Merge branch 'dev' into feat/dedup-search-param

Author: glin93

src/memos/api/handlers/memory_handler.py

@@ -23,6 +23,10 @@
     remove_embedding_recursive,
     sort_children_by_memory_type,
 )
+from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import (
+    cosine_similarity_matrix,
+    find_best_unrelated_subgroup,
+)
 
 
 if TYPE_CHECKING:
@@ -37,6 +41,7 @@ def handle_get_all_memories(
     mem_cube_id: str,
     memory_type: Literal["text_mem", "act_mem", "param_mem", "para_mem"],
     naive_mem_cube: Any,
+    embedder: Any,
 ) -> MemoryResponse:
     """
     Main handler for getting all memories.
@@ -59,6 +64,14 @@ def handle_get_all_memories(
             # Get all text memories from the graph database
             memories = naive_mem_cube.text_mem.get_all(user_name=mem_cube_id)
 
+            mems = [mem.get("memory", "") for mem in memories.get("nodes", [])]
+            embeddings = embedder.embed(mems)
+            similarity_matrix = cosine_similarity_matrix(embeddings)
+            selected_indices, _ = find_best_unrelated_subgroup(
+                embeddings, similarity_matrix, bar=0.9
+            )
+            memories["nodes"] = [memories["nodes"][i] for i in selected_indices]
+
             # Format and convert to tree structure
             memories_cleaned = remove_embedding_recursive(memories)
             custom_type_ratios = {

src/memos/api/product_models.py

@@ -1177,3 +1177,16 @@ class AllStatusResponse(BaseResponse[AllStatusResponseData]):
     """Response model for full scheduler status operations."""
 
     message: str = "Scheduler status summary retrieved successfully"
+
+
+# ─── Internal API Endpoints Models (for internal use) ───────────────────────────────────────────────────
+
+
+class GetUserNamesByMemoryIdsRequest(BaseRequest):
+    """Request model for getting user names by memory ids."""
+
+    memory_ids: list[str] = Field(..., description="Memory IDs")
+
+
+class GetUserNamesByMemoryIdsResponse(BaseResponse[dict[str, list[str]]]):
+    """Response model for getting user names by memory ids."""

src/memos/api/routers/server_router.py

@@ -36,13 +36,16 @@
     GetMemoryPlaygroundRequest,
     GetMemoryRequest,
     GetMemoryResponse,
+    GetUserNamesByMemoryIdsRequest,
+    GetUserNamesByMemoryIdsResponse,
     MemoryResponse,
     SearchResponse,
     StatusResponse,
     SuggestionRequest,
     SuggestionResponse,
     TaskQueueResponse,
 )
+from memos.graph_dbs.polardb import PolarDBGraphDB
 from memos.log import get_logger
 from memos.mem_scheduler.base_scheduler import BaseScheduler
 from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
@@ -83,6 +86,8 @@
 naive_mem_cube = components["naive_mem_cube"]
 redis_client = components["redis_client"]
 status_tracker = TaskStatusTracker(redis_client=redis_client)
+embedder = components["embedder"]
+graph_db = components["graph_db"]
 
 
 # =============================================================================
@@ -294,6 +299,7 @@ def get_all_memories(memory_req: GetMemoryPlaygroundRequest):
             ),
             memory_type=memory_req.memory_type or "text_mem",
             naive_mem_cube=naive_mem_cube,
+            embedder=embedder,
         )
 
 
@@ -327,3 +333,27 @@ def feedback_memories(feedback_req: APIFeedbackRequest):
     This endpoint uses the class-based FeedbackHandler for better code organization.
     """
     return feedback_handler.handle_feedback_memories(feedback_req)
+
+
+# =============================================================================
+# Other API Endpoints (for internal use)
+# =============================================================================
+
+
+@router.get(
+    "/get_user_names_by_memory_ids",
+    summary="Get user names by memory ids",
+    response_model=GetUserNamesByMemoryIdsResponse,
+)
+def get_user_names_by_memory_ids(memory_ids: GetUserNamesByMemoryIdsRequest):
+    """Get user names by memory ids."""
+    if not isinstance(graph_db, PolarDBGraphDB):
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "graph_db must be an instance of PolarDBGraphDB to use "
+                "get_user_names_by_memory_ids"
+                f"current graph_db is: {graph_db.__class__.__name__}"
+            ),
+        )
+    return graph_db.get_user_names_by_memory_ids(memory_ids=memory_ids)

src/memos/graph_dbs/polardb.py

@@ -4869,6 +4869,7 @@ def delete_node_by_prams(
         memory_ids: list[str] | None = None,
         file_ids: list[str] | None = None,
         filter: dict | None = None,
+        batch_size: int = 100,
     ) -> int:
         """
         Delete nodes by memory_ids, file_ids, or filter.
@@ -4898,31 +4899,6 @@ def delete_node_by_prams(
                     f"agtype_access_operator(VARIADIC ARRAY[properties, '\"user_name\"'::agtype]) = '\"{cube_id}\"'::agtype"
                 )
 
-        # Build WHERE conditions separately for memory_ids and file_ids
-        where_conditions = []
-
-        # Handle memory_ids: query properties.id
-        if memory_ids and len(memory_ids) > 0:
-            memory_id_conditions = []
-            for node_id in memory_ids:
-                memory_id_conditions.append(
-                    f"ag_catalog.agtype_access_operator(properties, '\"id\"'::agtype) = '\"{node_id}\"'::agtype"
-                )
-            if memory_id_conditions:
-                where_conditions.append(f"({' OR '.join(memory_id_conditions)})")
-
-        # Check if any file_id is in the file_ids array field (OR relationship)
-        if file_ids and len(file_ids) > 0:
-            file_id_conditions = []
-            for file_id in file_ids:
-                # Format: agtype_in_operator(agtype_access_operator(VARIADIC ARRAY[properties, '"file_ids"'::agtype]), '"file_id"'::agtype)
-                file_id_conditions.append(
-                    f"agtype_in_operator(agtype_access_operator(VARIADIC ARRAY[properties, '\"file_ids\"'::agtype]), '\"{file_id}\"'::agtype)"
-                )
-            if file_id_conditions:
-                # Use OR to match any file_id in the array
-                where_conditions.append(f"({' OR '.join(file_id_conditions)})")
-
         # Query nodes by filter if provided
         filter_ids = set()
         if filter:
@@ -4943,86 +4919,165 @@ def delete_node_by_prams(
                     "[delete_node_by_prams] Filter parsed to None, skipping filter query"
                 )
 
-        # If filter returned IDs, add condition for them
+        # Combine all IDs that need to be deleted
+        all_memory_ids = set()
+        if memory_ids:
+            all_memory_ids.update(memory_ids)
         if filter_ids:
-            filter_id_conditions = []
-            for node_id in filter_ids:
-                filter_id_conditions.append(
-                    f"ag_catalog.agtype_access_operator(properties, '\"id\"'::agtype) = '\"{node_id}\"'::agtype"
-                )
-            if filter_id_conditions:
-                where_conditions.append(f"({' OR '.join(filter_id_conditions)})")
+            all_memory_ids.update(filter_ids)
 
-        # If no conditions (except user_name), return 0
-        if not where_conditions:
+        # If no conditions to delete, return 0
+        if not all_memory_ids and not file_ids:
             logger.warning(
                 "[delete_node_by_prams] No nodes to delete (no memory_ids, file_ids, or filter provided)"
             )
             return 0
 
-        # Build WHERE clause
-        # First, combine memory_ids, file_ids, and filter conditions with OR (any condition can match)
-        data_conditions = " OR ".join([f"({cond})" for cond in where_conditions])
+        conn = None
+        total_deleted_count = 0
+        try:
+            conn = self._get_connection()
+            with conn.cursor() as cursor:
+                # Process memory_ids and filter_ids in batches
+                if all_memory_ids:
+                    memory_ids_list = list(all_memory_ids)
+                    total_batches = (len(memory_ids_list) + batch_size - 1) // batch_size
+                    logger.info(
+                        f"[delete_node_by_prams] memoryids Processing {len(memory_ids_list)} memory_ids in {total_batches} batches (batch_size={batch_size})"
+                    )
 
-        # Build final WHERE clause
-        # If user_name_conditions exist, combine with data_conditions using AND
-        # Otherwise, use only data_conditions
-        if user_name_conditions:
-            user_name_where = " OR ".join(user_name_conditions)
-            where_clause = f"({user_name_where}) AND ({data_conditions})"
-        else:
-            where_clause = f"({data_conditions})"
+                    for batch_idx in range(total_batches):
+                        batch_start = batch_idx * batch_size
+                        batch_end = min(batch_start + batch_size, len(memory_ids_list))
+                        batch_ids = memory_ids_list[batch_start:batch_end]
 
-        # Use SQL DELETE query for better performance
-        # First count matching nodes to get accurate count
-        count_query = f"""
-                SELECT COUNT(*)
-                FROM "{self.db_name}_graph"."Memory"
-                WHERE {where_clause}
-            """
-        logger.info(f"[delete_node_by_prams] count_query: {count_query}")
+                        # Build conditions for this batch
+                        batch_conditions = []
+                        for node_id in batch_ids:
+                            batch_conditions.append(
+                                f"ag_catalog.agtype_access_operator(properties, '\"id\"'::agtype) = '\"{node_id}\"'::agtype"
+                            )
+                        batch_where = f"({' OR '.join(batch_conditions)})"
 
-        # Then delete nodes
-        delete_query = f"""
-                DELETE FROM "{self.db_name}_graph"."Memory"
-                WHERE {where_clause}
-            """
+                        # Add user_name filter if provided
+                        if user_name_conditions:
+                            user_name_where = " OR ".join(user_name_conditions)
+                            where_clause = f"({user_name_where}) AND ({batch_where})"
+                        else:
+                            where_clause = batch_where
 
-        logger.info(
-            f"[delete_node_by_prams] Deleting nodes - memory_ids: {memory_ids}, file_ids: {file_ids}, filter: {filter}"
-        )
-        logger.info(f"[delete_node_by_prams] delete_query: {delete_query}")
+                        # Count before deletion
+                        count_query = f"""
+                            SELECT COUNT(*)
+                            FROM "{self.db_name}_graph"."Memory"
+                            WHERE {where_clause}
+                        """
+                        logger.info(
+                            f"[delete_node_by_prams] memoryids batch {batch_idx + 1}/{total_batches}: count_query: {count_query}"
+                        )
 
-        conn = None
-        deleted_count = 0
-        try:
-            conn = self._get_connection()
-            with conn.cursor() as cursor:
-                # Count nodes before deletion
-                cursor.execute(count_query)
-                count_result = cursor.fetchone()
-                expected_count = count_result[0] if count_result else 0
+                        cursor.execute(count_query)
+                        count_result = cursor.fetchone()
+                        expected_count = count_result[0] if count_result else 0
 
-                logger.info(
-                    f"[delete_node_by_prams] Found {expected_count} nodes matching the criteria"
-                )
+                        if expected_count == 0:
+                            logger.info(
+                                f"[delete_node_by_prams] memoryids Batch {batch_idx + 1}/{total_batches}: No nodes found, skipping"
+                            )
+                            continue
+
+                        # Delete batch
+                        delete_query = f"""
+                            DELETE FROM "{self.db_name}_graph"."Memory"
+                            WHERE {where_clause}
+                        """
+                        logger.info(
+                            f"[delete_node_by_prams] memoryids batch {batch_idx + 1}/{total_batches}: delete_query: {delete_query}"
+                        )
+
+                        logger.info(
+                            f"[delete_node_by_prams] memoryids Batch {batch_idx + 1}/{total_batches}: Executing delete query for {len(batch_ids)} nodes"
+                        )
+                        cursor.execute(delete_query)
+                        batch_deleted = cursor.rowcount
+                        total_deleted_count += batch_deleted
+
+                        logger.info(
+                            f"[delete_node_by_prams] memoryids Batch {batch_idx + 1}/{total_batches}: Deleted {batch_deleted} nodes (batch size: {len(batch_ids)})"
+                        )
+
+                # Process file_ids in batches
+                if file_ids:
+                    total_file_batches = (len(file_ids) + batch_size - 1) // batch_size
+                    logger.info(
+                        f"[delete_node_by_prams] Processing {len(file_ids)} file_ids in {total_file_batches} batches (batch_size={batch_size})"
+                    )
+
+                    for batch_idx in range(total_file_batches):
+                        batch_start = batch_idx * batch_size
+                        batch_end = min(batch_start + batch_size, len(file_ids))
+                        batch_file_ids = file_ids[batch_start:batch_end]
+
+                        # Build conditions for this batch
+                        batch_conditions = []
+                        for file_id in batch_file_ids:
+                            batch_conditions.append(
+                                f"agtype_in_operator(agtype_access_operator(VARIADIC ARRAY[properties, '\"file_ids\"'::agtype]), '\"{file_id}\"'::agtype)"
+                            )
+                        batch_where = f"({' OR '.join(batch_conditions)})"
+
+                        # Add user_name filter if provided
+                        if user_name_conditions:
+                            user_name_where = " OR ".join(user_name_conditions)
+                            where_clause = f"({user_name_where}) AND ({batch_where})"
+                        else:
+                            where_clause = batch_where
+
+                        # Count before deletion
+                        count_query = f"""
+                            SELECT COUNT(*)
+                            FROM "{self.db_name}_graph"."Memory"
+                            WHERE {where_clause}
+                        """
+
+                        logger.info(
+                            f"[delete_node_by_prams] File batch {batch_idx + 1}/{total_file_batches}: count_query: {count_query}"
+                        )
+                        cursor.execute(count_query)
+                        count_result = cursor.fetchone()
+                        expected_count = count_result[0] if count_result else 0
+
+                        if expected_count == 0:
+                            logger.info(
+                                f"[delete_node_by_prams] File batch {batch_idx + 1}/{total_file_batches}: No nodes found, skipping"
+                            )
+                            continue
+
+                        # Delete batch
+                        delete_query = f"""
+                            DELETE FROM "{self.db_name}_graph"."Memory"
+                            WHERE {where_clause}
+                        """
+                        cursor.execute(delete_query)
+                        batch_deleted = cursor.rowcount
+                        total_deleted_count += batch_deleted
+
+                        logger.info(
+                            f"[delete_node_by_prams] File batch {batch_idx + 1}/{total_file_batches}: delete_query: {delete_query}"
+                        )
 
-                # Delete nodes
-                cursor.execute(delete_query)
-                # Use rowcount to get actual deleted count
-                deleted_count = cursor.rowcount
                 elapsed_time = time.time() - batch_start_time
                 logger.info(
-                    f"[delete_node_by_prams] Deletion completed successfully in {elapsed_time:.2f}s, deleted {deleted_count} nodes"
+                    f"[delete_node_by_prams] Deletion completed successfully in {elapsed_time:.2f}s, total deleted {total_deleted_count} nodes"
                 )
         except Exception as e:
             logger.error(f"[delete_node_by_prams] Failed to delete nodes: {e}", exc_info=True)
             raise
         finally:
             self._return_connection(conn)
 
-        logger.info(f"[delete_node_by_prams] Successfully deleted {deleted_count} nodes")
-        return deleted_count
+        logger.info(f"[delete_node_by_prams] Successfully deleted {total_deleted_count} nodes")
+        return total_deleted_count
 
     @timed
     def get_user_names_by_memory_ids(self, memory_ids: list[str]) -> dict[str, list[str]]: