Merge remote-tracking branch 'origin' into litellm_user_promethus_metrics

.github/workflows/publish-migrations.yml

@@ -13,6 +13,7 @@ on:
 
 jobs:
   publish-migrations:
+    if: github.repository == 'BerriAI/litellm'
     runs-on: ubuntu-latest
     services:
       postgres:

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.
   <img src="https://contrib.rocks/image?repo=BerriAI/litellm" />
 </a>
 
-

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.
+:::

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/)
+

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

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) |
+
+<br />
+
+## 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())
+```

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

docs/my-website/docs/proxy/config_settings.md

@@ -581,6 +581,18 @@ router_settings:
 | FIREWORKS_AI_56_B_MOE | Size parameter for Fireworks AI 56B MOE model. Default is 56
 | FIREWORKS_AI_80_B | Size parameter for Fireworks AI 80B model. Default is 80
 | FIREWORKS_AI_176_B_MOE | Size parameter for Fireworks AI 176B MOE model. Default is 176
+| FOCUS_PROVIDER | Destination provider for Focus exports (e.g., `s3`). Defaults to `s3`.
+| FOCUS_FORMAT | Output format for Focus exports. Defaults to `parquet`.
+| FOCUS_FREQUENCY | Frequency for scheduled Focus exports (`hourly`, `daily`, or `interval`). Defaults to `hourly`.
+| FOCUS_CRON_OFFSET | Minute offset used when scheduling hourly/daily Focus exports. Defaults to `5` minutes.
+| FOCUS_INTERVAL_SECONDS | Interval (in seconds) for Focus exports when `frequency` is `interval`.
+| FOCUS_PREFIX | Object key prefix (or folder) used when uploading Focus export files. Defaults to `focus_exports`.
+| FOCUS_S3_BUCKET_NAME | S3 bucket to upload Focus export files when using the S3 destination.
+| FOCUS_S3_REGION_NAME | AWS region for the Focus export S3 bucket.
+| FOCUS_S3_ENDPOINT_URL | Custom endpoint for the Focus export S3 client (optional; useful for S3-compatible storage).
+| FOCUS_S3_ACCESS_KEY | AWS access key ID used by the Focus export S3 client.
+| FOCUS_S3_SECRET_KEY | AWS secret access key used by the Focus export S3 client.
+| FOCUS_S3_SESSION_TOKEN | AWS session token used by the Focus export S3 client (optional).
 | FUNCTION_DEFINITION_TOKEN_COUNT | Token count for function definitions. Default is 9
 | GALILEO_BASE_URL | Base URL for Galileo platform
 | GALILEO_PASSWORD | Password for Galileo authentication

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)

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",

litellm/__init__.py

@@ -134,6 +134,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
     "bitbucket",
     "gitlab",
     "cloudzero",
+    "focus",
     "posthog",
     "levo",
 ]

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",

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)}")

litellm/integrations/focus/database.py

@@ -0,0 +1,113 @@
+"""Database access helpers for Focus export."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import polars as pl
+
+
+class FocusLiteLLMDatabase:
+    """Retrieves LiteLLM usage data for Focus export workflows."""
+
+    def _ensure_prisma_client(self):
+        from litellm.proxy.proxy_server import prisma_client
+
+        if prisma_client is None:
+            raise RuntimeError(
+                "Database not connected. Connect a database to your proxy - "
+                "https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
+            )
+        return prisma_client
+
+    async def get_usage_data(
+        self,
+        *,
+        limit: Optional[int] = None,
+        start_time_utc: Optional[datetime] = None,
+        end_time_utc: Optional[datetime] = None,
+    ) -> pl.DataFrame:
+        """Return usage data for the requested window."""
+        client = self._ensure_prisma_client()
+
+        where_clauses: list[str] = []
+        query_params: list[Any] = []
+        placeholder_index = 1
+        if start_time_utc:
+            where_clauses.append(f"dus.updated_at >= ${placeholder_index}::timestamptz")
+            query_params.append(start_time_utc)
+            placeholder_index += 1
+        if end_time_utc:
+            where_clauses.append(f"dus.updated_at <= ${placeholder_index}::timestamptz")
+            query_params.append(end_time_utc)
+            placeholder_index += 1
+
+        where_clause = ""
+        if where_clauses:
+            where_clause = "WHERE " + " AND ".join(where_clauses)
+
+        limit_clause = ""
+        if limit is not None:
+            try:
+                limit_value = int(limit)
+            except (TypeError, ValueError) as exc:  # pragma: no cover - defensive guard
+                raise ValueError("limit must be an integer") from exc
+            if limit_value < 0:
+                raise ValueError("limit must be non-negative")
+            limit_clause = f" LIMIT ${placeholder_index}"
+            query_params.append(limit_value)
+
+        query = f"""
+        SELECT
+            dus.id,
+            dus.date,
+            dus.user_id,
+            dus.api_key,
+            dus.model,
+            dus.model_group,
+            dus.custom_llm_provider,
+            dus.prompt_tokens,
+            dus.completion_tokens,
+            dus.spend,
+            dus.api_requests,
+            dus.successful_requests,
+            dus.failed_requests,
+            dus.cache_creation_input_tokens,
+            dus.cache_read_input_tokens,
+            dus.created_at,
+            dus.updated_at,
+            vt.team_id,
+            vt.key_alias as api_key_alias,
+            tt.team_alias,
+            ut.user_email as user_email
+        FROM "LiteLLM_DailyUserSpend" dus
+        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}
+        ORDER BY dus.date DESC, dus.created_at DESC
+        {limit_clause}
+        """
+
+        try:
+            db_response = await client.db.query_raw(query, *query_params)
+            return pl.DataFrame(db_response, infer_schema_length=None)
+        except Exception as exc:
+            raise RuntimeError(f"Error retrieving usage data: {exc}") from exc
+
+    async def get_table_info(self) -> Dict[str, Any]:
+        """Return metadata about the spend table for diagnostics."""
+        client = self._ensure_prisma_client()
+
+        info_query = """
+        SELECT column_name, data_type, is_nullable
+        FROM information_schema.columns
+        WHERE table_name = 'LiteLLM_DailyUserSpend'
+        ORDER BY ordinal_position;
+        """
+        try:
+            columns_response = await client.db.query_raw(info_query)
+            return {"columns": columns_response, "table_name": "LiteLLM_DailyUserSpend"}
+        except Exception as exc:
+            raise RuntimeError(f"Error getting table info: {exc}") from exc

litellm/integrations/focus/destinations/__init__.py

@@ -0,0 +1,12 @@
+"""Destination implementations for Focus export."""
+
+from .base import FocusDestination, FocusTimeWindow
+from .factory import FocusDestinationFactory
+from .s3_destination import FocusS3Destination
+
+__all__ = [
+    "FocusDestination",
+    "FocusDestinationFactory",
+    "FocusTimeWindow",
+    "FocusS3Destination",
+]

litellm/integrations/focus/destinations/base.py

@@ -0,0 +1,30 @@
+"""Abstract destination interfaces for Focus export."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Protocol
+
+
+@dataclass(frozen=True)
+class FocusTimeWindow:
+    """Represents the span of data exported in a single batch."""
+
+    start_time: datetime
+    end_time: datetime
+    frequency: str
+
+
+class FocusDestination(Protocol):
+    """Protocol for anything that can receive Focus export files."""
+
+    async def deliver(
+        self,
+        *,
+        content: bytes,
+        time_window: FocusTimeWindow,
+        filename: str,
+    ) -> None:
+        """Persist the serialized export for the provided time window."""
+        ...

litellm/integrations/focus/destinations/factory.py

@@ -0,0 +1,59 @@
+"""Factory helpers for Focus export destinations."""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, Optional
+
+from .base import FocusDestination
+from .s3_destination import FocusS3Destination
+
+
+class FocusDestinationFactory:
+    """Builds destination instances based on provider/config settings."""
+
+    @staticmethod
+    def create(
+        *,
+        provider: str,
+        prefix: str,
+        config: Optional[Dict[str, Any]] = None,
+    ) -> FocusDestination:
+        """Return a destination implementation for the requested provider."""
+        provider_lower = provider.lower()
+        normalized_config = FocusDestinationFactory._resolve_config(
+            provider=provider_lower, overrides=config or {}
+        )
+        if provider_lower == "s3":
+            return FocusS3Destination(prefix=prefix, config=normalized_config)
+        raise NotImplementedError(
+            f"Provider '{provider}' not supported for Focus export"
+        )
+
+    @staticmethod
+    def _resolve_config(
+        *,
+        provider: str,
+        overrides: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        if provider == "s3":
+            resolved = {
+                "bucket_name": overrides.get("bucket_name")
+                or os.getenv("FOCUS_S3_BUCKET_NAME"),
+                "region_name": overrides.get("region_name")
+                or os.getenv("FOCUS_S3_REGION_NAME"),
+                "endpoint_url": overrides.get("endpoint_url")
+                or os.getenv("FOCUS_S3_ENDPOINT_URL"),
+                "aws_access_key_id": overrides.get("aws_access_key_id")
+                or os.getenv("FOCUS_S3_ACCESS_KEY"),
+                "aws_secret_access_key": overrides.get("aws_secret_access_key")
+                or os.getenv("FOCUS_S3_SECRET_KEY"),
+                "aws_session_token": overrides.get("aws_session_token")
+                or os.getenv("FOCUS_S3_SESSION_TOKEN"),
+            }
+            if not resolved.get("bucket_name"):
+                raise ValueError("FOCUS_S3_BUCKET_NAME must be provided for S3 exports")
+            return {k: v for k, v in resolved.items() if v is not None}
+        raise NotImplementedError(
+            f"Provider '{provider}' not supported for Focus export configuration"
+        )

litellm/integrations/focus/destinations/s3_destination.py

@@ -0,0 +1,74 @@
+"""S3 destination implementation for Focus export."""
+
+from __future__ import annotations
+
+import asyncio
+from datetime import timezone
+from typing import Any, Optional
+
+import boto3
+
+from .base import FocusDestination, FocusTimeWindow
+
+
+class FocusS3Destination(FocusDestination):
+    """Handles uploading serialized exports to S3 buckets."""
+
+    def __init__(
+        self,
+        *,
+        prefix: str,
+        config: Optional[dict[str, Any]] = None,
+    ) -> None:
+        config = config or {}
+        bucket_name = config.get("bucket_name")
+        if not bucket_name:
+            raise ValueError("bucket_name must be provided for S3 destination")
+        self.bucket_name = bucket_name
+        self.prefix = prefix.rstrip("/")
+        self.config = config
+
+    async def deliver(
+        self,
+        *,
+        content: bytes,
+        time_window: FocusTimeWindow,
+        filename: str,
+    ) -> None:
+        object_key = self._build_object_key(time_window=time_window, filename=filename)
+        await asyncio.to_thread(self._upload, content, object_key)
+
+    def _build_object_key(self, *, time_window: FocusTimeWindow, filename: str) -> str:
+        start_utc = time_window.start_time.astimezone(timezone.utc)
+        date_component = f"date={start_utc.strftime('%Y-%m-%d')}"
+        parts = [self.prefix, date_component]
+        if time_window.frequency == "hourly":
+            parts.append(f"hour={start_utc.strftime('%H')}")
+        key_prefix = "/".join(filter(None, parts))
+        return f"{key_prefix}/{filename}" if key_prefix else filename
+
+    def _upload(self, content: bytes, object_key: str) -> None:
+        client_kwargs: dict[str, Any] = {}
+        region_name = self.config.get("region_name")
+        if region_name:
+            client_kwargs["region_name"] = region_name
+        endpoint_url = self.config.get("endpoint_url")
+        if endpoint_url:
+            client_kwargs["endpoint_url"] = endpoint_url
+
+        session_kwargs: dict[str, Any] = {}
+        for key in (
+            "aws_access_key_id",
+            "aws_secret_access_key",
+            "aws_session_token",
+        ):
+            if self.config.get(key):
+                session_kwargs[key] = self.config[key]
+
+        s3_client = boto3.client("s3", **client_kwargs, **session_kwargs)
+        s3_client.put_object(
+            Bucket=self.bucket_name,
+            Key=object_key,
+            Body=content,
+            ContentType="application/octet-stream",
+        )

litellm/integrations/focus/export_engine.py

@@ -0,0 +1,124 @@
+"""Core export engine for Focus integrations (heavy dependencies)."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import polars as pl
+
+from litellm._logging import verbose_logger
+
+from .database import FocusLiteLLMDatabase
+from .destinations import FocusDestinationFactory, FocusTimeWindow
+from .serializers import FocusParquetSerializer, FocusSerializer
+from .transformer import FocusTransformer
+
+
+class FocusExportEngine:
+    """Engine that fetches, normalizes, and uploads Focus exports."""
+
+    def __init__(
+        self,
+        *,
+        provider: str,
+        export_format: str,
+        prefix: str,
+        destination_config: Optional[dict[str, Any]] = None,
+    ) -> None:
+        self.provider = provider
+        self.export_format = export_format
+        self.prefix = prefix
+        self._destination = FocusDestinationFactory.create(
+            provider=self.provider,
+            prefix=self.prefix,
+            config=destination_config,
+        )
+        self._serializer = self._init_serializer()
+        self._transformer = FocusTransformer()
+        self._database = FocusLiteLLMDatabase()
+
+    def _init_serializer(self) -> FocusSerializer:
+        if self.export_format != "parquet":
+            raise NotImplementedError("Only parquet export supported currently")
+        return FocusParquetSerializer()
+
+    async def dry_run_export_usage_data(self, limit: Optional[int]) -> Dict[str, Any]:
+        data = await self._database.get_usage_data(limit=limit)
+        normalized = self._transformer.transform(data)
+
+        usage_sample = data.head(min(50, len(data))).to_dicts()
+        normalized_sample = normalized.head(min(50, len(normalized))).to_dicts()
+
+        summary = {
+            "total_records": len(normalized),
+            "total_spend": self._sum_column(normalized, "spend"),
+            "total_tokens": self._sum_column(normalized, "total_tokens"),
+            "unique_teams": self._count_unique(normalized, "team_id"),
+            "unique_models": self._count_unique(normalized, "model"),
+        }
+
+        return {
+            "usage_data": usage_sample,
+            "normalized_data": normalized_sample,
+            "summary": summary,
+        }
+
+    async def export_window(
+        self,
+        *,
+        window: FocusTimeWindow,
+        limit: Optional[int],
+    ) -> None:
+        data = await self._database.get_usage_data(
+            limit=limit,
+            start_time_utc=window.start_time,
+            end_time_utc=window.end_time,
+        )
+        if data.is_empty():
+            verbose_logger.debug("Focus export: no usage data for window %s", window)
+            return
+
+        normalized = self._transformer.transform(data)
+        if normalized.is_empty():
+            verbose_logger.debug(
+                "Focus export: normalized data empty for window %s", window
+            )
+            return
+
+        await self._serialize_and_upload(normalized, window)
+
+    async def _serialize_and_upload(
+        self, frame: pl.DataFrame, window: FocusTimeWindow
+    ) -> None:
+        payload = self._serializer.serialize(frame)
+        if not payload:
+            verbose_logger.debug("Focus export: serializer returned empty payload")
+            return
+        await self._destination.deliver(
+            content=payload,
+            time_window=window,
+            filename=self._build_filename(),
+        )
+
+    def _build_filename(self) -> str:
+        if not self._serializer.extension:
+            raise ValueError("Serializer must declare a file extension")
+        return f"usage.{self._serializer.extension}"
+
+    @staticmethod
+    def _sum_column(frame: pl.DataFrame, column: str) -> float:
+        if frame.is_empty() or column not in frame.columns:
+            return 0.0
+        value = frame.select(pl.col(column).sum().alias("sum")).row(0)[0]
+        if value is None:
+            return 0.0
+        return float(value)
+
+    @staticmethod
+    def _count_unique(frame: pl.DataFrame, column: str) -> int:
+        if frame.is_empty() or column not in frame.columns:
+            return 0
+        value = frame.select(pl.col(column).n_unique().alias("unique")).row(0)[0]
+        if value is None:
+            return 0
+        return int(value)

litellm/integrations/focus/focus_logger.py

@@ -0,0 +1,211 @@
+"""Focus export logger orchestrating DB pull/transform/upload."""
+
+from __future__ import annotations
+
+import os
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.integrations.custom_logger import CustomLogger
+
+from .destinations import FocusTimeWindow
+
+if TYPE_CHECKING:
+    from apscheduler.schedulers.asyncio import AsyncIOScheduler
+    from .export_engine import FocusExportEngine
+else:
+    AsyncIOScheduler = Any
+
+FOCUS_USAGE_DATA_JOB_NAME = "focus_export_usage_data"
+DEFAULT_DRY_RUN_LIMIT = 500
+
+
+class FocusLogger(CustomLogger):
+    """Coordinates Focus export jobs across transformer/serializer/destination layers."""
+
+    def __init__(
+        self,
+        *,
+        provider: Optional[str] = None,
+        export_format: Optional[str] = None,
+        frequency: Optional[str] = None,
+        cron_offset_minute: Optional[int] = None,
+        interval_seconds: Optional[int] = None,
+        prefix: Optional[str] = None,
+        destination_config: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.provider = (provider or os.getenv("FOCUS_PROVIDER") or "s3").lower()
+        self.export_format = (
+            export_format or os.getenv("FOCUS_FORMAT") or "parquet"
+        ).lower()
+        self.frequency = (frequency or os.getenv("FOCUS_FREQUENCY") or "hourly").lower()
+        self.cron_offset_minute = (
+            cron_offset_minute
+            if cron_offset_minute is not None
+            else int(os.getenv("FOCUS_CRON_OFFSET", "5"))
+        )
+        raw_interval = (
+            interval_seconds
+            if interval_seconds is not None
+            else os.getenv("FOCUS_INTERVAL_SECONDS")
+        )
+        self.interval_seconds = int(raw_interval) if raw_interval is not None else None
+        env_prefix = os.getenv("FOCUS_PREFIX")
+        self.prefix: str = (
+            prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports")
+        )
+
+        self._destination_config = destination_config
+        self._engine: Optional["FocusExportEngine"] = None
+
+    def _ensure_engine(self) -> "FocusExportEngine":
+        """Instantiate the heavy export engine lazily."""
+        if self._engine is None:
+            from .export_engine import FocusExportEngine
+
+            self._engine = FocusExportEngine(
+                provider=self.provider,
+                export_format=self.export_format,
+                prefix=self.prefix,
+                destination_config=self._destination_config,
+            )
+        return self._engine
+
+    async def export_usage_data(
+        self,
+        *,
+        limit: Optional[int] = None,
+        start_time_utc: Optional[datetime] = None,
+        end_time_utc: Optional[datetime] = None,
+    ) -> None:
+        """Public hook to trigger export immediately."""
+        if bool(start_time_utc) ^ bool(end_time_utc):
+            raise ValueError(
+                "start_time_utc and end_time_utc must be provided together"
+            )
+
+        if start_time_utc and end_time_utc:
+            window = FocusTimeWindow(
+                start_time=start_time_utc,
+                end_time=end_time_utc,
+                frequency=self.frequency,
+            )
+        else:
+            window = self._compute_time_window(datetime.now(timezone.utc))
+        await self._export_window(window=window, limit=limit)
+
+    async def dry_run_export_usage_data(
+        self, limit: Optional[int] = DEFAULT_DRY_RUN_LIMIT
+    ) -> dict[str, Any]:
+        """Return transformed data without uploading."""
+        engine = self._ensure_engine()
+        return await engine.dry_run_export_usage_data(limit=limit)
+
+    async def initialize_focus_export_job(self) -> None:
+        """Entry point for scheduler jobs to run export cycle with locking."""
+        from litellm.proxy.proxy_server import proxy_logging_obj
+
+        pod_lock_manager = None
+        if proxy_logging_obj is not None:
+            writer = getattr(proxy_logging_obj, "db_spend_update_writer", None)
+            if writer is not None:
+                pod_lock_manager = getattr(writer, "pod_lock_manager", None)
+
+        if pod_lock_manager and pod_lock_manager.redis_cache:
+            acquired = await pod_lock_manager.acquire_lock(
+                cronjob_id=FOCUS_USAGE_DATA_JOB_NAME
+            )
+            if not acquired:
+                verbose_logger.debug("Focus export: unable to acquire pod lock")
+                return
+            try:
+                await self._run_scheduled_export()
+            finally:
+                await pod_lock_manager.release_lock(
+                    cronjob_id=FOCUS_USAGE_DATA_JOB_NAME
+                )
+        else:
+            await self._run_scheduled_export()
+
+    @staticmethod
+    async def init_focus_export_background_job(
+        scheduler: AsyncIOScheduler,
+    ) -> None:
+        """Register the export cron/interval job with the provided scheduler."""
+
+        focus_loggers: List[
+            CustomLogger
+        ] = litellm.logging_callback_manager.get_custom_loggers_for_type(
+            callback_type=FocusLogger
+        )
+        if not focus_loggers:
+            verbose_logger.debug(
+                "No Focus export logger registered; skipping scheduler"
+            )
+            return
+
+        focus_logger = cast(FocusLogger, focus_loggers[0])
+        trigger_kwargs = focus_logger._build_scheduler_trigger()
+        scheduler.add_job(
+            focus_logger.initialize_focus_export_job,
+            **trigger_kwargs,
+        )
+
+    def _build_scheduler_trigger(self) -> Dict[str, Any]:
+        """Return scheduler configuration for the selected frequency."""
+        if self.frequency == "interval":
+            seconds = self.interval_seconds or 60
+            return {"trigger": "interval", "seconds": seconds}
+
+        if self.frequency == "hourly":
+            minute = max(0, min(59, self.cron_offset_minute))
+            return {"trigger": "cron", "minute": minute, "second": 0}
+
+        if self.frequency == "daily":
+            total_minutes = max(0, self.cron_offset_minute)
+            hour = min(23, total_minutes // 60)
+            minute = min(59, total_minutes % 60)
+            return {"trigger": "cron", "hour": hour, "minute": minute, "second": 0}
+
+        raise ValueError(f"Unsupported frequency: {self.frequency}")
+
+    async def _run_scheduled_export(self) -> None:
+        """Execute the scheduled export for the configured window."""
+        window = self._compute_time_window(datetime.now(timezone.utc))
+        await self._export_window(window=window, limit=None)
+
+    async def _export_window(
+        self,
+        *,
+        window: FocusTimeWindow,
+        limit: Optional[int],
+    ) -> None:
+        engine = self._ensure_engine()
+        await engine.export_window(window=window, limit=limit)
+
+    def _compute_time_window(self, now: datetime) -> FocusTimeWindow:
+        """Derive the time window to export based on configured frequency."""
+        now_utc = now.astimezone(timezone.utc)
+        if self.frequency == "hourly":
+            end_time = now_utc.replace(minute=0, second=0, microsecond=0)
+            start_time = end_time - timedelta(hours=1)
+        elif self.frequency == "daily":
+            end_time = now_utc.replace(hour=0, minute=0, second=0, microsecond=0)
+            start_time = end_time - timedelta(days=1)
+        elif self.frequency == "interval":
+            interval = timedelta(seconds=self.interval_seconds or 60)
+            end_time = now_utc
+            start_time = end_time - interval
+        else:
+            raise ValueError(f"Unsupported frequency: {self.frequency}")
+        return FocusTimeWindow(
+            start_time=start_time,
+            end_time=end_time,
+            frequency=self.frequency,
+        )
+
+__all__ = ["FocusLogger"]

litellm/integrations/focus/schema.py

@@ -0,0 +1,50 @@
+"""Schema definitions for Focus export data."""
+
+from __future__ import annotations
+
+import polars as pl
+
+# see: https://focus.finops.org/focus-specification/v1-2/
+FOCUS_NORMALIZED_SCHEMA = pl.Schema(
+    [
+        ("BilledCost", pl.Decimal(18, 6)),
+        ("BillingAccountId", pl.String),
+        ("BillingAccountName", pl.String),
+        ("BillingCurrency", pl.String),
+        ("BillingPeriodStart", pl.Datetime(time_unit="us")),
+        ("BillingPeriodEnd", pl.Datetime(time_unit="us")),
+        ("ChargeCategory", pl.String),
+        ("ChargeClass", pl.String),
+        ("ChargeDescription", pl.String),
+        ("ChargeFrequency", pl.String),
+        ("ChargePeriodStart", pl.Datetime(time_unit="us")),
+        ("ChargePeriodEnd", pl.Datetime(time_unit="us")),
+        ("ConsumedQuantity", pl.Decimal(18, 6)),
+        ("ConsumedUnit", pl.String),
+        ("ContractedCost", pl.Decimal(18, 6)),
+        ("ContractedUnitPrice", pl.Decimal(18, 6)),
+        ("EffectiveCost", pl.Decimal(18, 6)),
+        ("InvoiceIssuerName", pl.String),
+        ("ListCost", pl.Decimal(18, 6)),
+        ("ListUnitPrice", pl.Decimal(18, 6)),
+        ("PricingCategory", pl.String),
+        ("PricingQuantity", pl.Decimal(18, 6)),
+        ("PricingUnit", pl.String),
+        ("ProviderName", pl.String),
+        ("PublisherName", pl.String),
+        ("RegionId", pl.String),
+        ("RegionName", pl.String),
+        ("ResourceId", pl.String),
+        ("ResourceName", pl.String),
+        ("ResourceType", pl.String),
+        ("ServiceCategory", pl.String),
+        ("ServiceSubcategory", pl.String),
+        ("ServiceName", pl.String),
+        ("SubAccountId", pl.String),
+        ("SubAccountName", pl.String),
+        ("SubAccountType", pl.String),
+        ("Tags", pl.Object),
+    ]
+)
+
+__all__ = ["FOCUS_NORMALIZED_SCHEMA"]

litellm/integrations/focus/serializers/__init__.py

@@ -0,0 +1,6 @@
+"""Serializer package exports for Focus integration."""
+
+from .base import FocusSerializer
+from .parquet import FocusParquetSerializer
+
+__all__ = ["FocusSerializer", "FocusParquetSerializer"]

litellm/integrations/focus/serializers/base.py

@@ -0,0 +1,18 @@
+"""Serializer abstractions for Focus export."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import polars as pl
+
+
+class FocusSerializer(ABC):
+    """Base serializer turning Focus frames into bytes."""
+
+    extension: str = ""
+
+    @abstractmethod
+    def serialize(self, frame: pl.DataFrame) -> bytes:
+        """Convert the normalized Focus frame into the chosen format."""
+        raise NotImplementedError

litellm/integrations/focus/serializers/parquet.py

@@ -0,0 +1,22 @@
+"""Parquet serializer for Focus export."""
+
+from __future__ import annotations
+
+import io
+
+import polars as pl
+
+from .base import FocusSerializer
+
+
+class FocusParquetSerializer(FocusSerializer):
+    """Serialize normalized Focus frames to Parquet bytes."""
+
+    extension = "parquet"
+
+    def serialize(self, frame: pl.DataFrame) -> bytes:
+        """Encode the provided frame as a parquet payload."""
+        target = frame if not frame.is_empty() else pl.DataFrame(schema=frame.schema)
+        buffer = io.BytesIO()
+        target.write_parquet(buffer, compression="snappy")
+        return buffer.getvalue()

litellm/integrations/focus/transformer.py

@@ -0,0 +1,90 @@
+"""Focus export data transformer."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+
+import polars as pl
+
+from .schema import FOCUS_NORMALIZED_SCHEMA
+
+
+class FocusTransformer:
+    """Transforms LiteLLM DB rows into Focus-compatible schema."""
+
+    schema = FOCUS_NORMALIZED_SCHEMA
+
+    def transform(self, frame: pl.DataFrame) -> pl.DataFrame:
+        """Return a normalized frame expected by downstream serializers."""
+        if frame.is_empty():
+            return pl.DataFrame(schema=self.schema)
+
+        # derive period start/end from usage date
+        frame = frame.with_columns(
+            pl.col("date")
+            .cast(pl.Utf8)
+            .str.strptime(pl.Datetime(time_unit="us"), format="%Y-%m-%d", strict=False)
+            .alias("usage_date"),
+        )
+        frame = frame.with_columns(
+            pl.col("usage_date").alias("ChargePeriodStart"),
+            (pl.col("usage_date") + timedelta(days=1)).alias("ChargePeriodEnd"),
+        )
+
+        def fmt(col):
+            return col.dt.strftime("%Y-%m-%dT%H:%M:%SZ")
+
+        DEC = pl.Decimal(18, 6)
+
+        def dec(col):
+            return col.cast(DEC)
+
+        none_str = pl.lit(None, dtype=pl.Utf8)
+        none_dec = pl.lit(None, dtype=pl.Decimal(18, 6))
+
+        return frame.select(
+            dec(pl.col("spend").fill_null(0.0)).alias("BilledCost"),
+            pl.col("api_key").cast(pl.String).alias("BillingAccountId"),
+            pl.col("api_key_alias").cast(pl.String).alias("BillingAccountName"),
+            pl.lit("API Key").alias("BillingAccountType"),
+            pl.lit("USD").alias("BillingCurrency"),
+            fmt(pl.col("ChargePeriodEnd")).alias("BillingPeriodEnd"),
+            fmt(pl.col("ChargePeriodStart")).alias("BillingPeriodStart"),
+            pl.lit("Usage").alias("ChargeCategory"),
+            none_str.alias("ChargeClass"),
+            pl.col("model").cast(pl.String).alias("ChargeDescription"),
+            pl.lit("Usage-Based").alias("ChargeFrequency"),
+            fmt(pl.col("ChargePeriodEnd")).alias("ChargePeriodEnd"),
+            fmt(pl.col("ChargePeriodStart")).alias("ChargePeriodStart"),
+            dec(pl.lit(1.0)).alias("ConsumedQuantity"),
+            pl.lit("Requests").alias("ConsumedUnit"),
+            dec(pl.col("spend").fill_null(0.0)).alias("ContractedCost"),
+            none_str.alias("ContractedUnitPrice"),
+            dec(pl.col("spend").fill_null(0.0)).alias("EffectiveCost"),
+            pl.col("custom_llm_provider").cast(pl.String).alias("InvoiceIssuerName"),
+            none_str.alias("InvoiceId"),
+            dec(pl.col("spend").fill_null(0.0)).alias("ListCost"),
+            none_dec.alias("ListUnitPrice"),
+            none_str.alias("AvailabilityZone"),
+            pl.lit("USD").alias("PricingCurrency"),
+            none_str.alias("PricingCategory"),
+            dec(pl.lit(1.0)).alias("PricingQuantity"),
+            none_dec.alias("PricingCurrencyContractedUnitPrice"),
+            dec(pl.col("spend").fill_null(0.0)).alias("PricingCurrencyEffectiveCost"),
+            none_dec.alias("PricingCurrencyListUnitPrice"),
+            pl.lit("Requests").alias("PricingUnit"),
+            pl.col("custom_llm_provider").cast(pl.String).alias("ProviderName"),
+            pl.col("custom_llm_provider").cast(pl.String).alias("PublisherName"),
+            none_str.alias("RegionId"),
+            none_str.alias("RegionName"),
+            pl.col("model").cast(pl.String).alias("ResourceId"),
+            pl.col("model").cast(pl.String).alias("ResourceName"),
+            pl.col("model").cast(pl.String).alias("ResourceType"),
+            pl.lit("AI and Machine Learning").alias("ServiceCategory"),
+            pl.lit("Generative AI").alias("ServiceSubcategory"),
+            pl.col("model_group").cast(pl.String).alias("ServiceName"),
+            pl.col("team_id").cast(pl.String).alias("SubAccountId"),
+            pl.col("team_alias").cast(pl.String).alias("SubAccountName"),
+            none_str.alias("SubAccountType"),
+            none_str.alias("Tags"),
+        )

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"
-    }
-}
+    "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"]
+  }
+}

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
 
@@ -262,6 +263,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",
@@ -354,6 +385,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
@@ -825,7 +875,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"
         )
@@ -845,7 +895,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", {})
@@ -856,7 +906,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 {}),
@@ -976,6 +1026,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.
@@ -1045,6 +1101,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],
@@ -1229,6 +1333,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
 
@@ -1249,7 +1369,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"
         )
@@ -1603,7 +1723,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"
@@ -1636,12 +1755,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
@@ -1784,6 +1902,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 = ""
@@ -2572,10 +2734,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))
@@ -2647,7 +2809,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",

litellm/litellm_core_utils/custom_logger_registry.py

@@ -18,6 +18,7 @@ from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLog
 from litellm.integrations.bitbucket import BitBucketPromptManager
 from litellm.integrations.braintrust_logging import BraintrustLogger
 from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
+from litellm.integrations.focus.focus_logger import FocusLogger
 from litellm.integrations.datadog.datadog import DataDogLogger
 from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
 from litellm.integrations.deepeval import DeepEvalLogger
@@ -93,6 +94,7 @@ class CustomLoggerRegistry:
         "bitbucket": BitBucketPromptManager,
         "gitlab": GitLabPromptManager,
         "cloudzero": CloudZeroLogger,
+        "focus": FocusLogger,
         "posthog": PostHogLogger,
     }
 

litellm/litellm_core_utils/litellm_logging.py

@@ -3756,6 +3756,15 @@ def _init_custom_logger_compatible_class(  # noqa: PLR0915
             cloudzero_logger = CloudZeroLogger()
             _in_memory_loggers.append(cloudzero_logger)
             return cloudzero_logger  # type: ignore
+        elif logging_integration == "focus":
+            from litellm.integrations.focus.focus_logger import FocusLogger
+
+            for callback in _in_memory_loggers:
+                if isinstance(callback, FocusLogger):
+                    return callback  # type: ignore
+            focus_logger = FocusLogger()
+            _in_memory_loggers.append(focus_logger)
+            return focus_logger  # type: ignore
         elif logging_integration == "deepeval":
             for callback in _in_memory_loggers:
                 if isinstance(callback, DeepEvalLogger):
@@ -4076,6 +4085,12 @@ def get_custom_logger_compatible_class(  # noqa: PLR0915
             for callback in _in_memory_loggers:
                 if isinstance(callback, CloudZeroLogger):
                     return callback
+        elif logging_integration == "focus":
+            from litellm.integrations.focus.focus_logger import FocusLogger
+
+            for callback in _in_memory_loggers:
+                if isinstance(callback, FocusLogger):
+                    return callback
         elif logging_integration == "deepeval":
             for callback in _in_memory_loggers:
                 if isinstance(callback, DeepEvalLogger):
@@ -4824,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(

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

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,

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

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

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(

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

litellm/llms/bedrock/count_tokens/handler.py

@@ -70,6 +70,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 +80,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)

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/"):

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(

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",

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,
+        )

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:

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,

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

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(
@@ -2104,6 +2104,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
     rotation_interval: Optional[str] = None  # How often to rotate (e.g., "30d", "90d")
     last_rotation_at: Optional[datetime] = None  # When this key was last rotated
     key_rotation_at: Optional[datetime] = None  # When this key should next be rotated
+    router_settings: Optional[Dict] = None  # Router settings for this key (Key > Team > Global precedence)
 
     model_config = ConfigDict(protected_namespaces=())
 

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={},

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,

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
@@ -440,6 +524,29 @@ class ProxyBaseLLMRequestProcessing:
             user_api_key_dict=user_api_key_dict, data=self.data, call_type=route_type  # type: ignore
         )
 
+        # Apply hierarchical router_settings (Key > Team > Global)
+        if llm_router is not None and proxy_config is not None:
+            from litellm.proxy.proxy_server import prisma_client
+            
+            router_settings = await proxy_config._get_hierarchical_router_settings(
+                user_api_key_dict=user_api_key_dict,
+                prisma_client=prisma_client,
+            )
+            
+            # If router_settings found (from key, team, or global), apply them
+            # This ensures key/team settings override global settings
+            if router_settings is not None and router_settings:
+                # Get model_list from current router
+                model_list = llm_router.get_model_list()
+                if model_list is not None:
+                    # Create user_config with model_list and router_settings
+                    # This creates a per-request router with the hierarchical settings
+                    user_config = {
+                        "model_list": model_list,
+                        **router_settings
+                    }
+                    self.data["user_config"] = user_config
+
         if "messages" in self.data and self.data["messages"]:
             logging_obj.update_messages(self.data["messages"])
 
@@ -647,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,
@@ -658,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,
@@ -900,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":
@@ -1155,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
 

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")
 

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(

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

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":

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
@@ -533,9 +533,9 @@ except ImportError:
 server_root_path = os.getenv("SERVER_ROOT_PATH", "")
 _license_check = LicenseCheck()
 premium_user: bool = _license_check.is_premium()
-premium_user_data: Optional["EnterpriseLicenseData"] = (
-    _license_check.airgapped_license_data
-)
+premium_user_data: Optional[
+    "EnterpriseLicenseData"
+] = _license_check.airgapped_license_data
 global_max_parallel_request_retries_env: Optional[str] = os.getenv(
     "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
 )
@@ -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]
 
 
@@ -1083,9 +1094,7 @@ try:
     # In non-root Docker, we restructure in /var/lib/litellm/ui.
     try:
         _restructure_ui_html_files(ui_path)
-        verbose_proxy_logger.info(
-            f"Restructured UI directory: {ui_path}"
-        )
+        verbose_proxy_logger.info(f"Restructured UI directory: {ui_path}")
     except PermissionError as e:
         verbose_proxy_logger.exception(
             f"Permission error while restructuring UI directory {ui_path}: {e}"
@@ -1171,9 +1180,9 @@ master_key: Optional[str] = None
 config_agents: Optional[List[AgentConfig]] = None
 otel_logging = False
 prisma_client: Optional[PrismaClient] = None
-shared_aiohttp_session: Optional["ClientSession"] = (
-    None  # Global shared session for connection reuse
-)
+shared_aiohttp_session: Optional[
+    "ClientSession"
+] = None  # Global shared session for connection reuse
 user_api_key_cache = DualCache(
     default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
 )
@@ -1181,9 +1190,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
     dual_cache=user_api_key_cache
 )
 litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
-redis_usage_cache: Optional[RedisCache] = (
-    None  # redis cache used for tracking spend, tpm/rpm limits
-)
+redis_usage_cache: Optional[
+    RedisCache
+] = None  # redis cache used for tracking spend, tpm/rpm limits
 polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False
 polling_cache_ttl: int = 3600  # Default 1 hour TTL for polling cache
 user_custom_auth = None
@@ -1522,9 +1531,9 @@ async def update_cache(  # noqa: PLR0915
         _id = "team_id:{}".format(team_id)
         try:
             # Fetch the existing cost for the given user
-            existing_spend_obj: Optional[LiteLLM_TeamTable] = (
-                await user_api_key_cache.async_get_cache(key=_id)
-            )
+            existing_spend_obj: Optional[
+                LiteLLM_TeamTable
+            ] = await user_api_key_cache.async_get_cache(key=_id)
             if existing_spend_obj is None:
                 # do nothing if team not in api key cache
                 return
@@ -1876,7 +1885,6 @@ class ProxyConfig:
                 "environment_variables" in config_to_save
                 and config_to_save["environment_variables"]
             ):
-
                 # decrypt the environment_variables - in case a caller function has already encrypted the environment_variables
                 decrypted_env_vars = self._decrypt_and_set_db_env_variables(
                     environment_variables=config_to_save["environment_variables"],
@@ -2794,21 +2802,21 @@ class ProxyConfig:
         verbose_proxy_logger.debug(f"_alerting_callbacks: {general_settings}")
         if _alerting_callbacks is None:
             return
-        
+
         # Ensure proxy_logging_obj.alerting is set for all alerting types
         _alerting_value = general_settings.get("alerting", None)
-        verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}")
+        verbose_proxy_logger.debug(
+            f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}"
+        )
         proxy_logging_obj.update_values(
             alerting=_alerting_value,
             alerting_threshold=general_settings.get("alerting_threshold", 600),
             alert_types=general_settings.get("alert_types", None),
-            alert_to_webhook_url=general_settings.get(
-                "alert_to_webhook_url", None
-            ),
+            alert_to_webhook_url=general_settings.get("alert_to_webhook_url", None),
             alerting_args=general_settings.get("alerting_args", None),
             redis_cache=redis_usage_cache,
         )
-        
+
         for _alert in _alerting_callbacks:
             if _alert == "slack":
                 # [OLD] v0 implementation - already handled by update_values above
@@ -3222,6 +3230,84 @@ class ProxyConfig:
             decrypted_variables[k] = decrypted_value
         return decrypted_variables
 
+    async def _get_hierarchical_router_settings(
+        self,
+        user_api_key_dict: Optional["UserAPIKeyAuth"],
+        prisma_client: Optional[PrismaClient],
+    ) -> Optional[dict]:
+        """
+        Get router_settings in priority order: Key > Team > Global
+        
+        Returns:
+            dict: Combined router_settings, or None if no settings found
+        """
+        if prisma_client is None:
+            return None
+        
+        import json
+        import yaml
+        
+        # 1. Try key-level router_settings
+        if user_api_key_dict is not None:
+            # Check if router_settings is available on the key object
+            key_router_settings_value = getattr(user_api_key_dict, "router_settings", None)
+            if key_router_settings_value is not None:
+                key_router_settings = None
+                if isinstance(key_router_settings_value, str):
+                    try:
+                        key_router_settings = yaml.safe_load(key_router_settings_value)
+                    except (yaml.YAMLError, json.JSONDecodeError):
+                        try:
+                            key_router_settings = json.loads(key_router_settings_value)
+                        except json.JSONDecodeError:
+                            pass
+                elif isinstance(key_router_settings_value, dict):
+                    key_router_settings = key_router_settings_value
+                
+                # If key has router_settings (non-empty dict), use it
+                if key_router_settings is not None and isinstance(key_router_settings, dict) and key_router_settings:
+                    return key_router_settings
+        
+        # 2. Try team-level router_settings
+        if user_api_key_dict is not None and user_api_key_dict.team_id is not None:
+            try:
+                team_obj = await prisma_client.db.litellm_teamtable.find_unique(
+                    where={"team_id": user_api_key_dict.team_id}
+                )
+                if team_obj is not None:
+                    team_router_settings_value = getattr(team_obj, "router_settings", None)
+                    if team_router_settings_value is not None:
+                        team_router_settings = None
+                        if isinstance(team_router_settings_value, str):
+                            try:
+                                team_router_settings = yaml.safe_load(team_router_settings_value)
+                            except (yaml.YAMLError, json.JSONDecodeError):
+                                try:
+                                    team_router_settings = json.loads(team_router_settings_value)
+                                except json.JSONDecodeError:
+                                    pass
+                        elif isinstance(team_router_settings_value, dict):
+                            team_router_settings = team_router_settings_value
+                        
+                        # If team has router_settings (non-empty dict), use it
+                        if team_router_settings is not None and isinstance(team_router_settings, dict) and team_router_settings:
+                            return team_router_settings
+            except Exception:
+                # If team lookup fails, continue to global settings
+                pass
+        
+        # 3. Try global router_settings
+        try:
+            db_router_settings = await prisma_client.db.litellm_config.find_first(
+                where={"param_name": "router_settings"}
+            )
+            if db_router_settings is not None and isinstance(db_router_settings.param_value, dict) and db_router_settings.param_value:
+                return db_router_settings.param_value
+        except Exception:
+            pass
+        
+        return None
+
     async def _add_router_settings_from_db_config(
         self,
         config_data: dict,
@@ -3279,7 +3365,7 @@ class ProxyConfig:
             proxy_logging_obj: ProxyLogging
         """
         _general_settings = config_data.get("general_settings", {})
-    
+
         if _general_settings is not None and "alerting" in _general_settings:
             if (
                 general_settings is not None
@@ -3294,7 +3380,8 @@ class ProxyConfig:
                 _merged_alerting = list(_yaml_alerting.union(_db_alerting))
                 # Preserve order: YAML values first, then DB values
                 _merged_alerting = list(general_settings["alerting"]) + [
-                    item for item in _general_settings["alerting"] 
+                    item
+                    for item in _general_settings["alerting"]
                     if item not in general_settings["alerting"]
                 ]
                 verbose_proxy_logger.debug(
@@ -3608,7 +3695,6 @@ class ProxyConfig:
             await self._init_vector_stores_in_db(prisma_client=prisma_client)
 
         if self._should_load_db_object(object_type="vector_store_indexes"):
-
             await self._init_vector_store_indexes_in_db(prisma_client=prisma_client)
 
         if self._should_load_db_object(object_type="mcp"):
@@ -3807,10 +3893,10 @@ class ProxyConfig:
         )
 
         try:
-            guardrails_in_db: List[Guardrail] = (
-                await GuardrailRegistry.get_all_guardrails_from_db(
-                    prisma_client=prisma_client
-                )
+            guardrails_in_db: List[
+                Guardrail
+            ] = await GuardrailRegistry.get_all_guardrails_from_db(
+                prisma_client=prisma_client
             )
             verbose_proxy_logger.debug(
                 "guardrails from the DB %s", str(guardrails_in_db)
@@ -4137,9 +4223,9 @@ async def initialize(  # noqa: PLR0915
         user_api_base = api_base
         dynamic_config[user_model]["api_base"] = api_base
     if api_version:
-        os.environ["AZURE_API_VERSION"] = (
-            api_version  # set this for azure - litellm can read this from the env
-        )
+        os.environ[
+            "AZURE_API_VERSION"
+        ] = api_version  # set this for azure - litellm can read this from the env
     if max_tokens:  # model-specific param
         dynamic_config[user_model]["max_tokens"] = max_tokens
     if temperature:  # model-specific param
@@ -4657,10 +4743,14 @@ class ProxyStartupEvent:
                     replace_existing=True,
                     misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
                 )
-                verbose_proxy_logger.info("Responses cost check job scheduled successfully")
+                verbose_proxy_logger.info(
+                    "Responses cost check job scheduled successfully"
+                )
 
             except Exception as e:
-                verbose_proxy_logger.debug(f"Failed to setup responses cost checking: {e}")
+                verbose_proxy_logger.debug(
+                    f"Failed to setup responses cost checking: {e}"
+                )
                 verbose_proxy_logger.debug(
                     "Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..."
                 )
@@ -4684,8 +4774,9 @@ class ProxyStartupEvent:
         """
         Initialize the spend tracking and other background jobs
         1. CloudZero Background Job
-        2. Prometheus Background Job
-        3. Key Rotation Background Job
+        2. Focus Background Job
+        3. Prometheus Background Job
+        4. Key Rotation Background Job
 
         Args:
             scheduler: The scheduler to add the background jobs to
@@ -4694,11 +4785,17 @@ class ProxyStartupEvent:
         # CloudZero Background Job
         ########################################################
         from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
+        from litellm.integrations.focus.focus_logger import FocusLogger
         from litellm.proxy.spend_tracking.cloudzero_endpoints import is_cloudzero_setup
 
         if await is_cloudzero_setup():
             await CloudZeroLogger.init_cloudzero_background_job(scheduler=scheduler)
 
+        ########################################################
+        # Focus Background Job
+        ########################################################
+        await FocusLogger.init_focus_export_background_job(scheduler=scheduler)
+
         ########################################################
         # Prometheus Background Job
         ########################################################
@@ -4829,6 +4926,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()
@@ -5940,7 +6045,6 @@ async def realtime_websocket_endpoint(
     ),
     user_api_key_dict=Depends(user_api_key_auth_websocket),
 ):
-
     await websocket.accept()
 
     # Only use explicit parameters, not all query params
@@ -6739,7 +6843,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,
@@ -9528,9 +9632,9 @@ async def get_config_list(
                             hasattr(sub_field_info, "description")
                             and sub_field_info.description is not None
                         ):
-                            nested_fields[idx].field_description = (
-                                sub_field_info.description
-                            )
+                            nested_fields[
+                                idx
+                            ].field_description = sub_field_info.description
                         idx += 1
 
                     _stored_in_db = None

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"],

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

litellm/router.py

@@ -4485,9 +4485,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
@@ -7664,6 +7676,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 (

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
@@ -188,6 +188,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",
 ]
 
 
@@ -222,6 +230,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,
@@ -451,6 +476,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)
@@ -475,11 +515,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)

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
@@ -7718,6 +7717,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 +7945,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

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,

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"

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"
+    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]

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

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")

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

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

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

tests/test_litellm/integrations/focus/test_focus_database.py

@@ -0,0 +1,74 @@
+"""Tests for FocusLiteLLMDatabase query construction."""
+
+from datetime import datetime, timezone
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from litellm.integrations.focus.database import FocusLiteLLMDatabase
+
+
+def _setup_db(monkeypatch: pytest.MonkeyPatch, query_return):
+    """Create a database instance with a stubbed prisma client."""
+    query_mock = AsyncMock(return_value=query_return)
+    mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=query_mock))
+    db = FocusLiteLLMDatabase()
+    monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
+    return db, query_mock
+
+
+@pytest.mark.asyncio
+async def test_should_parameterize_filters_and_limit(monkeypatch: pytest.MonkeyPatch):
+    start = datetime(2024, 1, 1, tzinfo=timezone.utc)
+    end = datetime(2024, 1, 2, tzinfo=timezone.utc)
+    db, query_mock = _setup_db(monkeypatch, [])
+
+    await db.get_usage_data(limit=25, 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, 25]
+
+
+@pytest.mark.asyncio
+async def test_should_execute_without_filters(monkeypatch: pytest.MonkeyPatch):
+    row = {
+        "id": 1,
+        "user_id": "user",
+        "date": datetime(2024, 1, 1, tzinfo=timezone.utc),
+    }
+    db, query_mock = _setup_db(monkeypatch, [row])
+
+    result = await db.get_usage_data()
+
+    query_text, *params = query_mock.await_args.args
+    assert "WHERE" not in query_text
+    assert "LIMIT $" not in query_text
+    assert params == []
+    assert result.height == 1
+    assert result["id"][0] == 1
+
+
+@pytest.mark.asyncio
+async def test_should_accept_string_timestamps(monkeypatch: pytest.MonkeyPatch):
+    db, query_mock = _setup_db(monkeypatch, [])
+
+    start = "2024-02-01T00:00:00+00:00"
+    end = "2024-02-02T00:00:00+00:00"
+    await db.get_usage_data(start_time_utc=start, end_time_utc=end)
+
+    _, *params = query_mock.await_args.args
+    assert params == [start, end]
+
+
+@pytest.mark.asyncio
+async def test_should_reject_invalid_limit(monkeypatch: pytest.MonkeyPatch):
+    db, query_mock = _setup_db(monkeypatch, [])
+
+    with pytest.raises(ValueError):
+        await db.get_usage_data(limit="invalid")
+
+    assert query_mock.await_count == 0

tests/test_litellm/integrations/focus/test_s3_destination.py

@@ -0,0 +1,100 @@
+"""Tests for FocusS3Destination behavior."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from types import SimpleNamespace
+from typing import Any, Dict
+
+import pytest
+
+import litellm.integrations.focus.destinations.s3_destination as s3_module
+from litellm.integrations.focus.destinations.base import FocusTimeWindow
+from litellm.integrations.focus.destinations.s3_destination import FocusS3Destination
+
+
+def _window(freq: str = "hourly", hour: int = 5) -> FocusTimeWindow:
+    start = datetime(2024, 1, 2, hour, tzinfo=timezone.utc)
+    end = start.replace(hour=hour + 1)
+    return FocusTimeWindow(start_time=start, end_time=end, frequency=freq)
+
+
+def test_should_require_bucket_name():
+    with pytest.raises(ValueError):
+        FocusS3Destination(prefix="focus", config={})
+
+
+def test_should_build_hourly_object_key():
+    dest = FocusS3Destination(prefix="exports/", config={"bucket_name": "bucket"})
+    key = dest._build_object_key(
+        time_window=_window(freq="hourly", hour=3), filename="data.snappy"
+    )
+    assert key == "exports/date=2024-01-02/hour=03/data.snappy"
+
+
+def test_should_build_daily_key_without_hour_segment():
+    dest = FocusS3Destination(prefix="", config={"bucket_name": "bucket"})
+    key = dest._build_object_key(
+        time_window=_window(freq="daily", hour=0), filename="daily.parquet"
+    )
+    assert key == "date=2024-01-02/daily.parquet"
+
+
+@pytest.mark.asyncio
+async def test_should_dispatch_upload_via_thread(monkeypatch: pytest.MonkeyPatch):
+    dest = FocusS3Destination(prefix="focus", config={"bucket_name": "bucket"})
+    captured: Dict[str, Any] = {}
+
+    async def fake_to_thread(func, *args, **kwargs):  # type: ignore[override]
+        captured["func"] = func
+        captured["args"] = args
+        captured["kwargs"] = kwargs
+
+    monkeypatch.setattr(s3_module.asyncio, "to_thread", fake_to_thread)
+
+    window = _window(freq="hourly", hour=1)
+    await dest.deliver(content=b"payload", time_window=window, filename="file.bin")
+
+    assert captured["func"] == dest._upload
+    assert captured["args"][0] == b"payload"
+    assert captured["args"][1].endswith("/file.bin")
+
+
+def test_should_upload_with_configured_client(monkeypatch: pytest.MonkeyPatch):
+    config = {
+        "bucket_name": "bucket",
+        "region_name": "us-east-2",
+        "endpoint_url": "http://localhost:4566",
+        "aws_access_key_id": "key",
+        "aws_secret_access_key": "secret",
+        "aws_session_token": "token",
+    }
+    dest = FocusS3Destination(prefix="focus", config=config)
+    captured: Dict[str, Any] = {}
+
+    def fake_client(service: str, **kwargs):
+        assert service == "s3"
+        captured["client_kwargs"] = kwargs
+
+        def put_object(**put_kwargs):
+            captured["put_kwargs"] = put_kwargs
+
+        return SimpleNamespace(put_object=put_object)
+
+    monkeypatch.setattr(s3_module.boto3, "client", fake_client)
+
+    dest._upload(content=b"payload", object_key="path/file.bin")
+
+    assert captured["client_kwargs"] == {
+        "region_name": "us-east-2",
+        "endpoint_url": "http://localhost:4566",
+        "aws_access_key_id": "key",
+        "aws_secret_access_key": "secret",
+        "aws_session_token": "token",
+    }
+    assert captured["put_kwargs"] == {
+        "Bucket": "bucket",
+        "Key": "path/file.bin",
+        "Body": b"payload",
+        "ContentType": "application/octet-stream",
+    }

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")

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"])

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

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 

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

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
+

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

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!")

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

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
@@ -349,6 +360,7 @@ async def test_proxy_admin_expired_key_from_cache():
             pass
 
 
+
 @pytest.mark.asyncio
 async def test_return_user_api_key_auth_obj_user_spend_and_budget():
     """

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())

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:

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):
         """

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
 
@@ -75,6 +76,84 @@ class TestProxyBaseLLMRequestProcessing:
             pytest.fail("litellm_call_id is not a valid UUID")
         assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"]
 
+    @pytest.mark.asyncio
+    async def test_should_apply_hierarchical_router_settings_to_user_config(
+        self, monkeypatch
+    ):
+        processing_obj = ProxyBaseLLMRequestProcessing(data={})
+        mock_request = MagicMock(spec=Request)
+        mock_request.headers = {}
+
+        async def mock_add_litellm_data_to_request(*args, **kwargs):
+            return {}
+
+        async def mock_common_processing_pre_call_logic(
+            user_api_key_dict, data, call_type
+        ):
+            data_copy = copy.deepcopy(data)
+            return data_copy
+
+        mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
+        mock_proxy_logging_obj.pre_call_hook = AsyncMock(
+            side_effect=mock_common_processing_pre_call_logic
+        )
+        monkeypatch.setattr(
+            litellm.proxy.common_request_processing,
+            "add_litellm_data_to_request",
+            mock_add_litellm_data_to_request,
+        )
+
+        mock_general_settings = {}
+        mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
+        mock_proxy_config = MagicMock(spec=ProxyConfig)
+        
+        mock_router_settings = {
+            "routing_strategy": "least-busy",
+            "timeout": 30.0,
+            "num_retries": 3,
+        }
+        mock_proxy_config._get_hierarchical_router_settings = AsyncMock(
+            return_value=mock_router_settings
+        )
+
+        mock_model_list = [
+            {"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo"}},
+            {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}},
+        ]
+        mock_llm_router = MagicMock()
+        mock_llm_router.get_model_list = MagicMock(return_value=mock_model_list)
+
+        mock_prisma_client = MagicMock()
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.prisma_client",
+            mock_prisma_client,
+        )
+
+        route_type = "acompletion"
+
+        returned_data, logging_obj = await processing_obj.common_processing_pre_call_logic(
+            request=mock_request,
+            general_settings=mock_general_settings,
+            user_api_key_dict=mock_user_api_key_dict,
+            proxy_logging_obj=mock_proxy_logging_obj,
+            proxy_config=mock_proxy_config,
+            route_type=route_type,
+            llm_router=mock_llm_router,
+        )
+
+        mock_proxy_config._get_hierarchical_router_settings.assert_called_once_with(
+            user_api_key_dict=mock_user_api_key_dict,
+            prisma_client=mock_prisma_client,
+        )
+        mock_llm_router.get_model_list.assert_called_once()
+
+        assert "user_config" in returned_data
+        user_config = returned_data["user_config"]
+        assert user_config["model_list"] == mock_model_list
+        assert user_config["routing_strategy"] == "least-busy"
+        assert user_config["timeout"] == 30.0
+        assert user_config["num_retries"] == 3
+
     @pytest.mark.asyncio
     async def test_stream_timeout_header_processing(self):
         """
@@ -602,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():
@@ -624,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
@@ -641,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
@@ -654,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 == []
@@ -665,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 = {
@@ -682,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():
@@ -702,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"
@@ -712,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",
             {},
@@ -729,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
@@ -742,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
@@ -773,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", {}
             )
 
@@ -810,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
@@ -827,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]}"

tests/test_litellm/proxy/test_proxy_server.py

@@ -3124,3 +3124,95 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch):
     assert result["general_settings"]["nested"]["key1"] == "updated_value1"
     assert result["general_settings"]["nested"]["key2"] == "value2"
     assert result["general_settings"]["nested"]["key3"] == "value3"
+
+
+@pytest.mark.asyncio
+async def test_get_hierarchical_router_settings():
+    """
+    Test _get_hierarchical_router_settings method's priority order: Key > Team > Global
+    """
+    from unittest.mock import AsyncMock, MagicMock
+
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    proxy_config = ProxyConfig()
+
+    # Test Case 1: Returns None when prisma_client is None
+    result = await proxy_config._get_hierarchical_router_settings(
+        user_api_key_dict=None,
+        prisma_client=None,
+    )
+    assert result is None
+
+    # Test Case 2: Returns key-level router_settings when available (as dict)
+    mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
+    mock_user_api_key_dict.router_settings = {"routing_strategy": "key-level", "timeout": 10}
+    mock_user_api_key_dict.team_id = None
+
+    mock_prisma_client = MagicMock()
+
+    result = await proxy_config._get_hierarchical_router_settings(
+        user_api_key_dict=mock_user_api_key_dict,
+        prisma_client=mock_prisma_client,
+    )
+    assert result == {"routing_strategy": "key-level", "timeout": 10}
+
+    # Test Case 3: Returns key-level router_settings when available (as YAML string)
+    mock_user_api_key_dict.router_settings = "routing_strategy: key-yaml\ntimeout: 20"
+    result = await proxy_config._get_hierarchical_router_settings(
+        user_api_key_dict=mock_user_api_key_dict,
+        prisma_client=mock_prisma_client,
+    )
+    assert result == {"routing_strategy": "key-yaml", "timeout": 20}
+
+    # Test Case 4: Falls back to team-level router_settings when key-level is not available
+    mock_user_api_key_dict.router_settings = None
+    mock_user_api_key_dict.team_id = "team-123"
+
+    mock_team_obj = MagicMock()
+    mock_team_obj.router_settings = {"routing_strategy": "team-level", "timeout": 30}
+
+    mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
+        return_value=mock_team_obj
+    )
+
+    result = await proxy_config._get_hierarchical_router_settings(
+        user_api_key_dict=mock_user_api_key_dict,
+        prisma_client=mock_prisma_client,
+    )
+    assert result == {"routing_strategy": "team-level", "timeout": 30}
+    mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with(
+        where={"team_id": "team-123"}
+    )
+
+    # Test Case 5: Falls back to global router_settings when neither key nor team settings are available
+    mock_user_api_key_dict.router_settings = None
+    mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
+
+    mock_db_config = MagicMock()
+    mock_db_config.param_value = {"routing_strategy": "global-level", "timeout": 40}
+
+    mock_prisma_client.db.litellm_config.find_first = AsyncMock(
+        return_value=mock_db_config
+    )
+
+    result = await proxy_config._get_hierarchical_router_settings(
+        user_api_key_dict=mock_user_api_key_dict,
+        prisma_client=mock_prisma_client,
+    )
+    assert result == {"routing_strategy": "global-level", "timeout": 40}
+    mock_prisma_client.db.litellm_config.find_first.assert_called_once_with(
+        where={"param_name": "router_settings"}
+    )
+
+    # Test Case 6: Returns None when no settings are found
+    mock_user_api_key_dict.router_settings = None
+    mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
+    mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
+
+    result = await proxy_config._get_hierarchical_router_settings(
+        user_api_key_dict=mock_user_api_key_dict,
+        prisma_client=mock_prisma_client,
+    )
+    assert result is None

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

tests/test_litellm/test_utils_custom.py

@@ -0,0 +1,42 @@
+import pytest
+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"
+    
+    # Mock anthropic
+    with patch("anthropic.AsyncAnthropic") as mock_cls:
+        mock_client = MagicMock()
+        mock_cls.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)
+        
+        # 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_cls.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_cls.assert_called_once()

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<string, Page> = {
+  "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,
+};

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",
+}

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<void> {
+  await page.goto(`/ui?page=${pageEnum}`);
+}

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();

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}(&|$)`));
+    });
   });
 }

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);
+    });
+  });
+});

ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx

@@ -84,10 +84,7 @@ interface ChatUIProps {
   };
 }
 
-const MCP_SUPPORTED_ENDPOINTS = new Set<EndpointType>([
-  EndpointType.CHAT,
-  EndpointType.RESPONSES,
-]);
+const MCP_SUPPORTED_ENDPOINTS = new Set<EndpointType>([EndpointType.CHAT, EndpointType.RESPONSES]);
 
 const ChatUI: React.FC<ChatUIProps> = ({
   accessToken,
@@ -131,7 +128,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
   });
   const [apiKey, setApiKey] = useState<string>(() => sessionStorage.getItem("apiKey") || "");
   const [customProxyBaseUrl, setCustomProxyBaseUrl] = useState<string>(
-    () => sessionStorage.getItem("customProxyBaseUrl") || ""
+    () => sessionStorage.getItem("customProxyBaseUrl") || "",
   );
   const [inputMessage, setInputMessage] = useState("");
   const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
@@ -215,7 +212,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
   const [temperature, setTemperature] = useState<number>(1.0);
   const [maxTokens, setMaxTokens] = useState<number>(2048);
   const [useAdvancedParams, setUseAdvancedParams] = useState<boolean>(false);
-  
+
   // Code Interpreter state (using custom hook)
   const codeInterpreter = useCodeInterpreter();
 
@@ -244,16 +241,15 @@ const ChatUI: React.FC<ChatUIProps> = ({
 
     try {
       const response = await listMCPTools(userApiKey, serverId);
-      setServerToolsMap(prev => ({
+      setServerToolsMap((prev) => ({
         ...prev,
-        [serverId]: response.tools || []
+        [serverId]: response.tools || [],
       }));
     } catch (error) {
       console.error(`Error fetching tools for server ${serverId}:`, error);
     }
   };
 
-
   useEffect(() => {
     if (isGetCodeModalVisible) {
       const code = generateCodeSnippet({
@@ -1007,7 +1003,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
             traceId,
             selectedVectorStores.length > 0 ? selectedVectorStores : undefined,
             selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
-            selectedMCPTools, // Pass the selected tools array
+            selectedMCPServers, // Pass the selected tools array
             customProxyBaseUrl || undefined,
           );
         } else if (endpointType === EndpointType.EMBEDDINGS) {
@@ -1205,9 +1201,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
                   icon={ApiOutlined}
                 />
                 {customProxyBaseUrl && (
-                  <Text className="text-xs text-gray-500 mt-1">
-                    API calls will be sent to: {customProxyBaseUrl}
-                  </Text>
+                  <Text className="text-xs text-gray-500 mt-1">API calls will be sent to: {customProxyBaseUrl}</Text>
                 )}
               </div>
 
@@ -1382,7 +1376,11 @@ const ChatUI: React.FC<ChatUIProps> = ({
                     optionLabelProp="label"
                   >
                     {agentInfo.map((agent) => (
-                      <Select.Option key={agent.agent_id} value={agent.agent_name} label={agent.agent_name || agent.agent_id}>
+                      <Select.Option
+                        key={agent.agent_id}
+                        value={agent.agent_name}
+                        label={agent.agent_name || agent.agent_id}
+                      >
                         <div className="flex flex-col py-1">
                           <span className="font-medium">{agent.agent_name || agent.agent_id}</span>
                           {agent.agent_card_params?.description && (
@@ -1416,10 +1414,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
               <div>
                 <Text className="font-medium block mb-2 text-gray-700 flex items-center">
                   <ToolOutlined className="mr-2" /> MCP Servers
-                  <Tooltip
-                    className="ml-1"
-                    title="Select MCP servers to use in your conversation."
-                  >
+                  <Tooltip className="ml-1" title="Select MCP servers to use in your conversation.">
                     <InfoCircleOutlined />
                   </Tooltip>
                 </Text>
@@ -1435,15 +1430,15 @@ const ChatUI: React.FC<ChatUIProps> = ({
                     } else {
                       setSelectedMCPServers(value);
                       // Clean up tool restrictions for removed servers
-                      setMCPServerToolRestrictions(prev => {
+                      setMCPServerToolRestrictions((prev) => {
                         const updated = { ...prev };
-                        Object.keys(updated).forEach(serverId => {
+                        Object.keys(updated).forEach((serverId) => {
                           if (!value.includes(serverId)) delete updated[serverId];
                         });
                         return updated;
                       });
                       // Load tools for newly selected servers
-                      value.forEach(serverId => {
+                      value.forEach((serverId) => {
                         if (!serverToolsMap[serverId]) {
                           loadServerTools(serverId);
                         }
@@ -1474,12 +1469,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
                       disabled={selectedMCPServers.includes("__all__")}
                     >
                       <div className="flex flex-col py-1">
-                        <span className="font-medium">
-                          {server.alias || server.server_name || server.server_id}
-                        </span>
-                        {server.description && (
-                          <span className="text-xs text-gray-500 mt-1">{server.description}</span>
-                        )}
+                        <span className="font-medium">{server.alias || server.server_name || server.server_id}</span>
+                        {server.description && <span className="text-xs text-gray-500 mt-1">{server.description}</span>}
                       </div>
                     </Select.Option>
                   ))}
@@ -1489,40 +1480,40 @@ const ChatUI: React.FC<ChatUIProps> = ({
                 {selectedMCPServers.length > 0 &&
                   !selectedMCPServers.includes("__all__") &&
                   MCP_SUPPORTED_ENDPOINTS.has(endpointType as EndpointType) && (
-                  <div className="mt-3 space-y-2">
-                    {selectedMCPServers.map(serverId => {
-                      const server = mcpServers.find(s => s.server_id === serverId);
-                      const tools = serverToolsMap[serverId] || [];
-                      if (tools.length === 0) return null;
+                    <div className="mt-3 space-y-2">
+                      {selectedMCPServers.map((serverId) => {
+                        const server = mcpServers.find((s) => s.server_id === serverId);
+                        const tools = serverToolsMap[serverId] || [];
+                        if (tools.length === 0) return null;
 
-                      return (
-                        <div key={serverId} className="border rounded p-2">
-                          <Text className="text-xs text-gray-600 mb-1">
-                            Limit tools for {server?.alias || server?.server_name || serverId}:
-                          </Text>
-                          <Select
-                            mode="multiple"
-                            size="small"
-                            style={{ width: "100%" }}
-                            placeholder="All tools (default)"
-                            value={mcpServerToolRestrictions[serverId] || []}
-                            onChange={(selectedTools) => {
-                              setMCPServerToolRestrictions(prev => ({
-                                ...prev,
-                                [serverId]: selectedTools
-                              }));
-                            }}
-                            options={tools.map(tool => ({
-                              value: tool.name,
-                              label: tool.name
-                            }))}
-                            maxTagCount={2}
-                          />
-                        </div>
-                      );
-                    })}
-                  </div>
-                )}
+                        return (
+                          <div key={serverId} className="border rounded p-2">
+                            <Text className="text-xs text-gray-600 mb-1">
+                              Limit tools for {server?.alias || server?.server_name || serverId}:
+                            </Text>
+                            <Select
+                              mode="multiple"
+                              size="small"
+                              style={{ width: "100%" }}
+                              placeholder="All tools (default)"
+                              value={mcpServerToolRestrictions[serverId] || []}
+                              onChange={(selectedTools) => {
+                                setMCPServerToolRestrictions((prev) => ({
+                                  ...prev,
+                                  [serverId]: selectedTools,
+                                }));
+                              }}
+                              options={tools.map((tool) => ({
+                                value: tool.name,
+                                label: tool.name,
+                              }))}
+                              maxTagCount={2}
+                            />
+                          </div>
+                        );
+                      })}
+                    </div>
+                  )}
               </div>
 
               <div>
@@ -2073,7 +2064,13 @@ const ChatUI: React.FC<ChatUIProps> = ({
                     )}
                     {/* Quick Code Interpreter toggle for Responses */}
                     {endpointType === EndpointType.RESPONSES && (
-                      <Tooltip title={codeInterpreter.enabled ? "Code Interpreter enabled (click to disable)" : "Enable Code Interpreter"}>
+                      <Tooltip
+                        title={
+                          codeInterpreter.enabled
+                            ? "Code Interpreter enabled (click to disable)"
+                            : "Enable Code Interpreter"
+                        }
+                      >
                         <button
                           className={`p-1.5 rounded-md transition-colors ${
                             codeInterpreter.enabled

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(<CodeInterpreterOutput code="print('hello')" accessToken="test-token" />);
+
+    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(<CodeInterpreterOutput code={code} accessToken="test-token" />);
+
+    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(
+      <CodeInterpreterOutput
+        code="import matplotlib.pyplot as plt"
+        annotations={annotations}
+        accessToken="test-token"
+      />,
+    );
+
+    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<Blob>((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(
+      <CodeInterpreterOutput
+        code="import matplotlib.pyplot as plt"
+        annotations={annotations}
+        accessToken="test-token"
+      />,
+    );
+
+    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(
+      <CodeInterpreterOutput
+        code="import matplotlib.pyplot as plt"
+        annotations={annotations}
+        accessToken="test-token"
+      />,
+    );
+
+    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(<CodeInterpreterOutput code="import pandas as pd" annotations={annotations} accessToken="test-token" />);
+
+    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(<CodeInterpreterOutput accessToken="test-token" />);
+
+    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(
+      <CodeInterpreterOutput
+        code="import matplotlib.pyplot as plt"
+        annotations={annotations}
+        accessToken="test-token"
+      />,
+    );
+
+    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(
+      <CodeInterpreterOutput
+        code="import matplotlib.pyplot as plt"
+        annotations={annotations}
+        accessToken="test-token"
+      />,
+    );
+
+    await waitFor(() => {
+      expect(consoleErrorSpy).toHaveBeenCalled();
+    });
+
+    consoleErrorSpy.mockRestore();
+  });
+});

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(<UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />);
+
+    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(
+      <UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />,
+    );
+
+    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(
+      <UnifiedSelector value="" options={options} loading={true} config={config} onChange={onChange} />,
+    );
+
+    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(<UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />);
+
+    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(
+      <UnifiedSelector value="option1" options={options} loading={false} config={config} onChange={onChange} />,
+    );
+
+    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(<UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />);
+
+    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(<UnifiedSelector value="" options={options} loading={true} config={config} onChange={onChange} />);
+
+    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(<UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />);
+
+    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(
+      <UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />,
+    );
+
+    const placeholder = container.querySelector(".ant-select-selection-placeholder");
+    expect(placeholder).toHaveTextContent(config.selectorPlaceholder);
+  });
+});

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);
+  });
+});

ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.test.tsx

@@ -119,8 +119,25 @@ describe("chat_completion", () => {
   it("should configure MCP tools per server with restrictions", async () => {
     const selectedMCPServers = ["server-1", "server-2"];
     const mcpServers = [
-      { server_id: "server-1", alias: "alpha", server_name: "Alpha" },
-      { server_id: "server-2", server_name: "Beta" },
+      {
+        server_id: "server-1",
+        alias: "alpha",
+        server_name: "Alpha",
+        url: "http://example.com",
+        created_at: "2024-01-01",
+        created_by: "test",
+        updated_at: "2024-01-01",
+        updated_by: "test",
+      },
+      {
+        server_id: "server-2",
+        server_name: "Beta",
+        url: "http://example.com",
+        created_at: "2024-01-01",
+        created_by: "test",
+        updated_at: "2024-01-01",
+        updated_by: "test",
+      },
     ];
     const mcpServerToolRestrictions = {
       "server-1": ["toolA", "toolB"],
@@ -132,41 +149,43 @@ describe("chat_completion", () => {
       mockUpdateUI,
       "gpt-4",
       "test-token",
-      undefined,
-      undefined,
-      undefined,
-      undefined,
-      undefined,
-      undefined,
-      undefined,
-      undefined,
+      undefined, // tags
+      undefined, // signal
+      undefined, // onReasoningContent
+      undefined, // onTimingData
+      undefined, // onUsageData
+      undefined, // traceId
+      undefined, // vector_store_ids
+      undefined, // guardrails
       selectedMCPServers,
-      undefined,
-      undefined,
-      undefined,
-      undefined,
-      undefined,
+      undefined, // onImageGenerated
+      undefined, // onSearchResults
+      undefined, // temperature
+      undefined, // max_tokens
+      undefined, // onTotalLatency
+      undefined, // customBaseUrl
       mcpServers,
       mcpServerToolRestrictions,
     );
 
     const callArgs = mockCreate.mock.calls[0][0];
     expect(callArgs.tool_choice).toBe("auto");
-    expect(callArgs.tools).toEqual([
-      {
-        type: "mcp",
-        server_label: "litellm",
-        server_url: "litellm_proxy/mcp/alpha",
-        require_approval: "never",
-        allowed_tools: ["toolA", "toolB"],
-      },
-      {
-        type: "mcp",
-        server_label: "litellm",
-        server_url: "litellm_proxy/mcp/Beta",
-        require_approval: "never",
-        allowed_tools: ["toolC"],
-      },
-    ]);
+    expect(callArgs.tools).toHaveLength(2);
+
+    // Check first tool
+    const firstTool = callArgs.tools[0];
+    expect(firstTool.type).toBe("mcp");
+    expect(firstTool.server_label).toBe("litellm");
+    expect(firstTool.server_url).toBe("litellm_proxy/mcp/alpha");
+    expect(firstTool.require_approval).toBe("never");
+    expect(firstTool.allowed_tools).toEqual(["toolA", "toolB"]);
+
+    // Check second tool
+    const secondTool = callArgs.tools[1];
+    expect(secondTool.type).toBe("mcp");
+    expect(secondTool.server_label).toBe("litellm");
+    expect(secondTool.server_url).toBe("litellm_proxy/mcp/Beta");
+    expect(secondTool.require_approval).toBe("never");
+    expect(secondTool.allowed_tools).toEqual(["toolC"]);
   });
 });

ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx

@@ -72,8 +72,25 @@ describe("responses_api", () => {
   it("should configure MCP tools per server with restrictions", async () => {
     const selectedMCPServers = ["server-1", "server-2"];
     const mcpServers = [
-      { server_id: "server-1", alias: "alpha", server_name: "Alpha" },
-      { server_id: "server-2", server_name: "Beta" },
+      {
+        server_id: "server-1",
+        alias: "alpha",
+        server_name: "Alpha",
+        url: "http://example.com",
+        created_at: "2024-01-01",
+        created_by: "test",
+        updated_at: "2024-01-01",
+        updated_by: "test",
+      },
+      {
+        server_id: "server-2",
+        server_name: "Beta",
+        url: "http://example.com",
+        created_at: "2024-01-01",
+        created_by: "test",
+        updated_at: "2024-01-01",
+        updated_by: "test",
+      },
     ];
     const mcpServerToolRestrictions: Record<string, string[]> = {
       "server-1": ["toolA"],
@@ -99,6 +116,7 @@ describe("responses_api", () => {
       undefined,
       undefined,
       undefined,
+      undefined,
       mcpServers,
       mcpServerToolRestrictions,
     );