diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 22d27bc3266..e8dd246a5cc 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -104,6 +104,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: """ @@ -338,8 +339,13 @@ class VectorStoreRegistry: break # 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: + # This ensures deleted vector stores are removed from cache. + # Config-defined vector stores are never in the database, so skip the check for them + 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( @@ -427,6 +433,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 | {vector_store_id} verbose_logger.debug( "all loaded vector stores = %s", 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..84e59d424fa 100644 --- a/tests/test_litellm/vector_stores/test_vector_store_registry.py +++ b/tests/test_litellm/vector_stores/test_vector_store_registry.py @@ -8,7 +8,7 @@ from fastapi.testclient import TestClient from datetime import datetime, timezone -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import litellm from litellm.types.vector_stores import LiteLLM_ManagedVectorStore @@ -182,3 +182,39 @@ def test_search_uses_registry_credentials(): assert getattr(called_params, "aws_region_name") == "us-east-1" finally: litellm.vector_store_registry = original_registry + + +@pytest.mark.asyncio +async def test_config_vector_store_survives_db_check(): + """Config-defined vector stores are not in the DB and must not be evicted when a DB is connected.""" + registry = VectorStoreRegistry() + registry.load_vector_stores_from_config( + [ + { + "vector_store_name": "config-kb", + "litellm_params": { + "vector_store_id": "CONFIGKB", + "custom_llm_provider": "bedrock", + }, + } + ] + ) + + db_store = LiteLLM_ManagedVectorStore( + vector_store_id="DBONLY", + custom_llm_provider="bedrock", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + registry.add_vector_store_to_registry(db_store) + + mock_prisma_client = MagicMock() + mock_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": ["CONFIGKB", "DBONLY"]}, + prisma_client=mock_prisma_client, + ) + + assert [vs.get("vector_store_id") for vs in result] == ["CONFIGKB"] + assert any(vs.get("vector_store_id") == "CONFIGKB" for vs in registry.vector_stores) + assert not any(vs.get("vector_store_id") == "DBONLY" for vs in registry.vector_stores)