Add dedup option to search pipeline

Author: [email protected]

src/memos/api/product_models.py

@@ -319,6 +319,15 @@ class APISearchRequest(BaseRequest):
         description="Number of textual memories to retrieve (top-K). Default: 10.",
     )
 
+    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."
+        ),
+    )
+
     pref_top_k: int = Field(
         6,
         ge=0,

src/memos/api/start_api.py

@@ -1,7 +1,7 @@
 import logging
 import os
 
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Literal, TypeVar
 
 from dotenv import load_dotenv
 from fastapi import FastAPI
@@ -145,6 +145,14 @@ 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):
@@ -349,6 +357,7 @@ 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,6 +551,7 @@ def search(
         internet_search: bool = False,
         moscube: bool = False,
         session_id: str | None = None,
+        dedup: str | None = None,
         **kwargs,
     ) -> MOSSearchResult:
         """
@@ -625,6 +626,7 @@ 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,6 +186,7 @@ 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,6 +161,7 @@ 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.
@@ -207,6 +208,7 @@ 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,6 +239,7 @@ def deep_search(
         user_name: str | None = None,
         **kwargs,
     ):
+        dedup = kwargs.get("dedup")
         previous_retrieval_phrases = [query]
         retrieved_memories = self.retrieve(
             query=query,
@@ -254,6 +255,7 @@ 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,9 +119,15 @@ def post_retrieve(
         info=None,
         search_tool_memory: bool = False,
         tool_mem_top_k: int = 6,
+        dedup: str | None = None,
         plugin=False,
     ):
-        deduped = self._deduplicate_results(retrieved_results)
+        if dedup == "no":
+            deduped = retrieved_results
+        elif dedup == "sim":
+            deduped = self._deduplicate_similar_results(retrieved_results)
+        else:
+            deduped = self._deduplicate_results(retrieved_results)
         final_results = self._sort_and_trim(
             deduped, top_k, plugin, search_tool_memory, tool_mem_top_k
         )
@@ -141,6 +147,7 @@ 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]:
         """
@@ -173,7 +180,11 @@ 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
+                query=query,
+                top_k=top_k,
+                search_filter=search_filter,
+                user_name=user_name,
+                dedup=dedup,
             )
         else:
             retrieved_results = self.retrieve(
@@ -202,6 +213,7 @@ 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.")
@@ -291,6 +303,7 @@ def _retrieve_simple(
         top_k: int,
         search_filter: dict | None = None,
         user_name: str | None = None,
+        dedup: str | None = None,
         **kwargs,
     ):
         """Retrieve from by keywords and embedding"""
@@ -710,14 +723,17 @@ def _retrieve_simple(
             user_name=user_name,
         )
         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)}"
-        )
+        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)}"
+            )
         return self.reranker.rerank(
             query=query,
             query_embedding=query_embeddings[0],
@@ -734,6 +750,26 @@ 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,6 +23,9 @@
     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 (
@@ -263,6 +266,7 @@ 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
@@ -328,6 +332,7 @@ 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
@@ -378,7 +383,24 @@ def _dedup_by_content(memories: list) -> list:
                 unique_memories.append(mem)
             return unique_memories
 
-        deduped_memories = _dedup_by_content(enhanced_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)
         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}")
@@ -463,6 +485,7 @@ 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]