fix(embedding): stabilize cache space switching

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

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")

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."""