(feat) Milvus - search vector store support + (fix) Passthrough Endpoints - support multi-part form data on passthrough (#16035)

* feat(milvus/): initial commit adding milvus vector store support to LiteLLM

allows querying milvus vector store through litellm

* feat(bedrock/vector_stores): support translating openai filters param to aws kb

adds filtering to aws kb

* feat(milvus/): add milvus vector store unified search support

allows calling milvus vector store in through chat completions

* docs(milvus_vector_stores.md): document new milvus vector search integration

* feat(pass_through_endpoints.py): support passing form data through to a passthrough endpoint

Closes LIT-1147

* fix: fix linting errors
This commit is contained in:
Krish Dholakia 2025-11-01 12:00:29 -07:00 committed by GitHub
parent 22d35e2552
commit b02be1ba70
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 1487 additions and 124 deletions

View file

@ -138,6 +138,125 @@ print(response.choices[0].message.content)
</Tabs>
## Filter Results
Filter by metadata attributes.
**Operators** (OpenAI-style, auto-translated):
- `eq`, `ne`, `gt`, `gte`, `lt`, `lte`, `in`, `nin`
**AWS operators** (use directly):
- `equals`, `notEquals`, `greaterThan`, `greaterThanOrEquals`, `lessThan`, `lessThanOrEquals`, `in`, `notIn`, `startsWith`, `listContains`, `stringContains`
<Tabs>
<TabItem value="single-filter" label="Single Filter">
```python
response = await litellm.acompletion(
model="anthropic/claude-3-5-sonnet",
messages=[{"role": "user", "content": "What are the latest updates?"}],
tools=[{
"type": "file_search",
"vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"],
"filters": {
"key": "category",
"value": "updates",
"operator": "eq"
}
}]
)
```
</TabItem>
<TabItem value="and-filters" label="AND">
```python
response = await litellm.acompletion(
model="anthropic/claude-3-5-sonnet",
messages=[{"role": "user", "content": "What are the policies?"}],
tools=[{
"type": "file_search",
"vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"],
"filters": {
"and": [
{"key": "category", "value": "policy", "operator": "eq"},
{"key": "year", "value": 2024, "operator": "gte"}
]
}
}]
)
```
</TabItem>
<TabItem value="or-filters" label="OR">
```python
response = await litellm.acompletion(
model="anthropic/claude-3-5-sonnet",
messages=[{"role": "user", "content": "Show me technical docs"}],
tools=[{
"type": "file_search",
"vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"],
"filters": {
"or": [
{"key": "category", "value": "api", "operator": "eq"},
{"key": "category", "value": "sdk", "operator": "eq"}
]
}
}]
)
```
</TabItem>
<TabItem value="advanced-filters" label="AWS Operators">
```python
response = await litellm.acompletion(
model="anthropic/claude-3-5-sonnet",
messages=[{"role": "user", "content": "Find docs"}],
tools=[{
"type": "file_search",
"vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"],
"filters": {
"and": [
{"key": "title", "value": "Guide", "operator": "stringContains"},
{"key": "tags", "value": "important", "operator": "listContains"}
]
}
}]
)
```
</TabItem>
<TabItem value="proxy-filters" label="Proxy">
```bash
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $LITELLM_API_KEY" \
-d '{
"model": "claude-3-5-sonnet",
"messages": [{"role": "user", "content": "What are our policies?"}],
"tools": [{
"type": "file_search",
"vector_store_ids": ["YOUR_KNOWLEDGE_BASE_ID"],
"filters": {
"and": [
{"key": "department", "value": "engineering", "operator": "eq"},
{"key": "type", "value": "policy", "operator": "eq"}
]
}
}]
}'
```
</TabItem>
</Tabs>
## Accessing Search Results
See how to access vector store search results in your response:

View file

@ -0,0 +1,221 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Milvus - Vector Store
Use Milvus as a vector store for RAG.
## Quick Start
You need three things:
1. A Milvus instance (cloud or self-hosted)
2. An embedding model (to convert your queries to vectors)
3. A Milvus collection with vector fields
## Usage
<Tabs>
<TabItem value="sdk" label="SDK">
### Basic Search
```python
from litellm import vector_stores
import os
# Set your credentials
os.environ["MILVUS_API_KEY"] = "your-milvus-api-key"
os.environ["MILVUS_API_BASE"] = "https://your-milvus-instance.milvus.io"
# Search the vector store
response = vector_stores.search(
vector_store_id="my-collection-name", # Your Milvus collection name
query="What is the capital of France?",
custom_llm_provider="milvus",
litellm_embedding_model="azure/text-embedding-3-large",
litellm_embedding_config={
"api_base": "your-embedding-endpoint",
"api_key": "your-embedding-api-key",
"api_version": "2025-09-01"
},
milvus_text_field="book_intro", # Field name that contains text content
api_key=os.getenv("MILVUS_API_KEY"),
)
print(response)
```
### Async Search
```python
from litellm import vector_stores
response = await vector_stores.asearch(
vector_store_id="my-collection-name",
query="What is the capital of France?",
custom_llm_provider="milvus",
litellm_embedding_model="azure/text-embedding-3-large",
litellm_embedding_config={
"api_base": "your-embedding-endpoint",
"api_key": "your-embedding-api-key",
"api_version": "2025-09-01"
},
milvus_text_field="book_intro",
api_key=os.getenv("MILVUS_API_KEY"),
)
print(response)
```
### Advanced Options
```python
from litellm import vector_stores
response = vector_stores.search(
vector_store_id="my-collection-name",
query="What is the capital of France?",
custom_llm_provider="milvus",
litellm_embedding_model="azure/text-embedding-3-large",
litellm_embedding_config={
"api_base": "your-embedding-endpoint",
"api_key": "your-embedding-api-key",
},
milvus_text_field="book_intro",
api_key=os.getenv("MILVUS_API_KEY"),
# Milvus-specific parameters
limit=10, # Number of results to return
offset=0, # Pagination offset
dbName="default", # Database name
annsField="book_intro_vector", # Vector field name
outputFields=["id", "book_intro", "title"], # Fields to return
filter='book_id > 0', # Metadata filter expression
searchParams={"metric_type": "L2", "params": {"nprobe": 10}}, # Search parameters
)
print(response)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
### Setup Config
Add this to your config.yaml:
```yaml
vector_store_registry:
- vector_store_name: "milvus-knowledgebase"
litellm_params:
vector_store_id: "my-collection-name"
custom_llm_provider: "milvus"
api_key: os.environ/MILVUS_API_KEY
api_base: https://your-milvus-instance.milvus.io
litellm_embedding_model: "azure/text-embedding-3-large"
litellm_embedding_config:
api_base: https://your-endpoint.cognitiveservices.azure.com/
api_key: os.environ/AZURE_API_KEY
api_version: "2025-09-01"
milvus_text_field: "book_intro"
# Optional Milvus parameters
annsField: "book_intro_vector"
limit: 10
```
### Start Proxy
```bash
litellm --config /path/to/config.yaml
```
### Search via API
```bash
curl -X POST 'http://0.0.0.0:4000/v1/vector_stores/my-collection-name/search' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"query": "What is the capital of France?"
}'
```
</TabItem>
</Tabs>
## Required Parameters
| Parameter | Type | Description |
|-----------|------|-------------|
| `vector_store_id` | string | Your Milvus collection name |
| `custom_llm_provider` | string | Set to `"milvus"` |
| `litellm_embedding_model` | string | Model to generate query embeddings (e.g., `"azure/text-embedding-3-large"`) |
| `litellm_embedding_config` | dict | Config for the embedding model (api_base, api_key, api_version) |
| `milvus_text_field` | string | Field name in your collection that contains text content |
| `api_key` | string | Your Milvus API key (or set `MILVUS_API_KEY` env var) |
| `api_base` | string | Your Milvus API base URL (or set `MILVUS_API_BASE` env var) |
## Optional Parameters
| Parameter | Type | Description |
|-----------|------|-------------|
| `dbName` | string | Database name (default: "default") |
| `annsField` | string | Vector field name to search (default: "book_intro_vector") |
| `limit` | integer | Maximum number of results to return |
| `offset` | integer | Pagination offset |
| `filter` | string | Filter expression for metadata filtering |
| `groupingField` | string | Field to group results by |
| `outputFields` | list | List of fields to return in results |
| `searchParams` | dict | Search parameters like metric type and search parameters |
| `partitionNames` | list | List of partition names to search |
| `consistencyLevel` | string | Consistency level for the search |
## Supported Features
| Feature | Status | Notes |
|---------|--------|-------|
| Logging | ✅ Supported | Full logging support available |
| 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 | |
## Response Format
The response follows the standard LiteLLM vector store format:
```json
{
"object": "vector_store.search_results.page",
"search_query": "What is the capital of France?",
"data": [
{
"score": 0.95,
"content": [
{
"text": "Paris is the capital of France...",
"type": "text"
}
],
"file_id": null,
"filename": null,
"attributes": {
"id": "123",
"title": "France Geography"
}
}
]
}
```
## How It Works
When you search:
1. LiteLLM converts your query to a vector using the embedding model you specified
2. It sends the vector to your Milvus instance via the `/v2/vectordb/entities/search` endpoint
3. Milvus finds the most similar documents in your collection using vector similarity search
4. Results come back with distance scores
The embedding model can be any model supported by LiteLLM - Azure OpenAI, OpenAI, Bedrock, etc.

View file

@ -12,7 +12,7 @@ Search a vector store for relevant chunks based on a query and file attributes f
| Cost Tracking | ✅ | Tracked per search operation |
| Logging | ✅ | Works across all integrations |
| End-user Tracking | ✅ | |
| Support LLM Providers | **OpenAI, Azure OpenAI, Bedrock, Vertex RAG Engine, Azure AI** | Full vector stores API support across providers |
| Support LLM Providers | **OpenAI, Azure OpenAI, Bedrock, Vertex RAG Engine, Azure AI, Milvus** | Full vector stores API support across providers |
## Usage
@ -134,6 +134,36 @@ print(response)
[See full Azure AI vector store documentation](../providers/azure_ai_vector_stores.md)
</TabItem>
<TabItem value="milvus-provider" label="Milvus Provider">
#### Using Milvus
```python showLineNumbers title="Search Vector Store - Milvus Provider"
import litellm
import os
# Set credentials
os.environ["MILVUS_API_KEY"] = "your-milvus-api-key"
os.environ["MILVUS_API_BASE"] = "https://your-milvus-instance.milvus.io"
response = await litellm.vector_stores.asearch(
vector_store_id="my-collection-name",
query="What is the capital of France?",
custom_llm_provider="milvus",
litellm_embedding_model="azure/text-embedding-3-large",
litellm_embedding_config={
"api_base": "your-embedding-endpoint",
"api_key": "your-embedding-api-key",
},
milvus_text_field="book_intro",
api_key=os.getenv("MILVUS_API_KEY"),
)
print(response)
```
[See full Milvus vector store documentation](../providers/milvus_vector_stores.md)
</TabItem>
</Tabs>

View file

@ -495,6 +495,7 @@ const sidebars = {
"providers/bedrock_vector_store",
]
},
"providers/milvus_vector_stores",
"providers/litellm_proxy",
"providers/meta_llama",
"providers/mistral",

View file

@ -5,6 +5,7 @@ import httpx
from litellm.types.router import GenericLiteLLMParams
from litellm.types.vector_stores import (
VECTOR_STORE_OPENAI_PARAMS,
VectorStoreCreateOptionalRequestParams,
VectorStoreCreateResponse,
VectorStoreSearchOptionalRequestParams,
@ -24,6 +25,20 @@ else:
class BaseVectorStoreConfig:
def get_supported_openai_params(
self, model: str
) -> List[VECTOR_STORE_OPENAI_PARAMS]:
return []
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
drop_params: bool,
) -> dict:
return optional_params
@abstractmethod
def transform_search_vector_store_request(
self,
@ -34,6 +49,7 @@ class BaseVectorStoreConfig:
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> Tuple[str, Dict]:
pass
@abstractmethod

View file

@ -13,6 +13,7 @@ from litellm.types.integrations.rag.bedrock_knowledgebase import (
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.vector_stores import (
VECTOR_STORE_OPENAI_PARAMS,
VectorStoreResultContent,
VectorStoreSearchOptionalRequestParams,
VectorStoreSearchResponse,
@ -32,6 +33,134 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
BaseVectorStoreConfig.__init__(self)
BaseAWSLLM.__init__(self)
def get_supported_openai_params(
self, model: str
) -> List[VECTOR_STORE_OPENAI_PARAMS]:
return ["filters", "max_num_results", "ranking_options"]
def _map_operator_to_aws(self, operator: str) -> str:
"""
Map OpenAI-style operators to AWS Bedrock operator names.
OpenAI uses: eq, ne, gt, gte, lt, lte, in, nin
AWS uses: equals, notEquals, greaterThan, greaterThanOrEquals, lessThan, lessThanOrEquals, in, notIn, startsWith, listContains, stringContains
"""
operator_mapping = {
"eq": "equals",
"ne": "notEquals",
"gt": "greaterThan",
"gte": "greaterThanOrEquals",
"lt": "lessThan",
"lte": "lessThanOrEquals",
"in": "in",
"nin": "notIn",
# AWS-specific operators (pass through)
"equals": "equals",
"notEquals": "notEquals",
"greaterThan": "greaterThan",
"greaterThanOrEquals": "greaterThanOrEquals",
"lessThan": "lessThan",
"lessThanOrEquals": "lessThanOrEquals",
"notIn": "notIn",
"startsWith": "startsWith",
"listContains": "listContains",
"stringContains": "stringContains",
}
return operator_mapping.get(operator, operator)
def _map_operator_filter(self, filter_dict: dict) -> dict:
"""
Map a single OpenAI operator filter to AWS KB format.
OpenAI format: {"key": <key>, "value": <value>, "operator": <operator>}
AWS KB format: {"operator": {"key": <key>, "value": <value>}}
"""
aws_operator = self._map_operator_to_aws(filter_dict["operator"])
return {
aws_operator: {
"key": filter_dict["key"],
"value": filter_dict["value"],
}
}
def _map_and_or_filters(self, value: dict) -> dict:
"""
Map OpenAI and/or filters to AWS KB format.
OpenAI format: {"and" | "or": [{"key": <key>, "value": <value>, "operator": <operator>}]}
AWS KB format: {"andAll" | "orAll": [{"operator": {"key": <key>, "value": <value>}}]}
Note: AWS requires andAll/orAll to have at least 2 elements.
For single filters, unwrap and return just the operator.
"""
aws_filters = {}
if "and" in value:
and_filters = value["and"]
# If only 1 filter, return just the operator (AWS requires andAll to have >=2 elements)
if len(and_filters) == 1:
return self._map_operator_filter(and_filters[0])
aws_filters["andAll"] = [
{
self._map_operator_to_aws(and_filters[i]["operator"]): {
"key": and_filters[i]["key"],
"value": and_filters[i]["value"],
}
}
for i in range(len(and_filters))
]
if "or" in value:
or_filters = value["or"]
# If only 1 filter, return just the operator (AWS requires orAll to have >=2 elements)
if len(or_filters) == 1:
return self._map_operator_filter(or_filters[0])
aws_filters["orAll"] = [
{
self._map_operator_to_aws(or_filters[i]["operator"]): {
"key": or_filters[i]["key"],
"value": or_filters[i]["value"],
}
}
for i in range(len(or_filters))
]
return aws_filters
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
drop_params: bool,
) -> dict:
for param, value in non_default_params.items():
if param == "max_num_results":
optional_params["numberOfResults"] = value
elif param == "filters" and value is not None:
# map the openai filters to the aws kb filters format
# openai filters = {"key": <key>, "value": <value>, "operator": <operator>} OR {"and" | "or": [{"key": <key>, "value": <value>, "operator": <operator>}]}
# aws kb filters = {"operator": {"<key>": <value>}} OR {"andAll | orAll": [{"operator": {"<key>": <value>}}]}
# 1. check if filter is in openai format
# 2. if it is, map it to the aws kb filters format
# 3. if it is not, assume it is in aws kb filters format and add it to the optional_params
aws_filters: Optional[Dict] = None
if isinstance(value, dict):
if "operator" in value.keys():
# Single operator - map directly (no wrapping needed)
aws_filters = self._map_operator_filter(value)
elif "and" in value.keys() or "or" in value.keys():
aws_filters = self._map_and_or_filters(value)
else:
# Assume it's already in AWS KB format
aws_filters = value
optional_params["filters"] = aws_filters
return optional_params
def validate_environment(
self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]
) -> dict:
@ -39,13 +168,13 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
headers.setdefault("Content-Type", "application/json")
return headers
def get_complete_url(
self, api_base: Optional[str], litellm_params: dict
) -> str:
def get_complete_url(self, api_base: Optional[str], litellm_params: dict) -> str:
aws_region_name = litellm_params.get("aws_region_name")
endpoint_url, _ = self.get_runtime_endpoint(
api_base=api_base,
aws_bedrock_runtime_endpoint=litellm_params.get("aws_bedrock_runtime_endpoint"),
aws_bedrock_runtime_endpoint=litellm_params.get(
"aws_bedrock_runtime_endpoint"
),
aws_region_name=self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=aws_region_name
),
@ -86,7 +215,9 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
# Create a properly typed retrieval configuration
typed_retrieval_config: BedrockKBRetrievalConfiguration = {}
if "vectorSearchConfiguration" in retrieval_config:
typed_retrieval_config["vectorSearchConfiguration"] = retrieval_config["vectorSearchConfiguration"]
typed_retrieval_config["vectorSearchConfiguration"] = retrieval_config[
"vectorSearchConfiguration"
]
request_body["retrievalConfiguration"] = typed_retrieval_config
litellm_logging_obj.model_call_details["query"] = query
@ -117,8 +248,12 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
source_uri = metadata.get("x-amz-bedrock-kb-source-uri", "") if metadata else ""
if source_uri:
return source_uri
chunk_id = metadata.get("x-amz-bedrock-kb-chunk-id", "unknown") if metadata else "unknown"
chunk_id = (
metadata.get("x-amz-bedrock-kb-chunk-id", "unknown")
if metadata
else "unknown"
)
return f"bedrock-kb-{chunk_id}"
def _get_filename_from_metadata(self, metadata: Dict[str, Any]) -> str:
@ -127,18 +262,26 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
Tries to extract filename from source URI, falls back to domain name or data source ID.
"""
source_uri = metadata.get("x-amz-bedrock-kb-source-uri", "") if metadata else ""
if source_uri:
try:
parsed_uri = urlparse(source_uri)
filename = parsed_uri.path.split('/')[-1] if parsed_uri.path and parsed_uri.path != '/' else parsed_uri.netloc
if not filename or filename == '/':
filename = (
parsed_uri.path.split("/")[-1]
if parsed_uri.path and parsed_uri.path != "/"
else parsed_uri.netloc
)
if not filename or filename == "/":
filename = parsed_uri.netloc
return filename
except Exception:
return source_uri
data_source_id = metadata.get("x-amz-bedrock-kb-data-source-id", "unknown") if metadata else "unknown"
data_source_id = (
metadata.get("x-amz-bedrock-kb-data-source-id", "unknown")
if metadata
else "unknown"
)
return f"bedrock-kb-document-{data_source_id}"
def _get_attributes_from_metadata(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
@ -161,13 +304,13 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
text = content.get("text") if content else None
if text is None:
continue
# Extract metadata and use helper functions
metadata = item.get("metadata", {}) or {}
file_id = self._get_file_id_from_metadata(metadata)
filename = self._get_filename_from_metadata(metadata)
attributes = self._get_attributes_from_metadata(metadata)
results.append(
VectorStoreSearchResult(
score=item.get("score"),

View file

@ -18,7 +18,7 @@ from litellm.constants import (
AIOHTTP_CONNECTOR_LIMIT,
AIOHTTP_KEEPALIVE_TIMEOUT,
AIOHTTP_TTL_DNS_CACHE,
DEFAULT_SSL_CIPHERS
DEFAULT_SSL_CIPHERS,
)
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.types.llms.custom_http import *
@ -141,11 +141,11 @@ def get_ssl_configuration(
if ssl_verify is not False:
custom_ssl_context = ssl.create_default_context(cafile=cafile)
# Optimize SSL handshake performance
# Set minimum TLS version to 1.2 for better performance
custom_ssl_context.minimum_version = ssl.TLSVersion.TLSv1_2
# Configure cipher suites for optimal performance
if ssl_security_level and isinstance(ssl_security_level, str):
# User provided custom cipher configuration (e.g., via SSL_SECURITY_LEVEL env var)
@ -753,7 +753,7 @@ class AsyncHTTPHandler:
keepalive_timeout=AIOHTTP_KEEPALIVE_TIMEOUT,
ttl_dns_cache=AIOHTTP_TTL_DNS_CACHE,
enable_cleanup_closed=True,
**connector_kwargs
**connector_kwargs,
),
trust_env=trust_env,
),
@ -1043,7 +1043,7 @@ class HTTPHandler:
if litellm.force_ipv4:
return HTTPTransport(local_address="0.0.0.0")
else:
return getattr(litellm, 'sync_transport', None)
return getattr(litellm, "sync_transport", None)
def get_async_httpx_client(

View file

@ -1593,8 +1593,6 @@ class BaseLLMHTTPHandler:
headers=headers or {},
)
# Transform the request
data = provider_config.transform_search_request(
query=query,
@ -1682,7 +1680,7 @@ class BaseLLMHTTPHandler:
query=query,
optional_params=optional_params,
)
# Get complete URL (pass data for providers that need request body for URL construction)
complete_url = provider_config.get_complete_url(
api_base=api_base,
@ -1704,6 +1702,7 @@ class BaseLLMHTTPHandler:
if client is None or not isinstance(client, AsyncHTTPHandler):
# For search providers, use special Search provider type
from litellm.types.llms.custom_http import httpxSpecialProvider
async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.Search
)
@ -1726,7 +1725,7 @@ class BaseLLMHTTPHandler:
response = await async_httpx_client.post(
url=complete_url,
headers=headers,
json=data,
json=data, # type: ignore
timeout=timeout,
)
except Exception as e:
@ -4079,16 +4078,16 @@ class BaseLLMHTTPHandler:
try:
# Use JSON when no files, otherwise use form data with files
if files and len(files) > 0:
# Use multipart/form-data when files are present
response = sync_httpx_client.post(
url=api_base,
headers=headers,
data=data,
files=files,
timeout=timeout,
)
# Use multipart/form-data when files are present
response = sync_httpx_client.post(
url=api_base,
headers=headers,
data=data,
files=files,
timeout=timeout,
)
# --- END MOCK VIDEO RESPONSE ---
# --- END MOCK VIDEO RESPONSE ---
else:
response = sync_httpx_client.post(
url=api_base,
@ -4355,7 +4354,7 @@ class BaseLLMHTTPHandler:
e=e,
provider_config=video_content_provider_config,
)
def video_remix_handler(
self,
video_id: str,
@ -4582,6 +4581,7 @@ class BaseLLMHTTPHandler:
else:
# For sync calls, we'll use the async handler in a sync context
import asyncio
return asyncio.run(
self.async_video_list_handler(
after=after,
@ -4682,7 +4682,7 @@ class BaseLLMHTTPHandler:
e=e,
provider_config=video_list_provider_config,
)
async def async_video_delete_handler(
self,
video_id: str,
@ -4820,12 +4820,14 @@ class BaseLLMHTTPHandler:
)
# Transform the request using the provider config
url, data = video_status_provider_config.transform_video_status_retrieve_request(
video_id=video_id,
model=model,
api_base=api_base,
litellm_params=litellm_params,
headers=headers,
url, data = (
video_status_provider_config.transform_video_status_retrieve_request(
video_id=video_id,
model=model,
api_base=api_base,
litellm_params=litellm_params,
headers=headers,
)
)
## LOGGING
@ -4845,10 +4847,12 @@ class BaseLLMHTTPHandler:
headers=headers,
)
return video_status_provider_config.transform_video_status_retrieve_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
return (
video_status_provider_config.transform_video_status_retrieve_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
)
)
except Exception as e:
@ -4898,12 +4902,14 @@ class BaseLLMHTTPHandler:
)
# Transform the request using the provider config
url, data = video_status_provider_config.transform_video_status_retrieve_request(
video_id=video_id,
model=model,
api_base=api_base,
litellm_params=litellm_params,
headers=headers,
url, data = (
video_status_provider_config.transform_video_status_retrieve_request(
video_id=video_id,
model=model,
api_base=api_base,
litellm_params=litellm_params,
headers=headers,
)
)
## LOGGING
@ -4923,10 +4929,12 @@ class BaseLLMHTTPHandler:
headers=headers,
)
return video_status_provider_config.transform_video_status_retrieve_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
return (
video_status_provider_config.transform_video_status_retrieve_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
)
)
except Exception as e:
@ -5006,6 +5014,7 @@ class BaseLLMHTTPHandler:
)
try:
response = await async_httpx_client.post(
url=url,
headers=headers,

View file

@ -0,0 +1,3 @@
from litellm.llms.milvus.vector_stores.transformation import MilvusVectorStoreConfig
__all__ = ["MilvusVectorStoreConfig"]

View file

@ -0,0 +1,252 @@
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
import litellm
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
VectorStoreCreateResponse,
VectorStoreResultContent,
VectorStoreSearchOptionalRequestParams,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
MILVUS_OPTIONAL_PARAMS = {
"dbName",
"annsField",
"limit",
"filter",
"offset",
"groupingField",
"outputFields",
"searchParams",
"partitionNames",
"consistencyLevel",
}
class MilvusVectorStoreConfig(BaseVectorStoreConfig):
"""
Configuration for Milvus Vector Store
This implementation uses the Azure AI Search API for vector store operations.
Supports vector search with embeddings generated via litellm.embeddings.
"""
def __init__(self):
super().__init__()
def validate_environment(
self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]
) -> dict:
api_key: Optional[str] = None
if litellm_params is not None:
api_key = litellm_params.api_key or get_secret_str("MILVUS_API_KEY")
if not api_key:
raise ValueError(
"MILVUS_API_KEY is not set. Either set it in the litellm_params or set the MILVUS_API_KEY environment variable."
)
headers.update({"Authorization": f"Bearer {api_key}"})
return headers
def map_openai_params(
self, non_default_params: dict, optional_params: dict, drop_params: bool
) -> dict:
for param, value in non_default_params.items():
if param in MILVUS_OPTIONAL_PARAMS:
optional_params[param] = value
return optional_params
def get_complete_url(
self,
api_base: Optional[str],
litellm_params: dict,
) -> str:
"""
Get the base endpoint for Milvus API
Expected format: https://{milvus_api_base}.milvus.io
"""
api_base = api_base or get_secret_str("MILVUS_API_BASE")
if not api_base:
raise ValueError(
"Milvus API base URL is required. Set MILVUS_API_BASE environment variable or pass api_base in litellm_params."
)
if api_base:
return api_base.rstrip("/")
return api_base
def transform_search_vector_store_request(
self,
vector_store_id: str,
query: Union[str, List[str]],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> Tuple[str, Dict[str, Any]]:
"""
Transform search request for Azure AI Search API
Generates embeddings using litellm.embeddings and constructs Azure AI Search request
"""
# Convert query to string if it's a list
if isinstance(query, list):
query = " ".join(query)
# Get embedding model from litellm_params (required)
embedding_model = litellm_params.get("litellm_embedding_model")
if not embedding_model:
raise ValueError(
"embedding_model is required in litellm_params for Milvus. You can call any litellm embedding model."
"Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'"
)
embedding_config = litellm_params.get("litellm_embedding_config", {})
if not embedding_config:
raise ValueError(
"embedding_config is required in litellm_params for Milvus. You can call any litellm embedding model."
"Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}"
)
# Get top_k (number of results to return)
# Generate embedding for the query using litellm.embeddings
try:
embedding_response = litellm.embedding(
model=embedding_model,
input=[query],
**embedding_config,
)
query_vector = embedding_response.data[0]["embedding"]
except Exception as e:
raise Exception(f"Failed to generate embedding for query: {str(e)}")
# Azure AI Search endpoint for search
index_name = vector_store_id # vector_store_id is the index name
url = f"{api_base}/v2/vectordb/entities/search"
# Build the request body for Azure AI Search with vector search
request_body = {
"collectionName": index_name,
"data": [query_vector],
"annsField": "book_intro_vector",
**vector_store_search_optional_params,
}
#########################################################
# Update logging object with details of the request
#########################################################
litellm_logging_obj.model_call_details["input"] = query
litellm_logging_obj.model_call_details["embedding_model"] = embedding_model
return url, request_body
def transform_search_vector_store_response(
self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj
) -> VectorStoreSearchResponse:
"""
Transform Azure AI Search API response to standard vector store search response
Handles the format from Azure AI Search which returns:
{
"value": [
{
"id": "...",
"content": "...",
"distance": 0.95,
}
]
}
"""
try:
response_json = response.json()
# Extract results from Azure AI Search API response
results = response_json.get("data", [])
# Try to get text_field from optional_params first, then litellm_params
optional_params = litellm_logging_obj.model_call_details.get(
"optional_params", {}
)
text_field = optional_params.get("milvus_text_field", "")
# Fallback to litellm_params if not in optional_params
if not text_field:
text_field = litellm_logging_obj.model_call_details.get(
"litellm_params", {}
).get("milvus_text_field", "")
# Transform results to standard format
search_results: List[VectorStoreSearchResult] = []
for result in results:
# Extract text content
text_content = result.get(text_field, "")
content = [
VectorStoreResultContent(
text=text_content,
type="text",
)
]
# Get the search score (distance from the query vector)
score = result.get("distance", 0.0)
# Build attributes with all available metadata
# Exclude system fields and already-processed fields
attributes = {}
for key, value in result.items():
if key not in ["id", "content", "distance", text_field]:
attributes[key] = value
result_obj = VectorStoreSearchResult(
score=score,
content=content,
file_id=None,
filename=None,
attributes=attributes,
)
search_results.append(result_obj)
return VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query=litellm_logging_obj.model_call_details.get("input", ""),
data=search_results,
)
except Exception as e:
raise self.get_error_class(
error_message=str(e),
status_code=response.status_code,
headers=response.headers,
)
def transform_create_vector_store_request(
self,
vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams,
api_base: str,
) -> Tuple[str, Dict]:
raise NotImplementedError
def transform_create_vector_store_response(
self, response: httpx.Response
) -> VectorStoreCreateResponse:
raise NotImplementedError

View file

@ -9,21 +9,11 @@ model_list:
api_key: os.environ/AZURE_API_KEY
api_version: "2025-09-01"
guardrails:
- guardrail_name: "enkryptai-guard"
litellm_params:
guardrail: enkryptai
mode: "pre_call"
api_key: os.environ/ENKRYPTAI_API_KEY
default_on: true
policy_name: "Sample Airline Guardrail"
detectors:
toxicity:
enabled: true
nsfw:
enabled: true
pii:
enabled: true
entities: ["email", "phone", "secrets"]
injection_attack:
enabled: true
general_settings:
master_key: sk-1234
pass_through_endpoints:
- path: "/v1/my-custom-path" # Route on LiteLLM Proxy
target: "http://0.0.0.0:8089/v1/my-custom-path" # Target endpoint
headers: # Headers to forward
Authorization: "bearer sk-1234"
forward_headers: true # Forward all incoming headers

View file

@ -441,10 +441,15 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
else:
form_data_dict[field_name] = field_value
# Remove content-type header - httpx will set it correctly with the new boundary
# when it creates the multipart body from files/data parameters
headers_copy = headers.copy()
headers_copy.pop("content-type", None)
response = await async_client.request(
method=request.method,
url=url,
headers=headers,
headers=headers_copy,
params=requested_query_params,
files=files,
data=form_data_dict,
@ -925,6 +930,69 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di
return metadata
async def _parse_request_data_by_content_type(
request: Request,
) -> Tuple[Optional[Any], Optional[Any], Optional[Any], Optional[Any]]:
"""
Parse request data based on content type.
Handles JSON, multipart/form-data, and URL-encoded form data.
Returns:
Tuple of (query_params_data, custom_body_data, file_data, stream)
"""
content_type = request.headers.get("content-type", "")
query_params_data = None
custom_body_data = None
file_data = None
stream = None
if "application/json" in content_type:
# ✅ Handle JSON
body = await request.json()
query_params_data = body.get("query_params")
custom_body_data = body.get("custom_body")
stream = body.get("stream")
elif "multipart/form-data" in content_type:
# ✅ Handle multipart form-data
form = await request.form()
if "query_params" in form:
form_value = form["query_params"]
if isinstance(form_value, str):
try:
query_params_data = json.loads(form_value)
except Exception:
query_params_data = form_value
else:
query_params_data = form_value
if "custom_body" in form:
form_value = form["custom_body"]
if isinstance(form_value, str):
try:
custom_body_data = json.loads(form_value)
except Exception:
custom_body_data = form_value
else:
custom_body_data = form_value
if "file" in form:
file_data = form["file"] # this is a Starlette UploadFile object
elif "application/x-www-form-urlencoded" in content_type:
# ✅ Handle URL-encoded form data
form = await request.form()
query_params_data = form.get("query_params")
custom_body_data = form.get("custom_body")
else:
# ✅ Fallback: maybe no body, just query params
query_params_data = dict(request.query_params) or None
return query_params_data, custom_body_data, file_data, stream
def create_pass_through_route(
endpoint,
target: str,
@ -968,11 +1036,6 @@ def create_pass_through_route(
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
query_params: Optional[dict] = None,
custom_body: Optional[dict] = None,
stream: Optional[
bool
] = None, # if pass-through endpoint is a streaming request
subpath: str = "", # captures sub-paths when include_subpath=True
):
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
@ -981,6 +1044,14 @@ def create_pass_through_route(
path = request.url.path
# Parse request data based on content type
(
query_params_data,
custom_body_data,
file_data,
stream,
) = await _parse_request_data_by_content_type(request)
if not InitPassThroughEndpointHelpers.is_registered_pass_through_route(
route=path
):
@ -1032,6 +1103,18 @@ def create_pass_through_route(
param_custom_headers if isinstance(param_custom_headers, dict) else {}
)
# Ensure query_params and custom_body are dicts or None
final_query_params = (
query_params_data
if isinstance(query_params_data, dict) or query_params_data is None
else None
)
final_custom_body = (
custom_body_data
if isinstance(custom_body_data, dict) or custom_body_data is None
else None
)
return await pass_through_request( # type: ignore
request=request,
target=full_target,
@ -1039,9 +1122,9 @@ def create_pass_through_route(
user_api_key_dict=user_api_key_dict,
forward_headers=cast(Optional[bool], param_forward_headers),
merge_query_params=cast(Optional[bool], param_merge_query_params),
query_params=query_params,
query_params=final_query_params,
stream=stream,
custom_body=custom_body,
custom_body=final_custom_body,
cost_per_request=cast(Optional[float], param_cost_per_request),
custom_llm_provider=custom_llm_provider,
)

View file

@ -209,6 +209,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
# Vector Store Params
vector_store_id: Optional[str] = None
milvus_text_field: Optional[str] = None
def __init__(
self,
@ -607,6 +608,7 @@ class SearchToolLiteLLMParams(TypedDict, total=False):
LiteLLM params for search tools.
Search tools don't require a 'model' field like regular deployments.
"""
search_provider: Required[SearchProvider]
api_key: Optional[str]
api_base: Optional[str]
@ -617,7 +619,7 @@ class SearchToolLiteLLMParams(TypedDict, total=False):
class SearchToolTypedDict(TypedDict):
"""
Configuration for a search tool in the router.
Example:
{
"search_tool_name": "litellm-search",
@ -627,6 +629,7 @@ class SearchToolTypedDict(TypedDict):
}
}
"""
search_tool_name: Required[str]
litellm_params: Required[SearchToolLiteLLMParams]

View file

@ -280,7 +280,7 @@ class CallTypes(str, Enum):
file_content = "file_content"
create_fine_tuning_job = "create_fine_tuning_job"
acreate_fine_tuning_job = "acreate_fine_tuning_job"
#########################################################
# Video Generation Call Types
#########################################################
@ -1179,8 +1179,6 @@ class StreamingChatCompletionChunk(OpenAIChatCompletionChunk):
super().__init__(**kwargs)
class ModelResponseBase(OpenAIObject):
id: str
"""A unique identifier for the completion."""
@ -2501,6 +2499,7 @@ class LlmProviders(str, Enum):
DEEPINFRA = "deepinfra"
PERPLEXITY = "perplexity"
MISTRAL = "mistral"
MILVUS = "milvus"
GROQ = "groq"
NVIDIA_NIM = "nvidia_nim"
CEREBRAS = "cerebras"
@ -2576,6 +2575,7 @@ class SearchProviders(str, Enum):
Enum for search provider types.
Separate from LlmProviders for semantic clarity.
"""
PERPLEXITY = "perplexity"
TAVILY = "tavily"
PARALLEL_AI = "parallel_ai"
@ -2822,13 +2822,13 @@ CostResponseTypes = Union[
class PriorityReservationDict(TypedDict, total=False):
"""
Dictionary format for priority reservation values.
Used in litellm.priority_reservation to specify how much capacity to reserve
for each priority level. Supports three formats:
1. Percentage-based: {"type": "percent", "value": 0.9} -> 90% of capacity
2. RPM-based: {"type": "rpm", "value": 900} -> 900 requests per minute
3. TPM-based: {"type": "tpm", "value": 900000} -> 900,000 tokens per minute
Attributes:
type: The type of value - "percent", "rpm", or "tpm". Defaults to "percent".
value: The numeric value. For percent (0.0-1.0), for rpm/tpm (absolute value).

View file

@ -76,6 +76,7 @@ class VectorStoreResultContent(TypedDict, total=False):
class VectorStoreSearchResult(TypedDict, total=False):
"""Result of a vector store search"""
score: Optional[float]
content: Optional[List[VectorStoreResultContent]]
file_id: Optional[str]
@ -92,44 +93,55 @@ class VectorStoreSearchResponse(TypedDict, total=False):
search_query: Optional[str]
data: Optional[List[VectorStoreSearchResult]]
class VectorStoreSearchOptionalRequestParams(TypedDict, total=False):
"""TypedDict for Optional parameters supported by the vector store search API."""
filters: Optional[Dict]
max_num_results: Optional[int]
max_num_results: Optional[int]
ranking_options: Optional[Dict]
rewrite_query: Optional[bool]
class VectorStoreSearchRequest(VectorStoreSearchOptionalRequestParams, total=False):
"""Request body for searching a vector store"""
query: Union[str, List[str]]
# Vector Store Creation Types
class VectorStoreExpirationPolicy(TypedDict, total=False):
"""The expiration policy for a vector store"""
anchor: Literal["last_active_at"] # Anchor timestamp after which the expiration policy applies
anchor: Literal[
"last_active_at"
] # Anchor timestamp after which the expiration policy applies
days: int # Number of days after anchor time that the vector store will expire
class VectorStoreAutoChunkingStrategy(TypedDict, total=False):
"""Auto chunking strategy configuration"""
type: Literal["auto"] # Always "auto"
class VectorStoreStaticChunkingStrategyConfig(TypedDict, total=False):
"""Static chunking strategy configuration"""
max_chunk_size_tokens: int # Maximum number of tokens per chunk
chunk_overlap_tokens: int # Number of tokens to overlap between chunks
class VectorStoreStaticChunkingStrategy(TypedDict, total=False):
"""Static chunking strategy"""
type: Literal["static"] # Always "static"
static: VectorStoreStaticChunkingStrategyConfig
class VectorStoreChunkingStrategy(TypedDict, total=False):
"""Union type for chunking strategies"""
# This can be either auto or static
type: Literal["auto", "static"]
static: Optional[VectorStoreStaticChunkingStrategyConfig]
@ -137,6 +149,7 @@ class VectorStoreChunkingStrategy(TypedDict, total=False):
class VectorStoreFileCounts(TypedDict, total=False):
"""File counts for a vector store"""
in_progress: int
completed: int
failed: int
@ -146,20 +159,27 @@ class VectorStoreFileCounts(TypedDict, total=False):
class VectorStoreCreateOptionalRequestParams(TypedDict, total=False):
"""TypedDict for Optional parameters supported by the vector store create API."""
name: Optional[str] # Name of the vector store
file_ids: Optional[List[str]] # List of File IDs that the vector store should use
expires_after: Optional[VectorStoreExpirationPolicy] # Expiration policy for the vector store
chunking_strategy: Optional[VectorStoreChunkingStrategy] # Chunking strategy for the files
expires_after: Optional[
VectorStoreExpirationPolicy
] # Expiration policy for the vector store
chunking_strategy: Optional[
VectorStoreChunkingStrategy
] # Chunking strategy for the files
metadata: Optional[Dict[str, str]] # Set of key-value pairs for metadata
class VectorStoreCreateRequest(VectorStoreCreateOptionalRequestParams, total=False):
"""Request body for creating a vector store"""
pass # All fields are optional for vector store creation
class VectorStoreCreateResponse(TypedDict, total=False):
"""Response after creating a vector store"""
id: str # ID of the vector store
object: Literal["vector_store"] # Always "vector_store"
created_at: int # Unix timestamp of when the vector store was created
@ -169,5 +189,15 @@ class VectorStoreCreateResponse(TypedDict, total=False):
status: Literal["expired", "in_progress", "completed"] # Status of the vector store
expires_after: Optional[VectorStoreExpirationPolicy] # Expiration policy
expires_at: Optional[int] # Unix timestamp of when the vector store expires
last_active_at: Optional[int] # Unix timestamp of when the vector store was last active
metadata: Optional[Dict[str, str]] # Metadata associated with the vector store
last_active_at: Optional[
int
] # Unix timestamp of when the vector store was last active
metadata: Optional[Dict[str, str]] # Metadata associated with the vector store
VECTOR_STORE_OPENAI_PARAMS = Literal[
"filters",
"max_num_results",
"ranking_options",
"rewrite_query",
]

View file

@ -5009,7 +5009,9 @@ def _get_model_info_helper( # noqa: PLR0915
"output_cost_per_token_above_200k_tokens", None
),
output_cost_per_second=_model_info.get("output_cost_per_second", None),
output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None),
output_cost_per_video_per_second=_model_info.get(
"output_cost_per_video_per_second", None
),
output_cost_per_image=_model_info.get("output_cost_per_image", None),
output_vector_size=_model_info.get("output_vector_size", None),
citation_cost_per_token=_model_info.get(
@ -7584,6 +7586,12 @@ class ProviderConfigManager:
)
return AzureAIVectorStoreConfig()
elif litellm.LlmProviders.MILVUS == provider:
from litellm.llms.milvus.vector_stores.transformation import (
MilvusVectorStoreConfig,
)
return MilvusVectorStoreConfig()
return None
@staticmethod
@ -7668,7 +7676,6 @@ class ProviderConfigManager:
return AzureVideoConfig()
return None
@staticmethod
def get_provider_realtime_config(
model: str,

View file

@ -430,7 +430,8 @@ def search(
# Get VectorStoreSearchOptionalRequestParams with only valid parameters
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams = (
VectorStoreRequestUtils.get_requested_vector_store_search_optional_param(
local_vars
local_vars,
vector_store_provider_config=vector_store_provider_config,
)
)
@ -445,6 +446,7 @@ def search(
litellm_params={
"litellm_call_id": litellm_call_id,
"vector_store_id": vector_store_id,
**litellm_params.model_dump(exclude_none=True),
},
custom_llm_provider=custom_llm_provider,
)

View file

@ -1,5 +1,6 @@
from typing import Any, Dict, cast, get_type_hints
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
VectorStoreSearchOptionalRequestParams,
@ -12,6 +13,7 @@ class VectorStoreRequestUtils:
@staticmethod
def get_requested_vector_store_search_optional_param(
params: Dict[str, Any],
vector_store_provider_config: BaseVectorStoreConfig,
) -> VectorStoreSearchOptionalRequestParams:
"""
Filter parameters to only include those defined in VectorStoreSearchOptionalRequestParams.
@ -27,7 +29,13 @@ class VectorStoreRequestUtils:
k: v for k, v in params.items() if k in valid_keys and v is not None
}
return cast(VectorStoreSearchOptionalRequestParams, filtered_params)
optional_params = vector_store_provider_config.map_openai_params(
non_default_params=params,
optional_params=filtered_params,
drop_params=False,
)
return cast(VectorStoreSearchOptionalRequestParams, optional_params)
@staticmethod
def get_requested_vector_store_create_optional_param(
@ -48,4 +56,3 @@ class VectorStoreRequestUtils:
}
return cast(VectorStoreCreateOptionalRequestParams, filtered_params)

View file

@ -120,6 +120,80 @@ async def test_make_multipart_http_request():
assert call_args["data"]["text_field"] == "test value"
@pytest.mark.asyncio
async def test_make_multipart_http_request_removes_content_type_header():
"""
Test that make_multipart_http_request removes the content-type header
to prevent boundary mismatch errors.
When forwarding multipart requests, the original content-type header contains
a boundary that doesn't match the new boundary httpx generates. This test
verifies that the content-type header is removed so httpx can set it correctly.
"""
# Mock request with form data
request = MagicMock(spec=Request)
request.method = "POST"
# Mock form data with both file and regular field
file_content = b"test file content"
file = BytesIO(file_content)
headers = Headers({"content-type": "text/plain"})
upload_file = UploadFile(file=file, filename="test.txt", headers=headers)
upload_file.read = AsyncMock(return_value=file_content)
form_data = {"file": upload_file, "key": "value"}
request.form = AsyncMock(return_value=form_data)
# Mock httpx client
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
async_client = MagicMock()
async_client.request = AsyncMock(return_value=mock_response)
# Headers with content-type containing old boundary (this is what causes the issue)
original_headers = {
"content-type": "multipart/form-data; boundary=--------------------------416423083260054165225918",
"user-agent": "PostmanRuntime/7.49.0",
"Authorization": "bearer sk-1234",
}
# Test the function
response = await HttpPassThroughEndpointHelpers.make_multipart_http_request(
request=request,
async_client=async_client,
url=httpx.URL("http://test.com"),
headers=original_headers,
requested_query_params={"param": "value"},
)
# Verify the response
assert response == mock_response
# Verify the client call
async_client.request.assert_called_once()
call_args = async_client.request.call_args[1]
# CRITICAL ASSERTION: content-type header should be removed
assert "content-type" not in call_args["headers"]
# Other headers should be preserved
assert call_args["headers"]["user-agent"] == "PostmanRuntime/7.49.0"
assert call_args["headers"]["Authorization"] == "bearer sk-1234"
# Verify other parameters are correct
assert call_args["method"] == "POST"
assert str(call_args["url"]) == "http://test.com"
assert isinstance(call_args["files"], dict)
assert isinstance(call_args["data"], dict)
assert call_args["data"]["key"] == "value"
assert call_args["params"] == {"param": "value"}
# Verify the original headers dict was not modified (copy was used)
assert "content-type" in original_headers
@pytest.mark.asyncio
async def test_pass_through_request_failure_handler():
"""
@ -1420,7 +1494,7 @@ async def test_pass_through_with_httpbin_redirect():
assert response.status_code == 200
# The response should be from the /get endpoint
response_content = response.body.decode("utf-8")
response_content = bytes(response.body).decode("utf-8")
# httpbin.org/get returns JSON with info about the request
assert '"url": "https://httpbin.org/get"' in response_content
@ -1699,11 +1773,11 @@ async def test_filter_endpoints_by_team_allowed_routes_partial_match():
async def test_bedrock_router_passthrough_metadata_initialization():
"""
Test that bedrock router passthrough properly initializes metadata for hooks.
This test verifies the fix for issue #15826 where metadata.headers and
This test verifies the fix for issue #15826 where metadata.headers and
litellm_params.proxy_server_request were missing for /bedrock passthrough
requests with router models.
The fix ensures router bedrock models use the same common processing path
as non-router models, which properly initializes all metadata structures.
"""
@ -1718,46 +1792,50 @@ async def test_bedrock_router_passthrough_metadata_initialization():
# Setup mock instance
mock_processor = MagicMock()
mock_processing_class.return_value = mock_processor
# Mock successful response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.aread = AsyncMock(return_value=b'{"content": [{"text": "Hello"}]}')
mock_response.aread = AsyncMock(
return_value=b'{"content": [{"text": "Hello"}]}'
)
mock_processor.base_passthrough_process_llm_request = AsyncMock(
return_value=mock_response
)
# Create mock request with headers
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://localhost:4000/bedrock/model/my-model/invoke"
mock_request.headers = Headers({
"content-type": "application/json",
"authorization": "Bearer sk-test-key",
"x-custom-header": "test-value"
})
mock_request.headers = Headers(
{
"content-type": "application/json",
"authorization": "Bearer sk-test-key",
"x-custom-header": "test-value",
}
)
mock_request.query_params = QueryParams({})
# Create mock user API key dict with all required fields
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.api_key = "sk-test-key"
mock_user_api_key_dict.key_alias = "test-alias"
mock_user_api_key_dict.user_id = "user-123"
mock_user_api_key_dict.team_id = "team-123"
# Mock other required dependencies
mock_router = MagicMock()
mock_proxy_logging = MagicMock()
mock_general_settings = {}
mock_proxy_config = MagicMock()
mock_select_data_generator = MagicMock()
request_body = {
"max_tokens": 100,
"messages": [{"role": "user", "content": "Hello"}],
"anthropic_version": "bedrock-2023-05-31"
"anthropic_version": "bedrock-2023-05-31",
}
# Call the function
result = await handle_bedrock_passthrough_router_model(
model="my-bedrock-model",
@ -1777,24 +1855,32 @@ async def test_bedrock_router_passthrough_metadata_initialization():
user_api_base=None,
version="1.0",
)
# Verify that ProxyBaseLLMRequestProcessing was instantiated
# This is the KEY assertion - router models now use the common processing path
mock_processing_class.assert_called_once()
# Verify that base_passthrough_process_llm_request was called
# This proves we're using the common processing path that initializes metadata
mock_processor.base_passthrough_process_llm_request.assert_called_once()
# Verify the call included all required parameters for proper metadata initialization
call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args[1]
# These are the critical parameters that ensure metadata is properly initialized:
assert call_kwargs["request"] == mock_request, "Request must be passed for header extraction"
assert call_kwargs["user_api_key_dict"] == mock_user_api_key_dict, "User API key dict needed for metadata"
assert call_kwargs["proxy_logging_obj"] == mock_proxy_logging, "Logging obj needed for hooks"
assert call_kwargs["llm_router"] == mock_router, "Router needed for model routing"
assert (
call_kwargs["request"] == mock_request
), "Request must be passed for header extraction"
assert (
call_kwargs["user_api_key_dict"] == mock_user_api_key_dict
), "User API key dict needed for metadata"
assert (
call_kwargs["proxy_logging_obj"] == mock_proxy_logging
), "Logging obj needed for hooks"
assert (
call_kwargs["llm_router"] == mock_router
), "Router needed for model routing"
assert call_kwargs["model"] == "my-bedrock-model", "Model name must be passed"
# Verify response was returned
assert result == mock_response

View file

@ -0,0 +1,361 @@
"""
Tests for Milvus Vector Store
"""
import json
import os
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm.vector_stores import asearch as vector_store_asearch
from litellm.vector_stores import search as vector_store_search
# Mock response from actual Milvus API
MOCK_MILVUS_SEARCH_RESPONSE = {
"code": 0,
"cost": 6,
"data": [
{
"book_id": 0,
"book_intro_text": "abababababa_0562efee-0f1f-4b6b-9ca3-1a160f124ad8",
"distance": 10.240219,
},
{
"book_id": 1,
"book_intro_text": "abababababa_9a13e8f3-bb1e-487f-b555-b8ae4b127243",
"distance": 10.240219,
},
{
"book_id": 2,
"book_intro_text": "abababababa_870f47f1-23ec-4364-ad30-6d364ba8ddb5",
"distance": 10.240219,
},
{
"book_id": 1000,
"book_intro_text": "abababababa_8ea2d76a-3fdf-49b3-8f16-a91638361bba",
"distance": 8.531628,
},
{
"book_id": 1001,
"book_intro_text": "abababababa_24758251-e740-4183-8649-2f742f676ca0",
"distance": 8.531628,
},
{
"book_id": 1002,
"book_intro_text": "abababababa_faa55789-220d-4ef1-b5bf-a72f2fbd061b",
"distance": 8.531628,
},
{
"book_id": 0,
"book_intro_text": "abababababa_0562efee-0f1f-4b6b-9ca3-1a160f124ad8",
"distance": 8.236887,
},
{
"book_id": 1,
"book_intro_text": "abababababa_9a13e8f3-bb1e-487f-b555-b8ae4b127243",
"distance": 8.236887,
},
{
"book_id": 2,
"book_intro_text": "abababababa_870f47f1-23ec-4364-ad30-6d364ba8ddb5",
"distance": 8.236887,
},
],
"topks": [3, 3, 3],
}
# Mock embedding response from OpenAI
MOCK_EMBEDDING_RESPONSE = MagicMock()
MOCK_EMBEDDING_RESPONSE.data = [
{
"embedding": [
0.023,
-0.019,
0.045,
-0.012,
0.067,
-0.034,
0.089,
-0.056,
]
* 128 # Simulate 1024-dimensional embedding
}
]
class TestMilvusVectorStore:
"""Test Milvus Vector Store with mocked responses"""
@pytest.mark.asyncio
async def test_basic_search_with_mock_async(self):
"""Test basic vector search with mocked backend response (async)"""
# Mock the HTTP response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = MOCK_MILVUS_SEARCH_RESPONSE
mock_response.text = json.dumps(MOCK_MILVUS_SEARCH_RESPONSE)
with patch("litellm.embedding") as mock_embedding:
mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post:
mock_post.return_value = mock_response
# Make the search request
response = await vector_store_asearch(
query="what is machine learning?",
vector_store_id="book_2",
custom_llm_provider="milvus",
api_base="https://in03-test.serverless.aws-eu-central-1.cloud.zilliz.com",
api_key="mock_milvus_api_key",
litellm_embedding_model="text-embedding-3-large",
litellm_embedding_config={
"api_key": "mock_openai_api_key",
},
outputFields=["book_intro_text"],
annsField="book_intro_vector",
milvus_text_field="book_intro_text",
)
print("Response:", json.dumps(response, indent=2, default=str))
# Verify embedding was called with correct parameters
mock_embedding.assert_called_once()
embedding_call_args = mock_embedding.call_args
assert embedding_call_args[1]["model"] == "text-embedding-3-large"
assert embedding_call_args[1]["input"] == ["what is machine learning?"]
assert embedding_call_args[1]["api_key"] == "mock_openai_api_key"
# Verify the API was called
mock_post.assert_called_once()
# Verify the request payload
call_args = mock_post.call_args
print(f"call_args: {call_args}")
print(f"call_args.kwargs: {call_args.kwargs}")
# The post method is called with 'data' parameter (JSON string) not 'json' parameter
request_data_str = call_args.kwargs.get("data")
if request_data_str:
request_data = json.loads(request_data_str)
else:
# Fallback: check for json kwarg or in args
request_data = call_args.kwargs.get("json")
if (
request_data is None
and len(call_args.args) > 0
and isinstance(call_args.args[0], dict)
):
request_data = call_args.args[0]
assert (
request_data is not None
), f"Could not extract request data. Call args: {call_args}"
print("Request data:", json.dumps(request_data, indent=2, default=str))
# Validate request structure
assert "collectionName" in request_data
assert request_data["collectionName"] == "book_2"
assert "data" in request_data
assert isinstance(request_data["data"], list)
assert len(request_data["data"]) == 1 # Single query vector
assert "annsField" in request_data
assert request_data["annsField"] == "book_intro_vector"
assert "outputFields" in request_data
assert request_data["outputFields"] == ["book_intro_text"]
# Verify the URL format
url = call_args.kwargs.get("url", "")
assert "v2/vectordb/entities/search" in url
# Validate the response structure (LiteLLM standard format)
assert response is not None
assert response["object"] == "vector_store.search_results.page" # type: ignore
assert "data" in response
assert len(response["data"]) == 9 # type: ignore # 9 results in mock response
# Validate first result
first_result = response["data"][0] # type: ignore
assert "score" in first_result
assert first_result["score"] == 10.240219 # type: ignore
assert "content" in first_result
assert "attributes" in first_result
# Validate content structure
assert len(first_result["content"]) > 0 # type: ignore
assert first_result["content"][0]["type"] == "text" # type: ignore
assert "text" in first_result["content"][0] # type: ignore
assert (
first_result["content"][0]["text"] # type: ignore
== "abababababa_0562efee-0f1f-4b6b-9ca3-1a160f124ad8"
)
# Validate attributes contain book_id but NOT book_intro_text (it's in content)
assert "book_id" in first_result["attributes"] # type: ignore
assert first_result["attributes"]["book_id"] == 0 # type: ignore
assert "book_intro_text" not in first_result["attributes"] # type: ignore # Should be in content, not attributes
def test_basic_search_with_mock_sync(self):
"""Test basic vector search with mocked backend response (sync)"""
# Mock the HTTP response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = MOCK_MILVUS_SEARCH_RESPONSE
mock_response.text = json.dumps(MOCK_MILVUS_SEARCH_RESPONSE)
with patch("litellm.embedding") as mock_embedding:
mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post"
) as mock_post:
mock_post.return_value = mock_response
# Make the search request
response = vector_store_search(
query="what is machine learning?",
vector_store_id="book_2",
custom_llm_provider="milvus",
api_base="https://in03-test.serverless.aws-eu-central-1.cloud.zilliz.com",
api_key="mock_milvus_api_key",
litellm_embedding_model="text-embedding-3-large",
litellm_embedding_config={
"api_key": "mock_openai_api_key",
},
outputFields=["book_intro_text"],
annsField="book_intro_vector",
milvus_text_field="book_intro_text",
)
print("Response:", json.dumps(response, indent=2, default=str))
# Verify embedding was called
mock_embedding.assert_called_once()
# Verify the API was called
mock_post.assert_called_once()
# Verify the request payload
call_args = mock_post.call_args
# The post method is called with 'data' parameter (JSON string) not 'json' parameter
request_data_str = call_args.kwargs.get("data")
if request_data_str:
request_data = json.loads(request_data_str)
else:
# Fallback: check for json kwarg or in args
request_data = call_args.kwargs.get("json")
if (
request_data is None
and len(call_args.args) > 0
and isinstance(call_args.args[0], dict)
):
request_data = call_args.args[0]
assert (
request_data is not None
), f"Could not extract request data. Call args: {call_args}"
# Validate request structure
assert "collectionName" in request_data
assert request_data["collectionName"] == "book_2"
assert "data" in request_data
assert isinstance(request_data["data"], list)
assert "annsField" in request_data
assert "outputFields" in request_data
# Validate the response structure
assert response is not None
assert response["object"] == "vector_store.search_results.page" # type: ignore
assert "data" in response # type: ignore
assert len(response["data"]) == 9 # type: ignore # 9 results in mock response
assert "search_query" in response # type: ignore
# Validate first few results
expected_results = [
{
"book_id": 0,
"text": "abababababa_0562efee-0f1f-4b6b-9ca3-1a160f124ad8",
"distance": 10.240219,
},
{
"book_id": 1,
"text": "abababababa_9a13e8f3-bb1e-487f-b555-b8ae4b127243",
"distance": 10.240219,
},
{
"book_id": 2,
"text": "abababababa_870f47f1-23ec-4364-ad30-6d364ba8ddb5",
"distance": 10.240219,
},
]
for idx, expected in enumerate(expected_results):
result = response["data"][idx] # type: ignore
assert "score" in result
assert result["score"] == expected["distance"] # type: ignore
assert "content" in result
assert len(result["content"]) > 0 # type: ignore
assert result["content"][0]["type"] == "text" # type: ignore
assert "text" in result["content"][0] # type: ignore
assert result["content"][0]["text"] == expected["text"] # type: ignore
assert "attributes" in result
assert result["attributes"]["book_id"] == expected["book_id"] # type: ignore
assert "book_intro_text" not in result["attributes"] # type: ignore # Should be in content, not attributes
# @pytest.mark.parametrize("sync_mode", [True, False])
# @pytest.mark.asyncio
# async def test_basic_search_vector_store(sync_mode):
# """Integration test with real Milvus API (requires credentials)"""
# litellm._turn_on_debug()
# litellm.set_verbose = True
# base_request_args = {
# "vector_store_id": "book_2",
# "custom_llm_provider": "milvus",
# "api_base": "https://in03-18505f064ffbc6f.serverless.aws-eu-central-1.cloud.zilliz.com",
# "litellm_embedding_model": "text-embedding-3-large",
# "litellm_embedding_config": {
# "api_key": os.getenv("OPENAI_API_KEY"),
# },
# "default_output_fields": [
# "book_intro_text"
# ], # field containing the text to return in the response
# "default_anns_field": "book_intro_vector",
# }
# default_query = base_request_args.pop("query", "Basic ping")
# print(f"base_request_args: {base_request_args}")
# try:
# if sync_mode:
# response = vector_store_search(query=default_query, **base_request_args)
# else:
# response = await vector_store_asearch(
# query=default_query, **base_request_args
# )
# except litellm.InternalServerError:
# pytest.skip("Skipping test due to litellm.InternalServerError")
# print("litellm response=", json.dumps(response, indent=4, default=str))
# assert len(response["data"]) > 0 # type: ignore
if __name__ == "__main__":
# Run tests
import asyncio
test = TestMilvusVectorStore()
print("Running async mock test...")
asyncio.run(test.test_basic_search_with_mock_async())
print("\nRunning sync mock test...")
test.test_basic_search_with_mock_sync()
print("\n✅ All mock tests passed!")