Merge branch 'dev' into dev

Author: CaralHsi

src/memos/graph_dbs/nebular.py

@@ -41,6 +41,20 @@ def _format_datetime(value: str | datetime) -> str:
     return str(value)
 
 
+def _normalize_datetime(val):
+    """
+    Normalize datetime to ISO 8601 UTC string with +00:00.
+    - If val is datetime object -> keep isoformat() (Neo4j)
+    - If val is string without timezone -> append +00:00 (Nebula)
+    - Otherwise just str()
+    """
+    if hasattr(val, "isoformat"):
+        return val.isoformat()
+    if isinstance(val, str) and not val.endswith(("+00:00", "Z", "+08:00")):
+        return val + "+08:00"
+    return str(val)
+
+
 class SessionPoolError(Exception):
     pass
 
@@ -62,6 +76,7 @@ def __init__(
         self.hosts = hosts
         self.user = user
         self.password = password
+        self.minsize = minsize
         self.maxsize = maxsize
         self.pool = Queue(maxsize)
         self.lock = Lock()
@@ -79,13 +94,13 @@ def _create_and_add_client(self):
         self.clients.append(client)
 
     def get_client(self, timeout: float = 5.0):
-        from nebulagraph_python import NebulaClient
-
         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
@@ -120,6 +135,25 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
         return _ClientContext(self)
 
+    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:
+                try:
+                    client.close()
+                except Exception:
+                    logger.error("Fail to close!!!")
+            self.clients.clear()
+            while not self.pool.empty():
+                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.")
+
 
 class NebulaGraphDB(BaseGraphDB):
     """
@@ -181,12 +215,27 @@ def __init__(self, config: NebulaGraphDBConfig):
 
     def execute_query(self, gql: str, timeout: float = 5.0, auto_set_db: bool = True):
         with self.pool.get() as client:
-            if auto_set_db and self.db_name:
-                client.execute(f"SESSION SET GRAPH `{self.db_name}`")
             try:
+                if auto_set_db and self.db_name:
+                    client.execute(f"SESSION SET GRAPH `{self.db_name}`")
                 return client.execute(gql, timeout=timeout)
-            except Exception:
+            except Exception as e:
                 logger.error(f"Fail to run gql {gql} trace: {traceback.format_exc()}")
+                if "Session not found" in str(e):
+                    logger.warning("[execute_query] Session expired, replacing client.")
+                    try:
+                        client.close()
+                    except Exception:
+                        logger.error("Fail to close!!!!!")
+                    finally:
+                        if client in self.pool.clients:
+                            self.pool.clients.remove(client)
+                    from nebulagraph_python import NebulaClient
+
+                    new_client = NebulaClient(self.pool.hosts, self.pool.user, self.pool.password)
+                    self.pool.clients.append(new_client)
+                    return new_client.execute(gql, timeout=timeout)
+                raise
 
     def close(self):
         self.pool.close()
@@ -923,9 +972,11 @@ def clear(self) -> None:
         except Exception as e:
             logger.error(f"[ERROR] Failed to clear database: {e}")
 
-    def export_graph(self) -> dict[str, Any]:
+    def export_graph(self, include_embedding: bool = False) -> dict[str, Any]:
         """
         Export all graph nodes and edges in a structured form.
+        Args:
+        include_embedding (bool): Whether to include the large embedding field.
 
         Returns:
             {
@@ -942,12 +993,41 @@ def export_graph(self) -> dict[str, Any]:
             edge_query += f' WHERE r.user_name = "{username}"'
 
         try:
-            full_node_query = f"{node_query} RETURN n"
-            node_result = self.execute_query(full_node_query)
+            if include_embedding:
+                return_fields = "n"
+            else:
+                return_fields = ",".join(
+                    [
+                        "n.id AS id",
+                        "n.memory AS memory",
+                        "n.user_name AS user_name",
+                        "n.user_id AS user_id",
+                        "n.session_id AS session_id",
+                        "n.status AS status",
+                        "n.key AS key",
+                        "n.confidence AS confidence",
+                        "n.tags AS tags",
+                        "n.created_at AS created_at",
+                        "n.updated_at AS updated_at",
+                        "n.memory_type AS memory_type",
+                        "n.sources AS sources",
+                        "n.source AS source",
+                        "n.node_type AS node_type",
+                        "n.visibility AS visibility",
+                        "n.usage AS usage",
+                        "n.background AS background",
+                    ]
+                )
+
+            full_node_query = f"{node_query} RETURN {return_fields}"
+            node_result = self.execute_query(full_node_query, timeout=20)
             nodes = []
+            logger.debug(f"Debugging: {node_result}")
             for row in node_result:
-                node_wrapper = row.values()[0].as_node()
-                props = node_wrapper.get_properties()
+                if include_embedding:
+                    props = row.values()[0].as_node().get_properties()
+                else:
+                    props = {k: v.value for k, v in row.items()}
 
                 node = self._parse_node(props)
                 nodes.append(node)
@@ -956,7 +1036,7 @@ def export_graph(self) -> dict[str, Any]:
 
         try:
             full_edge_query = f"{edge_query} RETURN a.id AS source, b.id AS target, type(r) as edge"
-            edge_result = self.execute_query(full_edge_query)
+            edge_result = self.execute_query(full_edge_query, timeout=20)
             edges = [
                 {
                     "source": row.values()[0].value,
@@ -1023,6 +1103,7 @@ def get_all_memory_items(self, scope: str) -> list[dict]:
                    MATCH (n@Memory)
                    {where_clause}
                    RETURN n
+                   LIMIT 100
                    """
         nodes = []
         try:
@@ -1065,7 +1146,7 @@ def get_structure_optimization_candidates(self, scope: str) -> list[dict]:
                 node_props = rec["n"].as_node().get_properties()
                 candidates.append(self._parse_node(node_props))
         except Exception as e:
-            logger.error(f"Failed : {e}")
+            logger.error(f"Failed : {e}, traceback: {traceback.format_exc()}")
         return candidates
 
     def drop_database(self) -> None:
@@ -1318,15 +1399,17 @@ def _parse_node(self, props: dict[str, Any]) -> dict[str, Any]:
         parsed = {k: self._parse_value(v) for k, v in props.items()}
 
         for tf in ("created_at", "updated_at"):
-            if tf in parsed and hasattr(parsed[tf], "isoformat"):
-                parsed[tf] = parsed[tf].isoformat()
+            if tf in parsed and parsed[tf] is not None:
+                parsed[tf] = _normalize_datetime(parsed[tf])
 
         node_id = parsed.pop("id")
         memory = parsed.pop("memory", "")
         parsed.pop("user_name", None)
         metadata = parsed
         metadata["type"] = metadata.pop("node_type")
-        metadata["embedding"] = metadata.pop(self.dim_field)
+
+        if self.dim_field in metadata:
+            metadata["embedding"] = metadata.pop(self.dim_field)
 
         return {"id": node_id, "memory": memory, "metadata": metadata}
 

src/memos/mem_os/utils/format_utils.py

@@ -570,15 +570,23 @@ def convert_graph_to_tree_forworkmem(
         else:
             other_roots.append(root_id)
 
-    def build_tree(node_id: str) -> dict[str, Any]:
-        """Recursively build tree structure"""
+    def build_tree(node_id: str, visited=None) -> dict[str, Any] | None:
+        """Recursively build tree structure with cycle detection"""
+        if visited is None:
+            visited = set()
+
+        if node_id in visited:
+            logger.warning(f"[build_tree] Detected cycle at node {node_id}, skipping.")
+            return None
+        visited.add(node_id)
+
         if node_id not in node_map:
             return None
 
         children_ids = children_map.get(node_id, [])
         children = []
         for child_id in children_ids:
-            child_tree = build_tree(child_id)
+            child_tree = build_tree(child_id, visited)
             if child_tree:
                 children.append(child_tree)
 

…al/tree_text_memory/organize/conflict.py …ual/tree_text_memory/organize/handler.pysrc/memos/memories/textual/tree_text_memory/organize/conflict.py renamed to src/memos/memories/textual/tree_text_memory/organize/handler.py src/memos/memories/textual/tree_text_memory/organize/conflict.py renamed to src/memos/memories/textual/tree_text_memory/organize/handler.py

@@ -1,6 +1,5 @@
 import json
 import re
-
 from datetime import datetime
 
 from dateutil import parser
@@ -11,82 +10,68 @@
 from memos.log import get_logger
 from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
 from memos.templates.tree_reorganize_prompts import (
-    CONFLICT_DETECTOR_PROMPT,
-    CONFLICT_RESOLVER_PROMPT,
+    MEMORY_RELATION_DETECTOR_PROMPT,
+    MEMORY_RELATION_RESOLVER_PROMPT,
 )
 
-
 logger = get_logger(__name__)
 
 
-class ConflictHandler:
+class NodeHandler:
     EMBEDDING_THRESHOLD: float = 0.8  # Threshold for embedding similarity to consider conflict
 
     def __init__(self, graph_store: Neo4jGraphDB, llm: BaseLLM, embedder: BaseEmbedder):
         self.graph_store = graph_store
         self.llm = llm
         self.embedder = embedder
 
-    def detect(
-        self, memory: TextualMemoryItem, top_k: int = 5, scope: str | None = None
-    ) -> list[tuple[TextualMemoryItem, TextualMemoryItem]]:
-        """
-        Detect conflicts by finding the most similar items in the graph database based on embedding, then use LLM to judge conflict.
-        Args:
-            memory: The memory item (should have an embedding attribute or field).
-            top_k: Number of top similar nodes to retrieve.
-            scope: Optional memory type filter.
-        Returns:
-            List of conflict pairs (each pair is a tuple: (memory, candidate)).
-        """
+    def detect(self, memory, top_k: int = 5, scope=None):
         # 1. Search for similar memories based on embedding
         embedding = memory.metadata.embedding
         embedding_candidates_info = self.graph_store.search_by_embedding(
-            embedding, top_k=top_k, scope=scope
+            embedding, top_k=top_k, scope=scope, threshold=self.EMBEDDING_THRESHOLD
         )
         # 2. Filter based on similarity threshold
         embedding_candidates_ids = [
-            info["id"]
-            for info in embedding_candidates_info
-            if info["score"] >= self.EMBEDDING_THRESHOLD and info["id"] != memory.id
+            info["id"] for info in embedding_candidates_info if info["id"] != memory.id
         ]
         # 3. Judge conflicts using LLM
         embedding_candidates = self.graph_store.get_nodes(embedding_candidates_ids)
-        conflict_pairs = []
+        detected_relationships = []
         for embedding_candidate in embedding_candidates:
             embedding_candidate = TextualMemoryItem.from_dict(embedding_candidate)
             prompt = [
-                {
-                    "role": "system",
-                    "content": "You are a conflict detector for memory items.",
-                },
                 {
                     "role": "user",
-                    "content": CONFLICT_DETECTOR_PROMPT.format(
-                        statement_1=memory.memory,
-                        statement_2=embedding_candidate.memory,
+                    "content": MEMORY_RELATION_DETECTOR_PROMPT.format(
+                        statement_1=memory.memory, statement_2=embedding_candidate.memory
                     ),
-                },
+                }
             ]
             result = self.llm.generate(prompt).strip()
-            if "yes" in result.lower():
-                conflict_pairs.append([memory, embedding_candidate])
-        if len(conflict_pairs):
-            conflict_text = "\n".join(
-                f'"{pair[0].memory!s}" <==CONFLICT==> "{pair[1].memory!s}"'
-                for pair in conflict_pairs
-            )
-            logger.warning(
-                f"Detected {len(conflict_pairs)} conflicts for memory {memory.id}\n {conflict_text}"
-            )
-        return conflict_pairs
+            if result == "contradictory":
+                logger.warning(
+                    f'detected "{memory.memory}" <==CONFLICT==> "{embedding_candidate.memory}"'
+                )
+                detected_relationships.append([memory, embedding_candidate, "contradictory"])
+            elif result == "redundant":
+                logger.warning(
+                    f'detected "{memory.memory}" <==REDUNDANT==> "{embedding_candidate.memory}"'
+                )
+                detected_relationships.append([memory, embedding_candidate, "redundant"])
+            elif result == "independent":
+                pass
+            else:
+                pass
+        return detected_relationships
 
-    def resolve(self, memory_a: TextualMemoryItem, memory_b: TextualMemoryItem) -> None:
+    def resolve(self, memory_a: TextualMemoryItem, memory_b: TextualMemoryItem, relation) -> None:
         """
         Resolve detected conflicts between two memory items using LLM fusion.
         Args:
             memory_a: The first conflicting memory item.
             memory_b: The second conflicting memory item.
+            relation: relation
         Returns:
             A fused TextualMemoryItem representing the resolved memory.
         """
@@ -96,13 +81,10 @@ def resolve(self, memory_a: TextualMemoryItem, memory_b: TextualMemoryItem) -> N
         metadata_1 = memory_a.metadata.model_dump_json(include=metadata_for_resolve)
         metadata_2 = memory_b.metadata.model_dump_json(include=metadata_for_resolve)
         prompt = [
-            {
-                "role": "system",
-                "content": "",
-            },
             {
                 "role": "user",
-                "content": CONFLICT_RESOLVER_PROMPT.format(
+                "content": MEMORY_RELATION_RESOLVER_PROMPT.format(
+                    relation=relation,
                     statement_1=memory_a.memory,
                     metadata_1=metadata_1,
                     statement_2=memory_b.memory,
@@ -119,7 +101,7 @@ def resolve(self, memory_a: TextualMemoryItem, memory_b: TextualMemoryItem) -> N
             # —————— 2.1 Can't resolve conflict, hard update by comparing timestamp ————
             if len(answer) <= 10 and "no" in answer.lower():
                 logger.warning(
-                    f"Conflict between {memory_a.id} and {memory_b.id} could not be resolved. "
+                    f"{relation} between {memory_a.id} and {memory_b.id} could not be resolved. "
                 )
                 self._hard_update(memory_a, memory_b)
             # —————— 2.2 Conflict resolved, update metadata and memory ————