Skip to content

Commit b18b9c7

Browse files
CaralHsiCaralHsi
committed
feat: add simple_openapi_memos_neo4j_community
1 parent 3c8f991 commit b18b9c7

File tree

1 file changed

+315
-0
lines changed

1 file changed

+315
-0
lines changed

‎examples/mem_os/simple_openapi_memos_neo4j_community.py‎

Lines changed: 315 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,315 @@
1+
import os
2+
import time
3+
import uuid
4+
5+
from datetime import datetime
6+
7+
from dotenv import load_dotenv
8+
9+
from memos.configs.mem_cube import GeneralMemCubeConfig
10+
from memos.configs.mem_os import MOSConfig
11+
from memos.mem_cube.general import GeneralMemCube
12+
from memos.mem_os.main import MOS
13+
14+
15+
load_dotenv()
16+
17+
# 1. Create MOS Config and set openai config
18+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to create MOS configuration...")
19+
start_time = time.time()
20+
21+
user_name = str(uuid.uuid4())
22+
print(user_name)
23+
24+
# 1.1 Set openai config
25+
openapi_config = {
26+
"model_name_or_path": "gpt-4o-mini",
27+
"temperature": 0.8,
28+
"max_tokens": 1024,
29+
"top_p": 0.9,
30+
"top_k": 50,
31+
"remove_think_prefix": True,
32+
"api_key": os.getenv("OPENAI_API_KEY", "sk-xxxxx"),
33+
"api_base": os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
34+
}
35+
embedder_config = {
36+
"backend": "universal_api",
37+
"config": {
38+
"provider": "openai",
39+
"api_key": os.getenv("OPENAI_API_KEY", "sk-xxxxx"),
40+
"model_name_or_path": "text-embedding-3-large",
41+
"base_url": os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
42+
},
43+
}
44+
EMBEDDING_DIMENSION = 3072
45+
46+
# 1.2 Set neo4j config
47+
neo4j_uri = os.getenv("NEO4J_URI", "bolt://localhost:7687")
48+
49+
# 1.3 Create MOS Config
50+
config = {
51+
"user_id": user_name,
52+
"chat_model": {
53+
"backend": "openai",
54+
"config": openapi_config,
55+
},
56+
"mem_reader": {
57+
"backend": "simple_struct",
58+
"config": {
59+
"llm": {
60+
"backend": "openai",
61+
"config": openapi_config,
62+
},
63+
"embedder": embedder_config,
64+
"chunker": {
65+
"backend": "sentence",
66+
"config": {
67+
"tokenizer_or_token_counter": "gpt2",
68+
"chunk_size": 512,
69+
"chunk_overlap": 128,
70+
"min_sentences_per_chunk": 1,
71+
},
72+
},
73+
},
74+
},
75+
"max_turns_window": 20,
76+
"top_k": 5,
77+
"enable_textual_memory": True,
78+
"enable_activation_memory": False,
79+
"enable_parametric_memory": False,
80+
}
81+
82+
mos_config = MOSConfig(**config)
83+
# you can set PRO_MODE to True to enable CoT enhancement mos_config.PRO_MODE = True
84+
mos = MOS(mos_config)
85+
86+
print(
87+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] MOS configuration created successfully, time elapsed: {time.time() - start_time:.2f}s\n"
88+
)
89+
90+
# 2. Initialize memory cube
91+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to initialize MemCube configuration...")
92+
start_time = time.time()
93+
94+
config = GeneralMemCubeConfig.model_validate(
95+
{
96+
"user_id": user_name,
97+
"cube_id": f"{user_name}",
98+
"text_mem": {
99+
"backend": "tree_text",
100+
"config": {
101+
"extractor_llm": {
102+
"backend": "openai",
103+
"config": openapi_config,
104+
},
105+
"dispatcher_llm": {
106+
"backend": "openai",
107+
"config": openapi_config,
108+
},
109+
"embedder": embedder_config,
110+
"graph_db": {
111+
"backend": "neo4j-community",
112+
"config": {
113+
"uri": neo4j_uri,
114+
"user": "neo4j",
115+
"password": "12345678",
116+
"db_name": "neo4j",
117+
"user_name": "alice",
118+
"use_multi_db": False,
119+
"auto_create": False,
120+
"embedding_dimension": EMBEDDING_DIMENSION,
121+
"vec_config": {
122+
"backend": "qdrant",
123+
"config": {
124+
"collection_name": "neo4j_vec_db",
125+
"vector_dimension": EMBEDDING_DIMENSION,
126+
"distance_metric": "cosine",
127+
"host": "localhost",
128+
"port": 6333,
129+
},
130+
},
131+
},
132+
},
133+
"reorganize": True,
134+
},
135+
},
136+
"act_mem": {},
137+
"para_mem": {},
138+
},
139+
)
140+
141+
print(
142+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] MemCube configuration initialization completed, time elapsed: {time.time() - start_time:.2f}s\n"
143+
)
144+
145+
# 3. Initialize the MemCube with the configuration
146+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to create MemCube instance...")
147+
start_time = time.time()
148+
149+
mem_cube = GeneralMemCube(config)
150+
try:
151+
mem_cube.dump(f"/tmp/{user_name}/")
152+
print(
153+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] MemCube created and saved successfully, time elapsed: {time.time() - start_time:.2f}s\n"
154+
)
155+
except Exception as e:
156+
print(
157+
f"❌ [{datetime.now().strftime('%H:%M:%S')}] MemCube save failed: {e}, time elapsed: {time.time() - start_time:.2f}s\n"
158+
)
159+
160+
# 4. Register the MemCube
161+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to register MemCube...")
162+
start_time = time.time()
163+
164+
mos.register_mem_cube(f"/tmp/{user_name}", mem_cube_id=user_name)
165+
166+
print(
167+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] MemCube registration completed, time elapsed: {time.time() - start_time:.2f}s\n"
168+
)
169+
170+
# 5. Add, get, search memory
171+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to add single memory...")
172+
start_time = time.time()
173+
174+
mos.add(memory_content="I like playing football.")
175+
176+
print(
177+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] Single memory added successfully, time elapsed: {time.time() - start_time:.2f}s"
178+
)
179+
180+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to get all memories...")
181+
start_time = time.time()
182+
183+
get_all_results = mos.get_all()
184+
185+
186+
# Filter out embedding fields, keeping only necessary fields
187+
def filter_memory_data(memories_data):
188+
filtered_data = {}
189+
for key, value in memories_data.items():
190+
if key == "text_mem":
191+
filtered_data[key] = []
192+
for mem_group in value:
193+
# Check if it's the new data structure (list of TextualMemoryItem objects)
194+
if "memories" in mem_group and isinstance(mem_group["memories"], list):
195+
# New data structure: directly a list of TextualMemoryItem objects
196+
filtered_memories = []
197+
for memory_item in mem_group["memories"]:
198+
# Create filtered dictionary
199+
filtered_item = {
200+
"id": memory_item.id,
201+
"memory": memory_item.memory,
202+
"metadata": {},
203+
}
204+
# Filter metadata, excluding embedding
205+
if hasattr(memory_item, "metadata") and memory_item.metadata:
206+
for attr_name in dir(memory_item.metadata):
207+
if not attr_name.startswith("_") and attr_name != "embedding":
208+
attr_value = getattr(memory_item.metadata, attr_name)
209+
if not callable(attr_value):
210+
filtered_item["metadata"][attr_name] = attr_value
211+
filtered_memories.append(filtered_item)
212+
213+
filtered_group = {
214+
"cube_id": mem_group.get("cube_id", ""),
215+
"memories": filtered_memories,
216+
}
217+
filtered_data[key].append(filtered_group)
218+
else:
219+
# Old data structure: dictionary with nodes and edges
220+
filtered_group = {
221+
"memories": {"nodes": [], "edges": mem_group["memories"].get("edges", [])}
222+
}
223+
for node in mem_group["memories"].get("nodes", []):
224+
filtered_node = {
225+
"id": node.get("id"),
226+
"memory": node.get("memory"),
227+
"metadata": {
228+
k: v
229+
for k, v in node.get("metadata", {}).items()
230+
if k != "embedding"
231+
},
232+
}
233+
filtered_group["memories"]["nodes"].append(filtered_node)
234+
filtered_data[key].append(filtered_group)
235+
else:
236+
filtered_data[key] = value
237+
return filtered_data
238+
239+
240+
filtered_results = filter_memory_data(get_all_results)
241+
print(f"Get all results after add memory: {filtered_results['text_mem'][0]['memories']}")
242+
243+
print(
244+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] Get all memories completed, time elapsed: {time.time() - start_time:.2f}s\n"
245+
)
246+
247+
# 6. Add messages
248+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to add conversation messages...")
249+
start_time = time.time()
250+
251+
messages = [
252+
{"role": "user", "content": "I like playing football."},
253+
{"role": "assistant", "content": "yes football is my favorite game."},
254+
]
255+
mos.add(messages)
256+
257+
print(
258+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] Conversation messages added successfully, time elapsed: {time.time() - start_time:.2f}s"
259+
)
260+
261+
print(
262+
f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to get all memories (after adding messages)..."
263+
)
264+
start_time = time.time()
265+
266+
get_all_results = mos.get_all()
267+
filtered_results = filter_memory_data(get_all_results)
268+
print(f"Get all results after add messages: {filtered_results}")
269+
270+
print(
271+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] Get all memories completed, time elapsed: {time.time() - start_time:.2f}s\n"
272+
)
273+
274+
# 7. Add document
275+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to add document...")
276+
start_time = time.time()
277+
## 7.1 add pdf for ./tmp/data if use doc mem mos.add(doc_path="./tmp/data/")
278+
start_time = time.time()
279+
280+
get_all_results = mos.get_all()
281+
filtered_results = filter_memory_data(get_all_results)
282+
print(f"Get all results after add doc: {filtered_results}")
283+
284+
print(
285+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] Get all memories completed, time elapsed: {time.time() - start_time:.2f}s\n"
286+
)
287+
288+
# 8. Search
289+
print(f"🚀 [{datetime.now().strftime('%H:%M:%S')}] Starting to search memories...")
290+
start_time = time.time()
291+
292+
search_results = mos.search(query="my favorite football game", user_id=user_name)
293+
filtered_search_results = filter_memory_data(search_results)
294+
print(f"Search results: {filtered_search_results}")
295+
296+
print(
297+
f"✅ [{datetime.now().strftime('%H:%M:%S')}] Memory search completed, time elapsed: {time.time() - start_time:.2f}s\n"
298+
)
299+
300+
# 9. Chat
301+
print(f"🎯 [{datetime.now().strftime('%H:%M:%S')}] Starting chat mode...")
302+
while True:
303+
user_input = input("👤 [You] ").strip()
304+
if user_input.lower() in ["quit", "exit"]:
305+
break
306+
307+
print()
308+
chat_start_time = time.time()
309+
response = mos.chat(user_input)
310+
chat_duration = time.time() - chat_start_time
311+
312+
print(f"🤖 [Assistant] {response}")
313+
print(f"⏱️ [Response time: {chat_duration:.2f}s]\n")
314+
315+
print("📢 [System] MemChat has stopped.")

0 commit comments

Comments
 (0)