Merge branch 'dev' into fix/pass-source-doc-id-in-completion-log

Author: glin93

src/memos/api/handlers/chat_handler.py

@@ -405,33 +405,26 @@ def generate_chat_response() -> Generator[str, None, None]:
                         async_mode="sync",
                     )
 
-                    # Use first readable cube ID for scheduler (backward compatibility)
-                    scheduler_cube_id = (
-                        readable_cube_ids[0] if readable_cube_ids else chat_req.user_id
-                    )
-                    self._send_message_to_scheduler(
-                        user_id=chat_req.user_id,
-                        mem_cube_id=scheduler_cube_id,
-                        query=chat_req.query,
-                        label=QUERY_TASK_LABEL,
-                    )
-
                     # ====== first search text mem with parse goal ======
                     search_req = APISearchPlaygroundRequest(
                         query=chat_req.query,
                         user_id=chat_req.user_id,
                         readable_cube_ids=readable_cube_ids,
-                        mode=chat_req.mode,
+                        mode="fast",
                         internet_search=False,
-                        top_k=chat_req.top_k,
+                        top_k=5,
                         chat_history=chat_req.history,
                         session_id=chat_req.session_id,
-                        include_preference=chat_req.include_preference,
+                        include_preference=False,
                         pref_top_k=chat_req.pref_top_k,
                         filter=chat_req.filter,
+                        search_tool_memory=False,
                         playground_search_goal_parser=False,
                     )
+                    start_time = time.time()
                     search_response = self.search_handler.handle_search_memories(search_req)
+                    end_time = time.time()
+                    self.logger.info(f"first search time: {end_time - start_time}")
 
                     yield f"data: {json.dumps({'type': 'status', 'data': '1'})}\n\n"
 
@@ -447,17 +440,19 @@ def generate_chat_response() -> Generator[str, None, None]:
 
                     # Prepare reference data (first search)
                     reference = prepare_reference_data(filtered_memories)
-                    # get preference string
-                    pref_string = search_response.data.get("pref_string", "")
 
                     yield f"data: {json.dumps({'type': 'reference', 'data': reference})}\n\n"
 
-                    # Prepare preference markdown string
-                    if chat_req.include_preference:
-                        pref_list = search_response.data.get("pref_mem") or []
-                        pref_memories = pref_list[0].get("memories", []) if pref_list else []
-                        pref_md_string = self._build_pref_md_string_for_playground(pref_memories)
-                        yield f"data: {json.dumps({'type': 'pref_md_string', 'data': pref_md_string})}\n\n"
+                    # Use first readable cube ID for scheduler (backward compatibility)
+                    scheduler_cube_id = (
+                        readable_cube_ids[0] if readable_cube_ids else chat_req.user_id
+                    )
+                    self._send_message_to_scheduler(
+                        user_id=chat_req.user_id,
+                        mem_cube_id=scheduler_cube_id,
+                        query=chat_req.query,
+                        label=QUERY_TASK_LABEL,
+                    )
 
                     # parse goal for internet search
                     searcher = self.dependencies.searcher
@@ -481,23 +476,28 @@ def generate_chat_response() -> Generator[str, None, None]:
                         # internet status
                         yield f"data: {json.dumps({'type': 'status', 'data': 'start_internet_search'})}\n\n"
 
-                    # ======  internet search with parse goal ======
+                    # ======  second deep search  ======
                     search_req = APISearchPlaygroundRequest(
                         query=parsed_goal.rephrased_query
                         or chat_req.query + (f"{parsed_goal.tags}" if parsed_goal.tags else ""),
                         user_id=chat_req.user_id,
                         readable_cube_ids=readable_cube_ids,
-                        mode=chat_req.mode,
-                        internet_search=chat_req.internet_search,
+                        mode="fast",
+                        internet_search=chat_req.internet_search or parsed_goal.internet_search,
                         top_k=chat_req.top_k,
                         chat_history=chat_req.history,
                         session_id=chat_req.session_id,
-                        include_preference=False,
+                        include_preference=chat_req.include_preference,
+                        pref_top_k=chat_req.pref_top_k,
                         filter=chat_req.filter,
                         search_memory_type="All",
+                        search_tool_memory=False,
                         playground_search_goal_parser=False,
                     )
+                    start_time = time.time()
                     search_response = self.search_handler.handle_search_memories(search_req)
+                    end_time = time.time()
+                    self.logger.info(f"second search time: {end_time - start_time}")
 
                     # Extract memories from search results (second search)
                     memories_list = []
@@ -516,12 +516,19 @@ def generate_chat_response() -> Generator[str, None, None]:
 
                     # Prepare remain reference data (second search)
                     reference = prepare_reference_data(filtered_memories)
+                    # get preference string
+                    pref_string = search_response.data.get("pref_string", "")
                     # get internet reference
                     internet_reference = self._get_internet_reference(
                         search_response.data.get("text_mem")[0]["memories"]
                     )
-
                     yield f"data: {json.dumps({'type': 'reference', 'data': reference})}\n\n"
+                    # Prepare preference markdown string
+                    if chat_req.include_preference:
+                        pref_list = search_response.data.get("pref_mem") or []
+                        pref_memories = pref_list[0].get("memories", []) if pref_list else []
+                        pref_md_string = self._build_pref_md_string_for_playground(pref_memories)
+                        yield f"data: {json.dumps({'type': 'pref_md_string', 'data': pref_md_string})}\n\n"
 
                     # Step 2: Build system prompt with memories
                     system_prompt = self._build_enhance_system_prompt(

src/memos/graph_dbs/polardb.py

@@ -3348,58 +3348,120 @@ def add_nodes_batch(
             with conn.cursor() as cursor:
                 # Process each group separately
                 for embedding_column, nodes_group in nodes_by_embedding_column.items():
-                    # Delete existing records first (batch delete)
-                    for node in nodes_group:
+                    # Batch delete existing records using IN clause
+                    ids_to_delete = [node["id"] for node in nodes_group]
+                    if ids_to_delete:
                         delete_query = f"""
                             DELETE FROM {self.db_name}_graph."Memory"
-                            WHERE id = ag_catalog._make_graph_id('{self.db_name}_graph'::name, 'Memory'::name, %s::text::cstring)
+                            WHERE id IN (
+                                SELECT ag_catalog._make_graph_id('{self.db_name}_graph'::name, 'Memory'::name, unnest(%s::text[])::cstring)
+                            )
                         """
-                        cursor.execute(delete_query, (node["id"],))
+                        cursor.execute(delete_query, (ids_to_delete,))
+
+                    # Batch get graph_ids for all nodes
+                    get_graph_ids_query = f"""
+                        SELECT
+                            id_val,
+                            ag_catalog._make_graph_id('{self.db_name}_graph'::name, 'Memory'::name, id_val::text::cstring) as graph_id
+                        FROM unnest(%s::text[]) as id_val
+                    """
+                    cursor.execute(get_graph_ids_query, (ids_to_delete,))
+                    graph_id_map = {row[0]: row[1] for row in cursor.fetchall()}
 
-                    # Insert nodes (batch insert using executemany for better performance)
+                    # Add graph_id to properties
                     for node in nodes_group:
-                        # Get graph_id for this node
-                        get_graph_id_query = f"""
-                            SELECT ag_catalog._make_graph_id('{self.db_name}_graph'::name, 'Memory'::name, %s::text::cstring)
-                        """
-                        cursor.execute(get_graph_id_query, (node["id"],))
-                        graph_id = cursor.fetchone()[0]
-                        node["properties"]["graph_id"] = str(graph_id)
-
-                        # Insert node
-                        if node["embedding_vector"]:
-                            insert_query = f"""
-                                INSERT INTO {self.db_name}_graph."Memory"(id, properties, {embedding_column})
-                                VALUES (
-                                    ag_catalog._make_graph_id('{self.db_name}_graph'::name, 'Memory'::name, %s::text::cstring),
-                                    %s,
-                                    %s
+                        graph_id = graph_id_map.get(node["id"])
+                        if graph_id:
+                            node["properties"]["graph_id"] = str(graph_id)
+
+                    # Batch insert using VALUES with multiple rows
+                    # Use psycopg2.extras.execute_values for efficient batch insert
+                    from psycopg2.extras import execute_values
+
+                    if embedding_column and any(node["embedding_vector"] for node in nodes_group):
+                        # Prepare data tuples for batch insert with embedding
+                        data_tuples = []
+                        for node in nodes_group:
+                            # Each tuple: (id, properties_json, embedding_json)
+                            data_tuples.append(
+                                (
+                                    node["id"],
+                                    json.dumps(node["properties"]),
+                                    json.dumps(node["embedding_vector"])
+                                    if node["embedding_vector"]
+                                    else None,
                                 )
-                            """
-                            logger.info(
-                                f"[add_nodes_batch] Inserting node insert_query={insert_query}"
                             )
-                            cursor.execute(
-                                insert_query,
+
+                        # Build the INSERT query template
+                        insert_query = f"""
+                            INSERT INTO {self.db_name}_graph."Memory"(id, properties, {embedding_column})
+                            VALUES %s
+                        """
+
+                        # Build the VALUES template for execute_values
+                        # Each row: (graph_id_function, agtype, vector)
+                        # Note: properties column is agtype, not jsonb
+                        template = f"""
+                            (
+                                ag_catalog._make_graph_id('{self.db_name}_graph'::name, 'Memory'::name, %s::text::cstring),
+                                %s::text::agtype,
+                                %s::vector
+                            )
+                        """
+                        logger.info(
+                            f"[add_nodes_batch] embedding_column Inserting insert_query:{insert_query}"
+                        )
+                        logger.info(
+                            f"[add_nodes_batch] embedding_column Inserting data_tuples:{data_tuples}"
+                        )
+
+                        # Execute batch insert
+                        execute_values(
+                            cursor,
+                            insert_query,
+                            data_tuples,
+                            template=template,
+                            page_size=100,  # Insert in batches of 100
+                        )
+                    else:
+                        # Prepare data tuples for batch insert without embedding
+                        data_tuples = []
+                        for node in nodes_group:
+                            # Each tuple: (id, properties_json)
+                            data_tuples.append(
                                 (
                                     node["id"],
                                     json.dumps(node["properties"]),
-                                    json.dumps(node["embedding_vector"]),
-                                ),
-                            )
-                        else:
-                            insert_query = f"""
-                                INSERT INTO {self.db_name}_graph."Memory"(id, properties)
-                                VALUES (
-                                    ag_catalog._make_graph_id('{self.db_name}_graph'::name, 'Memory'::name, %s::text::cstring),
-                                    %s
                                 )
-                            """
-                            cursor.execute(
-                                insert_query,
-                                (node["id"], json.dumps(node["properties"])),
                             )
 
+                        # Build the INSERT query template
+                        insert_query = f"""
+                            INSERT INTO {self.db_name}_graph."Memory"(id, properties)
+                            VALUES %s
+                        """
+
+                        # Build the VALUES template for execute_values
+                        # Note: properties column is agtype, not jsonb
+                        template = f"""
+                            (
+                                ag_catalog._make_graph_id('{self.db_name}_graph'::name, 'Memory'::name, %s::text::cstring),
+                                %s::text::agtype
+                            )
+                        """
+                        logger.info(f"[add_nodes_batch] Inserting insert_query:{insert_query}")
+                        logger.info(f"[add_nodes_batch] Inserting data_tuples:{data_tuples}")
+                        # Execute batch insert
+                        execute_values(
+                            cursor,
+                            insert_query,
+                            data_tuples,
+                            template=template,
+                            page_size=100,  # Insert in batches of 100
+                        )
+
                     logger.info(
                         f"[add_nodes_batch] Inserted {len(nodes_group)} nodes with embedding_column={embedding_column}"
                     )

src/memos/mem_reader/read_multi_modal/system_parser.py

@@ -1,5 +1,6 @@
 """Parser for system messages."""
 
+import ast
 import json
 import re
 import uuid
@@ -137,8 +138,14 @@ def parse_fine(
             tool_schema = json.loads(content)
             assert isinstance(tool_schema, list), "Tool schema must be a list[dict]"
         except json.JSONDecodeError:
-            logger.warning(f"[SystemParser] Failed to parse tool schema: {content}")
-            return []
+            try:
+                tool_schema = ast.literal_eval(content)
+                assert isinstance(tool_schema, list), "Tool schema must be a list[dict]"
+            except (ValueError, SyntaxError, AssertionError):
+                logger.warning(
+                    f"[SystemParser] Failed to parse tool schema with both JSON and ast.literal_eval: {content}"
+                )
+                return []
         except AssertionError:
             logger.warning(f"[SystemParser] Tool schema must be a list[dict]: {content}")
             return []

src/memos/memories/textual/tree.py

@@ -343,6 +343,17 @@ def delete_all(self) -> None:
             logger.error(f"An error occurred while deleting all memories: {e}")
             raise
 
+    def delete_by_filter(
+        self,
+        writable_cube_ids: list[str],
+        file_ids: list[str] | None = None,
+        filter: dict | None = None,
+    ) -> None:
+        """Delete memories by filter."""
+        self.graph_store.delete_node_by_prams(
+            writable_cube_ids=writable_cube_ids, file_ids=file_ids, filter=filter
+        )
+
     def load(self, dir: str) -> None:
         try:
             memory_file = os.path.join(dir, self.config.memory_filename)

src/memos/memories/textual/tree_text_memory/retrieve/searcher.py

@@ -701,15 +701,35 @@ def _sort_and_trim(
         """Sort results by score and trim to top_k"""
         final_items = []
         if search_tool_memory:
-            tool_results = [
+            tool_schema_results = [
                 (item, score)
                 for item, score in results
-                if item.metadata.memory_type in ["ToolSchemaMemory", "ToolTrajectoryMemory"]
+                if item.metadata.memory_type == "ToolSchemaMemory"
             ]
-            sorted_tool_results = sorted(tool_results, key=lambda pair: pair[1], reverse=True)[
-                :tool_mem_top_k
+            sorted_tool_schema_results = sorted(
+                tool_schema_results, key=lambda pair: pair[1], reverse=True
+            )[:tool_mem_top_k]
+            for item, score in sorted_tool_schema_results:
+                if plugin and round(score, 2) == 0.00:
+                    continue
+                meta_data = item.metadata.model_dump()
+                meta_data["relativity"] = score
+                final_items.append(
+                    TextualMemoryItem(
+                        id=item.id,
+                        memory=item.memory,
+                        metadata=SearchedTreeNodeTextualMemoryMetadata(**meta_data),
+                    )
+                )
+            tool_trajectory_results = [
+                (item, score)
+                for item, score in results
+                if item.metadata.memory_type == "ToolTrajectoryMemory"
             ]
-            for item, score in sorted_tool_results:
+            sorted_tool_trajectory_results = sorted(
+                tool_trajectory_results, key=lambda pair: pair[1], reverse=True
+            )[:tool_mem_top_k]
+            for item, score in sorted_tool_trajectory_results:
                 if plugin and round(score, 2) == 0.00:
                     continue
                 meta_data = item.metadata.model_dump()

src/memos/multi_mem_cube/single_cube.py

@@ -30,6 +30,7 @@
     SearchMode,
     UserContext,
 )
+from memos.utils import timed
 
 
 logger = get_logger(__name__)
@@ -198,6 +199,7 @@ def _get_search_mode(self, mode: str) -> str:
         """
         return mode
 
+    @timed
     def _search_text(
         self,
         search_req: APISearchRequest,
@@ -363,6 +365,7 @@ def _fine_search(
 
         return formatted_memories
 
+    @timed
     def _search_pref(
         self,
         search_req: APISearchRequest,
@@ -429,7 +432,7 @@ def _fast_search(
             top_k=search_req.top_k,
             mode=SearchMode.FAST,
             manual_close_internet=not search_req.internet_search,
-            momory_type=search_req.search_memory_type,
+            memory_type=search_req.search_memory_type,
             search_filter=search_filter,
             search_priority=search_priority,
             info={