fix:code

Author: fridayL

examples/mem_api/pipeline_test.py

@@ -7,7 +7,7 @@
 
 import time
 
-from typing import Optional, Any
+from typing import Any
 
 from dotenv import load_dotenv
 
@@ -91,7 +91,7 @@ def test_search_memories(
     mode: str = "fast",
     internet_search: bool = False,
     moscube: bool = False,
-    chat_history: Optional[list] = None,
+    chat_history: list | None = None,
 ) -> list[Any]:
     """
     Test searching memories from the system.
@@ -134,10 +134,6 @@ def test_search_memories(
         },
     )
 
-    elapsed_time = time.time() - time_start
-    logger.info(f"Search completed in {elapsed_time:.2f}s")
-    logger.info(f"Found {len(search_results)} results")
-
     # Print search results
     for idx, result in enumerate(search_results, 1):
         logger.info(f"\n  Result {idx}:")

src/memos/graph_dbs/nebular.py

@@ -4,7 +4,7 @@
 from contextlib import suppress
 from datetime import datetime
 from threading import Lock
-from typing import Optional, TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 import numpy as np
 
@@ -417,7 +417,7 @@ def create_index(
 
     @timed
     def remove_oldest_memory(
-        self, memory_type: str, keep_latest: int, user_name: Optional[str] = None
+        self, memory_type: str, keep_latest: int, user_name: str | None = None
     ) -> None:
         """
         Remove all WorkingMemory nodes except the latest `keep_latest` entries.
@@ -443,7 +443,7 @@ def remove_oldest_memory(
 
     @timed
     def add_node(
-        self, id: str, memory: str, metadata: dict[str, Any], user_name: Optional[str] = None
+        self, id: str, memory: str, metadata: dict[str, Any], user_name: str | None = None
     ) -> None:
         """
         Insert or update a Memory node in NebulaGraph.
@@ -477,7 +477,7 @@ def add_node(
             )
 
     @timed
-    def node_not_exist(self, scope: str, user_name: Optional[str] = None) -> int:
+    def node_not_exist(self, scope: str, user_name: str | None = None) -> int:
         user_name = user_name if user_name else self.config.user_name
         filter_clause = f'n.memory_type = "{scope}" AND n.user_name = "{user_name}"'
         query = f"""
@@ -495,7 +495,7 @@ def node_not_exist(self, scope: str, user_name: Optional[str] = None) -> int:
             raise
 
     @timed
-    def update_node(self, id: str, fields: dict[str, Any], user_name: Optional[str] = None) -> None:
+    def update_node(self, id: str, fields: dict[str, Any], user_name: str | None = None) -> None:
         """
         Update node fields in Nebular, auto-converting `created_at` and `updated_at` to datetime type if present.
         """
@@ -516,7 +516,7 @@ def update_node(self, id: str, fields: dict[str, Any], user_name: Optional[str]
         self.execute_query(query)
 
     @timed
-    def delete_node(self, id: str, user_name: Optional[str] = None) -> None:
+    def delete_node(self, id: str, user_name: str | None = None) -> None:
         """
         Delete a node from the graph.
         Args:
@@ -531,7 +531,7 @@ def delete_node(self, id: str, user_name: Optional[str] = None) -> None:
         self.execute_query(query)
 
     @timed
-    def add_edge(self, source_id: str, target_id: str, type: str, user_name: Optional[str] = None):
+    def add_edge(self, source_id: str, target_id: str, type: str, user_name: str | None = None):
         """
         Create an edge from source node to target node.
         Args:
@@ -555,7 +555,9 @@ def add_edge(self, source_id: str, target_id: str, type: str, user_name: Optiona
             logger.error(f"Failed to insert edge: {e}", exc_info=True)
 
     @timed
-    def delete_edge(self, source_id: str, target_id: str, type: str, user_name: Optional[str] = None) -> None:
+    def delete_edge(
+        self, source_id: str, target_id: str, type: str, user_name: str | None = None
+    ) -> None:
         """
         Delete a specific edge between two nodes.
         Args:
@@ -575,7 +577,7 @@ def delete_edge(self, source_id: str, target_id: str, type: str, user_name: Opti
         self.execute_query(query)
 
     @timed
-    def get_memory_count(self, memory_type: str, user_name: Optional[str] = None) -> int:
+    def get_memory_count(self, memory_type: str, user_name: str | None = None) -> int:
         user_name = user_name if user_name else self.config.user_name
         query = f"""
                 MATCH (n@Memory)
@@ -592,7 +594,7 @@ def get_memory_count(self, memory_type: str, user_name: Optional[str] = None) ->
             return -1
 
     @timed
-    def count_nodes(self, scope: str, user_name: Optional[str] = None) -> int:
+    def count_nodes(self, scope: str, user_name: str | None = None) -> int:
         user_name = user_name if user_name else self.config.user_name
         query = f"""
                 MATCH (n@Memory)
@@ -611,7 +613,7 @@ def edge_exists(
         target_id: str,
         type: str = "ANY",
         direction: str = "OUTGOING",
-        user_name: Optional[str] = None,
+        user_name: str | None = None,
     ) -> bool:
         """
         Check if an edge exists between two nodes.
@@ -654,7 +656,7 @@ def edge_exists(
     @timed
     # Graph Query & Reasoning
     def get_node(
-        self, id: str, include_embedding: bool = False, user_name: Optional[str] = None
+        self, id: str, include_embedding: bool = False, user_name: str | None = None
     ) -> dict[str, Any] | None:
         """
         Retrieve a Memory node by its unique ID.
@@ -691,7 +693,11 @@ def get_node(
 
     @timed
     def get_nodes(
-        self, ids: list[str], include_embedding: bool = False, user_name: Optional[str] = None, **kwargs
+        self,
+        ids: list[str],
+        include_embedding: bool = False,
+        user_name: str | None = None,
+        **kwargs,
     ) -> list[dict[str, Any]]:
         """
         Retrieve the metadata and memory of a list of nodes.
@@ -734,7 +740,7 @@ def get_nodes(
 
     @timed
     def get_edges(
-        self, id: str, type: str = "ANY", direction: str = "ANY", user_name: Optional[str] = None
+        self, id: str, type: str = "ANY", direction: str = "ANY", user_name: str | None = None
     ) -> list[dict[str, str]]:
         """
         Get edges connected to a node, with optional type and direction filter.
@@ -796,7 +802,7 @@ def get_neighbors_by_tag(
         top_k: int = 5,
         min_overlap: int = 1,
         include_embedding: bool = False,
-        user_name: Optional[str] = None,
+        user_name: str | None = None,
     ) -> list[dict[str, Any]]:
         """
         Find top-K neighbor nodes with maximum tag overlap.
@@ -857,7 +863,9 @@ def get_neighbors_by_tag(
         return result
 
     @timed
-    def get_children_with_embeddings(self, id: str, user_name: Optional[str] = None) -> list[dict[str, Any]]:
+    def get_children_with_embeddings(
+        self, id: str, user_name: str | None = None
+    ) -> list[dict[str, Any]]:
         user_name = user_name if user_name else self.config.user_name
         where_user = f"AND p.user_name = '{user_name}' AND c.user_name = '{user_name}'"
 
@@ -883,7 +891,7 @@ def get_subgraph(
         center_id: str,
         depth: int = 2,
         center_status: str = "activated",
-        user_name: Optional[str] = None,
+        user_name: str | None = None,
     ) -> dict[str, Any]:
         """
         Retrieve a local subgraph centered at a given node.
@@ -955,7 +963,7 @@ def search_by_embedding(
         status: str | None = None,
         threshold: float | None = None,
         search_filter: dict | None = None,
-        user_name: Optional[str] = None,
+        user_name: str | None = None,
         **kwargs,
     ) -> list[dict]:
         """
@@ -1034,7 +1042,9 @@ def search_by_embedding(
             return []
 
     @timed
-    def get_by_metadata(self, filters: list[dict[str, Any]], user_name: Optional[str] = None) -> list[str]:
+    def get_by_metadata(
+        self, filters: list[dict[str, Any]], user_name: str | None = None
+    ) -> list[str]:
         """
         1. ADD logic: "AND" vs "OR"(support logic combination);
         2. Support nested conditional expressions;
@@ -1102,7 +1112,7 @@ def get_grouped_counts(
         group_fields: list[str],
         where_clause: str = "",
         params: dict[str, Any] | None = None,
-        user_name: Optional[str] = None,
+        user_name: str | None = None,
     ) -> list[dict[str, Any]]:
         """
         Count nodes grouped by any fields.
@@ -1167,7 +1177,7 @@ def get_grouped_counts(
         return output
 
     @timed
-    def clear(self, user_name: Optional[str] = None) -> None:
+    def clear(self, user_name: str | None = None) -> None:
         """
         Clear the entire graph if the target database exists.
 
@@ -1185,7 +1195,7 @@ def clear(self, user_name: Optional[str] = None) -> None:
 
     @timed
     def export_graph(
-        self, include_embedding: bool = False, user_name: Optional[str] = None
+        self, include_embedding: bool = False, user_name: str | None = None
     ) -> dict[str, Any]:
         """
         Export all graph nodes and edges in a structured form.
@@ -1263,7 +1273,7 @@ def export_graph(
         return {"nodes": nodes, "edges": edges}
 
     @timed
-    def import_graph(self, data: dict[str, Any], user_name: Optional[str] = None) -> None:
+    def import_graph(self, data: dict[str, Any], user_name: str | None = None) -> None:
         """
         Import the entire graph from a serialized dictionary.
 
@@ -1301,7 +1311,7 @@ def import_graph(self, data: dict[str, Any], user_name: Optional[str] = None) ->
 
     @timed
     def get_all_memory_items(
-        self, scope: str, include_embedding: bool = False, user_name: Optional[str] = None
+        self, scope: str, include_embedding: bool = False, user_name: str | None = None
     ) -> (list)[dict]:
         """
         Retrieve all memory items of a specific memory_type.
@@ -1342,7 +1352,7 @@ def get_all_memory_items(
 
     @timed
     def get_structure_optimization_candidates(
-        self, scope: str, include_embedding: bool = False, user_name: Optional[str] = None
+        self, scope: str, include_embedding: bool = False, user_name: str | None = None
     ) -> list[dict]:
         """
         Find nodes that are likely candidates for structure optimization:

src/memos/memories/textual/simple_tree.py

@@ -1,7 +1,7 @@
 import time
 
 from datetime import datetime
-from typing import Optional, Any
+from typing import Any
 
 from memos.configs.memory import TreeTextMemoryConfig
 from memos.embedders.base import BaseEmbedder
@@ -77,7 +77,7 @@ def __init__(
         logger.info(f"time init: internet_retriever time is: {time.time() - time_start_ir}")
 
     def add(
-        self, memories: list[TextualMemoryItem | dict[str, Any]], user_name: Optional[str] = None
+        self, memories: list[TextualMemoryItem | dict[str, Any]], user_name: str | None = None
     ) -> list[str]:
         """Add memories.
         Args:
@@ -91,11 +91,11 @@ def add(
         return self.memory_manager.add(memories, user_name=user_name)
 
     def replace_working_memory(
-        self, memories: list[TextualMemoryItem], user_name: Optional[str] = None
+        self, memories: list[TextualMemoryItem], user_name: str | None = None
     ) -> None:
         self.memory_manager.replace_working_memory(memories, user_name=user_name)
 
-    def get_working_memory(self, user_name: Optional[str] = None) -> list[TextualMemoryItem]:
+    def get_working_memory(self, user_name: str | None = None) -> list[TextualMemoryItem]:
         working_memories = self.graph_store.get_all_memory_items(
             scope="WorkingMemory", user_name=user_name
         )
@@ -106,7 +106,7 @@ def get_working_memory(self, user_name: Optional[str] = None) -> list[TextualMem
         )
         return sorted_items
 
-    def get_current_memory_size(self, user_name: Optional[str] = None) -> dict[str, int]:
+    def get_current_memory_size(self, user_name: str | None = None) -> dict[str, int]:
         """
         Get the current size of each memory type.
         This delegates to the MemoryManager.
@@ -123,7 +123,7 @@ def search(
         manual_close_internet: bool = False,
         moscube: bool = False,
         search_filter: dict | None = None,
-        user_name: Optional[str] = None,
+        user_name: str | None = None,
     ) -> list[TextualMemoryItem]:
         """Search for memories based on a query.
         User query -> TaskGoalParser -> MemoryPathResolver ->

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

@@ -14,7 +14,6 @@
     GraphStructureReorganizer,
     QueueMessage,
 )
-from typing import Optional
 
 
 logger = get_logger(__name__)
@@ -52,7 +51,7 @@ def __init__(
         )
         self._merged_threshold = merged_threshold
 
-    def add(self, memories: list[TextualMemoryItem], user_name: Optional[str] = None) -> list[str]:
+    def add(self, memories: list[TextualMemoryItem], user_name: str | None = None) -> list[str]:
         """
         Add new memories in parallel to different memory types (WorkingMemory, LongTermMemory, UserMemory).
         """
@@ -81,7 +80,7 @@ def add(self, memories: list[TextualMemoryItem], user_name: Optional[str] = None
         return added_ids
 
     def replace_working_memory(
-        self, memories: list[TextualMemoryItem], user_name: Optional[str] = None
+        self, memories: list[TextualMemoryItem], user_name: str | None = None
     ) -> None:
         """
         Replace WorkingMemory
@@ -107,14 +106,14 @@ def replace_working_memory(
         )
         self._refresh_memory_size(user_name=user_name)
 
-    def get_current_memory_size(self, user_name: Optional[str] = None) -> dict[str, int]:
+    def get_current_memory_size(self, user_name: str | None = None) -> dict[str, int]:
         """
         Return the cached memory type counts.
         """
         self._refresh_memory_size(user_name=user_name)
         return self.current_memory_size
 
-    def _refresh_memory_size(self, user_name: Optional[str] = None) -> None:
+    def _refresh_memory_size(self, user_name: str | None = None) -> None:
         """
         Query the latest counts from the graph store and update internal state.
         """
@@ -124,7 +123,7 @@ def _refresh_memory_size(self, user_name: Optional[str] = None) -> None:
         self.current_memory_size = {record["memory_type"]: record["count"] for record in results}
         logger.info(f"[MemoryManager] Refreshed memory sizes: {self.current_memory_size}")
 
-    def _process_memory(self, memory: TextualMemoryItem, user_name: Optional[str] = None):
+    def _process_memory(self, memory: TextualMemoryItem, user_name: str | None = None):
         """
         Process and add memory to different memory types (WorkingMemory, LongTermMemory, UserMemory).
         This method runs asynchronously to process each memory item.
@@ -133,7 +132,7 @@ def _process_memory(self, memory: TextualMemoryItem, user_name: Optional[str] =
 
         # Add to WorkingMemory
         working_id = self._add_memory_to_db(memory, "WorkingMemory", user_name)
-        # ids.append(working_id)
+        ids.append(working_id)
 
         # Add to LongTermMemory and UserMemory
         if memory.metadata.memory_type in ["LongTermMemory", "UserMemory"]:
@@ -145,7 +144,7 @@ def _process_memory(self, memory: TextualMemoryItem, user_name: Optional[str] =
         return ids
 
     def _add_memory_to_db(
-        self, memory: TextualMemoryItem, memory_type: str, user_name: Optional[str] = None
+        self, memory: TextualMemoryItem, memory_type: str, user_name: str | None = None
     ) -> str:
         """
         Add a single memory item to the graph store, with FIFO logic for WorkingMemory.
@@ -161,7 +160,7 @@ def _add_memory_to_db(
         return working_memory.id
 
     def _add_to_graph_memory(
-        self, memory: TextualMemoryItem, memory_type: str, user_name: Optional[str] = None
+        self, memory: TextualMemoryItem, memory_type: str, user_name: str | None = None
     ):
         """
         Generalized method to add memory to a graph-based memory type (e.g., LongTermMemory, UserMemory).