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

memos/graph_dbs/neo4j.py

@@ -1,3 +1,4 @@
+import json
 import time
 
 from datetime import datetime
@@ -174,6 +175,12 @@ class Neo4jGraphDB(BaseGraphDB):
                 n.updated_at = datetime($updated_at),
                 n += $metadata
         """
+
+        # serialization
+        if metadata["sources"]:
+            for idx in range(len(metadata["sources"])):
+                metadata["sources"][idx] = json.dumps(metadata["sources"][idx])
+
         with self.driver.session(database=self.db_name) as session:
             session.run(
                 query,
@@ -323,12 +330,11 @@ class Neo4jGraphDB(BaseGraphDB):
             return result.single() is not None
 
     # Graph Query & Reasoning
-    def get_node(self, id: str, include_embedding: bool = True) -> dict[str, Any] | None:
+    def get_node(self, id: str, **kwargs) -> dict[str, Any] | None:
         """
         Retrieve the metadata and memory of a node.
         Args:
             id: Node identifier.
-            include_embedding (bool): Whether to include the large embedding field.
         Returns:
             Dictionary of node fields, or None if not found.
         """
@@ -345,12 +351,11 @@ class Neo4jGraphDB(BaseGraphDB):
             record = session.run(query, params).single()
             return self._parse_node(dict(record["n"])) if record else None
 
-    def get_nodes(self, ids: list[str], include_embedding: bool = True) -> list[dict[str, Any]]:
+    def get_nodes(self, ids: list[str], **kwargs) -> list[dict[str, Any]]:
         """
         Retrieve the metadata and memory of a list of nodes.
         Args:
             ids: List of Node identifier.
-            include_embedding (bool): Whether to include the large embedding field.
         Returns:
         list[dict]: Parsed node records containing 'id', 'memory', and 'metadata'.
 
@@ -367,7 +372,10 @@ class Neo4jGraphDB(BaseGraphDB):
 
         if not self.config.use_multi_db and self.config.user_name:
             where_user = " AND n.user_name = $user_name"
-            params["user_name"] = self.config.user_name
+            if kwargs.get("cube_name"):
+                params["user_name"] = kwargs["cube_name"]
+            else:
+                params["user_name"] = self.config.user_name
 
         query = f"MATCH (n:Memory) WHERE n.id IN $ids{where_user} RETURN n"
 
@@ -605,6 +613,8 @@ class Neo4jGraphDB(BaseGraphDB):
         scope: str | None = None,
         status: str | None = None,
         threshold: float | None = None,
+        search_filter: dict | None = None,
+        **kwargs,
     ) -> list[dict]:
         """
         Retrieve node IDs based on vector similarity.
@@ -616,6 +626,8 @@ class Neo4jGraphDB(BaseGraphDB):
             status (str, optional): Node status filter (e.g., 'active', 'archived').
                             If provided, restricts results to nodes with matching status.
             threshold (float, optional): Minimum similarity score threshold (0 ~ 1).
+            search_filter (dict, optional): Additional metadata filters for search results.
+                            Keys should match node properties, values are the expected values.
 
         Returns:
             list[dict]: A list of dicts with 'id' and 'score', ordered by similarity.
@@ -625,6 +637,7 @@ class Neo4jGraphDB(BaseGraphDB):
             - If scope is provided, it restricts results to nodes with matching memory_type.
             - If 'status' is provided, only nodes with the matching status will be returned.
             - If threshold is provided, only results with score >= threshold will be returned.
+            - If search_filter is provided, additional WHERE clauses will be added for metadata filtering.
             - Typical use case: restrict to 'status = activated' to avoid
             matching archived or merged nodes.
         """
@@ -637,6 +650,12 @@ class Neo4jGraphDB(BaseGraphDB):
         if not self.config.use_multi_db and self.config.user_name:
             where_clauses.append("node.user_name = $user_name")
 
+        # Add search_filter conditions
+        if search_filter:
+            for key, _ in search_filter.items():
+                param_name = f"filter_{key}"
+                where_clauses.append(f"node.{key} = ${param_name}")
+
         where_clause = ""
         if where_clauses:
             where_clause = "WHERE " + " AND ".join(where_clauses)
@@ -648,13 +667,23 @@ class Neo4jGraphDB(BaseGraphDB):
             RETURN node.id AS id, score
         """
 
-        parameters = {"embedding": vector, "k": top_k, "scope": scope}
+        parameters = {"embedding": vector, "k": top_k}
+
         if scope:
             parameters["scope"] = scope
         if status:
             parameters["status"] = status
         if not self.config.use_multi_db and self.config.user_name:
-            parameters["user_name"] = self.config.user_name
+            if kwargs.get("cube_name"):
+                parameters["user_name"] = kwargs["cube_name"]
+            else:
+                parameters["user_name"] = self.config.user_name
+
+        # Add search_filter parameters
+        if search_filter:
+            for key, value in search_filter.items():
+                param_name = f"filter_{key}"
+                parameters[param_name] = value
 
         with self.driver.session(database=self.db_name) as session:
             result = session.run(query, parameters)
@@ -833,7 +862,7 @@ class Neo4jGraphDB(BaseGraphDB):
             logger.error(f"[ERROR] Failed to clear database '{self.db_name}': {e}")
             raise
 
-    def export_graph(self, include_embedding: bool = True) -> dict[str, Any]:
+    def export_graph(self, **kwargs) -> dict[str, Any]:
         """
         Export all graph nodes and edges in a structured form.
 
@@ -914,13 +943,12 @@ class Neo4jGraphDB(BaseGraphDB):
                     target_id=edge["target"],
                 )
 
-    def get_all_memory_items(self, scope: str, include_embedding: bool = True) -> list[dict]:
+    def get_all_memory_items(self, scope: str, **kwargs) -> list[dict]:
         """
         Retrieve all memory items of a specific memory_type.
 
         Args:
             scope (str): Must be one of 'WorkingMemory', 'LongTermMemory', or 'UserMemory'.
-            include_embedding (bool): Whether to include the large embedding field.
         Returns:
 
         Returns:
@@ -946,9 +974,7 @@ class Neo4jGraphDB(BaseGraphDB):
             results = session.run(query, params)
             return [self._parse_node(dict(record["n"])) for record in results]
 
-    def get_structure_optimization_candidates(
-        self, scope: str, include_embedding: bool = True
-    ) -> list[dict]:
+    def get_structure_optimization_candidates(self, scope: str, **kwargs) -> list[dict]:
         """
         Find nodes that are likely candidates for structure optimization:
         - Isolated nodes, nodes with empty background, or nodes with exactly one child.
@@ -1109,4 +1135,14 @@ class Neo4jGraphDB(BaseGraphDB):
                 node[time_field] = node[time_field].isoformat()
         node.pop("user_name", None)
 
+        # serialization
+        if node["sources"]:
+            for idx in range(len(node["sources"])):
+                if not (
+                    isinstance(node["sources"][idx], str)
+                    and node["sources"][idx][0] == "{"
+                    and node["sources"][idx][0] == "}"
+                ):
+                    break
+                node["sources"][idx] = json.loads(node["sources"][idx])
         return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node}

memos/graph_dbs/neo4j_community.py

@@ -129,6 +129,8 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
         scope: str | None = None,
         status: str | None = None,
         threshold: float | None = None,
+        search_filter: dict | None = None,
+        **kwargs,
     ) -> list[dict]:
         """
         Retrieve node IDs based on vector similarity using external vector DB.
@@ -139,6 +141,7 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
             scope (str, optional): Memory type filter (e.g., 'WorkingMemory', 'LongTermMemory').
             status (str, optional): Node status filter (e.g., 'activated', 'archived').
             threshold (float, optional): Minimum similarity score threshold (0 ~ 1).
+            search_filter (dict, optional): Additional metadata filters to apply.
 
         Returns:
             list[dict]: A list of dicts with 'id' and 'score', ordered by similarity.
@@ -148,6 +151,7 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
             - If 'scope' is provided, it restricts results to nodes with matching memory_type.
             - If 'status' is provided, it further filters nodes by status.
             - If 'threshold' is provided, only results with score >= threshold will be returned.
+            - If 'search_filter' is provided, it applies additional metadata-based filtering.
             - The returned IDs can be used to fetch full node data from Neo4j if needed.
         """
         # Build VecDB filter
@@ -157,7 +161,14 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
         if status:
             vec_filter["status"] = status
         vec_filter["vector_sync"] = "success"
-        vec_filter["user_name"] = self.config.user_name
+        if kwargs.get("cube_name"):
+            vec_filter["user_name"] = kwargs["cube_name"]
+        else:
+            vec_filter["user_name"] = self.config.user_name
+
+        # Add search_filter conditions
+        if search_filter:
+            vec_filter.update(search_filter)
 
         # Perform vector search
         results = self.vec_db.search(query_vector=vector, top_k=top_k, filter=vec_filter)
@@ -169,13 +180,12 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
         # Return consistent format
         return [{"id": r.id, "score": r.score} for r in results]
 
-    def get_all_memory_items(self, scope: str) -> list[dict]:
+    def get_all_memory_items(self, scope: str, **kwargs) -> list[dict]:
         """
         Retrieve all memory items of a specific memory_type.
 
         Args:
             scope (str): Must be one of 'WorkingMemory', 'LongTermMemory', or 'UserMemory'.
-
         Returns:
             list[dict]: Full list of memory items under this scope.
         """

memos/llms/factory.py

@@ -9,6 +9,7 @@ from memos.llms.ollama import OllamaLLM
 from memos.llms.openai import AzureLLM, OpenAILLM
 from memos.llms.qwen import QwenLLM
 from memos.llms.vllm import VLLMLLM
+from memos.memos_tools.singleton import singleton_factory
 
 
 class LLMFactory(BaseLLM):
@@ -26,6 +27,7 @@ class LLMFactory(BaseLLM):
     }
 
     @classmethod
+    @singleton_factory()
     def from_config(cls, config_factory: LLMConfigFactory) -> BaseLLM:
         backend = config_factory.backend
         if backend not in cls.backend_to_class:

memos/llms/openai.py

@@ -1,4 +1,8 @@
+import hashlib
+import json
+
 from collections.abc import Generator
+from typing import ClassVar
 
 import openai
 
@@ -13,11 +17,44 @@ logger = get_logger(__name__)
 
 
 class OpenAILLM(BaseLLM):
-    """OpenAI LLM class."""
+    """OpenAI LLM class with singleton pattern."""
+
+    _instances: ClassVar[dict] = {}  # Class variable to store instances
+
+    def __new__(cls, config: OpenAILLMConfig) -> "OpenAILLM":
+        config_hash = cls._get_config_hash(config)
+
+        if config_hash not in cls._instances:
+            logger.info(f"Creating new OpenAI LLM instance for config hash: {config_hash}")
+            instance = super().__new__(cls)
+            cls._instances[config_hash] = instance
+        else:
+            logger.info(f"Reusing existing OpenAI LLM instance for config hash: {config_hash}")
+
+        return cls._instances[config_hash]
 
     def __init__(self, config: OpenAILLMConfig):
+        # Avoid duplicate initialization
+        if hasattr(self, "_initialized"):
+            return
+
         self.config = config
         self.client = openai.Client(api_key=config.api_key, base_url=config.api_base)
+        self._initialized = True
+        logger.info("OpenAI LLM instance initialized")
+
+    @classmethod
+    def _get_config_hash(cls, config: OpenAILLMConfig) -> str:
+        """Generate hash value of configuration"""
+        config_dict = config.model_dump()
+        config_str = json.dumps(config_dict, sort_keys=True)
+        return hashlib.md5(config_str.encode()).hexdigest()
+
+    @classmethod
+    def clear_cache(cls):
+        """Clear all cached instances"""
+        cls._instances.clear()
+        logger.info("OpenAI LLM instance cache cleared")
 
     def generate(self, messages: MessageList) -> str:
         """Generate a response from OpenAI LLM."""
@@ -71,15 +108,50 @@ class OpenAILLM(BaseLLM):
 
 
 class AzureLLM(BaseLLM):
-    """Azure OpenAI LLM class."""
+    """Azure OpenAI LLM class with singleton pattern."""
+
+    _instances: ClassVar[dict] = {}  # Class variable to store instances
+
+    def __new__(cls, config: AzureLLMConfig):
+        # Generate hash value of config as cache key
+        config_hash = cls._get_config_hash(config)
+
+        if config_hash not in cls._instances:
+            logger.info(f"Creating new Azure LLM instance for config hash: {config_hash}")
+            instance = super().__new__(cls)
+            cls._instances[config_hash] = instance
+        else:
+            logger.info(f"Reusing existing Azure LLM instance for config hash: {config_hash}")
+
+        return cls._instances[config_hash]
 
     def __init__(self, config: AzureLLMConfig):
+        # Avoid duplicate initialization
+        if hasattr(self, "_initialized"):
+            return
+
         self.config = config
         self.client = openai.AzureOpenAI(
             azure_endpoint=config.base_url,
             api_version=config.api_version,
             api_key=config.api_key,
         )
+        self._initialized = True
+        logger.info("Azure LLM instance initialized")
+
+    @classmethod
+    def _get_config_hash(cls, config: AzureLLMConfig) -> str:
+        """Generate hash value of configuration"""
+        # Convert config to dict and sort to ensure consistency
+        config_dict = config.model_dump()
+        config_str = json.dumps(config_dict, sort_keys=True)
+        return hashlib.md5(config_str.encode()).hexdigest()
+
+    @classmethod
+    def clear_cache(cls):
+        """Clear all cached instances"""
+        cls._instances.clear()
+        logger.info("Azure LLM instance cache cleared")
 
     def generate(self, messages: MessageList) -> str:
         """Generate a response from Azure OpenAI LLM."""

memos/llms/vllm.py

@@ -105,6 +105,7 @@ class VLLMLLM(BaseLLM):
                 "temperature": float(getattr(self.config, "temperature", 0.8)),
                 "max_tokens": int(getattr(self.config, "max_tokens", 1024)),
                 "top_p": float(getattr(self.config, "top_p", 0.9)),
+                "extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
             }
 
             response = self.client.chat.completions.create(**completion_kwargs)
@@ -142,6 +143,7 @@ class VLLMLLM(BaseLLM):
                 "max_tokens": int(getattr(self.config, "max_tokens", 1024)),
                 "top_p": float(getattr(self.config, "top_p", 0.9)),
                 "stream": True,  # Enable streaming
+                "extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
             }
 
             stream = self.client.chat.completions.create(**completion_kwargs)

memos/log.py

@@ -1,12 +1,20 @@
+import atexit
 import logging
+import os
+import threading
+import time
 
+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.context.context import get_current_api_path, get_current_trace_id
 
 
 # Load environment variables
@@ -26,19 +34,129 @@ 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 "trace-id"
+        except Exception:
+            record.trace_id = "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 = 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
+
+        # Only process INFO and ERROR level logs
+        if record.levelno < logging.INFO:  # Skip DEBUG and lower
+            return
+
+        try:
+            trace_id = get_current_trace_id() or "trace-id"
+            api_path = get_current_api_path()
+            if api_path is not None:
+                self._executor.submit(self._send_log_sync, record.getMessage(), trace_id, api_path)
+        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, api_path):
+        """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,
+                "action": api_path,
+                "current_time": round(time.time(), 3),
+            }
+
+            # 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:%(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": {
@@ -46,7 +164,7 @@ LOGGING_CONFIG = {
             "class": "logging.StreamHandler",
             "stream": stdout,
             "formatter": "no_datetime",
-            "filters": ["package_tree_filter"],
+            "filters": ["package_tree_filter", "trace_id_filter"],
         },
         "file": {
             "level": "DEBUG",
@@ -55,6 +173,12 @@ LOGGING_CONFIG = {
             "maxBytes": 1024**2 * 10,
             "backupCount": 10,
             "formatter": "standard",
+            "filters": ["trace_id_filter"],
+        },
+        "custom_logger": {
+            "level": "INFO",
+            "class": "memos.log.CustomLoggerRequestHandler",
+            "formatter": "simplified",
         },
     },
     "root": {  # Root logger handles all logs

memos/mem_cube/general.py

@@ -1,4 +1,5 @@
 import os
+import time
 
 from typing import Literal
 
@@ -23,11 +24,13 @@ class GeneralMemCube(BaseMemCube):
     def __init__(self, config: GeneralMemCubeConfig):
         """Initialize the MemCube with a configuration."""
         self.config = config
+        time_start = time.time()
         self._text_mem: BaseTextMemory | None = (
             MemoryFactory.from_config(config.text_mem)
             if config.text_mem.backend != "uninitialized"
             else None
         )
+        logger.info(f"init_text_mem in {time.time() - time_start} seconds")
         self._act_mem: BaseActMemory | None = (
             MemoryFactory.from_config(config.act_mem)
             if config.act_mem.backend != "uninitialized"
@@ -137,7 +140,6 @@ class GeneralMemCube(BaseMemCube):
         if default_config is not None:
             config = merge_config_with_default(config, default_config)
             logger.info(f"Applied default config to cube {config.cube_id}")
-
         mem_cube = GeneralMemCube(config)
         mem_cube.load(dir, memory_types)
         return mem_cube