diff --git a/.github/workflows/publish-migrations.yml b/.github/workflows/publish-migrations.yml index 8e5a67bcf85..a5187cb2f55 100644 --- a/.github/workflows/publish-migrations.yml +++ b/.github/workflows/publish-migrations.yml @@ -13,6 +13,7 @@ on: jobs: publish-migrations: + if: github.repository == 'BerriAI/litellm' runs-on: ubuntu-latest services: postgres: diff --git a/README.md b/README.md index a020bd80898..75a23faa5c1 100644 --- a/README.md +++ b/README.md @@ -262,6 +262,7 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature | Provider | `/chat/completions` | `/messages` | `/responses` | `/embeddings` | `/image/generations` | `/audio/transcriptions` | `/audio/speech` | `/moderations` | `/batches` | `/rerank` | |-------------------------------------------------------------------------------------|---------------------|-------------|--------------|---------------|----------------------|-------------------------|-----------------|----------------|-----------|-----------| +| [Abliteration (`abliteration`)](https://docs.litellm.ai/docs/providers/abliteration) | ✅ | | | | | | | | | | | [AI/ML API (`aiml`)](https://docs.litellm.ai/docs/providers/aiml) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | [AI21 (`ai21`)](https://docs.litellm.ai/docs/providers/ai21) | ✅ | ✅ | ✅ | | | | | | | | | [AI21 Chat (`ai21_chat`)](https://docs.litellm.ai/docs/providers/ai21) | ✅ | ✅ | ✅ | | | | | | | | @@ -455,4 +456,3 @@ All these checks must pass before your PR can be merged. - diff --git a/docs/my-website/docs/mcp_control.md b/docs/my-website/docs/mcp_control.md index a7d66a6b7fc..96c71ef9278 100644 --- a/docs/my-website/docs/mcp_control.md +++ b/docs/my-website/docs/mcp_control.md @@ -649,3 +649,16 @@ general_settings: ``` This is useful when you want discoverability for MCP offerings without granting additional execution privileges. + + +## Publish MCP Registry + +If you want other systems—for example external agent frameworks such as MCP-capable IDEs running outside your network—to automatically discover the MCP servers hosted on LiteLLM, you can expose a Model Context Protocol Registry endpoint. This registry lists the built-in LiteLLM MCP server and every server you have configured, using the [official MCP Registry spec](https://github.com/modelcontextprotocol/registry). + +1. Set `enable_mcp_registry: true` under `general_settings` in your proxy config (or DB settings) and restart the proxy. +2. LiteLLM will serve the registry at `GET /v1/mcp/registry.json`. +3. Each entry points to either `/mcp` (built-in server) or `/{mcp_server_name}/mcp` for your custom servers, so clients can connect directly using the advertised Streamable HTTP URL. + +:::note Permissions still apply +The registry only advertises server URLs. Actual access control is still enforced by LiteLLM when the client connects to `/mcp` or `/{server}/mcp`, so publishing the registry does not bypass per-key permissions. +::: diff --git a/docs/my-website/docs/observability/focus.md b/docs/my-website/docs/observability/focus.md new file mode 100644 index 00000000000..c282f4a220c --- /dev/null +++ b/docs/my-website/docs/observability/focus.md @@ -0,0 +1,93 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Focus Export (Experimental) + +:::caution Experimental feature +Focus Format export is under active development and currently considered experimental. +Interfaces, schema mappings, and configuration options may change as we iterate based on user feedback. +Please treat this integration as a preview and report any issues or suggestions to help us stabilize and improve the workflow. +::: + +LiteLLM can emit usage data in the [FinOps FOCUS format](https://focus.finops.org/focus-specification/v1-2/) and push artifacts (for example Parquet files) to destinations such as Amazon S3. This enables downstream cost-analysis tooling to ingest a standardised dataset directly from LiteLLM. + +LiteLLM currently conforms to the FinOps FOCUS v1.2 specification when emitting this dataset. + +## Overview + +| Property | Details | +|----------|---------| +| Destination | Export LiteLLM usage data in FOCUS format to managed storage (currently S3) | +| Callback name | `focus` | +| Supported operations | Automatic scheduled export | +| Data format | FOCUS Normalised Dataset (Parquet) | + +## Environment Variables + +### Common settings + +| Variable | Required | Description | +|----------|----------|-------------| +| `FOCUS_PROVIDER` | No | Destination provider (defaults to `s3`). | +| `FOCUS_FORMAT` | No | Output format (currently only `parquet`). | +| `FOCUS_FREQUENCY` | No | Export cadence. Prefer `hourly` or `daily` for production; `interval` is intended for short test loops. Defaults to `hourly`. | +| `FOCUS_CRON_OFFSET` | No | Minute offset used for hourly/daily cron triggers. Defaults to `5`. | +| `FOCUS_INTERVAL_SECONDS` | No | Interval (seconds) when `FOCUS_FREQUENCY="interval"`. | +| `FOCUS_PREFIX` | No | Object key prefix/folder. Defaults to `focus_exports`. | + +### S3 destination + +| Variable | Required | Description | +|----------|----------|-------------| +| `FOCUS_S3_BUCKET_NAME` | Yes | Destination bucket for exported files. | +| `FOCUS_S3_REGION_NAME` | No | AWS region for the bucket. | +| `FOCUS_S3_ENDPOINT_URL` | No | Custom endpoint (useful for S3-compatible storage). | +| `FOCUS_S3_ACCESS_KEY` | Yes | AWS access key for uploads. | +| `FOCUS_S3_SECRET_KEY` | Yes | AWS secret key for uploads. | +| `FOCUS_S3_SESSION_TOKEN` | No | AWS session token if using temporary credentials. | + +## Setup via Config + +### Configure environment variables + +```bash +export FOCUS_PROVIDER="s3" +export FOCUS_PREFIX="focus_exports" + +# S3 example +export FOCUS_S3_BUCKET_NAME="my-litellm-focus-bucket" +export FOCUS_S3_REGION_NAME="us-east-1" +export FOCUS_S3_ACCESS_KEY="AKIA..." +export FOCUS_S3_SECRET_KEY="..." +``` + +### Update LiteLLM config + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: sk-your-key + +litellm_settings: + callbacks: ["focus"] +``` + +### Start the proxy + +```bash +litellm --config /path/to/config.yaml +``` + +During boot LiteLLM registers the Focus logger and a background job that runs according to the configured frequency. + +## Planned Enhancements +- Add "Setup on UI" flow alongside the current configuration-based setup. +- Add GCS / Azure Blob to the Destination options. +- Support CSV output alongside Parquet. + +## Related Links + +- [Focus](https://focus.finops.org/) + diff --git a/docs/my-website/docs/observability/qualifire_integration.md b/docs/my-website/docs/observability/qualifire_integration.md new file mode 100644 index 00000000000..cf866f467bf --- /dev/null +++ b/docs/my-website/docs/observability/qualifire_integration.md @@ -0,0 +1,122 @@ +import Image from '@theme/IdealImage'; + +# Qualifire - LLM Evaluation, Guardrails & Observability + +[Qualifire](https://qualifire.ai/) provides real-time Agentic evaluations, guardrails and observability for production AI applications. + +**Key Features:** + +- **Evaluation** - Systematically assess AI behavior to detect hallucinations, jailbreaks, policy breaches, and other vulnerabilities +- **Guardrails** - Real-time interventions to prevent risks like brand damage, data leaks, and compliance breaches +- **Observability** - Complete tracing and logging for RAG pipelines, chatbots, and AI agents +- **Prompt Management** - Centralized prompt management with versioning and no-code studio + +:::tip + +Looking for Qualifire Guardrails? Check out the [Qualifire Guardrails Integration](../proxy/guardrails/qualifire.md) for real-time content moderation, prompt injection detection, PII checks, and more. + +::: + +## Pre-Requisites + +1. Create an account on [Qualifire](https://app.qualifire.ai/) +2. Get your API key and webhook URL from the Qualifire dashboard + +```bash +pip install litellm +``` + +## Quick Start + +Use just 2 lines of code to instantly log your responses **across all providers** with Qualifire. + +```python +litellm.callbacks = ["qualifire_eval"] +``` + +```python +import litellm +import os + +# Set Qualifire credentials +os.environ["QUALIFIRE_API_KEY"] = "your-qualifire-api-key" +os.environ["QUALIFIRE_WEBHOOK_URL"] = "https://your-qualifire-webhook-url" + +# LLM API Keys +os.environ['OPENAI_API_KEY'] = "your-openai-api-key" + +# Set qualifire_eval as a callback & LiteLLM will send the data to Qualifire +litellm.callbacks = ["qualifire_eval"] + +# OpenAI call +response = litellm.completion( + model="gpt-5", + messages=[ + {"role": "user", "content": "Hi 👋 - i'm openai"} + ] +) +``` + +## Using with LiteLLM Proxy + +1. Setup config.yaml + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + +litellm_settings: + callbacks: ["qualifire_eval"] + +general_settings: + master_key: "sk-1234" + +environment_variables: + QUALIFIRE_API_KEY: "your-qualifire-api-key" + QUALIFIRE_WEBHOOK_URL: "https://app.qualifire.ai/api/v1/webhooks/evaluations" +``` + +2. Start the proxy + +```bash +litellm --config config.yaml +``` + +3. Test it! + +```bash +curl -X POST 'http://0.0.0.0:4000/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-d '{ "model": "gpt-4o", "messages": [{"role": "user", "content": "Hi 👋 - i'm openai"}]}' +``` + +## Environment Variables + +| Variable | Description | +| ----------------------- | ------------------------------------------------------ | +| `QUALIFIRE_API_KEY` | Your Qualifire API key for authentication | +| `QUALIFIRE_WEBHOOK_URL` | The Qualifire webhook endpoint URL from your dashboard | + +## What Gets Logged? + +The [LiteLLM Standard Logging Payload](https://docs.litellm.ai/docs/proxy/logging_spec) is sent to your Qualifire endpoint on each successful LLM API call. + +This includes: + +- Request messages and parameters +- Response content and metadata +- Token usage statistics +- Latency metrics +- Model information +- Cost data + +Once data is in Qualifire, you can: + +- Run evaluations to detect hallucinations, toxicity, and policy violations +- Set up guardrails to block or modify responses in real-time +- View traces across your entire AI pipeline +- Track performance and quality metrics over time diff --git a/docs/my-website/docs/providers/abliteration.md b/docs/my-website/docs/providers/abliteration.md new file mode 100644 index 00000000000..a0fc7f39310 --- /dev/null +++ b/docs/my-website/docs/providers/abliteration.md @@ -0,0 +1,109 @@ +# Abliteration + +## Overview + +| Property | Details | +|-------|-------| +| Description | Abliteration provides an OpenAI-compatible `/chat/completions` endpoint. | +| Provider Route on LiteLLM | `abliteration/` | +| Link to Provider Doc | [Abliteration](https://abliteration.ai) | +| Base URL | `https://api.abliteration.ai/v1` | +| Supported Operations | [`/chat/completions`](#sample-usage) | + +
+ +## Required Variables + +```python showLineNumbers title="Environment Variables" +os.environ["ABLITERATION_API_KEY"] = "" # your Abliteration API key +``` + +## Sample Usage + +```python showLineNumbers title="Abliteration Completion" +import os +from litellm import completion + +os.environ["ABLITERATION_API_KEY"] = "" + +response = completion( + model="abliteration/abliterated-model", + messages=[{"role": "user", "content": "Hello from LiteLLM"}], +) + +print(response) +``` + +## Sample Usage - Streaming + +```python showLineNumbers title="Abliteration Streaming Completion" +import os +from litellm import completion + +os.environ["ABLITERATION_API_KEY"] = "" + +response = completion( + model="abliteration/abliterated-model", + messages=[{"role": "user", "content": "Stream a short reply"}], + stream=True, +) + +for chunk in response: + print(chunk) +``` + +## Usage with LiteLLM Proxy Server + +1. Add the model to your proxy config: + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: abliteration-chat + litellm_params: + model: abliteration/abliterated-model + api_key: os.environ/ABLITERATION_API_KEY +``` + +2. Start the proxy: + +```bash +litellm --config /path/to/config.yaml +``` + +## Direct API Usage (Bearer Token) + +Use the environment variable as a Bearer token against the OpenAI-compatible endpoint: +`https://api.abliteration.ai/v1/chat/completions`. + +```bash showLineNumbers title="cURL" +export ABLITERATION_API_KEY="" +curl https://api.abliteration.ai/v1/chat/completions \ + -H "Authorization: Bearer ${ABLITERATION_API_KEY}" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "abliterated-model", + "messages": [{"role": "user", "content": "Hello from Abliteration"}] + }' +``` + +```python showLineNumbers title="Python (requests)" +import os +import requests + +api_key = os.environ["ABLITERATION_API_KEY"] + +response = requests.post( + "https://api.abliteration.ai/v1/chat/completions", + headers={ + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + }, + json={ + "model": "abliterated-model", + "messages": [{"role": "user", "content": "Hello from Abliteration"}], + }, + timeout=60, +) + +print(response.json()) +``` diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 33ebf535d29..f46608aa57c 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -35,6 +35,8 @@ import json # !gcloud auth application-default login - run this to add vertex credentials to your env ## OR ## file_path = 'path/to/vertex_ai_service_account.json' +## OR ## +export VERTEXAI_API_KEY="your-api-key" # Load the JSON file with open(file_path, 'r') as file: @@ -47,7 +49,7 @@ vertex_credentials_json = json.dumps(vertex_credentials) response = completion( model="vertex_ai/gemini-2.5-pro", messages=[{ "content": "Hello, how are you?","role": "user"}], - vertex_credentials=vertex_credentials_json + vertex_credentials=vertex_credentials_json # Can remove this is added VERTEXAI_API_KEY in env ) ``` @@ -1329,15 +1331,41 @@ Here's how to use Vertex AI with the LiteLLM Proxy Server ## Authentication - vertex_project, vertex_location, etc. +LiteLLM supports two authentication methods for Vertex AI: + +1. **API Key Authentication** (Recommended for getting started) +2. **Service Account Credentials** (Recommended for production) + Set your vertex credentials via: - dynamic params OR - env vars +### **Authentication Method 1: -### **Dynamic Params** +The simplest way to authenticate with Vertex AI. You can set: +- `api_key` (str) - Your Vertex AI API key -You can set: +**Environment Variables:** +```bash +export VERTEXAI_API_KEY="your-api-key" +``` + +**Or pass as parameters:** +```python +from litellm import completion + +response = completion( + model="vertex_ai/gemini-2.0-flash-exp", + messages=[{"role": "user", "content": "Hello!"}], + api_key="your-vertex-api-key", + +) +``` + +### **Authentication Method 2: Service Account Credentials** + +For production environments with fine-grained access control. You can set: - `vertex_credentials` (str) - can be a json string or filepath to your vertex ai service account.json - `vertex_location` (str) - place where vertex model is deployed (us-central1, asia-southeast1, etc.). Some models support the global location, please see [Vertex AI documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#supported_models) - `vertex_project` Optional[str] - use if vertex project different from the one in vertex_credentials @@ -1392,7 +1420,16 @@ model_list: ### **Environment Variables** -You can set: +#### For API Key Authentication: + +- `VERTEXAI_API_KEY` or `VERTEX_API_KEY` - Your Vertex AI API key + +```bash +export VERTEXAI_API_KEY="your-vertex-api-key" +``` + +#### For Service Account Authentication: + - `GOOGLE_APPLICATION_CREDENTIALS` - store the filepath for your service_account.json in here (used by vertex sdk directly). - VERTEXAI_LOCATION - place where vertex model is deployed (us-central1, asia-southeast1, etc.) - VERTEXAI_PROJECT - Optional[str] - use if vertex project different from the one in vertex_credentials diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 68e1f629e8b..42494c1e5f5 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/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 diff --git a/docs/my-website/docs/proxy/guardrails/qualifire.md b/docs/my-website/docs/proxy/guardrails/qualifire.md index 66961c92d9d..850af37e47f 100644 --- a/docs/my-website/docs/proxy/guardrails/qualifire.md +++ b/docs/my-website/docs/proxy/guardrails/qualifire.md @@ -8,13 +8,7 @@ Use [Qualifire](https://qualifire.ai) to evaluate LLM outputs for quality, safet ## Quick Start -### 1. Install the Qualifire SDK - -```bash -pip install qualifire -``` - -### 2. Define Guardrails on your LiteLLM config.yaml +### 1. Define Guardrails on your LiteLLM config.yaml Define your guardrails under the `guardrails` section: @@ -61,13 +55,13 @@ guardrails: - `post_call` Run **after** LLM call, on **input & output** - `during_call` Run **during** LLM call, on **input**. Same as `pre_call` but runs in parallel as LLM call. Response not returned until guardrail check completes -### 3. Start LiteLLM Gateway +### 2. Start LiteLLM Gateway ```shell litellm --config config.yaml --detailed_debug ``` -### 4. Test request +### 3. Test request **[Langchain, OpenAI SDK Usage Examples](../proxy/user_keys#request-format)** @@ -142,7 +136,7 @@ guardrails: evaluation_id: eval_abc123 # Your evaluation ID from Qualifire dashboard ``` -When `evaluation_id` is provided, LiteLLM will use `invoke_evaluation()` instead of `evaluate()`, running the pre-configured evaluation from your dashboard. +When `evaluation_id` is provided, LiteLLM will use the invoke evaluation API endpoint instead of the evaluate endpoint, running the pre-configured evaluation from your dashboard. ## Available Checks @@ -213,19 +207,19 @@ guardrails: ### Parameter Reference -| Parameter | Type | Default | Description | -| ------------------------------ | ----------- | --------------------------- | -------------------------------------------------------- | -| `api_key` | `str` | `QUALIFIRE_API_KEY` env var | Your Qualifire API key | -| `api_base` | `str` | `None` | Custom API base URL (optional) | -| `evaluation_id` | `str` | `None` | Pre-configured evaluation ID from Qualifire dashboard | -| `prompt_injections` | `bool` | `true` (if no other checks) | Enable prompt injection detection | -| `hallucinations_check` | `bool` | `None` | Enable hallucination detection | -| `grounding_check` | `bool` | `None` | Enable grounding verification | -| `pii_check` | `bool` | `None` | Enable PII detection | -| `content_moderation_check` | `bool` | `None` | Enable content moderation | -| `tool_selection_quality_check` | `bool` | `None` | Enable tool selection quality check | -| `assertions` | `List[str]` | `None` | Custom assertions to validate | -| `on_flagged` | `str` | `"block"` | Action when content is flagged: `"block"` or `"monitor"` | +| Parameter | Type | Default | Description | +| ------------------------------ | ----------- | ---------------------------- | -------------------------------------------------------- | +| `api_key` | `str` | `QUALIFIRE_API_KEY` env var | Your Qualifire API key | +| `api_base` | `str` | `https://proxy.qualifire.ai` | Custom API base URL (optional) | +| `evaluation_id` | `str` | `None` | Pre-configured evaluation ID from Qualifire dashboard | +| `prompt_injections` | `bool` | `true` (if no other checks) | Enable prompt injection detection | +| `hallucinations_check` | `bool` | `None` | Enable hallucination detection | +| `grounding_check` | `bool` | `None` | Enable grounding verification | +| `pii_check` | `bool` | `None` | Enable PII detection | +| `content_moderation_check` | `bool` | `None` | Enable content moderation | +| `tool_selection_quality_check` | `bool` | `None` | Enable tool selection quality check | +| `assertions` | `List[str]` | `None` | Custom assertions to validate | +| `on_flagged` | `str` | `"block"` | Action when content is flagged: `"block"` or `"monitor"` | ### Default Behavior @@ -261,4 +255,3 @@ This evaluates whether the LLM selected the appropriate tools and provided corre - [Qualifire Documentation](https://docs.qualifire.ai) - [Qualifire Dashboard](https://app.qualifire.ai) -- [Qualifire Python SDK](https://github.com/qualifire-dev/qualifire-python-sdk) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 488fd616678..2e8ea07ab75 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -55,6 +55,7 @@ const sidebars = { "proxy/guardrails/test_playground", "proxy/guardrails/litellm_content_filter", ...[ + "proxy/guardrails/qualifire", "proxy/guardrails/aim_security", "proxy/guardrails/onyx_security", "proxy/guardrails/aporia_api", @@ -653,12 +654,13 @@ const sidebars = { "providers/bedrock_writer", "providers/bedrock_batches", "providers/aws_polly", - "providers/bedrock_vector_store", - ] - }, - "providers/litellm_proxy", - "providers/ai21", - "providers/aiml", + "providers/bedrock_vector_store", + ] + }, + "providers/litellm_proxy", + "providers/abliteration", + "providers/ai21", + "providers/aiml", "providers/aleph_alpha", "providers/amazon_nova", "providers/anyscale", diff --git a/litellm/__init__.py b/litellm/__init__.py index 0c019457027..4c43ae60aa6 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -134,6 +134,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "bitbucket", "gitlab", "cloudzero", + "focus", "posthog", "levo", ] diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 364fa3f5def..585de510e8b 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -225,10 +225,13 @@ class BraintrustLogger(CustomLogger): "id": litellm_call_id, "input": prompt["messages"], "metadata": standard_logging_object, - "tags": tags, "span_attributes": {"name": span_name, "type": "llm"}, } - + + # Braintrust cannot specify 'tags' for non-root spans + if dynamic_metadata.get("root_span_id") is None: + request_data["tags"] = tags + # Only add those that are not None (or falsy) for key, value in span_attributes.items(): if value: @@ -351,14 +354,37 @@ class BraintrustLogger(CustomLogger): # Allow metadata override for span name span_name = dynamic_metadata.get("span_name", "Chat Completion") + # Span parents is a special case + span_parents = dynamic_metadata.get("span_parents") + + # Convert comma-separated string to list if present + if span_parents: + span_parents = [s.strip() for s in span_parents.split(",") if s.strip()] + + # Add optional span attributes only if present + span_attributes = { + "span_id": dynamic_metadata.get("span_id"), + "root_span_id": dynamic_metadata.get("root_span_id"), + "span_parents": span_parents, + } + request_data = { "id": litellm_call_id, "input": prompt["messages"], "output": output, "metadata": standard_logging_object, - "tags": tags, "span_attributes": {"name": span_name, "type": "llm"}, } + + # Braintrust cannot specify 'tags' for non-root spans + if dynamic_metadata.get("root_span_id") is None: + request_data["tags"] = tags + + # Only add those that are not None (or falsy) + for key, value in span_attributes.items(): + if value: + request_data[key] = value + if choices is not None: request_data["output"] = [choice.dict() for choice in choices] else: @@ -367,9 +393,6 @@ class BraintrustLogger(CustomLogger): if metrics is not None: request_data["metrics"] = metrics - if metrics is not None: - request_data["metrics"] = metrics - try: await self.global_braintrust_http_handler.post( url=f"{self.api_base}/project_logs/{project_id}/insert", diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 2128b55bf83..71929398103 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -19,7 +19,7 @@ """Database connection and data extraction for LiteLLM.""" from datetime import datetime -from typing import Any, Dict, Optional +from typing import Any, Optional, List import polars as pl @@ -46,19 +46,9 @@ class LiteLLMDatabase: """Retrieve usage data from LiteLLM daily user spend table.""" client = self._ensure_prisma_client() - # Build WHERE clause for time filtering - where_conditions = [] - if start_time_utc: - where_conditions.append(f"dus.updated_at >= '{start_time_utc.isoformat()}'") - if end_time_utc: - where_conditions.append(f"dus.updated_at <= '{end_time_utc.isoformat()}'") - - where_clause = "" - if where_conditions: - where_clause = "WHERE " + " AND ".join(where_conditions) - - # Query to get user spend data with team information - query = f""" + # Query to get user spend data with team information. Use parameter binding to + # avoid SQL injection from user-supplied timestamps or limits. + query = """ SELECT dus.id, dus.date, @@ -85,163 +75,27 @@ class LiteLLMDatabase: LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id - {where_clause} + WHERE ($1::timestamptz IS NULL OR dus.updated_at >= $1::timestamptz) + AND ($2::timestamptz IS NULL OR dus.updated_at <= $2::timestamptz) ORDER BY dus.date DESC, dus.created_at DESC """ - if limit: - query += f" LIMIT {limit}" + params: List[Any] = [ + start_time_utc, + end_time_utc, + ] + + if limit is not None: + try: + params.append(int(limit)) + except (TypeError, ValueError): + raise ValueError("limit must be an integer") + query += " LIMIT $3" try: - db_response = await client.db.query_raw(query) + db_response = await client.db.query_raw(query, *params) # Convert the response to polars DataFrame with full schema inference # This prevents schema mismatch errors when data types vary across rows return pl.DataFrame(db_response, infer_schema_length=None) except Exception as e: raise Exception(f"Error retrieving usage data: {str(e)}") - - async def get_table_info(self) -> Dict[str, Any]: - """Get information about the daily user spend table.""" - client = self._ensure_prisma_client() - - try: - # Get row count from user spend table - user_count = await self._get_table_row_count("LiteLLM_DailyUserSpend") - - # Get column structure from user spend table - query = """ - SELECT column_name, data_type, is_nullable - FROM information_schema.columns - WHERE table_name = 'LiteLLM_DailyUserSpend' - ORDER BY ordinal_position; - """ - columns_response = await client.db.query_raw(query) - - return { - "columns": columns_response, - "row_count": user_count, - "table_name": "LiteLLM_DailyUserSpend", - } - except Exception as e: - raise Exception(f"Error getting table info: {str(e)}") - - async def _get_table_row_count(self, table_name: str) -> int: - """Get row count from specified table.""" - client = self._ensure_prisma_client() - - try: - query = f'SELECT COUNT(*) as count FROM "{table_name}"' - response = await client.db.query_raw(query) - - if response and len(response) > 0: - return response[0].get("count", 0) - return 0 - except Exception: - return 0 - - async def discover_all_tables(self) -> Dict[str, Any]: - """Discover all tables in the LiteLLM database and their schemas.""" - client = self._ensure_prisma_client() - - try: - # Get all LiteLLM tables - litellm_tables_query = """ - SELECT table_name - FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name LIKE 'LiteLLM_%' - ORDER BY table_name; - """ - tables_response = await client.db.query_raw(litellm_tables_query) - table_names = [row["table_name"] for row in tables_response] - - # Get detailed schema for each table - tables_info = {} - for table_name in table_names: - # Get column information - columns_query = """ - SELECT - column_name, - data_type, - is_nullable, - column_default, - character_maximum_length, - numeric_precision, - numeric_scale, - ordinal_position - FROM information_schema.columns - WHERE table_name = $1 - AND table_schema = 'public' - ORDER BY ordinal_position; - """ - columns_response = await client.db.query_raw(columns_query, table_name) - - # Get primary key information - pk_query = """ - SELECT a.attname - FROM pg_index i - JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) - WHERE i.indrelid = $1::regclass AND i.indisprimary; - """ - pk_response = await client.db.query_raw(pk_query, f'"{table_name}"') - primary_keys = ( - [row["attname"] for row in pk_response] if pk_response else [] - ) - - # Get foreign key information - fk_query = """ - SELECT - tc.constraint_name, - kcu.column_name, - ccu.table_name AS foreign_table_name, - ccu.column_name AS foreign_column_name - FROM information_schema.table_constraints AS tc - JOIN information_schema.key_column_usage AS kcu - ON tc.constraint_name = kcu.constraint_name - JOIN information_schema.constraint_column_usage AS ccu - ON ccu.constraint_name = tc.constraint_name - WHERE tc.constraint_type = 'FOREIGN KEY' - AND tc.table_name = $1; - """ - fk_response = await client.db.query_raw(fk_query, table_name) - foreign_keys = fk_response if fk_response else [] - - # Get indexes - indexes_query = """ - SELECT - i.relname AS index_name, - array_agg(a.attname ORDER BY a.attnum) AS column_names, - ix.indisunique AS is_unique - FROM pg_class t - JOIN pg_index ix ON t.oid = ix.indrelid - JOIN pg_class i ON i.oid = ix.indexrelid - JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey) - WHERE t.relname = $1 - AND t.relkind = 'r' - GROUP BY i.relname, ix.indisunique - ORDER BY i.relname; - """ - indexes_response = await client.db.query_raw(indexes_query, table_name) - indexes = indexes_response if indexes_response else [] - - # Get row count - try: - row_count = await self._get_table_row_count(table_name) - except Exception: - row_count = 0 - - tables_info[table_name] = { - "columns": columns_response, - "primary_keys": primary_keys, - "foreign_keys": foreign_keys, - "indexes": indexes, - "row_count": row_count, - } - - return { - "tables": tables_info, - "table_count": len(table_names), - "table_names": table_names, - } - except Exception as e: - raise Exception(f"Error discovering tables: {str(e)}") diff --git a/litellm/integrations/focus/__init__.py b/litellm/integrations/focus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py new file mode 100644 index 00000000000..298254670eb --- /dev/null +++ b/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 diff --git a/litellm/integrations/focus/destinations/__init__.py b/litellm/integrations/focus/destinations/__init__.py new file mode 100644 index 00000000000..233f1da0c9b --- /dev/null +++ b/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", +] diff --git a/litellm/integrations/focus/destinations/base.py b/litellm/integrations/focus/destinations/base.py new file mode 100644 index 00000000000..8042a7e23b9 --- /dev/null +++ b/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.""" + ... diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py new file mode 100644 index 00000000000..cb7696a11de --- /dev/null +++ b/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" + ) diff --git a/litellm/integrations/focus/destinations/s3_destination.py b/litellm/integrations/focus/destinations/s3_destination.py new file mode 100644 index 00000000000..c6d5554b438 --- /dev/null +++ b/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", + ) diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py new file mode 100644 index 00000000000..22ebce2a168 --- /dev/null +++ b/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) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py new file mode 100644 index 00000000000..ade1cf861b1 --- /dev/null +++ b/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"] diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py new file mode 100644 index 00000000000..ac2f33dad0a --- /dev/null +++ b/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"] diff --git a/litellm/integrations/focus/serializers/__init__.py b/litellm/integrations/focus/serializers/__init__.py new file mode 100644 index 00000000000..18187bf73e5 --- /dev/null +++ b/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"] diff --git a/litellm/integrations/focus/serializers/base.py b/litellm/integrations/focus/serializers/base.py new file mode 100644 index 00000000000..6da080dae81 --- /dev/null +++ b/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 diff --git a/litellm/integrations/focus/serializers/parquet.py b/litellm/integrations/focus/serializers/parquet.py new file mode 100644 index 00000000000..6b3dde5903d --- /dev/null +++ b/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() diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py new file mode 100644 index 00000000000..cac12b7be14 --- /dev/null +++ b/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"), + ) diff --git a/litellm/integrations/generic_api/generic_api_compatible_callbacks.json b/litellm/integrations/generic_api/generic_api_compatible_callbacks.json index 12dc4ae643c..13fe79ae671 100644 --- a/litellm/integrations/generic_api/generic_api_compatible_callbacks.json +++ b/litellm/integrations/generic_api/generic_api_compatible_callbacks.json @@ -1,28 +1,37 @@ { - "sample_callback": { - "event_types": ["llm_api_success", "llm_api_failure"], - "endpoint": "{{environment_variables.SAMPLE_CALLBACK_URL}}", - "headers": { - "Content-Type": "application/json", - "Authorization": "Bearer {{environment_variables.SAMPLE_CALLBACK_API_KEY}}" - }, - "environment_variables": ["SAMPLE_CALLBACK_URL", "SAMPLE_CALLBACK_API_KEY"] + "sample_callback": { + "event_types": ["llm_api_success", "llm_api_failure"], + "endpoint": "{{environment_variables.SAMPLE_CALLBACK_URL}}", + "headers": { + "Content-Type": "application/json", + "Authorization": "Bearer {{environment_variables.SAMPLE_CALLBACK_API_KEY}}" }, - "rubrik": { - "event_types": ["llm_api_success"], - "endpoint": "{{environment_variables.RUBRIK_WEBHOOK_URL}}", - "headers": { - "Content-Type": "application/json", - "Authorization": "Bearer {{environment_variables.RUBRIK_API_KEY}}" - }, - "environment_variables": ["RUBRIK_API_KEY", "RUBRIK_WEBHOOK_URL"] + "environment_variables": ["SAMPLE_CALLBACK_URL", "SAMPLE_CALLBACK_API_KEY"] + }, + "rubrik": { + "event_types": ["llm_api_success"], + "endpoint": "{{environment_variables.RUBRIK_WEBHOOK_URL}}", + "headers": { + "Content-Type": "application/json", + "Authorization": "Bearer {{environment_variables.RUBRIK_API_KEY}}" }, - "sumologic": { - "endpoint": "{{environment_variables.SUMOLOGIC_WEBHOOK_URL}}", - "headers": { - "Content-Type": "application/json" - }, - "environment_variables": ["SUMOLOGIC_WEBHOOK_URL"], - "log_format": "ndjson" - } -} \ No newline at end of file + "environment_variables": ["RUBRIK_API_KEY", "RUBRIK_WEBHOOK_URL"] + }, + "sumologic": { + "endpoint": "{{environment_variables.SUMOLOGIC_WEBHOOK_URL}}", + "headers": { + "Content-Type": "application/json" + }, + "environment_variables": ["SUMOLOGIC_WEBHOOK_URL"], + "log_format": "ndjson" + }, + "qualifire_eval": { + "event_types": ["llm_api_success"], + "endpoint": "{{environment_variables.QUALIFIRE_WEBHOOK_URL}}", + "headers": { + "Content-Type": "application/json", + "X-Qualifire-API-Key": "{{environment_variables.QUALIFIRE_API_KEY}}" + }, + "environment_variables": ["QUALIFIRE_API_KEY", "QUALIFIRE_WEBHOOK_URL"] + } +} diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index e11c8737a81..9eb3cd72b0b 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -45,6 +45,7 @@ def _get_cached_end_user_id_for_cost_tracking(): global _get_end_user_id_for_cost_tracking if _get_end_user_id_for_cost_tracking is None: from litellm.utils import get_end_user_id_for_cost_tracking + _get_end_user_id_for_cost_tracking = get_end_user_id_for_cost_tracking return _get_end_user_id_for_cost_tracking @@ -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", diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 47cbcb8aec9..a3c25ab65e9 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/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, } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ab55022f8ce..e0a799d8e5c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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( diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0c331e43038..d8e82199272 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1645,9 +1645,12 @@ def convert_to_anthropic_tool_result( ) elif content["type"] == "image_url": format = content["image_url"].get("format") if isinstance(content["image_url"], dict) else None - anthropic_content_list.append( - create_anthropic_image_param(content["image_url"], format=format) + _anthropic_image_param = create_anthropic_image_param(content["image_url"], format=format) + _anthropic_image_param = add_cache_control_to_content( + anthropic_content_element=_anthropic_image_param, + original_content_element=content, ) + anthropic_content_list.append(_anthropic_image_param) anthropic_content = anthropic_content_list anthropic_tool_result: Optional[AnthropicMessagesToolResultParam] = None diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 87f81d117f0..506b7fdfe5e 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -25,7 +25,24 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): return "gpt-5" in model or "gpt5_series" in model def get_supported_openai_params(self, model: str) -> List[str]: - return OpenAIGPT5Config.get_supported_openai_params(self, model=model) + """Get supported parameters for Azure OpenAI GPT-5 models. + + Azure OpenAI GPT-5.2 models support logprobs, unlike OpenAI's GPT-5. + This overrides the parent class to add logprobs support back for gpt-5.2. + + Reference: + - Tested with Azure OpenAI GPT-5.2 (api-version: 2025-01-01-preview) + - Azure returns logprobs successfully despite Microsoft's general + documentation stating reasoning models don't support it. + """ + params = OpenAIGPT5Config.get_supported_openai_params(self, model=model) + + # Only gpt-5.2 has been verified to support logprobs on Azure + if self.is_model_gpt_5_2_model(model): + azure_supported_params = ["logprobs", "top_logprobs"] + params.extend(azure_supported_params) + + return params def map_openai_params( self, diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 0d5494541ec..e9cea23ea4a 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -314,6 +314,12 @@ class BaseAWSLLM: if model.startswith("invoke/"): model = model.replace("invoke/", "", 1) + # Special case: Check for "nova" in model name first (before "amazon") + # This handles amazon.nova-* models which would otherwise match "amazon" (Titan) + if "nova" in model.lower(): + if "nova" in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL): + return cast(BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova") + _split_model = model.split(".")[0] if _split_model in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL): return cast(BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model) @@ -323,13 +329,9 @@ class BaseAWSLLM: if provider is not None: return provider - # check if provider == "nova" - if "nova" in model: - return "nova" - else: - for provider in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL): - if provider in model: - return provider + for provider in get_args(BEDROCK_INVOKE_PROVIDERS_LITERAL): + if provider in model: + return provider return None @staticmethod @@ -364,7 +366,7 @@ class BaseAWSLLM: elif provider == "qwen3" and "qwen3/" in model_id: model_id = BaseAWSLLM._get_model_id_from_model_with_spec( model_id, spec="qwen3" - ) + ) elif provider == "stability" and "stability/" in model_id: model_id = BaseAWSLLM._get_model_id_from_model_with_spec( model_id, spec="stability" @@ -416,7 +418,7 @@ class BaseAWSLLM: if "nova" in model.lower(): if "nova" in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL): return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, "nova") - + # Handle regional models like us.twelvelabs.marengo-embed-2-7-v1:0 if "." in model: parts = model.split(".") @@ -962,7 +964,9 @@ class BaseAWSLLM: return endpoint_url, proxy_endpoint_url def _select_default_endpoint_url( - self, endpoint_type: Optional[Literal["runtime", "agent", "agentcore"]], aws_region_name: str + self, + endpoint_type: Optional[Literal["runtime", "agent", "agentcore"]], + aws_region_name: str, ) -> str: """ Select the default endpoint url based on the endpoint type diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index c602b71fe05..cf8aee6954b 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -524,6 +524,12 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): if model.startswith("invoke/"): model = model.replace("invoke/", "", 1) + # Special case: Check for "nova" in model name first (before "amazon") + # This handles amazon.nova-* models which would otherwise match "amazon" (Titan) + if "nova" in model.lower(): + if "nova" in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): + return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova") + _split_model = model.split(".")[0] if _split_model in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model) @@ -533,10 +539,6 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): if provider is not None: return provider - # check if provider == "nova" - if "nova" in model: - return "nova" - for provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL): if provider in model: return provider diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9edfe320fb2..5cb51cf994f 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -15,7 +15,7 @@ import litellm from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) -from litellm.llms.base_llm.base_utils import BaseLLMModelInfo +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret @@ -359,6 +359,70 @@ def get_bedrock_tool_name(response_tool_name: str) -> str: return response_tool_name +# Cache the global regions list at module level +_BEDROCK_GLOBAL_REGIONS: Optional[List[str]] = None + + +def _get_all_bedrock_regions() -> List[str]: + """Get all Bedrock regions, cached at module level.""" + global _BEDROCK_GLOBAL_REGIONS + if _BEDROCK_GLOBAL_REGIONS is None: + _BEDROCK_GLOBAL_REGIONS = AmazonBedrockGlobalConfig().get_all_regions() + return _BEDROCK_GLOBAL_REGIONS + + +def get_bedrock_cross_region_inference_regions() -> List[str]: + """Abbreviations of regions AWS Bedrock supports for cross region inference.""" + return ["global", "us", "eu", "apac", "jp", "au", "us-gov"] + + +def extract_model_name_from_bedrock_arn(model: str) -> str: + """ + Extract the model name from an AWS Bedrock ARN. + Returns the string after the last '/' if 'arn' is in the input string. + """ + if "arn" in model.lower(): + return model.split("/")[-1] + return model + + +def strip_bedrock_routing_prefix(model: str) -> str: + """Strip LiteLLM routing prefixes from model name.""" + for prefix in ["bedrock/", "converse/", "invoke/", "openai/"]: + if model.startswith(prefix): + model = model.split("/", 1)[1] + return model + + +def get_bedrock_base_model(model: str) -> str: + """ + Get the base model from the given model name. + + Handle model names like: + - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1" + - "bedrock/converse/model" -> "model" + """ + model = strip_bedrock_routing_prefix(model) + model = extract_model_name_from_bedrock_arn(model) + + potential_region = model.split(".", 1)[0] + alt_potential_region = model.split("/", 1)[0] + + if potential_region in get_bedrock_cross_region_inference_regions(): + return model.split(".", 1)[1] + elif ( + alt_potential_region in _get_all_bedrock_regions() + and len(model.split("/", 1)) > 1 + ): + return model.split("/", 1)[1] + + return model + + +# Import after standalone functions to avoid circular imports +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter + + class BedrockModelInfo(BaseLLMModelInfo): global_config = AmazonBedrockGlobalConfig() all_global_regions = global_config.get_all_regions() @@ -394,76 +458,34 @@ class BedrockModelInfo(BaseLLMModelInfo): ) -> List[str]: return [] - @staticmethod - def extract_model_name_from_arn(model: str) -> str: + def get_token_counter(self) -> Optional[BaseTokenCounter]: """ - Extract the model name from an AWS Bedrock ARN. - Returns the string after the last '/' if 'arn' is in the input string. - - Args: - arn (str): The ARN string to parse + Factory method to create a Bedrock token counter. Returns: - str: The extracted model name if 'arn' is in the string, - otherwise returns the original string + BedrockTokenCounter instance for this provider. """ - if "arn" in model.lower(): - return model.split("/")[-1] - return model + return BedrockTokenCounter() + + @staticmethod + def extract_model_name_from_arn(model: str) -> str: + """Wrapper for standalone function. See extract_model_name_from_bedrock_arn().""" + return extract_model_name_from_bedrock_arn(model) @staticmethod def get_non_litellm_routing_model_name(model: str) -> str: - if model.startswith("bedrock/"): - model = model.split("/", 1)[1] - - if model.startswith("converse/"): - model = model.split("/", 1)[1] - - if model.startswith("invoke/"): - model = model.split("/", 1)[1] - - if model.startswith("openai/"): - model = model.split("/", 1)[1] - - return model + """Wrapper for standalone function. See strip_bedrock_routing_prefix().""" + return strip_bedrock_routing_prefix(model) @staticmethod def get_base_model(model: str) -> str: - """ - Get the base model from the given model name. - - Handle model names like - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1" - AND "meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1" - """ - - model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model) - model = BedrockModelInfo.extract_model_name_from_arn(model) - - potential_region = model.split(".", 1)[0] - - alt_potential_region = model.split("/", 1)[ - 0 - ] # in model cost map we store regional information like `/us-west-2/bedrock-model` - - if ( - potential_region - in BedrockModelInfo._supported_cross_region_inference_region() - ): - return model.split(".", 1)[1] - elif ( - alt_potential_region in BedrockModelInfo.all_global_regions - and len(model.split("/", 1)) > 1 - ): - return model.split("/", 1)[1] - - return model + """Wrapper for standalone function. See get_bedrock_base_model().""" + return get_bedrock_base_model(model) @staticmethod def _supported_cross_region_inference_region() -> List[str]: - """ - Abbreviations of regions AWS Bedrock supports for cross region inference - """ - return ["global", "us", "eu", "apac", "jp", "au", "us-gov"] + """Wrapper for standalone function. See get_bedrock_cross_region_inference_regions().""" + return get_bedrock_cross_region_inference_regions() @staticmethod def get_bedrock_route( diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py new file mode 100644 index 00000000000..b680bd046ef --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -0,0 +1,87 @@ +""" +Bedrock Token Counter implementation using the CountTokens API. +""" + +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_logger +from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.bedrock.common_utils import get_bedrock_base_model +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.types.utils import LlmProviders, TokenCountResponse + + +class BedrockTokenCounter(BaseTokenCounter): + """Token counter implementation for AWS Bedrock provider using the CountTokens API.""" + + def should_use_token_counting_api( + self, + custom_llm_provider: Optional[str] = None, + ) -> bool: + """ + Returns True if we should use the Bedrock CountTokens API for token counting. + """ + return custom_llm_provider == LlmProviders.BEDROCK.value + + async def count_tokens( + self, + model_to_use: str, + messages: Optional[List[Dict[str, Any]]], + contents: Optional[List[Dict[str, Any]]], + deployment: Optional[Dict[str, Any]] = None, + request_model: str = "", + ) -> Optional[TokenCountResponse]: + """ + Count tokens using AWS Bedrock's CountTokens API. + + This method calls the existing BedrockCountTokensHandler to make an API call + to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting. + + Args: + model_to_use: The model identifier + messages: The messages to count tokens for + contents: Alternative content format (not used for Bedrock) + deployment: Deployment configuration containing litellm_params + request_model: The original request model name + + Returns: + TokenCountResponse with token count, or None if counting fails + """ + if not messages: + return None + + deployment = deployment or {} + litellm_params = deployment.get("litellm_params", {}) + + # Build request data in the format expected by BedrockCountTokensHandler + request_data = { + "model": model_to_use, + "messages": messages, + } + + # Get the resolved model (strip prefixes like bedrock/, converse/, etc.) + resolved_model = get_bedrock_base_model(model_to_use) + + try: + handler = BedrockCountTokensHandler() + result = await handler.handle_count_tokens_request( + request_data=request_data, + litellm_params=litellm_params, + resolved_model=resolved_model, + ) + + # Transform response to TokenCountResponse + if result is not None: + return TokenCountResponse( + total_tokens=result.get("input_tokens", 0), + request_model=request_model, + model_used=model_to_use, + tokenizer_type="bedrock_api", + original_response=result, + ) + except Exception as e: + verbose_logger.warning( + f"Error calling Bedrock CountTokens API: {e}, falling back to default tokenizer" + ) + + return None diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index d4355c0c360..60ace7f3369 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/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) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index d46ed3aa452..b313cc9df3c 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -8,7 +8,7 @@ to AWS Bedrock's CountTokens API format and vice versa. from typing import Any, Dict, List from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM -from litellm.llms.bedrock.common_utils import BedrockModelInfo +from litellm.llms.bedrock.common_utils import get_bedrock_base_model class BedrockCountTokensConfig(BaseAWSLLM): @@ -141,7 +141,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): Complete endpoint URL for CountTokens API """ # Use existing LiteLLM function to get the base model ID (removes region prefix) - model_id = BedrockModelInfo.get_base_model(model) + model_id = get_bedrock_base_model(model) # Remove bedrock/ prefix if present if model_id.startswith("bedrock/"): diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index e53829d3329..30c5b4f17c5 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -150,6 +150,15 @@ def get_api_key_from_env() -> Optional[str]: return get_secret_str("GOOGLE_API_KEY") or get_secret_str("GEMINI_API_KEY") +def get_vertex_api_key_from_env() -> Optional[str]: + """ + Get API key from environment for Vertex AI. + Checks VERTEXAI_API_KEY and VERTEX_API_KEY environment variables. + This allows using Vertex AI with API keys instead of service account credentials. + """ + return get_secret_str("VERTEXAI_API_KEY") or get_secret_str("VERTEX_API_KEY") + + class GoogleAIStudioTokenCounter(BaseTokenCounter): """Token counter implementation for Google AI Studio provider.""" def should_use_token_counting_api( diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 206aee1359d..bda3684a8a8 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -61,6 +61,10 @@ "max_completion_tokens": "max_tokens" } }, + "abliteration": { + "base_url": "https://api.abliteration.ai/v1", + "api_key_env": "ABLITERATION_API_KEY" + }, "llamagate": { "base_url": "https://api.llamagate.dev/v1", "api_key_env": "LLAMAGATE_API_KEY", diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py new file mode 100644 index 00000000000..d1d0e911d16 --- /dev/null +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -0,0 +1,182 @@ +""" +OpenRouter Embedding API Configuration. + +This module provides the configuration for OpenRouter's Embedding API. +OpenRouter is OpenAI-compatible and supports embeddings via the /v1/embeddings endpoint. + +Docs: https://openrouter.ai/docs +""" +from typing import TYPE_CHECKING, Any, Optional + +import httpx + +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.types.llms.openai import AllEmbeddingInputValues +from litellm.types.utils import EmbeddingResponse +from litellm.utils import convert_to_model_response_object + +from ..common_utils import OpenRouterException + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): + """ + Configuration for OpenRouter's Embedding API. + + Reference: https://openrouter.ai/docs + """ + + def validate_environment( + self, + headers: dict, + model: str, + messages: list, + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate environment and set up headers for OpenRouter API. + + OpenRouter requires: + - Authorization header with Bearer token + - HTTP-Referer header (site URL) + - X-Title header (app name) + """ + from litellm import get_secret + + # Get OpenRouter-specific headers + openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai" + openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM" + + openrouter_headers = { + "HTTP-Referer": openrouter_site_url, + "X-Title": openrouter_app_name, + "Content-Type": "application/json", + } + + # Add Authorization header if api_key is provided + if api_key: + openrouter_headers["Authorization"] = f"Bearer {api_key}" + + # Merge with existing headers (user's extra_headers take priority) + merged_headers = {**openrouter_headers, **headers} + + return merged_headers + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete URL for OpenRouter Embedding API endpoint. + """ + # api_base is already set to https://openrouter.ai/api/v1 in main.py + # Remove trailing slashes + if api_base: + api_base = api_base.rstrip("/") + else: + api_base = "https://openrouter.ai/api/v1" + + # Return the embeddings endpoint + return f"{api_base}/embeddings" + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict, + headers: dict, + ) -> dict: + """ + Transform embedding request to OpenRouter format (OpenAI-compatible). + """ + # Ensure input is a list + if isinstance(input, str): + input = [input] + + # OpenRouter expects the full model name (e.g., google/gemini-embedding-001) + # Strip 'openrouter/' prefix if present + if model.startswith("openrouter/"): + model = model.replace("openrouter/", "", 1) + + return { + "model": model, + "input": input, + **optional_params, + } + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + request_data: dict, + optional_params: dict, + litellm_params: dict, + ) -> EmbeddingResponse: + """ + Transform embedding response from OpenRouter format (OpenAI-compatible). + """ + logging_obj.post_call(original_response=raw_response.text) + + # OpenRouter returns standard OpenAI-compatible embedding response + response_json = raw_response.json() + + return convert_to_model_response_object( + response_object=response_json, + model_response_object=model_response, + response_type="embedding", + ) + + def get_supported_openai_params(self, model: str) -> list: + """ + Get list of supported OpenAI parameters for OpenRouter embeddings. + """ + return [ + "timeout", + "dimensions", + "encoding_format", + "user", + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI parameters to OpenRouter format. + """ + for param, value in non_default_params.items(): + if param in self.get_supported_openai_params(model): + optional_params[param] = value + return optional_params + + def get_error_class( + self, error_message: str, status_code: int, headers: Any + ) -> Any: + """ + Get the error class for OpenRouter errors. + """ + return OpenRouterException( + message=error_message, + status_code=status_code, + headers=headers, + ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index ba1788a217f..91100cf7d7b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -480,20 +480,21 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): or tool_name == VertexToolName.CODE_EXECUTION.value ): # code_execution maintained for backwards compatibility code_execution = self.get_tool_value(tool, "codeExecution") - elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH.value: - googleSearch = self.get_tool_value( - tool, VertexToolName.GOOGLE_SEARCH.value - ) - elif ( - tool_name and tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value + elif tool_name and ( + tool_name == VertexToolName.GOOGLE_SEARCH.value + or tool_name == "google_search" ): - googleSearchRetrieval = self.get_tool_value( - tool, VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value - ) - elif tool_name and tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value: - enterpriseWebSearch = self.get_tool_value( - tool, VertexToolName.ENTERPRISE_WEB_SEARCH.value - ) + googleSearch = self.get_tool_value(tool, tool_name) + elif tool_name and ( + tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value + or tool_name == "google_search_retrieval" + ): + googleSearchRetrieval = self.get_tool_value(tool, tool_name) + elif tool_name and ( + tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value + or tool_name == "enterprise_web_search" + ): + enterpriseWebSearch = self.get_tool_value(tool, tool_name) elif tool_name and ( tool_name == VertexToolName.URL_CONTEXT.value or tool_name == "urlContext" @@ -1811,6 +1812,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): functions: Optional[ChatCompletionToolCallFunctionChunk] = None thinking_blocks: Optional[List[ChatCompletionThinkingBlock]] = None reasoning_content: Optional[str] = None + thought_signatures: Optional[Any] = None for idx, candidate in enumerate(_candidates): if "content" not in candidate: diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 826f151df35..a3606ff9deb 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -388,6 +388,10 @@ class VertexBase: Internal function. Returns the token and url for the call. Handles logic if it's google ai studio vs. vertex ai. + + For Vertex AI: + - If gemini_api_key is provided, use API key authentication (x-goog-api-key header) + - Otherwise, use service account credentials (OAuth2 Bearer token) Returns token, url @@ -400,7 +404,7 @@ class VertexBase: stream=stream, gemini_api_key=gemini_api_key, ) - auth_header = None # this field is not used for gemin + auth_header = None # this field is not used for gemini else: vertex_location = self.get_vertex_region( vertex_region=vertex_location, @@ -409,14 +413,32 @@ class VertexBase: ### SET RUNTIME ENDPOINT ### version = "v1beta1" if should_use_v1beta1_features is True else "v1" - url, endpoint = _get_vertex_url( - mode=mode, - model=model, - stream=stream, - vertex_project=vertex_project, - vertex_location=vertex_location, - vertex_api_version=version, - ) + + # Check if using API key authentication for Vertex AI + if gemini_api_key and not vertex_credentials: + # When using API key with Vertex AI, use the Google AI Studio endpoint + # This is because Vertex AI API keys work with generativelanguage.googleapis.com + verbose_logger.debug( + f"Using Vertex AI API key authentication for model: {model} - routing to Google AI Studio endpoint" + ) + url, endpoint = _get_gemini_url( + mode=mode, + model=model, + stream=stream, + gemini_api_key=gemini_api_key, + ) + # API key is already included in the URL by _get_gemini_url + auth_header = None + else: + # Use OAuth2 Bearer token authentication (traditional Vertex AI) + url, endpoint = _get_vertex_url( + mode=mode, + model=model, + stream=stream, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_api_version=version, + ) return self._check_custom_proxy( api_base=api_base, diff --git a/litellm/main.py b/litellm/main.py index e8a8b504d96..10e3bcac04b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -189,7 +189,7 @@ from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from .llms.custom_llm import CustomLLM, custom_chat_llm_router from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha, palm -from .llms.gemini.common_utils import get_api_key_from_env +from .llms.gemini.common_utils import get_api_key_from_env, get_vertex_api_key_from_env from .llms.groq.chat.handler import GroqChatCompletion from .llms.heroku.chat.transformation import HerokuChatConfig from .llms.huggingface.embedding.handler import HuggingFaceEmbedding @@ -3230,6 +3230,12 @@ def completion( # type: ignore # noqa: PLR0915 or get_secret("VERTEXAI_CREDENTIALS") ) + vertex_api_key = ( + api_key + or get_vertex_api_key_from_env() + or litellm.api_key + ) + api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE") new_params = safe_deep_copy(optional_params or {}) @@ -3271,7 +3277,7 @@ def completion( # type: ignore # noqa: PLR0915 vertex_location=vertex_ai_location, vertex_project=vertex_ai_project, vertex_credentials=vertex_credentials, - gemini_api_key=None, + gemini_api_key=vertex_api_key, # Support for Vertex AI API Key logging_obj=logging, acompletion=acompletion, timeout=timeout, @@ -4701,6 +4707,51 @@ def embedding( # noqa: PLR0915 litellm_params=litellm_params_dict, headers=headers, ) + elif custom_llm_provider == "openrouter": + api_base = ( + api_base + or litellm.api_base + or get_secret_str("OPENROUTER_API_BASE") + or "https://openrouter.ai/api/v1" + ) + + api_key = ( + api_key + or litellm.api_key + or litellm.openrouter_key + or get_secret("OPENROUTER_API_KEY") + or get_secret("OR_API_KEY") + ) + + openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai" + openrouter_app_name = get_secret("OR_APP_NAME") or "liteLLM" + + openrouter_headers = { + "HTTP-Referer": openrouter_site_url, + "X-Title": openrouter_app_name, + } + + _headers = headers or litellm.headers + if _headers: + openrouter_headers.update(_headers) + + headers = openrouter_headers + + response = base_llm_http_handler.embedding( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + litellm_params=litellm_params_dict, + headers=headers, + ) elif custom_llm_provider == "huggingface": api_key = ( api_key diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c3b4a1b847d..0250a5fa9be 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1942,7 +1942,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): description="connect to a postgres db - needed for generating temporary keys + tracking spend / key", ) database_connection_pool_limit: Optional[int] = Field( - 100, + 10, description="default connection pool for prisma client connecting to postgres db", ) database_connection_timeout: Optional[float] = Field( @@ -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=()) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 334362a0271..7de6b7fccfc 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -11,7 +11,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, - create_streaming_response, + create_response, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.types.utils import TokenCountResponse @@ -106,7 +106,7 @@ async def anthropic_response( # noqa: PLR0915 ) ) - return await create_streaming_response( + return await create_response( generator=selected_data_generator, media_type="text/event-stream", headers={}, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 44c2ec0b61a..bc0c164a0ad 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -138,7 +138,7 @@ def _apply_budget_limits_to_end_user_params( ) -> None: """ Helper function to apply budget limits to end user parameters. - + Args: end_user_params: Dictionary to update with budget parameters budget_info: Budget table object containing limits @@ -146,16 +146,14 @@ def _apply_budget_limits_to_end_user_params( """ if budget_info.tpm_limit is not None: end_user_params["end_user_tpm_limit"] = budget_info.tpm_limit - + if budget_info.rpm_limit is not None: end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit - + if budget_info.max_budget is not None: end_user_params["end_user_max_budget"] = budget_info.max_budget - - verbose_proxy_logger.debug( - f"Applied budget limits to end user {end_user_id}" - ) + + verbose_proxy_logger.debug(f"Applied budget limits to end user {end_user_id}") async def user_api_key_auth_websocket(websocket: WebSocket): @@ -170,12 +168,10 @@ async def user_api_key_auth_websocket(websocket: WebSocket): model = query_params.get("model") - async def return_body(): return _realtime_request_body(model) - - request.body = return_body # type: ignore + request.body = return_body # type: ignore authorization = websocket.headers.get("authorization") # If no Authorization header, try the api-key header @@ -586,7 +582,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_membership is not None else None ), - team_metadata=team_object.metadata if team_object is not None else None, + team_metadata=team_object.metadata + if team_object is not None + else None, ) # run through common checks _ = await common_checks( @@ -669,9 +667,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params["allowed_model_region"] = ( - _end_user_object.allowed_model_region - ) + end_user_params[ + "allowed_model_region" + ] = _end_user_object.allowed_model_region if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -753,7 +751,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, code=400, - param=api_key, + param=abbreviate_api_key(api_key=api_key), ) valid_token = update_valid_token_with_end_user_params( valid_token=valid_token, end_user_params=end_user_params @@ -994,7 +992,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Check 3. Check if user is in their team budget if valid_token.team_member_spend is not None: - if prisma_client is not None: _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" @@ -1055,7 +1052,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, code=400, - param=api_key, + param=abbreviate_api_key(api_key=api_key), ) # Check 4. Token Spend is under budget @@ -1216,8 +1213,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) - - @tracer.wrap() async def user_api_key_auth( request: Request, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 537b48f06ed..61bffde3aca 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -17,7 +17,7 @@ from typing import ( import httpx import orjson from fastapi import HTTPException, Request, status -from fastapi.responses import Response, StreamingResponse +from fastapi.responses import JSONResponse, Response, StreamingResponse import litellm from litellm._logging import verbose_proxy_logger @@ -96,16 +96,55 @@ async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional return None -async def create_streaming_response( +def _extract_error_from_sse_chunk(event_line: Union[str, bytes]) -> dict: + """ + Extract error dictionary from SSE format chunk. + + Args: + event_line: SSE format event line, e.g. "data: {"error": {...}}\n\n" + + Returns: + Error dictionary in OpenAI API format + """ + event_line = ( + event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line + ) + + # Default error format + default_error = { + "message": "Unknown error", + "type": "internal_server_error", + "param": None, + "code": "500", + } + + if event_line.startswith("data: "): + json_str = event_line[len("data: ") :].strip() + if not json_str or json_str == "[DONE]": + return default_error + + try: + data = orjson.loads(json_str) + if isinstance(data, dict) and "error" in data: + error_obj = data["error"] + if isinstance(error_obj, dict): + return error_obj + except (orjson.JSONDecodeError, json.JSONDecodeError): + pass + + return default_error + + +async def create_response( generator: AsyncGenerator[str, None], media_type: str, headers: dict, default_status_code: int = status.HTTP_200_OK, -) -> StreamingResponse: +) -> Union[StreamingResponse, JSONResponse]: """ - Creates a StreamingResponse by inspecting the first chunk for an error code. - The entire original generator content is streamed, but the HTTP status code - of the response is set based on the first chunk if it's a recognized error. + Create streaming response, checking if the first chunk is an error. + If the first chunk is an error, return a standard JSON error response. + Otherwise, return StreamingResponse and stream all content. """ first_chunk_value: Optional[str] = None final_status_code = default_status_code @@ -124,9 +163,27 @@ async def create_streaming_response( first_chunk_value ) if error_code_from_chunk is not None: + # First chunk is an error, stream hasn't really started yet + # Should return standard JSON error response instead of SSE format final_status_code = error_code_from_chunk verbose_proxy_logger.debug( - f"Error detected in first stream chunk. Status code set to: {final_status_code}" + f"Error detected in first stream chunk. Returning JSON error response with status code: {final_status_code}" + ) + + # Parse error content + error_dict = _extract_error_from_sse_chunk(first_chunk_value) + + # Consume and close generator (avoid resource leak) + try: + await generator.aclose() + except Exception: + pass + + # Return JSON format error response + return JSONResponse( + status_code=final_status_code, + content={"error": error_dict}, + headers=headers, ) except Exception as e: verbose_proxy_logger.debug(f"Error parsing first chunk value: {e}") @@ -237,7 +294,11 @@ class ProxyBaseLLMRequestProcessing: if response_cost is not None: try: # Convert response_cost to float if it's a string - cost_value = float(response_cost) if isinstance(response_cost, str) else response_cost + cost_value = ( + float(response_cost) + if isinstance(response_cost, str) + else response_cost + ) if cost_value > 0: updated_spend = current_spend + cost_value except (ValueError, TypeError): @@ -376,6 +437,16 @@ class ProxyBaseLLMRequestProcessing: ) -> Tuple[dict, LiteLLMLoggingObj]: start_time = datetime.now() # start before calling guardrail hooks + # Calculate request queue time if arrival_time is available + # Use start_time.timestamp() to avoid extra time.time() call for better performance + proxy_server_request = self.data.get("proxy_server_request", {}) + arrival_time = proxy_server_request.get("arrival_time") + queue_time_seconds = None + if arrival_time is not None: + # Convert start_time (datetime) to timestamp for calculation + processing_start_time = start_time.timestamp() + queue_time_seconds = processing_start_time - arrival_time + self.data = await add_litellm_data_to_request( data=self.data, request=request, @@ -385,6 +456,19 @@ class ProxyBaseLLMRequestProcessing: proxy_config=proxy_config, ) + # Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved + if queue_time_seconds is not None: + from litellm.proxy.litellm_pre_call_utils import _get_metadata_variable_name + + _metadata_variable_name = _get_metadata_variable_name(request) + if _metadata_variable_name not in self.data: + self.data[_metadata_variable_name] = {} + if not isinstance(self.data[_metadata_variable_name], dict): + self.data[_metadata_variable_name] = {} + self.data[_metadata_variable_name][ + "queue_time_seconds" + ] = queue_time_seconds + self.data["model"] = ( general_settings.get("completion_model", None) # server default or user_model # model name passed via cli args @@ -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 diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 406ddceabf5..c9c0cfe8f68 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -17,50 +17,141 @@ from litellm.secret_managers.main import str_to_bool class PrismaWrapper: + """ + Wrapper around Prisma client that handles RDS IAM token authentication. + + When iam_token_db_auth is enabled, this wrapper: + 1. Proactively refreshes IAM tokens before they expire (background task) + 2. Falls back to synchronous refresh if a token is found expired + 3. Uses proper locking to prevent race conditions during reconnection + + RDS IAM tokens are valid for 15 minutes. This wrapper refreshes them + 3 minutes before expiration to ensure uninterrupted database connectivity. + """ + + # Buffer time in seconds before token expiration to trigger refresh + # Refresh 3 minutes (180 seconds) before the token expires + TOKEN_REFRESH_BUFFER_SECONDS = 180 + + # Fallback refresh interval if token parsing fails (10 minutes) + FALLBACK_REFRESH_INTERVAL_SECONDS = 600 + def __init__(self, original_prisma: Any, iam_token_db_auth: bool): self._original_prisma = original_prisma self.iam_token_db_auth = iam_token_db_auth + # Background token refresh task management + self._token_refresh_task: Optional[asyncio.Task] = None + self._reconnection_lock = asyncio.Lock() + self._last_refresh_time: Optional[datetime] = None + + def _extract_token_from_db_url(self, db_url: Optional[str]) -> Optional[str]: + """ + Extract the token (password) from the DATABASE_URL. + + The token contains the AWS signature with X-Amz-Date and X-Amz-Expires parameters. + + Important: We must parse the URL while it's still encoded to preserve structure, + then decode the password portion. Otherwise the '?' in the token breaks URL parsing. + """ + if db_url is None: + return None + try: + # Parse URL while still encoded to preserve structure + parsed = urllib.parse.urlparse(db_url) + if parsed.password: + # Now decode just the password/token + return urllib.parse.unquote(parsed.password) + return None + except Exception: + return None + + def _parse_token_expiration(self, token: Optional[str]) -> Optional[datetime]: + """ + Parse the token to extract its expiration time. + + Returns the datetime when the token expires, or None if parsing fails. + """ + if token is None: + return None + + try: + # Token format: ...?X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=900&... + if "?" not in token: + return None + + query_string = token.split("?", 1)[1] + params = urllib.parse.parse_qs(query_string) + + expires_str = params.get("X-Amz-Expires", [None])[0] + date_str = params.get("X-Amz-Date", [None])[0] + + if not expires_str or not date_str: + return None + + token_created = datetime.strptime(date_str, "%Y%m%dT%H%M%SZ") + expires_in = int(expires_str) + + return token_created + timedelta(seconds=expires_in) + except Exception as e: + verbose_proxy_logger.debug(f"Failed to parse token expiration: {e}") + return None + + def _calculate_seconds_until_refresh(self) -> float: + """ + Calculate exactly how many seconds until we need to refresh the token. + + Uses precise timing: sleeps until (token_expiration - buffer_seconds). + For a 15-minute (900s) token with 180s buffer, this returns ~720s (12 min). + + Returns: + Number of seconds to sleep before the next refresh. + Returns 0 if token should be refreshed immediately. + Returns FALLBACK_REFRESH_INTERVAL_SECONDS if parsing fails. + """ + db_url = os.getenv("DATABASE_URL") + token = self._extract_token_from_db_url(db_url) + expiration_time = self._parse_token_expiration(token) + + if expiration_time is None: + # If we can't parse the token, use fallback interval + verbose_proxy_logger.debug( + f"Could not parse token expiration, using fallback interval of " + f"{self.FALLBACK_REFRESH_INTERVAL_SECONDS}s" + ) + return self.FALLBACK_REFRESH_INTERVAL_SECONDS + + # Calculate when we should refresh (expiration - buffer) + refresh_at = expiration_time - timedelta( + seconds=self.TOKEN_REFRESH_BUFFER_SECONDS + ) + + # How long until refresh time? + now = datetime.utcnow() + seconds_until_refresh = (refresh_at - now).total_seconds() + + # If already past refresh time, return 0 (refresh immediately) + return max(0, seconds_until_refresh) + def is_token_expired(self, token_url: Optional[str]) -> bool: + """Check if the token in the given URL is expired.""" if token_url is None: return True - # Decode the token URL to handle URL-encoded characters - decoded_url = urllib.parse.unquote(token_url) - # Parse the token URL - parsed_url = urllib.parse.urlparse(decoded_url) + token = self._extract_token_from_db_url(token_url) + expiration_time = self._parse_token_expiration(token) - # Parse the query parameters from the path component (if they exist there) - query_params = urllib.parse.parse_qs(parsed_url.query) + if expiration_time is None: + # If we can't parse the token, assume it's expired to trigger refresh + verbose_proxy_logger.debug( + "Could not parse token expiration, treating as expired" + ) + return True - # Get expiration time from the query parameters - expires = query_params.get("X-Amz-Expires", [None])[0] - if expires is None: - raise ValueError("X-Amz-Expires parameter is missing or invalid.") - - expires_int = int(expires) - - # Get the token's creation time from the X-Amz-Date parameter - token_time_str = query_params.get("X-Amz-Date", [""])[0] - if not token_time_str: - raise ValueError("X-Amz-Date parameter is missing or invalid.") - - # Ensure the token time string is parsed correctly - try: - token_time = datetime.strptime(token_time_str, "%Y%m%dT%H%M%SZ") - except ValueError as e: - raise ValueError(f"Invalid X-Amz-Date format: {e}") - - # Calculate the expiration time - expiration_time = token_time + timedelta(seconds=expires_int) - - # Current time in UTC - current_time = datetime.utcnow() - - # Check if the token is expired - return current_time > expiration_time + return datetime.utcnow() > expiration_time def get_rds_iam_token(self) -> Optional[str]: + """Generate a new RDS IAM token and update DATABASE_URL.""" if self.iam_token_db_auth: from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token @@ -74,7 +165,6 @@ class PrismaWrapper: db_host=db_host, db_port=db_port, db_user=db_user ) - # print(f"token: {token}") _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" if db_schema: _db_url += f"?schema={db_schema}" @@ -86,6 +176,7 @@ class PrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, http_client: Optional[Any] = None ): + """Disconnect and reconnect the Prisma client with a new database URL.""" from prisma import Prisma # type: ignore try: @@ -100,21 +191,159 @@ class PrismaWrapper: await self._original_prisma.connect() + async def start_token_refresh_task(self) -> None: + """ + Start the background token refresh task. + + This task proactively refreshes RDS IAM tokens before they expire, + preventing connection failures. Should be called after the initial + Prisma client connection is established. + """ + if not self.iam_token_db_auth: + verbose_proxy_logger.debug( + "IAM token auth not enabled, skipping token refresh task" + ) + return + + if self._token_refresh_task is not None: + verbose_proxy_logger.debug("Token refresh task already running") + return + + self._token_refresh_task = asyncio.create_task(self._token_refresh_loop()) + verbose_proxy_logger.info( + "Started RDS IAM token proactive refresh background task" + ) + + async def stop_token_refresh_task(self) -> None: + """ + Stop the background token refresh task gracefully. + + Should be called during application shutdown to clean up resources. + """ + if self._token_refresh_task is None: + return + + self._token_refresh_task.cancel() + try: + await self._token_refresh_task + except asyncio.CancelledError: + pass + self._token_refresh_task = None + verbose_proxy_logger.info("Stopped RDS IAM token refresh background task") + + async def _token_refresh_loop(self) -> None: + """ + Background loop that proactively refreshes RDS IAM tokens before expiration. + + Uses precise timing: calculates the exact sleep duration until the token + needs to be refreshed (expiration - 3 minute buffer), then refreshes. + This is more efficient than polling, requiring only 1 wake-up per token cycle. + """ + verbose_proxy_logger.info( + f"RDS IAM token refresh loop started. " + f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration." + ) + + while True: + try: + # Calculate exactly how long to sleep until next refresh + sleep_seconds = self._calculate_seconds_until_refresh() + + if sleep_seconds > 0: + verbose_proxy_logger.info( + f"RDS IAM token refresh scheduled in {sleep_seconds:.0f} seconds " + f"({sleep_seconds / 60:.1f} minutes)" + ) + await asyncio.sleep(sleep_seconds) + + # Refresh the token + verbose_proxy_logger.info("Proactively refreshing RDS IAM token...") + await self._safe_refresh_token() + + except asyncio.CancelledError: + verbose_proxy_logger.info("RDS IAM token refresh loop cancelled") + break + except Exception as e: + verbose_proxy_logger.error( + f"Error in RDS IAM token refresh loop: {e}. " + f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..." + ) + # On error, wait before retrying to avoid tight error loops + try: + await asyncio.sleep(self.FALLBACK_REFRESH_INTERVAL_SECONDS) + except asyncio.CancelledError: + break + + async def _safe_refresh_token(self) -> None: + """ + Refresh the RDS IAM token with proper locking to prevent race conditions. + + Uses an asyncio lock to ensure only one refresh operation happens at a time, + preventing multiple concurrent reconnection attempts. + """ + async with self._reconnection_lock: + new_db_url = self.get_rds_iam_token() + if new_db_url: + await self.recreate_prisma_client(new_db_url) + self._last_refresh_time = datetime.utcnow() + verbose_proxy_logger.info( + "RDS IAM token refreshed successfully. New token valid for ~15 minutes." + ) + else: + verbose_proxy_logger.error( + "Failed to generate new RDS IAM token during proactive refresh" + ) + def __getattr__(self, name: str): + """ + Proxy attribute access to the underlying Prisma client. + + If IAM token auth is enabled and the token is expired, this method + provides a synchronous fallback to refresh the token. However, this + should rarely be needed since the background task proactively refreshes + tokens before they expire. + + FIXED: Now properly waits for reconnection to complete before returning, + instead of the previous fire-and-forget pattern that caused the bug. + """ original_attr = getattr(self._original_prisma, name) + if self.iam_token_db_auth: db_url = os.getenv("DATABASE_URL") - if self.is_token_expired(db_url): - db_url = self.get_rds_iam_token() - loop = asyncio.get_event_loop() - if db_url: + # Check if token is expired (should be rare if background task is running) + if self.is_token_expired(db_url): + verbose_proxy_logger.warning( + "RDS IAM token expired in __getattr__ - proactive refresh may have failed. " + "Triggering synchronous fallback refresh..." + ) + + new_db_url = self.get_rds_iam_token() + if new_db_url: + loop = asyncio.get_event_loop() + if loop.is_running(): - asyncio.run_coroutine_threadsafe( - self.recreate_prisma_client(db_url), loop + # FIXED: Actually wait for the reconnection to complete! + # The previous code used fire-and-forget which caused the bug. + future = asyncio.run_coroutine_threadsafe( + self.recreate_prisma_client(new_db_url), loop ) + try: + # Wait up to 30 seconds for reconnection + future.result(timeout=30) + verbose_proxy_logger.info( + "Synchronous token refresh completed successfully" + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to refresh token synchronously: {e}" + ) + raise else: - asyncio.run(self.recreate_prisma_client(db_url)) + asyncio.run(self.recreate_prisma_client(new_db_url)) + + # Get the NEW attribute after reconnection + original_attr = getattr(self._original_prisma, name) else: raise ValueError("Failed to get RDS IAM token") diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index a6971b49f3b..87da11efad0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -5,6 +5,7 @@ # +-------------------------------------------------------------+ # Qualifire - Evaluate LLM outputs for quality, safety, and reliability +import json import os from typing import Any, Dict, List, Literal, Optional, Type @@ -15,12 +16,17 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs GUARDRAIL_NAME = "qualifire" +DEFAULT_QUALIFIRE_API_BASE = "https://proxy.qualifire.ai" class QualifireGuardrail(CustomGuardrail): @@ -44,7 +50,7 @@ class QualifireGuardrail(CustomGuardrail): Args: api_key: API key for Qualifire (or use QUALIFIRE_API_KEY env var) - api_base: Optional custom API base URL + api_base: Optional custom API base URL (defaults to https://api.qualifire.ai) evaluation_id: Pre-configured evaluation ID from Qualifire dashboard prompt_injections: Enable prompt injection detection (default if no other checks) hallucinations_check: Enable hallucination detection @@ -64,6 +70,7 @@ class QualifireGuardrail(CustomGuardrail): api_base or get_secret_str("QUALIFIRE_BASE_URL") or os.environ.get("QUALIFIRE_BASE_URL") + or DEFAULT_QUALIFIRE_API_BASE ) self.evaluation_id = evaluation_id self.prompt_injections = prompt_injections @@ -79,7 +86,11 @@ class QualifireGuardrail(CustomGuardrail): if not self._has_any_check_enabled() and not self.evaluation_id: self.prompt_injections = True - self._client = None + # Initialize async HTTP client for direct API calls + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + super().__init__(**kwargs) def _has_any_check_enabled(self) -> bool: @@ -96,43 +107,22 @@ class QualifireGuardrail(CustomGuardrail): ] ) - def _get_client(self): - """Lazy initialization of Qualifire client.""" - if self._client is None: - try: - from qualifire.client import Client - except ImportError: - raise ImportError( - "qualifire package is required for QualifireGuardrail. " - "Install it with: pip install qualifire" - ) - - client_kwargs: Dict[str, Any] = {} - if self.qualifire_api_key: - client_kwargs["api_key"] = self.qualifire_api_key - if self.qualifire_api_base: - client_kwargs["base_url"] = self.qualifire_api_base - - self._client = Client(**client_kwargs) - - return self._client - - def _convert_messages_to_qualifire_format( + def _convert_messages_to_api_format( self, messages: List[AllMessageValues] - ) -> List[Any]: + ) -> List[Dict[str, Any]]: """ - Convert LiteLLM messages to Qualifire's LLMMessage format. + Convert LiteLLM messages to Qualifire API format. Supports tool calls for tool_selection_quality_check. - """ - try: - from qualifire.types import LLMMessage, LLMToolCall - except ImportError: - raise ImportError( - "qualifire package is required for QualifireGuardrail. " - "Install it with: pip install qualifire" - ) - qualifire_messages = [] + Returns a list of dicts matching the API's ModelInvocationCanonicalMessage schema: + { + "role": "user" | "assistant" | "system" | "tool", + "content": "...", + "tool_call_id": "...", # optional + "tool_calls": [{"id": "...", "name": "...", "arguments": {...}}] # optional + } + """ + api_messages = [] for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") @@ -147,42 +137,86 @@ class QualifireGuardrail(CustomGuardrail): text_parts.append(part) content = "\n".join(text_parts) - llm_message_kwargs: Dict[str, Any] = { + api_message: Dict[str, Any] = { "role": role, "content": content if isinstance(content, str) else str(content), } + # Handle tool_call_id for tool response messages + tool_call_id = msg.get("tool_call_id") + if tool_call_id: + api_message["tool_call_id"] = tool_call_id + # Handle tool calls if present tool_calls = msg.get("tool_calls") if tool_calls and isinstance(tool_calls, list): - qualifire_tool_calls = [] + api_tool_calls = [] for tc in tool_calls: if isinstance(tc, dict): function_info = tc.get("function", {}) # Arguments can be a string (JSON) or dict args = function_info.get("arguments", {}) if isinstance(args, str): - import json - try: args = json.loads(args) except json.JSONDecodeError: args = {} - qualifire_tool_calls.append( - LLMToolCall( - id=tc.get("id") or "", - name=function_info.get("name") or "", - arguments=args if isinstance(args, dict) else {}, - ) + api_tool_calls.append( + { + "id": tc.get("id") or "", + "name": function_info.get("name") or "", + "arguments": args if isinstance(args, dict) else {}, + } ) - if qualifire_tool_calls: - llm_message_kwargs["tool_calls"] = qualifire_tool_calls + if api_tool_calls: + api_message["tool_calls"] = api_tool_calls - qualifire_messages.append(LLMMessage(**llm_message_kwargs)) + api_messages.append(api_message) - return qualifire_messages + return api_messages - def _check_if_flagged(self, result: Any) -> bool: + def _convert_tools_to_api_format( + self, tools: Optional[List[Any]] + ) -> Optional[List[Dict[str, Any]]]: + """ + Convert OpenAI-format tools to Qualifire API format. + + Returns a list of dicts matching the API's ModelInvocationToolDefinition schema: + { + "name": "...", + "description": "...", + "parameters": {...} + } + """ + if not tools: + return None + + api_tools = [] + for tool in tools: + if isinstance(tool, dict): + # Handle OpenAI function tool format + if tool.get("type") == "function": + function_def = tool.get("function", {}) + api_tools.append( + { + "name": function_def.get("name", ""), + "description": function_def.get("description", ""), + "parameters": function_def.get("parameters", {}), + } + ) + # Handle direct tool format + elif "name" in tool: + api_tools.append( + { + "name": tool.get("name", ""), + "description": tool.get("description", ""), + "parameters": tool.get("parameters", {}), + } + ) + + return api_tools if api_tools else None + + def _check_if_flagged(self, result: Dict[str, Any]) -> bool: """ Check if the Qualifire evaluation result indicates flagged content. @@ -190,65 +224,53 @@ class QualifireGuardrail(CustomGuardrail): A high score (close to 100) indicates GOOD content, low score indicates problems. """ # Check evaluation results for any flagged items - evaluation_results = getattr(result, "evaluationResults", None) or [] - if isinstance(result, dict): - evaluation_results = result.get("evaluationResults", []) or [] + evaluation_results = result.get("evaluationResults", []) or [] for eval_result in evaluation_results: - results: List[Any] = [] - if isinstance(eval_result, dict): - results = eval_result.get("results", []) or [] - else: - results = getattr(eval_result, "results", []) or [] - + results = eval_result.get("results", []) or [] for r in results: - flagged = ( - r.get("flagged") - if isinstance(r, dict) - else getattr(r, "flagged", False) - ) - if flagged: + if r.get("flagged"): return True return False - def _build_evaluate_kwargs( + def _build_evaluate_payload( self, - qualifire_messages: List[Any], + api_messages: List[Dict[str, Any]], output: Optional[str], assertions: Optional[List[str]], - available_tools: Optional[List[Any]], + available_tools: Optional[List[Dict[str, Any]]], ) -> Dict[str, Any]: - """Build kwargs dictionary for the evaluate call.""" - kwargs: Dict[str, Any] = {"messages": qualifire_messages} + """Build payload dictionary for the /api/evaluation/evaluate endpoint.""" + payload: Dict[str, Any] = {"messages": api_messages} if output is not None: - kwargs["output"] = output + payload["output"] = output # Add enabled checks if self.prompt_injections: - kwargs["prompt_injections"] = True + payload["prompt_injections"] = True if self.hallucinations_check: - kwargs["hallucinations_check"] = True + payload["hallucinations_check"] = True if self.grounding_check: - kwargs["grounding_check"] = True + payload["grounding_check"] = True if self.pii_check: - kwargs["pii_check"] = True + payload["pii_check"] = True if self.content_moderation_check: - kwargs["content_moderation_check"] = True + payload["content_moderation_check"] = True if self.tool_selection_quality_check: # Only enable tool_selection_quality_check if available_tools is provided if available_tools: - kwargs["tool_selection_quality_check"] = True - kwargs["available_tools"] = available_tools + payload["tool_selection_quality_check"] = True + payload["available_tools"] = available_tools else: verbose_proxy_logger.debug( "Qualifire Guardrail: tool_selection_quality_check enabled but no available_tools provided, skipping this check" ) if assertions: - kwargs["assertions"] = assertions + payload["assertions"] = assertions - return kwargs + return payload async def _run_qualifire_check( self, @@ -274,11 +296,17 @@ class QualifireGuardrail(CustomGuardrail): assertions = dynamic_params.get("assertions") or self.assertions on_flagged = dynamic_params.get("on_flagged") or self.on_flagged - try: - client = self._get_client() - qualifire_messages = self._convert_messages_to_qualifire_format(messages) + # Prepare headers + headers = { + "X-Qualifire-API-Key": self.qualifire_api_key or "", + "Content-Type": "application/json", + } - # Use invoke_evaluation if evaluation_id is provided + try: + # Convert messages to API format + api_messages = self._convert_messages_to_api_format(messages) + + # Use invoke endpoint if evaluation_id is provided if evaluation_id: # For invoke_evaluation, we need to extract input/output input_text = "" @@ -291,25 +319,47 @@ class QualifireGuardrail(CustomGuardrail): input_text = content break - result = client.invoke_evaluation( - evaluation_id=evaluation_id, - input=input_text, - output=output or "", - ) + payload = { + "evaluation_id": evaluation_id, + "input": input_text, + "output": output or "", + "messages": api_messages, + } + + # Convert tools if provided + api_tools = self._convert_tools_to_api_format(available_tools) + if api_tools: + payload["available_tools"] = api_tools + + url = f"{self.qualifire_api_base}/api/evaluation/invoke" else: - # Use evaluate with individual checks - kwargs = self._build_evaluate_kwargs( - qualifire_messages=qualifire_messages, + # Use evaluate endpoint with individual checks + api_tools = self._convert_tools_to_api_format(available_tools) + payload = self._build_evaluate_payload( + api_messages=api_messages, output=output, assertions=assertions, - available_tools=available_tools, + available_tools=api_tools, ) - result = client.evaluate(**kwargs) + url = f"{self.qualifire_api_base}/api/evaluation/evaluate" - # Convert result to dict for logging + verbose_proxy_logger.debug( + f"Qualifire Guardrail: Making request to {url}" + ) + + # Make the API request + response = await self.async_handler.post( + url=url, + headers=headers, + json=payload, + ) + response.raise_for_status() + result = response.json() + + # Extract response info for logging qualifire_response = { - "score": getattr(result, "score", None), - "status": getattr(result, "status", None), + "score": result.get("score"), + "status": result.get("status"), } verbose_proxy_logger.debug( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 7dc6741f023..1fbd8ee72c2 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -161,7 +161,6 @@ class KeyAndTeamLoggingSettings: @staticmethod def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): - if ( user_api_key_dict.team_metadata is not None and "logging" in user_api_key_dict.team_metadata @@ -174,12 +173,12 @@ def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> Optional[TeamCallbackMetadata]: callback_settings_obj: Optional[TeamCallbackMetadata] = None - key_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) - ) - team_dynamic_logging_settings: Optional[dict] = ( - KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) - ) + key_dynamic_logging_settings: Optional[ + dict + ] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + team_dynamic_logging_settings: Optional[ + dict + ] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) ######################################################################################### # Key-based callbacks ######################################################################################### @@ -462,7 +461,6 @@ class LiteLLMProxyRequestSetup: team_id=user_api_key_dict.team_id, ) # handles aliases, wildcards, etc. ): - _headers = LiteLLMProxyRequestSetup.add_headers_to_llm_call( headers, user_api_key_dict ) @@ -663,11 +661,11 @@ class LiteLLMProxyRequestSetup: ## KEY-LEVEL SPEND LOGS / TAGS if "tags" in key_metadata and key_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = ( - LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=key_metadata["tags"], - ) + data[_metadata_variable_name][ + "tags" + ] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=key_metadata["tags"], ) if "disable_global_guardrails" in key_metadata and isinstance( key_metadata["disable_global_guardrails"], bool @@ -815,11 +813,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915 # Init - Proxy Server Request # we do this as soon as entering so we track the original request ########################################################## + # Track arrival time for queue time metric + arrival_time = time.time() data["proxy_server_request"] = { "url": str(request.url), "method": request.method, "headers": _headers, "body": copy.copy(data), # use copy instead of deepcopy + "arrival_time": arrival_time, # Track when request arrived at proxy } safe_add_api_version_from_query_params(data, request) @@ -930,9 +931,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data[_metadata_variable_name]["litellm_api_version"] = version if general_settings is not None: - data[_metadata_variable_name]["global_max_parallel_requests"] = ( - general_settings.get("global_max_parallel_requests", None) - ) + data[_metadata_variable_name][ + "global_max_parallel_requests" + ] = general_settings.get("global_max_parallel_requests", None) ### KEY-LEVEL Controls key_metadata = user_api_key_dict.metadata diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 47793c8fc8e..d8816df010a 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -36,12 +36,19 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._experimental.mcp_server.utils import ( + get_server_prefix, validate_and_normalize_mcp_server_payload, ) router = APIRouter(prefix="/v1/mcp", tags=["mcp"]) + MCP_AVAILABLE: bool = True + TEMPORARY_MCP_SERVER_TTL_SECONDS = 300 +DEFAULT_MCP_REGISTRY_VERSION = "1.0.0" +LITELLM_MCP_SERVER_NAME = "litellm-mcp-server" +LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM" + try: importlib.import_module("mcp") except ImportError as e: @@ -57,6 +64,7 @@ if MCP_AVAILABLE: update_mcp_server, ) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, authorize_with_server, exchange_token_with_server, register_client_with_server, @@ -89,6 +97,66 @@ if MCP_AVAILABLE: server: MCPServer expires_at: datetime + def _is_public_registry_enabled() -> bool: + from litellm.proxy.proxy_server import ( + general_settings as proxy_general_settings, + ) + + return bool(proxy_general_settings.get("enable_mcp_registry")) + + def _build_registry_remote_url(base_url: str, path: str) -> str: + normalized_base = base_url.rstrip("/") + normalized_path = path if path.startswith("/") else f"/{path}" + return f"{normalized_base}{normalized_path}" + + def _build_mcp_registry_server_name(server: MCPServer) -> str: + if server.alias: + return server.alias + if server.server_name: + return server.server_name + return server.server_id + + def _build_mcp_registry_entry_for_server( + server: MCPServer, base_url: str + ) -> Dict[str, Any]: + server_name = _build_mcp_registry_server_name(server) + title = server_name + description = server_name + version = DEFAULT_MCP_REGISTRY_VERSION + + server_prefix = get_server_prefix(server) + if not server_prefix: + raise ValueError("MCP server prefix is missing") + remote_url = _build_registry_remote_url(base_url, f"/{server_prefix}/mcp") + + return { + "name": server_name, + "title": title, + "description": description, + "version": version, + "remotes": [ + { + "type": "streamable-http", + "url": remote_url, + } + ], + } + + def _build_builtin_registry_entry(base_url: str) -> Dict[str, Any]: + remote_url = _build_registry_remote_url(base_url, "/mcp") + return { + "name": LITELLM_MCP_SERVER_NAME, + "title": LITELLM_MCP_SERVER_NAME, + "description": LITELLM_MCP_SERVER_DESCRIPTION, + "version": DEFAULT_MCP_REGISTRY_VERSION, + "remotes": [ + { + "type": "streamable-http", + "url": remote_url, + } + ], + } + _temporary_mcp_servers: Dict[str, _TemporaryMCPServerEntry] = {} def _prune_expired_temporary_mcp_servers() -> None: @@ -302,15 +370,42 @@ if MCP_AVAILABLE: access_groups_list = sorted(list(access_groups)) return {"access_groups": access_groups_list} + @router.get( + "/registry.json", + tags=["mcp"], + description="MCP registry endpoint. Spec: https://github.com/modelcontextprotocol/registry", + ) + async def get_mcp_registry(request: Request): + if not _is_public_registry_enabled(): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="MCP registry is not enabled", + ) + + base_url = get_request_base_url(request) + registry_servers: List[Dict[str, Any]] = [] + registry_servers.append({"server": _build_builtin_registry_entry(base_url)}) + + registered_servers = list(global_mcp_server_manager.get_registry().values()) + registered_servers.sort(key=_build_mcp_registry_server_name) + + for server in registered_servers: + try: + entry = _build_mcp_registry_entry_for_server(server, base_url) + except Exception as e: + verbose_proxy_logger.debug( + f"Skipping MCP server {getattr(server, 'server_id', 'unknown')} in registry: {e}" + ) + continue + registry_servers.append({"server": entry}) + + return {"servers": registry_servers} + ## FastAPI Routes def _get_user_mcp_management_mode() -> UserMCPManagementMode: - proxy_general_settings: dict = {} - try: - from litellm.proxy.proxy_server import ( - general_settings as proxy_general_settings, - ) - except Exception: - pass + from litellm.proxy.proxy_server import ( + general_settings as proxy_general_settings, + ) mode = proxy_general_settings.get("user_mcp_management_mode") if mode == "view_all": diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d264b82b873..6e1079b886c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -229,7 +229,7 @@ from litellm.proxy.batches_endpoints.endpoints import router as batches_router from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, - create_streaming_response, + create_response, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.debug_utils import init_verbose_loggers @@ -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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f16c115fed3..171898b1631 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -131,6 +131,7 @@ else: unified_guardrail = UnifiedLLMGuardrails() +_anthropic_async_clients = {} def print_verbose(print_statement): """ @@ -961,6 +962,7 @@ class ProxyLogging: Updated data dictionary if guardrail passes, None if guardrail should be skipped """ from litellm.types.guardrails import GuardrailEventHooks + from litellm.integrations.prometheus import PrometheusLogger # Determine the event type based on call type event_type = GuardrailEventHooks.pre_call @@ -973,30 +975,62 @@ class ProxyLogging: guardrail_name = callback.guardrail_name - # Check if load balancing should be used - if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name): - response = await self._execute_guardrail_with_load_balancing( - guardrail_name=guardrail_name, - hook_type="pre_call", - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) - else: - # Single guardrail - execute directly - response = await self._execute_guardrail_hook( - callback=callback, - hook_type="pre_call", - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) + # Track timing and errors for prometheus metrics + # Use time.perf_counter() for more accurate duration measurements + guardrail_start_time = time.perf_counter() + status = "success" + error_type = None - # Process the response if one was returned - if response is not None: - data = await self.process_pre_call_hook_response( - response=response, data=data, call_type=call_type - ) + try: + # Check if load balancing should be used + if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name): + response = await self._execute_guardrail_with_load_balancing( + guardrail_name=guardrail_name, + hook_type="pre_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + else: + # Single guardrail - execute directly + response = await self._execute_guardrail_hook( + callback=callback, + hook_type="pre_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + + # Process the response if one was returned + if response is not None: + data = await self.process_pre_call_hook_response( + response=response, data=data, call_type=call_type + ) + + except Exception as e: + status = "error" + error_type = type(e).__name__ + # Re-raise the exception to maintain existing behavior + raise + finally: + # Record prometheus metrics + guardrail_end_time = time.perf_counter() + latency_seconds = guardrail_end_time - guardrail_start_time + + # Get guardrail name for metrics (fallback if not set) + metrics_guardrail_name = guardrail_name or getattr(callback, "guardrail_name", callback.__class__.__name__) or "unknown" + + # Find PrometheusLogger in callbacks and record metrics + for prom_callback in litellm.callbacks: + if isinstance(prom_callback, PrometheusLogger): + prom_callback._record_guardrail_metrics( + guardrail_name=metrics_guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type="pre_call", + ) + break return data @@ -4254,11 +4288,16 @@ async def count_tokens_with_anthropic_api( if anthropic_api_key and messages: # Call Anthropic API directly for more accurate token counting - client = anthropic.Anthropic(api_key=anthropic_api_key) + + # Use cached client if available to avoid socket exhaustion + if anthropic_api_key not in _anthropic_async_clients: + _anthropic_async_clients[anthropic_api_key] = anthropic.AsyncAnthropic(api_key=anthropic_api_key) + + client = _anthropic_async_clients[anthropic_api_key] # Call with explicit parameters to satisfy type checking # Type ignore for now since messages come from generic dict input - response = client.beta.messages.count_tokens( + response = await client.beta.messages.count_tokens( model=model_to_use, messages=messages, # type: ignore betas=["token-counting-2024-11-01"], diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index a92b5d25a37..7667d1bad84 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -443,11 +443,18 @@ class ResponseAPILoggingUtils: completion_tokens=0, total_tokens=0, ) - response_api_usage: ResponseAPIUsage = ( - ResponseAPIUsage(**usage_input) - if isinstance(usage_input, dict) - else usage_input - ) + response_api_usage: ResponseAPIUsage + if isinstance(usage_input, dict): + total_tokens = usage_input.get("total_tokens") + if total_tokens is None: + input_tokens = usage_input.get("input_tokens") + output_tokens = usage_input.get("output_tokens") + if input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + usage_input["total_tokens"] = total_tokens + response_api_usage = ResponseAPIUsage(**usage_input) + else: + response_api_usage = usage_input prompt_tokens: int = response_api_usage.input_tokens or 0 completion_tokens: int = response_api_usage.output_tokens or 0 prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None diff --git a/litellm/router.py b/litellm/router.py index 84b38b3985b..88b4087c1ee 100644 --- a/litellm/router.py +++ b/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 ( diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index bc8f06dcc40..fd9b722287d 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -1,7 +1,7 @@ import re from dataclasses import dataclass from enum import Enum -from typing import Dict, List, Literal, Optional, Tuple, Union +from typing import Dict, List, Literal, Optional, Tuple from pydantic import BaseModel, Field from typing_extensions import Annotated @@ -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) diff --git a/litellm/utils.py b/litellm/utils.py index 2260b2c7ba5..5ebf21f0644 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -47,7 +47,6 @@ from tiktoken import Encoding from tokenizers import Tokenizer import litellm - import litellm.litellm_core_utils # audio_utils.utils is lazy-loaded - only imported when needed for transcription calls import litellm.litellm_core_utils.json_validation_rule @@ -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 diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 673aab0990d..8432ce4e874 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -32,6 +32,23 @@ } }, "providers": { + "abliteration": { + "display_name": "Abliteration (`abliteration`)", + "url": "https://docs.litellm.ai/docs/providers/abliteration", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "aiml": { "display_name": "AI/ML API (`aiml`)", "url": "https://docs.litellm.ai/docs/providers/aiml", @@ -1559,7 +1576,7 @@ "chat_completions": true, "messages": true, "responses": true, - "embeddings": false, + "embeddings": true, "image_generations": false, "audio_transcriptions": false, "audio_speech": false, diff --git a/tests/litellm/llms/openai_like/test_abliteration_provider.py b/tests/litellm/llms/openai_like/test_abliteration_provider.py new file mode 100644 index 00000000000..8b8d443fc44 --- /dev/null +++ b/tests/litellm/llms/openai_like/test_abliteration_provider.py @@ -0,0 +1,50 @@ +""" +Unit tests for the Abliteration OpenAI-like provider. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +) + +from litellm.llms.openai_like.dynamic_config import create_config_class +from litellm.llms.openai_like.json_loader import JSONProviderRegistry + +ABLITERATION_BASE_URL = "https://api.abliteration.ai/v1" + + +def _get_config(): + provider = JSONProviderRegistry.get("abliteration") + assert provider is not None + config_class = create_config_class(provider) + return config_class() + + +def test_abliteration_provider_registered(): + provider = JSONProviderRegistry.get("abliteration") + assert provider is not None + assert provider.base_url == ABLITERATION_BASE_URL + assert provider.api_key_env == "ABLITERATION_API_KEY" + + +def test_abliteration_resolves_env_api_key(monkeypatch): + config = _get_config() + monkeypatch.setenv("ABLITERATION_API_KEY", "test-key") + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == ABLITERATION_BASE_URL + assert api_key == "test-key" + + +def test_abliteration_complete_url_appends_endpoint(): + config = _get_config() + url = config.get_complete_url( + api_base=ABLITERATION_BASE_URL, + api_key="test-key", + model="abliteration/abliterated-model", + optional_params={}, + litellm_params={}, + stream=False, + ) + assert url == f"{ABLITERATION_BASE_URL}/chat/completions" diff --git a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py index 6d005af28ac..20f48b6f393 100644 --- a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py @@ -7,6 +7,7 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path from litellm.llms.vertex_ai.gemini import transformation +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig from litellm.types.llms import openai from litellm.types import completion from litellm.types.llms.vertex_ai import RequestBody @@ -225,4 +226,65 @@ async def test__transform_request_body_image_config_with_image_size(): assert "generationConfig" in rb assert "imageConfig" in rb["generationConfig"] assert rb["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" - assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K" \ No newline at end of file + assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K" + + +def test_map_function_google_search_snake_case(): + """ + Test that google_search tool (snake_case) is properly mapped to googleSearch. + Fixes issue where tools=[{"google_search": {}}] was being stripped. + """ + config = VertexGeminiConfig() + optional_params = {} + + # Test snake_case google_search + tools = [{"google_search": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearch" in result[0] + assert result[0]["googleSearch"] == {} + + +def test_map_function_google_search_camel_case(): + """ + Test that googleSearch tool (camelCase) still works. + """ + config = VertexGeminiConfig() + optional_params = {} + + # Test camelCase googleSearch + tools = [{"googleSearch": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearch" in result[0] + assert result[0]["googleSearch"] == {} + + +def test_map_function_google_search_retrieval_snake_case(): + """ + Test that google_search_retrieval tool (snake_case) is properly mapped. + """ + config = VertexGeminiConfig() + optional_params = {} + + tools = [{"google_search_retrieval": {"dynamic_retrieval_config": {"mode": "MODE_DYNAMIC"}}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearchRetrieval" in result[0] + + +def test_map_function_enterprise_web_search_snake_case(): + """ + Test that enterprise_web_search tool (snake_case) is properly mapped. + """ + config = VertexGeminiConfig() + optional_params = {} + + tools = [{"enterprise_web_search": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "enterpriseWebSearch" in result[0] \ No newline at end of file diff --git a/tests/llm_translation/test_bedrock_common_utils.py b/tests/llm_translation/test_bedrock_common_utils.py new file mode 100644 index 00000000000..7b6a05b6988 --- /dev/null +++ b/tests/llm_translation/test_bedrock_common_utils.py @@ -0,0 +1,182 @@ +""" +Unit tests for litellm/llms/bedrock/common_utils.py + +Tests the standalone model name utility functions and BedrockTokenCounter. +""" + +import pytest + +from litellm.llms.bedrock.common_utils import ( + BedrockModelInfo, + extract_model_name_from_bedrock_arn, + get_bedrock_base_model, + get_bedrock_cross_region_inference_regions, + strip_bedrock_routing_prefix, +) +from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter + + +class TestStripBedrockRoutingPrefix: + """Tests for strip_bedrock_routing_prefix function.""" + + def test_strips_bedrock_prefix(self): + assert strip_bedrock_routing_prefix("bedrock/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_converse_prefix(self): + assert strip_bedrock_routing_prefix("converse/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_invoke_prefix(self): + assert strip_bedrock_routing_prefix("invoke/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_openai_prefix(self): + assert strip_bedrock_routing_prefix("openai/gpt-4") == "gpt-4" + + def test_strips_all_known_prefixes(self): + # Function strips all known prefixes iteratively + # bedrock/converse/model -> converse/model -> model + assert strip_bedrock_routing_prefix("bedrock/converse/claude-3") == "claude-3" + + def test_no_prefix_unchanged(self): + assert strip_bedrock_routing_prefix("claude-3-sonnet") == "claude-3-sonnet" + + def test_model_with_dots_unchanged(self): + assert ( + strip_bedrock_routing_prefix("anthropic.claude-3-sonnet-20240229-v1:0") + == "anthropic.claude-3-sonnet-20240229-v1:0" + ) + + +class TestExtractModelNameFromBedrockArn: + """Tests for extract_model_name_from_bedrock_arn function.""" + + def test_extracts_from_provisioned_model_arn(self): + arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model-id" + assert extract_model_name_from_bedrock_arn(arn) == "my-model-id" + + def test_extracts_from_foundation_model_arn(self): + arn = "arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2" + assert extract_model_name_from_bedrock_arn(arn) == "anthropic.claude-v2" + + def test_non_arn_unchanged(self): + model = "anthropic.claude-3-sonnet-20240229-v1:0" + assert extract_model_name_from_bedrock_arn(model) == model + + def test_case_insensitive_arn_detection(self): + arn = "ARN:aws:bedrock:us-east-1:123456789012:model/my-model" + assert extract_model_name_from_bedrock_arn(arn) == "my-model" + + +class TestGetBedrockCrossRegionInferenceRegions: + """Tests for get_bedrock_cross_region_inference_regions function.""" + + def test_returns_expected_regions(self): + regions = get_bedrock_cross_region_inference_regions() + assert "us" in regions + assert "eu" in regions + assert "global" in regions + assert "apac" in regions + + def test_returns_list(self): + regions = get_bedrock_cross_region_inference_regions() + assert isinstance(regions, list) + + +class TestGetBedrockBaseModel: + """Tests for get_bedrock_base_model function.""" + + def test_strips_bedrock_prefix(self): + assert get_bedrock_base_model("bedrock/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_converse_prefix(self): + assert get_bedrock_base_model("bedrock/converse/claude-3-sonnet") == "claude-3-sonnet" + + def test_strips_us_region_prefix(self): + # us.anthropic.model -> anthropic.model + assert ( + get_bedrock_base_model("us.anthropic.claude-3-sonnet-20240229-v1:0") + == "anthropic.claude-3-sonnet-20240229-v1:0" + ) + + def test_strips_eu_region_prefix(self): + assert ( + get_bedrock_base_model("eu.anthropic.claude-3-sonnet-20240229-v1:0") + == "anthropic.claude-3-sonnet-20240229-v1:0" + ) + + def test_extracts_from_arn(self): + arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model" + assert get_bedrock_base_model(arn) == "my-model" + + def test_model_without_prefix_unchanged(self): + model = "anthropic.claude-3-sonnet-20240229-v1:0" + assert get_bedrock_base_model(model) == model + + def test_combined_bedrock_and_region_prefix(self): + # bedrock/us.anthropic.model -> anthropic.model + assert ( + get_bedrock_base_model("bedrock/us.anthropic.claude-3-sonnet-20240229-v1:0") + == "anthropic.claude-3-sonnet-20240229-v1:0" + ) + + +class TestBedrockModelInfoWrappers: + """Tests that BedrockModelInfo methods correctly wrap standalone functions.""" + + def test_get_base_model_matches_standalone(self): + test_cases = [ + "bedrock/claude-3-sonnet", + "us.anthropic.claude-3-sonnet-20240229-v1:0", + "arn:aws:bedrock:us-east-1:123:model/my-model", + ] + for model in test_cases: + assert BedrockModelInfo.get_base_model(model) == get_bedrock_base_model(model) + + def test_extract_model_name_from_arn_matches_standalone(self): + arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model" + assert ( + BedrockModelInfo.extract_model_name_from_arn(arn) + == extract_model_name_from_bedrock_arn(arn) + ) + + def test_get_non_litellm_routing_model_name_matches_standalone(self): + model = "bedrock/converse/claude-3" + assert ( + BedrockModelInfo.get_non_litellm_routing_model_name(model) + == strip_bedrock_routing_prefix(model) + ) + + +class TestBedrockTokenCounter: + """Tests for BedrockTokenCounter class.""" + + def test_should_use_token_counting_api_for_bedrock(self): + counter = BedrockTokenCounter() + assert counter.should_use_token_counting_api("bedrock") is True + + def test_should_not_use_token_counting_api_for_other_providers(self): + counter = BedrockTokenCounter() + assert counter.should_use_token_counting_api("openai") is False + assert counter.should_use_token_counting_api("anthropic") is False + assert counter.should_use_token_counting_api(None) is False + + def test_get_token_counter_returns_bedrock_token_counter(self): + model_info = BedrockModelInfo() + token_counter = model_info.get_token_counter() + assert isinstance(token_counter, BedrockTokenCounter) + + @pytest.mark.asyncio + async def test_count_tokens_returns_none_for_empty_messages(self): + counter = BedrockTokenCounter() + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=None, + contents=None, + ) + assert result is None + + result = await counter.count_tokens( + model_to_use="anthropic.claude-3-sonnet", + messages=[], + contents=None, + ) + assert result is None diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index ec510b8f953..29c2d26981c 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -2835,6 +2835,34 @@ def test_bedrock_invoke_provider(): ) == "nova" ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider("amazon.nova-pro-v1:0") + == "nova" + ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider( + "amazon.nova-lite-v1:0" + ) + == "nova" + ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider( + "amazon.nova-micro-v1:0" + ) + == "nova" + ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider( + "amazon.nova-premier-v1:0" + ) + == "nova" + ) + assert ( + litellm.AmazonInvokeConfig().get_bedrock_invoke_provider( + "amazon.nova-2-lite-v1:0" + ) + == "nova" + ) def test_bedrock_description_param(): @@ -3488,7 +3516,9 @@ def test_bedrock_openai_imported_model(): url = mock_post.call_args.kwargs["url"] print(f"URL: {url}") assert "bedrock-runtime.us-east-1.amazonaws.com" in url - assert "arn:aws:bedrock:us-east-1:117159858402:imported-model/m4gc1mrfuddy" in url + assert ( + "arn:aws:bedrock:us-east-1:117159858402:imported-model/m4gc1mrfuddy" in url + ) assert "/invoke" in url # Validate request body follows OpenAI format @@ -3517,7 +3547,9 @@ def test_bedrock_openai_imported_model(): # Check image_url content assert user_msg["content"][1]["type"] == "image_url" assert "image_url" in user_msg["content"][1] - assert user_msg["content"][1]["image_url"]["url"].startswith("data:image/jpeg;base64,") + assert user_msg["content"][1]["image_url"]["url"].startswith( + "data:image/jpeg;base64," + ) assert user_msg["content"][2]["type"] == "image_url" assert "image_url" in user_msg["content"][2] @@ -3526,21 +3558,67 @@ def test_bedrock_openai_imported_model(): assert request_body["max_tokens"] == 300 assert request_body["temperature"] == 0.5 + +def test_bedrock_nova_provider_detection(): + """ + Test that Nova models are correctly detected even when prefixed with "amazon." + Regression test for issue #17910 where models like "amazon.nova-pro-v1:0" + were incorrectly identified as "amazon" (Titan) instead of "nova". + """ + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + # Test various Nova model formats + nova_test_cases = [ + ("us.amazon.nova-pro-v1:0", "nova"), + ("us.amazon.nova-lite-v1:0", "nova"), + ("us.amazon.nova-micro-v1:0", "nova"), + ("amazon.nova-pro-v1:0", "nova"), + ("amazon.nova-lite-v1:0", "nova"), + ("amazon.nova-micro-v1:0", "nova"), + ("amazon.nova-premier-v1:0", "nova"), + ("amazon.nova-2-lite-v1:0", "nova"), + ("bedrock/amazon.nova-pro-v1:0", "nova"), + ("bedrock/invoke/amazon.nova-pro-v1:0", "nova"), + ("amazon.Nova-pro-v1:0", "nova"), + ("amazon.NOVA-pro-v1:0", "nova"), + ] + + for model, expected in nova_test_cases: + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + assert ( + provider == expected + ), f"Failed for model: {model}, expected: {expected}, got: {provider}" + + # Verify that Amazon Titan models still return "amazon" + titan_test_cases = [ + ("amazon.titan-text-express-v1", "amazon"), + ("us.amazon.titan-text-lite-v1", "amazon"), + ] + + for model, expected in titan_test_cases: + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + assert ( + provider == expected + ), f"Failed for model: {model}, expected: {expected}, got: {provider}" + + def test_bedrock_openai_provider_detection(): """ Test that the OpenAI provider is correctly detected from model strings. """ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - + # Test various OpenAI model formats test_cases = [ "openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123", "bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/xyz789", ] - + for model in test_cases: provider = BaseAWSLLM.get_bedrock_invoke_provider(model) - assert provider == "openai", f"Failed for model: {model}, got provider: {provider}" + assert ( + provider == "openai" + ), f"Failed for model: {model}, got provider: {provider}" print(f"✓ Provider detection works for: {model}") @@ -3549,16 +3627,16 @@ def test_bedrock_openai_model_id_extraction(): Test that the model ID (ARN) is correctly extracted and encoded for OpenAI models. """ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - - model = "openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-model-123" - provider = BaseAWSLLM.get_bedrock_invoke_provider(model) - - model_id = BaseAWSLLM.get_bedrock_model_id( - model=model, - provider=provider, - optional_params={} + + model = ( + "openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-model-123" ) - + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + + model_id = BaseAWSLLM.get_bedrock_model_id( + model=model, provider=provider, optional_params={} + ) + # The ARN should be double URL encoded assert "arn" in model_id assert "imported-model" in model_id @@ -3570,20 +3648,17 @@ def test_bedrock_openai_convert_messages_to_prompt(): Test that convert_messages_to_prompt returns empty string for OpenAI models. """ from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM - + bedrock_llm = BedrockLLM() messages = [ {"role": "system", "content": "You are helpful"}, - {"role": "user", "content": "Hello"} + {"role": "user", "content": "Hello"}, ] - + prompt, chat_history = bedrock_llm.convert_messages_to_prompt( - model="test-model", - messages=messages, - provider="openai", - custom_prompt_dict={} + model="test-model", messages=messages, provider="openai", custom_prompt_dict={} ) - + # OpenAI models use messages directly, no prompt conversion assert prompt == "" assert chat_history is None @@ -3598,37 +3673,33 @@ def test_bedrock_openai_response_parsing(): from litellm import ModelResponse from unittest.mock import Mock import json - + bedrock_llm = BedrockLLM() - + # Mock OpenAI-style response openai_response = { "choices": [ { "message": { "content": "The capital of France is Paris.", - "role": "assistant" + "role": "assistant", }, "finish_reason": "stop", - "index": 0 + "index": 0, } ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 8, - "total_tokens": 18 - } + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, } - + mock_response = Mock() mock_response.json.return_value = openai_response mock_response.text = json.dumps(openai_response) mock_response.status_code = 200 mock_response.headers = {} - + model_response = ModelResponse() mock_logging = Mock() - + result = bedrock_llm.process_response( model="openai/arn:aws:bedrock:us-east-1:123:imported-model/test", response=mock_response, @@ -3640,18 +3711,18 @@ def test_bedrock_openai_response_parsing(): data={}, messages=[{"role": "user", "content": "What is the capital of France?"}], print_verbose=lambda x: None, - encoding=None + encoding=None, ) - + # Verify response content assert result.choices[0].message.content == "The capital of France is Paris." assert result.choices[0].finish_reason == "stop" - + # Verify usage assert result.usage.prompt_tokens == 10 assert result.usage.completion_tokens == 8 assert result.usage.total_tokens == 18 - + print("✓ OpenAI response parsing works correctly") @@ -3659,45 +3730,47 @@ def test_bedrock_openai_request_transformation(): """ Test that the request is correctly transformed for OpenAI models. """ - from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig - + from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig, + ) + config = AmazonInvokeConfig() - + model = "openai/arn:aws:bedrock:us-east-1:123:imported-model/test" messages = [ {"role": "system", "content": "You are helpful"}, - {"role": "user", "content": "Hello"} + {"role": "user", "content": "Hello"}, ] - + optional_params = { "max_tokens": 100, "temperature": 0.7, "top_p": 0.9, - "stream": False + "stream": False, } - + litellm_params = {} headers = {} - - with patch.object(config, 'get_bedrock_invoke_provider', return_value="openai"): + + with patch.object(config, "get_bedrock_invoke_provider", return_value="openai"): result = config.transform_request( model=model, messages=messages, optional_params=optional_params.copy(), litellm_params=litellm_params, - headers=headers + headers=headers, ) - + # Verify the request uses messages format (not prompt) assert "messages" in result assert len(result["messages"]) == 2 assert result["messages"][0]["role"] == "system" assert result["messages"][1]["role"] == "user" - + # Verify parameters are included assert "max_tokens" in result assert "temperature" in result - + print("✓ Request transformation works correctly") @@ -3705,20 +3778,22 @@ def test_bedrock_openai_parameter_filtering(): """ Test that only supported OpenAI parameters are included in the request. """ - from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig - + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, + ) + config = AmazonBedrockOpenAIConfig() model = "test-model" - + supported_params = config.get_supported_openai_params(model=model) - + # Verify common OpenAI parameters are supported assert "max_tokens" in supported_params assert "temperature" in supported_params assert "top_p" in supported_params assert "stream" in supported_params assert "stop" in supported_params - + print(f"✓ Parameter filtering supports: {len(supported_params)} parameters") print(f" Supported params: {supported_params}") @@ -3728,12 +3803,12 @@ def test_bedrock_openai_route_detection(): Test that the OpenAI route is correctly detected. """ from litellm.llms.bedrock.common_utils import BedrockModelInfo - + test_cases = [ ("openai/arn:aws:bedrock:us-east-1:123:imported-model/test", "openai"), ("bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test", "openai"), ] - + for model, expected_route in test_cases: route = BedrockModelInfo.get_bedrock_route(model) assert route == expected_route, f"Failed for model: {model}, got route: {route}" @@ -3745,15 +3820,30 @@ def test_bedrock_openai_explicit_route_check(): Test the explicit OpenAI route checker helper method. """ from litellm.llms.bedrock.common_utils import BedrockModelInfo - + # Test with openai/ prefix - assert BedrockModelInfo._explicit_openai_route("openai/arn:aws:bedrock:us-east-1:123:imported-model/test") is True - assert BedrockModelInfo._explicit_openai_route("bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test") is True - + assert ( + BedrockModelInfo._explicit_openai_route( + "openai/arn:aws:bedrock:us-east-1:123:imported-model/test" + ) + is True + ) + assert ( + BedrockModelInfo._explicit_openai_route( + "bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test" + ) + is True + ) + # Test without openai/ prefix assert BedrockModelInfo._explicit_openai_route("anthropic.claude-3-sonnet") is False - assert BedrockModelInfo._explicit_openai_route("arn:aws:bedrock:us-east-1:123:imported-model/test") is False - + assert ( + BedrockModelInfo._explicit_openai_route( + "arn:aws:bedrock:us-east-1:123:imported-model/test" + ) + is False + ) + print("✓ Explicit route check works correctly") @@ -3761,16 +3851,18 @@ def test_bedrock_openai_config_initialization(): """ Test that AmazonBedrockOpenAIConfig can be properly initialized. """ - from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig - + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, + ) + config = AmazonBedrockOpenAIConfig() - + # Verify it has the necessary methods - assert hasattr(config, 'get_supported_openai_params') - assert hasattr(config, 'transform_request') - assert hasattr(config, 'transform_response') - assert hasattr(config, 'map_openai_params') - + assert hasattr(config, "get_supported_openai_params") + assert hasattr(config, "transform_request") + assert hasattr(config, "transform_response") + assert hasattr(config, "map_openai_params") + print("✓ AmazonBedrockOpenAIConfig initializes correctly") @@ -3779,9 +3871,9 @@ def test_bedrock_openai_multiple_message_types(): Test that various message content types are handled correctly. """ from litellm.llms.custom_httpx.http_handler import HTTPHandler - + client = HTTPHandler() - + # Test with mixed content types messages = [ {"role": "system", "content": "You are helpful"}, @@ -3790,11 +3882,14 @@ def test_bedrock_openai_multiple_message_types(): "role": "user", "content": [ {"type": "text", "text": "Complex message with text"}, - {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,iVBORw0KGg"}} - ] - } + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,iVBORw0KGg"}, + }, + ], + }, ] - + with patch.object(client, "post") as mock_post: try: response = completion( @@ -3805,18 +3900,18 @@ def test_bedrock_openai_multiple_message_types(): ) except Exception as e: pass - + # Verify the request was made if mock_post.called: request_body = json.loads(mock_post.call_args.kwargs["data"]) - + # Verify messages are preserved assert "messages" in request_body assert len(request_body["messages"]) == 3 - + # Verify mixed content is handled assert isinstance(request_body["messages"][2]["content"], list) - + print("✓ Multiple message types handled correctly") @@ -3829,18 +3924,18 @@ def test_bedrock_openai_error_handling(): from litellm.llms.bedrock.common_utils import BedrockError from unittest.mock import Mock import json - + bedrock_llm = BedrockLLM() - + # Mock error response mock_response = Mock() mock_response.json.side_effect = Exception("Invalid JSON") mock_response.text = "Invalid response" mock_response.status_code = 422 - + model_response = ModelResponse() mock_logging = Mock() - + with pytest.raises(BedrockError) as exc_info: bedrock_llm.process_response( model="openai/arn:aws:bedrock:us-east-1:123:imported-model/test", @@ -3853,8 +3948,8 @@ def test_bedrock_openai_error_handling(): data={}, messages=[], print_verbose=lambda x: None, - encoding=None + encoding=None, ) - + assert exc_info.value.status_code == 422 print("✓ Error handling works correctly") diff --git a/tests/llm_translation/test_openrouter.py b/tests/llm_translation/test_openrouter.py index 839d08e12bf..105b05d3449 100644 --- a/tests/llm_translation/test_openrouter.py +++ b/tests/llm_translation/test_openrouter.py @@ -32,3 +32,18 @@ def test_completion_openrouter_image_generation(): .message.images[0]["image_url"]["url"] .startswith("data:image/png;base64,") ) + + +def test_openrouter_embedding(): + """Test OpenRouter embeddings support.""" + litellm._turn_on_debug() + resp = litellm.embedding( + model="openrouter/openai/text-embedding-3-small", + input=["Hello world", "How are you?"], + ) + print(resp) + assert resp is not None + assert len(resp.data) == 2 + assert resp.data[0]["embedding"] is not None + assert isinstance(resp.data[0]["embedding"], list) + assert len(resp.data[0]["embedding"]) > 0 diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py index 31a2f6cbf51..b0aac17e7d9 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero.py @@ -66,16 +66,77 @@ class TestCloudZeroHourlyExport: fake_client = MagicMock() fake_db = MagicMock() - async def query_raw_mock(query: str): - sql_context = pl.SQLContext( - LiteLLM_DailyUserSpend=spend_mock_data, - LiteLLM_VerificationToken=verification_mock_data, - LiteLLM_TeamTable=team_mock_data, - LiteLLM_UserTable=user_mock_data, - ) - result = sql_context.execute(query).collect() + async def query_raw_mock(query: str, *params): + start_time_utc = params[0] if len(params) > 0 else None + end_time_utc = params[1] if len(params) > 1 else None + limit = params[2] if len(params) > 2 else None - return result + spend_df = spend_mock_data.collect() + verification_df = verification_mock_data.collect().rename( + {"key_alias": "api_key_alias"} + ) + team_df = team_mock_data.collect() + user_df = user_mock_data.collect() + + joined = ( + spend_df.join( + verification_df, left_on="api_key", right_on="token", how="left" + ) + .join( + team_df, + left_on="team_id", + right_on="team_id", + how="left", + suffix="_team", + ) + .join( + user_df, + left_on="user_id", + right_on="user_id", + how="left", + suffix="_user", + ) + ) + + for duplicate_column in ("team_id_team", "user_id_user"): + if duplicate_column in joined.columns: + joined = joined.drop(duplicate_column) + + if start_time_utc is not None: + joined = joined.filter(pl.col("updated_at") >= start_time_utc) + if end_time_utc is not None: + joined = joined.filter(pl.col("updated_at") <= end_time_utc) + + joined = joined.select( + [ + "id", + "date", + "user_id", + "api_key", + "model", + "model_group", + "custom_llm_provider", + "prompt_tokens", + "completion_tokens", + "spend", + "api_requests", + "successful_requests", + "failed_requests", + "cache_creation_input_tokens", + "cache_read_input_tokens", + "created_at", + "updated_at", + "team_id", + "api_key_alias", + "team_alias", + "user_email", + ] + ).sort(["date", "created_at"], descending=[True, True]) + + if limit is not None: + joined = joined.head(int(limit)) + + return joined fake_db.query_raw = AsyncMock(side_effect=query_raw_mock) fake_client.db = fake_db diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py new file mode 100644 index 00000000000..89a5028011c --- /dev/null +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py @@ -0,0 +1,57 @@ +"""Tests for LiteLLM CloudZero database helper.""" + +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from litellm.integrations.cloudzero.database import LiteLLMDatabase + + +def _setup_db(monkeypatch: pytest.MonkeyPatch, query_return): + """Return a database instance with prisma client mocked out.""" + query_mock = AsyncMock(return_value=query_return) + mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=query_mock)) + db = LiteLLMDatabase() + monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client) + return db, query_mock + + +@pytest.mark.asyncio +async def test_get_usage_data_parameterized(monkeypatch: pytest.MonkeyPatch): + """Start/end filters and limit should be parameterized via placeholders.""" + start = datetime(2024, 5, 1, tzinfo=timezone.utc) + end = datetime(2024, 5, 2, tzinfo=timezone.utc) + db, query_mock = _setup_db(monkeypatch, []) + + await db.get_usage_data(limit=10, start_time_utc=start, end_time_utc=end) + + query_text, *params = query_mock.await_args.args + assert "dus.updated_at >= $1::timestamptz" in query_text + assert "dus.updated_at <= $2::timestamptz" in query_text + assert "LIMIT $3" in query_text + assert params == [start, end, 10] + + +@pytest.mark.asyncio +async def test_get_usage_data_handles_missing_filters(monkeypatch: pytest.MonkeyPatch): + """When no filters provided the params should be None placeholders.""" + db, query_mock = _setup_db(monkeypatch, []) + + await db.get_usage_data() + + query_text, *params = query_mock.await_args.args + assert "LIMIT $3" not in query_text + assert params == [None, None] + + +@pytest.mark.asyncio +async def test_get_usage_data_rejects_invalid_limit(monkeypatch: pytest.MonkeyPatch): + """limit must coerce to int or raise ValueError before hitting the DB.""" + db, query_mock = _setup_db(monkeypatch, []) + + with pytest.raises(ValueError): + await db.get_usage_data(limit="invalid") + + assert query_mock.await_count == 0 diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py new file mode 100644 index 00000000000..5ee98cc9dd0 --- /dev/null +++ b/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 diff --git a/tests/test_litellm/integrations/focus/test_s3_destination.py b/tests/test_litellm/integrations/focus/test_s3_destination.py new file mode 100644 index 00000000000..f915b2c56a3 --- /dev/null +++ b/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", + } diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index a322dfe9a2b..d7d7720ff43 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -530,7 +530,7 @@ class TestPassthroughCallTypeHandling: ) assert ( ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aembedding") - == "embeddings" + == "embedding" ) assert ( ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aresponses") diff --git a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py b/tests/test_litellm/integrations/test_prometheus_cache_metrics.py new file mode 100644 index 00000000000..660757673f6 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_cache_metrics.py @@ -0,0 +1,211 @@ +""" +Unit tests for cache Prometheus metrics. + +Run with: poetry run pytest tests/test_litellm/integrations/test_prometheus_cache_metrics.py -v +""" +import pytest +from unittest.mock import MagicMock, patch +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + +class TestPrometheusCacheMetrics: + """Tests for cache-related Prometheus metrics""" + + @pytest.fixture + def sample_enum_values(self): + """Create sample enum values for labels""" + return UserAPIKeyLabelValues( + end_user="test-end-user", + hashed_api_key="test-key-hash", + api_key_alias="test-key-alias", + team="test-team", + team_alias="test-team-alias", + user="test-user", + model="gpt-3.5-turbo", + ) + + def test_cache_metrics_defined_in_types(self): + """Test that cache metrics are defined in DEFINED_PROMETHEUS_METRICS""" + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + from typing import get_args + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + + assert "litellm_cache_hits_metric" in defined_metrics + assert "litellm_cache_misses_metric" in defined_metrics + assert "litellm_cached_tokens_metric" in defined_metrics + + def test_cache_metric_labels_defined(self): + """Test that cache metric labels are properly defined""" + from litellm.types.integrations.prometheus import PrometheusMetricLabels + + # Verify labels are defined for each cache metric + assert hasattr(PrometheusMetricLabels, "litellm_cache_hits_metric") + assert hasattr(PrometheusMetricLabels, "litellm_cache_misses_metric") + assert hasattr(PrometheusMetricLabels, "litellm_cached_tokens_metric") + + # Verify labels include expected keys + expected_labels = [ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + for label in expected_labels: + assert label in PrometheusMetricLabels.litellm_cache_hits_metric + assert label in PrometheusMetricLabels.litellm_cache_misses_metric + assert label in PrometheusMetricLabels.litellm_cached_tokens_metric + + def test_increment_cache_metrics_on_cache_hit(self, sample_enum_values): + """Test that cache hit increments the correct metrics""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + # Import the method directly and bind it to our mock + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=True + standard_logging_payload = { + "cache_hit": True, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method using unbound method approach + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify cache hits metric was incremented + mock_logger.litellm_cache_hits_metric.labels.assert_called() + mock_logger.litellm_cache_hits_metric.labels().inc.assert_called_once() + + # Verify cached tokens metric was incremented with total_tokens + mock_logger.litellm_cached_tokens_metric.labels.assert_called() + mock_logger.litellm_cached_tokens_metric.labels().inc.assert_called_once_with( + 100 + ) + + # Verify cache misses metric was NOT called + mock_logger.litellm_cache_misses_metric.labels.assert_not_called() + + def test_increment_cache_metrics_on_cache_miss(self, sample_enum_values): + """Test that cache miss increments the correct metrics""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=False + standard_logging_payload = { + "cache_hit": False, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify cache misses metric was incremented + mock_logger.litellm_cache_misses_metric.labels.assert_called() + mock_logger.litellm_cache_misses_metric.labels().inc.assert_called_once() + + # Verify cache hits and cached tokens metrics were NOT called + mock_logger.litellm_cache_hits_metric.labels.assert_not_called() + mock_logger.litellm_cached_tokens_metric.labels.assert_not_called() + + def test_increment_cache_metrics_when_cache_hit_is_none(self, sample_enum_values): + """Test that no metrics are incremented when cache_hit is None""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=None + standard_logging_payload = { + "cache_hit": None, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify NO metrics were called + mock_logger.litellm_cache_hits_metric.labels.assert_not_called() + mock_logger.litellm_cache_misses_metric.labels.assert_not_called() + mock_logger.litellm_cached_tokens_metric.labels.assert_not_called() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py b/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py new file mode 100644 index 00000000000..0743a9c7ba2 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_queue_guardrail_metrics.py @@ -0,0 +1,424 @@ +""" +Unit tests for prometheus queue time and guardrail metrics +""" +from datetime import datetime +from unittest.mock import MagicMock + +import pytest +from prometheus_client import REGISTRY + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + """Clean up prometheus registry between tests""" + # Clear the registry before each test + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + yield + # Clean up after test + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + + +class TestPrometheusQueueTimeMetric: + """Test request queue time metric recording""" + + def test_queue_time_metric_recorded_in_set_latency_metrics(self): + """Test that queue time metric is recorded when queue_time_seconds is present in metadata""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs with queue_time_seconds in metadata + queue_time_seconds = 0.5 + + kwargs = { + "litellm_params": {"metadata": {"queue_time_seconds": queue_time_seconds}}, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should be called + mock_metric.labels.assert_called() + # Check that observe was called on the queue time metric + assert mock_labeled_metric.observe.called + # Verify the observed value + observed_value = None + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0: + observed_value = call[0][0] + if observed_value == queue_time_seconds: + break + assert observed_value == queue_time_seconds + assert observed_value >= 0 + + def test_queue_time_metric_not_recorded_when_missing(self): + """Test that queue time metric is not recorded when queue_time_seconds is missing""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs without queue_time_seconds + kwargs = { + "litellm_params": {"metadata": {}}, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should not be called (queue_time_seconds is None) + # We check that observe was not called with queue_time_seconds + queue_time_called = False + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0 and call[0][0] == 0.5: # Our test queue time value + queue_time_called = True + break + assert ( + not queue_time_called + ), "Queue time metric should not be recorded when queue_time_seconds is missing" + + def test_queue_time_metric_not_recorded_when_negative(self): + """Test that queue time metric is not recorded when queue_time_seconds is negative""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock the metric + mock_metric = MagicMock() + mock_labeled_metric = MagicMock() + mock_metric.labels.return_value = mock_labeled_metric + prometheus_logger.litellm_request_queue_time_metric = mock_metric + + # Create mock kwargs with negative queue_time_seconds + kwargs = { + "litellm_params": { + "metadata": {"queue_time_seconds": -0.1} # Negative value + }, + "model": "gpt-3.5-turbo", + "start_time": datetime.now(), + "end_time": datetime.now(), + } + + enum_values = UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-3.5-turbo", + model_group="gpt-3.5-turbo", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-3.5-turbo", + litellm_model_name="gpt-3.5-turbo", + tags=[], + model_id="gpt-3.5-turbo", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + # Act + prometheus_logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-3.5-turbo", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # Assert - queue time metric should not be called for negative values + # We check that observe was not called with the negative value + negative_value_called = False + for call in mock_labeled_metric.observe.call_args_list: + if len(call[0]) > 0 and call[0][0] == -0.1: + negative_value_called = True + break + assert ( + not negative_value_called + ), "Queue time metric should not be recorded for negative values" + + +class TestPrometheusGuardrailMetrics: + """Test guardrail metrics recording""" + + def test_record_guardrail_metrics_success(self): + """Test recording guardrail metrics for successful execution""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + mock_errors_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + prometheus_logger.litellm_guardrail_errors_total = mock_errors_metric + + guardrail_name = "test_guardrail" + latency_seconds = 0.15 + status = "success" + error_type = None + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + + # Assert - latency metric should be recorded + mock_latency_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + error_type="none", + hook_type=hook_type, + ) + mock_latency_metric.labels.return_value.observe.assert_called_once_with( + latency_seconds + ) + + # Assert - requests metric should be incremented + mock_requests_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + hook_type=hook_type, + ) + mock_requests_metric.labels.return_value.inc.assert_called_once() + + # Assert - errors metric should NOT be called for success + mock_errors_metric.labels.assert_not_called() + + def test_record_guardrail_metrics_error(self): + """Test recording guardrail metrics for failed execution""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + mock_errors_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + prometheus_logger.litellm_guardrail_errors_total = mock_errors_metric + + guardrail_name = "test_guardrail" + latency_seconds = 0.2 + status = "error" + error_type = "ValueError" + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + + # Assert - latency metric should be recorded + mock_latency_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + mock_latency_metric.labels.return_value.observe.assert_called_once_with( + latency_seconds + ) + + # Assert - requests metric should be incremented + mock_requests_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + status=status, + hook_type=hook_type, + ) + mock_requests_metric.labels.return_value.inc.assert_called_once() + + # Assert - errors metric should be incremented + mock_errors_metric.labels.assert_called_once_with( + guardrail_name=guardrail_name, + error_type=error_type, + hook_type=hook_type, + ) + mock_errors_metric.labels.return_value.inc.assert_called_once() + + def test_record_guardrail_metrics_during_call_hook(self): + """Test recording guardrail metrics for during_call hook""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + + guardrail_name = "moderation_guardrail" + latency_seconds = 0.1 + status = "success" + hook_type = "during_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=None, + hook_type=hook_type, + ) + + # Assert - hook_type should be "during_call" + mock_latency_metric.labels.assert_called_once() + call_kwargs = mock_latency_metric.labels.call_args[1] + assert call_kwargs["hook_type"] == "during_call" + + def test_record_guardrail_metrics_handles_exception(self): + """Test that _record_guardrail_metrics handles exceptions gracefully""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metric to raise exception + mock_metric = MagicMock() + mock_metric.labels.side_effect = Exception("Test error") + prometheus_logger.litellm_guardrail_latency_metric = mock_metric + prometheus_logger.litellm_guardrail_requests_total = MagicMock() + + # Act & Assert - should not raise exception + try: + prometheus_logger._record_guardrail_metrics( + guardrail_name="test", + latency_seconds=0.1, + status="success", + error_type=None, + hook_type="pre_call", + ) + except Exception: + pytest.fail("_record_guardrail_metrics should handle exceptions gracefully") + + def test_record_guardrail_metrics_with_guardrail_name_attribute(self): + """Test that guardrail name is extracted from guardrail_name attribute if available""" + # Arrange + prometheus_logger = PrometheusLogger() + + # Mock metrics + mock_latency_metric = MagicMock() + mock_requests_metric = MagicMock() + + prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric + prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric + + guardrail_name = "custom_guardrail_name" + latency_seconds = 0.1 + status = "success" + hook_type = "pre_call" + + # Act + prometheus_logger._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=None, + hook_type=hook_type, + ) + + # Assert - guardrail_name should be used + mock_latency_metric.labels.assert_called_once() + call_kwargs = mock_latency_metric.labels.call_args[1] + assert call_kwargs["guardrail_name"] == guardrail_name diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 4914ec0bfb7..42a2b5d0971 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -620,26 +620,26 @@ def test_bedrock_tools_unpack_defs(): def test_bedrock_image_processor_content_type_fallback_url_extension(): """ - Test that _post_call_image_processing falls back to URL extension + Test that _post_call_image_processing falls back to URL extension when content-type is binary/octet-stream or application/octet-stream """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a simple PNG header (magic bytes) png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" png_content = png_header + b"\x00" * 100 # Add some padding mock_response.content = png_content - + # Test with .png URL image_url = "https://example.com/test-image.png" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/png" assert base64_bytes == base64.b64encode(png_content).decode("utf-8") @@ -650,22 +650,22 @@ def test_bedrock_image_processor_content_type_fallback_binary_detection(): when content-type is missing and URL extension is not recognized """ import base64 - + # Create mock response with no content-type mock_response = MagicMock() mock_response.headers.get.return_value = None - + # Create a JPEG header (magic bytes) jpeg_header = b"\xff\xd8\xff" jpeg_content = jpeg_header + b"\x00" * 100 # Add some padding mock_response.content = jpeg_content - + # Test with URL without extension image_url = "https://example.com/test-image-without-extension" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/jpeg" assert base64_bytes == base64.b64encode(jpeg_content).decode("utf-8") @@ -675,22 +675,22 @@ def test_bedrock_image_processor_content_type_fallback_application_octet_stream( Test that _post_call_image_processing handles application/octet-stream correctly """ import base64 - + # Create mock response with application/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "application/octet-stream" - + # Create a GIF header (magic bytes) gif_header = b"GIF8" + b"\x00" + b"a" gif_content = gif_header + b"\x00" * 100 # Add some padding mock_response.content = gif_content - + # Test with .gif URL image_url = "https://s3.amazonaws.com/bucket/image.gif" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/gif" assert base64_bytes == base64.b64encode(gif_content).decode("utf-8") @@ -700,22 +700,22 @@ def test_bedrock_image_processor_content_type_with_query_params(): Test that _post_call_image_processing correctly extracts extension from URL with query parameters """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a WebP header (magic bytes) webp_header = b"RIFF" + b"\x00\x00\x00\x00" + b"WEBP" webp_content = webp_header + b"\x00" * 100 # Add some padding mock_response.content = webp_content - + # Test with URL containing query parameters (common in S3 signed URLs) image_url = "https://s3.amazonaws.com/bucket/image.webp?AWSAccessKeyId=123&Expires=456&Signature=789" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/webp" assert base64_bytes == base64.b64encode(webp_content).decode("utf-8") @@ -725,21 +725,21 @@ def test_bedrock_image_processor_content_type_normal_header(): Test that _post_call_image_processing works normally when content-type is correctly set """ import base64 - + # Create mock response with correct content-type mock_response = MagicMock() mock_response.headers.get.return_value = "image/png" - + # Create a PNG header png_header = b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a" png_content = png_header + b"\x00" * 100 mock_response.content = png_content - + image_url = "https://example.com/test-image.png" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, image_url ) - + assert content_type == "image/png" assert base64_bytes == base64.b64encode(png_content).decode("utf-8") @@ -751,16 +751,16 @@ def test_bedrock_image_processor_content_type_fallback_failure(): # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create content with unrecognizable image format mock_response.content = b"\x00" * 100 - + # Test with URL without recognizable extension image_url = "https://example.com/unknown-file" - + with pytest.raises(ValueError) as excinfo: BedrockImageProcessor._post_call_image_processing(mock_response, image_url) - + assert "Unable to determine content type" in str(excinfo.value) @@ -771,18 +771,18 @@ def test_bedrock_image_processor_content_type_jpeg_variants(): # Create mock response with binary/octet-stream mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + jpeg_header = b"\xff\xd8\xff" jpeg_content = jpeg_header + b"\x00" * 100 mock_response.content = jpeg_content - + # Test with .jpg extension image_url_jpg = "https://example.com/photo.jpg" _, content_type_jpg = BedrockImageProcessor._post_call_image_processing( mock_response, image_url_jpg ) assert content_type_jpg == "image/jpeg" - + # Test with .jpeg extension image_url_jpeg = "https://example.com/photo.jpeg" _, content_type_jpeg = BedrockImageProcessor._post_call_image_processing( @@ -797,22 +797,22 @@ def test_bedrock_image_processor_content_type_pdf_document(): when content-type is binary/octet-stream """ import base64 - + # Create mock response with binary/octet-stream content-type mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + # Create a PDF header (magic bytes: %PDF) pdf_header = b"%PDF-1.4" pdf_content = pdf_header + b"\x00" * 100 mock_response.content = pdf_content - + # Test with .pdf URL pdf_url = "https://s3.amazonaws.com/bucket/document.pdf" base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, pdf_url ) - + assert content_type == "application/pdf" assert base64_bytes == base64.b64encode(pdf_content).decode("utf-8") @@ -822,12 +822,12 @@ def test_bedrock_image_processor_content_type_document_formats(): Test that _post_call_image_processing handles various document formats """ import base64 - + # Create mock response mock_response = MagicMock() mock_response.headers.get.return_value = "application/octet-stream" mock_response.content = b"\x00" * 100 - + # Test various document formats test_cases = [ ("https://example.com/doc.pdf", "application/pdf"), @@ -837,7 +837,7 @@ def test_bedrock_image_processor_content_type_document_formats(): ("https://example.com/page.html", "text/html"), ("https://example.com/readme.txt", "text/plain"), ] - + for url, expected_mime in test_cases: _, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, url @@ -850,21 +850,21 @@ def test_bedrock_image_processor_content_type_s3_pdf_with_query(): Test that _post_call_image_processing handles S3 PDF with query parameters """ import base64 - + # Create mock response mock_response = MagicMock() mock_response.headers.get.return_value = "binary/octet-stream" - + pdf_content = b"%PDF-1.4" + b"\x00" * 100 mock_response.content = pdf_content - + # S3 signed URL with query parameters s3_url = "https://my-bucket.s3.us-east-1.amazonaws.com/documents/report.pdf?AWSAccessKeyId=AKIAIOSFODNN7EXAMPLE&Expires=1234567890&Signature=abcdef123456" - + base64_bytes, content_type = BedrockImageProcessor._post_call_image_processing( mock_response, s3_url ) - + assert content_type == "application/pdf" assert base64_bytes == base64.b64encode(pdf_content).decode("utf-8") @@ -1139,6 +1139,170 @@ def test_bedrock_create_bedrock_block_different_document_formats(): assert block["document"]["format"] == format_type +def test_convert_to_anthropic_tool_result_image_with_cache_control(): + """ + Test that cache_control is properly applied to image content in tool results. + This tests the functionality added in the uncommitted changes where + add_cache_control_to_content is called for image_url content types. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + # Test with base64 image data URI + message = { + "role": "tool", + "tool_call_id": "call_test_123", + "content": [ + { + "type": "text", + "text": "Here is the image you requested:", + }, + { + "type": "image_url", + "image_url": "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQ", + "cache_control": {"type": "ephemeral"}, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify the result structure + assert result["type"] == "tool_result" + assert result["tool_use_id"] == "call_test_123" + assert isinstance(result["content"], list) + assert len(result["content"]) == 2 + + # Verify text content + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Here is the image you requested:" + + # Verify image content with cache_control + assert result["content"][1]["type"] == "image" + assert result["content"][1]["source"]["type"] == "base64" + assert result["content"][1]["source"]["media_type"] == "image/jpeg" + assert "cache_control" in result["content"][1] + assert result["content"][1]["cache_control"]["type"] == "ephemeral" + + +def test_convert_to_anthropic_tool_result_image_without_cache_control(): + """ + Test that images without cache_control in tool results work correctly. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_test_456", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAUA", + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify the result structure + assert result["type"] == "tool_result" + assert result["tool_use_id"] == "call_test_456" + assert isinstance(result["content"], list) + assert len(result["content"]) == 1 + + # Verify image content without cache_control (cache_control will be None if not set) + assert result["content"][0]["type"] == "image" + assert result["content"][0]["source"]["type"] == "base64" + assert result["content"][0]["source"]["media_type"] == "image/png" + assert result["content"][0].get("cache_control") is None + + +def test_convert_to_anthropic_tool_result_mixed_content_with_cache_control(): + """ + Test tool results with mixed content types (text and image) where only some have cache_control. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_test_789", + "content": [ + { + "type": "text", + "text": "First image:", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "image_url", + "image_url": "data:image/jpeg;base64,/9j/4AAQSkZJRg", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "text", + "text": "Second image (no cache):", + }, + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgo", + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + assert result["type"] == "tool_result" + assert isinstance(result["content"], list) + assert len(result["content"]) == 4 + + # First text with cache_control + assert result["content"][0]["type"] == "text" + assert result["content"][0]["cache_control"]["type"] == "ephemeral" + + # First image with cache_control + assert result["content"][1]["type"] == "image" + assert result["content"][1]["cache_control"]["type"] == "ephemeral" + + # Second text without cache_control (cache_control will be None if not set) + assert result["content"][2]["type"] == "text" + assert result["content"][2].get("cache_control") is None + + # Second image without cache_control (cache_control will be None if not set) + assert result["content"][3]["type"] == "image" + assert result["content"][3].get("cache_control") is None + + +def test_convert_to_anthropic_tool_result_image_url_as_http(): + """ + Test that HTTP/HTTPS URLs with cache_control are handled correctly. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "role": "tool", + "tool_call_id": "call_http_001", + "content": [ + { + "type": "image_url", + "image_url": "https://example.com/image.jpg", + "cache_control": {"type": "ephemeral"}, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + # Verify image is passed as URL reference with cache_control + assert result["content"][0]["type"] == "image" + assert result["content"][0]["source"]["type"] == "url" + assert result["content"][0]["source"]["url"] == "https://example.com/image.jpg" + assert result["content"][0]["cache_control"]["type"] == "ephemeral" def test_anthropic_messages_pt_server_tool_use_passthrough(): """ Test that anthropic_messages_pt passes through server_tool_use and diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 9b150fd89f4..bae0e5bbb4f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -393,6 +393,63 @@ def test_get_request_tags_from_metadata_and_litellm_metadata(): assert "User-Agent: litellm/1.0.0" in tags +def test_get_request_tags_does_not_mutate_original_tags(): + """ + Test that _get_request_tags does not mutate the original tags list in metadata. + + This is a regression test for a bug where calling _get_request_tags multiple times + would cause User-Agent tags to be duplicated because the function was mutating + the original tags list instead of creating a copy. + """ + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + # Create metadata with original tags + original_tags = ["custom-tag-1", "custom-tag-2"] + metadata = {"tags": original_tags} + litellm_params = {"metadata": metadata} + proxy_server_request = { + "headers": { + "user-agent": "AsyncOpenAI/Python 1.99.9", + } + } + + # Call _get_request_tags multiple times (simulating multiple callbacks) + tags1 = StandardLoggingPayloadSetup._get_request_tags( + litellm_params=litellm_params, + proxy_server_request=proxy_server_request, + ) + tags2 = StandardLoggingPayloadSetup._get_request_tags( + litellm_params=litellm_params, + proxy_server_request=proxy_server_request, + ) + tags3 = StandardLoggingPayloadSetup._get_request_tags( + litellm_params=litellm_params, + proxy_server_request=proxy_server_request, + ) + + # Verify the original tags list was NOT mutated + assert original_tags == ["custom-tag-1", "custom-tag-2"], ( + f"Original tags list was mutated: {original_tags}" + ) + assert metadata["tags"] == ["custom-tag-1", "custom-tag-2"], ( + f"metadata['tags'] was mutated: {metadata['tags']}" + ) + + # Verify each returned list has exactly 2 User-Agent tags (not duplicated) + user_agent_count_1 = len([t for t in tags1 if t.startswith("User-Agent:")]) + user_agent_count_2 = len([t for t in tags2 if t.startswith("User-Agent:")]) + user_agent_count_3 = len([t for t in tags3 if t.startswith("User-Agent:")]) + + assert user_agent_count_1 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_1}" + assert user_agent_count_2 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_2}" + assert user_agent_count_3 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_3}" + + # Verify all returned lists are independent (different objects) + assert tags1 is not tags2 + assert tags2 is not tags3 + assert tags1 is not original_tags + + def test_get_extra_header_tags(): """Test the _get_extra_header_tags method with various scenarios.""" import litellm diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index 91d664c3216..199a16d8590 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -204,3 +204,62 @@ def test_azure_gpt5_reasoning_effort_none_dropped(config: AzureOpenAIGPT5Config) ) assert "reasoning_effort" not in params or params.get("reasoning_effort") != "none" + +# Logprobs support tests for Azure GPT-5.2 +def test_azure_gpt5_2_supports_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.2 models support logprobs parameters. + + Only Azure OpenAI GPT-5.2 supports logprobs, unlike OpenAI's GPT-5 or Azure's gpt-5/gpt-5.1. + Tested with gpt-5.2 on api-version 2025-01-01-preview. + """ + supported_params = config.get_supported_openai_params(model="gpt-5.2") + assert "logprobs" in supported_params + assert "top_logprobs" in supported_params + + +def test_azure_gpt5_2_with_prefix_supports_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.2 with azure/ prefix supports logprobs parameters.""" + supported_params = config.get_supported_openai_params(model="azure/gpt-5.2") + assert "logprobs" in supported_params + assert "top_logprobs" in supported_params + + +def test_azure_gpt5_2_series_supports_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.2 with gpt5_series prefix supports logprobs.""" + supported_params = config.get_supported_openai_params(model="gpt5_series/gpt-5.2") + assert "logprobs" in supported_params + assert "top_logprobs" in supported_params + + +def test_azure_gpt5_2_logprobs_params_passed_through(config: AzureOpenAIGPT5Config): + """Test that logprobs parameters are correctly passed through to the API for gpt-5.2.""" + params = config.map_openai_params( + non_default_params={"logprobs": True, "top_logprobs": 5}, + optional_params={}, + model="azure/gpt-5.2", + drop_params=False, + api_version="2025-01-01-preview", + ) + assert params["logprobs"] is True + assert params["top_logprobs"] == 5 + + +def test_azure_gpt5_base_does_not_support_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5 (non-5.2) does not support logprobs parameters. + + Only gpt-5.2 has been verified to support logprobs on Azure. + """ + supported_params = config.get_supported_openai_params(model="gpt-5") + assert "logprobs" not in supported_params + assert "top_logprobs" not in supported_params + + +def test_azure_gpt5_1_does_not_support_logprobs(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.1 does not support logprobs parameters. + + Only gpt-5.2 has been verified to support logprobs on Azure. + """ + supported_params = config.get_supported_openai_params(model="gpt-5.1") + assert "logprobs" not in supported_params + assert "top_logprobs" not in supported_params + diff --git a/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py b/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py new file mode 100644 index 00000000000..714adc346db --- /dev/null +++ b/tests/test_litellm/llms/openrouter/test_openrouter_embedding_transformation.py @@ -0,0 +1,132 @@ +""" +Unit tests for OpenRouter embedding transformation logic. +""" +from litellm.llms.openrouter.embedding.transformation import ( + OpenrouterEmbeddingConfig, +) + + +def test_openrouter_embedding_supported_params(): + """Test that supported OpenAI params are correctly defined.""" + config = OpenrouterEmbeddingConfig() + supported = config.get_supported_openai_params("test-model") + + assert "timeout" in supported + assert "dimensions" in supported + assert "encoding_format" in supported + assert "user" in supported + + +def test_openrouter_embedding_transform_request(): + """Test request transformation logic.""" + config = OpenrouterEmbeddingConfig() + + # Test with string input + result = config.transform_embedding_request( + model="openrouter/google/text-embedding-004", + input="Hello world", + optional_params={}, + headers={}, + ) + + assert result["model"] == "google/text-embedding-004" + assert result["input"] == ["Hello world"] + + # Test with list input + result = config.transform_embedding_request( + model="google/text-embedding-004", + input=["Hello", "World"], + optional_params={"dimensions": 512}, + headers={}, + ) + + assert result["model"] == "google/text-embedding-004" + assert result["input"] == ["Hello", "World"] + assert result["dimensions"] == 512 + + +def test_openrouter_embedding_validate_environment(): + """Test environment validation and header setup.""" + config = OpenrouterEmbeddingConfig() + + # Test with API key + headers = config.validate_environment( + headers={"Custom-Header": "value"}, + model="test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-api-key", + ) + + # Should include OpenRouter-specific headers + assert "HTTP-Referer" in headers + assert "X-Title" in headers + # Should include Content-Type header + assert "Content-Type" in headers + assert headers["Content-Type"] == "application/json" + # Should include Authorization header + assert "Authorization" in headers + assert headers["Authorization"] == "Bearer test-api-key" + # Should preserve custom headers + assert headers["Custom-Header"] == "value" + + # Test without API key + headers_no_key = config.validate_environment( + headers={}, + model="test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + # Should still include OpenRouter headers but not Authorization + assert "HTTP-Referer" in headers_no_key + assert "X-Title" in headers_no_key + assert "Content-Type" in headers_no_key + assert "Authorization" not in headers_no_key + + +def test_openrouter_embedding_get_complete_url(): + """Test URL construction.""" + config = OpenrouterEmbeddingConfig() + + url = config.get_complete_url( + api_base="https://openrouter.ai/api/v1", + api_key="test-key", + model="test-model", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://openrouter.ai/api/v1/embeddings" + + # Test with trailing slash + url = config.get_complete_url( + api_base="https://openrouter.ai/api/v1/", + api_key="test-key", + model="test-model", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://openrouter.ai/api/v1/embeddings" + + +def test_openrouter_embedding_map_params(): + """Test parameter mapping.""" + config = OpenrouterEmbeddingConfig() + + result = config.map_openai_params( + non_default_params={"dimensions": 512, "timeout": 30, "unsupported": "value"}, + optional_params={}, + model="test-model", + drop_params=False, + ) + + # Supported params should be included + assert result["dimensions"] == 512 + assert result["timeout"] == 30 + # Unsupported params should not be included + assert "unsupported" not in result diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py new file mode 100644 index 00000000000..0a1ac7e2a54 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_gemini_unbound_local_error.py @@ -0,0 +1,38 @@ +import pytest +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm import ModelResponse + +def test_process_candidates_unbound_local_error_fix(): + # Setup + candidates = [ + { + "content": { + "role": "model" + # "parts" is missing intentionally to trigger the issue + }, + "finishReason": "STOP" + } + ] + model_response = ModelResponse() + + # Execution + try: + VertexGeminiConfig._process_candidates( + _candidates=candidates, + model_response=model_response, + standard_optional_params={}, + cumulative_tool_call_index=0 + ) + except UnboundLocalError as e: + pytest.fail(f"UnboundLocalError raised: {e}") + except Exception as e: + # Other exceptions might be okay if they are not UnboundLocalError, + # but ideally it should pass without error or raise a specific error if parts are required. + # However, the goal is to verify thought_signatures doesn't crash. + pass + + # Verify that we didn't crash with UnboundLocalError + +if __name__ == "__main__": + test_process_candidates_unbound_local_error_fix() + print("Test passed!") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 389c8446135..80d65991acb 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -13,6 +13,7 @@ sys.path.insert( import litellm from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.llms.vertex_ai.common_utils import _get_gemini_url def run_sync(coro): @@ -1048,3 +1049,139 @@ class TestVertexBase: MockCredentials.from_info.assert_called_once_with(json_obj) mock_creds.with_scopes.assert_called_once_with(scopes) assert result == "scoped_creds" + + def test_get_token_and_url_with_api_key(self): + """Test that API key authentication routes to Google AI Studio endpoint""" + vertex_base = VertexBase() + + # Test with API key and no credentials - should use Google AI Studio endpoint + auth_header, url = vertex_base._get_token_and_url( + model="gemini-2.0-flash-exp", + auth_header=None, + gemini_api_key="test-api-key-123", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, # No service account credentials + stream=False, + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="chat", + ) + + # Should route to Google AI Studio endpoint + assert "generativelanguage.googleapis.com" in url + assert "gemini-2.0-flash-exp" in url + assert "key=test-api-key-123" in url + assert auth_header is None # API key is in URL, not header + + def test_get_token_and_url_with_credentials(self): + """Test that service account credentials route to Vertex AI endpoint""" + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "mock-bearer-token" + mock_creds.expired = False + + with patch.object( + vertex_base, "_ensure_access_token", return_value=("mock-bearer-token", "test-project") + ): + # Test with credentials - should use Vertex AI endpoint + auth_header, url = vertex_base._get_token_and_url( + model="gemini-2.0-flash-exp", + auth_header="mock-bearer-token", + gemini_api_key=None, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials={"type": "service_account"}, + stream=False, + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="chat", + ) + + # Should route to Vertex AI endpoint + assert "aiplatform.googleapis.com" in url + assert "projects/test-project" in url + assert "locations/us-central1" in url + assert auth_header == "mock-bearer-token" + + def test_get_token_and_url_api_key_with_streaming(self): + """Test API key authentication with streaming enabled""" + vertex_base = VertexBase() + + auth_header, url = vertex_base._get_token_and_url( + model="gemini-2.0-flash-exp", + auth_header=None, + gemini_api_key="test-api-key-456", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + stream=True, # Streaming enabled + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="chat", + ) + + # Should route to Google AI Studio endpoint with streaming + assert "generativelanguage.googleapis.com" in url + assert "streamGenerateContent" in url + assert "key=test-api-key-456" in url + assert "alt=sse" in url + assert auth_header is None + + def test_get_token_and_url_api_key_priority(self): + """Test that credentials take priority over API key when both are provided""" + vertex_base = VertexBase() + + # When both API key and credentials are provided, credentials take priority + mock_creds = MagicMock() + mock_creds.token = "mock-bearer-token" + mock_creds.expired = False + + with patch.object( + vertex_base, "_ensure_access_token", return_value=("mock-bearer-token", "test-project") + ): + auth_header, url = vertex_base._get_token_and_url( + model="gemini-2.0-flash-exp", + auth_header="mock-bearer-token", + gemini_api_key="test-api-key-789", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials={"type": "service_account"}, # Credentials provided + stream=False, + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="chat", + ) + + # Should use Vertex AI endpoint with Bearer token (credentials take priority) + assert "aiplatform.googleapis.com" in url + assert auth_header == "mock-bearer-token" + + def test_get_token_and_url_with_embedding_mode(self): + """Test API key authentication with embedding mode""" + vertex_base = VertexBase() + + auth_header, url = vertex_base._get_token_and_url( + model="text-embedding-004", + auth_header=None, + gemini_api_key="test-embedding-key", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + stream=False, + custom_llm_provider="vertex_ai", + api_base=None, + should_use_v1beta1_features=False, + mode="embedding", + ) + + # Should route to Google AI Studio endpoint for embeddings + assert "generativelanguage.googleapis.com" in url + assert "embedContent" in url + assert "key=test-embedding-key" in url + assert auth_header is None \ No newline at end of file diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 46a7b0ff435..1f379f4371e 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -338,6 +338,17 @@ async def test_proxy_admin_expired_key_from_cache(): f"Exception message should mention 'Expired Key', got: {exc_info.value.message}" ) + # Verify that the param field does NOT leak the full API key (Issue #18731) + # The param should be abbreviated like "sk-...XXXX" not the full plaintext key + assert exc_info.value.param is not None, "Exception should have 'param' attribute" + assert exc_info.value.param != api_key, ( + f"SECURITY: Full API key should NOT be in param field! " + f"Got: {exc_info.value.param}, Expected abbreviated format like 'sk-...XXXX'" + ) + assert exc_info.value.param.startswith("sk-..."), ( + f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}" + ) + # Verify that cache deletion was called mock_delete_cache.assert_called_once() call_args = mock_delete_cache.call_args @@ -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(): """ diff --git a/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py b/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py new file mode 100644 index 00000000000..1492acb0794 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py @@ -0,0 +1,275 @@ +""" +Tests for the RDS IAM token proactive refresh implementation. + +Tests for GitHub Issue #16220: RDS IAM authentication connection failures after 15 minutes. + +The fix implements: +1. Proactive background token refresh (refreshes 3 min before expiration) +2. Precise sleep timing (1 wake-up per token cycle instead of polling) +3. Proper locking during reconnection +4. Fixed __getattr__ fallback that now waits for reconnection + +Run these tests: + poetry run pytest tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py -v -s +""" + +import asyncio +import os +import urllib.parse +from datetime import datetime, timedelta +from unittest.mock import MagicMock, patch + +import pytest + + +class TestPrismaWrapperTokenRefresh: + """Tests for the PrismaWrapper RDS IAM token refresh implementation.""" + + @pytest.fixture + def setup_env(self): + """Setup environment variables for testing.""" + os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "test_user" + os.environ["DATABASE_NAME"] = "test_db" + os.environ["IAM_TOKEN_DB_AUTH"] = "True" + yield + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + "IAM_TOKEN_DB_AUTH", + "DATABASE_SCHEMA", + ]: + os.environ.pop(key, None) + + def _generate_mock_token(self, expires_in_seconds: int = 900) -> str: + """Generate a mock IAM token with expiration info.""" + now = datetime.utcnow() + date_str = now.strftime("%Y%m%dT%H%M%SZ") + # Build the token like AWS does + token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires={expires_in_seconds}&X-Amz-Signature=abc123" + return urllib.parse.quote(token, safe="") + + def _set_database_url_with_token(self, expires_in_seconds: int = 900): + """Set DATABASE_URL with a mock token.""" + token = self._generate_mock_token(expires_in_seconds) + os.environ[ + "DATABASE_URL" + ] = f"postgresql://test_user:{token}@test-host:5432/test_db" + + @pytest.mark.asyncio + async def test_is_token_expired_fresh(self, setup_env): + """Test that fresh token is not detected as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + self._set_database_url_with_token(expires_in_seconds=900) + db_url = os.getenv("DATABASE_URL") + + assert wrapper.is_token_expired(db_url) is False + + @pytest.mark.asyncio + async def test_is_token_expired_old(self, setup_env): + """Test that old token is detected as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Create an expired token + old_date = datetime.utcnow() - timedelta(seconds=901) + date_str = old_date.strftime("%Y%m%dT%H%M%SZ") + token = ( + f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=900&X-Amz-Signature=abc" + ) + encoded_token = urllib.parse.quote(token, safe="") + db_url = f"postgresql://test_user:{encoded_token}@test-host:5432/test_db" + + assert wrapper.is_token_expired(db_url) is True + + @pytest.mark.asyncio + async def test_start_stop_token_refresh_task(self, setup_env): + """Test that token refresh task starts and stops correctly.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Set a valid token + self._set_database_url_with_token(expires_in_seconds=900) + + # Start the task + await wrapper.start_token_refresh_task() + assert wrapper._token_refresh_task is not None + assert not wrapper._token_refresh_task.done() + + # Stop the task + await wrapper.stop_token_refresh_task() + assert wrapper._token_refresh_task is None + + @pytest.mark.asyncio + async def test_start_task_not_enabled(self, setup_env): + """Test that task doesn't start when IAM auth is not enabled.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + # IAM auth disabled + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False) + + await wrapper.start_token_refresh_task() + assert wrapper._token_refresh_task is None + + @pytest.mark.asyncio + async def test_is_token_expired_null(self, setup_env): + """Test that None token is treated as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + assert wrapper.is_token_expired(None) is True + + +class TestTokenExpirationParsing: + """Tests for token expiration parsing utilities.""" + + def test_parse_token_expiration_valid(self): + """Test parsing expiration from a valid token.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Create a token with known expiration + token = "mock-token?X-Amz-Date=20240101T120000Z&X-Amz-Expires=900&X-Amz-Signature=abc" + + expiration = wrapper._parse_token_expiration(token) + + assert expiration is not None + expected = datetime(2024, 1, 1, 12, 0, 0) + timedelta(seconds=900) + assert expiration == expected + + def test_parse_token_expiration_invalid(self): + """Test that invalid token returns None.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Invalid tokens + assert wrapper._parse_token_expiration(None) is None + assert wrapper._parse_token_expiration("no-query-params") is None + assert wrapper._parse_token_expiration("?missing=params") is None + + +class TestBackgroundRefreshLoop: + """Tests for the background refresh loop timing.""" + + @pytest.fixture + def setup_env(self): + """Setup environment variables for testing.""" + os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "test_user" + os.environ["DATABASE_NAME"] = "test_db" + yield + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + ]: + os.environ.pop(key, None) + + @pytest.mark.asyncio + async def test_calculate_seconds_fallback_when_no_url(self, setup_env): + """Test that fallback is used when DATABASE_URL is not set.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Don't set DATABASE_URL + seconds = wrapper._calculate_seconds_until_refresh() + + # Should return fallback interval + assert seconds == wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS + + +# ============================================================================ +# DEMONSTRATION SCRIPT +# ============================================================================ + + +async def demonstrate_fix(): + """ + Demonstrates the fix for the RDS IAM token expiration bug. + + Shows how the proactive refresh prevents the 15-minute connection failure. + """ + # Import the actual implementation + try: + from litellm.proxy.db.prisma_client import PrismaWrapper + except ImportError: + return + + # Setup mock environment + os.environ["DATABASE_HOST"] = "mock-rds.region.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "iam_user" + os.environ["DATABASE_NAME"] = "litellm" + + # Create initial token (expires in 10 seconds for demo) + now = datetime.utcnow() + date_str = now.strftime("%Y%m%dT%H%M%SZ") + token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=10&X-Amz-Signature=abc123" + encoded_token = urllib.parse.quote(token, safe="") + os.environ[ + "DATABASE_URL" + ] = f"postgresql://iam_user:{encoded_token}@mock-rds:5432/litellm" + + # Create mock prisma client + mock_prisma = MagicMock() + + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Override buffer for faster demo + wrapper.TOKEN_REFRESH_BUFFER_SECONDS = 3 + wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS = 5 + _ = wrapper._calculate_seconds_until_refresh() # Verify calculation works + db_url = os.getenv("DATABASE_URL") + is_expired = wrapper.is_token_expired(db_url) + assert is_expired is False, "Fresh token should not be expired!" + + # Mock the _token_refresh_loop to prevent it from actually running + async def mock_loop(): + try: + await asyncio.sleep(1000) + except asyncio.CancelledError: + pass + + with patch.object(wrapper, "_token_refresh_loop", side_effect=mock_loop): + await wrapper.start_token_refresh_task() + await wrapper.stop_token_refresh_task() + + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + ]: + os.environ.pop(key, None) + + +if __name__ == "__main__": + asyncio.run(demonstrate_fix()) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py index 35ed49a84ed..fd72185d1e7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py @@ -2,7 +2,6 @@ Unit tests for Qualifire guardrail integration. """ -import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -75,9 +74,37 @@ class TestQualifireGuardrailInit: assert guardrail.on_flagged == "monitor" + def test_init_with_default_api_base(self): + """Test that default API base is set when not provided.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + DEFAULT_QUALIFIRE_API_BASE, + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + assert guardrail.qualifire_api_base == DEFAULT_QUALIFIRE_API_BASE + + def test_init_with_custom_api_base(self): + """Test initialization with custom API base URL.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + api_base="https://custom.qualifire.ai", + guardrail_name="test_guardrail", + ) + + assert guardrail.qualifire_api_base == "https://custom.qualifire.ai" + class TestQualifireGuardrailMessageConversion: - """Tests for message conversion to Qualifire format.""" + """Tests for message conversion to API format.""" def test_convert_simple_messages(self): """Test conversion of simple text messages.""" @@ -95,15 +122,13 @@ class TestQualifireGuardrailMessageConversion: {"role": "assistant", "content": "Hi there!"}, ] - # Create mock LLMMessage class - mock_llm_message = MagicMock() + result = guardrail._convert_messages_to_api_format(messages) - with patch( - "litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire.QualifireGuardrail._convert_messages_to_qualifire_format" - ) as mock_convert: - mock_convert.return_value = [mock_llm_message, mock_llm_message] - result = guardrail._convert_messages_to_qualifire_format(messages) - assert len(result) == 2 + assert len(result) == 2 + assert result[0]["role"] == "user" + assert result[0]["content"] == "Hello, world!" + assert result[1]["role"] == "assistant" + assert result[1]["content"] == "Hi there!" def test_convert_multimodal_messages(self): """Test conversion of multimodal messages with text parts.""" @@ -126,112 +151,258 @@ class TestQualifireGuardrailMessageConversion: }, ] - with patch( - "litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire.QualifireGuardrail._convert_messages_to_qualifire_format" - ) as mock_convert: - mock_convert.return_value = [MagicMock()] - result = guardrail._convert_messages_to_qualifire_format(messages) - assert len(result) == 1 + result = guardrail._convert_messages_to_api_format(messages) + + assert len(result) == 1 + assert result[0]["role"] == "user" + assert result[0]["content"] == "First part\nSecond part" + + def test_convert_messages_with_tool_calls(self): + """Test conversion of messages with tool calls.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + messages = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}', + }, + } + ], + }, + ] + + result = guardrail._convert_messages_to_api_format(messages) + + assert len(result) == 1 + assert result[0]["role"] == "assistant" + assert "tool_calls" in result[0] + assert len(result[0]["tool_calls"]) == 1 + assert result[0]["tool_calls"][0]["id"] == "call_123" + assert result[0]["tool_calls"][0]["name"] == "get_weather" + assert result[0]["tool_calls"][0]["arguments"] == {"location": "NYC"} -class TestQualifireGuardrailEvaluateKwargs: - """Tests for evaluate kwargs passed to Qualifire client.""" +class TestQualifireGuardrailToolConversion: + """Tests for tool definition conversion.""" + + def test_convert_openai_function_tools(self): + """Test conversion of OpenAI function tool format.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + result = guardrail._convert_tools_to_api_format(tools) + + assert result is not None + assert len(result) == 1 + assert result[0]["name"] == "get_weather" + assert result[0]["description"] == "Get weather for a location" + + def test_convert_empty_tools(self): + """Test that empty tools returns None.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + result = guardrail._convert_tools_to_api_format(None) + assert result is None + + result = guardrail._convert_tools_to_api_format([]) + assert result is None + + +class TestQualifireGuardrailAPICall: + """Tests for API call with httpx client.""" @pytest.mark.asyncio async def test_evaluate_called_with_prompt_injections(self): - """Test that evaluate is called with prompt_injections enabled.""" - # Mock the qualifire module and its types - mock_qualifire_types = MagicMock() - mock_llm_message = MagicMock() - mock_llm_tool_call = MagicMock() - mock_message_instance = MagicMock() - mock_llm_message.return_value = mock_message_instance - - mock_qualifire_types.LLMMessage = mock_llm_message - mock_qualifire_types.LLMToolCall = mock_llm_tool_call - - with patch.dict('sys.modules', {'qualifire': MagicMock(), 'qualifire.types': mock_qualifire_types}): - from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( - QualifireGuardrail, - ) + """Test that evaluate endpoint is called with prompt_injections enabled.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) - guardrail = QualifireGuardrail( - api_key="test_key", - prompt_injections=True, - guardrail_name="test_guardrail", - ) + guardrail = QualifireGuardrail( + api_key="test_key", + prompt_injections=True, + guardrail_name="test_guardrail", + ) - # Mock the client - mock_client = MagicMock() - mock_result = MagicMock() - mock_result.score = 100 - mock_result.status = "completed" - mock_result.evaluationResults = [] - mock_client.evaluate.return_value = mock_result - guardrail._client = mock_client + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) - messages = [{"role": "user", "content": "Hello, world!"}] + messages = [{"role": "user", "content": "Hello, world!"}] - await guardrail._run_qualifire_check( - messages=messages, output=None, dynamic_params={} - ) + await guardrail._run_qualifire_check( + messages=messages, output=None, dynamic_params={} + ) - # Verify evaluate was called with correct kwargs - mock_client.evaluate.assert_called_once() - call_kwargs = mock_client.evaluate.call_args[1] - assert call_kwargs["prompt_injections"] is True - assert "messages" in call_kwargs + # Verify the API was called + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + assert "json" in call_kwargs + payload = call_kwargs["json"] + assert payload["prompt_injections"] is True + assert "messages" in payload + assert call_kwargs["url"].endswith("/api/evaluation/evaluate") @pytest.mark.asyncio async def test_evaluate_called_with_multiple_checks(self): """Test that evaluate is called with multiple checks enabled.""" - # Mock the qualifire module and its types - mock_qualifire_types = MagicMock() - mock_llm_message = MagicMock() - mock_llm_tool_call = MagicMock() - mock_message_instance = MagicMock() - mock_llm_message.return_value = mock_message_instance - - mock_qualifire_types.LLMMessage = mock_llm_message - mock_qualifire_types.LLMToolCall = mock_llm_tool_call - - with patch.dict('sys.modules', {'qualifire': MagicMock(), 'qualifire.types': mock_qualifire_types}): - from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( - QualifireGuardrail, - ) + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) - guardrail = QualifireGuardrail( - api_key="test_key", - prompt_injections=True, - pii_check=True, - hallucinations_check=True, - assertions=["Output must be valid JSON"], - guardrail_name="test_guardrail", - ) + guardrail = QualifireGuardrail( + api_key="test_key", + prompt_injections=True, + pii_check=True, + hallucinations_check=True, + assertions=["Output must be valid JSON"], + guardrail_name="test_guardrail", + ) - # Mock the client - mock_client = MagicMock() - mock_result = MagicMock() - mock_result.score = 100 - mock_result.status = "completed" - mock_result.evaluationResults = [] - mock_client.evaluate.return_value = mock_result - guardrail._client = mock_client + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) - messages = [{"role": "user", "content": "Hello, world!"}] + messages = [{"role": "user", "content": "Hello, world!"}] - await guardrail._run_qualifire_check( - messages=messages, output="Test output", dynamic_params={} - ) + await guardrail._run_qualifire_check( + messages=messages, output="Test output", dynamic_params={} + ) - # Verify evaluate was called with correct kwargs - mock_client.evaluate.assert_called_once() - call_kwargs = mock_client.evaluate.call_args[1] - assert call_kwargs["prompt_injections"] is True - assert call_kwargs["pii_check"] is True - assert call_kwargs["hallucinations_check"] is True - assert call_kwargs["assertions"] == ["Output must be valid JSON"] - assert call_kwargs["output"] == "Test output" + # Verify the API was called with correct payload + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + payload = call_kwargs["json"] + assert payload["prompt_injections"] is True + assert payload["pii_check"] is True + assert payload["hallucinations_check"] is True + assert payload["assertions"] == ["Output must be valid JSON"] + assert payload["output"] == "Test output" + + @pytest.mark.asyncio + async def test_invoke_endpoint_used_with_evaluation_id(self): + """Test that invoke endpoint is used when evaluation_id is provided.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + evaluation_id="eval_123", + guardrail_name="test_guardrail", + ) + + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + messages = [{"role": "user", "content": "Hello, world!"}] + + await guardrail._run_qualifire_check( + messages=messages, output="Test output", dynamic_params={} + ) + + # Verify the invoke endpoint was called + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + assert call_kwargs["url"].endswith("/api/evaluation/invoke") + payload = call_kwargs["json"] + assert payload["evaluation_id"] == "eval_123" + assert payload["input"] == "Hello, world!" + assert payload["output"] == "Test output" + + @pytest.mark.asyncio + async def test_correct_headers_sent(self): + """Test that correct headers are sent with the API request.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="my_api_key", + guardrail_name="test_guardrail", + ) + + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + messages = [{"role": "user", "content": "Hello!"}] + + await guardrail._run_qualifire_check( + messages=messages, output=None, dynamic_params={} + ) + + call_kwargs = guardrail.async_handler.post.call_args[1] + headers = call_kwargs["headers"] + + assert headers["X-Qualifire-API-Key"] == "my_api_key" + assert headers["Content-Type"] == "application/json" class TestQualifireGuardrailCheckIfFlagged: @@ -248,12 +419,14 @@ class TestQualifireGuardrailCheckIfFlagged: guardrail_name="test_guardrail", ) - # Mock result with completed status and no flagged items - mock_result = MagicMock() - mock_result.status = "completed" - mock_result.evaluationResults = [] + # Result with completed status and no flagged items (dict format) + result = { + "status": "completed", + "score": 100, + "evaluationResults": [], + } - assert guardrail._check_if_flagged(mock_result) is False + assert guardrail._check_if_flagged(result) is False def test_check_if_flagged_returns_true_for_flagged_content(self): """Test that _check_if_flagged returns True when content is flagged.""" @@ -266,18 +439,25 @@ class TestQualifireGuardrailCheckIfFlagged: guardrail_name="test_guardrail", ) - # Mock result with flagged item - mock_inner_result = MagicMock() - mock_inner_result.flagged = True + # Result with flagged item (dict format matching API response) + result = { + "status": "completed", + "score": 15, + "evaluationResults": [ + { + "type": "prompt_injection", + "results": [ + { + "flagged": True, + "score": 0.15, + "reason": "Prompt injection detected", + } + ], + } + ], + } - mock_eval_result = MagicMock() - mock_eval_result.results = [mock_inner_result] - - mock_result = MagicMock() - mock_result.status = "completed" - mock_result.evaluationResults = [mock_eval_result] - - assert guardrail._check_if_flagged(mock_result) is True + assert guardrail._check_if_flagged(result) is True def test_check_if_flagged_returns_false_when_no_flagged_items(self): """Test that _check_if_flagged returns False when no items are flagged.""" @@ -291,17 +471,24 @@ class TestQualifireGuardrailCheckIfFlagged: ) # Result with evaluation results but nothing flagged - mock_inner_result = MagicMock() - mock_inner_result.flagged = False + result = { + "status": "completed", + "score": 95, + "evaluationResults": [ + { + "type": "prompt_injection", + "results": [ + { + "flagged": False, + "score": 0.95, + "reason": "No issues detected", + } + ], + } + ], + } - mock_eval_result = MagicMock() - mock_eval_result.results = [mock_inner_result] - - mock_result = MagicMock() - mock_result.status = "success" - mock_result.evaluationResults = [mock_eval_result] - - assert guardrail._check_if_flagged(mock_result) is False + assert guardrail._check_if_flagged(result) is False class TestQualifireGuardrailShouldRun: diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index f2bae2cb14a..bc223d15d5f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1,20 +1,20 @@ import json import os import sys -from litellm._uuid import uuid +import types from datetime import datetime, timedelta -from typing import List +from typing import List, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import FastAPI from fastapi.testclient import TestClient +from litellm._uuid import uuid sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from typing import Optional - from litellm.proxy._types import ( LiteLLM_MCPServerTable, LitellmUserRoles, @@ -118,6 +118,22 @@ def setup_mock_prisma_client( return mock_prisma_client +def create_mcp_router_test_client() -> TestClient: + from litellm.proxy.management_endpoints.mcp_management_endpoints import router + + app = FastAPI() + app.include_router(router) + return TestClient(app) + + +def patch_proxy_general_settings(settings: dict): + fake_proxy_server_module = types.SimpleNamespace(general_settings=settings) + return patch.dict( + sys.modules, + {"litellm.proxy.proxy_server": fake_proxy_server_module}, + ) + + class TestListMCPServers: """Test suite for list MCP servers functionality""" @@ -1082,6 +1098,55 @@ class TestHealthCheckServers: assert result[1]["server_id"] == "server-2" assert result[1]["status"] == "unhealthy" + +class TestMCPRegistryEndpoint: + def test_registry_returns_404_when_flag_missing(self): + client = create_mcp_router_test_client() + + with patch_proxy_general_settings({}): + response = client.get("/v1/mcp/registry.json") + + assert response.status_code == 404 + + def test_registry_returns_404_when_flag_false(self): + client = create_mcp_router_test_client() + + with patch_proxy_general_settings({"enable_mcp_registry": False}): + response = client.get("/v1/mcp/registry.json") + + assert response.status_code == 404 + + def test_registry_returns_entries_when_enabled(self): + client = create_mcp_router_test_client() + + mock_server = generate_mock_mcp_server_config_record( + server_id="server-123", + name="zapier", + url="https://zapier.example.com/mcp", + transport="http", + ) + + mock_manager = MagicMock() + mock_manager.get_registry.return_value = {mock_server.server_id: mock_server} + + with patch_proxy_general_settings({"enable_mcp_registry": True}), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ): + response = client.get("/v1/mcp/registry.json") + + assert response.status_code == 200 + data = response.json() + assert len(data["servers"]) == 2 # built-in + custom server + + builtin_entry = data["servers"][0]["server"] + assert builtin_entry["name"] == "litellm-mcp-server" + assert builtin_entry["remotes"][0]["url"].endswith("/mcp") + + custom_entry = data["servers"][1]["server"] + assert custom_entry["name"] == "zapier" + assert custom_entry["remotes"][0]["url"].endswith("/zapier/mcp") + @pytest.mark.asyncio async def test_health_check_specific_servers(self): """ diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index b5d44385698..3d1e9aece41 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import Request, status -from fastapi.responses import StreamingResponse +from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid @@ -11,9 +11,10 @@ from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ProxyConfig, + _extract_error_from_sse_chunk, _get_cost_breakdown_from_logging_obj, _parse_event_data_for_error, - create_streaming_response, + create_response, ) from litellm.proxy.utils import ProxyLogging @@ -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]}" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 53d89df5026..751a9033871 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/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 diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 96ac2e2c345..09628cd4a76 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -203,3 +203,22 @@ class TestResponseAPILoggingUtils: assert result.prompt_tokens == 0 assert result.completion_tokens == 20 assert result.total_tokens == 20 + + def test_transform_response_api_usage_calculates_total_from_input_and_output_tokens_if_available(self): + """Test transformation calculates total_tokens when it's None and input / output tokens are present""" + # Setup + usage = { + "input_tokens": 15, + "output_tokens": 25, + "total_tokens": None, + } + + # Execute + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + + # Assert + assert result.prompt_tokens == 15 + assert result.completion_tokens == 25 + assert result.total_tokens == 40 # 15 + 25 diff --git a/tests/test_litellm/test_utils_custom.py b/tests/test_litellm/test_utils_custom.py new file mode 100644 index 00000000000..292da4132b9 --- /dev/null +++ b/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() diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts b/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts new file mode 100644 index 00000000000..4a4bb64c8ed --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/menuMappings.ts @@ -0,0 +1,38 @@ +import { Page } from "./pages"; + +/** + * Maps sidebar menu item labels to their corresponding page enum values. + * This mapping is for the admin role. + */ +export const menuLabelToPage: Record = { + "Virtual Keys": Page.ApiKeys, + Playground: Page.LlmPlayground, + Models: Page.Models, + "Models + Endpoints": Page.Models, + Usage: Page.NewUsage, + Teams: Page.Teams, + "Internal Users": Page.Users, + "Internal User": Page.Users, // Legacy label support + Organizations: Page.Organizations, + "API Reference": Page.ApiRef, + "AI Hub": Page.ModelHubTable, + "Model Hub": Page.ModelHubTable, + Logs: Page.Logs, + Guardrails: Page.Guardrails, + // Settings submenu items + "Router Settings": Page.RouterSettings, + "Logging & Alerts": Page.LoggingAndAlerts, + "Admin Settings": Page.AdminPanel, + "Cost Tracking": Page.CostTracking, + "UI Theme": Page.UiTheme, + // Experimental submenu items + Caching: Page.Caching, + Prompts: Page.Prompts, + Budgets: Page.Budgets, + "API Playground": Page.TransformRequest, + "Tag Management": Page.TagManagement, + "Old Usage": Page.Usage, + // Tools submenu items + "MCP Servers": Page.McpServers, + "Vector Stores": Page.VectorStores, +}; diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts new file mode 100644 index 00000000000..3ea37718ab5 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/pages.ts @@ -0,0 +1,33 @@ +/** + * Enum for all page query parameters supported in the app. + * These values correspond to the `page` query parameter used in the URL. + */ +export enum Page { + ApiKeys = "api-keys", + Models = "models", + LlmPlayground = "llm-playground", + Users = "users", + Teams = "teams", + Organizations = "organizations", + AdminPanel = "admin-panel", + ApiRef = "api_ref", + LoggingAndAlerts = "logging-and-alerts", + Budgets = "budgets", + Guardrails = "guardrails", + Agents = "agents", + Prompts = "prompts", + TransformRequest = "transform-request", + RouterSettings = "router-settings", + UiTheme = "ui-theme", + CostTracking = "cost-tracking", + ModelHubTable = "model-hub-table", + Caching = "caching", + PassThroughSettings = "pass-through-settings", + Logs = "logs", + McpServers = "mcp-servers", + SearchTools = "search-tools", + TagManagement = "tag-management", + VectorStores = "vector-stores", + NewUsage = "new_usage", + Usage = "usage", +} diff --git a/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts new file mode 100644 index 00000000000..919e516b35b --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/helpers/navigation.ts @@ -0,0 +1,12 @@ +import { Page } from "../fixtures/pages"; +import { Page as PlaywrightPage } from "@playwright/test"; + +/** + * Navigates to a specific page using the page query parameter. + * Uses relative path which will be resolved against the baseURL configured in playwright.config.ts + * @param page - The Playwright page object + * @param pageEnum - The page enum value to navigate to + */ +export async function navigateToPage(page: PlaywrightPage, pageEnum: Page): Promise { + await page.goto(`/ui?page=${pageEnum}`); +} diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts index c0619cfa845..2ab782d5678 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts @@ -14,7 +14,7 @@ test.describe("Add Model", () => { await providerInputDropdown.fill("Anthropic"); await page.waitForTimeout(1000); await providerInputDropdown.press("Enter"); - await page.waitForTimeout(1000); + await page.waitForTimeout(2000); const providerModelsDropdown = page.locator(".ant-select-selection-overflow").first(); await providerModelsDropdown.click(); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts index c90be698ae1..ce07cc2b83d 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -1,6 +1,9 @@ import test, { expect } from "@playwright/test"; import { Role } from "../../fixtures/roles"; import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { menuLabelToPage } from "../../fixtures/menuMappings"; +import { navigateToPage } from "../../helpers/navigation"; const sidebarButtons = { [Role.ProxyAdmin]: [ @@ -9,9 +12,7 @@ const sidebarButtons = { "Models", "Usage", "Teams", - "Internal User", - "Settings", - "Experimental", + "Internal Users", "API Reference", "AI Hub", ], @@ -23,13 +24,36 @@ for (const { role, storage } of roles) { test.describe(`${role} sidebar`, () => { test.use({ storageState: storage }); - test("can see and navigate all sidebar buttons", async ({ page }) => { + test("should navigate to correct URL when clicking sidebar menu items from homepage", async ({ page }) => { await page.goto("/ui"); - for (const button of sidebarButtons[role as keyof typeof sidebarButtons]) { - const tab = page.getByRole("menuitem", { name: button }); + + for (const buttonLabel of sidebarButtons[role as keyof typeof sidebarButtons]) { + const expectedPage = menuLabelToPage[buttonLabel]; + + if (!expectedPage) { + throw new Error(`No page mapping found for menu label: ${buttonLabel}`); + } + + const tab = page.getByRole("menuitem", { name: buttonLabel }); await expect(tab).toBeVisible(); + await tab.click(); + + // Verify URL contains the correct page query parameter + await expect(page).toHaveURL(new RegExp(`[?&]page=${expectedPage}(&|$)`)); } }); + + test("should navigate directly to page using navigation helper", async ({ page }) => { + // Test direct navigation to verify the helper function works + await navigateToPage(page, Page.ApiKeys); + await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.ApiKeys}(&|$)`)); + + await navigateToPage(page, Page.Models); + await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.Models}(&|$)`)); + + await navigateToPage(page, Page.LlmPlayground); + await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.LlmPlayground}(&|$)`)); + }); }); } diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx new file mode 100644 index 00000000000..ecc4914b0ab --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx @@ -0,0 +1,187 @@ +import { describe, expect, it, vi, beforeEach } from "vitest"; +import { + convertImageToBase64, + createChatMultimodalMessage, + createChatDisplayMessage, + shouldShowChatAttachedImage, +} from "./ChatImageUtils"; +import { MessageType } from "./types"; + +describe("ChatImageUtils", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe("convertImageToBase64", () => { + it("should convert file to base64 data URI", async () => { + const file = new File(["test content"], "test.png", { type: "image/png" }); + const result = await convertImageToBase64(file); + expect(result).toMatch(/^data:image\/png;base64,/); + }); + + it("should handle different file types", async () => { + const jpegFile = new File(["jpeg content"], "test.jpg", { type: "image/jpeg" }); + const result = await convertImageToBase64(jpegFile); + expect(result).toMatch(/^data:image\/jpeg;base64,/); + }); + + it("should reject on file read error", async () => { + const file = new File(["test"], "test.png", { type: "image/png" }); + const originalReadAsDataURL = FileReader.prototype.readAsDataURL; + + FileReader.prototype.readAsDataURL = vi.fn(function (this: FileReader) { + setTimeout(() => { + if (this.onerror) { + this.onerror(new Error("Read error") as any); + } + }, 0); + }); + + await expect(convertImageToBase64(file)).rejects.toThrow(); + + FileReader.prototype.readAsDataURL = originalReadAsDataURL; + }); + }); + + describe("createChatMultimodalMessage", () => { + it("should create multimodal message with text and image", async () => { + const file = new File(["test content"], "test.png", { type: "image/png" }); + const inputMessage = "What is in this image?"; + + const result = await createChatMultimodalMessage(inputMessage, file); + + expect(result.role).toBe("user"); + expect(result.content).toHaveLength(2); + expect(result.content[0]).toEqual({ type: "text", text: inputMessage }); + expect(result.content[1]).toMatchObject({ + type: "image_url", + image_url: { + url: expect.stringMatching(/^data:image\/png;base64,/), + }, + }); + }); + + it("should include base64 data URI in image_url", async () => { + const file = new File(["test content"], "test.png", { type: "image/png" }); + const result = await createChatMultimodalMessage("test", file); + + const imageContent = result.content[1]; + expect(imageContent.type).toBe("image_url"); + if ("image_url" in imageContent && imageContent.image_url) { + expect(imageContent.image_url.url).toMatch(/^data:/); + } + }); + }); + + describe("createChatDisplayMessage", () => { + it("should create display message without file", () => { + const result = createChatDisplayMessage("Hello world", false); + + expect(result.role).toBe("user"); + expect(result.content).toBe("Hello world"); + expect(result.imagePreviewUrl).toBeUndefined(); + }); + + it("should create display message with PDF file", () => { + const filePreviewUrl = "blob:test-url"; + const result = createChatDisplayMessage("Read this", true, filePreviewUrl, "document.pdf"); + + expect(result.content).toBe("Read this [PDF attached]"); + expect(result.imagePreviewUrl).toBe(filePreviewUrl); + }); + + it("should create display message with image file", () => { + const filePreviewUrl = "blob:test-url"; + const result = createChatDisplayMessage("Look at this", true, filePreviewUrl, "photo.jpg"); + + expect(result.content).toBe("Look at this [Image attached]"); + expect(result.imagePreviewUrl).toBe(filePreviewUrl); + }); + + it("should create display message with file but no fileName", () => { + const filePreviewUrl = "blob:test-url"; + const result = createChatDisplayMessage("Check this", true, filePreviewUrl); + + expect(result.content).toBe("Check this "); + expect(result.imagePreviewUrl).toBe(filePreviewUrl); + }); + + it("should create display message with file but no preview URL", () => { + const result = createChatDisplayMessage("See this", true, undefined, "image.png"); + + expect(result.content).toBe("See this [Image attached]"); + expect(result.imagePreviewUrl).toBeUndefined(); + }); + }); + + describe("shouldShowChatAttachedImage", () => { + it("should return true for user message with image attachment", () => { + const message: MessageType = { + role: "user", + content: "Check this [Image attached]", + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(true); + }); + + it("should return true for user message with PDF attachment", () => { + const message: MessageType = { + role: "user", + content: "Read this [PDF attached]", + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(true); + }); + + it("should return false for assistant message", () => { + const message: MessageType = { + role: "assistant", + content: "Here is the image [Image attached]", + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + + it("should return false when content is not a string", () => { + const message: MessageType = { + role: "user", + content: [{ type: "input_text", text: "test" }], + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + + it("should return false when content does not include attachment marker", () => { + const message: MessageType = { + role: "user", + content: "Just regular text", + imagePreviewUrl: "blob:test-url", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + + it("should return false when imagePreviewUrl is missing", () => { + const message: MessageType = { + role: "user", + content: "Check this [Image attached]", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + + it("should return false when imagePreviewUrl is empty string", () => { + const message: MessageType = { + role: "user", + content: "Check this [Image attached]", + imagePreviewUrl: "", + }; + + expect(shouldShowChatAttachedImage(message)).toBe(false); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index 9918040ef64..f747b5b2599 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -84,10 +84,7 @@ interface ChatUIProps { }; } -const MCP_SUPPORTED_ENDPOINTS = new Set([ - EndpointType.CHAT, - EndpointType.RESPONSES, -]); +const MCP_SUPPORTED_ENDPOINTS = new Set([EndpointType.CHAT, EndpointType.RESPONSES]); const ChatUI: React.FC = ({ accessToken, @@ -131,7 +128,7 @@ const ChatUI: React.FC = ({ }); const [apiKey, setApiKey] = useState(() => sessionStorage.getItem("apiKey") || ""); const [customProxyBaseUrl, setCustomProxyBaseUrl] = useState( - () => sessionStorage.getItem("customProxyBaseUrl") || "" + () => sessionStorage.getItem("customProxyBaseUrl") || "", ); const [inputMessage, setInputMessage] = useState(""); const [chatHistory, setChatHistory] = useState(() => { @@ -215,7 +212,7 @@ const ChatUI: React.FC = ({ const [temperature, setTemperature] = useState(1.0); const [maxTokens, setMaxTokens] = useState(2048); const [useAdvancedParams, setUseAdvancedParams] = useState(false); - + // Code Interpreter state (using custom hook) const codeInterpreter = useCodeInterpreter(); @@ -244,16 +241,15 @@ const ChatUI: React.FC = ({ 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 = ({ 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 = ({ icon={ApiOutlined} /> {customProxyBaseUrl && ( - - API calls will be sent to: {customProxyBaseUrl} - + API calls will be sent to: {customProxyBaseUrl} )} @@ -1382,7 +1376,11 @@ const ChatUI: React.FC = ({ optionLabelProp="label" > {agentInfo.map((agent) => ( - +
{agent.agent_name || agent.agent_id} {agent.agent_card_params?.description && ( @@ -1416,10 +1414,7 @@ const ChatUI: React.FC = ({
MCP Servers - + @@ -1435,15 +1430,15 @@ const ChatUI: React.FC = ({ } 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 = ({ disabled={selectedMCPServers.includes("__all__")} >
- - {server.alias || server.server_name || server.server_id} - - {server.description && ( - {server.description} - )} + {server.alias || server.server_name || server.server_id} + {server.description && {server.description}}
))} @@ -1489,40 +1480,40 @@ const ChatUI: React.FC = ({ {selectedMCPServers.length > 0 && !selectedMCPServers.includes("__all__") && MCP_SUPPORTED_ENDPOINTS.has(endpointType as EndpointType) && ( -
- {selectedMCPServers.map(serverId => { - const server = mcpServers.find(s => s.server_id === serverId); - const tools = serverToolsMap[serverId] || []; - if (tools.length === 0) return null; +
+ {selectedMCPServers.map((serverId) => { + const server = mcpServers.find((s) => s.server_id === serverId); + const tools = serverToolsMap[serverId] || []; + if (tools.length === 0) return null; - return ( -
- - Limit tools for {server?.alias || server?.server_name || serverId}: - - { + setMCPServerToolRestrictions((prev) => ({ + ...prev, + [serverId]: selectedTools, + })); + }} + options={tools.map((tool) => ({ + value: tool.name, + label: tool.name, + }))} + maxTagCount={2} + /> +
+ ); + })} +
+ )}
@@ -2073,7 +2064,13 @@ const ChatUI: React.FC = ({ )} {/* Quick Code Interpreter toggle for Responses */} {endpointType === EndpointType.RESPONSES && ( - +