diff --git a/reme/components/as_embedding/__init__.py b/reme/components/as_embedding/__init__.py index 1da9c167..c61033e0 100644 --- a/reme/components/as_embedding/__init__.py +++ b/reme/components/as_embedding/__init__.py @@ -58,7 +58,7 @@ class BaseAsEmbedding(BaseComponent): self.backend or self.credential_cls.__name__, str(self.kwargs.get("model") or ""), str(self.dimensions), - self._endpoint(self.kwargs.get("credential")), + self._configured_endpoint(), ) @property @@ -79,6 +79,21 @@ class BaseAsEmbedding(BaseComponent): return str(value).rstrip("/") return "" + def _configured_endpoint(self) -> str: + """Resolve endpoint defaults without constructing the provider eagerly.""" + credential = self.kwargs.get("credential") + if not isinstance(credential, dict): + return self._endpoint(credential) + fields = getattr(self.credential_cls, "model_fields", {}) + for name in ("base_url", "host"): + value = credential.get(name) + field = fields.get(name) + if name not in credential and field is not None and not field.is_required(): + value = field.get_default(call_default_factory=True) + if value: + return str(value).rstrip("/") + return "" + async def __call__(self, inputs: list[Any], **kwargs) -> list[list[float]]: self._ensure_model() assert self.model is not None diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 3f6d137c..990b276e 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -117,7 +117,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): batch = misses[start : start + size] for idx, key, emb in await self._compute_batch(batch, **kwargs): results[idx] = emb - if vector_space_id == self.vector_space_id: + if vector_space_id == self.vector_space_id == self._cache_space: self._cache_put(key, emb) async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]: @@ -187,16 +187,25 @@ class LocalEmbeddingStore(BaseEmbeddingStore): if space == self._cache_space: return async with self._cache_space_lock: - space = self.vector_space_id - if space == self._cache_space: + while True: + space = self.vector_space_id + if space == self._cache_space: + return + previous = self._cache_space + snapshot = list(self._cache.items()) + if previous and self.enable_cache and snapshot: + await asyncio.to_thread(self._dump_sync, previous, snapshot) + if space != self.vector_space_id: + continue + dimensions = self.dimensions + cache: OrderedDict[str, np.ndarray] = OrderedDict() + if self.enable_cache and self._cache_path(space).exists(): + cache = await asyncio.to_thread(self._load_sync, space, dimensions) + if space != self.vector_space_id: + continue + self._cache = cache + self._cache_space = space return - previous = self._cache_space - if previous and self.enable_cache and self._cache: - await asyncio.to_thread(self._dump_sync, previous) - self._cache.clear() - self._cache_space = space - if self.enable_cache and self._cache_path(space).exists(): - await asyncio.to_thread(self._load_sync, space) def _cache_key(self, text: str) -> str: return hashlib.sha256(text.encode()).hexdigest() @@ -223,40 +232,41 @@ class LocalEmbeddingStore(BaseEmbeddingStore): async def load(self) -> None: self._cache.clear() - self._cache_space = self.vector_space_id - if not self.enable_cache or not self._cache_path(self._cache_space).exists(): - return - await asyncio.to_thread(self._load_sync, self._cache_space) + self._cache_space = "" + await self._sync_cache_space() - def _load_sync(self, vector_space_id: str) -> None: + def _load_sync(self, vector_space_id: str, dimensions: int) -> OrderedDict[str, np.ndarray]: path = self._cache_path(vector_space_id) + cache: OrderedDict[str, np.ndarray] = OrderedDict() try: with np.load(path) as data: for key, emb in zip(data["keys"], data["embeddings"]): - if len(emb) != self.dimensions: + if len(emb) != dimensions: continue - if len(self._cache) >= self.max_cache_size: + if len(cache) >= self.max_cache_size: break - self._cache[str(key)] = emb.astype(np.float16) + cache[str(key)] = emb.astype(np.float16) except Exception: self.logger.exception("Failed to load embedding cache, removing") path.unlink(missing_ok=True) - return - self.logger.info(f"Loaded {len(self._cache)} embeddings from {path}") + return cache + self.logger.info(f"Loaded {len(cache)} embeddings from {path}") + return cache async def dump(self) -> None: await self._sync_cache_space() - if not self.enable_cache or not self._cache: + snapshot = list(self._cache.items()) + if not self.enable_cache or not snapshot: return - await asyncio.to_thread(self._dump_sync, self._cache_space) + await asyncio.to_thread(self._dump_sync, self._cache_space, snapshot) - def _dump_sync(self, vector_space_id: str) -> None: + def _dump_sync(self, vector_space_id: str, cache: list[tuple[str, np.ndarray]]) -> None: path = self._cache_path(vector_space_id) path.parent.mkdir(parents=True, exist_ok=True) - keys = np.array(list(self._cache.keys()), dtype=str) - embeddings = np.stack(list(self._cache.values())) + keys = np.array([key for key, _ in cache], dtype=str) + embeddings = np.stack([embedding for _, embedding in cache]) try: np.savez(path, keys=keys, embeddings=embeddings) - self.logger.info(f"Saved {len(self._cache)} embeddings to {path}") + self.logger.info(f"Saved {len(cache)} embeddings to {path}") except Exception: self.logger.exception("Failed to save embedding cache") diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index fe432a8e..0650f194 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -7,7 +7,7 @@ from types import SimpleNamespace import numpy as np -from reme.components.as_embedding import OpenAIAsEmbedding +from reme.components.as_embedding import DashScopeAsEmbedding, OpenAIAsEmbedding from reme.components.embedding_store.base_embedding_store import BaseEmbeddingStore from reme.components.embedding_store.local_embedding_store import LocalEmbeddingStore from reme.schema import EmbNode @@ -255,6 +255,21 @@ def test_vector_space_id_is_stable_across_lazy_provider_construction(): assert embedding.vector_space_id == before +def test_vector_space_id_resolves_default_endpoint_before_lazy_construction(): + """Credential defaults must not change the cache namespace on the first request.""" + embedding = DashScopeAsEmbedding( + name="t_space_default_endpoint", + model="text-embedding-v3", + dimensions=1024, + credential={"api_key": "test"}, + ) + before = embedding.vector_space_id + + embedding._ensure_model() + + assert embedding.vector_space_id == before + + def test_cache_is_saved_and_restored_per_vector_space(monkeypatch, tmp_path): """Switching models persists the old cache and restores it when switched back.""" @@ -289,6 +304,66 @@ def test_cache_is_saved_and_restored_per_vector_space(monkeypatch, tmp_path): run(go()) +def test_cache_space_is_rechecked_after_async_load(monkeypatch, tmp_path): + """A provider switch during disk I/O must not publish the stale namespace.""" + + async def go(): + monkeypatch.setattr( + LocalEmbeddingStore, + "component_metadata_path", + property(lambda _self: tmp_path), + ) + embedding = OpenAIAsEmbedding(name="t_space_load_race", backend="openai", model="v3", dimensions=2) + store = LocalEmbeddingStore(name="t_local_load_race") + store.as_embedding = embedding + await store.load() + + embedding.model = FakeProviderModel("v4") + v4_space = embedding.vector_space_id + np.savez( + store._cache_path(v4_space), + keys=np.array([store._cache_key("hello")]), + embeddings=np.array([[4.0, 0.0]], dtype=np.float16), + ) + original_to_thread = asyncio.to_thread + + async def switch_during_load(func, *args): + if getattr(func, "__name__", "") == "_load_sync": + embedding.model = FakeProviderModel("v3") + return await original_to_thread(func, *args) + + monkeypatch.setattr(asyncio, "to_thread", switch_during_load) + await store._sync_cache_space() + + assert store._cache_space == embedding.vector_space_id + assert not store._cache + + run(go()) + + +def test_completed_request_only_writes_to_its_active_cache_space(): + """A v3 request must not populate v4 after the provider switches back to v3.""" + + async def go(): + embedding = OpenAIAsEmbedding(name="t_space_write_race", backend="openai", model="v3", dimensions=2) + store = LocalEmbeddingStore(name="t_local_write_race") + store.as_embedding = embedding + store._cache_space = embedding.vector_space_id + + async def compute_after_round_trip(_batch, **_kwargs): + embedding.model = FakeProviderModel("v4") + store._cache_space = embedding.vector_space_id + embedding.model = FakeProviderModel("v3") + return [(0, "key", np.array([3.0, 0.0], dtype=np.float16))] + + store._compute_batch = compute_after_round_trip + await store._fill_misses([(0, "text", "key")], [None]) + + assert "key" not in store._cache + + run(go()) + + def test_start_ignores_cache_file_without_vector_space_tag(monkeypatch, tmp_path): """An unattributable legacy cache is ignored without deleting derived data."""