Skip to content

Commit 5481f56

Browse files
authored
Feat: add chat complete for new server_api (#366)
* feat: add server api prd * feat: update memcube for api * feat: add run server api md and change user_id to user_id * fix: code format * fix:code * fix: fix code format * feat: remove ids * fix: working ids * feat: add_memreader config and change neo4j db user_name * feat: add chat model
1 parent de2b5c6 commit 5481f56

File tree

3 files changed

+484
-2
lines changed

3 files changed

+484
-2
lines changed

src/memos/api/product_models.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,23 @@ class APIADDRequest(BaseRequest):
197197
)
198198

199199

200+
class APIChatCompleteRequest(BaseRequest):
201+
"""Request model for chat operations."""
202+
203+
user_id: str = Field(..., description="User ID")
204+
query: str = Field(..., description="Chat query message")
205+
mem_cube_id: str | None = Field(None, description="Cube ID to use for chat")
206+
history: list[MessageDict] | None = Field(None, description="Chat history")
207+
internet_search: bool = Field(False, description="Whether to use internet search")
208+
moscube: bool = Field(True, description="Whether to use MemOSCube")
209+
base_prompt: str | None = Field(None, description="Base prompt to use for chat")
210+
top_k: int = Field(10, description="Number of results to return")
211+
threshold: float = Field(0.5, description="Threshold for filtering references")
212+
session_id: str | None = Field(
213+
"default_session", description="Session ID for soft-filtering memories"
214+
)
215+
216+
200217
class SuggestionRequest(BaseRequest):
201218
"""Request model for getting suggestion queries."""
202219

src/memos/api/routers/server_router.py

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
import os
2+
import traceback
23

34
from typing import Any
45

5-
from fastapi import APIRouter
6+
from fastapi import APIRouter, HTTPException
67

78
from memos.api.config import APIConfig
89
from memos.api.product_models import (
910
APIADDRequest,
11+
APIChatCompleteRequest,
1012
APISearchRequest,
1113
MemoryResponse,
1214
SearchResponse,
@@ -22,6 +24,7 @@
2224
from memos.llms.factory import LLMFactory
2325
from memos.log import get_logger
2426
from memos.mem_cube.navie import NaiveMemCube
27+
from memos.mem_os.product_server import MOSServer
2528
from memos.mem_reader.factory import MemReaderFactory
2629
from memos.memories.textual.tree_text_memory.organize.manager import MemoryManager
2730
from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory import (
@@ -126,7 +129,11 @@ def init_server():
126129
memory_size=_get_default_memory_size(default_cube_config),
127130
is_reorganize=getattr(default_cube_config.text_mem.config, "reorganize", False),
128131
)
129-
132+
mos_server = MOSServer(
133+
mem_reader=mem_reader,
134+
llm=llm,
135+
online_bot=False,
136+
)
130137
return (
131138
graph_db,
132139
mem_reader,
@@ -136,6 +143,7 @@ def init_server():
136143
internet_retriever,
137144
memory_manager,
138145
default_cube_config,
146+
mos_server,
139147
)
140148

141149

@@ -149,6 +157,7 @@ def init_server():
149157
internet_retriever,
150158
memory_manager,
151159
default_cube_config,
160+
mos_server,
152161
) = init_server()
153162

154163

@@ -280,3 +289,36 @@ def add_memories(add_req: APIADDRequest):
280289
message="Memory added successfully",
281290
data=response_data,
282291
)
292+
293+
294+
@router.post("/chat/complete", summary="Chat with MemOS (Complete Response)")
295+
def chat_complete(chat_req: APIChatCompleteRequest):
296+
"""Chat with MemOS for a specific user. Returns complete response (non-streaming)."""
297+
try:
298+
# Collect all responses from the generator
299+
naive_mem_cube = _create_naive_mem_cube()
300+
content, references = mos_server.chat(
301+
query=chat_req.query,
302+
user_id=chat_req.user_id,
303+
cube_id=chat_req.mem_cube_id,
304+
mem_cube=naive_mem_cube,
305+
history=chat_req.history,
306+
internet_search=chat_req.internet_search,
307+
moscube=chat_req.moscube,
308+
base_prompt=chat_req.base_prompt,
309+
top_k=chat_req.top_k,
310+
threshold=chat_req.threshold,
311+
session_id=chat_req.session_id,
312+
)
313+
314+
# Return the complete response
315+
return {
316+
"message": "Chat completed successfully",
317+
"data": {"response": content, "references": references},
318+
}
319+
320+
except ValueError as err:
321+
raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err
322+
except Exception as err:
323+
logger.error(f"Failed to start chat: {traceback.format_exc()}")
324+
raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err

0 commit comments

Comments
 (0)