TreeTextMemory: 树形明文记忆

from memos.configs.memory import TreeTextMemoryConfig

config = TreeTextMemoryConfig.from_json_file("examples/data/config/tree_config.json")

from memos.memories.textual.tree import TreeTextMemory

tree_memory = TreeTextMemory(config)

from memos.mem_reader.simple_struct import SimpleStructMemReader

reader = SimpleStructMemReader.from_json_file("examples/data/config/simple_struct_reader_config.json")

scene_data = [[
    {"role": "user", "content": "Tell me about your childhood."},
    {"role": "assistant", "content": "I loved playing in the garden with my dog."}
]]

memories = reader.get_memory(scene_data, type="chat", info={"user_id": "1234"})
for m_list in memories:
    tree_memory.add(m_list)

from memos.configs.mem_reader import MultiModalStructMemReaderConfig
from memos.mem_reader.multi_modal_struct import MultiModalStructMemReader

# 创建 MultiModal Reader 配置
multimodal_config = MultiModalStructMemReaderConfig(
    llm={
        "backend": "openai",
        "config": {
            "model_name_or_path": "gpt-4o-mini",
            "api_key": "your-api-key"
        }
    },
    embedder={
        "backend": "openai",
        "config": {
            "model_name_or_path": "text-embedding-3-small",
            "api_key": "your-api-key"
        }
    },
    chunker={
        "backend": "text_splitter",
        "config": {
            "chunk_size": 1000,
            "chunk_overlap": 200
        }
    },
    extractor_llm={
        "backend": "openai",
        "config": {
            "model_name_or_path": "gpt-4o-mini",
            "api_key": "your-api-key"
        }
    },
    # 可选：指定哪些域名直接返回 Markdown
    direct_markdown_hostnames=["github.com", "docs.python.org"]
)

# 初始化 MultiModal Reader
multimodal_reader = MultiModalStructMemReader(multimodal_config)

# ========================================
# 示例 1: 处理包含图片的对话
# ========================================
scene_with_image = [[
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "这是我家的花园"},
            {"type": "image_url", "image_url": {"url": "https://example.com/garden.jpg"}}
        ]
    },
    {
        "role": "assistant",
        "content": "你的花园很漂亮！"
    }
]]

memories = multimodal_reader.get_memory(
    scene_with_image,
    type="chat",
    info={"user_id": "1234", "session_id": "session_001"}
)
for m_list in memories:
    tree_memory.add(m_list)
print(f"✓ 已添加 {len(memories)} 条多模态记忆")

# ========================================
# 示例 2: 处理网页 URL
# ========================================
scene_with_url = [[
    {
        "role": "user",
        "content": "请分析这篇文章: https://example.com/article.html"
    },
    {
        "role": "assistant",
        "content": "我会帮你分析这篇文章"
    }
]]

url_memories = multimodal_reader.get_memory(
    scene_with_url,
    type="chat",
    info={"user_id": "1234", "session_id": "session_002"}
)
for m_list in url_memories:
    tree_memory.add(m_list)
print(f"✓ 已从 URL 提取并添加 {len(url_memories)} 条记忆")

# ========================================
# 示例 3: 处理本地文件
# ========================================
# 支持的文件类型: PDF, DOCX, TXT, Markdown, HTML 等
file_paths = [
    "./documents/report.pdf",
    "./documents/notes.md",
    "./documents/data.txt"
]

file_memories = multimodal_reader.get_memory(
    file_paths,
    type="doc",
    info={"user_id": "1234", "session_id": "session_003"}
)
for m_list in file_memories:
    tree_memory.add(m_list)
print(f"✓ 已从文件提取并添加 {len(file_memories)} 条记忆")

# ========================================
# 示例 4: 混合模式（文本 + 图片 + URL）
# ========================================
mixed_scene = [[
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "这是我的项目文档:"},
            {"type": "text", "text": "https://github.com/user/project/README.md"},
            {"type": "image_url", "image_url": {"url": "https://example.com/diagram.png"}}
        ]
    }
]]

mixed_memories = multimodal_reader.get_memory(
    mixed_scene,
    type="chat",
    info={"user_id": "1234", "session_id": "session_004"}
)
for m_list in mixed_memories:
    tree_memory.add(m_list)
print(f"✓ 已从混合内容提取并添加 {len(mixed_memories)} 条记忆")

results = tree_memory.search("Talk about the garden", top_k=5)
for i, node in enumerate(results):
    print(f"{i}: {node.memory}")

# 创建embedder
embedder = EmbedderFactory.from_config(
    EmbedderConfigFactory.model_validate({
        "backend": "ollama",
        "config": {"model_name_or_path": "nomic-embed-text:latest"},
    })
)

# 配置检索器（以 BochaAI 为例）
retriever_config = InternetRetrieverConfigFactory.model_validate({
    "backend": "bocha",
    "config": {
        "api_key": "sk-xxx",  # 替换为你的 BochaAI API Key
        "max_results": 5,
        "reader": {  # 自动分块的 Reader 配置
            "backend": "simple_struct",
            "config": ...,  # 你的mem-reader config
        },
    }
})

# 实例化检索器
retriever = InternetRetrieverFactory.from_config(retriever_config, embedder)

# 执行网页检索
results = retriever.retrieve_from_internet("Alibaba 2024 ESG report")

# 添加到记忆图中
for m in results:
    tree_memory.add(m)

{
  "internet_retriever": {
    "backend": "bocha",
    "config": {
      "api_key": "sk-xxx",
      "max_results": 5,
      "reader": {
        "backend": "simple_struct",
        "config": ...
      }
    }
  }
}

tree_memory.replace_working_memory(
    [{
        "memory": "User is discussing gardening tips.",
        "metadata": {"memory_type": "WorkingMemory"}
    }]
)

tree_memory.dump("tmp/tree_memories")
tree_memory.load("tmp/tree_memories")