mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
(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:
parent
22d35e2552
commit
b02be1ba70
20 changed files with 1487 additions and 124 deletions
|
|
@ -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:
|
||||
|
|
|
|||
221
docs/my-website/docs/providers/milvus_vector_stores.md
Normal file
221
docs/my-website/docs/providers/milvus_vector_stores.md
Normal 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.
|
||||
|
||||
|
|
@ -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>
|
||||
|
||||
|
|
|
|||
|
|
@ -495,6 +495,7 @@ const sidebars = {
|
|||
"providers/bedrock_vector_store",
|
||||
]
|
||||
},
|
||||
"providers/milvus_vector_stores",
|
||||
"providers/litellm_proxy",
|
||||
"providers/meta_llama",
|
||||
"providers/mistral",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
3
litellm/llms/milvus/vector_stores/__init__.py
Normal file
3
litellm/llms/milvus/vector_stores/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from litellm.llms.milvus.vector_stores.transformation import MilvusVectorStoreConfig
|
||||
|
||||
__all__ = ["MilvusVectorStoreConfig"]
|
||||
252
litellm/llms/milvus/vector_stores/transformation.py
Normal file
252
litellm/llms/milvus/vector_stores/transformation.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
361
tests/vector_store_tests/test_milvus_vector_store.py
Normal file
361
tests/vector_store_tests/test_milvus_vector_store.py
Normal 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!")
|
||||
Loading…
Add table
Reference in a new issue