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

memos/graph_dbs/nebular.py

@@ -1,10 +1,10 @@
+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,9 +15,28 @@ from memos.log import get_logger
 from memos.utils import timed
 
 
+if TYPE_CHECKING:
+    from nebulagraph_python import (
+        NebulaClient,
+    )
+
+
 logger = get_logger(__name__)
 
 
+_TRANSIENT_ERR_KEYS = (
+    "Session not found",
+    "Connection not established",
+    "timeout",
+    "deadline exceeded",
+    "Broken pipe",
+    "EOFError",
+    "socket closed",
+    "connection reset",
+    "connection refused",
+)
+
+
 @timed
 def _normalize(vec: list[float]) -> list[float]:
     v = np.asarray(vec, dtype=np.float32)
@@ -35,7 +54,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
@@ -61,145 +101,202 @@ def _normalize_datetime(val):
     return str(val)
 
 
-class SessionPoolError(Exception):
-    pass
-
-
-class SessionPool:
-    @require_python_package(
-        import_name="nebulagraph_python",
-        install_command="pip install ... @Tianxing",
-        install_link=".....",
-    )
-    def __init__(
-        self,
-        hosts: list[str],
-        user: str,
-        password: str,
-        minsize: int = 1,
-        maxsize: int = 10000,
-    ):
-        self.hosts = hosts
-        self.user = user
-        self.password = password
-        self.minsize = minsize
-        self.maxsize = maxsize
-        self.pool = Queue(maxsize)
-        self.lock = Lock()
-
-        self.clients = []
-
-        for _ in range(minsize):
-            self._create_and_add_client()
-
-    @timed
-    def _create_and_add_client(self):
-        from nebulagraph_python import NebulaClient
-
-        client = NebulaClient(self.hosts, self.user, self.password)
-        self.pool.put(client)
-        self.clients.append(client)
-
-    @timed
-    def get_client(self, timeout: float = 5.0):
-        try:
-            return self.pool.get(timeout=timeout)
-        except Empty:
-            with self.lock:
-                if len(self.clients) < self.maxsize:
-                    from nebulagraph_python import NebulaClient
-
-                    client = NebulaClient(self.hosts, self.user, self.password)
-                    self.clients.append(client)
-                    return client
-            raise RuntimeError("NebulaClientPool exhausted") from None
-
-    @timed
-    def return_client(self, client):
-        try:
-            client.execute("YIELD 1")
-            self.pool.put(client)
-        except Exception:
-            logger.info("[Pool] Client dead, replacing...")
-            self.replace_client(client)
-
-    @timed
-    def close(self):
-        for client in self.clients:
-            with suppress(Exception):
-                client.close()
-        self.clients.clear()
+class NebulaGraphDB(BaseGraphDB):
+    """
+    NebulaGraph-based implementation of a graph memory store.
+    """
 
-    @timed
-    def get(self):
-        """
-        Context manager: with pool.get() as client:
-        """
+    # ====== shared pool cache & refcount ======
+    # These are process-local; in a multi-process model each process will
+    # have its own cache.
+    _CLIENT_CACHE: ClassVar[dict[str, "NebulaClient"]] = {}
+    _CLIENT_REFCOUNT: ClassVar[dict[str, int]] = {}
+    _CLIENT_LOCK: ClassVar[Lock] = Lock()
+    _CLIENT_INIT_DONE: ClassVar[set[str]] = set()
+
+    @staticmethod
+    def _get_hosts_from_cfg(cfg: NebulaGraphDBConfig) -> list[str]:
+        hosts = getattr(cfg, "uri", None) or getattr(cfg, "hosts", None)
+        if isinstance(hosts, str):
+            return [hosts]
+        return list(hosts or [])
+
+    @staticmethod
+    def _make_client_key(cfg: NebulaGraphDBConfig) -> str:
+        hosts = NebulaGraphDB._get_hosts_from_cfg(cfg)
+        return "|".join(
+            [
+                "nebula-sync",
+                ",".join(hosts),
+                str(getattr(cfg, "user", "")),
+                str(getattr(cfg, "use_multi_db", False)),
+                str(getattr(cfg, "space", "")),
+            ]
+        )
 
-        class _ClientContext:
-            def __init__(self, outer):
-                self.outer = outer
-                self.client = None
+    @classmethod
+    def _bootstrap_admin(cls, cfg: NebulaGraphDBConfig, client: "NebulaClient") -> "NebulaGraphDB":
+        tmp = object.__new__(NebulaGraphDB)
+        tmp.config = cfg
+        tmp.db_name = cfg.space
+        tmp.user_name = getattr(cfg, "user_name", None)
+        tmp.embedding_dimension = getattr(cfg, "embedding_dimension", 3072)
+        tmp.default_memory_dimension = 3072
+        tmp.common_fields = {
+            "id",
+            "memory",
+            "user_name",
+            "user_id",
+            "session_id",
+            "status",
+            "key",
+            "confidence",
+            "tags",
+            "created_at",
+            "updated_at",
+            "memory_type",
+            "sources",
+            "source",
+            "node_type",
+            "visibility",
+            "usage",
+            "background",
+        }
+        tmp.base_fields = set(tmp.common_fields) - {"usage"}
+        tmp.heavy_fields = {"usage"}
+        tmp.dim_field = (
+            f"embedding_{tmp.embedding_dimension}"
+            if str(tmp.embedding_dimension) != str(tmp.default_memory_dimension)
+            else "embedding"
+        )
+        tmp.system_db_name = "system" if getattr(cfg, "use_multi_db", False) else cfg.space
+        tmp._client = client
+        tmp._owns_client = False
+        return tmp
+
+    @classmethod
+    def _get_or_create_shared_client(cls, cfg: NebulaGraphDBConfig) -> tuple[str, "NebulaClient"]:
+        from nebulagraph_python import (
+            ConnectionConfig,
+            NebulaClient,
+            SessionConfig,
+            SessionPoolConfig,
+        )
 
-            def __enter__(self):
-                self.client = self.outer.get_client()
-                return self.client
+        key = cls._make_client_key(cfg)
+        with cls._CLIENT_LOCK:
+            client = cls._CLIENT_CACHE.get(key)
+            if client is None:
+                # Connection setting
+                conn_conf: ConnectionConfig | None = getattr(cfg, "conn_config", None)
+                if conn_conf is None:
+                    conn_conf = ConnectionConfig.from_defults(
+                        cls._get_hosts_from_cfg(cfg),
+                        getattr(cfg, "ssl_param", None),
+                    )
+
+                sess_conf = SessionConfig(graph=getattr(cfg, "space", None))
+                pool_conf = SessionPoolConfig(
+                    size=int(getattr(cfg, "max_client", 1000)), wait_timeout=5000
+                )
 
-            def __exit__(self, exc_type, exc_val, exc_tb):
-                if self.client:
-                    self.outer.return_client(self.client)
+                client = NebulaClient(
+                    hosts=conn_conf.hosts,
+                    username=cfg.user,
+                    password=cfg.password,
+                    conn_config=conn_conf,
+                    session_config=sess_conf,
+                    session_pool_config=pool_conf,
+                )
+                cls._CLIENT_CACHE[key] = client
+                cls._CLIENT_REFCOUNT[key] = 0
+                logger.info(f"[NebulaGraphDBSync] Created shared NebulaClient key={key}")
 
-        return _ClientContext(self)
+            cls._CLIENT_REFCOUNT[key] = cls._CLIENT_REFCOUNT.get(key, 0) + 1
 
-    @timed
-    def reset_pool(self):
-        """⚠️ Emergency reset: Close all clients and clear the pool."""
-        logger.warning("[Pool] Resetting all clients. Existing sessions will be lost.")
-        with self.lock:
-            for client in self.clients:
+            if getattr(cfg, "auto_create", False) and key not in cls._CLIENT_INIT_DONE:
                 try:
-                    client.close()
-                except Exception:
-                    logger.error("Fail to close!!!")
-            self.clients.clear()
-            while not self.pool.empty():
+                    pass
+                finally:
+                    pass
+
+        if getattr(cfg, "auto_create", False) and key not in cls._CLIENT_INIT_DONE:
+            with cls._CLIENT_LOCK:
+                if key not in cls._CLIENT_INIT_DONE:
+                    admin = cls._bootstrap_admin(cfg, client)
+                    try:
+                        admin._ensure_database_exists()
+                        admin._create_basic_property_indexes()
+                        admin._create_vector_index(
+                            dimensions=int(
+                                admin.embedding_dimension or admin.default_memory_dimension
+                            ),
+                        )
+                        cls._CLIENT_INIT_DONE.add(key)
+                        logger.info("[NebulaGraphDBSync] One-time init done")
+                    except Exception:
+                        logger.exception("[NebulaGraphDBSync] One-time init failed")
+
+        return key, client
+
+    def _refresh_client(self):
+        """
+        refresh NebulaClient:
+        """
+        old_key = getattr(self, "_client_key", None)
+        if not old_key:
+            return
+
+        cls = self.__class__
+        with cls._CLIENT_LOCK:
+            try:
+                if old_key in cls._CLIENT_CACHE:
+                    try:
+                        cls._CLIENT_CACHE[old_key].close()
+                    except Exception as e:
+                        logger.warning(f"[refresh_client] close old client error: {e}")
+                    finally:
+                        cls._CLIENT_CACHE.pop(old_key, None)
+            finally:
+                cls._CLIENT_REFCOUNT[old_key] = 0
+
+            new_key, new_client = cls._get_or_create_shared_client(self.config)
+            self._client_key = new_key
+            self._client = new_client
+            logger.info(f"[NebulaGraphDBSync] client refreshed: {old_key} -> {new_key}")
+
+    @classmethod
+    def _release_shared_client(cls, key: str):
+        with cls._CLIENT_LOCK:
+            if key not in cls._CLIENT_CACHE:
+                return
+            cls._CLIENT_REFCOUNT[key] = max(0, cls._CLIENT_REFCOUNT.get(key, 0) - 1)
+            if cls._CLIENT_REFCOUNT[key] == 0:
                 try:
-                    self.pool.get_nowait()
-                except Empty:
-                    break
-            for _ in range(self.minsize):
-                self._create_and_add_client()
-        logger.info("[Pool] Pool has been reset successfully.")
-
-    @timed
-    def replace_client(self, client):
-        try:
-            client.close()
-        except Exception:
-            logger.error("Fail to close client")
-
-        if client in self.clients:
-            self.clients.remove(client)
-
-        from nebulagraph_python import NebulaClient
-
-        new_client = NebulaClient(self.hosts, self.user, self.password)
-        self.clients.append(new_client)
-
-        self.pool.put(new_client)
-
-        logger.info("[Pool] Replaced dead client with a new one.")
-        return new_client
-
-
-class NebulaGraphDB(BaseGraphDB):
-    """
-    NebulaGraph-based implementation of a graph memory store.
-    """
+                    cls._CLIENT_CACHE[key].close()
+                except Exception as e:
+                    logger.warning(f"[NebulaGraphDBSync] Error closing client: {e}")
+                finally:
+                    cls._CLIENT_CACHE.pop(key, None)
+                    cls._CLIENT_REFCOUNT.pop(key, None)
+                    logger.info(f"[NebulaGraphDBSync] Closed & removed client key={key}")
+
+    @classmethod
+    def close_all_shared_clients(cls):
+        with cls._CLIENT_LOCK:
+            for key, client in list(cls._CLIENT_CACHE.items()):
+                try:
+                    client.close()
+                except Exception as e:
+                    logger.warning(f"[NebulaGraphDBSync] Error closing client {key}: {e}")
+                finally:
+                    logger.info(f"[NebulaGraphDBSync] Closed client key={key}")
+            cls._CLIENT_CACHE.clear()
+            cls._CLIENT_REFCOUNT.clear()
 
     @require_python_package(
         import_name="nebulagraph_python",
-        install_command="pip install ... @Tianxing",
+        install_command="pip install nebulagraph-python>=5.1.1",
         install_link=".....",
     )
     def __init__(self, config: NebulaGraphDBConfig):
@@ -246,48 +343,65 @@ 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),
-        )
-
-        if config.auto_create:
-            self._ensure_database_exists()
-
-        self.execute_query(f"SESSION SET GRAPH `{self.db_name}`")
 
-        # Create only if not exists
-        self.create_index(dimensions=config.embedding_dimension)
+        # ---- NEW: pool acquisition strategy
+        # Get or create a shared pool from the class-level cache
+        self._client_key, self._client = self._get_or_create_shared_client(config)
+        self._owns_client = True
 
         logger.info("Connected to NebulaGraph successfully.")
 
     @timed
-    def execute_query(self, gql: str, timeout: float = 5.0, auto_set_db: bool = True):
-        with self.pool.get() as client:
-            try:
-                if auto_set_db and self.db_name:
-                    client.execute(f"SESSION SET GRAPH `{self.db_name}`")
-                return client.execute(gql, timeout=timeout)
+    def execute_query(self, gql: str, timeout: float = 60.0, auto_set_db: bool = True):
+        def _wrap_use_db(q: str) -> str:
+            if auto_set_db and self.db_name:
+                return f"USE `{self.db_name}`\n{q}"
+            return q
 
-            except Exception as e:
-                if "Session not found" in str(e) or "Connection not established" in str(e):
-                    logger.warning(f"[execute_query] {e!s}, replacing client...")
-                    self.pool.replace_client(client)
-                    return self.execute_query(gql, timeout, auto_set_db)
-                raise
+        try:
+            return self._client.execute(_wrap_use_db(gql), timeout=timeout)
+
+        except Exception as e:
+            emsg = str(e)
+            if any(k.lower() in emsg.lower() for k in _TRANSIENT_ERR_KEYS):
+                logger.warning(f"[execute_query] {e!s} → refreshing session pool and retry once...")
+                try:
+                    self._refresh_client()
+                    return self._client.execute(_wrap_use_db(gql), timeout=timeout)
+                except Exception:
+                    logger.exception("[execute_query] retry after refresh failed")
+                    raise
+            raise
 
     @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_client:
+            logger.debug("[NebulaGraphDBSync] close() skipped (injected client).")
+            return
+        if self._client_key:
+            self._release_shared_client(self._client_key)
+            self._client_key = None
+            self._client = 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 +480,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 +680,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 +691,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 +711,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 +729,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 +798,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 +808,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 +830,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 +845,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 +953,8 @@ class NebulaGraphDB(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.
@@ -851,6 +966,8 @@ class NebulaGraphDB(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.
@@ -860,6 +977,7 @@ class NebulaGraphDB(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.
         """
@@ -874,12 +992,22 @@ 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}"')
+
+                # Add search_filter conditions
+                if search_filter:
+                    for key, value in search_filter.items():
+                        if isinstance(value, str):
+                            where_clauses.append(f'n.{key} = "{value}"')
+                        else:
+                            where_clauses.append(f"n.{key} = {value}")
 
         where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
 
         gql = f"""
-               USE `{self.db_name}`
                MATCH (n@Memory)
                {where_clause}
                ORDER BY inner_product(n.{self.dim_field}, {gql_vector}) DESC
@@ -902,7 +1030,7 @@ class NebulaGraphDB(BaseGraphDB):
                 id_val = values[0].as_string()
                 score_val = values[1].as_double()
                 score_val = (score_val + 1) / 2  # align to neo4j, Normalized Cosine Score
-                if threshold is None or score_val <= threshold:
+                if threshold is None or score_val >= threshold:
                     output.append({"id": id_val, "score": score_val})
             return output
         except Exception as e:
@@ -936,20 +1064,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 +1273,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 +1336,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}")
@@ -1235,9 +1360,9 @@ class NebulaGraphDB(BaseGraphDB):
             where_clause += f' AND n.user_name = "{self.config.user_name}"'
 
         return_fields = self._build_return_fields(include_embedding)
+        return_fields += f", n.{self.dim_field} AS {self.dim_field}"
 
         query = f"""
-            USE `{self.db_name}`
             MATCH (n@Memory)
             WHERE {where_clause}
             OPTIONAL MATCH (n)-[@PARENT]->(c@Memory)
@@ -1247,14 +1372,16 @@ class NebulaGraphDB(BaseGraphDB):
         """
 
         candidates = []
+        node_ids = set()
         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()}
-                candidates.append(self._parse_node(props))
+                props = {k: v.value for k, v in row.items()}
+                node = self._parse_node(props)
+                node_id = node["id"]
+                if node_id not in node_ids:
+                    candidates.append(node)
+                    node_ids.add(node_id)
         except Exception as e:
             logger.error(f"Failed : {e}, traceback: {traceback.format_exc()}")
         return candidates
@@ -1408,18 +1535,19 @@ class NebulaGraphDB(BaseGraphDB):
                 logger.info(f"✅ Graph Type {graph_type_name} already include {self.dim_field}")
 
         create_graph = f"CREATE GRAPH IF NOT EXISTS `{self.db_name}` TYPED {graph_type_name}"
-        set_graph_working = f"SESSION SET GRAPH `{self.db_name}`"
-
         try:
             self.execute_query(create_graph, auto_set_db=False)
-            self.execute_query(set_graph_working)
             logger.info(f"✅ Graph ``{self.db_name}`` is now the working graph.")
         except Exception as e:
             logger.error(f"❌ Failed to create tag: {e} trace: {traceback.format_exc()}")
 
     @timed
     def _create_vector_index(
-        self, label: str, vector_property: str, dimensions: int, index_name: str
+        self,
+        label: str = "Memory",
+        vector_property: str = "embedding",
+        dimensions: int = 3072,
+        index_name: str = "memory_vector_index",
     ) -> None:
         """
         Create a vector index for the specified property in the label.
@@ -1555,6 +1683,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 +1692,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 +1715,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 +1751,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)