From 2a85c36fa9b261517a037ac01dae97717f65b137 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 15 Jul 2026 20:46:36 +0800 Subject: [PATCH] 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 | 19 ++- .../embedding_store/local_embedding_store.py | 4 +- .../components/file_store/local_file_store.py | 106 +++++++++---- .../keyword_index/base_keyword_index.py | 6 + reme/components/keyword_index/bm25_index.py | 8 +- tests/unit/test_as_embedding_lazy.py | 67 ++++++++ tests/unit/test_file_store_consistency.py | 144 ++++++++++++++++++ tests/unit/test_keyword_index.py | 1 + 8 files changed, 320 insertions(+), 35 deletions(-) create mode 100644 tests/unit/test_as_embedding_lazy.py diff --git a/reme/components/as_embedding/__init__.py b/reme/components/as_embedding/__init__.py index 1f897992..2a66167d 100644 --- a/reme/components/as_embedding/__init__.py +++ b/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 diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index adf8676a..bf58d0a4 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/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) diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index c7b17aed..0d5ee0aa 100644 --- a/reme/components/file_store/local_file_store.py +++ b/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(" 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: diff --git a/reme/components/keyword_index/base_keyword_index.py b/reme/components/keyword_index/base_keyword_index.py index bf3f3dc2..bb64b1b3 100644 --- a/reme/components/keyword_index/base_keyword_index.py +++ b/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: diff --git a/reme/components/keyword_index/bm25_index.py b/reme/components/keyword_index/bm25_index.py index 53a3c618..c049f9d6 100644 --- a/reme/components/keyword_index/bm25_index.py +++ b/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.""" diff --git a/tests/unit/test_as_embedding_lazy.py b/tests/unit/test_as_embedding_lazy.py new file mode 100644 index 00000000..dc81f85f --- /dev/null +++ b/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()) diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index 7954d8c6..d7b93153 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/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("