refactor(embedding): defer provider construction until first remote call (#361)

- Changed dimensions property to avoid forcing provider construction
- Added _ensure_model method to construct provider on demand
- Modified __call__ to ensure model exists before use
- Updated _start to defer provider initialization
- Removed eager health check during startup
- Added compact embedding serialization with base64 encoding
- Implemented batch processing for vector search with heap-based ranking
- Added document_ids property to keyword index interface
- Updated chunk persistence to handle legacy JSON embeddings
- Optimized memory usage by avoiding materialization of metadata in document_ids

reme/components/as_embedding/__init__.py

@@ -30,16 +30,26 @@ class BaseAsEmbedding(BaseComponent):
 
     @property
     def dimensions(self) -> int:
-        """Return the embedding dimension size."""
-        assert self.model is not None
-        return self.model.dimensions
+        """Return configured dimensions without forcing provider construction."""
+        if self.model is not None:
+            return self.model.dimensions
+        dimensions = self.kwargs.get("dimensions")
+        if dimensions is None:
+            raise RuntimeError("Embedding dimensions are required before provider initialization.")
+        return int(dimensions)
 
     async def __call__(self, inputs: list[Any], **kwargs) -> list[list[float]]:
+        self._ensure_model()
         assert self.model is not None
         response = await self.model(inputs, **kwargs)  # pylint: disable=not-callable
         return response.embeddings
 
     async def _start(self) -> None:
+        """Defer provider construction until the first remote embedding call."""
+        return None
+
+    def _ensure_model(self) -> None:
+        """Construct the provider on demand while keeping dimensions locally available."""
         if self.model is not None:
             return
 
@@ -50,7 +60,8 @@ class BaseAsEmbedding(BaseComponent):
         if model_cls is None:
             raise ValueError(f"{self.credential_cls.__name__} does not support embeddings.")
 
-        dimensions = kwargs.pop("dimensions")
+        dimensions = self.dimensions
+        kwargs.pop("dimensions", None)
         params_dict = kwargs.pop("parameters", None)
         parameters = model_cls.Parameters(**params_dict) if params_dict else None
 

reme/components/embedding_store/local_embedding_store.py

@@ -98,8 +98,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
 
     async def _fill_misses(self, misses: list[Miss], results: list[np.ndarray | None], **kwargs) -> None:
         size = self.max_batch_size
-        batches = [misses[i : i + size] for i in range(0, len(misses), size)]
-        for batch in batches:
+        for start in range(0, len(misses), size):
+            batch = misses[start : start + size]
             for idx, key, emb in await self._compute_batch(batch, **kwargs):
                 results[idx] = emb
                 self._cache_put(key, emb)

reme/components/file_store/local_file_store.py

@@ -1,6 +1,10 @@
 """In-memory file store with compressed JSONL persistence on close."""
 
+import base64
 import datetime
+import heapq
+import json
+from collections.abc import Iterable
 from contextlib import suppress
 
 import numpy as np
@@ -16,6 +20,9 @@ from ...utils import batch_cosine_similarity
 from ...utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst
 
 CachedEmbedding = tuple[str, np.ndarray]
+_EMBEDDING_F16_B64_FIELD = "_embedding_f16_b64"
+_EMBEDDING_F16_DTYPE = np.dtype("<f2")
+_VECTOR_SEARCH_BATCH_SIZE = 1024
 
 
 @R.register("local")
@@ -61,9 +68,6 @@ class LocalFileStore(BaseFileStore):
     async def _start(self) -> None:
         self.component_metadata_path.mkdir(parents=True, exist_ok=True)
         await super()._start()
-        if self.embedding_store is not None and not await self.embedding_store.health_check():
-            self.logger.warning(f"{self.name}: embedding unhealthy, vector disabled")
-            self.embedding_store = None
         await self.load()
 
     async def _close(self) -> None:
@@ -98,7 +102,7 @@ class LocalFileStore(BaseFileStore):
         chunk.embedding = None
         return True
 
-    def _drop_stale_embeddings(self, chunks: list[FileChunk], context: str) -> None:
+    def _drop_stale_embeddings(self, chunks: Iterable[FileChunk], context: str) -> None:
         for chunk in chunks:
             self._drop_stale_embedding(chunk, context)
 
@@ -112,7 +116,7 @@ class LocalFileStore(BaseFileStore):
             for line in read_jsonl_zst(self.chunks_path, self.encoding):
                 line = line.strip()
                 if line:
-                    chunk = FileChunk.model_validate_json(line)
+                    chunk = self._deserialize_chunk(line)
                     self.file_chunks[chunk.id] = chunk
             self.logger.info(f"Loaded {len(self.file_chunks)} chunks from {self.chunks_path}")
             self._invalidate_stale_embeddings()
@@ -121,11 +125,35 @@ class LocalFileStore(BaseFileStore):
         except Exception as e:
             self.logger.exception(f"Failed to load {self.chunks_path}: {e}")
 
+    @staticmethod
+    def _deserialize_chunk(line: str) -> FileChunk:
+        """Read compact vectors while retaining legacy JSON-list compatibility."""
+        payload = json.loads(line)
+        encoded = payload.pop(_EMBEDDING_F16_B64_FIELD, None)
+        if encoded is not None:
+            raw = base64.b64decode(encoded, validate=True)
+            if len(raw) % _EMBEDDING_F16_DTYPE.itemsize:
+                raise ValueError("Invalid float16 embedding byte length")
+            payload["embedding"] = np.frombuffer(raw, dtype=_EMBEDDING_F16_DTYPE)
+        return FileChunk.model_validate(payload)
+
+    @staticmethod
+    def _serialize_chunk(chunk: FileChunk) -> str:
+        """Serialize embeddings without expanding float16 values into Python floats."""
+        payload = chunk.model_dump(mode="json", exclude={"embedding"})
+        if chunk.embedding is not None:
+            embedding = np.asarray(chunk.embedding, dtype=_EMBEDDING_F16_DTYPE)
+            if embedding.ndim != 1:
+                raise ValueError("FileChunk embedding must be one-dimensional")
+            raw = np.ascontiguousarray(embedding).tobytes()
+            payload[_EMBEDDING_F16_B64_FIELD] = base64.b64encode(raw).decode("ascii")
+        return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
+
     def _invalidate_stale_embeddings(self) -> None:
         """Drop persisted embeddings whose dimension no longer matches the active model."""
         if self.embedding_store is None:
             return
-        self._drop_stale_embeddings(list(self.file_chunks.values()), "load")
+        self._drop_stale_embeddings(self.file_chunks.values(), "load")
 
     async def _backfill_missing_embeddings(self) -> None:
         """Embed persisted chunks that predate embedding being enabled."""
@@ -162,18 +190,14 @@ class LocalFileStore(BaseFileStore):
         if not docs:
             return
 
-        expected_ids = set(docs)
+        expected_ids = docs.keys()
         live_ids = None
         with suppress(Exception):
-            live_ids = set(getattr(self.keyword_index, "doc_meta", {}).keys())
+            live_ids = self.keyword_index.document_ids
 
         if live_ids == expected_ids:
             return
 
-        n_docs = getattr(self.keyword_index, "n_docs", None)
-        if live_ids is None and n_docs == len(expected_ids):
-            return
-
         self.logger.warning(f"{self.name}: keyword index mismatch with chunks; rebuilding {len(docs)} docs")
         await self.keyword_index.reset_index(docs)
 
@@ -181,7 +205,11 @@ class LocalFileStore(BaseFileStore):
         """Atomically rewrite the JSONL, then cascade dump into keyword_index and file_graph."""
         assert self.file_graph is not None
         try:
-            write_jsonl_zst(self.chunks_path, (c.model_dump_json() for c in self.file_chunks.values()), self.encoding)
+            write_jsonl_zst(
+                self.chunks_path,
+                (self._serialize_chunk(c) for c in self.file_chunks.values()),
+                self.encoding,
+            )
             self.logger.info(f"Saved {len(self.file_chunks)} chunks to {self.chunks_path}")
         except Exception as e:
             self.logger.exception(f"Failed to write {self.chunks_path}: {e}")
@@ -318,7 +346,7 @@ class LocalFileStore(BaseFileStore):
     # -- search ---------------------------------------------------------------
 
     async def vector_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
-        if self.embedding_store is None or not query:
+        if self.embedding_store is None or not query or limit <= 0:
             return []
 
         try:
@@ -334,23 +362,45 @@ class LocalFileStore(BaseFileStore):
             )
             return []
 
-        candidates = [
-            c
-            for c in self.file_chunks.values()
-            if self._embedding_dim_matches(c.embedding) and self._matches_search_filter(c, search_filter)
-        ]
-        if not candidates:
-            return []
+        top: list[tuple[float, int, FileChunk]] = []
+        candidates: list[FileChunk] = []
+        embeddings: list[np.ndarray] = []
+        order = 0
 
-        candidate_embeddings = np.stack([c.embedding for c in candidates])
-        similarities = batch_cosine_similarity(query_embedding.reshape(1, -1), candidate_embeddings)[0]
+        def score_batch() -> None:
+            nonlocal order
+            if not candidates:
+                return
+            matrix = np.stack(embeddings)
+            similarities = batch_cosine_similarity(query_embedding.reshape(1, -1), matrix)[0]
+            for candidate, similarity in zip(candidates, similarities):
+                score = float(similarity)
+                item = (score, -order, candidate)
+                if len(top) < limit:
+                    heapq.heappush(top, item)
+                elif item[:2] > top[0][:2]:
+                    heapq.heapreplace(top, item)
+                order += 1
+            candidates.clear()
+            embeddings.clear()
 
-        results = [
-            c.model_copy(update={"scores": {"vector": float(s), "score": float(s)}})
-            for c, s in zip(candidates, similarities)
+        for candidate in self.file_chunks.values():
+            if not self._embedding_dim_matches(candidate.embedding) or not self._matches_search_filter(
+                candidate,
+                search_filter,
+            ):
+                continue
+            candidates.append(candidate)
+            embeddings.append(candidate.embedding)
+            if len(candidates) >= _VECTOR_SEARCH_BATCH_SIZE:
+                score_batch()
+        score_batch()
+
+        ranked = sorted(top, key=lambda item: (-item[0], -item[1]))
+        return [
+            candidate.model_copy(update={"scores": {"vector": score, "score": score}})
+            for score, _neg_order, candidate in ranked
         ]
-        results.sort(key=lambda r: r.score, reverse=True)
-        return results[:limit]
 
     async def keyword_search(self, query: str, limit: int, search_filter: dict) -> list[FileChunk]:
         if not self.keyword_index:

reme/components/keyword_index/base_keyword_index.py

@@ -1,6 +1,7 @@
 """Abstract base class for keyword indexes (BM25 and other lexical backends)."""
 
 from abc import abstractmethod
+from collections.abc import Set
 
 from ..base_component import BaseComponent
 from ..tokenizer import BaseTokenizer
@@ -25,6 +26,11 @@ class BaseKeywordIndex(BaseComponent):
     async def _close(self) -> None:
         await self.dump()
 
+    @property
+    def document_ids(self) -> Set[str]:
+        """Return live document IDs without materializing per-document metadata."""
+        raise NotImplementedError(f"{type(self).__name__} does not expose live document IDs")
+
     def _tokenize(self, text: str) -> list[str]:
         """Tokenize a single text into a list of tokens."""
         if self.tokenizer is None:

reme/components/keyword_index/bm25_index.py

@@ -21,6 +21,7 @@ import math
 import pickle
 import re
 from collections import Counter
+from collections.abc import KeysView
 from pathlib import Path
 
 import numpy as np
@@ -96,7 +97,7 @@ class BM25Index(BaseKeywordIndex):
     @property
     def n_docs(self) -> int:
         """Number of live (non-deleted) documents."""
-        return 0 if self._deleted.size == 0 else int((~self._deleted).sum())
+        return len(self._doc_id_to_idx)
 
     @property
     def total_len(self) -> int:
@@ -109,6 +110,11 @@ class BM25Index(BaseKeywordIndex):
         n = self.n_docs
         return self.total_len / n if n > 0 else 0.0
 
+    @property
+    def document_ids(self) -> KeysView[str]:
+        """Return live document IDs without materializing per-document token metadata."""
+        return self._doc_id_to_idx.keys()
+
     @property
     def doc_meta(self) -> dict[str, dict]:
         """Per-live-doc length and unique token_id set, keyed by doc_id."""

tests/unit/test_as_embedding_lazy.py

@@ -0,0 +1,67 @@
+"""Tests for lazy AgentScope embedding provider construction."""
+
+import asyncio
+from types import SimpleNamespace
+
+from reme.components.as_embedding import BaseAsEmbedding
+
+
+class FakeModel:
+    """Minimal async embedding model used to observe construction."""
+
+    constructions = 0
+
+    class Parameters:
+        """Accept arbitrary provider parameters."""
+
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+    def __init__(self, credential, dimensions, parameters=None, **kwargs):
+        type(self).constructions += 1
+        self.credential = credential
+        self.dimensions = dimensions
+        self.parameters = parameters
+        self.kwargs = kwargs
+
+    async def __call__(self, inputs, **_kwargs):
+        return SimpleNamespace(embeddings=[[float(index)] * self.dimensions for index, _ in enumerate(inputs)])
+
+
+class FakeCredential:
+    """Credential boundary that resolves to ``FakeModel``."""
+
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+    @staticmethod
+    def get_embedding_model_class():
+        """Return the fake provider model class."""
+        return FakeModel
+
+
+class LazyAsEmbedding(BaseAsEmbedding):
+    """Concrete wrapper backed by the local fakes."""
+
+    credential_cls = FakeCredential
+
+
+def test_provider_is_constructed_once_on_first_call():
+    """Start and dimension inspection stay local; the first request builds the model."""
+
+    async def go():
+        FakeModel.constructions = 0
+        embedding = LazyAsEmbedding(dimensions=3, credential={"token": "test"}, parameters={"mode": "test"})
+
+        await embedding.start()
+        assert embedding.model is None
+        assert embedding.dimensions == 3
+        assert FakeModel.constructions == 0
+
+        assert await embedding(["first"]) == [[0.0, 0.0, 0.0]]
+        assert await embedding(["second"]) == [[0.0, 0.0, 0.0]]
+        assert FakeModel.constructions == 1
+
+        await embedding.close()
+
+    asyncio.run(go())

tests/unit/test_file_store_consistency.py

@@ -3,6 +3,8 @@
 # pylint: disable=protected-access
 
 import asyncio
+import base64
+import json
 import os
 import tempfile
 
@@ -10,7 +12,9 @@ import numpy as np
 import pytest
 
 from reme.components.file_store import FaissLocalFileStore, LocalFileStore
+from reme.components.file_store import local_file_store as local_file_store_module
 from reme.schema import FileChunk, FileNode
+from reme.utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst
 
 
 class temp_chdir:
@@ -72,6 +76,17 @@ class UnhealthyCountingEmbeddingStore(CountingFakeEmbeddingStore):
         return False
 
 
+class HealthCountingEmbeddingStore(FakeEmbeddingStore):
+    """Fake provider that records eager health checks."""
+
+    def __init__(self):
+        self.health_calls = 0
+
+    async def health_check(self, _timeout: float = 2.0) -> bool:
+        self.health_calls += 1
+        return True
+
+
 class WrongDimEmbeddingStore(FakeEmbeddingStore):
     """Fake embedding store that returns vectors with the wrong dimension."""
 
@@ -84,6 +99,23 @@ class WrongDimEmbeddingStore(FakeEmbeddingStore):
         return nodes
 
 
+class CountOnlyKeywordIndex:
+    """Keyword backend that knows its size but cannot expose document IDs."""
+
+    def __init__(self, n_docs: int):
+        self.n_docs = n_docs
+        self.reset_docs = None
+
+    @property
+    def document_ids(self):
+        """Signal that exact live IDs are unavailable."""
+        raise NotImplementedError
+
+    async def reset_index(self, docs):
+        """Record the documents requested for rebuilding."""
+        self.reset_docs = docs
+
+
 def run(coro):
     """Run an async test body."""
     return asyncio.run(coro)
@@ -120,6 +152,23 @@ def test_keyword_only_upsert_removes_old_chunks_and_docs():
     run(go())
 
 
+def test_start_does_not_health_check_embedding_without_backfill():
+    """Hot startup keeps local vector retrieval independent of provider health."""
+
+    async def go():
+        with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
+            store = LocalFileStore(name="t_lazy_embedding_health", embedding_store="")
+            embedding_store = HealthCountingEmbeddingStore()
+            store.embedding_store = embedding_store
+            await store.start()
+
+            assert embedding_store.health_calls == 0
+            assert store.embedding_store is embedding_store
+            await store.close()
+
+    run(go())
+
+
 def test_load_rebuilds_keyword_index_from_persisted_chunks_when_missing():
     """Loading persisted chunks repairs a missing keyword index."""
 
@@ -150,6 +199,101 @@ def test_load_rebuilds_keyword_index_from_persisted_chunks_when_missing():
     run(go())
 
 
+def test_keyword_sync_rebuilds_when_backend_only_exposes_matching_count():
+    """Matching counts cannot prove that a backend contains the expected IDs."""
+
+    async def go():
+        store = LocalFileStore(name="t_count_only_keyword", embedding_store="")
+        store.file_chunks = {
+            "expected": chunk("expected", "expected.md", "expected content"),
+        }
+        keyword_index = CountOnlyKeywordIndex(n_docs=1)
+        store.keyword_index = keyword_index
+
+        await store._sync_keyword_index_from_chunks()
+
+        assert keyword_index.reset_docs == {"expected": "expected content"}
+
+    run(go())
+
+
+def test_chunk_persistence_uses_compact_embedding_and_round_trips():
+    """Chunk persistence avoids JSON float lists while preserving float16 vectors."""
+
+    async def go():
+        with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
+            store = LocalFileStore(name="t_compact_embedding", embedding_store="")
+            await store.start()
+            original = chunk("a", "a.md", "alpha text", source="test")
+            original.embedding = np.array([0.25, -1.5, 3.0], dtype=np.float16)
+            store.file_chunks[original.id] = original
+            await store.dump()
+
+            payload = json.loads(next(read_jsonl_zst(store.chunks_path)))
+            assert "embedding" not in payload
+            assert isinstance(payload["_embedding_f16_b64"], str)
+            assert base64.b64decode(payload["_embedding_f16_b64"]) == original.embedding.astype("<f2").tobytes()
+
+            store.file_chunks.clear()
+            await store.load()
+            restored = store.file_chunks[original.id]
+            np.testing.assert_array_equal(restored.embedding, original.embedding)
+            assert restored.embedding.dtype == np.float16
+            assert restored.metadata == {"source": "test"}
+            await store.close()
+
+    run(go())
+
+
+def test_vector_search_batches_candidates_and_preserves_stable_ties(monkeypatch):
+    """Local vector search limits matrix size and retains insertion order for ties."""
+
+    async def go():
+        store = LocalFileStore(name="t_vector_batches", embedding_store="")
+        store.embedding_store = FakeEmbeddingStore()
+        for index in range(5):
+            candidate = chunk(str(index), f"{index}.md", "alpha")
+            candidate.embedding = np.array([1.0, 0.0], dtype=np.float16)
+            store.file_chunks[candidate.id] = candidate
+
+        batch_sizes = []
+        original_similarity = local_file_store_module.batch_cosine_similarity
+
+        def recording_similarity(query, matrix):
+            batch_sizes.append(len(matrix))
+            return original_similarity(query, matrix)
+
+        monkeypatch.setattr(local_file_store_module, "_VECTOR_SEARCH_BATCH_SIZE", 2)
+        monkeypatch.setattr(local_file_store_module, "batch_cosine_similarity", recording_similarity)
+
+        results = await store.vector_search("alpha", 3, {})
+
+        assert batch_sizes == [2, 2, 1]
+        assert [result.id for result in results] == ["0", "1", "2"]
+
+    run(go())
+
+
+def test_chunk_persistence_loads_legacy_json_embedding_list():
+    """Existing indexes with JSON float-list embeddings remain readable."""
+
+    async def go():
+        with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
+            store = LocalFileStore(name="t_legacy_embedding", embedding_store="")
+            await store.start()
+            original = chunk("legacy", "legacy.md", "legacy text")
+            original.embedding = np.array([0.5, 1.5], dtype=np.float16)
+            write_jsonl_zst(store.chunks_path, [original.model_dump_json()])
+
+            await store.load()
+            restored = store.file_chunks[original.id]
+            np.testing.assert_array_equal(restored.embedding, original.embedding)
+            assert restored.embedding.dtype == np.float16
+            await store.close()
+
+    run(go())
+
+
 def test_same_chunk_id_with_changed_text_gets_new_embedding():
     """Changing a chunk text refreshes its embedding."""
 

tests/unit/test_keyword_index.py

@@ -174,6 +174,7 @@ def test_add_single_doc():
             assert bm25.total_len == 2  # 'hello', 'world'
             assert bm25.avg_len == 2.0
             assert set(bm25.vocab) == {"hello", "world"}
+            assert bm25.document_ids == {"d1"}
             assert "d1" in bm25.doc_meta
             assert bm25.doc_meta["d1"]["len"] == 2
             await bm25.close()