From 8f80cd17c1f222a9184ae1dfe7fd3ff85c0690cd Mon Sep 17 00:00:00 2001
From: "jinli.yl" <jinli.yl@alibaba-inc.com>
Date: Wed, 25 Feb 2026 18:27:05 +0800
Subject: [PATCH] perf(core): optimize memory store similarity search
 performance

---
 reme/__init__.py                             |  2 +-
 reme/core/memory_store/local_memory_store.py | 39 +++++++++++++-------
 reme/core/utils/env_utils.py                 |  2 +-
 3 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/reme/__init__.py b/reme/__init__.py
index b02aaa94..21d11dda 100644
--- a/reme/__init__.py
+++ b/reme/__init__.py
@@ -20,7 +20,7 @@ __all__ = [
     "ReMeFs",
 ]
 
-__version__ = "0.3.0.0b4"
+__version__ = "0.3.0.0b5"
 
 """
 conda create -n fl_test2 python=3.10
diff --git a/reme/core/memory_store/local_memory_store.py b/reme/core/memory_store/local_memory_store.py
index bee6882b..989983df 100644
--- a/reme/core/memory_store/local_memory_store.py
+++ b/reme/core/memory_store/local_memory_store.py
@@ -5,12 +5,13 @@ import time
 from dataclasses import dataclass
 from pathlib import Path
 
+import numpy as np
 from loguru import logger
 
 from .base_memory_store import BaseMemoryStore
 from ..enumeration import MemorySource
 from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
-from ..utils.common_utils import cosine_similarity
+from ..utils.common_utils import batch_cosine_similarity
 
 
 @dataclass
@@ -305,25 +306,37 @@ class LocalMemoryStore(BaseMemoryStore):
             return []
 
         source_values = {s.value for s in sources} if sources else None
-        results = []
+
+        # Collect candidate chunks with embeddings
+        candidates: list[tuple[_ChunkRecord, list[float]]] = []
         for rec in self._chunks.values():
             if source_values and rec.source not in source_values:
                 continue
             if not rec.embedding:
                 continue
+            candidates.append((rec, rec.embedding))
 
-            similarity = cosine_similarity(query_embedding, rec.embedding)
-            results.append(
-                MemorySearchResult(
-                    path=rec.path,
-                    start_line=rec.start_line,
-                    end_line=rec.end_line,
-                    score=similarity,
-                    snippet=rec.text,
-                    source=MemorySource(rec.source),
-                    raw_metric=1.0 - similarity,  # distance equivalent
-                ),
+        if not candidates:
+            return []
+
+        # Build embedding matrix and compute similarities in batch
+        query_array = np.array([query_embedding])  # Shape: (1, emb_size)
+        chunk_embeddings = np.array([emb for _, emb in candidates])  # Shape: (n, emb_size)
+        similarities = batch_cosine_similarity(query_array, chunk_embeddings)[0]  # Shape: (n,)
+
+        # Build results
+        results = [
+            MemorySearchResult(
+                path=rec.path,
+                start_line=rec.start_line,
+                end_line=rec.end_line,
+                score=float(similarity),
+                snippet=rec.text,
+                source=MemorySource(rec.source),
+                raw_metric=1.0 - float(similarity),
             )
+            for (rec, _), similarity in zip(candidates, similarities)
+        ]
 
         results.sort(key=lambda r: r.score, reverse=True)
         return results[:limit]
diff --git a/reme/core/utils/env_utils.py b/reme/core/utils/env_utils.py
index 433039cd..34dc82d3 100644
--- a/reme/core/utils/env_utils.py
+++ b/reme/core/utils/env_utils.py
@@ -54,7 +54,7 @@ def load_env(path: str | Path | None = None, enable_log: bool = True) -> None:
             _ENV_LOADED = True
             return
 
-    logger.warning(".env file not found in search path")
+    # logger.info(".env file not found in search path")
 
 
 def reset_env_flag() -> None: