From 9a9292577ad599e71e5968b7c8a6558a1b31ef30 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 26 Jan 2026 19:10:02 -0800 Subject: [PATCH] test_rag_ingest_calls_save_vector_store_to_db --- .../test_vector_store_endpoints.py | 214 ++++++++++++++++++ 1 file changed, 214 insertions(+) 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 558fe18ae38..7a88b53d96d 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 @@ -1368,3 +1368,217 @@ async def test_new_vector_store_auto_resolves_embedding_config(): assert litellm_params_dict["litellm_embedding_config"]["api_key"] == "resolved-api-key" assert litellm_params_dict["litellm_embedding_config"]["api_base"] == "https://api.openai.com" assert litellm_params_dict["litellm_embedding_config"]["api_version"] == "2024-01-01" + + +@pytest.mark.asyncio +async def test_rag_ingest_calls_save_vector_store_to_db(): + """ + Test that rag_ingest endpoint calls _save_vector_store_to_db_from_rag_ingest + to persist vector stores in the database. + + This test verifies: + 1. The rag_ingest endpoint successfully ingests a document + 2. _save_vector_store_to_db_from_rag_ingest is called with correct parameters + 3. Vector store metadata is saved to the database + 4. File metadata is properly tracked + """ + from litellm.proxy.rag_endpoints.endpoints import ( + _save_vector_store_to_db_from_rag_ingest, + ) + + # Mock response from litellm.aingest + mock_ingest_response = { + "vector_store_id": "vs_test_123", + "file_id": "file_abc_456", + "status": "completed", + "chunks_created": 10, + } + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user API key + mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key.user_role = None + mock_user_api_key.user_id = "test-user-123" + + # Test ingest options with custom vector store name + ingest_options = { + "vector_store": { + "custom_llm_provider": "openai", + }, + "litellm_vector_store_params": { + "vector_store_name": "My Test Vector Store", + "vector_store_description": "A test vector store for unit testing", + }, + } + + # Test file data + file_data = ("test_document.pdf", b"fake pdf content", "application/pdf") + + # Mock database operations + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( + return_value=None # Vector store doesn't exist yet + ) + + captured_create_data = {} + + async def mock_create(*args, **kwargs): + captured_create_data.update(kwargs.get("data", {})) + mock_created_vs = MagicMock() + mock_created_vs.model_dump.return_value = kwargs.get("data", {}) + return mock_created_vs + + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( + side_effect=mock_create + ) + + mock_registry = MagicMock() + mock_registry.add_vector_store_to_registry = MagicMock() + + # Patch dependencies + with patch.object( + litellm, "vector_store_registry", mock_registry + ), patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client + ): + # Call the function directly + await _save_vector_store_to_db_from_rag_ingest( + response=mock_ingest_response, + ingest_options=ingest_options, + prisma_client=mock_prisma_client, + user_api_key_dict=mock_user_api_key, + file_data=file_data, + file_url=None, + ) + + # Verify database create was called + mock_prisma_client.db.litellm_managedvectorstorestable.create.assert_called_once() + + # Verify the created data has correct structure + assert captured_create_data["vector_store_id"] == "vs_test_123" + assert captured_create_data["custom_llm_provider"] == "openai" + assert captured_create_data["vector_store_name"] == "My Test Vector Store" + assert captured_create_data["vector_store_description"] == "A test vector store for unit testing" + + # Verify file metadata was included + import json + vector_store_metadata = json.loads(captured_create_data["vector_store_metadata"]) + assert "ingested_files" in vector_store_metadata + assert len(vector_store_metadata["ingested_files"]) == 1 + + file_entry = vector_store_metadata["ingested_files"][0] + assert file_entry["file_id"] == "file_abc_456" + assert file_entry["filename"] == "test_document.pdf" + assert file_entry["content_type"] == "application/pdf" + assert "ingested_at" in file_entry + + # Verify registry was updated + mock_registry.add_vector_store_to_registry.assert_called_once() + + +@pytest.mark.asyncio +async def test_rag_ingest_updates_existing_vector_store_metadata(): + """ + Test that rag_ingest appends file metadata to existing vector stores + instead of creating duplicates. + """ + from litellm.proxy.rag_endpoints.endpoints import ( + _save_vector_store_to_db_from_rag_ingest, + ) + + # Mock response from litellm.aingest + mock_ingest_response = { + "vector_store_id": "vs_existing_123", + "file_id": "file_new_789", + "status": "completed", + } + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user API key + mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key.user_role = None + mock_user_api_key.user_id = "test-user-123" + + # Test ingest options + ingest_options = { + "vector_store": { + "custom_llm_provider": "openai", + "vector_store_id": "vs_existing_123", + }, + } + + # Test file data + file_data = ("new_document.pdf", b"new content", "application/pdf") + + # Mock existing vector store in database + existing_metadata = { + "ingested_files": [ + { + "file_id": "file_old_456", + "filename": "old_document.pdf", + "ingested_at": "2024-01-01T00:00:00Z", + } + ] + } + + mock_existing_vs = MagicMock() + mock_existing_vs.vector_store_id = "vs_existing_123" + mock_existing_vs.custom_llm_provider = "openai" + mock_existing_vs.vector_store_metadata = existing_metadata + + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( + return_value=mock_existing_vs + ) + + captured_update_data = {} + + async def mock_update(*args, **kwargs): + captured_update_data.update(kwargs.get("data", {})) + mock_updated_vs = MagicMock() + mock_updated_vs.model_dump.return_value = kwargs.get("data", {}) + return mock_updated_vs + + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock( + side_effect=mock_update + ) + + mock_registry = MagicMock() + mock_registry.update_vector_store_in_registry = MagicMock() + + # Patch dependencies + with patch.object( + litellm, "vector_store_registry", mock_registry + ), patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client + ): + # Call the function + await _save_vector_store_to_db_from_rag_ingest( + response=mock_ingest_response, + ingest_options=ingest_options, + prisma_client=mock_prisma_client, + user_api_key_dict=mock_user_api_key, + file_data=file_data, + file_url=None, + ) + + # Verify update was called instead of create + mock_prisma_client.db.litellm_managedvectorstorestable.update.assert_called_once() + + # Verify the updated metadata includes both old and new files + import json + updated_metadata = json.loads(captured_update_data["vector_store_metadata"]) + assert "ingested_files" in updated_metadata + assert len(updated_metadata["ingested_files"]) == 2 + + # Check old file is still there + file_ids = [f["file_id"] for f in updated_metadata["ingested_files"]] + assert "file_old_456" in file_ids + assert "file_new_789" in file_ids + + # Verify registry was updated + mock_registry.update_vector_store_in_registry.assert_called_once()