MemoryOS 1.0.1__tar.gz → 1.1.2__tar.gz

{memoryos-1.0.1 → memoryos-1.1.2}/PKG-INFO

@@ -1,8 +1,9 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: MemoryOS
-Version: 1.0.1
+Version: 1.1.2
 Summary: Intelligence Begins with Memory
 License: Apache-2.0
+License-File: LICENSE
 Keywords: memory,llm,language model,memoryOS,agent,kv cache,lora
 Author: MemTensor
 Author-email: MemTensor@memtensor.cn
@@ -22,6 +23,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Provides-Extra: all
 Provides-Extra: mem-reader
 Provides-Extra: mem-scheduler
+Provides-Extra: mem-user
 Provides-Extra: tree-mem
 Requires-Dist: chonkie (>=1.0.7,<2.0.0) ; extra == "all"
 Requires-Dist: chonkie (>=1.0.7,<2.0.0) ; extra == "mem-reader"
@@ -35,6 +37,9 @@ Requires-Dist: ollama (>=0.4.8,<0.5.0)
 Requires-Dist: openai (>=1.77.0,<2.0.0)
 Requires-Dist: pika (>=1.3.2,<2.0.0) ; extra == "all"
 Requires-Dist: pika (>=1.3.2,<2.0.0) ; extra == "mem-scheduler"
+Requires-Dist: pymysql (>=1.1.0,<2.0.0)
+Requires-Dist: pymysql (>=1.1.0,<2.0.0) ; extra == "all"
+Requires-Dist: pymysql (>=1.1.0,<2.0.0) ; extra == "mem-user"
 Requires-Dist: python-dateutil (>=2.9.0.post0,<3.0.0)
 Requires-Dist: qdrant-client (>=1.14.2,<2.0.0) ; extra == "all"
 Requires-Dist: redis (>=6.2.0,<7.0.0) ; extra == "all"

{memoryos-1.0.1 → memoryos-1.1.2}/pyproject.toml

@@ -4,7 +4,7 @@
 ##############################################################################
 
 name = "MemoryOS"
-version = "1.0.1"
+version = "1.1.2"
 description = "Intelligence Begins with Memory"
 license = {text = "Apache-2.0"}
 readme = "README.md"
@@ -42,6 +42,7 @@ dependencies = [
     "tenacity (>=9.1.2,<10.0.0)",  # Error handling and retrying library
     "fastapi[all] (>=0.115.12,<0.116.0)",  # Web framework for building APIs
     "sqlalchemy (>=2.0.41,<3.0.0)",  # SQL toolkit
+    "pymysql (>=1.1.0,<2.0.0)",  # MySQL Python driver
     "scikit-learn (>=1.7.0,<2.0.0)",  # Machine learning
     "fastmcp (>=2.10.5,<3.0.0)",
     "python-dateutil (>=2.9.0.post0,<3.0.0)",
@@ -76,6 +77,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 +96,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 +165,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/"

{memoryos-1.0.1 → memoryos-1.1.2}/src/memos/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.0.1"
+__version__ = "1.1.2"
 
 from memos.configs.mem_cube import GeneralMemCubeConfig
 from memos.configs.mem_os import MOSConfig

memoryos-1.1.2/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/v1"
+        )
+        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_message(
+        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 = {"user_id": user_id, "conversation_id": 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_message(
+        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, "user_id": user_id, "conversation_id": 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_memory(
+        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,
+            "user_id": user_id,
+            "conversation_id": conversation_id,
+            "memory_limit_number": 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

{memoryos-1.0.1 → memoryos-1.1.2}/src/memos/api/config.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")),
@@ -100,8 +100,10 @@ class APIConfig:
                 "backend": "http_bge",
                 "config": {
                     "url": os.getenv("MOS_RERANKER_URL"),
-                    "model": "bge-reranker-v2-m3",
+                    "model": os.getenv("MOS_RERANKER_MODEL", "bge-reranker-v2-m3"),
                     "timeout": 10,
+                    "headers_extra": os.getenv("MOS_RERANKER_HEADERS_EXTRA"),
+                    "rerank_source": os.getenv("MOS_RERANK_SOURCE"),
                 },
             }
         else:
@@ -186,22 +188,22 @@ class APIConfig:
         return {
             "uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"),
             "user": os.getenv("NEO4J_USER", "neo4j"),
-            "db_name": os.getenv("NEO4J_DB_NAME", "shared-tree-textual-memory"),
+            "db_name": os.getenv("NEO4J_DB_NAME", "neo4j"),
             "password": os.getenv("NEO4J_PASSWORD", "12345678"),
             "user_name": f"memos{user_id.replace('-', '')}",
-            "auto_create": True,
+            "auto_create": False,
             "use_multi_db": False,
-            "embedding_dimension": int(os.getenv("EMBEDDING_DIMENSION", 3072)),
+            "embedding_dimension": int(os.getenv("EMBEDDING_DIMENSION", 1024)),
             "vec_config": {
                 # Pass nested config to initialize external vector DB
                 # If you use qdrant, please use Server instead of local mode.
                 "backend": "qdrant",
                 "config": {
                     "collection_name": "neo4j_vec_db",
-                    "vector_dimension": int(os.getenv("EMBEDDING_DIMENSION", 3072)),
+                    "vector_dimension": int(os.getenv("EMBEDDING_DIMENSION", 1024)),
                     "distance_metric": "cosine",
-                    "host": "localhost",
-                    "port": 6333,
+                    "host": os.getenv("QDRANT_HOST", "localhost"),
+                    "port": int(os.getenv("QDRANT_PORT", "6333")),
                 },
             },
         }
@@ -271,7 +273,7 @@ class APIConfig:
     def get_scheduler_config() -> dict[str, Any]:
         """Get scheduler configuration."""
         return {
-            "backend": "general_scheduler",
+            "backend": "optimized_scheduler",
             "config": {
                 "top_k": int(os.getenv("MOS_SCHEDULER_TOP_K", "10")),
                 "act_mem_update_interval": int(

memoryos-1.1.2/src/memos/api/context/dependencies.py

@@ -0,0 +1,50 @@
+import logging
+
+from memos.context.context import RequestContext, get_current_context
+
+
+logger = logging.getLogger(__name__)
+
+# Type alias for the RequestContext from context module
+G = RequestContext
+
+
+def get_g_object() -> G:
+    """
+    Get Flask g-like object for the current request.
+    Returns the context created by middleware.
+    """
+    ctx = get_current_context()
+    if ctx is None:
+        raise RuntimeError(
+            "No request context available. Make sure RequestContextMiddleware is properly configured."
+        )
+    return ctx
+
+
+def get_current_g() -> G | None:
+    """
+    Get the current request's g object from anywhere in the application.
+
+    Returns:
+        The current request's g object if available, None otherwise.
+    """
+    return get_current_context()
+
+
+def require_g() -> G:
+    """
+    Get the current request's g object, raising an error if not available.
+
+    Returns:
+        The current request's g object.
+
+    Raises:
+        RuntimeError: If called outside of a request context.
+    """
+    ctx = get_current_context()
+    if ctx is None:
+        raise RuntimeError(
+            "No request context available. This function must be called within a request handler."
+        )
+    return ctx

{memoryos-1.0.1 → memoryos-1.1.2}/src/memos/api/middleware/request_context.py

@@ -2,40 +2,25 @@
 Request context middleware for automatic trace_id injection.
 """
 
-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
-
+import memos.log
 
-logger = logging.getLogger(__name__)
+from memos.context.context import RequestContext, generate_trace_id, set_request_context
 
 
-def generate_trace_id() -> str:
-    """Generate a random trace_id."""
-    return os.urandom(16).hex()
+logger = memos.log.get_logger(__name__)
 
 
 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")
-    if trace_id:
-        return trace_id
-
-    trace_id = request.headers.get("x-trace-id")
-    if trace_id:
-        return trace_id
-
-    trace_id = request.headers.get("trace-id")
-    if trace_id:
-        return trace_id
-
+    for header in ["g-trace-id", "x-trace-id", "trace-id"]:
+        if trace_id := request.headers.get(header):
+            return trace_id
     return None
 
 
@@ -51,19 +36,12 @@ class RequestContextMiddleware(BaseHTTPMiddleware):
 
     async def dispatch(self, request: Request, call_next: Callable) -> Response:
         # Extract or generate trace_id
-        trace_id = extract_trace_id_from_headers(request)
-        if not trace_id:
-            trace_id = generate_trace_id()
+        trace_id = extract_trace_id_from_headers(request) or generate_trace_id()
 
         # Create and set request context
-        context = RequestContext(trace_id=trace_id)
+        context = RequestContext(trace_id=trace_id, api_path=request.url.path)
         set_request_context(context)
 
-        # Add request metadata to context
-        context.set("method", request.method)
-        context.set("path", request.url.path)
-        context.set("client_ip", request.client.host if request.client else None)
-
         # Log request start with parameters
         params_log = {}
 
@@ -71,16 +49,7 @@ class RequestContextMiddleware(BaseHTTPMiddleware):
         if request.query_params:
             params_log["query_params"] = dict(request.query_params)
 
-        # Get request body if it's available
-        try:
-            params_log = await request.json()
-        except Exception as e:
-            logger.error(f"Error getting request body: {e}")
-            # If body is not JSON or empty, ignore it
-
-        logger.info(
-            f"Request started: {request.method} {request.url.path} - Parameters: {params_log}"
-        )
+        logger.info(f"Request started: {request.method} {request.url.path}, {params_log}")
 
         # Process the request
         response = await call_next(request)

{memoryos-1.0.1 → memoryos-1.1.2}/src/memos/api/product_api.py

@@ -17,9 +17,7 @@ app = FastAPI(
     version="1.0.1",
 )
 
-# Add request context middleware (must be added first)
 app.add_middleware(RequestContextMiddleware)
-
 # Include routers
 app.include_router(product_router)
 
@@ -35,5 +33,6 @@ if __name__ == "__main__":
 
     parser = argparse.ArgumentParser()
     parser.add_argument("--port", type=int, default=8001)
+    parser.add_argument("--workers", type=int, default=32)
     args = parser.parse_args()
-    uvicorn.run(app, host="0.0.0.0", port=args.port)
+    uvicorn.run("memos.api.product_api:app", host="0.0.0.0", port=args.port, workers=args.workers)

{memoryos-1.0.1 → memoryos-1.1.2}/src/memos/api/product_models.py

@@ -1,26 +1,14 @@
 import uuid
 
-from typing import Generic, Literal, TypeAlias, TypeVar
+from typing import Generic, Literal, TypeVar
 
 from pydantic import BaseModel, Field
-from typing_extensions import TypedDict
 
+# Import message types from core types module
+from memos.types import MessageDict
 
-T = TypeVar("T")
-
-
-# ─── Message Types ──────────────────────────────────────────────────────────────
-
-# Chat message roles
-MessageRole: TypeAlias = Literal["user", "assistant", "system"]
-
-
-# Message structure
-class MessageDict(TypedDict):
-    """Typed dictionary for chat message dictionaries."""
 
-    role: MessageRole
-    content: str
+T = TypeVar("T")
 
 
 class BaseRequest(BaseModel):
@@ -42,6 +30,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")
 
@@ -85,6 +74,7 @@ class ChatRequest(BaseRequest):
     history: list[MessageDict] | None = Field(None, description="Chat history")
     internet_search: bool = Field(True, description="Whether to use internet search")
     moscube: bool = Field(False, description="Whether to use MemOSCube")
+    session_id: str | None = Field(None, description="Session ID for soft-filtering memories")
 
 
 class ChatCompleteRequest(BaseRequest):
@@ -99,6 +89,7 @@ class ChatCompleteRequest(BaseRequest):
     base_prompt: str | None = Field(None, description="Base prompt to use for chat")
     top_k: int = Field(10, description="Number of results to return")
     threshold: float = Field(0.5, description="Threshold for filtering references")
+    session_id: str | None = Field(None, description="Session ID for soft-filtering memories")
 
 
 class UserCreate(BaseRequest):
@@ -160,6 +151,7 @@ class MemoryCreateRequest(BaseRequest):
     mem_cube_id: str | None = Field(None, description="Cube ID")
     source: str | None = Field(None, description="Source of the memory")
     user_profile: bool = Field(False, description="User profile memory")
+    session_id: str | None = Field(None, description="Session id")
 
 
 class SearchRequest(BaseRequest):
@@ -169,6 +161,7 @@ class SearchRequest(BaseRequest):
     query: str = Field(..., description="Search query")
     mem_cube_id: str | None = Field(None, description="Cube ID to search in")
     top_k: int = Field(10, description="Number of results to return")
+    session_id: str | None = Field(None, description="Session ID for soft-filtering memories")
 
 
 class SuggestionRequest(BaseRequest):
@@ -177,3 +170,85 @@ 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."""
+
+    model_config = {"extra": "allow"}
+
+
+class MemoryDetail(BaseModel):
+    """Individual memory detail model based on actual API response."""
+
+    model_config = {"extra": "allow"}
+
+
+class GetMessagesData(BaseModel):
+    """Data model for get messages response based on actual API."""
+
+    message_detail_list: list[MessageDetail] = Field(
+        default_factory=list, alias="message_detail_list", 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="memory_detail_list", description="List of memory details"
+    )
+    message_detail_list: list[MessageDetail] | None = Field(
+        None, alias="message_detail_list", 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

{memoryos-1.0.1 → memoryos-1.1.2}/src/memos/api/routers/product_router.py

@@ -1,14 +1,11 @@
 import json
+import time
 import traceback
 
-from datetime import datetime
-from typing import Annotated
-
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
 
 from memos.api.config import APIConfig
-from memos.api.context.dependencies import G, get_g_object
 from memos.api.product_models import (
     BaseResponse,
     ChatCompleteRequest,
@@ -79,24 +76,19 @@ def set_config(config):
 
 
 @router.post("/users/register", summary="Register a new user", response_model=UserRegisterResponse)
-def register_user(user_req: UserRegisterRequest, g: Annotated[G, Depends(get_g_object)]):
+def register_user(user_req: UserRegisterRequest):
     """Register a new user with configuration and default cube."""
     try:
-        # Set request-related information in g object
-        g.user_id = user_req.user_id
-        g.action = "user_register"
-        g.timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-
-        logger.info(f"Starting user registration for user_id: {user_req.user_id}")
-        logger.info(f"Request trace_id: {g.trace_id}")
-        logger.info(f"Request timestamp: {g.timestamp}")
-
         # Get configuration for the user
+        time_start_register = time.time()
         user_config, default_mem_cube = APIConfig.create_user_config(
             user_name=user_req.user_id, user_id=user_req.user_id
         )
         logger.info(f"user_config: {user_config.model_dump(mode='json')}")
         logger.info(f"default_mem_cube: {default_mem_cube.config.model_dump(mode='json')}")
+        logger.info(
+            f"time register api : create user config time user_id: {user_req.user_id} time is: {time.time() - time_start_register}"
+        )
         mos_product = get_mos_product_instance()
 
         # Register user with default config and mem cube
@@ -106,8 +98,11 @@ 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,
+        )
+        logger.info(
+            f"time register api : register time user_id: {user_req.user_id} time is: {time.time() - time_start_register}"
         )
-
         if result["status"] == "success":
             return UserRegisterResponse(
                 message="User registered successfully",
@@ -194,6 +189,7 @@ def get_all_memories(memory_req: GetMemoryRequest):
 def create_memory(memory_req: MemoryCreateRequest):
     """Create a new memory for a specific user."""
     try:
+        time_start_add = time.time()
         mos_product = get_mos_product_instance()
         mos_product.add(
             user_id=memory_req.user_id,
@@ -203,6 +199,10 @@ def create_memory(memory_req: MemoryCreateRequest):
             mem_cube_id=memory_req.mem_cube_id,
             source=memory_req.source,
             user_profile=memory_req.user_profile,
+            session_id=memory_req.session_id,
+        )
+        logger.info(
+            f"time add api : add time user_id: {memory_req.user_id} time is: {time.time() - time_start_add}"
         )
         return SimpleResponse(message="Memory created successfully")
 
@@ -217,12 +217,17 @@ def create_memory(memory_req: MemoryCreateRequest):
 def search_memories(search_req: SearchRequest):
     """Search memories for a specific user."""
     try:
+        time_start_search = time.time()
         mos_product = get_mos_product_instance()
         result = mos_product.search(
             query=search_req.query,
             user_id=search_req.user_id,
             install_cube_ids=[search_req.mem_cube_id] if search_req.mem_cube_id else None,
             top_k=search_req.top_k,
+            session_id=search_req.session_id,
+        )
+        logger.info(
+            f"time search api : add time user_id: {search_req.user_id} time is: {time.time() - time_start_search}"
         )
         return SearchResponse(message="Search completed successfully", data=result)
 
@@ -250,6 +255,7 @@ def chat(chat_req: ChatRequest):
                     history=chat_req.history,
                     internet_search=chat_req.internet_search,
                     moscube=chat_req.moscube,
+                    session_id=chat_req.session_id,
                 )
 
             except Exception as e:
@@ -294,6 +300,7 @@ def chat_complete(chat_req: ChatCompleteRequest):
             base_prompt=chat_req.base_prompt,
             top_k=chat_req.top_k,
             threshold=chat_req.threshold,
+            session_id=chat_req.session_id,
         )
 
         # Return the complete response