Fix dedup handling in simple search

Author: [email protected]

src/memos/api/handlers/formatters_handler.py

@@ -29,7 +29,7 @@ def to_iter(running: Any) -> list[Any]:
     return list(running) if running else []
 
 
-def format_memory_item(memory_data: Any) -> dict[str, Any]:
+def format_memory_item(memory_data: Any, include_embedding: bool = False) -> dict[str, Any]:
     """
     Format a single memory item for API response.
 
@@ -47,7 +47,8 @@ def format_memory_item(memory_data: Any) -> dict[str, Any]:
     ref_id = f"[{memory_id.split('-')[0]}]"
 
     memory["ref_id"] = ref_id
-    memory["metadata"]["embedding"] = []
+    if not include_embedding:
+        memory["metadata"]["embedding"] = []
     memory["metadata"]["sources"] = []
     memory["metadata"]["usage"] = []
     memory["metadata"]["ref_id"] = ref_id

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,9 @@ 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._strip_embeddings(results)
 
         self.logger.info(
             f"[SearchHandler] Final search results: count={len(results)} results={results}"
@@ -63,6 +71,110 @@ 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]:
+        buckets = results.get("text_mem", [])
+        if not buckets:
+            return results
+
+        flat: list[tuple[int, dict[str, Any], float]] = []
+        for bucket_idx, bucket in enumerate(buckets):
+            for mem in bucket.get("memories", []):
+                score = mem.get("metadata", {}).get("relativity", 0.0)
+                flat.append((bucket_idx, mem, score))
+
+        if len(flat) <= 1:
+            return results
+
+        embeddings = self._extract_embeddings([mem for _, mem, _ in flat])
+        if embeddings is None:
+            documents = [mem.get("memory", "") for _, mem, _ in flat]
+            embeddings = self.searcher.embedder.embed(documents)
+
+        similarity_matrix = cosine_similarity_matrix(embeddings)
+
+        indices_by_bucket: dict[int, list[int]] = {i: [] for i in range(len(buckets))}
+        for flat_index, (bucket_idx, _, _) in enumerate(flat):
+            indices_by_bucket[bucket_idx].append(flat_index)
+
+        selected_global: list[int] = []
+        selected_by_bucket: dict[int, list[int]] = {i: [] for i in range(len(buckets))}
+
+        ordered_indices = sorted(range(len(flat)), key=lambda idx: flat[idx][2], reverse=True)
+        for idx in ordered_indices:
+            bucket_idx = flat[idx][0]
+            if len(selected_by_bucket[bucket_idx]) >= target_top_k:
+                continue
+            if self._is_unrelated(idx, selected_global, similarity_matrix, 0.85):
+                selected_by_bucket[bucket_idx].append(idx)
+                selected_global.append(idx)
+
+        for bucket_idx in range(len(buckets)):
+            if len(selected_by_bucket[bucket_idx]) >= min(
+                target_top_k, len(indices_by_bucket[bucket_idx])
+            ):
+                continue
+            remaining_indices = [
+                idx
+                for idx in indices_by_bucket.get(bucket_idx, [])
+                if idx not in selected_by_bucket[bucket_idx]
+            ]
+            if not remaining_indices:
+                continue
+            # Fill to target_top_k with the least-similar candidates to preserve diversity.
+            remaining_indices.sort(
+                key=lambda idx: self._max_similarity(idx, selected_global, similarity_matrix)
+            )
+            for idx in remaining_indices:
+                if len(selected_by_bucket[bucket_idx]) >= target_top_k:
+                    break
+                selected_by_bucket[bucket_idx].append(idx)
+                selected_global.append(idx)
+
+        for bucket_idx, bucket in enumerate(buckets):
+            selected_indices = selected_by_bucket.get(bucket_idx, [])
+            bucket["memories"] = [flat[i][1] for i in selected_indices[:target_top_k]]
+        return results
+
+    @staticmethod
+    def _is_unrelated(
+        index: int,
+        selected_indices: list[int],
+        similarity_matrix: list[list[float]],
+        similarity_threshold: float,
+    ) -> bool:
+        return all(similarity_matrix[index][j] <= similarity_threshold for j in selected_indices)
+
+    @staticmethod
+    def _max_similarity(
+        index: int, selected_indices: list[int], similarity_matrix: list[list[float]]
+    ) -> float:
+        if not selected_indices:
+            return 0.0
+        return max(similarity_matrix[index][j] for j in 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
+
+    @staticmethod
+    def _strip_embeddings(results: dict[str, Any]) -> None:
+        for bucket in results.get("text_mem", []):
+            for mem in bucket.get("memories", []):
+                metadata = mem.get("metadata", {})
+                if "embedding" in metadata:
+                    metadata["embedding"] = []
+        for bucket in results.get("tool_mem", []):
+            for mem in bucket.get("memories", []):
+                metadata = mem.get("metadata", {})
+                if "embedding" in metadata:
+                    metadata["embedding"] = []
+
     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

@@ -190,7 +190,10 @@ def mix_search_memories(
         )
         memories = merged_memories[: search_req.top_k]
 
-        formatted_memories = [format_textual_memory_item(item) for item in memories]
+        formatted_memories = [
+            format_textual_memory_item(item, include_embedding=search_req.dedup == "sim")
+            for item in memories
+        ]
         self.submit_memory_history_async_task(
             search_req=search_req,
             user_context=user_context,
@@ -234,7 +237,10 @@ def update_search_memories_to_redis(
                     mem_cube=self.mem_cube,
                     mode=SearchMode.FAST,
                 )
-                formatted_memories = [format_textual_memory_item(data) for data in memories]
+                formatted_memories = [
+                    format_textual_memory_item(data, include_embedding=search_req.dedup == "sim")
+                    for data in memories
+                ]
             else:
                 memories = [
                     TextualMemoryItem.from_dict(one) for one in memories_to_store["memories"]

src/memos/mem_scheduler/utils/api_utils.py

@@ -6,14 +6,15 @@
 from memos.memories.textual.tree import TextualMemoryItem
 
 
-def format_textual_memory_item(memory_data: Any) -> dict[str, Any]:
+def format_textual_memory_item(memory_data: Any, include_embedding: bool = False) -> dict[str, Any]:
     """Format a single memory item for API response."""
     memory = memory_data.model_dump()
     memory_id = memory["id"]
     ref_id = f"[{memory_id.split('-')[0]}]"
 
     memory["ref_id"] = ref_id
-    memory["metadata"]["embedding"] = []
+    if not include_embedding:
+        memory["metadata"]["embedding"] = []
     memory["metadata"]["sources"] = []
     memory["metadata"]["ref_id"] = ref_id
     memory["metadata"]["id"] = memory_id

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

@@ -124,8 +124,6 @@ def post_retrieve(
     ):
         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(
@@ -180,11 +178,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 +207,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,
+            dedup=dedup,
         )
 
         logger.info(f"[SEARCH] Done. Total {len(final_results)} results.")
@@ -296,50 +290,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 +673,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 +697,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