Merge remote-tracking branch 'upstream/dev' into dev

Author: Nyakult

src/memos/api/config.py

@@ -15,6 +15,7 @@
 
 from memos.configs.mem_cube import GeneralMemCubeConfig
 from memos.configs.mem_os import MOSConfig
+from memos.context.context import ContextThread
 from memos.mem_cube.general import GeneralMemCube
 
 
@@ -178,7 +179,6 @@ def start_watch_if_enabled(cls) -> None:
         if not enable:
             return
         interval = int(os.getenv("NACOS_WATCH_INTERVAL", "60"))
-        import threading
 
         def _loop() -> None:
             while True:
@@ -188,7 +188,7 @@ def _loop() -> None:
                     logger.error(f"❌ Nacos watch loop error: {e}")
                 time.sleep(interval)
 
-        threading.Thread(target=_loop, daemon=True).start()
+        ContextThread(target=_loop, daemon=True).start()
         logger.info(f"Nacos watch thread started (interval={interval}s).")
 
     @classmethod
@@ -861,9 +861,9 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
                             "reorganize": os.getenv("MOS_ENABLE_REORGANIZE", "false").lower()
                             == "true",
                             "memory_size": {
-                                "WorkingMemory": os.getenv("NEBULAR_WORKING_MEMORY", 20),
-                                "LongTermMemory": os.getenv("NEBULAR_LONGTERM_MEMORY", 1e6),
-                                "UserMemory": os.getenv("NEBULAR_USER_MEMORY", 1e6),
+                                "WorkingMemory": int(os.getenv("NEBULAR_WORKING_MEMORY", 20)),
+                                "LongTermMemory": int(os.getenv("NEBULAR_LONGTERM_MEMORY", 1e6)),
+                                "UserMemory": int(os.getenv("NEBULAR_USER_MEMORY", 1e6)),
                             },
                             "search_strategy": {
                                 "fast_graph": bool(os.getenv("FAST_GRAPH", "false") == "true"),
@@ -933,9 +933,9 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None:
                             == "true",
                             "internet_retriever": internet_config,
                             "memory_size": {
-                                "WorkingMemory": os.getenv("NEBULAR_WORKING_MEMORY", 20),
-                                "LongTermMemory": os.getenv("NEBULAR_LONGTERM_MEMORY", 1e6),
-                                "UserMemory": os.getenv("NEBULAR_USER_MEMORY", 1e6),
+                                "WorkingMemory": int(os.getenv("NEBULAR_WORKING_MEMORY", 20)),
+                                "LongTermMemory": int(os.getenv("NEBULAR_LONGTERM_MEMORY", 1e6)),
+                                "UserMemory": int(os.getenv("NEBULAR_USER_MEMORY", 1e6)),
                             },
                             "search_strategy": {
                                 "fast_graph": bool(os.getenv("FAST_GRAPH", "false") == "true"),

src/memos/api/middleware/request_context.py

@@ -72,10 +72,18 @@ async def dispatch(self, request: Request, call_next: Callable) -> Response:
             f"headers: {request.headers}"
         )
 
+        response = await call_next(request)
+        end_time = time.time()
+
         # Process the request
         try:
-            response = await call_next(request)
-            end_time = time.time()
+            if not response:
+                logger.error(
+                    f"Request Failed No Response, path: {request.url.path}, status: {response.status_code}, cost: {(end_time - start_time) * 1000:.2f}ms"
+                )
+
+                return response
+
             if response.status_code == 200:
                 logger.info(
                     f"Request completed: source: {self.source}, path: {request.url.path}, status: {response.status_code}, cost: {(end_time - start_time) * 1000:.2f}ms"
@@ -89,6 +97,5 @@ async def dispatch(self, request: Request, call_next: Callable) -> Response:
             logger.error(
                 f"Request Exception Error: source: {self.source}, path: {request.url.path}, error: {e}, cost: {(end_time - start_time) * 1000:.2f}ms"
             )
-            raise e
 
         return response

src/memos/api/routers/server_router.py

@@ -349,7 +349,7 @@ def search_memories(search_req: APISearchRequest):
         mem_cube_id=search_req.mem_cube_id,
         session_id=search_req.session_id or "default_session",
     )
-    logger.info(f"Search user_id is: {user_context.mem_cube_id}")
+    logger.info(f"Search Req is: {search_req}")
     memories_result: MOSSearchResult = {
         "text_mem": [],
         "act_mem": [],
@@ -502,6 +502,9 @@ def add_memories(add_req: APIADDRequest):
         mem_cube_id=add_req.mem_cube_id,
         session_id=add_req.session_id or "default_session",
     )
+
+    logger.info(f"Add Req is: {add_req}")
+
     target_session_id = add_req.session_id
     if not target_session_id:
         target_session_id = "default_session"

src/memos/graph_dbs/polardb.py

@@ -1740,9 +1740,11 @@ def get_grouped_counts(
         for field in group_fields:
             alias = field.replace(".", "_")
             return_fields.append(
-                f"ag_catalog.agtype_access_operator(properties, '\"{field}\"'::agtype) AS {alias}"
+                f"ag_catalog.agtype_access_operator(properties, '\"{field}\"'::agtype)::text AS {alias}"
+            )
+            group_by_fields.append(
+                f"ag_catalog.agtype_access_operator(properties, '\"{field}\"'::agtype)::text"
             )
-            group_by_fields.append(alias)
 
         # Full SQL query construction
         query = f"""
@@ -1751,7 +1753,6 @@ def get_grouped_counts(
             {where_clause}
             GROUP BY {", ".join(group_by_fields)}
         """
-
         conn = self._get_connection()
         try:
             with conn.cursor() as cursor:
@@ -1772,7 +1773,7 @@ def get_grouped_counts(
                         else:
                             group_values[field] = str(value)
                     count_value = row[-1]  # Last column is count
-                    output.append({**group_values, "count": count_value})
+                    output.append({**group_values, "count": int(count_value)})
 
                 return output
 

src/memos/mem_reader/simple_struct.py

@@ -67,9 +67,18 @@ def detect_lang(text):
     try:
         if not text or not isinstance(text, str):
             return "en"
+        cleaned_text = text
+        # remove role and timestamp
+        cleaned_text = re.sub(
+            r"\b(user|assistant|query|answer)\s*:", "", cleaned_text, flags=re.IGNORECASE
+        )
+        cleaned_text = re.sub(r"\[[\d\-:\s]+\]", "", cleaned_text)
+
+        # extract chinese characters
         chinese_pattern = r"[\u4e00-\u9fff\u3400-\u4dbf\U00020000-\U0002a6df\U0002a700-\U0002b73f\U0002b740-\U0002b81f\U0002b820-\U0002ceaf\uf900-\ufaff]"
-        chinese_chars = re.findall(chinese_pattern, text)
-        if len(chinese_chars) / len(re.sub(r"[\s\d\W]", "", text)) > 0.3:
+        chinese_chars = re.findall(chinese_pattern, cleaned_text)
+        text_without_special = re.sub(r"[\s\d\W]", "", cleaned_text)
+        if text_without_special and len(chinese_chars) / len(text_without_special) > 0.3:
             return "zh"
         return "en"
     except Exception:
@@ -466,15 +475,11 @@ def get_scene_data_info(self, scene_data: list, type: str) -> list[str]:
         if type == "chat":
             for items in scene_data:
                 result = []
-                for item in items:
-                    # Convert dictionary to string
-                    if "chat_time" in item:
-                        result.append(item)
-                    else:
-                        result.append(item)
+                for i, item in enumerate(items):
+                    result.append(item)
                     if len(result) >= 10:
                         results.append(result)
-                        context = copy.deepcopy(result[-2:])
+                        context = copy.deepcopy(result[-2:]) if i + 1 < len(items) else []
                         result = context
                 if result:
                     results.append(result)

src/memos/mem_scheduler/base_scheduler.py

@@ -12,6 +12,7 @@
 from sqlalchemy.engine import Engine
 
 from memos.configs.mem_scheduler import AuthConfig, BaseSchedulerConfig
+from memos.context.context import ContextThread
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
@@ -689,7 +690,7 @@ def start(self) -> None:
             logger.info("Message consumer process started")
         else:
             # Default to thread mode
-            self._consumer_thread = threading.Thread(
+            self._consumer_thread = ContextThread(
                 target=self._message_consumer,
                 daemon=True,
                 name="MessageConsumerThread",

src/memos/mem_scheduler/general_modules/task_threads.py

@@ -5,6 +5,7 @@
 from concurrent.futures import as_completed
 from typing import Any, TypeVar
 
+from memos.context.context import ContextThread
 from memos.log import get_logger
 from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
 
@@ -138,7 +139,7 @@ def worker(task_name: str, func: Callable, args: tuple):
 
             # Start all threads
             for task_name, (func, args) in tasks.items():
-                thread = threading.Thread(
+                thread = ContextThread(
                     target=worker, args=(task_name, func, args), name=f"task-{task_name}"
                 )
                 threads[task_name] = thread
@@ -283,7 +284,7 @@ def run_race(
 
         # Create and start threads for each task
         for task_name, task_func in tasks.items():
-            thread = threading.Thread(
+            thread = ContextThread(
                 target=self.worker, args=(task_func, task_name), name=f"race-{task_name}"
             )
             self.threads[task_name] = thread

src/memos/mem_scheduler/monitors/dispatcher_monitor.py

@@ -4,7 +4,7 @@
 from time import perf_counter
 
 from memos.configs.mem_scheduler import BaseSchedulerConfig
-from memos.context.context import ContextThreadPoolExecutor
+from memos.context.context import ContextThread, ContextThreadPoolExecutor
 from memos.log import get_logger
 from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
 from memos.mem_scheduler.general_modules.dispatcher import SchedulerDispatcher
@@ -340,7 +340,7 @@ def start(self) -> bool:
             return False
 
         self._running = True
-        self._monitor_thread = threading.Thread(
+        self._monitor_thread = ContextThread(
             target=self._monitor_loop, name="threadpool_monitor", daemon=True
         )
         self._monitor_thread.start()

src/memos/mem_scheduler/webservice_modules/rabbitmq_service.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 
 from memos.configs.mem_scheduler import AuthConfig, RabbitMQConfig
+from memos.context.context import ContextThread
 from memos.dependency import require_python_package
 from memos.log import get_logger
 from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
@@ -96,7 +97,7 @@ def initialize_rabbitmq(
             )
 
             # Start IOLoop in dedicated thread
-            self._io_loop_thread = threading.Thread(
+            self._io_loop_thread = ContextThread(
                 target=self.rabbitmq_connection.ioloop.start, daemon=True
             )
             self._io_loop_thread.start()

src/memos/mem_scheduler/webservice_modules/redis_service.py

@@ -1,12 +1,12 @@
 import asyncio
 import os
 import subprocess
-import threading
 import time
 
 from collections.abc import Callable
 from typing import Any
 
+from memos.context.context import ContextThread
 from memos.dependency import require_python_package
 from memos.log import get_logger
 from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
@@ -41,7 +41,7 @@ def __init__(self):
         self.query_list_capacity = 1000
 
         self._redis_listener_running = False
-        self._redis_listener_thread: threading.Thread | None = None
+        self._redis_listener_thread: ContextThread | None = None
         self._redis_listener_loop: asyncio.AbstractEventLoop | None = None
 
     @property
@@ -336,7 +336,7 @@ def redis_start_listening(self, handler: Callable | None = None):
         if handler is None:
             handler = self.redis_consume_message_stream
 
-        self._redis_listener_thread = threading.Thread(
+        self._redis_listener_thread = ContextThread(
             target=self._redis_run_listener_async,
             args=(handler,),
             daemon=True,