diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 2c875eb5e46..48539efc158 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -93,6 +93,40 @@ class VectorStoreIndexRegistry: return vector_stores_from_db +def _build_managed_vector_store_from_config_entry(vector_store_config: dict) -> LiteLLM_ManagedVectorStore: + """Construct a LiteLLM_ManagedVectorStore from one config.yaml entry. + + Extracted from load_vector_stores_from_config to keep that method under the + C901 complexity budget; also enforces the two required params (vector_store_id + and custom_llm_provider) in one place with clear error messages. + """ + litellm_vector_store_config = LiteLLM_VectorStoreConfig(**vector_store_config) + vector_store_name = litellm_vector_store_config.get("vector_store_name") + vector_store_litellm_params: dict[str, Any] = litellm_vector_store_config.get("litellm_params") or {} + + vector_store_id = vector_store_litellm_params.get("vector_store_id") + if vector_store_id is None: + raise ValueError( + f"vector_store_id is required for initializing vector store, got vector_store_id={vector_store_id}" + ) + custom_llm_provider = vector_store_litellm_params.get("custom_llm_provider") + if custom_llm_provider is None: + raise ValueError( + f"custom_llm_provider is required for initializing vector store, got custom_llm_provider={custom_llm_provider}" + ) + + return LiteLLM_ManagedVectorStore( + vector_store_id=vector_store_id, + custom_llm_provider=custom_llm_provider, + litellm_params=vector_store_litellm_params, + vector_store_name=vector_store_name, + vector_store_description=vector_store_litellm_params.get("vector_store_description"), + vector_store_metadata=vector_store_litellm_params.get("vector_store_metadata"), + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + + class VectorStoreRegistry: def __init__(self, vector_stores: list[LiteLLM_ManagedVectorStore] | None = None): # Never bind a mutable default: the previous `= []` default caused every @@ -394,6 +428,14 @@ class VectorStoreRegistry: vector_store_ids.extend(tool["vector_store_ids"]) return vector_store_ids + def _drop_config_entries_removed_from_new_config(self, new_config_ids: set[str]) -> None: + """Evict any previously-config-loaded store whose id isn't in the new config, + so config removals stop being protected from list-endpoint reconciliation.""" + for stale_id in self.config_loaded_vector_store_ids - new_config_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) + def load_vector_stores_from_config(self, vector_stores_config: list[dict]): """ Loads vector stores from the litellm proxy config.yaml. @@ -407,44 +449,14 @@ class VectorStoreRegistry: (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) + self._drop_config_entries_removed_from_new_config(new_config_ids) for vector_store_config in vector_stores_config: - # cast to VectorStoreConfig - litellm_vector_store_config = LiteLLM_VectorStoreConfig(**vector_store_config) - vector_store_name = litellm_vector_store_config.get("vector_store_name") - vector_store_litellm_params: dict[str, Any] = litellm_vector_store_config.get("litellm_params") or {} - - vector_store_id = vector_store_litellm_params.get("vector_store_id") - if vector_store_id is None: - raise ValueError( - f"vector_store_id is required for initializing vector store, got vector_store_id={vector_store_id}" - ) - custom_llm_provider = vector_store_litellm_params.get("custom_llm_provider") - if custom_llm_provider is None: - raise ValueError( - f"custom_llm_provider is required for initializing vector store, got custom_llm_provider={custom_llm_provider}" - ) - - litellm_managed_vector_store = LiteLLM_ManagedVectorStore( - vector_store_id=vector_store_id, - custom_llm_provider=custom_llm_provider, - litellm_params=vector_store_litellm_params, - vector_store_name=vector_store_name, - vector_store_description=vector_store_litellm_params.get("vector_store_description"), - vector_store_metadata=vector_store_litellm_params.get("vector_store_metadata"), - created_at=datetime.now(timezone.utc), - updated_at=datetime.now(timezone.utc), - ) + managed = _build_managed_vector_store_from_config_entry(vector_store_config) # 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) + self.add_vector_store_to_registry(managed) + self.config_loaded_vector_store_ids.add(managed["vector_store_id"]) verbose_logger.debug( "all loaded vector stores = %s", diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 49ca2db5059..5b195f18cbf 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3162,7 +3162,9 @@ async def test_list_vector_stores_returns_config_sourced_and_leaves_registry_int ), ) - returned_ids = [vs["vector_store_id"] for vs in response.data] + # LiteLLM_ManagedVectorStoreListResponse is a TypedDict — use dict access. + response_data = response["data"] if isinstance(response, dict) else response.data + returned_ids = [vs["vector_store_id"] for vs in response_data] assert "vs_from_config" in returned_ids # And, critically, the registry must NOT have been mutated — # config-loaded stores stay resident for downstream routing.