feat:add time step (#326)

* fix time

* add timed step for llm

* fix:ci code

* fix：test

* remove

Author: fridayL

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,18 +76,9 @@ 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
         user_config, default_mem_cube = APIConfig.create_user_config(
             user_name=user_req.user_id, user_id=user_req.user_id
@@ -195,6 +183,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,
@@ -206,6 +195,9 @@ def create_memory(memory_req: MemoryCreateRequest):
             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")
 
     except ValueError as err:
@@ -219,6 +211,7 @@ 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,
@@ -227,6 +220,9 @@ def search_memories(search_req: SearchRequest):
             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)
 
     except ValueError as err:

src/memos/mem_os/core.py

@@ -662,6 +662,8 @@ def add(
             "messages_or_doc_path or memory_content or doc_path must be provided."
         )
         # TODO: asure that session_id is a valid string
+        time_start = time.time()
+
         target_session_id = session_id if session_id else self.session_id
         target_user_id = user_id if user_id is not None else self.user_id
         if mem_cube_id is None:
@@ -674,14 +676,25 @@ def add(
             mem_cube_id = accessible_cubes[0].cube_id  # TODO not only first
         else:
             self._validate_cube_access(target_user_id, mem_cube_id)
+        logger.info(
+            f"time add: get mem_cube_id time user_id: {target_user_id} time is: {time.time() - time_start}"
+        )
 
+        time_start_0 = time.time()
         if mem_cube_id not in self.mem_cubes:
             raise ValueError(f"MemCube '{mem_cube_id}' is not loaded. Please register.")
+        logger.info(
+            f"time add: get mem_cube_id check in mem_cubes time user_id: {target_user_id} time is: {time.time() - time_start_0}"
+        )
+        time_start_1 = time.time()
         if (
             (messages is not None)
             and self.config.enable_textual_memory
             and self.mem_cubes[mem_cube_id].text_mem
         ):
+            logger.info(
+                f"time add: messages is not None and enable_textual_memory and text_mem is not None time user_id: {target_user_id} time is: {time.time() - time_start_1}"
+            )
             if self.mem_cubes[mem_cube_id].config.text_mem.backend != "tree_text":
                 add_memory = []
                 metadata = TextualMemoryMetadata(
@@ -694,12 +707,15 @@ def add(
                 self.mem_cubes[mem_cube_id].text_mem.add(add_memory)
             else:
                 messages_list = [messages]
+                time_start_2 = time.time()
                 memories = self.mem_reader.get_memory(
                     messages_list,
                     type="chat",
                     info={"user_id": target_user_id, "session_id": target_session_id},
                 )
-
+                logger.info(
+                    f"time add: get mem_reader time user_id: {target_user_id} time is: {time.time() - time_start_2}"
+                )
                 mem_ids = []
                 for mem in memories:
                     mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)

src/memos/mem_os/product.py

@@ -2,7 +2,6 @@
 import json
 import os
 import random
-import sys
 import time
 
 from collections.abc import Generator
@@ -47,6 +46,7 @@
     get_memos_prompt,
 )
 from memos.types import MessageList
+from memos.utils import timed
 
 
 logger = get_logger(__name__)
@@ -258,6 +258,7 @@ def _preload_user_cubes(
         except Exception as e:
             logger.error(f"Error pre-loading cubes for user {user_id}: {e}", exc_info=True)
 
+    @timed
     def _load_user_cubes(
         self, user_id: str, default_cube_config: GeneralMemCubeConfig | None = None
     ) -> None:
@@ -289,6 +290,7 @@ def _load_user_cubes(
                         )
                 except Exception as e:
                     logger.error(f"Failed to load cube {cube.cube_id} for user {user_id}: {e}")
+        logger.info(f"load user {user_id} cubes successfully")
 
     def _ensure_user_instance(self, user_id: str, max_instances: int | None = None) -> None:
         """
@@ -1320,14 +1322,6 @@ def search(
         # Load user cubes if not already loaded
         time_start = time.time()
         self._load_user_cubes(user_id, self.default_cube_config)
-        try:
-            dict_size = sys.getsizeof(self.mem_cubes)
-            size_mb = dict_size / (1024 * 1024)
-            logger.info(
-                f"now search memcubes_size is : len is {len(self.mem_cubes)} and {size_mb}MB"
-            )
-        except Exception as e:
-            logger.warning(f"Failed to get memcubes size: {e}, ignore it")
         load_user_cubes_time_end = time.time()
         logger.info(
             f"time search: load_user_cubes time user_id: {user_id} time is: {load_user_cubes_time_end - time_start}"
@@ -1376,12 +1370,6 @@ def add(
 
         # Load user cubes if not already loaded
         self._load_user_cubes(user_id, self.default_cube_config)
-        try:
-            dict_size = sys.getsizeof(self.mem_cubes)
-            size_mb = dict_size / (1024 * 1024)
-            logger.info(f"now add memcubes_size is : {len is len(self.mem_cubes)} and {size_mb}MB")
-        except Exception as e:
-            logger.warning(f"Failed to get memcubes size: {e}, ignore it")
         result = super().add(
             messages, memory_content, doc_path, mem_cube_id, user_id, session_id=session_id
         )

src/memos/mem_reader/simple_struct.py

@@ -27,6 +27,7 @@
     SIMPLE_STRUCT_MEM_READER_PROMPT,
     SIMPLE_STRUCT_MEM_READER_PROMPT_ZH,
 )
+from memos.utils import timed
 
 
 logger = log.get_logger(__name__)
@@ -126,6 +127,7 @@ def __init__(self, config: SimpleStructMemReaderConfig):
         self.embedder = EmbedderFactory.from_config(config.embedder)
         self.chunker = ChunkerFactory.from_config(config.chunker)
 
+    @timed
     def _process_chat_data(self, scene_data_info, info):
         mem_list = []
         for item in scene_data_info: