Performance optimization for searching memories

src/memos/api/handlers/search_handler.py

@@ -16,7 +16,7 @@
 from memos.api.handlers.formatters_handler import rerank_knowledge_mem
 from memos.api.product_models import APISearchRequest, SearchResponse
 from memos.dream.contextualization import CONTEXT_MEMORY_TYPE
-from memos.log import get_logger
+from memos.log import get_logger, summarize_search_request, summarize_search_results
 from memos.memories.textual.tree_text_memory.retrieve.retrieve_utils import (
     cosine_similarity_matrix,
 )
@@ -32,6 +32,8 @@
 _ENV_CONTEXT_RECALL = "MEMOS_DREAM_CONTEXT_RECALL"
 _ENV_CONTEXT_RECALL_TOP_K = "MEMOS_DREAM_CONTEXT_RECALL_TOP_K"
 _DEFAULT_CONTEXT_RECALL_TOP_K = 2
+_ENV_MMR_CANDIDATE_PRUNING = "MEMOS_MMR_CANDIDATE_PRUNING_ENABLED"
+_MMR_CANDIDATE_MULTIPLIER = 2
 
 
 def _env_enabled(name: str, default: str = "off") -> bool:
@@ -77,7 +79,10 @@ def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse
         Returns:
             SearchResponse with formatted results
         """
-        self.logger.info(f"[SearchHandler] Search Req is: {search_req}")
+        self.logger.info(
+            "[SearchHandler] Search request summary: %s",
+            summarize_search_request(search_req),
+        )
 
         # Use deepcopy to avoid modifying the original request object
         search_req_local = copy.deepcopy(search_req)
@@ -137,7 +142,8 @@ def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse
             results = hooked_results
 
         self.logger.info(
-            f"[SearchHandler] Final search results: count={len(results)} results={results}"
+            "[SearchHandler] Final search result summary: %s",
+            summarize_search_results(results),
         )
 
         return SearchResponse(
@@ -348,6 +354,12 @@ def _mmr_dedup_text_memories(
         if not text_buckets and not pref_buckets:
             return results
 
+        self._prune_mmr_candidates_by_bucket(
+            results,
+            text_top_k=text_top_k,
+            pref_top_k=pref_top_k,
+        )
+
         # Flatten all memories with their type and scores
         # flat structure: (memory_type, bucket_idx, mem, score)
         flat: list[tuple[str, int, dict[str, Any], float]] = []
@@ -552,6 +564,87 @@ def _mmr_dedup_text_memories(
 
         return results
 
+    def _prune_mmr_candidates_by_bucket(
+        self,
+        results: dict[str, Any],
+        *,
+        text_top_k: int,
+        pref_top_k: int,
+    ) -> dict[str, Any]:
+        """Keep at most twice the final quota for each memory type in each bucket."""
+        if not _env_enabled(_ENV_MMR_CANDIDATE_PRUNING, "off"):
+            return results
+
+        total_before = 0
+        total_after = 0
+        bucket_counts: list[tuple[int, int]] = []
+
+        for result_key, target_top_k in (
+            ("text_mem", text_top_k),
+            ("pref_mem", pref_top_k),
+        ):
+            buckets = results.get(result_key)
+            if not isinstance(buckets, list):
+                continue
+
+            candidate_limit = max(0, int(target_top_k)) * _MMR_CANDIDATE_MULTIPLIER
+            for bucket in buckets:
+                memories = bucket.get("memories") if isinstance(bucket, dict) else None
+                if not isinstance(memories, list):
+                    continue
+
+                before_count = len(memories)
+                total_before += before_count
+                memories_by_type: dict[str, list[dict[str, Any]]] = {}
+                for memory in memories:
+                    if not isinstance(memory, dict):
+                        continue
+                    metadata = memory.get("metadata")
+                    memory_type = (
+                        metadata.get("memory_type") if isinstance(metadata, dict) else None
+                    )
+                    memories_by_type.setdefault(str(memory_type or result_key), []).append(memory)
+
+                selected: list[dict[str, Any]] = []
+                if candidate_limit > 0:
+                    for typed_memories in memories_by_type.values():
+                        selected.extend(
+                            sorted(
+                                typed_memories,
+                                key=self._mmr_candidate_score,
+                                reverse=True,
+                            )[:candidate_limit]
+                        )
+                    selected.sort(key=self._mmr_candidate_score, reverse=True)
+
+                bucket["memories"] = selected
+                if "total_nodes" in bucket:
+                    bucket["total_nodes"] = len(selected)
+                total_after += len(selected)
+                bucket_counts.append((before_count, len(selected)))
+
+        self.logger.info(
+            "[SearchHandler] MMR candidate pruning: multiplier=%s before=%s after=%s "
+            "dropped=%s bucket_counts=%s",
+            _MMR_CANDIDATE_MULTIPLIER,
+            total_before,
+            total_after,
+            total_before - total_after,
+            bucket_counts,
+        )
+        return results
+
+    @staticmethod
+    def _mmr_candidate_score(memory: dict[str, Any]) -> float:
+        metadata = memory.get("metadata")
+        if not isinstance(metadata, dict):
+            return 0.0
+        score = metadata.get("score", metadata.get("relativity", 0.0))
+        try:
+            return float(score or 0.0)
+        except (TypeError, ValueError):
+            return 0.0
+
     @staticmethod
     def _is_unrelated(
         index: int,

src/memos/embedders/ark.py

@@ -1,6 +1,6 @@
 from memos.configs.embedder import ArkEmbedderConfig
 from memos.dependency import require_python_package
-from memos.embedders.base import BaseEmbedder
+from memos.embedders.base import BaseEmbedder, log_embedding_call
 from memos.log import get_logger
 
 
@@ -35,6 +35,7 @@ def __init__(self, config: ArkEmbedderConfig):
         # Initialize ark client
         self.client = Ark(api_key=self.config.api_key, base_url=self.config.api_base)
 
+    @log_embedding_call
     def embed(self, texts: list[str]) -> list[list[float]]:
         """
         Generate embeddings for the given texts.

src/memos/embedders/base.py

@@ -1,8 +1,63 @@
+import functools
 import re
+import time
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, TypeVar, cast
 
 from memos.configs.embedder import BaseEmbedderConfig
+from memos.log import get_logger, text_hash
+
+
+logger = get_logger(__name__)
+EmbeddingCallable = TypeVar("EmbeddingCallable", bound=Callable[..., Any])
+
+
+def log_embedding_call(func: EmbeddingCallable) -> EmbeddingCallable:
+    """Log embedding request dimensions and timing without text or vectors."""
+
+    @functools.wraps(func)
+    def wrapper(self, texts, *args, **kwargs):
+        normalized_texts = [texts] if isinstance(texts, str) else list(texts or [])
+        text_lengths = [len(str(text or "")) for text in normalized_texts]
+        config = getattr(self, "config", None)
+        model = getattr(config, "model_name_or_path", None) or "unknown"
+        backup_model = getattr(config, "backup_model_name_or_path", None) or "none"
+        backup_enabled = bool(getattr(self, "use_backup_client", False))
+        started_at = time.perf_counter()
+        status = "success"
+        error_type = None
+        try:
+            return func(self, texts, *args, **kwargs)
+        except Exception as exc:
+            status = "failed"
+            error_type = type(exc).__name__
+            raise
+        finally:
+            elapsed_ms = (time.perf_counter() - started_at) * 1000
+            log_message = (
+                "Embedding request model=%s backup_model=%s backup_enabled=%s "
+                "batch_size=%d total_chars=%d max_chars=%d text_hash=%s "
+                "elapsed_ms=%.2f status=%s"
+            )
+            log_values = (
+                model,
+                backup_model,
+                backup_enabled,
+                len(normalized_texts),
+                sum(text_lengths),
+                max(text_lengths, default=0),
+                text_hash(normalized_texts),
+                elapsed_ms,
+                status,
+            )
+            if error_type is None:
+                logger.info(log_message, *log_values)
+            else:
+                logger.info(log_message + " error_type=%s", *log_values, error_type)
+
+    return cast("EmbeddingCallable", wrapper)
 
 
 def _count_tokens_for_embedding(text: str) -> int: