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

memos/memories/textual/item.py

@@ -1,13 +1,48 @@
 """Defines memory item types for textual memory."""
 
+import json
 import uuid
 
 from datetime import datetime
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 
+ALLOWED_ROLES = {"user", "assistant", "system"}
+
+
+class SourceMessage(BaseModel):
+    """
+    Purpose: **memory provenance / traceability**.
+
+    Capture the minimal, reproducible origin context of a memory item so it can be
+    audited, traced, rolled back, or de-duplicated later.
+
+    Fields & conventions:
+        - type: Source kind (e.g., "chat", "doc", "web", "file", "system", ...).
+            If not provided, upstream logic may infer it:
+            presence of `role` ⇒ "chat"; otherwise ⇒ "doc".
+        - role: Conversation role ("user" | "assistant" | "system") when the
+            source is a chat turn.
+        - content: Minimal reproducible snippet from the source. If omitted,
+            upstream may fall back to `doc_path` / `url` / `message_id`.
+        - chat_time / message_id / doc_path: Locators for precisely pointing back
+            to the original record (timestamp, message id, document path).
+        - Extra fields: Allowed (`model_config.extra="allow"`) to carry arbitrary
+            provenance attributes (e.g., url, page, offset, span, local_confidence).
+    """
+
+    type: str | None = "chat"
+    role: Literal["user", "assistant", "system"] | None = None
+    chat_time: str | None = None
+    message_id: str | None = None
+    content: str | None = None
+    doc_path: str | None = None
+
+    model_config = ConfigDict(extra="allow")
+
+
 class TextualMemoryMetadata(BaseModel):
     """Metadata for a memory item.
 
@@ -62,7 +97,7 @@ class TreeNodeTextualMemoryMetadata(TextualMemoryMetadata):
     memory_type: Literal["WorkingMemory", "LongTermMemory", "UserMemory", "OuterMemory"] = Field(
         default="WorkingMemory", description="Memory lifecycle type."
     )
-    sources: list[str] | None = Field(
+    sources: list[SourceMessage] | None = Field(
         default=None, description="Multiple origins of the memory (e.g., URLs, notes)."
     )
     embedding: list[float] | None = Field(
@@ -74,8 +109,8 @@ class TreeNodeTextualMemoryMetadata(TextualMemoryMetadata):
         description="The timestamp of the first creation to the memory. Useful "
         "for tracking memory initialization. Format: ISO 8601.",
     )
-    usage: list[str] | None = Field(
-        default=[],
+    usage: list[str] = Field(
+        default_factory=list,
         description="Usage history of this node",
     )
     background: str | None = Field(
@@ -83,12 +118,40 @@ class TreeNodeTextualMemoryMetadata(TextualMemoryMetadata):
         description="background of this node",
     )
 
-    @field_validator("sources")
+    @field_validator("sources", mode="before")
     @classmethod
-    def validate_sources(cls, v):
-        if v is not None and not isinstance(v, list):
-            raise ValueError("Sources must be a list of strings.")
-        return v
+    def coerce_sources(cls, v):
+        if v is None:
+            return v
+        if not isinstance(v, list):
+            raise TypeError("sources must be a list")
+        out = []
+        for item in v:
+            if isinstance(item, SourceMessage):
+                out.append(item)
+
+            elif isinstance(item, dict):
+                d = dict(item)
+                if d.get("type") is None:
+                    d["type"] = "chat" if d.get("role") in ALLOWED_ROLES else "doc"
+                out.append(SourceMessage(**d))
+
+            elif isinstance(item, str):
+                try:
+                    parsed = json.loads(item)
+                except Exception:
+                    parsed = None
+
+                if isinstance(parsed, dict):
+                    if parsed.get("type") is None:
+                        parsed["type"] = "chat" if parsed.get("role") in ALLOWED_ROLES else "doc"
+                    out.append(SourceMessage(**parsed))
+                else:
+                    out.append(SourceMessage(type="doc", content=item))
+
+            else:
+                out.append(SourceMessage(type="doc", content=str(item)))
+        return out
 
     def __str__(self) -> str:
         """Pretty string representation of the metadata."""
@@ -114,19 +177,17 @@ class TextualMemoryItem(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     memory: str
     metadata: (
-        TextualMemoryMetadata
+        SearchedTreeNodeTextualMemoryMetadata
         | TreeNodeTextualMemoryMetadata
-        | SearchedTreeNodeTextualMemoryMetadata
+        | TextualMemoryMetadata
     ) = Field(default_factory=TextualMemoryMetadata)
 
     model_config = ConfigDict(extra="forbid")
 
+    @field_validator("id")
     @classmethod
-    def validate_id(cls, v):
-        try:
-            uuid.UUID(v)
-        except ValueError as e:
-            raise ValueError("Invalid UUID format") from e
+    def _validate_id(cls, v: str) -> str:
+        uuid.UUID(v)
         return v
 
     @classmethod
@@ -136,6 +197,24 @@ class TextualMemoryItem(BaseModel):
     def to_dict(self) -> dict:
         return self.model_dump(exclude_none=True)
 
+    @field_validator("metadata", mode="before")
+    @classmethod
+    def _coerce_metadata(cls, v: Any):
+        if isinstance(
+            v,
+            SearchedTreeNodeTextualMemoryMetadata
+            | TreeNodeTextualMemoryMetadata
+            | TextualMemoryMetadata,
+        ):
+            return v
+        if isinstance(v, dict):
+            if v.get("relativity") is not None:
+                return SearchedTreeNodeTextualMemoryMetadata(**v)
+            if any(k in v for k in ("sources", "memory_type", "embedding", "background", "usage")):
+                return TreeNodeTextualMemoryMetadata(**v)
+            return TextualMemoryMetadata(**v)
+        return v
+
     def __str__(self) -> str:
         """Pretty string representation of the memory item."""
         return f"<ID: {self.id} | Memory: {self.memory} | Metadata: {self.metadata!s}>"

memos/memories/textual/naive.py

@@ -115,7 +115,7 @@ class NaiveTextMemory(BaseTextMemory):
                 self.memories[i] = memory_dict
                 break
 
-    def search(self, query: str, top_k: int) -> list[TextualMemoryItem]:
+    def search(self, query: str, top_k: int, **kwargs) -> list[TextualMemoryItem]:
         """Search for memories based on a query."""
         sims = [
             (memory, len(set(query.split()) & set(memory["memory"].split())))

memos/memories/textual/tree.py

@@ -2,6 +2,7 @@ import json
 import os
 import shutil
 import tempfile
+import time
 
 from datetime import datetime
 from pathlib import Path
@@ -32,15 +33,28 @@ class TreeTextMemory(BaseTextMemory):
 
     def __init__(self, config: TreeTextMemoryConfig):
         """Initialize memory with the given configuration."""
+        time_start = time.time()
         self.config: TreeTextMemoryConfig = config
         self.extractor_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
             config.extractor_llm
         )
+        logger.info(f"time init: extractor_llm time is: {time.time() - time_start}")
+
+        time_start_ex = time.time()
         self.dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
             config.dispatcher_llm
         )
+        logger.info(f"time init: dispatcher_llm time is: {time.time() - time_start_ex}")
+
+        time_start_em = time.time()
         self.embedder: OllamaEmbedder = EmbedderFactory.from_config(config.embedder)
+        logger.info(f"time init: embedder time is: {time.time() - time_start_em}")
+
+        time_start_gs = time.time()
         self.graph_store: Neo4jGraphDB = GraphStoreFactory.from_config(config.graph_db)
+        logger.info(f"time init: graph_store time is: {time.time() - time_start_gs}")
+
+        time_start_rr = time.time()
         if config.reranker is None:
             default_cfg = RerankerConfigFactory.model_validate(
                 {
@@ -54,9 +68,10 @@ class TreeTextMemory(BaseTextMemory):
             self.reranker = RerankerFactory.from_config(default_cfg)
         else:
             self.reranker = RerankerFactory.from_config(config.reranker)
-
+        logger.info(f"time init: reranker time is: {time.time() - time_start_rr}")
         self.is_reorganize = config.reorganize
 
+        time_start_mm = time.time()
         self.memory_manager: MemoryManager = MemoryManager(
             self.graph_store,
             self.embedder,
@@ -69,7 +84,8 @@ class TreeTextMemory(BaseTextMemory):
             },
             is_reorganize=self.is_reorganize,
         )
-
+        logger.info(f"time init: memory_manager time is: {time.time() - time_start_mm}")
+        time_start_ir = time.time()
         # Create internet retriever if configured
         self.internet_retriever = None
         if config.internet_retriever is not None:
@@ -81,6 +97,7 @@ class TreeTextMemory(BaseTextMemory):
             )
         else:
             logger.info("No internet retriever configured")
+        logger.info(f"time init: internet_retriever time is: {time.time() - time_start_ir}")
 
     def add(self, memories: list[TextualMemoryItem | dict[str, Any]]) -> list[str]:
         """Add memories.
@@ -122,6 +139,7 @@ class TreeTextMemory(BaseTextMemory):
         memory_type: str = "All",
         manual_close_internet: bool = False,
         moscube: bool = False,
+        search_filter: dict | None = None,
     ) -> list[TextualMemoryItem]:
         """Search for memories based on a query.
         User query -> TaskGoalParser -> MemoryPathResolver ->
@@ -136,6 +154,12 @@ class TreeTextMemory(BaseTextMemory):
             memory_type (str): Type restriction for search.
             ['All', 'WorkingMemory', 'LongTermMemory', 'UserMemory']
             manual_close_internet (bool): If True, the internet retriever will be closed by this search, it high priority than config.
+            moscube (bool): whether you use moscube to answer questions
+            search_filter (dict, optional): Optional metadata filters for search results.
+                - Keys correspond to memory metadata fields (e.g., "user_id", "session_id").
+                - Values are exact-match conditions.
+                Example: {"user_id": "123", "session_id": "abc"}
+                If None, no additional filtering is applied.
         Returns:
             list[TextualMemoryItem]: List of matching memories.
         """
@@ -160,7 +184,7 @@ class TreeTextMemory(BaseTextMemory):
                 internet_retriever=self.internet_retriever,
                 moscube=moscube,
             )
-        return searcher.search(query, top_k, info, mode, memory_type)
+        return searcher.search(query, top_k, info, mode, memory_type, search_filter)
 
     def get_relevant_subgraph(
         self, query: str, top_k: int = 5, depth: int = 2, center_status: str = "activated"

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

@@ -1,5 +1,6 @@
 import json
 import re
+
 from datetime import datetime
 
 from dateutil import parser
@@ -14,6 +15,7 @@ from memos.templates.tree_reorganize_prompts import (
     MEMORY_RELATION_RESOLVER_PROMPT,
 )
 
+
 logger = get_logger(__name__)
 
 
@@ -50,12 +52,12 @@ class NodeHandler:
             ]
             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"])

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

@@ -1,8 +1,10 @@
+import traceback
 import uuid
 
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import as_completed
 from datetime import datetime
 
+from memos.context.context import ContextThreadPoolExecutor
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.neo4j import Neo4jGraphDB
 from memos.llms.factory import AzureLLM, OllamaLLM, OpenAILLM
@@ -55,24 +57,35 @@ class MemoryManager:
         """
         added_ids: list[str] = []
 
-        with ThreadPoolExecutor(max_workers=8) as executor:
+        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)
                 except Exception as e:
                     logger.exception("Memory processing error: ", exc_info=e)
 
-        self.graph_store.remove_oldest_memory(
-            memory_type="WorkingMemory", keep_latest=self.memory_size["WorkingMemory"]
-        )
-        self.graph_store.remove_oldest_memory(
-            memory_type="LongTermMemory", keep_latest=self.memory_size["LongTermMemory"]
-        )
-        self.graph_store.remove_oldest_memory(
-            memory_type="UserMemory", keep_latest=self.memory_size["UserMemory"]
-        )
+        try:
+            self.graph_store.remove_oldest_memory(
+                memory_type="WorkingMemory", keep_latest=self.memory_size["WorkingMemory"]
+            )
+        except Exception:
+            logger.warning(f"Remove WorkingMemory error: {traceback.format_exc()}")
+
+        try:
+            self.graph_store.remove_oldest_memory(
+                memory_type="LongTermMemory", keep_latest=self.memory_size["LongTermMemory"]
+            )
+        except Exception:
+            logger.warning(f"Remove LongTermMemory error: {traceback.format_exc()}")
+
+        try:
+            self.graph_store.remove_oldest_memory(
+                memory_type="UserMemory", keep_latest=self.memory_size["UserMemory"]
+            )
+        except Exception:
+            logger.warning(f"Remove UserMemory error: {traceback.format_exc()}")
 
         self._refresh_memory_size()
         return added_ids
@@ -82,12 +95,12 @@ class MemoryManager:
         Replace WorkingMemory
         """
         working_memory_top_k = memories[: self.memory_size["WorkingMemory"]]
-        with ThreadPoolExecutor(max_workers=8) as executor:
+        with ContextThreadPoolExecutor(max_workers=8) as executor:
             futures = [
                 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:
@@ -102,6 +115,7 @@ class MemoryManager:
         """
         Return the cached memory type counts.
         """
+        self._refresh_memory_size()
         return self.current_memory_size
 
     def _refresh_memory_size(self) -> None:

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

@@ -46,7 +46,7 @@ class RelationAndReasoningDetector:
                     "sequence_links": [],
                     "aggregate_nodes": [],
                 }
-
+            """
             nearest = self.graph_store.get_neighbors_by_tag(
                 tags=node.metadata.tags,
                 exclude_ids=exclude_ids,
@@ -55,7 +55,6 @@ class RelationAndReasoningDetector:
             )
             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"])

memos/memories/textual/tree_text_memory/organize/reorganizer.py

@@ -4,19 +4,20 @@ import time
 import traceback
 
 from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import as_completed
 from queue import PriorityQueue
 from typing import Literal
 
 import numpy as np
 
+from memos.context.context import ContextThreadPoolExecutor
 from memos.dependency import require_python_package
 from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.item import GraphDBEdge, GraphDBNode
 from memos.graph_dbs.neo4j import Neo4jGraphDB
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
-from memos.memories.textual.item import TreeNodeTextualMemoryMetadata
+from memos.memories.textual.item import SourceMessage, TreeNodeTextualMemoryMetadata
 from memos.memories.textual.tree_text_memory.organize.handler import NodeHandler
 from memos.memories.textual.tree_text_memory.organize.relation_reason_detector import (
     RelationAndReasoningDetector,
@@ -27,6 +28,22 @@ from memos.templates.tree_reorganize_prompts import LOCAL_SUBCLUSTER_PROMPT, REO
 logger = get_logger(__name__)
 
 
+def build_summary_parent_node(cluster_nodes):
+    normalized_sources = []
+    for n in cluster_nodes:
+        sm = SourceMessage(
+            type="chat",
+            role=None,
+            chat_time=None,
+            message_id=None,
+            content=n.memory,
+            # extra
+            node_id=n.id,
+        )
+        normalized_sources.append(sm)
+    return normalized_sources
+
+
 class QueueMessage:
     def __init__(
         self,
@@ -51,6 +68,15 @@ class QueueMessage:
         return op_priority[self.op] < op_priority[other.op]
 
 
+def extract_first_to_last_brace(text: str):
+    start = text.find("{")
+    end = text.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        return "", None
+    json_str = text[start : end + 1]
+    return json_str, json.loads(json_str)
+
+
 class GraphStructureReorganizer:
     def __init__(
         self, graph_store: Neo4jGraphDB, llm: BaseLLM, embedder: OllamaEmbedder, is_reorganize: bool
@@ -87,6 +113,7 @@ class GraphStructureReorganizer:
         1) queue is empty
         2) any running structure optimization is done
         """
+        deadline = time.time() + 600
         if not self.is_reorganize:
             return
 
@@ -96,6 +123,9 @@ class GraphStructureReorganizer:
 
         while any(self._is_optimizing.values()):
             logger.debug(f"Waiting for structure optimizer to finish... {self._is_optimizing}")
+            if time.time() > deadline:
+                logger.error(f"Wait timed out; flags={self._is_optimizing}")
+                break
             time.sleep(1)
         logger.debug("Structure optimizer is now idle.")
 
@@ -129,6 +159,9 @@ class GraphStructureReorganizer:
 
         logger.info("Structure optimizer schedule started.")
         while not getattr(self, "_stop_scheduler", False):
+            if any(self._is_optimizing.values()):
+                time.sleep(1)
+                continue
             if self._reorganize_needed:
                 logger.info("[Reorganizer] Triggering optimize_structure due to new nodes.")
                 self.optimize_structure(scope="LongTermMemory")
@@ -176,6 +209,7 @@ class GraphStructureReorganizer:
         local_tree_threshold: int = 10,
         min_cluster_size: int = 4,
         min_group_size: int = 20,
+        max_duration_sec: int = 600,
     ):
         """
         Periodically reorganize the graph:
@@ -183,8 +217,20 @@ class GraphStructureReorganizer:
         2. Summarize each cluster.
         3. Create parent nodes and build local PARENT trees.
         """
+        # --- Total time watch dog: check functions ---
+        start_ts = time.time()
+
+        def _check_deadline(where: str):
+            if time.time() - start_ts > max_duration_sec:
+                logger.error(
+                    f"[GraphStructureReorganize] {scope} surpass {max_duration_sec}s，time "
+                    f"over at {where}"
+                )
+                return True
+            return False
+
         if self._is_optimizing[scope]:
-            logger.info(f"Already optimizing for {scope}. Skipping.")
+            logger.info(f"[GraphStructureReorganize] Already optimizing for {scope}. Skipping.")
             return
 
         if self.graph_store.node_not_exist(scope):
@@ -198,32 +244,35 @@ class GraphStructureReorganizer:
             )
 
             logger.debug(
-                f"Num of scope in self.graph_store is {self.graph_store.get_memory_count(scope)}"
+                f"[GraphStructureReorganize] Num of scope in self.graph_store is"
+                f" {self.graph_store.get_memory_count(scope)}"
             )
             # Load candidate nodes
+            if _check_deadline("[GraphStructureReorganize] Before loading candidates"):
+                return
             raw_nodes = self.graph_store.get_structure_optimization_candidates(scope)
             nodes = [GraphDBNode(**n) for n in raw_nodes]
 
             if not nodes:
                 logger.info("[GraphStructureReorganize] No nodes to optimize. Skipping.")
                 return
-
             if len(nodes) < min_group_size:
                 logger.info(
                     f"[GraphStructureReorganize] Only {len(nodes)} candidate nodes found. Not enough to reorganize. Skipping."
                 )
                 return
 
-            logger.info(f"[GraphStructureReorganize] Loaded {len(nodes)} nodes.")
-
             # Step 2: Partition nodes
+            if _check_deadline("[GraphStructureReorganize] Before partition"):
+                return
             partitioned_groups = self._partition(nodes)
-
             logger.info(
                 f"[GraphStructureReorganize] Partitioned into {len(partitioned_groups)} clusters."
             )
 
-            with ThreadPoolExecutor(max_workers=4) as executor:
+            if _check_deadline("[GraphStructureReorganize] Before submit partition task"):
+                return
+            with ContextThreadPoolExecutor(max_workers=4) as executor:
                 futures = []
                 for cluster_nodes in partitioned_groups:
                     futures.append(
@@ -237,14 +286,17 @@ class GraphStructureReorganizer:
                     )
 
                 for f in as_completed(futures):
+                    if _check_deadline("[GraphStructureReorganize] Waiting clusters..."):
+                        for x in futures:
+                            x.cancel()
+                        return
                     try:
                         f.result()
                     except Exception as e:
                         logger.warning(
-                            f"[Reorganize] Cluster processing "
-                            f"failed: {e}, cluster_nodes: {cluster_nodes}, trace: {traceback.format_exc()}"
+                            f"[GraphStructureReorganize] Cluster processing failed: {e}, trace: {traceback.format_exc()}"
                         )
-                logger.info("[GraphStructure Reorganize] Structure optimization finished.")
+            logger.info("[GraphStructure Reorganize] Structure optimization finished.")
 
         finally:
             self._is_optimizing[scope] = False
@@ -282,7 +334,7 @@ class GraphStructureReorganizer:
         nodes_to_check = cluster_nodes
         exclude_ids = [n.id for n in nodes_to_check]
 
-        with ThreadPoolExecutor(max_workers=4) as executor:
+        with ContextThreadPoolExecutor(max_workers=4) as executor:
             futures = []
             for node in nodes_to_check:
                 futures.append(
@@ -294,7 +346,7 @@ class GraphStructureReorganizer:
                     )
                 )
 
-            for f in as_completed(futures):
+            for f in as_completed(futures, timeout=300):
                 results = f.result()
 
                 # 1) Add pairwise relations
@@ -331,11 +383,11 @@ class GraphStructureReorganizer:
                     for child_id in agg_node.metadata.sources:
                         self.graph_store.add_edge(agg_node.id, child_id, "AGGREGATE_TO")
 
-            logger.info("[Reorganizer] Cluster relation/reasoning done.")
+        logger.info("[Reorganizer] Cluster relation/reasoning done.")
 
     def _local_subcluster(
-        self, cluster_nodes: list[GraphDBNode], max_length: int = 8000
-    ) -> (list)[list[GraphDBNode]]:
+        self, cluster_nodes: list[GraphDBNode], max_length: int = 15000
+    ) -> list[list[GraphDBNode]]:
         """
         Use LLM to split a large cluster into semantically coherent sub-clusters.
         """
@@ -350,7 +402,7 @@ class GraphStructureReorganizer:
 
         joined_scene = "\n".join(scene_lines)
         if len(joined_scene) > max_length:
-            logger.warning(f"Sub-cluster too long: {joined_scene}")
+            logger.warning("Sub-cluster too long")
         prompt = LOCAL_SUBCLUSTER_PROMPT.replace("{joined_scene}", joined_scene[:max_length])
 
         messages = [{"role": "user", "content": prompt}]
@@ -499,17 +551,17 @@ class GraphStructureReorganizer:
         parent_node = GraphDBNode(
             memory=parent_value,
             metadata=TreeNodeTextualMemoryMetadata(
-                user_id="",  # TODO: summarized node: no user_id
-                session_id="",  # TODO: summarized node: no session_id
+                user_id=None,
+                session_id=None,
                 memory_type=scope,
                 status="activated",
                 key=parent_key,
                 tags=parent_tags,
                 embedding=embedding,
                 usage=[],
-                sources=[n.id for n in cluster_nodes],
+                sources=build_summary_parent_node(cluster_nodes),
                 background=parent_background,
-                confidence=0.99,
+                confidence=0.66,
                 type="topic",
             ),
         )
@@ -518,7 +570,7 @@ class GraphStructureReorganizer:
     def _parse_json_result(self, response_text):
         try:
             response_text = response_text.replace("```", "").replace("json", "")
-            response_json = json.loads(response_text)
+            response_json = extract_first_to_last_brace(response_text)[1]
             return response_json
         except json.JSONDecodeError as e:
             logger.warning(