feat: add api client (#316)

Co-authored-by: harvey_xiang <[email protected]>

Author: CarltonXiang

src/memos/api/client.py

@@ -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/product_models.py

@@ -178,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