1
import json
2
import os
3

4
from dotenv import load_dotenv
5
from openai import OpenAI
6
from pathlib import Path
7

8
from mem0 import Memory
9
from mem0.configs.base import MemoryConfig
10
from sentence_transformers import SentenceTransformer
11

12
def build_memory() -> Memory:
13
    llm_model = os.getenv("LLM_MODEL")
14
    embedding_path = os.getenv("EMBEDDING_PATH")
15
    openai_api_key = os.getenv("OPENAI_API_KEY")
16
    openai_base_url = os.getenv("OPENAI_BASE_URL")
17

18
    if not openai_api_key:
19
        raise RuntimeError("Missing OPENAI_API_KEY. Add it to .env before running mem0_demo.py")
20
    if not llm_model:
21
        raise RuntimeError("Missing LLM_MODEL. Add it to .env before running mem0_demo.py")
22
    if not embedding_path:
23
        raise RuntimeError("Missing EMBEDDING_PATH. Add it to .env before running mem0_demo.py")
24

25
    model_path = Path(embedding_path)
26
    if not model_path.exists():
27
        raise RuntimeError(f"EMBEDDING_PATH does not exist: {model_path}")
28

29
    embedding_model = SentenceTransformer(str(model_path))
30
    embedding_dims = embedding_model.get_embedding_dimension()
31

32
    config = MemoryConfig(
33
        llm={
34
            "provider": "openai",
35
            "config": {
36
                "model": llm_model,
37
                "api_key": openai_api_key,
38
                "openai_base_url": openai_base_url,
39
            },
40
        },
41
        embedder={
42
            "provider": "huggingface",
43
            "config": {
44
                "model": str(model_path),
45
                "embedding_dims": embedding_dims,
46
            },
47
        },
48
        vector_store={
49
            "provider": "qdrant",
50
            "config": {
51
                "collection_name": "mem0",
52
                "path": ".mem0/qdrant",
53
                "embedding_model_dims": embedding_dims,
54
            },
55
        },
56
    )
57
    return Memory(config=config)
58

59

60
USER_ID = "whyulooksad"
61
SYSTEM_PROMPT = "你是一个带长期记忆的 AI 助手。回答时优先参考历史记忆，回答要简洁准确。"
62

63

64
def close_memory(memory) -> None:
65
    memory.close()
66

67
    vector_store = getattr(memory, "vector_store", None)
68
    if vector_store and getattr(vector_store, "client", None):
69
        vector_store.client.close()
70

71
    telemetry_store = getattr(memory, "_telemetry_vector_store", None)
72
    if telemetry_store and getattr(telemetry_store, "client", None):
73
        telemetry_store.client.close()
74

75

76
def build_llm_client() -> tuple[OpenAI, str]:
77
    api_key = os.getenv("OPENAI_API_KEY")
78
    base_url = os.getenv("OPENAI_BASE_URL")
79
    model = os.getenv("LLM_MODEL")
80

81
    if not api_key:
82
        raise RuntimeError("Missing OPENAI_API_KEY in .env")
83
    if not base_url:
84
        raise RuntimeError("Missing OPENAI_BASE_URL in .env")
85
    if not model:
86
        raise RuntimeError("Missing LLM_MODEL in .env")
87

88
    return OpenAI(api_key=api_key, base_url=base_url), model
89

90

91
def format_memories(search_result: dict) -> str:
92
    items = search_result.get("results", [])
93
    if not items:
94
        return "暂无历史记忆"
95
    return "\n".join(f"- {item['memory']}" for item in items)
96

97

98
def ask_agent(client: OpenAI, model: str, memory, user_input: str) -> None:
99
    memories = memory.search(user_input, user_id=USER_ID, limit=5)
100
    memory_text = format_memories(memories)
101

102
    messages = [
103
        {"role": "system", "content": SYSTEM_PROMPT},
104
        {
105
            "role": "user",
106
            "content": f"历史记忆：\n{memory_text}\n\n用户问题：{user_input}",
107
        },
108
    ]
109

110
    print("=== 检索到的记忆 ===")
111
    print(json.dumps(memories, ensure_ascii=False, indent=2))
112

113
    print("\n=== 注入 Prompt ===")
114
    print(messages[1]["content"])
115

116
    response = client.chat.completions.create(
117
        model=model,
118
        messages=messages,
119
        temperature=0.2,
120
    )
121
    answer = response.choices[0].message.content or ""
122

123
    print("\n=== Agent 回复 ===")
124
    print(answer)
125

126
    add_result = memory.add(
127
        [
128
            {"role": "user", "content": user_input},
129
            {"role": "assistant", "content": answer},
130
        ],
131
        user_id=USER_ID,
132
        metadata={"source": "agent-chat", "topic": "agent-memory-demo"},
133
    )
134
    print("\n=== 写回记忆 ===")
135
    print(json.dumps(add_result, ensure_ascii=False, indent=2))
136

137

138
def main() -> None:
139
    load_dotenv()
140
    memory = build_memory()
141
    client, model = build_llm_client()
142
    try:
143
        ask_agent(client, model, memory, "结合你的历史记忆，概括一下我平时研究的技术方向。")
144
    finally:
145
        close_memory(memory)
146

147

148
if __name__ == "__main__":
149
    main()

1
# 1. 添加记忆
2
mem.add("内容", user_id="xxx")
3

4
# 2. 语义搜索
5
mem.search("问题", user_id="xxx")
6

7
# 3. 获取该用户所有记忆
8
mem.get_all(user_id="xxx")
9

10
# 4. 手动更新记忆
11
mem.update(memory_id="xxx", data="新内容")
12

13
# 5. 手动删除记忆
14
mem.delete(memory_id="xxx")