创建你的第一个记忆

from memos import log
from memos.configs.mem_reader import SimpleStructMemReaderConfig
from memos.configs.memory import TreeTextMemoryConfig
from memos.mem_reader.simple_struct import SimpleStructMemReader
from memos.memories.textual.tree import TreeTextMemory

logger = log.get_logger(__name__)

# 1. 初始化 TreeTextMemory（记忆仓库）
tree_config = TreeTextMemoryConfig.from_json_file(
    "examples/data/config/tree_config_shared_database.json"
)
my_tree_textual_memory = TreeTextMemory(tree_config)

# ⚠️ 注意：这里为了演示方便清空了旧数据。生产环境千万别这么干！
my_tree_textual_memory.delete_all()

# 2. 初始化 SimpleStructMemReader（信息提取器）
reader_config = SimpleStructMemReaderConfig.from_json_file(
    "examples/data/config/simple_struct_reader_config.json"
)
reader = SimpleStructMemReader(reader_config)

scene_data = [
    [
        {
            "role": "user",
            "chat_time": "3 May 2025",
            "content": "This week I’ve been feeling a bit off, especially when playing tennis. My body just doesn’t feel right.",
        },
        {
            "role": "assistant",
            "chat_time": "3 May 2025",
            "content": "It sounds like you've been having some physical discomfort lately...",
        },
        # ... (中间省略了几轮吐槽) ...
        {
            "role": "user",
            "chat_time": "3 May 2025",
            "content": "I think it might be due to stress and lack of sleep recently...",
        },
    ]
]

# 1. 提取 (Extract)
# Reader 会调用 LLM 分析对话，返回一个记忆列表
memory = reader.get_memory(
    scene_data, 
    type="chat", 
    info={"user_id": "1234", "session_id": "2222"}
)

# 2. 存储 (Add)
for m_list in memory:
    added_ids = my_tree_textual_memory.add(m_list)
    
    # 看看存进去了啥
    for i, id in enumerate(added_ids):
        print(f"存入第 {i} 条记忆: " + my_tree_textual_memory.get(id).memory)
    
    # 等待后台整理完成（建立索引需要一点时间）
    my_tree_textual_memory.memory_manager.wait_reorganizer()

# 稍微等一下索引构建
import time
time.sleep(2)

init_time = time.time()

# 试着搜一下关于“童年”的事（假设之前的对话里包含相关内容）
# 或者搜 "Why is the user feeling bad?" 试试
results = my_tree_textual_memory.search(
    "Talk about the user's childhood story?",
    top_k=10,
    info={
        "query": "Talk about the user's childhood story?",
        "user_id": "111",
        "session_id": "2234",
    },
)

for i, r in enumerate(results):
    print(f"搜到的第 {i} 条结果: {r.memory}")

print(f"搜索耗时: {round(time.time() - init_time)}s")

# 开启 Fine 模式
results_fine_search = my_tree_textual_memory.search(
    "Recent news in the first city you've mentioned.",
    top_k=10,
    mode="fine", # 关键在这里
    info={
        "query": "Recent news in NewYork",
        "user_id": "111",
        "session_id": "2234",
        "chat_history": [
            {"role": "user", "content": "I want to know three beautiful cities"},
            {"role": "assistant", "content": "New York, London, and Shanghai"},
        ],
    },
)

for i, r in enumerate(results_fine_search):
    print(f"Fine Search 结果: {r.memory}")

# 构造文档数据
doc_data = [
    {
        "type": "file",
        "file": {
            "filename": "tennis_rule.txt",
            "path": "./tennis_rule.txt", # 确保文件存在
            # 或者直接提供 content: "file_data": "..."
        }
    }
]

# 告诉 Reader 这是 "doc" 类型
doc_memories = reader.get_memory(
    doc_data, 
    type="doc", 
    info={"user_id": "1234", "session_id": "docs_import"}
)

# 存入记忆
for m in doc_memories:
    my_tree_textual_memory.add(m)

tool_scene = [
    [
        {"role": "user", "content": "What's the weather in Beijing?"},
        {
            "role": "assistant", 
            "content": "", 
            "tool_calls": [{"id": "call_1", "function": {"name": "get_weather", "arguments": "{'city': 'Beijing'}"}}]
        },
        {
            "role": "tool", 
            "tool_call_id": "call_1", 
            "content": "Sunny, 25°C"
        }
    ]
]

# Reader 会自动理解这是工具交互
tool_memories = reader.get_memory(tool_scene, type="chat", info={"user_id": "1234"})

from memos.memories.textual.simple_preference import SimplePreferenceTextMemory
# 注意：初始化需要配置 VectorDB, Embedder 等，这里仅作示意
# pref_memory = SimplePreferenceTextMemory(...)

# 从对话中自动提取偏好
pref_memories = pref_memory.get_memory(chat_data, type="chat", info=...)

# 存入偏好
pref_memory.add(pref_memories)

# 搜索偏好
prefs = pref_memory.search("What is the user's UI preference?", top_k=1)
print(prefs[0].memory) # 输出: "User prefers dark mode"

from memos.mem_feedback.simple_feedback import SimpleMemFeedback

# 初始化反馈模块
# feedback_module = SimpleMemFeedback(...)

# 处理用户反馈
# 假设用户说："Actually, I started playing tennis in 2020, not 2018."
feedback_module.process_feedback({
    "user_id": "1234",
    "feedback_content": "Actually, I started playing tennis in 2020, not 2018.",
    "chat_history": [...], # 提供上下文
    "feedback_time": "Now"
})

# 反馈模块会自动在后台更新 Graph 数据库中的节点和关系

# 关闭后台线程
my_tree_textual_memory.memory_manager.close()

# 备份一下记忆
my_tree_textual_memory.dump("tmp/my_tree_textual_memory")

# 删库跑路（仅限测试环境！）
my_tree_textual_memory.drop()