diff --git a/docs/my-website/docs/providers/milvus_vector_stores.md b/docs/my-website/docs/providers/milvus_vector_stores.md index c1a3042051b..84f16fbc74a 100644 --- a/docs/my-website/docs/providers/milvus_vector_stores.md +++ b/docs/my-website/docs/providers/milvus_vector_stores.md @@ -178,7 +178,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/vector_stores/my-collection-name/search' \ | Guardrails | ❌ Not Yet Supported | Guardrails are not currently supported for vector stores | | Cost Tracking | ✅ Supported | Cost is $0 for Milvus searches | | Unified API | ✅ Supported | Call via OpenAI compatible `/v1/vector_stores/search` endpoint | -| Passthrough | ❌ Not yet supported | | +| Passthrough | ✅ Supported | Use native Milvus API format | ## Response Format @@ -208,6 +208,313 @@ The response follows the standard LiteLLM vector store format: } ``` +## Passthrough API (Native Milvus Format) + +Use this to allow developers to **create** and **search** vector stores using the native Milvus API format, without giving them the Milvus credentials. + +This is for the proxy only. + +### Admin Flow + +#### 1. Add the vector store to LiteLLM + +```yaml +model_list: + - model_name: embedding-model + litellm_params: + model: azure/text-embedding-3-large + api_base: https://your-endpoint.cognitiveservices.azure.com/ + api_key: os.environ/AZURE_API_KEY + api_version: "2025-09-01" + +vector_store_registry: + - vector_store_name: "milvus-store" + litellm_params: + vector_store_id: "can-be-anything" # vector store id can be anything for the purpose of passthrough api + custom_llm_provider: "milvus" + api_key: os.environ/MILVUS_API_KEY + api_base: https://your-milvus-instance.milvus.io + +general_settings: + database_url: "postgresql://user:password@host:port/database" + master_key: "sk-1234" +``` + +Add your vector store credentials to LiteLLM. + +#### 2. Start the proxy + +```bash +litellm --config /path/to/config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +#### 3. Create a virtual index + +```bash +curl -L -X POST 'http://0.0.0.0:4000/v1/indexes' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-d '{ + "index_name": "dall-e-6", + "litellm_params": { + "vector_store_index": "real-collection-name", + "vector_store_name": "milvus-store" + } +}' +``` + +This is a virtual index, which the developer can use to create and search vector stores. + +#### 4. Create a key with the vector store permissions + +```bash +curl -L -X POST 'http://0.0.0.0:4000/key/generate' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-d '{ + "allowed_vector_store_indexes": [{"index_name": "dall-e-6", "index_permissions": ["write", "read"]}], + "models": ["embedding-model"] +}' +``` + +Give the key access to the virtual index and the embedding model. + +**Expected response** + +```json +{ + "key": "sk-my-virtual-key" +} +``` + +### Developer Flow + +#### 1. Create a collection with schema + +Note: Use the `/milvus` endpoint for the passthrough api that uses the `milvus` provider in your config. + +```python +from milvus_rest_client import MilvusRESTClient, DataType +import random +import time + +# Configuration +uri = "http://0.0.0.0:4000/milvus" # IMPORTANT: Use the '/milvus' endpoint for passthrough +token = "sk-my-virtual-key" +collection_name = "dall-e-6" # Virtual index name + +# Initialize client +milvus_client = MilvusRESTClient(uri=uri, token=token) +print(f"Connected to DB: {uri} successfully") + +# Check if the collection exists and drop if it does +check_collection = milvus_client.has_collection(collection_name) +if check_collection: + milvus_client.drop_collection(collection_name) + print(f"Dropped the existing collection {collection_name} successfully") + +# Define schema +dim = 64 # Vector dimension + +print("Start to create the collection schema") +schema = milvus_client.create_schema() +schema.add_field( + "book_id", DataType.INT64, is_primary=True, description="customized primary id" +) +schema.add_field("word_count", DataType.INT64, description="word count") +schema.add_field( + "book_intro", DataType.FLOAT_VECTOR, dim=dim, description="book introduction" +) + +# Prepare index parameters +print("Start to prepare index parameters with default AUTOINDEX") +index_params = milvus_client.prepare_index_params() +index_params.add_index("book_intro", metric_type="L2") + +# Create collection +print(f"Start to create example collection: {collection_name}") +milvus_client.create_collection( + collection_name, schema=schema, index_params=index_params +) +collection_property = milvus_client.describe_collection(collection_name) +print("Collection details: %s" % collection_property) +``` + +#### 2. Insert data into the collection + +```python +# Insert data with customized ids +nb = 1000 +insert_rounds = 2 +start = 0 # first primary key id +total_rt = 0 # total response time for insert + +print( + f"Start to insert {nb*insert_rounds} entities into example collection: {collection_name}" +) +for i in range(insert_rounds): + vector = [random.random() for _ in range(dim)] + rows = [ + {"book_id": i, "word_count": random.randint(1, 100), "book_intro": vector} + for i in range(start, start + nb) + ] + t0 = time.time() + milvus_client.insert(collection_name, rows) + ins_rt = time.time() - t0 + start += nb + total_rt += ins_rt +print(f"Insert completed in {round(total_rt, 4)} seconds") + +# Flush the collection +print("Start to flush") +start_flush = time.time() +milvus_client.flush(collection_name) +end_flush = time.time() +print(f"Flush completed in {round(end_flush - start_flush, 4)} seconds") +``` + +#### 3. Search the collection + +```python +# Search configuration +nq = 3 # Number of query vectors +search_params = {"metric_type": "L2", "params": {"level": 2}} +limit = 2 # Number of results to return + +# Perform searches +for i in range(5): + search_vectors = [[random.random() for _ in range(dim)] for _ in range(nq)] + t0 = time.time() + results = milvus_client.search( + collection_name, + data=search_vectors, + limit=limit, + search_params=search_params, + anns_field="book_intro", + ) + t1 = time.time() + print(f"Search {i} results: {results}") + print(f"Search {i} latency: {round(t1-t0, 4)} seconds") +``` + +#### Complete Example + +Here's a full working example: + +```python +from milvus_rest_client import MilvusRESTClient, DataType +import random +import time + +# ---------------------------- +# 🔐 CONFIGURATION +# ---------------------------- +uri = "http://0.0.0.0:4000/milvus" # IMPORTANT: Use the '/milvus' endpoint +token = "sk-my-virtual-key" +collection_name = "dall-e-6" # Your virtual index name + +# ---------------------------- +# 📋 STEP 1 — Initialize Client +# ---------------------------- +milvus_client = MilvusRESTClient(uri=uri, token=token) +print(f"✅ Connected to DB: {uri} successfully") + +# ---------------------------- +# 🗑️ STEP 2 — Drop Existing Collection (if needed) +# ---------------------------- +check_collection = milvus_client.has_collection(collection_name) +if check_collection: + milvus_client.drop_collection(collection_name) + print(f"🗑️ Dropped the existing collection {collection_name} successfully") + +# ---------------------------- +# 📐 STEP 3 — Create Collection Schema +# ---------------------------- +dim = 64 # Vector dimension + +print("📐 Creating the collection schema") +schema = milvus_client.create_schema() +schema.add_field( + "book_id", DataType.INT64, is_primary=True, description="customized primary id" +) +schema.add_field("word_count", DataType.INT64, description="word count") +schema.add_field( + "book_intro", DataType.FLOAT_VECTOR, dim=dim, description="book introduction" +) + +# ---------------------------- +# 🔍 STEP 4 — Create Index +# ---------------------------- +print("🔍 Preparing index parameters with default AUTOINDEX") +index_params = milvus_client.prepare_index_params() +index_params.add_index("book_intro", metric_type="L2") + +# ---------------------------- +# 🏗️ STEP 5 — Create Collection +# ---------------------------- +print(f"🏗️ Creating collection: {collection_name}") +milvus_client.create_collection( + collection_name, schema=schema, index_params=index_params +) +collection_property = milvus_client.describe_collection(collection_name) +print(f"✅ Collection created: {collection_property}") + +# ---------------------------- +# 📤 STEP 6 — Insert Data +# ---------------------------- +nb = 1000 +insert_rounds = 2 +start = 0 +total_rt = 0 + +print(f"📤 Inserting {nb*insert_rounds} entities into collection") +for i in range(insert_rounds): + vector = [random.random() for _ in range(dim)] + rows = [ + {"book_id": i, "word_count": random.randint(1, 100), "book_intro": vector} + for i in range(start, start + nb) + ] + t0 = time.time() + milvus_client.insert(collection_name, rows) + ins_rt = time.time() - t0 + start += nb + total_rt += ins_rt +print(f"✅ Insert completed in {round(total_rt, 4)} seconds") + +# ---------------------------- +# 💾 STEP 7 — Flush Collection +# ---------------------------- +print("💾 Flushing collection") +start_flush = time.time() +milvus_client.flush(collection_name) +end_flush = time.time() +print(f"✅ Flush completed in {round(end_flush - start_flush, 4)} seconds") + +# ---------------------------- +# 🔍 STEP 8 — Search +# ---------------------------- +nq = 3 +search_params = {"metric_type": "L2", "params": {"level": 2}} +limit = 2 + +print(f"🔍 Performing {5} search operations") +for i in range(5): + search_vectors = [[random.random() for _ in range(dim)] for _ in range(nq)] + t0 = time.time() + results = milvus_client.search( + collection_name, + data=search_vectors, + limit=limit, + search_params=search_params, + anns_field="book_intro", + ) + t1 = time.time() + print(f"✅ Search {i} results: {results}") + print(f" Search {i} latency: {round(t1-t0, 4)} seconds") +``` + ## How It Works When you search: diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 1f57cd9f159..9ef9812dd49 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -12,4 +12,10 @@ vector_store_registry: vector_store_id: "litellm-docs_1761094140318" custom_llm_provider: "vertex_ai/search_api" vertex_project: "test-vector-store-db" - vertex_location: "global" \ No newline at end of file + vertex_location: "global" + - vector_store_name: "milvus-litellm-website-knowledgebase" + litellm_params: + vector_store_id: "can-be-anything" + custom_llm_provider: "milvus" + api_base: os.environ/MILVUS_API_BASE + api_key: os.environ/MILVUS_API_KEY \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 974e25615e4..0a6240b0f9f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -346,6 +346,7 @@ class LiteLLMRoutes(enum.Enum): "/eu.assemblyai", "/vllm", "/mistral", + "/milvus", ] ######################################################### diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index cd8f6c38000..b7a82bb8d60 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, + _safe_set_request_parsed_body, get_form_data, get_request_body, ) @@ -39,6 +40,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( is_allowed_to_call_vector_store_endpoint, ) from litellm.secret_managers.main import get_secret_str +from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter @@ -161,12 +163,12 @@ async def llm_passthrough_factory_proxy_route( endpoint=endpoint, target=str(updated_url), custom_headers=auth_headers, + is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, fastapi_response, user_api_key_dict, - stream=is_streaming_request, # type: ignore ) return received_value @@ -230,13 +232,12 @@ async def gemini_proxy_route( endpoint=endpoint, target=str(updated_url), custom_llm_provider="gemini", + is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, fastapi_response, user_api_key_dict, - query_params=merged_params, # type: ignore - stream=is_streaming_request, # type: ignore ) return received_value @@ -283,12 +284,12 @@ async def cohere_proxy_route( endpoint=endpoint, target=str(updated_url), custom_headers={"Authorization": "Bearer {}".format(cohere_api_key)}, + is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, fastapi_response, user_api_key_dict, - stream=is_streaming_request, # type: ignore ) return received_value @@ -412,12 +413,141 @@ async def mistral_proxy_route( endpoint=endpoint, target=str(updated_url), custom_headers={"Authorization": "Bearer {}".format(mistral_api_key)}, + is_streaming_request=is_streaming_request, + ) # dynamically construct pass-through endpoint based on incoming path + received_value = await endpoint_func( + request, + fastapi_response, + user_api_key_dict, + ) + + return received_value + + +@router.api_route( + "/milvus/{endpoint:path}", + methods=["GET", "POST", "PUT", "DELETE", "PATCH"], + tags=["Milvus Pass-through", "pass-through"], +) +async def milvus_proxy_route( + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Enable using Milvus `/vectors` endpoint as a pass-through endpoint. + """ + + provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=LlmProviders.MILVUS + ) + if not provider_config: + raise HTTPException( + status_code=500, + detail="Unable to find Milvus vector store config.", + ) + + # check if managed vector store index is used + request_body = await get_request_body(request) + + # check collectionName + collection_name = cast(Optional[str], request_body.get("collectionName")) + extra_headers = {} + base_target_url: Optional[str] = None + if not collection_name: + raise HTTPException( + status_code=400, + detail=f"Collection name is required. Got {request_body}", + ) + + if not litellm.vector_store_index_registry or not litellm.vector_store_registry: + raise HTTPException( + status_code=500, + detail="Unable to find Milvus vector store index registry or vector store registry.", + ) + + # check if vector store index + is_vector_store_index = litellm.vector_store_index_registry.is_vector_store_index( + vector_store_index_name=collection_name + ) + + if not is_vector_store_index: + raise HTTPException( + status_code=400, + detail=f"Collection {collection_name} is not a litellm managed vector store index. Only litellm managed vector store indexes are supported.", + ) + + is_allowed_to_call_vector_store_endpoint( + index_name=collection_name, + provider=LlmProviders.MILVUS, + request=request, + user_api_key_dict=user_api_key_dict, + ) + # get the vector store name from index registry + + index_object = ( + ( + litellm.vector_store_index_registry.get_vector_store_index_by_name( + vector_store_index_name=collection_name + ) + ) + if litellm.vector_store_index_registry is not None + else None + ) + if index_object is None: + raise Exception(f"Vector store index not found for {collection_name}") + + vector_store_name = index_object.litellm_params.vector_store_name + vector_store_index = index_object.litellm_params.vector_store_index + + request_body["collectionName"] = vector_store_index + + # Update the request object with the modified collection name + _safe_set_request_parsed_body(request, request_body) + + vector_store = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry_by_name( + vector_store_name=vector_store_name + ) + if vector_store is None: + raise Exception(f"Vector store not found for {vector_store_name}") + litellm_params = vector_store.get("litellm_params") or {} + auth_credentials = provider_config.get_auth_credentials( + litellm_params=litellm_params + ) + + extra_headers = auth_credentials.get("headers") or {} + + litellm_params = vector_store.get("litellm_params") or {} + + base_target_url = provider_config.get_complete_url( + api_base=litellm_params.get("api_base"), litellm_params=litellm_params + ) + + if base_target_url is None: + raise Exception( + f"api_base not found in vector store configuration for {vector_store_name}" + ) + + encoded_endpoint = httpx.URL(endpoint).path + + # Ensure endpoint starts with '/' for proper URL construction + if not encoded_endpoint.startswith("/"): + encoded_endpoint = "/" + encoded_endpoint + + # Construct the full target URL using httpx + base_url = httpx.URL(base_target_url) + updated_url = base_url.copy_with(path=encoded_endpoint) + ## CREATE PASS-THROUGH + endpoint_func = create_pass_through_route( + endpoint=endpoint, + target=str(updated_url), + custom_headers=extra_headers, ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, fastapi_response, user_api_key_dict, - stream=is_streaming_request, # type: ignore ) return received_value @@ -475,12 +605,12 @@ async def anthropic_proxy_route( target=str(updated_url), custom_headers={"x-api-key": "{}".format(anthropic_api_key)}, _forward_headers=True, + is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, fastapi_response, user_api_key_dict, - stream=is_streaming_request, # type: ignore ) return received_value @@ -898,12 +1028,12 @@ async def bedrock_proxy_route( endpoint=endpoint, target=str(prepped.url), custom_headers=prepped.headers, # type: ignore + is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, fastapi_response, user_api_key_dict, - stream=is_streaming_request, # type: ignore custom_body=data, # type: ignore query_params={}, # type: ignore ) @@ -981,12 +1111,12 @@ async def assemblyai_proxy_route( endpoint=endpoint, target=str(updated_url), custom_headers={"Authorization": "{}".format(assemblyai_api_key)}, + is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, - stream=is_streaming_request, # type: ignore ) return received_value @@ -1687,13 +1817,12 @@ class BaseOpenAIPassThroughHandler: custom_headers=BaseOpenAIPassThroughHandler._assemble_headers( api_key=api_key, request=request, extra_headers=extra_headers ), + is_streaming_request=is_streaming_request, # type: ignore ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, fastapi_response, user_api_key_dict, - stream=is_streaming_request, # type: ignore - query_params=dict(request.query_params), # type: ignore ) return received_value diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 1d88ad0edf4..81c205095c1 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -2,7 +2,7 @@ from typing import Any, Dict, Literal, Optional from fastapi import HTTPException, Request -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -73,6 +73,11 @@ def is_allowed_to_call_vector_store_endpoint( 1. Creating a vector store index 2. Reading a vector store index (Search / List / Get) """ + if ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ): + return True # check what allowed permissions are for the key key_metadata = user_api_key_dict.metadata team_metadata = user_api_key_dict.team_metadata diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 84e53a903a5..ea1017e1d5a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -21,6 +21,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( bedrock_llm_proxy_route, create_pass_through_route, llm_passthrough_factory_proxy_route, + milvus_proxy_route, vertex_discovery_proxy_route, vertex_proxy_route, vllm_proxy_route, @@ -179,11 +180,8 @@ class TestBaseOpenAIPassThroughHandler: print("Verifying endpoint_func call parameters...") mock_endpoint_func.assert_awaited_once() assert mock_endpoint_func.await_args is not None - call_kwargs = mock_endpoint_func.await_args[1] - print(f"stream parameter: {call_kwargs['stream']}") - print(f"query_params: {call_kwargs['query_params']}") - assert call_kwargs["stream"] is False - assert call_kwargs["query_params"] == {"model": "gpt-4"} + # The endpoint_func is called with request, fastapi_response, user_api_key_dict + # No longer checking for stream and query_params as they're handled differently class TestVertexAIPassThroughHandler: @@ -291,6 +289,7 @@ class TestVertexAIPassThroughHandler: endpoint=endpoint, target=f"https://{test_location}-aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {test_token}"}, + is_streaming_request=False, ) @pytest.mark.asyncio @@ -389,6 +388,7 @@ class TestVertexAIPassThroughHandler: endpoint=endpoint, target=f"https://aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {test_token}"}, + is_streaming_request=False, ) @pytest.mark.parametrize( @@ -481,6 +481,7 @@ class TestVertexAIPassThroughHandler: endpoint=endpoint, target=f"https://{default_location}-aiplatform.googleapis.com/v1/projects/{default_project}/locations/{default_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {default_credentials}"}, + is_streaming_request=False, ) @pytest.mark.asyncio @@ -559,6 +560,7 @@ class TestVertexAIPassThroughHandler: endpoint=endpoint, target=f"https://{test_location}-aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"authorization": f"Bearer {test_token}"}, + is_streaming_request=False, ) @pytest.mark.asyncio @@ -1134,17 +1136,17 @@ class TestBedrockLLMProxyRoute: } mock_llm_router = Mock() - + # Mock ProxyBaseLLMRequestProcessing to raise the httpx error with patch( "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.base_passthrough_process_llm_request", new_callable=AsyncMock, - side_effect=mock_http_error + side_effect=mock_http_error, ): mock_user_api_key_dict = Mock() mock_user_api_key_dict.api_key = "test-key" mock_user_api_key_dict.allowed_model_region = None - + mock_proxy_logging_obj = Mock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() @@ -1226,6 +1228,7 @@ class TestLLMPassthroughFactoryProxyRoute: endpoint="/chat/completions", target="https://example.com/v1/chat/completions", custom_headers={"x-api-key": "dummy"}, + is_streaming_request=False, ) mock_endpoint_func.assert_awaited_once() @@ -1293,3 +1296,453 @@ class TestVLLMProxyRoute: assert result == "factory_success" mock_factory_route.assert_awaited_once() + + +class TestMilvusProxyRoute: + """ + Test cases for Milvus passthrough endpoint + """ + + @pytest.mark.asyncio + async def test_milvus_proxy_route_success(self): + """ + Test successful Milvus proxy route with valid managed vector store index + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + collection_name = "dall-e-6" + vector_store_name = "milvus-store-1" + vector_store_index = "collection_123" + api_base = "http://localhost:19530" + + # Mock request + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.headers = {"content-type": "application/json"} + mock_request.url = MagicMock() + mock_request.url.path = "/milvus/vectors/search" + + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + # Mock vector store index object + mock_index_object = MagicMock() + mock_index_object.litellm_params.vector_store_name = vector_store_name + mock_index_object.litellm_params.vector_store_index = vector_store_index + + # Mock vector store + mock_vector_store = { + "litellm_params": { + "api_base": api_base, + "api_key": "test-milvus-key", + } + } + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"collectionName": collection_name, "data": [[0.1, 0.2]]}, + ) as mock_get_body, patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config" + ) as mock_get_config, patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint" + ) as mock_is_allowed, patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body" + ) as mock_safe_set, patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_create_route, patch.object( + litellm, "vector_store_index_registry" + ) as mock_index_registry, patch.object( + litellm, "vector_store_registry" + ) as mock_vector_registry: + # Setup mocks + mock_provider_config = MagicMock() + mock_provider_config.get_auth_credentials.return_value = { + "headers": {"Authorization": "Bearer test-token"} + } + mock_provider_config.get_complete_url.return_value = api_base + mock_get_config.return_value = mock_provider_config + + mock_index_registry.is_vector_store_index.return_value = True + mock_index_registry.get_vector_store_index_by_name.return_value = ( + mock_index_object + ) + + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = ( + mock_vector_store + ) + + mock_endpoint_func = AsyncMock( + return_value={"results": [{"id": 1, "distance": 0.5}]} + ) + mock_create_route.return_value = mock_endpoint_func + + # Call the route + result = await milvus_proxy_route( + endpoint="vectors/search", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify calls + mock_get_body.assert_called_once() + mock_index_registry.is_vector_store_index.assert_called_once_with( + vector_store_index_name=collection_name + ) + mock_is_allowed.assert_called_once() + mock_safe_set.assert_called_once() + + # Verify collection name was updated to the actual index + set_body_call_args = mock_safe_set.call_args[0] + assert set_body_call_args[1]["collectionName"] == vector_store_index + + # Verify create_pass_through_route was called with correct URL + mock_create_route.assert_called_once() + create_route_args = mock_create_route.call_args[1] + assert "vectors/search" in create_route_args["target"] + assert create_route_args["custom_headers"] == { + "Authorization": "Bearer test-token" + } + + # Verify endpoint function was called + mock_endpoint_func.assert_awaited_once() + assert result == {"results": [{"id": 1, "distance": 0.5}]} + + @pytest.mark.asyncio + async def test_milvus_proxy_route_missing_collection_name(self): + """ + Test that missing collection name raises HTTPException + """ + from fastapi import HTTPException + + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + mock_request = MagicMock(spec=Request) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"data": [[0.1, 0.2]]}, # No collectionName + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config" + ) as mock_get_config: + mock_get_config.return_value = MagicMock() + + with pytest.raises(HTTPException) as exc_info: + await milvus_proxy_route( + endpoint="vectors/search", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert exc_info.value.status_code == 400 + assert "Collection name is required" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_milvus_proxy_route_no_provider_config(self): + """ + Test that missing provider config raises HTTPException + """ + from fastapi import HTTPException + + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + mock_request = MagicMock(spec=Request) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config", + return_value=None, + ): + with pytest.raises(HTTPException) as exc_info: + await milvus_proxy_route( + endpoint="vectors/search", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert exc_info.value.status_code == 500 + assert "Unable to find Milvus vector store config" in str( + exc_info.value.detail + ) + + @pytest.mark.asyncio + async def test_milvus_proxy_route_no_index_registry(self): + """ + Test that missing index registry raises HTTPException + """ + from fastapi import HTTPException + + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + collection_name = "test-collection" + + mock_request = MagicMock(spec=Request) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"collectionName": collection_name}, + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config" + ) as mock_get_config, patch.object( + litellm, "vector_store_index_registry", None + ): + mock_get_config.return_value = MagicMock() + + with pytest.raises(HTTPException) as exc_info: + await milvus_proxy_route( + endpoint="vectors/search", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert exc_info.value.status_code == 500 + assert "Unable to find Milvus vector store index registry" in str( + exc_info.value.detail + ) + + @pytest.mark.asyncio + async def test_milvus_proxy_route_not_managed_index(self): + """ + Test that non-managed vector store index raises HTTPException + """ + from fastapi import HTTPException + + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + collection_name = "unmanaged-collection" + + mock_request = MagicMock(spec=Request) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"collectionName": collection_name}, + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config" + ) as mock_get_config, patch.object( + litellm, "vector_store_index_registry" + ) as mock_index_registry, patch.object( + litellm, "vector_store_registry", MagicMock() + ): + mock_get_config.return_value = MagicMock() + mock_index_registry.is_vector_store_index.return_value = False + + with pytest.raises(HTTPException) as exc_info: + await milvus_proxy_route( + endpoint="vectors/search", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert exc_info.value.status_code == 400 + assert ( + f"Collection {collection_name} is not a litellm managed vector store index" + in str(exc_info.value.detail) + ) + + @pytest.mark.asyncio + async def test_milvus_proxy_route_vector_store_not_found(self): + """ + Test that missing vector store raises Exception + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + collection_name = "test-collection" + vector_store_name = "missing-store" + vector_store_index = "collection_123" + + mock_request = MagicMock(spec=Request) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + mock_index_object = MagicMock() + mock_index_object.litellm_params.vector_store_name = vector_store_name + mock_index_object.litellm_params.vector_store_index = vector_store_index + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"collectionName": collection_name}, + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config" + ) as mock_get_config, patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint" + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body" + ), patch.object( + litellm, "vector_store_index_registry" + ) as mock_index_registry, patch.object( + litellm, "vector_store_registry" + ) as mock_vector_registry: + mock_get_config.return_value = MagicMock() + mock_index_registry.is_vector_store_index.return_value = True + mock_index_registry.get_vector_store_index_by_name.return_value = ( + mock_index_object + ) + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = ( + None + ) + + with pytest.raises(Exception) as exc_info: + await milvus_proxy_route( + endpoint="vectors/search", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert f"Vector store not found for {vector_store_name}" in str( + exc_info.value + ) + + @pytest.mark.asyncio + async def test_milvus_proxy_route_no_api_base(self): + """ + Test that missing api_base raises Exception + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + collection_name = "test-collection" + vector_store_name = "milvus-store-1" + vector_store_index = "collection_123" + + mock_request = MagicMock(spec=Request) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + mock_index_object = MagicMock() + mock_index_object.litellm_params.vector_store_name = vector_store_name + mock_index_object.litellm_params.vector_store_index = vector_store_index + + mock_vector_store = {"litellm_params": {}} # No api_base + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"collectionName": collection_name}, + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config" + ) as mock_get_config, patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint" + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body" + ), patch.object( + litellm, "vector_store_index_registry" + ) as mock_index_registry, patch.object( + litellm, "vector_store_registry" + ) as mock_vector_registry: + mock_provider_config = MagicMock() + mock_provider_config.get_auth_credentials.return_value = {"headers": {}} + mock_provider_config.get_complete_url.return_value = None + mock_get_config.return_value = mock_provider_config + + mock_index_registry.is_vector_store_index.return_value = True + mock_index_registry.get_vector_store_index_by_name.return_value = ( + mock_index_object + ) + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = ( + mock_vector_store + ) + + with pytest.raises(Exception) as exc_info: + await milvus_proxy_route( + endpoint="vectors/search", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + assert ( + f"api_base not found in vector store configuration for {vector_store_name}" + in str(exc_info.value) + ) + + @pytest.mark.asyncio + async def test_milvus_proxy_route_endpoint_without_leading_slash(self): + """ + Test that endpoint without leading slash is handled correctly + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + collection_name = "test-collection" + vector_store_name = "milvus-store-1" + vector_store_index = "collection_123" + api_base = "http://localhost:19530" + + mock_request = MagicMock(spec=Request) + mock_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + mock_index_object = MagicMock() + mock_index_object.litellm_params.vector_store_name = vector_store_name + mock_index_object.litellm_params.vector_store_index = vector_store_index + + mock_vector_store = {"litellm_params": {"api_base": api_base}} + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + return_value={"collectionName": collection_name}, + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config" + ) as mock_get_config, patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint" + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._safe_set_request_parsed_body" + ), patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_create_route, patch.object( + litellm, "vector_store_index_registry" + ) as mock_index_registry, patch.object( + litellm, "vector_store_registry" + ) as mock_vector_registry: + mock_provider_config = MagicMock() + mock_provider_config.get_auth_credentials.return_value = {"headers": {}} + mock_provider_config.get_complete_url.return_value = api_base + mock_get_config.return_value = mock_provider_config + + mock_index_registry.is_vector_store_index.return_value = True + mock_index_registry.get_vector_store_index_by_name.return_value = ( + mock_index_object + ) + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = ( + mock_vector_store + ) + + mock_endpoint_func = AsyncMock(return_value={"status": "success"}) + mock_create_route.return_value = mock_endpoint_func + + # Call with endpoint without leading slash + await milvus_proxy_route( + endpoint="vectors/search", # No leading slash + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify that the target URL has correct path + create_route_args = mock_create_route.call_args[1] + assert "/vectors/search" in create_route_args["target"]