From 141f2b8b65fdbd522782d8f62a43e53b54c12831 Mon Sep 17 00:00:00 2001
From: Ziyang Guo <121015044+RerankerGuo@users.noreply.github.com>
Date: Thu, 2 Jul 2026 15:58:09 +0800
Subject: [PATCH] fix: batch missing search embeddings

---
 src/memos/api/handlers/search_handler.py      | 12 ++++--
 .../test_search_handler_embedding_batches.py  | 41 +++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)
 create mode 100644 tests/api/test_search_handler_embedding_batches.py

diff --git a/src/memos/api/handlers/search_handler.py b/src/memos/api/handlers/search_handler.py
index 03e6977ad..821d8cd2d 100644
--- a/src/memos/api/handlers/search_handler.py
+++ b/src/memos/api/handlers/search_handler.py
@@ -32,6 +32,7 @@
 _ENV_CONTEXT_RECALL = "MEMOS_DREAM_CONTEXT_RECALL"
 _ENV_CONTEXT_RECALL_TOP_K = "MEMOS_DREAM_CONTEXT_RECALL_TOP_K"
 _DEFAULT_CONTEXT_RECALL_TOP_K = 2
+_MISSING_EMBEDDING_BATCH_SIZE = 10
 
 
 def _env_enabled(name: str, default: str = "off") -> bool:
@@ -590,10 +591,13 @@ def _extract_embeddings(self, memories: list[dict[str, Any]]) -> list[list[float
             missing_documents.append(mem.get("memory", ""))
 
         if missing_indices:
-            computed = self.searcher.embedder.embed(missing_documents)
-            for idx, embedding in zip(missing_indices, computed, strict=False):
-                embeddings[idx] = embedding
-                memories[idx]["metadata"]["embedding"] = embedding
+            for start in range(0, len(missing_documents), _MISSING_EMBEDDING_BATCH_SIZE):
+                batch_documents = missing_documents[start : start + _MISSING_EMBEDDING_BATCH_SIZE]
+                batch_indices = missing_indices[start : start + _MISSING_EMBEDDING_BATCH_SIZE]
+                computed = self.searcher.embedder.embed(batch_documents)
+                for idx, embedding in zip(batch_indices, computed, strict=False):
+                    embeddings[idx] = embedding
+                    memories[idx]["metadata"]["embedding"] = embedding
 
         return embeddings
 
diff --git a/tests/api/test_search_handler_embedding_batches.py b/tests/api/test_search_handler_embedding_batches.py
new file mode 100644
index 000000000..c9e44926e
--- /dev/null
+++ b/tests/api/test_search_handler_embedding_batches.py
@@ -0,0 +1,41 @@
+from memos.api.handlers.base_handler import HandlerDependencies
+from memos.api.handlers.search_handler import SearchHandler
+
+
+class BatchLimitedEmbedder:
+    def __init__(self, *, limit: int):
+        self.limit = limit
+        self.calls: list[list[str]] = []
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        self.calls.append(list(texts))
+        if len(texts) > self.limit:
+            raise AssertionError(f"batch too large: {len(texts)}")
+        return [[float(len(text)), 0.0] for text in texts]
+
+
+def _handler(embedder: BatchLimitedEmbedder) -> SearchHandler:
+    searcher = type("FakeSearcher", (), {"embedder": embedder})()
+    return SearchHandler(
+        HandlerDependencies(
+            naive_mem_cube=object(),
+            mem_scheduler=object(),
+            searcher=searcher,
+            deepsearch_agent=object(),
+        )
+    )
+
+
+def test_extract_embeddings_batches_missing_documents():
+    embedder = BatchLimitedEmbedder(limit=10)
+    handler = _handler(embedder)
+    memories = [
+        {"memory": f"memory {idx}", "metadata": {}}
+        for idx in range(25)
+    ]
+
+    embeddings = handler._extract_embeddings(memories)
+
+    assert [len(call) for call in embedder.calls] == [10, 10, 5]
+    assert len(embeddings) == 25
+    assert all(mem["metadata"]["embedding"] for mem in memories)