Fix dedup handling in simple search

Author: [email protected]

src/memos/api/handlers/search_handler.py

@@ -5,9 +5,14 @@
 using dependency injection for better modularity and testability.
 """
 
+from typing import Any
+
 from memos.api.handlers.base_handler import BaseHandler, HandlerDependencies
 from memos.api.product_models import APISearchRequest, SearchResponse
 from memos.log import get_logger
+from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import (
+    cosine_similarity_matrix,
+)
 from memos.multi_mem_cube.composite_cube import CompositeCubeView
 from memos.multi_mem_cube.single_cube import SingleCubeView
 from memos.multi_mem_cube.views import MemCubeView
@@ -53,6 +58,8 @@ def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse
         cube_view = self._build_cube_view(search_req)
 
         results = cube_view.search_memories(search_req)
+        if search_req.dedup == "sim":
+            results = self._dedup_text_memories(results, search_req.top_k)
 
         self.logger.info(
             f"[SearchHandler] Final search results: count={len(results)} results={results}"
@@ -63,6 +70,48 @@ def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse
             data=results,
         )
 
+    def _dedup_text_memories(self, results: dict[str, Any], target_top_k: int) -> dict[str, Any]:
+        for bucket in results.get("text_mem", []):
+            memories = bucket.get("memories", [])
+            if len(memories) <= 1:
+                continue
+            embeddings = self._extract_embeddings(memories)
+            if embeddings is None:
+                documents = [mem.get("memory", "") for mem in memories]
+                embeddings = self.searcher.embedder.embed(documents)
+            similarity_matrix = cosine_similarity_matrix(embeddings)
+            selected_indices = self._select_unrelated_indices(similarity_matrix, 0.85)
+            if len(selected_indices) < min(target_top_k, len(memories)):
+                selected = set(selected_indices)
+                for i in range(len(memories)):
+                    if i in selected:
+                        continue
+                    selected_indices.append(i)
+                    if len(selected_indices) >= target_top_k:
+                        break
+            bucket["memories"] = [memories[i] for i in selected_indices[:target_top_k]]
+        return results
+
+    @staticmethod
+    def _select_unrelated_indices(
+        similarity_matrix: list[list[float]], similarity_threshold: float
+    ) -> list[int]:
+        selected_indices: list[int] = []
+        for i in range(len(similarity_matrix)):
+            if all(similarity_matrix[i][j] <= similarity_threshold for j in selected_indices):
+                selected_indices.append(i)
+        return selected_indices
+
+    @staticmethod
+    def _extract_embeddings(memories: list[dict[str, Any]]) -> list[list[float]] | None:
+        embeddings: list[list[float]] = []
+        for mem in memories:
+            embedding = mem.get("metadata", {}).get("embedding")
+            if not embedding:
+                return None
+            embeddings.append(embedding)
+        return embeddings
+
     def _resolve_cube_ids(self, search_req: APISearchRequest) -> list[str]:
         """
         Normalize target cube ids from search_req.

src/memos/api/start_api.py

@@ -1,7 +1,7 @@
 import logging
 import os
 
-from typing import Any, Generic, Literal, TypeVar
+from typing import Any, Generic, TypeVar
 
 from dotenv import load_dotenv
 from fastapi import FastAPI
@@ -145,14 +145,6 @@ class SearchRequest(BaseRequest):
         description="List of cube IDs to search in",
         json_schema_extra={"example": ["cube123", "cube456"]},
     )
-    dedup: Literal["no", "sim"] | None = Field(
-        None,
-        description=(
-            "Optional dedup option for textual memories. "
-            "Use 'no' for no dedup, 'sim' for similarity dedup. "
-            "If None, default exact-text dedup is applied."
-        ),
-    )
 
 
 class MemCubeRegister(BaseRequest):
@@ -357,7 +349,6 @@ async def search_memories(search_req: SearchRequest):
         query=search_req.query,
         user_id=search_req.user_id,
         install_cube_ids=search_req.install_cube_ids,
-        dedup=search_req.dedup,
     )
     return SearchResponse(message="Search completed successfully", data=result)
 

src/memos/mem_os/core.py

@@ -551,7 +551,6 @@ def search(
         internet_search: bool = False,
         moscube: bool = False,
         session_id: str | None = None,
-        dedup: str | None = None,
         **kwargs,
     ) -> MOSSearchResult:
         """
@@ -626,7 +625,6 @@ def search_textual_memory(cube_id, cube):
                         },
                         moscube=moscube,
                         search_filter=search_filter,
-                        dedup=dedup,
                     )
                     search_time_end = time.time()
                     logger.info(

src/memos/mem_scheduler/optimized_scheduler.py

@@ -186,7 +186,6 @@ def mix_search_memories(
             info=info,
             search_tool_memory=search_req.search_tool_memory,
             tool_mem_top_k=search_req.tool_mem_top_k,
-            dedup=search_req.dedup,
         )
         memories = merged_memories[: search_req.top_k]
 

src/memos/memories/textual/tree.py

@@ -161,7 +161,6 @@ def search(
         user_name: str | None = None,
         search_tool_memory: bool = False,
         tool_mem_top_k: int = 6,
-        dedup: str | None = None,
         **kwargs,
     ) -> list[TextualMemoryItem]:
         """Search for memories based on a query.
@@ -208,7 +207,6 @@ def search(
             user_name=user_name,
             search_tool_memory=search_tool_memory,
             tool_mem_top_k=tool_mem_top_k,
-            dedup=dedup,
             **kwargs,
         )
 

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

@@ -239,7 +239,6 @@ def deep_search(
         user_name: str | None = None,
         **kwargs,
     ):
-        dedup = kwargs.get("dedup")
         previous_retrieval_phrases = [query]
         retrieved_memories = self.retrieve(
             query=query,
@@ -255,7 +254,6 @@ def deep_search(
             top_k=top_k,
             user_name=user_name,
             info=info,
-            dedup=dedup,
         )
         if len(memories) == 0:
             logger.warning("Requirements not met; returning memories as-is.")

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

@@ -119,15 +119,9 @@ def post_retrieve(
         info=None,
         search_tool_memory: bool = False,
         tool_mem_top_k: int = 6,
-        dedup: str | None = None,
         plugin=False,
     ):
-        if dedup == "no":
-            deduped = retrieved_results
-        elif dedup == "sim":
-            deduped = self._deduplicate_similar_results(retrieved_results)
-        else:
-            deduped = self._deduplicate_results(retrieved_results)
+        deduped = self._deduplicate_results(retrieved_results)
         final_results = self._sort_and_trim(
             deduped, top_k, plugin, search_tool_memory, tool_mem_top_k
         )
@@ -147,7 +141,6 @@ def search(
         user_name: str | None = None,
         search_tool_memory: bool = False,
         tool_mem_top_k: int = 6,
-        dedup: str | None = None,
         **kwargs,
     ) -> list[TextualMemoryItem]:
         """
@@ -180,11 +173,7 @@ def search(
         if kwargs.get("plugin", False):
             logger.info(f"[SEARCH] Retrieve from plugin: {query}")
             retrieved_results = self._retrieve_simple(
-                query=query,
-                top_k=top_k,
-                search_filter=search_filter,
-                user_name=user_name,
-                dedup=dedup,
+                query=query, top_k=top_k, search_filter=search_filter, user_name=user_name
             )
         else:
             retrieved_results = self.retrieve(
@@ -213,7 +202,6 @@ def search(
             plugin=kwargs.get("plugin", False),
             search_tool_memory=search_tool_memory,
             tool_mem_top_k=tool_mem_top_k,
-            dedup=None if kwargs.get("plugin", False) and dedup == "sim" else dedup,
         )
 
         logger.info(f"[SEARCH] Done. Total {len(final_results)} results.")
@@ -296,50 +284,6 @@ def _parse_task(
 
         return parsed_goal, query_embedding, context, query
 
-    @timed
-    def _retrieve_simple(
-        self,
-        query: str,
-        top_k: int,
-        search_filter: dict | None = None,
-        user_name: str | None = None,
-        dedup: str | None = None,
-        **kwargs,
-    ):
-        """Retrieve from by keywords and embedding"""
-        query_words = []
-        if self.tokenizer:
-            query_words = self.tokenizer.tokenize_mixed(query)
-        else:
-            query_words = query.strip().split()
-        query_words = [query, *query_words]
-        logger.info(f"[SIMPLESEARCH] Query words: {query_words}")
-        query_embeddings = self.embedder.embed(query_words)
-
-        items = self.graph_retriever.retrieve_from_mixed(
-            top_k=top_k * 2,
-            memory_scope=None,
-            query_embedding=query_embeddings,
-            search_filter=search_filter,
-            user_name=user_name,
-            use_fast_graph=self.use_fast_graph,
-        )
-        logger.info(f"[SIMPLESEARCH] Items count: {len(items)}")
-        documents = [getattr(item, "memory", "") for item in items]
-        documents_embeddings = self.embedder.embed(documents)
-        similarity_matrix = cosine_similarity_matrix(documents_embeddings)
-        selected_indices, _ = find_best_unrelated_subgroup(documents, similarity_matrix)
-        selected_items = [items[i] for i in selected_indices]
-        logger.info(
-            f"[SIMPLESEARCH] after unrelated subgroup selection items count: {len(selected_items)}"
-        )
-        return self.reranker.rerank(
-            query=query,
-            query_embedding=query_embeddings[0],
-            graph_results=selected_items,
-            top_k=top_k,
-        )
-
     @timed
     def _retrieve_paths(
         self,
@@ -723,17 +667,14 @@ def _retrieve_simple(
             user_name=user_name,
         )
         logger.info(f"[SIMPLESEARCH] Items count: {len(items)}")
-        if dedup == "no":
-            selected_items = items
-        else:
-            documents = [getattr(item, "memory", "") for item in items]
-            documents_embeddings = self.embedder.embed(documents)
-            similarity_matrix = cosine_similarity_matrix(documents_embeddings)
-            selected_indices, _ = find_best_unrelated_subgroup(documents, similarity_matrix)
-            selected_items = [items[i] for i in selected_indices]
-            logger.info(
-                f"[SIMPLESEARCH] after unrelated subgroup selection items count: {len(selected_items)}"
-            )
+        documents = [getattr(item, "memory", "") for item in items]
+        documents_embeddings = self.embedder.embed(documents)
+        similarity_matrix = cosine_similarity_matrix(documents_embeddings)
+        selected_indices, _ = find_best_unrelated_subgroup(documents, similarity_matrix)
+        selected_items = [items[i] for i in selected_indices]
+        logger.info(
+            f"[SIMPLESEARCH] after unrelated subgroup selection items count: {len(selected_items)}"
+        )
         return self.reranker.rerank(
             query=query,
             query_embedding=query_embeddings[0],
@@ -750,26 +691,6 @@ def _deduplicate_results(self, results):
                 deduped[item.memory] = (item, score)
         return list(deduped.values())
 
-    @timed
-    def _deduplicate_similar_results(
-        self, results: list[tuple[TextualMemoryItem, float]], similarity_threshold: float = 0.85
-    ):
-        """Deduplicate results by semantic similarity while keeping higher scores."""
-        if len(results) <= 1:
-            return results
-
-        sorted_results = sorted(results, key=lambda pair: pair[1], reverse=True)
-        documents = [getattr(item, "memory", "") for item, _ in sorted_results]
-        embeddings = self.embedder.embed(documents)
-        similarity_matrix = cosine_similarity_matrix(embeddings)
-
-        selected_indices: list[int] = []
-        for i in range(len(sorted_results)):
-            if all(similarity_matrix[i][j] <= similarity_threshold for j in selected_indices):
-                selected_indices.append(i)
-
-        return [sorted_results[i] for i in selected_indices]
-
     @timed
     def _sort_and_trim(
         self, results, top_k, plugin=False, search_tool_memory=False, tool_mem_top_k=6

src/memos/multi_mem_cube/single_cube.py

@@ -23,9 +23,6 @@
     MEM_READ_TASK_LABEL,
     PREF_ADD_TASK_LABEL,
 )
-from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import (
-    cosine_similarity_matrix,
-)
 from memos.multi_mem_cube.views import MemCubeView
 from memos.templates.mem_reader_prompts import PROMPT_MAPPING
 from memos.types.general_types import (
@@ -266,7 +263,6 @@ def _deep_search(
             moscube=search_req.moscube,
             search_filter=search_filter,
             info=info,
-            dedup=search_req.dedup,
         )
         formatted_memories = [format_memory_item(data) for data in enhanced_memories]
         return formatted_memories
@@ -332,7 +328,6 @@ def _fine_search(
             top_k=search_req.top_k,
             user_name=user_context.mem_cube_id,
             info=info,
-            dedup=search_req.dedup,
         )
 
         # Enhance with query
@@ -383,24 +378,7 @@ def _dedup_by_content(memories: list) -> list:
                 unique_memories.append(mem)
             return unique_memories
 
-        def _dedup_by_similarity(memories: list) -> list:
-            if len(memories) <= 1:
-                return memories
-            documents = [getattr(mem, "memory", "") for mem in memories]
-            embeddings = self.searcher.embedder.embed(documents)
-            similarity_matrix = cosine_similarity_matrix(embeddings)
-            selected_indices = []
-            for i in range(len(memories)):
-                if all(similarity_matrix[i][j] <= 0.85 for j in selected_indices):
-                    selected_indices.append(i)
-            return [memories[i] for i in selected_indices]
-
-        if search_req.dedup == "no":
-            deduped_memories = enhanced_memories
-        elif search_req.dedup == "sim":
-            deduped_memories = _dedup_by_similarity(enhanced_memories)
-        else:
-            deduped_memories = _dedup_by_content(enhanced_memories)
+        deduped_memories = _dedup_by_content(enhanced_memories)
         formatted_memories = [format_memory_item(data) for data in deduped_memories]
 
         logger.info(f"Found {len(formatted_memories)} memories for user {search_req.user_id}")
@@ -485,7 +463,6 @@ def _fast_search(
             plugin=plugin,
             search_tool_memory=search_req.search_tool_memory,
             tool_mem_top_k=search_req.tool_mem_top_k,
-            dedup=search_req.dedup,
         )
 
         formatted_memories = [format_memory_item(data) for data in search_results]