MemoryOS 1.0.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

memos/log.py

@@ -1,12 +1,19 @@
+import atexit
 import logging
+import os
+import threading
 
 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
+from memos.api.context.context_thread import ContextThreadPoolExecutor
 
 
 # Load environment variables
@@ -26,27 +33,126 @@ 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
+                    cls._instance._executor = None
+                    cls._instance._session = None
+                    cls._instance._is_shutting_down = None
+        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 = ContextThreadPoolExecutor(
+                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
+
+        try:
+            trace_id = get_current_trace_id() or "no-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"""
+        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"
+        },
+        "simplified": {
+            "format": "%(asctime)s | %(trace_id)s | %(levelname)s | %(filename)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"],
+            "formatter": "simplified",
+            "filters": ["package_tree_filter", "trace_id_filter"],
         },
         "file": {
             "level": "DEBUG",
@@ -54,12 +160,18 @@ LOGGING_CONFIG = {
             "filename": _setup_logfile(),
             "maxBytes": 1024**2 * 10,
             "backupCount": 10,
-            "formatter": "standard",
+            "formatter": "simplified",
+            "filters": ["trace_id_filter"],
+        },
+        "custom_logger": {
+            "level": selected_log_level,
+            "class": "memos.log.CustomLoggerRequestHandler",
+            "formatter": "simplified",
         },
     },
     "root": {  # Root logger handles all logs
-        "level": logging.DEBUG if settings.DEBUG else logging.INFO,
-        "handlers": ["console", "file"],
+        "level": selected_log_level,
+        "handlers": ["console", "file", "custom_logger"],
     },
     "loggers": {
         "memos": {

memos/mem_os/core.py

@@ -352,6 +352,7 @@ class MOSCore:
         self,
         memories: list[TextualMemoryItem] | list[str] | None = None,
         base_prompt: str | None = None,
+        **kwargs,
     ) -> str:
         """Build system prompt with optional memories context."""
         if base_prompt is None:
@@ -545,6 +546,8 @@ class MOSCore:
         top_k: int | None = None,
         mode: Literal["fast", "fine"] = "fast",
         internet_search: bool = False,
+        moscube: bool = False,
+        **kwargs,
     ) -> MOSSearchResult:
         """
         Search for textual memories across all registered MemCubes.
@@ -602,6 +605,7 @@ class MOSCore:
                         "session_id": self.session_id,
                         "chat_history": chat_history.chat_history,
                     },
+                    moscube=moscube,
                 )
                 result["text_mem"].append({"cube_id": mem_cube_id, "memories": memories})
                 logger.info(
@@ -659,7 +663,7 @@ class MOSCore:
             if self.mem_cubes[mem_cube_id].config.text_mem.backend != "tree_text":
                 add_memory = []
                 metadata = TextualMemoryMetadata(
-                    user_id=self.user_id, session_id=self.session_id, source="conversation"
+                    user_id=target_user_id, session_id=self.session_id, source="conversation"
                 )
                 for message in messages:
                     add_memory.append(
@@ -956,6 +960,30 @@ class MOSCore:
         self.mem_cubes[mem_cube_id].dump(dump_dir)
         logger.info(f"MemCube {mem_cube_id} dumped to {dump_dir}")
 
+    def load(
+        self,
+        load_dir: str,
+        user_id: str | None = None,
+        mem_cube_id: str | None = None,
+        memory_types: list[Literal["text_mem", "act_mem", "para_mem"]] | None = None,
+    ) -> None:
+        """Dump the MemCube to a dictionary.
+        Args:
+            load_dir (str): The directory to load the MemCube from.
+            user_id (str, optional): The identifier of the user to load the MemCube from.
+                If None, the default user is used.
+            mem_cube_id (str, optional): The identifier of the MemCube to load.
+                If None, the default MemCube for the user is used.
+        """
+        target_user_id = user_id if user_id is not None else self.user_id
+        accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
+        if not mem_cube_id:
+            mem_cube_id = accessible_cubes[0].cube_id
+        if mem_cube_id not in self.mem_cubes:
+            raise ValueError(f"MemCube with ID {mem_cube_id} does not exist. please regiester")
+        self.mem_cubes[mem_cube_id].load(load_dir, memory_types=memory_types)
+        logger.info(f"MemCube {mem_cube_id} loaded from {load_dir}")
+
     def get_user_info(self) -> dict[str, Any]:
         """Get current user information including accessible cubes.
 
@@ -971,7 +999,7 @@ class MOSCore:
         return {
             "user_id": user.user_id,
             "user_name": user.user_name,
-            "role": user.role.value,
+            "role": user.role.value if hasattr(user.role, "value") else user.role,
             "created_at": user.created_at.isoformat(),
             "accessible_cubes": [
                 {