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