From 55905ed59566b9defd63e7710702c7a464c8b200 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 28 Aug 2026 11:31:11 +0800 Subject: [PATCH] fix(embedding): retry vector space changes per request --- docs/en/plugin_management.md | 7 +- docs/zh/plugin_management.md | 6 +- plugins/auto-fin/README.md | 6 +- plugins/auto-fin/README_ZH.md | 6 +- .../embedding_store/local_embedding_store.py | 62 ++++++++------- tests/unit/test_local_embedding_store.py | 79 ++++++++++++++----- 6 files changed, 98 insertions(+), 68 deletions(-) diff --git a/docs/en/plugin_management.md b/docs/en/plugin_management.md index 46a729be..31b8400a 100644 --- a/docs/en/plugin_management.md +++ b/docs/en/plugin_management.md @@ -173,11 +173,8 @@ curl -s http://127.0.0.1:2333/auto_fin \ When the application uses an MCP service, service-enabled plugin Jobs appear as MCP tools instead. -To add the plugin to another application config, select it explicitly: - -```bash -reme start config=demo plugins='["auto-fin"]' -``` +Custom application configs must provide the plugin's runtime dependencies, including an `agent_wrapper.default` and +the `search` and `read` Jobs used by Auto Fin. ## Uninstall a plugin diff --git a/docs/zh/plugin_management.md b/docs/zh/plugin_management.md index 0ab75aa8..f9f735f8 100644 --- a/docs/zh/plugin_management.md +++ b/docs/zh/plugin_management.md @@ -167,11 +167,7 @@ curl -s http://127.0.0.1:2333/auto_fin \ 当应用使用 MCP service 时，允许对外服务的插件 Job 会显示为 MCP tool。 -如果需要将插件叠加到其他应用配置，则显式选择该配置： - -```bash -reme start config=demo plugins='["auto-fin"]' -``` +自定义应用配置需要提供插件的运行依赖，包括 `agent_wrapper.default`，以及 Auto Fin 使用的 `search` 和 `read` Jobs。 ## 卸载插件 diff --git a/plugins/auto-fin/README.md b/plugins/auto-fin/README.md index 804e0fdb..8012ce81 100644 --- a/plugins/auto-fin/README.md +++ b/plugins/auto-fin/README.md @@ -59,11 +59,7 @@ reme start plugins='["auto-fin"]' \ service.backend=http ``` -To add Auto Fin to another application instead, select that config explicitly, for example: - -```bash -reme start config=demo plugins='["auto-fin"]' -``` +Custom application configs must provide `agent_wrapper.default` and the `search` and `read` Jobs used by Auto Fin. ## Pipeline diff --git a/plugins/auto-fin/README_ZH.md b/plugins/auto-fin/README_ZH.md index a100804a..e7811bb4 100644 --- a/plugins/auto-fin/README_ZH.md +++ b/plugins/auto-fin/README_ZH.md @@ -54,11 +54,7 @@ reme start plugins='["auto-fin"]' \ service.backend=http ``` -如果需要将 Auto Fin 叠加到其他应用，则显式选择相应配置，例如： - -```bash -reme start config=demo plugins='["auto-fin"]' -``` +自定义应用配置需要提供 `agent_wrapper.default`，以及 Auto Fin 使用的 `search` 和 `read` Jobs。 ## 流程 diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 795a81eb..5c9b557b 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -104,12 +104,25 @@ 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: - await self._fill_misses(misses, results, **kwargs) - return results + for attempt in range(1, _MAX_VECTOR_SPACE_ATTEMPTS + 1): + await self._sync_cache_space() + vector_space_id = self._cache_space + results, misses = self._partition_by_cache(texts) + stable = not misses or await self._fill_misses(misses, results, vector_space_id, **kwargs) + if stable and vector_space_id == self.vector_space_id == self._cache_space: + return results + if attempt == _MAX_VECTOR_SPACE_ATTEMPTS: + self.logger.warning( + f"Embedding vector space kept changing while computing a request; " + f"discarding all result(s) after {attempt} attempts", + ) + else: + self.logger.info( + f"Embedding vector space changed while computing a request; " + f"discarding all result(s) and retrying ({attempt}/{_MAX_VECTOR_SPACE_ATTEMPTS})", + ) + return [None] * len(texts) # -- Batching -- @@ -125,31 +138,26 @@ class LocalEmbeddingStore(BaseEmbeddingStore): misses.append((idx, text, key)) return results, misses - async def _fill_misses(self, misses: list[Miss], results: list[np.ndarray | None], **kwargs) -> None: + async def _fill_misses( + self, + misses: list[Miss], + results: list[np.ndarray | None], + vector_space_id: str, + **kwargs, + ) -> bool: + """Fill every miss only while the request remains in one vector space.""" size = self.max_batch_size for start in range(0, len(misses), size): + if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space: + return False batch = misses[start : start + size] - for attempt in range(1, _MAX_VECTOR_SPACE_ATTEMPTS + 1): - await self._sync_cache_space() - vector_space_id = self._cache_space - computed = await self._compute_batch(batch, **kwargs) - if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space: - if attempt == _MAX_VECTOR_SPACE_ATTEMPTS: - self.logger.warning( - f"Embedding vector space kept changing while computing a batch; " - f"discarding {len(computed)} stale result(s) after {attempt} attempts", - ) - else: - self.logger.info( - f"Embedding vector space changed while computing a batch; " - f"discarding {len(computed)} stale result(s) and retrying " - f"({attempt}/{_MAX_VECTOR_SPACE_ATTEMPTS})", - ) - continue - for idx, key, emb in computed: - results[idx] = emb - self._cache_put(key, emb) - break + computed = await self._compute_batch(batch, **kwargs) + if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space: + return False + for idx, key, emb in computed: + results[idx] = emb + self._cache_put(key, emb) + return True async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]: texts = [text for _, text, _ in batch] diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index f0d449d1..757e184b 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -509,41 +509,79 @@ def test_cache_space_is_rechecked_after_async_load(monkeypatch, tmp_path): run(go()) -def test_completed_request_retries_after_vector_space_changes(): - """A request completed by the old provider must not escape into the new vector space.""" +def test_whole_request_retries_after_vector_space_changes_between_batches(): + """Completed batches must be discarded when a later batch changes vector space.""" async def go(): - embedding = OpenAIAsEmbedding(name="t_space_write_race", backend="openai", model="v3", dimensions=2) - store = LocalEmbeddingStore(name="t_local_write_race") + embedding = FakeAsEmbedding() + embedding.vector_space_id = "v3" + store = LocalEmbeddingStore(name="t_local_write_race", max_batch_size=1, enable_cache=False) store.as_embedding = embedding store._cache_space = embedding.vector_space_id calls = 0 - async def compute_after_round_trip(_batch, **_kwargs): + async def switch_during_second_batch(batch, **_kwargs): nonlocal calls calls += 1 - if calls == 1: - embedding.model = FakeProviderModel("v4") - return [(0, "key", np.array([3.0, 0.0], dtype=np.float16))] - return [(0, "key", np.array([4.0, 0.0], dtype=np.float16))] + if calls == 2: + embedding.vector_space_id = "v4" + idx, _text, key = batch[0] + version = 3.0 if calls < 3 else 4.0 + return [(idx, key, np.array([version, 0.0], dtype=np.float16))] - store._compute_batch = compute_after_round_trip - results = [None] - await store._fill_misses([(0, "text", "key")], results) + store._compute_batch = switch_during_second_batch + results = await store.get_embeddings(["first", "second"]) - assert calls == 2 + assert calls == 4 assert store._cache_space == embedding.vector_space_id - np.testing.assert_array_equal(results[0], np.array([4.0, 0.0], dtype=np.float16)) - np.testing.assert_array_equal(store._cache["key"], np.array([4.0, 0.0], dtype=np.float16)) + for result in results: + np.testing.assert_array_equal(result, np.array([4.0, 0.0], dtype=np.float16)) run(go()) -def test_completed_request_stops_retrying_when_vector_space_keeps_changing(): - """Continuous configuration churn must leave the batch empty instead of blocking forever.""" +def test_whole_request_rereads_cache_after_vector_space_changes(monkeypatch, tmp_path): + """A cache hit from the old space must not survive a later provider switch.""" async def go(): - embedding = OpenAIAsEmbedding(name="t_space_write_churn", backend="openai", model="v3", dimensions=2) + monkeypatch.setattr( + LocalEmbeddingStore, + "component_metadata_path", + property(lambda _self: tmp_path), + ) + embedding = FakeAsEmbedding() + embedding.vector_space_id = "v3" + store = LocalEmbeddingStore(name="t_local_cache_race", enable_cache=True) + store.as_embedding = embedding + store._cache_space = embedding.vector_space_id + first_key = store._cache_key("first") + store._cache[first_key] = np.array([3.0, 0.0], dtype=np.float16) + calls = 0 + + async def switch_on_miss(batch, **_kwargs): + nonlocal calls + calls += 1 + if calls == 1: + embedding.vector_space_id = "v4" + version = 3.0 if calls == 1 else 4.0 + return [(idx, key, np.array([version, 0.0], dtype=np.float16)) for idx, _text, key in batch] + + store._compute_batch = switch_on_miss + results = await store.get_embeddings(["first", "second"]) + + assert calls == 2 + for result in results: + np.testing.assert_array_equal(result, np.array([4.0, 0.0], dtype=np.float16)) + + run(go()) + + +def test_whole_request_stops_retrying_when_vector_space_keeps_changing(): + """Continuous configuration churn must discard the whole request instead of blocking forever.""" + + async def go(): + embedding = FakeAsEmbedding() + embedding.vector_space_id = "v3" store = LocalEmbeddingStore(name="t_local_write_churn") store.as_embedding = embedding store._cache_space = embedding.vector_space_id @@ -552,12 +590,11 @@ def test_completed_request_stops_retrying_when_vector_space_keeps_changing(): async def change_space_every_time(_batch, **_kwargs): nonlocal calls calls += 1 - embedding.model = FakeProviderModel(f"v{calls + 3}") + embedding.vector_space_id = f"v{calls + 3}" return [(0, "key", np.array([float(calls), 0.0], dtype=np.float16))] store._compute_batch = change_space_every_time - results = [None] - await store._fill_misses([(0, "text", "key")], results) + results = await store.get_embeddings(["text"]) assert calls == 3 assert results == [None]