Merge branch 'feat/reranker' of github.com:CaralHsi/MemOSRealPublic into feat/reranker

Author: CaralHsi

src/memos/api/context/dependencies.py

@@ -1,5 +1,4 @@
 import logging
-import os
 
 from fastapi import Depends, Header, Request
 
@@ -25,13 +24,6 @@ def get_trace_id_from_header(
     return g_trace_id or x_trace_id or trace_id
 
 
-def generate_trace_id() -> str:
-    """
-    Get a random trace_id.
-    """
-    return os.urandom(16).hex()
-
-
 def get_request_context(
     request: Request, trace_id: str | None = Depends(get_trace_id_from_header)
 ) -> RequestContext:
@@ -65,9 +57,6 @@ def get_g_object(trace_id: str | None = Depends(get_trace_id_from_header)) -> G:
     This creates a RequestContext and sets it globally for access
     throughout the request lifecycle.
     """
-    if trace_id is None:
-        trace_id = generate_trace_id()
-
     g = RequestContext(trace_id=trace_id)
     set_request_context(g)
     logger.info(f"Request g object created with trace_id: {g.trace_id}")

src/memos/api/middleware/request_context.py

@@ -0,0 +1,81 @@
+"""
+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
+
+
+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")
+    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
+
+    return None
+
+
+class RequestContextMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to automatically inject request context for every HTTP request.
+
+    This middleware:
+    1. Extracts trace_id from headers or generates a new one
+    2. Creates a RequestContext and sets it globally
+    3. Ensures the context is available throughout the request lifecycle
+    """
+
+    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()
+
+        # Create and set request context
+        context = RequestContext(trace_id=trace_id)
+        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
+        logger.info(f"Request started: {request.method} {request.url.path} - trace_id: {trace_id}")
+
+        # Process the request
+        response = await call_next(request)
+
+        # Log request completion
+        logger.info(
+            f"Request completed: {request.method} {request.url.path} - trace_id: {trace_id} - status: {response.status_code}"
+        )
+
+        # Add trace_id to response headers for debugging
+        response.headers["x-trace-id"] = trace_id
+
+        return response

src/memos/api/product_api.py

@@ -3,6 +3,7 @@
 from fastapi import FastAPI
 
 from memos.api.exceptions import APIExceptionHandler
+from memos.api.middleware.request_context import RequestContextMiddleware
 from memos.api.routers.product_router import router as product_router
 
 
@@ -16,6 +17,9 @@
     version="1.0.0",
 )
 
+# Add request context middleware (must be added first)
+app.add_middleware(RequestContextMiddleware)
+
 # Include routers
 app.include_router(product_router)
 

src/memos/api/product_models.py

@@ -97,6 +97,8 @@ class ChatCompleteRequest(BaseRequest):
     internet_search: bool = Field(False, description="Whether to use internet search")
     moscube: bool = Field(False, description="Whether to use MemOSCube")
     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")
 
 
 class UserCreate(BaseRequest):

src/memos/api/routers/product_router.py

@@ -1,5 +1,5 @@
 import json
-import logging
+from memos.log import get_logger
 import traceback
 
 from datetime import datetime
@@ -30,7 +30,7 @@
 from memos.memos_tools.notification_service import get_error_bot_function, get_online_bot_function
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 router = APIRouter(prefix="/product", tags=["Product API"])
 
@@ -284,18 +284,23 @@ def chat_complete(chat_req: ChatCompleteRequest):
         mos_product = get_mos_product_instance()
 
         # Collect all responses from the generator
-        content = mos_product.chat(
+        content, references = mos_product.chat(
             query=chat_req.query,
             user_id=chat_req.user_id,
             cube_id=chat_req.mem_cube_id,
             history=chat_req.history,
             internet_search=chat_req.internet_search,
             moscube=chat_req.moscube,
             base_prompt=chat_req.base_prompt,
+            top_k=chat_req.top_k,
+            threshold=chat_req.threshold,
         )
 
         # Return the complete response
-        return {"message": "Chat completed successfully", "data": {"response": content}}
+        return {
+            "message": "Chat completed successfully",
+            "data": {"response": content, "references": references},
+        }
 
     except ValueError as err:
         raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err

src/memos/api/start_api.py

@@ -9,6 +9,7 @@
 from fastapi.responses import JSONResponse, RedirectResponse
 from pydantic import BaseModel, Field
 
+from memos.api.middleware.request_context import RequestContextMiddleware
 from memos.configs.mem_os import MOSConfig
 from memos.mem_os.main import MOS
 from memos.mem_user.user_manager import UserManager, UserRole
@@ -78,6 +79,8 @@ def get_mos_instance():
     version="1.0.0",
 )
 
+app.add_middleware(RequestContextMiddleware)
+
 
 class BaseRequest(BaseModel):
     """Base model for all requests."""

src/memos/log.py

@@ -1,12 +1,19 @@
+import atexit
 import logging
+import os
+import threading
 
+from concurrent.futures import ThreadPoolExecutor
 from logging.config import dictConfig
 from pathlib import Path
 from sys import stdout
 
+import requests
+
 from dotenv import load_dotenv
 
 from memos import settings
+from memos.api.context.context import get_current_trace_id
 
 
 # Load environment variables
@@ -26,27 +33,125 @@ def _setup_logfile() -> Path:
     return logfile
 
 
+class TraceIDFilter(logging.Filter):
+    """add trace_id to the log record"""
+
+    def filter(self, record):
+        try:
+            trace_id = get_current_trace_id()
+            record.trace_id = trace_id if trace_id else "no-trace-id"
+        except Exception:
+            record.trace_id = "no-trace-id"
+        return True
+
+
+class CustomLoggerRequestHandler(logging.Handler):
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self):
+        """Initialize handler with minimal setup"""
+        if not self._initialized:
+            super().__init__()
+            workers = int(os.getenv("CUSTOM_LOGGER_WORKERS", "2"))
+            self._executor = ThreadPoolExecutor(
+                max_workers=workers, thread_name_prefix="log_sender"
+            )
+            self._is_shutting_down = threading.Event()
+            self._session = requests.Session()
+            self._initialized = True
+            atexit.register(self._cleanup)
+
+    def emit(self, record):
+        """Process log records of INFO or ERROR level (non-blocking)"""
+        if os.getenv("CUSTOM_LOGGER_URL") is None or self._is_shutting_down.is_set():
+            return
+
+        if record.levelno in (logging.INFO, logging.ERROR):
+            try:
+                trace_id = (
+                    get_current_trace_id()
+                )  # TODO: get trace_id from request context instead of get_current_trace_id
+                if trace_id:
+                    self._executor.submit(self._send_log_sync, record.getMessage(), trace_id)
+            except Exception as e:
+                if not self._is_shutting_down.is_set():
+                    print(f"Error sending log: {e}")
+
+    def _send_log_sync(self, message, trace_id):
+        """Send log message synchronously in a separate thread"""
+        print(f"send_log_sync: {message} {trace_id}")
+        try:
+            logger_url = os.getenv("CUSTOM_LOGGER_URL")
+            token = os.getenv("CUSTOM_LOGGER_TOKEN")
+
+            headers = {"Content-Type": "application/json"}
+            post_content = {"message": message, "trace_id": trace_id}
+
+            # Add auth token if exists
+            if token:
+                headers["Authorization"] = f"Bearer {token}"
+
+            # Add traceId to headers for consistency
+            headers["traceId"] = trace_id
+
+            # Add custom attributes from env
+            for key, value in os.environ.items():
+                if key.startswith("CUSTOM_LOGGER_ATTRIBUTE_"):
+                    attribute_key = key[len("CUSTOM_LOGGER_ATTRIBUTE_") :].lower()
+                    post_content[attribute_key] = value
+
+            self._session.post(logger_url, headers=headers, json=post_content, timeout=5)
+        except Exception:
+            # Silently ignore errors to avoid affecting main application
+            pass
+
+    def _cleanup(self):
+        """Clean up resources during program exit"""
+        if not self._initialized:
+            return
+
+        self._is_shutting_down.set()
+        try:
+            self._executor.shutdown(wait=False)
+            self._session.close()
+        except Exception as e:
+            print(f"Error during cleanup: {e}")
+
+    def close(self):
+        """Override close to prevent premature shutdown"""
+
+
 LOGGING_CONFIG = {
     "version": 1,
     "disable_existing_loggers": False,
     "formatters": {
         "standard": {
-            "format": "%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
+            "format": "%(asctime)s [%(trace_id)s] - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
         },
         "no_datetime": {
-            "format": "%(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
+            "format": "[%(trace_id)s] - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
         },
     },
     "filters": {
-        "package_tree_filter": {"()": "logging.Filter", "name": settings.LOG_FILTER_TREE_PREFIX}
+        "package_tree_filter": {"()": "logging.Filter", "name": settings.LOG_FILTER_TREE_PREFIX},
+        "trace_id_filter": {"()": "memos.log.TraceIDFilter"},
     },
     "handlers": {
         "console": {
             "level": selected_log_level,
             "class": "logging.StreamHandler",
             "stream": stdout,
             "formatter": "no_datetime",
-            "filters": ["package_tree_filter"],
+            "filters": ["package_tree_filter", "trace_id_filter"],
         },
         "file": {
             "level": "DEBUG",
@@ -55,6 +160,7 @@ def _setup_logfile() -> Path:
             "maxBytes": 1024**2 * 10,
             "backupCount": 10,
             "formatter": "standard",
+            "filters": ["trace_id_filter"],
         },
     },
     "root": {  # Root logger handles all logs