feat: update add context

Author: fridayL

src/memos/api/context/context.py

@@ -0,0 +1,147 @@
+"""
+Global request context management for trace_id and request-scoped data.
+
+This module provides optional trace_id functionality that can be enabled
+when using the API components. It uses ContextVar to ensure thread safety
+and request isolation.
+"""
+
+import uuid
+
+from collections.abc import Callable
+from contextvars import ContextVar
+from typing import Any
+
+
+# Global context variable for request-scoped data
+_request_context: ContextVar[dict[str, Any] | None] = ContextVar("request_context", default=None)
+
+
+class RequestContext:
+    """
+    Request-scoped context object that holds trace_id and other request data.
+
+    This provides a Flask g-like object for FastAPI applications.
+    """
+
+    def __init__(self, trace_id: str | None = None):
+        self.trace_id = trace_id or str(uuid.uuid4())
+        self._data: dict[str, Any] = {}
+
+    def set(self, key: str, value: Any) -> None:
+        """Set a value in the context."""
+        self._data[key] = value
+
+    def get(self, key: str, default: Any | None = None) -> Any:
+        """Get a value from the context."""
+        return self._data.get(key, default)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name.startswith("_") or name == "trace_id":
+            super().__setattr__(name, value)
+        else:
+            if not hasattr(self, "_data"):
+                super().__setattr__(name, value)
+            else:
+                self._data[name] = value
+
+    def __getattr__(self, name: str) -> Any:
+        if hasattr(self, "_data") and name in self._data:
+            return self._data[name]
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert context to dictionary."""
+        return {"trace_id": self.trace_id, "data": self._data.copy()}
+
+
+def set_request_context(context: RequestContext) -> None:
+    """
+    Set the current request context.
+
+    This is typically called by the API dependency injection system.
+    """
+    _request_context.set(context.to_dict())
+
+
+def get_current_trace_id() -> str | None:
+    """
+    Get the current request's trace_id.
+
+    Returns:
+        The trace_id if available, None otherwise.
+    """
+    context = _request_context.get()
+    if context:
+        return context.get("trace_id")
+    return None
+
+
+def get_current_context() -> RequestContext | None:
+    """
+    Get the current request context.
+
+    Returns:
+        The current RequestContext if available, None otherwise.
+    """
+    context_dict = _request_context.get()
+    if context_dict:
+        ctx = RequestContext(trace_id=context_dict.get("trace_id"))
+        ctx._data = context_dict.get("data", {}).copy()
+        return ctx
+    return None
+
+
+def require_context() -> RequestContext:
+    """
+    Get the current request context, raising an error if not available.
+
+    Returns:
+        The current RequestContext.
+
+    Raises:
+        RuntimeError: If called outside of a request context.
+    """
+    context = get_current_context()
+    if context is None:
+        raise RuntimeError(
+            "No request context available. This function must be called within a request handler."
+        )
+    return context
+
+
+# Type for trace_id getter function
+TraceIdGetter = Callable[[], str | None]
+
+# Global variable to hold the trace_id getter function
+_trace_id_getter: TraceIdGetter | None = None
+
+
+def set_trace_id_getter(getter: TraceIdGetter) -> None:
+    """
+    Set a custom trace_id getter function.
+
+    This allows the logging system to retrieve trace_id without importing
+    API-specific modules.
+    """
+    global _trace_id_getter
+    _trace_id_getter = getter
+
+
+def get_trace_id_for_logging() -> str | None:
+    """
+    Get trace_id for logging purposes.
+
+    This function is used by the logging system and will use either
+    the custom getter function or fall back to the default context.
+    """
+    if _trace_id_getter:
+        try:
+            return _trace_id_getter()
+        except Exception:
+            pass
+    return get_current_trace_id()
+
+
+# Initialize the default trace_id getter
+set_trace_id_getter(get_current_trace_id)

src/memos/api/context/dependencies.py

@@ -0,0 +1,90 @@
+import logging
+
+from fastapi import Depends, Header, Request
+
+from memos.api.context.context import RequestContext, set_request_context
+
+
+logger = logging.getLogger(__name__)
+
+# Type alias for the RequestContext from context module
+G = RequestContext
+
+
+def get_trace_id_from_header(
+    trace_id: str | None = Header(None, alias="trace-id"),
+    x_trace_id: str | None = Header(None, alias="x-trace-id"),
+    g_trace_id: str | None = Header(None, alias="g-trace-id"),
+) -> str | None:
+    """
+    Extract trace_id from various possible headers.
+
+    Priority: g-trace-id > x-trace-id > trace-id
+    """
+    return g_trace_id or x_trace_id or trace_id
+
+
+def get_request_context(
+    request: Request, trace_id: str | None = Depends(get_trace_id_from_header)
+) -> RequestContext:
+    """
+    Get request context object with trace_id and request metadata.
+
+    This function creates a RequestContext and automatically sets it
+    in the global context for use throughout the request lifecycle.
+    """
+    # Create context object
+    ctx = RequestContext(trace_id=trace_id)
+
+    # Set the context globally for this request
+    set_request_context(ctx)
+
+    # Log request start
+    logger.info(f"Request started with trace_id: {ctx.trace_id}")
+
+    # Add request metadata to context
+    ctx.set("method", request.method)
+    ctx.set("path", request.url.path)
+    ctx.set("client_ip", request.client.host if request.client else None)
+
+    return ctx
+
+
+def get_g_object(trace_id: str | None = Depends(get_trace_id_from_header)) -> G:
+    """
+    Get Flask g-like object for the current request.
+
+    This creates a RequestContext and sets it globally for access
+    throughout the request lifecycle.
+    """
+    g = RequestContext(trace_id=trace_id)
+    set_request_context(g)
+    logger.info(f"Request g object created with trace_id: {g.trace_id}")
+    return g
+
+
+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.
+    """
+    from memos.context import get_current_context
+
+    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.
+    """
+    from memos.context import require_context
+
+    return require_context()

src/memos/api/routers/product_router.py

@@ -2,10 +2,14 @@
 import logging
 import traceback
 
-from fastapi import APIRouter, HTTPException
+from datetime import datetime
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, 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,
     ChatRequest,
@@ -64,9 +68,18 @@ async def set_config(config):
 
 
 @router.post("/users/register", summary="Register a new user", response_model=UserRegisterResponse)
-async def register_user(user_req: UserRegisterRequest):
+async def register_user(user_req: UserRegisterRequest, g: Annotated[G, Depends(get_g_object)]):
     """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
         user_config, default_mem_cube = APIConfig.create_user_config(
             user_name=user_req.user_id, user_id=user_req.user_id

src/memos/mem_os/product.py

@@ -498,6 +498,14 @@ def _send_message_to_scheduler(
             )
             self.mem_scheduler.submit_messages(messages=[message_item])
 
+    def _filter_memories_by_threshold(
+        self, memories: list[TextualMemoryItem], threshold: float = 0.20
+    ) -> list[TextualMemoryItem]:
+        """
+        Filter memories by threshold.
+        """
+        return [memory for memory in memories if memory.metadata.relativity >= threshold]
+
     def register_mem_cube(
         self,
         mem_cube_name_or_path_or_object: str | GeneralMemCube,
@@ -689,6 +697,7 @@ def chat_with_references(
         user_id: str,
         cube_id: str | None = None,
         history: MessageList | None = None,
+        top_k: int = 10,
     ) -> Generator[str, None, None]:
         """
         Chat with LLM with memory references and streaming output.
@@ -710,10 +719,11 @@ def chat_with_references(
         time_start = time.time()
         memories_list = []
         memories_result = super().search(
-            query, user_id, install_cube_ids=[cube_id] if cube_id else None, top_k=10
+            query, user_id, install_cube_ids=[cube_id] if cube_id else None, top_k=top_k
         )["text_mem"]
         if memories_result:
             memories_list = memories_result[0]["memories"]
+            memories_list = self._filter_memories_by_threshold(memories_list)
 
         # Build custom system prompt with relevant memories
         system_prompt = self._build_system_prompt(user_id, memories_list)
@@ -766,7 +776,7 @@ def chat_with_references(
         # Initialize buffer for streaming
         buffer = ""
         full_response = ""
-
+        token_count = 0
         # Use tiktoken for proper token-based chunking
         if self.config.chat_model.backend not in ["huggingface", "vllm"]:
             # For non-huggingface backends, we need to collect the full response first
@@ -779,6 +789,7 @@ def chat_with_references(
         for chunk in response_stream:
             if chunk in ["<think>", "</think>"]:
                 continue
+            token_count += 1
             buffer += chunk
             full_response += chunk
 
@@ -809,7 +820,8 @@ def chat_with_references(
 
         yield f"data: {json.dumps({'type': 'reference', 'data': reference})}\n\n"
         total_time = round(float(time_end - time_start), 1)
-        yield f"data: {json.dumps({'type': 'time', 'data': {'total_time': total_time, 'speed_improvement': '23%'}})}\n\n"
+        speed_improvement = round(float((len(system_prompt) / 2) * 0.0048 + 44.5), 1)
+        yield f"data: {json.dumps({'type': 'time', 'data': {'total_time': total_time, 'speed_improvement': f'{speed_improvement}%'}})}\n\n"
         logger.info(f"user_id: {user_id}, cube_id: {cube_id}, current_messages: {current_messages}")
         logger.info(f"user_id: {user_id}, cube_id: {cube_id}, full_response: {full_response}")
         self._send_message_to_scheduler(
@@ -935,6 +947,7 @@ def get_subgraph(
         user_id: str,
         query: str,
         mem_cube_ids: list[str] | None = None,
+        top_k: int = 20,
     ) -> list[dict[str, Any]]:
         """Get all memory items for a user.
 
@@ -950,7 +963,7 @@ def get_subgraph(
         # Load user cubes if not already loaded
         self._load_user_cubes(user_id, self.default_cube_config)
         memory_list = self._get_subgraph(
-            query=query, mem_cube_id=mem_cube_ids[0], user_id=user_id, top_k=20
+            query=query, mem_cube_id=mem_cube_ids[0], user_id=user_id, top_k=top_k
         )["text_mem"]
         reformat_memory_list = []
         for memory in memory_list:

src/memos/mem_os/utils/format_utils.py

@@ -533,7 +533,7 @@ def convert_graph_to_tree_forworkmem(
         node_name = extract_node_name(memory)
         memory_key = node.get("metadata", {}).get("key", node_name)
         usage = node.get("metadata", {}).get("usage", [])
-        frequency = len(usage)
+        frequency = len(usage) if len(usage) < 100 else 100
         node_map[node["id"]] = {
             "id": node["id"],
             "value": memory,