diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index c3e0055e38..4ac5582d06 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -38,6 +38,10 @@ spec: serviceAccountName: {{ include "litellm.serviceAccountName" . }} securityContext: {{- toYaml .Values.podSecurityContext | nindent 8 }} + {{- with .Values.extraInitContainers }} + initContainers: + {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: {{ include "litellm.name" . }} securityContext: diff --git a/deploy/charts/litellm-helm/templates/migrations-job.yaml b/deploy/charts/litellm-helm/templates/migrations-job.yaml index f8893a47af..3459fa12d1 100644 --- a/deploy/charts/litellm-helm/templates/migrations-job.yaml +++ b/deploy/charts/litellm-helm/templates/migrations-job.yaml @@ -35,6 +35,10 @@ spec: {{- toYaml . | nindent 8 }} {{- end }} serviceAccountName: {{ include "litellm.serviceAccountName" . }} + {{- with .Values.migrationJob.extraInitContainers }} + initContainers: + {{- toYaml . | nindent 8 }} + {{- end }} containers: - name: prisma-migrations image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default (printf "main-%s" .Chart.AppVersion) }}" diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index b75ce64037..cea25974bb 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -281,6 +281,7 @@ migrationJob: # cpu: 100m # memory: 100Mi extraContainers: [] + extraInitContainers: [] # Hook configuration hooks: diff --git a/docs/my-website/docs/providers/gemini.md b/docs/my-website/docs/providers/gemini.md index 23a02f7365..b9ad7820dd 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/docs/my-website/docs/providers/gemini.md @@ -1840,6 +1840,57 @@ content = response.get('choices', [{}])[0].get('message', {}).get('content') print(content) ``` +## gemini-robotics-er-1.5-preview Usage + +```python +from litellm import api_base +from openai import OpenAI +import os +import base64 + +client = OpenAI(base_url="http://0.0.0.0:4000", api_key="sk-12345") +base64_image = base64.b64encode(open("closeup-object-on-table-many-260nw-1216144471.webp", "rb").read()).decode() + +import json +import re +tools = [{"codeExecution": {}}] +response = client.chat.completions.create( + model="gemini/gemini-robotics-er-1.5-preview", + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Point to no more than 10 items in the image. The label returned should be an identifying name for the object detected. The answer should follow the json format: [{\"point\": [y, x], \"label\": }, ...]. The points are in [y, x] format normalized to 0-1000." + }, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"} + } + ] + } + ], + tools=tools +) + +# Extract JSON from markdown code block if present +content = response.choices[0].message.content +# Look for triple-backtick JSON block +match = re.search(r'```json\s*(.*?)\s*```', content, re.DOTALL) +if match: + json_str = match.group(1) +else: + json_str = content + +try: + data = json.loads(json_str) + print(json.dumps(data, indent=2)) +except Exception as e: + print("Error parsing response as JSON:", e) + print("Response content:", content) +``` + ## Usage - PDF / Videos / etc. Files ### Inline Data (e.g. audio stream) diff --git a/docs/my-website/docs/providers/sarvam.md b/docs/my-website/docs/providers/sarvam.md new file mode 100644 index 0000000000..d77e9c0c75 --- /dev/null +++ b/docs/my-website/docs/providers/sarvam.md @@ -0,0 +1,89 @@ +# Sarvam.ai + +LiteLLM supports all the text models from [Sarvam ai](https://docs.sarvam.ai/api-reference-docs/chat/chat-completions) + +## Usage + +```python +import os +from litellm import completion + +# Set your Sarvam API key +os.environ["SARVAM_API_KEY"] = "" + +messages = [{"role": "user", "content": "Hello"}] + +response = completion( + model="sarvam/sarvam-m", + messages=messages, +) +print(response) +``` + +## Usage with LiteLLM Proxy Server + +Here's how to call a Sarvam.ai model with the LiteLLM Proxy Server + +1. **Modify the `config.yaml`:** + + ```yaml + model_list: + - model_name: my-model + litellm_params: + model: sarvam/ # add sarvam/ prefix to route as Sarvam provider + api_key: api-key # api key to send your model + ``` + +2. **Start the proxy:** + + ```bash + $ litellm --config /path/to/config.yaml + ``` + +3. **Send a request to LiteLLM Proxy Server:** + + + + + + ```python + import openai + + client = openai.OpenAI( + api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys + base_url="http://0.0.0.0:4000" # litellm-proxy-base url + ) + + response = client.chat.completions.create( + model="my-model", + messages=[ + { + "role": "user", + "content": "what llm are you" + } + ], + ) + + print(response) + ``` + + + + + ```shell + curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "model": "my-model", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] + }' + ``` + + + diff --git a/docs/my-website/docs/proxy/guardrails/quick_start.md b/docs/my-website/docs/proxy/guardrails/quick_start.md index cb6379d49f..ddb215fcb6 100644 --- a/docs/my-website/docs/proxy/guardrails/quick_start.md +++ b/docs/my-website/docs/proxy/guardrails/quick_start.md @@ -405,14 +405,10 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ ## **Proxy Admin Controls** -### ✨ Monitoring Guardrails +### Monitoring Guardrails Monitor which guardrails were executed and whether they passed or failed. e.g. guardrail going rogue and failing requests we don't intend to fail -:::info - -✨ This is an Enterprise only feature [Get a free trial](https://www.litellm.ai/enterprise#trial) - ::: #### Setup diff --git a/docs/my-website/docs/rag_ingest.md b/docs/my-website/docs/rag_ingest.md index 1133b85f20..7adc2d70b5 100644 --- a/docs/my-website/docs/rag_ingest.md +++ b/docs/my-website/docs/rag_ingest.md @@ -5,7 +5,7 @@ All-in-one document ingestion pipeline: **Upload → Chunk → Embed → Vector | Feature | Supported | |---------|-----------| | Logging | Yes | -| Supported Providers | `openai`, `bedrock`, `vertex_ai`, `gemini` | +| Supported Providers | `openai`, `bedrock`, `vertex_ai`, `gemini`, `s3_vectors` | :::tip After ingesting documents, use [/rag/query](./rag_query.md) to search and generate responses with your ingested content. @@ -75,6 +75,31 @@ curl -X POST "http://localhost:4000/v1/rag/ingest" \ }" ``` +### AWS S3 Vectors + +```bash showLineNumbers title="Ingest to S3 Vectors" +curl -X POST "http://localhost:4000/v1/rag/ingest" \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d "{ + \"file\": { + \"filename\": \"document.txt\", + \"content\": \"$(base64 -i document.txt)\", + \"content_type\": \"text/plain\" + }, + \"ingest_options\": { + \"embedding\": { + \"model\": \"text-embedding-3-small\" + }, + \"vector_store\": { + \"custom_llm_provider\": \"s3_vectors\", + \"vector_bucket_name\": \"my-embeddings\", + \"aws_region_name\": \"us-west-2\" + } + } + }" +``` + ## Response ```json @@ -265,6 +290,57 @@ When `vector_store_id` is omitted, LiteLLM automatically creates: 4. Install: `pip install 'google-cloud-aiplatform>=1.60.0'` ::: +### vector_store (AWS S3 Vectors) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `custom_llm_provider` | string | - | `"s3_vectors"` | +| `vector_bucket_name` | string | **required** | S3 vector bucket name | +| `index_name` | string | auto-create | Vector index name | +| `dimension` | integer | auto-detect | Vector dimension (auto-detected from embedding model) | +| `distance_metric` | string | `cosine` | Distance metric: `cosine` or `euclidean` | +| `non_filterable_metadata_keys` | array | `["source_text"]` | Metadata keys excluded from filtering | +| `aws_region_name` | string | `us-west-2` | AWS region | +| `aws_access_key_id` | string | env | AWS access key | +| `aws_secret_access_key` | string | env | AWS secret key | + +:::info S3 Vectors Auto-Creation +When `index_name` is omitted, LiteLLM automatically creates: +- S3 vector bucket (if it doesn't exist) +- Vector index with auto-detected dimensions from your embedding model + +**Dimension Auto-Detection**: The vector dimension is automatically detected by making a test embedding request to your specified model. No need to manually specify dimensions! + +**Supported Embedding Models**: Works with any LiteLLM-supported embedding model (OpenAI, Cohere, Bedrock, Azure, etc.) +::: + +**Example with auto-detection:** +```json +{ + "embedding": { + "model": "text-embedding-3-small" // Dimension auto-detected as 1536 + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "my-embeddings" + } +} +``` + +**Example with custom embedding provider:** +```json +{ + "embedding": { + "model": "cohere/embed-english-v3.0" // Dimension auto-detected as 1024 + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "my-embeddings", + "distance_metric": "cosine" + } +} +``` + ## Input Examples ### File (Base64) diff --git a/docs/my-website/docs/routing.md b/docs/my-website/docs/routing.md index 47967775e1..b5ece7237a 100644 --- a/docs/my-website/docs/routing.md +++ b/docs/my-website/docs/routing.md @@ -828,7 +828,12 @@ asyncio.run(router_acompletion()) ``` - + +## Traffic Mirroring / Silent Experiments + +Traffic mirroring allows you to "mimic" production traffic to a secondary (silent) model for evaluation purposes. The silent model's response is gathered in the background and does not affect the latency or result of the primary request. + +[**See detailed guide on A/B Testing - Traffic Mirroring here**](./traffic_mirroring.md) ## Basic Reliability diff --git a/docs/my-website/docs/traffic_mirroring.md b/docs/my-website/docs/traffic_mirroring.md new file mode 100644 index 0000000000..3bdcb0f161 --- /dev/null +++ b/docs/my-website/docs/traffic_mirroring.md @@ -0,0 +1,83 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# A/B Testing - Traffic Mirroring + +Traffic mirroring allows you to "mimic" production traffic to a secondary (silent) model for evaluation purposes. The silent model's response is gathered in the background and does not affect the latency or result of the primary request. + +This is useful for: +- Testing a new model's performance on production prompts before switching. +- Comparing costs and latency between different providers. +- Debugging issues by mirroring traffic to a more verbose model. + +## Quick Start + +To enable traffic mirroring, add `silent_model` to the `litellm_params` of a deployment. + + + + +```python +from litellm import Router + +model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "azure/chatgpt-v-2", + "api_key": "...", + "silent_model": "gpt-4" # 👈 Mirror traffic to gpt-4 + }, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "..." + }, + } +] + +router = Router(model_list=model_list) + +# The request to "gpt-3.5-turbo" will trigger a background call to "gpt-4" +response = await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "How does traffic mirroring work?"}] +) +``` + + + + +Add `silent_model` to your `config.yaml`: + +```yaml +model_list: + - model_name: primary-model + litellm_params: + model: azure/gpt-35-turbo + api_key: os.environ/AZURE_API_KEY + silent_model: evaluation-model # 👈 Mirror traffic here + - model_name: evaluation-model + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY +``` + + + + +## How it works +1. **Request Received**: A request is made to a model group (e.g. `primary-model`). +2. **Deployment Picked**: LiteLLM picks a deployment from the group. +3. **Primary Call**: LiteLLM makes the call to the primary deployment. +4. **Mirroring**: If `silent_model` is present, LiteLLM triggers a background call to that model. + - For **Sync** calls: Uses a shared thread pool. + - For **Async** calls: Uses `asyncio.create_task`. +5. **Isolation**: The background call uses a `deepcopy` of the original request parameters and sets `metadata["is_silent_experiment"] = True`. It also strips out logging IDs to prevent collisions in usage tracking. + +## Key Features +- **Latency Isolation**: The primary request returns as soon as it's ready. The background (silent) call does not block. +- **Unified Logging**: Background calls are processed via the Router, meaning they are automatically logged to your configured observability tools (Langfuse, S3, etc.). +- **Evaluation**: Use the `is_silent_experiment: True` flag in your logs to filter and compare results between the primary and mirrored calls. diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 85573599e4..ef7573e9ce 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -364,6 +364,7 @@ const sidebars = { label: "Load Balancing, Routing, Fallbacks", href: "https://docs.litellm.ai/docs/routing-load-balancing", }, + "traffic_mirroring", { type: "category", label: "Logging, Alerting, Metrics", @@ -775,6 +776,7 @@ const sidebars = { "providers/oci", "providers/ollama", "providers/openrouter", + "providers/sarvam", "providers/ovhcloud", "providers/perplexity", "providers/petals", diff --git a/enterprise/litellm_enterprise/proxy/auth/route_checks.py b/enterprise/litellm_enterprise/proxy/auth/route_checks.py index 6cce781faf..6f7cf9143f 100644 --- a/enterprise/litellm_enterprise/proxy/auth/route_checks.py +++ b/enterprise/litellm_enterprise/proxy/auth/route_checks.py @@ -36,7 +36,7 @@ class EnterpriseRouteChecks: if not premium_user: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"🚨🚨🚨 DISABLING LLM API ENDPOINTS is an Enterprise feature\n🚨 {CommonProxyErrors.not_premium_user.value}", + detail=f"🚨🚨🚨 DISABLING ADMIN ENDPOINTS is an Enterprise feature\n🚨 {CommonProxyErrors.not_premium_user.value}", ) return get_secret_bool("DISABLE_ADMIN_ENDPOINTS") is True diff --git a/litellm/constants.py b/litellm/constants.py index 0525bf3843..0c9d1d941b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1327,6 +1327,13 @@ COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int( DEFAULT_CHUNK_SIZE = int(os.getenv("DEFAULT_CHUNK_SIZE", 1000)) DEFAULT_CHUNK_OVERLAP = int(os.getenv("DEFAULT_CHUNK_OVERLAP", 200)) +########################### S3 Vectors RAG Constants ########################### +S3_VECTORS_DEFAULT_DIMENSION = int(os.getenv("S3_VECTORS_DEFAULT_DIMENSION", 1024)) +S3_VECTORS_DEFAULT_DISTANCE_METRIC = str( + os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine") +) +S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS = ["source_text"] + ########################### Microsoft SSO Constants ########################### MICROSOFT_USER_EMAIL_ATTRIBUTE = str( os.getenv("MICROSOFT_USER_EMAIL_ATTRIBUTE", "userPrincipalName") diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 5d0bedb776..7137a4e422 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -31,15 +31,19 @@ def _process_image_response(response: Response, url: str) -> str: f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) - image_bytes = response.content + # Stream download with size checking to prevent downloading huge files + max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) + image_bytes = bytearray() + bytes_downloaded = 0 - # Check actual size after download if Content-Length was not available - if content_length is None: - size_mb = len(image_bytes) / (1024 * 1024) - if size_mb > MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: + for chunk in response.iter_bytes(chunk_size=8192): + bytes_downloaded += len(chunk) + if bytes_downloaded > max_bytes: + size_mb = bytes_downloaded / (1024 * 1024) raise litellm.ImageFetchError( f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) + image_bytes.extend(chunk) base64_image = base64.b64encode(image_bytes).decode("utf-8") diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 89f2094d5d..935fd53c19 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -5,8 +5,8 @@ import httpx from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( - BaseVectorStoreAuthCredentials, VECTOR_STORE_OPENAI_PARAMS, + BaseVectorStoreAuthCredentials, VectorStoreCreateOptionalRequestParams, VectorStoreCreateResponse, VectorStoreIndexEndpoints, @@ -64,6 +64,30 @@ class BaseVectorStoreConfig: pass + async def atransform_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]: + """ + Optional async version of transform_search_vector_store_request. + If not implemented, the handler will fall back to the sync version. + Providers that need to make async calls (e.g., generating embeddings) should override this. + """ + # Default implementation: call the sync version + return self.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=litellm_logging_obj, + litellm_params=litellm_params, + ) + @abstractmethod def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 642d15fe3e..1de1c40c43 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1163,7 +1163,7 @@ class BaseAWSLLM: def _sign_request( self, - service_name: Literal["bedrock", "sagemaker", "bedrock-agentcore"], + service_name: Literal["bedrock", "sagemaker", "bedrock-agentcore", "s3vectors"], headers: dict, optional_params: dict, request_data: dict, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 89b42f5e94..65d237bdbd 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -797,7 +797,7 @@ class BedrockEventStreamDecoderBase: def get_anthropic_beta_from_headers(headers: dict) -> List[str]: """ Extract anthropic-beta header values and convert them to a list. - Supports comma-separated values from user headers. + Supports both JSON array format and comma-separated values from user headers. Used by both converse and invoke transformations for consistent handling of anthropic-beta headers that should be passed to AWS Bedrock. @@ -812,8 +812,25 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: if not anthropic_beta_header: return [] - # Split comma-separated values and strip whitespace - return [beta.strip() for beta in anthropic_beta_header.split(",")] + # If it's already a list, return it + if isinstance(anthropic_beta_header, list): + return anthropic_beta_header + + # Try to parse as JSON array first (e.g., '["interleaved-thinking-2025-05-14", "claude-code-20250219"]') + if isinstance(anthropic_beta_header, str): + anthropic_beta_header = anthropic_beta_header.strip() + if anthropic_beta_header.startswith("[") and anthropic_beta_header.endswith("]"): + try: + parsed = json.loads(anthropic_beta_header) + if isinstance(parsed, list): + return [str(beta).strip() for beta in parsed] + except json.JSONDecodeError: + pass # Fall through to comma-separated parsing + + # Fall back to comma-separated values + return [beta.strip() for beta in anthropic_beta_header.split(",")] + + return [] class CommonBatchFilesUtils: diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index a7065caece..fc4220c054 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -162,6 +162,49 @@ class AmazonAnthropicClaudeMessagesConfig( return any(pattern in model_lower for pattern in supported_patterns) + def _is_claude_opus_4_5(self, model: str) -> bool: + """ + Check if the model is Claude Opus 4.5. + + Args: + model: The model name + + Returns: + True if the model is Claude Opus 4.5 + """ + model_lower = model.lower() + opus_4_5_patterns = [ + "opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5", + ] + return any(pattern in model_lower for pattern in opus_4_5_patterns) + + def _supports_tool_search_on_bedrock(self, model: str) -> bool: + """ + Check if the model supports tool search on Bedrock. + + On Amazon Bedrock, server-side tool search is supported on Claude Opus 4.5 + and Claude Sonnet 4.5 with the tool-search-tool-2025-10-19 beta header. + + Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool + + Args: + model: The model name + + Returns: + True if the model supports tool search on Bedrock + """ + model_lower = model.lower() + + # Supported models for tool search on Bedrock + supported_patterns = [ + # Opus 4.5 + "opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5", + # Sonnet 4.5 + "sonnet-4.5", "sonnet_4.5", "sonnet-4-5", "sonnet_4_5", + ] + + return any(pattern in model_lower for pattern in supported_patterns) + def _filter_unsupported_beta_headers_for_bedrock( self, model: str, beta_set: set ) -> None: @@ -169,25 +212,33 @@ class AmazonAnthropicClaudeMessagesConfig( Remove beta headers that are not supported on Bedrock for the given model. Extended thinking beta headers are only supported on specific Claude 4+ models. - Advanced tool use headers are not supported on Bedrock Invoke API. + Advanced tool use headers are not supported on Bedrock Invoke API, but need to be + translated to Bedrock-specific headers for models that support tool search + (Claude Opus 4.5, Sonnet 4.5). This prevents 400 "invalid beta flag" errors on Bedrock. Note: Bedrock Invoke API fails with a 400 error when unsupported beta headers are sent, returning: {"message":"invalid beta flag"} + Translation for models supporting tool search (Opus 4.5, Sonnet 4.5): + - advanced-tool-use-2025-11-20 -> tool-search-tool-2025-10-19 + tool-examples-2025-10-29 + Args: model: The model name beta_set: The set of beta headers to filter in-place """ beta_headers_to_remove = set() + has_advanced_tool_use = False - # 1. Filter out beta headers that are universally unsupported on Bedrock Invoke + # 1. Filter out beta headers that are universally unsupported on Bedrock Invoke and track if advanced-tool-use header is present for beta in beta_set: for unsupported_pattern in self.UNSUPPORTED_BEDROCK_INVOKE_BETA_PATTERNS: if unsupported_pattern in beta.lower(): beta_headers_to_remove.add(beta) + has_advanced_tool_use = True break - + + # 2. Filter out extended thinking headers for models that don't support them extended_thinking_patterns = [ "extended-thinking", @@ -204,6 +255,14 @@ class AmazonAnthropicClaudeMessagesConfig( for beta in beta_headers_to_remove: beta_set.discard(beta) + # 3. Translate advanced-tool-use to Bedrock-specific headers for models that support tool search + # Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + # Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool + if has_advanced_tool_use and self._supports_tool_search_on_bedrock(model): + beta_set.add("tool-search-tool-2025-10-19") + beta_set.add("tool-examples-2025-10-29") + + def _get_tool_search_beta_header_for_bedrock( self, model: str, @@ -256,7 +315,7 @@ class AmazonAnthropicClaudeMessagesConfig( Ref: https://aws.amazon.com/blogs/machine-learning/structured-data-response-with-amazon-bedrock-prompt-engineering-and-tool-use/ """ import json - + # Extract schema from output_format schema = output_format.get("schema") if not schema: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6a87967c3a..d2ea7e872a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7033,17 +7033,31 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - ( - url, - request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), - ) + # Check if provider has async transform method + if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"): + ( + url, + request_body, + ) = await vector_store_provider_config.atransform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + ) + else: + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + ) all_optional_params: Dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) headers, signed_json_body = vector_store_provider_config.sign_request( diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index a75ecd8cc7..34ea7b03dd 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -323,4 +323,12 @@ class GroqChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): status_code=error.get("code"), message=error.get("message"), body=error ) + # Map Groq's 'reasoning' field to LiteLLM's 'reasoning_content' field + # Groq returns delta.reasoning, but LiteLLM expects delta.reasoning_content + choices = chunk.get("choices", []) + for choice in choices: + delta = choice.get("delta", {}) + if "reasoning" in delta: + delta["reasoning_content"] = delta.pop("reasoning") + return super().chunk_parser(chunk) diff --git a/litellm/llms/s3_vectors/__init__.py b/litellm/llms/s3_vectors/__init__.py new file mode 100644 index 0000000000..e8367949c3 --- /dev/null +++ b/litellm/llms/s3_vectors/__init__.py @@ -0,0 +1 @@ +# S3 Vectors LLM integration diff --git a/litellm/llms/s3_vectors/vector_stores/__init__.py b/litellm/llms/s3_vectors/vector_stores/__init__.py new file mode 100644 index 0000000000..ac24b4a38d --- /dev/null +++ b/litellm/llms/s3_vectors/vector_stores/__init__.py @@ -0,0 +1 @@ +# S3 Vectors vector store integration diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py new file mode 100644 index 0000000000..df81a78289 --- /dev/null +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -0,0 +1,254 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +import httpx + +from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.types.router import GenericLiteLLMParams +from litellm.types.vector_stores import ( + VECTOR_STORE_OPENAI_PARAMS, + BaseVectorStoreAuthCredentials, + VectorStoreIndexEndpoints, + VectorStoreResultContent, + VectorStoreSearchOptionalRequestParams, + VectorStoreSearchResponse, + VectorStoreSearchResult, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): + """Vector store configuration for AWS S3 Vectors.""" + + def __init__(self) -> None: + BaseVectorStoreConfig.__init__(self) + BaseAWSLLM.__init__(self) + + def get_auth_credentials( + self, litellm_params: dict + ) -> BaseVectorStoreAuthCredentials: + return {} + + def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: + return { + "read": [("POST", "/QueryVectors")], + "write": [], + } + + def get_supported_openai_params( + self, model: str + ) -> List[VECTOR_STORE_OPENAI_PARAMS]: + return ["max_num_results"] + + 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["maxResults"] = value + return optional_params + + def validate_environment( + self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] + ) -> dict: + headers = headers or {} + headers.setdefault("Content-Type", "application/json") + return headers + + def get_complete_url(self, api_base: Optional[str], litellm_params: dict) -> str: + aws_region_name = litellm_params.get("aws_region_name") + if not aws_region_name: + raise ValueError("aws_region_name is required for S3 Vectors") + return f"https://s3vectors.{aws_region_name}.api.aws" + + 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]: + """Sync version - generates embedding synchronously.""" + # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name + # If not in that format, try to construct it from litellm_params + bucket_name: str + index_name: str + + if ":" in vector_store_id: + bucket_name, index_name = vector_store_id.split(":", 1) + else: + # Try to get bucket_name from litellm_params + bucket_name_from_params = litellm_params.get("vector_bucket_name") + if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): + raise ValueError( + "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " + "or vector_bucket_name must be provided in litellm_params" + ) + bucket_name = bucket_name_from_params + index_name = vector_store_id + + if isinstance(query, list): + query = " ".join(query) + + # Generate embedding for the query + embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") + + import litellm as litellm_module + embedding_response = litellm_module.embedding(model=embedding_model, input=[query]) + query_embedding = embedding_response.data[0]["embedding"] + + url = f"{api_base}/QueryVectors" + + request_body: Dict[str, Any] = { + "vectorBucketName": bucket_name, + "indexName": index_name, + "queryVector": {"float32": query_embedding}, + "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 + "returnDistance": True, + "returnMetadata": True, + } + + litellm_logging_obj.model_call_details["query"] = query + return url, request_body + + async def atransform_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]: + """Async version - generates embedding asynchronously.""" + # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name + # If not in that format, try to construct it from litellm_params + bucket_name: str + index_name: str + + if ":" in vector_store_id: + bucket_name, index_name = vector_store_id.split(":", 1) + else: + # Try to get bucket_name from litellm_params + bucket_name_from_params = litellm_params.get("vector_bucket_name") + if not bucket_name_from_params or not isinstance(bucket_name_from_params, str): + raise ValueError( + "vector_store_id must be in format 'bucket_name:index_name' for S3 Vectors, " + "or vector_bucket_name must be provided in litellm_params" + ) + bucket_name = bucket_name_from_params + index_name = vector_store_id + + if isinstance(query, list): + query = " ".join(query) + + # Generate embedding for the query asynchronously + embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") + + import litellm as litellm_module + embedding_response = await litellm_module.aembedding(model=embedding_model, input=[query]) + query_embedding = embedding_response.data[0]["embedding"] + + url = f"{api_base}/QueryVectors" + + request_body: Dict[str, Any] = { + "vectorBucketName": bucket_name, + "indexName": index_name, + "queryVector": {"float32": query_embedding}, + "topK": vector_store_search_optional_params.get("max_num_results", 5), # Default to 5 + "returnDistance": True, + "returnMetadata": True, + } + + litellm_logging_obj.model_call_details["query"] = query + return url, request_body + + def sign_request( + self, + headers: dict, + optional_params: Dict, + request_data: Dict, + api_base: str, + api_key: Optional[str] = None, + ) -> Tuple[dict, Optional[bytes]]: + return self._sign_request( + service_name="s3vectors", + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + api_key=api_key, + ) + + def transform_search_vector_store_response( + self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj + ) -> VectorStoreSearchResponse: + try: + response_data = response.json() + results: List[VectorStoreSearchResult] = [] + + for item in response_data.get("vectors", []) or []: + metadata = item.get("metadata", {}) or {} + source_text = metadata.get("source_text", "") + + if not source_text: + continue + + # Extract file information from metadata + chunk_index = metadata.get("chunk_index", "0") + file_id = f"s3-vectors-chunk-{chunk_index}" + filename = metadata.get("filename", f"document-{chunk_index}") + + # S3 Vectors returns distance, convert to similarity score (0-1) + # Lower distance = higher similarity + # We'll normalize using 1 / (1 + distance) to get a 0-1 score + distance = item.get("distance") + score = None + if distance is not None: + # Convert distance to similarity score between 0 and 1 + # For cosine distance: similarity = 1 - distance + # For euclidean: use 1 / (1 + distance) + # Assuming cosine distance here + score = max(0.0, min(1.0, 1.0 - float(distance))) + + results.append( + VectorStoreSearchResult( + score=score, + content=[VectorStoreResultContent(text=source_text, type="text")], + file_id=file_id, + filename=filename, + attributes=metadata, + ) + ) + + return VectorStoreSearchResponse( + object="vector_store.search_results.page", + search_query=litellm_logging_obj.model_call_details.get("query", ""), + data=results, + ) + except Exception as e: + raise self.get_error_class( + error_message=str(e), + status_code=response.status_code, + headers=response.headers, + ) + + # Vector store creation is not yet implemented + def transform_create_vector_store_request( + self, + vector_store_create_optional_params, + api_base: str, + ) -> Tuple[str, Dict]: + raise NotImplementedError + + def transform_create_vector_store_response(self, response: httpx.Response): + raise NotImplementedError diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cf6b0e7e3f..8cb488c170 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10232,6 +10232,48 @@ "mode": "completion", "output_cost_per_token": 5e-07 }, + "deepseek-v3-2-251201": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 98304, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "glm-4-7-251222": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "kimi-k2-thinking-251104": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 229376, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "doubao-embedding": { "input_cost_per_token": 0.0, "litellm_provider": "volcengine", @@ -13480,6 +13522,43 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-robotics-er-1.5-preview": { + "cache_read_input_token_cost": 0, + "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_reasoning_token": 2.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-robotics-er-1-5-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true + }, "gemini-2.5-computer-use-preview-10-2025": { "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, diff --git a/litellm/proxy/_experimental/out/assets/logos/s3_vector.png b/litellm/proxy/_experimental/out/assets/logos/s3_vector.png new file mode 100644 index 0000000000..15a1a456e1 Binary files /dev/null and b/litellm/proxy/_experimental/out/assets/logos/s3_vector.png differ diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 61d9044f92..e0b056d450 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1016,8 +1016,10 @@ async def _get_fuzzy_user_object( ) if response is None and user_email is not None: + # Use case-insensitive query to handle emails with different casing + # This matches the pattern used in _check_duplicate_user_email response = await prisma_client.db.litellm_usertable.find_first( - where={"user_email": user_email}, + where={"user_email": {"equals": user_email, "mode": "insensitive"}}, include={"organization_memberships": True}, ) diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index bba5b87024..5c41b371ca 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -264,7 +264,7 @@ async def create_interaction( general_settings=general_settings, proxy_config=proxy_config, select_data_generator=select_data_generator, - model=data.get("model"), + model=data.get("model") or data.get("agent"), user_model=user_model, user_temperature=user_temperature, user_request_timeout=user_request_timeout, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d0afb62e8b..63db2d72fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1771,6 +1771,113 @@ async def _add_team_members_to_team( return updated_team, updated_users, updated_team_memberships +async def _validate_and_populate_member_user_info( + member: Member, + prisma_client: PrismaClient, +) -> Member: + """ + Validate and populate user_email/user_id for a member. + + Logic: + 1. If both user_email and user_id are provided, verify they belong to the same user (use user_email as source of truth) + 2. If only user_email is provided, populate user_id from DB + 3. If only user_id is provided, populate user_email from DB (if user exists) + 4. If only user_id is provided and doesn't exist, allow it to pass with user_email as None (will be upserted later) + 5. If user_email and user_id mismatch, throw error + + Returns a Member with user_email and user_id populated (user_email may be None if only user_id provided and user doesn't exist). + """ + if member.user_email is None and member.user_id is None: + raise HTTPException( + status_code=400, + detail={"error": "Either user_id or user_email must be provided"}, + ) + + # Case 1: Both user_email and user_id provided - verify they match + if member.user_email is not None and member.user_id is not None: + # Use user_email as source of truth + # Check for multiple users with same email first + users_by_email = await prisma_client.get_data( + key_val={"user_email": member.user_email}, + table_name="user", + query_type="find_all", + ) + + if users_by_email is None or ( + isinstance(users_by_email, list) and len(users_by_email) == 0 + ): + # User doesn't exist yet - this is fine, will be created later + return member + + if isinstance(users_by_email, list) and len(users_by_email) > 1: + raise HTTPException( + status_code=400, + detail={ + "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." + }, + ) + + # Get the single user + user_by_email = users_by_email[0] + + # Verify the user_id matches + if user_by_email.user_id != member.user_id: + raise HTTPException( + status_code=400, + detail={ + "error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user." + }, + ) + + # Both match, return as is + return member + + # Case 2: Only user_email provided - populate user_id from DB + if member.user_email is not None and member.user_id is None: + user_by_email = await prisma_client.db.litellm_usertable.find_first( + where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} + ) + + if user_by_email is None: + # User doesn't exist yet - this is fine, will be created later + return member + + # Check for multiple users with same email + users_by_email = await prisma_client.get_data( + key_val={"user_email": member.user_email}, + table_name="user", + query_type="find_all", + ) + + if users_by_email and isinstance(users_by_email, list) and len(users_by_email) > 1: + raise HTTPException( + status_code=400, + detail={ + "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." + }, + ) + + # Populate user_id + member.user_id = user_by_email.user_id + return member + + # Case 3: Only user_id provided - populate user_email from DB if user exists + if member.user_id is not None and member.user_email is None: + user_by_id = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": member.user_id} + ) + + if user_by_id is None: + # User doesn't exist yet - allow it to pass with user_email as None + # Will be upserted later with just user_id and null email + return member + + # Populate user_email + member.user_email = user_by_id.user_email + return member + + return member + @router.post( "/team/member_add", tags=["team management"], @@ -1846,6 +1953,19 @@ async def team_member_add( complete_team_data=complete_team_data, ) + # Validate and populate user_email/user_id for members before processing + if isinstance(data.member, Member): + await _validate_and_populate_member_user_info( + member=data.member, + prisma_client=prisma_client, + ) + elif isinstance(data.member, List): + for m in data.member: + await _validate_and_populate_member_user_info( + member=m, + prisma_client=prisma_client, + ) + updated_team, updated_users, updated_team_memberships = ( await _add_team_members_to_team( data=data, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 1b827411e5..4048b3731c 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -93,6 +93,26 @@ else: router = APIRouter() +def normalize_email(email: Optional[str]) -> Optional[str]: + """ + Normalize email address to lowercase for consistent storage and comparison. + + Email addresses should be treated as case-insensitive for SSO purposes, + even though RFC 5321 technically allows case-sensitive local parts. + This prevents issues where SSO providers return emails with different casing + than what's stored in the database. + + Args: + email: Email address to normalize, can be None + + Returns: + Lowercased email address, or None if input is None + """ + if email is None: + return None + return email.lower() if isinstance(email, str) else email + + def determine_role_from_groups( user_groups: List[str], role_mappings: "RoleMappings", @@ -395,7 +415,7 @@ def generic_response_convertor( display_name=get_nested_value( response, generic_user_display_name_attribute_name ), - email=get_nested_value(response, generic_user_email_attribute_name), + email=normalize_email(get_nested_value(response, generic_user_email_attribute_name)), first_name=get_nested_value(response, generic_user_first_name_attribute_name), last_name=get_nested_value(response, generic_user_last_name_attribute_name), provider=get_nested_value(response, generic_provider_attribute_name), @@ -731,7 +751,7 @@ async def get_user_info_from_db( if _id is not None and isinstance(_id, str): potential_user_ids.append(_id) - user_email = ( + user_email = normalize_email( getattr(result, "email", None) if not isinstance(result, dict) else result.get("email", None) @@ -806,8 +826,8 @@ def _build_sso_user_update_data( Returns: dict: Update data containing user_email and optionally user_role if valid - """ - update_data: dict = {"user_email": user_email} + """ + update_data: dict = {"user_email": normalize_email(user_email)} # Get SSO role from result and include if valid sso_role = getattr(result, "user_role", None) @@ -1316,7 +1336,7 @@ async def insert_sso_user( new_user_request = NewUserRequest( user_id=user_defined_values["user_id"], - user_email=user_defined_values["user_email"], + user_email=normalize_email(user_defined_values["user_email"]), user_role=user_defined_values["user_role"], # type: ignore max_budget=user_defined_values["max_budget"], budget_duration=user_defined_values["budget_duration"], @@ -1981,7 +2001,7 @@ class SSOAuthenticationHandler: """ Gets the user email and id from the OpenID result after validating the email domain """ - user_email: Optional[str] = getattr(result, "email", None) + user_email: Optional[str] = normalize_email(getattr(result, "email", None)) user_id: Optional[str] = ( getattr(result, "id", None) if result is not None else None ) @@ -2020,7 +2040,7 @@ class SSOAuthenticationHandler: "GENERIC_USER_ROLE_ATTRIBUTE", "role" ) user_id = getattr(result, "id", None) - user_email = getattr(result, "email", None) + user_email = normalize_email(getattr(result, "email", None)) if user_role is None: _role_from_attr = getattr(result, generic_user_role_attribute_name, None) # type: ignore if _role_from_attr is not None: @@ -2413,7 +2433,7 @@ class MicrosoftSSOHandler: response = response or {} verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}") openid_response = CustomOpenID( - email=response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail"), + email=normalize_email(response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail")), display_name=response.get(MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE), provider="microsoft", id=response.get(MICROSOFT_USER_ID_ATTRIBUTE), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 183c25ed46..12e34e221d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7807,6 +7807,7 @@ async def _apply_search_filter_to_models( size: int, prisma_client: Optional[Any], proxy_config: Any, + sort_by: Optional[str] = None, ) -> Tuple[List[Dict[str, Any]], Optional[int]]: """ Apply search filter to models, querying database for additional matching models. @@ -7818,6 +7819,7 @@ async def _apply_search_filter_to_models( size: Page size prisma_client: Prisma client for database queries proxy_config: Proxy config for decrypting models + sort_by: Optional sort field - if provided, fetch all matching models instead of paginating at DB level Returns: Tuple of (filtered_models, total_count). total_count is None if not searching. @@ -7881,17 +7883,15 @@ async def _apply_search_filter_to_models( # Calculate total count for search results search_total_count = router_models_count + db_models_total_count - # Fetch database models if we need more for the current page - if router_models_count < models_needed_for_page: - models_to_fetch = min( - models_needed_for_page - router_models_count, db_models_total_count - ) - - if models_to_fetch > 0: + # If sorting is requested, we need to fetch ALL matching models to sort correctly + # Otherwise, we can optimize by only fetching what's needed for the current page + if sort_by: + # Fetch all matching database models for sorting + if db_models_total_count > 0: db_models_raw = ( await prisma_client.db.litellm_proxymodeltable.find_many( where=db_where_condition, - take=models_to_fetch, + take=db_models_total_count, # Fetch all matching models ) ) @@ -7902,6 +7902,28 @@ async def _apply_search_filter_to_models( ) if decrypted_models: db_models.extend(decrypted_models) + else: + # Fetch database models if we need more for the current page + if router_models_count < models_needed_for_page: + models_to_fetch = min( + models_needed_for_page - router_models_count, db_models_total_count + ) + + if models_to_fetch > 0: + db_models_raw = ( + await prisma_client.db.litellm_proxymodeltable.find_many( + where=db_where_condition, + take=models_to_fetch, + ) + ) + + # Convert database models to router format + for db_model in db_models_raw: + decrypted_models = proxy_config.decrypt_model_list_from_db( + [db_model] + ) + if decrypted_models: + db_models.extend(decrypted_models) except Exception as e: verbose_proxy_logger.exception( f"Error querying database models with search: {str(e)}" @@ -7917,6 +7939,80 @@ async def _apply_search_filter_to_models( return filtered_models, search_total_count +def _sort_models( + all_models: List[Dict[str, Any]], + sort_by: Optional[str], + sort_order: str = "asc", +) -> List[Dict[str, Any]]: + """ + Sort models by the specified field and order. + + Args: + all_models: List of models to sort + sort_by: Field to sort by (model_name, created_at, updated_at, costs, status) + sort_order: Sort order (asc or desc) + + Returns: + Sorted list of models + """ + if not sort_by or sort_by not in ["model_name", "created_at", "updated_at", "costs", "status"]: + return all_models + + reverse = sort_order.lower() == "desc" + + def get_sort_key(model: Dict[str, Any]) -> Any: + model_info = model.get("model_info", {}) + + if sort_by == "model_name": + return model.get("model_name", "").lower() + + elif sort_by == "created_at": + created_at = model_info.get("created_at") + if created_at is None: + # Put None values at the end for asc, at the start for desc + return (datetime.max if not reverse else datetime.min) + if isinstance(created_at, str): + try: + return datetime.fromisoformat(created_at.replace("Z", "+00:00")) + except (ValueError, AttributeError): + return datetime.min if not reverse else datetime.max + return created_at + + elif sort_by == "updated_at": + updated_at = model_info.get("updated_at") + if updated_at is None: + return (datetime.max if not reverse else datetime.min) + if isinstance(updated_at, str): + try: + return datetime.fromisoformat(updated_at.replace("Z", "+00:00")) + except (ValueError, AttributeError): + return datetime.min if not reverse else datetime.max + return updated_at + + elif sort_by == "costs": + input_cost = model_info.get("input_cost_per_token", 0) or 0 + output_cost = model_info.get("output_cost_per_token", 0) or 0 + total_cost = input_cost + output_cost + # Put 0 or None costs at the end for asc, at the start for desc + if total_cost == 0: + return (float("inf") if not reverse else float("-inf")) + return total_cost + + elif sort_by == "status": + # False (config) comes before True (db) for asc + db_model = model_info.get("db_model", False) + return db_model + + return None + + try: + sorted_models = sorted(all_models, key=get_sort_key, reverse=reverse) + return sorted_models + except Exception as e: + verbose_proxy_logger.exception(f"Error sorting models by {sort_by}: {str(e)}") + return all_models + + def _paginate_models_response( all_models: List[Dict[str, Any]], page: int, @@ -8078,6 +8174,55 @@ async def _filter_models_by_team_id( return filtered_models +async def _find_model_by_id( + model_id: str, + search: Optional[str], + llm_router, + prisma_client, + proxy_config, +) -> tuple[list, Optional[int]]: + """Find a model by its ID and optionally filter by search term.""" + found_model = None + + # First, search in config + if llm_router is not None: + found_model = llm_router.get_model_info(id=model_id) + if found_model: + found_model = copy.deepcopy(found_model) + + # If not found in config, search in database + if found_model is None: + try: + db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( + where={"model_id": model_id} + ) + if db_model: + # Convert database model to router format + decrypted_models = proxy_config.decrypt_model_list_from_db( + [db_model] + ) + if decrypted_models: + found_model = decrypted_models[0] + except Exception as e: + verbose_proxy_logger.exception( + f"Error querying database for modelId {model_id}: {str(e)}" + ) + + # If model found, verify search filter if provided + if found_model is not None: + if search is not None and search.strip(): + search_lower = search.lower().strip() + model_name = found_model.get("model_name", "") + if search_lower not in model_name.lower(): + # Model found but doesn't match search filter + found_model = None + + # Set all_models to the found model or empty list + all_models = [found_model] if found_model is not None else [] + search_total_count: Optional[int] = len(all_models) + return all_models, search_total_count + + @router.get( "/v2/model/info", description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true", @@ -8109,6 +8254,14 @@ async def model_info_v2( None, description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids", ), + sortBy: Optional[str] = fastapi.Query( + None, + description="Field to sort by. Options: model_name, created_at, updated_at, costs, status", + ), + sortOrder: Optional[str] = fastapi.Query( + "asc", + description="Sort order. Options: asc, desc", + ), ): """ BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now. @@ -8136,44 +8289,13 @@ async def model_info_v2( # If modelId is provided, search for the specific model if modelId is not None: - found_model = None - - # First, search in config - if llm_router is not None: - found_model = llm_router.get_model_info(id=modelId) - if found_model: - found_model = copy.deepcopy(found_model) - - # If not found in config, search in database - if found_model is None: - try: - db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( - where={"model_id": modelId} - ) - if db_model: - # Convert database model to router format - decrypted_models = proxy_config.decrypt_model_list_from_db( - [db_model] - ) - if decrypted_models: - found_model = decrypted_models[0] - except Exception as e: - verbose_proxy_logger.exception( - f"Error querying database for modelId {modelId}: {str(e)}" - ) - - # If model found, verify search filter if provided - if found_model is not None: - if search is not None and search.strip(): - search_lower = search.lower().strip() - model_name = found_model.get("model_name", "") - if search_lower not in model_name.lower(): - # Model found but doesn't match search filter - found_model = None - - # Set all_models to the found model or empty list - all_models = [found_model] if found_model is not None else [] - search_total_count: Optional[int] = len(all_models) + all_models, search_total_count = await _find_model_by_id( + model_id=modelId, + search=search, + llm_router=llm_router, + prisma_client=prisma_client, + proxy_config=proxy_config, + ) else: # Normal flow when modelId is not provided all_models = copy.deepcopy(llm_router.model_list) @@ -8193,6 +8315,7 @@ async def model_info_v2( size=size, prisma_client=prisma_client, proxy_config=proxy_config, + sort_by=sortBy, ) if user_models_only: @@ -8236,6 +8359,20 @@ async def model_info_v2( if modelId is not None: search_total_count = len(all_models) + # Apply sorting before pagination + if sortBy: + # Validate sortOrder + if sortOrder and sortOrder.lower() not in ["asc", "desc"]: + raise HTTPException( + status_code=400, + detail=f"Invalid sortOrder: {sortOrder}. Must be 'asc' or 'desc'", + ) + all_models = _sort_models( + all_models=all_models, + sort_by=sortBy, + sort_order=sortOrder or "asc", + ) + verbose_proxy_logger.debug("all_models: %s", all_models) return _paginate_models_response( diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 79f182817f..2c34457c02 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -129,6 +129,14 @@ async def _save_vector_store_to_db_from_rag_ingest( litellm_vector_store_params = ingest_options.get("litellm_vector_store_params", {}) custom_vector_store_name = litellm_vector_store_params.get("vector_store_name") custom_vector_store_description = litellm_vector_store_params.get("vector_store_description") + + # Extract provider-specific params from vector_store_config to save as litellm_params + # This ensures params like aws_region_name, embedding_model, etc. are available for search + provider_specific_params = {} + excluded_keys = {"custom_llm_provider", "vector_store_id"} + for key, value in vector_store_config.items(): + if key not in excluded_keys and value is not None: + provider_specific_params[key] = value # Build file metadata entry using helper file_entry = _build_file_metadata_entry( @@ -167,6 +175,7 @@ async def _save_vector_store_to_db_from_rag_ingest( vector_store_name=vector_store_name, vector_store_description=vector_store_description, vector_store_metadata=initial_metadata, + litellm_params=provider_specific_params if provider_specific_params else None, ) verbose_proxy_logger.info( diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 236f58bb82..af441ee43e 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -221,8 +221,9 @@ async def route_request( "aretrieve_container_file_content", ]: return getattr(llm_router, f"{route_type}")(**data) - # Interactions API: get/delete/cancel don't need model routing + # Interactions API: create with agent, get/delete/cancel don't need model routing if route_type in [ + "acreate_interaction", "aget_interaction", "adelete_interaction", "acancel_interaction", diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index 20059487b4..3daa767188 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.rag.ingestion.file_parsers import extract_text_from_pdf from litellm.rag.text_splitters import RecursiveCharacterTextSplitter from litellm.types.rag import RAGIngestOptions, RAGIngestResponse @@ -193,11 +194,23 @@ class BaseRAGIngestion(ABC): if text: text_to_chunk = text elif file_content and not ocr_was_used: + # Try UTF-8 decode first try: text_to_chunk = file_content.decode("utf-8") except UnicodeDecodeError: - verbose_logger.debug("Binary file detected, skipping text chunking") - return [] + # Check if it's a PDF and try to extract text + if file_content.startswith(b"%PDF"): + verbose_logger.debug("PDF detected, attempting text extraction") + text_to_chunk = extract_text_from_pdf(file_content) + if not text_to_chunk: + verbose_logger.debug( + "PDF text extraction failed. Install 'pypdf' or 'PyPDF2' for PDF support, " + "or enable OCR with a vision model." + ) + return [] + else: + verbose_logger.debug("Binary file detected, skipping text chunking") + return [] if not text_to_chunk: return [] diff --git a/litellm/rag/ingestion/file_parsers/__init__.py b/litellm/rag/ingestion/file_parsers/__init__.py new file mode 100644 index 0000000000..5be68cdbb7 --- /dev/null +++ b/litellm/rag/ingestion/file_parsers/__init__.py @@ -0,0 +1,9 @@ +""" +File parsers for RAG ingestion. + +Provides text extraction utilities for various file formats. +""" + +from .pdf_parser import extract_text_from_pdf + +__all__ = ["extract_text_from_pdf"] diff --git a/litellm/rag/ingestion/file_parsers/pdf_parser.py b/litellm/rag/ingestion/file_parsers/pdf_parser.py new file mode 100644 index 0000000000..9a533ccf13 --- /dev/null +++ b/litellm/rag/ingestion/file_parsers/pdf_parser.py @@ -0,0 +1,70 @@ +""" +PDF text extraction utilities. + +Provides text extraction from PDF files using pypdf or PyPDF2. +""" + +from typing import Optional + +from litellm._logging import verbose_logger + + +def extract_text_from_pdf(file_content: bytes) -> Optional[str]: + """ + Extract text from PDF using pypdf if available. + + Args: + file_content: Raw PDF bytes + + Returns: + Extracted text or None if extraction fails + """ + try: + from io import BytesIO + + # Try pypdf first (most common) + try: + from pypdf import PdfReader as PypdfReader + + pdf_file = BytesIO(file_content) + reader = PypdfReader(pdf_file) + + text_parts = [] + for page in reader.pages: + text = page.extract_text() + if text: + text_parts.append(text) + + if text_parts: + extracted_text = "\n\n".join(text_parts) + verbose_logger.debug(f"Extracted {len(extracted_text)} characters from PDF using pypdf") + return extracted_text + + except ImportError: + verbose_logger.debug("pypdf not available, trying PyPDF2") + + # Fallback to PyPDF2 + try: + from PyPDF2 import PdfReader as PyPDF2Reader + + pdf_file = BytesIO(file_content) + reader = PyPDF2Reader(pdf_file) + + text_parts = [] + for page in reader.pages: + text = page.extract_text() + if text: + text_parts.append(text) + + if text_parts: + extracted_text = "\n\n".join(text_parts) + verbose_logger.debug(f"Extracted {len(extracted_text)} characters from PDF using PyPDF2") + return extracted_text + + except ImportError: + verbose_logger.debug("PyPDF2 not available, PDF extraction requires OCR or pypdf/PyPDF2 library") + + except Exception as e: + verbose_logger.debug(f"PDF text extraction failed: {e}") + + return None diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py new file mode 100644 index 0000000000..e6c166a101 --- /dev/null +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -0,0 +1,573 @@ +""" +S3 Vectors-specific RAG Ingestion implementation. + +S3 Vectors is AWS's native vector storage service that provides: +- Purpose-built vector buckets for storing and querying vectors +- Vector indexes with configurable dimensions and distance metrics +- Metadata filtering for semantic search + +This implementation: +1. Auto-creates vector buckets and indexes if not provided +2. Uses LiteLLM's embedding API (supports any provider) +3. Uses httpx + AWS SigV4 signing (no boto3 dependency for S3 Vectors APIs) +4. Stores vectors with metadata using PutVectors API +""" + +from __future__ import annotations + +import hashlib +import uuid +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import litellm +from litellm._logging import verbose_logger +from litellm.constants import ( + S3_VECTORS_DEFAULT_DIMENSION, + S3_VECTORS_DEFAULT_DISTANCE_METRIC, + S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, +) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion + +if TYPE_CHECKING: + from litellm import Router + from litellm.types.rag import RAGIngestOptions + + +class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): + """ + S3 Vectors RAG ingestion using httpx + AWS SigV4 signing. + + Workflow: + 1. Auto-create vector bucket if needed (CreateVectorBucket API) + 2. Auto-create vector index if needed (CreateVectorIndex API) + 3. Generate embeddings using LiteLLM (supports any provider) + 4. Store vectors with PutVectors API + + Configuration: + - vector_bucket_name: S3 vector bucket name (required) + - index_name: Vector index name (auto-creates if not provided) + - dimension: Vector dimension (default: S3_VECTORS_DEFAULT_DIMENSION) + - distance_metric: "cosine" or "euclidean" (default: S3_VECTORS_DEFAULT_DISTANCE_METRIC) + - non_filterable_metadata_keys: List of metadata keys to exclude from filtering + """ + + def __init__( + self, + ingest_options: "RAGIngestOptions", + router: Optional["Router"] = None, + ): + BaseRAGIngestion.__init__(self, ingest_options=ingest_options, router=router) + BaseAWSLLM.__init__(self) + + # Extract config + self.vector_bucket_name = self.vector_store_config["vector_bucket_name"] + self.index_name = self.vector_store_config.get("index_name") + self.distance_metric = self.vector_store_config.get( + "distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC + ) + self.non_filterable_metadata_keys = self.vector_store_config.get( + "non_filterable_metadata_keys", + S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, + ) + + # Get dimension from config (will be auto-detected on first use if not provided) + self.dimension = self._get_dimension_from_config() + + # Get AWS region using BaseAWSLLM method + _aws_region = self.vector_store_config.get("aws_region_name") + self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( + aws_region_name=str(_aws_region) if _aws_region else None + ) + + # Create httpx client (similar to s3_v2.py) + ssl_verify = self._get_ssl_verify( + ssl_verify=self.vector_store_config.get("ssl_verify") + ) + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.RAG, + params={"ssl_verify": ssl_verify} if ssl_verify is not None else None, + ) + + # Track if infrastructure is initialized + self._config_initialized = False + + async def _get_dimension_from_embedding_request(self) -> int: + """ + Auto-detect dimension by making a test embedding request. + + Makes a single embedding request with a test string to determine + the output dimension of the embedding model. + """ + if not self.embedding_config or "model" not in self.embedding_config: + return S3_VECTORS_DEFAULT_DIMENSION + + try: + model_name = self.embedding_config["model"] + verbose_logger.debug( + f"Auto-detecting dimension by making test embedding request to {model_name}" + ) + + # Make a test embedding request + test_input = "test" + if self.router: + response = await self.router.aembedding(model=model_name, input=[test_input]) + else: + response = await litellm.aembedding(model=model_name, input=[test_input]) + + # Get dimension from the response + if response.data and len(response.data) > 0: + dimension = len(response.data[0]["embedding"]) + verbose_logger.debug( + f"Auto-detected dimension {dimension} for embedding model {model_name}" + ) + return dimension + except Exception as e: + verbose_logger.warning( + f"Could not auto-detect dimension from embedding model: {e}. " + f"Using default dimension of {S3_VECTORS_DEFAULT_DIMENSION}." + ) + + return S3_VECTORS_DEFAULT_DIMENSION + + def _get_dimension_from_config(self) -> Optional[int]: + """ + Get vector dimension from config if explicitly provided. + + Returns None if dimension should be auto-detected. + """ + if "dimension" in self.vector_store_config: + return int(self.vector_store_config["dimension"]) + return None + + async def _ensure_config_initialized(self): + """Lazily initialize S3 Vectors infrastructure.""" + if self._config_initialized: + return + + # Auto-detect dimension if not provided + if self.dimension is None: + self.dimension = await self._get_dimension_from_embedding_request() + + # Ensure vector bucket exists + await self._ensure_vector_bucket_exists() + + # Ensure vector index exists + if not self.index_name: + # Auto-generate index name + unique_id = uuid.uuid4().hex[:8] + self.index_name = f"litellm-index-{unique_id}" + + await self._ensure_vector_index_exists() + + self._config_initialized = True + + async def _sign_and_execute_request( + self, + method: str, + url: str, + data: Optional[str] = None, + headers: Optional[Dict[str, str]] = None, + ) -> Any: + """ + Helper to sign and execute AWS API requests using httpx + SigV4. + + Pattern from litellm/integrations/s3_v2.py + """ + try: + import requests + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError( + "Missing botocore to call S3 Vectors. Run 'pip install boto3'." + ) + + # Get AWS credentials using BaseAWSLLM's get_credentials method + credentials = self.get_credentials( + aws_access_key_id=self.vector_store_config.get("aws_access_key_id"), + aws_secret_access_key=self.vector_store_config.get("aws_secret_access_key"), + aws_session_token=self.vector_store_config.get("aws_session_token"), + aws_region_name=self.aws_region_name, + aws_session_name=self.vector_store_config.get("aws_session_name"), + aws_profile_name=self.vector_store_config.get("aws_profile_name"), + aws_role_name=self.vector_store_config.get("aws_role_name"), + aws_web_identity_token=self.vector_store_config.get("aws_web_identity_token"), + aws_sts_endpoint=self.vector_store_config.get("aws_sts_endpoint"), + aws_external_id=self.vector_store_config.get("aws_external_id"), + ) + + # Prepare headers + if headers is None: + headers = {} + + if data: + headers["Content-Type"] = "application/json" + # Calculate SHA256 hash of the content + content_hash = hashlib.sha256(data.encode("utf-8")).hexdigest() + headers["x-amz-content-sha256"] = content_hash + else: + # For requests without body, use hash of empty string + headers["x-amz-content-sha256"] = hashlib.sha256(b"").hexdigest() + + # Prepare the request + req = requests.Request(method, url, data=data, headers=headers) + prepped = req.prepare() + + # Sign the request + aws_request = AWSRequest( + method=prepped.method, + url=prepped.url, + data=prepped.body, + headers=prepped.headers, + ) + SigV4Auth(credentials, "s3vectors", self.aws_region_name).add_auth(aws_request) + + # Prepare the signed headers + signed_headers = dict(aws_request.headers.items()) + + # Make the request using specific method (pattern from s3_v2.py) + method_upper = method.upper() + if method_upper == "PUT": + response = await self.async_httpx_client.put( + url, data=data, headers=signed_headers + ) + elif method_upper == "POST": + response = await self.async_httpx_client.post( + url, data=data, headers=signed_headers + ) + elif method_upper == "GET": + response = await self.async_httpx_client.get(url, headers=signed_headers) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + + return response + + async def _ensure_vector_bucket_exists(self): + """Create vector bucket if it doesn't exist using GetVectorBucket and CreateVectorBucket APIs.""" + verbose_logger.debug( + f"Ensuring S3 vector bucket exists: {self.vector_bucket_name}" + ) + + # Validate bucket name (AWS S3 naming rules) + if len(self.vector_bucket_name) < 3: + raise ValueError( + f"Invalid vector_bucket_name '{self.vector_bucket_name}': " + f"AWS S3 bucket names must be at least 3 characters long. " + f"Please provide a valid bucket name (e.g., 'my-vector-bucket')." + ) + if not self.vector_bucket_name.replace("-", "").replace(".", "").isalnum(): + raise ValueError( + f"Invalid vector_bucket_name '{self.vector_bucket_name}': " + f"AWS S3 bucket names can only contain lowercase letters, numbers, hyphens, and periods. " + f"Please provide a valid bucket name (e.g., 'my-vector-bucket')." + ) + + # Try to get bucket info using GetVectorBucket API + get_url = f"https://s3vectors.{self.aws_region_name}.api.aws/GetVectorBucket" + get_body = safe_dumps({"vectorBucketName": self.vector_bucket_name}) + + try: + response = await self._sign_and_execute_request("POST", get_url, data=get_body) + if response.status_code == 200: + verbose_logger.debug(f"Vector bucket {self.vector_bucket_name} exists") + return + except Exception as e: + verbose_logger.debug( + f"Bucket check failed (may not exist): {e}, attempting to create" + ) + + # Create vector bucket using CreateVectorBucket API + try: + verbose_logger.debug(f"Creating vector bucket: {self.vector_bucket_name}") + create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateVectorBucket" + create_body = safe_dumps({ + "vectorBucketName": self.vector_bucket_name + }) + + response = await self._sign_and_execute_request("POST", create_url, data=create_body) + + if response.status_code in (200, 201): + verbose_logger.info(f"Created vector bucket: {self.vector_bucket_name}") + elif response.status_code == 409: + # Bucket already exists (ConflictException) + verbose_logger.debug( + f"Vector bucket {self.vector_bucket_name} already exists" + ) + else: + verbose_logger.error(f"CreateVectorBucket failed: {response.status_code} - {response.text}") + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error creating vector bucket: {e}") + raise + + async def _ensure_vector_index_exists(self): + """Create vector index if it doesn't exist using GetIndex and CreateIndex APIs.""" + verbose_logger.debug( + f"Ensuring vector index exists: {self.vector_bucket_name}/{self.index_name}" + ) + + # Try to get index info using GetIndex API + get_url = f"https://s3vectors.{self.aws_region_name}.api.aws/GetIndex" + get_body = safe_dumps({ + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name + }) + + try: + response = await self._sign_and_execute_request("POST", get_url, data=get_body) + if response.status_code == 200: + verbose_logger.debug(f"Vector index {self.index_name} exists") + return + except Exception as e: + verbose_logger.debug( + f"Index check failed (may not exist): {e}, attempting to create" + ) + + # Create vector index using CreateIndex API + try: + verbose_logger.debug( + f"Creating vector index: {self.index_name} with dimension={self.dimension}, metric={self.distance_metric}" + ) + + # Prepare index configuration per AWS API docs + index_config = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + "dataType": "float32", + "dimension": self.dimension, + "distanceMetric": self.distance_metric, + } + + if self.non_filterable_metadata_keys: + index_config["metadataConfiguration"] = { + "nonFilterableMetadataKeys": self.non_filterable_metadata_keys + } + + create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateIndex" + response = await self._sign_and_execute_request( + "POST", create_url, data=safe_dumps(index_config) + ) + + if response.status_code in (200, 201): + verbose_logger.info(f"Created vector index: {self.index_name}") + elif response.status_code == 409: + verbose_logger.debug(f"Vector index {self.index_name} already exists") + else: + verbose_logger.error(f"CreateIndex failed: {response.status_code} - {response.text}") + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error creating vector index: {e}") + raise + + async def _put_vectors(self, vectors: List[Dict[str, Any]]): + """ + Call PutVectors API to store vectors in S3 Vectors. + + Args: + vectors: List of vector objects with keys: "key", "data", "metadata" + """ + verbose_logger.debug( + f"Storing {len(vectors)} vectors in {self.vector_bucket_name}/{self.index_name}" + ) + + url = f"https://s3vectors.{self.aws_region_name}.api.aws/PutVectors" + + # Prepare request body per AWS API docs + request_body = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + "vectors": vectors + } + + try: + response = await self._sign_and_execute_request( + "POST", url, data=safe_dumps(request_body) + ) + + if response.status_code in (200, 201): + verbose_logger.info( + f"Successfully stored {len(vectors)} vectors in index {self.index_name}" + ) + else: + verbose_logger.error( + f"PutVectors failed with status {response.status_code}: {response.text}" + ) + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error storing vectors: {e}") + raise + + async def embed( + self, + chunks: List[str], + ) -> Optional[List[List[float]]]: + """ + Generate embeddings using LiteLLM's embedding API. + + Supports any embedding provider (OpenAI, Bedrock, Cohere, etc.) + """ + if not chunks: + return None + + # Use embedding config from ingest_options or default + if not self.embedding_config: + verbose_logger.warning( + "No embedding config provided, using default text-embedding-3-small" + ) + self.embedding_config = {"model": "text-embedding-3-small"} + + embedding_model = self.embedding_config.get("model", "text-embedding-3-small") + + verbose_logger.debug( + f"Generating embeddings for {len(chunks)} chunks using {embedding_model}" + ) + + # Convert to list to ensure type compatibility + input_chunks: List[str] = list(chunks) + + if self.router: + response = await self.router.aembedding(model=embedding_model, input=input_chunks) + else: + response = await litellm.aembedding(model=embedding_model, input=input_chunks) + + return [item["embedding"] for item in response.data] + + async def store( + self, + file_content: Optional[bytes], + filename: Optional[str], + content_type: Optional[str], + chunks: List[str], + embeddings: Optional[List[List[float]]], + ) -> Tuple[Optional[str], Optional[str]]: + """ + Store vectors in S3 Vectors using PutVectors API. + + Steps: + 1. Ensure vector bucket exists (auto-create if needed) + 2. Ensure vector index exists (auto-create if needed) + 3. Prepare vector data with metadata + 4. Call PutVectors API with httpx + SigV4 signing + + Args: + file_content: Raw file bytes (not used for S3 Vectors) + filename: Name of the file + content_type: MIME type (not used for S3 Vectors) + chunks: Text chunks + embeddings: Vector embeddings + + Returns: + Tuple of (index_name, filename) + """ + # Ensure infrastructure exists + await self._ensure_config_initialized() + + if not embeddings or not chunks: + error_msg = ( + "No text content could be extracted from the file for embedding. " + "Possible causes:\n" + " 1. PDF files require OCR - add 'ocr' config with a vision model (e.g., 'anthropic/claude-3-5-sonnet-20241022')\n" + " 2. Binary files cannot be processed - convert to text first\n" + " 3. File is empty or contains no extractable text\n" + "For PDFs, either enable OCR or use a PDF extraction library to convert to text before ingestion." + ) + verbose_logger.error(error_msg) + raise ValueError(error_msg) + + # Prepare vectors for PutVectors API + vectors = [] + for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): + # Build metadata dict + metadata: Dict[str, str] = { + "source_text": chunk, # Non-filterable (for reference) + "chunk_index": str(i), # Filterable + } + + if filename: + metadata["filename"] = filename # Filterable + + vector_obj = { + "key": f"{filename}_{i}" if filename else f"chunk_{i}", + "data": {"float32": embedding}, + "metadata": metadata, + } + + vectors.append(vector_obj) + + # Call PutVectors API + await self._put_vectors(vectors) + + # Return vector_store_id in format bucket_name:index_name for S3 Vectors search compatibility + vector_store_id = f"{self.vector_bucket_name}:{self.index_name}" + return vector_store_id, filename + + async def query_vector_store( + self, vector_store_id: str, query: str, top_k: int = 5 + ) -> Optional[Dict[str, Any]]: + """ + Query S3 Vectors using QueryVectors API. + + Args: + vector_store_id: Index name + query: Query text + top_k: Number of results to return + + Returns: + Query results with vectors and metadata + """ + verbose_logger.debug(f"Querying index {vector_store_id} with query: {query}") + + # Generate query embedding + if not self.embedding_config: + self.embedding_config = {"model": "text-embedding-3-small"} + + embedding_model = self.embedding_config.get("model", "text-embedding-3-small") + + response = await litellm.aembedding(model=embedding_model, input=[query]) + query_embedding = response.data[0]["embedding"] + + # Call QueryVectors API + url = f"https://s3vectors.{self.aws_region_name}.api.aws/QueryVectors" + + request_body = { + "vectorBucketName": self.vector_bucket_name, + "indexName": vector_store_id, + "queryVector": {"float32": query_embedding}, + "topK": top_k, + "returnDistance": True, + "returnMetadata": True, + } + + try: + response = await self._sign_and_execute_request( + "POST", url, data=safe_dumps(request_body) + ) + + if response.status_code == 200: + results = response.json() + verbose_logger.debug(f"Query returned {len(results.get('vectors', []))} results") + + # Check if query terms appear in results + if results.get("vectors"): + for result in results["vectors"]: + metadata = result.get("metadata", {}) + source_text = metadata.get("source_text", "") + if query.lower() in source_text.lower(): + return results + + # Return results even if exact match not found + return results + else: + verbose_logger.error( + f"QueryVectors failed with status {response.status_code}: {response.text}" + ) + return None + except Exception as e: + verbose_logger.exception(f"Error querying vectors: {e}") + return None diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 620097f83d..571f78f2f7 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -31,6 +31,7 @@ from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion +from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion from litellm.rag.rag_query import RAGQuery from litellm.types.rag import ( RAGIngestOptions, @@ -48,6 +49,7 @@ INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = { "openai": OpenAIRAGIngestion, "bedrock": BedrockRAGIngestion, "gemini": GeminiRAGIngestion, + "s3_vectors": S3VectorsRAGIngestion, } diff --git a/litellm/router.py b/litellm/router.py index 9fe37efa3a..09d71b6b49 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -58,6 +58,7 @@ from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_metadata_variable_name_from_kwargs, ) +from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer @@ -1250,10 +1251,25 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + # Check for silent model experiment + # Make a local copy of litellm_params to avoid mutating the Router's state + litellm_params = deployment["litellm_params"].copy() + silent_model = litellm_params.pop("silent_model", None) + if silent_model is not None: + # Mirroring traffic to a secondary model + # Use shared thread pool for background calls + executor.submit( + self._silent_experiment_completion, + silent_model, + messages, + **kwargs, + ) + + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + kwargs.pop("silent_model", None) # Ensure it's not in kwargs either # No copy needed - data is only read and spread into new dict below - data = deployment["litellm_params"] + data = litellm_params.copy() # Use the local copy of litellm_params model_name = data["model"] potential_model_client = self._get_client( deployment=deployment, kwargs=kwargs @@ -1274,15 +1290,14 @@ class Router: if not self.has_model_id(model): self.routing_strategy_pre_call_checks(deployment=deployment) - response = litellm.completion( - **{ - **data, - "messages": messages, - "caching": self.cache_responses, - "client": model_client, - **kwargs, - } - ) + input_kwargs = { + **data, + "messages": messages, + "caching": self.cache_responses, + "client": model_client, + **kwargs, + } + response = litellm.completion(**input_kwargs) verbose_router_logger.info( f"litellm.completion(model={model_name})\033[32m 200 OK\033[0m" ) @@ -1309,6 +1324,56 @@ class Router: self._set_deployment_num_retries_on_exception(e, deployment) raise e + def _get_silent_experiment_kwargs(self, **kwargs) -> dict: + """ + Prepare kwargs for a silent experiment by ensuring isolation from the primary call. + """ + # Copy kwargs to ensure isolation + silent_kwargs = copy.deepcopy(kwargs) + if "metadata" not in silent_kwargs: + silent_kwargs["metadata"] = {} + + silent_kwargs["metadata"]["is_silent_experiment"] = True + + # Pop logging objects and call IDs to ensure a fresh logging context + # This prevents collisions in the Proxy's database (spend_logs) + silent_kwargs.pop("litellm_call_id", None) + silent_kwargs.pop("litellm_logging_obj", None) + silent_kwargs.pop("standard_logging_object", None) + silent_kwargs.pop("proxy_server_request", None) + + return silent_kwargs + + def _silent_experiment_completion( + self, silent_model: str, messages: List[Any], **kwargs + ): + """ + Run a silent experiment in the background (thread). + """ + try: + # Prevent infinite recursion if silent model also has a silent model + if kwargs.get("metadata", {}).get("is_silent_experiment", False): + return + + messages = copy.deepcopy(messages) + + verbose_router_logger.info( + f"Starting silent experiment for model {silent_model}" + ) + + silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) + + # Trigger the silent request + self.completion( + model=silent_model, + messages=cast(List[Dict[str, str]], messages), + **silent_kwargs, + ) + except Exception as e: + verbose_router_logger.error( + f"Silent experiment failed for model {silent_model}: {str(e)}" + ) + # fmt: off @overload @@ -1517,6 +1582,36 @@ class Router: return FallbackStreamWrapper(stream_with_fallbacks()) + async def _silent_experiment_acompletion( + self, silent_model: str, messages: List[Any], **kwargs + ): + """ + Run a silent experiment in the background. + """ + try: + # Prevent infinite recursion if silent model also has a silent model + if kwargs.get("metadata", {}).get("is_silent_experiment", False): + return + + messages = copy.deepcopy(messages) + + verbose_router_logger.info( + f"Starting silent experiment for model {silent_model}" + ) + + silent_kwargs = self._get_silent_experiment_kwargs(**kwargs) + + # Trigger the silent request + await self.acompletion( + model=silent_model, + messages=cast(List[AllMessageValues], messages), + **silent_kwargs, + ) + except Exception as e: + verbose_router_logger.error( + f"Silent experiment failed for model {silent_model}: {str(e)}" + ) + async def _acompletion( # noqa: PLR0915 self, model: str, messages: List[Dict[str, str]], **kwargs ) -> Union[ModelResponse, CustomStreamWrapper,]: @@ -1563,9 +1658,27 @@ class Router: self._track_deployment_metrics( deployment=deployment, parent_otel_span=parent_otel_span ) + + # Check for silent model experiment + # Make a local copy of litellm_params to avoid mutating the Router's state + litellm_params = deployment["litellm_params"].copy() + silent_model = litellm_params.pop("silent_model", None) + + if silent_model is not None: + # Mirroring traffic to a secondary model + # This is a silent experiment, so we don't want to block the primary request + asyncio.create_task( + self._silent_experiment_acompletion( + silent_model=silent_model, + messages=messages, # Use messages instead of *args + **kwargs, + ) + ) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + kwargs.pop("silent_model", None) # Ensure it's not in kwargs either # No copy needed - data is only read and spread into new dict below - data = deployment["litellm_params"] + data = litellm_params.copy() # Use the local copy of litellm_params model_name = data["model"] @@ -1582,6 +1695,7 @@ class Router: "client": model_client, **kwargs, } + input_kwargs.pop("silent_model", None) _response = litellm.acompletion(**input_kwargs) diff --git a/litellm/types/rag.py b/litellm/types/rag.py index fe237a1343..cae0770868 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -129,9 +129,53 @@ class VertexAIVectorStoreOptions(TypedDict, total=False): import_timeout: Optional[int] # Timeout in seconds (default: 600) +class S3VectorsVectorStoreOptions(TypedDict, total=False): + """ + AWS S3 Vectors configuration. + + Example (auto-create): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings"} + + Example (use existing): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings", + "index_name": "my-index"} + + Example (with credentials): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings", + "litellm_credential_name": "my-aws-creds"} + + Auto-creation creates: S3 vector bucket and vector index (if not provided). + Embeddings are generated using LiteLLM's embedding API (supports any provider). + """ + + custom_llm_provider: Literal["s3_vectors"] + vector_bucket_name: str # Required - S3 vector bucket name + index_name: Optional[str] # Vector index name (auto-creates if not provided) + + # Index configuration (for auto-creation) + dimension: Optional[int] # Vector dimension (auto-detected from embedding model, or default: 1024) + distance_metric: Optional[Literal["cosine", "euclidean"]] # Default: cosine + non_filterable_metadata_keys: Optional[List[str]] # Keys excluded from filtering (e.g., ["source_text"]) + + # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) + litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list + + # AWS auth (uses BaseAWSLLM) + aws_access_key_id: Optional[str] + aws_secret_access_key: Optional[str] + aws_session_token: Optional[str] + aws_region_name: Optional[str] # default: us-west-2 + aws_role_name: Optional[str] + aws_session_name: Optional[str] + aws_profile_name: Optional[str] + aws_web_identity_token: Optional[str] + aws_sts_endpoint: Optional[str] + aws_external_id: Optional[str] + + # Union type for vector store options RAGIngestVectorStoreOptions = Union[ - OpenAIVectorStoreOptions, BedrockVectorStoreOptions, VertexAIVectorStoreOptions + OpenAIVectorStoreOptions, BedrockVectorStoreOptions, VertexAIVectorStoreOptions, S3VectorsVectorStoreOptions ] diff --git a/litellm/types/router.py b/litellm/types/router.py index 8ea7a20753..43943d9e07 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -404,6 +404,10 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): aws_access_key_id: Optional[str] aws_secret_access_key: Optional[str] aws_region_name: Optional[str] + ## AWS S3 VECTORS ## + vector_bucket_name: Optional[str] + index_name: Optional[str] + embedding_model: Optional[str] ## IBM WATSONX ## watsonx_region_name: Optional[str] ## CUSTOM PRICING ## diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2d67e13e92..49c903502d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4,20 +4,22 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union from openai._models import BaseModel as OpenAIObject -from openai.types.audio.transcription_create_params import FileTypes as FileTypes # type: ignore +from openai.types.audio.transcription_create_params import ( + FileTypes as FileTypes, # type: ignore +) from openai.types.chat.chat_completion import ChatCompletion as ChatCompletion from openai.types.completion_usage import ( CompletionTokensDetails, CompletionUsage, PromptTokensDetails, ) +from openai.types.moderation import Categories as Categories from openai.types.moderation import ( - Categories as Categories, CategoryAppliedInputTypes as CategoryAppliedInputTypes, - CategoryScores as CategoryScores, ) +from openai.types.moderation import CategoryScores as CategoryScores +from openai.types.moderation_create_response import Moderation as Moderation from openai.types.moderation_create_response import ( - Moderation as Moderation, ModerationCreateResponse as ModerationCreateResponse, ) from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator @@ -3075,6 +3077,7 @@ class LlmProviders(str, Enum): LLAMA = "meta_llama" NSCALE = "nscale" PG_VECTOR = "pg_vector" + S3_VECTORS = "s3_vectors" HELICONE = "helicone" HYPERBOLIC = "hyperbolic" RECRAFT = "recraft" diff --git a/litellm/utils.py b/litellm/utils.py index bb95be05b5..d7fb4855a4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8452,6 +8452,12 @@ class ProviderConfigManager: ) return RAGFlowVectorStoreConfig() + elif litellm.LlmProviders.S3_VECTORS == provider: + from litellm.llms.s3_vectors.vector_stores.transformation import ( + S3VectorsVectorStoreConfig, + ) + + return S3VectorsVectorStoreConfig() return None @staticmethod @@ -8699,9 +8705,9 @@ class ProviderConfigManager: """ Get Search configuration for a given provider. """ + from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig - from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.firecrawl.search.transformation import FirecrawlSearchConfig from litellm.llms.google_pse.search.transformation import GooglePSESearchConfig from litellm.llms.linkup.search.transformation import LinkupSearchConfig diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 421f2b8b4d..bd090d071c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -13521,6 +13521,43 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-robotics-er-1.5-preview": { + "cache_read_input_token_cost": 0, + "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "output_cost_per_reasoning_token": 2.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-robotics-er-1-5-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true + }, "gemini-2.5-computer-use-preview-10-2025": { "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, diff --git a/requirements.txt b/requirements.txt index 36328a081a..f250fb213e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -68,6 +68,7 @@ jsonschema>=4.23.0,<5.0.0 # validating json schema - aligned with openapi-core + websockets==15.0.1 # for realtime API soundfile==0.12.1 # for audio file processing openapi-core==0.21.0 # for OpenAPI compliance tests +pypdf>=6.6.2 # for PDF text extraction in RAG ingestion ######################## # LITELLM ENTERPRISE DEPENDENCIES diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py b/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py index ce77ff69e0..706e3b7187 100644 --- a/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py +++ b/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py @@ -180,3 +180,68 @@ class TestEnterpriseRouteChecks: # Should not raise exception since management routes are enabled EnterpriseRouteChecks.should_call_route("/config/update") + + +class TestEnterpriseRouteChecksErrorMessages: + """Test that error messages correctly identify which feature requires Enterprise license""" + + @patch("litellm.secret_managers.main.get_secret_bool") + @patch("litellm.proxy.proxy_server.premium_user", False) + def test_disable_llm_api_endpoints_error_message(self, mock_get_secret_bool): + """ + Test that when DISABLE_LLM_API_ENDPOINTS is set without Enterprise license, + the error message correctly mentions 'LLM API ENDPOINTS' + """ + with patch.dict(os.environ, {"DISABLE_LLM_API_ENDPOINTS": "true"}): + with pytest.raises(HTTPException) as exc_info: + EnterpriseRouteChecks.is_llm_api_route_disabled() + + assert exc_info.value.status_code == 500 + assert "DISABLING LLM API ENDPOINTS is an Enterprise feature" in str( + exc_info.value.detail + ) + + @patch("litellm.secret_managers.main.get_secret_bool") + @patch("litellm.proxy.proxy_server.premium_user", False) + def test_disable_admin_endpoints_error_message(self, mock_get_secret_bool): + """ + Test that when DISABLE_ADMIN_ENDPOINTS is set without Enterprise license, + the error message correctly mentions 'ADMIN ENDPOINTS' (not 'LLM API ENDPOINTS') + + This is a regression test for a bug where the error message incorrectly said + 'DISABLING LLM API ENDPOINTS' when the actual issue was DISABLE_ADMIN_ENDPOINTS. + """ + with patch.dict(os.environ, {"DISABLE_ADMIN_ENDPOINTS": "true"}): + with pytest.raises(HTTPException) as exc_info: + EnterpriseRouteChecks.is_management_routes_disabled() + + assert exc_info.value.status_code == 500 + assert "DISABLING ADMIN ENDPOINTS is an Enterprise feature" in str( + exc_info.value.detail + ) + # Ensure it does NOT mention LLM API ENDPOINTS (the old buggy message) + assert "LLM API ENDPOINTS" not in str(exc_info.value.detail) + + @patch("litellm.secret_managers.main.get_secret_bool") + @patch("litellm.proxy.proxy_server.premium_user", True) + def test_disable_llm_api_endpoints_with_premium_user(self, mock_get_secret_bool): + """ + Test that premium users can use DISABLE_LLM_API_ENDPOINTS without error + """ + mock_get_secret_bool.return_value = True + with patch.dict(os.environ, {"DISABLE_LLM_API_ENDPOINTS": "true"}): + # Should not raise exception for premium users + result = EnterpriseRouteChecks.is_llm_api_route_disabled() + assert result is True + + @patch("litellm.secret_managers.main.get_secret_bool") + @patch("litellm.proxy.proxy_server.premium_user", True) + def test_disable_admin_endpoints_with_premium_user(self, mock_get_secret_bool): + """ + Test that premium users can use DISABLE_ADMIN_ENDPOINTS without error + """ + mock_get_secret_bool.return_value = True + with patch.dict(os.environ, {"DISABLE_ADMIN_ENDPOINTS": "true"}): + # Should not raise exception for premium users + result = EnterpriseRouteChecks.is_management_routes_disabled() + assert result is True diff --git a/tests/llm_translation/test_groq.py b/tests/llm_translation/test_groq.py index 33c1a42587..b82cffd418 100644 --- a/tests/llm_translation/test_groq.py +++ b/tests/llm_translation/test_groq.py @@ -11,7 +11,10 @@ import pytest import litellm from base_llm_unit_tests import BaseLLMChatTest -from litellm.llms.groq.chat.transformation import GroqChatConfig +from litellm.llms.groq.chat.transformation import ( + GroqChatConfig, + GroqChatCompletionStreamingHandler, +) class TestGroq(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: @@ -164,3 +167,124 @@ class TestGroqStructuredOutputs: if "tools" in result: tool_names = [t.get("function", {}).get("name") for t in result["tools"]] assert "json_tool_call" not in tool_names + + +class TestGroqReasoning: + """ + Tests for Groq reasoning field mapping. + + Groq returns 'reasoning' field in delta, but LiteLLM expects 'reasoning_content'. + """ + + def test_reasoning_field_mapping_in_streaming_chunks(self): + """ + Test that Groq's 'reasoning' field in streaming chunks is properly mapped + to LiteLLM's 'reasoning_content' field. + """ + handler = GroqChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Simulate a chunk with reasoning field as returned by Groq + groq_chunk = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1769511767, + "model": "qwen/qwen3-32b", + "choices": [ + { + "delta": { + "reasoning": "This is reasoning content", + "role": None, + }, + "finish_reason": None, + "index": 0, + } + ], + } + + # Parse the chunk + parsed_chunk = handler.chunk_parser(groq_chunk) + + # Verify that reasoning was mapped to reasoning_content + assert parsed_chunk.choices[0].delta.reasoning_content == "This is reasoning content" + # Verify that the original 'reasoning' field was removed + assert not hasattr(parsed_chunk.choices[0].delta, "reasoning") + + def test_reasoning_field_not_present(self): + """ + Test that chunks without reasoning field still work correctly. + """ + handler = GroqChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Simulate a chunk without reasoning field + groq_chunk = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1769511767, + "model": "qwen/qwen3-32b", + "choices": [ + { + "delta": { + "content": "Regular content", + "role": "assistant", + }, + "finish_reason": None, + "index": 0, + } + ], + } + + # Parse the chunk + parsed_chunk = handler.chunk_parser(groq_chunk) + + # Verify that content is present + assert parsed_chunk.choices[0].delta.content == "Regular content" + assert parsed_chunk.choices[0].delta.role == "assistant" + # Verify that reasoning_content is not set (it should be deleted by Delta.__init__) + assert not hasattr(parsed_chunk.choices[0].delta, "reasoning_content") + + def test_reasoning_with_tool_calls(self): + """ + Test that reasoning field is properly mapped even when tool_calls are present. + """ + handler = GroqChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + # Simulate a chunk with both reasoning and tool_calls + groq_chunk = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1769511767, + "model": "qwen/qwen3-32b", + "choices": [ + { + "delta": { + "reasoning": "Reasoning before tool call", + "tool_calls": [ + { + "index": 0, + "id": "call_123", + "function": {"name": "test_function", "arguments": "{}"}, + "type": "function", + } + ], + }, + "finish_reason": None, + "index": 0, + } + ], + } + + # Parse the chunk + parsed_chunk = handler.chunk_parser(groq_chunk) + + # Verify that reasoning was mapped to reasoning_content + assert parsed_chunk.choices[0].delta.reasoning_content == "Reasoning before tool call" + # Verify tool_calls are still present + assert parsed_chunk.choices[0].delta.tool_calls is not None + assert len(parsed_chunk.choices[0].delta.tool_calls) == 1 + assert parsed_chunk.choices[0].delta.tool_calls[0]["function"]["name"] == "test_function" diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index 24c005443a..05c6e4984a 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -584,7 +584,7 @@ async def test_get_fuzzy_user_object(): ) assert result == test_user mock_prisma.db.litellm_usertable.find_first.assert_called_with( - where={"user_email": "test@example.com"}, + where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}}, include={"organization_memberships": True}, ) @@ -612,7 +612,7 @@ async def test_get_fuzzy_user_object(): ) assert result == test_user mock_prisma.db.litellm_usertable.find_first.assert_called_with( - where={"user_email": "test@example.com"}, + where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}}, include={"organization_memberships": True}, ) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 839d09b52d..d864a0d986 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1117,7 +1117,10 @@ async def test_create_team_member_add(prisma_client, new_member_method): ) as mock_litellm_usertable, patch( "litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache", new=AsyncMock(return_value=team_obj), - ) as mock_team_obj: + ) as mock_team_obj, patch( + "litellm.proxy.proxy_server.prisma_client.get_data", + new=AsyncMock(return_value=[]), + ) as mock_get_data: mock_client = AsyncMock( return_value=LiteLLM_UserTable( @@ -1126,6 +1129,10 @@ async def test_create_team_member_add(prisma_client, new_member_method): ) mock_litellm_usertable.upsert = mock_client mock_litellm_usertable.find_many = AsyncMock(return_value=None) + # Mock find_first for user_email validation (returns None for new users) + mock_litellm_usertable.find_first = AsyncMock(return_value=None) + # Mock find_unique for user_id validation (returns None for new users) + mock_litellm_usertable.find_unique = AsyncMock(return_value=None) team_mock_client = AsyncMock() original_val = getattr( litellm.proxy.proxy_server.prisma_client.db, "litellm_teamtable" @@ -1299,7 +1306,10 @@ async def test_create_team_member_add_team_admin( ) as mock_litellm_usertable, patch( "litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache", new=AsyncMock(return_value=team_obj), - ) as mock_team_obj: + ) as mock_team_obj, patch( + "litellm.proxy.proxy_server.prisma_client.get_data", + new=AsyncMock(return_value=[]), + ) as mock_get_data: mock_client = AsyncMock( return_value=LiteLLM_UserTable( user_id="1234", max_budget=100, user_email="1234" @@ -1307,6 +1317,10 @@ async def test_create_team_member_add_team_admin( ) mock_litellm_usertable.upsert = mock_client mock_litellm_usertable.find_many = AsyncMock(return_value=None) + # Mock find_first for user_email validation (returns None for new users) + mock_litellm_usertable.find_first = AsyncMock(return_value=None) + # Mock find_unique for user_id validation (returns None for new users) + mock_litellm_usertable.find_unique = AsyncMock(return_value=None) team_mock_client = AsyncMock() original_val = getattr( diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index b15d75a414..9c2939b2da 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -66,6 +66,41 @@ class LargeImageClient: ) +class StreamingLargeImageClient: + """ + Client that streams a large image to test streaming download protection. + This simulates a huge file without actually creating it all in memory. + """ + + def __init__(self, size_mb=100, include_content_length=False): + self.size_mb = size_mb + self.include_content_length = include_content_length + + def get(self, url, follow_redirects=True): + size_bytes = int(self.size_mb * 1024 * 1024) + headers = {"Content-Type": "image/jpeg"} + if self.include_content_length: + headers["Content-Length"] = str(size_bytes) + + # Create a generator that yields chunks without creating the whole file in memory + def generate_chunks(total_size, chunk_size=8192): + bytes_sent = 0 + while bytes_sent < total_size: + chunk = b"x" * min(chunk_size, total_size - bytes_sent) + bytes_sent += len(chunk) + yield chunk + + # Create response with streaming content + response = Response( + status_code=200, + headers=headers, + request=Request("GET", url), + ) + # Mock the iter_bytes method to return our generator + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) + return response + + def test_image_exceeds_size_limit_with_content_length(monkeypatch): """ Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected when Content-Length header is present. @@ -83,6 +118,7 @@ def test_image_exceeds_size_limit_with_content_length(monkeypatch): def test_image_exceeds_size_limit_without_content_length(monkeypatch): """ Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected even without Content-Length header. + This uses the old non-streaming mock for backward compatibility. """ monkeypatch.setattr( litellm, "module_level_client", LargeImageClient(size_mb=100, include_content_length=False) @@ -94,6 +130,29 @@ def test_image_exceeds_size_limit_without_content_length(monkeypatch): assert "exceeds maximum allowed size" in str(excinfo.value) +def test_streaming_download_protects_against_huge_files(monkeypatch): + """ + Test that streaming download aborts early when file exceeds size limit, + preventing memory exhaustion from huge files (e.g., petabyte-sized files). + + This test verifies that the streaming implementation doesn't download the entire + file into memory before checking size. Instead, it should abort as soon as the + limit is exceeded during streaming. + """ + # Simulate a 1GB file - far larger than the 50MB default limit + client = StreamingLargeImageClient(size_mb=1024, include_content_length=False) + monkeypatch.setattr(litellm, "module_level_client", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + convert_url_to_base64("https://example.com/huge-image.jpg") + + # Verify the error message shows it was caught during streaming + assert "exceeds maximum allowed size" in str(excinfo.value) + + # The error should be raised after downloading just slightly more than the limit + # not after downloading the full 1GB + + class SmallImageClient: """ Client that returns a small valid image. @@ -124,6 +183,26 @@ def test_image_within_size_limit(monkeypatch): assert result.startswith("data:image/jpeg;base64,") +def test_streaming_download_handles_petabyte_file(monkeypatch): + """ + Test that streaming download can handle extremely large file URLs (e.g., petabyte-sized) + without attempting to download the entire file or causing memory exhaustion. + + This simulates what happens if a malicious actor or misconfiguration provides + a URL to an extremely large file. + """ + # Simulate a 1 petabyte file (1,000,000 GB) + # Without streaming protection, this would cause OOM or hang indefinitely + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) + monkeypatch.setattr(litellm, "module_level_client", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + convert_url_to_base64("https://example.com/petabyte-file.jpg") + + # Should fail fast without downloading anywhere near 1 petabyte + assert "exceeds maximum allowed size" in str(excinfo.value) + + def test_image_size_limit_disabled(monkeypatch): """ Test that setting MAX_IMAGE_URL_DOWNLOAD_SIZE_MB to 0 disables all image URL downloads. diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 1cb84d32c1..98b392a353 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -279,3 +279,188 @@ def test_output_format_with_no_schema(): # Content should remain as string (not converted to list) assert isinstance(last_user_message["content"], str) assert last_user_message["content"] == "Hello" + + +def test_advanced_tool_use_header_translation_for_opus_4_5(): + """ + Test that advanced-tool-use-2025-11-20 header is translated to Bedrock-specific headers + for Claude Opus 4.5. + + Regression test for: Claude Code sends advanced-tool-use header which needs to be + translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for Bedrock + Invoke API on Claude Opus 4.5. + + Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Simulate advanced-tool-use header from Claude Code + headers = { + "anthropic-beta": "advanced-tool-use-2025-11-20" + } + + # Test with Claude Opus 4.5 + result = config.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-5-20250514-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + # Verify advanced-tool-use header was removed + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + assert "advanced-tool-use-2025-11-20" not in beta_headers, \ + "advanced-tool-use header should be removed for Bedrock" + + # Verify Bedrock-specific headers were added + assert "tool-search-tool-2025-10-19" in beta_headers, \ + "tool-search-tool-2025-10-19 should be added for Opus 4.5" + assert "tool-examples-2025-10-29" in beta_headers, \ + "tool-examples-2025-10-29 should be added for Opus 4.5" + + +def test_advanced_tool_use_header_filtered_for_non_opus_4_5(): + """ + Test that advanced-tool-use-2025-11-20 header is filtered out for non-Opus 4.5 models + without adding Bedrock-specific headers. + + The translation to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 should + only happen for Claude Opus 4.5. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Simulate advanced-tool-use header from Claude Code + headers = { + "anthropic-beta": "advanced-tool-use-2025-11-20" + } + + # Test with Claude Sonnet 4.5 (not Opus 4.5) + result = config.transform_anthropic_messages_request( + model="anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + # Verify advanced-tool-use header was removed + beta_headers = result.get("anthropic_beta", []) + assert "advanced-tool-use-2025-11-20" not in beta_headers, \ + "advanced-tool-use header should be removed for Bedrock" + + # Verify Bedrock-specific headers were NOT added (only for Opus 4.5) + assert "tool-search-tool-2025-10-19" not in beta_headers, \ + "tool-search-tool should not be added for non-Opus 4.5 models" + assert "tool-examples-2025-10-29" not in beta_headers, \ + "tool-examples should not be added for non-Opus 4.5 models" + + +def test_advanced_tool_use_header_translation_with_multiple_beta_headers(): + """ + Test that advanced-tool-use header translation works correctly when multiple + beta headers are present. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + messages = [ + {"role": "user", "content": "What's the weather like?"} + ] + + anthropic_messages_optional_request_params = { + "max_tokens": 100, + } + + # Multiple beta headers including advanced-tool-use + headers = { + "anthropic-beta": "claude-code-20250219,advanced-tool-use-2025-11-20,interleaved-thinking-2025-05-14" + } + + # Test with Claude Opus 4.5 + result = config.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-5-20250514-v1:0", + messages=messages, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + litellm_params={}, + headers=headers, + ) + + beta_headers = result.get("anthropic_beta", []) + + # Verify advanced-tool-use was removed + assert "advanced-tool-use-2025-11-20" not in beta_headers + + # Verify Bedrock-specific headers were added + assert "tool-search-tool-2025-10-19" in beta_headers + assert "tool-examples-2025-10-29" in beta_headers + + # Verify other beta headers are preserved + assert "claude-code-20250219" in beta_headers + assert "interleaved-thinking-2025-05-14" in beta_headers + + +def test_opus_4_5_model_detection(): + """ + Test that the _is_claude_opus_4_5 method correctly identifies Opus 4.5 models + with various naming conventions. + """ + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + + # Test various Opus 4.5 naming patterns + opus_4_5_models = [ + "anthropic.claude-opus-4-5-20250514-v1:0", + "anthropic.claude-opus-4.5-20250514-v1:0", + "anthropic.claude-opus_4_5-20250514-v1:0", + "anthropic.claude-opus_4.5-20250514-v1:0", + "us.anthropic.claude-opus-4-5-20250514-v1:0", + "ANTHROPIC.CLAUDE-OPUS-4-5-20250514-V1:0", # Case insensitive + ] + + for model in opus_4_5_models: + assert config._is_claude_opus_4_5(model), \ + f"Should detect {model} as Opus 4.5" + + # Test non-Opus 4.5 models + non_opus_4_5_models = [ + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-opus-4-20250514-v1:0", # Opus 4, not 4.5 + "anthropic.claude-opus-4-1-20250514-v1:0", # Opus 4.1, not 4.5 + "anthropic.claude-haiku-4-5-20251001-v1:0", + ] + + for model in non_opus_4_5_models: + assert not config._is_claude_opus_4_5(model), \ + f"Should not detect {model} as Opus 4.5" diff --git a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py index b9324e4966..e7b6de29b6 100644 --- a/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py +++ b/tests/test_litellm/llms/bedrock/test_anthropic_beta_support.py @@ -390,3 +390,103 @@ class TestAnthropicBetaHeaderSupport: "anthropic_beta SHOULD be added for Anthropic models with cross-region prefix." ) assert "context-1m-2025-08-07" in additional_fields["anthropic_beta"] + + def test_messages_advanced_tool_use_translation_opus_4_5(self): + """Test that advanced-tool-use header is translated to Bedrock-specific headers for Opus 4.5. + + Regression test for: Claude Code sends advanced-tool-use-2025-11-20 header which needs + to be translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for + Bedrock Invoke API on Claude Opus 4.5. + + Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-opus-4-5-20250514-v1:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers, ( + "advanced-tool-use-2025-11-20 should be removed for Bedrock Invoke API" + ) + + # Bedrock-specific headers should be added for Opus 4.5 + assert "tool-search-tool-2025-10-19" in beta_headers, ( + "tool-search-tool-2025-10-19 should be added for Opus 4.5" + ) + assert "tool-examples-2025-10-29" in beta_headers, ( + "tool-examples-2025-10-29 should be added for Opus 4.5" + ) + + def test_messages_advanced_tool_use_translation_sonnet_4_5(self): + """Test that advanced-tool-use header is translated to Bedrock-specific headers for Sonnet 4.5. + + Regression test for: Claude Code sends advanced-tool-use-2025-11-20 header which needs + to be translated to tool-search-tool-2025-10-19 and tool-examples-2025-10-29 for + Bedrock Invoke API on Claude Sonnet 4.5. + + Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + assert "anthropic_beta" in result + beta_headers = result["anthropic_beta"] + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers, ( + "advanced-tool-use-2025-11-20 should be removed for Bedrock Invoke API" + ) + + # Bedrock-specific headers should be added for Sonnet 4.5 + assert "tool-search-tool-2025-10-19" in beta_headers, ( + "tool-search-tool-2025-10-19 should be added for Sonnet 4.5" + ) + assert "tool-examples-2025-10-29" in beta_headers, ( + "tool-examples-2025-10-29 should be added for Sonnet 4.5" + ) + + def test_messages_advanced_tool_use_filtered_unsupported_model(self): + """Test that advanced-tool-use header is filtered out for models that don't support tool search. + + The translation to Bedrock-specific headers should only happen for models that + support tool search on Bedrock (Opus 4.5, Sonnet 4.5). + For other models, the advanced-tool-use header should just be removed. + """ + config = AmazonAnthropicClaudeMessagesConfig() + headers = {"anthropic-beta": "advanced-tool-use-2025-11-20"} + + # Test with Claude 3.5 Sonnet (does NOT support tool search on Bedrock) + result = config.transform_anthropic_messages_request( + model="us.anthropic.claude-3-5-sonnet-20241022-v2:0", + messages=[{"role": "user", "content": "Test"}], + anthropic_messages_optional_request_params={"max_tokens": 100}, + litellm_params={}, + headers=headers + ) + + beta_headers = result.get("anthropic_beta", []) + + # advanced-tool-use should be removed + assert "advanced-tool-use-2025-11-20" not in beta_headers + + # Bedrock-specific headers should NOT be added for unsupported models + assert "tool-search-tool-2025-10-19" not in beta_headers + assert "tool-examples-2025-10-29" not in beta_headers diff --git a/tests/test_litellm/llms/s3_vectors/__init__.py b/tests/test_litellm/llms/s3_vectors/__init__.py new file mode 100644 index 0000000000..d4b0c4d855 --- /dev/null +++ b/tests/test_litellm/llms/s3_vectors/__init__.py @@ -0,0 +1 @@ +# S3 Vectors tests diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py b/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py new file mode 100644 index 0000000000..231735c1de --- /dev/null +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py @@ -0,0 +1 @@ +# S3 Vectors vector store tests diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py new file mode 100644 index 0000000000..3a84da3154 --- /dev/null +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -0,0 +1,115 @@ +from unittest.mock import MagicMock, Mock + +import httpx +import pytest + +from litellm.llms.s3_vectors.vector_stores.transformation import ( + S3VectorsVectorStoreConfig, +) +from litellm.types.vector_stores import VectorStoreSearchResponse + + +class TestS3VectorsVectorStoreConfig: + def test_init(self): + """Test that S3VectorsVectorStoreConfig initializes correctly""" + config = S3VectorsVectorStoreConfig() + assert config is not None + + def test_get_supported_openai_params(self): + """Test that supported OpenAI params are returned""" + config = S3VectorsVectorStoreConfig() + params = config.get_supported_openai_params("test-model") + assert "max_num_results" in params + + def test_get_complete_url(self): + """Test URL generation for S3 Vectors""" + config = S3VectorsVectorStoreConfig() + litellm_params = {"aws_region_name": "us-west-2"} + url = config.get_complete_url(None, litellm_params) + assert url == "https://s3vectors.us-west-2.api.aws" + + def test_get_complete_url_missing_region(self): + """Test that missing region raises error""" + config = S3VectorsVectorStoreConfig() + litellm_params = {} + with pytest.raises(ValueError, match="aws_region_name is required"): + config.get_complete_url(None, litellm_params) + + @pytest.mark.skip(reason="Requires embedding API call, tested in integration tests") + def test_transform_search_request(self): + """Test search request transformation""" + # This test requires making an actual embedding API call + # It's better tested in integration tests + pass + + def test_transform_search_request_invalid_vector_store_id(self): + """Test that invalid vector_store_id format raises error""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + + with pytest.raises( + ValueError, match="vector_store_id must be in format 'bucket_name:index_name'" + ): + config.transform_search_vector_store_request( + vector_store_id="invalid-format", + query="test query", + vector_store_search_optional_params={}, + api_base="https://s3vectors.us-west-2.api.aws", + litellm_logging_obj=mock_logging_obj, + litellm_params={}, + ) + + def test_transform_search_response(self): + """Test search response transformation""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {"query": "test query"} + + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "vectors": [ + { + "distance": 0.05, # S3 Vectors returns distance, not score + "metadata": { + "source_text": "This is test content", + "chunk_index": "0", + "filename": "test.pdf", + }, + }, + { + "distance": 0.15, + "metadata": { + "source_text": "More test content", + "chunk_index": "1", + }, + }, + ] + } + mock_response.status_code = 200 + mock_response.headers = {} + + result = config.transform_search_vector_store_response( + mock_response, mock_logging_obj + ) + + # VectorStoreSearchResponse is a TypedDict, so check structure instead of isinstance + assert result["object"] == "vector_store.search_results.page" + assert result["search_query"] == "test query" + assert len(result["data"]) == 2 + # Score should be 1 - distance (cosine similarity) + assert result["data"][0]["score"] == 0.95 # 1 - 0.05 + assert result["data"][0]["content"][0]["text"] == "This is test content" + assert result["data"][0]["filename"] == "test.pdf" + assert result["data"][1]["score"] == 0.85 # 1 - 0.15 + assert result["data"][1]["content"][0]["text"] == "More test content" + + def test_map_openai_params(self): + """Test OpenAI parameter mapping""" + config = S3VectorsVectorStoreConfig() + non_default_params = {"max_num_results": 5} + optional_params = {} + + result = config.map_openai_params(non_default_params, optional_params, False) + + assert result["maxResults"] == 5 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 3df0dc881e..af0591e5f5 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -28,6 +28,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _can_object_call_vector_stores, + _get_fuzzy_user_object, _get_team_db_check, _virtual_key_max_budget_alert_check, _virtual_key_soft_budget_check, @@ -1331,3 +1332,42 @@ async def test_virtual_key_max_budget_alert_check_scenarios( assert ( alert_triggered == expect_alert ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}" + + +@pytest.mark.asyncio +async def test_get_fuzzy_user_object_case_insensitive_email(): + """Test that _get_fuzzy_user_object uses case-insensitive email lookup""" + # Setup mock Prisma client + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_usertable = MagicMock() + + # Mock user data with mixed case email + test_user = LiteLLM_UserTable( + user_id="test_123", + sso_user_id=None, + user_email="Test@Example.com", # Mixed case in DB + organization_memberships=[], + max_budget=None, + ) + + # Test: SSO ID not found, find by email with different casing + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user) + + # Search with lowercase email (different from DB) + result = await _get_fuzzy_user_object( + prisma_client=mock_prisma, + sso_user_id=None, + user_email="test@example.com", # Lowercase search + ) + + # Verify user was found despite case difference + assert result == test_user + + # Verify the query used case-insensitive mode + mock_prisma.db.litellm_usertable.find_first.assert_called_once() + call_args = mock_prisma.db.litellm_usertable.find_first.call_args + assert call_args.kwargs["where"]["user_email"]["equals"] == "test@example.com" + assert call_args.kwargs["where"]["user_email"]["mode"] == "insensitive" + assert call_args.kwargs["include"] == {"organization_memberships": True} diff --git a/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py b/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py new file mode 100644 index 0000000000..1063f59afb --- /dev/null +++ b/tests/test_litellm/proxy/google_endpoints/test_interactions_agent_param.py @@ -0,0 +1,75 @@ +""" +Test for interactions endpoint agent parameter handling. + +Tests that the /v1beta/interactions endpoint correctly extracts +the `agent` parameter as a fallback when `model` is not provided. +""" + +import pytest + + +class TestInteractionsAgentParameter: + """Test agent parameter handling in interactions endpoint.""" + + def test_agent_parameter_fallback_logic(self): + """ + Test the core logic: model or agent extraction. + + This tests the fix in endpoints.py line ~267: + model=data.get("model") or data.get("agent") + """ + # Case 1: Only agent provided (Deep Research use case) + data = { + "agent": "deep-research-pro-preview-12-2025", + "input": "Research quantum computing", + "background": True, + } + model = data.get("model") or data.get("agent") + assert model == "deep-research-pro-preview-12-2025" + + # Case 2: Only model provided (normal use case) + data = { + "model": "gemini-2.5-flash", + "input": "Hello world", + } + model = data.get("model") or data.get("agent") + assert model == "gemini-2.5-flash" + + # Case 3: Both provided (model takes precedence) + data = { + "model": "gemini-2.5-flash", + "agent": "deep-research-pro-preview-12-2025", + "input": "Test", + } + model = data.get("model") or data.get("agent") + assert model == "gemini-2.5-flash" + + # Case 4: Neither provided + data = { + "input": "Test", + } + model = data.get("model") or data.get("agent") + assert model is None + + def test_route_type_in_skip_model_routing_list(self): + """ + Test that acreate_interaction is in the list of routes + that skip model-based routing. + + This tests the fix in route_llm_request.py. + """ + # The list of routes that skip model routing for interactions + skip_model_routing_routes = [ + "acreate_interaction", + "aget_interaction", + "adelete_interaction", + "acancel_interaction", + ] + + # acreate_interaction should be in the list (this is the fix) + assert "acreate_interaction" in skip_model_routing_routes + + # All interaction routes should be covered + assert "aget_interaction" in skip_model_routing_routes + assert "adelete_interaction" in skip_model_routing_routes + assert "acancel_interaction" in skip_model_routing_routes diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a97ed93a85..467ee3661d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, + _validate_and_populate_member_user_info, delete_team, router, team_member_add_duplication_check, @@ -5466,3 +5467,122 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) ) and mock_db_client.db.litellm_verificationtoken.find_many.called: # If it was called, that's unexpected for admin users assert False, "API keys should not be fetched for team admin users" + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_both_provided_match(): + """ + Test _validate_and_populate_member_user_info when both user_email and user_id + are provided and they match the same user in the database. + """ + # Create member with both user_email and user_id + member = Member(user_email="test@example.com", user_id="user-123", role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user object that matches both email and user_id + mock_user = MagicMock() + mock_user.user_id = "user-123" + mock_user.user_email = "test@example.com" + + # Mock get_data to return single user matching email + mock_prisma_client.get_data = AsyncMock(return_value=[mock_user]) + + # Call the function + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify result matches input (both already provided and match) + assert result.user_email == "test@example.com" + assert result.user_id == "user-123" + + # Verify get_data was called with correct parameters + mock_prisma_client.get_data.assert_called_once_with( + key_val={"user_email": "test@example.com"}, + table_name="user", + query_type="find_all", + ) + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_only_email_provided(): + """ + Test _validate_and_populate_member_user_info when only user_email is provided. + Should populate user_id from database. + """ + # Create member with only user_email + member = Member(user_email="test@example.com", user_id=None, role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock user object from find_first + mock_user_find_first = MagicMock() + mock_user_find_first.user_id = "user-456" + mock_user_find_first.user_email = "test@example.com" + + # Mock find_first to return the user + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + return_value=mock_user_find_first + ) + + # Mock get_data to return single user (no duplicates) + mock_prisma_client.get_data = AsyncMock(return_value=[mock_user_find_first]) + + # Call the function + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify user_id was populated + assert result.user_email == "test@example.com" + assert result.user_id == "user-456" + + # Verify find_first was called with correct parameters + mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( + where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}} + ) + + # Verify get_data was called to check for duplicates + mock_prisma_client.get_data.assert_called_once_with( + key_val={"user_email": "test@example.com"}, + table_name="user", + query_type="find_all", + ) + + +@pytest.mark.asyncio +async def test_validate_and_populate_member_user_info_only_user_id_not_found(): + """ + Test _validate_and_populate_member_user_info when only user_id is provided + but the user doesn't exist in the database. Should allow it to pass with + user_email as None (will be upserted later). + """ + # Create member with only user_id + member = Member(user_email=None, user_id="nonexistent-user", role="user") + + # Mock prisma client + mock_prisma_client = MagicMock() + + # Mock find_unique to return None (user not found) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + # Call the function - should NOT raise an exception + result = await _validate_and_populate_member_user_info( + member=member, + prisma_client=mock_prisma_client, + ) + + # Verify the result - should return member with user_id set and user_email as None + assert result.user_id == "nonexistent-user" + assert result.user_email is None + assert result.role == "user" + + # Verify find_unique was called with correct parameters + mock_prisma_client.db.litellm_usertable.find_unique.assert_called_once_with( + where={"user_id": "nonexistent-user"} + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index f983af2d0b..5e9078ea87 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -24,6 +24,7 @@ from litellm.proxy.management_endpoints.ui_sso import ( GoogleSSOHandler, MicrosoftSSOHandler, SSOAuthenticationHandler, + normalize_email, ) from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, @@ -667,6 +668,85 @@ def test_build_sso_user_update_data_without_role(): assert "user_role" not in update_data +def test_normalize_email(): + """ + Test that normalize_email correctly lowercases email addresses and handles edge cases. + """ + # Test with lowercase email + assert normalize_email("test@example.com") == "test@example.com" + + # Test with uppercase email + assert normalize_email("TEST@EXAMPLE.COM") == "test@example.com" + + # Test with mixed case email + assert normalize_email("Test.User@Example.COM") == "test.user@example.com" + + # Test with None + assert normalize_email(None) is None + + # Test with empty string + assert normalize_email("") == "" + + +def test_build_sso_user_update_data_normalizes_email(): + """ + Test that _build_sso_user_update_data normalizes email addresses to lowercase. + """ + from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import _build_sso_user_update_data + + sso_result = CustomOpenID( + id="test-user-789", + email="Test.User@Example.COM", + display_name="Test User", + provider="microsoft", + team_ids=[], + user_role=None, + ) + + update_data = _build_sso_user_update_data( + result=sso_result, + user_email="Test.User@Example.COM", + user_id="test-user-789", + ) + + # Email should be normalized to lowercase + assert update_data["user_email"] == "test.user@example.com" + assert "user_role" not in update_data + + +def test_generic_response_convertor_normalizes_email(): + """ + Test that generic_response_convertor normalizes email addresses. + """ + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + mock_response = { + "preferred_username": "user123", + "email": "Test.User@Example.COM", + "sub": "Test User", + "first_name": "Test", + "last_name": "User", + "provider": "generic", + } + + # Mock JWT handler + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + result = generic_response_convertor( + response=mock_response, + jwt_handler=mock_jwt_handler, + sso_jwt_handler=None, + role_mappings=None, + ) + + # Email should be normalized to lowercase + assert result.email == "test.user@example.com" + assert result.id == "user123" + assert result.display_name == "Test User" + + @pytest.mark.asyncio async def test_upsert_sso_user_updates_role_for_existing_user(): """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2f6eccaf04..22a60f220e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4013,6 +4013,279 @@ async def test_model_info_v2_filter_by_team_id(monkeypatch): app.dependency_overrides = original_overrides +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sort_by,sort_order,expected_order", + [ + # Test model_name sorting + ("model_name", "asc", ["a-model", "b-model", "z-model"]), + ("model_name", "desc", ["z-model", "b-model", "a-model"]), + # Test created_at sorting + ("created_at", "asc", ["old-model", "mid-model", "new-model"]), + ("created_at", "desc", ["new-model", "mid-model", "old-model"]), + # Test updated_at sorting + ("updated_at", "asc", ["old-updated", "mid-updated", "new-updated"]), + ("updated_at", "desc", ["new-updated", "mid-updated", "old-updated"]), + # Test costs sorting + ("costs", "asc", ["low-cost", "mid-cost", "high-cost"]), + ("costs", "desc", ["high-cost", "mid-cost", "low-cost"]), + # Test status sorting (False/config models come before True/db models in asc) + ("status", "asc", ["config-model-1", "config-model-2", "db-model"]), + ("status", "desc", ["db-model", "config-model-1", "config-model-2"]), + ], +) +async def test_model_info_v2_sorting(monkeypatch, sort_by, sort_order, expected_order): + """ + Test sorting functionality for /v2/model/info endpoint. + Tests all sortBy fields (model_name, created_at, updated_at, costs, status) + with both asc and desc sort orders. + """ + from datetime import datetime, timedelta + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + # Create base time for date comparisons + base_time = datetime(2024, 1, 1, 12, 0, 0) + + # Create mock models with different values for each sort field + mock_models = [] + + if sort_by == "model_name": + # Models with different names + mock_models = [ + { + "model_name": "z-model", + "litellm_params": {"model": "z-model"}, + "model_info": {"id": "z-model"}, + }, + { + "model_name": "a-model", + "litellm_params": {"model": "a-model"}, + "model_info": {"id": "a-model"}, + }, + { + "model_name": "b-model", + "litellm_params": {"model": "b-model"}, + "model_info": {"id": "b-model"}, + }, + ] + elif sort_by == "created_at": + # Models with different created_at timestamps + mock_models = [ + { + "model_name": "new-model", + "litellm_params": {"model": "new-model"}, + "model_info": { + "id": "new-model", + "created_at": (base_time + timedelta(days=3)).isoformat(), + }, + }, + { + "model_name": "old-model", + "litellm_params": {"model": "old-model"}, + "model_info": { + "id": "old-model", + "created_at": (base_time - timedelta(days=3)).isoformat(), + }, + }, + { + "model_name": "mid-model", + "litellm_params": {"model": "mid-model"}, + "model_info": { + "id": "mid-model", + "created_at": base_time.isoformat(), + }, + }, + ] + elif sort_by == "updated_at": + # Models with different updated_at timestamps + mock_models = [ + { + "model_name": "new-updated", + "litellm_params": {"model": "new-updated"}, + "model_info": { + "id": "new-updated", + "updated_at": (base_time + timedelta(days=3)).isoformat(), + }, + }, + { + "model_name": "old-updated", + "litellm_params": {"model": "old-updated"}, + "model_info": { + "id": "old-updated", + "updated_at": (base_time - timedelta(days=3)).isoformat(), + }, + }, + { + "model_name": "mid-updated", + "litellm_params": {"model": "mid-updated"}, + "model_info": { + "id": "mid-updated", + "updated_at": base_time.isoformat(), + }, + }, + ] + elif sort_by == "costs": + # Models with different costs (input_cost + output_cost) + mock_models = [ + { + "model_name": "high-cost", + "litellm_params": {"model": "high-cost"}, + "model_info": { + "id": "high-cost", + "input_cost_per_token": 0.00005, + "output_cost_per_token": 0.00015, + }, + }, + { + "model_name": "low-cost", + "litellm_params": {"model": "low-cost"}, + "model_info": { + "id": "low-cost", + "input_cost_per_token": 0.00001, + "output_cost_per_token": 0.00003, + }, + }, + { + "model_name": "mid-cost", + "litellm_params": {"model": "mid-cost"}, + "model_info": { + "id": "mid-cost", + "input_cost_per_token": 0.00003, + "output_cost_per_token": 0.00007, + }, + }, + ] + elif sort_by == "status": + # Models with different db_model status (False = config, True = db) + mock_models = [ + { + "model_name": "db-model", + "litellm_params": {"model": "db-model"}, + "model_info": {"id": "db-model", "db_model": True}, + }, + { + "model_name": "config-model-1", + "litellm_params": {"model": "config-model-1"}, + "model_info": {"id": "config-model-1", "db_model": False}, + }, + { + "model_name": "config-model-2", + "litellm_params": {"model": "config-model-2"}, + "model_info": {"id": "config-model-2", "db_model": False}, + }, + ] + + # Mock llm_router + mock_router = MagicMock() + mock_router.model_list = mock_models + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock proxy_config.get_config + mock_get_config = AsyncMock(return_value={}) + + # Mock user authentication + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.api_key = "test-key" + mock_user_api_key_dict.team_models = [] + mock_user_api_key_dict.models = [] + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + # Override auth dependency + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_api_key_dict + + client = TestClient(app) + try: + # Test sorting with specified sortBy and sortOrder + response = client.get( + "/v2/model/info", params={"sortBy": sort_by, "sortOrder": sort_order} + ) + assert response.status_code == 200 + data = response.json() + assert len(data["data"]) == len(expected_order) + + # Verify models are in expected order + actual_order = [m["model_name"] for m in data["data"]] + assert actual_order == expected_order, ( + f"Sorting failed for sortBy={sort_by}, sortOrder={sort_order}. " + f"Expected: {expected_order}, Got: {actual_order}" + ) + + finally: + app.dependency_overrides = original_overrides + + +@pytest.mark.asyncio +async def test_model_info_v2_sorting_invalid_sort_order(monkeypatch): + """ + Test that invalid sortOrder values return a 400 error. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth + + # Create mock models + mock_models = [ + { + "model_name": "test-model", + "litellm_params": {"model": "test-model"}, + "model_info": {"id": "test-model"}, + } + ] + + # Mock llm_router + mock_router = MagicMock() + mock_router.model_list = mock_models + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock proxy_config.get_config + mock_get_config = AsyncMock(return_value={}) + + # Mock user authentication + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.api_key = "test-key" + mock_user_api_key_dict.team_models = [] + mock_user_api_key_dict.models = [] + + # Apply monkeypatches + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + + # Override auth dependency + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_api_key_dict + + client = TestClient(app) + try: + # Test invalid sortOrder + response = client.get( + "/v2/model/info", params={"sortBy": "model_name", "sortOrder": "invalid"} + ) + assert response.status_code == 400 + data = response.json() + assert "Invalid sortOrder" in data["detail"] + + finally: + app.dependency_overrides = original_overrides + + @pytest.mark.asyncio async def test_apply_search_filter_to_models(monkeypatch): """ diff --git a/tests/test_litellm/test_router_silent_experiment.py b/tests/test_litellm/test_router_silent_experiment.py new file mode 100644 index 0000000000..9b82cde13c --- /dev/null +++ b/tests/test_litellm/test_router_silent_experiment.py @@ -0,0 +1,157 @@ +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +from litellm.router import Router + + +@pytest.mark.asyncio +async def test_router_silent_experiment_acompletion(): + """ + Test that silent_model triggers a background acompletion call + and that the silent_model parameter is stripped from both calls. + """ + model_list = [ + { + "model_name": "primary-model", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "fake-key", + "silent_model": "silent-model", + }, + }, + { + "model_name": "silent-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key", + }, + }, + ] + + router = Router(model_list=model_list) + + # Mock litellm.acompletion + mock_acompletion = MagicMock() + # Create a future that resolves to a ModelResponse + mock_response = litellm.ModelResponse(choices=[{"message": {"content": "hello"}}]) + future = asyncio.Future() + future.set_result(mock_response) + mock_acompletion.return_value = future + + with patch("litellm.acompletion", mock_acompletion): + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert response.choices[0].message.content == "hello" + + # Give the background task a moment to trigger (it's an asyncio task) + await asyncio.sleep(0.1) + + # Should have 2 calls: one for primary, one for silent + assert mock_acompletion.call_count == 2 + + # Check call arguments + call_args_list = mock_acompletion.call_args_list + + # Verify no silent_model in any call to litellm.acompletion + for call in call_args_list: + args, kwargs = call + assert "silent_model" not in kwargs + if "metadata" in kwargs: + # One call should have is_silent_experiment=True + pass + + # Find the silent call + silent_call = next( + ( + c + for c in call_args_list + if c[1].get("metadata", {}).get("is_silent_experiment") is True + ), + None, + ) + assert silent_call is not None + assert silent_call[1]["model"] == "openai/gpt-4" + + # Find the primary call + primary_call = next( + ( + c + for c in call_args_list + if not c[1].get("metadata", {}).get("is_silent_experiment") + ), + None, + ) + assert primary_call is not None + assert primary_call[1]["model"] == "openai/gpt-3.5-turbo" + + +def test_router_silent_experiment_completion(): + """ + Test that silent_model triggers a background completion call (sync) + and that the silent_model parameter is stripped. + """ + model_list = [ + { + "model_name": "primary-model", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "fake-key", + "silent_model": "silent-model", + }, + }, + { + "model_name": "silent-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key", + }, + }, + ] + + router = Router(model_list=model_list) + + # Mock litellm.completion + mock_completion = MagicMock() + mock_response = litellm.ModelResponse(choices=[{"message": {"content": "hello"}}]) + mock_completion.return_value = mock_response + + with patch("litellm.completion", mock_completion): + response = router.completion( + model="primary-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert response.choices[0].message.content == "hello" + + # The sync background call uses a thread pool. We might need to wait a bit. + import time + + time.sleep(0.5) + + # Should have 2 calls + assert mock_completion.call_count == 2 + + call_args_list = mock_completion.call_args_list + + # Verify no silent_model in any call + for call in call_args_list: + args, kwargs = call + assert "silent_model" not in kwargs + + # Find the silent call + silent_call = next( + ( + c + for c in call_args_list + if c[1].get("metadata", {}).get("is_silent_experiment") is True + ), + None, + ) + assert silent_call is not None + assert silent_call[1]["model"] == "openai/gpt-4" diff --git a/tests/vector_store_tests/rag/test_rag_s3_vectors.py b/tests/vector_store_tests/rag/test_rag_s3_vectors.py new file mode 100644 index 0000000000..cd8a362a7b --- /dev/null +++ b/tests/vector_store_tests/rag/test_rag_s3_vectors.py @@ -0,0 +1,107 @@ +""" +S3 Vectors RAG ingestion tests. + +Requires environment variables: +- AWS_ACCESS_KEY_ID +- AWS_SECRET_ACCESS_KEY +- AWS_REGION_NAME (optional, defaults to us-west-2) + +Optional: +- S3_VECTOR_BUCKET_NAME (optional, auto-generates if not set) +""" + +import os +import sys +from typing import Any, Dict, Optional + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm.types.rag import RAGIngestOptions +from tests.vector_store_tests.rag.base_rag_tests import BaseRAGTest + + +class TestRAGS3Vectors(BaseRAGTest): + """Test RAG Ingest with AWS S3 Vectors.""" + + @pytest.fixture(autouse=True) + def check_env_vars(self): + """Check required environment variables before each test.""" + aws_key = os.environ.get("AWS_ACCESS_KEY_ID") + aws_secret = os.environ.get("AWS_SECRET_ACCESS_KEY") + + if not aws_key or not aws_secret: + pytest.skip("Skipping S3 Vectors test: AWS credentials required") + + def get_base_ingest_options(self) -> RAGIngestOptions: + """ + Return S3 Vectors-specific ingest options. + + Chunking is configured via chunking_strategy (unified interface). + Embeddings are generated using LiteLLM's embedding API. + """ + vector_bucket_name = os.environ.get( + "S3_VECTOR_BUCKET_NAME", "test-litellm-vectors" + ) + aws_region = os.environ.get("AWS_REGION_NAME", "us-west-2") + + return { + "chunking_strategy": { + "chunk_size": 512, + "chunk_overlap": 100, + }, + "embedding": { + "model": "text-embedding-3-small" # Can use any LiteLLM-supported model + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": vector_bucket_name, + "index_name": "test-index", + # dimension is auto-detected from embedding model (text-embedding-3-small = 1536) + "distance_metric": "cosine", + "non_filterable_metadata_keys": ["source_text"], + "aws_region_name": aws_region, + }, + } + + async def query_vector_store( + self, + vector_store_id: str, + query: str, + ) -> Optional[Dict[str, Any]]: + """Query S3 Vectors index.""" + try: + # Import the ingestion class to use its query method + from litellm.rag.ingestion.s3_vectors_ingestion import ( + S3VectorsRAGIngestion, + ) + except ImportError: + pytest.skip("S3 Vectors ingestion not available") + + vector_bucket_name = os.environ.get( + "S3_VECTOR_BUCKET_NAME", "test-litellm-vectors" + ) + aws_region = os.environ.get("AWS_REGION_NAME", "us-west-2") + + # Create ingestion instance to use query method + ingest_options = { + "embedding": {"model": "text-embedding-3-small"}, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": vector_bucket_name, + "aws_region_name": aws_region, + }, + } + + ingestion = S3VectorsRAGIngestion(ingest_options=ingest_options) + + # Query the index + results = await ingestion.query_vector_store( + vector_store_id=vector_store_id, + query=query, + top_k=5, + ) + + return results diff --git a/tests/vector_store_tests/test_s3_vectors_vector_store.py b/tests/vector_store_tests/test_s3_vectors_vector_store.py new file mode 100644 index 0000000000..a7a1568c1c --- /dev/null +++ b/tests/vector_store_tests/test_s3_vectors_vector_store.py @@ -0,0 +1,42 @@ +from base_vector_store_test import BaseVectorStoreTest +import os +import pytest + + +class TestS3VectorsVectorStore(BaseVectorStoreTest): + @pytest.fixture(autouse=True) + def check_env_vars(self): + """Check if required environment variables are set""" + required_vars = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"] + missing_vars = [var for var in required_vars if not os.getenv(var)] + if missing_vars: + pytest.skip(f"Missing required environment variables: {', '.join(missing_vars)}") + + def get_base_request_args(self) -> dict: + """ + Must return the base request args for searching. + For S3 Vectors, vector_store_id should be in format: bucket_name:index_name + """ + return { + "custom_llm_provider": "s3_vectors", + "vector_store_id": os.getenv( + "S3_VECTORS_VECTOR_STORE_ID", "test-litellm-vectors:test-index" + ), + "query": "What is machine learning?", + "aws_region_name": os.getenv("AWS_REGION_NAME", "us-west-2"), + "aws_access_key_id": os.getenv("AWS_ACCESS_KEY_ID"), + "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"), + } + + def get_base_create_vector_store_args(self) -> dict: + """ + Vector store creation is not yet implemented for S3 Vectors. + This test will be skipped. + """ + return {} + + @pytest.mark.parametrize("sync_mode", [True, False]) + @pytest.mark.asyncio + async def test_basic_create_vector_store(self, sync_mode): + """S3 Vectors doesn't support vector store creation via this API yet""" + pytest.skip("Vector store creation not yet implemented for S3 Vectors") diff --git a/ui/litellm-dashboard/public/assets/logos/s3_vector.png b/ui/litellm-dashboard/public/assets/logos/s3_vector.png new file mode 100644 index 0000000000..15a1a456e1 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/s3_vector.png differ diff --git a/ui/litellm-dashboard/src/components/molecules/models/columns.tsx b/ui/litellm-dashboard/src/components/molecules/models/columns.tsx index 43813852f8..1fc08e502a 100644 --- a/ui/litellm-dashboard/src/components/molecules/models/columns.tsx +++ b/ui/litellm-dashboard/src/components/molecules/models/columns.tsx @@ -17,297 +17,300 @@ export const columns = ( expandedRows: Set, setExpandedRows: (expandedRows: Set) => void, ): ColumnDef[] => [ - { - header: () => Model ID, - accessorKey: "model_info.id", - cell: ({ row }) => { - const model = row.original; - return ( - -
setSelectedModelId(model.model_info.id)} - > - {model.model_info.id} -
- - ); - }, - }, - { - header: () => Model Information, - accessorKey: "model_name", - size: 250, // Fixed column width - cell: ({ row }) => { - const model = row.original; - const displayName = getDisplayModelName(row.original) || "-"; - const tooltipContent = ( -
-
- Provider: {model.provider || "-"} -
-
- Public Model Name: {displayName} -
-
- LiteLLM Model Name: {model.litellm_model_name || "-"} -
-
- ); - - return ( - -
- {/* Provider Icon */} -
- {model.provider ? ( - - ) : ( -
-
- )} -
- - {/* Model Names Container */} -
- {/* Public Model Name */} -
{displayName}
- {/* LiteLLM Model Name */} -
- {model.litellm_model_name || "-"} -
-
-
- - ); - }, - }, - { - header: () => Credentials, - accessorKey: "litellm_credential_name", - size: 180, // Fixed column width - cell: ({ row }) => { - const model = row.original; - const credentialName = model.litellm_params?.litellm_credential_name; - - return credentialName ? ( - -
- - - {credentialName} - -
- - ) : ( -
- - No credentials -
- ); - }, - }, - { - header: () => Created By, - accessorKey: "model_info.created_by", - sortingFn: "datetime", - size: 160, // Fixed column width - cell: ({ row }) => { - const model = row.original; - const isConfigModel = !model.model_info?.db_model; - const createdBy = model.model_info.created_by; - const createdAt = model.model_info.created_at ? new Date(model.model_info.created_at).toLocaleDateString() : null; - - return ( -
- {/* Created By - Primary */} -
- {isConfigModel ? "Defined in config" : createdBy || "Unknown"} -
- {/* Created At - Secondary */} -
- {isConfigModel ? "-" : createdAt || "Unknown date"} -
-
- ); - }, - }, - { - header: () => Updated At, - accessorKey: "model_info.updated_at", - sortingFn: "datetime", - cell: ({ row }) => { - const model = row.original; - return ( - - {model.model_info.updated_at ? new Date(model.model_info.updated_at).toLocaleDateString() : "-"} - - ); - }, - }, - { - header: () => Costs, - accessorKey: "input_cost", - size: 120, // Fixed column width - cell: ({ row }) => { - const model = row.original; - const inputCost = model.input_cost; - const outputCost = model.output_cost; - - // If both costs are missing or undefined, show "-" - if (!inputCost && !outputCost) { + { + header: () => Model ID, + accessorKey: "model_info.id", + enableSorting: false, + cell: ({ row }) => { + const model = row.original; return ( -
- - + +
setSelectedModelId(model.model_info.id)} + > + {model.model_info.id} +
+ + ); + }, + }, + { + header: () => Model Information, + accessorKey: "model_name", + size: 250, // Fixed column width + cell: ({ row }) => { + const model = row.original; + const displayName = getDisplayModelName(row.original) || "-"; + const tooltipContent = ( +
+
+ Provider: {model.provider || "-"} +
+
+ Public Model Name: {displayName} +
+
+ LiteLLM Model Name: {model.litellm_model_name || "-"} +
); - } - return ( - -
- {/* Input Cost - Primary */} - {inputCost &&
In: ${inputCost}
} - {/* Output Cost - Secondary */} - {outputCost &&
Out: ${outputCost}
} -
- - ); - }, - }, - { - header: () => Team ID, - accessorKey: "model_info.team_id", - cell: ({ row }) => { - const model = row.original; - return model.model_info.team_id ? ( -
- - + return ( + +
+ {/* Provider Icon */} +
+ {model.provider ? ( + + ) : ( +
-
+ )} +
+ + {/* Model Names Container */} +
+ {/* Public Model Name */} +
{displayName}
+ {/* LiteLLM Model Name */} +
+ {model.litellm_model_name || "-"} +
+
+
-
- ) : ( - "-" - ); + ); + }, }, - }, - { - header: () => Model Access Group, - accessorKey: "model_info.model_access_group", - enableSorting: false, - cell: ({ row }) => { - const model = row.original; - const accessGroups = model.model_info.access_groups; + { + header: () => Credentials, + accessorKey: "litellm_credential_name", + enableSorting: false, + size: 180, // Fixed column width + cell: ({ row }) => { + const model = row.original; + const credentialName = model.litellm_params?.litellm_credential_name; - if (!accessGroups || accessGroups.length === 0) { - return "-"; - } + return credentialName ? ( + +
+ + + {credentialName} + +
+ + ) : ( +
+ + No credentials +
+ ); + }, + }, + { + header: () => Created By, + accessorKey: "model_info.created_by", + sortingFn: "datetime", + size: 160, // Fixed column width + cell: ({ row }) => { + const model = row.original; + const isConfigModel = !model.model_info?.db_model; + const createdBy = model.model_info.created_by; + const createdAt = model.model_info.created_at ? new Date(model.model_info.created_at).toLocaleDateString() : null; - const modelId = model.model_info.id; - const isExpanded = expandedRows.has(modelId); - const shouldShowExpandButton = accessGroups.length > 1; - - const toggleExpanded = () => { - const newExpanded = new Set(expandedRows); - if (isExpanded) { - newExpanded.delete(modelId); - } else { - newExpanded.add(modelId); - } - setExpandedRows(newExpanded); - }; - - return ( -
- - {accessGroups[0]} - - - {(isExpanded || (!shouldShowExpandButton && accessGroups.length === 2)) && - accessGroups.slice(1).map((group: string, index: number) => ( - - {group} - - ))} - - {shouldShowExpandButton && ( - - )} -
- ); + {isConfigModel ? "Defined in config" : createdBy || "Unknown"} +
+ {/* Created At - Secondary */} +
+ {isConfigModel ? "-" : createdAt || "Unknown date"} +
+ + ); + }, }, - }, - { - header: () => Status, - accessorKey: "model_info.db_model", - cell: ({ row }) => { - const model = row.original; - return ( -
Updated At, + accessorKey: "model_info.updated_at", + sortingFn: "datetime", + cell: ({ row }) => { + const model = row.original; + return ( + + {model.model_info.updated_at ? new Date(model.model_info.updated_at).toLocaleDateString() : "-"} + + ); + }, + }, + { + header: () => Costs, + accessorKey: "input_cost", + size: 120, // Fixed column width + cell: ({ row }) => { + const model = row.original; + const inputCost = model.input_cost; + const outputCost = model.output_cost; + + // If both costs are missing or undefined, show "-" + if (!inputCost && !outputCost) { + return ( +
+ - +
+ ); + } + + return ( + +
+ {/* Input Cost - Primary */} + {inputCost &&
In: ${inputCost}
} + {/* Output Cost - Secondary */} + {outputCost &&
Out: ${outputCost}
} +
+ + ); + }, + }, + { + header: () => Team ID, + accessorKey: "model_info.team_id", + enableSorting: false, + cell: ({ row }) => { + const model = row.original; + return model.model_info.team_id ? ( +
+ + + +
+ ) : ( + "-" + ); + }, + }, + { + header: () => Model Access Group, + accessorKey: "model_info.model_access_group", + enableSorting: false, + cell: ({ row }) => { + const model = row.original; + const accessGroups = model.model_info.access_groups; + + if (!accessGroups || accessGroups.length === 0) { + return "-"; + } + + const modelId = model.model_info.id; + const isExpanded = expandedRows.has(modelId); + const shouldShowExpandButton = accessGroups.length > 1; + + const toggleExpanded = () => { + const newExpanded = new Set(expandedRows); + if (isExpanded) { + newExpanded.delete(modelId); + } else { + newExpanded.add(modelId); + } + setExpandedRows(newExpanded); + }; + + return ( +
+ + {accessGroups[0]} + + + {(isExpanded || (!shouldShowExpandButton && accessGroups.length === 2)) && + accessGroups.slice(1).map((group: string, index: number) => ( + + {group} + + ))} + + {shouldShowExpandButton && ( + + )} +
+ ); + }, + }, + { + header: () => Status, + accessorKey: "model_info.db_model", + cell: ({ row }) => { + const model = row.original; + return ( +
- {model.model_info.db_model ? "DB Model" : "Config Model"} -
- ); + > + {model.model_info.db_model ? "DB Model" : "Config Model"} +
+ ); + }, }, - }, - { - id: "actions", - header: () => Actions, - cell: ({ row }) => { - const model = row.original; - const canEditModel = userRole === "Admin" || model.model_info?.created_by === userID; - const isConfigModel = !model.model_info?.db_model; - return ( -
- {isConfigModel ? ( - - - - ) : ( - - { - if (canEditModel) { - setSelectedModelId(model.model_info.id); - } - }} - className={!canEditModel ? "opacity-50 cursor-not-allowed" : "cursor-pointer hover:text-red-600"} - /> - - )} -
- ); + { + id: "actions", + header: () => Actions, + cell: ({ row }) => { + const model = row.original; + const canEditModel = userRole === "Admin" || model.model_info?.created_by === userID; + const isConfigModel = !model.model_info?.db_model; + return ( +
+ {isConfigModel ? ( + + + + ) : ( + + { + if (canEditModel) { + setSelectedModelId(model.model_info.id); + } + }} + className={!canEditModel ? "opacity-50 cursor-not-allowed" : "cursor-pointer hover:text-red-600"} + /> + + )} +
+ ); + }, }, - }, -]; + ]; diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index c78f355ff1..8e4fda0ba5 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -219,7 +219,7 @@ const Navbar: React.FC = ({ style={{ animationDuration: "2s" }} title="Happy Holidays!" > - 🎄 + ❄️ diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index a1625b6ffb..e0dccbe6e8 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6954,7 +6954,8 @@ export const ragIngestCall = async ( customLlmProvider: string, vectorStoreId?: string, vectorStoreName?: string, - vectorStoreDescription?: string + vectorStoreDescription?: string, + providerSpecificParams?: Record ): Promise => { try { let url = proxyBaseUrl ? `${proxyBaseUrl}/rag/ingest` : `/rag/ingest`; @@ -6967,6 +6968,7 @@ export const ragIngestCall = async ( vector_store: { custom_llm_provider: customLlmProvider, ...(vectorStoreId && { vector_store_id: vectorStoreId }), + ...(providerSpecificParams && providerSpecificParams), }, }, }; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx index db2975781b..a8ba55b91f 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.test.tsx @@ -22,17 +22,51 @@ vi.mock("../vector_store_providers", () => ({ BEDROCK: "Amazon Bedrock", OPENAI: "OpenAI", AZURE_OPENAI: "Azure OpenAI", + S3Vectors: "AWS S3 Vectors", }, vectorStoreProviderMap: { BEDROCK: "bedrock", OPENAI: "openai", AZURE_OPENAI: "azure_openai", + S3Vectors: "s3_vectors", }, vectorStoreProviderLogoMap: { "Amazon Bedrock": "https://example.com/bedrock.png", "OpenAI": "https://example.com/openai.png", "Azure OpenAI": "https://example.com/azure.png", + "AWS S3 Vectors": "https://example.com/aws.png", }, + getProviderSpecificFields: vi.fn((provider: string) => { + if (provider === "s3_vectors") { + return [ + { + name: "vector_bucket_name", + label: "Vector Bucket Name", + tooltip: "S3 bucket name for vector storage", + placeholder: "my-vector-bucket", + required: true, + type: "text", + }, + { + name: "aws_region_name", + label: "AWS Region", + tooltip: "AWS region", + placeholder: "us-west-2", + required: true, + type: "text", + }, + { + name: "embedding_model", + label: "Embedding Model", + tooltip: "Embedding model to use", + placeholder: "text-embedding-3-small", + required: true, + type: "select", + }, + ]; + } + return []; + }), })); describe("CreateVectorStore", () => { @@ -43,9 +77,9 @@ describe("CreateVectorStore", () => { it("should render the component successfully", () => { render(); - expect(screen.getByText("Create Vector Store")).toBeInTheDocument(); + expect(screen.getAllByText("Create Vector Store").length).toBeGreaterThan(0); expect(screen.getByText("Step 1: Upload Documents")).toBeInTheDocument(); - expect(screen.getByText("Step 2: Select Provider")).toBeInTheDocument(); + expect(screen.getByText("Step 2: Configure Vector Store")).toBeInTheDocument(); }); it("should display upload area with correct text", () => { @@ -123,7 +157,15 @@ describe("CreateVectorStore", () => { }); await waitFor(() => { - expect(mockRagIngestCall).toHaveBeenCalledWith("test-token", expect.any(File), "bedrock", undefined); + expect(mockRagIngestCall).toHaveBeenCalledWith( + "test-token", + expect.any(File), + "bedrock", + undefined, + undefined, + undefined, + {} + ); }); }); @@ -163,4 +205,72 @@ describe("CreateVectorStore", () => { expect(screen.getByText("Vector Store Created Successfully")).toBeInTheDocument(); }); }); + + it("should display S3 Vectors provider-specific fields when selected", async () => { + render(); + + // Find and click the provider dropdown + const providerSelect = screen.getByRole("combobox"); + + await act(async () => { + fireEvent.mouseDown(providerSelect); + }); + + // Wait for dropdown options to appear + await waitFor(() => { + const s3Option = screen.queryByText("AWS S3 Vectors"); + if (s3Option) { + fireEvent.click(s3Option); + } + }); + + // Check if S3-specific fields are displayed + await waitFor(() => { + expect(screen.queryByText("Vector Bucket Name")).toBeInTheDocument(); + expect(screen.queryByText("AWS Region")).toBeInTheDocument(); + expect(screen.queryByText("Embedding Model")).toBeInTheDocument(); + }); + }); + + it("should validate S3 Vectors required fields before submission", async () => { + render(); + + // Upload a file first + const file = new File(["test content"], "test.pdf", { type: "application/pdf" }); + const uploadInput = document.querySelector('input[type="file"]') as HTMLInputElement; + + await act(async () => { + if (uploadInput) { + fireEvent.change(uploadInput, { target: { files: [file] } }); + } + }); + + await waitFor(() => { + expect(screen.getByText("Uploaded Documents (1)")).toBeInTheDocument(); + }); + + // Select S3 Vectors provider + const providerSelect = screen.getByRole("combobox"); + + await act(async () => { + fireEvent.mouseDown(providerSelect); + }); + + await waitFor(() => { + const s3Option = screen.queryByText("AWS S3 Vectors"); + if (s3Option) { + fireEvent.click(s3Option); + } + }); + + // Try to create without filling required fields + const createButton = screen.getByRole("button", { name: /Create Vector Store/i }); + + await act(async () => { + fireEvent.click(createButton); + }); + + // Should show validation warning (mocked message.warning would be called) + // The actual validation happens in the component + }); }); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx index 685e37ac73..97162aa132 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx @@ -10,8 +10,11 @@ import { VectorStoreProviders, vectorStoreProviderLogoMap, vectorStoreProviderMap, + getProviderSpecificFields, + VectorStoreFieldConfig, } from "../vector_store_providers"; import NotificationsManager from "../molecules/notifications_manager"; +import S3VectorsConfig from "./S3VectorsConfig"; const { Dragger } = Upload; @@ -28,6 +31,7 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu const [vectorStoreName, setVectorStoreName] = useState(""); const [vectorStoreDescription, setVectorStoreDescription] = useState(""); const [ingestResults, setIngestResults] = useState([]); + const [providerParams, setProviderParams] = useState>({}); const uploadProps: UploadProps = { name: "file", @@ -92,6 +96,27 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu return; } + // Validate provider-specific required fields + const requiredFields = getProviderSpecificFields(selectedProvider).filter((field) => field.required); + for (const field of requiredFields) { + if (!providerParams[field.name]) { + message.warning(`Please provide ${field.label}`); + return; + } + } + + // S3 Vectors specific validation + if (selectedProvider === "s3_vectors") { + if (providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3) { + message.warning("Vector bucket name must be at least 3 characters"); + return; + } + if (providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3) { + message.warning("Index name must be at least 3 characters if provided"); + return; + } + } + if (!accessToken) { message.error("No access token available"); return; @@ -118,7 +143,8 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu selectedProvider, vectorStoreId, // Use the same vector store ID for subsequent uploads vectorStoreName || undefined, - vectorStoreDescription || undefined + vectorStoreDescription || undefined, + providerParams ); // Store the vector store ID from the first successful ingest @@ -298,6 +324,74 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu })} + + {/* S3 Vectors Configuration */} + {selectedProvider === "s3_vectors" && ( + + )} + + {/* Other Provider-specific fields */} + {selectedProvider !== "s3_vectors" && + getProviderSpecificFields(selectedProvider).map((field: VectorStoreFieldConfig) => { + if (field.type === "select") { + // For embedding model selection, we'd need to fetch available models + // For now, provide a text input as fallback + return ( + + {field.label}{" "} + + + + + } + required={field.required} + > + + setProviderParams((prev) => ({ ...prev, [field.name]: e.target.value })) + } + placeholder={field.placeholder} + size="large" + className="rounded-md" + /> + + ); + } + + return ( + + {field.label}{" "} + + + + + } + required={field.required} + > + + setProviderParams((prev) => ({ ...prev, [field.name]: e.target.value })) + } + placeholder={field.placeholder} + size="large" + className="rounded-md" + /> + + ); + })}
diff --git a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx index 761bc2b7df..507ad0537d 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.test.tsx @@ -62,7 +62,7 @@ describe("DocumentsTable", () => { render(); expect(screen.getByText(/1000.00 KB/)).toBeInTheDocument(); - expect(screen.getByText(/2.00 MB/)).toBeInTheDocument(); + expect(screen.getByText(/1.95 MB/)).toBeInTheDocument(); expect(screen.getByText(/500.00 KB/)).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx new file mode 100644 index 0000000000..28f885a9dd --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx @@ -0,0 +1,203 @@ +import { render, screen, fireEvent, waitFor, act } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import S3VectorsConfig from "./S3VectorsConfig"; +import * as fetchModels from "../playground/llm_calls/fetch_models"; + +// Mock fetchAvailableModels +vi.mock("../playground/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn(), +})); + +describe("S3VectorsConfig", () => { + const mockOnParamsChange = vi.fn(); + const defaultProps = { + accessToken: "test-token", + providerParams: {}, + onParamsChange: mockOnParamsChange, + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render the component successfully", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + expect(screen.getByText("AWS S3 Vectors Setup")).toBeInTheDocument(); + expect(screen.getByText("Vector Bucket Name")).toBeInTheDocument(); + expect(screen.getByText("Index Name")).toBeInTheDocument(); + expect(screen.getByText("AWS Region")).toBeInTheDocument(); + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should display setup instructions", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + expect( + screen.getByText(/AWS S3 Vectors allows you to store and query vector embeddings directly in S3/) + ).toBeInTheDocument(); + expect(screen.getByText(/Vector buckets and indexes will be automatically created/)).toBeInTheDocument(); + expect(screen.getByText(/Vector dimensions are auto-detected/)).toBeInTheDocument(); + }); + + it("should fetch embedding models on mount", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + { model_group: "gpt-4", mode: "chat" }, + ]; + + const fetchSpy = vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + await waitFor(() => { + expect(fetchSpy).toHaveBeenCalledWith("test-token"); + }); + }); + + it("should filter and display only embedding models", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + { model_group: "gpt-4", mode: "chat" }, + { model_group: "gpt-3.5-turbo", mode: "chat" }, + ]; + + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + // Wait for models to load + await waitFor(() => { + expect(fetchModels.fetchAvailableModels).toHaveBeenCalled(); + }); + + // The component should filter to only embedding models internally + // We can verify this by checking the component loaded successfully + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should call onParamsChange when vector bucket name changes", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + const bucketInput = screen.getByPlaceholderText("my-vector-bucket (min 3 chars)"); + + await act(async () => { + fireEvent.change(bucketInput, { target: { value: "test-bucket" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + vector_bucket_name: "test-bucket", + }); + }); + + it("should call onParamsChange when AWS region changes", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + render(); + + const regionInput = screen.getByPlaceholderText("us-west-2"); + + await act(async () => { + fireEvent.change(regionInput, { target: { value: "us-east-1" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + aws_region_name: "us-east-1", + }); + }); + + it("should call onParamsChange when embedding model is selected", async () => { + const mockModels = [ + { model_group: "text-embedding-3-small", mode: "embedding" }, + { model_group: "text-embedding-3-large", mode: "embedding" }, + ]; + + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue(mockModels); + + render(); + + await waitFor(() => { + expect(fetchModels.fetchAvailableModels).toHaveBeenCalled(); + }); + + // Find the Select component and trigger change directly + const selectElement = screen.getByRole("combobox"); + + await act(async () => { + // Simulate selecting a value by firing the change event + fireEvent.change(selectElement, { target: { value: "text-embedding-3-small" } }); + }); + + // The component should handle the selection + expect(screen.getByText("Embedding Model")).toBeInTheDocument(); + }); + + it("should preserve existing params when updating a field", async () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + const existingParams = { + vector_bucket_name: "existing-bucket", + aws_region_name: "us-west-2", + }; + + render(); + + const indexInput = screen.getByPlaceholderText("my-vector-index (optional, min 3 chars)"); + + await act(async () => { + fireEvent.change(indexInput, { target: { value: "my-index" } }); + }); + + expect(mockOnParamsChange).toHaveBeenCalledWith({ + vector_bucket_name: "existing-bucket", + aws_region_name: "us-west-2", + index_name: "my-index", + }); + }); + + it("should display existing param values", () => { + vi.spyOn(fetchModels, "fetchAvailableModels").mockResolvedValue([]); + + const existingParams = { + vector_bucket_name: "my-bucket", + index_name: "my-index", + aws_region_name: "eu-west-1", + embedding_model: "text-embedding-3-small", + }; + + render(); + + expect(screen.getByDisplayValue("my-bucket")).toBeInTheDocument(); + expect(screen.getByDisplayValue("my-index")).toBeInTheDocument(); + expect(screen.getByDisplayValue("eu-west-1")).toBeInTheDocument(); + }); + + it("should handle model fetch error gracefully", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + vi.spyOn(fetchModels, "fetchAvailableModels").mockRejectedValue(new Error("Failed to fetch models")); + + render(); + + await waitFor(() => { + expect(consoleErrorSpy).toHaveBeenCalledWith("Error fetching embedding models:", expect.any(Error)); + }); + + consoleErrorSpy.mockRestore(); + }); + + it("should not fetch models if accessToken is null", () => { + const fetchSpy = vi.spyOn(fetchModels, "fetchAvailableModels"); + + render(); + + expect(fetchSpy).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx new file mode 100644 index 0000000000..e54d3a163c --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx @@ -0,0 +1,192 @@ +import React, { useState, useEffect } from "react"; +import { Alert, Form, Input, Select, Tooltip } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; + +interface S3VectorsConfigProps { + accessToken: string | null; + providerParams: Record; + onParamsChange: (params: Record) => void; +} + +const S3VectorsConfig: React.FC = ({ + accessToken, + providerParams, + onParamsChange, +}) => { + const [embeddingModels, setEmbeddingModels] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + + useEffect(() => { + if (!accessToken) return; + + const loadModels = async () => { + setIsLoadingModels(true); + try { + const models = await fetchAvailableModels(accessToken); + // Filter for embedding models only + const embeddingOnly = models.filter((model) => model.mode === "embedding"); + setEmbeddingModels(embeddingOnly); + } catch (error) { + console.error("Error fetching embedding models:", error); + } finally { + setIsLoadingModels(false); + } + }; + + loadModels(); + }, [accessToken]); + + const handleFieldChange = (fieldName: string, value: string) => { + onParamsChange({ + ...providerParams, + [fieldName]: value, + }); + }; + + return ( + <> + {/* S3 Vectors Setup Instructions */} + +

AWS S3 Vectors allows you to store and query vector embeddings directly in S3:

+
    +
  • Vector buckets and indexes will be automatically created if they don't exist
    • +
  • Vector dimensions are auto-detected from your selected embedding model
    • +
  • Ensure your AWS credentials have permissions for S3 Vectors operations
    • +
  • + Learn more:{" "} + + AWS S3 Vectors Documentation + +
    • +
+
+ } + type="info" + showIcon + style={{ marginBottom: "16px" }} + /> + + {/* Vector Bucket Name */} + + Vector Bucket Name{" "} + + + + + } + required + validateStatus={ + providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3 + ? "error" + : undefined + } + help={ + providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3 + ? "Bucket name must be at least 3 characters" + : undefined + } + > + handleFieldChange("vector_bucket_name", e.target.value)} + placeholder="my-vector-bucket (min 3 chars)" + size="large" + className="rounded-md" + /> + + + {/* Index Name (Optional) */} + + Index Name{" "} + + + + + } + validateStatus={ + providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3 + ? "error" + : undefined + } + help={ + providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3 + ? "Index name must be at least 3 characters if provided" + : undefined + } + > + handleFieldChange("index_name", e.target.value)} + placeholder="my-vector-index (optional, min 3 chars)" + size="large" + className="rounded-md" + /> + + + {/* AWS Region */} + + AWS Region{" "} + + + + + } + required + > + handleFieldChange("aws_region_name", e.target.value)} + placeholder="us-west-2" + size="large" + className="rounded-md" + /> + + + {/* Embedding Model */} + + Embedding Model{" "} + + + + + } + required + > +