feat: optimize memory search deduplication and fix parsing bugs

- Tune similarity threshold to 0.92 for 'dedup=sim' to preserve subtle semantic nuances.
- Implement recall expansion (5x Top-K) when deduplicating to ensure output diversity.
- Remove aggressive filling logic to strictly enforce the similarity threshold.
- Fix attribute error in MultiModalStructMemReader by correctly importing parse_json_result.
- Replace fragile eval() with robust parse_json_result in TaskGoalParser to handle JSON booleans.

Author: [email protected]

src/memos/api/handlers/search_handler.py

@@ -55,12 +55,19 @@ def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse
         """
         self.logger.info(f"[SearchHandler] Search Req is: {search_req}")
 
+        # Increase recall pool if deduplication is enabled to ensure diversity
+        original_top_k = search_req.top_k
+        if search_req.dedup == "sim":
+            search_req.top_k = original_top_k * 5
+
         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)
+            results = self._dedup_text_memories(results, original_top_k)
             self._strip_embeddings(results)
+            # Restore original top_k for downstream logic or response metadata
+            search_req.top_k = original_top_k
 
         self.logger.info(
             f"[SearchHandler] Final search results: count={len(results)} results={results}"
@@ -104,35 +111,18 @@ def _dedup_text_memories(self, results: dict[str, Any], target_top_k: int) -> di
             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):
+            # Use 0.92 threshold strictly
+            if self._is_unrelated(idx, selected_global, similarity_matrix, 0.92):
                 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)
+        # Removed the 'filling' logic that was pulling back similar items.
+        # Now it will only return items that truly pass the 0.92 threshold,
+        # up to target_top_k.
 
         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]]
+            bucket["memories"] = [flat[i][1] for i in selected_indices]
         return results
 
     @staticmethod

src/memos/mem_reader/multi_modal_struct.py

@@ -10,6 +10,7 @@
 from memos.mem_reader.read_multi_modal import MultiModalParser, detect_lang
 from memos.mem_reader.read_multi_modal.base import _derive_key
 from memos.mem_reader.simple_struct import PROMPT_DICT, SimpleStructMemReader
+from memos.mem_reader.utils import parse_json_result
 from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
 from memos.templates.tool_mem_prompts import TOOL_TRAJECTORY_PROMPT_EN, TOOL_TRAJECTORY_PROMPT_ZH
 from memos.types import MessagesType
@@ -377,7 +378,7 @@ def _get_llm_response(
         messages = [{"role": "user", "content": prompt}]
         try:
             response_text = self.llm.generate(messages)
-            response_json = self.parse_json_result(response_text)
+            response_json = parse_json_result(response_text)
         except Exception as e:
             logger.error(f"[LLM] Exception during chat generation: {e}")
             response_json = {

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

@@ -4,7 +4,6 @@
 from pathlib import Path
 from typing import Any
 
-
 import numpy as np
 
 from memos.dependency import require_python_package

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

@@ -5,7 +5,10 @@
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
 from memos.memories.textual.tree_text_memory.retrieve.retrieval_mid_structs import ParsedTaskGoal
-from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import FastTokenizer
+from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import (
+    FastTokenizer,
+    parse_json_result,
+)
 from memos.memories.textual.tree_text_memory.retrieve.utils import TASK_PARSE_PROMPT
 
 
@@ -111,8 +114,10 @@ def _parse_response(self, response: str, **kwargs) -> ParsedTaskGoal:
         for attempt_times in range(attempts):
             try:
                 context = kwargs.get("context", "")
-                response = response.replace("```", "").replace("json", "").strip()
-                response_json = eval(response)
+                response_json = parse_json_result(response)
+                if not response_json:
+                    raise ValueError("Parsed JSON is empty")
+
                 return ParsedTaskGoal(
                     memories=response_json.get("memories", []),
                     keys=response_json.get("keys", []),
@@ -123,6 +128,8 @@ def _parse_response(self, response: str, **kwargs) -> ParsedTaskGoal:
                     context=context,
                 )
             except Exception as e:
-                raise ValueError(
-                    f"Failed to parse LLM output: {e}\nRaw response:\n{response} retried: {attempt_times + 1}/{attempts + 1}"
-                ) from e
+                if attempt_times == attempts - 1:
+                    raise ValueError(
+                        f"Failed to parse LLM output: {e}\nRaw response:\n{response} retried: {attempt_times + 1}/{attempts}"
+                    ) from e
+                continue