From 50b419b011547b4ce7f578581bf4b9471074609d Mon Sep 17 00:00:00 2001 From: fridayL Date: Wed, 16 Jul 2025 06:57:09 +0000 Subject: [PATCH 01/14] feat: update config --- src/memos/api/config.py | 51 +++++++++------- src/memos/mem_cube/utils.py | 113 +++++++++++++++++------------------- src/memos/mem_os/product.py | 3 + 3 files changed, 85 insertions(+), 82 deletions(-) diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 90acb054a..ce650d0f0 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -89,12 +89,39 @@ def get_activation_vllm_config() -> dict[str, Any]: }, } + @staticmethod + def get_embedder_config() -> dict[str, Any]: + """Get embedder configuration.""" + embedder_backend = os.getenv("MOS_EMBEDDER_BACKEND", "ollama") + + if embedder_backend == "universal_api": + return { + "backend": "universal_api", + "config": { + "provider": os.getenv("MOS_EMBEDDER_PROVIDER", "openai"), + "api_key": os.getenv("OPENAI_API_KEY", "sk-xxxx"), + "model_name_or_path": os.getenv("MOS_EMBEDDER_MODEL", "text-embedding-3-large"), + "base_url": os.getenv("OPENAI_API_BASE", "http://openai.com"), + }, + } + else: # ollama + return { + "backend": "ollama", + "config": { + "model_name_or_path": os.getenv( + "MOS_EMBEDDER_MODEL", "nomic-embed-text:latest" + ), + "api_base": os.getenv("OLLAMA_API_BASE", "http://localhost:11434"), + }, + } + @staticmethod def get_neo4j_config() -> dict[str, Any]: """Get Neo4j configuration.""" return { "uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"), "user": os.getenv("NEO4J_USER", "neo4j"), + "db_name": "default", "password": os.getenv("NEO4J_PASSWORD", "12345678"), "auto_create": True, } @@ -157,13 +184,7 @@ def get_product_default_config() -> dict[str, Any]: "backend": "openai", "config": openai_config, }, - "embedder": { - "backend": "ollama", - "config": { - "model_name_or_path": "nomic-embed-text:latest", - "api_base": os.getenv("OLLAMA_API_BASE", "http://localhost:11434"), - }, - }, + "embedder": APIConfig.get_embedder_config(), "chunker": { "backend": "sentence", "config": { @@ -252,13 +273,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General "backend": "openai", "config": openai_config, }, - "embedder": { - "backend": "ollama", - "config": { - "model_name_or_path": "nomic-embed-text:latest", - "api_base": os.getenv("OLLAMA_API_BASE", "http://localhost:11434"), - }, - }, + "embedder": APIConfig.get_embedder_config(), "chunker": { "backend": "sentence", "config": { @@ -353,13 +368,7 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None: "backend": "neo4j", "config": neo4j_config, }, - "embedder": { - "backend": "ollama", - "config": { - "model_name_or_path": "nomic-embed-text:latest", - "api_base": os.getenv("OLLAMA_API_BASE", "http://localhost:11434"), - }, - }, + "embedder": APIConfig.get_embedder_config(), "reorganize": os.getenv("MOS_ENABLE_REORGANIZE", "false").lower() == "true", }, }, diff --git a/src/memos/mem_cube/utils.py b/src/memos/mem_cube/utils.py index d0c3c5bd7..4a3b830c5 100644 --- a/src/memos/mem_cube/utils.py +++ b/src/memos/mem_cube/utils.py @@ -3,8 +3,6 @@ import subprocess import tempfile -from typing import Any - from memos.configs.mem_cube import GeneralMemCubeConfig @@ -50,77 +48,70 @@ def merge_config_with_default( GeneralMemCubeConfig: Merged configuration """ - def deep_merge_dicts( - existing: dict[str, Any], default: dict[str, Any], preserve_keys: set[str] | None = None - ) -> dict[str, Any]: - """Recursively merge dictionaries, preserving specified keys from existing dict.""" - if preserve_keys is None: - preserve_keys = set() - - result = copy.deepcopy(existing) - - for key, default_value in default.items(): - if key in preserve_keys: - # Preserve existing value for critical keys - continue - - if key in result and isinstance(result[key], dict) and isinstance(default_value, dict): - # Recursively merge nested dictionaries - result[key] = deep_merge_dicts(result[key], default_value, preserve_keys) - elif key not in result or result[key] is None: - # Use default value if key doesn't exist or is None - result[key] = copy.deepcopy(default_value) - # For non-dict values, keep existing value unless it's None - - return result - # Convert configs to dictionaries existing_dict = existing_config.model_dump(mode="json") default_dict = default_config.model_dump(mode="json") - # Merge text_mem config + logger.info( + f"Starting config merge for user {existing_config.user_id}, cube {existing_config.cube_id}" + ) + + # Define fields that should be preserved from existing config + preserve_fields = {"user_id", "cube_id", "config_filename", "model_schema"} + + # Preserve graph_db from existing config if it exists, but merge some fields + preserved_graph_db = None if "text_mem" in existing_dict and "text_mem" in default_dict: existing_text_config = existing_dict["text_mem"].get("config", {}) default_text_config = default_dict["text_mem"].get("config", {}) - # Handle nested graph_db config specially if "graph_db" in existing_text_config and "graph_db" in default_text_config: - existing_graph_config = existing_text_config["graph_db"].get("config", {}) - default_graph_config = default_text_config["graph_db"].get("config", {}) - - # Merge graph_db config, preserving critical keys - merged_graph_config = deep_merge_dicts( - existing_graph_config, - default_graph_config, - preserve_keys={"uri", "user", "password", "db_name", "auto_create"}, - ) - - # Update the configs - existing_text_config["graph_db"]["config"] = merged_graph_config - default_text_config["graph_db"]["config"] = merged_graph_config - - # Merge other text_mem config fields - merged_text_config = deep_merge_dicts(existing_text_config, default_text_config) - existing_dict["text_mem"]["config"] = merged_text_config - - # Merge act_mem config - if "act_mem" in existing_dict and "act_mem" in default_dict: - existing_act_config = existing_dict["act_mem"].get("config", {}) - default_act_config = default_dict["act_mem"].get("config", {}) - merged_act_config = deep_merge_dicts(existing_act_config, default_act_config) - existing_dict["act_mem"]["config"] = merged_act_config - - # Merge para_mem config - if "para_mem" in existing_dict and "para_mem" in default_dict: - existing_para_config = existing_dict["para_mem"].get("config", {}) - default_para_config = default_dict["para_mem"].get("config", {}) - merged_para_config = deep_merge_dicts(existing_para_config, default_para_config) - existing_dict["para_mem"]["config"] = merged_para_config + existing_graph_config = existing_text_config["graph_db"]["config"] + default_graph_config = default_text_config["graph_db"]["config"] + + # Define graph_db fields to preserve (user-specific) + preserve_graph_fields = { + "uri", + "user", + "password", + "db_name", + "auto_create", + "user_name", + } + + # Create merged graph_db config + merged_graph_config = copy.deepcopy(existing_graph_config) + for key, value in default_graph_config.items(): + if key not in preserve_graph_fields: + merged_graph_config[key] = value + logger.debug( + f"Updated graph_db field '{key}': {existing_graph_config.get(key)} -> {value}" + ) + + preserved_graph_db = { + "backend": existing_text_config["graph_db"]["backend"], + "config": merged_graph_config, + } + + # Use default config as base + merged_dict = copy.deepcopy(default_dict) + + # Restore preserved fields from existing config + for field in preserve_fields: + if field in existing_dict: + merged_dict[field] = existing_dict[field] + logger.debug(f"Preserved field '{field}': {existing_dict[field]}") + + # Restore graph_db if it was preserved + if preserved_graph_db and "text_mem" in merged_dict: + merged_dict["text_mem"]["config"]["graph_db"] = preserved_graph_db + logger.debug(f"Preserved graph_db with merged config: {preserved_graph_db}") # Create new config from merged dictionary - merged_config = GeneralMemCubeConfig.model_validate(existing_dict) + merged_config = GeneralMemCubeConfig.model_validate(merged_dict) + logger.info( - f"Merged cube config for user {merged_config.user_id}, cube {merged_config.cube_id}" + f"Successfully merged cube config for user {merged_config.user_id}, cube {merged_config.cube_id}" ) return merged_config diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index 220b7deb8..15463a6ac 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -544,6 +544,9 @@ def register_mem_cube( ) # Register the MemCube + logger.info( + f"Registering MemCube {mem_cube_id} with cube config {mem_cube.config.model_dump(mode='json')}" + ) self.mem_cubes[mem_cube_id] = mem_cube def user_register( From bceb1414828d10aad7007d9bacbc7d5ec42dafb7 Mon Sep 17 00:00:00 2001 From: fridayL Date: Wed, 16 Jul 2025 07:55:55 +0000 Subject: [PATCH 02/14] fix:dim --- src/memos/api/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/memos/api/config.py b/src/memos/api/config.py index ce650d0f0..dab4f3437 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -102,6 +102,7 @@ def get_embedder_config() -> dict[str, Any]: "api_key": os.getenv("OPENAI_API_KEY", "sk-xxxx"), "model_name_or_path": os.getenv("MOS_EMBEDDER_MODEL", "text-embedding-3-large"), "base_url": os.getenv("OPENAI_API_BASE", "http://openai.com"), + "embedding_dimension": 3072, }, } else: # ollama From 1ad36c2c8fd4aa72ad13489f188277a8404b4ff5 Mon Sep 17 00:00:00 2001 From: fridayL Date: Wed, 16 Jul 2025 08:01:15 +0000 Subject: [PATCH 03/14] change dim --- src/memos/api/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/memos/api/config.py b/src/memos/api/config.py index dab4f3437..9dcbd516c 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -102,7 +102,6 @@ def get_embedder_config() -> dict[str, Any]: "api_key": os.getenv("OPENAI_API_KEY", "sk-xxxx"), "model_name_or_path": os.getenv("MOS_EMBEDDER_MODEL", "text-embedding-3-large"), "base_url": os.getenv("OPENAI_API_BASE", "http://openai.com"), - "embedding_dimension": 3072, }, } else: # ollama @@ -125,6 +124,7 @@ def get_neo4j_config() -> dict[str, Any]: "db_name": "default", "password": os.getenv("NEO4J_PASSWORD", "12345678"), "auto_create": True, + "embedding_dimension": 3072, } @staticmethod From 9a71fc0c46c3ef1f0756a7a320ca5ec01444a55a Mon Sep 17 00:00:00 2001 From: fridayL Date: Wed, 16 Jul 2025 09:38:53 +0000 Subject: [PATCH 04/14] fix:change default db --- src/memos/api/config.py | 49 ++++++++++++++++++++++--------------- src/memos/mem_cube/utils.py | 9 +++++++ src/memos/mem_os/product.py | 2 +- 3 files changed, 39 insertions(+), 21 deletions(-) diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 9dcbd516c..feb267ccb 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -116,14 +116,37 @@ def get_embedder_config() -> dict[str, Any]: } @staticmethod - def get_neo4j_config() -> dict[str, Any]: + def get_neo4j_config(user_id: str | None = None) -> dict[str, Any]: + """Get Neo4j configuration.""" + if os.getenv("MOS_NEO4J_SHARED_DB", "false").lower() == "true": + return APIConfig.get_neo4j_shared_config(user_id) + else: + return APIConfig.get_noshared_neo4j_config(user_id) + + @staticmethod + def get_noshared_neo4j_config(user_id) -> dict[str, Any]: """Get Neo4j configuration.""" return { "uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"), "user": os.getenv("NEO4J_USER", "neo4j"), - "db_name": "default", + "db_name": f"memos{user_id.replace('-', '')}", "password": os.getenv("NEO4J_PASSWORD", "12345678"), "auto_create": True, + "use_multi_db": True, + "embedding_dimension": 3072, + } + + @staticmethod + def get_neo4j_shared_config(user_id: str | None = None) -> dict[str, Any]: + """Get Neo4j configuration.""" + return { + "uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"), + "user": os.getenv("NEO4J_USER", "neo4j"), + "db_name": os.getenv("NEO4J_DB_NAME", "shared-tree-textual-memory"), + "password": os.getenv("NEO4J_PASSWORD", "12345678"), + "user_name": f"memos{user_id.replace('-', '')}", + "auto_create": True, + "use_multi_db": False, "embedding_dimension": 3072, } @@ -251,7 +274,7 @@ def get_start_default_config() -> dict[str, Any]: def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, GeneralMemCube]: """Create configuration for a specific user.""" openai_config = APIConfig.get_openai_config() - neo4j_config = APIConfig.get_neo4j_config() + neo4j_config = APIConfig.get_neo4j_config(user_id) qwen_config = APIConfig.qwen_config() vllm_config = APIConfig.vllm_config() backend = os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai") @@ -314,23 +337,9 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General "dispatcher_llm": {"backend": "openai", "config": openai_config}, "graph_db": { "backend": "neo4j", - "config": { - "uri": neo4j_config["uri"], - "user": neo4j_config["user"], - "password": neo4j_config["password"], - "db_name": os.getenv( - "NEO4J_DB_NAME", f"memos{user_id.replace('-', '')}" - ), # , replace with - "auto_create": neo4j_config["auto_create"], - }, - }, - "embedder": { - "backend": "ollama", - "config": { - "model_name_or_path": "nomic-embed-text:latest", - "api_base": os.getenv("OLLAMA_API_BASE", "http://localhost:11434"), - }, + "config": neo4j_config, }, + "embedder": APIConfig.get_embedder_config(), }, }, "act_mem": {} @@ -354,7 +363,7 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None: return None openai_config = APIConfig.get_openai_config() - neo4j_config = APIConfig.get_neo4j_config() + neo4j_config = APIConfig.get_neo4j_config(user_id="default") return GeneralMemCubeConfig.model_validate( { diff --git a/src/memos/mem_cube/utils.py b/src/memos/mem_cube/utils.py index 4a3b830c5..0e7afaf39 100644 --- a/src/memos/mem_cube/utils.py +++ b/src/memos/mem_cube/utils.py @@ -77,6 +77,7 @@ def merge_config_with_default( "db_name", "auto_create", "user_name", + "use_multi_db", } # Create merged graph_db config @@ -87,6 +88,14 @@ def merge_config_with_default( logger.debug( f"Updated graph_db field '{key}': {existing_graph_config.get(key)} -> {value}" ) + if not default_graph_config.get("use_multi_db", True): + # set original use_multi_db to False if default_graph_config.use_multi_db is False + if merged_graph_config.get("use_multi_db", True): + merged_graph_config["use_multi_db"] = False + merged_graph_config["user_name"] = merged_graph_config.get("db_name") + merged_graph_config["db_name"] = default_graph_config.get("db_name") + else: + logger.info("use_multi_db is already False, no need to change") preserved_graph_db = { "backend": existing_text_config["graph_db"]["backend"], diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index 15463a6ac..737e05a8a 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -772,7 +772,7 @@ def chat_with_references( current_messages = [ {"role": "system", "content": system_prompt}, *chat_history.chat_history, - {"role": "user", "content": query + "/nothink"}, + {"role": "user", "content": query}, ] # Generate response with custom prompt From 58addaf8e7eb9015c99ef943d92791d0e6dd6afe Mon Sep 17 00:00:00 2001 From: fridayL Date: Wed, 16 Jul 2025 10:57:53 +0000 Subject: [PATCH 05/14] =?UTF-8?q?fix=EF=BC=9Adelay?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/memos/api/routers/product_router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/memos/api/routers/product_router.py b/src/memos/api/routers/product_router.py index 96feba668..92acb38a4 100644 --- a/src/memos/api/routers/product_router.py +++ b/src/memos/api/routers/product_router.py @@ -222,7 +222,7 @@ async def generate_chat_response(): history=chat_req.history, ): yield chunk - await asyncio.sleep(0.05) # 50ms delay between chunks + await asyncio.sleep(0.00001) # 50ms delay between chunks except Exception as e: logger.error(f"Error in chat stream: {e}") error_data = f"data: {json.dumps({'type': 'error', 'content': str(traceback.format_exc())})}\n\n" From 1dc8d6f8acb805de803ff849e5ae22a0771dd231 Mon Sep 17 00:00:00 2001 From: fridayL Date: Wed, 16 Jul 2025 11:45:23 +0000 Subject: [PATCH 06/14] fix:len --- src/memos/mem_os/product.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index 737e05a8a..2cab6fb95 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -882,7 +882,7 @@ def chat_with_references( mem_cube_id=cube_id, ) # Keep chat history under 30 messages by removing oldest conversation pair - if len(self.chat_history_manager[user_id].chat_history) > 30: + if len(self.chat_history_manager[user_id].chat_history) > 10: self.chat_history_manager[user_id].chat_history.pop(0) # Remove oldest user message self.chat_history_manager[user_id].chat_history.pop( 0 From 3eccccbdb142e7c868b62e848631d8d968dbd3d4 Mon Sep 17 00:00:00 2001 From: fridayL Date: Thu, 17 Jul 2025 03:24:56 +0000 Subject: [PATCH 07/14] fix:change recently mem size --- src/memos/mem_os/product.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index 2cab6fb95..d1ae206e3 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -666,7 +666,7 @@ def get_suggestion_query(self, user_id: str, language: str = "zh") -> list[str]: "query": ["query1", "query2", "query3"] }} """ - text_mem_result = super().search("my recently memories", user_id=user_id, top_k=10)[ + text_mem_result = super().search("my recently memories", user_id=user_id, top_k=3)[ "text_mem" ] if text_mem_result: From bb5ee43069ae943691c700917e87029d89a89ac9 Mon Sep 17 00:00:00 2001 From: fridayL Date: Thu, 17 Jul 2025 07:02:02 +0000 Subject: [PATCH 08/14] fix:dup node error --- src/memos/mem_os/product.py | 7 +- src/memos/mem_os/utils/format_utils.py | 190 ++++++++++++++++++++++++- 2 files changed, 195 insertions(+), 2 deletions(-) diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index d1ae206e3..a8223da85 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -17,6 +17,7 @@ from memos.mem_os.core import MOSCore from memos.mem_os.utils.format_utils import ( convert_graph_to_tree_forworkmem, + ensure_unique_tree_ids, filter_nodes_by_tree_ids, remove_embedding_recursive, sort_children_by_memory_type, @@ -920,8 +921,10 @@ def get_all( "UserMemory": 0.40, } tree_result, node_type_count = convert_graph_to_tree_forworkmem( - memories, target_node_count=150, type_ratios=custom_type_ratios + memories, target_node_count=200, type_ratios=custom_type_ratios ) + # Ensure all node IDs are unique in the tree structure + tree_result = ensure_unique_tree_ids(tree_result) memories_filtered = filter_nodes_by_tree_ids(tree_result, memories) children = tree_result["children"] children_sort = sort_children_by_memory_type(children) @@ -1011,6 +1014,8 @@ def get_subgraph( tree_result, node_type_count = convert_graph_to_tree_forworkmem( memories, target_node_count=150, type_ratios=custom_type_ratios ) + # Ensure all node IDs are unique in the tree structure + tree_result = ensure_unique_tree_ids(tree_result) memories_filtered = filter_nodes_by_tree_ids(tree_result, memories) children = tree_result["children"] children_sort = sort_children_by_memory_type(children) diff --git a/src/memos/mem_os/utils/format_utils.py b/src/memos/mem_os/utils/format_utils.py index 9722779f0..344e3b905 100644 --- a/src/memos/mem_os/utils/format_utils.py +++ b/src/memos/mem_os/utils/format_utils.py @@ -493,7 +493,7 @@ def sample_nodes_by_importance( # Modified main function to use new sampling strategy def convert_graph_to_tree_forworkmem( json_data: dict[str, Any], - target_node_count: int = 150, + target_node_count: int = 200, type_ratios: dict[str, float] | None = None, ) -> dict[str, Any]: """ @@ -1152,3 +1152,191 @@ def convert_activation_memory_summary(act_mem_items: list[KVCacheItem]) -> dict[ "total_parameters": total_parameters, "summary": f"Activation memory contains {total_items} items with {total_layers} layers and approximately {total_parameters:,} parameters", } + + +def detect_and_remove_duplicate_ids(tree_node: dict[str, Any]) -> dict[str, Any]: + """ + Detect and remove duplicate IDs in tree structure by skipping duplicate nodes. + First occurrence of each ID is kept, subsequent duplicates are removed. + + Args: + tree_node: Tree node (dictionary format) + + Returns: + dict: Fixed tree node with duplicate nodes removed + """ + used_ids = set() + removed_count = 0 + + def remove_duplicates_recursive( + node: dict[str, Any], parent_path: str = "" + ) -> dict[str, Any] | None: + """Recursively remove duplicate IDs by skipping duplicate nodes""" + nonlocal removed_count + + if not isinstance(node, dict): + return node + + # Create node copy + fixed_node = node.copy() + + # Handle current node ID + current_id = fixed_node.get("id", "") + if current_id in used_ids and current_id not in ["root", "WorkingMemory"]: + # Skip this duplicate node + print(f"Skipping duplicate node: {current_id} (path: {parent_path})") + removed_count += 1 + return None # Return None to indicate this node should be removed + else: + used_ids.add(current_id) + + # Recursively process child nodes + if "children" in fixed_node and isinstance(fixed_node["children"], list): + fixed_children = [] + for i, child in enumerate(fixed_node["children"]): + child_path = f"{parent_path}/{fixed_node.get('node_name', 'unknown')}[{i}]" + fixed_child = remove_duplicates_recursive(child, child_path) + if fixed_child is not None: # Only add non-None children + fixed_children.append(fixed_child) + fixed_node["children"] = fixed_children + + return fixed_node + + result = remove_duplicates_recursive(tree_node) + if result is not None: + print(f"Removed {removed_count} duplicate nodes") + return result + else: + # If root node itself was removed (shouldn't happen), return empty root + return { + "id": "root", + "node_name": "root", + "value": "root", + "memory_type": "Root", + "children": [], + } + + +def validate_tree_structure(tree_node: dict[str, Any]) -> dict[str, Any]: + """ + Validate tree structure integrity, including ID uniqueness check + + Args: + tree_node: Tree node (dictionary format) + + Returns: + dict: Validation result containing error messages and fix suggestions + """ + validation_result = { + "is_valid": True, + "errors": [], + "warnings": [], + "total_nodes": 0, + "unique_ids": set(), + "duplicate_ids": set(), + "missing_ids": set(), + "invalid_structure": [], + } + + def validate_recursive(node: dict[str, Any], path: str = "", depth: int = 0): + """Recursively validate tree structure""" + if not isinstance(node, dict): + validation_result["errors"].append(f"Node is not a dictionary: {path}") + validation_result["is_valid"] = False + return + + validation_result["total_nodes"] += 1 + + # Check required fields + if "id" not in node: + validation_result["errors"].append(f"Node missing ID field: {path}") + validation_result["missing_ids"].add(path) + validation_result["is_valid"] = False + else: + node_id = node["id"] + if node_id in validation_result["unique_ids"]: + validation_result["errors"].append(f"Duplicate node ID: {node_id} (path: {path})") + validation_result["duplicate_ids"].add(node_id) + validation_result["is_valid"] = False + else: + validation_result["unique_ids"].add(node_id) + + # Check other required fields + required_fields = ["node_name", "value", "memory_type"] + for field in required_fields: + if field not in node: + validation_result["warnings"].append(f"Node missing field '{field}': {path}") + + # Recursively validate child nodes + if "children" in node: + if not isinstance(node["children"], list): + validation_result["errors"].append(f"Children field is not a list: {path}") + validation_result["is_valid"] = False + else: + for i, child in enumerate(node["children"]): + child_path = f"{path}/children[{i}]" + validate_recursive(child, child_path, depth + 1) + + # Check depth limit + if depth > 20: + validation_result["warnings"].append(f"Tree depth too deep ({depth}): {path}") + + validate_recursive(tree_node) + + # Generate fix suggestions + if validation_result["duplicate_ids"]: + validation_result["fix_suggestion"] = ( + "Use detect_and_fix_duplicate_ids() function to fix duplicate IDs" + ) + + return validation_result + + +def ensure_unique_tree_ids(tree_result: dict[str, Any]) -> dict[str, Any]: + """ + Ensure all node IDs in tree structure are unique by removing duplicate nodes, + this is a post-processing function for convert_graph_to_tree_forworkmem + + Args: + tree_result: Tree structure returned by convert_graph_to_tree_forworkmem + + Returns: + dict: Fixed tree structure with duplicate nodes removed + """ + print("šŸ” Starting duplicate ID check in tree structure...") + + # First validate tree structure + validation = validate_tree_structure(tree_result) + + if validation["is_valid"]: + print("Tree structure validation passed, no duplicate IDs found") + return tree_result + + # Report issues + print(f"Found {len(validation['errors'])} errors:") + for error in validation["errors"][:5]: # Only show first 5 errors + print(f" - {error}") + + if len(validation["errors"]) > 5: + print(f" ... and {len(validation['errors']) - 5} more errors") + + print("Statistics:") + print(f" - Total nodes: {validation['total_nodes']}") + print(f" - Unique IDs: {len(validation['unique_ids'])}") + print(f" - Duplicate IDs: {len(validation['duplicate_ids'])}") + + # Remove duplicate nodes + print(" Starting duplicate node removal...") + fixed_tree = detect_and_remove_duplicate_ids(tree_result) + + # Validate again + post_validation = validate_tree_structure(fixed_tree) + if post_validation["is_valid"]: + print("Removal completed, tree structure is now valid") + print(f"Final node count: {post_validation['total_nodes']}") + else: + print("Issues remain after removal, please check code logic") + for error in post_validation["errors"][:3]: + print(f" - {error}") + + return fixed_tree From 13639a360a900c54fa95a1379f0f4dcab71fcda1 Mon Sep 17 00:00:00 2001 From: fridayL Date: Fri, 18 Jul 2025 13:03:09 +0000 Subject: [PATCH 09/14] fix: remove mock_data --- src/memos/mem_os/product.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index a8223da85..5579d6139 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -36,8 +36,6 @@ load_dotenv() CUBE_PATH = os.getenv("MOS_CUBE_PATH", "/tmp/data/") -with open("./tmp/fake_data.json") as f: - MOCK_DATA = json.loads(f.read()) class MOSProduct(MOSCore): From ea63a1c204759dd83cae78639ff9bf560a9bd84f Mon Sep 17 00:00:00 2001 From: fridayL Date: Tue, 22 Jul 2025 03:30:55 +0000 Subject: [PATCH 10/14] fix: change config --- src/memos/api/config.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 79ba59a2e..a709a74a1 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -368,7 +368,6 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General }, "embedder": APIConfig.get_embedder_config(), }, - "embedder": APIConfig.get_embedder_config(), }, "act_mem": {} if os.getenv("ENABLE_ACTIVATION_MEMORY", "false").lower() == "false" From 81c87ca547057f147bb13a8a204a6bf0a90d2e8c Mon Sep 17 00:00:00 2001 From: fridayL Date: Tue, 22 Jul 2025 03:49:06 +0000 Subject: [PATCH 11/14] feat: reorganize code --- src/memos/api/config.py | 81 ++++++++++++----------------------------- 1 file changed, 23 insertions(+), 58 deletions(-) diff --git a/src/memos/api/config.py b/src/memos/api/config.py index a709a74a1..61e73f9ae 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -350,8 +350,15 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General default_config = MOSConfig(**config_dict) - if os.getenv("NEO4J_BACKEND", "neo4j_community").lower() == "neo4j_community": - neo4j_community_config = APIConfig.get_neo4j_community_config(user_id) + neo4j_community_config = APIConfig.get_neo4j_community_config(user_id) + neo4j_config = APIConfig.get_neo4j_config(user_id) + + graph_db_backend_map = { + "neo4j-community": neo4j_community_config, + "neo4j": neo4j_config, + } + graph_db_backend = os.getenv("NEO4J_BACKEND", "neo4j-community").lower() + if graph_db_backend in graph_db_backend_map: # Create MemCube config default_cube_config = GeneralMemCubeConfig.model_validate( { @@ -363,8 +370,8 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General "extractor_llm": {"backend": "openai", "config": openai_config}, "dispatcher_llm": {"backend": "openai", "config": openai_config}, "graph_db": { - "backend": "neo4j-community", - "config": neo4j_community_config, + "backend": graph_db_backend, + "config": graph_db_backend_map[graph_db_backend], }, "embedder": APIConfig.get_embedder_config(), }, @@ -376,30 +383,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General } ) else: - neo4j_config = APIConfig.get_neo4j_config(user_id) - # Create MemCube config - default_cube_config = GeneralMemCubeConfig.model_validate( - { - "user_id": user_id, - "cube_id": f"{user_name}_default_cube", - "text_mem": { - "backend": "tree_text", - "config": { - "extractor_llm": {"backend": "openai", "config": openai_config}, - "dispatcher_llm": {"backend": "openai", "config": openai_config}, - "graph_db": { - "backend": "neo4j", - "config": neo4j_config, - }, - "embedder": APIConfig.get_embedder_config(), - }, - }, - "act_mem": {} - if os.getenv("ENABLE_ACTIVATION_MEMORY", "false").lower() == "false" - else APIConfig.get_activation_vllm_config(), - "para_mem": {}, - } - ) + raise ValueError(f"Invalid Neo4j backend: {graph_db_backend}") default_mem_cube = GeneralMemCube(default_cube_config) return default_config, default_mem_cube @@ -415,9 +399,14 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None: return None openai_config = APIConfig.get_openai_config() - - if os.getenv("NEO4J_BACKEND", "neo4j_community").lower() == "neo4j_community": - neo4j_community_config = APIConfig.get_neo4j_community_config(user_id="default") + neo4j_community_config = APIConfig.get_neo4j_community_config(user_id="default") + neo4j_config = APIConfig.get_neo4j_config(user_id="default") + graph_db_backend_map = { + "neo4j-community": neo4j_community_config, + "neo4j": neo4j_config, + } + graph_db_backend = os.getenv("NEO4J_BACKEND", "neo4j-community").lower() + if graph_db_backend in graph_db_backend_map: return GeneralMemCubeConfig.model_validate( { "user_id": "default", @@ -428,8 +417,8 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None: "extractor_llm": {"backend": "openai", "config": openai_config}, "dispatcher_llm": {"backend": "openai", "config": openai_config}, "graph_db": { - "backend": "neo4j-community", - "config": neo4j_community_config, + "backend": graph_db_backend, + "config": graph_db_backend_map[graph_db_backend], }, "embedder": APIConfig.get_embedder_config(), "reorganize": os.getenv("MOS_ENABLE_REORGANIZE", "false").lower() @@ -443,28 +432,4 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None: } ) else: - neo4j_config = APIConfig.get_neo4j_config(user_id="default") - return GeneralMemCubeConfig.model_validate( - { - "user_id": "default", - "cube_id": "default_cube", - "text_mem": { - "backend": "tree_text", - "config": { - "extractor_llm": {"backend": "openai", "config": openai_config}, - "dispatcher_llm": {"backend": "openai", "config": openai_config}, - "graph_db": { - "backend": "neo4j", - "config": neo4j_config, - }, - "embedder": APIConfig.get_embedder_config(), - "reorganize": os.getenv("MOS_ENABLE_REORGANIZE", "false").lower() - == "true", - }, - }, - "act_mem": {} - if os.getenv("ENABLE_ACTIVATION_MEMORY", "false").lower() == "false" - else APIConfig.get_activation_vllm_config(), - "para_mem": {}, - } - ) + raise ValueError(f"Invalid Neo4j backend: {graph_db_backend}") From 5b51bec9bf172242aca2a2704f50c339ae354836 Mon Sep 17 00:00:00 2001 From: fridayL Date: Tue, 22 Jul 2025 08:36:27 +0000 Subject: [PATCH 12/14] add: add json parse for en --- src/memos/mem_os/product.py | 7 ++++--- src/memos/mem_os/utils/format_utils.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index 5579d6139..9c235d062 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -16,6 +16,7 @@ from memos.mem_cube.general import GeneralMemCube from memos.mem_os.core import MOSCore from memos.mem_os.utils.format_utils import ( + clean_json_response, convert_graph_to_tree_forworkmem, ensure_unique_tree_ids, filter_nodes_by_tree_ids, @@ -657,7 +658,7 @@ def get_suggestion_query(self, user_id: str, language: str = "zh") -> list[str]: you should generate some suggestion query, the query should be user what to query, user recently memories is: {memories} - please generate 3 suggestion query in English, + if the user recently memories is empty, please generate 3 suggestion query in English, output should be a json format, the key is "query", the value is a list of suggestion query. example: @@ -674,8 +675,8 @@ def get_suggestion_query(self, user_id: str, language: str = "zh") -> list[str]: memories = "" message_list = [{"role": "system", "content": suggestion_prompt.format(memories=memories)}] response = self.chat_llm.generate(message_list) - response_json = json.loads(response) - + clean_response = clean_json_response(response) + response_json = json.loads(clean_response) return response_json["query"] def chat( diff --git a/src/memos/mem_os/utils/format_utils.py b/src/memos/mem_os/utils/format_utils.py index 344e3b905..c6167f0bd 100644 --- a/src/memos/mem_os/utils/format_utils.py +++ b/src/memos/mem_os/utils/format_utils.py @@ -1340,3 +1340,16 @@ def ensure_unique_tree_ids(tree_result: dict[str, Any]) -> dict[str, Any]: print(f" - {error}") return fixed_tree + + +def clean_json_response(response: str) -> str: + """ + Remove markdown JSON code block formatting from LLM response. + + Args: + response: Raw response string that may contain ```json and ``` + + Returns: + str: Clean JSON string without markdown formatting + """ + return response.replace("```json", "").replace("```", "").strip() From 5dece52c4f9105198844787e18cdf5ecfda36b92 Mon Sep 17 00:00:00 2001 From: fridayL Date: Tue, 22 Jul 2025 08:49:15 +0000 Subject: [PATCH 13/14] fix:change user_id --- src/memos/mem_os/product.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index 9c235d062..ccb0e52f9 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -764,11 +764,10 @@ def chat_with_references( system_prompt = self._build_system_prompt(user_id, memories_list) # Get chat history - target_user_id = user_id if user_id is not None else self.user_id - if target_user_id not in self.chat_history_manager: - self._register_chat_history(target_user_id) + if user_id not in self.chat_history_manager: + self._register_chat_history(user_id) - chat_history = self.chat_history_manager[target_user_id] + chat_history = self.chat_history_manager[user_id] current_messages = [ {"role": "system", "content": system_prompt}, *chat_history.chat_history, From 2b7dae81e0812797bd0b8b435bf546eac0b4c7cd Mon Sep 17 00:00:00 2001 From: fridayL Date: Tue, 22 Jul 2025 09:28:35 +0000 Subject: [PATCH 14/14] fix: logger info --- src/memos/mem_os/core.py | 10 +- src/memos/mem_os/utils/format_utils.py | 164 +++++++++++++------------ 2 files changed, 85 insertions(+), 89 deletions(-) diff --git a/src/memos/mem_os/core.py b/src/memos/mem_os/core.py index d09380332..906420e27 100644 --- a/src/memos/mem_os/core.py +++ b/src/memos/mem_os/core.py @@ -646,14 +646,8 @@ def add( ) else: messages_list = [ - [ - {"role": "user", "content": memory_content}, - { - "role": "assistant", - "content": "", - }, # add by str to keep the format,assistant role is empty - ] - ] + [{"role": "user", "content": memory_content}] + ] # for only user-str input and convert message memories = self.mem_reader.get_memory( messages_list, type="chat", diff --git a/src/memos/mem_os/utils/format_utils.py b/src/memos/mem_os/utils/format_utils.py index c6167f0bd..8465b44f0 100644 --- a/src/memos/mem_os/utils/format_utils.py +++ b/src/memos/mem_os/utils/format_utils.py @@ -239,10 +239,10 @@ def sample_nodes_with_type_balance( "MetaMemory": 0.05, # 5% } - print( + logger.info( f"Starting type-balanced sampling, original nodes: {len(nodes)}, target nodes: {target_count}" ) - print(f"Target type ratios: {type_ratios}") + logger.info(f"Target type ratios: {type_ratios}") # Analyze current node type distribution current_type_counts = {} @@ -255,7 +255,7 @@ def sample_nodes_with_type_balance( nodes_by_type[memory_type] = [] nodes_by_type[memory_type].append(node) - print(f"Current type distribution: {current_type_counts}") + logger.info(f"Current type distribution: {current_type_counts}") # Calculate target node count for each type type_targets = {} @@ -290,7 +290,7 @@ def sample_nodes_with_type_balance( ) type_targets[memory_type] = type_targets.get(memory_type, 0) + additional - print(f"Target node count for each type: {type_targets}") + logger.info(f"Target node count for each type: {type_targets}") # Perform subtree quality sampling for each type selected_nodes = [] @@ -300,16 +300,18 @@ def sample_nodes_with_type_balance( continue type_nodes = nodes_by_type[memory_type] - print(f"\n--- Processing {memory_type} type: {len(type_nodes)} -> {target_for_type} ---") + logger.info( + f"\n--- Processing {memory_type} type: {len(type_nodes)} -> {target_for_type} ---" + ) if len(type_nodes) <= target_for_type: selected_nodes.extend(type_nodes) - print(f" Select all: {len(type_nodes)} nodes") + logger.info(f" Select all: {len(type_nodes)} nodes") else: # Use enhanced subtree quality sampling type_selected = sample_by_enhanced_subtree_quality(type_nodes, edges, target_for_type) selected_nodes.extend(type_selected) - print(f" Sampled selection: {len(type_selected)} nodes") + logger.info(f" Sampled selection: {len(type_selected)} nodes") # Filter edges selected_node_ids = {node["id"] for node in selected_nodes} @@ -319,8 +321,8 @@ def sample_nodes_with_type_balance( if edge["source"] in selected_node_ids and edge["target"] in selected_node_ids ] - print(f"\nFinal selected nodes: {len(selected_nodes)}") - print(f"Final edges: {len(filtered_edges)}") + logger.info(f"\nFinal selected nodes: {len(selected_nodes)}") + logger.info(f"Final edges: {len(filtered_edges)}") # Verify final type distribution final_type_counts = {} @@ -328,11 +330,11 @@ def sample_nodes_with_type_balance( memory_type = node.get("metadata", {}).get("memory_type", "Unknown") final_type_counts[memory_type] = final_type_counts.get(memory_type, 0) + 1 - print(f"Final type distribution: {final_type_counts}") + logger.info(f"Final type distribution: {final_type_counts}") for memory_type, count in final_type_counts.items(): percentage = count / len(selected_nodes) * 100 target_percentage = type_ratios.get(memory_type, 0) * 100 - print( + logger.info( f" {memory_type}: {count} nodes ({percentage:.1f}%, target: {target_percentage:.1f}%)" ) @@ -358,9 +360,9 @@ def sample_by_enhanced_subtree_quality( subtree_analysis.items(), key=lambda x: x[1]["quality_score"], reverse=True ) - print(" Subtree quality ranking:") + logger.info(" Subtree quality ranking:") for i, (root_id, analysis) in enumerate(sorted_subtrees[:5]): - print( + logger.info( f" #{i + 1} Root node {root_id}: Quality={analysis['quality_score']:.2f}, " f"Depth={analysis['max_depth']}, Branches={analysis['branch_nodes']}, " f"Leaves={analysis['leaf_count']}, Max Width={analysis['max_width']}" @@ -386,7 +388,7 @@ def sample_by_enhanced_subtree_quality( if node: selected_nodes.append(node) selected_node_ids.add(node_id) - print(f" Select entire subtree {root_id}: +{len(new_nodes)} nodes") + logger.info(f" Select entire subtree {root_id}: +{len(new_nodes)} nodes") else: # Subtree too large, need partial selection if analysis["quality_score"] > 5: # Only partial selection for high-quality subtrees @@ -398,7 +400,7 @@ def sample_by_enhanced_subtree_quality( selected_nodes.extend(partial_selection) for node in partial_selection: selected_node_ids.add(node["id"]) - print( + logger.info( f" Partial selection of subtree {root_id}: +{len(partial_selection)} nodes" ) @@ -411,7 +413,7 @@ def sample_by_enhanced_subtree_quality( remaining_count = target_count - len(selected_nodes) additional = sample_nodes_by_importance(remaining_nodes, edges, remaining_count) selected_nodes.extend(additional) - print(f" Supplementary selection: +{len(additional)} nodes") + logger.info(f" Supplementary selection: +{len(additional)} nodes") return selected_nodes @@ -502,8 +504,8 @@ def convert_graph_to_tree_forworkmem( original_nodes = json_data.get("nodes", []) original_edges = json_data.get("edges", []) - print(f"Original node count: {len(original_nodes)}") - print(f"Target node count: {target_node_count}") + logger.info(f"Original node count: {len(original_nodes)}") + logger.info(f"Target node count: {target_node_count}") filter_original_edges = [] for original_edge in original_edges: if original_edge["type"] == "PARENT": @@ -633,7 +635,7 @@ def build_tree(node_id: str) -> dict[str, Any]: def print_tree_structure(node: dict[str, Any], level: int = 0, max_level: int = 5): - """Print the first few layers of tree structure for easy viewing""" + """logger.info the first few layers of tree structure for easy viewing""" if level > max_level: return @@ -647,21 +649,21 @@ def print_tree_structure(node: dict[str, Any], level: int = 0, max_level: int = children = node.get("children", []) if children: # Intermediate node, display name, type and child count - print(f"{indent}- {node_name} [{memory_type}] ({len(children)} children)") - print(f"{indent} ID: {node_id}") + logger.info(f"{indent}- {node_name} [{memory_type}] ({len(children)} children)") + logger.info(f"{indent} ID: {node_id}") display_value = node_value[:80] + "..." if len(node_value) > 80 else node_value - print(f"{indent} Value: {display_value}") + logger.info(f"{indent} Value: {display_value}") if level < max_level: for child in children: print_tree_structure(child, level + 1, max_level) elif level == max_level: - print(f"{indent} ... (expansion limited)") + logger.info(f"{indent} ... (expansion limited)") else: # Leaf node, display name, type and value display_value = node_value[:80] + "..." if len(node_value) > 80 else node_value - print(f"{indent}- {node_name} [{memory_type}]: {display_value}") - print(f"{indent} ID: {node_id}") + logger.info(f"{indent}- {node_name} [{memory_type}]: {display_value}") + logger.info(f"{indent} ID: {node_id}") def analyze_final_tree_quality(tree_data: dict[str, Any]) -> dict: @@ -856,107 +858,107 @@ def count_subtree(subnode, subdepth): def print_tree_analysis(tree_data: dict[str, Any]): - """Print enhanced tree analysis results""" + """logger.info enhanced tree analysis results""" stats = analyze_final_tree_quality(tree_data) - print("\n" + "=" * 60) - print("šŸŒ³ Enhanced Tree Structure Quality Analysis Report") - print("=" * 60) + logger.info("\n" + "=" * 60) + logger.info("šŸŒ³ Enhanced Tree Structure Quality Analysis Report") + logger.info("=" * 60) # Basic statistics - print("\nšŸ“Š Basic Statistics:") - print(f" Total nodes: {stats['total_nodes']}") - print(f" Max depth: {stats['max_depth']}") - print( + logger.info("\nšŸ“Š Basic Statistics:") + logger.info(f" Total nodes: {stats['total_nodes']}") + logger.info(f" Max depth: {stats['max_depth']}") + logger.info( f" Leaf nodes: {stats['total_leaves']} ({stats['total_leaves'] / stats['total_nodes'] * 100:.1f}%)" ) - print( + logger.info( f" Branch nodes: {stats['total_branches']} ({stats['total_branches'] / stats['total_nodes'] * 100:.1f}%)" ) # Structure quality assessment structure = stats.get("structure_quality", {}) if structure: - print("\nšŸ—ļø Structure Quality Assessment:") - print( + logger.info("\nšŸ—ļø Structure Quality Assessment:") + logger.info( f" Branch density: {structure['branch_density']:.3f} ({'āœ… Good' if 0.2 <= structure['branch_density'] <= 0.6 else 'āš ļø Needs improvement'})" ) - print( + logger.info( f" Leaf ratio: {structure['leaf_ratio']:.3f} ({'āœ… Good' if 0.3 <= structure['leaf_ratio'] <= 0.7 else 'āš ļø Needs improvement'})" ) - print(f" Max width: {structure['max_width']}") - print( + logger.info(f" Max width: {structure['max_width']}") + logger.info( f" Depth-width ratio: {structure['depth_width_ratio']:.2f} ({'āœ… Good' if structure['depth_width_ratio'] <= 3 else 'āš ļø Too thin'})" ) - print( + logger.info( f" Overall balance: {'āœ… Good' if structure['is_well_balanced'] else 'āš ļø Needs improvement'}" ) # Single chain analysis chain_analysis = stats.get("chain_analysis", {}) if chain_analysis: - print("\nšŸ”— Single Chain Structure Analysis:") - print(f" Longest chain: {chain_analysis.get('max_chain_length', 0)} layers") - print(f" Single chain subtrees: {chain_analysis.get('single_chain_subtrees', 0)}") - print( + logger.info("\nšŸ”— Single Chain Structure Analysis:") + logger.info(f" Longest chain: {chain_analysis.get('max_chain_length', 0)} layers") + logger.info(f" Single chain subtrees: {chain_analysis.get('single_chain_subtrees', 0)}") + logger.info( f" Single chain subtree ratio: {chain_analysis.get('chain_subtree_ratio', 0) * 100:.1f}%" ) if chain_analysis.get("max_chain_length", 0) > 5: - print(" āš ļø Warning: Overly long single chain structure may affect display") + logger.info(" āš ļø Warning: Overly long single chain structure may affect display") elif chain_analysis.get("chain_subtree_ratio", 0) > 0.3: - print( + logger.info( " āš ļø Warning: Too many single chain subtrees, suggest increasing branch structure" ) else: - print(" āœ… Single chain structure well controlled") + logger.info(" āœ… Single chain structure well controlled") # Type diversity type_div = stats.get("type_diversity", {}) if type_div: - print("\nšŸŽØ Type Diversity Analysis:") - print(f" Total types: {type_div['total_types']}") - print(f" Diversity index: {type_div['shannon_diversity']:.3f}") - print(f" Normalized diversity: {type_div['normalized_diversity']:.3f}") - print(f" Distribution balance: {type_div['distribution_balance']:.3f}") + logger.info("\nšŸŽØ Type Diversity Analysis:") + logger.info(f" Total types: {type_div['total_types']}") + logger.info(f" Diversity index: {type_div['shannon_diversity']:.3f}") + logger.info(f" Normalized diversity: {type_div['normalized_diversity']:.3f}") + logger.info(f" Distribution balance: {type_div['distribution_balance']:.3f}") # Type distribution - print("\nšŸ“‹ Type Distribution Details:") + logger.info("\nšŸ“‹ Type Distribution Details:") for mem_type, count in sorted(stats["by_type"].items(), key=lambda x: x[1], reverse=True): percentage = count / stats["total_nodes"] * 100 - print(f" {mem_type}: {count} nodes ({percentage:.1f}%)") + logger.info(f" {mem_type}: {count} nodes ({percentage:.1f}%)") # Depth distribution - print("\nšŸ“ Depth Distribution:") + logger.info("\nšŸ“ Depth Distribution:") for depth in sorted(stats["by_depth"].keys()): count = stats["by_depth"][depth] - print(f" Depth {depth}: {count} nodes") + logger.info(f" Depth {depth}: {count} nodes") # Major subtree analysis if stats["subtrees"]: - print("\nšŸŒ² Major Subtree Analysis (sorted by quality):") + logger.info("\nšŸŒ² Major Subtree Analysis (sorted by quality):") sorted_subtrees = sorted( stats["subtrees"], key=lambda x: x.get("quality_score", 0), reverse=True ) for i, subtree in enumerate(sorted_subtrees[:8]): # Show first 8 quality = subtree.get("quality_score", 0) - print(f" #{i + 1} {subtree['root']} [{subtree['type']}]:") - print(f" Quality score: {quality:.2f}") - print( + logger.info(f" #{i + 1} {subtree['root']} [{subtree['type']}]:") + logger.info(f" Quality score: {quality:.2f}") + logger.info( f" Structure: Depth={subtree['depth']}, Branches={subtree['branches']}, Leaves={subtree['leaves']}" ) - print( + logger.info( f" Density: Branch density={subtree.get('branch_density', 0):.3f}, Leaf ratio={subtree.get('leaf_ratio', 0):.3f}" ) if quality > 15: - print(" āœ… High quality subtree") + logger.info(" āœ… High quality subtree") elif quality > 8: - print(" šŸŸ” Medium quality subtree") + logger.info(" šŸŸ” Medium quality subtree") else: - print(" šŸ”“ Low quality subtree") + logger.info(" šŸ”“ Low quality subtree") - print("\n" + "=" * 60) + logger.info("\n" + "=" * 60) def remove_embedding_recursive(memory_info: dict) -> Any: @@ -1184,7 +1186,7 @@ def remove_duplicates_recursive( current_id = fixed_node.get("id", "") if current_id in used_ids and current_id not in ["root", "WorkingMemory"]: # Skip this duplicate node - print(f"Skipping duplicate node: {current_id} (path: {parent_path})") + logger.info(f"Skipping duplicate node: {current_id} (path: {parent_path})") removed_count += 1 return None # Return None to indicate this node should be removed else: @@ -1204,7 +1206,7 @@ def remove_duplicates_recursive( result = remove_duplicates_recursive(tree_node) if result is not None: - print(f"Removed {removed_count} duplicate nodes") + logger.info(f"Removed {removed_count} duplicate nodes") return result else: # If root node itself was removed (shouldn't happen), return empty root @@ -1303,41 +1305,41 @@ def ensure_unique_tree_ids(tree_result: dict[str, Any]) -> dict[str, Any]: Returns: dict: Fixed tree structure with duplicate nodes removed """ - print("šŸ” Starting duplicate ID check in tree structure...") + logger.info("šŸ” Starting duplicate ID check in tree structure...") # First validate tree structure validation = validate_tree_structure(tree_result) if validation["is_valid"]: - print("Tree structure validation passed, no duplicate IDs found") + logger.info("Tree structure validation passed, no duplicate IDs found") return tree_result # Report issues - print(f"Found {len(validation['errors'])} errors:") + logger.info(f"Found {len(validation['errors'])} errors:") for error in validation["errors"][:5]: # Only show first 5 errors - print(f" - {error}") + logger.info(f" - {error}") if len(validation["errors"]) > 5: - print(f" ... and {len(validation['errors']) - 5} more errors") + logger.info(f" ... and {len(validation['errors']) - 5} more errors") - print("Statistics:") - print(f" - Total nodes: {validation['total_nodes']}") - print(f" - Unique IDs: {len(validation['unique_ids'])}") - print(f" - Duplicate IDs: {len(validation['duplicate_ids'])}") + logger.info("Statistics:") + logger.info(f" - Total nodes: {validation['total_nodes']}") + logger.info(f" - Unique IDs: {len(validation['unique_ids'])}") + logger.info(f" - Duplicate IDs: {len(validation['duplicate_ids'])}") # Remove duplicate nodes - print(" Starting duplicate node removal...") + logger.info(" Starting duplicate node removal...") fixed_tree = detect_and_remove_duplicate_ids(tree_result) # Validate again post_validation = validate_tree_structure(fixed_tree) if post_validation["is_valid"]: - print("Removal completed, tree structure is now valid") - print(f"Final node count: {post_validation['total_nodes']}") + logger.info("Removal completed, tree structure is now valid") + logger.info(f"Final node count: {post_validation['total_nodes']}") else: - print("Issues remain after removal, please check code logic") + logger.info("Issues remain after removal, please check code logic") for error in post_validation["errors"][:3]: - print(f" - {error}") + logger.info(f" - {error}") return fixed_tree