diff --git a/.github/workflows/publish-migrations.yml b/.github/workflows/publish-migrations.yml index 8e5a67bcf85..a5187cb2f55 100644 --- a/.github/workflows/publish-migrations.yml +++ b/.github/workflows/publish-migrations.yml @@ -13,6 +13,7 @@ on: jobs: publish-migrations: + if: github.repository == 'BerriAI/litellm' runs-on: ubuntu-latest services: postgres: diff --git a/README.md b/README.md index a020bd80898..75a23faa5c1 100644 --- a/README.md +++ b/README.md @@ -262,6 +262,7 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature | Provider | `/chat/completions` | `/messages` | `/responses` | `/embeddings` | `/image/generations` | `/audio/transcriptions` | `/audio/speech` | `/moderations` | `/batches` | `/rerank` | |-------------------------------------------------------------------------------------|---------------------|-------------|--------------|---------------|----------------------|-------------------------|-----------------|----------------|-----------|-----------| +| [Abliteration (`abliteration`)](https://docs.litellm.ai/docs/providers/abliteration) | ✅ | | | | | | | | | | | [AI/ML API (`aiml`)](https://docs.litellm.ai/docs/providers/aiml) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | [AI21 (`ai21`)](https://docs.litellm.ai/docs/providers/ai21) | ✅ | ✅ | ✅ | | | | | | | | | [AI21 Chat (`ai21_chat`)](https://docs.litellm.ai/docs/providers/ai21) | ✅ | ✅ | ✅ | | | | | | | | @@ -455,4 +456,3 @@ All these checks must pass before your PR can be merged. - diff --git a/docs/my-website/docs/mcp_control.md b/docs/my-website/docs/mcp_control.md index a7d66a6b7fc..96c71ef9278 100644 --- a/docs/my-website/docs/mcp_control.md +++ b/docs/my-website/docs/mcp_control.md @@ -649,3 +649,16 @@ general_settings: ``` This is useful when you want discoverability for MCP offerings without granting additional execution privileges. + + +## Publish MCP Registry + +If you want other systems—for example external agent frameworks such as MCP-capable IDEs running outside your network—to automatically discover the MCP servers hosted on LiteLLM, you can expose a Model Context Protocol Registry endpoint. This registry lists the built-in LiteLLM MCP server and every server you have configured, using the [official MCP Registry spec](https://github.com/modelcontextprotocol/registry). + +1. Set `enable_mcp_registry: true` under `general_settings` in your proxy config (or DB settings) and restart the proxy. +2. LiteLLM will serve the registry at `GET /v1/mcp/registry.json`. +3. Each entry points to either `/mcp` (built-in server) or `/{mcp_server_name}/mcp` for your custom servers, so clients can connect directly using the advertised Streamable HTTP URL. + +:::note Permissions still apply +The registry only advertises server URLs. Actual access control is still enforced by LiteLLM when the client connects to `/mcp` or `/{server}/mcp`, so publishing the registry does not bypass per-key permissions. +::: diff --git a/docs/my-website/docs/observability/focus.md b/docs/my-website/docs/observability/focus.md new file mode 100644 index 00000000000..c282f4a220c --- /dev/null +++ b/docs/my-website/docs/observability/focus.md @@ -0,0 +1,93 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Focus Export (Experimental) + +:::caution Experimental feature +Focus Format export is under active development and currently considered experimental. +Interfaces, schema mappings, and configuration options may change as we iterate based on user feedback. +Please treat this integration as a preview and report any issues or suggestions to help us stabilize and improve the workflow. +::: + +LiteLLM can emit usage data in the [FinOps FOCUS format](https://focus.finops.org/focus-specification/v1-2/) and push artifacts (for example Parquet files) to destinations such as Amazon S3. This enables downstream cost-analysis tooling to ingest a standardised dataset directly from LiteLLM. + +LiteLLM currently conforms to the FinOps FOCUS v1.2 specification when emitting this dataset. + +## Overview + +| Property | Details | +|----------|---------| +| Destination | Export LiteLLM usage data in FOCUS format to managed storage (currently S3) | +| Callback name | `focus` | +| Supported operations | Automatic scheduled export | +| Data format | FOCUS Normalised Dataset (Parquet) | + +## Environment Variables + +### Common settings + +| Variable | Required | Description | +|----------|----------|-------------| +| `FOCUS_PROVIDER` | No | Destination provider (defaults to `s3`). | +| `FOCUS_FORMAT` | No | Output format (currently only `parquet`). | +| `FOCUS_FREQUENCY` | No | Export cadence. Prefer `hourly` or `daily` for production; `interval` is intended for short test loops. Defaults to `hourly`. | +| `FOCUS_CRON_OFFSET` | No | Minute offset used for hourly/daily cron triggers. Defaults to `5`. | +| `FOCUS_INTERVAL_SECONDS` | No | Interval (seconds) when `FOCUS_FREQUENCY="interval"`. | +| `FOCUS_PREFIX` | No | Object key prefix/folder. Defaults to `focus_exports`. | + +### S3 destination + +| Variable | Required | Description | +|----------|----------|-------------| +| `FOCUS_S3_BUCKET_NAME` | Yes | Destination bucket for exported files. | +| `FOCUS_S3_REGION_NAME` | No | AWS region for the bucket. | +| `FOCUS_S3_ENDPOINT_URL` | No | Custom endpoint (useful for S3-compatible storage). | +| `FOCUS_S3_ACCESS_KEY` | Yes | AWS access key for uploads. | +| `FOCUS_S3_SECRET_KEY` | Yes | AWS secret key for uploads. | +| `FOCUS_S3_SESSION_TOKEN` | No | AWS session token if using temporary credentials. | + +## Setup via Config + +### Configure environment variables + +```bash +export FOCUS_PROVIDER="s3" +export FOCUS_PREFIX="focus_exports" + +# S3 example +export FOCUS_S3_BUCKET_NAME="my-litellm-focus-bucket" +export FOCUS_S3_REGION_NAME="us-east-1" +export FOCUS_S3_ACCESS_KEY="AKIA..." +export FOCUS_S3_SECRET_KEY="..." +``` + +### Update LiteLLM config + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: sk-your-key + +litellm_settings: + callbacks: ["focus"] +``` + +### Start the proxy + +```bash +litellm --config /path/to/config.yaml +``` + +During boot LiteLLM registers the Focus logger and a background job that runs according to the configured frequency. + +## Planned Enhancements +- Add "Setup on UI" flow alongside the current configuration-based setup. +- Add GCS / Azure Blob to the Destination options. +- Support CSV output alongside Parquet. + +## Related Links + +- [Focus](https://focus.finops.org/) + diff --git a/docs/my-website/docs/observability/qualifire_integration.md b/docs/my-website/docs/observability/qualifire_integration.md new file mode 100644 index 00000000000..cf866f467bf --- /dev/null +++ b/docs/my-website/docs/observability/qualifire_integration.md @@ -0,0 +1,122 @@ +import Image from '@theme/IdealImage'; + +# Qualifire - LLM Evaluation, Guardrails & Observability + +[Qualifire](https://qualifire.ai/) provides real-time Agentic evaluations, guardrails and observability for production AI applications. + +**Key Features:** + +- **Evaluation** - Systematically assess AI behavior to detect hallucinations, jailbreaks, policy breaches, and other vulnerabilities +- **Guardrails** - Real-time interventions to prevent risks like brand damage, data leaks, and compliance breaches +- **Observability** - Complete tracing and logging for RAG pipelines, chatbots, and AI agents +- **Prompt Management** - Centralized prompt management with versioning and no-code studio + +:::tip + +Looking for Qualifire Guardrails? Check out the [Qualifire Guardrails Integration](../proxy/guardrails/qualifire.md) for real-time content moderation, prompt injection detection, PII checks, and more. + +::: + +## Pre-Requisites + +1. Create an account on [Qualifire](https://app.qualifire.ai/) +2. Get your API key and webhook URL from the Qualifire dashboard + +```bash +pip install litellm +``` + +## Quick Start + +Use just 2 lines of code to instantly log your responses **across all providers** with Qualifire. + +```python +litellm.callbacks = ["qualifire_eval"] +``` + +```python +import litellm +import os + +# Set Qualifire credentials +os.environ["QUALIFIRE_API_KEY"] = "your-qualifire-api-key" +os.environ["QUALIFIRE_WEBHOOK_URL"] = "https://your-qualifire-webhook-url" + +# LLM API Keys +os.environ['OPENAI_API_KEY'] = "your-openai-api-key" + +# Set qualifire_eval as a callback & LiteLLM will send the data to Qualifire +litellm.callbacks = ["qualifire_eval"] + +# OpenAI call +response = litellm.completion( + model="gpt-5", + messages=[ + {"role": "user", "content": "Hi 👋 - i'm openai"} + ] +) +``` + +## Using with LiteLLM Proxy + +1. Setup config.yaml + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + +litellm_settings: + callbacks: ["qualifire_eval"] + +general_settings: + master_key: "sk-1234" + +environment_variables: + QUALIFIRE_API_KEY: "your-qualifire-api-key" + QUALIFIRE_WEBHOOK_URL: "https://app.qualifire.ai/api/v1/webhooks/evaluations" +``` + +2. Start the proxy + +```bash +litellm --config config.yaml +``` + +3. Test it! + +```bash +curl -X POST 'http://0.0.0.0:4000/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-d '{ "model": "gpt-4o", "messages": [{"role": "user", "content": "Hi 👋 - i'm openai"}]}' +``` + +## Environment Variables + +| Variable | Description | +| ----------------------- | ------------------------------------------------------ | +| `QUALIFIRE_API_KEY` | Your Qualifire API key for authentication | +| `QUALIFIRE_WEBHOOK_URL` | The Qualifire webhook endpoint URL from your dashboard | + +## What Gets Logged? + +The [LiteLLM Standard Logging Payload](https://docs.litellm.ai/docs/proxy/logging_spec) is sent to your Qualifire endpoint on each successful LLM API call. + +This includes: + +- Request messages and parameters +- Response content and metadata +- Token usage statistics +- Latency metrics +- Model information +- Cost data + +Once data is in Qualifire, you can: + +- Run evaluations to detect hallucinations, toxicity, and policy violations +- Set up guardrails to block or modify responses in real-time +- View traces across your entire AI pipeline +- Track performance and quality metrics over time diff --git a/docs/my-website/docs/providers/abliteration.md b/docs/my-website/docs/providers/abliteration.md new file mode 100644 index 00000000000..a0fc7f39310 --- /dev/null +++ b/docs/my-website/docs/providers/abliteration.md @@ -0,0 +1,109 @@ +# Abliteration + +## Overview + +| Property | Details | +|-------|-------| +| Description | Abliteration provides an OpenAI-compatible `/chat/completions` endpoint. | +| Provider Route on LiteLLM | `abliteration/` | +| Link to Provider Doc | [Abliteration](https://abliteration.ai) | +| Base URL | `https://api.abliteration.ai/v1` | +| Supported Operations | [`/chat/completions`](#sample-usage) | + +
+ +## Required Variables + +```python showLineNumbers title="Environment Variables" +os.environ["ABLITERATION_API_KEY"] = "" # your Abliteration API key +``` + +## Sample Usage + +```python showLineNumbers title="Abliteration Completion" +import os +from litellm import completion + +os.environ["ABLITERATION_API_KEY"] = "" + +response = completion( + model="abliteration/abliterated-model", + messages=[{"role": "user", "content": "Hello from LiteLLM"}], +) + +print(response) +``` + +## Sample Usage - Streaming + +```python showLineNumbers title="Abliteration Streaming Completion" +import os +from litellm import completion + +os.environ["ABLITERATION_API_KEY"] = "" + +response = completion( + model="abliteration/abliterated-model", + messages=[{"role": "user", "content": "Stream a short reply"}], + stream=True, +) + +for chunk in response: + print(chunk) +``` + +## Usage with LiteLLM Proxy Server + +1. Add the model to your proxy config: + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: abliteration-chat + litellm_params: + model: abliteration/abliterated-model + api_key: os.environ/ABLITERATION_API_KEY +``` + +2. Start the proxy: + +```bash +litellm --config /path/to/config.yaml +``` + +## Direct API Usage (Bearer Token) + +Use the environment variable as a Bearer token against the OpenAI-compatible endpoint: +`https://api.abliteration.ai/v1/chat/completions`. + +```bash showLineNumbers title="cURL" +export ABLITERATION_API_KEY="" +curl https://api.abliteration.ai/v1/chat/completions \ + -H "Authorization: Bearer ${ABLITERATION_API_KEY}" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "abliterated-model", + "messages": [{"role": "user", "content": "Hello from Abliteration"}] + }' +``` + +```python showLineNumbers title="Python (requests)" +import os +import requests + +api_key = os.environ["ABLITERATION_API_KEY"] + +response = requests.post( + "https://api.abliteration.ai/v1/chat/completions", + headers={ + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + }, + json={ + "model": "abliterated-model", + "messages": [{"role": "user", "content": "Hello from Abliteration"}], + }, + timeout=60, +) + +print(response.json()) +``` diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 33ebf535d29..f46608aa57c 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -35,6 +35,8 @@ import json # !gcloud auth application-default login - run this to add vertex credentials to your env ## OR ## file_path = 'path/to/vertex_ai_service_account.json' +## OR ## +export VERTEXAI_API_KEY="your-api-key" # Load the JSON file with open(file_path, 'r') as file: @@ -47,7 +49,7 @@ vertex_credentials_json = json.dumps(vertex_credentials) response = completion( model="vertex_ai/gemini-2.5-pro", messages=[{ "content": "Hello, how are you?","role": "user"}], - vertex_credentials=vertex_credentials_json + vertex_credentials=vertex_credentials_json # Can remove this is added VERTEXAI_API_KEY in env ) ``` @@ -1329,15 +1331,41 @@ Here's how to use Vertex AI with the LiteLLM Proxy Server ## Authentication - vertex_project, vertex_location, etc. +LiteLLM supports two authentication methods for Vertex AI: + +1. **API Key Authentication** (Recommended for getting started) +2. **Service Account Credentials** (Recommended for production) + Set your vertex credentials via: - dynamic params OR - env vars +### **Authentication Method 1: -### **Dynamic Params** +The simplest way to authenticate with Vertex AI. You can set: +- `api_key` (str) - Your Vertex AI API key -You can set: +**Environment Variables:** +```bash +export VERTEXAI_API_KEY="your-api-key" +``` + +**Or pass as parameters:** +```python +from litellm import completion + +response = completion( + model="vertex_ai/gemini-2.0-flash-exp", + messages=[{"role": "user", "content": "Hello!"}], + api_key="your-vertex-api-key", + +) +``` + +### **Authentication Method 2: Service Account Credentials** + +For production environments with fine-grained access control. You can set: - `vertex_credentials` (str) - can be a json string or filepath to your vertex ai service account.json - `vertex_location` (str) - place where vertex model is deployed (us-central1, asia-southeast1, etc.). Some models support the global location, please see [Vertex AI documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#supported_models) - `vertex_project` Optional[str] - use if vertex project different from the one in vertex_credentials @@ -1392,7 +1420,16 @@ model_list: ### **Environment Variables** -You can set: +#### For API Key Authentication: + +- `VERTEXAI_API_KEY` or `VERTEX_API_KEY` - Your Vertex AI API key + +```bash +export VERTEXAI_API_KEY="your-vertex-api-key" +``` + +#### For Service Account Authentication: + - `GOOGLE_APPLICATION_CREDENTIALS` - store the filepath for your service_account.json in here (used by vertex sdk directly). - VERTEXAI_LOCATION - place where vertex model is deployed (us-central1, asia-southeast1, etc.) - VERTEXAI_PROJECT - Optional[str] - use if vertex project different from the one in vertex_credentials diff --git a/docs/my-website/docs/proxy/guardrails/qualifire.md b/docs/my-website/docs/proxy/guardrails/qualifire.md index 66961c92d9d..850af37e47f 100644 --- a/docs/my-website/docs/proxy/guardrails/qualifire.md +++ b/docs/my-website/docs/proxy/guardrails/qualifire.md @@ -8,13 +8,7 @@ Use [Qualifire](https://qualifire.ai) to evaluate LLM outputs for quality, safet ## Quick Start -### 1. Install the Qualifire SDK - -```bash -pip install qualifire -``` - -### 2. Define Guardrails on your LiteLLM config.yaml +### 1. Define Guardrails on your LiteLLM config.yaml Define your guardrails under the `guardrails` section: @@ -61,13 +55,13 @@ guardrails: - `post_call` Run **after** LLM call, on **input & output** - `during_call` Run **during** LLM call, on **input**. Same as `pre_call` but runs in parallel as LLM call. Response not returned until guardrail check completes -### 3. Start LiteLLM Gateway +### 2. Start LiteLLM Gateway ```shell litellm --config config.yaml --detailed_debug ``` -### 4. Test request +### 3. Test request **[Langchain, OpenAI SDK Usage Examples](../proxy/user_keys#request-format)** @@ -142,7 +136,7 @@ guardrails: evaluation_id: eval_abc123 # Your evaluation ID from Qualifire dashboard ``` -When `evaluation_id` is provided, LiteLLM will use `invoke_evaluation()` instead of `evaluate()`, running the pre-configured evaluation from your dashboard. +When `evaluation_id` is provided, LiteLLM will use the invoke evaluation API endpoint instead of the evaluate endpoint, running the pre-configured evaluation from your dashboard. ## Available Checks @@ -213,19 +207,19 @@ guardrails: ### Parameter Reference -| Parameter | Type | Default | Description | -| ------------------------------ | ----------- | --------------------------- | -------------------------------------------------------- | -| `api_key` | `str` | `QUALIFIRE_API_KEY` env var | Your Qualifire API key | -| `api_base` | `str` | `None` | Custom API base URL (optional) | -| `evaluation_id` | `str` | `None` | Pre-configured evaluation ID from Qualifire dashboard | -| `prompt_injections` | `bool` | `true` (if no other checks) | Enable prompt injection detection | -| `hallucinations_check` | `bool` | `None` | Enable hallucination detection | -| `grounding_check` | `bool` | `None` | Enable grounding verification | -| `pii_check` | `bool` | `None` | Enable PII detection | -| `content_moderation_check` | `bool` | `None` | Enable content moderation | -| `tool_selection_quality_check` | `bool` | `None` | Enable tool selection quality check | -| `assertions` | `List[str]` | `None` | Custom assertions to validate | -| `on_flagged` | `str` | `"block"` | Action when content is flagged: `"block"` or `"monitor"` | +| Parameter | Type | Default | Description | +| ------------------------------ | ----------- | ---------------------------- | -------------------------------------------------------- | +| `api_key` | `str` | `QUALIFIRE_API_KEY` env var | Your Qualifire API key | +| `api_base` | `str` | `https://proxy.qualifire.ai` | Custom API base URL (optional) | +| `evaluation_id` | `str` | `None` | Pre-configured evaluation ID from Qualifire dashboard | +| `prompt_injections` | `bool` | `true` (if no other checks) | Enable prompt injection detection | +| `hallucinations_check` | `bool` | `None` | Enable hallucination detection | +| `grounding_check` | `bool` | `None` | Enable grounding verification | +| `pii_check` | `bool` | `None` | Enable PII detection | +| `content_moderation_check` | `bool` | `None` | Enable content moderation | +| `tool_selection_quality_check` | `bool` | `None` | Enable tool selection quality check | +| `assertions` | `List[str]` | `None` | Custom assertions to validate | +| `on_flagged` | `str` | `"block"` | Action when content is flagged: `"block"` or `"monitor"` | ### Default Behavior @@ -261,4 +255,3 @@ This evaluates whether the LLM selected the appropriate tools and provided corre - [Qualifire Documentation](https://docs.qualifire.ai) - [Qualifire Dashboard](https://app.qualifire.ai) -- [Qualifire Python SDK](https://github.com/qualifire-dev/qualifire-python-sdk) diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index 5fe8f17d7b0..a27b6dcf083 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -67,7 +67,7 @@ Set `litellm.turn_off_message_logging=True` This will prevent the messages and r -**1. Setup config.yaml ** +**1. Setup config.yaml** ```yaml model_list: - model_name: gpt-3.5-turbo diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 488fd616678..2e8ea07ab75 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -55,6 +55,7 @@ const sidebars = { "proxy/guardrails/test_playground", "proxy/guardrails/litellm_content_filter", ...[ + "proxy/guardrails/qualifire", "proxy/guardrails/aim_security", "proxy/guardrails/onyx_security", "proxy/guardrails/aporia_api", @@ -653,12 +654,13 @@ const sidebars = { "providers/bedrock_writer", "providers/bedrock_batches", "providers/aws_polly", - "providers/bedrock_vector_store", - ] - }, - "providers/litellm_proxy", - "providers/ai21", - "providers/aiml", + "providers/bedrock_vector_store", + ] + }, + "providers/litellm_proxy", + "providers/abliteration", + "providers/ai21", + "providers/aiml", "providers/aleph_alpha", "providers/amazon_nova", "providers/anyscale", diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index b7523ef8c16..1dc805a6b54 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -130,6 +130,9 @@ class GenerateContentHelper: api_key=litellm_params.api_key, ) + if litellm_params.custom_llm_provider is None: + litellm_params.custom_llm_provider = custom_llm_provider + # get provider config generate_content_provider_config: Optional[ BaseGoogleGenAIGenerateContentConfig @@ -407,6 +410,9 @@ async def agenerate_content_stream( # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: + if "stream" in kwargs: + kwargs.pop("stream", None) + # Use the adapter to convert to completion format return ( await GenerateContentToCompletionHandler.async_generate_content_handler( @@ -490,6 +496,9 @@ def generate_content_stream( # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: + if "stream" in kwargs: + kwargs.pop("stream", None) + # Use the adapter to convert to completion format return GenerateContentToCompletionHandler.generate_content_handler( model=model, diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 364fa3f5def..585de510e8b 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -225,10 +225,13 @@ class BraintrustLogger(CustomLogger): "id": litellm_call_id, "input": prompt["messages"], "metadata": standard_logging_object, - "tags": tags, "span_attributes": {"name": span_name, "type": "llm"}, } - + + # Braintrust cannot specify 'tags' for non-root spans + if dynamic_metadata.get("root_span_id") is None: + request_data["tags"] = tags + # Only add those that are not None (or falsy) for key, value in span_attributes.items(): if value: @@ -351,14 +354,37 @@ class BraintrustLogger(CustomLogger): # Allow metadata override for span name span_name = dynamic_metadata.get("span_name", "Chat Completion") + # Span parents is a special case + span_parents = dynamic_metadata.get("span_parents") + + # Convert comma-separated string to list if present + if span_parents: + span_parents = [s.strip() for s in span_parents.split(",") if s.strip()] + + # Add optional span attributes only if present + span_attributes = { + "span_id": dynamic_metadata.get("span_id"), + "root_span_id": dynamic_metadata.get("root_span_id"), + "span_parents": span_parents, + } + request_data = { "id": litellm_call_id, "input": prompt["messages"], "output": output, "metadata": standard_logging_object, - "tags": tags, "span_attributes": {"name": span_name, "type": "llm"}, } + + # Braintrust cannot specify 'tags' for non-root spans + if dynamic_metadata.get("root_span_id") is None: + request_data["tags"] = tags + + # Only add those that are not None (or falsy) + for key, value in span_attributes.items(): + if value: + request_data[key] = value + if choices is not None: request_data["output"] = [choice.dict() for choice in choices] else: @@ -367,9 +393,6 @@ class BraintrustLogger(CustomLogger): if metrics is not None: request_data["metrics"] = metrics - if metrics is not None: - request_data["metrics"] = metrics - try: await self.global_braintrust_http_handler.post( url=f"{self.api_base}/project_logs/{project_id}/insert", diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 2128b55bf83..71929398103 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -19,7 +19,7 @@ """Database connection and data extraction for LiteLLM.""" from datetime import datetime -from typing import Any, Dict, Optional +from typing import Any, Optional, List import polars as pl @@ -46,19 +46,9 @@ class LiteLLMDatabase: """Retrieve usage data from LiteLLM daily user spend table.""" client = self._ensure_prisma_client() - # Build WHERE clause for time filtering - where_conditions = [] - if start_time_utc: - where_conditions.append(f"dus.updated_at >= '{start_time_utc.isoformat()}'") - if end_time_utc: - where_conditions.append(f"dus.updated_at <= '{end_time_utc.isoformat()}'") - - where_clause = "" - if where_conditions: - where_clause = "WHERE " + " AND ".join(where_conditions) - - # Query to get user spend data with team information - query = f""" + # Query to get user spend data with team information. Use parameter binding to + # avoid SQL injection from user-supplied timestamps or limits. + query = """ SELECT dus.id, dus.date, @@ -85,163 +75,27 @@ class LiteLLMDatabase: LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id - {where_clause} + WHERE ($1::timestamptz IS NULL OR dus.updated_at >= $1::timestamptz) + AND ($2::timestamptz IS NULL OR dus.updated_at <= $2::timestamptz) ORDER BY dus.date DESC, dus.created_at DESC """ - if limit: - query += f" LIMIT {limit}" + params: List[Any] = [ + start_time_utc, + end_time_utc, + ] + + if limit is not None: + try: + params.append(int(limit)) + except (TypeError, ValueError): + raise ValueError("limit must be an integer") + query += " LIMIT $3" try: - db_response = await client.db.query_raw(query) + db_response = await client.db.query_raw(query, *params) # Convert the response to polars DataFrame with full schema inference # This prevents schema mismatch errors when data types vary across rows return pl.DataFrame(db_response, infer_schema_length=None) except Exception as e: raise Exception(f"Error retrieving usage data: {str(e)}") - - async def get_table_info(self) -> Dict[str, Any]: - """Get information about the daily user spend table.""" - client = self._ensure_prisma_client() - - try: - # Get row count from user spend table - user_count = await self._get_table_row_count("LiteLLM_DailyUserSpend") - - # Get column structure from user spend table - query = """ - SELECT column_name, data_type, is_nullable - FROM information_schema.columns - WHERE table_name = 'LiteLLM_DailyUserSpend' - ORDER BY ordinal_position; - """ - columns_response = await client.db.query_raw(query) - - return { - "columns": columns_response, - "row_count": user_count, - "table_name": "LiteLLM_DailyUserSpend", - } - except Exception as e: - raise Exception(f"Error getting table info: {str(e)}") - - async def _get_table_row_count(self, table_name: str) -> int: - """Get row count from specified table.""" - client = self._ensure_prisma_client() - - try: - query = f'SELECT COUNT(*) as count FROM "{table_name}"' - response = await client.db.query_raw(query) - - if response and len(response) > 0: - return response[0].get("count", 0) - return 0 - except Exception: - return 0 - - async def discover_all_tables(self) -> Dict[str, Any]: - """Discover all tables in the LiteLLM database and their schemas.""" - client = self._ensure_prisma_client() - - try: - # Get all LiteLLM tables - litellm_tables_query = """ - SELECT table_name - FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name LIKE 'LiteLLM_%' - ORDER BY table_name; - """ - tables_response = await client.db.query_raw(litellm_tables_query) - table_names = [row["table_name"] for row in tables_response] - - # Get detailed schema for each table - tables_info = {} - for table_name in table_names: - # Get column information - columns_query = """ - SELECT - column_name, - data_type, - is_nullable, - column_default, - character_maximum_length, - numeric_precision, - numeric_scale, - ordinal_position - FROM information_schema.columns - WHERE table_name = $1 - AND table_schema = 'public' - ORDER BY ordinal_position; - """ - columns_response = await client.db.query_raw(columns_query, table_name) - - # Get primary key information - pk_query = """ - SELECT a.attname - FROM pg_index i - JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) - WHERE i.indrelid = $1::regclass AND i.indisprimary; - """ - pk_response = await client.db.query_raw(pk_query, f'"{table_name}"') - primary_keys = ( - [row["attname"] for row in pk_response] if pk_response else [] - ) - - # Get foreign key information - fk_query = """ - SELECT - tc.constraint_name, - kcu.column_name, - ccu.table_name AS foreign_table_name, - ccu.column_name AS foreign_column_name - FROM information_schema.table_constraints AS tc - JOIN information_schema.key_column_usage AS kcu - ON tc.constraint_name = kcu.constraint_name - JOIN information_schema.constraint_column_usage AS ccu - ON ccu.constraint_name = tc.constraint_name - WHERE tc.constraint_type = 'FOREIGN KEY' - AND tc.table_name = $1; - """ - fk_response = await client.db.query_raw(fk_query, table_name) - foreign_keys = fk_response if fk_response else [] - - # Get indexes - indexes_query = """ - SELECT - i.relname AS index_name, - array_agg(a.attname ORDER BY a.attnum) AS column_names, - ix.indisunique AS is_unique - FROM pg_class t - JOIN pg_index ix ON t.oid = ix.indrelid - JOIN pg_class i ON i.oid = ix.indexrelid - JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey) - WHERE t.relname = $1 - AND t.relkind = 'r' - GROUP BY i.relname, ix.indisunique - ORDER BY i.relname; - """ - indexes_response = await client.db.query_raw(indexes_query, table_name) - indexes = indexes_response if indexes_response else [] - - # Get row count - try: - row_count = await self._get_table_row_count(table_name) - except Exception: - row_count = 0 - - tables_info[table_name] = { - "columns": columns_response, - "primary_keys": primary_keys, - "foreign_keys": foreign_keys, - "indexes": indexes, - "row_count": row_count, - } - - return { - "tables": tables_info, - "table_count": len(table_names), - "table_names": table_names, - } - except Exception as e: - raise Exception(f"Error discovering tables: {str(e)}") diff --git a/litellm/integrations/generic_api/generic_api_compatible_callbacks.json b/litellm/integrations/generic_api/generic_api_compatible_callbacks.json index 12dc4ae643c..13fe79ae671 100644 --- a/litellm/integrations/generic_api/generic_api_compatible_callbacks.json +++ b/litellm/integrations/generic_api/generic_api_compatible_callbacks.json @@ -1,28 +1,37 @@ { - "sample_callback": { - "event_types": ["llm_api_success", "llm_api_failure"], - "endpoint": "{{environment_variables.SAMPLE_CALLBACK_URL}}", - "headers": { - "Content-Type": "application/json", - "Authorization": "Bearer {{environment_variables.SAMPLE_CALLBACK_API_KEY}}" - }, - "environment_variables": ["SAMPLE_CALLBACK_URL", "SAMPLE_CALLBACK_API_KEY"] + "sample_callback": { + "event_types": ["llm_api_success", "llm_api_failure"], + "endpoint": "{{environment_variables.SAMPLE_CALLBACK_URL}}", + "headers": { + "Content-Type": "application/json", + "Authorization": "Bearer {{environment_variables.SAMPLE_CALLBACK_API_KEY}}" }, - "rubrik": { - "event_types": ["llm_api_success"], - "endpoint": "{{environment_variables.RUBRIK_WEBHOOK_URL}}", - "headers": { - "Content-Type": "application/json", - "Authorization": "Bearer {{environment_variables.RUBRIK_API_KEY}}" - }, - "environment_variables": ["RUBRIK_API_KEY", "RUBRIK_WEBHOOK_URL"] + "environment_variables": ["SAMPLE_CALLBACK_URL", "SAMPLE_CALLBACK_API_KEY"] + }, + "rubrik": { + "event_types": ["llm_api_success"], + "endpoint": "{{environment_variables.RUBRIK_WEBHOOK_URL}}", + "headers": { + "Content-Type": "application/json", + "Authorization": "Bearer {{environment_variables.RUBRIK_API_KEY}}" }, - "sumologic": { - "endpoint": "{{environment_variables.SUMOLOGIC_WEBHOOK_URL}}", - "headers": { - "Content-Type": "application/json" - }, - "environment_variables": ["SUMOLOGIC_WEBHOOK_URL"], - "log_format": "ndjson" - } -} \ No newline at end of file + "environment_variables": ["RUBRIK_API_KEY", "RUBRIK_WEBHOOK_URL"] + }, + "sumologic": { + "endpoint": "{{environment_variables.SUMOLOGIC_WEBHOOK_URL}}", + "headers": { + "Content-Type": "application/json" + }, + "environment_variables": ["SUMOLOGIC_WEBHOOK_URL"], + "log_format": "ndjson" + }, + "qualifire_eval": { + "event_types": ["llm_api_success"], + "endpoint": "{{environment_variables.QUALIFIRE_WEBHOOK_URL}}", + "headers": { + "Content-Type": "application/json", + "X-Qualifire-API-Key": "{{environment_variables.QUALIFIRE_API_KEY}}" + }, + "environment_variables": ["QUALIFIRE_API_KEY", "QUALIFIRE_WEBHOOK_URL"] + } +} diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index e4aca5ced04..c32a7b75c51 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -45,6 +45,7 @@ def _get_cached_end_user_id_for_cost_tracking(): global _get_end_user_id_for_cost_tracking if _get_end_user_id_for_cost_tracking is None: from litellm.utils import get_end_user_id_for_cost_tracking + _get_end_user_id_for_cost_tracking = get_end_user_id_for_cost_tracking return _get_end_user_id_for_cost_tracking @@ -238,6 +239,36 @@ class PrometheusLogger(CustomLogger): ), buckets=LATENCY_BUCKETS, ) + + # Request queue time metric + self.litellm_request_queue_time_metric = self._histogram_factory( + "litellm_request_queue_time_seconds", + "Time spent in request queue before processing starts (seconds)", + labelnames=self.get_labels_for_metric( + "litellm_request_queue_time_seconds" + ), + buckets=LATENCY_BUCKETS, + ) + + # Guardrail metrics + self.litellm_guardrail_latency_metric = self._histogram_factory( + "litellm_guardrail_latency_seconds", + "Latency (seconds) for guardrail execution", + labelnames=["guardrail_name", "status", "error_type", "hook_type"], + buckets=LATENCY_BUCKETS, + ) + + self.litellm_guardrail_errors_total = self._counter_factory( + "litellm_guardrail_errors_total", + "Total number of errors encountered during guardrail execution", + labelnames=["guardrail_name", "error_type", "hook_type"], + ) + + self.litellm_guardrail_requests_total = self._counter_factory( + "litellm_guardrail_requests_total", + "Total number of guardrail invocations", + labelnames=["guardrail_name", "status", "hook_type"], + ) # llm api provider budget metrics self.litellm_provider_remaining_budget_metric = self._gauge_factory( "litellm_provider_remaining_budget_metric", @@ -330,6 +361,25 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_requests_metric"), ) + # Cache metrics + self.litellm_cache_hits_metric = self._counter_factory( + name="litellm_cache_hits_metric", + documentation="Total number of LiteLLM cache hits", + labelnames=self.get_labels_for_metric("litellm_cache_hits_metric"), + ) + + self.litellm_cache_misses_metric = self._counter_factory( + name="litellm_cache_misses_metric", + documentation="Total number of LiteLLM cache misses", + labelnames=self.get_labels_for_metric("litellm_cache_misses_metric"), + ) + + self.litellm_cached_tokens_metric = self._counter_factory( + name="litellm_cached_tokens_metric", + documentation="Total tokens served from LiteLLM cache", + labelnames=self.get_labels_for_metric("litellm_cached_tokens_metric"), + ) + except Exception as e: print_verbose(f"Got exception on init prometheus client {str(e)}") raise e @@ -801,7 +851,7 @@ class PrometheusLogger(CustomLogger): litellm_params = kwargs.get("litellm_params", {}) or {} _metadata = litellm_params.get("metadata", {}) get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() - + end_user_id = get_end_user_id_for_cost_tracking( litellm_params, service_type="prometheus" ) @@ -821,7 +871,7 @@ class PrometheusLogger(CustomLogger): user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[ "metadata" ].get("user_api_key_auth_metadata") - + # Include top-level metadata fields (excluding nested dictionaries) # This allows accessing fields like requester_ip_address from top-level metadata top_level_metadata = standard_logging_payload.get("metadata", {}) @@ -832,7 +882,7 @@ class PrometheusLogger(CustomLogger): for k, v in top_level_metadata.items() if not isinstance(v, dict) # Exclude nested dicts to avoid conflicts } - + combined_metadata: Dict[str, Any] = { **top_level_fields, # Include top-level fields first **(_requester_metadata if _requester_metadata else {}), @@ -951,6 +1001,12 @@ class PrometheusLogger(CustomLogger): kwargs, start_time, end_time, enum_values, output_tokens ) + # cache metrics + self._increment_cache_metrics( + standard_logging_payload=standard_logging_payload, # type: ignore + enum_values=enum_values, + ) + if ( standard_logging_payload["stream"] is True ): # log successful streaming requests from logging event hook. @@ -1020,6 +1076,54 @@ class PrometheusLogger(CustomLogger): standard_logging_payload["completion_tokens"] ) + def _increment_cache_metrics( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + ): + """ + Increment cache-related Prometheus metrics based on cache hit/miss status. + + Args: + standard_logging_payload: Contains cache_hit field (True/False/None) + enum_values: Label values for Prometheus metrics + """ + cache_hit = standard_logging_payload.get("cache_hit") + + # Only track if cache_hit has a definite value (True or False) + if cache_hit is None: + return + + if cache_hit is True: + # Increment cache hits counter + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_cache_hits_metric" + ), + enum_values=enum_values, + ) + self.litellm_cache_hits_metric.labels(**_labels).inc() + + # Increment cached tokens counter + total_tokens = standard_logging_payload.get("total_tokens", 0) + if total_tokens > 0: + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_cached_tokens_metric" + ), + enum_values=enum_values, + ) + self.litellm_cached_tokens_metric.labels(**_labels).inc(total_tokens) + else: + # cache_hit is False - increment cache misses counter + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_cache_misses_metric" + ), + enum_values=enum_values, + ) + self.litellm_cache_misses_metric.labels(**_labels).inc() + async def _increment_remaining_budget_metrics( self, user_api_team: Optional[str], @@ -1188,6 +1292,22 @@ class PrometheusLogger(CustomLogger): total_time_seconds ) + # request queue time (time from arrival to processing start) + _litellm_params = kwargs.get("litellm_params", {}) or {} + queue_time_seconds = _litellm_params.get("metadata", {}).get( + "queue_time_seconds" + ) + if queue_time_seconds is not None and queue_time_seconds >= 0: + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_request_queue_time_seconds" + ), + enum_values=enum_values, + ) + self.litellm_request_queue_time_metric.labels(**_labels).observe( + queue_time_seconds + ) + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): from litellm.types.utils import StandardLoggingPayload @@ -1208,7 +1328,7 @@ class PrometheusLogger(CustomLogger): litellm_params = kwargs.get("litellm_params", {}) or {} get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() - + end_user_id = get_end_user_id_for_cost_tracking( litellm_params, service_type="prometheus" ) @@ -1562,7 +1682,6 @@ class PrometheusLogger(CustomLogger): api_provider=llm_provider or "", ) if exception is not None: - _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( metric_name="litellm_deployment_failure_responses" @@ -1595,12 +1714,11 @@ class PrometheusLogger(CustomLogger): enum_values: UserAPIKeyLabelValues, output_tokens: float = 1.0, ): - try: verbose_logger.debug("setting remaining tokens requests metric") - standard_logging_payload: Optional[StandardLoggingPayload] = ( - request_kwargs.get("standard_logging_object") - ) + standard_logging_payload: Optional[ + StandardLoggingPayload + ] = request_kwargs.get("standard_logging_object") if standard_logging_payload is None: return @@ -1743,6 +1861,50 @@ class PrometheusLogger(CustomLogger): ) return + def _record_guardrail_metrics( + self, + guardrail_name: str, + latency_seconds: float, + status: str, + error_type: Optional[str], + hook_type: str, + ): + """ + Record guardrail metrics for prometheus. + + Args: + guardrail_name: Name of the guardrail + latency_seconds: Execution latency in seconds + status: "success" or "error" + error_type: Type of error if any, None otherwise + hook_type: "pre_call", "during_call", or "post_call" + """ + try: + # Record latency + self.litellm_guardrail_latency_metric.labels( + guardrail_name=guardrail_name, + status=status, + error_type=error_type or "none", + hook_type=hook_type, + ).observe(latency_seconds) + + # Record request count + self.litellm_guardrail_requests_total.labels( + guardrail_name=guardrail_name, + status=status, + hook_type=hook_type, + ).inc() + + # Record error count if there was an error + if status == "error" and error_type: + self.litellm_guardrail_errors_total.labels( + guardrail_name=guardrail_name, + error_type=error_type, + hook_type=hook_type, + ).inc() + except Exception as e: + verbose_logger.debug(f"Error recording guardrail metrics: {str(e)}") + @staticmethod def _get_exception_class_name(exception: Exception) -> str: exception_class_name = "" @@ -2380,10 +2542,10 @@ class PrometheusLogger(CustomLogger): from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES from litellm.integrations.custom_logger import CustomLogger - prometheus_loggers: List[CustomLogger] = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=PrometheusLogger - ) + prometheus_loggers: List[ + CustomLogger + ] = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=PrometheusLogger ) # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them verbose_logger.debug("found %s prometheus loggers", len(prometheus_loggers)) @@ -2455,7 +2617,7 @@ def prometheus_label_factory( if UserAPIKeyLabelNames.END_USER.value in filtered_labels: get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() - + filtered_labels["end_user"] = get_end_user_id_for_cost_tracking( litellm_params={"user_api_key_end_user_id": enum_values.end_user}, service_type="prometheus", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4902c1f8343..e0a799d8e5c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4839,9 +4839,9 @@ class StandardLoggingPayloadSetup: metadata = litellm_params.get("metadata") or {} litellm_metadata = litellm_params.get("litellm_metadata") or {} if metadata.get("tags", []): - request_tags = metadata.get("tags", []) + request_tags = metadata.get("tags", []).copy() elif litellm_metadata.get("tags", []): - request_tags = litellm_metadata.get("tags", []) + request_tags = litellm_metadata.get("tags", []).copy() else: request_tags = [] user_agent_tags = StandardLoggingPayloadSetup._get_user_agent_tags( diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0c331e43038..d8e82199272 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1645,9 +1645,12 @@ def convert_to_anthropic_tool_result( ) elif content["type"] == "image_url": format = content["image_url"].get("format") if isinstance(content["image_url"], dict) else None - anthropic_content_list.append( - create_anthropic_image_param(content["image_url"], format=format) + _anthropic_image_param = create_anthropic_image_param(content["image_url"], format=format) + _anthropic_image_param = add_cache_control_to_content( + anthropic_content_element=_anthropic_image_param, + original_content_element=content, ) + anthropic_content_list.append(_anthropic_image_param) anthropic_content = anthropic_content_list anthropic_tool_result: Optional[AnthropicMessagesToolResultParam] = None diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 87f81d117f0..506b7fdfe5e 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -25,7 +25,24 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): return "gpt-5" in model or "gpt5_series" in model def get_supported_openai_params(self, model: str) -> List[str]: - return OpenAIGPT5Config.get_supported_openai_params(self, model=model) + """Get supported parameters for Azure OpenAI GPT-5 models. + + Azure OpenAI GPT-5.2 models support logprobs, unlike OpenAI's GPT-5. + This overrides the parent class to add logprobs support back for gpt-5.2. + + Reference: + - Tested with Azure OpenAI GPT-5.2 (api-version: 2025-01-01-preview) + - Azure returns logprobs successfully despite Microsoft's general + documentation stating reasoning models don't support it. + """ + params = OpenAIGPT5Config.get_supported_openai_params(self, model=model) + + # Only gpt-5.2 has been verified to support logprobs on Azure + if self.is_model_gpt_5_2_model(model): + azure_supported_params = ["logprobs", "top_logprobs"] + params.extend(azure_supported_params) + + return params def map_openai_params( self, diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 0d5494541ec..e9cea23ea4a 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -314,6 +314,12 @@ class BaseAWSLLM: if model.startswith("invoke/"): model = model.replace("invoke/", "", 1) + # Special case: Check for "nova" in model name first (before "amazon") + # This handles amazon.nova-* models which would otherwise match "amazon" (Titan) + if "nova" in model.lower(): + if "nova" in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL): + return cast(BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova") + _split_model = model.split(".")[0] if _split_model in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL): return cast(BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model) @@ -323,13 +329,9 @@ class BaseAWSLLM: if provider is not None: return provider - # check if provider == "nova" - if "nova" in model: - return "nova" - else: - for provider in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL): - if provider in model: - return provider + for provider in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL): + if provider in model: + return provider return None @staticmethod @@ -364,7 +366,7 @@ class BaseAWSLLM: elif provider == "qwen3" and "qwen3/" in model_id: model_id = BaseAWSLLM._get_model_id_from_model_with_spec( model_id, spec="qwen3" - ) + ) elif provider == "stability" and "stability/" in model_id: model_id = BaseAWSLLM._get_model_id_from_model_with_spec( model_id, spec="stability" @@ -416,7 +418,7 @@ class BaseAWSLLM: if "nova" in model.lower(): if "nova" in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL): return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, "nova") - + # Handle regional models like us.twelvelabs.marengo-embed-2-7-v1:0 if "." in model: parts = model.split(".") @@ -962,7 +964,9 @@ class BaseAWSLLM: return endpoint_url, proxy_endpoint_url def _select_default_endpoint_url( - self, endpoint_type: Optional[Literal["runtime", "agent", "agentcore"]], aws_region_name: str + self, + endpoint_type: Optional[Literal["runtime", "agent", "agentcore"]], + aws_region_name: str, ) -> str: """ Select the default endpoint url based on the endpoint type diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index c602b71fe05..cf8aee6954b 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -524,6 +524,12 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): if model.startswith("invoke/"): model = model.replace("invoke/", "", 1) + # Special case: Check for "nova" in model name first (before "amazon") + # This handles amazon.nova-* models which would otherwise match "amazon" (Titan) + if "nova" in model.lower(): + if "nova" in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): + return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova") + _split_model = model.split(".")[0] if _split_model in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model) @@ -533,10 +539,6 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): if provider is not None: return provider - # check if provider == "nova" - if "nova" in model: - return "nova" - for provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): if provider in model: return provider diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9edfe320fb2..5cb51cf994f 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -15,7 +15,7 @@ import litellm from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) -from litellm.llms.base_llm.base_utils import BaseLLMModelInfo +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret @@ -359,6 +359,70 @@ def get_bedrock_tool_name(response_tool_name: str) -> str: return response_tool_name +# Cache the global regions list at module level +_BEDROCK_GLOBAL_REGIONS: Optional[List[str]] = None + + +def _get_all_bedrock_regions() -> List[str]: + """Get all Bedrock regions, cached at module level.""" + global _BEDROCK_GLOBAL_REGIONS + if _BEDROCK_GLOBAL_REGIONS is None: + _BEDROCK_GLOBAL_REGIONS = AmazonBedrockGlobalConfig().get_all_regions() + return _BEDROCK_GLOBAL_REGIONS + + +def get_bedrock_cross_region_inference_regions() -> List[str]: + """Abbreviations of regions AWS Bedrock supports for cross region inference.""" + return ["global", "us", "eu", "apac", "jp", "au", "us-gov"] + + +def extract_model_name_from_bedrock_arn(model: str) -> str: + """ + Extract the model name from an AWS Bedrock ARN. + Returns the string after the last '/' if 'arn' is in the input string. + """ + if "arn" in model.lower(): + return model.split("/")[-1] + return model + + +def strip_bedrock_routing_prefix(model: str) -> str: + """Strip LiteLLM routing prefixes from model name.""" + for prefix in ["bedrock/", "converse/", "invoke/", "openai/"]: + if model.startswith(prefix): + model = model.split("/", 1)[1] + return model + + +def get_bedrock_base_model(model: str) -> str: + """ + Get the base model from the given model name. + + Handle model names like: + - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1" + - "bedrock/converse/model" -> "model" + """ + model = strip_bedrock_routing_prefix(model) + model = extract_model_name_from_bedrock_arn(model) + + potential_region = model.split(".", 1)[0] + alt_potential_region = model.split("/", 1)[0] + + if potential_region in get_bedrock_cross_region_inference_regions(): + return model.split(".", 1)[1] + elif ( + alt_potential_region in _get_all_bedrock_regions() + and len(model.split("/", 1)) > 1 + ): + return model.split("/", 1)[1] + + return model + + +# Import after standalone functions to avoid circular imports +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter + + class BedrockModelInfo(BaseLLMModelInfo): global_config = AmazonBedrockGlobalConfig() all_global_regions = global_config.get_all_regions() @@ -394,76 +458,34 @@ class BedrockModelInfo(BaseLLMModelInfo): ) -> List[str]: return [] - @staticmethod - def extract_model_name_from_arn(model: str) -> str: + def get_token_counter(self) -> Optional[BaseTokenCounter]: """ - Extract the model name from an AWS Bedrock ARN. - Returns the string after the last '/' if 'arn' is in the input string. - - Args: - arn (str): The ARN string to parse + Factory method to create a Bedrock token counter. Returns: - str: The extracted model name if 'arn' is in the string, - otherwise returns the original string + BedrockTokenCounter instance for this provider. """ - if "arn" in model.lower(): - return model.split("/")[-1] - return model + return BedrockTokenCounter() + + @staticmethod + def extract_model_name_from_arn(model: str) -> str: + """Wrapper for standalone function. See extract_model_name_from_bedrock_arn().""" + return extract_model_name_from_bedrock_arn(model) @staticmethod def get_non_litellm_routing_model_name(model: str) -> str: - if model.startswith("bedrock/"): - model = model.split("/", 1)[1] - - if model.startswith("converse/"): - model = model.split("/", 1)[1] - - if model.startswith("invoke/"): - model = model.split("/", 1)[1] - - if model.startswith("openai/"): - model = model.split("/", 1)[1] - - return model + """Wrapper for standalone function. See strip_bedrock_routing_prefix().""" + return strip_bedrock_routing_prefix(model) @staticmethod def get_base_model(model: str) -> str: - """ - Get the base model from the given model name. - - Handle model names like - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1" - AND "meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1" - """ - - model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model) - model = BedrockModelInfo.extract_model_name_from_arn(model) - - potential_region = model.split(".", 1)[0] - - alt_potential_region = model.split("/", 1)[ - 0 - ] # in model cost map we store regional information like `/us-west-2/bedrock-model` - - if ( - potential_region - in BedrockModelInfo._supported_cross_region_inference_region() - ): - return model.split(".", 1)[1] - elif ( - alt_potential_region in BedrockModelInfo.all_global_regions - and len(model.split("/", 1)) > 1 - ): - return model.split("/", 1)[1] - - return model + """Wrapper for standalone function. See get_bedrock_base_model().""" + return get_bedrock_base_model(model) @staticmethod def _supported_cross_region_inference_region() -> List[str]: - """ - Abbreviations of regions AWS Bedrock supports for cross region inference - """ - return ["global", "us", "eu", "apac", "jp", "au", "us-gov"] + """Wrapper for standalone function. See get_bedrock_cross_region_inference_regions().""" + return get_bedrock_cross_region_inference_regions() @staticmethod def get_bedrock_route( diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py new file mode 100644 index 00000000000..b680bd046ef --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -0,0 +1,87 @@ +""" +Bedrock Token Counter implementation using the CountTokens API. +""" + +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_logger +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.bedrock.common_utils import get_bedrock_base_model +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.types.utils import LlmProviders, TokenCountResponse + + +class BedrockTokenCounter(BaseTokenCounter): + """Token counter implementation for AWS Bedrock provider using the CountTokens API.""" + + def should_use_token_counting_api( + self, + custom_llm_provider: Optional[str] = None, + ) -> bool: + """ + Returns True if we should use the Bedrock CountTokens API for token counting. + """ + return custom_llm_provider == LlmProviders.BEDROCK.value + + async def count_tokens( + self, + model_to_use: str, + messages: Optional[List[Dict[str, Any]]], + contents: Optional[List[Dict[str, Any]]], + deployment: Optional[Dict[str, Any]] = None, + request_model: str = "", + ) -> Optional[TokenCountResponse]: + """ + Count tokens using AWS Bedrock's CountTokens API. + + This method calls the existing BedrockCountTokensHandler to make an API call + to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting. + + Args: + model_to_use: The model identifier + messages: The messages to count tokens for + contents: Alternative content format (not used for Bedrock) + deployment: Deployment configuration containing litellm_params + request_model: The original request model name + + Returns: + TokenCountResponse with token count, or None if counting fails + """ + if not messages: + return None + + deployment = deployment or {} + litellm_params = deployment.get("litellm_params", {}) + + # Build request data in the format expected by BedrockCountTokensHandler + request_data = { + "model": model_to_use, + "messages": messages, + } + + # Get the resolved model (strip prefixes like bedrock/, converse/, etc.) + resolved_model = get_bedrock_base_model(model_to_use) + + try: + handler = BedrockCountTokensHandler() + result = await handler.handle_count_tokens_request( + request_data=request_data, + litellm_params=litellm_params, + resolved_model=resolved_model, + ) + + # Transform response to TokenCountResponse + if result is not None: + return TokenCountResponse( + total_tokens=result.get("input_tokens", 0), + request_model=request_model, + model_used=model_to_use, + tokenizer_type="bedrock_api", + original_response=result, + ) + except Exception as e: + verbose_logger.warning( + f"Error calling Bedrock CountTokens API: {e}, falling back to default tokenizer" + ) + + return None diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index d4355c0c360..e8366165b65 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -6,10 +6,9 @@ Simplified handler leveraging existing LiteLLM Bedrock infrastructure. from typing import Any, Dict -from fastapi import HTTPException - import litellm from litellm._logging import verbose_logger +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig from litellm.llms.custom_httpx.http_handler import get_async_httpx_client @@ -70,6 +69,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): verbose_logger.debug(f"Making request to: {endpoint_url}") # Use existing _sign_request method from BaseAWSLLM + # Extract api_key for bearer token auth if provided + api_key = litellm_params.get("api_key", None) headers = {"Content-Type": "application/json"} signed_headers, signed_body = self._sign_request( service_name="bedrock", @@ -78,6 +79,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): request_data=bedrock_request, api_base=endpoint_url, model=resolved_model, + api_key=api_key, ) async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) @@ -94,9 +96,9 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): if response.status_code != 200: error_text = response.text verbose_logger.error(f"AWS Bedrock error: {error_text}") - raise HTTPException( - status_code=400, - detail={"error": f"AWS Bedrock error: {error_text}"}, + raise BedrockError( + status_code=response.status_code, + message=f"AWS Bedrock error: {error_text}", ) bedrock_response = response.json() @@ -112,12 +114,12 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): return final_response - except HTTPException: - # Re-raise HTTP exceptions as-is + except BedrockError: + # Re-raise Bedrock exceptions as-is raise except Exception as e: verbose_logger.error(f"Error in CountTokens handler: {str(e)}") - raise HTTPException( + raise BedrockError( status_code=500, - detail={"error": f"CountTokens processing error: {str(e)}"}, + message=f"CountTokens processing error: {str(e)}", ) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index d46ed3aa452..b313cc9df3c 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -8,7 +8,7 @@ to AWS Bedrock's CountTokens API format and vice versa. from typing import Any, Dict, List from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM -from litellm.llms.bedrock.common_utils import BedrockModelInfo +from litellm.llms.bedrock.common_utils import get_bedrock_base_model class BedrockCountTokensConfig(BaseAWSLLM): @@ -141,7 +141,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): Complete endpoint URL for CountTokens API """ # Use existing LiteLLM function to get the base model ID (removes region prefix) - model_id = BedrockModelInfo.get_base_model(model) + model_id = get_bedrock_base_model(model) # Remove bedrock/ prefix if present if model_id.startswith("bedrock/"): diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index e53829d3329..30c5b4f17c5 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -150,6 +150,15 @@ def get_api_key_from_env() -> Optional[str]: return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") +def get_vertex_api_key_from_env() -> Optional[str]: + """ + Get API key from environment for Vertex AI. + Checks VERTEXAI_API_KEY and VERTEX_API_KEY environment variables. + This allows using Vertex AI with API keys instead of service account credentials. + """ + return get_secret_str("VERTEXAI_API_KEY") or get_secret_str("VERTEX_API_KEY") + + class GoogleAIStudioTokenCounter(BaseTokenCounter): """Token counter implementation for Google AI Studio provider.""" def should_use_token_counting_api( diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 206aee1359d..bda3684a8a8 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -61,6 +61,10 @@ "max_completion_tokens": "max_tokens" } }, + "abliteration": { + "base_url": "https://api.abliteration.ai/v1", + "api_key_env": "ABLITERATION_API_KEY" + }, "llamagate": { "base_url": "https://api.llamagate.dev/v1", "api_key_env": "LLAMAGATE_API_KEY", diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py new file mode 100644 index 00000000000..d1d0e911d16 --- /dev/null +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -0,0 +1,182 @@ +""" +OpenRouter Embedding API Configuration. + +This module provides the configuration for OpenRouter's Embedding API. +OpenRouter is OpenAI-compatible and supports embeddings via the /v1/embeddings endpoint. + +Docs: https://openrouter.ai/docs +""" +from typing import TYPE_CHECKING, Any, Optional + +import httpx + +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.types.llms.openai import AllEmbeddingInputValues +from litellm.types.utils import EmbeddingResponse +from litellm.utils import convert_to_model_response_object + +from ..common_utils import OpenRouterException + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): + """ + Configuration for OpenRouter's Embedding API. + + Reference: https://openrouter.ai/docs + """ + + def validate_environment( + self, + headers: dict, + model: str, + messages: list, + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate environment and set up headers for OpenRouter API. + + OpenRouter requires: + - Authorization header with Bearer token + - HTTP-Referer header (site URL) + - X-Title header (app name) + """ + from litellm import get_secret + + # Get OpenRouter-specific headers + openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai" + openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM" + + openrouter_headers = { + "HTTP-Referer": openrouter_site_url, + "X-Title": openrouter_app_name, + "Content-Type": "application/json", + } + + # Add Authorization header if api_key is provided + if api_key: + openrouter_headers["Authorization"] = f"Bearer {api_key}" + + # Merge with existing headers (user's extra_headers take priority) + merged_headers = {**openrouter_headers, **headers} + + return merged_headers + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete URL for OpenRouter Embedding API endpoint. + """ + # api_base is already set to https://openrouter.ai/api/v1 in main.py + # Remove trailing slashes + if api_base: + api_base = api_base.rstrip("/") + else: + api_base = "https://openrouter.ai/api/v1" + + # Return the embeddings endpoint + return f"{api_base}/embeddings" + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict, + headers: dict, + ) -> dict: + """ + Transform embedding request to OpenRouter format (OpenAI-compatible). + """ + # Ensure input is a list + if isinstance(input, str): + input = [input] + + # OpenRouter expects the full model name (e.g., google/gemini-embedding-001) + # Strip 'openrouter/' prefix if present + if model.startswith("openrouter/"): + model = model.replace("openrouter/", "", 1) + + return { + "model": model, + "input": input, + **optional_params, + } + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + request_data: dict, + optional_params: dict, + litellm_params: dict, + ) -> EmbeddingResponse: + """ + Transform embedding response from OpenRouter format (OpenAI-compatible). + """ + logging_obj.post_call(original_response=raw_response.text) + + # OpenRouter returns standard OpenAI-compatible embedding response + response_json = raw_response.json() + + return convert_to_model_response_object( + response_object=response_json, + model_response_object=model_response, + response_type="embedding", + ) + + def get_supported_openai_params(self, model: str) -> list: + """ + Get list of supported OpenAI parameters for OpenRouter embeddings. + """ + return [ + "timeout", + "dimensions", + "encoding_format", + "user", + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI parameters to OpenRouter format. + """ + for param, value in non_default_params.items(): + if param in self.get_supported_openai_params(model): + optional_params[param] = value + return optional_params + + def get_error_class( + self, error_message: str, status_code: int, headers: Any + ) -> Any: + """ + Get the error class for OpenRouter errors. + """ + return OpenRouterException( + message=error_message, + status_code=status_code, + headers=headers, + ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index ba1788a217f..91100cf7d7b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -480,20 +480,21 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): or tool_name == VertexToolName.CODE_EXECUTION.value ): # code_execution maintained for backwards compatibility code_execution = self.get_tool_value(tool, "codeExecution") - elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH.value: - googleSearch = self.get_tool_value( - tool, VertexToolName.GOOGLE_SEARCH.value - ) - elif ( - tool_name and tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value + elif tool_name and ( + tool_name == VertexToolName.GOOGLE_SEARCH.value + or tool_name == "google_search" ): - googleSearchRetrieval = self.get_tool_value( - tool, VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value - ) - elif tool_name and tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value: - enterpriseWebSearch = self.get_tool_value( - tool, VertexToolName.ENTERPRISE_WEB_SEARCH.value - ) + googleSearch = self.get_tool_value(tool, tool_name) + elif tool_name and ( + tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value + or tool_name == "google_search_retrieval" + ): + googleSearchRetrieval = self.get_tool_value(tool, tool_name) + elif tool_name and ( + tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value + or tool_name == "enterprise_web_search" + ): + enterpriseWebSearch = self.get_tool_value(tool, tool_name) elif tool_name and ( tool_name == VertexToolName.URL_CONTEXT.value or tool_name == "urlContext" @@ -1811,6 +1812,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): functions: Optional[ChatCompletionToolCallFunctionChunk] = None thinking_blocks: Optional[List[ChatCompletionThinkingBlock]] = None reasoning_content: Optional[str] = None + thought_signatures: Optional[Any] = None for idx, candidate in enumerate(_candidates): if "content" not in candidate: diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 826f151df35..a3606ff9deb 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -388,6 +388,10 @@ class VertexBase: Internal function. Returns the token and url for the call. Handles logic if it's google ai studio vs. vertex ai. + + For Vertex AI: + - If gemini_api_key is provided, use API key authentication (x-goog-api-key header) + - Otherwise, use service account credentials (OAuth2 Bearer token) Returns token, url @@ -400,7 +404,7 @@ class VertexBase: stream=stream, gemini_api_key=gemini_api_key, ) - auth_header = None # this field is not used for gemin + auth_header = None # this field is not used for gemini else: vertex_location = self.get_vertex_region( vertex_region=vertex_location, @@ -409,14 +413,32 @@ class VertexBase: ### SET RUNTIME ENDPOINT ### version = "v1beta1" if should_use_v1beta1_features is True else "v1" - url, endpoint = _get_vertex_url( - mode=mode, - model=model, - stream=stream, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_api_version=version, - ) + + # Check if using API key authentication for Vertex AI + if gemini_api_key and not vertex_credentials: + # When using API key with Vertex AI, use the Google AI Studio endpoint + # This is because Vertex AI API keys work with generativelanguage.googleapis.com + verbose_logger.debug( + f"Using Vertex AI API key authentication for model: {model} - routing to Google AI Studio endpoint" + ) + url, endpoint = _get_gemini_url( + mode=mode, + model=model, + stream=stream, + gemini_api_key=gemini_api_key, + ) + # API key is already included in the URL by _get_gemini_url + auth_header = None + else: + # Use OAuth2 Bearer token authentication (traditional Vertex AI) + url, endpoint = _get_vertex_url( + mode=mode, + model=model, + stream=stream, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version=version, + ) return self._check_custom_proxy( api_base=api_base, diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 186d858321a..c7e6a77b96f 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -7,13 +7,14 @@ WatsonX follows the OpenAI spec for audio transcription. from typing import Any, Dict, List, Optional import litellm +from httpx import Response from litellm.litellm_core_utils.audio_utils.utils import process_audio_file from litellm.types.llms.openai import ( AllMessageValues, OpenAIAudioTranscriptionOptionalParams, ) from litellm.types.llms.watsonx import WatsonXAudioTranscriptionRequestBody -from litellm.types.utils import FileTypes +from litellm.types.utils import FileTypes, TranscriptionResponse from ...base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, @@ -156,3 +157,48 @@ class IBMWatsonXAudioTranscriptionConfig( url = f"{url}?version={api_version}" return url + + def transform_audio_transcription_response( + self, + raw_response: Response, + ) -> TranscriptionResponse: + """ + Transform the audio transcription response from WatsonX. + + WatsonX may include a 'model' field in the response, which needs to be + removed before creating the TranscriptionResponse object. + """ + try: + raw_response_json = raw_response.json() + except Exception as e: + raise ValueError( + f"Error transforming response to json: {str(e)}\nResponse: {raw_response.text}" + ) + + # Extract only valid fields for TranscriptionResponse.__init__() + # TranscriptionResponse only accepts 'text' and 'usage' in __init__() + text = raw_response_json.get("text") + usage = raw_response_json.get("usage") + + # Create response with only valid fields + response_kwargs = {} + if text is not None: + response_kwargs["text"] = text + if usage is not None: + response_kwargs["usage"] = usage + + if not response_kwargs: + raise ValueError( + "Invalid response format. Received response does not match the expected format. Got: ", + raw_response_json, + ) + + response = TranscriptionResponse(**response_kwargs) + + # Add other fields using dictionary-style assignment (like duration, task, etc.) + # Skip fields that TranscriptionResponse doesn't accept in __init__() + for key, value in raw_response_json.items(): + if key not in ["text", "usage", "model"]: # text/usage already set, model should be excluded + response[key] = value + + return response diff --git a/litellm/main.py b/litellm/main.py index e8a8b504d96..10e3bcac04b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -189,7 +189,7 @@ from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from .llms.custom_llm import CustomLLM, custom_chat_llm_router from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha, palm -from .llms.gemini.common_utils import get_api_key_from_env +from .llms.gemini.common_utils import get_api_key_from_env, get_vertex_api_key_from_env from .llms.groq.chat.handler import GroqChatCompletion from .llms.heroku.chat.transformation import HerokuChatConfig from .llms.huggingface.embedding.handler import HuggingFaceEmbedding @@ -3230,6 +3230,12 @@ def completion( # type: ignore # noqa: PLR0915 or get_secret("VERTEXAI_CREDENTIALS") ) + vertex_api_key = ( + api_key + or get_vertex_api_key_from_env() + or litellm.api_key + ) + api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE") new_params = safe_deep_copy(optional_params or {}) @@ -3271,7 +3277,7 @@ def completion( # type: ignore # noqa: PLR0915 vertex_location=vertex_ai_location, vertex_project=vertex_ai_project, vertex_credentials=vertex_credentials, - gemini_api_key=None, + gemini_api_key=vertex_api_key, # Support for Vertex AI API Key logging_obj=logging, acompletion=acompletion, timeout=timeout, @@ -4701,6 +4707,51 @@ def embedding( # noqa: PLR0915 litellm_params=litellm_params_dict, headers=headers, ) + elif custom_llm_provider == "openrouter": + api_base = ( + api_base + or litellm.api_base + or get_secret_str("OPENROUTER_API_BASE") + or "https://openrouter.ai/api/v1" + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.openrouter_key + or get_secret("OPENROUTER_API_KEY") + or get_secret("OR_API_KEY") + ) + + openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai" + openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM" + + openrouter_headers = { + "HTTP-Referer": openrouter_site_url, + "X-Title": openrouter_app_name, + } + + _headers = headers or litellm.headers + if _headers: + openrouter_headers.update(_headers) + + headers = openrouter_headers + + response = base_llm_http_handler.embedding( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + litellm_params=litellm_params_dict, + headers=headers, + ) elif custom_llm_provider == "huggingface": api_key = ( api_key diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3a548e203c5..1029f2241a1 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -551,6 +551,7 @@ class MCPServerManager: allowed_tools=getattr(mcp_server, "allowed_tools", None), disallowed_tools=getattr(mcp_server, "disallowed_tools", None), allow_all_keys=mcp_server.allow_all_keys, + updated_at=getattr(mcp_server, "updated_at", None), ) return new_server @@ -697,9 +698,7 @@ class MCPServerManager: results = await asyncio.gather(*tasks) # Flatten results into single list - list_tools_result: List[MCPTool] = [ - tool for tools in results for tool in tools - ] + list_tools_result: List[MCPTool] = [tool for tools in results for tool in tools] verbose_logger.info( f"Successfully fetched {len(list_tools_result)} tools total from all servers" @@ -2059,7 +2058,8 @@ class MCPServerManager: return None - async def _add_mcp_servers_from_db_to_in_memory_registry(self): + async def reload_servers_from_database(self): + """Re-synchronize the in-memory MCP server registry with the database.""" from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers from litellm.proxy.management_endpoints.mcp_management_endpoints import ( get_prisma_client_or_throw, @@ -2074,15 +2074,34 @@ class MCPServerManager: db_mcp_servers = await get_all_mcp_servers(prisma_client) verbose_logger.info(f"Found {len(db_mcp_servers)} MCP servers in database") - # ensure the global_mcp_server_manager is up to date with the db + previous_registry = self.registry + new_registry: Dict[str, MCPServer] = {} + for server in db_mcp_servers: + existing_server = previous_registry.get(server.server_id) + + if ( + existing_server is not None + and existing_server.updated_at is not None + and server.updated_at is not None + and existing_server.updated_at == server.updated_at + ): + # Re-use existing server instance to avoid re-running build_mcp_server_from_table() + # which can perform network discovery for OAuth2 servers. + new_registry[server.server_id] = existing_server + continue + verbose_logger.debug( - f"Adding server to registry: {server.server_id} ({server.server_name})" + f"Building server from DB: {server.server_id} ({server.server_name})" ) - await self.add_server(server) + new_registry[server.server_id] = await self.build_mcp_server_from_table( + server + ) + + self.registry = new_registry verbose_logger.debug( - f"Registry now contains {len(self.get_registry())} servers" + "MCP registry refreshed (%s servers in registry)", len(new_registry) ) def get_mcp_servers_from_ids(self, server_ids: List[str]) -> List[MCPServer]: @@ -2369,13 +2388,6 @@ class MCPServerManager: servers.append(self._build_mcp_server_table(server)) return servers - async def reload_servers_from_database(self): - """ - Public method to reload all MCP servers from database into registry. - This can be called from management endpoints to ensure registry is up to date. - """ - await self._add_mcp_servers_from_db_to_in_memory_registry() - async def get_all_mcp_servers_with_health_unfiltered( self, server_ids: Optional[List[str]] = None ) -> List[LiteLLM_MCPServerTable]: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index be1e4dbcdc9..09c8bb562f7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1942,7 +1942,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): description="connect to a postgres db - needed for generating temporary keys + tracking spend / key", ) database_connection_pool_limit: Optional[int] = Field( - 100, + 10, description="default connection pool for prisma client connecting to postgres db", ) database_connection_timeout: Optional[float] = Field( diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 334362a0271..7de6b7fccfc 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -11,7 +11,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, - create_streaming_response, + create_response, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.types.utils import TokenCountResponse @@ -106,7 +106,7 @@ async def anthropic_response( # noqa: PLR0915 ) ) - return await create_streaming_response( + return await create_response( generator=selected_data_generator, media_type="text/event-stream", headers={}, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index de4973ecc69..26778ece60e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -147,6 +147,7 @@ async def common_checks( # 3.1. If organization is in budget await _organization_max_budget_check( valid_token=valid_token, + team_object=team_object, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, @@ -2310,61 +2311,86 @@ async def _team_max_budget_check( async def _organization_max_budget_check( valid_token: Optional[UserAPIKeyAuth], + team_object: Optional[LiteLLM_TeamTable], prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache, proxy_logging_obj: ProxyLogging, ): """ Check if the organization is over its max budget. + + This function checks the organization budget using: + 1. First, tries to use valid_token.org_id (if key has organization_id set) + 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does) + + This ensures organization budget checks work even when keys don't have organization_id + set directly, as long as their team belongs to an organization. Raises: BudgetExceededError if the organization is over its max budget. Triggers a budget alert if the organization is over its max budget. """ - # Only check if token has organization info and organization_max_budget is set - if ( - valid_token is None - or valid_token.org_id is None - or valid_token.organization_max_budget is None - or valid_token.organization_max_budget <= 0 - ): + if valid_token is None or prisma_client is None: return - # Get organization object to check current spend - if prisma_client is not None: - org_table = await get_org_object( - org_id=valid_token.org_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + # Determine organization_id: first try from token, then fallback to team + org_id: Optional[str] = None + if valid_token.org_id is not None: + org_id = valid_token.org_id + elif team_object is not None and team_object.organization_id is not None: + org_id = team_object.organization_id + + # If no organization_id found, skip the check + if org_id is None: + return + + # Get organization object with budget table to check current spend and max budget + try: + org_table = await prisma_client.db.litellm_organizationtable.find_unique( + where={"organization_id": org_id}, + include={"litellm_budget_table": True}, + ) + except Exception: + # If organization lookup fails, skip the check + return + + if org_table is None: + return + + # Get max_budget from organization's budget table + org_max_budget: Optional[float] = None + if org_table.litellm_budget_table is not None: + org_max_budget = org_table.litellm_budget_table.max_budget + + # Only check if organization has a valid max_budget set + if org_max_budget is None or org_max_budget <= 0: + return + + # Check if organization spend exceeds max budget + if org_table.spend >= org_max_budget: + # Trigger budget alert + call_info = CallInfo( + token=valid_token.token, + spend=org_table.spend, + max_budget=org_max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=org_id, + event_group=Litellm_EntityType.ORGANIZATION, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="organization_budget", + user_info=call_info, + ) ) - if ( - org_table is not None - and org_table.spend >= valid_token.organization_max_budget - ): - # Trigger budget alert - call_info = CallInfo( - token=valid_token.token, - spend=org_table.spend, - max_budget=valid_token.organization_max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.ORGANIZATION, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="organization_budget", - user_info=call_info, - ) - ) - - raise litellm.BudgetExceededError( - current_cost=org_table.spend, - max_budget=valid_token.organization_max_budget, - message=f"Budget has been exceeded! Organization={valid_token.org_id} Current cost: {org_table.spend}, Max budget: {valid_token.organization_max_budget}", - ) + raise litellm.BudgetExceededError( + current_cost=org_table.spend, + max_budget=org_max_budget, + message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}", + ) async def _tag_max_budget_check( @@ -2601,4 +2627,4 @@ def _can_object_call_vector_stores( code=status.HTTP_401_UNAUTHORIZED, ) - return True + return True \ No newline at end of file diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 7a71af1da5c..797540deaa4 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -426,38 +426,65 @@ def get_key_model_rpm_limit( user_api_key_dict: UserAPIKeyAuth, ) -> Optional[Dict[str, int]]: """ - Get the model rpm limit for a given api key - - check key metadata - - check key model max budget - - check team metadata + Get the model rpm limit for a given api key. + + Priority order (returns first found): + 1. Key metadata (model_rpm_limit) + 2. Key model_max_budget (rpm_limit per model) + 3. Team metadata (model_rpm_limit) """ + # 1. Check key metadata first (takes priority) if user_api_key_dict.metadata: - if "model_rpm_limit" in user_api_key_dict.metadata: - return user_api_key_dict.metadata["model_rpm_limit"] - elif user_api_key_dict.model_max_budget: + result = user_api_key_dict.metadata.get("model_rpm_limit") + if result: + return result + + # 2. Check model_max_budget + if user_api_key_dict.model_max_budget: model_rpm_limit: Dict[str, Any] = {} for model, budget in user_api_key_dict.model_max_budget.items(): - if "rpm_limit" in budget and budget["rpm_limit"] is not None: + if isinstance(budget, dict) and budget.get("rpm_limit") is not None: model_rpm_limit[model] = budget["rpm_limit"] - return model_rpm_limit - elif user_api_key_dict.team_metadata: - if "model_rpm_limit" in user_api_key_dict.team_metadata: - return user_api_key_dict.team_metadata["model_rpm_limit"] + if model_rpm_limit: + return model_rpm_limit + + # 3. Fallback to team metadata + if user_api_key_dict.team_metadata: + return user_api_key_dict.team_metadata.get("model_rpm_limit") + return None def get_key_model_tpm_limit( user_api_key_dict: UserAPIKeyAuth, ) -> Optional[Dict[str, int]]: + """ + Get the model tpm limit for a given api key. + + Priority order (returns first found): + 1. Key metadata (model_tpm_limit) + 2. Key model_max_budget (tpm_limit per model) + 3. Team metadata (model_tpm_limit) + """ + # 1. Check key metadata first (takes priority) if user_api_key_dict.metadata: - if "model_tpm_limit" in user_api_key_dict.metadata: - return user_api_key_dict.metadata["model_tpm_limit"] - elif user_api_key_dict.model_max_budget: - if "tpm_limit" in user_api_key_dict.model_max_budget: - return user_api_key_dict.model_max_budget["tpm_limit"] - elif user_api_key_dict.team_metadata: - if "model_tpm_limit" in user_api_key_dict.team_metadata: - return user_api_key_dict.team_metadata["model_tpm_limit"] + result = user_api_key_dict.metadata.get("model_tpm_limit") + if result: + return result + + # 2. Check model_max_budget (iterate per-model like RPM does) + if user_api_key_dict.model_max_budget: + model_tpm_limit: Dict[str, Any] = {} + for model, budget in user_api_key_dict.model_max_budget.items(): + if isinstance(budget, dict) and budget.get("tpm_limit") is not None: + model_tpm_limit[model] = budget["tpm_limit"] + if model_tpm_limit: + return model_tpm_limit + + # 3. Fallback to team metadata + if user_api_key_dict.team_metadata: + return user_api_key_dict.team_metadata.get("model_tpm_limit") + return None @@ -469,7 +496,8 @@ def get_model_rate_limit_from_metadata( if getattr(user_api_key_dict, metadata_accessor_key): return getattr(user_api_key_dict, metadata_accessor_key).get(rate_limit_key) return None - + + def get_team_model_rpm_limit( user_api_key_dict: UserAPIKeyAuth, ) -> Optional[Dict[str, int]]: diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 8cc33ce6cdd..5be44f479b8 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -188,7 +188,7 @@ async def authenticate_user( # noqa: PLR0915 _user_row = cast( Optional[LiteLLM_UserTable], await prisma_client.db.litellm_usertable.find_first( - where={"user_email": {"equals": username}} + where={"user_email": {"equals": username, "mode": "insensitive"}} ), ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9b53d9a3a80..401aa7fd443 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -138,7 +138,7 @@ def _apply_budget_limits_to_end_user_params( ) -> None: """ Helper function to apply budget limits to end user parameters. - + Args: end_user_params: Dictionary to update with budget parameters budget_info: Budget table object containing limits @@ -146,16 +146,14 @@ def _apply_budget_limits_to_end_user_params( """ if budget_info.tpm_limit is not None: end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit - + if budget_info.rpm_limit is not None: end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit - + if budget_info.max_budget is not None: end_user_params["end_user_max_budget"] = budget_info.max_budget - - verbose_proxy_logger.debug( - f"Applied budget limits to end user {end_user_id}" - ) + + verbose_proxy_logger.debug(f"Applied budget limits to end user {end_user_id}") async def user_api_key_auth_websocket(websocket: WebSocket): @@ -170,12 +168,10 @@ async def user_api_key_auth_websocket(websocket: WebSocket): model = query_params.get("model") - async def return_body(): return _realtime_request_body(model) - - request.body = return_body # type: ignore + request.body = return_body # type: ignore authorization = websocket.headers.get("authorization") # If no Authorization header, try the api-key header @@ -586,7 +582,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_membership is not None else None ), - team_metadata=team_object.metadata if team_object is not None else None, + team_metadata=team_object.metadata + if team_object is not None + else None, ) # run through common checks _ = await common_checks( @@ -669,9 +667,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params["allowed_model_region"] = ( - _end_user_object.allowed_model_region - ) + end_user_params[ + "allowed_model_region" + ] = _end_user_object.allowed_model_region if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -753,7 +751,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, code=400, - param=api_key, + param=abbreviate_api_key(api_key=api_key), ) valid_token = update_valid_token_with_end_user_params( valid_token=valid_token, end_user_params=end_user_params @@ -994,7 +992,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Check 3. Check if user is in their team budget if valid_token.team_member_spend is not None: - if prisma_client is not None: _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" @@ -1055,7 +1052,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, code=400, - param=api_key, + param=abbreviate_api_key(api_key=api_key), ) # Check 4. Token Spend is under budget @@ -1216,8 +1213,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) - - @tracer.wrap() async def user_api_key_auth( request: Request, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e9ce10ccf31..61bffde3aca 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -17,7 +17,7 @@ from typing import ( import httpx import orjson from fastapi import HTTPException, Request, status -from fastapi.responses import Response, StreamingResponse +from fastapi.responses import JSONResponse, Response, StreamingResponse import litellm from litellm._logging import verbose_proxy_logger @@ -96,16 +96,55 @@ async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional return None -async def create_streaming_response( +def _extract_error_from_sse_chunk(event_line: Union[str, bytes]) -> dict: + """ + Extract error dictionary from SSE format chunk. + + Args: + event_line: SSE format event line, e.g. "data: {"error": {...}}\n\n" + + Returns: + Error dictionary in OpenAI API format + """ + event_line = ( + event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line + ) + + # Default error format + default_error = { + "message": "Unknown error", + "type": "internal_server_error", + "param": None, + "code": "500", + } + + if event_line.startswith("data: "): + json_str = event_line[len("data: ") :].strip() + if not json_str or json_str == "[DONE]": + return default_error + + try: + data = orjson.loads(json_str) + if isinstance(data, dict) and "error" in data: + error_obj = data["error"] + if isinstance(error_obj, dict): + return error_obj + except (orjson.JSONDecodeError, json.JSONDecodeError): + pass + + return default_error + + +async def create_response( generator: AsyncGenerator[str, None], media_type: str, headers: dict, default_status_code: int = status.HTTP_200_OK, -) -> StreamingResponse: +) -> Union[StreamingResponse, JSONResponse]: """ - Creates a StreamingResponse by inspecting the first chunk for an error code. - The entire original generator content is streamed, but the HTTP status code - of the response is set based on the first chunk if it's a recognized error. + Create streaming response, checking if the first chunk is an error. + If the first chunk is an error, return a standard JSON error response. + Otherwise, return StreamingResponse and stream all content. """ first_chunk_value: Optional[str] = None final_status_code = default_status_code @@ -124,9 +163,27 @@ async def create_streaming_response( first_chunk_value ) if error_code_from_chunk is not None: + # First chunk is an error, stream hasn't really started yet + # Should return standard JSON error response instead of SSE format final_status_code = error_code_from_chunk verbose_proxy_logger.debug( - f"Error detected in first stream chunk. Status code set to: {final_status_code}" + f"Error detected in first stream chunk. Returning JSON error response with status code: {final_status_code}" + ) + + # Parse error content + error_dict = _extract_error_from_sse_chunk(first_chunk_value) + + # Consume and close generator (avoid resource leak) + try: + await generator.aclose() + except Exception: + pass + + # Return JSON format error response + return JSONResponse( + status_code=final_status_code, + content={"error": error_dict}, + headers=headers, ) except Exception as e: verbose_proxy_logger.debug(f"Error parsing first chunk value: {e}") @@ -237,7 +294,11 @@ class ProxyBaseLLMRequestProcessing: if response_cost is not None: try: # Convert response_cost to float if it's a string - cost_value = float(response_cost) if isinstance(response_cost, str) else response_cost + cost_value = ( + float(response_cost) + if isinstance(response_cost, str) + else response_cost + ) if cost_value > 0: updated_spend = current_spend + cost_value except (ValueError, TypeError): @@ -376,6 +437,16 @@ class ProxyBaseLLMRequestProcessing: ) -> Tuple[dict, LiteLLMLoggingObj]: start_time = datetime.now() # start before calling guardrail hooks + # Calculate request queue time if arrival_time is available + # Use start_time.timestamp() to avoid extra time.time() call for better performance + proxy_server_request = self.data.get("proxy_server_request", {}) + arrival_time = proxy_server_request.get("arrival_time") + queue_time_seconds = None + if arrival_time is not None: + # Convert start_time (datetime) to timestamp for calculation + processing_start_time = start_time.timestamp() + queue_time_seconds = processing_start_time - arrival_time + self.data = await add_litellm_data_to_request( data=self.data, request=request, @@ -385,6 +456,19 @@ class ProxyBaseLLMRequestProcessing: proxy_config=proxy_config, ) + # Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved + if queue_time_seconds is not None: + from litellm.proxy.litellm_pre_call_utils import _get_metadata_variable_name + + _metadata_variable_name = _get_metadata_variable_name(request) + if _metadata_variable_name not in self.data: + self.data[_metadata_variable_name] = {} + if not isinstance(self.data[_metadata_variable_name], dict): + self.data[_metadata_variable_name] = {} + self.data[_metadata_variable_name][ + "queue_time_seconds" + ] = queue_time_seconds + self.data["model"] = ( general_settings.get("completion_model", None) # server default or user_model # model name passed via cli args @@ -670,7 +754,7 @@ class ProxyBaseLLMRequestProcessing: proxy_logging_obj=proxy_logging_obj, ) ) - return await create_streaming_response( + return await create_response( generator=selected_data_generator, media_type="text/event-stream", headers=custom_headers, @@ -681,7 +765,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, request_data=self.data, ) - return await create_streaming_response( + return await create_response( generator=selected_data_generator, media_type="text/event-stream", headers=custom_headers, @@ -923,11 +1007,11 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def _get_pre_call_type( route_type: Literal["acompletion", "aembedding", "aresponses", "allm_passthrough_route"], - ) -> Literal["completion", "embeddings", "responses", "allm_passthrough_route"]: + ) -> Literal["completion", "embedding", "responses", "allm_passthrough_route"]: if route_type == "acompletion": return "completion" elif route_type == "aembedding": - return "embeddings" + return "embedding" elif route_type == "aresponses": return "responses" elif route_type == "allm_passthrough_route": @@ -1178,9 +1262,9 @@ class ProxyBaseLLMRequestProcessing: # Add cache-related fields to **params (handled by Usage.__init__) if cache_creation_input_tokens is not None: - usage_kwargs["cache_creation_input_tokens"] = ( - cache_creation_input_tokens - ) + usage_kwargs[ + "cache_creation_input_tokens" + ] = cache_creation_input_tokens if cache_read_input_tokens is not None: usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 406ddceabf5..c9c0cfe8f68 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -17,50 +17,141 @@ from litellm.secret_managers.main import str_to_bool class PrismaWrapper: + """ + Wrapper around Prisma client that handles RDS IAM token authentication. + + When iam_token_db_auth is enabled, this wrapper: + 1. Proactively refreshes IAM tokens before they expire (background task) + 2. Falls back to synchronous refresh if a token is found expired + 3. Uses proper locking to prevent race conditions during reconnection + + RDS IAM tokens are valid for 15 minutes. This wrapper refreshes them + 3 minutes before expiration to ensure uninterrupted database connectivity. + """ + + # Buffer time in seconds before token expiration to trigger refresh + # Refresh 3 minutes (180 seconds) before the token expires + TOKEN_REFRESH_BUFFER_SECONDS = 180 + + # Fallback refresh interval if token parsing fails (10 minutes) + FALLBACK_REFRESH_INTERVAL_SECONDS = 600 + def __init__(self, original_prisma: Any, iam_token_db_auth: bool): self._original_prisma = original_prisma self.iam_token_db_auth = iam_token_db_auth + # Background token refresh task management + self._token_refresh_task: Optional[asyncio.Task] = None + self._reconnection_lock = asyncio.Lock() + self._last_refresh_time: Optional[datetime] = None + + def _extract_token_from_db_url(self, db_url: Optional[str]) -> Optional[str]: + """ + Extract the token (password) from the DATABASE_URL. + + The token contains the AWS signature with X-Amz-Date and X-Amz-Expires parameters. + + Important: We must parse the URL while it's still encoded to preserve structure, + then decode the password portion. Otherwise the '?' in the token breaks URL parsing. + """ + if db_url is None: + return None + try: + # Parse URL while still encoded to preserve structure + parsed = urllib.parse.urlparse(db_url) + if parsed.password: + # Now decode just the password/token + return urllib.parse.unquote(parsed.password) + return None + except Exception: + return None + + def _parse_token_expiration(self, token: Optional[str]) -> Optional[datetime]: + """ + Parse the token to extract its expiration time. + + Returns the datetime when the token expires, or None if parsing fails. + """ + if token is None: + return None + + try: + # Token format: ...?X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=900&... + if "?" not in token: + return None + + query_string = token.split("?", 1)[1] + params = urllib.parse.parse_qs(query_string) + + expires_str = params.get("X-Amz-Expires", [None])[0] + date_str = params.get("X-Amz-Date", [None])[0] + + if not expires_str or not date_str: + return None + + token_created = datetime.strptime(date_str, "%Y%m%dT%H%M%SZ") + expires_in = int(expires_str) + + return token_created + timedelta(seconds=expires_in) + except Exception as e: + verbose_proxy_logger.debug(f"Failed to parse token expiration: {e}") + return None + + def _calculate_seconds_until_refresh(self) -> float: + """ + Calculate exactly how many seconds until we need to refresh the token. + + Uses precise timing: sleeps until (token_expiration - buffer_seconds). + For a 15-minute (900s) token with 180s buffer, this returns ~720s (12 min). + + Returns: + Number of seconds to sleep before the next refresh. + Returns 0 if token should be refreshed immediately. + Returns FALLBACK_REFRESH_INTERVAL_SECONDS if parsing fails. + """ + db_url = os.getenv("DATABASE_URL") + token = self._extract_token_from_db_url(db_url) + expiration_time = self._parse_token_expiration(token) + + if expiration_time is None: + # If we can't parse the token, use fallback interval + verbose_proxy_logger.debug( + f"Could not parse token expiration, using fallback interval of " + f"{self.FALLBACK_REFRESH_INTERVAL_SECONDS}s" + ) + return self.FALLBACK_REFRESH_INTERVAL_SECONDS + + # Calculate when we should refresh (expiration - buffer) + refresh_at = expiration_time - timedelta( + seconds=self.TOKEN_REFRESH_BUFFER_SECONDS + ) + + # How long until refresh time? + now = datetime.utcnow() + seconds_until_refresh = (refresh_at - now).total_seconds() + + # If already past refresh time, return 0 (refresh immediately) + return max(0, seconds_until_refresh) + def is_token_expired(self, token_url: Optional[str]) -> bool: + """Check if the token in the given URL is expired.""" if token_url is None: return True - # Decode the token URL to handle URL-encoded characters - decoded_url = urllib.parse.unquote(token_url) - # Parse the token URL - parsed_url = urllib.parse.urlparse(decoded_url) + token = self._extract_token_from_db_url(token_url) + expiration_time = self._parse_token_expiration(token) - # Parse the query parameters from the path component (if they exist there) - query_params = urllib.parse.parse_qs(parsed_url.query) + if expiration_time is None: + # If we can't parse the token, assume it's expired to trigger refresh + verbose_proxy_logger.debug( + "Could not parse token expiration, treating as expired" + ) + return True - # Get expiration time from the query parameters - expires = query_params.get("X-Amz-Expires", [None])[0] - if expires is None: - raise ValueError("X-Amz-Expires parameter is missing or invalid.") - - expires_int = int(expires) - - # Get the token's creation time from the X-Amz-Date parameter - token_time_str = query_params.get("X-Amz-Date", [""])[0] - if not token_time_str: - raise ValueError("X-Amz-Date parameter is missing or invalid.") - - # Ensure the token time string is parsed correctly - try: - token_time = datetime.strptime(token_time_str, "%Y%m%dT%H%M%SZ") - except ValueError as e: - raise ValueError(f"Invalid X-Amz-Date format: {e}") - - # Calculate the expiration time - expiration_time = token_time + timedelta(seconds=expires_int) - - # Current time in UTC - current_time = datetime.utcnow() - - # Check if the token is expired - return current_time > expiration_time + return datetime.utcnow() > expiration_time def get_rds_iam_token(self) -> Optional[str]: + """Generate a new RDS IAM token and update DATABASE_URL.""" if self.iam_token_db_auth: from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token @@ -74,7 +165,6 @@ class PrismaWrapper: db_host=db_host, db_port=db_port, db_user=db_user ) - # print(f"token: {token}") _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" if db_schema: _db_url += f"?schema={db_schema}" @@ -86,6 +176,7 @@ class PrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, http_client: Optional[Any] = None ): + """Disconnect and reconnect the Prisma client with a new database URL.""" from prisma import Prisma # type: ignore try: @@ -100,21 +191,159 @@ class PrismaWrapper: await self._original_prisma.connect() + async def start_token_refresh_task(self) -> None: + """ + Start the background token refresh task. + + This task proactively refreshes RDS IAM tokens before they expire, + preventing connection failures. Should be called after the initial + Prisma client connection is established. + """ + if not self.iam_token_db_auth: + verbose_proxy_logger.debug( + "IAM token auth not enabled, skipping token refresh task" + ) + return + + if self._token_refresh_task is not None: + verbose_proxy_logger.debug("Token refresh task already running") + return + + self._token_refresh_task = asyncio.create_task(self._token_refresh_loop()) + verbose_proxy_logger.info( + "Started RDS IAM token proactive refresh background task" + ) + + async def stop_token_refresh_task(self) -> None: + """ + Stop the background token refresh task gracefully. + + Should be called during application shutdown to clean up resources. + """ + if self._token_refresh_task is None: + return + + self._token_refresh_task.cancel() + try: + await self._token_refresh_task + except asyncio.CancelledError: + pass + self._token_refresh_task = None + verbose_proxy_logger.info("Stopped RDS IAM token refresh background task") + + async def _token_refresh_loop(self) -> None: + """ + Background loop that proactively refreshes RDS IAM tokens before expiration. + + Uses precise timing: calculates the exact sleep duration until the token + needs to be refreshed (expiration - 3 minute buffer), then refreshes. + This is more efficient than polling, requiring only 1 wake-up per token cycle. + """ + verbose_proxy_logger.info( + f"RDS IAM token refresh loop started. " + f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration." + ) + + while True: + try: + # Calculate exactly how long to sleep until next refresh + sleep_seconds = self._calculate_seconds_until_refresh() + + if sleep_seconds > 0: + verbose_proxy_logger.info( + f"RDS IAM token refresh scheduled in {sleep_seconds:.0f} seconds " + f"({sleep_seconds / 60:.1f} minutes)" + ) + await asyncio.sleep(sleep_seconds) + + # Refresh the token + verbose_proxy_logger.info("Proactively refreshing RDS IAM token...") + await self._safe_refresh_token() + + except asyncio.CancelledError: + verbose_proxy_logger.info("RDS IAM token refresh loop cancelled") + break + except Exception as e: + verbose_proxy_logger.error( + f"Error in RDS IAM token refresh loop: {e}. " + f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..." + ) + # On error, wait before retrying to avoid tight error loops + try: + await asyncio.sleep(self.FALLBACK_REFRESH_INTERVAL_SECONDS) + except asyncio.CancelledError: + break + + async def _safe_refresh_token(self) -> None: + """ + Refresh the RDS IAM token with proper locking to prevent race conditions. + + Uses an asyncio lock to ensure only one refresh operation happens at a time, + preventing multiple concurrent reconnection attempts. + """ + async with self._reconnection_lock: + new_db_url = self.get_rds_iam_token() + if new_db_url: + await self.recreate_prisma_client(new_db_url) + self._last_refresh_time = datetime.utcnow() + verbose_proxy_logger.info( + "RDS IAM token refreshed successfully. New token valid for ~15 minutes." + ) + else: + verbose_proxy_logger.error( + "Failed to generate new RDS IAM token during proactive refresh" + ) + def __getattr__(self, name: str): + """ + Proxy attribute access to the underlying Prisma client. + + If IAM token auth is enabled and the token is expired, this method + provides a synchronous fallback to refresh the token. However, this + should rarely be needed since the background task proactively refreshes + tokens before they expire. + + FIXED: Now properly waits for reconnection to complete before returning, + instead of the previous fire-and-forget pattern that caused the bug. + """ original_attr = getattr(self._original_prisma, name) + if self.iam_token_db_auth: db_url = os.getenv("DATABASE_URL") - if self.is_token_expired(db_url): - db_url = self.get_rds_iam_token() - loop = asyncio.get_event_loop() - if db_url: + # Check if token is expired (should be rare if background task is running) + if self.is_token_expired(db_url): + verbose_proxy_logger.warning( + "RDS IAM token expired in __getattr__ - proactive refresh may have failed. " + "Triggering synchronous fallback refresh..." + ) + + new_db_url = self.get_rds_iam_token() + if new_db_url: + loop = asyncio.get_event_loop() + if loop.is_running(): - asyncio.run_coroutine_threadsafe( - self.recreate_prisma_client(db_url), loop + # FIXED: Actually wait for the reconnection to complete! + # The previous code used fire-and-forget which caused the bug. + future = asyncio.run_coroutine_threadsafe( + self.recreate_prisma_client(new_db_url), loop ) + try: + # Wait up to 30 seconds for reconnection + future.result(timeout=30) + verbose_proxy_logger.info( + "Synchronous token refresh completed successfully" + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to refresh token synchronously: {e}" + ) + raise else: - asyncio.run(self.recreate_prisma_client(db_url)) + asyncio.run(self.recreate_prisma_client(new_db_url)) + + # Get the NEW attribute after reconnection + original_attr = getattr(self._original_prisma, name) else: raise ValueError("Failed to get RDS IAM token") diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index a6971b49f3b..87da11efad0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -5,6 +5,7 @@ # +-------------------------------------------------------------+ # Qualifire - Evaluate LLM outputs for quality, safety, and reliability +import json import os from typing import Any, Dict, List, Literal, Optional, Type @@ -15,12 +16,17 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs GUARDRAIL_NAME = "qualifire" +DEFAULT_QUALIFIRE_API_BASE = "https://proxy.qualifire.ai" class QualifireGuardrail(CustomGuardrail): @@ -44,7 +50,7 @@ class QualifireGuardrail(CustomGuardrail): Args: api_key: API key for Qualifire (or use QUALIFIRE_API_KEY env var) - api_base: Optional custom API base URL + api_base: Optional custom API base URL (defaults to https://api.qualifire.ai) evaluation_id: Pre-configured evaluation ID from Qualifire dashboard prompt_injections: Enable prompt injection detection (default if no other checks) hallucinations_check: Enable hallucination detection @@ -64,6 +70,7 @@ class QualifireGuardrail(CustomGuardrail): api_base or get_secret_str("QUALIFIRE_BASE_URL") or os.environ.get("QUALIFIRE_BASE_URL") + or DEFAULT_QUALIFIRE_API_BASE ) self.evaluation_id = evaluation_id self.prompt_injections = prompt_injections @@ -79,7 +86,11 @@ class QualifireGuardrail(CustomGuardrail): if not self._has_any_check_enabled() and not self.evaluation_id: self.prompt_injections = True - self._client = None + # Initialize async HTTP client for direct API calls + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + super().__init__(**kwargs) def _has_any_check_enabled(self) -> bool: @@ -96,43 +107,22 @@ class QualifireGuardrail(CustomGuardrail): ] ) - def _get_client(self): - """Lazy initialization of Qualifire client.""" - if self._client is None: - try: - from qualifire.client import Client - except ImportError: - raise ImportError( - "qualifire package is required for QualifireGuardrail. " - "Install it with: pip install qualifire" - ) - - client_kwargs: Dict[str, Any] = {} - if self.qualifire_api_key: - client_kwargs["api_key"] = self.qualifire_api_key - if self.qualifire_api_base: - client_kwargs["base_url"] = self.qualifire_api_base - - self._client = Client(**client_kwargs) - - return self._client - - def _convert_messages_to_qualifire_format( + def _convert_messages_to_api_format( self, messages: List[AllMessageValues] - ) -> List[Any]: + ) -> List[Dict[str, Any]]: """ - Convert LiteLLM messages to Qualifire's LLMMessage format. + Convert LiteLLM messages to Qualifire API format. Supports tool calls for tool_selection_quality_check. - """ - try: - from qualifire.types import LLMMessage, LLMToolCall - except ImportError: - raise ImportError( - "qualifire package is required for QualifireGuardrail. " - "Install it with: pip install qualifire" - ) - qualifire_messages = [] + Returns a list of dicts matching the API's ModelInvocationCanonicalMessage schema: + { + "role": "user" | "assistant" | "system" | "tool", + "content": "...", + "tool_call_id": "...", # optional + "tool_calls": [{"id": "...", "name": "...", "arguments": {...}}] # optional + } + """ + api_messages = [] for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") @@ -147,42 +137,86 @@ class QualifireGuardrail(CustomGuardrail): text_parts.append(part) content = "\n".join(text_parts) - llm_message_kwargs: Dict[str, Any] = { + api_message: Dict[str, Any] = { "role": role, "content": content if isinstance(content, str) else str(content), } + # Handle tool_call_id for tool response messages + tool_call_id = msg.get("tool_call_id") + if tool_call_id: + api_message["tool_call_id"] = tool_call_id + # Handle tool calls if present tool_calls = msg.get("tool_calls") if tool_calls and isinstance(tool_calls, list): - qualifire_tool_calls = [] + api_tool_calls = [] for tc in tool_calls: if isinstance(tc, dict): function_info = tc.get("function", {}) # Arguments can be a string (JSON) or dict args = function_info.get("arguments", {}) if isinstance(args, str): - import json - try: args = json.loads(args) except json.JSONDecodeError: args = {} - qualifire_tool_calls.append( - LLMToolCall( - id=tc.get("id") or "", - name=function_info.get("name") or "", - arguments=args if isinstance(args, dict) else {}, - ) + api_tool_calls.append( + { + "id": tc.get("id") or "", + "name": function_info.get("name") or "", + "arguments": args if isinstance(args, dict) else {}, + } ) - if qualifire_tool_calls: - llm_message_kwargs["tool_calls"] = qualifire_tool_calls + if api_tool_calls: + api_message["tool_calls"] = api_tool_calls - qualifire_messages.append(LLMMessage(**llm_message_kwargs)) + api_messages.append(api_message) - return qualifire_messages + return api_messages - def _check_if_flagged(self, result: Any) -> bool: + def _convert_tools_to_api_format( + self, tools: Optional[List[Any]] + ) -> Optional[List[Dict[str, Any]]]: + """ + Convert OpenAI-format tools to Qualifire API format. + + Returns a list of dicts matching the API's ModelInvocationToolDefinition schema: + { + "name": "...", + "description": "...", + "parameters": {...} + } + """ + if not tools: + return None + + api_tools = [] + for tool in tools: + if isinstance(tool, dict): + # Handle OpenAI function tool format + if tool.get("type") == "function": + function_def = tool.get("function", {}) + api_tools.append( + { + "name": function_def.get("name", ""), + "description": function_def.get("description", ""), + "parameters": function_def.get("parameters", {}), + } + ) + # Handle direct tool format + elif "name" in tool: + api_tools.append( + { + "name": tool.get("name", ""), + "description": tool.get("description", ""), + "parameters": tool.get("parameters", {}), + } + ) + + return api_tools if api_tools else None + + def _check_if_flagged(self, result: Dict[str, Any]) -> bool: """ Check if the Qualifire evaluation result indicates flagged content. @@ -190,65 +224,53 @@ class QualifireGuardrail(CustomGuardrail): A high score (close to 100) indicates GOOD content, low score indicates problems. """ # Check evaluation results for any flagged items - evaluation_results = getattr(result, "evaluationResults", None) or [] - if isinstance(result, dict): - evaluation_results = result.get("evaluationResults", []) or [] + evaluation_results = result.get("evaluationResults", []) or [] for eval_result in evaluation_results: - results: List[Any] = [] - if isinstance(eval_result, dict): - results = eval_result.get("results", []) or [] - else: - results = getattr(eval_result, "results", []) or [] - + results = eval_result.get("results", []) or [] for r in results: - flagged = ( - r.get("flagged") - if isinstance(r, dict) - else getattr(r, "flagged", False) - ) - if flagged: + if r.get("flagged"): return True return False - def _build_evaluate_kwargs( + def _build_evaluate_payload( self, - qualifire_messages: List[Any], + api_messages: List[Dict[str, Any]], output: Optional[str], assertions: Optional[List[str]], - available_tools: Optional[List[Any]], + available_tools: Optional[List[Dict[str, Any]]], ) -> Dict[str, Any]: - """Build kwargs dictionary for the evaluate call.""" - kwargs: Dict[str, Any] = {"messages": qualifire_messages} + """Build payload dictionary for the /api/evaluation/evaluate endpoint.""" + payload: Dict[str, Any] = {"messages": api_messages} if output is not None: - kwargs["output"] = output + payload["output"] = output # Add enabled checks if self.prompt_injections: - kwargs["prompt_injections"] = True + payload["prompt_injections"] = True if self.hallucinations_check: - kwargs["hallucinations_check"] = True + payload["hallucinations_check"] = True if self.grounding_check: - kwargs["grounding_check"] = True + payload["grounding_check"] = True if self.pii_check: - kwargs["pii_check"] = True + payload["pii_check"] = True if self.content_moderation_check: - kwargs["content_moderation_check"] = True + payload["content_moderation_check"] = True if self.tool_selection_quality_check: # Only enable tool_selection_quality_check if available_tools is provided if available_tools: - kwargs["tool_selection_quality_check"] = True - kwargs["available_tools"] = available_tools + payload["tool_selection_quality_check"] = True + payload["available_tools"] = available_tools else: verbose_proxy_logger.debug( "Qualifire Guardrail: tool_selection_quality_check enabled but no available_tools provided, skipping this check" ) if assertions: - kwargs["assertions"] = assertions + payload["assertions"] = assertions - return kwargs + return payload async def _run_qualifire_check( self, @@ -274,11 +296,17 @@ class QualifireGuardrail(CustomGuardrail): assertions = dynamic_params.get("assertions") or self.assertions on_flagged = dynamic_params.get("on_flagged") or self.on_flagged - try: - client = self._get_client() - qualifire_messages = self._convert_messages_to_qualifire_format(messages) + # Prepare headers + headers = { + "X-Qualifire-API-Key": self.qualifire_api_key or "", + "Content-Type": "application/json", + } - # Use invoke_evaluation if evaluation_id is provided + try: + # Convert messages to API format + api_messages = self._convert_messages_to_api_format(messages) + + # Use invoke endpoint if evaluation_id is provided if evaluation_id: # For invoke_evaluation, we need to extract input/output input_text = "" @@ -291,25 +319,47 @@ class QualifireGuardrail(CustomGuardrail): input_text = content break - result = client.invoke_evaluation( - evaluation_id=evaluation_id, - input=input_text, - output=output or "", - ) + payload = { + "evaluation_id": evaluation_id, + "input": input_text, + "output": output or "", + "messages": api_messages, + } + + # Convert tools if provided + api_tools = self._convert_tools_to_api_format(available_tools) + if api_tools: + payload["available_tools"] = api_tools + + url = f"{self.qualifire_api_base}/api/evaluation/invoke" else: - # Use evaluate with individual checks - kwargs = self._build_evaluate_kwargs( - qualifire_messages=qualifire_messages, + # Use evaluate endpoint with individual checks + api_tools = self._convert_tools_to_api_format(available_tools) + payload = self._build_evaluate_payload( + api_messages=api_messages, output=output, assertions=assertions, - available_tools=available_tools, + available_tools=api_tools, ) - result = client.evaluate(**kwargs) + url = f"{self.qualifire_api_base}/api/evaluation/evaluate" - # Convert result to dict for logging + verbose_proxy_logger.debug( + f"Qualifire Guardrail: Making request to {url}" + ) + + # Make the API request + response = await self.async_handler.post( + url=url, + headers=headers, + json=payload, + ) + response.raise_for_status() + result = response.json() + + # Extract response info for logging qualifire_response = { - "score": getattr(result, "score", None), - "status": getattr(result, "status", None), + "score": result.get("score"), + "status": result.get("status"), } verbose_proxy_logger.debug( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c416527990e..4d17cca22ad 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -167,7 +167,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.token_increment_script = None self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) - + # Batch rate limiter (lazy loaded) self._batch_rate_limiter: Optional[Any] = None @@ -1013,7 +1013,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) # Fail safe: enforce limits if we can't check return True - + def get_rate_limiter_for_call_type(self, call_type: str) -> Optional[Any]: """Get the rate limiter for the call type.""" if call_type == "acreate_batch": @@ -1095,9 +1095,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): now = self._get_current_time().timestamp() reset_time = now + self.window_size - reset_time_formatted = datetime.fromtimestamp( - reset_time - ).strftime("%Y-%m-%d %H:%M:%S UTC") + reset_time_formatted = datetime.fromtimestamp(reset_time).strftime( + "%Y-%m-%d %H:%M:%S UTC" + ) remaining_display = max(0, status["limit_remaining"]) rate_limit_type = status["rate_limit_type"] @@ -1137,7 +1137,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Check if the call type has a specific rate limiter # eg. for Batch APIs we need to use the batch rate limiter to read the input file and count the tokens and requests ######################################################### - call_type_specific_rate_limiter = self.get_rate_limiter_for_call_type(call_type=call_type) + call_type_specific_rate_limiter = self.get_rate_limiter_for_call_type( + call_type=call_type + ) if call_type_specific_rate_limiter: return await call_type_specific_rate_limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -1233,26 +1235,58 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return pipeline_operations - def _get_total_tokens_from_usage(self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"]) -> int: - # Get total tokens from response + def _get_total_tokens_from_usage( + self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"] + ) -> int: + """ + Get total tokens from response usage for rate limiting. + + For 'input' and 'total' rate limit types, cached tokens are excluded + because providers like AWS Bedrock don't count cached tokens toward + rate limits. This aligns LiteLLM's TPM calculation with provider behavior. + """ total_tokens = 0 - # spot fix for /responses api + cached_tokens = 0 + if usage: if isinstance(usage, Usage): if rate_limit_type == "output": - total_tokens = usage.completion_tokens + total_tokens = usage.completion_tokens or 0 elif rate_limit_type == "input": - total_tokens = usage.prompt_tokens + total_tokens = usage.prompt_tokens or 0 elif rate_limit_type == "total": - total_tokens = usage.total_tokens + total_tokens = usage.total_tokens or 0 + + # Get cached tokens to exclude from input/total + if rate_limit_type in ("input", "total"): + if ( + hasattr(usage, "prompt_tokens_details") + and usage.prompt_tokens_details is not None + ): + cached_tokens = ( + getattr(usage.prompt_tokens_details, "cached_tokens", 0) + or 0 + ) + elif isinstance(usage, dict): - # Responses API usage comes as a dict in ResponsesAPIResponse + # Responses API usage comes as a dict if rate_limit_type == "output": - total_tokens = usage.get("completion_tokens", 0) + total_tokens = usage.get("completion_tokens", 0) or 0 elif rate_limit_type == "input": - total_tokens = usage.get("prompt_tokens", 0) + total_tokens = usage.get("prompt_tokens", 0) or 0 elif rate_limit_type == "total": - total_tokens = usage.get("total_tokens", 0) + total_tokens = usage.get("total_tokens", 0) or 0 + + # Get cached tokens from dict + if rate_limit_type in ("input", "total"): + prompt_details = usage.get("prompt_tokens_details") or {} + if isinstance(prompt_details, dict): + cached_tokens = prompt_details.get("cached_tokens", 0) or 0 + + # Subtract cached tokens for input/total (providers don't count them) + if cached_tokens > 0: + total_tokens = max(0, total_tokens - cached_tokens) + return total_tokens async def _execute_token_increment_script( @@ -1336,6 +1370,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def get_rate_limit_type(self) -> Literal["output", "input", "total"]: from litellm.proxy.proxy_server import general_settings + specified_rate_limit_type = general_settings.get( "token_rate_limit_type", "total" ) @@ -1381,9 +1416,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): user_api_key_organization_id = standard_logging_metadata.get( "user_api_key_org_id" ) - user_api_key_end_user_id = kwargs.get("user") or standard_logging_metadata.get( - "user_api_key_end_user_id" - ) + user_api_key_end_user_id = kwargs.get( + "user" + ) or standard_logging_metadata.get("user_api_key_end_user_id") model_group = get_model_group_from_litellm_kwargs(kwargs) # Get total tokens from response @@ -1393,7 +1428,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): response_obj, BaseLiteLLMOpenAIResponseObject ): _usage = getattr(response_obj, "usage", None) - total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type) + total_tokens = self._get_total_tokens_from_usage( + usage=_usage, rate_limit_type=rate_limit_type + ) # Create pipeline operations for TPM increments pipeline_operations: List[RedisPipelineIncrementOperation] = [] @@ -1518,9 +1555,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): from litellm.types.caching import RedisPipelineIncrementOperation try: - litellm_parent_otel_span: Union[ - Span, None - ] = _get_parent_otel_span_from_kwargs(kwargs) + litellm_parent_otel_span: Union[Span, None] = ( + _get_parent_otel_span_from_kwargs(kwargs) + ) # Get metadata from standard_logging_object - this correctly handles both # 'metadata' and 'litellm_metadata' fields from litellm_params standard_logging_object = kwargs.get("standard_logging_object") or {} @@ -1555,7 +1592,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): f"Error in rate limit failure event: {str(e)}" ) - async def async_post_call_success_hook( self, data: dict, user_api_key_dict: UserAPIKeyAuth, response ): diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 5b5723efc3d..3f844f21eb0 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -161,7 +161,6 @@ class KeyAndTeamLoggingSettings: @staticmethod def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): - if ( user_api_key_dict.team_metadata is not None and "logging" in user_api_key_dict.team_metadata @@ -174,12 +173,12 @@ def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> Optional[TeamCallbackMetadata]: callback_settings_obj: Optional[TeamCallbackMetadata] = None - key_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) - ) - team_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) - ) + key_dynamic_logging_settings: Optional[ + dict + ] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + team_dynamic_logging_settings: Optional[ + dict + ] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) ######################################################################################### # Key-based callbacks ######################################################################################### @@ -462,7 +461,6 @@ class LiteLLMProxyRequestSetup: team_id=user_api_key_dict.team_id, ) # handles aliases, wildcards, etc. ): - _headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call( headers, user_api_key_dict ) @@ -663,11 +661,11 @@ class LiteLLMProxyRequestSetup: ## KEY-LEVEL SPEND LOGS / TAGS if "tags" in key_metadata and key_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = ( - LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=key_metadata["tags"], - ) + data[_metadata_variable_name][ + "tags" + ] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=key_metadata["tags"], ) if "disable_global_guardrails" in key_metadata and isinstance( key_metadata["disable_global_guardrails"], bool @@ -815,11 +813,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915 # Init - Proxy Server Request # we do this as soon as entering so we track the original request ########################################################## + # Track arrival time for queue time metric + arrival_time = time.time() data["proxy_server_request"] = { "url": str(request.url), "method": request.method, "headers": _headers, "body": copy.copy(data), # use copy instead of deepcopy + "arrival_time": arrival_time, # Track when request arrived at proxy } safe_add_api_version_from_query_params(data, request) @@ -930,9 +931,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data[_metadata_variable_name]["litellm_api_version"] = version if general_settings is not None: - data[_metadata_variable_name]["global_max_parallel_requests"] = ( - general_settings.get("global_max_parallel_requests", None) - ) + data[_metadata_variable_name][ + "global_max_parallel_requests" + ] = general_settings.get("global_max_parallel_requests", None) ### KEY-LEVEL Controls key_metadata = user_api_key_dict.metadata diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 47793c8fc8e..d8816df010a 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -36,12 +36,19 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._experimental.mcp_server.utils import ( + get_server_prefix, validate_and_normalize_mcp_server_payload, ) router = APIRouter(prefix="/v1/mcp", tags=["mcp"]) + MCP_AVAILABLE: bool = True + TEMPORARY_MCP_SERVER_TTL_SECONDS = 300 +DEFAULT_MCP_REGISTRY_VERSION = "1.0.0" +LITELLM_MCP_SERVER_NAME = "litellm-mcp-server" +LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM" + try: importlib.import_module("mcp") except ImportError as e: @@ -57,6 +64,7 @@ if MCP_AVAILABLE: update_mcp_server, ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, authorize_with_server, exchange_token_with_server, register_client_with_server, @@ -89,6 +97,66 @@ if MCP_AVAILABLE: server: MCPServer expires_at: datetime + def _is_public_registry_enabled() -> bool: + from litellm.proxy.proxy_server import ( + general_settings as proxy_general_settings, + ) + + return bool(proxy_general_settings.get("enable_mcp_registry")) + + def _build_registry_remote_url(base_url: str, path: str) -> str: + normalized_base = base_url.rstrip("/") + normalized_path = path if path.startswith("/") else f"/{path}" + return f"{normalized_base}{normalized_path}" + + def _build_mcp_registry_server_name(server: MCPServer) -> str: + if server.alias: + return server.alias + if server.server_name: + return server.server_name + return server.server_id + + def _build_mcp_registry_entry_for_server( + server: MCPServer, base_url: str + ) -> Dict[str, Any]: + server_name = _build_mcp_registry_server_name(server) + title = server_name + description = server_name + version = DEFAULT_MCP_REGISTRY_VERSION + + server_prefix = get_server_prefix(server) + if not server_prefix: + raise ValueError("MCP server prefix is missing") + remote_url = _build_registry_remote_url(base_url, f"/{server_prefix}/mcp") + + return { + "name": server_name, + "title": title, + "description": description, + "version": version, + "remotes": [ + { + "type": "streamable-http", + "url": remote_url, + } + ], + } + + def _build_builtin_registry_entry(base_url: str) -> Dict[str, Any]: + remote_url = _build_registry_remote_url(base_url, "/mcp") + return { + "name": LITELLM_MCP_SERVER_NAME, + "title": LITELLM_MCP_SERVER_NAME, + "description": LITELLM_MCP_SERVER_DESCRIPTION, + "version": DEFAULT_MCP_REGISTRY_VERSION, + "remotes": [ + { + "type": "streamable-http", + "url": remote_url, + } + ], + } + _temporary_mcp_servers: Dict[str, _TemporaryMCPServerEntry] = {} def _prune_expired_temporary_mcp_servers() -> None: @@ -302,15 +370,42 @@ if MCP_AVAILABLE: access_groups_list = sorted(list(access_groups)) return {"access_groups": access_groups_list} + @router.get( + "/registry.json", + tags=["mcp"], + description="MCP registry endpoint. Spec: https://github.com/modelcontextprotocol/registry", + ) + async def get_mcp_registry(request: Request): + if not _is_public_registry_enabled(): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="MCP registry is not enabled", + ) + + base_url = get_request_base_url(request) + registry_servers: List[Dict[str, Any]] = [] + registry_servers.append({"server": _build_builtin_registry_entry(base_url)}) + + registered_servers = list(global_mcp_server_manager.get_registry().values()) + registered_servers.sort(key=_build_mcp_registry_server_name) + + for server in registered_servers: + try: + entry = _build_mcp_registry_entry_for_server(server, base_url) + except Exception as e: + verbose_proxy_logger.debug( + f"Skipping MCP server {getattr(server, 'server_id', 'unknown')} in registry: {e}" + ) + continue + registry_servers.append({"server": entry}) + + return {"servers": registry_servers} + ## FastAPI Routes def _get_user_mcp_management_mode() -> UserMCPManagementMode: - proxy_general_settings: dict = {} - try: - from litellm.proxy.proxy_server import ( - general_settings as proxy_general_settings, - ) - except Exception: - pass + from litellm.proxy.proxy_server import ( + general_settings as proxy_general_settings, + ) mode = proxy_general_settings.get("user_mcp_management_mode") if mode == "view_all": diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 167160c72d1..4d4c41a3dc0 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -4,6 +4,7 @@ ROUTER SETTINGS MANAGEMENT Endpoints for accessing router configuration and metadata GET /router/settings - Get router configuration including available routing strategies +GET /router/fields - Get router settings field definitions without values (for UI rendering) """ import inspect @@ -37,6 +38,15 @@ class RouterSettingsResponse(BaseModel): ) +class RouterFieldsResponse(BaseModel): + fields: List[RouterSettingsField] = Field( + description="List of all configurable router settings with metadata (without field values)" + ) + routing_strategy_descriptions: Dict[str, str] = Field( + description="Descriptions for each routing strategy option" + ) + + def _get_routing_strategies_from_router_class() -> List[str]: """ Dynamically extract routing strategies from the Router class __init__ method. @@ -120,3 +130,53 @@ async def get_router_settings( ) raise + +@router.get( + "/router/fields", + tags=["Router Settings"], + dependencies=[Depends(user_api_key_auth)], + response_model=RouterFieldsResponse, +) +async def get_router_fields( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get router settings field definitions without values. + + Returns only the field metadata (type, description, default, options) without + populating field_value. This is useful for UI components that need to know + what fields to render, but will get the actual values from a different endpoint. + + Returns: + - fields: List of all configurable router settings with their metadata (type, description, default, options) + The routing_strategy field includes available options extracted from the Router class + Note: field_value will be None for all fields + - routing_strategy_descriptions: Descriptions for each routing strategy option + """ + try: + # Get available routing strategies dynamically from Router class + available_routing_strategies = _get_routing_strategies_from_router_class() + + # Get router settings fields from types file + router_fields = [field.model_copy(deep=True) for field in ROUTER_SETTINGS_FIELDS] + + # Populate routing_strategy field with available options + for field in router_fields: + if field.field_name == "routing_strategy": + field.options = available_routing_strategies + break + + # Ensure field_value is None for all fields (don't populate values) + for field in router_fields: + field.field_value = None + + return RouterFieldsResponse( + fields=router_fields, + routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS, + ) + except Exception as e: + verbose_proxy_logger.error( + f"Error fetching router fields: {str(e)}" + ) + raise + diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 84550092d2e..d9798dae690 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -776,6 +776,7 @@ async def handle_bedrock_count_tokens( - /v1/messages/count_tokens - /v1/messages/count-tokens """ + from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.proxy.proxy_server import llm_router @@ -822,6 +823,12 @@ async def handle_bedrock_count_tokens( return result + except BedrockError as e: + # Convert BedrockError to HTTPException for FastAPI + verbose_proxy_logger.error(f"BedrockError in handle_bedrock_count_tokens: {str(e)}") + raise HTTPException( + status_code=e.status_code, detail={"error": e.message} + ) except HTTPException: # Re-raise HTTP exceptions as-is raise diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e58ae7c643..83d4ab5657e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -229,7 +229,7 @@ from litellm.proxy.batches_endpoints.endpoints import router as batches_router from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, - create_streaming_response, + create_response, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.debug_utils import init_verbose_loggers @@ -658,7 +658,7 @@ async def _initialize_shared_aiohttp_session(): @asynccontextmanager -async def proxy_startup_event(app: FastAPI): +async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session import json @@ -788,6 +788,17 @@ async def proxy_startup_event(app: FastAPI): except Exception as e: verbose_proxy_logger.error(f"Error closing shared aiohttp session: {e}") + # Shutdown event - stop RDS IAM token refresh background task + if ( + prisma_client is not None + and hasattr(prisma_client, "db") + and hasattr(prisma_client.db, "stop_token_refresh_task") + ): + try: + await prisma_client.db.stop_token_refresh_task() + except Exception as e: + verbose_proxy_logger.error(f"Error stopping token refresh task: {e}") + await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] @@ -3972,7 +3983,7 @@ class ProxyConfig: ) try: - await global_mcp_server_manager._add_mcp_servers_from_db_to_in_memory_registry() + await global_mcp_server_manager.reload_servers_from_database() except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db - {}".format( @@ -4109,6 +4120,23 @@ class ProxyConfig: return [] +async def _reload_mcp_servers_job(): + """Background job entrypoint for MCP registry refreshes.""" + if proxy_config._should_load_db_object(object_type="mcp") is False: + return + + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + await global_mcp_server_manager.reload_servers_from_database() + except Exception as e: + verbose_proxy_logger.exception( + "Failed to reload MCP servers from database: %s", str(e) + ) + + proxy_config = ProxyConfig() @@ -4646,6 +4674,18 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) await proxy_config.get_credentials(prisma_client=prisma_client) + + from litellm.proxy._experimental.mcp_server.utils import is_mcp_available + + if is_mcp_available(): + scheduler.add_job( + _reload_mcp_servers_job, + "interval", + seconds=30, + id="reload_mcp_servers_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) await cls._initialize_slack_alerting_jobs( scheduler=scheduler, general_settings=general_settings, @@ -4915,6 +4955,14 @@ class ProxyStartupEvent: await prisma_client.connect() + ## Start RDS IAM token refresh background task if enabled ## + # This proactively refreshes IAM tokens before they expire, + # preventing the 15-minute connection failure bug (#16220) + if hasattr(prisma_client, "db") and hasattr( + prisma_client.db, "start_token_refresh_task" + ): + await prisma_client.db.start_token_refresh_task() + ## Add necessary views to proxy ## asyncio.create_task( prisma_client.check_view_exists() @@ -6824,7 +6872,7 @@ async def run_thread( if ( "stream" in data and data["stream"] is True ): # use generate_responses to stream responses - return await create_streaming_response( + return await create_response( generator=async_assistants_data_generator( user_api_key_dict=user_api_key_dict, response=response, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f16c115fed3..171898b1631 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -131,6 +131,7 @@ else: unified_guardrail = UnifiedLLMGuardrails() +_anthropic_async_clients = {} def print_verbose(print_statement): """ @@ -961,6 +962,7 @@ class ProxyLogging: Updated data dictionary if guardrail passes, None if guardrail should be skipped """ from litellm.types.guardrails import GuardrailEventHooks + from litellm.integrations.prometheus import PrometheusLogger # Determine the event type based on call type event_type = GuardrailEventHooks.pre_call @@ -973,30 +975,62 @@ class ProxyLogging: guardrail_name = callback.guardrail_name - # Check if load balancing should be used - if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name): - response = await self._execute_guardrail_with_load_balancing( - guardrail_name=guardrail_name, - hook_type="pre_call", - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) - else: - # Single guardrail - execute directly - response = await self._execute_guardrail_hook( - callback=callback, - hook_type="pre_call", - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) + # Track timing and errors for prometheus metrics + # Use time.perf_counter() for more accurate duration measurements + guardrail_start_time = time.perf_counter() + status = "success" + error_type = None - # Process the response if one was returned - if response is not None: - data = await self.process_pre_call_hook_response( - response=response, data=data, call_type=call_type - ) + try: + # Check if load balancing should be used + if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name): + response = await self._execute_guardrail_with_load_balancing( + guardrail_name=guardrail_name, + hook_type="pre_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + else: + # Single guardrail - execute directly + response = await self._execute_guardrail_hook( + callback=callback, + hook_type="pre_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + + # Process the response if one was returned + if response is not None: + data = await self.process_pre_call_hook_response( + response=response, data=data, call_type=call_type + ) + + except Exception as e: + status = "error" + error_type = type(e).__name__ + # Re-raise the exception to maintain existing behavior + raise + finally: + # Record prometheus metrics + guardrail_end_time = time.perf_counter() + latency_seconds = guardrail_end_time - guardrail_start_time + + # Get guardrail name for metrics (fallback if not set) + metrics_guardrail_name = guardrail_name or getattr(callback, "guardrail_name", callback.__class__.__name__) or "unknown" + + # Find PrometheusLogger in callbacks and record metrics + for prom_callback in litellm.callbacks: + if isinstance(prom_callback, PrometheusLogger): + prom_callback._record_guardrail_metrics( + guardrail_name=metrics_guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type="pre_call", + ) + break return data @@ -4254,11 +4288,16 @@ async def count_tokens_with_anthropic_api( if anthropic_api_key and messages: # Call Anthropic API directly for more accurate token counting - client = anthropic.Anthropic(api_key=anthropic_api_key) + + # Use cached client if available to avoid socket exhaustion + if anthropic_api_key not in _anthropic_async_clients: + _anthropic_async_clients[anthropic_api_key] = anthropic.AsyncAnthropic(api_key=anthropic_api_key) + + client = _anthropic_async_clients[anthropic_api_key] # Call with explicit parameters to satisfy type checking # Type ignore for now since messages come from generic dict input - response = client.beta.messages.count_tokens( + response = await client.beta.messages.count_tokens( model=model_to_use, messages=messages, # type: ignore betas=["token-counting-2024-11-01"], diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index a92b5d25a37..7667d1bad84 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -443,11 +443,18 @@ class ResponseAPILoggingUtils: completion_tokens=0, total_tokens=0, ) - response_api_usage: ResponseAPIUsage = ( - ResponseAPIUsage(**usage_input) - if isinstance(usage_input, dict) - else usage_input - ) + response_api_usage: ResponseAPIUsage + if isinstance(usage_input, dict): + total_tokens = usage_input.get("total_tokens") + if total_tokens is None: + input_tokens = usage_input.get("input_tokens") + output_tokens = usage_input.get("output_tokens") + if input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + usage_input["total_tokens"] = total_tokens + response_api_usage = ResponseAPIUsage(**usage_input) + else: + response_api_usage = usage_input prompt_tokens: int = response_api_usage.input_tokens or 0 completion_tokens: int = response_api_usage.output_tokens or 0 prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None diff --git a/litellm/router.py b/litellm/router.py index 84b38b3985b..364e6719300 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -386,9 +386,9 @@ class Router: ) # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### - cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = ( - "local" # default to an in-memory cache - ) + cache_type: Literal[ + "local", "redis", "redis-semantic", "s3", "disk" + ] = "local" # default to an in-memory cache redis_cache = None cache_config: Dict[str, Any] = {} @@ -430,9 +430,9 @@ class Router: self.default_max_parallel_requests = default_max_parallel_requests self.provider_default_deployment_ids: List[str] = [] self.pattern_router = PatternMatchRouter() - self.team_pattern_routers: Dict[str, PatternMatchRouter] = ( - {} - ) # {"TEAM_ID": PatternMatchRouter} + self.team_pattern_routers: Dict[ + str, PatternMatchRouter + ] = {} # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} # Initialize model_group_alias early since it's used in set_model_list @@ -613,9 +613,9 @@ class Router: ) ) - self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = ( - model_group_retry_policy - ) + self.model_group_retry_policy: Optional[ + Dict[str, RetryPolicy] + ] = model_group_retry_policy self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None if allowed_fails_policy is not None: @@ -722,7 +722,10 @@ class Router: valid_strategy_strings = ["simple-shuffle"] + [s.value for s in RoutingStrategy] if routing_strategy is not None: - is_valid_string = isinstance(routing_strategy, str) and routing_strategy in valid_strategy_strings + is_valid_string = ( + isinstance(routing_strategy, str) + and routing_strategy in valid_strategy_strings + ) is_valid_enum = isinstance(routing_strategy, RoutingStrategy) if not is_valid_string and not is_valid_enum: raise ValueError( @@ -1071,7 +1074,7 @@ class Router: self.delete_container = self.factory_function( delete_container, call_type="delete_container" ) - + # Auto-register JSON-generated container file endpoints for name, func in container_file_endpoints.items(): setattr(self, name, self.factory_function(func, call_type=name)) # type: ignore[arg-type] @@ -1500,10 +1503,7 @@ class Router: async def _acompletion( self, model: str, messages: List[Dict[str, str]], **kwargs - ) -> Union[ - ModelResponse, - CustomStreamWrapper, - ]: + ) -> Union[ModelResponse, CustomStreamWrapper,]: """ - Get an available deployment - call it with a semaphore over the call @@ -3021,7 +3021,9 @@ class Router: kwargs["original_generic_function"] = original_function kwargs["original_function"] = self._aguardrail_helper self._update_kwargs_before_fallbacks( - model=guardrail_name, kwargs=kwargs, metadata_variable_name="litellm_metadata" + model=guardrail_name, + kwargs=kwargs, + metadata_variable_name="litellm_metadata", ) verbose_router_logger.debug( f"Inside aguardrail() - guardrail_name: {guardrail_name}; kwargs: {kwargs}" @@ -3314,8 +3316,7 @@ class Router: kwargs["model"] = model kwargs["input"] = input kwargs["original_function"] = self._embedding - kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) - kwargs.setdefault("metadata", {}).update({"model_group": model}) + self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) response = self.function_with_fallbacks(**kwargs) return response except Exception as e: @@ -3617,9 +3618,9 @@ class Router: healthy_deployments=healthy_deployments, responses=responses ) returned_response = cast(OpenAIFileObject, responses[0]) - returned_response._hidden_params["model_file_id_mapping"] = ( - model_file_id_mapping - ) + returned_response._hidden_params[ + "model_file_id_mapping" + ] = model_file_id_mapping return returned_response except Exception as e: verbose_router_logger.exception( @@ -4366,11 +4367,11 @@ class Router: if isinstance(e, litellm.ContextWindowExceededError): if context_window_fallbacks is not None: - context_window_fallback_model_group: Optional[List[str]] = ( - self._get_fallback_model_group_from_fallbacks( - fallbacks=context_window_fallbacks, - model_group=model_group, - ) + context_window_fallback_model_group: Optional[ + List[str] + ] = self._get_fallback_model_group_from_fallbacks( + fallbacks=context_window_fallbacks, + model_group=model_group, ) if context_window_fallback_model_group is None: raise original_exception @@ -4402,11 +4403,11 @@ class Router: e.message += "\n{}".format(error_message) elif isinstance(e, litellm.ContentPolicyViolationError): if content_policy_fallbacks is not None: - content_policy_fallback_model_group: Optional[List[str]] = ( - self._get_fallback_model_group_from_fallbacks( - fallbacks=content_policy_fallbacks, - model_group=model_group, - ) + content_policy_fallback_model_group: Optional[ + List[str] + ] = self._get_fallback_model_group_from_fallbacks( + fallbacks=content_policy_fallbacks, + model_group=model_group, ) if content_policy_fallback_model_group is None: raise original_exception @@ -4485,9 +4486,21 @@ class Router: if hasattr(original_exception, "message"): # add the available fallbacks to the exception - original_exception.message += ". Received Model Group={}\nAvailable Model Group Fallbacks={}".format( # type: ignore - model_group, - fallback_model_group, + deployment_info = "" + if kwargs is not None: + metadata = kwargs.get('metadata', {}) + if metadata and 'deployment' in metadata: + deployment_info = f"\nUsed Deployment: {metadata['deployment']}" + if 'model_info' in metadata: + model_info = metadata['model_info'] + if isinstance(model_info, dict): + deployment_info += f"\nDeployment ID: {model_info.get('id', 'unknown')}" + + original_exception.message += ( # type: ignore + f". Received Model Group={model_group}" + f"\nAvailable Model Group Fallbacks={fallback_model_group}" + f"{deployment_info}" + f"\n\n💡 Tip: If using wildcard patterns (e.g., 'openai/*'), ensure all matching deployments have credentials with access to this model." ) if len(fallback_failure_exception_str) > 0: original_exception.message += ( # type: ignore @@ -5669,26 +5682,26 @@ class Router: """ from litellm.router_strategy.auto_router.auto_router import AutoRouter - auto_router_config_path: Optional[str] = ( - deployment.litellm_params.auto_router_config_path - ) + auto_router_config_path: Optional[ + str + ] = deployment.litellm_params.auto_router_config_path auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config if auto_router_config_path is None and auto_router_config is None: raise ValueError( "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" ) - default_model: Optional[str] = ( - deployment.litellm_params.auto_router_default_model - ) + default_model: Optional[ + str + ] = deployment.litellm_params.auto_router_default_model if default_model is None: raise ValueError( "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" ) - embedding_model: Optional[str] = ( - deployment.litellm_params.auto_router_embedding_model - ) + embedding_model: Optional[ + str + ] = deployment.litellm_params.auto_router_embedding_model if embedding_model is None: raise ValueError( "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" @@ -6235,9 +6248,9 @@ class Router: # Add custom_llm_provider if deployment.litellm_params.custom_llm_provider: - credentials["custom_llm_provider"] = ( - deployment.litellm_params.custom_llm_provider - ) + credentials[ + "custom_llm_provider" + ] = deployment.litellm_params.custom_llm_provider elif "/" in deployment.litellm_params.model: # Extract provider from "provider/model" format credentials["custom_llm_provider"] = deployment.litellm_params.model.split( @@ -6931,42 +6944,44 @@ class Router: """ return candidate_id in self.model_id_to_deployment_index_map - def resolve_model_name_from_model_id(self, model_id: Optional[str]) -> Optional[str]: + def resolve_model_name_from_model_id( + self, model_id: Optional[str] + ) -> Optional[str]: """ Resolve model_name from model_id. - + This method attempts to find the correct model_name to use with the router so that litellm_params can be automatically injected from the model config. - + Strategy: 1. First, check if model_id directly matches a model_name or deployment ID 2. If not, search through router's model_list to find a match by litellm_params.model 3. Return the model_name if found, None otherwise - + Args: model_id: The model_id extracted from decoded video_id (could be model_name or litellm_params.model value) - + Returns: model_name if found, None otherwise. If None, the request will fall through to normal flow using environment variables. """ if not model_id: return None - + # Strategy 1: Check if model_id directly matches a model_name or deployment ID if model_id in self.model_names or self.has_model_id(model_id): return model_id - + # Strategy 2: Search through router's model_list to find by litellm_params.model all_models = self.get_model_list(model_name=None) if not all_models: return None - + for deployment in all_models: litellm_params = deployment.get("litellm_params", {}) actual_model = litellm_params.get("model") - + # Match by exact match or by checking if actual_model ends with /model_id or :model_id # e.g., model_id="veo-2.0-generate-001" matches actual_model="vertex_ai/veo-2.0-generate-001" matches = ( @@ -6974,12 +6989,12 @@ class Router: or (actual_model and actual_model.endswith(f"/{model_id}")) or (actual_model and actual_model.endswith(f":{model_id}")) ) - + if matches: model_name = deployment.get("model_name") if model_name: return model_name - + # No match found return None @@ -7664,6 +7679,10 @@ class Router: ) if pattern_deployments: + verbose_router_logger.debug( + f"Pattern match for model='{model}': Found {len(pattern_deployments)} deployments. " + f"Deployment IDs: {[d.get('model_info', {}).get('id', 'unknown') for d in pattern_deployments]}" + ) return model, pattern_deployments if ( @@ -7769,14 +7788,18 @@ class Router: request_kwargs=request_kwargs, ) - verbose_router_logger.debug(f"healthy_deployments after team filter: {healthy_deployments}") + verbose_router_logger.debug( + f"healthy_deployments after team filter: {healthy_deployments}" + ) healthy_deployments = filter_web_search_deployments( healthy_deployments=healthy_deployments, request_kwargs=request_kwargs, ) - verbose_router_logger.debug(f"healthy_deployments after web search filter: {healthy_deployments}") + verbose_router_logger.debug( + f"healthy_deployments after web search filter: {healthy_deployments}" + ) if isinstance(healthy_deployments, dict): return healthy_deployments diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 6a254fc8252..88dee19ae5a 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -1,7 +1,7 @@ import re from dataclasses import dataclass from enum import Enum -from typing import Dict, List, Literal, Optional, Tuple, Union +from typing import Dict, List, Literal, Optional, Tuple from pydantic import BaseModel, Field from typing_extensions import Annotated @@ -185,6 +185,14 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_redis_daily_spend_update_queue_size", "litellm_in_memory_spend_update_queue_size", "litellm_redis_spend_update_queue_size", + "litellm_request_queue_time_seconds", + "litellm_guardrail_latency_seconds", + "litellm_guardrail_errors_total", + "litellm_guardrail_requests_total", + # Cache metrics + "litellm_cache_hits_metric", + "litellm_cache_misses_metric", + "litellm_cached_tokens_metric", ] @@ -219,6 +227,23 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, ] + litellm_request_queue_time_seconds = [ + UserAPIKeyLabelNames.END_USER.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.USER.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + ] + + # Guardrail metrics - these use custom labels (guardrail_name, status, error_type, hook_type) + # which are not part of UserAPIKeyLabelNames + litellm_guardrail_latency_seconds: List[str] = [] + litellm_guardrail_errors_total: List[str] = [] + litellm_guardrail_requests_total: List[str] = [] + litellm_proxy_total_requests_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, @@ -436,6 +461,21 @@ class PrometheusMetricLabels: litellm_redis_spend_update_queue_size: List[str] = [] + # Cache metrics - track cache hits, misses, and tokens served from cache + _cache_metric_labels = [ + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.END_USER.value, + UserAPIKeyLabelNames.USER.value, + ] + + litellm_cache_hits_metric = _cache_metric_labels + litellm_cache_misses_metric = _cache_metric_labels + litellm_cached_tokens_metric = _cache_metric_labels + @staticmethod def get_labels(label_name: DEFINED_PROMETHEUS_METRICS) -> List[str]: default_labels = getattr(PrometheusMetricLabels, label_name) @@ -460,11 +500,6 @@ class PrometheusMetricLabels: return default_labels + custom_labels -from typing import List, Optional - -from pydantic import BaseModel, Field - - class UserAPIKeyLabelValues(BaseModel): end_user: Annotated[ Optional[str], Field(..., alias=UserAPIKeyLabelNames.END_USER.value) diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 96fd79f466b..94f33ffb297 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,3 +1,4 @@ +from datetime import datetime from typing import Any, Dict, List, Optional from pydantic import BaseModel, ConfigDict @@ -50,4 +51,5 @@ class MCPServer(BaseModel): env: Optional[Dict[str, str]] = None access_groups: Optional[List[str]] = None allow_all_keys: bool = False + updated_at: Optional[datetime] = None model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/litellm/utils.py b/litellm/utils.py index 2260b2c7ba5..42d4a1ac372 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -47,7 +47,6 @@ from tiktoken import Encoding from tokenizers import Tokenizer import litellm - import litellm.litellm_core_utils # audio_utils.utils is lazy-loaded - only imported when needed for transcription calls import litellm.litellm_core_utils.json_validation_rule @@ -2895,6 +2894,7 @@ def get_optional_params_image_gen( litellm.drop_params is True or drop_params is True ) and k not in supported_params: # drop the unsupported non-default values non_default_params.pop(k, None) + passed_params.pop(k, None) elif k not in supported_params: raise UnsupportedParamsError( status_code=500, @@ -7410,15 +7410,184 @@ def validate_chat_completion_tool_choice( class ProviderConfigManager: + # Dictionary mapping for O(1) provider lookup + # Stores tuples of (factory_function, needs_model_parameter) + # This is initialized lazily on first access to avoid circular imports + _PROVIDER_CONFIG_MAP: Optional[dict[LlmProviders, tuple[Callable, bool]]] = None + + @staticmethod + def _build_provider_config_map() -> dict[LlmProviders, tuple[Callable, bool]]: + """Build the provider-to-config mapping dictionary. + + Returns a dict mapping provider to (factory_function, needs_model_parameter). + This avoids expensive inspect.signature() calls at runtime. + """ + return { + # Most common providers first for readability + # Format: (factory_function, needs_model_parameter: bool) + LlmProviders.OPENAI: (lambda: litellm.OpenAIGPTConfig(), False), + LlmProviders.ANTHROPIC: (lambda: litellm.AnthropicConfig(), False), + LlmProviders.AZURE: (lambda model: ProviderConfigManager._get_azure_config(model), True), + LlmProviders.AZURE_AI: (lambda model: ProviderConfigManager._get_azure_ai_config(model), True), + LlmProviders.VERTEX_AI: (lambda model: ProviderConfigManager._get_vertex_ai_config(model), True), + LlmProviders.BEDROCK: (lambda model: ProviderConfigManager._get_bedrock_config(model), True), + LlmProviders.COHERE: (lambda model: ProviderConfigManager._get_cohere_config(model), True), + LlmProviders.COHERE_CHAT: (lambda model: ProviderConfigManager._get_cohere_config(model), True), + # Simple provider mappings (no model parameter needed) + LlmProviders.DEEPSEEK: (lambda: litellm.DeepSeekChatConfig(), False), + LlmProviders.GROQ: (lambda: litellm.GroqChatConfig(), False), + LlmProviders.BYTEZ: (lambda: litellm.BytezChatConfig(), False), + LlmProviders.DATABRICKS: (lambda: litellm.DatabricksConfig(), False), + LlmProviders.XAI: (lambda: litellm.XAIChatConfig(), False), + LlmProviders.ZAI: (lambda: litellm.ZAIChatConfig(), False), + LlmProviders.LAMBDA_AI: (lambda: litellm.LambdaAIChatConfig(), False), + LlmProviders.LLAMA: (lambda: litellm.LlamaAPIConfig(), False), + LlmProviders.TEXT_COMPLETION_OPENAI: (lambda: litellm.OpenAITextCompletionConfig(), False), + LlmProviders.SNOWFLAKE: (lambda: litellm.SnowflakeConfig(), False), + LlmProviders.CLARIFAI: (lambda: litellm.ClarifaiConfig(), False), + LlmProviders.ANTHROPIC_TEXT: (lambda: litellm.AnthropicTextConfig(), False), + LlmProviders.VERTEX_AI_BETA: (lambda: litellm.VertexGeminiConfig(), False), + LlmProviders.CLOUDFLARE: (lambda: litellm.CloudflareChatConfig(), False), + LlmProviders.SAGEMAKER_CHAT: (lambda: litellm.SagemakerChatConfig(), False), + LlmProviders.SAGEMAKER: (lambda: litellm.SagemakerConfig(), False), + LlmProviders.FIREWORKS_AI: (lambda: litellm.FireworksAIConfig(), False), + LlmProviders.FRIENDLIAI: (lambda: litellm.FriendliaiChatConfig(), False), + LlmProviders.WATSONX: (lambda: litellm.IBMWatsonXChatConfig(), False), + LlmProviders.WATSONX_TEXT: (lambda: litellm.IBMWatsonXAIConfig(), False), + LlmProviders.EMPOWER: (lambda: litellm.EmpowerChatConfig(), False), + LlmProviders.MINIMAX: (lambda: litellm.MinimaxChatConfig(), False), + LlmProviders.GITHUB: (lambda: litellm.GithubChatConfig(), False), + LlmProviders.COMPACTIFAI: (lambda: litellm.CompactifAIChatConfig(), False), + LlmProviders.GITHUB_COPILOT: (lambda: litellm.GithubCopilotConfig(), False), + LlmProviders.GIGACHAT: (lambda: litellm.GigaChatConfig(), False), + LlmProviders.RAGFLOW: (lambda: litellm.RAGFlowConfig(), False), + LlmProviders.CUSTOM: (lambda: litellm.OpenAILikeChatConfig(), False), + LlmProviders.CUSTOM_OPENAI: (lambda: litellm.OpenAILikeChatConfig(), False), + LlmProviders.OPENAI_LIKE: (lambda: litellm.OpenAILikeChatConfig(), False), + LlmProviders.AIOHTTP_OPENAI: (lambda: litellm.AiohttpOpenAIChatConfig(), False), + LlmProviders.HOSTED_VLLM: (lambda: litellm.HostedVLLMChatConfig(), False), + LlmProviders.LLAMAFILE: (lambda: litellm.LlamafileChatConfig(), False), + LlmProviders.LM_STUDIO: (lambda: litellm.LMStudioChatConfig(), False), + LlmProviders.GALADRIEL: (lambda: litellm.GaladrielChatConfig(), False), + LlmProviders.REPLICATE: (lambda: litellm.ReplicateConfig(), False), + LlmProviders.HUGGINGFACE: (lambda: litellm.HuggingFaceChatConfig(), False), + LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIConfig(), False), + LlmProviders.OPENROUTER: (lambda: litellm.OpenrouterConfig(), False), + LlmProviders.VERCEL_AI_GATEWAY: (lambda: litellm.VercelAIGatewayConfig(), False), + LlmProviders.COMETAPI: (lambda: litellm.CometAPIConfig(), False), + LlmProviders.DATAROBOT: (lambda: litellm.DataRobotConfig(), False), + LlmProviders.GEMINI: (lambda: litellm.GoogleAIStudioGeminiConfig(), False), + LlmProviders.AI21: (lambda: litellm.AI21ChatConfig(), False), + LlmProviders.AI21_CHAT: (lambda: litellm.AI21ChatConfig(), False), + LlmProviders.AZURE_TEXT: (lambda: litellm.AzureOpenAITextConfig(), False), + LlmProviders.NLP_CLOUD: (lambda: litellm.NLPCloudConfig(), False), + LlmProviders.OOBABOOGA: (lambda: litellm.OobaboogaConfig(), False), + LlmProviders.OLLAMA_CHAT: (lambda: litellm.OllamaChatConfig(), False), + LlmProviders.DEEPINFRA: (lambda: litellm.DeepInfraConfig(), False), + LlmProviders.PERPLEXITY: (lambda: litellm.PerplexityChatConfig(), False), + LlmProviders.MISTRAL: (lambda: litellm.MistralConfig(), False), + LlmProviders.CODESTRAL: (lambda: litellm.MistralConfig(), False), + LlmProviders.NVIDIA_NIM: (lambda: litellm.NvidiaNimConfig(), False), + LlmProviders.CEREBRAS: (lambda: litellm.CerebrasConfig(), False), + LlmProviders.BASETEN: (lambda: litellm.BasetenConfig(), False), + LlmProviders.VOLCENGINE: (lambda: litellm.VolcEngineConfig(), False), + LlmProviders.TEXT_COMPLETION_CODESTRAL: (lambda: litellm.CodestralTextCompletionConfig(), False), + LlmProviders.SAMBANOVA: (lambda: litellm.SambanovaConfig(), False), + LlmProviders.MARITALK: (lambda: litellm.MaritalkConfig(), False), + LlmProviders.VLLM: (lambda: litellm.VLLMConfig(), False), + LlmProviders.OLLAMA: (lambda: litellm.OllamaConfig(), False), + LlmProviders.PREDIBASE: (lambda: litellm.PredibaseConfig(), False), + LlmProviders.TRITON: (lambda: litellm.TritonConfig(), False), + LlmProviders.PETALS: (lambda: litellm.PetalsConfig(), False), + LlmProviders.SAP_GENERATIVE_AI_HUB: (lambda: litellm.GenAIHubOrchestrationConfig(), False), + LlmProviders.FEATHERLESS_AI: (lambda: litellm.FeatherlessAIConfig(), False), + LlmProviders.NOVITA: (lambda: litellm.NovitaConfig(), False), + LlmProviders.NEBIUS: (lambda: litellm.NebiusConfig(), False), + LlmProviders.WANDB: (lambda: litellm.WandbConfig(), False), + LlmProviders.DASHSCOPE: (lambda: litellm.DashScopeChatConfig(), False), + LlmProviders.MOONSHOT: (lambda: litellm.MoonshotChatConfig(), False), + LlmProviders.DOCKER_MODEL_RUNNER: (lambda: litellm.DockerModelRunnerChatConfig(), False), + LlmProviders.V0: (lambda: litellm.V0ChatConfig(), False), + LlmProviders.MORPH: (lambda: litellm.MorphChatConfig(), False), + LlmProviders.LITELLM_PROXY: (lambda: litellm.LiteLLMProxyChatConfig(), False), + LlmProviders.GRADIENT_AI: (lambda: litellm.GradientAIConfig(), False), + LlmProviders.NSCALE: (lambda: litellm.NscaleConfig(), False), + LlmProviders.HEROKU: (lambda: litellm.HerokuChatConfig(), False), + LlmProviders.OCI: (lambda: litellm.OCIChatConfig(), False), + LlmProviders.HYPERBOLIC: (lambda: litellm.HyperbolicChatConfig(), False), + LlmProviders.OVHCLOUD: (lambda: litellm.OVHCloudChatConfig(), False), + LlmProviders.AMAZON_NOVA: (lambda: litellm.AmazonNovaChatConfig(), False), + LlmProviders.LANGGRAPH: (lambda: ProviderConfigManager._get_langgraph_config(), False), + } + + @staticmethod + def _get_azure_config(model: str) -> BaseConfig: + """Get Azure config based on model type.""" + if litellm.AzureOpenAIO1Config().is_o_series_model(model=model): + return litellm.AzureOpenAIO1Config() + if litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=model): + return litellm.AzureOpenAIGPT5Config() + return litellm.AzureOpenAIConfig() + + @staticmethod + def _get_azure_ai_config(model: str) -> BaseConfig: + """Get Azure AI config based on model type.""" + if "claude" in model.lower(): + return litellm.AzureAnthropicConfig() + return litellm.AzureAIStudioConfig() + + @staticmethod + def _get_vertex_ai_config(model: str) -> BaseConfig: + """Get Vertex AI config based on model type.""" + if "gemini" in model: + return litellm.VertexGeminiConfig() + elif "claude" in model: + return litellm.VertexAIAnthropicConfig() + elif "gpt-oss" in model: + from litellm.llms.vertex_ai.vertex_ai_partner_models.gpt_oss.transformation import ( + VertexAIGPTOSSTransformation, + ) + return VertexAIGPTOSSTransformation() + elif model in litellm.vertex_mistral_models: + if "codestral" in model: + return litellm.CodestralTextCompletionConfig() + return litellm.MistralConfig() + elif model in litellm.vertex_ai_ai21_models: + return litellm.VertexAIAi21Config() + else: + return litellm.VertexAILlama3Config() + + @staticmethod + def _get_bedrock_config(model: str) -> BaseConfig: + """Get Bedrock config based on model.""" + from litellm.llms.bedrock.common_utils import get_bedrock_chat_config + return get_bedrock_chat_config(model=model) + + @staticmethod + def _get_cohere_config(model: str) -> BaseConfig: + """Get Cohere config based on route.""" + CohereModelInfo = getattr(sys.modules[__name__], 'CohereModelInfo') + route = CohereModelInfo.get_cohere_route(model) + if route == "v2": + return litellm.CohereV2ChatConfig() + return litellm.CohereChatConfig() + + @staticmethod + def _get_langgraph_config() -> BaseConfig: + """Get LangGraph config.""" + from litellm.llms.langgraph.chat.transformation import LangGraphConfig + return LangGraphConfig() + @staticmethod def get_provider_chat_config( # noqa: PLR0915 model: str, provider: LlmProviders ) -> Optional[BaseConfig]: """ Returns the provider config for a given provider. + + Uses O(1) dictionary lookup for fast provider resolution. """ - - # Check JSON providers FIRST + # Check JSON providers FIRST (these override standard mappings) from litellm.llms.openai_like.dynamic_config import create_config_class from litellm.llms.openai_like.json_loader import JSONProviderRegistry @@ -7428,244 +7597,29 @@ class ProviderConfigManager: raise ValueError(f"Provider {provider.value} not found") return create_config_class(provider_config)() - if ( - provider == LlmProviders.OPENAI - and litellm.openaiOSeriesConfig.is_model_o_series_model(model=model) - ): - return litellm.openaiOSeriesConfig - elif ( - provider == LlmProviders.OPENAI - and litellm.OpenAIGPT5Config.is_model_gpt_5_model(model=model) - ): - return litellm.OpenAIGPT5Config() - elif litellm.LlmProviders.DEEPSEEK == provider: - return litellm.DeepSeekChatConfig() - elif litellm.LlmProviders.GROQ == provider: - return litellm.GroqChatConfig() - elif litellm.LlmProviders.BYTEZ == provider: - return litellm.BytezChatConfig() - elif litellm.LlmProviders.DATABRICKS == provider: - return litellm.DatabricksConfig() - elif litellm.LlmProviders.XAI == provider: - return litellm.XAIChatConfig() - elif litellm.LlmProviders.ZAI == provider: - return litellm.ZAIChatConfig() - elif litellm.LlmProviders.LAMBDA_AI == provider: - return litellm.LambdaAIChatConfig() - elif litellm.LlmProviders.LLAMA == provider: - return litellm.LlamaAPIConfig() - elif litellm.LlmProviders.TEXT_COMPLETION_OPENAI == provider: - return litellm.OpenAITextCompletionConfig() - elif ( - litellm.LlmProviders.COHERE_CHAT == provider - or litellm.LlmProviders.COHERE == provider - ): - CohereModelInfo = getattr(sys.modules[__name__], 'CohereModelInfo') - route = CohereModelInfo.get_cohere_route(model) - if route == "v2": - return litellm.CohereV2ChatConfig() - else: + # Handle OpenAI special cases (O-series and GPT-5 models) + if provider == LlmProviders.OPENAI: + if litellm.openaiOSeriesConfig.is_model_o_series_model(model=model): + return litellm.openaiOSeriesConfig + if litellm.OpenAIGPT5Config.is_model_gpt_5_model(model=model): + return litellm.OpenAIGPT5Config() - return litellm.CohereChatConfig() - elif litellm.LlmProviders.SNOWFLAKE == provider: - return litellm.SnowflakeConfig() - elif litellm.LlmProviders.CLARIFAI == provider: - return litellm.ClarifaiConfig() - elif litellm.LlmProviders.ANTHROPIC == provider: - return litellm.AnthropicConfig() - elif litellm.LlmProviders.ANTHROPIC_TEXT == provider: - return litellm.AnthropicTextConfig() - elif litellm.LlmProviders.VERTEX_AI_BETA == provider: - return litellm.VertexGeminiConfig() - elif litellm.LlmProviders.VERTEX_AI == provider: - if "gemini" in model: - return litellm.VertexGeminiConfig() - elif "claude" in model: - return litellm.VertexAIAnthropicConfig() - elif "gpt-oss" in model: - from litellm.llms.vertex_ai.vertex_ai_partner_models.gpt_oss.transformation import ( - VertexAIGPTOSSTransformation, - ) + # Initialize provider config map lazily (avoids circular imports) + if ProviderConfigManager._PROVIDER_CONFIG_MAP is None: + ProviderConfigManager._PROVIDER_CONFIG_MAP = ProviderConfigManager._build_provider_config_map() - return VertexAIGPTOSSTransformation() - elif model in litellm.vertex_mistral_models: - if "codestral" in model: - return litellm.CodestralTextCompletionConfig() - else: - return litellm.MistralConfig() - elif model in litellm.vertex_ai_ai21_models: - return litellm.VertexAIAi21Config() - else: # use generic openai-like param mapping - return litellm.VertexAILlama3Config() - elif litellm.LlmProviders.CLOUDFLARE == provider: - return litellm.CloudflareChatConfig() - elif litellm.LlmProviders.SAGEMAKER_CHAT == provider: - return litellm.SagemakerChatConfig() - elif litellm.LlmProviders.SAGEMAKER == provider: - return litellm.SagemakerConfig() - elif litellm.LlmProviders.FIREWORKS_AI == provider: - return litellm.FireworksAIConfig() - elif litellm.LlmProviders.FRIENDLIAI == provider: - return litellm.FriendliaiChatConfig() - elif litellm.LlmProviders.WATSONX == provider: - return litellm.IBMWatsonXChatConfig() - elif litellm.LlmProviders.WATSONX_TEXT == provider: - return litellm.IBMWatsonXAIConfig() - elif litellm.LlmProviders.EMPOWER == provider: - return litellm.EmpowerChatConfig() - elif litellm.LlmProviders.MINIMAX == provider: - return litellm.MinimaxChatConfig() - elif litellm.LlmProviders.GITHUB == provider: - return litellm.GithubChatConfig() - elif litellm.LlmProviders.COMPACTIFAI == provider: - return litellm.CompactifAIChatConfig() - elif litellm.LlmProviders.GITHUB_COPILOT == provider: - return litellm.GithubCopilotConfig() - elif litellm.LlmProviders.GIGACHAT == provider: - return litellm.GigaChatConfig() - elif litellm.LlmProviders.RAGFLOW == provider: - return litellm.RAGFlowConfig() - elif ( - litellm.LlmProviders.CUSTOM == provider - or litellm.LlmProviders.CUSTOM_OPENAI == provider - or litellm.LlmProviders.OPENAI_LIKE == provider - ): - return litellm.OpenAILikeChatConfig() - elif litellm.LlmProviders.AIOHTTP_OPENAI == provider: - return litellm.AiohttpOpenAIChatConfig() - elif litellm.LlmProviders.HOSTED_VLLM == provider: - return litellm.HostedVLLMChatConfig() - elif litellm.LlmProviders.LLAMAFILE == provider: - return litellm.LlamafileChatConfig() - elif litellm.LlmProviders.LM_STUDIO == provider: - return litellm.LMStudioChatConfig() - elif litellm.LlmProviders.GALADRIEL == provider: - return litellm.GaladrielChatConfig() - elif litellm.LlmProviders.REPLICATE == provider: - return litellm.ReplicateConfig() - elif litellm.LlmProviders.HUGGINGFACE == provider: - return litellm.HuggingFaceChatConfig() - elif litellm.LlmProviders.TOGETHER_AI == provider: - return litellm.TogetherAIConfig() - elif litellm.LlmProviders.OPENROUTER == provider: - return litellm.OpenrouterConfig() - elif litellm.LlmProviders.VERCEL_AI_GATEWAY == provider: - return litellm.VercelAIGatewayConfig() - elif litellm.LlmProviders.COMETAPI == provider: - return litellm.CometAPIConfig() - elif litellm.LlmProviders.DATAROBOT == provider: - return litellm.DataRobotConfig() - elif litellm.LlmProviders.GEMINI == provider: - return litellm.GoogleAIStudioGeminiConfig() - elif ( - litellm.LlmProviders.AI21 == provider - or litellm.LlmProviders.AI21_CHAT == provider - ): - return litellm.AI21ChatConfig() - elif litellm.LlmProviders.AZURE == provider: - if litellm.AzureOpenAIO1Config().is_o_series_model(model=model): - return litellm.AzureOpenAIO1Config() - if litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=model): - return litellm.AzureOpenAIGPT5Config() - return litellm.AzureOpenAIConfig() - elif litellm.LlmProviders.AZURE_AI == provider: - if "claude" in model.lower(): - return litellm.AzureAnthropicConfig() - return litellm.AzureAIStudioConfig() - elif litellm.LlmProviders.AZURE_TEXT == provider: - return litellm.AzureOpenAITextConfig() - elif litellm.LlmProviders.HOSTED_VLLM == provider: - return litellm.HostedVLLMChatConfig() - elif litellm.LlmProviders.NLP_CLOUD == provider: - return litellm.NLPCloudConfig() - elif litellm.LlmProviders.OOBABOOGA == provider: - return litellm.OobaboogaConfig() - elif litellm.LlmProviders.OLLAMA_CHAT == provider: - return litellm.OllamaChatConfig() - elif litellm.LlmProviders.DEEPINFRA == provider: - return litellm.DeepInfraConfig() - elif litellm.LlmProviders.PERPLEXITY == provider: - return litellm.PerplexityChatConfig() - elif ( - litellm.LlmProviders.MISTRAL == provider - or litellm.LlmProviders.CODESTRAL == provider - ): - return litellm.MistralConfig() - elif litellm.LlmProviders.NVIDIA_NIM == provider: - return litellm.NvidiaNimConfig() - elif litellm.LlmProviders.CEREBRAS == provider: - return litellm.CerebrasConfig() - elif litellm.LlmProviders.BASETEN == provider: - return litellm.BasetenConfig() - elif litellm.LlmProviders.VOLCENGINE == provider: - return litellm.VolcEngineConfig() - elif litellm.LlmProviders.TEXT_COMPLETION_CODESTRAL == provider: - return litellm.CodestralTextCompletionConfig() - elif litellm.LlmProviders.SAMBANOVA == provider: - return litellm.SambanovaConfig() - elif litellm.LlmProviders.MARITALK == provider: - return litellm.MaritalkConfig() - elif litellm.LlmProviders.CLOUDFLARE == provider: - return litellm.CloudflareChatConfig() - elif litellm.LlmProviders.ANTHROPIC_TEXT == provider: - return litellm.AnthropicTextConfig() - elif litellm.LlmProviders.VLLM == provider: - return litellm.VLLMConfig() - elif litellm.LlmProviders.OLLAMA == provider: - return litellm.OllamaConfig() - elif litellm.LlmProviders.PREDIBASE == provider: - return litellm.PredibaseConfig() - elif litellm.LlmProviders.TRITON == provider: - return litellm.TritonConfig() - elif litellm.LlmProviders.PETALS == provider: - return litellm.PetalsConfig() - elif litellm.LlmProviders.SAP_GENERATIVE_AI_HUB == provider: - return litellm.GenAIHubOrchestrationConfig() - elif litellm.LlmProviders.FEATHERLESS_AI == provider: - return litellm.FeatherlessAIConfig() - elif litellm.LlmProviders.NOVITA == provider: - return litellm.NovitaConfig() - elif litellm.LlmProviders.NEBIUS == provider: - return litellm.NebiusConfig() - elif litellm.LlmProviders.WANDB == provider: - return litellm.WandbConfig() - elif litellm.LlmProviders.DASHSCOPE == provider: - return litellm.DashScopeChatConfig() - elif litellm.LlmProviders.MOONSHOT == provider: - return litellm.MoonshotChatConfig() - elif litellm.LlmProviders.DOCKER_MODEL_RUNNER == provider: - return litellm.DockerModelRunnerChatConfig() - elif litellm.LlmProviders.V0 == provider: - return litellm.V0ChatConfig() - elif litellm.LlmProviders.MORPH == provider: - return litellm.MorphChatConfig() - elif litellm.LlmProviders.BEDROCK == provider: - from litellm.llms.bedrock.common_utils import get_bedrock_chat_config + # O(1) dictionary lookup + config_entry = ProviderConfigManager._PROVIDER_CONFIG_MAP.get(provider) + if config_entry is None: + return None - return get_bedrock_chat_config(model=model) - elif litellm.LlmProviders.LITELLM_PROXY == provider: - return litellm.LiteLLMProxyChatConfig() - elif litellm.LlmProviders.OPENAI == provider: - return litellm.OpenAIGPTConfig() - elif litellm.LlmProviders.GRADIENT_AI == provider: - return litellm.GradientAIConfig() - elif litellm.LlmProviders.NSCALE == provider: - return litellm.NscaleConfig() - elif litellm.LlmProviders.HEROKU == provider: - return litellm.HerokuChatConfig() - elif litellm.LlmProviders.OCI == provider: - return litellm.OCIChatConfig() - elif litellm.LlmProviders.HYPERBOLIC == provider: - return litellm.HyperbolicChatConfig() - elif litellm.LlmProviders.OVHCLOUD == provider: - return litellm.OVHCloudChatConfig() - elif litellm.LlmProviders.AMAZON_NOVA == provider: - return litellm.AmazonNovaChatConfig() - elif litellm.LlmProviders.LANGGRAPH == provider: - from litellm.llms.langgraph.chat.transformation import LangGraphConfig - - return LangGraphConfig() - return None + # Unpack factory function and whether it needs model parameter + # This avoids expensive inspect.signature() calls at runtime + config_factory, needs_model = config_entry + if needs_model: + return config_factory(model) # type: ignore + else: + return config_factory() # type: ignore @staticmethod def get_provider_embedding_config( @@ -7718,6 +7672,11 @@ class ProviderConfigManager: return litellm.CometAPIEmbeddingConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: return litellm.GithubCopilotEmbeddingConfig() + elif litellm.LlmProviders.OPENROUTER == provider: + from litellm.llms.openrouter.embedding.transformation import ( + OpenrouterEmbeddingConfig, + ) + return OpenrouterEmbeddingConfig() elif litellm.LlmProviders.GIGACHAT == provider: return litellm.GigaChatEmbeddingConfig() elif litellm.LlmProviders.SAGEMAKER == provider: @@ -7941,6 +7900,8 @@ class ProviderConfigManager: return litellm.LemonadeChatConfig() elif LlmProviders.CLARIFAI == provider: return litellm.ClarifaiConfig() + elif LlmProviders.BEDROCK == provider: + return litellm.llms.bedrock.common_utils.BedrockModelInfo() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3c2c20b4dce..2f34d2a7d3d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -410,8 +410,8 @@ "max_input_tokens": 8172, "max_tokens": 8172, "mode": "embedding", - "input_cost_per_token": 1.35e-7, - "input_cost_per_image": 6e-5, + "input_cost_per_token": 1.35e-07, + "input_cost_per_image": 6e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "output_cost_per_token": 0.0, @@ -1398,8 +1398,8 @@ "mode": "chat" }, "azure_ai/gpt-oss-120b": { - "input_cost_per_token": 1.5e-7, - "output_cost_per_token": 6e-7, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, "max_output_tokens": 131072, @@ -2077,7 +2077,7 @@ "litellm_provider": "azure", "max_input_tokens": 4097, "max_output_tokens": 4096, - "max_tokens": 4097, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, "supports_function_calling": true, @@ -2090,7 +2090,7 @@ "litellm_provider": "azure", "max_input_tokens": 4097, "max_output_tokens": 4096, - "max_tokens": 4097, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, "supports_function_calling": true, @@ -2869,7 +2869,7 @@ "/v1/audio/transcriptions" ] }, - "azure/gpt-5.1-2025-11-13": { + "azure/gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -2905,7 +2905,7 @@ "supports_service_tier": true, "supports_vision": true }, - "azure/gpt-5.1-chat-2025-11-13": { + "azure/gpt-5.1-chat-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -2940,7 +2940,7 @@ "supports_tool_choice": false, "supports_vision": true }, - "azure/gpt-5.1-codex-2025-11-13": { + "azure/gpt-5.1-codex-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, "input_cost_per_token": 1.25e-06, @@ -3298,7 +3298,7 @@ "litellm_provider": "azure", "max_input_tokens": 272000, "max_output_tokens": 128000, - "max_tokens": 400000, + "max_tokens": 128000, "mode": "responses", "output_cost_per_token": 0.00012, "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/foundry-models/concepts/models-sold-directly-by-azure?pivots=azure-openai&tabs=global-standard-aoai%2Cstandard-chat-completions%2Cglobal-standard#gpt-5", @@ -3623,7 +3623,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1.68e-04, + "output_cost_per_token": 0.000168, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -3654,7 +3654,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1.68e-04, + "output_cost_per_token": 0.000168, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -4297,13 +4297,13 @@ "output_cost_per_token": 0.0 }, "azure/speech/azure-tts": { - "input_cost_per_character": 15e-06, + "input_cost_per_character": 1.5e-05, "litellm_provider": "azure", "mode": "audio_speech", "source": "https://azure.microsoft.com/en-us/pricing/calculator/" }, "azure/speech/azure-tts-hd": { - "input_cost_per_character": 30e-06, + "input_cost_per_character": 3e-05, "litellm_provider": "azure", "mode": "audio_speech", "source": "https://azure.microsoft.com/en-us/pricing/calculator/" @@ -5197,7 +5197,7 @@ }, "azure_ai/mistral-document-ai-2505": { "litellm_provider": "azure_ai", - "ocr_cost_per_page": 3e-3, + "ocr_cost_per_page": 0.003, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" @@ -5206,7 +5206,7 @@ }, "azure_ai/doc-intelligence/prebuilt-read": { "litellm_provider": "azure_ai", - "ocr_cost_per_page": 1.5e-3, + "ocr_cost_per_page": 0.0015, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" @@ -5215,7 +5215,7 @@ }, "azure_ai/doc-intelligence/prebuilt-layout": { "litellm_provider": "azure_ai", - "ocr_cost_per_page": 1e-2, + "ocr_cost_per_page": 0.01, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" @@ -5224,7 +5224,7 @@ }, "azure_ai/doc-intelligence/prebuilt-document": { "litellm_provider": "azure_ai", - "ocr_cost_per_page": 1e-2, + "ocr_cost_per_page": 0.01, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" @@ -5298,12 +5298,12 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, - "azure_ai/deepseek-v3.2": { + "azure_ai/deepseek-v3.2": { "input_cost_per_token": 5.8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 8192, + "max_tokens": 163840, "mode": "chat", "output_cost_per_token": 1.68e-06, "supports_assistant_prefill": true, @@ -5317,7 +5317,7 @@ "litellm_provider": "azure_ai", "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 8192, + "max_tokens": 163840, "mode": "chat", "output_cost_per_token": 1.68e-06, "supports_assistant_prefill": true, @@ -5452,7 +5452,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-non-reasoning": { - "input_cost_per_token": 0.43e-06, + "input_cost_per_token": 4.3e-07, "output_cost_per_token": 1.73e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -5465,7 +5465,7 @@ "supports_web_search": true }, "azure_ai/grok-4-fast-reasoning": { - "input_cost_per_token": 0.43e-06, + "input_cost_per_token": 4.3e-07, "output_cost_per_token": 1.73e-06, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -5623,7 +5623,7 @@ "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, - "max_tokens": 16384, + "max_tokens": 4096, "mode": "completion", "output_cost_per_token": 4e-07 }, @@ -7038,7 +7038,7 @@ "litellm_provider": "anthropic", "max_input_tokens": 1000000, "max_output_tokens": 64000, - "max_tokens": 1000000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_above_200k_tokens": 2.25e-05, @@ -7821,7 +7821,7 @@ "litellm_provider": "deepseek", "max_input_tokens": 131072, "max_output_tokens": 65536, - "max_tokens": 131072, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1.7e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", @@ -7842,7 +7842,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 1000000, "max_output_tokens": 16384, - "max_tokens": 1000000, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.alibabacloud.com/help/en/model-studio/models", @@ -7883,7 +7883,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 997952, "max_output_tokens": 32768, - "max_tokens": 1000000, + "max_tokens": 32768, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -7913,7 +7913,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 30720, "max_output_tokens": 8192, - "max_tokens": 32768, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 6.4e-06, "source": "https://www.alibabacloud.com/help/en/model-studio/models", @@ -7926,7 +7926,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 129024, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://www.alibabacloud.com/help/en/model-studio/models", @@ -7939,7 +7939,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 129024, "max_output_tokens": 8192, - "max_tokens": 131072, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://www.alibabacloud.com/help/en/model-studio/models", @@ -7952,7 +7952,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 129024, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", "output_cost_per_reasoning_token": 4e-06, "output_cost_per_token": 1.2e-06, @@ -7966,7 +7966,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 129024, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", "output_cost_per_reasoning_token": 4e-06, "output_cost_per_token": 1.2e-06, @@ -7979,7 +7979,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 997952, "max_output_tokens": 32768, - "max_tokens": 1000000, + "max_tokens": 32768, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8010,7 +8010,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 997952, "max_output_tokens": 32768, - "max_tokens": 1000000, + "max_tokens": 32768, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8041,7 +8041,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 997952, "max_output_tokens": 32768, - "max_tokens": 1000000, + "max_tokens": 32768, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8073,7 +8073,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 129024, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", "output_cost_per_reasoning_token": 5e-07, "output_cost_per_token": 2e-07, @@ -8087,7 +8087,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 1000000, "max_output_tokens": 8192, - "max_tokens": 1000000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 2e-07, "source": "https://www.alibabacloud.com/help/en/model-studio/models", @@ -8100,7 +8100,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 1000000, "max_output_tokens": 16384, - "max_tokens": 1000000, + "max_tokens": 16384, "mode": "chat", "output_cost_per_reasoning_token": 5e-07, "output_cost_per_token": 2e-07, @@ -8114,7 +8114,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 1000000, "max_output_tokens": 16384, - "max_tokens": 1000000, + "max_tokens": 16384, "mode": "chat", "output_cost_per_reasoning_token": 5e-07, "output_cost_per_token": 2e-07, @@ -8127,7 +8127,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 129024, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8138,7 +8138,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 997952, "max_output_tokens": 65536, - "max_tokens": 1000000, + "max_tokens": 65536, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8187,7 +8187,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 997952, "max_output_tokens": 65536, - "max_tokens": 1000000, + "max_tokens": 65536, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8232,7 +8232,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 997952, "max_output_tokens": 65536, - "max_tokens": 1000000, + "max_tokens": 65536, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8281,7 +8281,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 997952, "max_output_tokens": 65536, - "max_tokens": 1000000, + "max_tokens": 65536, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8326,7 +8326,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 258048, "max_output_tokens": 65536, - "max_tokens": 262144, + "max_tokens": 65536, "mode": "chat", "source": "https://www.alibabacloud.com/help/en/model-studio/models", "supports_function_calling": true, @@ -8364,7 +8364,7 @@ "litellm_provider": "dashscope", "max_input_tokens": 98304, "max_output_tokens": 8192, - "max_tokens": 131072, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.4e-06, "source": "https://www.alibabacloud.com/help/en/model-studio/models", @@ -8393,7 +8393,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 128000, - "max_tokens": 200000, + "max_tokens": 128000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8412,7 +8412,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 64000, - "max_tokens": 200000, + "max_tokens": 64000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8431,7 +8431,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, - "max_tokens": 200000, + "max_tokens": 32000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8450,7 +8450,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, - "max_tokens": 200000, + "max_tokens": 32000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8469,7 +8469,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 64000, - "max_tokens": 200000, + "max_tokens": 64000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8488,7 +8488,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 64000, - "max_tokens": 200000, + "max_tokens": 64000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8507,7 +8507,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 64000, - "max_tokens": 200000, + "max_tokens": 64000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8526,7 +8526,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 64000, - "max_tokens": 200000, + "max_tokens": 64000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8545,7 +8545,7 @@ "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65535, - "max_tokens": 1048576, + "max_tokens": 65535, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8562,7 +8562,7 @@ "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, - "max_tokens": 1048576, + "max_tokens": 65536, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8594,7 +8594,7 @@ "litellm_provider": "databricks", "max_input_tokens": 400000, "max_output_tokens": 128000, - "max_tokens": 400000, + "max_tokens": 128000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8609,7 +8609,7 @@ "litellm_provider": "databricks", "max_input_tokens": 400000, "max_output_tokens": 128000, - "max_tokens": 400000, + "max_tokens": 128000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8624,7 +8624,7 @@ "litellm_provider": "databricks", "max_input_tokens": 400000, "max_output_tokens": 128000, - "max_tokens": 400000, + "max_tokens": 128000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8639,7 +8639,7 @@ "litellm_provider": "databricks", "max_input_tokens": 400000, "max_output_tokens": 128000, - "max_tokens": 400000, + "max_tokens": 128000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8747,7 +8747,7 @@ "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 128000, - "max_tokens": 200000, + "max_tokens": 128000, "metadata": { "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, @@ -8846,7 +8846,7 @@ "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, - "max_tokens": 16384, + "max_tokens": 4096, "mode": "completion", "output_cost_per_token": 2e-06 }, @@ -10107,7 +10107,7 @@ "litellm_provider": "deepseek", "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 8192, + "max_tokens": 163840, "mode": "chat", "output_cost_per_token": 4e-07, "supports_assistant_prefill": true, @@ -10121,7 +10121,7 @@ "litellm_provider": "bedrock_converse", "max_input_tokens": 163840, "max_output_tokens": 81920, - "max_tokens": 163840, + "max_tokens": 81920, "mode": "chat", "output_cost_per_token": 1.68e-06, "supports_function_calling": true, @@ -10202,14 +10202,14 @@ "mode": "search", "tiered_pricing": [ { - "input_cost_per_query": 5e-03, + "input_cost_per_query": 0.005, "max_results_range": [ 0, 25 ] }, { - "input_cost_per_query": 25e-03, + "input_cost_per_query": 0.025, "max_results_range": [ 26, 100 @@ -10222,70 +10222,70 @@ "mode": "search", "tiered_pricing": [ { - "input_cost_per_query": 1.66e-03, + "input_cost_per_query": 0.00166, "max_results_range": [ 1, 10 ] }, { - "input_cost_per_query": 3.32e-03, + "input_cost_per_query": 0.00332, "max_results_range": [ 11, 20 ] }, { - "input_cost_per_query": 4.98e-03, + "input_cost_per_query": 0.00498, "max_results_range": [ 21, 30 ] }, { - "input_cost_per_query": 6.64e-03, + "input_cost_per_query": 0.00664, "max_results_range": [ 31, 40 ] }, { - "input_cost_per_query": 8.3e-03, + "input_cost_per_query": 0.0083, "max_results_range": [ 41, 50 ] }, { - "input_cost_per_query": 9.96e-03, + "input_cost_per_query": 0.00996, "max_results_range": [ 51, 60 ] }, { - "input_cost_per_query": 11.62e-03, + "input_cost_per_query": 0.01162, "max_results_range": [ 61, 70 ] }, { - "input_cost_per_query": 13.28e-03, + "input_cost_per_query": 0.01328, "max_results_range": [ 71, 80 ] }, { - "input_cost_per_query": 14.94e-03, + "input_cost_per_query": 0.01494, "max_results_range": [ 81, 90 ] }, { - "input_cost_per_query": 16.6e-03, + "input_cost_per_query": 0.0166, "max_results_range": [ 91, 100 @@ -10297,7 +10297,7 @@ } }, "perplexity/search": { - "input_cost_per_query": 5e-03, + "input_cost_per_query": 0.005, "litellm_provider": "perplexity", "mode": "search" }, @@ -10395,7 +10395,7 @@ "supports_embedding_image_input": true }, "embed-multilingual-light-v3.0": { - "input_cost_per_token": 1e-04, + "input_cost_per_token": 0.0001, "litellm_provider": "cohere", "max_input_tokens": 1024, "max_tokens": 1024, @@ -10689,7 +10689,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.3e-07, "supports_function_calling": true, @@ -10700,7 +10700,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.9e-07, "supports_function_calling": true, @@ -10711,7 +10711,7 @@ "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-06, "supports_function_calling": true, @@ -10817,14 +10817,14 @@ "litellm_provider": "featherless_ai", "max_input_tokens": 32768, "max_output_tokens": 4096, - "max_tokens": 32768, + "max_tokens": 4096, "mode": "chat" }, "featherless_ai/featherless-ai/Qwerky-QwQ-32B": { "litellm_provider": "featherless_ai", "max_input_tokens": 32768, "max_output_tokens": 4096, - "max_tokens": 32768, + "max_tokens": 4096, "mode": "chat" }, "fireworks-ai-4.1b-to-16b": { @@ -11031,7 +11031,7 @@ "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/glm-4p6": { - "input_cost_per_token": 0.55e-06, + "input_cost_per_token": 5.5e-07, "output_cost_per_token": 2.19e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 202800, @@ -11077,7 +11077,7 @@ "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2.5e-06, "source": "https://fireworks.ai/models/fireworks/kimi-k2-instruct", @@ -11090,7 +11090,7 @@ "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, - "max_tokens": 262144, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.5e-06, "source": "https://app.fireworks.ai/models/fireworks/kimi-k2-instruct-0905", @@ -11337,7 +11337,7 @@ "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, - "max_tokens": 16384, + "max_tokens": 4096, "mode": "completion", "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 2e-07 @@ -11348,7 +11348,7 @@ "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, - "max_tokens": 16384, + "max_tokens": 4096, "mode": "completion", "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 1e-06 @@ -11638,7 +11638,7 @@ "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 8192, "max_output_tokens": 2048, - "max_tokens": 8192, + "max_tokens": 2048, "mode": "chat", "output_cost_per_character": 3.75e-07, "output_cost_per_token": 1.5e-06, @@ -11655,7 +11655,7 @@ "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 8192, "max_output_tokens": 2048, - "max_tokens": 8192, + "max_tokens": 2048, "mode": "chat", "output_cost_per_character": 3.75e-07, "output_cost_per_token": 1.5e-06, @@ -12585,10 +12585,10 @@ "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, - "max_tokens": 65536, + "max_tokens": 32768, "mode": "image_generation", "output_cost_per_image": 0.134, - "output_cost_per_image_token": 1.2e-04, + "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", @@ -14112,7 +14112,7 @@ "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", @@ -14162,7 +14162,7 @@ "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_pdf_size_mb": 30, - "max_tokens": 8192, + "max_tokens": 65536, "max_video_length": 1, "max_videos_per_prompt": 10, "mode": "chat", @@ -14388,10 +14388,10 @@ "litellm_provider": "gemini", "max_input_tokens": 65536, "max_output_tokens": 32768, - "max_tokens": 65536, + "max_tokens": 32768, "mode": "image_generation", "output_cost_per_image": 0.134, - "output_cost_per_image_token": 1.2e-04, + "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, "rpm": 1000, "tpm": 4000000, @@ -15524,7 +15524,7 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.40, + "output_cost_per_second": 0.4, "source": "https://ai.google.dev/gemini-api/docs/video", "supported_modalities": [ "text" @@ -15552,7 +15552,7 @@ "max_input_tokens": 1024, "max_tokens": 1024, "mode": "video_generation", - "output_cost_per_second": 0.40, + "output_cost_per_second": 0.4, "source": "https://ai.google.dev/gemini-api/docs/video", "supported_modalities": [ "text" @@ -16056,11 +16056,11 @@ "supports_vision": true }, "gpt-3.5-turbo": { - "input_cost_per_token": 0.5e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, - "max_tokens": 4097, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-06, "supports_function_calling": true, @@ -16073,7 +16073,7 @@ "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, - "max_tokens": 16385, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-06, "supports_function_calling": true, @@ -16087,7 +16087,7 @@ "litellm_provider": "openai", "max_input_tokens": 4097, "max_output_tokens": 4096, - "max_tokens": 4097, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, "supports_prompt_caching": true, @@ -16099,7 +16099,7 @@ "litellm_provider": "openai", "max_input_tokens": 4097, "max_output_tokens": 4096, - "max_tokens": 4097, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, "supports_function_calling": true, @@ -16113,7 +16113,7 @@ "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, - "max_tokens": 16385, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, "supports_function_calling": true, @@ -16127,7 +16127,7 @@ "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, - "max_tokens": 16385, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 4e-06, "supports_prompt_caching": true, @@ -16139,7 +16139,7 @@ "litellm_provider": "openai", "max_input_tokens": 16385, "max_output_tokens": 4096, - "max_tokens": 16385, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 4e-06, "supports_prompt_caching": true, @@ -17191,7 +17191,7 @@ "supports_pdf_input": true }, "high/1024-x-1536/gpt-image-1.5": { - "input_cost_per_image": 0.20, + "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", "supported_endpoints": [ @@ -17202,7 +17202,7 @@ "supports_pdf_input": true }, "high/1536-x-1024/gpt-image-1.5": { - "input_cost_per_image": 0.20, + "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", "supported_endpoints": [ @@ -17356,7 +17356,7 @@ "supports_pdf_input": true }, "high/1024-x-1536/gpt-image-1.5-2025-12-16": { - "input_cost_per_image": 0.20, + "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", "supported_endpoints": [ @@ -17367,7 +17367,7 @@ "supports_pdf_input": true }, "high/1536-x-1024/gpt-image-1.5-2025-12-16": { - "input_cost_per_image": 0.20, + "input_cost_per_image": 0.2, "litellm_provider": "openai", "mode": "image_generation", "supported_endpoints": [ @@ -17704,7 +17704,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1.68e-04, + "output_cost_per_token": 0.000168, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -17735,7 +17735,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 1.68e-04, + "output_cost_per_token": 0.000168, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -17767,7 +17767,7 @@ "max_output_tokens": 272000, "max_tokens": 272000, "mode": "responses", - "output_cost_per_token": 1.2e-04, + "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, "supported_endpoints": [ "/v1/batch", @@ -17800,7 +17800,7 @@ "max_output_tokens": 272000, "max_tokens": 272000, "mode": "responses", - "output_cost_per_token": 1.2e-04, + "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, "supported_endpoints": [ "/v1/batch", @@ -18503,7 +18503,7 @@ "lemonade/Qwen3-Coder-30B-A3B-Instruct-GGUF": { "input_cost_per_token": 0, "litellm_provider": "lemonade", - "max_tokens": 262144, + "max_tokens": 32768, "max_input_tokens": 262144, "max_output_tokens": 32768, "mode": "chat", @@ -18515,7 +18515,7 @@ "lemonade/gpt-oss-20b-mxfp4-GGUF": { "input_cost_per_token": 0, "litellm_provider": "lemonade", - "max_tokens": 131072, + "max_tokens": 32768, "max_input_tokens": 131072, "max_output_tokens": 32768, "mode": "chat", @@ -18527,7 +18527,7 @@ "lemonade/gpt-oss-120b-mxfp-GGUF": { "input_cost_per_token": 0, "litellm_provider": "lemonade", - "max_tokens": 131072, + "max_tokens": 32768, "max_input_tokens": 131072, "max_output_tokens": 32768, "mode": "chat", @@ -18539,7 +18539,7 @@ "lemonade/Gemma-3-4b-it-GGUF": { "input_cost_per_token": 0, "litellm_provider": "lemonade", - "max_tokens": 128000, + "max_tokens": 8192, "max_input_tokens": 128000, "max_output_tokens": 8192, "mode": "chat", @@ -18551,7 +18551,7 @@ "lemonade/Qwen3-4B-Instruct-2507-GGUF": { "input_cost_per_token": 0, "litellm_provider": "lemonade", - "max_tokens": 262144, + "max_tokens": 32768, "max_input_tokens": 262144, "max_output_tokens": 32768, "mode": "chat", @@ -18688,11 +18688,11 @@ "groq/moonshotai/kimi-k2-instruct-0905": { "input_cost_per_token": 1e-06, "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 0.5e-06, + "cache_read_input_token_cost": 5e-07, "litellm_provider": "groq", "max_input_tokens": 262144, "max_output_tokens": 16384, - "max_tokens": 278528, + "max_tokens": 16384, "mode": "chat", "supports_function_calling": true, "supports_response_schema": true, @@ -19353,7 +19353,7 @@ "litellm_provider": "lambda_ai", "max_input_tokens": 131072, "max_output_tokens": 8192, - "max_tokens": 131072, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1e-07, "supports_function_calling": true, @@ -19366,7 +19366,7 @@ "litellm_provider": "lambda_ai", "max_input_tokens": 16384, "max_output_tokens": 8192, - "max_tokens": 16384, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1e-07, "supports_function_calling": true, @@ -19702,7 +19702,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.6e-05, "supports_function_calling": true, @@ -19713,7 +19713,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 2048, - "max_tokens": 128000, + "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 9.9e-07, "supports_function_calling": true, @@ -19724,7 +19724,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 2048, - "max_tokens": 128000, + "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 2.2e-07, "supports_function_calling": true, @@ -19735,7 +19735,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 3.5e-07, "supports_function_calling": true, @@ -19747,7 +19747,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1e-07, "supports_function_calling": true, @@ -19758,7 +19758,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-07, "supports_function_calling": true, @@ -19769,7 +19769,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, "supports_function_calling": true, @@ -19851,7 +19851,7 @@ "litellm_provider": "meta_llama", "max_input_tokens": 128000, "max_output_tokens": 4028, - "max_tokens": 128000, + "max_tokens": 4028, "mode": "chat", "source": "https://llama.developer.meta.com/docs/models", "supported_modalities": [ @@ -19867,7 +19867,7 @@ "litellm_provider": "meta_llama", "max_input_tokens": 128000, "max_output_tokens": 4028, - "max_tokens": 128000, + "max_tokens": 4028, "mode": "chat", "source": "https://llama.developer.meta.com/docs/models", "supported_modalities": [ @@ -19883,7 +19883,7 @@ "litellm_provider": "meta_llama", "max_input_tokens": 1000000, "max_output_tokens": 4028, - "max_tokens": 128000, + "max_tokens": 4028, "mode": "chat", "source": "https://llama.developer.meta.com/docs/models", "supported_modalities": [ @@ -19900,7 +19900,7 @@ "litellm_provider": "meta_llama", "max_input_tokens": 10000000, "max_output_tokens": 4028, - "max_tokens": 128000, + "max_tokens": 4028, "mode": "chat", "source": "https://llama.developer.meta.com/docs/models", "supported_modalities": [ @@ -19932,7 +19932,7 @@ ] }, "minimax/speech-02-turbo": { - "input_cost_per_character": 0.00006, + "input_cost_per_character": 6e-05, "litellm_provider": "minimax", "mode": "audio_speech", "supported_endpoints": [ @@ -19948,7 +19948,7 @@ ] }, "minimax/speech-2.6-turbo": { - "input_cost_per_character": 0.00006, + "input_cost_per_character": 6e-05, "litellm_provider": "minimax", "mode": "audio_speech", "supported_endpoints": [ @@ -20278,8 +20278,8 @@ }, "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", - "ocr_cost_per_page": 1e-3, - "annotation_cost_per_page": 3e-3, + "ocr_cost_per_page": 0.001, + "annotation_cost_per_page": 0.003, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" @@ -20288,8 +20288,8 @@ }, "mistral/mistral-ocr-2505-completion": { "litellm_provider": "mistral", - "ocr_cost_per_page": 1e-3, - "annotation_cost_per_page": 3e-3, + "ocr_cost_per_page": 0.001, + "annotation_cost_per_page": 0.003, "mode": "ocr", "supported_endpoints": [ "/v1/ocr" @@ -20349,14 +20349,14 @@ "mode": "embedding" }, "mistral/codestral-embed": { - "input_cost_per_token": 0.15e-06, + "input_cost_per_token": 1.5e-07, "litellm_provider": "mistral", "max_input_tokens": 8192, "max_tokens": 8192, "mode": "embedding" }, "mistral/codestral-embed-2505": { - "input_cost_per_token": 0.15e-06, + "input_cost_per_token": 1.5e-07, "litellm_provider": "mistral", "max_input_tokens": 8192, "max_tokens": 8192, @@ -20757,28 +20757,28 @@ "supports_vision": true }, "moonshot/kimi-k2-thinking": { - "cache_read_input_token_cost": 1.5e-7, - "input_cost_per_token": 6e-7, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 6e-07, "litellm_provider": "moonshot", "max_input_tokens": 262144, "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 2.5e-6, + "output_cost_per_token": 2.5e-06, "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", "supports_function_calling": true, "supports_tool_choice": true, "supports_web_search": true }, "moonshot/kimi-k2-thinking-turbo": { - "cache_read_input_token_cost": 1.5e-7, - "input_cost_per_token": 1.15e-6, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.15e-06, "litellm_provider": "moonshot", "max_input_tokens": 262144, "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8e-6, + "output_cost_per_token": 8e-06, "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", "supports_function_calling": true, "supports_tool_choice": true, @@ -21650,7 +21650,7 @@ "litellm_provider": "oci", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.068e-05, "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", @@ -21662,7 +21662,7 @@ "litellm_provider": "oci", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", @@ -21674,7 +21674,7 @@ "litellm_provider": "oci", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 7.2e-07, "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", @@ -21686,7 +21686,7 @@ "litellm_provider": "oci", "max_input_tokens": 512000, "max_output_tokens": 4000, - "max_tokens": 512000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 7.2e-07, "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", @@ -21698,7 +21698,7 @@ "litellm_provider": "oci", "max_input_tokens": 192000, "max_output_tokens": 4000, - "max_tokens": 192000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 7.2e-07, "source": "https://www.oracle.com/artificial-intelligence/generative-ai/generative-ai-service/pricing", @@ -21770,7 +21770,7 @@ "litellm_provider": "oci", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", @@ -21782,7 +21782,7 @@ "litellm_provider": "oci", "max_input_tokens": 256000, "max_output_tokens": 4000, - "max_tokens": 256000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", @@ -21794,7 +21794,7 @@ "litellm_provider": "oci", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", @@ -21806,7 +21806,7 @@ "litellm_provider": "ollama", "max_input_tokens": 32768, "max_output_tokens": 8192, - "max_tokens": 32768, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.0, "supports_function_calling": false @@ -21844,7 +21844,7 @@ "litellm_provider": "ollama", "max_input_tokens": 32768, "max_output_tokens": 8192, - "max_tokens": 32768, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.0, "supports_function_calling": true @@ -21864,12 +21864,12 @@ "litellm_provider": "ollama", "max_input_tokens": 32768, "max_output_tokens": 8192, - "max_tokens": 32768, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.0, "supports_function_calling": true }, - "ollama/deepseek-v3.1:671b-cloud" : { + "ollama/deepseek-v3.1:671b-cloud": { "input_cost_per_token": 0.0, "litellm_provider": "ollama", "max_input_tokens": 163840, @@ -21879,7 +21879,7 @@ "output_cost_per_token": 0.0, "supports_function_calling": true }, - "ollama/gpt-oss:120b-cloud" : { + "ollama/gpt-oss:120b-cloud": { "input_cost_per_token": 0.0, "litellm_provider": "ollama", "max_input_tokens": 131072, @@ -21889,7 +21889,7 @@ "output_cost_per_token": 0.0, "supports_function_calling": true }, - "ollama/gpt-oss:20b-cloud" : { + "ollama/gpt-oss:20b-cloud": { "input_cost_per_token": 0.0, "litellm_provider": "ollama", "max_input_tokens": 131072, @@ -21904,7 +21904,7 @@ "litellm_provider": "ollama", "max_input_tokens": 32768, "max_output_tokens": 8192, - "max_tokens": 32768, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.0, "supports_function_calling": true @@ -21968,7 +21968,7 @@ "litellm_provider": "ollama", "max_input_tokens": 8192, "max_output_tokens": 8192, - "max_tokens": 32768, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.0, "supports_function_calling": true @@ -22026,7 +22026,7 @@ "litellm_provider": "ollama", "max_input_tokens": 65536, "max_output_tokens": 8192, - "max_tokens": 65536, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.0, "supports_function_calling": true @@ -22084,7 +22084,7 @@ "litellm_provider": "openai", "max_input_tokens": 32768, "max_output_tokens": 0, - "max_tokens": 32768, + "max_tokens": 0, "mode": "moderation", "output_cost_per_token": 0.0 }, @@ -22093,7 +22093,7 @@ "litellm_provider": "openai", "max_input_tokens": 32768, "max_output_tokens": 0, - "max_tokens": 32768, + "max_tokens": 0, "mode": "moderation", "output_cost_per_token": 0.0 }, @@ -22102,7 +22102,7 @@ "litellm_provider": "openai", "max_input_tokens": 32768, "max_output_tokens": 0, - "max_tokens": 32768, + "max_tokens": 0, "mode": "moderation", "output_cost_per_token": 0.0 }, @@ -22156,7 +22156,7 @@ "input_cost_per_token": 1.102e-05, "litellm_provider": "openrouter", "max_output_tokens": 8191, - "max_tokens": 100000, + "max_tokens": 8191, "mode": "chat", "output_cost_per_token": 3.268e-05, "supports_tool_choice": true @@ -22296,7 +22296,7 @@ "input_cost_per_token": 1.63e-06, "litellm_provider": "openrouter", "max_output_tokens": 8191, - "max_tokens": 100000, + "max_tokens": 8191, "mode": "chat", "output_cost_per_token": 5.51e-06, "supports_tool_choice": true @@ -22491,7 +22491,7 @@ "litellm_provider": "openrouter", "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 8192, + "max_tokens": 163840, "mode": "chat", "output_cost_per_token": 8e-07, "supports_assistant_prefill": true, @@ -22506,7 +22506,7 @@ "litellm_provider": "openrouter", "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 8192, + "max_tokens": 163840, "mode": "chat", "output_cost_per_token": 4e-07, "supports_assistant_prefill": true, @@ -22521,7 +22521,7 @@ "litellm_provider": "openrouter", "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 8192, + "max_tokens": 163840, "mode": "chat", "output_cost_per_token": 4e-07, "supports_assistant_prefill": true, @@ -22535,7 +22535,7 @@ "litellm_provider": "openrouter", "max_input_tokens": 66000, "max_output_tokens": 4096, - "max_tokens": 8192, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2.8e-07, "supports_prompt_caching": true, @@ -22693,51 +22693,51 @@ "supports_web_search": true }, "openrouter/google/gemini-3-flash-preview": { - "cache_read_input_token_cost": 5e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 5e-07, - "litellm_provider": "openrouter", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_reasoning_token": 3e-06, - "output_cost_per_token": 3e-06, - "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "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, - "tpm": 800000 + "cache_read_input_token_cost": 5e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 5e-07, + "litellm_provider": "openrouter", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 3e-06, + "output_cost_per_token": 3e-06, + "rpm": 2000, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "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, + "tpm": 800000 }, "openrouter/google/gemini-pro-1.5": { "input_cost_per_image": 0.00265, @@ -22868,13 +22868,13 @@ "supports_tool_choice": true }, "openrouter/minimax/minimax-m2": { - "input_cost_per_token": 2.55e-7, + "input_cost_per_token": 2.55e-07, "litellm_provider": "openrouter", "max_input_tokens": 204800, "max_output_tokens": 204800, - "max_tokens": 32768, + "max_tokens": 204800, "mode": "chat", - "output_cost_per_token": 1.02e-6, + "output_cost_per_token": 1.02e-06, "supports_function_calling": true, "supports_prompt_caching": false, "supports_reasoning": true, @@ -22900,7 +22900,7 @@ "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, - "max_tokens": 262144, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 6e-07, "supports_function_calling": true, @@ -23285,7 +23285,7 @@ "litellm_provider": "openrouter", "max_input_tokens": 400000, "max_output_tokens": 128000, - "max_tokens": 400000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.4e-05, "supports_function_calling": true, @@ -23301,7 +23301,7 @@ "litellm_provider": "openrouter", "max_input_tokens": 128000, "max_output_tokens": 16384, - "max_tokens": 128000, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.4e-05, "supports_function_calling": true, @@ -23315,9 +23315,9 @@ "litellm_provider": "openrouter", "max_input_tokens": 400000, "max_output_tokens": 128000, - "max_tokens": 400000, + "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.68e-04, + "output_cost_per_token": 0.000168, "supports_function_calling": true, "supports_prompt_caching": false, "supports_reasoning": true, @@ -23474,20 +23474,20 @@ "litellm_provider": "openrouter", "max_input_tokens": 8192, "max_output_tokens": 2048, - "max_tokens": 8192, + "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 6.3e-07, "supports_tool_choice": true, "supports_vision": true }, "openrouter/qwen/qwen3-coder": { - "input_cost_per_token": 2.2e-7, + "input_cost_per_token": 2.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262100, "max_output_tokens": 262100, "max_tokens": 262100, "mode": "chat", - "output_cost_per_token": 9.5e-7, + "output_cost_per_token": 9.5e-07, "source": "https://openrouter.ai/qwen/qwen3-coder", "supports_tool_choice": true, "supports_function_calling": true @@ -23530,7 +23530,7 @@ "litellm_provider": "openrouter", "max_input_tokens": 2000000, "max_output_tokens": 30000, - "max_tokens": 2000000, + "max_tokens": 30000, "mode": "chat", "output_cost_per_token": 0, "source": "https://openrouter.ai/x-ai/grok-4-fast:free", @@ -23540,26 +23540,26 @@ "supports_web_search": false }, "openrouter/z-ai/glm-4.6": { - "input_cost_per_token": 4.0e-7, + "input_cost_per_token": 4e-07, "litellm_provider": "openrouter", "max_input_tokens": 202800, "max_output_tokens": 131000, - "max_tokens": 202800, + "max_tokens": 131000, "mode": "chat", - "output_cost_per_token": 1.75e-6, + "output_cost_per_token": 1.75e-06, "source": "https://openrouter.ai/z-ai/glm-4.6", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true }, "openrouter/z-ai/glm-4.6:exacto": { - "input_cost_per_token": 4.5e-7, + "input_cost_per_token": 4.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 202800, "max_output_tokens": 131000, - "max_tokens": 202800, + "max_tokens": 131000, "mode": "chat", - "output_cost_per_token": 1.9e-6, + "output_cost_per_token": 1.9e-06, "source": "https://openrouter.ai/z-ai/glm-4.6:exacto", "supports_function_calling": true, "supports_reasoning": true, @@ -24107,7 +24107,7 @@ "litellm_provider": "publicai", "max_input_tokens": 8192, "max_output_tokens": 4096, - "max_tokens": 8192, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24119,7 +24119,7 @@ "litellm_provider": "publicai", "max_input_tokens": 8192, "max_output_tokens": 4096, - "max_tokens": 8192, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24131,7 +24131,7 @@ "litellm_provider": "publicai", "max_input_tokens": 8192, "max_output_tokens": 4096, - "max_tokens": 8192, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24143,7 +24143,7 @@ "litellm_provider": "publicai", "max_input_tokens": 16384, "max_output_tokens": 4096, - "max_tokens": 16384, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24155,7 +24155,7 @@ "litellm_provider": "publicai", "max_input_tokens": 8192, "max_output_tokens": 4096, - "max_tokens": 8192, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24167,7 +24167,7 @@ "litellm_provider": "publicai", "max_input_tokens": 32768, "max_output_tokens": 4096, - "max_tokens": 32768, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24179,7 +24179,7 @@ "litellm_provider": "publicai", "max_input_tokens": 32768, "max_output_tokens": 4096, - "max_tokens": 32768, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24191,7 +24191,7 @@ "litellm_provider": "publicai", "max_input_tokens": 32768, "max_output_tokens": 4096, - "max_tokens": 32768, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24204,7 +24204,7 @@ "litellm_provider": "publicai", "max_input_tokens": 32768, "max_output_tokens": 4096, - "max_tokens": 32768, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", @@ -24217,7 +24217,7 @@ "litellm_provider": "bedrock_converse", "max_input_tokens": 262000, "max_output_tokens": 65536, - "max_tokens": 262144, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1.8e-06, "supports_function_calling": true, @@ -24229,7 +24229,7 @@ "litellm_provider": "bedrock_converse", "max_input_tokens": 262144, "max_output_tokens": 131072, - "max_tokens": 262144, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 8.8e-07, "supports_function_calling": true, @@ -24241,9 +24241,9 @@ "litellm_provider": "bedrock_converse", "max_input_tokens": 262144, "max_output_tokens": 131072, - "max_tokens": 262144, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 6.0e-07, + "output_cost_per_token": 6e-07, "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true @@ -24253,9 +24253,9 @@ "litellm_provider": "bedrock_converse", "max_input_tokens": 131072, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 6.0e-07, + "output_cost_per_token": 6e-07, "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true @@ -24756,12 +24756,11 @@ "supports_reasoning": true, "source": "https://cloud.sambanova.ai/plans/pricing" }, - "snowflake/claude-3-5-sonnet": { "litellm_provider": "snowflake", "max_input_tokens": 18000, "max_output_tokens": 8192, - "max_tokens": 18000, + "max_tokens": 8192, "mode": "chat", "supports_computer_use": true }, @@ -24769,7 +24768,7 @@ "litellm_provider": "snowflake", "max_input_tokens": 32768, "max_output_tokens": 8192, - "max_tokens": 32768, + "max_tokens": 8192, "mode": "chat", "supports_reasoning": true }, @@ -24777,293 +24776,339 @@ "litellm_provider": "snowflake", "max_input_tokens": 8000, "max_output_tokens": 8192, - "max_tokens": 8000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/jamba-1.5-large": { "litellm_provider": "snowflake", "max_input_tokens": 256000, "max_output_tokens": 8192, - "max_tokens": 256000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/jamba-1.5-mini": { "litellm_provider": "snowflake", "max_input_tokens": 256000, "max_output_tokens": 8192, - "max_tokens": 256000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/jamba-instruct": { "litellm_provider": "snowflake", "max_input_tokens": 256000, "max_output_tokens": 8192, - "max_tokens": 256000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama2-70b-chat": { "litellm_provider": "snowflake", "max_input_tokens": 4096, "max_output_tokens": 8192, - "max_tokens": 4096, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama3-70b": { "litellm_provider": "snowflake", "max_input_tokens": 8000, "max_output_tokens": 8192, - "max_tokens": 8000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama3-8b": { "litellm_provider": "snowflake", "max_input_tokens": 8000, "max_output_tokens": 8192, - "max_tokens": 8000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama3.1-405b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama3.1-70b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama3.1-8b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama3.2-1b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama3.2-3b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/llama3.3-70b": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/mistral-7b": { "litellm_provider": "snowflake", "max_input_tokens": 32000, "max_output_tokens": 8192, - "max_tokens": 32000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/mistral-large": { "litellm_provider": "snowflake", "max_input_tokens": 32000, "max_output_tokens": 8192, - "max_tokens": 32000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/mistral-large2": { "litellm_provider": "snowflake", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/mixtral-8x7b": { "litellm_provider": "snowflake", "max_input_tokens": 32000, "max_output_tokens": 8192, - "max_tokens": 32000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/reka-core": { "litellm_provider": "snowflake", "max_input_tokens": 32000, "max_output_tokens": 8192, - "max_tokens": 32000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/reka-flash": { "litellm_provider": "snowflake", "max_input_tokens": 100000, "max_output_tokens": 8192, - "max_tokens": 100000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/snowflake-arctic": { "litellm_provider": "snowflake", "max_input_tokens": 4096, "max_output_tokens": 8192, - "max_tokens": 4096, + "max_tokens": 8192, "mode": "chat" }, "snowflake/snowflake-llama-3.1-405b": { "litellm_provider": "snowflake", "max_input_tokens": 8000, "max_output_tokens": 8192, - "max_tokens": 8000, + "max_tokens": 8192, "mode": "chat" }, "snowflake/snowflake-llama-3.3-70b": { "litellm_provider": "snowflake", "max_input_tokens": 8000, "max_output_tokens": 8192, - "max_tokens": 8000, + "max_tokens": 8192, "mode": "chat" }, "stability/sd3": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.065, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability/sd3-large": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.065, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability/sd3-large-turbo": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.04, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability/sd3-medium": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.035, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability/sd3.5-large": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.065, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability/sd3.5-large-turbo": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.04, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability/sd3.5-medium": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.035, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability/stable-image-ultra": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.08, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability/inpaint": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.005, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/outpaint": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.004, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/erase": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.005, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/search-and-replace": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.005, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/search-and-recolor": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.005, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/remove-background": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.005, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/replace-background-and-relight": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.008, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/sketch": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.005, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/structure": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.005, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/style": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.005, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/style-transfer": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.008, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/fast": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.002, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/conservative": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.04, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/creative": { "litellm_provider": "stability", "mode": "image_edit", "output_cost_per_image": 0.06, - "supported_endpoints": ["/v1/images/edits"] + "supported_endpoints": [ + "/v1/images/edits" + ] }, "stability/stable-image-core": { "litellm_provider": "stability", "mode": "image_generation", "output_cost_per_image": 0.03, - "supported_endpoints": ["/v1/images/generations"] + "supported_endpoints": [ + "/v1/images/generations" + ] }, "stability.sd3-5-large-v1:0": { "litellm_provider": "bedrock", @@ -25090,13 +25135,13 @@ "litellm_provider": "bedrock", "max_input_tokens": 77, "mode": "image_edit", - "output_cost_per_image": 0.40 + "output_cost_per_image": 0.4 }, "stability.stable-creative-upscale-v1:0": { "litellm_provider": "bedrock", "max_input_tokens": 77, "mode": "image_edit", - "output_cost_per_image": 0.60 + "output_cost_per_image": 0.6 }, "stability.stable-fast-upscale-v1:0": { "litellm_provider": "bedrock", @@ -25204,12 +25249,12 @@ "output_cost_per_pixel": 0.0 }, "linkup/search": { - "input_cost_per_query": 5.87e-03, + "input_cost_per_query": 0.00587, "litellm_provider": "linkup", "mode": "search" }, "linkup/search-deep": { - "input_cost_per_query": 58.67e-03, + "input_cost_per_query": 0.05867, "litellm_provider": "linkup", "mode": "search" }, @@ -25388,7 +25433,7 @@ "litellm_provider": "openai", "max_input_tokens": 32768, "max_output_tokens": 0, - "max_tokens": 32768, + "max_tokens": 0, "mode": "moderation", "output_cost_per_token": 0.0 }, @@ -25397,7 +25442,7 @@ "litellm_provider": "openai", "max_input_tokens": 32768, "max_output_tokens": 0, - "max_tokens": 32768, + "max_tokens": 0, "mode": "moderation", "output_cost_per_token": 0.0 }, @@ -25406,7 +25451,7 @@ "litellm_provider": "openai", "max_input_tokens": 32768, "max_output_tokens": 0, - "max_tokens": 32768, + "max_tokens": 0, "mode": "moderation", "output_cost_per_token": 0.0 }, @@ -25842,7 +25887,7 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.6": { - "input_cost_per_token": 0.6e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, @@ -25925,7 +25970,7 @@ "source": "https://aws.amazon.com/polly/pricing/" }, "aws_polly/long-form": { - "input_cost_per_character": 1e-04, + "input_cost_per_character": 0.0001, "litellm_provider": "aws_polly", "mode": "audio_speech", "supported_endpoints": [ @@ -26357,7 +26402,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.6e-05, "supports_function_calling": true, @@ -26368,7 +26413,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 2048, - "max_tokens": 128000, + "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 9.9e-07, "supports_function_calling": true, @@ -26379,7 +26424,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 2048, - "max_tokens": 128000, + "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 2.2e-07, "supports_function_calling": true, @@ -26390,7 +26435,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 3.5e-07, "supports_function_calling": true, @@ -26402,7 +26447,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1e-07, "supports_function_calling": true, @@ -26413,7 +26458,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-07, "supports_function_calling": true, @@ -26424,7 +26469,7 @@ "litellm_provider": "bedrock", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06, "supports_function_calling": true, @@ -26489,7 +26534,7 @@ "litellm_provider": "bedrock_converse", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-06, "supports_function_calling": true, @@ -26542,7 +26587,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 40960, "max_output_tokens": 16384, - "max_tokens": 40960, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2.4e-07 }, @@ -26551,7 +26596,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 40960, "max_output_tokens": 16384, - "max_tokens": 40960, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07 }, @@ -26560,7 +26605,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 40960, "max_output_tokens": 16384, - "max_tokens": 40960, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3e-07 }, @@ -26569,7 +26614,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 40960, "max_output_tokens": 16384, - "max_tokens": 40960, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 3e-07 }, @@ -26578,7 +26623,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 262144, "max_output_tokens": 66536, - "max_tokens": 262144, + "max_tokens": 66536, "mode": "chat", "output_cost_per_token": 1.6e-06 }, @@ -26587,7 +26632,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 300000, "max_output_tokens": 8192, - "max_tokens": 300000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.4e-07 }, @@ -26596,7 +26641,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.4e-07 }, @@ -26605,7 +26650,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 300000, "max_output_tokens": 8192, - "max_tokens": 300000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 3.2e-06 }, @@ -26625,7 +26670,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 4096, - "max_tokens": 200000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.25e-06 }, @@ -26636,7 +26681,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 4096, - "max_tokens": 200000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 7.5e-05 }, @@ -26647,7 +26692,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 8192, - "max_tokens": 200000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 4e-06 }, @@ -26658,7 +26703,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 8192, - "max_tokens": 200000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.5e-05 }, @@ -26669,7 +26714,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 64000, - "max_tokens": 200000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05 }, @@ -26680,7 +26725,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 32000, - "max_tokens": 200000, + "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 7.5e-05 }, @@ -26691,7 +26736,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 64000, - "max_tokens": 200000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-05 }, @@ -26700,7 +26745,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 256000, "max_output_tokens": 8000, - "max_tokens": 256000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1e-05 }, @@ -26709,7 +26754,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-07 }, @@ -26718,7 +26763,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1e-05 }, @@ -26736,7 +26781,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 2.19e-06 }, @@ -26754,7 +26799,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 9e-07 }, @@ -26763,7 +26808,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_tokens": 1048576, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 6e-07 }, @@ -26772,7 +26817,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 8192, - "max_tokens": 1048576, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 3e-07 }, @@ -26781,7 +26826,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1000000, "max_output_tokens": 65536, - "max_tokens": 1000000, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 2.5e-06 }, @@ -26790,7 +26835,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 65536, - "max_tokens": 1048576, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05 }, @@ -26835,7 +26880,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 32000, "max_output_tokens": 16384, - "max_tokens": 32000, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-06 }, @@ -26862,7 +26907,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07 }, @@ -26871,7 +26916,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 131000, "max_output_tokens": 131072, - "max_tokens": 131000, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 8e-08 }, @@ -26880,7 +26925,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.6e-07 }, @@ -26889,7 +26934,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1e-07 }, @@ -26898,7 +26943,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.5e-07 }, @@ -26907,7 +26952,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07 }, @@ -26916,7 +26961,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 8192, - "max_tokens": 128000, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 7.2e-07 }, @@ -26925,7 +26970,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 131072, "max_output_tokens": 8192, - "max_tokens": 131072, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 6e-07 }, @@ -26934,7 +26979,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 131072, "max_output_tokens": 8192, - "max_tokens": 131072, + "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 3e-07 }, @@ -26943,7 +26988,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 256000, "max_output_tokens": 4000, - "max_tokens": 256000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 9e-07 }, @@ -26970,7 +27015,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 64000, - "max_tokens": 128000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 5e-06 }, @@ -26979,7 +27024,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 64000, - "max_tokens": 128000, + "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 1.5e-06 }, @@ -26988,7 +27033,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 4e-08 }, @@ -26997,7 +27042,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1e-07 }, @@ -27015,7 +27060,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 32000, "max_output_tokens": 4000, - "max_tokens": 32000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 6e-06 }, @@ -27033,7 +27078,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 32000, "max_output_tokens": 4000, - "max_tokens": 32000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 3e-07 }, @@ -27042,7 +27087,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 65536, "max_output_tokens": 2048, - "max_tokens": 65536, + "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 1.2e-06 }, @@ -27051,7 +27096,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.5e-07 }, @@ -27060,7 +27105,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 4000, - "max_tokens": 128000, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 6e-06 }, @@ -27069,7 +27114,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 131072, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 2.2e-06 }, @@ -27078,7 +27123,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 32768, "max_output_tokens": 16384, - "max_tokens": 32768, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.2e-06 }, @@ -27087,7 +27132,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 32768, "max_output_tokens": 16384, - "max_tokens": 32768, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1.9e-06 }, @@ -27096,7 +27141,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 16385, "max_output_tokens": 4096, - "max_tokens": 16385, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-06 }, @@ -27105,7 +27150,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 8192, "max_output_tokens": 4096, - "max_tokens": 8192, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 2e-06 }, @@ -27114,7 +27159,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 4096, - "max_tokens": 128000, + "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 3e-05 }, @@ -27125,7 +27170,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1047576, "max_output_tokens": 32768, - "max_tokens": 1047576, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-06 }, @@ -27136,7 +27181,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1047576, "max_output_tokens": 32768, - "max_tokens": 1047576, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.6e-06 }, @@ -27147,7 +27192,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1047576, "max_output_tokens": 32768, - "max_tokens": 1047576, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-07 }, @@ -27158,7 +27203,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 16384, - "max_tokens": 128000, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05 }, @@ -27169,7 +27214,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 16384, - "max_tokens": 128000, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07 }, @@ -27180,7 +27225,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 100000, - "max_tokens": 200000, + "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 6e-05 }, @@ -27191,7 +27236,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 100000, - "max_tokens": 200000, + "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06 }, @@ -27202,7 +27247,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 100000, - "max_tokens": 200000, + "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06 }, @@ -27213,7 +27258,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 100000, - "max_tokens": 200000, + "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06 }, @@ -27249,7 +27294,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 127000, "max_output_tokens": 8000, - "max_tokens": 127000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1e-06 }, @@ -27258,7 +27303,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 200000, "max_output_tokens": 8000, - "max_tokens": 200000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.5e-05 }, @@ -27267,7 +27312,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 127000, "max_output_tokens": 8000, - "max_tokens": 127000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 5e-06 }, @@ -27276,7 +27321,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 127000, "max_output_tokens": 8000, - "max_tokens": 127000, + "max_tokens": 8000, "mode": "chat", "output_cost_per_token": 8e-06 }, @@ -27285,7 +27330,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 32000, - "max_tokens": 128000, + "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 1.5e-05 }, @@ -27294,7 +27339,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 32768, - "max_tokens": 128000, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.5e-05 }, @@ -27303,7 +27348,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 131072, "max_output_tokens": 4000, - "max_tokens": 131072, + "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1e-05 }, @@ -27375,7 +27420,7 @@ "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 128000, "max_output_tokens": 96000, - "max_tokens": 128000, + "max_tokens": 96000, "mode": "chat", "output_cost_per_token": 1.1e-06 }, @@ -27394,7 +27439,7 @@ "supports_tool_choice": true }, "vertex_ai/chirp": { - "input_cost_per_character": 30e-06, + "input_cost_per_character": 3e-05, "litellm_provider": "vertex_ai", "mode": "audio_speech", "source": "https://cloud.google.com/text-to-speech/pricing", @@ -27938,7 +27983,7 @@ "litellm_provider": "vertex_ai-deepseek_models", "max_input_tokens": 163840, "max_output_tokens": 32768, - "max_tokens": 163840, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 5.4e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", @@ -27957,7 +28002,7 @@ "litellm_provider": "vertex_ai-deepseek_models", "max_input_tokens": 163840, "max_output_tokens": 32768, - "max_tokens": 163840, + "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.68e-06, "output_cost_per_token_batches": 8.4e-07, @@ -28042,10 +28087,10 @@ "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, - "max_tokens": 65536, + "max_tokens": 32768, "mode": "image_generation", "output_cost_per_image": 0.134, - "output_cost_per_image_token": 1.2e-04, + "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" @@ -28154,7 +28199,7 @@ "litellm_provider": "vertex_ai-llama_models", "max_input_tokens": 128000, "max_output_tokens": 2048, - "max_tokens": 128000, + "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 1.6e-05, "source": "https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-3.2-90b-vision-instruct-maas", @@ -28167,7 +28212,7 @@ "litellm_provider": "vertex_ai-llama_models", "max_input_tokens": 128000, "max_output_tokens": 2048, - "max_tokens": 128000, + "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 0.0, "source": "https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-3.2-90b-vision-instruct-maas", @@ -28180,7 +28225,7 @@ "litellm_provider": "vertex_ai-llama_models", "max_input_tokens": 128000, "max_output_tokens": 2048, - "max_tokens": 128000, + "max_tokens": 2048, "metadata": { "notes": "VertexAI states that The Llama 3.1 API service for llama-3.1-70b-instruct-maas and llama-3.1-8b-instruct-maas are in public preview and at no cost." }, @@ -28196,7 +28241,7 @@ "litellm_provider": "vertex_ai-llama_models", "max_input_tokens": 128000, "max_output_tokens": 2048, - "max_tokens": 128000, + "max_tokens": 2048, "metadata": { "notes": "VertexAI states that The Llama 3.2 API service is at no cost during public preview, and will be priced as per dollar-per-1M-tokens at GA." }, @@ -28494,7 +28539,7 @@ "vertex_ai/mistral-ocr-2505": { "litellm_provider": "vertex_ai", "mode": "ocr", - "ocr_cost_per_page": 5e-4, + "ocr_cost_per_page": 0.0005, "supported_endpoints": [ "/v1/ocr" ], @@ -28505,7 +28550,7 @@ "mode": "ocr", "input_cost_per_token": 3e-07, "output_cost_per_token": 1.2e-06, - "ocr_cost_per_page": 3e-04, + "ocr_cost_per_page": 0.0003, "source": "https://cloud.google.com/vertex-ai/pricing" }, "vertex_ai/openai/gpt-oss-120b-maas": { @@ -28993,13 +29038,13 @@ "mode": "chat" }, "watsonx/ibm/granite-3-8b-instruct": { - "input_cost_per_token": 0.2e-06, + "input_cost_per_token": 2e-07, "litellm_provider": "watsonx", "max_input_tokens": 8192, "max_output_tokens": 1024, - "max_tokens": 8192, + "max_tokens": 1024, "mode": "chat", - "output_cost_per_token": 0.2e-06, + "output_cost_per_token": 2e-07, "supports_audio_input": false, "supports_audio_output": false, "supports_function_calling": true, @@ -29015,9 +29060,9 @@ "litellm_provider": "watsonx", "max_input_tokens": 131072, "max_output_tokens": 16384, - "max_tokens": 131072, + "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 10e-06, + "output_cost_per_token": 1e-05, "supports_audio_input": false, "supports_audio_output": false, "supports_function_calling": true, @@ -29056,8 +29101,8 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.6e-06, - "output_cost_per_token": 0.6e-06, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29068,8 +29113,8 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.6e-06, - "output_cost_per_token": 0.6e-06, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29080,8 +29125,8 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.6e-06, - "output_cost_per_token": 0.6e-06, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29092,8 +29137,8 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.2e-06, - "output_cost_per_token": 0.2e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29104,8 +29149,8 @@ "max_tokens": 20480, "max_input_tokens": 20480, "max_output_tokens": 20480, - "input_cost_per_token": 0.06e-06, - "output_cost_per_token": 0.25e-06, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 2.5e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29116,8 +29161,8 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.1e-06, - "output_cost_per_token": 0.1e-06, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29128,8 +29173,8 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.2e-06, - "output_cost_per_token": 0.2e-06, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 2e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29140,8 +29185,8 @@ "max_tokens": 512, "max_input_tokens": 512, "max_output_tokens": 512, - "input_cost_per_token": 0.38e-06, - "output_cost_per_token": 0.38e-06, + "input_cost_per_token": 3.8e-07, + "output_cost_per_token": 3.8e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29152,8 +29197,8 @@ "max_tokens": 512, "max_input_tokens": 512, "max_output_tokens": 512, - "input_cost_per_token": 0.38e-06, - "output_cost_per_token": 0.38e-06, + "input_cost_per_token": 3.8e-07, + "output_cost_per_token": 3.8e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29164,8 +29209,8 @@ "max_tokens": 512, "max_input_tokens": 512, "max_output_tokens": 512, - "input_cost_per_token": 0.38e-06, - "output_cost_per_token": 0.38e-06, + "input_cost_per_token": 3.8e-07, + "output_cost_per_token": 3.8e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29176,8 +29221,8 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.1e-06, - "output_cost_per_token": 0.1e-06, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29188,8 +29233,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.35e-06, - "output_cost_per_token": 0.35e-06, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 3.5e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29200,8 +29245,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.1e-06, - "output_cost_per_token": 0.1e-06, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29212,8 +29257,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.15e-06, - "output_cost_per_token": 0.15e-06, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 1.5e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29236,8 +29281,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.71e-06, - "output_cost_per_token": 0.71e-06, + "input_cost_per_token": 7.1e-07, + "output_cost_per_token": 7.1e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29248,7 +29293,7 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.35e-06, + "input_cost_per_token": 3.5e-07, "output_cost_per_token": 1.4e-06, "litellm_provider": "watsonx", "mode": "chat", @@ -29260,8 +29305,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.35e-06, - "output_cost_per_token": 0.35e-06, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 3.5e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29273,7 +29318,7 @@ "max_input_tokens": 128000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, - "output_cost_per_token": 10e-06, + "output_cost_per_token": 1e-05, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29284,8 +29329,8 @@ "max_tokens": 32000, "max_input_tokens": 32000, "max_output_tokens": 32000, - "input_cost_per_token": 0.1e-06, - "output_cost_per_token": 0.3e-06, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29296,8 +29341,8 @@ "max_tokens": 32000, "max_input_tokens": 32000, "max_output_tokens": 32000, - "input_cost_per_token": 0.1e-06, - "output_cost_per_token": 0.3e-06, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": true, @@ -29308,8 +29353,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.35e-06, - "output_cost_per_token": 0.35e-06, + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 3.5e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29320,8 +29365,8 @@ "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "input_cost_per_token": 0.15e-06, - "output_cost_per_token": 0.6e-06, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, "litellm_provider": "watsonx", "mode": "chat", "supports_function_calling": false, @@ -29611,15 +29656,15 @@ }, "xai/grok-4-fast-reasoning": { "litellm_provider": "xai", - "max_input_tokens": 2e6, - "max_output_tokens": 2e6, - "max_tokens": 2e6, + "max_input_tokens": 2000000.0, + "max_output_tokens": 2000000.0, + "max_tokens": 2000000.0, "mode": "chat", - "input_cost_per_token": 0.2e-06, - "input_cost_per_token_above_128k_tokens": 0.4e-06, - "output_cost_per_token": 0.5e-06, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, + "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, - "cache_read_input_token_cost": 0.05e-06, + "cache_read_input_token_cost": 5e-08, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_tool_choice": true, @@ -29627,14 +29672,14 @@ }, "xai/grok-4-fast-non-reasoning": { "litellm_provider": "xai", - "max_input_tokens": 2e6, - "max_output_tokens": 2e6, - "cache_read_input_token_cost": 0.05e-06, - "max_tokens": 2e6, + "max_input_tokens": 2000000.0, + "max_output_tokens": 2000000.0, + "cache_read_input_token_cost": 5e-08, + "max_tokens": 2000000.0, "mode": "chat", - "input_cost_per_token": 0.2e-06, - "input_cost_per_token_above_128k_tokens": 0.4e-06, - "output_cost_per_token": 0.5e-06, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, + "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, @@ -29650,7 +29695,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_128k_tokens": 30e-06, + "output_cost_per_token_above_128k_tokens": 3e-05, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_tool_choice": true, @@ -29665,22 +29710,22 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_128k_tokens": 30e-06, + "output_cost_per_token_above_128k_tokens": 3e-05, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_tool_choice": true, "supports_web_search": true }, "xai/grok-4-1-fast": { - "cache_read_input_token_cost": 0.05e-06, - "input_cost_per_token": 0.2e-06, - "input_cost_per_token_above_128k_tokens": 0.4e-06, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", - "max_input_tokens": 2e6, - "max_output_tokens": 2e6, - "max_tokens": 2e6, + "max_input_tokens": 2000000.0, + "max_output_tokens": 2000000.0, + "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 0.5e-06, + "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, @@ -29692,15 +29737,15 @@ "supports_web_search": true }, "xai/grok-4-1-fast-reasoning": { - "cache_read_input_token_cost": 0.05e-06, - "input_cost_per_token": 0.2e-06, - "input_cost_per_token_above_128k_tokens": 0.4e-06, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", - "max_input_tokens": 2e6, - "max_output_tokens": 2e6, - "max_tokens": 2e6, + "max_input_tokens": 2000000.0, + "max_output_tokens": 2000000.0, + "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 0.5e-06, + "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, @@ -29712,15 +29757,15 @@ "supports_web_search": true }, "xai/grok-4-1-fast-reasoning-latest": { - "cache_read_input_token_cost": 0.05e-06, - "input_cost_per_token": 0.2e-06, - "input_cost_per_token_above_128k_tokens": 0.4e-06, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", - "max_input_tokens": 2e6, - "max_output_tokens": 2e6, - "max_tokens": 2e6, + "max_input_tokens": 2000000.0, + "max_output_tokens": 2000000.0, + "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 0.5e-06, + "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning", "supports_audio_input": true, @@ -29732,15 +29777,15 @@ "supports_web_search": true }, "xai/grok-4-1-fast-non-reasoning": { - "cache_read_input_token_cost": 0.05e-06, - "input_cost_per_token": 0.2e-06, - "input_cost_per_token_above_128k_tokens": 0.4e-06, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", - "max_input_tokens": 2e6, - "max_output_tokens": 2e6, - "max_tokens": 2e6, + "max_input_tokens": 2000000.0, + "max_output_tokens": 2000000.0, + "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 0.5e-06, + "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, @@ -29751,15 +29796,15 @@ "supports_web_search": true }, "xai/grok-4-1-fast-non-reasoning-latest": { - "cache_read_input_token_cost": 0.05e-06, - "input_cost_per_token": 0.2e-06, - "input_cost_per_token_above_128k_tokens": 0.4e-06, + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_128k_tokens": 4e-07, "litellm_provider": "xai", - "max_input_tokens": 2e6, - "max_output_tokens": 2e6, - "max_tokens": 2e6, + "max_input_tokens": 2000000.0, + "max_output_tokens": 2000000.0, + "max_tokens": 2000000.0, "mode": "chat", - "output_cost_per_token": 0.5e-06, + "output_cost_per_token": 5e-07, "output_cost_per_token_above_128k_tokens": 1e-06, "source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning", "supports_audio_input": true, @@ -29942,7 +29987,7 @@ "source": "https://docs.z.ai/guides/overview/pricing" }, "vertex_ai/search_api": { - "input_cost_per_query": 1.5e-03, + "input_cost_per_query": 0.0015, "litellm_provider": "vertex_ai", "mode": "vector_store" }, @@ -29954,7 +29999,7 @@ "openai/sora-2": { "litellm_provider": "openai", "mode": "video_generation", - "output_cost_per_video_per_second": 0.10, + "output_cost_per_video_per_second": 0.1, "source": "https://platform.openai.com/docs/api-reference/videos", "supported_modalities": [ "text", @@ -29971,7 +30016,7 @@ "openai/sora-2-pro": { "litellm_provider": "openai", "mode": "video_generation", - "output_cost_per_video_per_second": 0.30, + "output_cost_per_video_per_second": 0.3, "source": "https://platform.openai.com/docs/api-reference/videos", "supported_modalities": [ "text", @@ -29988,7 +30033,7 @@ "azure/sora-2": { "litellm_provider": "azure", "mode": "video_generation", - "output_cost_per_video_per_second": 0.10, + "output_cost_per_video_per_second": 0.1, "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", "supported_modalities": [ "text" @@ -30004,7 +30049,7 @@ "azure/sora-2-pro": { "litellm_provider": "azure", "mode": "video_generation", - "output_cost_per_video_per_second": 0.30, + "output_cost_per_video_per_second": 0.3, "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", "supported_modalities": [ "text" @@ -30020,7 +30065,7 @@ "azure/sora-2-pro-high-res": { "litellm_provider": "azure", "mode": "video_generation", - "output_cost_per_video_per_second": 0.50, + "output_cost_per_video_per_second": 0.5, "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", "supported_modalities": [ "text" @@ -32191,6 +32236,1100 @@ "litellm_provider": "fireworks_ai", "mode": "chat" }, + "novita/deepseek/deepseek-v3.2": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.69e-03, + "output_cost_per_token": 4e-03, + "max_input_tokens": 163840, + "max_output_tokens": 65536, + "max_tokens": 65536, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 1.345e-03, + "input_cost_per_token_cache_hit": 1.345e-03, + "supports_reasoning": true + }, + "novita/minimax/minimax-m2.1": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3e-03, + "output_cost_per_token": 1.2e-02, + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 3e-04, + "input_cost_per_token_cache_hit": 3e-04, + "supports_reasoning": true + }, + "novita/zai-org/glm-4.7": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 6e-03, + "output_cost_per_token": 2.2e-02, + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 1.1e-03, + "input_cost_per_token_cache_hit": 1.1e-03, + "supports_reasoning": true + }, + "novita/xiaomimimo/mimo-v2-flash": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1e-03, + "output_cost_per_token": 3e-03, + "max_input_tokens": 262144, + "max_output_tokens": 32000, + "max_tokens": 32000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 2e-04, + "input_cost_per_token_cache_hit": 2e-04, + "supports_reasoning": true + }, + "novita/zai-org/autoglm-phone-9b-multilingual": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3.5e-04, + "output_cost_per_token": 1.38e-03, + "max_input_tokens": 65536, + "max_output_tokens": 65536, + "max_tokens": 65536, + "supports_vision": true, + "supports_system_messages": true + }, + "novita/moonshotai/kimi-k2-thinking": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.8e-03, + "output_cost_per_token": 2e-02, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/minimax/minimax-m2": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.4e-03, + "output_cost_per_token": 9.6e-03, + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "cache_read_input_token_cost": 2.4e-04, + "input_cost_per_token_cache_hit": 2.4e-04, + "supports_reasoning": true + }, + "novita/paddlepaddle/paddleocr-vl": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.6e-04, + "output_cost_per_token": 1.6e-04, + "max_input_tokens": 16384, + "max_output_tokens": 16384, + "max_tokens": 16384, + "supports_vision": true, + "supports_system_messages": true + }, + "novita/deepseek/deepseek-v3.2-exp": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.16e-03, + "output_cost_per_token": 3.28e-03, + "max_input_tokens": 163840, + "max_output_tokens": 65536, + "max_tokens": 65536, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/qwen/qwen3-vl-235b-a22b-thinking": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 7.84e-03, + "output_cost_per_token": 3.16e-02, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_vision": true, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/zai-org/glm-4.6v": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3e-03, + "output_cost_per_token": 9e-03, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 5.5e-04, + "input_cost_per_token_cache_hit": 5.5e-04, + "supports_reasoning": true + }, + "novita/zai-org/glm-4.6": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.4e-03, + "output_cost_per_token": 1.76e-02, + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 8.8e-04, + "input_cost_per_token_cache_hit": 8.8e-04, + "supports_reasoning": true + }, + "novita/qwen/qwen3-next-80b-a3b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.2e-03, + "output_cost_per_token": 1.2e-02, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/qwen/qwen3-next-80b-a3b-thinking": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.2e-03, + "output_cost_per_token": 1.2e-02, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/deepseek/deepseek-ocr": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.4e-04, + "output_cost_per_token": 2.4e-04, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/deepseek/deepseek-v3.1-terminus": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.16e-03, + "output_cost_per_token": 8e-03, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 1.08e-03, + "input_cost_per_token_cache_hit": 1.08e-03, + "supports_reasoning": true + }, + "novita/qwen/qwen3-vl-235b-a22b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.4e-03, + "output_cost_per_token": 1.2e-02, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/qwen/qwen3-max": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.688e-02, + "output_cost_per_token": 6.76e-02, + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/skywork/r1v4-lite": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2e-03, + "output_cost_per_token": 6e-03, + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/deepseek/deepseek-v3.1": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.16e-03, + "output_cost_per_token": 8e-03, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 1.08e-03, + "input_cost_per_token_cache_hit": 1.08e-03, + "supports_reasoning": true + }, + "novita/moonshotai/kimi-k2-0905": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.8e-03, + "output_cost_per_token": 2e-02, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/qwen/qwen3-coder-480b-a35b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.4e-03, + "output_cost_per_token": 1.04e-02, + "max_input_tokens": 262144, + "max_output_tokens": 65536, + "max_tokens": 65536, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/qwen/qwen3-coder-30b-a3b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 7e-04, + "output_cost_per_token": 2.7e-03, + "max_input_tokens": 160000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/openai/gpt-oss-120b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4e-04, + "output_cost_per_token": 2e-03, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/moonshotai/kimi-k2-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.56e-03, + "output_cost_per_token": 1.84e-02, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/deepseek/deepseek-v3-0324": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.16e-03, + "output_cost_per_token": 8.96e-03, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 1.08e-03, + "input_cost_per_token_cache_hit": 1.08e-03 + }, + "novita/zai-org/glm-4.5": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.8e-03, + "output_cost_per_token": 1.76e-02, + "max_input_tokens": 131072, + "max_output_tokens": 98304, + "max_tokens": 98304, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "cache_read_input_token_cost": 8.8e-04, + "input_cost_per_token_cache_hit": 8.8e-04, + "supports_reasoning": true + }, + "novita/qwen/qwen3-235b-a22b-thinking-2507": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.4e-03, + "output_cost_per_token": 2.4e-02, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/meta-llama/llama-3.1-8b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2e-04, + "output_cost_per_token": 5e-04, + "max_input_tokens": 16384, + "max_output_tokens": 16384, + "max_tokens": 16384, + "supports_system_messages": true + }, + "novita/google/gemma-3-12b-it": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4e-04, + "output_cost_per_token": 8e-04, + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/zai-org/glm-4.5v": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.8e-03, + "output_cost_per_token": 1.44e-02, + "max_input_tokens": 65536, + "max_output_tokens": 16384, + "max_tokens": 16384, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 8.8e-04, + "input_cost_per_token_cache_hit": 8.8e-04, + "supports_reasoning": true + }, + "novita/openai/gpt-oss-20b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3.2e-04, + "output_cost_per_token": 1.2e-03, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/qwen/qwen3-235b-a22b-instruct-2507": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 7.2e-04, + "output_cost_per_token": 4.64e-03, + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/deepseek/deepseek-r1-distill-qwen-14b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.2e-03, + "output_cost_per_token": 1.2e-03, + "max_input_tokens": 32768, + "max_output_tokens": 16384, + "max_tokens": 16384, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/meta-llama/llama-3.3-70b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.08e-03, + "output_cost_per_token": 3.2e-03, + "max_input_tokens": 131072, + "max_output_tokens": 120000, + "max_tokens": 120000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true + }, + "novita/qwen/qwen-2.5-72b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3.04e-03, + "output_cost_per_token": 3.2e-03, + "max_input_tokens": 32000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/mistralai/mistral-nemo": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3.2e-04, + "output_cost_per_token": 1.36e-03, + "max_input_tokens": 60288, + "max_output_tokens": 16000, + "max_tokens": 16000, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/minimaxai/minimax-m1-80k": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.4e-03, + "output_cost_per_token": 1.76e-02, + "max_input_tokens": 1000000, + "max_output_tokens": 40000, + "max_tokens": 40000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/deepseek/deepseek-r1-0528": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5.6e-03, + "output_cost_per_token": 2e-02, + "max_input_tokens": 163840, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "cache_read_input_token_cost": 2.8e-03, + "input_cost_per_token_cache_hit": 2.8e-03, + "supports_reasoning": true + }, + "novita/deepseek/deepseek-r1-distill-qwen-32b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.4e-03, + "output_cost_per_token": 2.4e-03, + "max_input_tokens": 64000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/meta-llama/llama-3-8b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3.2e-04, + "output_cost_per_token": 3.2e-04, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_system_messages": true + }, + "novita/microsoft/wizardlm-2-8x22b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.96e-03, + "output_cost_per_token": 4.96e-03, + "max_input_tokens": 65535, + "max_output_tokens": 8000, + "max_tokens": 8000, + "supports_system_messages": true + }, + "novita/deepseek/deepseek-r1-0528-qwen3-8b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 4.8e-04, + "output_cost_per_token": 7.2e-04, + "max_input_tokens": 128000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/deepseek/deepseek-r1-distill-llama-70b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 6.4e-03, + "output_cost_per_token": 6.4e-03, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/meta-llama/llama-3-70b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5.1e-03, + "output_cost_per_token": 7.4e-03, + "max_input_tokens": 8192, + "max_output_tokens": 8000, + "max_tokens": 8000, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/qwen/qwen3-235b-a22b-fp8": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.6e-03, + "output_cost_per_token": 6.4e-03, + "max_input_tokens": 40960, + "max_output_tokens": 20000, + "max_tokens": 20000, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/meta-llama/llama-4-maverick-17b-128e-instruct-fp8": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.6e-03, + "output_cost_per_token": 7.2e-03, + "max_input_tokens": 1048576, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_vision": true, + "supports_system_messages": true + }, + "novita/meta-llama/llama-4-scout-17b-16e-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 8e-04, + "output_cost_per_token": 4e-03, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "supports_vision": true, + "supports_system_messages": true + }, + "novita/nousresearch/hermes-2-pro-llama-3-8b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.4e-03, + "output_cost_per_token": 1.4e-03, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/qwen/qwen2.5-vl-72b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 6.4e-03, + "output_cost_per_token": 6.4e-03, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_vision": true, + "supports_system_messages": true + }, + "novita/sao10k/l3-70b-euryale-v2.1": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.48e-02, + "output_cost_per_token": 1.48e-02, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true + }, + "novita/baidu/ernie-4.5-21B-a3b-thinking": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5.6e-04, + "output_cost_per_token": 2.24e-03, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/sao10k/l3-8b-lunaris": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5e-04, + "output_cost_per_token": 5e-04, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/baichuan/baichuan-m2-32b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5.6e-04, + "output_cost_per_token": 5.6e-04, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/thudm/glm-4.1v-9b-thinking": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.8e-04, + "output_cost_per_token": 1.104e-03, + "max_input_tokens": 65536, + "max_output_tokens": 8000, + "max_tokens": 8000, + "supports_vision": true, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/baidu/ernie-4.5-vl-424b-a47b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3.36e-03, + "output_cost_per_token": 1e-02, + "max_input_tokens": 123000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "supports_vision": true, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/baidu/ernie-4.5-300b-a47b-paddle": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.24e-03, + "output_cost_per_token": 8.8e-03, + "max_input_tokens": 123000, + "max_output_tokens": 12000, + "max_tokens": 12000, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/deepseek/deepseek-prover-v2-671b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5.6e-03, + "output_cost_per_token": 2e-02, + "max_input_tokens": 160000, + "max_output_tokens": 160000, + "max_tokens": 160000, + "supports_system_messages": true + }, + "novita/qwen/qwen3-32b-fp8": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 8e-04, + "output_cost_per_token": 3.6e-03, + "max_input_tokens": 40960, + "max_output_tokens": 20000, + "max_tokens": 20000, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/qwen/qwen3-30b-a3b-fp8": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 7.2e-04, + "output_cost_per_token": 3.6e-03, + "max_input_tokens": 40960, + "max_output_tokens": 20000, + "max_tokens": 20000, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/google/gemma-3-27b-it": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 9.52e-04, + "output_cost_per_token": 1.6e-03, + "max_input_tokens": 98304, + "max_output_tokens": 16384, + "max_tokens": 16384, + "supports_vision": true, + "supports_system_messages": true + }, + "novita/deepseek/deepseek-v3-turbo": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3.2e-03, + "output_cost_per_token": 1.04e-02, + "max_input_tokens": 64000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true + }, + "novita/deepseek/deepseek-r1-turbo": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5.6e-03, + "output_cost_per_token": 2e-02, + "max_input_tokens": 64000, + "max_output_tokens": 16000, + "max_tokens": 16000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/Sao10K/L3-8B-Stheno-v3.2": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5e-04, + "output_cost_per_token": 5e-04, + "max_input_tokens": 8192, + "max_output_tokens": 32000, + "max_tokens": 32000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true + }, + "novita/gryphe/mythomax-l2-13b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 7.2e-04, + "output_cost_per_token": 7.2e-04, + "max_input_tokens": 4096, + "max_output_tokens": 3200, + "max_tokens": 3200, + "supports_system_messages": true + }, + "novita/baidu/ernie-4.5-vl-28b-a3b-thinking": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 3.9e-03, + "output_cost_per_token": 3.9e-03, + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true, + "supports_reasoning": true + }, + "novita/qwen/qwen3-vl-8b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 6.4e-04, + "output_cost_per_token": 4e-03, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/zai-org/glm-4.5-air": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.04e-03, + "output_cost_per_token": 6.8e-03, + "max_input_tokens": 131072, + "max_output_tokens": 98304, + "max_tokens": 98304, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/qwen/qwen3-vl-30b-a3b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.6e-03, + "output_cost_per_token": 5.6e-03, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/qwen/qwen3-vl-30b-a3b-thinking": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.6e-03, + "output_cost_per_token": 8e-03, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/qwen/qwen-mt-plus": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2e-03, + "output_cost_per_token": 6e-03, + "max_input_tokens": 16384, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_system_messages": true + }, + "novita/baidu/ernie-4.5-vl-28b-a3b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.12e-03, + "output_cost_per_token": 4.48e-03, + "max_input_tokens": 30000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/baidu/ernie-4.5-21B-a3b": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5.6e-04, + "output_cost_per_token": 2.24e-03, + "max_input_tokens": 120000, + "max_output_tokens": 8000, + "max_tokens": 8000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true + }, + "novita/qwen/qwen3-8b-fp8": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.8e-04, + "output_cost_per_token": 1.104e-03, + "max_input_tokens": 128000, + "max_output_tokens": 20000, + "max_tokens": 20000, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/qwen/qwen3-4b-fp8": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.4e-04, + "output_cost_per_token": 2.4e-04, + "max_input_tokens": 128000, + "max_output_tokens": 20000, + "max_tokens": 20000, + "supports_system_messages": true, + "supports_reasoning": true + }, + "novita/qwen/qwen2.5-7b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 5.6e-04, + "output_cost_per_token": 5.6e-04, + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true, + "supports_response_schema": true + }, + "novita/meta-llama/llama-3.2-3b-instruct": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 2.4e-04, + "output_cost_per_token": 4e-04, + "max_input_tokens": 32768, + "max_output_tokens": 32000, + "max_tokens": 32000, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true + }, + "novita/sao10k/l31-70b-euryale-v2.2": { + "litellm_provider": "novita", + "mode": "chat", + "input_cost_per_token": 1.48e-02, + "output_cost_per_token": 1.48e-02, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_tool_choice": true, + "supports_system_messages": true + }, + "novita/qwen/qwen3-embedding-0.6b": { + "litellm_provider": "novita", + "mode": "embedding", + "input_cost_per_token": 5.6e-04, + "output_cost_per_token": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768 + }, + "novita/qwen/qwen3-embedding-8b": { + "litellm_provider": "novita", + "mode": "embedding", + "input_cost_per_token": 5.6e-04, + "output_cost_per_token": 0, + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 4096 + }, + "novita/baai/bge-m3": { + "litellm_provider": "novita", + "mode": "embedding", + "input_cost_per_token": 1e-04, + "output_cost_per_token": 1e-04, + "max_input_tokens": 8192, + "max_output_tokens": 96000, + "max_tokens": 96000 + }, + "novita/qwen/qwen3-reranker-8b": { + "litellm_provider": "novita", + "mode": "rerank", + "input_cost_per_token": 4e-04, + "output_cost_per_token": 4e-04, + "max_input_tokens": 32768, + "max_output_tokens": 4096, + "max_tokens": 4096 + }, + "novita/baai/bge-reranker-v2-m3": { + "litellm_provider": "novita", + "mode": "rerank", + "input_cost_per_token": 1e-04, + "output_cost_per_token": 1e-04, + "max_input_tokens": 8000, + "max_output_tokens": 8000, + "max_tokens": 8000 + }, "llamagate/llama-3.1-8b": { "max_tokens": 8192, "max_input_tokens": 131072, @@ -32367,4 +33506,3 @@ "mode": "embedding" } } - diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 673aab0990d..8432ce4e874 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -32,6 +32,23 @@ } }, "providers": { + "abliteration": { + "display_name": "Abliteration (`abliteration`)", + "url": "https://docs.litellm.ai/docs/providers/abliteration", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "aiml": { "display_name": "AI/ML API (`aiml`)", "url": "https://docs.litellm.ai/docs/providers/aiml", @@ -1559,7 +1576,7 @@ "chat_completions": true, "messages": true, "responses": true, - "embeddings": false, + "embeddings": true, "image_generations": false, "audio_transcriptions": false, "audio_speech": false, diff --git a/tests/litellm/llms/openai_like/test_abliteration_provider.py b/tests/litellm/llms/openai_like/test_abliteration_provider.py new file mode 100644 index 00000000000..8b8d443fc44 --- /dev/null +++ b/tests/litellm/llms/openai_like/test_abliteration_provider.py @@ -0,0 +1,50 @@ +""" +Unit tests for the Abliteration OpenAI-like provider. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +) + +from litellm.llms.openai_like.dynamic_config import create_config_class +from litellm.llms.openai_like.json_loader import JSONProviderRegistry + +ABLITERATION_BASE_URL = "https://api.abliteration.ai/v1" + + +def _get_config(): + provider = JSONProviderRegistry.get("abliteration") + assert provider is not None + config_class = create_config_class(provider) + return config_class() + + +def test_abliteration_provider_registered(): + provider = JSONProviderRegistry.get("abliteration") + assert provider is not None + assert provider.base_url == ABLITERATION_BASE_URL + assert provider.api_key_env == "ABLITERATION_API_KEY" + + +def test_abliteration_resolves_env_api_key(monkeypatch): + config = _get_config() + monkeypatch.setenv("ABLITERATION_API_KEY", "test-key") + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == ABLITERATION_BASE_URL + assert api_key == "test-key" + + +def test_abliteration_complete_url_appends_endpoint(): + config = _get_config() + url = config.get_complete_url( + api_base=ABLITERATION_BASE_URL, + api_key="test-key", + model="abliteration/abliterated-model", + optional_params={}, + litellm_params={}, + stream=False, + ) + assert url == f"{ABLITERATION_BASE_URL}/chat/completions" diff --git a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py index 6d005af28ac..20f48b6f393 100644 --- a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py @@ -7,6 +7,7 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path from litellm.llms.vertex_ai.gemini import transformation +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig from litellm.types.llms import openai from litellm.types import completion from litellm.types.llms.vertex_ai import RequestBody @@ -225,4 +226,65 @@ async def test__transform_request_body_image_config_with_image_size(): assert "generationConfig" in rb assert "imageConfig" in rb["generationConfig"] assert rb["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" - assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K" \ No newline at end of file + assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K" + + +def test_map_function_google_search_snake_case(): + """ + Test that google_search tool (snake_case) is properly mapped to googleSearch. + Fixes issue where tools=[{"google_search": {}}] was being stripped. + """ + config = VertexGeminiConfig() + optional_params = {} + + # Test snake_case google_search + tools = [{"google_search": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearch" in result[0] + assert result[0]["googleSearch"] == {} + + +def test_map_function_google_search_camel_case(): + """ + Test that googleSearch tool (camelCase) still works. + """ + config = VertexGeminiConfig() + optional_params = {} + + # Test camelCase googleSearch + tools = [{"googleSearch": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearch" in result[0] + assert result[0]["googleSearch"] == {} + + +def test_map_function_google_search_retrieval_snake_case(): + """ + Test that google_search_retrieval tool (snake_case) is properly mapped. + """ + config = VertexGeminiConfig() + optional_params = {} + + tools = [{"google_search_retrieval": {"dynamic_retrieval_config": {"mode": "MODE_DYNAMIC"}}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearchRetrieval" in result[0] + + +def test_map_function_enterprise_web_search_snake_case(): + """ + Test that enterprise_web_search tool (snake_case) is properly mapped. + """ + config = VertexGeminiConfig() + optional_params = {} + + tools = [{"enterprise_web_search": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "enterpriseWebSearch" in result[0] \ No newline at end of file diff --git a/tests/llm_translation/test_bedrock_common_utils.py b/tests/llm_translation/test_bedrock_common_utils.py new file mode 100644 index 00000000000..7b6a05b6988 --- /dev/null +++ b/tests/llm_translation/test_bedrock_common_utils.py @@ -0,0 +1,182 @@ +""" +Unit tests for litellm/llms/bedrock/common_utils.py + +Tests the standalone model name utility functions and BedrockTokenCounter. +""" + +import pytest + +from litellm.llms.bedrock.common_utils import ( + BedrockModelInfo, + extract_model_name_from_bedrock_arn, + get_bedrock_base_model, + get_bedrock_cross_region_inference_regions, + strip_bedrock_routing_prefix, +) +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter + + +class TestStripBedrockRoutingPrefix: + """Tests for strip_bedrock_routing_prefix function.""" + + def test_strips_bedrock_prefix(self): + assert strip_bedrock_routing_prefix("bedrock/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_converse_prefix(self): + assert strip_bedrock_routing_prefix("converse/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_invoke_prefix(self): + assert strip_bedrock_routing_prefix("invoke/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_openai_prefix(self): + assert strip_bedrock_routing_prefix("openai/gpt-4") == "gpt-4" + + def test_strips_all_known_prefixes(self): + # Function strips all known prefixes iteratively + # bedrock/converse/model -> converse/model -> model + assert strip_bedrock_routing_prefix("bedrock/converse/claude-3") == "claude-3" + + def test_no_prefix_unchanged(self): + assert strip_bedrock_routing_prefix("claude-3-sonnet") == "claude-3-sonnet" + + def test_model_with_dots_unchanged(self): + assert ( + strip_bedrock_routing_prefix("anthropic.claude-3-sonnet-20240229-v1:0") + == "anthropic.claude-3-sonnet-20240229-v1:0" + ) + + +class TestExtractModelNameFromBedrockArn: + """Tests for extract_model_name_from_bedrock_arn function.""" + + def test_extracts_from_provisioned_model_arn(self): + arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model-id" + assert extract_model_name_from_bedrock_arn(arn) == "my-model-id" + + def test_extracts_from_foundation_model_arn(self): + arn = "arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2" + assert extract_model_name_from_bedrock_arn(arn) == "anthropic.claude-v2" + + def test_non_arn_unchanged(self): + model = "anthropic.claude-3-sonnet-20240229-v1:0" + assert extract_model_name_from_bedrock_arn(model) == model + + def test_case_insensitive_arn_detection(self): + arn = "ARN:aws:bedrock:us-east-1:123456789012:model/my-model" + assert extract_model_name_from_bedrock_arn(arn) == "my-model" + + +class TestGetBedrockCrossRegionInferenceRegions: + """Tests for get_bedrock_cross_region_inference_regions function.""" + + def test_returns_expected_regions(self): + regions = get_bedrock_cross_region_inference_regions() + assert "us" in regions + assert "eu" in regions + assert "global" in regions + assert "apac" in regions + + def test_returns_list(self): + regions = get_bedrock_cross_region_inference_regions() + assert isinstance(regions, list) + + +class TestGetBedrockBaseModel: + """Tests for get_bedrock_base_model function.""" + + def test_strips_bedrock_prefix(self): + assert get_bedrock_base_model("bedrock/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_converse_prefix(self): + assert get_bedrock_base_model("bedrock/converse/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_us_region_prefix(self): + # us.anthropic.model -> anthropic.model + assert ( + get_bedrock_base_model("us.anthropic.claude-3-sonnet-20240229-v1:0") + == "anthropic.claude-3-sonnet-20240229-v1:0" + ) + + def test_strips_eu_region_prefix(self): + assert ( + get_bedrock_base_model("eu.anthropic.claude-3-sonnet-20240229-v1:0") + == "anthropic.claude-3-sonnet-20240229-v1:0" + ) + + def test_extracts_from_arn(self): + arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model" + assert get_bedrock_base_model(arn) == "my-model" + + def test_model_without_prefix_unchanged(self): + model = "anthropic.claude-3-sonnet-20240229-v1:0" + assert get_bedrock_base_model(model) == model + + def test_combined_bedrock_and_region_prefix(self): + # bedrock/us.anthropic.model -> anthropic.model + assert ( + get_bedrock_base_model("bedrock/us.anthropic.claude-3-sonnet-20240229-v1:0") + == "anthropic.claude-3-sonnet-20240229-v1:0" + ) + + +class TestBedrockModelInfoWrappers: + """Tests that BedrockModelInfo methods correctly wrap standalone functions.""" + + def test_get_base_model_matches_standalone(self): + test_cases = [ + "bedrock/claude-3-sonnet", + "us.anthropic.claude-3-sonnet-20240229-v1:0", + "arn:aws:bedrock:us-east-1:123:model/my-model", + ] + for model in test_cases: + assert BedrockModelInfo.get_base_model(model) == get_bedrock_base_model(model) + + def test_extract_model_name_from_arn_matches_standalone(self): + arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model" + assert ( + BedrockModelInfo.extract_model_name_from_arn(arn) + == extract_model_name_from_bedrock_arn(arn) + ) + + def test_get_non_litellm_routing_model_name_matches_standalone(self): + model = "bedrock/converse/claude-3" + assert ( + BedrockModelInfo.get_non_litellm_routing_model_name(model) + == strip_bedrock_routing_prefix(model) + ) + + +class TestBedrockTokenCounter: + """Tests for BedrockTokenCounter class.""" + + def test_should_use_token_counting_api_for_bedrock(self): + counter = BedrockTokenCounter() + assert counter.should_use_token_counting_api("bedrock") is True + + def test_should_not_use_token_counting_api_for_other_providers(self): + counter = BedrockTokenCounter() + assert counter.should_use_token_counting_api("openai") is False + assert counter.should_use_token_counting_api("anthropic") is False + assert counter.should_use_token_counting_api(None) is False + + def test_get_token_counter_returns_bedrock_token_counter(self): + model_info = BedrockModelInfo() + token_counter = model_info.get_token_counter() + assert isinstance(token_counter, BedrockTokenCounter) + + @pytest.mark.asyncio + async def test_count_tokens_returns_none_for_empty_messages(self): + counter = BedrockTokenCounter() + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=None, + contents=None, + ) + assert result is None + + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=[], + contents=None, + ) + assert result is None diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index ec510b8f953..29c2d26981c 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -2835,6 +2835,34 @@ def test_bedrock_invoke_provider(): ) == "nova" ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider("amazon.nova-pro-v1:0") + == "nova" + ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider( + "amazon.nova-lite-v1:0" + ) + == "nova" + ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider( + "amazon.nova-micro-v1:0" + ) + == "nova" + ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider( + "amazon.nova-premier-v1:0" + ) + == "nova" + ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider( + "amazon.nova-2-lite-v1:0" + ) + == "nova" + ) def test_bedrock_description_param(): @@ -3488,7 +3516,9 @@ def test_bedrock_openai_imported_model(): url = mock_post.call_args.kwargs["url"] print(f"URL: {url}") assert "bedrock-runtime.us-east-1.amazonaws.com" in url - assert "arn:aws:bedrock:us-east-1:117159858402:imported-model/m4gc1mrfuddy" in url + assert ( + "arn:aws:bedrock:us-east-1:117159858402:imported-model/m4gc1mrfuddy" in url + ) assert "/invoke" in url # Validate request body follows OpenAI format @@ -3517,7 +3547,9 @@ def test_bedrock_openai_imported_model(): # Check image_url content assert user_msg["content"][1]["type"] == "image_url" assert "image_url" in user_msg["content"][1] - assert user_msg["content"][1]["image_url"]["url"].startswith("data:image/jpeg;base64,") + assert user_msg["content"][1]["image_url"]["url"].startswith( + "data:image/jpeg;base64," + ) assert user_msg["content"][2]["type"] == "image_url" assert "image_url" in user_msg["content"][2] @@ -3526,21 +3558,67 @@ def test_bedrock_openai_imported_model(): assert request_body["max_tokens"] == 300 assert request_body["temperature"] == 0.5 + +def test_bedrock_nova_provider_detection(): + """ + Test that Nova models are correctly detected even when prefixed with "amazon." + Regression test for issue #17910 where models like "amazon.nova-pro-v1:0" + were incorrectly identified as "amazon" (Titan) instead of "nova". + """ + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + # Test various Nova model formats + nova_test_cases = [ + ("us.amazon.nova-pro-v1:0", "nova"), + ("us.amazon.nova-lite-v1:0", "nova"), + ("us.amazon.nova-micro-v1:0", "nova"), + ("amazon.nova-pro-v1:0", "nova"), + ("amazon.nova-lite-v1:0", "nova"), + ("amazon.nova-micro-v1:0", "nova"), + ("amazon.nova-premier-v1:0", "nova"), + ("amazon.nova-2-lite-v1:0", "nova"), + ("bedrock/amazon.nova-pro-v1:0", "nova"), + ("bedrock/invoke/amazon.nova-pro-v1:0", "nova"), + ("amazon.Nova-pro-v1:0", "nova"), + ("amazon.NOVA-pro-v1:0", "nova"), + ] + + for model, expected in nova_test_cases: + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + assert ( + provider == expected + ), f"Failed for model: {model}, expected: {expected}, got: {provider}" + + # Verify that Amazon Titan models still return "amazon" + titan_test_cases = [ + ("amazon.titan-text-express-v1", "amazon"), + ("us.amazon.titan-text-lite-v1", "amazon"), + ] + + for model, expected in titan_test_cases: + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + assert ( + provider == expected + ), f"Failed for model: {model}, expected: {expected}, got: {provider}" + + def test_bedrock_openai_provider_detection(): """ Test that the OpenAI provider is correctly detected from model strings. """ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - + # Test various OpenAI model formats test_cases = [ "openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123", "bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/xyz789", ] - + for model in test_cases: provider = BaseAWSLLM.get_bedrock_invoke_provider(model) - assert provider == "openai", f"Failed for model: {model}, got provider: {provider}" + assert ( + provider == "openai" + ), f"Failed for model: {model}, got provider: {provider}" print(f"✓ Provider detection works for: {model}") @@ -3549,16 +3627,16 @@ def test_bedrock_openai_model_id_extraction(): Test that the model ID (ARN) is correctly extracted and encoded for OpenAI models. """ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - - model = "openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-model-123" - provider = BaseAWSLLM.get_bedrock_invoke_provider(model) - - model_id = BaseAWSLLM.get_bedrock_model_id( - model=model, - provider=provider, - optional_params={} + + model = ( + "openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-model-123" ) - + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + + model_id = BaseAWSLLM.get_bedrock_model_id( + model=model, provider=provider, optional_params={} + ) + # The ARN should be double URL encoded assert "arn" in model_id assert "imported-model" in model_id @@ -3570,20 +3648,17 @@ def test_bedrock_openai_convert_messages_to_prompt(): Test that convert_messages_to_prompt returns empty string for OpenAI models. """ from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM - + bedrock_llm = BedrockLLM() messages = [ {"role": "system", "content": "You are helpful"}, - {"role": "user", "content": "Hello"} + {"role": "user", "content": "Hello"}, ] - + prompt, chat_history = bedrock_llm.convert_messages_to_prompt( - model="test-model", - messages=messages, - provider="openai", - custom_prompt_dict={} + model="test-model", messages=messages, provider="openai", custom_prompt_dict={} ) - + # OpenAI models use messages directly, no prompt conversion assert prompt == "" assert chat_history is None @@ -3598,37 +3673,33 @@ def test_bedrock_openai_response_parsing(): from litellm import ModelResponse from unittest.mock import Mock import json - + bedrock_llm = BedrockLLM() - + # Mock OpenAI-style response openai_response = { "choices": [ { "message": { "content": "The capital of France is Paris.", - "role": "assistant" + "role": "assistant", }, "finish_reason": "stop", - "index": 0 + "index": 0, } ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 8, - "total_tokens": 18 - } + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, } - + mock_response = Mock() mock_response.json.return_value = openai_response mock_response.text = json.dumps(openai_response) mock_response.status_code = 200 mock_response.headers = {} - + model_response = ModelResponse() mock_logging = Mock() - + result = bedrock_llm.process_response( model="openai/arn:aws:bedrock:us-east-1:123:imported-model/test", response=mock_response, @@ -3640,18 +3711,18 @@ def test_bedrock_openai_response_parsing(): data={}, messages=[{"role": "user", "content": "What is the capital of France?"}], print_verbose=lambda x: None, - encoding=None + encoding=None, ) - + # Verify response content assert result.choices[0].message.content == "The capital of France is Paris." assert result.choices[0].finish_reason == "stop" - + # Verify usage assert result.usage.prompt_tokens == 10 assert result.usage.completion_tokens == 8 assert result.usage.total_tokens == 18 - + print("✓ OpenAI response parsing works correctly") @@ -3659,45 +3730,47 @@ def test_bedrock_openai_request_transformation(): """ Test that the request is correctly transformed for OpenAI models. """ - from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig - + from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig, + ) + config = AmazonInvokeConfig() - + model = "openai/arn:aws:bedrock:us-east-1:123:imported-model/test" messages = [ {"role": "system", "content": "You are helpful"}, - {"role": "user", "content": "Hello"} + {"role": "user", "content": "Hello"}, ] - + optional_params = { "max_tokens": 100, "temperature": 0.7, "top_p": 0.9, - "stream": False + "stream": False, } - + litellm_params = {} headers = {} - - with patch.object(config, 'get_bedrock_invoke_provider', return_value="openai"): + + with patch.object(config, "get_bedrock_invoke_provider", return_value="openai"): result = config.transform_request( model=model, messages=messages, optional_params=optional_params.copy(), litellm_params=litellm_params, - headers=headers + headers=headers, ) - + # Verify the request uses messages format (not prompt) assert "messages" in result assert len(result["messages"]) == 2 assert result["messages"][0]["role"] == "system" assert result["messages"][1]["role"] == "user" - + # Verify parameters are included assert "max_tokens" in result assert "temperature" in result - + print("✓ Request transformation works correctly") @@ -3705,20 +3778,22 @@ def test_bedrock_openai_parameter_filtering(): """ Test that only supported OpenAI parameters are included in the request. """ - from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig - + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, + ) + config = AmazonBedrockOpenAIConfig() model = "test-model" - + supported_params = config.get_supported_openai_params(model=model) - + # Verify common OpenAI parameters are supported assert "max_tokens" in supported_params assert "temperature" in supported_params assert "top_p" in supported_params assert "stream" in supported_params assert "stop" in supported_params - + print(f"✓ Parameter filtering supports: {len(supported_params)} parameters") print(f" Supported params: {supported_params}") @@ -3728,12 +3803,12 @@ def test_bedrock_openai_route_detection(): Test that the OpenAI route is correctly detected. """ from litellm.llms.bedrock.common_utils import BedrockModelInfo - + test_cases = [ ("openai/arn:aws:bedrock:us-east-1:123:imported-model/test", "openai"), ("bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test", "openai"), ] - + for model, expected_route in test_cases: route = BedrockModelInfo.get_bedrock_route(model) assert route == expected_route, f"Failed for model: {model}, got route: {route}" @@ -3745,15 +3820,30 @@ def test_bedrock_openai_explicit_route_check(): Test the explicit OpenAI route checker helper method. """ from litellm.llms.bedrock.common_utils import BedrockModelInfo - + # Test with openai/ prefix - assert BedrockModelInfo._explicit_openai_route("openai/arn:aws:bedrock:us-east-1:123:imported-model/test") is True - assert BedrockModelInfo._explicit_openai_route("bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test") is True - + assert ( + BedrockModelInfo._explicit_openai_route( + "openai/arn:aws:bedrock:us-east-1:123:imported-model/test" + ) + is True + ) + assert ( + BedrockModelInfo._explicit_openai_route( + "bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test" + ) + is True + ) + # Test without openai/ prefix assert BedrockModelInfo._explicit_openai_route("anthropic.claude-3-sonnet") is False - assert BedrockModelInfo._explicit_openai_route("arn:aws:bedrock:us-east-1:123:imported-model/test") is False - + assert ( + BedrockModelInfo._explicit_openai_route( + "arn:aws:bedrock:us-east-1:123:imported-model/test" + ) + is False + ) + print("✓ Explicit route check works correctly") @@ -3761,16 +3851,18 @@ def test_bedrock_openai_config_initialization(): """ Test that AmazonBedrockOpenAIConfig can be properly initialized. """ - from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig - + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, + ) + config = AmazonBedrockOpenAIConfig() - + # Verify it has the necessary methods - assert hasattr(config, 'get_supported_openai_params') - assert hasattr(config, 'transform_request') - assert hasattr(config, 'transform_response') - assert hasattr(config, 'map_openai_params') - + assert hasattr(config, "get_supported_openai_params") + assert hasattr(config, "transform_request") + assert hasattr(config, "transform_response") + assert hasattr(config, "map_openai_params") + print("✓ AmazonBedrockOpenAIConfig initializes correctly") @@ -3779,9 +3871,9 @@ def test_bedrock_openai_multiple_message_types(): Test that various message content types are handled correctly. """ from litellm.llms.custom_httpx.http_handler import HTTPHandler - + client = HTTPHandler() - + # Test with mixed content types messages = [ {"role": "system", "content": "You are helpful"}, @@ -3790,11 +3882,14 @@ def test_bedrock_openai_multiple_message_types(): "role": "user", "content": [ {"type": "text", "text": "Complex message with text"}, - {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,iVBORw0KGg"}} - ] - } + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,iVBORw0KGg"}, + }, + ], + }, ] - + with patch.object(client, "post") as mock_post: try: response = completion( @@ -3805,18 +3900,18 @@ def test_bedrock_openai_multiple_message_types(): ) except Exception as e: pass - + # Verify the request was made if mock_post.called: request_body = json.loads(mock_post.call_args.kwargs["data"]) - + # Verify messages are preserved assert "messages" in request_body assert len(request_body["messages"]) == 3 - + # Verify mixed content is handled assert isinstance(request_body["messages"][2]["content"], list) - + print("✓ Multiple message types handled correctly") @@ -3829,18 +3924,18 @@ def test_bedrock_openai_error_handling(): from litellm.llms.bedrock.common_utils import BedrockError from unittest.mock import Mock import json - + bedrock_llm = BedrockLLM() - + # Mock error response mock_response = Mock() mock_response.json.side_effect = Exception("Invalid JSON") mock_response.text = "Invalid response" mock_response.status_code = 422 - + model_response = ModelResponse() mock_logging = Mock() - + with pytest.raises(BedrockError) as exc_info: bedrock_llm.process_response( model="openai/arn:aws:bedrock:us-east-1:123:imported-model/test", @@ -3853,8 +3948,8 @@ def test_bedrock_openai_error_handling(): data={}, messages=[], print_verbose=lambda x: None, - encoding=None + encoding=None, ) - + assert exc_info.value.status_code == 422 print("✓ Error handling works correctly") diff --git a/tests/llm_translation/test_openrouter.py b/tests/llm_translation/test_openrouter.py index 839d08e12bf..105b05d3449 100644 --- a/tests/llm_translation/test_openrouter.py +++ b/tests/llm_translation/test_openrouter.py @@ -32,3 +32,18 @@ def test_completion_openrouter_image_generation(): .message.images[0]["image_url"]["url"] .startswith("data:image/png;base64,") ) + + +def test_openrouter_embedding(): + """Test OpenRouter embeddings support.""" + litellm._turn_on_debug() + resp = litellm.embedding( + model="openrouter/openai/text-embedding-3-small", + input=["Hello world", "How are you?"], + ) + print(resp) + assert resp is not None + assert len(resp.data) == 2 + assert resp.data[0]["embedding"] is not None + assert isinstance(resp.data[0]["embedding"], list) + assert len(resp.data[0]["embedding"]) > 0 diff --git a/tests/router_unit_tests/test_router_embedding_headers.py b/tests/router_unit_tests/test_router_embedding_headers.py new file mode 100644 index 00000000000..6d480792b7c --- /dev/null +++ b/tests/router_unit_tests/test_router_embedding_headers.py @@ -0,0 +1,372 @@ +""" +Test suite for router embedding method header propagation. + +This tests the fix for the issue where the embedding method was not +propagating proxy model configuration headers to the LLM API calls. + +The fix ensures that router.embedding() calls _update_kwargs_before_fallbacks() +just like router.completion() does, which properly sets up metadata and allows +default_litellm_params (including headers) to be propagated. +""" +import os +import sys +from unittest.mock import MagicMock, patch, AsyncMock + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm import Router + + +class TestRouterEmbeddingHeaders: + """Test that embedding methods properly propagate headers from router configuration.""" + + def test_embedding_calls_update_kwargs_before_fallbacks(self): + """ + Test that router.embedding() calls _update_kwargs_before_fallbacks. + + This ensures that metadata is properly set up before the fallback mechanism, + which is necessary for header propagation to work correctly. + """ + model_list = [ + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fake-key", + }, + } + ] + + router = Router(model_list=model_list) + + # Mock the _update_kwargs_before_fallbacks method to verify it's called + with patch.object( + router, + "_update_kwargs_before_fallbacks", + wraps=router._update_kwargs_before_fallbacks, + ) as mock_update: + with patch("litellm.embedding") as mock_litellm_embedding: + mock_litellm_embedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2, 0.3]}] + ) + + router.embedding(model="text-embedding-ada-002", input=["test input"]) + + # Verify _update_kwargs_before_fallbacks was called + mock_update.assert_called_once() + call_kwargs = mock_update.call_args[1] + assert call_kwargs["model"] == "text-embedding-ada-002" + assert "kwargs" in call_kwargs + + @pytest.mark.asyncio + async def test_aembedding_calls_update_kwargs_before_fallbacks(self): + """ + Test that router.aembedding() calls _update_kwargs_before_fallbacks. + + This ensures consistency between sync and async embedding methods. + """ + model_list = [ + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fake-key", + }, + } + ] + + router = Router(model_list=model_list) + + # Mock the _update_kwargs_before_fallbacks method to verify it's called + with patch.object( + router, + "_update_kwargs_before_fallbacks", + wraps=router._update_kwargs_before_fallbacks, + ) as mock_update: + with patch( + "litellm.aembedding", new_callable=AsyncMock + ) as mock_litellm_aembedding: + mock_litellm_aembedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2, 0.3]}] + ) + + await router.aembedding( + model="text-embedding-ada-002", input=["test input"] + ) + + # Verify _update_kwargs_before_fallbacks was called + mock_update.assert_called_once() + call_kwargs = mock_update.call_args[1] + assert call_kwargs["model"] == "text-embedding-ada-002" + assert "kwargs" in call_kwargs + + def test_embedding_propagates_default_litellm_params(self): + """ + Test that embedding calls properly propagate default_litellm_params including headers. + + This is the main fix - ensuring that headers set in default_litellm_params + are included in the embedding request. + """ + custom_headers = {"X-Custom-Header": "test-value", "X-API-Version": "v2"} + + model_list = [ + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fake-key", + }, + } + ] + + # Create router with default_litellm_params containing headers + router = Router( + model_list=model_list, + default_litellm_params={ + "headers": custom_headers, + "metadata": {"test_key": "test_value"}, + }, + ) + + with patch("litellm.embedding") as mock_litellm_embedding: + mock_litellm_embedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2, 0.3]}] + ) + + router.embedding(model="text-embedding-ada-002", input=["test input"]) + + # Verify that litellm.embedding was called with the headers + mock_litellm_embedding.assert_called_once() + call_kwargs = mock_litellm_embedding.call_args[1] + + # Check that headers were included + assert "headers" in call_kwargs + assert call_kwargs["headers"] == custom_headers + + # Check that metadata was properly set up + assert "metadata" in call_kwargs + assert "model_group" in call_kwargs["metadata"] + assert call_kwargs["metadata"]["model_group"] == "text-embedding-ada-002" + + @pytest.mark.asyncio + async def test_aembedding_propagates_default_litellm_params(self): + """ + Test that async embedding calls properly propagate default_litellm_params including headers. + """ + custom_headers = {"X-Custom-Header": "test-value", "X-API-Version": "v2"} + + model_list = [ + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fake-key", + }, + } + ] + + # Create router with default_litellm_params containing headers + router = Router( + model_list=model_list, + default_litellm_params={ + "headers": custom_headers, + "metadata": {"test_key": "test_value"}, + }, + ) + + with patch( + "litellm.aembedding", new_callable=AsyncMock + ) as mock_litellm_aembedding: + mock_litellm_aembedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2, 0.3]}] + ) + + await router.aembedding( + model="text-embedding-ada-002", input=["test input"] + ) + + # Verify that litellm.aembedding was called with the headers + mock_litellm_aembedding.assert_called_once() + call_kwargs = mock_litellm_aembedding.call_args[1] + + # Check that headers were included + assert "headers" in call_kwargs + assert call_kwargs["headers"] == custom_headers + + # Check that metadata was properly set up + assert "metadata" in call_kwargs + assert "model_group" in call_kwargs["metadata"] + assert call_kwargs["metadata"]["model_group"] == "text-embedding-ada-002" + + def test_embedding_metadata_includes_model_group(self): + """ + Test that embedding calls include model_group in metadata. + + The _update_kwargs_before_fallbacks method should set this up. + """ + model_list = [ + { + "model_name": "test-embedding-model", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fake-key", + }, + } + ] + + router = Router(model_list=model_list) + + with patch("litellm.embedding") as mock_litellm_embedding: + mock_litellm_embedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2, 0.3]}] + ) + + router.embedding(model="test-embedding-model", input=["test input"]) + + call_kwargs = mock_litellm_embedding.call_args[1] + + # Verify metadata contains model_group + assert "metadata" in call_kwargs + assert "model_group" in call_kwargs["metadata"] + assert call_kwargs["metadata"]["model_group"] == "test-embedding-model" + + def test_embedding_sets_num_retries_from_router(self): + """ + Test that embedding calls inherit num_retries from router configuration. + + This is set by _update_kwargs_before_fallbacks. + """ + model_list = [ + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fake-key", + }, + } + ] + + # Create router with num_retries set + router = Router(model_list=model_list, num_retries=3) + + with patch("litellm.embedding") as mock_litellm_embedding: + mock_litellm_embedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2, 0.3]}] + ) + + router.embedding(model="text-embedding-ada-002", input=["test input"]) + + # Verify num_retries was not set in the call (it's handled by function_with_fallbacks) + # The important thing is that it was set in kwargs before being passed to function_with_fallbacks + # We verify this indirectly by checking that _update_kwargs_before_fallbacks was called + mock_litellm_embedding.assert_called_once() + + def test_embedding_sets_litellm_trace_id(self): + """ + Test that embedding calls include a litellm_trace_id. + + This is generated and set by _update_kwargs_before_fallbacks. + """ + model_list = [ + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fake-key", + }, + } + ] + + router = Router(model_list=model_list) + + with patch("litellm.embedding") as mock_litellm_embedding: + mock_litellm_embedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2, 0.3]}] + ) + + router.embedding(model="text-embedding-ada-002", input=["test input"]) + + call_kwargs = mock_litellm_embedding.call_args[1] + + # Verify litellm_trace_id was set + assert "litellm_trace_id" in call_kwargs + assert isinstance(call_kwargs["litellm_trace_id"], str) + assert len(call_kwargs["litellm_trace_id"]) > 0 + + def test_embedding_consistency_with_completion(self): + """ + Test that embedding and completion methods handle kwargs similarly. + + Both should call _update_kwargs_before_fallbacks to ensure consistent behavior. + """ + custom_headers = {"X-Test": "value"} + + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "fake-key", + }, + }, + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fake-key", + }, + }, + ] + + router = Router( + model_list=model_list, default_litellm_params={"headers": custom_headers} + ) + + # Test completion + with patch("litellm.completion") as mock_completion: + mock_completion.return_value = MagicMock() + + router.completion( + model="gpt-3.5-turbo", messages=[{"role": "user", "content": "test"}] + ) + + completion_kwargs = mock_completion.call_args[1] + + # Test embedding + with patch("litellm.embedding") as mock_embedding: + mock_embedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2, 0.3]}] + ) + + router.embedding(model="text-embedding-ada-002", input=["test input"]) + + embedding_kwargs = mock_embedding.call_args[1] + + # Both should have headers from default_litellm_params + assert "headers" in completion_kwargs + assert "headers" in embedding_kwargs + assert completion_kwargs["headers"] == custom_headers + assert embedding_kwargs["headers"] == custom_headers + + # Both should have metadata with model_group + assert "metadata" in completion_kwargs + assert "metadata" in embedding_kwargs + assert "model_group" in completion_kwargs["metadata"] + assert "model_group" in embedding_kwargs["metadata"] + + # Both should have litellm_trace_id + assert "litellm_trace_id" in completion_kwargs + assert "litellm_trace_id" in embedding_kwargs + + +if __name__ == "__main__": + # Run a simple test + test = TestRouterEmbeddingHeaders() + test.test_embedding_calls_update_kwargs_before_fallbacks() + test.test_embedding_propagates_default_litellm_params() + test.test_embedding_metadata_includes_model_group() + test.test_embedding_sets_litellm_trace_id() + test.test_embedding_consistency_with_completion() + print("All tests passed!") # noqa: T201 diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py new file mode 100644 index 00000000000..ab2071714a9 --- /dev/null +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -0,0 +1,355 @@ +""" +Integration tests for router embedding method with various configurations. + +These tests simulate real-world scenarios where headers and configuration +need to be properly propagated through the router to the LLM API. +""" +import os +import sys +from unittest.mock import MagicMock, patch, AsyncMock + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm import Router + + +class TestRouterEmbeddingIntegration: + """Integration tests for embedding with router configuration.""" + + def test_embedding_with_deployment_specific_headers(self): + """ + Test that deployment-specific headers are propagated. + + This simulates a scenario where different deployments have + different header requirements (e.g., different API versions). + """ + model_list = [ + { + "model_name": "embedding-deployment-1", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "key-1", + "headers": {"X-Deployment": "deployment-1"}, + }, + }, + { + "model_name": "embedding-deployment-2", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "key-2", + "headers": {"X-Deployment": "deployment-2"}, + }, + }, + ] + + router = Router(model_list=model_list) + + # Test first deployment + with patch("litellm.embedding") as mock_embedding: + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) + + router.embedding(model="embedding-deployment-1", input=["test"]) + + call_kwargs = mock_embedding.call_args[1] + assert call_kwargs["api_key"] == "key-1" + + # Test second deployment + with patch("litellm.embedding") as mock_embedding: + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) + + router.embedding(model="embedding-deployment-2", input=["test"]) + + call_kwargs = mock_embedding.call_args[1] + assert call_kwargs["api_key"] == "key-2" + + def test_embedding_with_router_and_deployment_headers_merge(self): + """ + Test that router-level headers are propagated. + + When no request headers are provided, router default headers should be used. + """ + model_list = [ + { + "model_name": "test-embedding", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "test-key", + }, + } + ] + + router = Router( + model_list=model_list, + default_litellm_params={ + "headers": { + "X-Router-Header": "router-value", + "X-Common-Header": "router-common", + } + }, + ) + + # Test: No request headers - router headers should be used + with patch("litellm.embedding") as mock_embedding: + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) + + router.embedding( + model="test-embedding", + input=["test"], + ) + + call_kwargs = mock_embedding.call_args[1] + + # Router headers should be present + assert "headers" in call_kwargs + assert call_kwargs["headers"]["X-Router-Header"] == "router-value" + assert call_kwargs["headers"]["X-Common-Header"] == "router-common" + + def test_embedding_metadata_propagation(self): + """ + Test that metadata is properly set up and propagated. + + This is important for logging, tracking, and debugging. + """ + model_list = [ + { + "model_name": "test-embedding", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "test-key", + }, + } + ] + + router = Router( + model_list=model_list, + default_litellm_params={ + "metadata": {"environment": "test", "service": "embedding-service"} + }, + ) + + with patch("litellm.embedding") as mock_embedding: + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) + + router.embedding( + model="test-embedding", + input=["test"], + metadata={"request_id": "req-123"}, # Additional metadata from request + ) + + call_kwargs = mock_embedding.call_args[1] + + # Check metadata contains all expected fields + assert "metadata" in call_kwargs + metadata = call_kwargs["metadata"] + + # From _update_kwargs_before_fallbacks + assert "model_group" in metadata + assert metadata["model_group"] == "test-embedding" + + # From default_litellm_params + assert "environment" in metadata + assert metadata["environment"] == "test" + assert "service" in metadata + assert metadata["service"] == "embedding-service" + + # From request + assert "request_id" in metadata + assert metadata["request_id"] == "req-123" + + @pytest.mark.asyncio + async def test_async_embedding_with_multiple_retries(self): + """ + Test that async embedding properly uses num_retries from router config. + + This ensures the fix works with the retry mechanism. + """ + model_list = [ + { + "model_name": "test-embedding", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "test-key", + }, + } + ] + + router = Router(model_list=model_list, num_retries=2) + + with patch("litellm.aembedding", new_callable=AsyncMock) as mock_aembedding: + mock_aembedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) + + await router.aembedding(model="test-embedding", input=["test"]) + + # The call should succeed + mock_aembedding.assert_called_once() + + def test_embedding_with_timeout_from_router(self): + """ + Test that timeout settings from router config are propagated. + """ + model_list = [ + { + "model_name": "test-embedding", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "test-key", + }, + } + ] + + router = Router(model_list=model_list, timeout=30.0) + + with patch("litellm.embedding") as mock_embedding: + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) + + router.embedding(model="test-embedding", input=["test"]) + + call_kwargs = mock_embedding.call_args[1] + + # Timeout should be set from router config + assert "timeout" in call_kwargs + assert call_kwargs["timeout"] == 30.0 + + def test_embedding_with_multiple_deployments_load_balancing(self): + """ + Test that headers are correctly propagated when router load balances + between multiple deployments. + """ + model_list = [ + { + "model_name": "shared-embedding-model", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "key-1", + }, + }, + { + "model_name": "shared-embedding-model", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "key-2", + }, + }, + ] + + router = Router( + model_list=model_list, + default_litellm_params={"headers": {"X-Shared-Header": "shared-value"}}, + ) + + # Make multiple calls and verify headers are always present + for i in range(5): + with patch("litellm.embedding") as mock_embedding: + mock_embedding.return_value = MagicMock( + data=[{"embedding": [0.1, 0.2]}] + ) + + router.embedding(model="shared-embedding-model", input=[f"test {i}"]) + + call_kwargs = mock_embedding.call_args[1] + + # Headers should always be present regardless of which deployment is chosen + assert "headers" in call_kwargs + assert call_kwargs["headers"]["X-Shared-Header"] == "shared-value" + + @pytest.mark.asyncio + async def test_embedding_with_fallback_configuration(self): + """ + Test that headers are propagated correctly when using fallback models. + """ + model_list = [ + { + "model_name": "primary-embedding", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "primary-key", + }, + }, + { + "model_name": "fallback-embedding", + "litellm_params": { + "model": "text-embedding-ada-002", + "api_key": "fallback-key", + }, + }, + ] + + router = Router( + model_list=model_list, + fallbacks=[{"primary-embedding": ["fallback-embedding"]}], + default_litellm_params={"headers": {"X-Fallback-Test": "test-value"}}, + ) + + # Simulate primary failing, fallback succeeding + with patch("litellm.aembedding", new_callable=AsyncMock) as mock_aembedding: + call_count = 0 + + async def side_effect(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + # First call (primary) fails + raise Exception("Primary failed") + else: + # Second call (fallback) succeeds + return MagicMock(data=[{"embedding": [0.1, 0.2]}]) + + mock_aembedding.side_effect = side_effect + + await router.aembedding(model="primary-embedding", input=["test"]) + + # Both calls should have headers + assert mock_aembedding.call_count == 2 + + # Check that both calls had headers + for call_obj in mock_aembedding.call_args_list: + call_kwargs = call_obj[1] + assert "headers" in call_kwargs + assert call_kwargs["headers"]["X-Fallback-Test"] == "test-value" + + def test_embedding_with_custom_provider_headers(self): + """ + Test that provider-specific headers are correctly propagated. + + Some providers require specific headers for API versioning, features, etc. + """ + model_list = [ + { + "model_name": "azure-embedding", + "litellm_params": { + "model": "azure/text-embedding-ada-002", + "api_key": "azure-key", + "api_base": "https://example.openai.azure.com", + "api_version": "2024-02-01", + }, + } + ] + + router = Router( + model_list=model_list, + default_litellm_params={ + "headers": {"X-Custom-Azure-Header": "azure-value"} + }, + ) + + with patch("litellm.embedding") as mock_embedding: + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) + + router.embedding(model="azure-embedding", input=["test"]) + + call_kwargs = mock_embedding.call_args[1] + + # Verify Azure-specific params are present + assert call_kwargs["api_base"] == "https://example.openai.azure.com" + assert call_kwargs["api_version"] == "2024-02-01" + + # Verify custom headers are present + assert "headers" in call_kwargs + assert call_kwargs["headers"]["X-Custom-Azure-Header"] == "azure-value" + + +if __name__ == "__main__": + # Run tests + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py index 31a2f6cbf51..b0aac17e7d9 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py @@ -66,16 +66,77 @@ class TestCloudZeroHourlyExport: fake_client = MagicMock() fake_db = MagicMock() - async def query_raw_mock(query: str): - sql_context = pl.SQLContext( - LiteLLM_DailyUserSpend=spend_mock_data, - LiteLLM_VerificationToken=verification_mock_data, - LiteLLM_TeamTable=team_mock_data, - LiteLLM_UserTable=user_mock_data, - ) - result = sql_context.execute(query).collect() + async def query_raw_mock(query: str, *params): + start_time_utc = params[0] if len(params) > 0 else None + end_time_utc = params[1] if len(params) > 1 else None + limit = params[2] if len(params) > 2 else None - return result + spend_df = spend_mock_data.collect() + verification_df = verification_mock_data.collect().rename( + {"key_alias": "api_key_alias"} + ) + team_df = team_mock_data.collect() + user_df = user_mock_data.collect() + + joined = ( + spend_df.join( + verification_df, left_on="api_key", right_on="token", how="left" + ) + .join( + team_df, + left_on="team_id", + right_on="team_id", + how="left", + suffix="_team", + ) + .join( + user_df, + left_on="user_id", + right_on="user_id", + how="left", + suffix="_user", + ) + ) + + for duplicate_column in ("team_id_team", "user_id_user"): + if duplicate_column in joined.columns: + joined = joined.drop(duplicate_column) + + if start_time_utc is not None: + joined = joined.filter(pl.col("updated_at") >= start_time_utc) + if end_time_utc is not None: + joined = joined.filter(pl.col("updated_at") <= end_time_utc) + + joined = joined.select( + [ + "id", + "date", + "user_id", + "api_key", + "model", + "model_group", + "custom_llm_provider", + "prompt_tokens", + "completion_tokens", + "spend", + "api_requests", + "successful_requests", + "failed_requests", + "cache_creation_input_tokens", + "cache_read_input_tokens", + "created_at", + "updated_at", + "team_id", + "api_key_alias", + "team_alias", + "user_email", + ] + ).sort(["date", "created_at"], descending=[True, True]) + + if limit is not None: + joined = joined.head(int(limit)) + + return joined fake_db.query_raw = AsyncMock(side_effect=query_raw_mock) fake_client.db = fake_db diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py new file mode 100644 index 00000000000..89a5028011c --- /dev/null +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py @@ -0,0 +1,57 @@ +"""Tests for LiteLLM CloudZero database helper.""" + +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from litellm.integrations.cloudzero.database import LiteLLMDatabase + + +def _setup_db(monkeypatch: pytest.MonkeyPatch, query_return): + """Return a database instance with prisma client mocked out.""" + query_mock = AsyncMock(return_value=query_return) + mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=query_mock)) + db = LiteLLMDatabase() + monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client) + return db, query_mock + + +@pytest.mark.asyncio +async def test_get_usage_data_parameterized(monkeypatch: pytest.MonkeyPatch): + """Start/end filters and limit should be parameterized via placeholders.""" + start = datetime(2024, 5, 1, tzinfo=timezone.utc) + end = datetime(2024, 5, 2, tzinfo=timezone.utc) + db, query_mock = _setup_db(monkeypatch, []) + + await db.get_usage_data(limit=10, start_time_utc=start, end_time_utc=end) + + query_text, *params = query_mock.await_args.args + assert "dus.updated_at >= $1::timestamptz" in query_text + assert "dus.updated_at <= $2::timestamptz" in query_text + assert "LIMIT $3" in query_text + assert params == [start, end, 10] + + +@pytest.mark.asyncio +async def test_get_usage_data_handles_missing_filters(monkeypatch: pytest.MonkeyPatch): + """When no filters provided the params should be None placeholders.""" + db, query_mock = _setup_db(monkeypatch, []) + + await db.get_usage_data() + + query_text, *params = query_mock.await_args.args + assert "LIMIT $3" not in query_text + assert params == [None, None] + + +@pytest.mark.asyncio +async def test_get_usage_data_rejects_invalid_limit(monkeypatch: pytest.MonkeyPatch): + """limit must coerce to int or raise ValueError before hitting the DB.""" + db, query_mock = _setup_db(monkeypatch, []) + + with pytest.raises(ValueError): + await db.get_usage_data(limit="invalid") + + assert query_mock.await_count == 0 diff --git a/tests/test_litellm/integrations/focus/test_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py similarity index 100% rename from tests/test_litellm/integrations/focus/test_database.py rename to tests/test_litellm/integrations/focus/test_focus_database.py diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index a322dfe9a2b..d7d7720ff43 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -530,7 +530,7 @@ class TestPassthroughCallTypeHandling: ) assert ( ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aembedding") - == "embeddings" + == "embedding" ) assert ( ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aresponses") diff --git a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py b/tests/test_litellm/integrations/test_prometheus_cache_metrics.py new file mode 100644 index 00000000000..660757673f6 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_cache_metrics.py @@ -0,0 +1,211 @@ +""" +Unit tests for cache Prometheus metrics. + +Run with: poetry run pytest tests/test_litellm/integrations/test_prometheus_cache_metrics.py -v +""" +import pytest +from unittest.mock import MagicMock, patch +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + +class TestPrometheusCacheMetrics: + """Tests for cache-related Prometheus metrics""" + + @pytest.fixture + def sample_enum_values(self): + """Create sample enum values for labels""" + return UserAPIKeyLabelValues( + end_user="test-end-user", + hashed_api_key="test-key-hash", + api_key_alias="test-key-alias", + team="test-team", + team_alias="test-team-alias", + user="test-user", + model="gpt-3.5-turbo", + ) + + def test_cache_metrics_defined_in_types(self): + """Test that cache metrics are defined in DEFINED_PROMETHEUS_METRICS""" + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + from typing import get_args + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + + assert "litellm_cache_hits_metric" in defined_metrics + assert "litellm_cache_misses_metric" in defined_metrics + assert "litellm_cached_tokens_metric" in defined_metrics + + def test_cache_metric_labels_defined(self): + """Test that cache metric labels are properly defined""" + from litellm.types.integrations.prometheus import PrometheusMetricLabels + + # Verify labels are defined for each cache metric + assert hasattr(PrometheusMetricLabels, "litellm_cache_hits_metric") + assert hasattr(PrometheusMetricLabels, "litellm_cache_misses_metric") + assert hasattr(PrometheusMetricLabels, "litellm_cached_tokens_metric") + + # Verify labels include expected keys + expected_labels = [ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + for label in expected_labels: + assert label in PrometheusMetricLabels.litellm_cache_hits_metric + assert label in PrometheusMetricLabels.litellm_cache_misses_metric + assert label in PrometheusMetricLabels.litellm_cached_tokens_metric + + def test_increment_cache_metrics_on_cache_hit(self, sample_enum_values): + """Test that cache hit increments the correct metrics""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + # Import the method directly and bind it to our mock + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=True + standard_logging_payload = { + "cache_hit": True, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method using unbound method approach + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify cache hits metric was incremented + mock_logger.litellm_cache_hits_metric.labels.assert_called() + mock_logger.litellm_cache_hits_metric.labels().inc.assert_called_once() + + # Verify cached tokens metric was incremented with total_tokens + mock_logger.litellm_cached_tokens_metric.labels.assert_called() + mock_logger.litellm_cached_tokens_metric.labels().inc.assert_called_once_with( + 100 + ) + + # Verify cache misses metric was NOT called + mock_logger.litellm_cache_misses_metric.labels.assert_not_called() + + def test_increment_cache_metrics_on_cache_miss(self, sample_enum_values): + """Test that cache miss increments the correct metrics""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=False + standard_logging_payload = { + "cache_hit": False, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify cache misses metric was incremented + mock_logger.litellm_cache_misses_metric.labels.assert_called() + mock_logger.litellm_cache_misses_metric.labels().inc.assert_called_once() + + # Verify cache hits and cached tokens metrics were NOT called + mock_logger.litellm_cache_hits_metric.labels.assert_not_called() + mock_logger.litellm_cached_tokens_metric.labels.assert_not_called() + + def test_increment_cache_metrics_when_cache_hit_is_none(self, sample_enum_values): + """Test that no metrics are incremented when cache_hit is None""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=None + standard_logging_payload = { + "cache_hit": None, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify NO metrics were called + mock_logger.litellm_cache_hits_metric.labels.assert_not_called() + mock_logger.litellm_cache_misses_metric.labels.assert_not_called() + mock_logger.litellm_cached_tokens_metric.labels.assert_not_called() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py b/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py new file mode 100644 index 00000000000..0743a9c7ba2 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py @@ -0,0 +1,424 @@ +""" +Unit tests for prometheus queue time and guardrail metrics +""" +from datetime import datetime +from unittest.mock import MagicMock + +import pytest +from prometheus_client import REGISTRY + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + """Clean up prometheus registry between tests""" + # Clear the registry before each test + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + yield + # Clean up after test + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + + +class TestPrometheusQueueTimeMetric: + """Test request queue time metric recording""" + + def test_queue_time_metric_recorded_in_set_latency_metrics(self): + """Test that queue time metric is recorded when queue_time_seconds is present in metadata""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs with queue_time_seconds in metadata + queue_time_seconds = 0.5 + + kwargs = { + "litellm_params": {"metadata": {"queue_time_seconds": queue_time_seconds}}, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should be called + mock_metric.labels.assert_called() + # Check that observe was called on the queue time metric + assert mock_labeled_metric.observe.called + # Verify the observed value + observed_value = None + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0: + observed_value = call[0][0] + if observed_value == queue_time_seconds: + break + assert observed_value == queue_time_seconds + assert observed_value >= 0 + + def test_queue_time_metric_not_recorded_when_missing(self): + """Test that queue time metric is not recorded when queue_time_seconds is missing""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs without queue_time_seconds + kwargs = { + "litellm_params": {"metadata": {}}, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should not be called (queue_time_seconds is None) + # We check that observe was not called with queue_time_seconds + queue_time_called = False + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0 and call[0][0] == 0.5: # Our test queue time value + queue_time_called = True + break + assert ( + not queue_time_called + ), "Queue time metric should not be recorded when queue_time_seconds is missing" + + def test_queue_time_metric_not_recorded_when_negative(self): + """Test that queue time metric is not recorded when queue_time_seconds is negative""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs with negative queue_time_seconds + kwargs = { + "litellm_params": { + "metadata": {"queue_time_seconds": -0.1} # Negative value + }, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should not be called for negative values + # We check that observe was not called with the negative value + negative_value_called = False + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0 and call[0][0] == -0.1: + negative_value_called = True + break + assert ( + not negative_value_called + ), "Queue time metric should not be recorded for negative values" + + +class TestPrometheusGuardrailMetrics: + """Test guardrail metrics recording""" + + def test_record_guardrail_metrics_success(self): + """Test recording guardrail metrics for successful execution""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + mock_errors_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + prometheus_logger.litellm_guardrail_errors_total = mock_errors_metric + + guardrail_name = "test_guardrail" + latency_seconds = 0.15 + status = "success" + error_type = None + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + + # Assert - latency metric should be recorded + mock_latency_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + error_type="none", + hook_type=hook_type, + ) + mock_latency_metric.labels.return_value.observe.assert_called_once_with( + latency_seconds + ) + + # Assert - requests metric should be incremented + mock_requests_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + hook_type=hook_type, + ) + mock_requests_metric.labels.return_value.inc.assert_called_once() + + # Assert - errors metric should NOT be called for success + mock_errors_metric.labels.assert_not_called() + + def test_record_guardrail_metrics_error(self): + """Test recording guardrail metrics for failed execution""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + mock_errors_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + prometheus_logger.litellm_guardrail_errors_total = mock_errors_metric + + guardrail_name = "test_guardrail" + latency_seconds = 0.2 + status = "error" + error_type = "ValueError" + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + + # Assert - latency metric should be recorded + mock_latency_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + mock_latency_metric.labels.return_value.observe.assert_called_once_with( + latency_seconds + ) + + # Assert - requests metric should be incremented + mock_requests_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + hook_type=hook_type, + ) + mock_requests_metric.labels.return_value.inc.assert_called_once() + + # Assert - errors metric should be incremented + mock_errors_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + error_type=error_type, + hook_type=hook_type, + ) + mock_errors_metric.labels.return_value.inc.assert_called_once() + + def test_record_guardrail_metrics_during_call_hook(self): + """Test recording guardrail metrics for during_call hook""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + + guardrail_name = "moderation_guardrail" + latency_seconds = 0.1 + status = "success" + hook_type = "during_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=None, + hook_type=hook_type, + ) + + # Assert - hook_type should be "during_call" + mock_latency_metric.labels.assert_called_once() + call_kwargs = mock_latency_metric.labels.call_args[1] + assert call_kwargs["hook_type"] == "during_call" + + def test_record_guardrail_metrics_handles_exception(self): + """Test that _record_guardrail_metrics handles exceptions gracefully""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metric to raise exception + mock_metric = MagicMock() + mock_metric.labels.side_effect = Exception("Test error") + prometheus_logger.litellm_guardrail_latency_metric = mock_metric + prometheus_logger.litellm_guardrail_requests_total = MagicMock() + + # Act & Assert - should not raise exception + try: + prometheus_logger._record_guardrail_metrics( + guardrail_name="test", + latency_seconds=0.1, + status="success", + error_type=None, + hook_type="pre_call", + ) + except Exception: + pytest.fail("_record_guardrail_metrics should handle exceptions gracefully") + + def test_record_guardrail_metrics_with_guardrail_name_attribute(self): + """Test that guardrail name is extracted from guardrail_name attribute if available""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + + guardrail_name = "custom_guardrail_name" + latency_seconds = 0.1 + status = "success" + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=None, + hook_type=hook_type, + ) + + # Assert - guardrail_name should be used + mock_latency_metric.labels.assert_called_once() + call_kwargs = mock_latency_metric.labels.call_args[1] + assert call_kwargs["guardrail_name"] == guardrail_name diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 4914ec0bfb7..42a2b5d0971 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -620,26 +620,26 @@ def test_bedrock_tools_unpack_defs(): def test_bedrock_image_processor_content_type_fallback_url_extension(): """ - Test that _post_call_image_processing falls back to URL extension + Test that _post_call_image_processing falls back to URL extension when content-type is binary/octet-stream or application/octet-stream """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a simple PNG header (magic bytes) png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" png_content = png_header + b"\x00" * 100 # Add some padding mock_response.content = png_content - + # Test with .png URL image_url = "https://example.com/test-image.png" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/png" assert base64_bytes == base64.b64encode(png_content).decode("utf-8") @@ -650,22 +650,22 @@ def test_bedrock_image_processor_content_type_fallback_binary_detection(): when content-type is missing and URL extension is not recognized """ import base64 - + # Create mock response with no content-type mock_response = MagicMock() mock_response.headers.get.return_value = None - + # Create a JPEG header (magic bytes) jpeg_header = b"\xff\xd8\xff" jpeg_content = jpeg_header + b"\x00" * 100 # Add some padding mock_response.content = jpeg_content - + # Test with URL without extension image_url = "https://example.com/test-image-without-extension" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/jpeg" assert base64_bytes == base64.b64encode(jpeg_content).decode("utf-8") @@ -675,22 +675,22 @@ def test_bedrock_image_processor_content_type_fallback_application_octet_stream( Test that _post_call_image_processing handles application/octet-stream correctly """ import base64 - + # Create mock response with application/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "application/octet-stream" - + # Create a GIF header (magic bytes) gif_header = b"GIF8" + b"\x00" + b"a" gif_content = gif_header + b"\x00" * 100 # Add some padding mock_response.content = gif_content - + # Test with .gif URL image_url = "https://s3.amazonaws.com/bucket/image.gif" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/gif" assert base64_bytes == base64.b64encode(gif_content).decode("utf-8") @@ -700,22 +700,22 @@ def test_bedrock_image_processor_content_type_with_query_params(): Test that _post_call_image_processing correctly extracts extension from URL with query parameters """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a WebP header (magic bytes) webp_header = b"RIFF" + b"\x00\x00\x00\x00" + b"WEBP" webp_content = webp_header + b"\x00" * 100 # Add some padding mock_response.content = webp_content - + # Test with URL containing query parameters (common in S3 signed URLs) image_url = "https://s3.amazonaws.com/bucket/image.webp?AWSAccessKeyId=123&Expires=456&Signature=789" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/webp" assert base64_bytes == base64.b64encode(webp_content).decode("utf-8") @@ -725,21 +725,21 @@ def test_bedrock_image_processor_content_type_normal_header(): Test that _post_call_image_processing works normally when content-type is correctly set """ import base64 - + # Create mock response with correct content-type mock_response = MagicMock() mock_response.headers.get.return_value = "image/png" - + # Create a PNG header png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" png_content = png_header + b"\x00" * 100 mock_response.content = png_content - + image_url = "https://example.com/test-image.png" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/png" assert base64_bytes == base64.b64encode(png_content).decode("utf-8") @@ -751,16 +751,16 @@ def test_bedrock_image_processor_content_type_fallback_failure(): # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create content with unrecognizable image format mock_response.content = b"\x00" * 100 - + # Test with URL without recognizable extension image_url = "https://example.com/unknown-file" - + with pytest.raises(ValueError) as excinfo: BedrockImageProcessor._post_call_image_processing(mock_response, image_url) - + assert "Unable to determine content type" in str(excinfo.value) @@ -771,18 +771,18 @@ def test_bedrock_image_processor_content_type_jpeg_variants(): # Create mock response with binary/octet-stream mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + jpeg_header = b"\xff\xd8\xff" jpeg_content = jpeg_header + b"\x00" * 100 mock_response.content = jpeg_content - + # Test with .jpg extension image_url_jpg = "https://example.com/photo.jpg" _, content_type_jpg = BedrockImageProcessor._post_call_image_processing( mock_response, image_url_jpg ) assert content_type_jpg == "image/jpeg" - + # Test with .jpeg extension image_url_jpeg = "https://example.com/photo.jpeg" _, content_type_jpeg = BedrockImageProcessor._post_call_image_processing( @@ -797,22 +797,22 @@ def test_bedrock_image_processor_content_type_pdf_document(): when content-type is binary/octet-stream """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a PDF header (magic bytes: %PDF) pdf_header = b"%PDF-1.4" pdf_content = pdf_header + b"\x00" * 100 mock_response.content = pdf_content - + # Test with .pdf URL pdf_url = "https://s3.amazonaws.com/bucket/document.pdf" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, pdf_url ) - + assert content_type == "application/pdf" assert base64_bytes == base64.b64encode(pdf_content).decode("utf-8") @@ -822,12 +822,12 @@ def test_bedrock_image_processor_content_type_document_formats(): Test that _post_call_image_processing handles various document formats """ import base64 - + # Create mock response mock_response = MagicMock() mock_response.headers.get.return_value = "application/octet-stream" mock_response.content = b"\x00" * 100 - + # Test various document formats test_cases = [ ("https://example.com/doc.pdf", "application/pdf"), @@ -837,7 +837,7 @@ def test_bedrock_image_processor_content_type_document_formats(): ("https://example.com/page.html", "text/html"), ("https://example.com/readme.txt", "text/plain"), ] - + for url, expected_mime in test_cases: _, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, url @@ -850,21 +850,21 @@ def test_bedrock_image_processor_content_type_s3_pdf_with_query(): Test that _post_call_image_processing handles S3 PDF with query parameters """ import base64 - + # Create mock response mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + pdf_content = b"%PDF-1.4" + b"\x00" * 100 mock_response.content = pdf_content - + # S3 signed URL with query parameters s3_url = "https://my-bucket.s3.us-east-1.amazonaws.com/documents/report.pdf?AWSAccessKeyId=AKIAIOSFODNN7EXAMPLE&Expires=1234567890&Signature=abcdef123456" - + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, s3_url ) - + assert content_type == "application/pdf" assert base64_bytes == base64.b64encode(pdf_content).decode("utf-8") @@ -1139,6 +1139,170 @@ def test_bedrock_create_bedrock_block_different_document_formats(): assert block["document"]["format"] == format_type +def test_convert_to_anthropic_tool_result_image_with_cache_control(): + """ + Test that cache_control is properly applied to image content in tool results. + This tests the functionality added in the uncommitted changes where + add_cache_control_to_content is called for image_url content types. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + # Test with base64 image data URI + message = { + "role": "tool", + "tool_call_id": "call_test_123", + "content": [ + { + "type": "text", + "text": "Here is the image you requested:", + }, + { + "type": "image_url", + "image_url": "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQ", + "cache_control": {"type": "ephemeral"}, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify the result structure + assert result["type"] == "tool_result" + assert result["tool_use_id"] == "call_test_123" + assert isinstance(result["content"], list) + assert len(result["content"]) == 2 + + # Verify text content + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Here is the image you requested:" + + # Verify image content with cache_control + assert result["content"][1]["type"] == "image" + assert result["content"][1]["source"]["type"] == "base64" + assert result["content"][1]["source"]["media_type"] == "image/jpeg" + assert "cache_control" in result["content"][1] + assert result["content"][1]["cache_control"]["type"] == "ephemeral" + + +def test_convert_to_anthropic_tool_result_image_without_cache_control(): + """ + Test that images without cache_control in tool results work correctly. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_test_456", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAUA", + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify the result structure + assert result["type"] == "tool_result" + assert result["tool_use_id"] == "call_test_456" + assert isinstance(result["content"], list) + assert len(result["content"]) == 1 + + # Verify image content without cache_control (cache_control will be None if not set) + assert result["content"][0]["type"] == "image" + assert result["content"][0]["source"]["type"] == "base64" + assert result["content"][0]["source"]["media_type"] == "image/png" + assert result["content"][0].get("cache_control") is None + + +def test_convert_to_anthropic_tool_result_mixed_content_with_cache_control(): + """ + Test tool results with mixed content types (text and image) where only some have cache_control. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_test_789", + "content": [ + { + "type": "text", + "text": "First image:", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "image_url", + "image_url": "data:image/jpeg;base64,/9j/4AAQSkZJRg", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "text", + "text": "Second image (no cache):", + }, + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgo", + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + assert result["type"] == "tool_result" + assert isinstance(result["content"], list) + assert len(result["content"]) == 4 + + # First text with cache_control + assert result["content"][0]["type"] == "text" + assert result["content"][0]["cache_control"]["type"] == "ephemeral" + + # First image with cache_control + assert result["content"][1]["type"] == "image" + assert result["content"][1]["cache_control"]["type"] == "ephemeral" + + # Second text without cache_control (cache_control will be None if not set) + assert result["content"][2]["type"] == "text" + assert result["content"][2].get("cache_control") is None + + # Second image without cache_control (cache_control will be None if not set) + assert result["content"][3]["type"] == "image" + assert result["content"][3].get("cache_control") is None + + +def test_convert_to_anthropic_tool_result_image_url_as_http(): + """ + Test that HTTP/HTTPS URLs with cache_control are handled correctly. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_http_001", + "content": [ + { + "type": "image_url", + "image_url": "https://example.com/image.jpg", + "cache_control": {"type": "ephemeral"}, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify image is passed as URL reference with cache_control + assert result["content"][0]["type"] == "image" + assert result["content"][0]["source"]["type"] == "url" + assert result["content"][0]["source"]["url"] == "https://example.com/image.jpg" + assert result["content"][0]["cache_control"]["type"] == "ephemeral" def test_anthropic_messages_pt_server_tool_use_passthrough(): """ Test that anthropic_messages_pt passes through server_tool_use and diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 9b150fd89f4..bae0e5bbb4f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -393,6 +393,63 @@ def test_get_request_tags_from_metadata_and_litellm_metadata(): assert "User-Agent: litellm/1.0.0" in tags +def test_get_request_tags_does_not_mutate_original_tags(): + """ + Test that _get_request_tags does not mutate the original tags list in metadata. + + This is a regression test for a bug where calling _get_request_tags multiple times + would cause User-Agent tags to be duplicated because the function was mutating + the original tags list instead of creating a copy. + """ + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + # Create metadata with original tags + original_tags = ["custom-tag-1", "custom-tag-2"] + metadata = {"tags": original_tags} + litellm_params = {"metadata": metadata} + proxy_server_request = { + "headers": { + "user-agent": "AsyncOpenAI/Python 1.99.9", + } + } + + # Call _get_request_tags multiple times (simulating multiple callbacks) + tags1 = StandardLoggingPayloadSetup._get_request_tags( + litellm_params=litellm_params, + proxy_server_request=proxy_server_request, + ) + tags2 = StandardLoggingPayloadSetup._get_request_tags( + litellm_params=litellm_params, + proxy_server_request=proxy_server_request, + ) + tags3 = StandardLoggingPayloadSetup._get_request_tags( + litellm_params=litellm_params, + proxy_server_request=proxy_server_request, + ) + + # Verify the original tags list was NOT mutated + assert original_tags == ["custom-tag-1", "custom-tag-2"], ( + f"Original tags list was mutated: {original_tags}" + ) + assert metadata["tags"] == ["custom-tag-1", "custom-tag-2"], ( + f"metadata['tags'] was mutated: {metadata['tags']}" + ) + + # Verify each returned list has exactly 2 User-Agent tags (not duplicated) + user_agent_count_1 = len([t for t in tags1 if t.startswith("User-Agent:")]) + user_agent_count_2 = len([t for t in tags2 if t.startswith("User-Agent:")]) + user_agent_count_3 = len([t for t in tags3 if t.startswith("User-Agent:")]) + + assert user_agent_count_1 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_1}" + assert user_agent_count_2 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_2}" + assert user_agent_count_3 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_3}" + + # Verify all returned lists are independent (different objects) + assert tags1 is not tags2 + assert tags2 is not tags3 + assert tags1 is not original_tags + + def test_get_extra_header_tags(): """Test the _get_extra_header_tags method with various scenarios.""" import litellm diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index 91d664c3216..199a16d8590 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -204,3 +204,62 @@ def test_azure_gpt5_reasoning_effort_none_dropped(config: AzureOpenAIGPT5Config) ) assert "reasoning_effort" not in params or params.get("reasoning_effort") != "none" + +# Logprobs support tests for Azure GPT-5.2 +def test_azure_gpt5_2_supports_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.2 models support logprobs parameters. + + Only Azure OpenAI GPT-5.2 supports logprobs, unlike OpenAI's GPT-5 or Azure's gpt-5/gpt-5.1. + Tested with gpt-5.2 on api-version 2025-01-01-preview. + """ + supported_params = config.get_supported_openai_params(model="gpt-5.2") + assert "logprobs" in supported_params + assert "top_logprobs" in supported_params + + +def test_azure_gpt5_2_with_prefix_supports_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.2 with azure/ prefix supports logprobs parameters.""" + supported_params = config.get_supported_openai_params(model="azure/gpt-5.2") + assert "logprobs" in supported_params + assert "top_logprobs" in supported_params + + +def test_azure_gpt5_2_series_supports_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.2 with gpt5_series prefix supports logprobs.""" + supported_params = config.get_supported_openai_params(model="gpt5_series/gpt-5.2") + assert "logprobs" in supported_params + assert "top_logprobs" in supported_params + + +def test_azure_gpt5_2_logprobs_params_passed_through(config: AzureOpenAIGPT5Config): + """Test that logprobs parameters are correctly passed through to the API for gpt-5.2.""" + params = config.map_openai_params( + non_default_params={"logprobs": True, "top_logprobs": 5}, + optional_params={}, + model="azure/gpt-5.2", + drop_params=False, + api_version="2025-01-01-preview", + ) + assert params["logprobs"] is True + assert params["top_logprobs"] == 5 + + +def test_azure_gpt5_base_does_not_support_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5 (non-5.2) does not support logprobs parameters. + + Only gpt-5.2 has been verified to support logprobs on Azure. + """ + supported_params = config.get_supported_openai_params(model="gpt-5") + assert "logprobs" not in supported_params + assert "top_logprobs" not in supported_params + + +def test_azure_gpt5_1_does_not_support_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.1 does not support logprobs parameters. + + Only gpt-5.2 has been verified to support logprobs on Azure. + """ + supported_params = config.get_supported_openai_params(model="gpt-5.1") + assert "logprobs" not in supported_params + assert "top_logprobs" not in supported_params + diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index 998510efcd9..987eb5bf998 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -164,3 +164,90 @@ def test_azure_image_generation_headers_without_api_key(): # Verify api-key is added when api_key is valid assert "api-key" in default_headers_with_key assert default_headers_with_key["api-key"] == "valid-key-123" + + +def test_azure_image_generation_drop_params_response_format(): + """ + Test that unsupported params like response_format are dropped when drop_params=True. + + Azure gpt-image-1.5 doesn't support response_format parameter. When drop_params=True, + this parameter should be completely removed and not appear in the final request body, + including not being added to extra_body. + + This test verifies the fix where: + 1. Unsupported params are removed from non_default_params in _check_valid_arg + 2. Unsupported params are also removed from passed_params to prevent them from + being re-added via extra_body in add_provider_specific_params_to_optional_params + + Without the fix, response_format would be added to extra_body and cause Azure to + return a 400 Bad Request error due to strict schema validation. + """ + from litellm.llms.openai.image_generation.gpt_transformation import ( + GPTImageGenerationConfig, + ) + + # Test with gpt-image-1.5 which doesn't support response_format + config = GPTImageGenerationConfig() + supported_params = config.get_supported_openai_params(model="gpt-image-1.5") + + # Verify response_format is NOT in supported params for gpt-image-1.5 + assert "response_format" not in supported_params + assert "n" in supported_params + assert "size" in supported_params + + # Test get_optional_params_image_gen with drop_params=True + optional_params = get_optional_params_image_gen( + model="gpt-image-1.5", + n=1, + size="1024x1024", + response_format="b64_json", # This should be dropped + custom_llm_provider="azure", + provider_config=config, + drop_params=True, + ) + + # Verify response_format is NOT in optional_params + assert "response_format" not in optional_params, ( + "response_format should be dropped from optional_params" + ) + + # Verify response_format is NOT in extra_body either + if "extra_body" in optional_params: + assert "response_format" not in optional_params["extra_body"], ( + "response_format should not be in extra_body" + ) + + # Verify supported params ARE in optional_params + assert "n" in optional_params + assert optional_params["n"] == 1 + assert "size" in optional_params + assert optional_params["size"] == "1024x1024" + + +def test_azure_image_generation_drop_params_false_raises_error(): + """ + Test that unsupported params raise an error when drop_params=False. + + This verifies that the error handling still works correctly when drop_params + is not enabled. + """ + from litellm.exceptions import UnsupportedParamsError + from litellm.llms.openai.image_generation.gpt_transformation import ( + GPTImageGenerationConfig, + ) + + config = GPTImageGenerationConfig() + + # Test that passing unsupported param with drop_params=False raises error + with pytest.raises(UnsupportedParamsError) as exc_info: + optional_params = get_optional_params_image_gen( + model="gpt-image-1.5", + n=1, + response_format="b64_json", # Unsupported param + custom_llm_provider="azure", + provider_config=config, + drop_params=False, + ) + + # Verify the error message mentions the unsupported parameter + assert "response_format" in str(exc_info.value) diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py b/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py new file mode 100644 index 00000000000..714adc346db --- /dev/null +++ b/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py @@ -0,0 +1,132 @@ +""" +Unit tests for OpenRouter embedding transformation logic. +""" +from litellm.llms.openrouter.embedding.transformation import ( + OpenrouterEmbeddingConfig, +) + + +def test_openrouter_embedding_supported_params(): + """Test that supported OpenAI params are correctly defined.""" + config = OpenrouterEmbeddingConfig() + supported = config.get_supported_openai_params("test-model") + + assert "timeout" in supported + assert "dimensions" in supported + assert "encoding_format" in supported + assert "user" in supported + + +def test_openrouter_embedding_transform_request(): + """Test request transformation logic.""" + config = OpenrouterEmbeddingConfig() + + # Test with string input + result = config.transform_embedding_request( + model="openrouter/google/text-embedding-004", + input="Hello world", + optional_params={}, + headers={}, + ) + + assert result["model"] == "google/text-embedding-004" + assert result["input"] == ["Hello world"] + + # Test with list input + result = config.transform_embedding_request( + model="google/text-embedding-004", + input=["Hello", "World"], + optional_params={"dimensions": 512}, + headers={}, + ) + + assert result["model"] == "google/text-embedding-004" + assert result["input"] == ["Hello", "World"] + assert result["dimensions"] == 512 + + +def test_openrouter_embedding_validate_environment(): + """Test environment validation and header setup.""" + config = OpenrouterEmbeddingConfig() + + # Test with API key + headers = config.validate_environment( + headers={"Custom-Header": "value"}, + model="test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-api-key", + ) + + # Should include OpenRouter-specific headers + assert "HTTP-Referer" in headers + assert "X-Title" in headers + # Should include Content-Type header + assert "Content-Type" in headers + assert headers["Content-Type"] == "application/json" + # Should include Authorization header + assert "Authorization" in headers + assert headers["Authorization"] == "Bearer test-api-key" + # Should preserve custom headers + assert headers["Custom-Header"] == "value" + + # Test without API key + headers_no_key = config.validate_environment( + headers={}, + model="test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + # Should still include OpenRouter headers but not Authorization + assert "HTTP-Referer" in headers_no_key + assert "X-Title" in headers_no_key + assert "Content-Type" in headers_no_key + assert "Authorization" not in headers_no_key + + +def test_openrouter_embedding_get_complete_url(): + """Test URL construction.""" + config = OpenrouterEmbeddingConfig() + + url = config.get_complete_url( + api_base="https://openrouter.ai/api/v1", + api_key="test-key", + model="test-model", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://openrouter.ai/api/v1/embeddings" + + # Test with trailing slash + url = config.get_complete_url( + api_base="https://openrouter.ai/api/v1/", + api_key="test-key", + model="test-model", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://openrouter.ai/api/v1/embeddings" + + +def test_openrouter_embedding_map_params(): + """Test parameter mapping.""" + config = OpenrouterEmbeddingConfig() + + result = config.map_openai_params( + non_default_params={"dimensions": 512, "timeout": 30, "unsupported": "value"}, + optional_params={}, + model="test-model", + drop_params=False, + ) + + # Supported params should be included + assert result["dimensions"] == 512 + assert result["timeout"] == 30 + # Unsupported params should not be included + assert "unsupported" not in result diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py new file mode 100644 index 00000000000..0a1ac7e2a54 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py @@ -0,0 +1,38 @@ +import pytest +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm import ModelResponse + +def test_process_candidates_unbound_local_error_fix(): + # Setup + candidates = [ + { + "content": { + "role": "model" + # "parts" is missing intentionally to trigger the issue + }, + "finishReason": "STOP" + } + ] + model_response = ModelResponse() + + # Execution + try: + VertexGeminiConfig._process_candidates( + _candidates=candidates, + model_response=model_response, + standard_optional_params={}, + cumulative_tool_call_index=0 + ) + except UnboundLocalError as e: + pytest.fail(f"UnboundLocalError raised: {e}") + except Exception as e: + # Other exceptions might be okay if they are not UnboundLocalError, + # but ideally it should pass without error or raise a specific error if parts are required. + # However, the goal is to verify thought_signatures doesn't crash. + pass + + # Verify that we didn't crash with UnboundLocalError + +if __name__ == "__main__": + test_process_candidates_unbound_local_error_fix() + print("Test passed!") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 389c8446135..80d65991acb 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -13,6 +13,7 @@ sys.path.insert( import litellm from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.llms.vertex_ai.common_utils import _get_gemini_url def run_sync(coro): @@ -1048,3 +1049,139 @@ class TestVertexBase: MockCredentials.from_info.assert_called_once_with(json_obj) mock_creds.with_scopes.assert_called_once_with(scopes) assert result == "scoped_creds" + + def test_get_token_and_url_with_api_key(self): + """Test that API key authentication routes to Google AI Studio endpoint""" + vertex_base = VertexBase() + + # Test with API key and no credentials - should use Google AI Studio endpoint + auth_header, url = vertex_base._get_token_and_url( + model="gemini-2.0-flash-exp", + auth_header=None, + gemini_api_key="test-api-key-123", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, # No service account credentials + stream=False, + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="chat", + ) + + # Should route to Google AI Studio endpoint + assert "generativelanguage.googleapis.com" in url + assert "gemini-2.0-flash-exp" in url + assert "key=test-api-key-123" in url + assert auth_header is None # API key is in URL, not header + + def test_get_token_and_url_with_credentials(self): + """Test that service account credentials route to Vertex AI endpoint""" + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "mock-bearer-token" + mock_creds.expired = False + + with patch.object( + vertex_base, "_ensure_access_token", return_value=("mock-bearer-token", "test-project") + ): + # Test with credentials - should use Vertex AI endpoint + auth_header, url = vertex_base._get_token_and_url( + model="gemini-2.0-flash-exp", + auth_header="mock-bearer-token", + gemini_api_key=None, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials={"type": "service_account"}, + stream=False, + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="chat", + ) + + # Should route to Vertex AI endpoint + assert "aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "locations/us-central1" in url + assert auth_header == "mock-bearer-token" + + def test_get_token_and_url_api_key_with_streaming(self): + """Test API key authentication with streaming enabled""" + vertex_base = VertexBase() + + auth_header, url = vertex_base._get_token_and_url( + model="gemini-2.0-flash-exp", + auth_header=None, + gemini_api_key="test-api-key-456", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + stream=True, # Streaming enabled + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="chat", + ) + + # Should route to Google AI Studio endpoint with streaming + assert "generativelanguage.googleapis.com" in url + assert "streamGenerateContent" in url + assert "key=test-api-key-456" in url + assert "alt=sse" in url + assert auth_header is None + + def test_get_token_and_url_api_key_priority(self): + """Test that credentials take priority over API key when both are provided""" + vertex_base = VertexBase() + + # When both API key and credentials are provided, credentials take priority + mock_creds = MagicMock() + mock_creds.token = "mock-bearer-token" + mock_creds.expired = False + + with patch.object( + vertex_base, "_ensure_access_token", return_value=("mock-bearer-token", "test-project") + ): + auth_header, url = vertex_base._get_token_and_url( + model="gemini-2.0-flash-exp", + auth_header="mock-bearer-token", + gemini_api_key="test-api-key-789", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials={"type": "service_account"}, # Credentials provided + stream=False, + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="chat", + ) + + # Should use Vertex AI endpoint with Bearer token (credentials take priority) + assert "aiplatform.googleapis.com" in url + assert auth_header == "mock-bearer-token" + + def test_get_token_and_url_with_embedding_mode(self): + """Test API key authentication with embedding mode""" + vertex_base = VertexBase() + + auth_header, url = vertex_base._get_token_and_url( + model="text-embedding-004", + auth_header=None, + gemini_api_key="test-embedding-key", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + stream=False, + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="embedding", + ) + + # Should route to Google AI Studio endpoint for embeddings + assert "generativelanguage.googleapis.com" in url + assert "embedContent" in url + assert "key=test-embedding-key" in url + assert auth_header is None \ No newline at end of file diff --git a/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py b/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py index e36a494998b..fd5f8f3eff8 100644 --- a/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/watsonx/audio_transcription/test_watsonx_audio_transcription_transformation.py @@ -14,6 +14,10 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) import litellm +from litellm.llms.watsonx.audio_transcription.transformation import ( + IBMWatsonXAudioTranscriptionConfig, +) +from litellm.types.utils import TranscriptionResponse class TestWatsonXAudioTranscription: @@ -189,3 +193,72 @@ class TestWatsonXAudioTranscription: # Verify file is sent separately files = captured_request.get("files", {}) assert "file" in files + + def test_transform_audio_transcription_response_removes_model_field(self): + """ + Test that transform_audio_transcription_response removes the 'model' field + from WatsonX response before creating TranscriptionResponse. + + This test ensures that when WatsonX returns a response with a 'model' field, + it is removed before creating the TranscriptionResponse object, since + TranscriptionResponse doesn't accept a 'model' parameter. + """ + handler = IBMWatsonXAudioTranscriptionConfig() + + # Mock response with 'model' field (as WatsonX may return) + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello, this is a test transcription.", + "model": "whisper-large-v3-turbo", # This field should be removed + "duration": 5.5, + } + mock_response.text = '{"text": "Hello, this is a test transcription.", "model": "whisper-large-v3-turbo", "duration": 5.5}' + + # This should not raise a TypeError - model field should be removed + result = handler.transform_audio_transcription_response(mock_response) + + # Verify the result is a TranscriptionResponse + assert isinstance(result, TranscriptionResponse) + + # Verify the text is correct + assert result.text == "Hello, this is a test transcription." + + # Verify duration is set via dictionary assignment + assert result["duration"] == 5.5 + + # Verify the model field is NOT in the serialized result + # Check via model_dump() or dict() to ensure it's not in the output + try: + result_dict = result.model_dump() + except AttributeError: + # Fallback for pydantic v1 + result_dict = result.dict() + + # The 'model' field should not be in the result + assert "model" not in result_dict, "Model field should be removed from response" + + def test_transform_audio_transcription_response_without_model_field(self): + """ + Test that transform_audio_transcription_response works correctly + when WatsonX response doesn't include a 'model' field. + """ + handler = IBMWatsonXAudioTranscriptionConfig() + + # Mock response without 'model' field + mock_response = MagicMock() + mock_response.json.return_value = { + "text": "Hello, this is a test transcription.", + "duration": 5.5, + } + mock_response.text = '{"text": "Hello, this is a test transcription.", "duration": 5.5}' + + result = handler.transform_audio_transcription_response(mock_response) + + # Verify the result is a TranscriptionResponse + assert isinstance(result, TranscriptionResponse) + + # Verify the text is correct + assert result.text == "Hello, this is a test transcription." + + # Verify duration is set via dictionary assignment + assert result["duration"] == 5.5 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 8062243dfdd..f1558ac5791 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,4 +1,5 @@ import asyncio +from datetime import datetime, timedelta from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -7,7 +8,12 @@ from fastapi import HTTPException from mcp import ReadResourceResult, Resource from mcp.types import Prompt, ResourceTemplate, TextResourceContents -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_MCPServerTable, + MCPTransport, + UserAPIKeyAuth, +) +from litellm.types.mcp_server.mcp_server_manager import MCPServer @pytest.mark.asyncio @@ -1688,3 +1694,99 @@ def test_filter_tools_by_allowed_tools(): assert len(filtered_tools) == 2 assert filtered_tools[0].name == "my_api_mcp-getpetbyid" assert filtered_tools[1].name == "my_api_mcp-findpetsbystatus" + + +def _make_db_mcp_server(server_id: str, updated_at: datetime) -> LiteLLM_MCPServerTable: + return LiteLLM_MCPServerTable( + server_id=server_id, + server_name="server", + alias="server", + url="https://example.com", + transport=MCPTransport.http, + created_at=updated_at, + updated_at=updated_at, + mcp_info={}, + ) + + +class TestMCPServerManagerReload: + @pytest.mark.asyncio + async def test_reuses_existing_server_when_updated_at_matches(self): + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + timestamp = datetime.utcnow() + existing_server = MCPServer( + server_id="server-1", + name="server", + transport=MCPTransport.http, + updated_at=timestamp, + ) + manager.registry = {existing_server.server_id: existing_server} + + db_row = _make_db_mcp_server("server-1", timestamp) + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_all_mcp_servers", + new=AsyncMock(return_value=[db_row]), + ) as mock_get_all, patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=object(), + ), patch.object( + manager, "build_mcp_server_from_table", AsyncMock() + ) as mock_build: + await manager.reload_servers_from_database() + + mock_get_all.assert_awaited_once() + mock_build.assert_not_awaited() + assert manager.registry["server-1"] is existing_server + + @pytest.mark.asyncio + async def test_rebuilds_server_when_updated_at_changes(self): + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + timestamp = datetime.utcnow() + existing_server = MCPServer( + server_id="server-1", + name="server", + transport=MCPTransport.http, + updated_at=timestamp, + ) + manager.registry = {existing_server.server_id: existing_server} + + new_timestamp = timestamp + timedelta(minutes=5) + db_row = _make_db_mcp_server("server-1", new_timestamp) + rebuilt_server = MCPServer( + server_id="server-1", + name="server", + transport=MCPTransport.http, + updated_at=new_timestamp, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_all_mcp_servers", + new=AsyncMock(return_value=[db_row]), + ) as mock_get_all, patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=object(), + ), patch.object( + manager, + "build_mcp_server_from_table", + AsyncMock(return_value=rebuilt_server), + ) as mock_build: + await manager.reload_servers_from_database() + + mock_get_all.assert_awaited_once() + mock_build.assert_awaited_once_with(db_row) + assert manager.registry["server-1"] is rebuilt_server diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py new file mode 100644 index 00000000000..b1bef63933e --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -0,0 +1,131 @@ +""" +Unit tests for auth_utils functions related to rate limiting. +""" + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import ( + get_key_model_rpm_limit, + get_key_model_tpm_limit, +) + + +class TestGetKeyModelRpmLimit: + """Tests for get_key_model_rpm_limit function.""" + + def test_returns_key_metadata_when_present(self): + """Key metadata takes priority over team metadata.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 50}}, + ) + result = get_key_model_rpm_limit(user_api_key_dict) + assert result == {"gpt-4": 100} + + def test_falls_back_to_team_metadata_when_key_has_other_metadata(self): + """Should fall back to team metadata when key metadata exists but has no model_rpm_limit.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + metadata={ + "some_other_key": "value" + }, # Has metadata, but not model_rpm_limit + team_metadata={"model_rpm_limit": {"gpt-4": 50}}, + ) + result = get_key_model_rpm_limit(user_api_key_dict) + assert result == {"gpt-4": 50} + + def test_extracts_from_model_max_budget(self): + """Should extract rpm_limit from model_max_budget when metadata is empty.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + model_max_budget={ + "gpt-4": {"rpm_limit": 100, "tpm_limit": 1000}, + "gpt-3.5-turbo": {"rpm_limit": 200}, + }, + ) + result = get_key_model_rpm_limit(user_api_key_dict) + assert result == {"gpt-4": 100, "gpt-3.5-turbo": 200} + + def test_skips_models_without_rpm_limit(self): + """Should skip models that don't have rpm_limit in model_max_budget.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + model_max_budget={ + "gpt-4": {"rpm_limit": 100}, + "gpt-3.5-turbo": {"tpm_limit": 1000}, # No rpm_limit + }, + ) + result = get_key_model_rpm_limit(user_api_key_dict) + assert result == {"gpt-4": 100} + + def test_returns_none_when_no_limits_configured(self): + """Should return None when no rate limits are configured.""" + user_api_key_dict = UserAPIKeyAuth(api_key="sk-123") + result = get_key_model_rpm_limit(user_api_key_dict) + assert result is None + + +class TestGetKeyModelTpmLimit: + """Tests for get_key_model_tpm_limit function.""" + + def test_returns_key_metadata_when_present(self): + """Key metadata takes priority over team metadata.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + metadata={"model_tpm_limit": {"gpt-4": 10000}}, + team_metadata={"model_tpm_limit": {"gpt-4": 5000}}, + ) + result = get_key_model_tpm_limit(user_api_key_dict) + assert result == {"gpt-4": 10000} + + def test_falls_back_to_team_metadata_when_key_has_other_metadata(self): + """Should fall back to team metadata when key metadata exists but has no model_tpm_limit.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + metadata={ + "some_other_key": "value" + }, # Has metadata, but not model_tpm_limit + team_metadata={"model_tpm_limit": {"gpt-4": 5000}}, + ) + result = get_key_model_tpm_limit(user_api_key_dict) + assert result == {"gpt-4": 5000} + + def test_extracts_from_model_max_budget(self): + """Should extract tpm_limit from model_max_budget when metadata is empty.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + model_max_budget={ + "gpt-4": {"tpm_limit": 10000, "rpm_limit": 100}, + "gpt-3.5-turbo": {"tpm_limit": 20000}, + }, + ) + result = get_key_model_tpm_limit(user_api_key_dict) + assert result == {"gpt-4": 10000, "gpt-3.5-turbo": 20000} + + def test_skips_models_without_tpm_limit(self): + """Should skip models that don't have tpm_limit in model_max_budget.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + model_max_budget={ + "gpt-4": {"tpm_limit": 10000}, + "gpt-3.5-turbo": {"rpm_limit": 100}, # No tpm_limit + }, + ) + result = get_key_model_tpm_limit(user_api_key_dict) + assert result == {"gpt-4": 10000} + + def test_returns_none_when_no_limits_configured(self): + """Should return None when no rate limits are configured.""" + user_api_key_dict = UserAPIKeyAuth(api_key="sk-123") + result = get_key_model_tpm_limit(user_api_key_dict) + assert result is None + + def test_model_max_budget_priority_over_team(self): + """model_max_budget should take priority over team_metadata.""" + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + model_max_budget={"gpt-4": {"tpm_limit": 10000}}, + team_metadata={"model_tpm_limit": {"gpt-4": 5000}}, + ) + result = get_key_model_tpm_limit(user_api_key_dict) + assert result == {"gpt-4": 10000} diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index c04b8114939..e7b27908c14 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -248,6 +248,83 @@ async def test_authenticate_user_wrong_password(): assert "Invalid credentials" in exc_info.value.message +@pytest.mark.asyncio +async def test_authenticate_user_email_case_insensitive_login(): + """Test that email lookup is case-insensitive during login""" + master_key = "sk-1234" + stored_email = "testemail@test.com" + login_email_mixed_case = "testEmail@test.com" + correct_password = "correct-password" + hashed_password = hash_token(token=correct_password) + + # `LiteLLM_UserTable` does not define a `password` field, but `authenticate_user()` + # expects `user_row.password` to exist (invite-link login). Use a simple object. + mock_user = MagicMock() + mock_user.user_id = "test-user-123" + mock_user.user_email = stored_email + mock_user.password = hashed_password + mock_user.user_role = LitellmUserRoles.INTERNAL_USER + + def mock_find_first(**kwargs): + where = kwargs.get("where", {}) + user_email = where.get("user_email", {}) + if user_email.get("mode") != "insensitive": + return None + if str(user_email.get("equals", "")).lower() == stored_email.lower(): + return mock_user + return None + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + side_effect=mock_find_first + ) + + with patch.dict( + os.environ, + { + "DATABASE_URL": "postgresql://test:test@localhost/test", + "UI_USERNAME": "admin", + "UI_PASSWORD": "admin-password", + }, + ): + with patch( + "litellm.proxy.auth.login_utils.expire_previous_ui_session_tokens", + new_callable=AsyncMock, + return_value=None, + ): + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.side_effect = [ + {"token": "token-1"}, + {"token": "token-2"}, + ] + + result_mixed = await authenticate_user( + username=login_email_mixed_case, + password=correct_password, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + result_lower = await authenticate_user( + username=stored_email, + password=correct_password, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert result_mixed.user_id == result_lower.user_id == "test-user-123" + assert result_mixed.user_email == result_lower.user_email == stored_email + + calls = mock_prisma_client.db.litellm_usertable.find_first.await_args_list + assert len(calls) == 2 + for call, expected_username in zip(calls, [login_email_mixed_case, stored_email]): + where = call.kwargs["where"] + assert where["user_email"]["equals"] == expected_username + assert where["user_email"]["mode"] == "insensitive" + + @pytest.mark.asyncio async def test_authenticate_user_database_required_for_admin(): """Test that database is required for admin login""" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index fcc8c1f0f2e..5f49db66089 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -338,6 +338,17 @@ async def test_proxy_admin_expired_key_from_cache(): f"Exception message should mention 'Expired Key', got: {exc_info.value.message}" ) + # Verify that the param field does NOT leak the full API key (Issue #18731) + # The param should be abbreviated like "sk-...XXXX" not the full plaintext key + assert exc_info.value.param is not None, "Exception should have 'param' attribute" + assert exc_info.value.param != api_key, ( + f"SECURITY: Full API key should NOT be in param field! " + f"Got: {exc_info.value.param}, Expected abbreviated format like 'sk-...XXXX'" + ) + assert exc_info.value.param.startswith("sk-..."), ( + f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}" + ) + # Verify that cache deletion was called mock_delete_cache.assert_called_once() call_args = mock_delete_cache.call_args @@ -347,3 +358,4 @@ async def test_proxy_admin_expired_key_from_cache(): finally: # Clean up - restore original values if needed pass + diff --git a/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py b/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py new file mode 100644 index 00000000000..1492acb0794 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py @@ -0,0 +1,275 @@ +""" +Tests for the RDS IAM token proactive refresh implementation. + +Tests for GitHub Issue #16220: RDS IAM authentication connection failures after 15 minutes. + +The fix implements: +1. Proactive background token refresh (refreshes 3 min before expiration) +2. Precise sleep timing (1 wake-up per token cycle instead of polling) +3. Proper locking during reconnection +4. Fixed __getattr__ fallback that now waits for reconnection + +Run these tests: + poetry run pytest tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py -v -s +""" + +import asyncio +import os +import urllib.parse +from datetime import datetime, timedelta +from unittest.mock import MagicMock, patch + +import pytest + + +class TestPrismaWrapperTokenRefresh: + """Tests for the PrismaWrapper RDS IAM token refresh implementation.""" + + @pytest.fixture + def setup_env(self): + """Setup environment variables for testing.""" + os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "test_user" + os.environ["DATABASE_NAME"] = "test_db" + os.environ["IAM_TOKEN_DB_AUTH"] = "True" + yield + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + "IAM_TOKEN_DB_AUTH", + "DATABASE_SCHEMA", + ]: + os.environ.pop(key, None) + + def _generate_mock_token(self, expires_in_seconds: int = 900) -> str: + """Generate a mock IAM token with expiration info.""" + now = datetime.utcnow() + date_str = now.strftime("%Y%m%dT%H%M%SZ") + # Build the token like AWS does + token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires={expires_in_seconds}&X-Amz-Signature=abc123" + return urllib.parse.quote(token, safe="") + + def _set_database_url_with_token(self, expires_in_seconds: int = 900): + """Set DATABASE_URL with a mock token.""" + token = self._generate_mock_token(expires_in_seconds) + os.environ[ + "DATABASE_URL" + ] = f"postgresql://test_user:{token}@test-host:5432/test_db" + + @pytest.mark.asyncio + async def test_is_token_expired_fresh(self, setup_env): + """Test that fresh token is not detected as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + self._set_database_url_with_token(expires_in_seconds=900) + db_url = os.getenv("DATABASE_URL") + + assert wrapper.is_token_expired(db_url) is False + + @pytest.mark.asyncio + async def test_is_token_expired_old(self, setup_env): + """Test that old token is detected as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Create an expired token + old_date = datetime.utcnow() - timedelta(seconds=901) + date_str = old_date.strftime("%Y%m%dT%H%M%SZ") + token = ( + f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=900&X-Amz-Signature=abc" + ) + encoded_token = urllib.parse.quote(token, safe="") + db_url = f"postgresql://test_user:{encoded_token}@test-host:5432/test_db" + + assert wrapper.is_token_expired(db_url) is True + + @pytest.mark.asyncio + async def test_start_stop_token_refresh_task(self, setup_env): + """Test that token refresh task starts and stops correctly.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Set a valid token + self._set_database_url_with_token(expires_in_seconds=900) + + # Start the task + await wrapper.start_token_refresh_task() + assert wrapper._token_refresh_task is not None + assert not wrapper._token_refresh_task.done() + + # Stop the task + await wrapper.stop_token_refresh_task() + assert wrapper._token_refresh_task is None + + @pytest.mark.asyncio + async def test_start_task_not_enabled(self, setup_env): + """Test that task doesn't start when IAM auth is not enabled.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + # IAM auth disabled + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False) + + await wrapper.start_token_refresh_task() + assert wrapper._token_refresh_task is None + + @pytest.mark.asyncio + async def test_is_token_expired_null(self, setup_env): + """Test that None token is treated as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + assert wrapper.is_token_expired(None) is True + + +class TestTokenExpirationParsing: + """Tests for token expiration parsing utilities.""" + + def test_parse_token_expiration_valid(self): + """Test parsing expiration from a valid token.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Create a token with known expiration + token = "mock-token?X-Amz-Date=20240101T120000Z&X-Amz-Expires=900&X-Amz-Signature=abc" + + expiration = wrapper._parse_token_expiration(token) + + assert expiration is not None + expected = datetime(2024, 1, 1, 12, 0, 0) + timedelta(seconds=900) + assert expiration == expected + + def test_parse_token_expiration_invalid(self): + """Test that invalid token returns None.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Invalid tokens + assert wrapper._parse_token_expiration(None) is None + assert wrapper._parse_token_expiration("no-query-params") is None + assert wrapper._parse_token_expiration("?missing=params") is None + + +class TestBackgroundRefreshLoop: + """Tests for the background refresh loop timing.""" + + @pytest.fixture + def setup_env(self): + """Setup environment variables for testing.""" + os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "test_user" + os.environ["DATABASE_NAME"] = "test_db" + yield + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + ]: + os.environ.pop(key, None) + + @pytest.mark.asyncio + async def test_calculate_seconds_fallback_when_no_url(self, setup_env): + """Test that fallback is used when DATABASE_URL is not set.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Don't set DATABASE_URL + seconds = wrapper._calculate_seconds_until_refresh() + + # Should return fallback interval + assert seconds == wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS + + +# ============================================================================ +# DEMONSTRATION SCRIPT +# ============================================================================ + + +async def demonstrate_fix(): + """ + Demonstrates the fix for the RDS IAM token expiration bug. + + Shows how the proactive refresh prevents the 15-minute connection failure. + """ + # Import the actual implementation + try: + from litellm.proxy.db.prisma_client import PrismaWrapper + except ImportError: + return + + # Setup mock environment + os.environ["DATABASE_HOST"] = "mock-rds.region.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "iam_user" + os.environ["DATABASE_NAME"] = "litellm" + + # Create initial token (expires in 10 seconds for demo) + now = datetime.utcnow() + date_str = now.strftime("%Y%m%dT%H%M%SZ") + token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=10&X-Amz-Signature=abc123" + encoded_token = urllib.parse.quote(token, safe="") + os.environ[ + "DATABASE_URL" + ] = f"postgresql://iam_user:{encoded_token}@mock-rds:5432/litellm" + + # Create mock prisma client + mock_prisma = MagicMock() + + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Override buffer for faster demo + wrapper.TOKEN_REFRESH_BUFFER_SECONDS = 3 + wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS = 5 + _ = wrapper._calculate_seconds_until_refresh() # Verify calculation works + db_url = os.getenv("DATABASE_URL") + is_expired = wrapper.is_token_expired(db_url) + assert is_expired is False, "Fresh token should not be expired!" + + # Mock the _token_refresh_loop to prevent it from actually running + async def mock_loop(): + try: + await asyncio.sleep(1000) + except asyncio.CancelledError: + pass + + with patch.object(wrapper, "_token_refresh_loop", side_effect=mock_loop): + await wrapper.start_token_refresh_task() + await wrapper.stop_token_refresh_task() + + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + ]: + os.environ.pop(key, None) + + +if __name__ == "__main__": + asyncio.run(demonstrate_fix()) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py index 35ed49a84ed..fd72185d1e7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py @@ -2,7 +2,6 @@ Unit tests for Qualifire guardrail integration. """ -import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -75,9 +74,37 @@ class TestQualifireGuardrailInit: assert guardrail.on_flagged == "monitor" + def test_init_with_default_api_base(self): + """Test that default API base is set when not provided.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + DEFAULT_QUALIFIRE_API_BASE, + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + assert guardrail.qualifire_api_base == DEFAULT_QUALIFIRE_API_BASE + + def test_init_with_custom_api_base(self): + """Test initialization with custom API base URL.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + api_base="https://custom.qualifire.ai", + guardrail_name="test_guardrail", + ) + + assert guardrail.qualifire_api_base == "https://custom.qualifire.ai" + class TestQualifireGuardrailMessageConversion: - """Tests for message conversion to Qualifire format.""" + """Tests for message conversion to API format.""" def test_convert_simple_messages(self): """Test conversion of simple text messages.""" @@ -95,15 +122,13 @@ class TestQualifireGuardrailMessageConversion: {"role": "assistant", "content": "Hi there!"}, ] - # Create mock LLMMessage class - mock_llm_message = MagicMock() + result = guardrail._convert_messages_to_api_format(messages) - with patch( - "litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire.QualifireGuardrail._convert_messages_to_qualifire_format" - ) as mock_convert: - mock_convert.return_value = [mock_llm_message, mock_llm_message] - result = guardrail._convert_messages_to_qualifire_format(messages) - assert len(result) == 2 + assert len(result) == 2 + assert result[0]["role"] == "user" + assert result[0]["content"] == "Hello, world!" + assert result[1]["role"] == "assistant" + assert result[1]["content"] == "Hi there!" def test_convert_multimodal_messages(self): """Test conversion of multimodal messages with text parts.""" @@ -126,112 +151,258 @@ class TestQualifireGuardrailMessageConversion: }, ] - with patch( - "litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire.QualifireGuardrail._convert_messages_to_qualifire_format" - ) as mock_convert: - mock_convert.return_value = [MagicMock()] - result = guardrail._convert_messages_to_qualifire_format(messages) - assert len(result) == 1 + result = guardrail._convert_messages_to_api_format(messages) + + assert len(result) == 1 + assert result[0]["role"] == "user" + assert result[0]["content"] == "First part\nSecond part" + + def test_convert_messages_with_tool_calls(self): + """Test conversion of messages with tool calls.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + messages = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}', + }, + } + ], + }, + ] + + result = guardrail._convert_messages_to_api_format(messages) + + assert len(result) == 1 + assert result[0]["role"] == "assistant" + assert "tool_calls" in result[0] + assert len(result[0]["tool_calls"]) == 1 + assert result[0]["tool_calls"][0]["id"] == "call_123" + assert result[0]["tool_calls"][0]["name"] == "get_weather" + assert result[0]["tool_calls"][0]["arguments"] == {"location": "NYC"} -class TestQualifireGuardrailEvaluateKwargs: - """Tests for evaluate kwargs passed to Qualifire client.""" +class TestQualifireGuardrailToolConversion: + """Tests for tool definition conversion.""" + + def test_convert_openai_function_tools(self): + """Test conversion of OpenAI function tool format.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + result = guardrail._convert_tools_to_api_format(tools) + + assert result is not None + assert len(result) == 1 + assert result[0]["name"] == "get_weather" + assert result[0]["description"] == "Get weather for a location" + + def test_convert_empty_tools(self): + """Test that empty tools returns None.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + result = guardrail._convert_tools_to_api_format(None) + assert result is None + + result = guardrail._convert_tools_to_api_format([]) + assert result is None + + +class TestQualifireGuardrailAPICall: + """Tests for API call with httpx client.""" @pytest.mark.asyncio async def test_evaluate_called_with_prompt_injections(self): - """Test that evaluate is called with prompt_injections enabled.""" - # Mock the qualifire module and its types - mock_qualifire_types = MagicMock() - mock_llm_message = MagicMock() - mock_llm_tool_call = MagicMock() - mock_message_instance = MagicMock() - mock_llm_message.return_value = mock_message_instance - - mock_qualifire_types.LLMMessage = mock_llm_message - mock_qualifire_types.LLMToolCall = mock_llm_tool_call - - with patch.dict('sys.modules', {'qualifire': MagicMock(), 'qualifire.types': mock_qualifire_types}): - from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( - QualifireGuardrail, - ) + """Test that evaluate endpoint is called with prompt_injections enabled.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) - guardrail = QualifireGuardrail( - api_key="test_key", - prompt_injections=True, - guardrail_name="test_guardrail", - ) + guardrail = QualifireGuardrail( + api_key="test_key", + prompt_injections=True, + guardrail_name="test_guardrail", + ) - # Mock the client - mock_client = MagicMock() - mock_result = MagicMock() - mock_result.score = 100 - mock_result.status = "completed" - mock_result.evaluationResults = [] - mock_client.evaluate.return_value = mock_result - guardrail._client = mock_client + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) - messages = [{"role": "user", "content": "Hello, world!"}] + messages = [{"role": "user", "content": "Hello, world!"}] - await guardrail._run_qualifire_check( - messages=messages, output=None, dynamic_params={} - ) + await guardrail._run_qualifire_check( + messages=messages, output=None, dynamic_params={} + ) - # Verify evaluate was called with correct kwargs - mock_client.evaluate.assert_called_once() - call_kwargs = mock_client.evaluate.call_args[1] - assert call_kwargs["prompt_injections"] is True - assert "messages" in call_kwargs + # Verify the API was called + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + assert "json" in call_kwargs + payload = call_kwargs["json"] + assert payload["prompt_injections"] is True + assert "messages" in payload + assert call_kwargs["url"].endswith("/api/evaluation/evaluate") @pytest.mark.asyncio async def test_evaluate_called_with_multiple_checks(self): """Test that evaluate is called with multiple checks enabled.""" - # Mock the qualifire module and its types - mock_qualifire_types = MagicMock() - mock_llm_message = MagicMock() - mock_llm_tool_call = MagicMock() - mock_message_instance = MagicMock() - mock_llm_message.return_value = mock_message_instance - - mock_qualifire_types.LLMMessage = mock_llm_message - mock_qualifire_types.LLMToolCall = mock_llm_tool_call - - with patch.dict('sys.modules', {'qualifire': MagicMock(), 'qualifire.types': mock_qualifire_types}): - from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( - QualifireGuardrail, - ) + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) - guardrail = QualifireGuardrail( - api_key="test_key", - prompt_injections=True, - pii_check=True, - hallucinations_check=True, - assertions=["Output must be valid JSON"], - guardrail_name="test_guardrail", - ) + guardrail = QualifireGuardrail( + api_key="test_key", + prompt_injections=True, + pii_check=True, + hallucinations_check=True, + assertions=["Output must be valid JSON"], + guardrail_name="test_guardrail", + ) - # Mock the client - mock_client = MagicMock() - mock_result = MagicMock() - mock_result.score = 100 - mock_result.status = "completed" - mock_result.evaluationResults = [] - mock_client.evaluate.return_value = mock_result - guardrail._client = mock_client + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) - messages = [{"role": "user", "content": "Hello, world!"}] + messages = [{"role": "user", "content": "Hello, world!"}] - await guardrail._run_qualifire_check( - messages=messages, output="Test output", dynamic_params={} - ) + await guardrail._run_qualifire_check( + messages=messages, output="Test output", dynamic_params={} + ) - # Verify evaluate was called with correct kwargs - mock_client.evaluate.assert_called_once() - call_kwargs = mock_client.evaluate.call_args[1] - assert call_kwargs["prompt_injections"] is True - assert call_kwargs["pii_check"] is True - assert call_kwargs["hallucinations_check"] is True - assert call_kwargs["assertions"] == ["Output must be valid JSON"] - assert call_kwargs["output"] == "Test output" + # Verify the API was called with correct payload + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + payload = call_kwargs["json"] + assert payload["prompt_injections"] is True + assert payload["pii_check"] is True + assert payload["hallucinations_check"] is True + assert payload["assertions"] == ["Output must be valid JSON"] + assert payload["output"] == "Test output" + + @pytest.mark.asyncio + async def test_invoke_endpoint_used_with_evaluation_id(self): + """Test that invoke endpoint is used when evaluation_id is provided.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + evaluation_id="eval_123", + guardrail_name="test_guardrail", + ) + + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + messages = [{"role": "user", "content": "Hello, world!"}] + + await guardrail._run_qualifire_check( + messages=messages, output="Test output", dynamic_params={} + ) + + # Verify the invoke endpoint was called + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + assert call_kwargs["url"].endswith("/api/evaluation/invoke") + payload = call_kwargs["json"] + assert payload["evaluation_id"] == "eval_123" + assert payload["input"] == "Hello, world!" + assert payload["output"] == "Test output" + + @pytest.mark.asyncio + async def test_correct_headers_sent(self): + """Test that correct headers are sent with the API request.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="my_api_key", + guardrail_name="test_guardrail", + ) + + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + messages = [{"role": "user", "content": "Hello!"}] + + await guardrail._run_qualifire_check( + messages=messages, output=None, dynamic_params={} + ) + + call_kwargs = guardrail.async_handler.post.call_args[1] + headers = call_kwargs["headers"] + + assert headers["X-Qualifire-API-Key"] == "my_api_key" + assert headers["Content-Type"] == "application/json" class TestQualifireGuardrailCheckIfFlagged: @@ -248,12 +419,14 @@ class TestQualifireGuardrailCheckIfFlagged: guardrail_name="test_guardrail", ) - # Mock result with completed status and no flagged items - mock_result = MagicMock() - mock_result.status = "completed" - mock_result.evaluationResults = [] + # Result with completed status and no flagged items (dict format) + result = { + "status": "completed", + "score": 100, + "evaluationResults": [], + } - assert guardrail._check_if_flagged(mock_result) is False + assert guardrail._check_if_flagged(result) is False def test_check_if_flagged_returns_true_for_flagged_content(self): """Test that _check_if_flagged returns True when content is flagged.""" @@ -266,18 +439,25 @@ class TestQualifireGuardrailCheckIfFlagged: guardrail_name="test_guardrail", ) - # Mock result with flagged item - mock_inner_result = MagicMock() - mock_inner_result.flagged = True + # Result with flagged item (dict format matching API response) + result = { + "status": "completed", + "score": 15, + "evaluationResults": [ + { + "type": "prompt_injection", + "results": [ + { + "flagged": True, + "score": 0.15, + "reason": "Prompt injection detected", + } + ], + } + ], + } - mock_eval_result = MagicMock() - mock_eval_result.results = [mock_inner_result] - - mock_result = MagicMock() - mock_result.status = "completed" - mock_result.evaluationResults = [mock_eval_result] - - assert guardrail._check_if_flagged(mock_result) is True + assert guardrail._check_if_flagged(result) is True def test_check_if_flagged_returns_false_when_no_flagged_items(self): """Test that _check_if_flagged returns False when no items are flagged.""" @@ -291,17 +471,24 @@ class TestQualifireGuardrailCheckIfFlagged: ) # Result with evaluation results but nothing flagged - mock_inner_result = MagicMock() - mock_inner_result.flagged = False + result = { + "status": "completed", + "score": 95, + "evaluationResults": [ + { + "type": "prompt_injection", + "results": [ + { + "flagged": False, + "score": 0.95, + "reason": "No issues detected", + } + ], + } + ], + } - mock_eval_result = MagicMock() - mock_eval_result.results = [mock_inner_result] - - mock_result = MagicMock() - mock_result.status = "success" - mock_result.evaluationResults = [mock_eval_result] - - assert guardrail._check_if_flagged(mock_result) is False + assert guardrail._check_if_flagged(result) is False class TestQualifireGuardrailShouldRun: diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index b76957dbf39..134fc84965f 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -247,7 +247,9 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller ) @pytest.mark.flaky(reruns=3) @pytest.mark.asyncio -async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_controller): +async def test_normal_router_call_tpm_v3( + monkeypatch, rate_limit_object, time_controller +): """ Test normal router call with parallel request limiter v3 for TPM rate limiting """ @@ -394,8 +396,10 @@ async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_co # Manually increment the token counter to simulate token usage from previous call # This simulates what would happen after a successful call - await local_cache.async_increment_cache(key=counter_key, value=15, ttl=2) # Use up most of our 10 token limit - + await local_cache.async_increment_cache( + key=counter_key, value=15, ttl=2 + ) # Use up most of our 10 token limit + # Make another request to test rate limiting - this should fail as we've consumed tokens with pytest.raises(HTTPException) as exc_info: await parallel_request_handler.async_pre_call_hook( @@ -535,7 +539,9 @@ async def test_async_log_failure_event_v3(): ) # Mock kwargs with user_api_key via standard_logging_object - mock_kwargs = {"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}} + mock_kwargs = { + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}} + } # Capture pipeline operations captured_ops = [] @@ -785,7 +791,7 @@ async def test_tpm_api_key_rate_limits_v3(): tpm_limit_per_model=tpms, models=[], ) - + user_api_key_dict.metadata["model_tpm_limit"] = tpms user_api_key_dict.metadata["model_rpm_limit"] = rpms @@ -804,32 +810,45 @@ async def test_tpm_api_key_rate_limits_v3(): # Return Error response to ensure HTTPException return { "overall_code": "OVER_LIMIT", - "statuses": [{'code': 'OK', 'current_limit': 2, 'limit_remaining': 1, 'rate_limit_type': 'requests', 'descriptor_key': 'model_per_key'}, - {'code': 'OVER_LIMIT', 'current_limit': 2, 'limit_remaining': -18, 'rate_limit_type': 'tokens', 'descriptor_key': 'model_per_key'}] + "statuses": [ + { + "code": "OK", + "current_limit": 2, + "limit_remaining": 1, + "rate_limit_type": "requests", + "descriptor_key": "model_per_key", + }, + { + "code": "OVER_LIMIT", + "current_limit": 2, + "limit_remaining": -18, + "rate_limit_type": "tokens", + "descriptor_key": "model_per_key", + }, + ], } - + parallel_request_handler.should_rate_limit = mock_should_rate_limit - + # Test the pre-call hook error = None try: - await parallel_request_handler.async_pre_call_hook( + await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, data={"model": model}, call_type="", ) except HTTPException as e: - error=e + error = e assert e.status_code == 429 assert "rate_limit_type" in e.headers assert e.headers.get("rate_limit_type") == "tokens" assert "retry-after" in e.headers - - + assert error is not None, "An Exception must be thrown" assert captured_descriptors is not None, "Rate limit descriptors should be captured" - + model_per_key_descriptor = None for descriptor in captured_descriptors: if descriptor["key"] == "model_per_key": @@ -837,9 +856,15 @@ async def test_tpm_api_key_rate_limits_v3(): break assert model_per_key_descriptor is not None, "Api-Key descriptor should be present" - assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", "Api-Key value should combine api_key and model" - assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set" - assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set" + assert ( + model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}" + ), "Api-Key value should combine api_key and model" + assert ( + model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit + ), "Api-Key RPM limit should be set" + assert ( + model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit + ), "Api-Key TPM limit should be set" @pytest.mark.asyncio @@ -861,7 +886,7 @@ async def test_rpm_api_key_rate_limits_v3(): tpm_limit_per_model=tpms, models=[], ) - + user_api_key_dict.metadata["model_tpm_limit"] = tpms user_api_key_dict.metadata["model_rpm_limit"] = rpms @@ -880,31 +905,45 @@ async def test_rpm_api_key_rate_limits_v3(): # Return Error response to ensure HTTPException return { "overall_code": "OVER_LIMIT", - "statuses": [{'code': 'OVER_LIMIT', 'current_limit': 2, 'limit_remaining': -2, 'rate_limit_type': 'requests', 'descriptor_key': 'model_per_key'}, - {'code': 'OK', 'current_limit': 2, 'limit_remaining': 2, 'rate_limit_type': 'tokens', 'descriptor_key': 'model_per_key'}] + "statuses": [ + { + "code": "OVER_LIMIT", + "current_limit": 2, + "limit_remaining": -2, + "rate_limit_type": "requests", + "descriptor_key": "model_per_key", + }, + { + "code": "OK", + "current_limit": 2, + "limit_remaining": 2, + "rate_limit_type": "tokens", + "descriptor_key": "model_per_key", + }, + ], } - + parallel_request_handler.should_rate_limit = mock_should_rate_limit - + # Test the pre-call hook error = None try: - await parallel_request_handler.async_pre_call_hook( + await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, data={"model": model}, call_type="", ) except HTTPException as e: - error=e + error = e assert e.status_code == 429 assert "rate_limit_type" in e.headers assert e.headers.get("rate_limit_type") == "requests" assert "retry-after" in e.headers - + assert error is not None, "An Exception must be thrown" assert captured_descriptors is not None, "Rate limit descriptors should be captured" - + model_per_key_descriptor = None for descriptor in captured_descriptors: if descriptor["key"] == "model_per_key": @@ -912,9 +951,16 @@ async def test_rpm_api_key_rate_limits_v3(): break assert model_per_key_descriptor is not None, "Api-Key descriptor should be present" - assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", "Api-Key value should combine api_key and model" - assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set" - assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set" + assert ( + model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}" + ), "Api-Key value should combine api_key and model" + assert ( + model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit + ), "Api-Key RPM limit should be set" + assert ( + model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit + ), "Api-Key TPM limit should be set" + @pytest.mark.asyncio async def test_team_member_rate_limits_v3(): @@ -925,7 +971,7 @@ async def test_team_member_rate_limits_v3(): _api_key = hash_token(_api_key) _team_id = "team_123" _user_id = "user_456" - + user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, team_id=_team_id, @@ -933,7 +979,7 @@ async def test_team_member_rate_limits_v3(): team_member_rpm_limit=10, team_member_tpm_limit=1000, ) - + local_cache = DualCache() parallel_request_handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) @@ -947,15 +993,12 @@ async def test_team_member_rate_limits_v3(): nonlocal captured_descriptors captured_descriptors = descriptors # Return OK response to avoid HTTPException - return { - "overall_code": "OK", - "statuses": [] - } + return {"overall_code": "OK", "statuses": []} parallel_request_handler.should_rate_limit = mock_should_rate_limit # Test the pre-call hook - + await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, @@ -965,24 +1008,32 @@ async def test_team_member_rate_limits_v3(): # Verify team member descriptor was created assert captured_descriptors is not None, "Rate limit descriptors should be captured" - + team_member_descriptor = None for descriptor in captured_descriptors: if descriptor["key"] == "team_member": team_member_descriptor = descriptor break - - assert team_member_descriptor is not None, "Team member descriptor should be present" - assert team_member_descriptor["value"] == f"{_team_id}:{_user_id}", "Team member value should combine team_id and user_id" - assert team_member_descriptor["rate_limit"]["requests_per_unit"] == 10, "Team member RPM limit should be set" - assert team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000, "Team member TPM limit should be set" + + assert ( + team_member_descriptor is not None + ), "Team member descriptor should be present" + assert ( + team_member_descriptor["value"] == f"{_team_id}:{_user_id}" + ), "Team member value should combine team_id and user_id" + assert ( + team_member_descriptor["rate_limit"]["requests_per_unit"] == 10 + ), "Team member RPM limit should be set" + assert ( + team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000 + ), "Team member TPM limit should be set" @pytest.mark.asyncio async def test_dynamic_rate_limiting_v3(): """ Test that dynamic rate limiting only enforces limits when model has failures. - + When rpm_limit_type is set to "dynamic": - If model has no failures, rate limits should NOT be enforced (allow exceeding) - If model has failures above threshold, rate limits SHOULD be enforced @@ -990,75 +1041,75 @@ async def test_dynamic_rate_limiting_v3(): _api_key = "sk-12345" _api_key_hash = hash_token(_api_key) model = "gpt-3.5-turbo" - + # Set a low RPM limit to make testing easier user_api_key_dict = UserAPIKeyAuth( api_key=_api_key_hash, rpm_limit=2, metadata={"rpm_limit_type": "dynamic"}, ) - + local_cache = DualCache() parallel_request_handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) - + # Mock should_rate_limit to track if limits are enforced captured_descriptors = [] - + async def mock_should_rate_limit(descriptors, **kwargs): captured_descriptors.clear() captured_descriptors.extend(descriptors) return {"overall_code": "OK", "statuses": []} - + parallel_request_handler.should_rate_limit = mock_should_rate_limit - + # Test 1: No failures - rate limits should NOT be enforced (rpm_limit should be None) async def mock_check_no_failures(*args, **kwargs): return False - + parallel_request_handler._check_model_has_recent_failures = mock_check_no_failures - + await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, data={"model": model}, call_type="", ) - + # Find the API key descriptor api_key_descriptor = None for descriptor in captured_descriptors: if descriptor["key"] == "api_key": api_key_descriptor = descriptor break - + assert api_key_descriptor is not None, "API key descriptor should be present" assert ( api_key_descriptor["rate_limit"]["requests_per_unit"] is None ), "RPM limit should be None when dynamic mode and no failures" - + # Test 2: With failures - rate limits SHOULD be enforced (rpm_limit should be set) async def mock_check_with_failures(*args, **kwargs): return True - + parallel_request_handler._check_model_has_recent_failures = mock_check_with_failures captured_descriptors.clear() - + await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, data={"model": model}, call_type="", ) - + # Find the API key descriptor again api_key_descriptor = None for descriptor in captured_descriptors: if descriptor["key"] == "api_key": api_key_descriptor = descriptor break - + assert api_key_descriptor is not None, "API key descriptor should be present" assert ( api_key_descriptor["rate_limit"]["requests_per_unit"] == 2 @@ -1069,17 +1120,17 @@ async def test_dynamic_rate_limiting_v3(): async def test_async_increment_tokens_with_ttl_preservation(): """ Test TTL preservation functionality for token increment operations. - + This test verifies that: 1. Keys are created with proper TTL on first increment 2. TTL is preserved on subsequent increments (not reset) 3. Both TTL and non-TTL operations work correctly in the same call - + Environment variables required: - REDIS_HOST: Redis server hostname - REDIS_PORT: Redis server port - REDIS_PASSWORD: Redis password (optional) - + Test scenario: 1. First call: Create keys with TTL=60s and TTL=None 2. Wait 2 seconds @@ -1094,38 +1145,40 @@ async def test_async_increment_tokens_with_ttl_preservation(): # Skip test if Redis environment variables are not set redis_host = os.getenv("REDIS_HOST") - redis_port = os.getenv("REDIS_PORT") + redis_port = os.getenv("REDIS_PORT") redis_password = os.getenv("REDIS_PASSWORD") - + if not redis_host or not redis_port: pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set") - + # Setup Redis cache redis_cache = RedisCache( host=redis_host, port=int(redis_port), password=redis_password, ) - + local_cache = DualCache(redis_cache=redis_cache) parallel_request_handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) - + # Verify Redis connection is working try: await redis_cache.ping() except Exception as e: pytest.skip(f"Redis connection failed: {str(e)}") - + # Verify the TTL preservation script is registered if parallel_request_handler.token_increment_script is None: - pytest.skip("Token increment script not available - Redis Lua scripting may not be supported") - + pytest.skip( + "Token increment script not available - Redis Lua scripting may not be supported" + ) + # Test keys - use hash tags to ensure they map to same Redis cluster slot test_key_with_ttl = "{test_ttl}:with_ttl" test_key_without_ttl = "{test_ttl}:without_ttl" - + try: # Clean up any existing test keys try: @@ -1134,88 +1187,108 @@ async def test_async_increment_tokens_with_ttl_preservation(): except Exception: # Keys might not exist, ignore cleanup errors pass - + # First increment: Create operations with mixed TTL scenarios pipeline_operations_first = [ RedisPipelineIncrementOperation( - key=test_key_with_ttl, - increment_value=10.0, - ttl=60 + key=test_key_with_ttl, increment_value=10.0, ttl=60 ), RedisPipelineIncrementOperation( - key=test_key_without_ttl, - increment_value=5.0, - ttl=None # No TTL - ) + key=test_key_without_ttl, increment_value=5.0, ttl=None # No TTL + ), ] - + # Execute first increment await parallel_request_handler.async_increment_tokens_with_ttl_preservation( pipeline_operations=pipeline_operations_first ) - + # Small delay to ensure Redis has processed the commands await asyncio.sleep(0.1) - + # Verify keys exist and check initial TTL ttl_after_first = await redis_cache.async_get_ttl(test_key_with_ttl) - value_after_first_with_ttl = await redis_cache.async_get_cache(test_key_with_ttl) - value_after_first_without_ttl = await redis_cache.async_get_cache(test_key_without_ttl) - - assert value_after_first_with_ttl == 10.0, f"First increment should set value to 10.0, got {value_after_first_with_ttl}" - assert value_after_first_without_ttl == 5.0, "First increment should set value to 5.0" - assert ttl_after_first is not None and ttl_after_first > 0, "Key with TTL should have positive TTL after first increment" + value_after_first_with_ttl = await redis_cache.async_get_cache( + test_key_with_ttl + ) + value_after_first_without_ttl = await redis_cache.async_get_cache( + test_key_without_ttl + ) + + assert ( + value_after_first_with_ttl == 10.0 + ), f"First increment should set value to 10.0, got {value_after_first_with_ttl}" + assert ( + value_after_first_without_ttl == 5.0 + ), "First increment should set value to 5.0" + assert ( + ttl_after_first is not None and ttl_after_first > 0 + ), "Key with TTL should have positive TTL after first increment" assert ttl_after_first <= 60, "TTL should not exceed the set value" - + # Check TTL for key without TTL (should be None, meaning no expiry) ttl_no_ttl_key = await redis_cache.async_get_ttl(test_key_without_ttl) - assert ttl_no_ttl_key is None, "Key without TTL should have no expiry (None from async_get_ttl)" - + assert ( + ttl_no_ttl_key is None + ), "Key without TTL should have no expiry (None from async_get_ttl)" + # Wait a moment to ensure TTL decreases await asyncio.sleep(2) - + # Second increment: Same operations to test TTL preservation pipeline_operations_second = [ RedisPipelineIncrementOperation( - key=test_key_with_ttl, - increment_value=15.0, - ttl=60 # Same TTL value + key=test_key_with_ttl, increment_value=15.0, ttl=60 # Same TTL value ), RedisPipelineIncrementOperation( - key=test_key_without_ttl, - increment_value=7.0, - ttl=None # No TTL - ) + key=test_key_without_ttl, increment_value=7.0, ttl=None # No TTL + ), ] - + # Execute second increment await parallel_request_handler.async_increment_tokens_with_ttl_preservation( pipeline_operations=pipeline_operations_second ) - + # Small delay to ensure Redis has processed the commands await asyncio.sleep(0.1) - + # Verify TTL preservation and value updates ttl_after_second = await redis_cache.async_get_ttl(test_key_with_ttl) - value_after_second_with_ttl = await redis_cache.async_get_cache(test_key_with_ttl) - value_after_second_without_ttl = await redis_cache.async_get_cache(test_key_without_ttl) - - assert value_after_second_with_ttl == 25.0, "Second increment should update value to 25.0" - assert value_after_second_without_ttl == 12.0, "Second increment should update value to 12.0" - + value_after_second_with_ttl = await redis_cache.async_get_cache( + test_key_with_ttl + ) + value_after_second_without_ttl = await redis_cache.async_get_cache( + test_key_without_ttl + ) + + assert ( + value_after_second_with_ttl == 25.0 + ), "Second increment should update value to 25.0" + assert ( + value_after_second_without_ttl == 12.0 + ), "Second increment should update value to 12.0" + # Critical test: TTL should be preserved (not reset to 60) assert ttl_after_second is not None, "TTL should still exist" - assert ttl_after_second < ttl_after_first, "TTL should have decreased (not been reset)" + assert ( + ttl_after_second < ttl_after_first + ), "TTL should have decreased (not been reset)" assert ttl_after_second > 0, "TTL should still be positive" - + # TTL should not be close to the original 60 seconds (proving it wasn't reset) - assert ttl_after_second < 59, "TTL should be significantly less than original, proving preservation" - + assert ( + ttl_after_second < 59 + ), "TTL should be significantly less than original, proving preservation" + # Key without TTL should still have no expiry - ttl_no_ttl_key_after_second = await redis_cache.async_get_ttl(test_key_without_ttl) - assert ttl_no_ttl_key_after_second is None, "Key without TTL should still have no expiry" - + ttl_no_ttl_key_after_second = await redis_cache.async_get_ttl( + test_key_without_ttl + ) + assert ( + ttl_no_ttl_key_after_second is None + ), "Key without TTL should still have no expiry" + finally: # Clean up test keys try: @@ -1224,7 +1297,7 @@ async def test_async_increment_tokens_with_ttl_preservation(): except Exception: # Ignore cleanup errors pass - + # Properly close Redis connections to prevent warnings try: await redis_cache.disconnect() @@ -1239,115 +1312,125 @@ async def test_async_increment_tokens_fallback_behavior(): Test fallback behavior when Lua script is not available. """ from litellm.types.caching import RedisPipelineIncrementOperation - + local_cache = DualCache() parallel_request_handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) - + # Mock the token_increment_script to None to simulate unavailable script parallel_request_handler.token_increment_script = None - + # Mock the fallback method fallback_called = False - original_method = parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline - + original_method = ( + parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline + ) + async def mock_fallback(*args, **kwargs): nonlocal fallback_called fallback_called = True return await original_method(*args, **kwargs) - - parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = mock_fallback - + + parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( + mock_fallback + ) + # Test operations pipeline_operations = [ RedisPipelineIncrementOperation( - key="test_fallback_key", - increment_value=10.0, - ttl=60 + key="test_fallback_key", increment_value=10.0, ttl=60 ) ] - + # Execute increment await parallel_request_handler.async_increment_tokens_with_ttl_preservation( pipeline_operations=pipeline_operations ) - + # Verify fallback was called - assert fallback_called, "Fallback method should be called when Lua script is not available" + assert ( + fallback_called + ), "Fallback method should be called when Lua script is not available" # Redis Cluster Compatibility Tests def test_group_keys_by_hash_tag_regular_redis(): """ Test that keys are correctly grouped for regular Redis (non-cluster). - + For regular Redis, all keys should be grouped together under a single group. """ local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) - + # Test keys with different hash tags test_keys = [ "{api_key:sk-123}:window", - "{api_key:sk-123}:requests", + "{api_key:sk-123}:requests", "{api_key:sk-123}:tokens", "{user:user-456}:window", "{user:user-456}:requests", "{team:team-789}:window", "{team:team-789}:tokens", - "no_hash_tag_key" + "no_hash_tag_key", ] - + # Group the keys (should be single group for regular Redis) groups = handler._group_keys_by_hash_tag(test_keys) - + # Verify all keys are in single group for regular Redis assert len(groups) == 1, f"Expected 1 group for regular Redis, got {len(groups)}" assert "all_keys" in groups, "Expected 'all_keys' group for regular Redis" - assert set(groups["all_keys"]) == set(test_keys), "All keys should be in single group" + assert set(groups["all_keys"]) == set( + test_keys + ), "All keys should be in single group" def test_group_keys_by_hash_tag_redis_cluster(): """ Test that keys are correctly grouped by Redis cluster slots when using Redis cluster. - + This ensures that keys are grouped by their slot number for cluster compatibility. """ from unittest.mock import patch - + local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) - + # Mock _is_redis_cluster to return True - with patch.object(handler, '_is_redis_cluster', return_value=True): + with patch.object(handler, "_is_redis_cluster", return_value=True): # Test keys with different hash tags test_keys = [ "{api_key:sk-123}:window", - "{api_key:sk-123}:requests", + "{api_key:sk-123}:requests", "{user:user-456}:window", "{user:user-456}:requests", ] - + # Group the keys (should be grouped by slot for Redis cluster) groups = handler._group_keys_by_hash_tag(test_keys) - + # Verify keys are grouped by slot assert len(groups) >= 1, "Should have at least 1 slot group" - + # All group keys should start with "slot_" for group_key in groups.keys(): - assert group_key.startswith("slot_"), f"Group key {group_key} should start with 'slot_'" - + assert group_key.startswith( + "slot_" + ), f"Group key {group_key} should start with 'slot_'" + # Verify all original keys are present across groups all_grouped_keys = [] for group_keys in groups.values(): all_grouped_keys.extend(group_keys) - assert set(all_grouped_keys) == set(test_keys), "All keys should be present in groups" + assert set(all_grouped_keys) == set( + test_keys + ), "All keys should be present in groups" def test_keyslot_for_redis_cluster(): @@ -1358,16 +1441,16 @@ def test_keyslot_for_redis_cluster(): handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) - + # Test basic key slot1 = handler.keyslot_for_redis_cluster("user:1000") assert 0 <= slot1 < 16384, "Slot should be in valid range" - + # Test key with hash tag slot2 = handler.keyslot_for_redis_cluster("foo{bar}baz") slot3 = handler.keyslot_for_redis_cluster("{bar}") assert slot2 == slot3, "Keys with same hash tag should have same slot" - + # Test keys with same hash tag should have same slot slot4 = handler.keyslot_for_redis_cluster("{api_key:sk-123}:requests") slot5 = handler.keyslot_for_redis_cluster("{api_key:sk-123}:window") @@ -1379,67 +1462,70 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): """ Test that the Redis batch rate limiter script execution handles cluster compatibility by grouping keys and falling back gracefully on errors. - + This simulates the Redis cluster error scenario and verifies fallback behavior. """ from unittest.mock import AsyncMock, patch - + local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) - + # Mock _is_redis_cluster to return True for this test - with patch.object(handler, '_is_redis_cluster', return_value=True): + with patch.object(handler, "_is_redis_cluster", return_value=True): # Mock script that simulates Redis cluster slot conflict mock_script = AsyncMock() mock_script.side_effect = [ - Exception("EVALSHA - all keys must map to the same key slot"), # First group fails - [1234, 1, 1234, 2] # Second group succeeds + Exception( + "EVALSHA - all keys must map to the same key slot" + ), # First group fails + [1234, 1, 1234, 2], # Second group succeeds ] handler.batch_rate_limiter_script = mock_script - + # Mock in-memory fallback (returns 2 values for 2 keys: window_start, counter) handler.in_memory_cache_sliding_window = AsyncMock(return_value=[1234, 1]) - + # Test keys from different hash tags (would fail in cluster without grouping) test_keys = [ "{api_key:sk-123}:window", "{api_key:sk-123}:requests", - "{user:user-456}:window", - "{user:user-456}:requests" + "{user:user-456}:window", + "{user:user-456}:requests", ] - + # Execute the method results = await handler._execute_redis_batch_rate_limiter_script( - keys_to_fetch=test_keys, - now_int=1234 + keys_to_fetch=test_keys, now_int=1234 ) - + # Verify results: 2 from fallback + 4 from successful script = 6 total assert len(results) == 6, f"Expected 6 results, got {len(results)}" - + # Verify script was called twice (once per slot group) assert mock_script.call_count == 2 - + # Verify fallback was called for the failed group handler.in_memory_cache_sliding_window.assert_called_once() - + # Verify the calls were made with grouped keys call_args_list = mock_script.call_args_list - + # Both calls should have keys, but we can't predict exact grouping without knowing slots # Just verify that keys were grouped and calls were made assert len(call_args_list) == 2, "Should have made 2 script calls" - + # Verify all keys were processed all_processed_keys = [] for call_args in call_args_list: - all_processed_keys.extend(call_args[1]['keys']) - + all_processed_keys.extend(call_args[1]["keys"]) + # Should have processed all keys (some might be duplicated due to fallback) unique_processed_keys = set(all_processed_keys) - assert len(unique_processed_keys) >= 2, "Should have processed at least some keys" + assert ( + len(unique_processed_keys) >= 2 + ), "Should have processed at least some keys" @pytest.mark.asyncio @@ -1485,23 +1571,23 @@ async def test_multiple_rate_limits_per_descriptor(): "current_limit": 2, "limit_remaining": 1, "rate_limit_type": "requests", - "descriptor_key": "api_key" + "descriptor_key": "api_key", }, { "code": "OK", "current_limit": 10, "limit_remaining": 8, "rate_limit_type": "tokens", - "descriptor_key": "api_key" + "descriptor_key": "api_key", }, { "code": "OVER_LIMIT", "current_limit": 1, "limit_remaining": -1, "rate_limit_type": "max_parallel_requests", - "descriptor_key": "api_key" - } - ] + "descriptor_key": "api_key", + }, + ], } parallel_request_handler.should_rate_limit = mock_should_rate_limit @@ -1560,9 +1646,9 @@ async def test_missing_descriptor_fallback(): "current_limit": 2, "limit_remaining": -1, "rate_limit_type": "requests", - "descriptor_key": "nonexistent_key" # This won't match any descriptor + "descriptor_key": "nonexistent_key", # This won't match any descriptor } - ] + ], } parallel_request_handler.should_rate_limit = mock_should_rate_limit @@ -1597,14 +1683,17 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): # Mock general_settings to return empty dict (no token_rate_limit_type set) import litellm.proxy.proxy_server as proxy_server - original_settings = getattr(proxy_server, 'general_settings', {}) - monkeypatch.setattr(proxy_server, 'general_settings', {}) + + original_settings = getattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "general_settings", {}) try: result = parallel_request_handler.get_rate_limit_type() - assert result == "total", f"Default rate limit type should be 'total', got '{result}'" + assert ( + result == "total" + ), f"Default rate limit type should be 'total', got '{result}'" finally: - monkeypatch.setattr(proxy_server, 'general_settings', original_settings) + monkeypatch.setattr(proxy_server, "general_settings", original_settings) @pytest.mark.asyncio @@ -1619,14 +1708,19 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): # Mock general_settings to return an invalid token_rate_limit_type import litellm.proxy.proxy_server as proxy_server - original_settings = getattr(proxy_server, 'general_settings', {}) - monkeypatch.setattr(proxy_server, 'general_settings', {'token_rate_limit_type': 'invalid_type'}) + + original_settings = getattr(proxy_server, "general_settings", {}) + monkeypatch.setattr( + proxy_server, "general_settings", {"token_rate_limit_type": "invalid_type"} + ) try: result = parallel_request_handler.get_rate_limit_type() - assert result == "total", f"Invalid rate limit type should fall back to 'total', got '{result}'" + assert ( + result == "total" + ), f"Invalid rate limit type should fall back to 'total', got '{result}'" finally: - monkeypatch.setattr(proxy_server, 'general_settings', original_settings) + monkeypatch.setattr(proxy_server, "general_settings", original_settings) @pytest.mark.parametrize( @@ -1638,7 +1732,9 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): ], ) @pytest.mark.asyncio -async def test_async_log_success_event_with_dict_usage(monkeypatch, token_rate_limit_type, expected_field): +async def test_async_log_success_event_with_dict_usage( + monkeypatch, token_rate_limit_type, expected_field +): """ Test that async_log_success_event correctly handles usage as a dict (Responses API format). @@ -1664,13 +1760,13 @@ async def test_async_log_success_event_with_dict_usage(monkeypatch, token_rate_l # Create a mock response object with usage as a dict (Responses API format) from litellm.types.utils import BaseLiteLLMOpenAIResponseObject - + # Use spec to make isinstance checks work correctly with MagicMock mock_response = MagicMock(spec=BaseLiteLLMOpenAIResponseObject) mock_response.usage = { "prompt_tokens": 25, "completion_tokens": 35, - "total_tokens": 60 + "total_tokens": 60, } # Create mock kwargs for the success event @@ -1760,7 +1856,10 @@ async def test_async_log_success_event_with_dict_usage_missing_fields(monkeypatc # total_tokens is missing } from litellm.types.utils import BaseLiteLLMOpenAIResponseObject - mock_response.__class__ = type('MockResponse', (BaseLiteLLMOpenAIResponseObject,), {}) + + mock_response.__class__ = type( + "MockResponse", (BaseLiteLLMOpenAIResponseObject,), {} + ) # Create mock kwargs for the success event mock_kwargs = { @@ -1805,7 +1904,9 @@ async def test_async_log_success_event_with_dict_usage_missing_fields(monkeypatc assert tpm_operation is not None, "Should have a TPM increment operation" # Should default to 0 when field is missing - assert tpm_operation["increment_value"] == 0, "Should default to 0 when completion_tokens is missing" + assert ( + tpm_operation["increment_value"] == 0 + ), "Should default to 0 when completion_tokens is missing" @pytest.mark.asyncio @@ -1813,68 +1914,154 @@ async def test_execute_token_increment_script_cluster_compatibility(): """ Test that token increment script execution handles Redis cluster compatibility by grouping operations by slot. - + This ensures token increments work correctly in cluster environments. """ from typing import List from unittest.mock import AsyncMock, patch from litellm.types.caching import RedisPipelineIncrementOperation - + local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) ) - + # Mock _is_redis_cluster to return True for this test - with patch.object(handler, '_is_redis_cluster', return_value=True): + with patch.object(handler, "_is_redis_cluster", return_value=True): # Mock script mock_script = AsyncMock() handler.token_increment_script = mock_script - + # Create pipeline operations with different hash tags pipeline_operations: List[RedisPipelineIncrementOperation] = [ + {"key": "{api_key:sk-123}:tokens", "increment_value": 100, "ttl": 60}, { - "key": "{api_key:sk-123}:tokens", - "increment_value": 100, - "ttl": 60 - }, - { - "key": "{api_key:sk-123}:max_parallel_requests", + "key": "{api_key:sk-123}:max_parallel_requests", "increment_value": -1, - "ttl": 60 + "ttl": 60, }, - { - "key": "{user:user-456}:tokens", - "increment_value": 50, - "ttl": 60 - } + {"key": "{user:user-456}:tokens", "increment_value": 50, "ttl": 60}, ] - + # Execute the method await handler._execute_token_increment_script(pipeline_operations) - + # Verify script was called (at least once, possibly more depending on slot grouping) assert mock_script.call_count >= 1, "Script should be called at least once" - + call_args_list = mock_script.call_args_list - + # Verify all operations were processed all_processed_keys = [] for call_args in call_args_list: - all_processed_keys.extend(call_args[1]['keys']) - + all_processed_keys.extend(call_args[1]["keys"]) + # Should have processed all 3 keys expected_keys = { "{api_key:sk-123}:tokens", "{api_key:sk-123}:max_parallel_requests", - "{user:user-456}:tokens" + "{user:user-456}:tokens", } - assert set(all_processed_keys) == expected_keys, "All operation keys should be processed" - + assert ( + set(all_processed_keys) == expected_keys + ), "All operation keys should be processed" + # Verify args structure is correct for each call for call_args in call_args_list: - keys = call_args[1]['keys'] - args = call_args[1]['args'] + keys = call_args[1]["keys"] + args = call_args[1]["args"] # Each key should have 2 args (increment_value, ttl) - assert len(args) == len(keys) * 2, f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" + assert ( + len(args) == len(keys) * 2 + ), f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" + + +class TestGetTotalTokensFromUsageCacheExclusion: + """ + Tests for _get_total_tokens_from_usage cache token exclusion. + + Issue: AWS Bedrock and similar providers exclude cache tokens from TPM calculation, + but LiteLLM was including them, causing up to 10x difference in rate limiting. + """ + + @pytest.fixture + def handler(self): + """Create a handler instance for testing.""" + local_cache = DualCache() + return _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + + def test_excludes_cached_tokens_from_total(self, handler): + """Cached tokens should be excluded from total token count.""" + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800), + ) + + # Total should be 1500 - 800 = 700 + result = handler._get_total_tokens_from_usage(usage, "total") + assert result == 700, f"Expected 700 (1500 - 800 cached), got {result}" + + def test_excludes_cached_tokens_from_input(self, handler): + """Cached tokens should be excluded from input token count.""" + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800), + ) + + # Input should be 1000 - 800 = 200 + result = handler._get_total_tokens_from_usage(usage, "input") + assert result == 200, f"Expected 200 (1000 - 800 cached), got {result}" + + def test_does_not_exclude_cached_tokens_from_output(self, handler): + """Cached tokens should NOT affect output token count.""" + from litellm.types.utils import PromptTokensDetailsWrapper + + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800), + ) + + # Output tokens should be unchanged + result = handler._get_total_tokens_from_usage(usage, "output") + assert result == 500, f"Expected 500 (no change for output), got {result}" + + def test_handles_no_cached_tokens(self, handler): + """Should work correctly when no cached tokens present.""" + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + ) + + result = handler._get_total_tokens_from_usage(usage, "total") + assert result == 1500, f"Expected 1500 (no cache), got {result}" + + def test_handles_dict_usage_with_cached_tokens(self, handler): + """Should handle dict usage format (Responses API) with cached tokens.""" + usage = { + "prompt_tokens": 1000, + "completion_tokens": 500, + "total_tokens": 1500, + "prompt_tokens_details": {"cached_tokens": 600}, + } + + result = handler._get_total_tokens_from_usage(usage, "total") + assert result == 900, f"Expected 900 (1500 - 600 cached), got {result}" + + def test_handles_none_usage(self, handler): + """Should handle None usage gracefully.""" + result = handler._get_total_tokens_from_usage(None, "total") + assert result == 0, f"Expected 0 for None usage, got {result}" diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index f2bae2cb14a..bc223d15d5f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1,20 +1,20 @@ import json import os import sys -from litellm._uuid import uuid +import types from datetime import datetime, timedelta -from typing import List +from typing import List, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import FastAPI from fastapi.testclient import TestClient +from litellm._uuid import uuid sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from typing import Optional - from litellm.proxy._types import ( LiteLLM_MCPServerTable, LitellmUserRoles, @@ -118,6 +118,22 @@ def setup_mock_prisma_client( return mock_prisma_client +def create_mcp_router_test_client() -> TestClient: + from litellm.proxy.management_endpoints.mcp_management_endpoints import router + + app = FastAPI() + app.include_router(router) + return TestClient(app) + + +def patch_proxy_general_settings(settings: dict): + fake_proxy_server_module = types.SimpleNamespace(general_settings=settings) + return patch.dict( + sys.modules, + {"litellm.proxy.proxy_server": fake_proxy_server_module}, + ) + + class TestListMCPServers: """Test suite for list MCP servers functionality""" @@ -1082,6 +1098,55 @@ class TestHealthCheckServers: assert result[1]["server_id"] == "server-2" assert result[1]["status"] == "unhealthy" + +class TestMCPRegistryEndpoint: + def test_registry_returns_404_when_flag_missing(self): + client = create_mcp_router_test_client() + + with patch_proxy_general_settings({}): + response = client.get("/v1/mcp/registry.json") + + assert response.status_code == 404 + + def test_registry_returns_404_when_flag_false(self): + client = create_mcp_router_test_client() + + with patch_proxy_general_settings({"enable_mcp_registry": False}): + response = client.get("/v1/mcp/registry.json") + + assert response.status_code == 404 + + def test_registry_returns_entries_when_enabled(self): + client = create_mcp_router_test_client() + + mock_server = generate_mock_mcp_server_config_record( + server_id="server-123", + name="zapier", + url="https://zapier.example.com/mcp", + transport="http", + ) + + mock_manager = MagicMock() + mock_manager.get_registry.return_value = {mock_server.server_id: mock_server} + + with patch_proxy_general_settings({"enable_mcp_registry": True}), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ): + response = client.get("/v1/mcp/registry.json") + + assert response.status_code == 200 + data = response.json() + assert len(data["servers"]) == 2 # built-in + custom server + + builtin_entry = data["servers"][0]["server"] + assert builtin_entry["name"] == "litellm-mcp-server" + assert builtin_entry["remotes"][0]["url"].endswith("/mcp") + + custom_entry = data["servers"][1]["server"] + assert custom_entry["name"] == "zapier" + assert custom_entry["remotes"][0]["url"].endswith("/zapier/mcp") + @pytest.mark.asyncio async def test_health_check_specific_servers(self): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py new file mode 100644 index 00000000000..1f5473e75d4 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py @@ -0,0 +1,71 @@ +""" +Tests for router settings management endpoints. + +Tests the GET endpoints for router settings and router fields. +""" +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../../..") +) + +from litellm.proxy.proxy_server import app + +client = TestClient(app) + + +class TestRouterSettingsEndpoints: + """Test suite for router settings endpoints""" + + @pytest.mark.asyncio + async def test_get_router_fields_success(self): + """ + Test GET /router/fields endpoint successfully returns field definitions without values. + """ + # Make request to router fields endpoint + response = client.get( + "/router/fields", + headers={"Authorization": "Bearer sk-1234"} + ) + + # Verify response + assert response.status_code == 200 + + response_data = response.json() + + # Verify response structure + assert "fields" in response_data + assert "routing_strategy_descriptions" in response_data + + # Verify fields is a list + assert isinstance(response_data["fields"], list) + assert len(response_data["fields"]) > 0 + + # Verify each field has required properties and field_value is None + for field in response_data["fields"]: + assert "field_name" in field + assert "field_type" in field + assert "field_description" in field + assert "field_default" in field + assert "ui_field_name" in field + assert "field_value" in field + assert field["field_value"] is None # Ensure field_value is None + + # Verify routing_strategy_descriptions is a dict + assert isinstance(response_data["routing_strategy_descriptions"], dict) + assert len(response_data["routing_strategy_descriptions"]) > 0 + + # Verify routing_strategy field has options populated + routing_strategy_field = next( + (f for f in response_data["fields"] if f["field_name"] == "routing_strategy"), + None + ) + assert routing_strategy_field is not None + assert "options" in routing_strategy_field + assert isinstance(routing_strategy_field["options"], list) + assert len(routing_strategy_field["options"]) > 0 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index bab95feec06..3d1e9aece41 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import Request, status -from fastapi.responses import StreamingResponse +from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid @@ -11,9 +11,10 @@ from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ProxyConfig, + _extract_error_from_sse_chunk, _get_cost_breakdown_from_logging_obj, _parse_event_data_for_error, - create_streaming_response, + create_response, ) from litellm.proxy.utils import ProxyLogging @@ -680,21 +681,27 @@ class TestCommonRequestProcessingHelpers: assert await _parse_event_data_for_error(event_line) == expected_code async def test_create_streaming_response_first_chunk_is_error(self): + """ + Test that when the first chunk is an error, a JSON error response is returned + instead of an SSE streaming response + """ async def mock_generator(): yield 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n' yield 'data: {"content": "more data"}\n\n' yield "data: [DONE]\n\n" - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {} ) + # Should return JSONResponse instead of StreamingResponse + assert isinstance(response, JSONResponse) assert response.status_code == status.HTTP_403_FORBIDDEN - content = await self.consume_stream(response) - assert content == [ - 'data: {"error": {"code": 403, "message": "forbidden"}}\n\n', - 'data: {"content": "more data"}\n\n', - "data: [DONE]\n\n", - ] + # Verify the response is in standard JSON error format + import json + body = json.loads(response.body.decode()) + assert "error" in body + assert body["error"]["code"] == 403 + assert body["error"]["message"] == "forbidden" async def test_create_streaming_response_first_chunk_not_error(self): async def mock_generator(): @@ -702,7 +709,7 @@ class TestCommonRequestProcessingHelpers: yield 'data: {"content": "second part"}\n\n' yield "data: [DONE]\n\n" - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {} ) assert response.status_code == status.HTTP_200_OK @@ -719,7 +726,7 @@ class TestCommonRequestProcessingHelpers: yield # Implicitly raises StopAsyncIteration - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {} ) assert response.status_code == status.HTTP_200_OK @@ -732,7 +739,7 @@ class TestCommonRequestProcessingHelpers: mock_gen = AsyncMock() mock_gen.__anext__.side_effect = StopAsyncIteration - response = await create_streaming_response(mock_gen, "text/event-stream", {}) + response = await create_response(mock_gen, "text/event-stream", {}) assert response.status_code == status.HTTP_200_OK content = await self.consume_stream(response) assert content == [] @@ -743,7 +750,7 @@ class TestCommonRequestProcessingHelpers: mock_gen = AsyncMock() mock_gen.__anext__.side_effect = ValueError("Test error from generator") - response = await create_streaming_response(mock_gen, "text/event-stream", {}) + response = await create_response(mock_gen, "text/event-stream", {}) assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR content = await self.consume_stream(response) expected_error_data = { @@ -760,19 +767,24 @@ class TestCommonRequestProcessingHelpers: assert content[1] == "data: [DONE]\n\n" async def test_create_streaming_response_first_chunk_error_string_code(self): + """ + Test that when the first chunk contains a string error code, a JSON error response is returned + """ async def mock_generator(): yield 'data: {"error": {"code": "429", "message": "too many requests"}}\n\n' yield "data: [DONE]\n\n" - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {} ) + assert isinstance(response, JSONResponse) assert response.status_code == status.HTTP_429_TOO_MANY_REQUESTS - content = await self.consume_stream(response) - assert content == [ - 'data: {"error": {"code": "429", "message": "too many requests"}}\n\n', - "data: [DONE]\n\n", - ] + # Verify the response is in standard JSON error format + import json + body = json.loads(response.body.decode()) + assert "error" in body + assert body["error"]["code"] == "429" + assert body["error"]["message"] == "too many requests" async def test_create_streaming_response_custom_headers(self): async def mock_generator(): @@ -780,7 +792,7 @@ class TestCommonRequestProcessingHelpers: yield "data: [DONE]\n\n" custom_headers = {"X-Custom-Header": "TestValue"} - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", custom_headers ) assert response.headers["x-custom-header"] == "TestValue" @@ -790,7 +802,7 @@ class TestCommonRequestProcessingHelpers: yield 'data: {"content": "data"}\n\n' yield "data: [DONE]\n\n" - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {}, @@ -807,7 +819,7 @@ class TestCommonRequestProcessingHelpers: async def mock_generator(): yield "data: [DONE]\n\n" - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {} ) assert response.status_code == status.HTTP_200_OK # Default status @@ -820,7 +832,7 @@ class TestCommonRequestProcessingHelpers: yield 'data: {"content": "actual data"}\n\n' yield "data: [DONE]\n\n" - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {} ) assert response.status_code == status.HTTP_200_OK # Default status @@ -851,7 +863,7 @@ class TestCommonRequestProcessingHelpers: # Patch the tracer in the common_request_processing module with patch("litellm.proxy.common_request_processing.tracer", mock_tracer): - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {} ) @@ -888,7 +900,10 @@ class TestCommonRequestProcessingHelpers: ), f"Call {i} should have operation name 'streaming.chunk.yield', got {args[0]}" async def test_create_streaming_response_dd_trace_with_error_chunk(self): - """Test that dd trace is applied even when the first chunk contains an error""" + """ + Test that when the first chunk contains an error, JSONResponse is returned + and tracing is not triggered (since it's not a streaming response) + """ from unittest.mock import patch # Create a mock tracer @@ -905,28 +920,107 @@ class TestCommonRequestProcessingHelpers: # Patch the tracer in the common_request_processing module with patch("litellm.proxy.common_request_processing.tracer", mock_tracer): - response = await create_streaming_response( + response = await create_response( mock_generator(), "text/event-stream", {} ) - # Even with error, status should be set to error code but tracing should still work + # Should return JSONResponse instead of StreamingResponse + assert isinstance(response, JSONResponse) assert response.status_code == 400 - # Consume the stream to trigger the tracer calls - content = await self.consume_stream(response) + # Verify the response is in standard JSON error format + import json + body = json.loads(response.body.decode()) + assert "error" in body + assert body["error"]["code"] == 400 + assert body["error"]["message"] == "bad request" - # Verify all chunks are present - assert len(content) == 3 + # Since JSONResponse is returned instead of StreamingResponse, streaming tracing should not be triggered + # tracer.trace should not be called + assert mock_tracer.trace.call_count == 0 - # Verify that tracer.trace was called for each chunk - assert mock_tracer.trace.call_count == 3 - # Verify that each call was made with the correct operation name - actual_calls = mock_tracer.trace.call_args_list - assert len(actual_calls) == 3 +class TestExtractErrorFromSSEChunk: + """Tests for _extract_error_from_sse_chunk function""" + + def test_extract_error_from_sse_chunk_with_valid_error(self): + """Test extracting error information from a standard SSE chunk""" + chunk = 'data: {"error": {"code": 403, "message": "forbidden", "type": "auth_error", "param": "api_key"}}\n\n' + error = _extract_error_from_sse_chunk(chunk) + + assert error["code"] == 403 + assert error["message"] == "forbidden" + assert error["type"] == "auth_error" + assert error["param"] == "api_key" + + def test_extract_error_from_sse_chunk_with_string_code(self): + """Test error code as string type""" + chunk = 'data: {"error": {"code": "429", "message": "too many requests"}}\n\n' + error = _extract_error_from_sse_chunk(chunk) + + assert error["code"] == "429" + assert error["message"] == "too many requests" + + def test_extract_error_from_sse_chunk_with_bytes(self): + """Test input as bytes type""" + chunk = b'data: {"error": {"code": 500, "message": "internal error"}}\n\n' + error = _extract_error_from_sse_chunk(chunk) + + assert error["code"] == 500 + assert error["message"] == "internal error" + + def test_extract_error_from_sse_chunk_with_done(self): + """Test [DONE] marker should return default error""" + chunk = "data: [DONE]\n\n" + error = _extract_error_from_sse_chunk(chunk) + + assert error["message"] == "Unknown error" + assert error["type"] == "internal_server_error" + assert error["code"] == "500" + assert error["param"] is None + + def test_extract_error_from_sse_chunk_without_error_field(self): + """Test missing error field should return default error""" + chunk = 'data: {"content": "some content"}\n\n' + error = _extract_error_from_sse_chunk(chunk) + + assert error["message"] == "Unknown error" + assert error["type"] == "internal_server_error" + assert error["code"] == "500" + + def test_extract_error_from_sse_chunk_with_invalid_json(self): + """Test invalid JSON should return default error""" + chunk = 'data: {invalid json}\n\n' + error = _extract_error_from_sse_chunk(chunk) + + assert error["message"] == "Unknown error" + assert error["type"] == "internal_server_error" + assert error["code"] == "500" + + def test_extract_error_from_sse_chunk_without_data_prefix(self): + """Test missing 'data:' prefix should return default error""" + chunk = '{"error": {"code": 400, "message": "bad request"}}\n\n' + error = _extract_error_from_sse_chunk(chunk) + + assert error["message"] == "Unknown error" + assert error["type"] == "internal_server_error" + assert error["code"] == "500" + + def test_extract_error_from_sse_chunk_with_empty_string(self): + """Test empty string should return default error""" + chunk = "" + error = _extract_error_from_sse_chunk(chunk) + + assert error["message"] == "Unknown error" + assert error["type"] == "internal_server_error" + assert error["code"] == "500" + + def test_extract_error_from_sse_chunk_with_minimal_error(self): + """Test minimal error object""" + chunk = 'data: {"error": {"message": "error occurred"}}\n\n' + error = _extract_error_from_sse_chunk(chunk) + + assert error["message"] == "error occurred" + # Other fields should be obtained from the original error object (if exists) + - for i, call in enumerate(actual_calls): - args, kwargs = call - assert ( - args[0] == "streaming.chunk.yield" - ), f"Call {i} should have operation name 'streaming.chunk.yield', got {args[0]}" diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 96ac2e2c345..09628cd4a76 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -203,3 +203,22 @@ class TestResponseAPILoggingUtils: assert result.prompt_tokens == 0 assert result.completion_tokens == 20 assert result.total_tokens == 20 + + def test_transform_response_api_usage_calculates_total_from_input_and_output_tokens_if_available(self): + """Test transformation calculates total_tokens when it's None and input / output tokens are present""" + # Setup + usage = { + "input_tokens": 15, + "output_tokens": 25, + "total_tokens": None, + } + + # Execute + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + + # Assert + assert result.prompt_tokens == 15 + assert result.completion_tokens == 25 + assert result.total_tokens == 40 # 15 + 25 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index cd76c438ded..bfa162e019b 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -749,6 +749,57 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): raise AssertionError(error_message) +def test_max_tokens_consistency(): + """ + Test that max_tokens == max_output_tokens for all models. + + According to the spec in model_prices_and_context_window.json: + - max_tokens is a LEGACY parameter + - It should be set to max_output_tokens if the provider specifies it + + This test ensures consistency across all model definitions. + """ + import json + from pathlib import Path + + # Load the model configuration + config_path = Path(__file__).parent.parent.parent / "model_prices_and_context_window.json" + with open(config_path, 'r') as f: + models = json.load(f) + + inconsistencies = [] + + for model_name, config in models.items(): + # Skip the sample_spec + if model_name == "sample_spec": + continue + + # Check if both max_tokens and max_output_tokens exist + if isinstance(config, dict): + max_tokens = config.get('max_tokens') + max_output_tokens = config.get('max_output_tokens') + + # Only validate if both exist + if max_tokens is not None and max_output_tokens is not None: + if max_tokens != max_output_tokens: + inconsistencies.append({ + 'model': model_name, + 'max_tokens': max_tokens, + 'max_output_tokens': max_output_tokens + }) + + if inconsistencies: + error_msg = f"\n\n❌ Found {len(inconsistencies)} models with max_tokens != max_output_tokens:\n\n" + for item in inconsistencies[:10]: # Show first 10 + error_msg += f" {item['model']}: max_tokens={item['max_tokens']}, max_output_tokens={item['max_output_tokens']}\n" + + if len(inconsistencies) > 10: + error_msg += f"\n ... and {len(inconsistencies) - 10} more\n" + + error_msg += "\nTo fix these inconsistencies, run: poetry run python fix_max_tokens_inconsistencies.py" + raise AssertionError(error_msg) + + def test_get_model_info_gemini(): """ Tests if ALL gemini models have 'tpm' and 'rpm' in the model info diff --git a/tests/test_litellm/test_utils_custom.py b/tests/test_litellm/test_utils_custom.py new file mode 100644 index 00000000000..3e924e9c719 --- /dev/null +++ b/tests/test_litellm/test_utils_custom.py @@ -0,0 +1,45 @@ +import pytest +import sys +from unittest.mock import MagicMock, patch, AsyncMock +from litellm.proxy.utils import count_tokens_with_anthropic_api, _anthropic_async_clients + +@pytest.mark.asyncio +async def test_count_tokens_caching(): + """ + Test that count_tokens_with_anthropic_api caches the client. + """ + # Clear cache + _anthropic_async_clients.clear() + + api_key = "sk-ant-test-key" + messages = [{"role": "user", "content": "hello"}] + model = "claude-3-opus-20240229" + + # Create a mock anthropic module + mock_anthropic = MagicMock() + mock_client = MagicMock() + mock_anthropic.AsyncAnthropic.return_value = mock_client + + # Mock response + mock_response = MagicMock() + mock_response.input_tokens = 10 + + # Setup async return for count_tokens + mock_client.beta.messages.count_tokens = AsyncMock(return_value=mock_response) + + # Patch sys.modules to ensure our mock is used when anthropic is imported + with patch.dict(sys.modules, {"anthropic": mock_anthropic}): + # First call + with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}): + await count_tokens_with_anthropic_api(model, messages) + + assert api_key in _anthropic_async_clients + assert _anthropic_async_clients[api_key] == mock_client + mock_anthropic.AsyncAnthropic.assert_called_once() # Should be called once + + # Second call + with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}): + await count_tokens_with_anthropic_api(model, messages) + + # Should still be called once (cached) + mock_anthropic.AsyncAnthropic.assert_called_once() diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts b/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts new file mode 100644 index 00000000000..4a4bb64c8ed --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts @@ -0,0 +1,38 @@ +import { Page } from "./pages"; + +/** + * Maps sidebar menu item labels to their corresponding page enum values. + * This mapping is for the admin role. + */ +export const menuLabelToPage: Record = { + "Virtual Keys": Page.ApiKeys, + Playground: Page.LlmPlayground, + Models: Page.Models, + "Models + Endpoints": Page.Models, + Usage: Page.NewUsage, + Teams: Page.Teams, + "Internal Users": Page.Users, + "Internal User": Page.Users, // Legacy label support + Organizations: Page.Organizations, + "API Reference": Page.ApiRef, + "AI Hub": Page.ModelHubTable, + "Model Hub": Page.ModelHubTable, + Logs: Page.Logs, + Guardrails: Page.Guardrails, + // Settings submenu items + "Router Settings": Page.RouterSettings, + "Logging & Alerts": Page.LoggingAndAlerts, + "Admin Settings": Page.AdminPanel, + "Cost Tracking": Page.CostTracking, + "UI Theme": Page.UiTheme, + // Experimental submenu items + Caching: Page.Caching, + Prompts: Page.Prompts, + Budgets: Page.Budgets, + "API Playground": Page.TransformRequest, + "Tag Management": Page.TagManagement, + "Old Usage": Page.Usage, + // Tools submenu items + "MCP Servers": Page.McpServers, + "Vector Stores": Page.VectorStores, +}; diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts new file mode 100644 index 00000000000..3ea37718ab5 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts @@ -0,0 +1,33 @@ +/** + * Enum for all page query parameters supported in the app. + * These values correspond to the `page` query parameter used in the URL. + */ +export enum Page { + ApiKeys = "api-keys", + Models = "models", + LlmPlayground = "llm-playground", + Users = "users", + Teams = "teams", + Organizations = "organizations", + AdminPanel = "admin-panel", + ApiRef = "api_ref", + LoggingAndAlerts = "logging-and-alerts", + Budgets = "budgets", + Guardrails = "guardrails", + Agents = "agents", + Prompts = "prompts", + TransformRequest = "transform-request", + RouterSettings = "router-settings", + UiTheme = "ui-theme", + CostTracking = "cost-tracking", + ModelHubTable = "model-hub-table", + Caching = "caching", + PassThroughSettings = "pass-through-settings", + Logs = "logs", + McpServers = "mcp-servers", + SearchTools = "search-tools", + TagManagement = "tag-management", + VectorStores = "vector-stores", + NewUsage = "new_usage", + Usage = "usage", +} diff --git a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts new file mode 100644 index 00000000000..919e516b35b --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts @@ -0,0 +1,12 @@ +import { Page } from "../fixtures/pages"; +import { Page as PlaywrightPage } from "@playwright/test"; + +/** + * Navigates to a specific page using the page query parameter. + * Uses relative path which will be resolved against the baseURL configured in playwright.config.ts + * @param page - The Playwright page object + * @param pageEnum - The page enum value to navigate to + */ +export async function navigateToPage(page: PlaywrightPage, pageEnum: Page): Promise { + await page.goto(`/ui?page=${pageEnum}`); +} diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts index c0619cfa845..2ab782d5678 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts @@ -14,7 +14,7 @@ test.describe("Add Model", () => { await providerInputDropdown.fill("Anthropic"); await page.waitForTimeout(1000); await providerInputDropdown.press("Enter"); - await page.waitForTimeout(1000); + await page.waitForTimeout(2000); const providerModelsDropdown = page.locator(".ant-select-selection-overflow").first(); await providerModelsDropdown.click(); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts index c90be698ae1..ce07cc2b83d 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -1,6 +1,9 @@ import test, { expect } from "@playwright/test"; import { Role } from "../../fixtures/roles"; import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { menuLabelToPage } from "../../fixtures/menuMappings"; +import { navigateToPage } from "../../helpers/navigation"; const sidebarButtons = { [Role.ProxyAdmin]: [ @@ -9,9 +12,7 @@ const sidebarButtons = { "Models", "Usage", "Teams", - "Internal User", - "Settings", - "Experimental", + "Internal Users", "API Reference", "AI Hub", ], @@ -23,13 +24,36 @@ for (const { role, storage } of roles) { test.describe(`${role} sidebar`, () => { test.use({ storageState: storage }); - test("can see and navigate all sidebar buttons", async ({ page }) => { + test("should navigate to correct URL when clicking sidebar menu items from homepage", async ({ page }) => { await page.goto("/ui"); - for (const button of sidebarButtons[role as keyof typeof sidebarButtons]) { - const tab = page.getByRole("menuitem", { name: button }); + + for (const buttonLabel of sidebarButtons[role as keyof typeof sidebarButtons]) { + const expectedPage = menuLabelToPage[buttonLabel]; + + if (!expectedPage) { + throw new Error(`No page mapping found for menu label: ${buttonLabel}`); + } + + const tab = page.getByRole("menuitem", { name: buttonLabel }); await expect(tab).toBeVisible(); + await tab.click(); + + // Verify URL contains the correct page query parameter + await expect(page).toHaveURL(new RegExp(`[?&]page=${expectedPage}(&|$)`)); } }); + + test("should navigate directly to page using navigation helper", async ({ page }) => { + // Test direct navigation to verify the helper function works + await navigateToPage(page, Page.ApiKeys); + await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.ApiKeys}(&|$)`)); + + await navigateToPage(page, Page.Models); + await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.Models}(&|$)`)); + + await navigateToPage(page, Page.LlmPlayground); + await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.LlmPlayground}(&|$)`)); + }); }); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.test.ts new file mode 100644 index 00000000000..fe4680cedeb --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.test.ts @@ -0,0 +1,387 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { RouterFieldsResponse, useRouterFields } from "./useRouterFields"; + +// Mock the networking module +vi.mock("@/components/networking", () => ({ + proxyBaseUrl: null, +})); + +// Mock useAuthorized hook +const mockUseAuthorized = vi.fn(); +vi.mock("../useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock global fetch +const mockFetch = vi.fn(); +global.fetch = mockFetch; + +// Mock console methods to avoid noise in tests +vi.spyOn(console, "log").mockImplementation(() => {}); +vi.spyOn(console, "error").mockImplementation(() => {}); + +// Mock data +const mockRouterFieldsResponse: RouterFieldsResponse = { + fields: [ + { + field_name: "routing_strategy", + field_type: "String", + field_description: "Routing strategy to use for load balancing across deployments", + field_default: "simple-shuffle", + options: ["simple-shuffle", "least-busy", "latency-based-routing"], + ui_field_name: "Routing Strategy", + link: null, + }, + { + field_name: "num_retries", + field_type: "Integer", + field_description: "Number of retries for failed requests", + field_default: 0, + options: null, + ui_field_name: "Number of Retries", + link: null, + }, + ], + routing_strategy_descriptions: { + "simple-shuffle": "Randomly picks a deployment from the list. Simple and fast.", + "least-busy": "Routes to the deployment with the lowest number of ongoing requests.", + "latency-based-routing": "Routes to the deployment with the lowest latency over a sliding window.", + }, +}; + +describe("useRouterFields", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should render", () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockRouterFieldsResponse, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + expect(result.current).toBeDefined(); + }); + + it("should return router fields data when query is successful", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockRouterFieldsResponse, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockRouterFieldsResponse); + expect(result.current.error).toBeNull(); + expect(mockFetch).toHaveBeenCalledTimes(1); + expect(mockFetch).toHaveBeenCalledWith("/router/fields", { + method: "GET", + headers: { + Authorization: "Bearer test-access-token", + "Content-Type": "application/json", + }, + }); + }); + + it("should handle error when fetch fails", async () => { + const errorMessage = "Failed to fetch router fields"; + const errorResponse = { error: errorMessage }; + + mockFetch.mockResolvedValueOnce({ + ok: false, + json: async () => errorResponse, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toBeDefined(); + expect(result.current.data).toBeUndefined(); + expect(mockFetch).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(mockFetch).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(mockFetch).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", () => { + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: null, + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(mockFetch).not.toHaveBeenCalled(); + }); + + it("should handle network error", async () => { + const networkError = new Error("Network error"); + mockFetch.mockRejectedValueOnce(networkError); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toBeDefined(); + expect(result.current.data).toBeUndefined(); + }); + + it("should use relative URL when proxyBaseUrl is null", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockRouterFieldsResponse, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + // When proxyBaseUrl is null, should use relative URL + expect(mockFetch).toHaveBeenCalledWith("/router/fields", { + method: "GET", + headers: { + Authorization: "Bearer test-access-token", + "Content-Type": "application/json", + }, + }); + }); + + it("should handle error response with different error formats", async () => { + const errorFormats = [ + { error: { message: "Error message" } }, + { message: "Error message" }, + { detail: "Error detail" }, + { error: "Error string" }, + { unknown: "format" }, + ]; + + for (const errorFormat of errorFormats) { + vi.clearAllMocks(); + mockFetch.mockResolvedValueOnce({ + ok: false, + json: async () => errorFormat, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toBeDefined(); + } + }); + + it("should return empty fields array when API returns empty fields", async () => { + const emptyResponse: RouterFieldsResponse = { + fields: [], + routing_strategy_descriptions: {}, + }; + + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => emptyResponse, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data?.fields).toEqual([]); + expect(result.current.data?.routing_strategy_descriptions).toEqual({}); + }); + + it("should have correct query configuration", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockRouterFieldsResponse, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + // Verify the query was called + expect(mockFetch).toHaveBeenCalledTimes(1); + + // The hook should have the expected properties from useQuery + expect(result.current).toHaveProperty("data"); + expect(result.current).toHaveProperty("isLoading"); + expect(result.current).toHaveProperty("isError"); + expect(result.current).toHaveProperty("isSuccess"); + expect(result.current).toHaveProperty("error"); + }); + + it("should handle fields with null options", async () => { + const responseWithNullOptions: RouterFieldsResponse = { + fields: [ + { + field_name: "timeout", + field_type: "Float", + field_description: "Timeout for requests in seconds", + field_default: null, + options: null, + ui_field_name: "Timeout", + link: null, + }, + ], + routing_strategy_descriptions: {}, + }; + + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => responseWithNullOptions, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data?.fields[0].options).toBeNull(); + }); + + it("should handle fields with link property", async () => { + const responseWithLink: RouterFieldsResponse = { + fields: [ + { + field_name: "enable_tag_filtering", + field_type: "Boolean", + field_description: "Enable tag-based routing", + field_default: false, + options: null, + ui_field_name: "Enable Tag Filtering", + link: "https://docs.litellm.ai/docs/proxy/tag_routing", + }, + ], + routing_strategy_descriptions: {}, + }; + + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => responseWithLink, + }); + + const { result } = renderHook(() => useRouterFields(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data?.fields[0].link).toBe("https://docs.litellm.ai/docs/proxy/tag_routing"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.ts new file mode 100644 index 00000000000..589508c5ddb --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.ts @@ -0,0 +1,69 @@ +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { proxyBaseUrl } from "@/components/networking"; + +export interface RouterSettingsField { + field_name: string; + field_type: string; + field_description: string; + field_default: any; + options: string[] | null; + ui_field_name: string; + link: string | null; +} + +export interface RouterFieldsResponse { + fields: RouterSettingsField[]; + routing_strategy_descriptions: Record; +} + +const routerFieldsKeys = createQueryKeys("routerFields"); + +const deriveErrorMessage = (errorData: any): string => { + return ( + (errorData?.error && (errorData.error.message || errorData.error)) || + errorData?.message || + errorData?.detail || + errorData?.error || + JSON.stringify(errorData) + ); +}; + +const getRouterFields = async (accessToken: string): Promise => { + try { + const url = proxyBaseUrl ? `${proxyBaseUrl}/router/fields` : `/router/fields`; + + console.log("Fetching router fields from:", url); + + const response = await fetch(url, { + method: "GET", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + throw new Error(errorMessage); + } + + const data: RouterFieldsResponse = await response.json(); + console.log("Fetched router fields:", data); + return data; + } catch (error) { + console.error("Failed to fetch router fields:", error); + throw error; + } +}; + +export const useRouterFields = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: routerFieldsKeys.detail("fields"), + queryFn: async () => await getRouterFields(accessToken!), + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx new file mode 100644 index 00000000000..ecc4914b0ab --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx @@ -0,0 +1,187 @@ +import { describe, expect, it, vi, beforeEach } from "vitest"; +import { + convertImageToBase64, + createChatMultimodalMessage, + createChatDisplayMessage, + shouldShowChatAttachedImage, +} from "./ChatImageUtils"; +import { MessageType } from "./types"; + +describe("ChatImageUtils", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe("convertImageToBase64", () => { + it("should convert file to base64 data URI", async () => { + const file = new File(["test content"], "test.png", { type: "image/png" }); + const result = await convertImageToBase64(file); + expect(result).toMatch(/^data:image\/png;base64,/); + }); + + it("should handle different file types", async () => { + const jpegFile = new File(["jpeg content"], "test.jpg", { type: "image/jpeg" }); + const result = await convertImageToBase64(jpegFile); + expect(result).toMatch(/^data:image\/jpeg;base64,/); + }); + + it("should reject on file read error", async () => { + const file = new File(["test"], "test.png", { type: "image/png" }); + const originalReadAsDataURL = FileReader.prototype.readAsDataURL; + + FileReader.prototype.readAsDataURL = vi.fn(function (this: FileReader) { + setTimeout(() => { + if (this.onerror) { + this.onerror(new Error("Read error") as any); + } + }, 0); + }); + + await expect(convertImageToBase64(file)).rejects.toThrow(); + + FileReader.prototype.readAsDataURL = originalReadAsDataURL; + }); + }); + + describe("createChatMultimodalMessage", () => { + it("should create multimodal message with text and image", async () => { + const file = new File(["test content"], "test.png", { type: "image/png" }); + const inputMessage = "What is in this image?"; + + const result = await createChatMultimodalMessage(inputMessage, file); + + expect(result.role).toBe("user"); + expect(result.content).toHaveLength(2); + expect(result.content[0]).toEqual({ type: "text", text: inputMessage }); + expect(result.content[1]).toMatchObject({ + type: "image_url", + image_url: { + url: expect.stringMatching(/^data:image\/png;base64,/), + }, + }); + }); + + it("should include base64 data URI in image_url", async () => { + const file = new File(["test content"], "test.png", { type: "image/png" }); + const result = await createChatMultimodalMessage("test", file); + + const imageContent = result.content[1]; + expect(imageContent.type).toBe("image_url"); + if ("image_url" in imageContent && imageContent.image_url) { + expect(imageContent.image_url.url).toMatch(/^data:/); + } + }); + }); + + describe("createChatDisplayMessage", () => { + it("should create display message without file", () => { + const result = createChatDisplayMessage("Hello world", false); + + expect(result.role).toBe("user"); + expect(result.content).toBe("Hello world"); + expect(result.imagePreviewUrl).toBeUndefined(); + }); + + it("should create display message with PDF file", () => { + const filePreviewUrl = "blob:test-url"; + const result = createChatDisplayMessage("Read this", true, filePreviewUrl, "document.pdf"); + + expect(result.content).toBe("Read this [PDF attached]"); + expect(result.imagePreviewUrl).toBe(filePreviewUrl); + }); + + it("should create display message with image file", () => { + const filePreviewUrl = "blob:test-url"; + const result = createChatDisplayMessage("Look at this", true, filePreviewUrl, "photo.jpg"); + + expect(result.content).toBe("Look at this [Image attached]"); + expect(result.imagePreviewUrl).toBe(filePreviewUrl); + }); + + it("should create display message with file but no fileName", () => { + const filePreviewUrl = "blob:test-url"; + const result = createChatDisplayMessage("Check this", true, filePreviewUrl); + + expect(result.content).toBe("Check this "); + expect(result.imagePreviewUrl).toBe(filePreviewUrl); + }); + + it("should create display message with file but no preview URL", () => { + const result = createChatDisplayMessage("See this", true, undefined, "image.png"); + + expect(result.content).toBe("See this [Image attached]"); + expect(result.imagePreviewUrl).toBeUndefined(); + }); + }); + + describe("shouldShowChatAttachedImage", () => { + it("should return true for user message with image attachment", () => { + const message: MessageType = { + role: "user", + content: "Check this [Image attached]", + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(true); + }); + + it("should return true for user message with PDF attachment", () => { + const message: MessageType = { + role: "user", + content: "Read this [PDF attached]", + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(true); + }); + + it("should return false for assistant message", () => { + const message: MessageType = { + role: "assistant", + content: "Here is the image [Image attached]", + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + + it("should return false when content is not a string", () => { + const message: MessageType = { + role: "user", + content: [{ type: "input_text", text: "test" }], + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + + it("should return false when content does not include attachment marker", () => { + const message: MessageType = { + role: "user", + content: "Just regular text", + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + + it("should return false when imagePreviewUrl is missing", () => { + const message: MessageType = { + role: "user", + content: "Check this [Image attached]", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + + it("should return false when imagePreviewUrl is empty string", () => { + const message: MessageType = { + role: "user", + content: "Check this [Image attached]", + imagePreviewUrl: "", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx new file mode 100644 index 00000000000..d87a74f4641 --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx @@ -0,0 +1,326 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; +import CodeInterpreterOutput from "./CodeInterpreterOutput"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => "https://example.com"), +})); + +global.fetch = vi.fn(); + +describe("CodeInterpreterOutput", () => { + beforeEach(() => { + vi.clearAllMocks(); + URL.createObjectURL = vi.fn((blob) => `blob:${blob}`); + URL.revokeObjectURL = vi.fn(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("should render", () => { + render(); + + expect(screen.getByText("Python Code Executed")).toBeInTheDocument(); + }); + + it("should display code in syntax highlighter", async () => { + const user = userEvent.setup(); + const code = "print('hello world')"; + const { container } = render(); + + expect(screen.getByText("Python Code Executed")).toBeInTheDocument(); + + const collapseHeader = screen.getByRole("button"); + await user.click(collapseHeader); + + await waitFor(() => { + const codeElement = container.querySelector("code.language-python"); + expect(codeElement).toBeInTheDocument(); + expect(codeElement?.textContent).toContain(code); + }); + }); + + it("should fetch and display images from annotations", async () => { + const mockBlob = new Blob(["image data"], { type: "image/png" }); + const mockResponse = { + ok: true, + blob: vi.fn().mockResolvedValue(mockBlob), + }; + + (global.fetch as any).mockResolvedValue(mockResponse); + + const annotations = [ + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-1", + filename: "chart.png", + start_index: 0, + end_index: 10, + }, + ]; + + render( + , + ); + + await waitFor(() => { + expect(global.fetch).toHaveBeenCalledWith( + "https://example.com/v1/containers/container-1/files/file-1/content", + expect.objectContaining({ + headers: { + Authorization: "Bearer test-token", + }, + }), + ); + }); + + await waitFor(() => { + expect(screen.getByText("chart.png")).toBeInTheDocument(); + }); + }); + + it("should show loading state while fetching images", async () => { + const mockBlob = new Blob(["image data"], { type: "image/png" }); + let resolveBlob: (value: Blob) => void; + const blobPromise = new Promise((resolve) => { + resolveBlob = resolve; + }); + + const mockResponse = { + ok: true, + blob: vi.fn().mockReturnValue(blobPromise), + }; + + (global.fetch as any).mockResolvedValue(mockResponse); + + const annotations = [ + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-1", + filename: "chart.png", + start_index: 0, + end_index: 10, + }, + ]; + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Loading image...")).toBeInTheDocument(); + }); + + resolveBlob!(mockBlob); + await waitFor(() => { + expect(screen.queryByText("Loading image...")).not.toBeInTheDocument(); + }); + }); + + it("should handle download for image files", async () => { + const user = userEvent.setup(); + const mockBlob = new Blob(["image data"], { type: "image/png" }); + const mockResponse = { + ok: true, + blob: vi.fn().mockResolvedValue(mockBlob), + }; + + (global.fetch as any).mockResolvedValue(mockResponse); + + const annotations = [ + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-1", + filename: "chart.png", + start_index: 0, + end_index: 10, + }, + ]; + + const createElementSpy = vi.spyOn(document, "createElement"); + const appendChildSpy = vi.spyOn(document.body, "appendChild"); + const removeChildSpy = vi.spyOn(document.body, "removeChild"); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("chart.png")).toBeInTheDocument(); + }); + + const downloadButton = screen.getByText("Download"); + await user.click(downloadButton); + + await waitFor(() => { + expect(global.fetch).toHaveBeenCalledWith( + "https://example.com/v1/containers/container-1/files/file-1/content", + expect.objectContaining({ + headers: { + Authorization: "Bearer test-token", + }, + }), + ); + }); + + createElementSpy.mockRestore(); + appendChildSpy.mockRestore(); + removeChildSpy.mockRestore(); + }); + + it("should handle download for non-image files", async () => { + const user = userEvent.setup(); + const mockBlob = new Blob(["file data"], { type: "text/plain" }); + const mockResponse = { + ok: true, + blob: vi.fn().mockResolvedValue(mockBlob), + }; + + (global.fetch as any).mockResolvedValue(mockResponse); + + const annotations = [ + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-1", + filename: "data.csv", + start_index: 0, + end_index: 10, + }, + ]; + + render(); + + await waitFor(() => { + expect(screen.getByText("data.csv")).toBeInTheDocument(); + }); + + const downloadButton = screen.getByText("data.csv").closest("button"); + expect(downloadButton).toBeInTheDocument(); + if (downloadButton) { + await user.click(downloadButton); + } + + await waitFor(() => { + expect(global.fetch).toHaveBeenCalledWith( + "https://example.com/v1/containers/container-1/files/file-1/content", + expect.objectContaining({ + headers: { + Authorization: "Bearer test-token", + }, + }), + ); + }); + }); + + it("should return null when no code and no annotations", () => { + const { container } = render(); + + expect(container.firstChild).toBeNull(); + }); + + it("should handle multiple image formats", async () => { + const mockBlob = new Blob(["image data"], { type: "image/png" }); + const mockResponse = { + ok: true, + blob: vi.fn().mockResolvedValue(mockBlob), + }; + + (global.fetch as any).mockResolvedValue(mockResponse); + + const annotations = [ + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-1", + filename: "image.png", + start_index: 0, + end_index: 10, + }, + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-2", + filename: "image.jpg", + start_index: 0, + end_index: 10, + }, + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-3", + filename: "image.jpeg", + start_index: 0, + end_index: 10, + }, + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-4", + filename: "image.gif", + start_index: 0, + end_index: 10, + }, + ]; + + render( + , + ); + + await waitFor(() => { + expect(global.fetch).toHaveBeenCalledTimes(4); + }); + }); + + it("should handle fetch errors gracefully", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + (global.fetch as any).mockRejectedValue(new Error("Network error")); + + const annotations = [ + { + type: "container_file_citation" as const, + container_id: "container-1", + file_id: "file-1", + filename: "chart.png", + start_index: 0, + end_index: 10, + }, + ]; + + render( + , + ); + + await waitFor(() => { + expect(consoleErrorSpy).toHaveBeenCalled(); + }); + + consoleErrorSpy.mockRestore(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/UnifiedSelector.test.tsx b/ui/litellm-dashboard/src/components/playground/compareUI/components/UnifiedSelector.test.tsx new file mode 100644 index 00000000000..1c6951fc421 --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/compareUI/components/UnifiedSelector.test.tsx @@ -0,0 +1,161 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import { UnifiedSelector } from "./UnifiedSelector"; +import { EndpointId, ENDPOINT_CONFIGS } from "../endpoint_config"; + +describe("UnifiedSelector", () => { + it("should render", () => { + const onChange = vi.fn(); + const options = [ + { value: "option1", label: "Option 1" }, + { value: "option2", label: "Option 2" }, + ]; + const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]; + + render(); + + const select = screen.getByRole("combobox"); + expect(select).toBeInTheDocument(); + }); + + it("should display placeholder when not loading", () => { + const onChange = vi.fn(); + const options = [{ value: "option1", label: "Option 1" }]; + const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]; + + const { container } = render( + , + ); + + const placeholder = container.querySelector(".ant-select-selection-placeholder"); + expect(placeholder).toHaveTextContent(config.selectorPlaceholder); + }); + + it("should display loading placeholder when loading", () => { + const onChange = vi.fn(); + const options = [{ value: "option1", label: "Option 1" }]; + const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]; + + const { container } = render( + , + ); + + const placeholder = container.querySelector(".ant-select-selection-placeholder"); + expect(placeholder).toHaveTextContent(`Loading ${config.selectorLabel.toLowerCase()}s...`); + }); + + it("should call onChange when option is selected", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const options = [ + { value: "option1", label: "Option 1" }, + { value: "option2", label: "Option 2" }, + ]; + const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]; + + render(); + + const select = screen.getByRole("combobox"); + await user.click(select); + + await waitFor(() => { + const option = screen.getByText("Option 1"); + expect(option).toBeInTheDocument(); + }); + + const option = screen.getByText("Option 1"); + await user.click(option); + + await waitFor(() => { + expect(onChange).toHaveBeenCalled(); + }); + const callArgs = onChange.mock.calls[0]; + expect(callArgs[0]).toBe("option1"); + }); + + it("should display selected value", () => { + const onChange = vi.fn(); + const options = [ + { value: "option1", label: "Option 1" }, + { value: "option2", label: "Option 2" }, + ]; + const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]; + + const { container } = render( + , + ); + + const selectedValue = container.querySelector(".ant-select-selection-item"); + expect(selectedValue).toHaveTextContent("Option 1"); + }); + + it("should filter options by search input", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const options = [ + { value: "option1", label: "Option One" }, + { value: "option2", label: "Option Two" }, + { value: "option3", label: "Different" }, + ]; + const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]; + + render(); + + const select = screen.getByRole("combobox"); + await user.click(select); + await user.type(select, "One"); + + await waitFor(() => { + expect(screen.getByText("Option One")).toBeInTheDocument(); + expect(screen.queryByText("Option Two")).not.toBeInTheDocument(); + expect(screen.queryByText("Different")).not.toBeInTheDocument(); + }); + }); + + it("should show loading spinner in notFoundContent when loading", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const options: { value: string; label: string }[] = []; + const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]; + + render(); + + const select = screen.getByRole("combobox"); + await user.click(select); + + await waitFor(() => { + const spin = document.querySelector(".ant-spin"); + expect(spin).toBeInTheDocument(); + }); + }); + + it("should show no options message when not loading and no options", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const options: { value: string; label: string }[] = []; + const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]; + + render(); + + const select = screen.getByRole("combobox"); + await user.click(select); + + await waitFor(() => { + expect(screen.getByText(`No ${config.selectorLabel.toLowerCase()}s available`)).toBeInTheDocument(); + }); + }); + + it("should work with agent endpoint config", () => { + const onChange = vi.fn(); + const options = [{ value: "agent1", label: "Agent One" }]; + const config = ENDPOINT_CONFIGS[EndpointId.A2A_AGENTS]; + + const { container } = render( + , + ); + + const placeholder = container.querySelector(".ant-select-selection-placeholder"); + expect(placeholder).toHaveTextContent(config.selectorPlaceholder); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.test.ts b/ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.test.ts new file mode 100644 index 00000000000..67ecf32fc2a --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.test.ts @@ -0,0 +1,104 @@ +import { describe, expect, it } from "vitest"; +import { + EndpointId, + ENDPOINT_CONFIGS, + getAvailableEndpoints, + getEndpointConfig, + isAgentEndpoint, + isModelEndpoint, + modelOptionsToSelectorOptions, + agentOptionsToSelectorOptions, + getSelectionFieldName, + getComparisonSelection, + hasValidSelection, +} from "./endpoint_config"; +import { Agent } from "../llm_calls/fetch_agents"; + +describe("endpoint_config", () => { + it("should export EndpointId constants", () => { + expect(EndpointId.CHAT_COMPLETIONS).toBe("/v1/chat/completions"); + expect(EndpointId.A2A_AGENTS).toBe("/a2a"); + }); + + it("should have endpoint configs for all endpoint IDs", () => { + expect(ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS]).toBeDefined(); + expect(ENDPOINT_CONFIGS[EndpointId.A2A_AGENTS]).toBeDefined(); + expect(ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS].selectorType).toBe("model"); + expect(ENDPOINT_CONFIGS[EndpointId.A2A_AGENTS].selectorType).toBe("agent"); + }); + + it("should get available endpoints", () => { + const endpoints = getAvailableEndpoints(); + expect(endpoints).toHaveLength(2); + expect(endpoints).toContainEqual({ + value: EndpointId.CHAT_COMPLETIONS, + label: "/v1/chat/completions", + }); + expect(endpoints).toContainEqual({ + value: EndpointId.A2A_AGENTS, + label: "/a2a (Agents)", + }); + }); + + it("should get endpoint config by ID", () => { + const config = getEndpointConfig(EndpointId.CHAT_COMPLETIONS); + expect(config.id).toBe(EndpointId.CHAT_COMPLETIONS); + expect(config.selectorType).toBe("model"); + expect(config.selectorLabel).toBe("Model"); + }); + + it("should check if endpoint is agent endpoint", () => { + expect(isAgentEndpoint(EndpointId.A2A_AGENTS)).toBe(true); + expect(isAgentEndpoint(EndpointId.CHAT_COMPLETIONS)).toBe(false); + }); + + it("should check if endpoint is model endpoint", () => { + expect(isModelEndpoint(EndpointId.CHAT_COMPLETIONS)).toBe(true); + expect(isModelEndpoint(EndpointId.A2A_AGENTS)).toBe(false); + }); + + it("should convert model options to selector options", () => { + const models = ["gpt-4", "gpt-3.5-turbo", "claude-3"]; + const options = modelOptionsToSelectorOptions(models); + expect(options).toHaveLength(3); + expect(options[0]).toEqual({ value: "gpt-4", label: "gpt-4" }); + expect(options[1]).toEqual({ value: "gpt-3.5-turbo", label: "gpt-3.5-turbo" }); + expect(options[2]).toEqual({ value: "claude-3", label: "claude-3" }); + }); + + it("should convert agent options to selector options", () => { + const agents: Agent[] = [ + { agent_id: "agent-1", agent_name: "Agent One" }, + { agent_id: "agent-2", agent_name: "Agent Two" }, + { agent_id: "agent-3", agent_name: undefined as any }, + ]; + const options = agentOptionsToSelectorOptions(agents); + expect(options).toHaveLength(3); + expect(options[0]).toEqual({ value: "Agent One", label: "Agent One" }); + expect(options[1]).toEqual({ value: "Agent Two", label: "Agent Two" }); + expect(options[2]).toEqual({ value: undefined, label: "agent-3" }); + }); + + it("should get selection field name based on endpoint", () => { + expect(getSelectionFieldName(EndpointId.CHAT_COMPLETIONS)).toBe("model"); + expect(getSelectionFieldName(EndpointId.A2A_AGENTS)).toBe("agent"); + }); + + it("should get comparison selection based on endpoint", () => { + const comparison = { model: "gpt-4", agent: "agent-1" }; + expect(getComparisonSelection(comparison, EndpointId.CHAT_COMPLETIONS)).toBe("gpt-4"); + expect(getComparisonSelection(comparison, EndpointId.A2A_AGENTS)).toBe("agent-1"); + }); + + it("should check if comparison has valid selection", () => { + const comparisonWithModel = { model: "gpt-4", agent: "" }; + const comparisonWithAgent = { model: "", agent: "agent-1" }; + const comparisonEmpty = { model: "", agent: "" }; + const comparisonWhitespace = { model: " ", agent: "" }; + + expect(hasValidSelection(comparisonWithModel, EndpointId.CHAT_COMPLETIONS)).toBe(true); + expect(hasValidSelection(comparisonWithAgent, EndpointId.A2A_AGENTS)).toBe(true); + expect(hasValidSelection(comparisonEmpty, EndpointId.CHAT_COMPLETIONS)).toBe(false); + expect(hasValidSelection(comparisonWhitespace, EndpointId.CHAT_COMPLETIONS)).toBe(false); + }); +});