From 57b08e8f7a30b5c504102ee7d0e35dff005ae9c5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 14:36:07 -0700 Subject: [PATCH] fix(vector-stores): preserve config registry entries --- litellm/proxy/_lazy_openapi_snapshot.json | 4 +- .../management_endpoints.py | 10 ++- .../vector_stores/vector_store_registry.py | 10 ++- .../test_vector_store_endpoints.py | 78 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 5 files changed, 96 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 316bfb8cf92..3599b8c8ada 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -38535,7 +38535,7 @@ "paths": { "/v1/vector_store/list": { "get": { - "description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nDatabase is the source of truth - deleted stores are removed from memory, updated stores sync to memory.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)", + "description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nConfig entries remain authoritative; database-backed entries sync from the database.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)", "operationId": "list_vector_stores_v1_vector_store_list_get", "parameters": [ { @@ -38686,7 +38686,7 @@ }, "/vector_store/list": { "get": { - "description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nDatabase is the source of truth - deleted stores are removed from memory, updated stores sync to memory.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)", + "description": "List all available vector stores with optional filtering and pagination.\nCombines both in-memory vector stores and those stored in the database.\nConfig entries remain authoritative; database-backed entries sync from the database.\n\nParameters:\n- page: int - Page number for pagination (default: 1)\n- page_size: int - Number of items per page (default: 100)", "operationId": "list_vector_stores_vector_store_list_get", "parameters": [ { diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index ec266cb87ca..c0b31b258bb 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -369,7 +369,7 @@ async def list_vector_stores( """ List all available vector stores with optional filtering and pagination. Combines both in-memory vector stores and those stored in the database. - Database is the source of truth - deleted stores are removed from memory, updated stores sync to memory. + Config entries remain authoritative; database-backed entries sync from the database. Parameters: - page: int - Page number for pagination (default: 1) @@ -396,6 +396,7 @@ async def list_vector_stores( # Process in-memory vector stores if litellm.vector_store_registry is not None: in_memory_vector_stores: Final = copy.deepcopy(litellm.vector_store_registry.vector_stores) + config_vector_store_ids: Final = litellm.vector_store_registry.config_vector_store_ids vector_stores_to_delete_from_memory: Final[list[str]] = [] @@ -404,8 +405,9 @@ async def list_vector_stores( if not vector_store_id: continue - # If vector store is in memory but NOT in database, it was deleted - if vector_store_id not in db_vector_store_ids: + if vector_store_id in config_vector_store_ids: + vector_store_map[vector_store_id] = vector_store + elif vector_store_id not in db_vector_store_ids: verbose_proxy_logger.info( "Vector store %s exists in memory but not in database - marking for deletion from cache", vector_store_id, @@ -424,7 +426,7 @@ async def list_vector_stores( # 2. Update in-memory registry with database versions (for updates) for vector_store in vector_stores_from_db: vector_store_id = vector_store.get("vector_store_id", None) - if vector_store_id: + if vector_store_id and vector_store_id not in config_vector_store_ids: litellm.vector_store_registry.update_vector_store_in_registry( vector_store_id=vector_store_id, updated_data=vector_store ) diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 5580e867001..cf0fb0787af 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -105,6 +105,7 @@ class VectorStoreRegistry: def __init__(self, vector_stores: list[LiteLLM_ManagedVectorStore] = []): self.vector_stores: list[LiteLLM_ManagedVectorStore] = vector_stores self.vector_store_ids_to_vector_store_map: dict[str, LiteLLM_ManagedVectorStore] = {} + self.config_vector_store_ids: frozenset[str] = frozenset() def _extract_tool_params(self, tool: dict) -> VectorStoreToolParams: """ @@ -342,7 +343,11 @@ class VectorStoreRegistry: # Verify vector store still exists in database (if we have DB access) # This ensures deleted vector stores are removed from cache - if vector_store is not None and prisma_client is not None: + if ( + vector_store is not None + and prisma_client is not None + and vector_store_id not in self.config_vector_store_ids + ): try: # Check if it still exists in database db_vector_store = await ManagedVectorStoresRepository(prisma_client).table.find_unique( @@ -409,7 +414,7 @@ class VectorStoreRegistry: 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: + if not isinstance(vector_store_id, str) or not vector_store_id: raise ValueError( f"vector_store_id is required for initializing vector store, got vector_store_id={vector_store_id}" ) @@ -433,6 +438,7 @@ class VectorStoreRegistry: updated_at=datetime.now(timezone.utc), ) self.vector_stores.append(litellm_managed_vector_store) + self.config_vector_store_ids = self.config_vector_store_ids.union((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 04ec48447b7..33170f97a2b 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 @@ -21,6 +21,7 @@ from litellm.proxy.vector_store_endpoints.endpoints import ( from litellm.proxy.vector_store_endpoints.management_endpoints import ( _check_vector_store_access, create_vector_store_in_db, + list_vector_stores, new_vector_store, ) from litellm.proxy.vector_store_endpoints.utils import ( @@ -799,6 +800,83 @@ async def test_config_loaded_milvus_grpc_connection_is_trusted(): assert result[MILVUS_ADMIN_CONFIGURED_CONNECTION] is True +@pytest.mark.asyncio +async def test_list_vector_stores_preserves_config_source_and_evicts_stale_database_cache(): + registry = VectorStoreRegistry() + registry.load_vector_stores_from_config( + [ + { + "vector_store_name": "configured", + "litellm_params": { + "vector_store_id": "configured", + "custom_llm_provider": "milvus", + "milvus_transport": "grpc", + "api_base": "https://configured-milvus:19530", + "litellm_embedding_model": "team-embedding-alias", + }, + } + ] + ) + registry.add_vector_store_to_registry( + LiteLLM_ManagedVectorStore( + vector_store_id="stale-database-cache", + custom_llm_provider="openai", + ) + ) + + with ( + patch.object( # test-quality-ok: installs the isolated registry used by the endpoint + litellm, "vector_store_registry", registry + ), + patch( # test-quality-ok: models a proxy without persisted vector stores + "litellm.proxy.proxy_server.prisma_client", None + ), + patch( # test-quality-ok: isolates registry synchronization from feature entitlement + "litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user", + new=AsyncMock(), + ), + patch( # test-quality-ok: isolates registry synchronization from per-user access policy + "litellm.proxy.vector_store_endpoints.management_endpoints._check_vector_store_access", + new=AsyncMock(return_value=True), + ), + ): + response = await list_vector_stores( + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert registry.config_vector_store_ids == frozenset(("configured",)) + assert [store["vector_store_id"] for store in response["data"]] == ["configured"] + assert registry.get_litellm_managed_vector_store_from_registry("configured") is not None + assert registry.get_litellm_managed_vector_store_from_registry("stale-database-cache") is None + + +@pytest.mark.asyncio +async def test_db_fallback_does_not_evict_config_source(): + registry = VectorStoreRegistry() + registry.load_vector_stores_from_config( + [ + { + "vector_store_name": "configured", + "litellm_params": { + "vector_store_id": "configured", + "custom_llm_provider": "openai", + }, + } + ] + ) + prisma_client = MagicMock() + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=None) + + result = await registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params={"vector_store_ids": ["configured"]}, + prisma_client=prisma_client, + ) + + assert [store["vector_store_id"] for store in result] == ["configured"] + assert registry.get_litellm_managed_vector_store_from_registry("configured") is not None + prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_awaited() + + def test_admin_persistence_strips_forged_marker_and_adds_server_marker(): params = prepare_milvus_connection_for_persistence( custom_llm_provider="milvus/probe", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e668645a4ec..d2693acc42b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -20285,7 +20285,7 @@ export interface paths { * List Vector Stores * @description List all available vector stores with optional filtering and pagination. * Combines both in-memory vector stores and those stored in the database. - * Database is the source of truth - deleted stores are removed from memory, updated stores sync to memory. + * Config entries remain authoritative; database-backed entries sync from the database. * * Parameters: * - page: int - Page number for pagination (default: 1) @@ -21586,7 +21586,7 @@ export interface paths { * List Vector Stores * @description List all available vector stores with optional filtering and pagination. * Combines both in-memory vector stores and those stored in the database. - * Database is the source of truth - deleted stores are removed from memory, updated stores sync to memory. + * Config entries remain authoritative; database-backed entries sync from the database. * * Parameters: * - page: int - Page number for pagination (default: 1)