scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Milvus - Passthrough API support - adds create + read vector store support via passthrough API's (#16170)

Browse source 
* feat(llm_passthrough_endpoints.py): support milvus passthrough api

* fix(llm_passthrough_endpoints.py): move streaming request value to the top of the function

* docs: document new milvus vector store passthrough flow
This commit is contained in:
Krish Dholakia 2025-11-02 09:47:58 -08:00 • committed by GitHub
parent 6ed76ff809
commit 07d2a27f14
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 922 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

309
docs/my-website/docs/providers/milvus_vector_stores.md
View file

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

8
litellm/proxy/_new_secret_config.yaml
View file

@ -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"
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

1
litellm/proxy/_types.py
View file

@ -346,6 +346,7 @@ class LiteLLMRoutes(enum.Enum):
"/eu.assemblyai",
"/vllm",
"/mistral",
"/milvus",
]
#########################################################

149
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

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

7
litellm/proxy/vector_store_endpoints/utils.py
View file

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

469
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -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"]