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

memos/configs/reranker.py

@@ -0,0 +1,18 @@
+# memos/configs/reranker.py
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class RerankerConfigFactory(BaseModel):
+    """
+    {
+      "backend": "http_bge" | "cosine_local" | "noop",
+      "config": { ... backend-specific ... }
+    }
+    """
+
+    backend: str = Field(..., description="Reranker backend id")
+    config: dict[str, Any] = Field(default_factory=dict, description="Backend-specific options")

memos/graph_dbs/base.py

@@ -81,7 +81,9 @@ class BaseGraphDB(ABC):
         """
 
     @abstractmethod
-    def get_nodes(self, id: str, include_embedding: bool = False) -> dict[str, Any] | None:
+    def get_nodes(
+        self, id: str, include_embedding: bool = False, **kwargs
+    ) -> dict[str, Any] | None:
         """
         Retrieve the metadata and memory of a list of nodes.
         Args:
@@ -141,7 +143,7 @@ class BaseGraphDB(ABC):
 
     # Search / recall operations
     @abstractmethod
-    def search_by_embedding(self, vector: list[float], top_k: int = 5) -> list[dict]:
+    def search_by_embedding(self, vector: list[float], top_k: int = 5, **kwargs) -> list[dict]:
         """
         Retrieve node IDs based on vector similarity.
 

memos/graph_dbs/nebular.py

@@ -1,10 +1,11 @@
+import json
 import traceback
 
 from contextlib import suppress
 from datetime import datetime
 from queue import Empty, Queue
 from threading import Lock
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 import numpy as np
 
@@ -15,6 +16,10 @@ from memos.log import get_logger
 from memos.utils import timed
 
 
+if TYPE_CHECKING:
+    from nebulagraph_python.client.pool import NebulaPool
+
+
 logger = get_logger(__name__)
 
 
@@ -35,7 +40,28 @@ def _compose_node(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]]:
 
 @timed
 def _escape_str(value: str) -> str:
-    return value.replace('"', '\\"')
+    out = []
+    for ch in value:
+        code = ord(ch)
+        if ch == "\\":
+            out.append("\\\\")
+        elif ch == '"':
+            out.append('\\"')
+        elif ch == "\n":
+            out.append("\\n")
+        elif ch == "\r":
+            out.append("\\r")
+        elif ch == "\t":
+            out.append("\\t")
+        elif ch == "\b":
+            out.append("\\b")
+        elif ch == "\f":
+            out.append("\\f")
+        elif code < 0x20 or code in (0x2028, 0x2029):
+            out.append(f"\\u{code:04x}")
+        else:
+            out.append(ch)
+    return "".join(out)
 
 
 @timed
@@ -197,6 +223,94 @@ class NebulaGraphDB(BaseGraphDB):
     NebulaGraph-based implementation of a graph memory store.
     """
 
+    # ====== shared pool cache & refcount ======
+    # These are process-local; in a multi-process model each process will
+    # have its own cache.
+    _POOL_CACHE: ClassVar[dict[str, "NebulaPool"]] = {}
+    _POOL_REFCOUNT: ClassVar[dict[str, int]] = {}
+    _POOL_LOCK: ClassVar[Lock] = Lock()
+
+    @staticmethod
+    def _make_pool_key(cfg: NebulaGraphDBConfig) -> str:
+        """
+        Build a cache key that captures all connection-affecting options.
+        Keep this key stable and include fields that change the underlying pool behavior.
+        """
+        # NOTE: Do not include tenant-like or query-scope-only fields here.
+        # Only include things that affect the actual TCP/auth/session pool.
+        return "|".join(
+            [
+                "nebula",
+                str(getattr(cfg, "uri", "")),
+                str(getattr(cfg, "user", "")),
+                str(getattr(cfg, "password", "")),
+                # pool sizing / tls / timeouts if you have them in config:
+                str(getattr(cfg, "max_client", 1000)),
+                # multi-db mode can impact how we use sessions; keep it to be safe
+                str(getattr(cfg, "use_multi_db", False)),
+            ]
+        )
+
+    @classmethod
+    def _get_or_create_shared_pool(cls, cfg: NebulaGraphDBConfig):
+        """
+        Get a shared NebulaPool from cache or create one if missing.
+        Thread-safe with a lock; maintains a simple refcount.
+        """
+        key = cls._make_pool_key(cfg)
+
+        with cls._POOL_LOCK:
+            pool = cls._POOL_CACHE.get(key)
+            if pool is None:
+                # Create a new pool and put into cache
+                pool = SessionPool(
+                    hosts=cfg.get("uri"),
+                    user=cfg.get("user"),
+                    password=cfg.get("password"),
+                    minsize=1,
+                    maxsize=cfg.get("max_client", 1000),
+                )
+                cls._POOL_CACHE[key] = pool
+                cls._POOL_REFCOUNT[key] = 0
+                logger.info(f"[NebulaGraphDB] Created new shared NebulaPool for key={key}")
+
+            # Increase refcount for the caller
+            cls._POOL_REFCOUNT[key] = cls._POOL_REFCOUNT.get(key, 0) + 1
+            return key, pool
+
+    @classmethod
+    def _release_shared_pool(cls, key: str):
+        """
+        Decrease refcount for the given pool key; only close when refcount hits zero.
+        """
+        with cls._POOL_LOCK:
+            if key not in cls._POOL_CACHE:
+                return
+            cls._POOL_REFCOUNT[key] = max(0, cls._POOL_REFCOUNT.get(key, 0) - 1)
+            if cls._POOL_REFCOUNT[key] == 0:
+                try:
+                    cls._POOL_CACHE[key].close()
+                except Exception as e:
+                    logger.warning(f"[NebulaGraphDB] Error closing shared pool: {e}")
+                finally:
+                    cls._POOL_CACHE.pop(key, None)
+                    cls._POOL_REFCOUNT.pop(key, None)
+                    logger.info(f"[NebulaGraphDB] Closed and removed shared pool key={key}")
+
+    @classmethod
+    def close_all_shared_pools(cls):
+        """Force close all cached pools. Call this on graceful shutdown."""
+        with cls._POOL_LOCK:
+            for key, pool in list(cls._POOL_CACHE.items()):
+                try:
+                    pool.close()
+                except Exception as e:
+                    logger.warning(f"[NebulaGraphDB] Error closing pool key={key}: {e}")
+                finally:
+                    logger.info(f"[NebulaGraphDB] Closed pool key={key}")
+            cls._POOL_CACHE.clear()
+            cls._POOL_REFCOUNT.clear()
+
     @require_python_package(
         import_name="nebulagraph_python",
         install_command="pip install ... @Tianxing",
@@ -246,20 +360,21 @@ class NebulaGraphDB(BaseGraphDB):
             "usage",
             "background",
         }
+        self.base_fields = set(self.common_fields) - {"usage"}
+        self.heavy_fields = {"usage"}
         self.dim_field = (
             f"embedding_{self.embedding_dimension}"
             if (str(self.embedding_dimension) != str(self.default_memory_dimension))
             else "embedding"
         )
         self.system_db_name = "system" if config.use_multi_db else config.space
-        self.pool = SessionPool(
-            hosts=config.get("uri"),
-            user=config.get("user"),
-            password=config.get("password"),
-            minsize=1,
-            maxsize=config.get("max_client", 1000),
-        )
 
+        # ---- NEW: pool acquisition strategy
+        # Get or create a shared pool from the class-level cache
+        self._pool_key, self.pool = self._get_or_create_shared_pool(config)
+        self._owns_pool = True  # We manage refcount for this instance
+
+        # auto-create graph type / graph / index if needed
         if config.auto_create:
             self._ensure_database_exists()
 
@@ -271,7 +386,7 @@ class NebulaGraphDB(BaseGraphDB):
         logger.info("Connected to NebulaGraph successfully.")
 
     @timed
-    def execute_query(self, gql: str, timeout: float = 5.0, auto_set_db: bool = True):
+    def execute_query(self, gql: str, timeout: float = 10.0, auto_set_db: bool = True):
         with self.pool.get() as client:
             try:
                 if auto_set_db and self.db_name:
@@ -287,7 +402,25 @@ class NebulaGraphDB(BaseGraphDB):
 
     @timed
     def close(self):
-        self.pool.close()
+        """
+        Close the connection resource if this instance owns it.
+
+        - If pool was injected (`shared_pool`), do nothing.
+        - If pool was acquired via shared cache, decrement refcount and close
+          when the last owner releases it.
+        """
+        if not self._owns_pool:
+            logger.debug("[NebulaGraphDB] close() skipped (injected pool).")
+            return
+        if self._pool_key:
+            self._release_shared_pool(self._pool_key)
+            self._pool_key = None
+            self.pool = None
+
+    # NOTE: __del__ is best-effort; do not rely on GC order.
+    def __del__(self):
+        with suppress(Exception):
+            self.close()
 
     @timed
     def create_index(
@@ -366,12 +499,10 @@ class NebulaGraphDB(BaseGraphDB):
             filter_clause = f'n.memory_type = "{scope}" AND n.user_name = "{self.config.user_name}"'
         else:
             filter_clause = f'n.memory_type = "{scope}"'
-        return_fields = ", ".join(f"n.{field} AS {field}" for field in self.common_fields)
-
         query = f"""
         MATCH (n@Memory)
         WHERE {filter_clause}
-        RETURN {return_fields}
+        RETURN n.id AS id
         LIMIT 1
         """
 
@@ -568,10 +699,7 @@ class NebulaGraphDB(BaseGraphDB):
         try:
             result = self.execute_query(gql)
             for row in result:
-                if include_embedding:
-                    props = row.values()[0].as_node().get_properties()
-                else:
-                    props = {k: v.value for k, v in row.items()}
+                props = {k: v.value for k, v in row.items()}
                 node = self._parse_node(props)
                 return node
 
@@ -582,7 +710,9 @@ class NebulaGraphDB(BaseGraphDB):
             return None
 
     @timed
-    def get_nodes(self, ids: list[str], include_embedding: bool = False) -> list[dict[str, Any]]:
+    def get_nodes(
+        self, ids: list[str], include_embedding: bool = False, **kwargs
+    ) -> list[dict[str, Any]]:
         """
         Retrieve the metadata and memory of a list of nodes.
         Args:
@@ -600,7 +730,10 @@ class NebulaGraphDB(BaseGraphDB):
 
         where_user = ""
         if not self.config.use_multi_db and self.config.user_name:
-            where_user = f" AND n.user_name = '{self.config.user_name}'"
+            if kwargs.get("cube_name"):
+                where_user = f" AND n.user_name = '{kwargs['cube_name']}'"
+            else:
+                where_user = f" AND n.user_name = '{self.config.user_name}'"
 
         # Safe formatting of the ID list
         id_list = ",".join(f'"{_id}"' for _id in ids)
@@ -615,10 +748,7 @@ class NebulaGraphDB(BaseGraphDB):
         try:
             results = self.execute_query(query)
             for row in results:
-                if include_embedding:
-                    props = row.values()[0].as_node().get_properties()
-                else:
-                    props = {k: v.value for k, v in row.items()}
+                props = {k: v.value for k, v in row.items()}
                 nodes.append(self._parse_node(props))
         except Exception as e:
             logger.error(
@@ -687,6 +817,7 @@ class NebulaGraphDB(BaseGraphDB):
         exclude_ids: list[str],
         top_k: int = 5,
         min_overlap: int = 1,
+        include_embedding: bool = False,
     ) -> list[dict[str, Any]]:
         """
         Find top-K neighbor nodes with maximum tag overlap.
@@ -696,6 +827,7 @@ class NebulaGraphDB(BaseGraphDB):
             exclude_ids: Node IDs to exclude (e.g., local cluster).
             top_k: Max number of neighbors to return.
             min_overlap: Minimum number of overlapping tags required.
+            include_embedding: with/without embedding
 
         Returns:
             List of dicts with node details and overlap count.
@@ -717,12 +849,13 @@ class NebulaGraphDB(BaseGraphDB):
         where_clause = " AND ".join(where_clauses)
         tag_list_literal = "[" + ", ".join(f'"{_escape_str(t)}"' for t in tags) + "]"
 
+        return_fields = self._build_return_fields(include_embedding)
         query = f"""
             LET tag_list = {tag_list_literal}
 
             MATCH (n@Memory)
             WHERE {where_clause}
-            RETURN n,
+            RETURN {return_fields},
                size( filter( n.tags, t -> t IN tag_list ) ) AS overlap_count
             ORDER BY overlap_count DESC
             LIMIT {top_k}
@@ -731,9 +864,8 @@ class NebulaGraphDB(BaseGraphDB):
         result = self.execute_query(query)
         neighbors: list[dict[str, Any]] = []
         for r in result:
-            node_props = r["n"].as_node().get_properties()
-            parsed = self._parse_node(node_props)  # --> {id, memory, metadata}
-
+            props = {k: v.value for k, v in r.items() if k != "overlap_count"}
+            parsed = self._parse_node(props)
             parsed["overlap_count"] = r["overlap_count"].value
             neighbors.append(parsed)
 
@@ -840,6 +972,7 @@ class NebulaGraphDB(BaseGraphDB):
         scope: str | None = None,
         status: str | None = None,
         threshold: float | None = None,
+        **kwargs,
     ) -> list[dict]:
         """
         Retrieve node IDs based on vector similarity.
@@ -874,7 +1007,10 @@ class NebulaGraphDB(BaseGraphDB):
         if status:
             where_clauses.append(f'n.status = "{status}"')
         if not self.config.use_multi_db and self.config.user_name:
-            where_clauses.append(f'n.user_name = "{self.config.user_name}"')
+            if kwargs.get("cube_name"):
+                where_clauses.append(f'n.user_name = "{kwargs["cube_name"]}"')
+            else:
+                where_clauses.append(f'n.user_name = "{self.config.user_name}"')
 
         where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
 
@@ -936,20 +1072,12 @@ class NebulaGraphDB(BaseGraphDB):
         """
         where_clauses = []
 
-        def _escape_value(value):
-            if isinstance(value, str):
-                return f'"{value}"'
-            elif isinstance(value, list):
-                return "[" + ", ".join(_escape_value(v) for v in value) + "]"
-            else:
-                return str(value)
-
         for _i, f in enumerate(filters):
             field = f["field"]
             op = f.get("op", "=")
             value = f["value"]
 
-            escaped_value = _escape_value(value)
+            escaped_value = self._format_value(value)
 
             # Build WHERE clause
             if op == "=":
@@ -1153,28 +1281,36 @@ class NebulaGraphDB(BaseGraphDB):
             data: A dictionary containing all nodes and edges to be loaded.
         """
         for node in data.get("nodes", []):
-            id, memory, metadata = _compose_node(node)
+            try:
+                id, memory, metadata = _compose_node(node)
 
-            if not self.config.use_multi_db and self.config.user_name:
-                metadata["user_name"] = self.config.user_name
+                if not self.config.use_multi_db and self.config.user_name:
+                    metadata["user_name"] = self.config.user_name
 
-            metadata = self._prepare_node_metadata(metadata)
-            metadata.update({"id": id, "memory": memory})
-            properties = ", ".join(f"{k}: {self._format_value(v, k)}" for k, v in metadata.items())
-            node_gql = f"INSERT OR IGNORE (n@Memory {{{properties}}})"
-            self.execute_query(node_gql)
+                metadata = self._prepare_node_metadata(metadata)
+                metadata.update({"id": id, "memory": memory})
+                properties = ", ".join(
+                    f"{k}: {self._format_value(v, k)}" for k, v in metadata.items()
+                )
+                node_gql = f"INSERT OR IGNORE (n@Memory {{{properties}}})"
+                self.execute_query(node_gql)
+            except Exception as e:
+                logger.error(f"Fail to load node: {node}, error: {e}")
 
         for edge in data.get("edges", []):
-            source_id, target_id = edge["source"], edge["target"]
-            edge_type = edge["type"]
-            props = ""
-            if not self.config.use_multi_db and self.config.user_name:
-                props = f'{{user_name: "{self.config.user_name}"}}'
-            edge_gql = f'''
-               MATCH (a@Memory {{id: "{source_id}"}}), (b@Memory {{id: "{target_id}"}})
-               INSERT OR IGNORE (a) -[e@{edge_type} {props}]-> (b)
-           '''
-            self.execute_query(edge_gql)
+            try:
+                source_id, target_id = edge["source"], edge["target"]
+                edge_type = edge["type"]
+                props = ""
+                if not self.config.use_multi_db and self.config.user_name:
+                    props = f'{{user_name: "{self.config.user_name}"}}'
+                edge_gql = f'''
+                   MATCH (a@Memory {{id: "{source_id}"}}), (b@Memory {{id: "{target_id}"}})
+                   INSERT OR IGNORE (a) -[e@{edge_type} {props}]-> (b)
+               '''
+                self.execute_query(edge_gql)
+            except Exception as e:
+                logger.error(f"Fail to load edge: {edge}, error: {e}")
 
     @timed
     def get_all_memory_items(self, scope: str, include_embedding: bool = False) -> (list)[dict]:
@@ -1208,10 +1344,7 @@ class NebulaGraphDB(BaseGraphDB):
         try:
             results = self.execute_query(query)
             for row in results:
-                if include_embedding:
-                    props = row.values()[0].as_node().get_properties()
-                else:
-                    props = {k: v.value for k, v in row.items()}
+                props = {k: v.value for k, v in row.items()}
                 nodes.append(self._parse_node(props))
         except Exception as e:
             logger.error(f"Failed to get memories: {e}")
@@ -1250,10 +1383,7 @@ class NebulaGraphDB(BaseGraphDB):
         try:
             results = self.execute_query(query)
             for row in results:
-                if include_embedding:
-                    props = row.values()[0].as_node().get_properties()
-                else:
-                    props = {k: v.value for k, v in row.items()}
+                props = {k: v.value for k, v in row.items()}
                 candidates.append(self._parse_node(props))
         except Exception as e:
             logger.error(f"Failed : {e}, traceback: {traceback.format_exc()}")
@@ -1555,6 +1685,7 @@ class NebulaGraphDB(BaseGraphDB):
         # Normalize embedding type
         embedding = metadata.get("embedding")
         if embedding and isinstance(embedding, list):
+            metadata.pop("embedding")
             metadata[self.dim_field] = _normalize([float(x) for x in embedding])
 
         return metadata
@@ -1563,12 +1694,22 @@ class NebulaGraphDB(BaseGraphDB):
     def _format_value(self, val: Any, key: str = "") -> str:
         from nebulagraph_python.py_data_types import NVector
 
+        # None
+        if val is None:
+            return "NULL"
+        # bool
+        if isinstance(val, bool):
+            return "true" if val else "false"
+        # str
         if isinstance(val, str):
             return f'"{_escape_str(val)}"'
+        # num
         elif isinstance(val, (int | float)):
             return str(val)
+        # time
         elif isinstance(val, datetime):
             return f'datetime("{val.isoformat()}")'
+        # list
         elif isinstance(val, list):
             if key == self.dim_field:
                 dim = len(val)
@@ -1576,13 +1717,18 @@ class NebulaGraphDB(BaseGraphDB):
                 return f"VECTOR<{dim}, FLOAT>([{joined}])"
             else:
                 return f"[{', '.join(self._format_value(v) for v in val)}]"
+        # NVector
         elif isinstance(val, NVector):
             if key == self.dim_field:
                 dim = len(val)
                 joined = ",".join(str(float(x)) for x in val)
                 return f"VECTOR<{dim}, FLOAT>([{joined}])"
-        elif val is None:
-            return "NULL"
+            else:
+                logger.warning("Invalid NVector")
+        # dict
+        if isinstance(val, dict):
+            j = json.dumps(val, ensure_ascii=False, separators=(",", ":"))
+            return f'"{_escape_str(j)}"'
         else:
             return f'"{_escape_str(str(val))}"'
 
@@ -1607,6 +1753,7 @@ class NebulaGraphDB(BaseGraphDB):
         return filtered_metadata
 
     def _build_return_fields(self, include_embedding: bool = False) -> str:
+        fields = set(self.base_fields)
         if include_embedding:
-            return "n"
-        return ", ".join(f"n.{field} AS {field}" for field in self.common_fields)
+            fields.add(self.dim_field)
+        return ", ".join(f"n.{f} AS {f}" for f in fields)

memos/graph_dbs/neo4j.py

@@ -323,12 +323,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 +344,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 +365,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 +606,7 @@ class Neo4jGraphDB(BaseGraphDB):
         scope: str | None = None,
         status: str | None = None,
         threshold: float | None = None,
+        **kwargs,
     ) -> list[dict]:
         """
         Retrieve node IDs based on vector similarity.
@@ -654,7 +656,10 @@ class Neo4jGraphDB(BaseGraphDB):
         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
 
         with self.driver.session(database=self.db_name) as session:
             result = session.run(query, parameters)
@@ -833,7 +838,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 +919,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 +950,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.

memos/graph_dbs/neo4j_community.py

@@ -129,6 +129,7 @@ class Neo4jCommunityGraphDB(Neo4jGraphDB):
         scope: str | None = None,
         status: str | None = None,
         threshold: float | None = None,
+        **kwargs,
     ) -> list[dict]:
         """
         Retrieve node IDs based on vector similarity using external vector DB.
@@ -157,7 +158,10 @@ 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
 
         # Perform vector search
         results = self.vec_db.search(query_vector=vector, top_k=top_k, filter=vec_filter)
@@ -169,13 +173,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/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)