diff --git a/docs/my-website/docs/providers/bedrock_vector_store.md b/docs/my-website/docs/providers/bedrock_vector_store.md index 39e1aec5ab8..5fae0c76c11 100644 --- a/docs/my-website/docs/providers/bedrock_vector_store.md +++ b/docs/my-website/docs/providers/bedrock_vector_store.md @@ -138,6 +138,125 @@ print(response.choices[0].message.content) +## 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` + + + + +```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" + } + }] +) +``` + + + + + +```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"} + ] + } + }] +) +``` + + + + + +```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"} + ] + } + }] +) +``` + + + + + +```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"} + ] + } + }] +) +``` + + + + + +```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"} + ] + } + }] + }' +``` + + + + ## Accessing Search Results See how to access vector store search results in your response: diff --git a/docs/my-website/docs/providers/milvus_vector_stores.md b/docs/my-website/docs/providers/milvus_vector_stores.md new file mode 100644 index 00000000000..c1a3042051b --- /dev/null +++ b/docs/my-website/docs/providers/milvus_vector_stores.md @@ -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 + + + + +### 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) +``` + + + + + +### 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?" +}' +``` + + + + +## 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. + diff --git a/docs/my-website/docs/vector_stores/search.md b/docs/my-website/docs/vector_stores/search.md index 5d0a2b737b9..2ffc8ef12e5 100644 --- a/docs/my-website/docs/vector_stores/search.md +++ b/docs/my-website/docs/vector_stores/search.md @@ -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) + + + + +#### 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) + diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 3b993e4620c..1ed2840d29e 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -495,6 +495,7 @@ const sidebars = { "providers/bedrock_vector_store", ] }, + "providers/milvus_vector_stores", "providers/litellm_proxy", "providers/meta_llama", "providers/mistral", diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 9d7ba7d61a8..daca48fc323 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -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 diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index c05b6ba3fb1..c485defe8a8 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -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": , "value": , "operator": } + AWS KB format: {"operator": {"key": , "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": , "value": , "operator": }]} + AWS KB format: {"andAll" | "orAll": [{"operator": {"key": , "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": , "value": , "operator": } OR {"and" | "or": [{"key": , "value": , "operator": }]} + # aws kb filters = {"operator": {"": }} OR {"andAll | orAll": [{"operator": {"": }}]} + # 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"), diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index cc451e5fa95..37b4af306a1 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -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( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 0a01e9cd68a..d3039dbefdb 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/litellm/llms/milvus/vector_stores/__init__.py b/litellm/llms/milvus/vector_stores/__init__.py new file mode 100644 index 00000000000..c20f5fa94b7 --- /dev/null +++ b/litellm/llms/milvus/vector_stores/__init__.py @@ -0,0 +1,3 @@ +from litellm.llms.milvus.vector_stores.transformation import MilvusVectorStoreConfig + +__all__ = ["MilvusVectorStoreConfig"] diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py new file mode 100644 index 00000000000..895d7bf37a0 --- /dev/null +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -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 diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 318cb56619b..8544956db92 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 64dea058a27..72e139dea52 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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, ) diff --git a/litellm/types/router.py b/litellm/types/router.py index 354e97d14d7..3801d5bb785 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b6ed7c38b9b..eceea4c652e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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). diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 58a2e7f9006..144b1853382 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -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 \ No newline at end of file + 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", +] diff --git a/litellm/utils.py b/litellm/utils.py index 0d4803ac881..8c0b2b4bb3b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index a0d0e338edc..ba5c1ffd3ac 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -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, ) diff --git a/litellm/vector_stores/utils.py b/litellm/vector_stores/utils.py index b7eb7790add..9430c68ccc9 100644 --- a/litellm/vector_stores/utils.py +++ b/litellm/vector_stores/utils.py @@ -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) - diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 3431d168529..ab0faa615b9 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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 diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py new file mode 100644 index 00000000000..ece07377786 --- /dev/null +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -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!")