Fix dedup handling in simple search

Author: [email protected]

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

@@ -296,50 +296,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,
@@ -699,6 +655,7 @@ def _retrieve_simple(
         top_k: int,
         search_filter: dict | None = None,
         user_name: str | None = None,
+        dedup: str | None = None,
         **kwargs,
     ):
         """
@@ -721,10 +678,16 @@ def _retrieve_simple(
             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)}")
         if dedup == "no":
             selected_items = items
+        elif dedup == "sim":
+            selected_items = self.deduplicate_similar_items(items)
+            logger.info(
+                f"[SIMPLESEARCH] after similarity dedup items count: {len(selected_items)}"
+            )
         else:
             documents = [getattr(item, "memory", "") for item in items]
             documents_embeddings = self.embedder.embed(documents)
@@ -763,12 +726,34 @@ def _deduplicate_similar_results(
         embeddings = self.embedder.embed(documents)
         similarity_matrix = cosine_similarity_matrix(embeddings)
 
+        selected_indices = self._select_unrelated_indices(
+            similarity_matrix, similarity_threshold
+        )
+        return [sorted_results[i] for i in selected_indices]
+
+    def deduplicate_similar_items(
+        self, items: list[TextualMemoryItem], similarity_threshold: float = 0.85
+    ) -> list[TextualMemoryItem]:
+        """Deduplicate memory items by semantic similarity while preserving order."""
+        if len(items) <= 1:
+            return items
+        documents = [getattr(item, "memory", "") for item in items]
+        embeddings = self.embedder.embed(documents)
+        similarity_matrix = cosine_similarity_matrix(embeddings)
+        selected_indices = self._select_unrelated_indices(
+            similarity_matrix, similarity_threshold
+        )
+        return [items[i] for i in selected_indices]
+
+    @staticmethod
+    def _select_unrelated_indices(
+        similarity_matrix: list[list[float]], similarity_threshold: float
+    ) -> list[int]:
         selected_indices: list[int] = []
-        for i in range(len(sorted_results)):
+        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 [sorted_results[i] for i in selected_indices]
+        return selected_indices
 
     @timed
     def _sort_and_trim(

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 (
@@ -383,22 +380,10 @@ 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)
+            deduped_memories = self.searcher.deduplicate_similar_items(enhanced_memories)
         else:
             deduped_memories = _dedup_by_content(enhanced_memories)
         formatted_memories = [format_memory_item(data) for data in deduped_memories]