From 21057931a9a0a2ae6dc7b3c8194a84d82cb9a941 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:42:16 +0800 Subject: [PATCH] fix(embedding): isolate caches by vector space (#442) * fix(embedding): isolate caches by vector space * fix(embedding): stabilize cache space switching * Revert "fix(embedding): stabilize cache space switching" This reverts commit 74193c9a0aa0725b45ff14a627ae4ce8b7beb4c5. * fix(embedding): include resolved OpenAI endpoint in cache ID * fix(embedding): stabilize cache space switching * fix(embedding): isolate Ollama endpoint caches --- reme/components/as_embedding/__init__.py | 90 ++++++ .../embedding_store/local_embedding_store.py | 97 ++++-- tests/unit/test_as_embedding_lazy.py | 10 +- tests/unit/test_local_embedding_store.py | 305 +++++++++++++++++- 4 files changed, 474 insertions(+), 28 deletions(-) diff --git a/reme/components/as_embedding/__init__.py b/reme/components/as_embedding/__init__.py index 2a66167d..908249ef 100644 --- a/reme/components/as_embedding/__init__.py +++ b/reme/components/as_embedding/__init__.py @@ -1,5 +1,7 @@ """AgentScope embedding model wrappers.""" +import hashlib +import os from typing import Any from agentscope.credential import ( @@ -23,6 +25,9 @@ class BaseAsEmbedding(BaseComponent): component_type = ComponentEnum.AS_EMBEDDING credential_cls: type[CredentialBase] + endpoint_fields = ("base_url", "host") + endpoint_env: str | None = None + default_endpoint = "" def __init__(self, **kwargs) -> None: super().__init__(**kwargs) @@ -38,6 +43,78 @@ class BaseAsEmbedding(BaseComponent): raise RuntimeError("Embedding dimensions are required before provider initialization.") return int(dimensions) + @property + def vector_space(self) -> tuple[str, ...]: + """Return the fields that make vectors from two embedding setups incompatible. + + The wrapper identifies the provider consistently before and after lazy model + construction. Model details come from the live provider when one has been + injected at runtime, otherwise they come from the configured kwargs. + """ + if self.model is not None: + return ( + self.backend or self.credential_cls.__name__, + str(getattr(self.model, "model", self.kwargs.get("model") or "")), + str(self.dimensions), + self._model_endpoint(), + ) + return ( + self.backend or self.credential_cls.__name__, + str(self.kwargs.get("model") or ""), + str(self.dimensions), + self._configured_endpoint(), + ) + + @property + def vector_space_id(self) -> str: + """Return a short digest of :attr:`vector_space` for naming persisted vectors. + + Cache consumers use this digest to avoid reusing vectors produced by a + different embedding setup. + """ + return hashlib.sha256("\x1f".join(self.vector_space).encode()).hexdigest()[:12] + + @staticmethod + def _endpoint(credential: Any) -> str: + """Read the provider endpoint from a credential object or a raw kwargs dict.""" + for field in ("base_url", "host"): + value = credential.get(field) if isinstance(credential, dict) else getattr(credential, field, None) + if value: + return str(value).rstrip("/") + return "" + + def _configured_endpoint(self) -> str: + """Resolve the provider endpoint without constructing it eagerly.""" + credential = self.kwargs.get("credential") + if isinstance(credential, dict): + fields = getattr(self.credential_cls, "model_fields", {}) + for name in self.endpoint_fields: + if name in credential: + value = credential[name] + if value is not None: + return str(value).rstrip("/") + continue + field = fields.get(name) + if field is not None and not field.is_required(): + value = field.get_default(call_default_factory=True) + if value: + return str(value).rstrip("/") + else: + endpoint = self._endpoint(credential) + if endpoint: + return endpoint + if self.endpoint_env: + endpoint = os.environ.get(self.endpoint_env) + if endpoint: + return endpoint.rstrip("/") + return self.default_endpoint + + def _model_endpoint(self) -> str: + """Return the endpoint used by the constructed provider.""" + assert self.model is not None + endpoint = self._endpoint(getattr(self.model, "credential", self.kwargs.get("credential"))) + return endpoint or self._configured_endpoint() + async def __call__(self, inputs: list[Any], **kwargs) -> list[list[float]]: self._ensure_model() assert self.model is not None @@ -78,6 +155,16 @@ class OpenAIAsEmbedding(BaseAsEmbedding): """OpenAI embedding model wrapper.""" credential_cls = OpenAICredential + endpoint_fields = ("base_url",) + endpoint_env = "OPENAI_BASE_URL" + default_endpoint = "https://api.openai.com/v1" + + def _model_endpoint(self) -> str: + """Read the endpoint the OpenAI client actually resolved.""" + assert self.model is not None + client = getattr(self.model, "client", None) + base_url = getattr(client, "base_url", None) + return str(base_url).rstrip("/") if base_url is not None else super()._model_endpoint() @R.register("dashscope") @@ -106,6 +193,9 @@ class OllamaAsEmbedding(BaseAsEmbedding): """Ollama embedding model wrapper.""" credential_cls = OllamaCredential + endpoint_fields = ("host",) + endpoint_env = "OLLAMA_HOST" + default_endpoint = "http://127.0.0.1:11434" __all__ = [ diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 052a0851..990b276e 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -35,7 +35,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore): self.enable_cache = enable_cache self.cache_version = cache_version self._cache: OrderedDict[str, np.ndarray] = OrderedDict() - self._key_suffix: bytes = b"" + self._cache_space: str = "" + self._cache_space_lock = asyncio.Lock() @property def dimensions(self) -> int: @@ -43,13 +44,25 @@ class LocalEmbeddingStore(BaseEmbeddingStore): assert self.as_embedding is not None, "embedding component not bound" return self.as_embedding.dimensions + @property + def vector_space_id(self) -> str: + """Return the digest of the vector space the bound provider currently produces.""" + assert self.as_embedding is not None, "embedding component not bound" + return self.as_embedding.vector_space_id + @property def cache_path(self) -> Path: - """Return the path to the disk cache file.""" - return self.component_metadata_path / f"{self.name}_{self.cache_version}.npz" + """Return the disk cache file for the current vector space. + + Each vector space owns its own file, so switching the embedding model cannot + read or overwrite vectors that belong to a different model. + """ + return self._cache_path(self.vector_space_id) + + def _cache_path(self, vector_space_id: str) -> Path: + return self.component_metadata_path / f"{self.name}_{self.cache_version}_{vector_space_id}.npz" async def _start(self) -> None: - self._key_suffix = f"|{self.dimensions}".encode() await self.load() async def _close(self) -> None: @@ -76,6 +89,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): # -- Public API -- async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]: + await self._sync_cache_space() texts = [self._truncate(t) for t in input_text] results, misses = self._partition_by_cache(texts) if misses: @@ -97,12 +111,14 @@ class LocalEmbeddingStore(BaseEmbeddingStore): return results, misses async def _fill_misses(self, misses: list[Miss], results: list[np.ndarray | None], **kwargs) -> None: + vector_space_id = self._cache_space size = self.max_batch_size 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) + 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]]: texts = [text for _, text, _ in batch] @@ -165,8 +181,34 @@ class LocalEmbeddingStore(BaseEmbeddingStore): # -- Cache -- + async def _sync_cache_space(self) -> None: + """Persist the previous space and restore the newly active space.""" + space = self.vector_space_id + if space == self._cache_space: + return + async with self._cache_space_lock: + 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 + def _cache_key(self, text: str) -> str: - return hashlib.sha256(text.encode() + self._key_suffix).hexdigest() + return hashlib.sha256(text.encode()).hexdigest() def _cache_get(self, key: str) -> np.ndarray | None: if not self.enable_cache or key not in self._cache: @@ -190,36 +232,41 @@ class LocalEmbeddingStore(BaseEmbeddingStore): async def load(self) -> None: self._cache.clear() - if not self.enable_cache or not self.cache_path.exists(): - return - await asyncio.to_thread(self._load_sync) + self._cache_space = "" + await self._sync_cache_space() - def _load_sync(self) -> 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(self.cache_path) as data: + 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") - self.cache_path.unlink(missing_ok=True) - return - self.logger.info(f"Loaded {len(self._cache)} embeddings from {self.cache_path}") + path.unlink(missing_ok=True) + return cache + self.logger.info(f"Loaded {len(cache)} embeddings from {path}") + return cache async def dump(self) -> None: - if not self.enable_cache or not self._cache: + await self._sync_cache_space() + snapshot = list(self._cache.items()) + if not self.enable_cache or not snapshot: return - await asyncio.to_thread(self._dump_sync) + await asyncio.to_thread(self._dump_sync, self._cache_space, snapshot) - def _dump_sync(self) -> None: - self.cache_path.parent.mkdir(parents=True, exist_ok=True) - keys = np.array(list(self._cache.keys()), dtype=str) - embeddings = np.stack(list(self._cache.values())) + 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([key for key, _ in cache], dtype=str) + embeddings = np.stack([embedding for _, embedding in cache]) try: - np.savez(self.cache_path, keys=keys, embeddings=embeddings) - self.logger.info(f"Saved {len(self._cache)} embeddings to {self.cache_path}") + np.savez(path, keys=keys, embeddings=embeddings) + 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_as_embedding_lazy.py b/tests/unit/test_as_embedding_lazy.py index dc81f85f..14a4bb27 100644 --- a/tests/unit/test_as_embedding_lazy.py +++ b/tests/unit/test_as_embedding_lazy.py @@ -51,14 +51,22 @@ def test_provider_is_constructed_once_on_first_call(): async def go(): FakeModel.constructions = 0 - embedding = LazyAsEmbedding(dimensions=3, credential={"token": "test"}, parameters={"mode": "test"}) + embedding = LazyAsEmbedding( + backend="fake", + model="fake-model", + dimensions=3, + credential={"token": "test"}, + parameters={"mode": "test"}, + ) await embedding.start() assert embedding.model is None assert embedding.dimensions == 3 assert FakeModel.constructions == 0 + vector_space_id = embedding.vector_space_id assert await embedding(["first"]) == [[0.0, 0.0, 0.0]] + assert embedding.vector_space_id == vector_space_id assert await embedding(["second"]) == [[0.0, 0.0, 0.0]] assert FakeModel.constructions == 1 diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 98d6b93a..8ff851f4 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -1,11 +1,13 @@ -"""Regression tests for LocalEmbeddingStore dimension handling.""" +"""Regression tests for LocalEmbeddingStore dimension and vector space handling.""" # pylint: disable=protected-access import asyncio +from types import SimpleNamespace import numpy as np +from reme.components.as_embedding import DashScopeAsEmbedding, OllamaAsEmbedding, 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 @@ -15,6 +17,7 @@ class FakeAsEmbedding: """Fake AgentScope embedding component.""" dimensions = 2 + vector_space_id = "fakespace000" async def __call__(self, texts: list[str], **_kwargs): return [[1.0] if text == "bad" else [1.0, 0.0] for text in texts] @@ -24,11 +27,26 @@ class BadHealthAsEmbedding: """Fake provider whose health probe returns the wrong dimension.""" dimensions = 2 + vector_space_id = "fakespace000" async def __call__(self, _texts: list[str], **_kwargs): return [[1.0]] +class FakeProviderModel: + """Stand-in for a constructed AgentScope embedding model object.""" + + def __init__( + self, + model: str, + dimensions: int = 2, + base_url: str = "https://api.openai.com/v1", + ): + self.model = model + self.dimensions = dimensions + self.credential = SimpleNamespace(base_url=base_url) + + class InsufficientQuotaError(Exception): """OpenAI-compatible quota error used without importing the provider SDK.""" @@ -86,7 +104,6 @@ def test_compute_batch_rejects_embeddings_with_wrong_dimension(): async def go(): store = LocalEmbeddingStore(name="t_local_embedding_dim") store.as_embedding = FakeAsEmbedding() - store._key_suffix = f"|{store.dimensions}".encode() results = await store._compute_batch( [ @@ -178,3 +195,287 @@ def test_insufficient_quota_does_not_retry_without_opt_in(monkeypatch): assert not sleeps run(go()) + + +def test_vector_space_id_separates_models_of_equal_dimension(): + """Two models of the same width must not claim the same vector space.""" + common = {"backend": "openai", "dimensions": 1024, "credential": {"base_url": "https://example.com/v1"}} + v3 = OpenAIAsEmbedding(name="t_space_v3", model="text-embedding-v3", **common) + v4 = OpenAIAsEmbedding(name="t_space_v4", model="text-embedding-v4", **common) + + assert v3.dimensions == v4.dimensions + assert v3.vector_space_id != v4.vector_space_id + + +def test_vector_space_id_separates_endpoints_of_one_model_name(): + """The same model name served by two endpoints is two vector spaces.""" + common = {"backend": "openai", "model": "text-embedding-v4", "dimensions": 1024} + official = OpenAIAsEmbedding(name="t_space_a", credential={"base_url": "https://example.com/v1"}, **common) + self_hosted = OpenAIAsEmbedding(name="t_space_b", credential={"base_url": "http://127.0.0.1:8000/v1"}, **common) + + assert official.vector_space_id != self_hosted.vector_space_id + + +def test_vector_space_id_ignores_trailing_slash_and_api_key(): + """Cosmetic and secret credential changes must not invalidate stored vectors.""" + common = {"backend": "openai", "model": "text-embedding-v4", "dimensions": 1024} + first = OpenAIAsEmbedding( + name="t_space_c", + credential={"base_url": "https://example.com/v1", "api_key": "key-one"}, + **common, + ) + second = OpenAIAsEmbedding( + name="t_space_d", + credential={"base_url": "https://example.com/v1/", "api_key": "key-two"}, + **common, + ) + + assert first.vector_space_id == second.vector_space_id + + +def test_vector_space_id_follows_a_model_swapped_in_after_start(): + """A provider replaced at runtime must win over the original kwargs.""" + embedding = OpenAIAsEmbedding(name="t_space_swap", backend="openai", model="v3", dimensions=2) + before = embedding.vector_space_id + + # Mirrors Application.update_component("as_embedding", "default", model=). + embedding.model = FakeProviderModel("v4") + + assert embedding.vector_space_id != before + + +def test_vector_space_id_is_stable_across_lazy_provider_construction(): + """Constructing the configured provider must not look like a model switch.""" + embedding = OpenAIAsEmbedding( + name="t_space_lazy", + backend="openai", + model="v3", + dimensions=2, + credential={"base_url": "https://example.com/v1"}, + ) + before = embedding.vector_space_id + + embedding.model = FakeProviderModel("v3", base_url="https://example.com/v1") + + 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_openai_vector_space_id_uses_sdk_resolved_endpoint(monkeypatch): + """OPENAI_BASE_URL must separate caches and stay stable after lazy construction.""" + clients = [] + ids = [] + try: + for endpoint in ("https://provider-a.example/v1", "https://provider-b.example/v1"): + monkeypatch.setenv("OPENAI_BASE_URL", endpoint) + embedding = OpenAIAsEmbedding( + name="t_space_openai_env", + backend="openai", + model="text-embedding-3-small", + dimensions=1536, + credential={"api_key": "test"}, + ) + before = embedding.vector_space_id + + embedding._ensure_model() + clients.append(embedding.model.client) + + assert str(embedding.model.client.base_url).rstrip("/") == endpoint + assert embedding.vector_space_id == before + ids.append(before) + + assert ids[0] != ids[1] + finally: + for client in clients: + run(client.close()) + + +def test_openai_vector_space_id_uses_sdk_default_endpoint(monkeypatch): + """The SDK default URL must not change the namespace on first construction.""" + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + embedding = OpenAIAsEmbedding( + name="t_space_openai_default", + backend="openai", + model="text-embedding-3-small", + dimensions=1536, + credential={"api_key": "test"}, + ) + before = embedding.vector_space_id + + embedding._ensure_model() + try: + assert str(embedding.model.client.base_url).rstrip("/") == "https://api.openai.com/v1" + assert embedding.vector_space_id == before + finally: + run(embedding.model.client.close()) + + +def test_ollama_vector_space_id_uses_sdk_resolved_endpoint(monkeypatch): + """OLLAMA_HOST must separate caches and stay stable after lazy construction.""" + ids = [] + for endpoint in ("http://provider-a.example:11434", "http://provider-b.example:11434"): + monkeypatch.setenv("OLLAMA_HOST", endpoint) + embedding = OllamaAsEmbedding( + name="t_space_ollama_env", + backend="ollama", + model="nomic-embed-text", + dimensions=768, + credential={}, + ) + before = embedding.vector_space_id + + embedding._ensure_model() + + assert embedding.model.host is None + assert embedding.vector_space_id == before + ids.append(before) + + assert ids[0] != ids[1] + + +def test_ollama_vector_space_id_uses_sdk_default_endpoint(monkeypatch): + """The Ollama SDK default URL must remain stable after lazy construction.""" + monkeypatch.delenv("OLLAMA_HOST", raising=False) + embedding = OllamaAsEmbedding( + name="t_space_ollama_default", + backend="ollama", + model="nomic-embed-text", + dimensions=768, + credential={}, + ) + before = embedding.vector_space_id + + embedding._ensure_model() + + assert embedding.vector_space[-1] == "http://127.0.0.1:11434" + 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.""" + + async def go(): + monkeypatch.setattr( + LocalEmbeddingStore, + "component_metadata_path", + property(lambda _self: tmp_path), + ) + embedding = OpenAIAsEmbedding(name="t_space_store", backend="openai", model="v3", dimensions=2) + store = LocalEmbeddingStore(name="t_local_space") + store.as_embedding = embedding + await store.load() + + key = store._cache_key("hello") + store._cache_put(key, np.array([1.0, 0.0], dtype=np.float16)) + v3_path = store.cache_path + + embedding.model = FakeProviderModel("v4") + await store._sync_cache_space() + + assert store._cache_key("hello") == key + assert store.cache_path != v3_path + assert store._cache_get(store._cache_key("hello")) is None + assert v3_path.exists() + + embedding.model = FakeProviderModel("v3") + await store._sync_cache_space() + + np.testing.assert_array_equal(store._cache_get(key), np.array([1.0, 0.0], dtype=np.float16)) + + 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.""" + + async def go(): + monkeypatch.setattr( + LocalEmbeddingStore, + "component_metadata_path", + property(lambda _self: tmp_path), + ) + store = LocalEmbeddingStore(name="t_local_untagged") + store.as_embedding = FakeAsEmbedding() + untagged = tmp_path / f"{store.name}_{store.cache_version}.npz" + untagged.write_bytes(b"vectors from an unknown model") + + await store._start() + + assert untagged.exists() + assert not store._cache + + run(go())