From 99fd96687f15edb1e1f76db431568492ab24674b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 5 Dec 2025 11:46:14 +0530 Subject: [PATCH] Fix vector store configuration synchronization failure --- .../proxy/vector_stores/endpoints.py | 42 +-- .../vector_store_pre_call_hook.py | 15 +- .../vector_stores/vector_store_registry.py | 119 +++++++++ .../test_vector_store_endpoints.py | 243 ++++++++++++++++++ 4 files changed, 400 insertions(+), 19 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py index fdb1dba372f..21933165217 100644 --- a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py +++ b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py @@ -141,28 +141,36 @@ async def list_vector_stores( """ from litellm.proxy.proxy_server import prisma_client - seen_vector_store_ids = set() - try: - # Get in-memory vector stores - in_memory_vector_stores: List[LiteLLM_ManagedVectorStore] = [] - if litellm.vector_store_registry is not None: - in_memory_vector_stores = copy.deepcopy( - litellm.vector_store_registry.vector_stores - ) - - # Get vector stores from database + # Get vector stores from database (source of truth) + # Only return what's in the database to ensure consistency across instances vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db( prisma_client=prisma_client ) + + # Also clean up in-memory registry to remove any deleted vector stores + if litellm.vector_store_registry is not None: + db_vector_store_ids = { + vs.get("vector_store_id") + for vs in vector_stores_from_db + if vs.get("vector_store_id") + } + # Remove any in-memory vector stores that no longer exist in database + vector_stores_to_remove = [] + for vs in litellm.vector_store_registry.vector_stores: + vs_id = vs.get("vector_store_id") + if vs_id and vs_id not in db_vector_store_ids: + vector_stores_to_remove.append(vs_id) + for vs_id in vector_stores_to_remove: + litellm.vector_store_registry.delete_vector_store_from_registry( + vector_store_id=vs_id + ) + verbose_proxy_logger.debug( + f"Removed deleted vector store {vs_id} from in-memory registry" + ) - # Combine in-memory and database vector stores - combined_vector_stores: List[LiteLLM_ManagedVectorStore] = [] - for vector_store in in_memory_vector_stores + vector_stores_from_db: - vector_store_id = vector_store.get("vector_store_id", None) - if vector_store_id not in seen_vector_store_ids: - combined_vector_stores.append(vector_store) - seen_vector_store_ids.add(vector_store_id) + # Use database as single source of truth for listing + combined_vector_stores: List[LiteLLM_ManagedVectorStore] = vector_stores_from_db total_count = len(combined_vector_stores) total_pages = (total_count + page_size - 1) // page_size diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 236935778d6..218581a41ad 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -74,9 +74,20 @@ class VectorStorePreCallHook(CustomLogger): if litellm.vector_store_registry is None: return model, messages, non_default_params + # Get prisma_client for database fallback + prisma_client = None + try: + from litellm.proxy.proxy_server import prisma_client as _prisma_client + prisma_client = _prisma_client + except ImportError: + pass + + # Use database fallback to ensure synchronization across instances vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = ( - litellm.vector_store_registry.pop_vector_stores_to_run( - non_default_params=non_default_params, tools=tools + await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params=non_default_params, + tools=tools, + prisma_client=prisma_client ) ) diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 78c8d7cf2ec..cf0bf89d701 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -233,6 +233,36 @@ class VectorStoreRegistry: return vector_store return None + async def get_litellm_managed_vector_store_from_registry_or_db( + self, vector_store_id: str, prisma_client: Optional[PrismaClient] = None + ) -> Optional[LiteLLM_ManagedVectorStore]: + """ + Returns the vector store from the registry, falling back to database if not found. + This ensures synchronization across multiple instances. + """ + # First check in-memory registry + vector_store = self.get_litellm_managed_vector_store_from_registry(vector_store_id) + if vector_store is not None: + return vector_store + + # Fall back to database if not found in memory + if prisma_client is not None: + try: + vector_stores_from_db = await self._get_vector_stores_from_db( + prisma_client=prisma_client + ) + for db_vector_store in vector_stores_from_db: + if db_vector_store.get("vector_store_id") == vector_store_id: + # Add to in-memory registry for future use + self.add_vector_store_to_registry(vector_store=db_vector_store) + return db_vector_store + except Exception as e: + verbose_logger.debug( + f"Error fetching vector store from database: {str(e)}" + ) + + return None + def get_litellm_managed_vector_store_from_registry_by_name( self, vector_store_name: str ) -> Optional[LiteLLM_ManagedVectorStore]: @@ -289,6 +319,95 @@ class VectorStoreRegistry: return vector_stores_to_run + async def pop_vector_stores_to_run_with_db_fallback( + self, + non_default_params: Dict, + tools: Optional[List[Dict]] = None, + prisma_client: Optional[PrismaClient] = None + ) -> List[LiteLLM_ManagedVectorStore]: + """ + Pops the vector stores to run with their tool parameters merged. + Falls back to database if vector stores are not found in memory. + This ensures synchronization across multiple instances. + + Primary function to use for vector store pre call hook. + + Args: + non_default_params: Parameters dict to pop vector_store_ids from + tools: Optional list of tools to extract vector store params from + prisma_client: Optional database client for fallback lookup + + Returns: + List of vector stores with tool parameters merged into litellm_params + """ + # Pop vector_store_ids from params + vector_store_ids: List[str] = non_default_params.pop("vector_store_ids", None) or [] + + # Extract params from tools and collect IDs + params_by_id = self.get_and_pop_recognised_vector_store_tools( + tools=tools, + vector_store_ids=vector_store_ids + ) + + vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = [] + + for vector_store_id in vector_store_ids: + vector_store = None + + # First check in-memory registry + for vs in self.vector_stores: + if vs.get("vector_store_id") == vector_store_id: + vector_store = vs + 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: + try: + # Check if it still exists in database + db_vector_store = await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store_id} + ) + if db_vector_store is None: + # Vector store was deleted from database, remove from cache + verbose_logger.debug( + f"Vector store {vector_store_id} found in memory but deleted from database, removing from cache" + ) + self.delete_vector_store_from_registry(vector_store_id=vector_store_id) + vector_store = None + except Exception as e: + verbose_logger.debug( + f"Error verifying vector store {vector_store_id} in database: {str(e)}" + ) + + # Fall back to database if not found in memory (or was deleted) + if vector_store is None and prisma_client is not None: + try: + vector_store = await self.get_litellm_managed_vector_store_from_registry_or_db( + vector_store_id=vector_store_id, + prisma_client=prisma_client + ) + except Exception as e: + verbose_logger.debug( + f"Error fetching vector store {vector_store_id} from database: {str(e)}" + ) + + if vector_store is not None: + # Create a copy to avoid modifying the registry + vector_store_copy = vector_store.copy() + + # Merge tool params if they exist + if vector_store_id in params_by_id: + existing_params = vector_store_copy.get("litellm_params", {}) or {} + tool_params_dict = params_by_id[vector_store_id].to_dict() + # Tool params take precedence over existing params + tool_params_dict.update(existing_params) + vector_store_copy["litellm_params"] = tool_params_dict + + vector_stores_to_run.append(vector_store_copy) + + return vector_stores_to_run + def _get_vector_store_ids_from_tool_calls( self, tools: Optional[List[Dict]] = None, vector_store_ids: List[str] = [] ) -> List[str]: 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 badbef42d6c..b98354032fe 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 @@ -1,5 +1,6 @@ import os import sys +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -802,3 +803,245 @@ class TestVectorStoreManagementEndpointsExist: f"Expected endpoint {method} {path} not found in registered routes. " f"Available routes: {app_routes}" ) + + +@pytest.mark.asyncio +async def test_vector_store_synchronization_across_instances(): + """ + Test that vector stores are properly synchronized across multiple instances. + + This test simulates the scenario where: + 1. Instance 1 creates a vector store (writes to DB, updates its own cache) + 2. Instance 2 should be able to find it (via database fallback) + 3. Instance 1 deletes the vector store (removes from DB, updates its own cache) + 4. Instance 2 should not show it in the list (database is source of truth) + """ + from datetime import datetime, timezone + from unittest.mock import AsyncMock, MagicMock + + from litellm.types.vector_stores import ( + LiteLLM_ManagedVectorStore, + VectorStoreDeleteRequest, + ) + from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + + # Simulate two instances with separate in-memory registries + instance_1_registry = VectorStoreRegistry(vector_stores=[]) + instance_2_registry = VectorStoreRegistry(vector_stores=[]) + + # Mock database that both instances share + mock_db_vector_stores = [] + + async def mock_find_unique(where): + """Mock find_unique for checking if vector store exists""" + vector_store_id = where.get("vector_store_id") + for vs in mock_db_vector_stores: + if vs.get("vector_store_id") == vector_store_id: + # Create a simple object that dict() can convert + class MockVectorStore: + def __init__(self, data): + for key, value in data.items(): + setattr(self, key, value) + self._data = data + + def __iter__(self): + return iter(self._data.items()) + return MockVectorStore(vs) + return None + + async def mock_find_many(order=None): + """Mock find_many for listing vector stores""" + # Return objects that can be converted to dict using dict() + # The _get_vector_stores_from_db uses dict(vector_store), so we need to make it work + result = [] + for vs in mock_db_vector_stores: + # Create a simple object that dict() can convert + class MockVectorStore: + def __init__(self, data): + for key, value in data.items(): + setattr(self, key, value) + self._data = data + + def __iter__(self): + return iter(self._data.items()) + result.append(MockVectorStore(vs)) + return result + + async def mock_create(data): + """Mock create for adding vector store to DB""" + vector_store = data.copy() + mock_db_vector_stores.append(vector_store) + mock_obj = MagicMock() + mock_obj.model_dump.return_value = vector_store + for key, value in vector_store.items(): + setattr(mock_obj, key, value) + return mock_obj + + async def mock_delete(where): + """Mock delete for removing vector store from DB""" + vector_store_id = where.get("vector_store_id") + mock_db_vector_stores[:] = [ + vs for vs in mock_db_vector_stores + if vs.get("vector_store_id") != vector_store_id + ] + return None + + # Create mock prisma client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( + side_effect=mock_find_unique + ) + mock_prisma_client.db.litellm_managedvectorstorestable.find_many = AsyncMock( + side_effect=mock_find_many + ) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( + side_effect=mock_create + ) + mock_prisma_client.db.litellm_managedvectorstorestable.delete = AsyncMock( + side_effect=mock_delete + ) + + # Test vector store data + test_vector_store_id = "test-sync-store-001" + test_vector_store: LiteLLM_ManagedVectorStore = { + "vector_store_id": test_vector_store_id, + "custom_llm_provider": "bedrock", + "vector_store_name": "Test Sync Store", + "vector_store_description": "Testing synchronization", + "litellm_params": { + "vector_store_id": test_vector_store_id, + "custom_llm_provider": "bedrock", + "region_name": "us-east-1" + }, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + + # Step 1: Create vector store on Instance 1 + # (Simulate what happens in new_vector_store endpoint) + await mock_prisma_client.db.litellm_managedvectorstorestable.create( + data=test_vector_store + ) + instance_1_registry.add_vector_store_to_registry(vector_store=test_vector_store) + + # Verify it's in Instance 1's memory + assert instance_1_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) is not None, "Vector store should be in Instance 1's memory" + + # Verify it's in the database + db_store = await mock_prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": test_vector_store_id} + ) + assert db_store is not None, "Vector store should be in database" + + # Step 2: Instance 2 should be able to find it via database fallback + # (Simulate what happens in pop_vector_stores_to_run_with_db_fallback) + found_store = await instance_2_registry.get_litellm_managed_vector_store_from_registry_or_db( + vector_store_id=test_vector_store_id, + prisma_client=mock_prisma_client + ) + assert found_store is not None, "Instance 2 should find vector store from database" + assert found_store.get("vector_store_id") == test_vector_store_id + + # Verify it's now cached in Instance 2's memory + assert instance_2_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) is not None, "Vector store should now be cached in Instance 2's memory" + + # Step 3: Test that Instance 2 can list vector stores from database + # (Simulate what happens in list_vector_stores endpoint - using DB as source of truth) + vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db( + prisma_client=mock_prisma_client + ) + + # Verify vector store appears in the database list + vector_store_ids = [vs.get("vector_store_id") for vs in vector_stores_from_db] + assert test_vector_store_id in vector_store_ids, ( + "Instance 2 should see vector store from database" + ) + + # Verify the list endpoint logic: only show DB stores (filter out stale cache) + # This simulates what list_vector_stores does + db_vector_store_ids = { + vs.get("vector_store_id") + for vs in vector_stores_from_db + if vs.get("vector_store_id") + } + + # Instance 2's in-memory cache should only contain stores that exist in DB + # (This is what the list endpoint cleanup does) + for vs in list(instance_2_registry.vector_stores): + vs_id = vs.get("vector_store_id") + if vs_id and vs_id not in db_vector_store_ids: + instance_2_registry.delete_vector_store_from_registry(vector_store_id=vs_id) + + # After cleanup, instance 2 should still have the vector store (it's in DB) + assert instance_2_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) is not None, "Instance 2 should still have vector store (it exists in DB)" + + # Step 4: Delete vector store on Instance 1 + # (Simulate what happens in delete_vector_store endpoint) + await mock_prisma_client.db.litellm_managedvectorstorestable.delete( + where={"vector_store_id": test_vector_store_id} + ) + instance_1_registry.delete_vector_store_from_registry( + vector_store_id=test_vector_store_id + ) + + # Verify it's removed from Instance 1's memory + assert instance_1_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) is None, "Vector store should be removed from Instance 1's memory" + + # Verify it's removed from database + db_store_after_delete = await mock_prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": test_vector_store_id} + ) + assert db_store_after_delete is None, "Vector store should be removed from database" + + # Step 5: Instance 2 should NOT show it in the list (database is source of truth) + # The list endpoint logic should clean up stale cache entries + vector_stores_from_db_after_delete = await VectorStoreRegistry._get_vector_stores_from_db( + prisma_client=mock_prisma_client + ) + + # Verify vector store does NOT appear in the database list + vector_store_ids_after_delete = [vs.get("vector_store_id") for vs in vector_stores_from_db_after_delete] + assert test_vector_store_id not in vector_store_ids_after_delete, ( + "Deleted vector store should not be in database" + ) + + # Simulate list endpoint cleanup logic + db_vector_store_ids_after_delete = { + vs.get("vector_store_id") + for vs in vector_stores_from_db_after_delete + if vs.get("vector_store_id") + } + + # Remove any in-memory vector stores that no longer exist in database + for vs in list(instance_2_registry.vector_stores): + vs_id = vs.get("vector_store_id") + if vs_id and vs_id not in db_vector_store_ids_after_delete: + instance_2_registry.delete_vector_store_from_registry(vector_store_id=vs_id) + + # Verify it was removed from Instance 2's cache + assert instance_2_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id + ) is None, ( + "Deleted vector store should be removed from Instance 2's cache" + ) + + # Step 6: Test that using a deleted vector store fails gracefully + # (Simulate what happens in pop_vector_stores_to_run_with_db_fallback) + non_default_params = {"vector_store_ids": [test_vector_store_id]} + vector_stores_to_run = await instance_2_registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params=non_default_params, + tools=None, + prisma_client=mock_prisma_client + ) + + assert len(vector_stores_to_run) == 0, ( + "Deleted vector store should not be returned when trying to use it" + )