Skip to content

Commit 88993de

Browse files
authored
feat: simplify simple tree (#461)
* feat: simplify simple tree * feat: add product_api examples * feat: modify online bot * feat: modify notification * feat: time * format: dingding report
1 parent 8a63675 commit 88993de

File tree

7 files changed

+277
-302
lines changed

7 files changed

+277
-302
lines changed

examples/api/__init__.py

Whitespace-only changes.

examples/api/product_api.py

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Simulate full MemOS Product API workflow:
4+
1. Register user
5+
2. Add memory
6+
3. Search memory
7+
4. Chat (stream)
8+
"""
9+
10+
import json
11+
12+
import requests
13+
14+
15+
BASE_URL = "http://0.0.0.0:8001/product"
16+
HEADERS = {"Content-Type": "application/json"}
17+
18+
index = "24"
19+
USER_ID = f"memos_user_id_{index}"
20+
USER_NAME = f"memos_user_alice_{index}"
21+
MEM_CUBE_ID = f"memos_cube_id_{index}"
22+
SESSION_ID = f"memos_session_id_{index}"
23+
SESSION_ID2 = f"memos_session_id_{index}_s2"
24+
25+
26+
def register_user():
27+
url = f"{BASE_URL}/users/register"
28+
data = {
29+
"user_id": USER_ID,
30+
"user_name": USER_NAME,
31+
"interests": "memory,retrieval,test",
32+
"mem_cube_id": MEM_CUBE_ID,
33+
}
34+
print(f"[*] Registering user {USER_ID} ...")
35+
resp = requests.post(url, headers=HEADERS, data=json.dumps(data), timeout=30)
36+
print(resp.status_code, resp.text)
37+
return resp.json()
38+
39+
40+
def add_memory():
41+
url = f"{BASE_URL}/add"
42+
data = {
43+
"user_id": USER_ID,
44+
"memory_content": "今天我在测试 MemOS 的记忆添加与检索流程。",
45+
"messages": [{"role": "user", "content": "我今天在做系统测试"}],
46+
"doc_path": None,
47+
"mem_cube_id": MEM_CUBE_ID,
48+
"source": "test_script",
49+
"user_profile": False,
50+
"session_id": SESSION_ID,
51+
}
52+
print("[*] Adding memory ...")
53+
resp = requests.post(url, headers=HEADERS, data=json.dumps(data), timeout=30)
54+
print(resp.status_code, resp.text)
55+
return resp.json()
56+
57+
58+
def search_memory(query="系统测试"):
59+
url = f"{BASE_URL}/search"
60+
data = {
61+
"user_id": USER_ID,
62+
"query": query,
63+
"mem_cube_id": MEM_CUBE_ID,
64+
"top_k": 5,
65+
"session_id": SESSION_ID,
66+
}
67+
print("[*] Searching memory ...")
68+
resp = requests.post(url, headers=HEADERS, data=json.dumps(data), timeout=30)
69+
print(resp.status_code, resp.text)
70+
return resp.json()
71+
72+
73+
def chat_stream(query: str, session_id: str, history: list | None = None):
74+
url = f"{BASE_URL}/chat"
75+
data = {
76+
"user_id": USER_ID,
77+
"query": query,
78+
"mem_cube_id": MEM_CUBE_ID,
79+
"history": history,
80+
"internet_search": False,
81+
"moscube": False,
82+
"session_id": session_id,
83+
}
84+
85+
print("[*] Starting streaming chat ...")
86+
87+
with requests.post(url, headers=HEADERS, data=json.dumps(data), stream=True) as resp:
88+
for raw_line in resp.iter_lines():
89+
if not raw_line:
90+
continue
91+
line = raw_line.decode("utf-8", errors="ignore")
92+
93+
payload = line.removeprefix("data: ").strip()
94+
if payload == "[DONE]":
95+
print("[done]")
96+
break
97+
98+
try:
99+
msg = json.loads(payload)
100+
msg_type = msg.get("type")
101+
msg_data = msg.get("data") or msg.get("content")
102+
103+
if msg_type == "text":
104+
print(msg_data, end="", flush=True)
105+
elif msg_type == "reference":
106+
print(f"\n[参考记忆] {msg_data}")
107+
elif msg_type == "status":
108+
pass
109+
elif msg_type == "suggestion":
110+
print(f"\n[建议] {msg_data}")
111+
elif msg_type == "end":
112+
print("\n[✅ Chat End]")
113+
else:
114+
print(f"\n[{msg_type}] {msg_data}")
115+
except Exception:
116+
try:
117+
print(payload.encode("latin-1").decode("utf-8"), end="")
118+
except Exception:
119+
print(payload)
120+
121+
122+
if __name__ == "__main__":
123+
print("===== STEP 1: Register User =====")
124+
register_user()
125+
126+
print("\n===== STEP 2: Add Memory =====")
127+
add_memory()
128+
129+
print("\n===== STEP 3: Search Memory =====")
130+
search_memory()
131+
132+
print("\n===== STEP 4: Stream Chat =====")
133+
chat_stream("我很开心,我今天吃了好吃的拉面", SESSION_ID, history=[])
134+
chat_stream(
135+
"我刚和你说什么",
136+
SESSION_ID,
137+
history=[
138+
{"role": "user", "content": "我很开心,我今天吃了好吃的拉面"},
139+
{"role": "assistant", "content": "🉑"},
140+
],
141+
)
142+
143+
print("\n===== STEP 4: Stream Chat =====")
144+
chat_stream("我刚和你说什么了呢", SESSION_ID2, history=[])

src/memos/mem_os/product.py

Lines changed: 48 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,34 @@ def _extract_references_from_response(self, response: str) -> tuple[str, list[di
563563
logger.error(f"Error extracting references from response: {e}", exc_info=True)
564564
return response, []
565565

566+
def _extract_struct_data_from_history(self, chat_data: list[dict]) -> dict:
567+
"""
568+
get struct message from chat-history
569+
# TODO: @xcy make this more general
570+
"""
571+
system_content = ""
572+
memory_content = ""
573+
chat_history = []
574+
575+
for item in chat_data:
576+
role = item.get("role")
577+
content = item.get("content", "")
578+
if role == "system":
579+
parts = content.split("# Memories", 1)
580+
system_content = parts[0].strip()
581+
if len(parts) > 1:
582+
memory_content = "# Memories" + parts[1].strip()
583+
elif role in ("user", "assistant"):
584+
chat_history.append({"role": role, "content": content})
585+
586+
if chat_history and chat_history[-1]["role"] == "assistant":
587+
if len(chat_history) >= 2 and chat_history[-2]["role"] == "user":
588+
chat_history = chat_history[:-2]
589+
else:
590+
chat_history = chat_history[:-1]
591+
592+
return {"system": system_content, "memory": memory_content, "chat_history": chat_history}
593+
566594
def _chunk_response_with_tiktoken(
567595
self, response: str, chunk_size: int = 5
568596
) -> Generator[str, None, None]:
@@ -640,23 +668,26 @@ async def _post_chat_processing(
640668
clean_response, extracted_references = self._extract_references_from_response(
641669
full_response
642670
)
671+
struct_message = self._extract_struct_data_from_history(current_messages)
643672
logger.info(f"Extracted {len(extracted_references)} references from response")
644673

645674
# Send chat report notifications asynchronously
646675
if self.online_bot:
676+
logger.info("Online Bot Open!")
647677
try:
648678
from memos.memos_tools.notification_utils import (
649679
send_online_bot_notification_async,
650680
)
651681

652682
# Prepare notification data
653-
chat_data = {
654-
"query": query,
655-
"user_id": user_id,
656-
"cube_id": cube_id,
657-
"system_prompt": system_prompt,
658-
"full_response": full_response,
659-
}
683+
chat_data = {"query": query, "user_id": user_id, "cube_id": cube_id}
684+
chat_data.update(
685+
{
686+
"memory": struct_message["memory"],
687+
"chat_history": struct_message["chat_history"],
688+
"full_response": full_response,
689+
}
690+
)
660691

661692
system_data = {
662693
"references": extracted_references,
@@ -720,6 +751,7 @@ def _start_post_chat_processing(
720751
"""
721752
Asynchronous processing of logs, notifications and memory additions, handle synchronous and asynchronous environments
722753
"""
754+
logger.info("Start post_chat_processing...")
723755

724756
def run_async_in_thread():
725757
"""Running asynchronous tasks in a new thread"""
@@ -1046,14 +1078,20 @@ def chat(
10461078
memories_list = new_memories_list
10471079

10481080
system_prompt = super()._build_system_prompt(memories_list, base_prompt)
1049-
history_info = []
1050-
if history:
1081+
if history is not None:
1082+
# Use the provided history (even if it's empty)
10511083
history_info = history[-20:]
1084+
else:
1085+
# Fall back to internal chat_history
1086+
if user_id not in self.chat_history_manager:
1087+
self._register_chat_history(user_id, session_id)
1088+
history_info = self.chat_history_manager[user_id].chat_history[-20:]
10521089
current_messages = [
10531090
{"role": "system", "content": system_prompt},
10541091
*history_info,
10551092
{"role": "user", "content": query},
10561093
]
1094+
logger.info("Start to get final answer...")
10571095
response = self.chat_llm.generate(current_messages)
10581096
time_end = time.time()
10591097
self._start_post_chat_processing(
@@ -1129,7 +1167,7 @@ def chat_with_references(
11291167
self._register_chat_history(user_id, session_id)
11301168

11311169
chat_history = self.chat_history_manager[user_id]
1132-
if history:
1170+
if history is not None:
11331171
chat_history.chat_history = history[-20:]
11341172
current_messages = [
11351173
{"role": "system", "content": system_prompt},

0 commit comments

Comments
 (0)