From 30cdbd4d9e04046836db54f317cd665da995212c Mon Sep 17 00:00:00 2001 From: KK-MCP Date: Tue, 25 Aug 2026 13:27:53 -0700 Subject: [PATCH] address review: reload-safe markers, tests in CI path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - vector_store_registry: load_vector_stores_from_config now reconciles markers on reload — any previously-config-loaded id absent from the new config is dropped from both the registry and the marker set, so stores removed from config.yaml stop being protected on the next reload (per Greptile review). - Use add_vector_store_to_registry (idempotent by id) inside the load path so reloads don't duplicate entries. - Move new tests into tests/test_litellm/vector_stores/, which is already invoked by the "vector-stores / Run tests" shard, fixing the assert-ci-coverage failure. Adds coverage for reload eviction and reload idempotency. --- .../vector_stores/vector_store_registry.py | 23 ++++- .../test_vector_store_registry.py | 84 +++++++++++++++++++ .../test_config_loaded_vector_stores.py | 57 ------------- 3 files changed, 105 insertions(+), 59 deletions(-) delete mode 100644 tests/vector_store_tests/test_config_loaded_vector_stores.py diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index d9e0e9faa09..2a7a0b7edd4 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -392,8 +392,25 @@ class VectorStoreRegistry: def load_vector_stores_from_config(self, vector_stores_config: list[dict]): """ - Loads vector stores from the litellm proxy config.yaml + Loads vector stores from the litellm proxy config.yaml. + + Safe to call on config reload: any previously-config-loaded store whose + id is not in the new config is removed from the registry and the marker + set, so a store removed from config.yaml stops being protected from the + list-endpoint reconciliation on the very next reload. """ + new_config_ids: set[str] = { + (cfg.get("litellm_params") or {}).get("vector_store_id") + for cfg in vector_stores_config + } + new_config_ids.discard(None) # ids missing here will re-raise below with a clearer error + + stale_ids: Final = self.config_loaded_vector_store_ids - new_config_ids + for stale_id in stale_ids: + # delete_vector_store_from_registry also discards from + # config_loaded_vector_store_ids, keeping the two in sync. + self.delete_vector_store_from_registry(stale_id) + for vector_store_config in vector_stores_config: # cast to VectorStoreConfig litellm_vector_store_config = LiteLLM_VectorStoreConfig(**vector_store_config) @@ -421,7 +438,9 @@ class VectorStoreRegistry: created_at=datetime.now(timezone.utc), updated_at=datetime.now(timezone.utc), ) - self.vector_stores.append(litellm_managed_vector_store) + # add_vector_store_to_registry is idempotent by id, so a reload + # that re-lists an existing store won't create duplicates. + self.add_vector_store_to_registry(litellm_managed_vector_store) self.config_loaded_vector_store_ids.add(vector_store_id) verbose_logger.debug( diff --git a/tests/test_litellm/vector_stores/test_vector_store_registry.py b/tests/test_litellm/vector_stores/test_vector_store_registry.py index f19c3706845..60c44516e64 100644 --- a/tests/test_litellm/vector_stores/test_vector_store_registry.py +++ b/tests/test_litellm/vector_stores/test_vector_store_registry.py @@ -182,3 +182,87 @@ def test_search_uses_registry_credentials(): assert getattr(called_params, "aws_region_name") == "us-east-1" finally: litellm.vector_store_registry = original_registry + + +# --------------------------------------------------------------------------- +# config-loaded marker bookkeeping (fixes: /vector_store/list used to prune +# config.yaml-sourced entries because they have no DB row). +# --------------------------------------------------------------------------- + + +def _cfg_entry(vs_id: str) -> dict: + return { + "vector_store_name": f"name-{vs_id}", + "litellm_params": { + "vector_store_id": vs_id, + "custom_llm_provider": "bedrock", + }, + } + + +def test_load_from_config_tracks_marker_ids(): + registry = VectorStoreRegistry() + assert registry.config_loaded_vector_store_ids == set() + + registry.load_vector_stores_from_config([_cfg_entry("vs_a"), _cfg_entry("vs_b")]) + + assert {vs["vector_store_id"] for vs in registry.vector_stores} == {"vs_a", "vs_b"} + assert registry.config_loaded_vector_store_ids == {"vs_a", "vs_b"} + + +def test_delete_drops_config_marker(): + """An explicit delete of a config-loaded id must also clear the marker + so a later reload can re-load the entry without being incorrectly protected.""" + registry = VectorStoreRegistry() + registry.load_vector_stores_from_config([_cfg_entry("vs_a")]) + assert "vs_a" in registry.config_loaded_vector_store_ids + + registry.delete_vector_store_from_registry("vs_a") + + assert registry.vector_stores == [] + assert "vs_a" not in registry.config_loaded_vector_store_ids + + +def test_delete_of_db_only_store_does_not_touch_marker_set(): + registry = VectorStoreRegistry() + registry.load_vector_stores_from_config([_cfg_entry("vs_config")]) + registry.add_vector_store_to_registry( + LiteLLM_ManagedVectorStore( + vector_store_id="vs_from_db", + custom_llm_provider="openai", + vector_store_name="db-store", + litellm_credential_name=None, + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + ) + + registry.delete_vector_store_from_registry("vs_from_db") + + assert [vs["vector_store_id"] for vs in registry.vector_stores] == ["vs_config"] + assert registry.config_loaded_vector_store_ids == {"vs_config"} + + +def test_reload_removes_entries_dropped_from_config(): + """Reloading config with a store removed should evict it from both the + registry and the marker set — otherwise a stale marker would permanently + protect the deleted store from the list-endpoint reconciliation.""" + registry = VectorStoreRegistry() + registry.load_vector_stores_from_config([_cfg_entry("vs_a"), _cfg_entry("vs_b")]) + assert registry.config_loaded_vector_store_ids == {"vs_a", "vs_b"} + + # Simulate config reload with vs_b removed from config.yaml. + registry.load_vector_stores_from_config([_cfg_entry("vs_a")]) + + assert registry.config_loaded_vector_store_ids == {"vs_a"} + assert [vs["vector_store_id"] for vs in registry.vector_stores] == ["vs_a"] + + +def test_reload_is_idempotent_for_unchanged_config(): + """Re-loading the same config must not create duplicate registry entries.""" + registry = VectorStoreRegistry() + registry.load_vector_stores_from_config([_cfg_entry("vs_a")]) + registry.load_vector_stores_from_config([_cfg_entry("vs_a")]) + + assert [vs["vector_store_id"] for vs in registry.vector_stores] == ["vs_a"] + assert registry.config_loaded_vector_store_ids == {"vs_a"} diff --git a/tests/vector_store_tests/test_config_loaded_vector_stores.py b/tests/vector_store_tests/test_config_loaded_vector_stores.py deleted file mode 100644 index 93e75a11948..00000000000 --- a/tests/vector_store_tests/test_config_loaded_vector_stores.py +++ /dev/null @@ -1,57 +0,0 @@ -""" -Tests for the config-sourced vector store fix. - -Regression: `/vector_store/list` used to prune config.yaml-loaded stores from -the in-memory registry because they have no DB row. The fix tracks -`config_loaded_vector_store_ids` on the registry and exempts them from the -"in memory but not in DB" reconciliation branch. -""" -from litellm.vector_stores.vector_store_registry import VectorStoreRegistry - - -def _config_entry(vs_id: str) -> dict: - return { - "vector_store_name": f"name-{vs_id}", - "litellm_params": { - "vector_store_id": vs_id, - "custom_llm_provider": "bedrock", - }, - } - - -def test_load_from_config_tracks_id(): - registry = VectorStoreRegistry() - assert registry.config_loaded_vector_store_ids == set() - - registry.load_vector_stores_from_config([_config_entry("vs_a"), _config_entry("vs_b")]) - - assert {vs["vector_store_id"] for vs in registry.vector_stores} == {"vs_a", "vs_b"} - assert registry.config_loaded_vector_store_ids == {"vs_a", "vs_b"} - - -def test_delete_drops_config_marker(): - """After a config-loaded store is explicitly deleted, the marker must go - with it so a later reload can re-create the entry without confusion.""" - registry = VectorStoreRegistry() - registry.load_vector_stores_from_config([_config_entry("vs_a")]) - assert "vs_a" in registry.config_loaded_vector_store_ids - - registry.delete_vector_store_from_registry("vs_a") - - assert registry.vector_stores == [] - assert "vs_a" not in registry.config_loaded_vector_store_ids - - -def test_delete_of_db_only_store_does_not_touch_marker_set(): - registry = VectorStoreRegistry() - registry.load_vector_stores_from_config([_config_entry("vs_config")]) - # Simulate a DB-sourced entry being added via /vector_store/new - registry.add_vector_store_to_registry( - {"vector_store_id": "vs_from_db", "custom_llm_provider": "openai"} - ) - - registry.delete_vector_store_from_registry("vs_from_db") - - assert [vs["vector_store_id"] for vs in registry.vector_stores] == ["vs_config"] - # Config marker for the untouched entry is intact. - assert registry.config_loaded_vector_store_ids == {"vs_config"}