feat: nebula&reorganize update (#322)

* feat: update nebula to nebula 5.1.1

* fix: bug in nebula and manager

* feat: update product

* test: update

* fix: bug in finding reorganize node

* fix: duplicate bug

* chore: delelte useless annotation

* feat: update nebula init

* feat: update reorganize

Author: CaralHsi

src/memos/graph_dbs/nebular.py

@@ -24,6 +24,19 @@
 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)
@@ -99,6 +112,7 @@ class NebulaGraphDB(BaseGraphDB):
     _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]:
@@ -115,13 +129,53 @@ def _make_client_key(cfg: NebulaGraphDBConfig) -> str:
                 "nebula-sync",
                 ",".join(hosts),
                 str(getattr(cfg, "user", "")),
-                str(getattr(cfg, "password", "")),
                 str(getattr(cfg, "use_multi_db", False)),
+                str(getattr(cfg, "space", "")),
             ]
         )
 
     @classmethod
-    def _get_or_create_shared_client(cls, cfg: NebulaGraphDBConfig) -> (tuple)[str, "NebulaClient"]:
+    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,
@@ -159,7 +213,60 @@ def _get_or_create_shared_client(cls, cfg: NebulaGraphDBConfig) -> (tuple)[str,
                 logger.info(f"[NebulaGraphDBSync] Created shared NebulaClient key={key}")
 
             cls._CLIENT_REFCOUNT[key] = cls._CLIENT_REFCOUNT.get(key, 0) + 1
-            return key, client
+
+            if getattr(cfg, "auto_create", False) and key not in cls._CLIENT_INIT_DONE:
+                try:
+                    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(
+                            label="Memory",
+                            vector_property=admin.dim_field,
+                            dimensions=int(
+                                admin.embedding_dimension or admin.default_memory_dimension
+                            ),
+                            index_name="memory_vector_index",
+                        )
+                        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):
@@ -253,32 +360,27 @@ def __init__(self, config: NebulaGraphDBConfig):
         self._client_key, self._client = self._get_or_create_shared_client(config)
         self._owns_client = True
 
-        # auto-create graph type / graph / index if needed
-        if getattr(config, "auto_create", False):
-            self._ensure_database_exists()
-
-        # Create only if not exists
-        self.create_index(dimensions=config.embedding_dimension)
         logger.info("Connected to NebulaGraph successfully.")
 
     @timed
     def execute_query(self, gql: str, timeout: float = 60.0, auto_set_db: bool = True):
-        try:
+        def _wrap_use_db(q: str) -> str:
             if auto_set_db and self.db_name:
-                gql = f"""USE `{self.db_name}`
-                       {gql}"""
-            return self._client.execute(gql, timeout=timeout)
+                return f"USE `{self.db_name}`\n{q}"
+            return q
+
+        try:
+            return self._client.execute(_wrap_use_db(gql), timeout=timeout)
+
         except Exception as e:
             emsg = str(e)
-            if "Session not found" in emsg or "Connection not established" in emsg:
-                logger.warning(f"[execute_query] {e!s}, retry once...")
+            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:
-                    if auto_set_db and self.db_name:
-                        gql = f"""USE `{self.db_name}`
-                                  {gql}"""
-                    return self._client.execute(gql, timeout=timeout)
+                    self._refresh_client()
+                    return self._client.execute(_wrap_use_db(gql), timeout=timeout)
                 except Exception:
-                    logger.exception("[execute_query] retry failed")
+                    logger.exception("[execute_query] retry after refresh failed")
                     raise
             raise
 
@@ -931,7 +1033,7 @@ def search_by_embedding(
                 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:
@@ -1261,6 +1363,7 @@ def get_structure_optimization_candidates(
             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"""
             MATCH (n@Memory)
@@ -1272,11 +1375,16 @@ def get_structure_optimization_candidates(
         """
 
         candidates = []
+        node_ids = set()
         try:
             results = self.execute_query(query)
             for row in results:
                 props = {k: v.value for k, v in row.items()}
-                candidates.append(self._parse_node(props))
+                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

src/memos/memories/textual/tree_text_memory/organize/handler.py

@@ -52,12 +52,12 @@ def detect(self, memory, top_k: int = 5, scope=None):
             ]
             result = self.llm.generate(prompt).strip()
             if result == "contradictory":
-                logger.warning(
+                logger.info(
                     f'detected "{memory.memory}" <==CONFLICT==> "{embedding_candidate.memory}"'
                 )
                 detected_relationships.append([memory, embedding_candidate, "contradictory"])
             elif result == "redundant":
-                logger.warning(
+                logger.info(
                     f'detected "{memory.memory}" <==REDUNDANT==> "{embedding_candidate.memory}"'
                 )
                 detected_relationships.append([memory, embedding_candidate, "redundant"])

src/memos/memories/textual/tree_text_memory/organize/manager.py

@@ -58,7 +58,7 @@ def add(self, memories: list[TextualMemoryItem]) -> list[str]:
 
         with ContextThreadPoolExecutor(max_workers=8) as executor:
             futures = {executor.submit(self._process_memory, m): m for m in memories}
-            for future in as_completed(futures):
+            for future in as_completed(futures, timeout=60):
                 try:
                     ids = future.result()
                     added_ids.extend(ids)
@@ -88,7 +88,7 @@ def replace_working_memory(self, memories: list[TextualMemoryItem]) -> None:
                 executor.submit(self._add_memory_to_db, memory, "WorkingMemory")
                 for memory in working_memory_top_k
             ]
-            for future in as_completed(futures):
+            for future in as_completed(futures, timeout=60):
                 try:
                     future.result()
                 except Exception as e:

src/memos/memories/textual/tree_text_memory/organize/relation_reason_detector.py

@@ -46,7 +46,7 @@ def process_node(self, node: GraphDBNode, exclude_ids: list[str], top_k: int = 5
                     "sequence_links": [],
                     "aggregate_nodes": [],
                 }
-
+            """
             nearest = self.graph_store.get_neighbors_by_tag(
                 tags=node.metadata.tags,
                 exclude_ids=exclude_ids,
@@ -55,7 +55,6 @@ def process_node(self, node: GraphDBNode, exclude_ids: list[str], top_k: int = 5
             )
             nearest = [GraphDBNode(**cand_data) for cand_data in nearest]
 
-            """
             # 1) Pairwise relations (including CAUSE/CONDITION/CONFLICT)
             pairwise = self._detect_pairwise_causal_condition_relations(node, nearest)
             results["relations"].extend(pairwise["relations"])