Merge pull request #17525 from BerriAI/litellm_fix_in_memory_vector_store

Fix vector store configuration synchronization failure

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

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
                 )
             )
 

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]:

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"
+    )