Merge branch 'dev' into yjydev

Author: fridayL

poetry.lock

@@ -76,6 +76,11 @@ mem-scheduler = [
     "pika (>=1.3.2,<2.0.0)",  # RabbitMQ client
 ]
 
+# MemUser (MySQL support)
+mem-user = [
+    "pymysql (>=1.1.0,<2.0.0)",  # MySQL client for SQLAlchemy
+]
+
 # MemReader
 mem-reader = [
     "chonkie (>=1.0.7,<2.0.0)",  # Sentence chunking library
@@ -90,6 +95,7 @@ all = [
     "schedule (>=1.2.2,<2.0.0)",
     "redis (>=6.2.0,<7.0.0)",
     "pika (>=1.3.2,<2.0.0)",
+    "pymysql (>=1.1.0,<2.0.0)",
     "chonkie (>=1.0.7,<2.0.0)",
     "markitdown[docx,pdf,pptx,xls,xlsx] (>=0.1.1,<0.2.0)",
 
@@ -158,6 +164,10 @@ python-dotenv = "^1.1.1"
 langgraph = "^0.5.1"
 langmem = "^0.0.27"
 
+
+[tool.poetry.group.mem-user.dependencies]
+pymysql = "^1.1.2"
+
 [[tool.poetry.source]]
 name = "mirrors"
 url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple/"

pyproject.toml

@@ -0,0 +1,109 @@
+import json
+import os
+
+from typing import Any
+
+import requests
+
+from memos.api.product_models import MemOSAddResponse, MemOSGetMessagesResponse, MemOSSearchResponse
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+MAX_RETRY_COUNT = 3
+
+
+class MemOSClient:
+    """MemOS API client"""
+
+    def __init__(self, api_key: str | None = None, base_url: str | None = None):
+        self.base_url = (
+            base_url or os.getenv("MEMOS_BASE_URL") or "https://memos.memtensor.cn/api/openmem"
+        )
+        api_key = api_key or os.getenv("MEMOS_API_KEY")
+
+        if not api_key:
+            raise ValueError("MemOS API key is required")
+
+        self.headers = {"Content-Type": "application/json", "Authorization": f"Token {api_key}"}
+
+    def _validate_required_params(self, **params):
+        """Validate required parameters - if passed, they must not be empty"""
+        for param_name, param_value in params.items():
+            if not param_value:
+                raise ValueError(f"{param_name} is required")
+
+    def get_messages(
+        self, user_id: str, conversation_id: str | None = None
+    ) -> MemOSGetMessagesResponse:
+        """Get messages"""
+        # Validate required parameters
+        self._validate_required_params(user_id=user_id)
+
+        url = f"{self.base_url}/get/message"
+        payload = {"userId": user_id, "conversationId": conversation_id}
+        for retry in range(MAX_RETRY_COUNT):
+            try:
+                response = requests.post(
+                    url, data=json.dumps(payload), headers=self.headers, timeout=30
+                )
+                response.raise_for_status()
+                response_data = response.json()
+                return MemOSGetMessagesResponse(**response_data)
+            except Exception as e:
+                logger.error(f"Failed to get messages (retry {retry + 1}/3): {e}")
+                if retry == MAX_RETRY_COUNT - 1:
+                    raise
+
+    def add(
+        self, messages: list[dict[str, Any]], user_id: str, conversation_id: str
+    ) -> MemOSAddResponse:
+        """Add memories"""
+        # Validate required parameters
+        self._validate_required_params(
+            messages=messages, user_id=user_id, conversation_id=conversation_id
+        )
+
+        url = f"{self.base_url}/add/message"
+        payload = {"messages": messages, "userId": user_id, "conversationId": conversation_id}
+        for retry in range(MAX_RETRY_COUNT):
+            try:
+                response = requests.post(
+                    url, data=json.dumps(payload), headers=self.headers, timeout=30
+                )
+                response.raise_for_status()
+                response_data = response.json()
+                return MemOSAddResponse(**response_data)
+            except Exception as e:
+                logger.error(f"Failed to add memory (retry {retry + 1}/3): {e}")
+                if retry == MAX_RETRY_COUNT - 1:
+                    raise
+
+    def search(
+        self, query: str, user_id: str, conversation_id: str, memory_limit_number: int = 6
+    ) -> MemOSSearchResponse:
+        """Search memories"""
+        # Validate required parameters
+        self._validate_required_params(query=query, user_id=user_id)
+
+        url = f"{self.base_url}/search/memory"
+        payload = {
+            "query": query,
+            "userId": user_id,
+            "conversationId": conversation_id,
+            "memoryLimitNumber": memory_limit_number,
+        }
+
+        for retry in range(MAX_RETRY_COUNT):
+            try:
+                response = requests.post(
+                    url, data=json.dumps(payload), headers=self.headers, timeout=30
+                )
+                response.raise_for_status()
+                response_data = response.json()
+                return MemOSSearchResponse(**response_data)
+            except Exception as e:
+                logger.error(f"Failed to search memory (retry {retry + 1}/3): {e}")
+                if retry == MAX_RETRY_COUNT - 1:
+                    raise

src/memos/api/client.py

@@ -21,7 +21,7 @@ class APIConfig:
     def get_openai_config() -> dict[str, Any]:
         """Get OpenAI configuration."""
         return {
-            "model_name_or_path": os.getenv("MOS_OPENAI_MODEL", "gpt-4o-mini"),
+            "model_name_or_path": os.getenv("MOS_CHAT_MODEL", "gpt-4o-mini"),
             "temperature": float(os.getenv("MOS_CHAT_TEMPERATURE", "0.8")),
             "max_tokens": int(os.getenv("MOS_MAX_TOKENS", "1024")),
             "top_p": float(os.getenv("MOS_TOP_P", "0.9")),

src/memos/api/config.py

@@ -6,6 +6,7 @@
 and request isolation.
 """
 
+import os
 import uuid
 
 from collections.abc import Callable
@@ -117,6 +118,11 @@ def require_context() -> RequestContext:
 _trace_id_getter: TraceIdGetter | None = None
 
 
+def generate_trace_id() -> str:
+    """Generate a random trace_id."""
+    return os.urandom(16).hex()
+
+
 def set_trace_id_getter(getter: TraceIdGetter) -> None:
     """
     Set a custom trace_id getter function.

src/memos/api/context/context.py

@@ -3,25 +3,19 @@
 """
 
 import logging
-import os
 
 from collections.abc import Callable
 
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.requests import Request
 from starlette.responses import Response
 
-from memos.api.context.context import RequestContext, set_request_context
+from memos.api.context.context import RequestContext, generate_trace_id, set_request_context
 
 
 logger = logging.getLogger(__name__)
 
 
-def generate_trace_id() -> str:
-    """Generate a random trace_id."""
-    return os.urandom(16).hex()
-
-
 def extract_trace_id_from_headers(request: Request) -> str | None:
     """Extract trace_id from various possible headers with priority: g-trace-id > x-trace-id > trace-id."""
     trace_id = request.headers.get("g-trace-id")

src/memos/api/middleware/request_context.py

@@ -42,6 +42,7 @@ class UserRegisterRequest(BaseRequest):
     user_id: str = Field(
         default_factory=lambda: str(uuid.uuid4()), description="User ID for registration"
     )
+    mem_cube_id: str | None = Field(None, description="Cube ID for registration")
     user_name: str | None = Field(None, description="User name for registration")
     interests: str | None = Field(None, description="User interests")
 
@@ -177,3 +178,104 @@ class SuggestionRequest(BaseRequest):
     user_id: str = Field(..., description="User ID")
     language: Literal["zh", "en"] = Field("zh", description="Language for suggestions")
     message: list[MessageDict] | None = Field(None, description="List of messages to store.")
+
+
+# ─── MemOS Client Response Models ──────────────────────────────────────────────
+
+
+class MessageDetail(BaseModel):
+    """Individual message detail model based on actual API response."""
+
+    role: str = Field(..., description="Message role (user/assistant)")
+    content: str = Field(..., description="Message content")
+    create_time: int | None = Field(
+        None, alias="createTime", description="Message creation timestamp"
+    )
+    update_time: int | None = Field(
+        None, alias="updateTime", description="Message update timestamp"
+    )
+
+
+class MemoryDetail(BaseModel):
+    """Individual memory detail model based on actual API response."""
+
+    id: str = Field(..., description="Memory ID")
+    memory_key: str = Field(..., alias="memoryKey", description="Memory key/title")
+    memory_value: str = Field(..., alias="memoryValue", description="Memory content")
+    memory_type: str = Field(
+        ..., alias="memoryType", description="Memory type (e.g., WorkingMemory)"
+    )
+    memory_time: int | None = Field(None, alias="memoryTime", description="Memory timestamp")
+    conversation_id: str = Field(..., alias="conversationId", description="Conversation ID")
+    status: str = Field(..., description="Memory status (e.g., activated)")
+    confidence: float = Field(..., description="Memory confidence score")
+    tags: list[str] = Field(default_factory=list, description="Memory tags")
+    update_time: int = Field(..., alias="updateTime", description="Last update timestamp")
+    relativity: float = Field(..., description="Memory relativity/similarity score")
+
+
+class GetMessagesData(BaseModel):
+    """Data model for get messages response based on actual API."""
+
+    message_detail_list: list[MessageDetail] = Field(
+        default_factory=list, alias="messageDetailList", description="List of message details"
+    )
+
+
+class SearchMemoryData(BaseModel):
+    """Data model for search memory response based on actual API."""
+
+    memory_detail_list: list[MemoryDetail] = Field(
+        default_factory=list, alias="memoryDetailList", description="List of memory details"
+    )
+    message_detail_list: list[MessageDetail] | None = Field(
+        None, alias="messageDetailList", description="List of message details (usually None)"
+    )
+
+
+class AddMessageData(BaseModel):
+    """Data model for add message response based on actual API."""
+
+    success: bool = Field(..., description="Operation success status")
+
+
+# ─── MemOS Response Models (Similar to OpenAI ChatCompletion) ──────────────────
+
+
+class MemOSGetMessagesResponse(BaseModel):
+    """Response model for get messages operation based on actual API."""
+
+    code: int = Field(..., description="Response status code")
+    message: str = Field(..., description="Response message")
+    data: GetMessagesData = Field(..., description="Messages data")
+
+    @property
+    def messages(self) -> list[MessageDetail]:
+        """Convenient access to message list."""
+        return self.data.message_detail_list
+
+
+class MemOSSearchResponse(BaseModel):
+    """Response model for search memory operation based on actual API."""
+
+    code: int = Field(..., description="Response status code")
+    message: str = Field(..., description="Response message")
+    data: SearchMemoryData = Field(..., description="Search results data")
+
+    @property
+    def memories(self) -> list[MemoryDetail]:
+        """Convenient access to memory list."""
+        return self.data.memory_detail_list
+
+
+class MemOSAddResponse(BaseModel):
+    """Response model for add message operation based on actual API."""
+
+    code: int = Field(..., description="Response status code")
+    message: str = Field(..., description="Response message")
+    data: AddMessageData = Field(..., description="Add operation data")
+
+    @property
+    def success(self) -> bool:
+        """Convenient access to success status."""
+        return self.data.success

src/memos/api/product_models.py

@@ -106,6 +106,7 @@ def register_user(user_req: UserRegisterRequest, g: Annotated[G, Depends(get_g_o
             interests=user_req.interests,
             config=user_config,
             default_mem_cube=default_mem_cube,
+            mem_cube_id=user_req.mem_cube_id,
         )
 
         if result["status"] == "success":