diff --git a/.circleci/config.yml b/.circleci/config.yml index fbbb6deeba8..8709f730c23 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -4100,6 +4100,63 @@ jobs: path: playwright-report destination: playwright-report + prisma_schema_sync: + machine: + image: ubuntu-2204:2023.10.1 + resource_class: xlarge + working_directory: ~/project + steps: + - checkout + - setup_google_dns + - attach_workspace: + at: ~/project + - run: + name: Load Docker Database Image + command: | + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database + - run: + name: Install Neon CLI + command: | + npm i -g neonctl + - run: + name: Install curl and dockerize + command: | + sudo apt-get update + sudo apt-get install -y curl + sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz + sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz + sudo rm dockerize-linux-amd64-v0.6.1.tar.gz + - run: + name: Sync schema on base e2e database + command: | + BASE_DATABASE_URL=$(neon connection-string \ + --project-id $NEON_PROJECT_ID \ + --api-key $NEON_API_KEY \ + --branch br-fancy-paper-ad1olsb3 \ + --database-name yuneng-trial-db \ + --role neondb_owner) + docker run -d \ + -p 4000:4000 \ + -e DATABASE_URL=$BASE_DATABASE_URL \ + -e LITELLM_MASTER_KEY="sk-1234" \ + --name schema-sync \ + -v $(pwd)/litellm/proxy/example_config_yaml/simple_config.yaml:/app/config.yaml \ + litellm-docker-database:ci \ + --config /app/config.yaml \ + --port 4000 \ + --use_prisma_db_push + - run: + name: Start outputting logs + command: docker logs -f schema-sync + background: true + - run: + name: Wait for proxy to be ready (schema sync complete) + command: dockerize -wait http://localhost:4000 -timeout 5m + - run: + name: Stop schema sync container + command: docker stop schema-sync + test_nonroot_image: machine: image: ubuntu-2204:2023.10.1 @@ -4298,6 +4355,15 @@ workflows: only: - main - /litellm_.*/ + - prisma_schema_sync: + context: e2e_ui_tests + requires: + - build_docker_database_image + filters: + branches: + only: + - main + - /litellm_.*/ - e2e_ui_testing: name: e2e_ui_testing_chromium browser: chromium @@ -4305,6 +4371,7 @@ workflows: requires: - ui_build - build_docker_database_image + - prisma_schema_sync filters: branches: only: @@ -4317,6 +4384,7 @@ workflows: requires: - ui_build - build_docker_database_image + - prisma_schema_sync filters: branches: only: diff --git a/docs/my-website/docs/ocr.md b/docs/my-website/docs/ocr.md index 93cb74ee69f..cea6fce1254 100644 --- a/docs/my-website/docs/ocr.md +++ b/docs/my-website/docs/ocr.md @@ -61,6 +61,52 @@ async def test_async_ocr(): asyncio.run(test_async_ocr()) ``` +### Using Local Files + +LiteLLM can read local files directly — no manual base64 encoding needed: + +```python +from litellm import ocr + +# OCR with a local PDF file path +response = ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "file", + "file": "/path/to/document.pdf" + } +) + +# OCR with a file object +response = ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "file", + "file": open("document.pdf", "rb") + } +) + +# OCR with raw bytes +with open("document.pdf", "rb") as f: + pdf_bytes = f.read() + +response = ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "file", + "file": pdf_bytes, + "mime_type": "application/pdf" # recommended for raw bytes (auto-detected from extension for file paths) + } +) +``` + +The `file` field accepts: +- **File path** (`str` or `pathlib.Path`) — LiteLLM reads the file and detects the MIME type from the extension +- **File object** (binary file-like object) — e.g. `open("doc.pdf", "rb")` +- **Raw bytes** (`bytes`) — use `mime_type` to specify the content type + +LiteLLM automatically converts file inputs to base64 data URIs internally, so all providers work seamlessly. + ### Using Base64 Encoded Documents ```python @@ -121,7 +167,7 @@ litellm --config /path/to/config.yaml # RUNNING on http://0.0.0.0:4000 ``` -Test request +**Test request — JSON body** ```bash curl http://0.0.0.0:4000/v1/ocr \ @@ -136,6 +182,27 @@ curl http://0.0.0.0:4000/v1/ocr \ }' ``` +**Test request — multipart file upload** + +Upload a file directly using multipart form data. No need to base64-encode the file yourself. + +```bash +curl http://0.0.0.0:4000/v1/ocr \ + -H "Authorization: Bearer sk-1234" \ + -F "model=mistral-ocr" \ + -F "file=@/path/to/document.pdf" +``` + +You can also pass optional parameters as additional form fields: + +```bash +curl http://0.0.0.0:4000/v1/ocr \ + -H "Authorization: Bearer sk-1234" \ + -F "model=mistral-ocr" \ + -F "file=@screenshot.png" \ + -F 'pages=[0,1,2]' \ + -F "include_image_base64=true" +``` ## **Request/Response Format** @@ -168,10 +235,12 @@ See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilitie | Parameter | Type | Required | Description | |-----------|------|----------|-------------| | `model` | string | Yes | The OCR model to use (e.g., `"mistral/mistral-ocr-latest"`) | -| `document` | object | Yes | Document to process. Must contain `type` and URL field | -| `document.type` | string | Yes | Either `"document_url"` for PDFs/docs or `"image_url"` for images | -| `document.document_url` | string | Conditional | URL to the document (required if `type` is `"document_url"`) | -| `document.image_url` | string | Conditional | URL to the image (required if `type` is `"image_url"`) | +| `document` | object | Yes | Document to process. Must contain `type` and the corresponding field | +| `document.type` | string | Yes | `"document_url"` for PDFs/docs, `"image_url"` for images, or `"file"` for local files | +| `document.document_url` | string | Conditional | URL or data URI to the document (required if `type` is `"document_url"`) | +| `document.image_url` | string | Conditional | URL or data URI to the image (required if `type` is `"image_url"`) | +| `document.file` | string/bytes/file | Conditional | File path, bytes, or file-like object (required if `type` is `"file"`) | +| `document.mime_type` | string | No | Explicit MIME type for file inputs (auto-detected from extension if not provided) | | `pages` | array | No | List of specific page indices to process (0-indexed) | | `include_image_base64` | boolean | No | Whether to include extracted images as base64 strings | | `image_limit` | integer | No | Maximum number of images to return | @@ -179,7 +248,7 @@ See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilitie #### Document Format Examples -**For PDFs and documents:** +**For PDFs and documents (URL):** ```json { "type": "document_url", @@ -187,7 +256,7 @@ See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilitie } ``` -**For images:** +**For images (URL):** ```json { "type": "image_url", @@ -203,6 +272,21 @@ See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilitie } ``` +**For local files (SDK):** +```python +{"type": "file", "file": "/path/to/document.pdf"} +{"type": "file", "file": open("image.png", "rb")} +{"type": "file", "file": pdf_bytes, "mime_type": "application/pdf"} +``` + +**For file uploads (Proxy — multipart form):** +```bash +curl http://0.0.0.0:4000/v1/ocr \ + -H "Authorization: Bearer sk-1234" \ + -F "model=mistral-ocr" \ + -F "file=@document.pdf" +``` + ### Response Format The response follows Mistral's OCR format with the following structure: diff --git a/docs/my-website/docs/pass_through/assembly_ai.md b/docs/my-website/docs/pass_through/assembly_ai.md index 4606640c5c4..c7c70639e7e 100644 --- a/docs/my-website/docs/pass_through/assembly_ai.md +++ b/docs/my-website/docs/pass_through/assembly_ai.md @@ -1,31 +1,36 @@ -# Assembly AI +# AssemblyAI -Pass-through endpoints for Assembly AI - call Assembly AI endpoints, in native format (no translation). +Pass-through endpoints for AssemblyAI - call AssemblyAI endpoints, in native format (no translation). -| Feature | Supported | Notes | +| Feature | Supported | Notes | |-------|-------|-------| | Cost Tracking | ✅ | works across all integrations | | Logging | ✅ | works across all integrations | -Supports **ALL** Assembly AI Endpoints +Supports **ALL** AssemblyAI Endpoints -[**See All Assembly AI Endpoints**](https://www.assemblyai.com/docs/api-reference) +[**See All AssemblyAI Endpoints**](https://www.assemblyai.com/docs/api-reference) - +## Supported Routes + +| AssemblyAI Service | LiteLLM Route | AssemblyAI Base URL | +|-------------------|---------------|---------------------| +| Speech-to-Text (US) | `/assemblyai/*` | `api.assemblyai.com` | +| Speech-to-Text (EU) | `/eu.assemblyai/*` | `eu.api.assemblyai.com` | ## Quick Start -Let's call the Assembly AI [`/v2/transcripts` endpoint](https://www.assemblyai.com/docs/api-reference/transcripts) +Let's call the AssemblyAI [`/v2/transcripts` endpoint](https://www.assemblyai.com/docs/api-reference/transcripts) -1. Add Assembly AI API Key to your environment +1. Add AssemblyAI API Key to your environment ```bash export ASSEMBLYAI_API_KEY="" ``` -2. Start LiteLLM Proxy +2. Start LiteLLM Proxy ```bash litellm @@ -33,53 +38,157 @@ litellm # RUNNING on http://0.0.0.0:4000 ``` -3. Test it! +3. Test it! -Let's call the Assembly AI `/v2/transcripts` endpoint +Let's call the AssemblyAI [`/v2/transcripts` endpoint](https://www.assemblyai.com/docs/api-reference/transcripts). Includes commented-out [Speech Understanding](https://www.assemblyai.com/docs/speech-understanding) features you can toggle on. ```python import assemblyai as aai -LITELLM_VIRTUAL_KEY = "sk-1234" # -LITELLM_PROXY_BASE_URL = "http://0.0.0.0:4000/assemblyai" # /assemblyai +aai.settings.base_url = "http://0.0.0.0:4000/assemblyai" # /assemblyai +aai.settings.api_key = "Bearer sk-1234" # Bearer -aai.settings.api_key = f"Bearer {LITELLM_VIRTUAL_KEY}" -aai.settings.base_url = LITELLM_PROXY_BASE_URL +# Use a publicly-accessible URL +audio_file = "https://assembly.ai/wildfires.mp3" -# URL of the file to transcribe -FILE_URL = "https://assembly.ai/wildfires.mp3" +# Or use a local file: +# audio_file = "./example.mp3" -# You can also transcribe a local file by passing in a file path -# FILE_URL = './path/to/file.mp3' +config = aai.TranscriptionConfig( + speech_models=["universal-3-pro", "universal-2"], + language_detection=True, + speaker_labels=True, + # Speech understanding features + # sentiment_analysis=True, + # entity_detection=True, + # auto_chapters=True, + # summarization=True, + # summary_type=aai.SummarizationType.bullets, + # redact_pii=True, + # content_safety=True, +) -transcriber = aai.Transcriber() -transcript = transcriber.transcribe(FILE_URL) -print(transcript) -print(transcript.id) +transcript = aai.Transcriber().transcribe(audio_file, config=config) + +if transcript.status == aai.TranscriptStatus.error: + raise RuntimeError(f"Transcription failed: {transcript.error}") + +print(f"\nFull Transcript:\n\n{transcript.text}") + +# Optionally print speaker diarization results +# for utterance in transcript.utterances: +# print(f"Speaker {utterance.speaker}: {utterance.text}") ``` -## Calling Assembly AI EU endpoints +4. [Prompting with Universal-3 Pro](https://www.assemblyai.com/docs/speech-to-text/prompting) (optional) -If you want to send your request to the Assembly AI EU endpoint, you can do so by setting the `LITELLM_PROXY_BASE_URL` to `/eu.assemblyai` +```python +import assemblyai as aai + +aai.settings.base_url = "http://0.0.0.0:4000/assemblyai" # /assemblyai +aai.settings.api_key = "Bearer sk-1234" # Bearer + +audio_file = "https://assemblyaiassets.com/audios/verbatim.mp3" + +config = aai.TranscriptionConfig( + speech_models=["universal-3-pro", "universal-2"], + language_detection=True, + prompt="Produce a transcript suitable for conversational analysis. Every disfluency is meaningful data. Include: fillers (um, uh, er, ah, hmm, mhm, like, you know, I mean), repetitions (I I, the the), restarts (I was- I went), stutters (th-that, b-but, no-not), and informal speech (gonna, wanna, gotta)", +) + +transcript = aai.Transcriber().transcribe(audio_file, config) + +print(transcript.text) +``` + +## Calling AssemblyAI EU endpoints + +If you want to send your request to the AssemblyAI EU endpoint, you can do so by setting the `LITELLM_PROXY_BASE_URL` to `/eu.assemblyai` ```python import assemblyai as aai -LITELLM_VIRTUAL_KEY = "sk-1234" # -LITELLM_PROXY_BASE_URL = "http://0.0.0.0:4000/eu.assemblyai" # /eu.assemblyai +aai.settings.base_url = "http://0.0.0.0:4000/eu.assemblyai" # /eu.assemblyai +aai.settings.api_key = "Bearer sk-1234" # Bearer -aai.settings.api_key = f"Bearer {LITELLM_VIRTUAL_KEY}" -aai.settings.base_url = LITELLM_PROXY_BASE_URL +# Use a publicly-accessible URL +audio_file = "https://assembly.ai/wildfires.mp3" -# URL of the file to transcribe -FILE_URL = "https://assembly.ai/wildfires.mp3" - -# You can also transcribe a local file by passing in a file path -# FILE_URL = './path/to/file.mp3' +# Or use a local file: +# audio_file = "./path/to/file.mp3" transcriber = aai.Transcriber() -transcript = transcriber.transcribe(FILE_URL) +transcript = transcriber.transcribe(audio_file) print(transcript) print(transcript.id) ``` + +## LLM Gateway + +Use AssemblyAI's [LLM Gateway](https://www.assemblyai.com/docs/llm-gateway) as an OpenAI-compatible provider — a unified API for Claude, GPT, and Gemini models with full LiteLLM logging, guardrails, and cost tracking support. + +[**See Available Models**](https://www.assemblyai.com/docs/llm-gateway#available-models) + +### Usage + +#### LiteLLM Python SDK + +```python +import litellm +import os + +os.environ["ASSEMBLYAI_API_KEY"] = "your-assemblyai-api-key" + +response = litellm.completion( + model="assemblyai/claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "What is the capital of France?"}] +) + +print(response.choices[0].message.content) +``` + +#### LiteLLM Proxy + +1. Config + +```yaml +model_list: + - model_name: assemblyai/* + litellm_params: + model: assemblyai/* + api_key: os.environ/ASSEMBLYAI_API_KEY +``` + +2. Start proxy + +```bash +litellm --config config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +3. Test it! + +```python +import requests + +headers = { + "authorization": "Bearer sk-1234" # Bearer +} + +response = requests.post( + "http://0.0.0.0:4000/v1/chat/completions", + headers=headers, + json={ + "model": "assemblyai/claude-sonnet-4-5-20250929", + "messages": [ + {"role": "user", "content": "What is the capital of France?"} + ], + "max_tokens": 1000 + } +) + +result = response.json() +print(result["choices"][0]["message"]["content"]) +``` diff --git a/docs/my-website/docs/providers/anthropic.md b/docs/my-website/docs/providers/anthropic.md index de5a4dc610c..428cfda4128 100644 --- a/docs/my-website/docs/providers/anthropic.md +++ b/docs/my-website/docs/providers/anthropic.md @@ -4,6 +4,7 @@ import TabItem from '@theme/TabItem'; # Anthropic LiteLLM supports all anthropic models. +- `claude-opus-4-6-20260205` - `claude-sonnet-4-5-20250929` - `claude-opus-4-5-20251101` - `claude-opus-4-1-20250805` @@ -415,7 +416,10 @@ print(response) | Model Name | Function Call | |------------------|--------------------------------------------| +| claude-opus-4-6 | `completion('claude-opus-4-6-20260205', messages)` | `os.environ['ANTHROPIC_API_KEY']` | | claude-sonnet-4-5 | `completion('claude-sonnet-4-5-20250929', messages)` | `os.environ['ANTHROPIC_API_KEY']` | +| claude-opus-4-5 | `completion('claude-opus-4-5-20251101', messages)` | `os.environ['ANTHROPIC_API_KEY']` | +| claude-opus-4-1 | `completion('claude-opus-4-1-20250805', messages)` | `os.environ['ANTHROPIC_API_KEY']` | | claude-opus-4 | `completion('claude-opus-4-20250514', messages)` | `os.environ['ANTHROPIC_API_KEY']` | | claude-sonnet-4 | `completion('claude-sonnet-4-20250514', messages)` | `os.environ['ANTHROPIC_API_KEY']` | | claude-3.7 | `completion('claude-3-7-sonnet-20250219', messages)` | `os.environ['ANTHROPIC_API_KEY']` | diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index e546ed97656..bb07216a295 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -660,7 +660,7 @@ Same as [Anthropic API response](../providers/anthropic#usage---thinking--reason LiteLLM supports Anthropic's beta features on AWS Bedrock through the `anthropic-beta` header. This enables access to experimental features like: -- **1M Context Window** - Up to 1 million tokens of context (Claude Sonnet 4) +- **1M Context Window** - Up to 1 million tokens of context (Claude Opus 4.6, Sonnet 4.5, Sonnet 4) - **Computer Use Tools** - AI that can interact with computer interfaces - **Token-Efficient Tools** - More efficient tool usage patterns - **Extended Output** - Up to 128K output tokens @@ -670,7 +670,7 @@ LiteLLM supports Anthropic's beta features on AWS Bedrock through the `anthropic | Beta Feature | Header Value | Compatible Models | Description | |--------------|-------------|------------------|-------------| -| 1M Context Window | `context-1m-2025-08-07` | Claude Sonnet 4 | Enable 1 million token context window | +| 1M Context Window | `context-1m-2025-08-07` | Claude Opus 4.6, Sonnet 4.5, Sonnet 4 | Enable 1 million token context window | | Computer Use (Latest) | `computer-use-2025-01-24` | Claude 3.7 Sonnet | Latest computer use tools | | Computer Use (Legacy) | `computer-use-2024-10-22` | Claude 3.5 Sonnet v2 | Computer use tools for Claude 3.5 | | Token-Efficient Tools | `token-efficient-tools-2025-02-19` | Claude 3.7 Sonnet | More efficient tool usage | diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index b694549cf40..decffb18833 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -196,6 +196,7 @@ router_settings: | disable_end_user_cost_tracking_prometheus_only | boolean | If true, turns off end user cost tracking on prometheus metrics only. | | key_generation_settings | object | Restricts who can generate keys. [Further docs](./virtual_keys.md#restricting-key-generation) | | disable_add_transform_inline_image_block | boolean | For Fireworks AI models - if true, turns off the auto-add of `#transform=inline` to the url of the image_url, if the model is not a vision model. | +| use_chat_completions_url_for_anthropic_messages | boolean | If true, routes OpenAI `/v1/messages` requests through chat/completions instead of the Responses API. Can also be set via env var `LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES=true`. | | disable_hf_tokenizer_download | boolean | If true, it defaults to using the openai tokenizer for all models (including huggingface models). | | enable_json_schema_validation | boolean | If true, enables json schema validation for all requests. | | disable_copilot_system_to_assistant | boolean | **DEPRECATED** - GitHub Copilot API supports system prompts. | diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 6f98265e40a..2764a6f0d4f 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -330,6 +330,22 @@ model_list: health_check_timeout: 10 # 👈 OVERRIDE HEALTH CHECK TIMEOUT ``` +## Health Check Max Tokens + +By default, health checks use `max_tokens=1` to minimize cost and latency. For wildcard models, the default is `max_tokens=10`. + +You can override this per-model by setting `health_check_max_tokens` in the `model_info` section of your config.yaml. + +```yaml +model_list: + - model_name: openai/gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + model_info: + health_check_max_tokens: 5 # 👈 OVERRIDE HEALTH CHECK MAX TOKENS +``` + ## `/health/readiness` Unprotected endpoint for checking if proxy is ready to accept requests diff --git a/docs/my-website/docs/proxy/prometheus.md b/docs/my-website/docs/proxy/prometheus.md index 18a139d1d29..d8f0d83b59d 100644 --- a/docs/my-website/docs/proxy/prometheus.md +++ b/docs/my-website/docs/proxy/prometheus.md @@ -113,6 +113,31 @@ litellm_settings: ``` +## Pod Health Metrics + +Use these to measure per-pod queue depth and diagnose latency that occurs **before** LiteLLM starts processing a request. + +| Metric Name | Type | Description | +|---|---|---| +| `litellm_in_flight_requests` | Gauge | Number of HTTP requests currently in-flight on this uvicorn worker. Tracks the pod's queue depth in real time. With multiple workers, values are summed across all live workers (`livesum`). | + +### When to use this + +LiteLLM measures latency from when its handler starts. If a request waits in uvicorn's event loop before the handler runs, that wait is invisible to LiteLLM's own logs. `litellm_in_flight_requests` shows how loaded the pod was at any point in time. + +``` +high in_flight_requests + high ALB TargetResponseTime → pod overloaded, scale out +low in_flight_requests + high ALB TargetResponseTime → delay is pre-ASGI (event loop blocking) +``` + +You can also check the current value directly without Prometheus: + +```bash +curl http://localhost:4000/health/backlog \ + -H "Authorization: Bearer sk-..." +# {"in_flight_requests": 47} +``` + ## Proxy Level Tracking Metrics Use this to track overall LiteLLM Proxy usage. diff --git a/docs/my-website/docs/troubleshoot/latency_overhead.md b/docs/my-website/docs/troubleshoot/latency_overhead.md index cfb2cb43a7e..dd7f012dcde 100644 --- a/docs/my-website/docs/troubleshoot/latency_overhead.md +++ b/docs/my-website/docs/troubleshoot/latency_overhead.md @@ -2,9 +2,41 @@ Use this guide when you see unexpected latency overhead between LiteLLM proxy and the LLM provider. +## The Invisible Latency Gap + +LiteLLM measures latency from when its handler starts. If a request waits in uvicorn's event loop **before** the handler runs, that wait is invisible to LiteLLM's own logs. + +``` +T=0 Request arrives at load balancer + [queue wait — LiteLLM never logs this] +T=10 LiteLLM handler starts → timer begins +T=20 Response sent + +LiteLLM logs: 10s User experiences: 20s +``` + +To measure the pre-handler wait, poll `/health/backlog` on each pod: + +```bash +curl http://localhost:4000/health/backlog \ + -H "Authorization: Bearer sk-..." +# {"in_flight_requests": 47} +``` + +Or scrape the `litellm_in_flight_requests` Prometheus gauge at `/metrics`. + +| `in_flight_requests` | ALB `TargetResponseTime` | Diagnosis | +|---|---|---| +| High | High | Pod overloaded → scale out | +| Low | High | Delay is pre-ASGI — check for sync blocking code or event loop saturation | +| High | Normal | Pod is busy but healthy, no queue buildup | + +If you're on **AWS ALB**, correlate `litellm_in_flight_requests` spikes with ALB's `TargetResponseTime` CloudWatch metric. The gap between what ALB reports and what LiteLLM logs is the invisible wait. + ## Quick Checklist -1. **Collect the `x-litellm-overhead-duration-ms` response header** — this tells you LiteLLM's total overhead on every request. Start here. +1. **Check `in_flight_requests` on each pod** via `/health/backlog` or the `litellm_in_flight_requests` Prometheus gauge — this tells you if requests are queuing before LiteLLM starts processing. Start here for unexplained latency. +2. **Collect the `x-litellm-overhead-duration-ms` response header** — this tells you LiteLLM's total overhead on every request. 2. **Is DEBUG logging enabled?** This is the #1 cause of latency with large payloads. 3. **Are you sending large base64 payloads?** (images, PDFs) — see [Large Payload Overhead](#large-payload-overhead). 4. **Enable detailed timing headers** to pinpoint where time is spent. diff --git a/docs/my-website/release_notes/v1.81.14.md b/docs/my-website/release_notes/v1.81.14.md index b3a0018b162..3a133f092ae 100644 --- a/docs/my-website/release_notes/v1.81.14.md +++ b/docs/my-website/release_notes/v1.81.14.md @@ -489,6 +489,71 @@ graph LR --- +## Security + +We run [Grype](https://github.com/anchore/grype) and [Trivy](https://github.com/aquasecurity/trivy) security scans on every LiteLLM Docker image. Here's the vulnerability report for this release across all published images: + +### Docker Image Scan Summary + +| Image | Critical | High | Medium | Low | +|-------|----------|------|--------|-----| +| `ghcr.io/berriai/litellm:main-latest` | **0** ✅ | 4 unique CVEs | 4 | 1 | +| `ghcr.io/berriai/litellm-ee:main-latest` | **0** ✅ | 4 unique CVEs | 4 | 1 | +| `ghcr.io/berriai/litellm-non_root:main-latest` | **1** | 11 unique CVEs | 6 | 2 | +| `ghcr.io/berriai/litellm-database:main-latest` | **1** | 7 unique CVEs | 5 | 1 | +| `ghcr.io/berriai/litellm-spend_logs:main-latest` | **4** | 35 matches | 40 | 10 | + +:::note +Vulnerability counts are based on full image scans including build-time tooling. High match counts are often inflated by packages like `minimatch` appearing at multiple versions; the unique CVE counts above reflect the actual distinct vulnerabilities. +::: + +### Critical Severity + +**1. Node.js Critical (non-root, database, spend_logs images):** +Node.js 24.12.0 is used **only** for the Admin UI build and Prisma client generation — it is **not** part of the LiteLLM Python application runtime. + +| Package | Vulnerability | Description | Fix Version | +|---------|---------------|-------------|-------------| +| `node` | CVE-2025-55130 | Node.js critical vulnerability | 20.20.0 | + +**2. OpenSSL & Go Critical (spend_logs image only):** +The `spend_logs` image contains additional vulnerabilities in the underlying Go modules and system libraries. + +| Package | Vulnerability | Description | Fix Version | +|---------|---------------|-------------|-------------| +| `libcrypto3`, `libssl3` | CVE-2025-15467 | OpenSSL critical vulnerability | 3.3.6-r0 | +| `stdlib` (Go) | CVE-2025-68121 | Go standard library critical vulnerability | 1.24.13+ | + +### High Severity + +All high-severity vulnerabilities are in **npm/Node.js build-time dependencies** or system-level libraries — they are **not** in the LiteLLM Python application code. + +**Present in all images:** + +| Package | Vulnerability | Description | Fix Version | +|---------|---------------|-------------|-------------| +| `minimatch` | CVE-2026-26996 | DoS via specially crafted glob patterns | 10.2.1+ / 9.0.6+ | +| `minimatch` | CVE-2026-27903 | DoS due to unbounded recursive backtracking | 10.2.3+ / 9.0.7+ | +| `minimatch` | CVE-2026-27904 | DoS via catastrophic backtracking in glob expressions | 10.2.3+ / 9.0.7+ | +| `tar` | CVE-2026-26960 / GHSA-83g3-92jg-28cx | Arbitrary file read/write via malicious archive hardlinks | 7.5.8 | + +### Medium Severity (all images) + +| Package | Vulnerability | Status | +|---------|---------------|--------| +| `pypdf` 6.7.2 | GHSA-x7hp-r3qg-r3cj | Fix available in 6.7.3 | +| Python 3.13 | CVE-2025-15366, CVE-2025-15367, CVE-2025-12781 | No upstream fix available | + +### Recommendations + +- **LiteLLM Main & EE images** (`litellm:main-latest`, `litellm-ee:main-latest`) have the best security posture with **0 critical vulnerabilities**. +- All HIGH/CRITICAL findings in the main images relate to build-time Node.js/npm tooling, not the Python runtime. +- We are actively monitoring upstream Python and system library fixes for remaining medium-severity vulnerabilities. + +To report a security vulnerability, email support@berri.ai with details and steps to reproduce. + +--- + ## Documentation Updates - Add OpenAI Agents SDK with LiteLLM guide - [PR #21311](https://github.com/BerriAI/litellm/pull/21311) diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49-py3-none-any.whl new file mode 100644 index 00000000000..e44b58f8e63 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49.tar.gz new file mode 100644 index 00000000000..2c8549ad069 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql new file mode 100644 index 00000000000..594ab9ac1a2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "agent_id" TEXT; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql new file mode 100644 index 00000000000..e2a3694e8ef --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql @@ -0,0 +1,18 @@ +-- CreateTable +CREATE TABLE "LiteLLM_ClaudeCodePluginTable" ( + "id" TEXT NOT NULL, + "name" TEXT NOT NULL, + "version" TEXT, + "description" TEXT, + "manifest_json" TEXT, + "files_json" TEXT DEFAULT '{}', + "enabled" BOOLEAN NOT NULL DEFAULT true, + "created_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + + CONSTRAINT "LiteLLM_ClaudeCodePluginTable_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_ClaudeCodePluginTable_name_key" ON "LiteLLM_ClaudeCodePluginTable"("name"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 440c9c1d829..2717480c7ef 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -300,7 +300,7 @@ model LiteLLM_MCPServerTable { token_url String? registration_url String? allow_all_keys Boolean @default(false) - available_on_public_internet Boolean @default(false) + available_on_public_internet Boolean @default(true) } // Generate Tokens for Proxy @@ -390,6 +390,7 @@ model LiteLLM_DeletedVerificationToken { config Json @default("{}") user_id String? team_id String? + agent_id String? project_id String? permissions Json @default("{}") max_parallel_requests Int? @@ -504,6 +505,7 @@ model LiteLLM_SpendLogs { agent_id String? proxy_server_request Json? @default("{}") @@index([startTime]) + @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) } diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index bd57b248cdb..968536712dc 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.48" +version = "0.4.49" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.48" +version = "0.4.49" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/__init__.py b/litellm/__init__.py index 6e42f2c1ea5..50fa0e76755 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -197,6 +197,9 @@ telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) +use_chat_completions_url_for_anthropic_messages: bool = bool( + os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) +) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API retry = True ### AUTH ### api_key: Optional[str] = None diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index 5dc16a224c7..331aa8f51cd 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -3,22 +3,29 @@ Add the event loop to the cache key, to prevent event loop closed errors. """ import asyncio +from typing import Set from .in_memory_cache import InMemoryCache class LLMClientCache(InMemoryCache): + # Background tasks must be stored to prevent garbage collection, which would + # trigger "coroutine was never awaited" warnings. See: + # https://docs.python.org/3/library/asyncio-task.html#creating-tasks + # Intentionally shared across all instances as a global task registry. + _background_tasks: Set[asyncio.Task] = set() + def _remove_key(self, key: str) -> None: """Close async clients before evicting them to prevent connection pool leaks.""" value = self.cache_dict.get(key) super()._remove_key(key) if value is not None: - close_fn = getattr(value, "aclose", None) or getattr( - value, "close", None - ) + close_fn = getattr(value, "aclose", None) or getattr(value, "close", None) if close_fn and asyncio.iscoroutinefunction(close_fn): try: - asyncio.get_running_loop().create_task(close_fn()) + task = asyncio.get_running_loop().create_task(close_fn()) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) except RuntimeError: pass elif close_fn and callable(close_fn): diff --git a/litellm/constants.py b/litellm/constants.py index b1a0021bcc6..4c38ecd74b5 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -49,6 +49,14 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int( ) DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) +# Maximum wall-clock seconds a streaming response is allowed to run. +# Streams exceeding this duration are terminated with a Timeout error. +# None (default) = no limit. Set env var to a number of seconds to enable globally. +_max_stream_duration_env = os.getenv("LITELLM_MAX_STREAMING_DURATION_SECONDS", None) +LITELLM_MAX_STREAMING_DURATION_SECONDS = ( + float(_max_stream_duration_env) if _max_stream_duration_env is not None else None +) + # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. # Set to 0 to disable truncation. @@ -185,9 +193,9 @@ _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client fo # Aiohttp connection pooling - prevents memory leaks from unbounded connection growth # Set to 0 for unlimited (not recommended for production) -AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 300)) +AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000)) AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int( - os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 50) + os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500) ) AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index eb027334606..b36d4ef877c 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -955,7 +955,8 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore generated_content: str = "", is_pre_first_chunk: bool = False, ): - self.status_code = 503 # Service Unavailable + original_status = getattr(original_exception, "status_code", None) + self.status_code = int(original_status) if original_status is not None else 503 self.message = f"litellm.MidStreamFallbackError: {message}" self.model = model self.llm_provider = llm_provider @@ -978,7 +979,14 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore else: self.response = response - # Call the parent constructor + # Save the original attributes before they are overridden by ServiceUnavailableError + _saved_response = self.response + _saved_request = getattr(self.response, "request", None) or httpx.Request( + method="POST", url=f"https://{llm_provider}.com/v1/" + ) + _saved_message = self.message + + # Call the parent constructor (which hardcodes status_code=503 and modifies the response object) super().__init__( message=self.message, llm_provider=llm_provider, @@ -988,6 +996,13 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore max_retries=self.max_retries, num_retries=self.num_retries, ) + + # Restore the propagated status and original response/request objects + self.status_code = int(original_status) if original_status is not None else 503 + self.response = _saved_response + self.request = _saved_request + self.message = _saved_message + self.args = (_saved_message,) def __str__(self): _message = self.message diff --git a/litellm/images/main.py b/litellm/images/main.py index 6c4c502a7b0..236266af6ad 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -483,6 +483,7 @@ def image_generation( # noqa: PLR0915 organization=organization, aimg_generation=aimg_generation, client=client, + headers=headers, ) elif custom_llm_provider == "bedrock": if model is None: diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 08db77e8571..7a08432b9a1 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2686,6 +2686,8 @@ class PrometheusLogger(CustomLogger): if team_info: team_object.budget_reset_at = team_info.budget_reset_at + if team_object.max_budget is None and team_info.max_budget is not None: + team_object.max_budget = team_info.max_budget return team_object @@ -2903,6 +2905,8 @@ class PrometheusLogger(CustomLogger): if user_info: user_object.budget_reset_at = user_info.budget_reset_at + if user_object.max_budget is None and user_info.max_budget is not None: + user_object.max_budget = user_info.max_budget return user_object diff --git a/litellm/litellm_core_utils/dd_tracing.py b/litellm/litellm_core_utils/dd_tracing.py index ce784ecf6a8..ae4f46c38bd 100644 --- a/litellm/litellm_core_utils/dd_tracing.py +++ b/litellm/litellm_core_utils/dd_tracing.py @@ -5,7 +5,7 @@ If the ddtrace package is not installed, the tracer will be a no-op. """ from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, Any, Optional, Union from litellm.secret_managers.main import get_secret_bool @@ -76,3 +76,48 @@ if should_use_dd_tracer: tracer = NullTracer() else: tracer = NullTracer() + + +def get_active_span() -> Optional[Any]: + """ + Return the active Datadog span, checking current span first and then root span. + """ + try: + current_span_fn = getattr(tracer, "current_span", None) + if callable(current_span_fn): + current_span = current_span_fn() + if current_span is not None: + return current_span + + current_root_span_fn = getattr(tracer, "current_root_span", None) + if callable(current_root_span_fn): + return current_root_span_fn() + except Exception: + return None + return None + + +def set_active_span_tag(tag_key: str, tag_value: str) -> bool: + """ + Best-effort helper to set a tag on the active Datadog span. + + Returns: + bool: True if a span tag was set, False otherwise. + """ + if not tag_key or tag_value is None: + return False + + span = get_active_span() + if span is None: + return False + + try: + if hasattr(span, "set_tag_str"): + span.set_tag_str(tag_key, str(tag_value)) + return True + if hasattr(span, "set_tag"): + span.set_tag(tag_key, str(tag_value)) + return True + except Exception: + return False + return False diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 47a27c8ef5b..9e972f1910b 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -14,7 +14,6 @@ TEST_PDF_URL = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9U class HealthCheckHelpers: - @staticmethod async def ahealth_check_wildcard_models( model: str, @@ -44,7 +43,9 @@ class HealthCheckHelpers: model_params["model"] = cheapest_models[0] model_params["litellm_logging_obj"] = litellm_logging_obj model_params["fallbacks"] = fallback_models - model_params["max_tokens"] = 10 # gpt-5-nano throws errors for max_tokens=1 + model_params["max_tokens"] = model_params.get( + "max_tokens", 10 + ) # gpt-5-nano throws errors for max_tokens=1 await acompletion(**model_params) return {} @@ -130,7 +131,7 @@ class HealthCheckHelpers: Callable, ]: """ - Returns a dictionary of mode handlers for health check calls. + Returns a dictionary of mode handlers for health check calls. Mode Handlers are Callables that need to be run for execution of the health check call. @@ -215,4 +216,4 @@ class HealthCheckHelpers: "document_url": TEST_PDF_URL, }, ), - } \ No newline at end of file + } diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index a9fd0f4ea8a..bf0b2709365 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -8,9 +8,11 @@ from litellm._logging import verbose_logger from litellm.types.utils import ( CacheCreationTokenDetails, CallTypes, + CompletionTokensDetailsWrapper, ImageResponse, ModelInfo, PassthroughCallTypes, + PromptTokensDetailsWrapper, ServiceTier, Usage, ) @@ -767,6 +769,64 @@ def generic_cost_per_token( # noqa: PLR0915 return prompt_cost, completion_cost +def calculate_image_response_cost_from_usage( + model: str, + image_response: ImageResponse, + custom_llm_provider: str, +) -> Optional[float]: + """ + Calculate image generation cost from usage metadata when available. + + Returns: + Optional[float]: total cost from token usage, or None when usage metadata + is missing/incomplete and caller should fall back to flat per-image pricing. + """ + usage = image_response.usage + if usage is None: + return None + + prompt_tokens = usage.input_tokens + completion_tokens = usage.output_tokens + total_tokens = usage.total_tokens + + if prompt_tokens is None or completion_tokens is None or total_tokens is None: + return None + + # ImageResponse may carry a default zeroed usage object even when provider + # usage metadata is absent. Treat this as missing usage and fall back. + if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0: + return None + + input_tokens_details = getattr(usage, "input_tokens_details", None) + prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None + if input_tokens_details is not None: + prompt_tokens_details = PromptTokensDetailsWrapper( + text_tokens=getattr(input_tokens_details, "text_tokens", None), + image_tokens=getattr(input_tokens_details, "image_tokens", None), + cached_tokens=0, + ) + + normalized_usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + prompt_tokens_details=prompt_tokens_details, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=0, + image_tokens=completion_tokens, + reasoning_tokens=0, + audio_tokens=0, + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=normalized_usage, + custom_llm_provider=custom_llm_provider, + ) + return prompt_cost + completion_cost + + class CostCalculatorUtils: @staticmethod def _call_type_has_image_response(call_type: str) -> bool: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index ba415af9a5a..796223ff8e1 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1766,6 +1766,7 @@ def convert_function_to_anthropic_tool_invoke( def convert_to_anthropic_tool_invoke( tool_calls: List[ChatCompletionAssistantToolCall], web_search_results: Optional[List[Any]] = None, + tool_results: Optional[List[Any]] = None, ) -> List[Union[AnthropicMessagesToolUseParam, Dict[str, Any]]]: """ OpenAI tool invokes: @@ -1840,12 +1841,18 @@ def convert_to_anthropic_tool_invoke( } anthropic_tool_invoke.append(_anthropic_server_tool_use) - # Add corresponding web_search_tool_result if available + # Add corresponding tool result if available. + # Check both web_search_results (web_search_tool_result / web_fetch_tool_result) + # and tool_results (bash_code_execution_tool_result, etc.) + _all_tool_results: List[Any] = [] if web_search_results: - for result in web_search_results: - if result.get("tool_use_id") == tool_id: - anthropic_tool_invoke.append(result) - break + _all_tool_results.extend(web_search_results) + if tool_results: + _all_tool_results.extend(tool_results) + for result in _all_tool_results: + if result.get("tool_use_id") == tool_id: + anthropic_tool_invoke.append(result) + break else: # Regular tool_use sanitized_tool_id = _sanitize_anthropic_tool_use_id(tool_id) @@ -2472,9 +2479,10 @@ def anthropic_messages_pt( # noqa: PLR0915 # Pass through as-is since these are Anthropic-native content types elif m.get("type", "") == "server_tool_use": assistant_content.append(m) # type: ignore - # handle tool_search_tool_result blocks + # handle all *_tool_result blocks (tool_search_tool_result, + # web_search_tool_result, bash_code_execution_tool_result, etc.) # Pass through as-is since these are Anthropic-native content types - elif m.get("type", "") == "tool_search_tool_result": + elif m.get("type", "").endswith("_tool_result"): assistant_content.append(m) # type: ignore elif ( "content" in assistant_content_block @@ -2504,7 +2512,8 @@ def anthropic_messages_pt( # noqa: PLR0915 if ( assistant_tool_calls is not None ): # support assistant tool invoke conversion - # Get web_search_results from provider_specific_fields for server_tool_use reconstruction + # Get web_search_results and tool_results from provider_specific_fields + # for server_tool_use reconstruction. # Fixes: https://github.com/BerriAI/litellm/issues/17737 _provider_specific_fields_raw = assistant_content_block.get( "provider_specific_fields" @@ -2517,9 +2526,11 @@ def anthropic_messages_pt( # noqa: PLR0915 _web_search_results = _provider_specific_fields.get( "web_search_results" ) + _tool_results = _provider_specific_fields.get("tool_results") tool_invoke_results = convert_to_anthropic_tool_invoke( assistant_tool_calls, web_search_results=_web_search_results, + tool_results=_tool_results, ) # Prevent "tool_use ids must be unique" errors by filtering duplicates diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 8df41aea4a3..294f9c485c1 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -72,6 +72,9 @@ class RealTimeStreaming: self.request_data: Dict = request_data or {} # Violation counter for end_session_after_n_fails support self._violation_count: int = 0 + # When a text message is blocked, hold the guardrail reason so the next + # response.create can be rewritten to include the failure context. + self._pending_guardrail_message: Optional[str] = None def _should_store_message( self, @@ -230,9 +233,9 @@ class RealTimeStreaming: message, self.model, self.session_configuration_request ) for msg in transformed: - await self.backend_ws.send(msg) + await self.backend_ws.send(msg) # type: ignore[union-attr] else: - await self.backend_ws.send(message) + await self.backend_ws.send(message) # type: ignore[union-attr] def _has_realtime_guardrails(self) -> bool: """Return True if any callback is registered for realtime guardrail event types.""" @@ -261,18 +264,12 @@ class RealTimeStreaming: When this returns True, we inject a session.update to disable the LLM's auto-response so the guardrail can gate it first. - """ - from litellm.integrations.custom_guardrail import CustomGuardrail - from litellm.types.guardrails import GuardrailEventHooks - return any( - isinstance(cb, CustomGuardrail) - and cb.should_run_guardrail( - data=self.request_data, - event_type=GuardrailEventHooks.realtime_input_transcription, - ) - for cb in litellm.callbacks - ) + Must match the same hook criteria as run_realtime_guardrails() so that + any guardrail that would actually check the transcript also disables + auto-response before the transcript arrives. + """ + return self._has_realtime_guardrails() async def run_realtime_guardrails( self, @@ -335,18 +332,35 @@ class RealTimeStreaming: # Use realtime_violation_message if configured; fall back to guardrail error text. error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg - # Return the error directly to the WebSocket consumer. + # Cancel any in-progress LLM response (e.g. VAD auto-response). + await self._send_to_backend(json.dumps({"type": "response.cancel"})) + # Send the policy violation hint (shows as small gray status text in UI). await self.websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "guardrail_violation", - "message": error_msg, - "code": "content_policy_violation", - }, - } - ) + json.dumps({ + "type": "error", + "error": { + "type": "guardrail_violation", + "message": error_msg, + "code": "content_policy_violation", + }, + }) + ) + # Ask the LLM to voice the exact guardrail message so the + # user hears it as audio in voice sessions (not just text). + guardrail_prompt = ( + f"Say exactly the following message to the user, word for word, " + f"do not add anything else: {error_msg}" + ) + await self._send_to_backend(json.dumps({ + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": guardrail_prompt}], + }, + })) + await self._send_to_backend( + json.dumps({"type": "response.create"}) ) self._violation_count += 1 @@ -362,7 +376,7 @@ class RealTimeStreaming: "[realtime guardrail] ending session after violation %d", self._violation_count, ) - await self.backend_ws.close() + await self.backend_ws.close() # type: ignore[union-attr] verbose_logger.warning( "[realtime guardrail] BLOCKED transcript (violation %d): %r", @@ -502,11 +516,11 @@ class RealTimeStreaming: try: while True: try: - raw_response = await self.backend_ws.recv( + raw_response = await self.backend_ws.recv( # type: ignore[union-attr] decode=False ) # improves performance except TypeError: - raw_response = await self.backend_ws.recv() # type: ignore[assignment] + raw_response = await self.backend_ws.recv() # type: ignore[union-attr, assignment] if self.provider_config: try: @@ -559,7 +573,17 @@ class RealTimeStreaming: combined_text ) if blocked: - continue # don't forward to backend + # Store the guardrail reason so the next response.create + # (sent automatically by the client) is rewritten to + # include it as response instructions. + self._pending_guardrail_message = combined_text + continue # don't forward the original blocked message + + if msg_type == "response.create" and self._pending_guardrail_message: + # The guardrail already sent the synthetic AI bubble — drop this + # response.create so OpenAI doesn't generate an additional response. + self._pending_guardrail_message = None + continue except (json.JSONDecodeError, AttributeError): pass @@ -573,9 +597,9 @@ class RealTimeStreaming: ) for msg in message: - await self.backend_ws.send(msg) + await self.backend_ws.send(msg) # type: ignore[union-attr] else: - await self.backend_ws.send(message) + await self.backend_ws.send(message) # type: ignore[union-attr] except Exception as e: verbose_logger.debug(f"Error in client ack messages: {e}") diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ccd5c1dd8f5..3b75a56fcc9 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -96,6 +96,7 @@ class CustomStreamWrapper: self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False + self._stream_created_time: float = time.time() litellm_params: GenericLiteLLMParams = GenericLiteLLMParams( **self.logging_obj.model_call_details.get("litellm_params", {}) @@ -161,6 +162,20 @@ class CustomStreamWrapper: self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) self.created: Optional[int] = None + def _check_max_streaming_duration(self) -> None: + """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" + from litellm.constants import LITELLM_MAX_STREAMING_DURATION_SECONDS + + if LITELLM_MAX_STREAMING_DURATION_SECONDS is None: + return + elapsed = time.time() - self._stream_created_time + if elapsed > LITELLM_MAX_STREAMING_DURATION_SECONDS: + raise litellm.Timeout( + message=f"Stream exceeded max streaming duration of {LITELLM_MAX_STREAMING_DURATION_SECONDS}s (elapsed {elapsed:.1f}s)", + model=self.model or "", + llm_provider=self.custom_llm_provider or "", + ) + def __iter__(self) -> Iterator["ModelResponseStream"]: return self @@ -1236,27 +1251,27 @@ class CustomStreamWrapper: else: completion_obj["content"] = str(chunk) elif self.custom_llm_provider == "petals": - if len(self.completion_stream) == 0: + if self.completion_stream is None or len(self.completion_stream) == 0: if self.received_finish_reason is not None: raise StopIteration else: self.received_finish_reason = "stop" chunk_size = 30 - new_chunk = self.completion_stream[:chunk_size] + new_chunk = self.completion_stream[:chunk_size] # type: ignore[index] completion_obj["content"] = new_chunk - self.completion_stream = self.completion_stream[chunk_size:] + self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index] elif self.custom_llm_provider == "palm": # fake streaming response_obj = {} - if len(self.completion_stream) == 0: + if self.completion_stream is None or len(self.completion_stream) == 0: if self.received_finish_reason is not None: raise StopIteration else: self.received_finish_reason = "stop" chunk_size = 30 - new_chunk = self.completion_stream[:chunk_size] + new_chunk = self.completion_stream[:chunk_size] # type: ignore[index] completion_obj["content"] = new_chunk - self.completion_stream = self.completion_stream[chunk_size:] + self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index] elif self.custom_llm_provider == "triton": response_obj = self.handle_triton_stream(chunk) completion_obj["content"] = response_obj["text"] @@ -1743,6 +1758,7 @@ class CustomStreamWrapper: and self.custom_llm_provider == "cached_response" ): cache_hit = True + self._check_max_streaming_duration() try: if self.completion_stream is None: self.fetch_sync_stream() @@ -1755,7 +1771,7 @@ class CustomStreamWrapper: ): chunk = self.completion_stream else: - chunk = next(self.completion_stream) + chunk = next(self.completion_stream) # type: ignore[arg-type] if chunk is not None and chunk != b"": print_verbose( f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk.decode('utf-8', errors='replace') if isinstance(chunk, bytes) else chunk}; custom_llm_provider: {self.custom_llm_provider}" @@ -1917,12 +1933,13 @@ class CustomStreamWrapper: and self.custom_llm_provider == "cached_response" ): cache_hit = True + self._check_max_streaming_duration() try: if self.completion_stream is None: await self.fetch_stream() if is_async_iterable(self.completion_stream): - async for chunk in self.completion_stream: + async for chunk in self.completion_stream: # type: ignore[union-attr] if chunk == "None" or chunk is None: continue # skip None chunks @@ -1951,22 +1968,24 @@ class CustomStreamWrapper: self.rules.post_call_rules( input=self.response_uptil_now, model=self.model ) - # Store a shallow copy so usage stripping below - # does not mutate the stored chunk. - self.chunks.append(processed_chunk.model_copy()) - # Add mcp_list_tools to first chunk if present if not self.sent_first_chunk: processed_chunk = self._add_mcp_list_tools_to_first_chunk(processed_chunk) self.sent_first_chunk = True - if ( + + _has_usage = ( hasattr(processed_chunk, "usage") and getattr(processed_chunk, "usage", None) is not None - ): + ) + + if _has_usage: + # Store a copy ONLY when usage stripping below will mutate + # the chunk. For non-usage chunks (vast majority), store + # directly to avoid expensive model_copy() per chunk. + self.chunks.append(processed_chunk.model_copy()) + # Strip usage from the outgoing chunk so it's not sent twice # (once in the chunk, once in _hidden_params). - # Create a new object without usage, matching sync behavior. - # The copy in self.chunks retains usage for calculate_total_usage(). obj_dict = processed_chunk.model_dump() if "usage" in obj_dict: del obj_dict["usage"] @@ -1978,6 +1997,9 @@ class CustomStreamWrapper: ) if is_empty: continue + else: + # No usage data — safe to store directly without copying + self.chunks.append(processed_chunk) # add usage as hidden param if self.sent_last_chunk is True and self.stream_options is None: @@ -2004,7 +2026,7 @@ class CustomStreamWrapper: ): chunk = self.completion_stream else: - chunk = next(self.completion_stream) + chunk = next(self.completion_stream) # type: ignore[arg-type] if chunk is not None and chunk != b"": processed_chunk = self.chunk_creator(chunk=chunk) if processed_chunk is None: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 7e5a4f22a7f..5b215c1fe54 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -25,8 +25,24 @@ from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client from ..adapters.handler import LiteLLMMessagesToCompletionTransformationHandler +from ..responses_adapters.handler import LiteLLMMessagesToResponsesAPIHandler from .utils import AnthropicMessagesRequestUtils, mock_response +# Providers that are routed directly to the OpenAI Responses API instead of +# going through chat/completions. +_RESPONSES_API_PROVIDERS = frozenset({"openai"}) + + +def _should_route_to_responses_api(custom_llm_provider: Optional[str]) -> bool: + """Return True when the provider should use the Responses API path. + + Set ``litellm.use_chat_completions_url_for_anthropic_messages = True`` to + opt out and route OpenAI/Azure requests through chat/completions instead. + """ + if litellm.use_chat_completions_url_for_anthropic_messages: + return False + return custom_llm_provider in _RESPONSES_API_PROVIDERS + ####### ENVIRONMENT VARIABLES ################### # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() @@ -282,29 +298,34 @@ def anthropic_messages_handler( ) ) if anthropic_messages_provider_config is None: - # Handle non-Anthropic models using the adapter - return ( - LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( - max_tokens=max_tokens, - messages=messages, - model=model, - metadata=metadata, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - _is_async=is_async, - api_key=api_key, - api_base=api_base, - client=client, - custom_llm_provider=custom_llm_provider, - **kwargs, + # Route to Responses API for OpenAI / Azure, chat/completions for everything else. + _shared_kwargs = dict( + max_tokens=max_tokens, + messages=messages, + model=model, + metadata=metadata, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + _is_async=is_async, + api_key=api_key, + api_base=api_base, + client=client, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + if _should_route_to_responses_api(custom_llm_provider): + return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( + **_shared_kwargs ) + return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( + **_shared_kwargs ) if custom_llm_provider is None: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py new file mode 100644 index 00000000000..6ad3c7b0164 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py @@ -0,0 +1,3 @@ +from .transformation import LiteLLMAnthropicToResponsesAPIAdapter + +__all__ = ["LiteLLMAnthropicToResponsesAPIAdapter"] diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py new file mode 100644 index 00000000000..c268d6c5be8 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -0,0 +1,229 @@ +""" +Handler for the Anthropic v1/messages -> OpenAI Responses API path. + +Used when the target model is an OpenAI or Azure model. +""" + +from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Union + +import litellm +from litellm.types.llms.anthropic import AnthropicMessagesRequest +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, +) +from litellm.types.llms.openai import ResponsesAPIResponse + +from .streaming_iterator import AnthropicResponsesStreamWrapper +from .transformation import LiteLLMAnthropicToResponsesAPIAdapter + +_ADAPTER = LiteLLMAnthropicToResponsesAPIAdapter() + + +def _build_responses_kwargs( + *, + max_tokens: int, + messages: List[Dict], + model: str, + context_management: Optional[Dict] = None, + metadata: Optional[Dict] = None, + output_config: Optional[Dict] = None, + stop_sequences: Optional[List[str]] = None, + stream: Optional[bool] = False, + system: Optional[str] = None, + temperature: Optional[float] = None, + thinking: Optional[Dict] = None, + tool_choice: Optional[Dict] = None, + tools: Optional[List[Dict]] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + output_format: Optional[Dict] = None, + extra_kwargs: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """ + Build the kwargs dict to pass directly to litellm.responses() / litellm.aresponses(). + """ + # Build a typed AnthropicMessagesRequest for the adapter + request_data: Dict[str, Any] = {"model": model, "messages": messages, "max_tokens": max_tokens} + if context_management: + request_data["context_management"] = context_management + if output_config: + request_data["output_config"] = output_config + if metadata: + request_data["metadata"] = metadata + if system: + request_data["system"] = system + if temperature is not None: + request_data["temperature"] = temperature + if thinking: + request_data["thinking"] = thinking + if tool_choice: + request_data["tool_choice"] = tool_choice + if tools: + request_data["tools"] = tools + if top_p is not None: + request_data["top_p"] = top_p + if output_format: + request_data["output_format"] = output_format + + anthropic_request = AnthropicMessagesRequest(**request_data) + responses_kwargs = _ADAPTER.translate_request(anthropic_request) + + if stream: + responses_kwargs["stream"] = True + + # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.) + excluded = {"anthropic_messages"} + for key, value in (extra_kwargs or {}).items(): + if key == "litellm_logging_obj" and value is not None: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObject, + ) + from litellm.types.utils import CallTypes + + if isinstance(value, LiteLLMLoggingObject): + # Reclassify as acompletion so the success handler doesn't try to + # validate the Responses API event as an AnthropicResponse. + # (Mirrors the pattern used in LiteLLMMessagesToCompletionTransformationHandler.) + setattr(value, "call_type", CallTypes.acompletion.value) + responses_kwargs[key] = value + elif key not in excluded and key not in responses_kwargs and value is not None: + responses_kwargs[key] = value + + return responses_kwargs + + +class LiteLLMMessagesToResponsesAPIHandler: + """ + Handles Anthropic /v1/messages requests for OpenAI / Azure models by + calling litellm.responses() / litellm.aresponses() directly and translating + the response back to Anthropic format. + """ + + @staticmethod + async def async_anthropic_messages_handler( + max_tokens: int, + messages: List[Dict], + model: str, + context_management: Optional[Dict] = None, + metadata: Optional[Dict] = None, + output_config: Optional[Dict] = None, + stop_sequences: Optional[List[str]] = None, + stream: Optional[bool] = False, + system: Optional[str] = None, + temperature: Optional[float] = None, + thinking: Optional[Dict] = None, + tool_choice: Optional[Dict] = None, + tools: Optional[List[Dict]] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + output_format: Optional[Dict] = None, + **kwargs, + ) -> Union[AnthropicMessagesResponse, AsyncIterator]: + responses_kwargs = _build_responses_kwargs( + max_tokens=max_tokens, + messages=messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + extra_kwargs=kwargs, + ) + + result = await litellm.aresponses(**responses_kwargs) + + if stream: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) + return wrapper.async_anthropic_sse_wrapper() + + if not isinstance(result, ResponsesAPIResponse): + raise ValueError(f"Expected ResponsesAPIResponse, got {type(result)}") + + return _ADAPTER.translate_response(result) + + @staticmethod + def anthropic_messages_handler( + max_tokens: int, + messages: List[Dict], + model: str, + context_management: Optional[Dict] = None, + metadata: Optional[Dict] = None, + output_config: Optional[Dict] = None, + stop_sequences: Optional[List[str]] = None, + stream: Optional[bool] = False, + system: Optional[str] = None, + temperature: Optional[float] = None, + thinking: Optional[Dict] = None, + tool_choice: Optional[Dict] = None, + tools: Optional[List[Dict]] = None, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + output_format: Optional[Dict] = None, + _is_async: bool = False, + **kwargs, + ) -> Union[ + AnthropicMessagesResponse, + AsyncIterator[Any], + Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any]]], + ]: + if _is_async: + return LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=max_tokens, + messages=messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + **kwargs, + ) + + # Sync path + responses_kwargs = _build_responses_kwargs( + max_tokens=max_tokens, + messages=messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + extra_kwargs=kwargs, + ) + + result = litellm.responses(**responses_kwargs) + + if stream: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) + return wrapper.async_anthropic_sse_wrapper() + + if not isinstance(result, ResponsesAPIResponse): + raise ValueError(f"Expected ResponsesAPIResponse, got {type(result)}") + + return _ADAPTER.translate_response(result) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py new file mode 100644 index 00000000000..0e6268e82f3 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -0,0 +1,265 @@ +# What is this? +## Translates OpenAI call to Anthropic `/v1/messages` format +import json +import traceback +from collections import deque +from typing import Any, AsyncIterator, Dict + +from litellm import verbose_logger +from litellm._uuid import uuid + + +class AnthropicResponsesStreamWrapper: + """ + Wraps a Responses API streaming iterator and re-emits events in Anthropic SSE format. + + Responses API event flow (relevant subset): + response.created -> message_start + response.output_item.added -> content_block_start (if message/function_call) + response.output_text.delta -> content_block_delta (text_delta) + response.reasoning_summary_text.delta -> content_block_delta (thinking_delta) + response.function_call_arguments.delta -> content_block_delta (input_json_delta) + response.output_item.done -> content_block_stop + response.completed -> message_delta + message_stop + """ + + def __init__( + self, + responses_stream: Any, + model: str, + ) -> None: + self.responses_stream = responses_stream + self.model = model + self._message_id: str = f"msg_{uuid.uuid4()}" + self._current_block_index: int = -1 + # Map item_id -> content_block_index so we can stop the right block later + self._item_id_to_block_index: Dict[str, int] = {} + # Track open function_call items by item_id so we can emit tool_use start + self._pending_tool_ids: Dict[str, str] = {} # item_id -> call_id / name accumulator + self._sent_message_start = False + self._sent_message_stop = False + self._chunk_queue: deque = deque() + + def _make_message_start(self) -> Dict[str, Any]: + return { + "type": "message_start", + "message": { + "id": self._message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": self.model, + "stop_reason": None, + "stop_sequence": None, + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, + }, + } + + def _next_block_index(self) -> int: + self._current_block_index += 1 + return self._current_block_index + + def _process_event(self, event: Any) -> None: + """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" + event_type = getattr(event, "type", None) + if event_type is None and isinstance(event, dict): + event_type = event.get("type") + + if event_type is None: + return + + # ---- message_start ---- + if event_type == "response.created": + self._sent_message_start = True + self._chunk_queue.append(self._make_message_start()) + return + + # ---- content_block_start for a new output message item ---- + if event_type == "response.output_item.added": + item = getattr(event, "item", None) or (event.get("item") if isinstance(event, dict) else None) + if item is None: + return + item_type = getattr(item, "type", None) or (item.get("type") if isinstance(item, dict) else None) + item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) + + if item_type == "message": + block_idx = self._next_block_index() + if item_id: + self._item_id_to_block_index[item_id] = block_idx + self._chunk_queue.append({ + "type": "content_block_start", + "index": block_idx, + "content_block": {"type": "text", "text": ""}, + }) + elif item_type == "function_call": + call_id = getattr(item, "call_id", None) or (item.get("call_id") if isinstance(item, dict) else None) or "" + name = getattr(item, "name", None) or (item.get("name") if isinstance(item, dict) else None) or "" + block_idx = self._next_block_index() + if item_id: + self._item_id_to_block_index[item_id] = block_idx + self._pending_tool_ids[item_id] = call_id + self._chunk_queue.append({ + "type": "content_block_start", + "index": block_idx, + "content_block": { + "type": "tool_use", + "id": call_id, + "name": name, + "input": {}, + }, + }) + elif item_type == "reasoning": + block_idx = self._next_block_index() + if item_id: + self._item_id_to_block_index[item_id] = block_idx + self._chunk_queue.append({ + "type": "content_block_start", + "index": block_idx, + "content_block": {"type": "thinking", "thinking": ""}, + }) + return + + # ---- text delta ---- + if event_type == "response.output_text.delta": + item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) + delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") + block_idx = self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id else self._current_block_index + self._chunk_queue.append({ + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "text_delta", "text": delta}, + }) + return + + # ---- reasoning summary text delta ---- + if event_type == "response.reasoning_summary_text.delta": + item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) + delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") + block_idx = self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id else self._current_block_index + self._chunk_queue.append({ + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "thinking_delta", "thinking": delta}, + }) + return + + # ---- function call arguments delta ---- + if event_type == "response.function_call_arguments.delta": + item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) + delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") + block_idx = self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id else self._current_block_index + self._chunk_queue.append({ + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "input_json_delta", "partial_json": delta}, + }) + return + + # ---- output item done -> content_block_stop ---- + if event_type == "response.output_item.done": + item = getattr(event, "item", None) or (event.get("item") if isinstance(event, dict) else None) + item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None + block_idx = self._item_id_to_block_index.get(item_id, self._current_block_index) if item_id else self._current_block_index + self._chunk_queue.append({ + "type": "content_block_stop", + "index": block_idx, + }) + return + + # ---- response completed -> message_delta + message_stop ---- + if event_type in ("response.completed", "response.failed", "response.incomplete"): + response_obj = getattr(event, "response", None) or (event.get("response") if isinstance(event, dict) else None) + stop_reason = "end_turn" + input_tokens = 0 + output_tokens = 0 + cache_creation_tokens = 0 + cache_read_tokens = 0 + + if response_obj is not None: + status = getattr(response_obj, "status", None) + if status == "incomplete": + stop_reason = "max_tokens" + usage = getattr(response_obj, "usage", None) + if usage is not None: + input_tokens = getattr(usage, "input_tokens", 0) or 0 + output_tokens = getattr(usage, "output_tokens", 0) or 0 + cache_creation_tokens = getattr(usage, "input_tokens_details", None) + cache_read_tokens = getattr(usage, "output_tokens_details", None) + # Prefer direct cache fields if present + cache_creation_tokens = getattr(usage, "cache_creation_input_tokens", 0) or 0 + cache_read_tokens = getattr(usage, "cache_read_input_tokens", 0) or 0 + + # Check if tool_use was in the output to override stop_reason + if response_obj is not None: + output = getattr(response_obj, "output", []) or [] + for out_item in output: + out_type = getattr(out_item, "type", None) or (out_item.get("type") if isinstance(out_item, dict) else None) + if out_type == "function_call": + stop_reason = "tool_use" + break + + usage_delta: Dict[str, Any] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + if cache_creation_tokens: + usage_delta["cache_creation_input_tokens"] = cache_creation_tokens + if cache_read_tokens: + usage_delta["cache_read_input_tokens"] = cache_read_tokens + + self._chunk_queue.append({ + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": usage_delta, + }) + self._chunk_queue.append({"type": "message_stop"}) + self._sent_message_stop = True + return + + def __aiter__(self) -> "AnthropicResponsesStreamWrapper": + return self + + async def __anext__(self) -> Dict[str, Any]: + # Return any queued chunks first + if self._chunk_queue: + return self._chunk_queue.popleft() + + # Emit message_start if not yet done (fallback if response.created wasn't fired) + if not self._sent_message_start: + self._sent_message_start = True + self._chunk_queue.append(self._make_message_start()) + return self._chunk_queue.popleft() + + # Consume the upstream stream + try: + async for event in self.responses_stream: + self._process_event(event) + if self._chunk_queue: + return self._chunk_queue.popleft() + except StopAsyncIteration: + pass + except Exception as e: + verbose_logger.error( + f"AnthropicResponsesStreamWrapper error: {e}\n{traceback.format_exc()}" + ) + + # Drain any remaining queued chunks + if self._chunk_queue: + return self._chunk_queue.popleft() + + raise StopAsyncIteration + + async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]: + """Yield SSE-encoded bytes for each Anthropic event chunk.""" + async for chunk in self: + if isinstance(chunk, dict): + event_type: str = str(chunk.get("type", "message")) + payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n" + yield payload.encode() + else: + yield chunk diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py new file mode 100644 index 00000000000..c2752272905 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -0,0 +1,450 @@ +""" +Transformation layer: Anthropic /v1/messages <-> OpenAI Responses API. + +This module owns all format conversions for the direct v1/messages -> Responses API +path used for OpenAI and Azure models. +""" + +import json +from typing import Any, Dict, List, Optional, Union, cast + +from litellm.types.llms.anthropic import ( + AllAnthropicToolsValues, + AnthopicMessagesAssistantMessageParam, + AnthropicFinishReason, + AnthropicMessagesRequest, + AnthropicMessagesToolChoice, + AnthropicMessagesUserMessageParam, + AnthropicResponseContentBlockText, + AnthropicResponseContentBlockThinking, + AnthropicResponseContentBlockToolUse, +) +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + AnthropicUsage, +) +from litellm.types.llms.openai import ResponsesAPIResponse + + +class LiteLLMAnthropicToResponsesAPIAdapter: + """ + Converts Anthropic /v1/messages requests to OpenAI Responses API format and + converts Responses API responses back to Anthropic format. + """ + + # ------------------------------------------------------------------ # + # Request translation: Anthropic -> Responses API # + # ------------------------------------------------------------------ # + + @staticmethod + def _translate_anthropic_image_source_to_url(source: dict) -> Optional[str]: + """Convert Anthropic image source to a URL string.""" + source_type = source.get("type") + if source_type == "base64": + media_type = source.get("media_type", "image/jpeg") + data = source.get("data", "") + return f"data:{media_type};base64,{data}" if data else None + elif source_type == "url": + return source.get("url") + return None + + def translate_messages_to_responses_input( + self, + messages: List[ + Union[ + AnthropicMessagesUserMessageParam, + AnthopicMessagesAssistantMessageParam, + ] + ], + ) -> List[Dict[str, Any]]: + """ + Convert Anthropic messages list to Responses API `input` items. + + Mapping: + user text -> message(role=user, input_text) + user image -> message(role=user, input_image) + user tool_result -> function_call_output + assistant text -> message(role=assistant, output_text) + assistant tool_use -> function_call + """ + input_items: List[Dict[str, Any]] = [] + + for m in messages: + role = m["role"] + content = m.get("content") + + if role == "user": + if isinstance(content, str): + input_items.append({ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": content}], + }) + elif isinstance(content, list): + user_parts: List[Dict[str, Any]] = [] + for block in content: + if not isinstance(block, dict): + continue + btype = block.get("type") + if btype == "text": + user_parts.append({"type": "input_text", "text": block.get("text", "")}) + elif btype == "image": + url = self._translate_anthropic_image_source_to_url(block.get("source", {})) + if url: + user_parts.append({"type": "input_image", "image_url": url}) + elif btype == "tool_result": + tool_use_id = block.get("tool_use_id", "") + inner = block.get("content") + if inner is None: + output_text = "" + elif isinstance(inner, str): + output_text = inner + elif isinstance(inner, list): + parts = [ + c.get("text", "") + for c in inner + if isinstance(c, dict) and c.get("type") == "text" + ] + output_text = "\n".join(parts) + else: + output_text = str(inner) + # tool_result is a top-level item, not inside the message + input_items.append({ + "type": "function_call_output", + "call_id": tool_use_id, + "output": output_text, + }) + if user_parts: + input_items.append({ + "type": "message", + "role": "user", + "content": user_parts, + }) + + elif role == "assistant": + if isinstance(content, str): + input_items.append({ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": content}], + }) + elif isinstance(content, list): + asst_parts: List[Dict[str, Any]] = [] + for block in content: + if not isinstance(block, dict): + continue + btype = block.get("type") + if btype == "text": + asst_parts.append({"type": "output_text", "text": block.get("text", "")}) + elif btype == "tool_use": + # tool_use becomes a top-level function_call item + input_items.append({ + "type": "function_call", + "call_id": block.get("id", ""), + "name": block.get("name", ""), + "arguments": json.dumps(block.get("input", {})), + }) + elif btype == "thinking": + thinking_text = block.get("thinking", "") + if thinking_text: + asst_parts.append({"type": "output_text", "text": thinking_text}) + if asst_parts: + input_items.append({ + "type": "message", + "role": "assistant", + "content": asst_parts, + }) + + return input_items + + def translate_tools_to_responses_api( + self, + tools: List[AllAnthropicToolsValues], + ) -> List[Dict[str, Any]]: + """Convert Anthropic tool definitions to Responses API function tools.""" + result: List[Dict[str, Any]] = [] + for tool in tools: + tool_dict = cast(Dict[str, Any], tool) + tool_type = tool_dict.get("type", "") + tool_name = tool_dict.get("name", "") + # web_search tool + if (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search": + result.append({"type": "web_search_preview"}) + continue + func_tool: Dict[str, Any] = {"type": "function", "name": tool_name} + if "description" in tool_dict: + func_tool["description"] = tool_dict["description"] + if "input_schema" in tool_dict: + func_tool["parameters"] = tool_dict["input_schema"] + result.append(func_tool) + return result + + @staticmethod + def translate_tool_choice_to_responses_api( + tool_choice: AnthropicMessagesToolChoice, + ) -> Dict[str, Any]: + """Convert Anthropic tool_choice to Responses API tool_choice.""" + tc_type = tool_choice.get("type") + if tc_type == "any": + return {"type": "required"} + elif tc_type == "tool": + return {"type": "function", "name": tool_choice.get("name", "")} + return {"type": "auto"} + + @staticmethod + def translate_context_management_to_responses_api( + context_management: Dict[str, Any], + ) -> Optional[List[Dict[str, Any]]]: + """ + Convert Anthropic context_management dict to OpenAI Responses API array format. + + Anthropic format: {"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]} + OpenAI format: [{"type": "compaction", "compact_threshold": 150000}] + """ + if not isinstance(context_management, dict): + return None + + edits = context_management.get("edits", []) + if not isinstance(edits, list): + return None + + result: List[Dict[str, Any]] = [] + for edit in edits: + if not isinstance(edit, dict): + continue + edit_type = edit.get("type", "") + if edit_type == "compact_20260112": + entry: Dict[str, Any] = {"type": "compaction"} + trigger = edit.get("trigger") + if isinstance(trigger, dict) and trigger.get("value") is not None: + entry["compact_threshold"] = int(trigger["value"]) + result.append(entry) + + return result if result else None + + @staticmethod + def translate_thinking_to_reasoning(thinking: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """ + Convert Anthropic thinking param to Responses API reasoning param. + + thinking.budget_tokens maps to reasoning effort: + >= 10000 -> high, >= 5000 -> medium, >= 2000 -> low, < 2000 -> minimal + """ + if not isinstance(thinking, dict) or thinking.get("type") != "enabled": + return None + budget = thinking.get("budget_tokens", 0) + if budget >= 10000: + effort = "high" + elif budget >= 5000: + effort = "medium" + elif budget >= 2000: + effort = "low" + else: + effort = "minimal" + return {"effort": effort, "summary": "detailed"} + + def translate_request( + self, + anthropic_request: AnthropicMessagesRequest, + ) -> Dict[str, Any]: + """ + Translate a full Anthropic /v1/messages request dict to + litellm.responses() / litellm.aresponses() kwargs. + """ + model: str = anthropic_request["model"] + messages_list = cast( + List[Union[AnthropicMessagesUserMessageParam, AnthopicMessagesAssistantMessageParam]], + anthropic_request["messages"], + ) + + responses_kwargs: Dict[str, Any] = { + "model": model, + "input": self.translate_messages_to_responses_input(messages_list), + } + + # system -> instructions + system = anthropic_request.get("system") + if system: + if isinstance(system, str): + responses_kwargs["instructions"] = system + elif isinstance(system, list): + text_parts = [ + b.get("text", "") + for b in system + if isinstance(b, dict) and b.get("type") == "text" + ] + responses_kwargs["instructions"] = "\n".join(filter(None, text_parts)) + + # max_tokens -> max_output_tokens + max_tokens = anthropic_request.get("max_tokens") + if max_tokens: + responses_kwargs["max_output_tokens"] = max_tokens + + # temperature / top_p passed through + if "temperature" in anthropic_request: + responses_kwargs["temperature"] = anthropic_request["temperature"] + if "top_p" in anthropic_request: + responses_kwargs["top_p"] = anthropic_request["top_p"] + + # tools + tools = anthropic_request.get("tools") + if tools: + responses_kwargs["tools"] = self.translate_tools_to_responses_api( + cast(List[AllAnthropicToolsValues], tools) + ) + + # tool_choice + tool_choice = anthropic_request.get("tool_choice") + if tool_choice: + responses_kwargs["tool_choice"] = self.translate_tool_choice_to_responses_api( + cast(AnthropicMessagesToolChoice, tool_choice) + ) + + # thinking -> reasoning + thinking = anthropic_request.get("thinking") + if isinstance(thinking, dict): + reasoning = self.translate_thinking_to_reasoning(thinking) + if reasoning: + responses_kwargs["reasoning"] = reasoning + + # output_format / output_config.format -> text format + # output_format: {"type": "json_schema", "schema": {...}} + # output_config: {"format": {"type": "json_schema", "schema": {...}}} + output_format = anthropic_request.get("output_format") + output_config = anthropic_request.get("output_config") + if not isinstance(output_format, dict) and isinstance(output_config, dict): + output_format = output_config.get("format") + if isinstance(output_format, dict) and output_format.get("type") == "json_schema": + schema = output_format.get("schema") + if schema: + responses_kwargs["text"] = { + "format": { + "type": "json_schema", + "name": "structured_output", + "schema": schema, + "strict": True, + } + } + + # context_management: Anthropic dict -> OpenAI array + context_management = anthropic_request.get("context_management") + if isinstance(context_management, dict): + openai_cm = self.translate_context_management_to_responses_api(context_management) + if openai_cm is not None: + responses_kwargs["context_management"] = openai_cm + + # metadata user_id -> user + metadata = anthropic_request.get("metadata") + if isinstance(metadata, dict) and "user_id" in metadata: + responses_kwargs["user"] = str(metadata["user_id"])[:64] + + return responses_kwargs + + # ------------------------------------------------------------------ # + # Response translation: Responses API -> Anthropic # + # ------------------------------------------------------------------ # + + def translate_response( + self, + response: ResponsesAPIResponse, + ) -> AnthropicMessagesResponse: + """ + Translate an OpenAI ResponsesAPIResponse to AnthropicMessagesResponse. + """ + from openai.types.responses import ( + ResponseFunctionToolCall, + ResponseOutputMessage, + ResponseReasoningItem, + ) + + from litellm.types.llms.openai import ResponseAPIUsage + + content: List[Dict[str, Any]] = [] + stop_reason: AnthropicFinishReason = "end_turn" + + for item in response.output: + if isinstance(item, ResponseReasoningItem): + for summary in item.summary: + text = getattr(summary, "text", "") + if text: + content.append( + AnthropicResponseContentBlockThinking( + type="thinking", + thinking=text, + signature=None, + ).model_dump() + ) + + elif isinstance(item, ResponseOutputMessage): + for part in item.content: + if getattr(part, "type", None) == "output_text": + content.append( + AnthropicResponseContentBlockText( + type="text", text=getattr(part, "text", "") + ).model_dump() + ) + + elif isinstance(item, ResponseFunctionToolCall): + try: + input_data = json.loads(item.arguments) if item.arguments else {} + except (json.JSONDecodeError, TypeError): + input_data = {} + content.append( + AnthropicResponseContentBlockToolUse( + type="tool_use", + id=item.call_id or item.id, + name=item.name, + input=input_data, + ).model_dump() + ) + stop_reason = "tool_use" + + elif isinstance(item, dict): + item_type = item.get("type") + if item_type == "message": + for part in item.get("content", []): + if isinstance(part, dict) and part.get("type") == "output_text": + content.append( + AnthropicResponseContentBlockText( + type="text", text=part.get("text", "") + ).model_dump() + ) + elif item_type == "function_call": + try: + input_data = json.loads(item.get("arguments", "{}")) + except (json.JSONDecodeError, TypeError): + input_data = {} + content.append( + AnthropicResponseContentBlockToolUse( + type="tool_use", + id=item.get("call_id") or item.get("id", ""), + name=item.get("name", ""), + input=input_data, + ).model_dump() + ) + stop_reason = "tool_use" + + # status -> stop_reason override + if response.status == "incomplete": + stop_reason = "max_tokens" + + # usage + raw_usage: Optional[ResponseAPIUsage] = response.usage + input_tokens = int(getattr(raw_usage, "input_tokens", 0) or 0) + output_tokens = int(getattr(raw_usage, "output_tokens", 0) or 0) + + anthropic_usage = AnthropicUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + + return AnthropicMessagesResponse( + id=response.id, + type="message", + role="assistant", + model=response.model or "unknown-model", + stop_sequence=None, + usage=anthropic_usage, # type: ignore + content=content, # type: ignore + stop_reason=stop_reason, + ) diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index fb13332c464..29929a2bf62 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -15,7 +15,9 @@ else: LiteLLMLoggingObj = Any -# DocumentType for OCR - Mistral format document dict +# DocumentType for OCR - providers always receive a dict with +# type="document_url" or type="image_url" (str values only). +# File-type inputs are preprocessed to this format in litellm/ocr/main.py. DocumentType = Dict[str, str] @@ -141,9 +143,13 @@ class BaseOCRConfig: Transform OCR request to provider-specific format. Override in provider-specific implementations. + Note: By the time this method is called, any file-type documents have already + been converted to document_url/image_url format with base64 data URIs by + the preprocessing in litellm/ocr/main.py. + Args: model: Model name - document: Document to process (Mistral format dict, or file path, bytes, etc.) + document: Document to process - always a dict with type="document_url" or type="image_url" optional_params: Optional parameters for the request headers: Request headers diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 60a93b169c8..ec5b942ec1b 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -68,7 +68,7 @@ def make_sync_call( model_response=model_response, json_mode=json_mode ) else: - decoder = AWSEventStreamDecoder(model=model) + decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) # LOGGING diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index d4fd0606302..d210f294c64 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1217,15 +1217,15 @@ class AmazonConverseConfig(BaseConfig): # Handle parallel_tool_calls configuration parallel_tool_use_config = additional_request_params.pop("_parallel_tool_use_config", None) - if parallel_tool_use_config is not None: - # Merge the tool_choice config from parallel_tool_calls into additional_request_params + if parallel_tool_use_config is not None and is_claude_4_5_on_bedrock(model): for key, value in parallel_tool_use_config.items(): if key in additional_request_params and isinstance(additional_request_params[key], dict) and isinstance(value, dict): - # Merge dictionaries additional_request_params[key].update(value) else: additional_request_params[key] = value + additional_request_params.pop("parallel_tool_calls", None) + # Only set the topK value in for models that support it additional_request_params.update( self._handle_top_k_value(model, inference_params) @@ -1779,6 +1779,92 @@ class AmazonConverseConfig(BaseConfig): return content_str, tools, reasoningContentBlocks, citationsContentBlocks + @staticmethod + def _unwrap_bedrock_properties(json_str: str) -> str: + """ + Unwrap Bedrock's response_format JSON structure. + + If the JSON has a single "properties" key, extract its value. + Otherwise, return the original string. + + Args: + json_str: JSON string to unwrap + + Returns: + Unwrapped JSON string or original if unwrapping not needed + """ + try: + response_data = json.loads(json_str) + if ( + isinstance(response_data, dict) + and "properties" in response_data + and len(response_data) == 1 + ): + response_data = response_data["properties"] + return json.dumps(response_data) + except json.JSONDecodeError: + pass + return json_str + + @staticmethod + def _filter_json_mode_tools( + json_mode: Optional[bool], + tools: List[ChatCompletionToolCallChunk], + chat_completion_message: ChatCompletionResponseMessage, + ) -> Optional[List[ChatCompletionToolCallChunk]]: + """ + When json_mode is True, Bedrock may return the internal `json_tool_call` + tool alongside real user-defined tools. This method handles 3 scenarios: + + 1. Only json_tool_call present -> convert to text content, return None + 2. Mixed json_tool_call + real -> filter out json_tool_call, return real tools + 3. No json_tool_call / no json_mode -> return tools as-is + """ + if not json_mode or not tools: + return tools if tools else None + + json_tool_indices = [ + i + for i, t in enumerate(tools) + if t["function"].get("name") == RESPONSE_FORMAT_TOOL_NAME + ] + + if not json_tool_indices: + # No json_tool_call found, return tools unchanged + return tools + + if len(json_tool_indices) == len(tools): + # All tools are json_tool_call — convert first one to content + verbose_logger.debug( + "Processing JSON tool call response for response_format" + ) + json_mode_content_str: Optional[str] = tools[0]["function"].get( + "arguments" + ) + if json_mode_content_str is not None: + json_mode_content_str = AmazonConverseConfig._unwrap_bedrock_properties( + json_mode_content_str + ) + chat_completion_message["content"] = json_mode_content_str + return None + + # Mixed: filter out json_tool_call, keep real tools. + # Preserve the json_tool_call content as message text so the structured + # output from response_format is not silently lost. + first_idx = json_tool_indices[0] + json_mode_args = tools[first_idx]["function"].get("arguments") + if json_mode_args is not None: + json_mode_args = AmazonConverseConfig._unwrap_bedrock_properties( + json_mode_args + ) + existing = chat_completion_message.get("content") or "" + chat_completion_message["content"] = ( + existing + json_mode_args if existing else json_mode_args + ) + + real_tools = [t for i, t in enumerate(tools) if i not in json_tool_indices] + return real_tools if real_tools else None + def _transform_response( # noqa: PLR0915 self, model: str, @@ -1801,7 +1887,7 @@ class AmazonConverseConfig(BaseConfig): additional_args={"complete_input_dict": data}, ) - json_mode: Optional[bool] = optional_params.pop("json_mode", None) + json_mode: Optional[bool] = optional_params.get("json_mode", None) ## RESPONSE OBJECT try: completion_response = ConverseResponseBlock(**response.json()) # type: ignore @@ -1885,37 +1971,13 @@ class AmazonConverseConfig(BaseConfig): self._transform_thinking_blocks(reasoningContentBlocks) ) chat_completion_message["content"] = content_str - if ( - json_mode is True - and tools is not None - and len(tools) == 1 - and tools[0]["function"].get("name") == RESPONSE_FORMAT_TOOL_NAME - ): - verbose_logger.debug( - "Processing JSON tool call response for response_format" - ) - json_mode_content_str: Optional[str] = tools[0]["function"].get("arguments") - if json_mode_content_str is not None: - # Bedrock returns the response wrapped in a "properties" object - # We need to extract the actual content from this wrapper - try: - response_data = json.loads(json_mode_content_str) - - # If Bedrock wrapped the response in "properties", extract the content - if ( - isinstance(response_data, dict) - and "properties" in response_data - and len(response_data) == 1 - ): - response_data = response_data["properties"] - json_mode_content_str = json.dumps(response_data) - except json.JSONDecodeError: - # If parsing fails, use the original response - pass - - chat_completion_message["content"] = json_mode_content_str - elif tools: - chat_completion_message["tool_calls"] = tools + filtered_tools = self._filter_json_mode_tools( + json_mode=json_mode, + tools=tools, + chat_completion_message=chat_completion_message, + ) + if filtered_tools: + chat_completion_message["tool_calls"] = filtered_tools ## CALCULATING USAGE - bedrock returns usage in the headers usage = self._transform_usage(completion_response["usage"]) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 1c58a11eebe..88f7341ed08 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -22,6 +22,7 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.caching.caching import InMemoryCache +from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.logging_utils import track_llm_api_timing @@ -252,7 +253,7 @@ async def make_call( response.aiter_bytes(chunk_size=stream_chunk_size) ) else: - decoder = AWSEventStreamDecoder(model=model) + decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) completion_stream = decoder.aiter_bytes( response.aiter_bytes(chunk_size=stream_chunk_size) ) @@ -346,7 +347,7 @@ def make_sync_call( response.iter_bytes(chunk_size=stream_chunk_size) ) else: - decoder = AWSEventStreamDecoder(model=model) + decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) completion_stream = decoder.iter_bytes( response.iter_bytes(chunk_size=stream_chunk_size) ) @@ -1282,7 +1283,7 @@ def get_response_stream_shape(): class AWSEventStreamDecoder: - def __init__(self, model: str) -> None: + def __init__(self, model: str, json_mode: Optional[bool] = False) -> None: from botocore.parsers import EventStreamJSONParser self.model = model @@ -1290,6 +1291,8 @@ class AWSEventStreamDecoder: self.content_blocks: List[ContentBlockDeltaEvent] = [] self.tool_calls_index: Optional[int] = None self.response_id: Optional[str] = None + self.json_mode = json_mode + self._current_tool_name: Optional[str] = None def check_empty_tool_call_args(self) -> bool: """ @@ -1391,6 +1394,16 @@ class AWSEventStreamDecoder: response_tool_name = get_bedrock_tool_name( response_tool_name=_response_tool_name ) + self._current_tool_name = response_tool_name + + # When json_mode is True, suppress the internal json_tool_call + # and convert its content to text in delta events instead + if ( + self.json_mode is True + and response_tool_name == RESPONSE_FORMAT_TOOL_NAME + ): + return tool_use, provider_specific_fields, thinking_blocks + self.tool_calls_index = ( 0 if self.tool_calls_index is None else self.tool_calls_index + 1 ) @@ -1445,19 +1458,27 @@ class AWSEventStreamDecoder: if "text" in delta_obj: text = delta_obj["text"] elif "toolUse" in delta_obj: - tool_use = { - "id": None, - "type": "function", - "function": { - "name": None, - "arguments": delta_obj["toolUse"]["input"], - }, - "index": ( - self.tool_calls_index - if self.tool_calls_index is not None - else index - ), - } + # When json_mode is True and this is the internal json_tool_call, + # convert tool input to text content instead of tool call arguments + if ( + self.json_mode is True + and self._current_tool_name == RESPONSE_FORMAT_TOOL_NAME + ): + text = delta_obj["toolUse"]["input"] + else: + tool_use = { + "id": None, + "type": "function", + "function": { + "name": None, + "arguments": delta_obj["toolUse"]["input"], + }, + "index": ( + self.tool_calls_index + if self.tool_calls_index is not None + else index + ), + } elif "reasoningContent" in delta_obj: provider_specific_fields = { "reasoningContent": delta_obj["reasoningContent"], @@ -1494,6 +1515,17 @@ class AWSEventStreamDecoder: ) -> Optional[ChatCompletionToolCallChunk]: """Handle stop/contentBlockIndex event in converse chunk parsing.""" tool_use: Optional[ChatCompletionToolCallChunk] = None + + # If the ending block was the internal json_tool_call, skip emitting + # the empty-args tool chunk and reset tracking state + if ( + self.json_mode is True + and self._current_tool_name == RESPONSE_FORMAT_TOOL_NAME + ): + self._current_tool_name = None + return tool_use + + self._current_tool_name = None is_empty = self.check_empty_tool_call_args() if is_empty: tool_use = { diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b09a36be60f..d6fdc58099f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,4 +1,5 @@ import json +import ssl from typing import ( TYPE_CHECKING, Any, @@ -4659,6 +4660,8 @@ class BaseLLMHTTPHandler: api_key: Optional[str] = None, client: Optional[Any] = None, timeout: Optional[float] = None, + user_api_key_dict: Optional[Any] = None, + litellm_metadata: Optional[Dict[str, Any]] = None, ): import websockets from websockets.asyncio.client import ClientConnection @@ -4672,6 +4675,11 @@ class BaseLLMHTTPHandler: try: ssl_context = get_shared_realtime_ssl_context() + if url.startswith("wss://") and ssl_context is False: + # Keep TLS for wss:// while honoring SSL_VERIFY=False semantics. + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE async with websockets.connect( # type: ignore url, additional_headers=headers, @@ -4686,12 +4694,17 @@ class BaseLLMHTTPHandler: if _session_config: await backend_ws.send(_session_config) + _request_data: Dict[str, Any] = {} + if litellm_metadata: + _request_data["litellm_metadata"] = litellm_metadata realtime_streaming = RealTimeStreaming( websocket, cast(ClientConnection, backend_ws), logging_obj, provider_config, model, + user_api_key_dict=user_api_key_dict, + request_data=_request_data, ) if _session_config: realtime_streaming.session_configuration_request = _session_config diff --git a/litellm/llms/gemini/image_generation/cost_calculator.py b/litellm/llms/gemini/image_generation/cost_calculator.py index 0a9ca2e5276..941ab0d50f7 100644 --- a/litellm/llms/gemini/image_generation/cost_calculator.py +++ b/litellm/llms/gemini/image_generation/cost_calculator.py @@ -5,6 +5,9 @@ Google AI Image Generation Cost Calculator from typing import Any import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_image_response_cost_from_usage, +) from litellm.types.utils import ImageResponse @@ -13,13 +16,22 @@ def cost_calculator( image_response: Any, ) -> float: """ - Vertex AI Image Generation Cost Calculator + Google AI Image Generation Cost Calculator """ _model_info = litellm.get_model_info( model=model, custom_llm_provider="gemini", ) + if isinstance(image_response, ImageResponse): + token_based_cost = calculate_image_response_cost_from_usage( + model=model, + image_response=image_response, + custom_llm_provider="gemini", + ) + if token_based_cost is not None: + return token_based_cost + output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 num_images: int = 0 if isinstance(image_response, ImageResponse): diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 2e0e678e69f..d9465c95e3b 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -867,6 +867,52 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) returned_message: List[OpenAIRealtimeEvents] = [] + # Handle transcription events that arrive independently from model + # content. Gemini sends inputTranscription / outputTranscription + # inside serverContent, separately from modelTurn / turnComplete. + server_content = json_message.get("serverContent") + if isinstance(server_content, dict): + input_tx = server_content.get("inputTranscription") + if isinstance(input_tx, dict) and input_tx.get("text"): + returned_message.append( + cast(OpenAIRealtimeEvents, { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_{}".format(uuid.uuid4()), + "transcript": input_tx["text"], + "item_id": "item_{}".format(uuid.uuid4()), + "content_index": 0, + }) + ) + + output_tx = server_content.get("outputTranscription") + if isinstance(output_tx, dict) and output_tx.get("text"): + returned_message.append( + cast(OpenAIRealtimeEvents, { + "type": "response.audio_transcript.delta", + "event_id": "event_{}".format(uuid.uuid4()), + "delta": output_tx["text"], + "item_id": current_output_item_id or "item_{}".format(uuid.uuid4()), + "response_id": current_response_id or "resp_{}".format(uuid.uuid4()), + "output_index": 0, + "content_index": 0, + }) + ) + + # If serverContent only contained transcription(s) and no model + # content, return early — the main loop would fail on unknown keys. + _model_content_keys = {"modelTurn", "turnComplete", "interrupted", "generationComplete"} + if not any(k in server_content for k in _model_content_keys): + return { + "response": returned_message, + "current_output_item_id": current_output_item_id, + "current_response_id": current_response_id, + "current_delta_chunks": current_delta_chunks, + "current_conversation_id": current_conversation_id, + "current_item_chunks": current_item_chunks, + "current_delta_type": current_delta_type, + "session_configuration_request": session_configuration_request, + } + for key, value in json_message.items(): # Check if this key or any nested key matches our mapping openai_event = self.map_openai_event( @@ -974,6 +1020,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): setup_config: BidiGenerateContentSetup = { "model": f"models/{model}", "generationConfig": {"responseModalities": response_modalities}, + # Return input transcript so guardrails can inspect user speech. + "inputAudioTranscription": {}, } if output_audio_transcription: setup_config["outputAudioTranscription"] = {} diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index c7524925bd0..7020f796bb7 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1401,6 +1401,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=None, max_retries=None, organization: Optional[str] = None, + headers: Optional[dict] = None, ): response = None try: @@ -1414,6 +1415,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=client, ) + if headers: + data["extra_headers"] = headers response = await openai_aclient.images.generate(**data, timeout=timeout) # type: ignore stringified_response = response.model_dump() ## LOGGING @@ -1446,6 +1449,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=None, aimg_generation=None, organization: Optional[str] = None, + headers: Optional[dict] = None, ) -> ImageResponse: data = {} try: @@ -1455,7 +1459,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError(status_code=422, message="max retries must be an int") if aimg_generation is True: - return self.aimage_generation(data=data, prompt=prompt, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries, organization=organization) # type: ignore + return self.aimage_generation(data=data, prompt=prompt, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries, organization=organization, headers=headers) # type: ignore openai_client: OpenAI = self._get_openai_client( # type: ignore is_async=False, @@ -1480,6 +1484,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) ## COMPLETION CALL + if headers: + data["extra_headers"] = headers _response = openai_client.images.generate(**data, timeout=timeout) # type: ignore response = _response.model_dump() diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 1b1b1c2f8cc..b3125d4ad38 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -90,5 +90,9 @@ "headers": { "api-subscription-key": "{api_key}" } + }, + "assemblyai": { + "base_url": "https://llm-gateway.assemblyai.com/v1", + "api_key_env": "ASSEMBLYAI_API_KEY" } } diff --git a/litellm/llms/vertex_ai/image_generation/cost_calculator.py b/litellm/llms/vertex_ai/image_generation/cost_calculator.py index 646c6080a2e..012de5498cb 100644 --- a/litellm/llms/vertex_ai/image_generation/cost_calculator.py +++ b/litellm/llms/vertex_ai/image_generation/cost_calculator.py @@ -3,6 +3,9 @@ Vertex AI Image Generation Cost Calculator """ import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_image_response_cost_from_usage, +) from litellm.types.utils import ImageResponse @@ -18,6 +21,14 @@ def cost_calculator( custom_llm_provider="vertex_ai", ) + token_based_cost = calculate_image_response_cost_from_usage( + model=model, + image_response=image_response, + custom_llm_provider="vertex_ai", + ) + if token_based_cost is not None: + return token_based_cost + output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 num_images: int = 0 if image_response.data: diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index eaa9844f108..5eae143175b 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -124,6 +124,8 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): "silenceDurationMs": 800, } }, + # Return input transcript so guardrails can inspect user speech. + "inputAudioTranscription": {}, # Return output transcript so clients can read what the model said. "outputAudioTranscription": {}, } diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 57563fc0bcc..f52288ea72a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14194,6 +14194,38 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-3.1-flash-image-preview": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true + }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -19178,6 +19210,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "gpt-audio-1.5": { + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "gpt-audio-2025-08-28": { "input_cost_per_audio_token": 3.2e-05, "input_cost_per_token": 2.5e-06, @@ -20895,6 +20960,38 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "gpt-realtime-1.5": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "openai", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, @@ -25060,6 +25157,25 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "openrouter/anthropic/claude-opus-4.6": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, "cache_creation_input_token_cost": 3.75e-06, @@ -26072,6 +26188,42 @@ "supports_prompt_caching": true, "supports_computer_use": false }, + "openrouter/openrouter/auto": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "max_input_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true + }, + "openrouter/openrouter/free": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "max_input_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, + "openrouter/openrouter/bodybuilder": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "chat" + }, "ovhcloud/DeepSeek-R1-Distill-Llama-70B": { "input_cost_per_token": 6.7e-07, "litellm_provider": "ovhcloud", @@ -26618,8 +26770,8 @@ "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", - "supports_function_calling": true, - "supports_tool_choice": true + "supports_function_calling": false, + "supports_tool_choice": false }, "publicai/swiss-ai/apertus-70b-instruct": { "input_cost_per_token": 0.0, @@ -26630,8 +26782,8 @@ "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", - "supports_function_calling": true, - "supports_tool_choice": true + "supports_function_calling": false, + "supports_tool_choice": false }, "publicai/aisingapore/Gemma-SEA-LION-v4-27B-IT": { "input_cost_per_token": 0.0, @@ -31545,6 +31697,19 @@ "output_cost_per_token_batches": 6e-06, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, + "vertex_ai/gemini-3.1-flash-image-preview": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" + }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -32946,6 +33111,7 @@ "supports_web_search": true }, "xai/grok-2-vision-1212": { + "deprecation_date": "2026-02-28", "input_cost_per_image": 2e-06, "input_cost_per_token": 2e-06, "litellm_provider": "xai", @@ -33050,6 +33216,7 @@ }, "xai/grok-3-mini": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2026-02-28", "input_cost_per_token": 3e-07, "litellm_provider": "xai", "max_input_tokens": 131072, @@ -33066,6 +33233,7 @@ }, "xai/grok-3-mini-beta": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2026-02-28", "input_cost_per_token": 3e-07, "litellm_provider": "xai", "max_input_tokens": 131072, diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 5acab8cbf2c..47cff8a2c0c 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -2,8 +2,14 @@ Main OCR function for LiteLLM. """ import asyncio +import base64 import contextvars +import mimetypes +import os +import re from functools import partial +from io import IOBase +from pathlib import Path from typing import Any, Coroutine, Dict, Optional, Union import httpx @@ -25,7 +31,7 @@ base_llm_http_handler = BaseLLMHTTPHandler() @client async def aocr( model: str, - document: Dict[str, str], + document: Dict[str, Any], api_key: Optional[str] = None, api_base: Optional[str] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, @@ -35,26 +41,27 @@ async def aocr( ) -> OCRResponse: """ Async OCR function. - + Args: model: Model name (e.g., "mistral/mistral-ocr-latest") document: Document to process in Mistral format: - {"type": "document_url", "document_url": "https://..."} for PDFs/docs or - {"type": "image_url", "image_url": "https://..."} for images + {"type": "document_url", "document_url": "https://..."} for PDFs/docs, + {"type": "image_url", "image_url": "https://..."} for images, or + {"type": "file", "file": } for local files api_key: Optional API key api_base: Optional API base URL timeout: Optional timeout custom_llm_provider: Optional custom LLM provider extra_headers: Optional extra headers **kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit) - + Returns: OCRResponse in Mistral OCR format with pages, model, usage_info, etc. - + Example: ```python import litellm - + # OCR with PDF response = await litellm.aocr( model="mistral/mistral-ocr-latest", @@ -64,7 +71,7 @@ async def aocr( }, include_image_base64=True ) - + # OCR with image response = await litellm.aocr( model="mistral/mistral-ocr-latest", @@ -73,7 +80,7 @@ async def aocr( "image_url": "https://example.com/image.png" } ) - + # OCR with base64 encoded PDF response = await litellm.aocr( model="mistral/mistral-ocr-latest", @@ -82,6 +89,12 @@ async def aocr( "document_url": f"data:application/pdf;base64,{base64_pdf}" } ) + + # OCR with local file + response = await litellm.aocr( + model="mistral/mistral-ocr-latest", + document={"type": "file", "file": "/path/to/document.pdf"} + ) ``` """ local_vars = locals() @@ -135,7 +148,7 @@ async def aocr( @client def ocr( model: str, - document: Dict[str, str], + document: Dict[str, Any], api_key: Optional[str] = None, api_base: Optional[str] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, @@ -145,26 +158,27 @@ def ocr( ) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: """ Synchronous OCR function. - + Args: model: Model name (e.g., "mistral/mistral-ocr-latest") document: Document to process in Mistral format: - {"type": "document_url", "document_url": "https://..."} for PDFs/docs or - {"type": "image_url", "image_url": "https://..."} for images + {"type": "document_url", "document_url": "https://..."} for PDFs/docs, + {"type": "image_url", "image_url": "https://..."} for images, or + {"type": "file", "file": } for local files api_key: Optional API key api_base: Optional API base URL timeout: Optional timeout custom_llm_provider: Optional custom LLM provider extra_headers: Optional extra headers **kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit) - + Returns: OCRResponse in Mistral OCR format with pages, model, usage_info, etc. - + Example: ```python import litellm - + # OCR with PDF response = litellm.ocr( model="mistral/mistral-ocr-latest", @@ -174,7 +188,7 @@ def ocr( }, include_image_base64=True ) - + # OCR with image response = litellm.ocr( model="mistral/mistral-ocr-latest", @@ -183,7 +197,7 @@ def ocr( "image_url": "https://example.com/image.png" } ) - + # OCR with base64 encoded PDF response = litellm.ocr( model="mistral/mistral-ocr-latest", @@ -192,7 +206,13 @@ def ocr( "document_url": f"data:application/pdf;base64,{base64_pdf}" } ) - + + # OCR with local file + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={"type": "file", "file": "/path/to/document.pdf"} + ) + # Access pages for page in response.pages: print(f"Page {page.index}: {page.markdown}") @@ -203,24 +223,38 @@ def ocr( litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("aocr", False) is True - - # Validate document parameter format (Mistral spec) - if not isinstance(document, dict): - raise ValueError(f"document must be a dict with 'type' and URL field, got {type(document)}") - - doc_type = document.get("type") - if doc_type not in ["document_url", "image_url"]: - raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url' or 'image_url'") - model, custom_llm_provider, dynamic_api_key, dynamic_api_base = ( - litellm.get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - api_key=api_key, + # Validate document parameter format + if not isinstance(document, dict): + raise ValueError( + f"document must be a dict with 'type' and URL/file field, got {type(document)}" ) + + doc_type = document.get("type") + + # Handle file type: convert to document_url/image_url with base64 data URI + if doc_type == "file": + document = convert_file_document_to_url_document(document) + doc_type = document.get("type") + + if doc_type not in ["document_url", "image_url"]: + raise ValueError( + f"Invalid document type: {doc_type}. " + "Must be 'document_url', 'image_url', or 'file'" + ) + + ( + model, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, ) - + # Update with dynamic values if available if dynamic_api_key: api_key = dynamic_api_key @@ -228,11 +262,11 @@ def ocr( api_base = dynamic_api_base # Get provider config - ocr_provider_config: Optional[BaseOCRConfig] = ( - ProviderConfigManager.get_provider_ocr_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + ocr_provider_config: Optional[ + BaseOCRConfig + ] = ProviderConfigManager.get_provider_ocr_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) if ocr_provider_config is None: @@ -246,21 +280,21 @@ def ocr( # Get litellm params using GenericLiteLLMParams (same as responses API) litellm_params = GenericLiteLLMParams(**kwargs) - + # Extract OCR-specific parameters from kwargs supported_params = ocr_provider_config.get_supported_ocr_params(model=model) non_default_params = {} for param in supported_params: if param in kwargs: non_default_params[param] = kwargs.pop(param) - + # Map parameters to provider-specific format optional_params = ocr_provider_config.map_ocr_params( non_default_params=non_default_params, optional_params={}, model=model, ) - + verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}") # Pre Call logging @@ -300,3 +334,111 @@ def ocr( extra_kwargs=kwargs, ) + +################################################# +# Public utilities — used by the SDK and the proxy +################################################# + +_MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$") + +_MIME_TYPE_MAP = { + ".pdf": "application/pdf", + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + ".bmp": "image/bmp", +} + + +def get_mime_type(file_path: str) -> str: + """ + Determine MIME type from file path extension. + + Falls back to mimetypes.guess_type, then to 'application/octet-stream'. + """ + ext = os.path.splitext(file_path)[1].lower() + mime = _MIME_TYPE_MAP.get(ext) + if mime: + return mime + guessed, _ = mimetypes.guess_type(file_path) + return guessed or "application/octet-stream" + + +def convert_file_document_to_url_document(document: Dict[str, Any]) -> Dict[str, str]: + """ + Convert a file-type document dict to a document_url-type document dict + with an inline base64 data URI. + + Accepts document dicts like: + {"type": "file", "file": "/path/to/document.pdf"} # file path string + {"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path + {"type": "file", "file": } # file-like object (BinaryIO) + {"type": "file", "file": b"raw bytes"} # raw bytes + + Returns: + {"type": "document_url", "document_url": "data:;base64,"} + or {"type": "image_url", "image_url": "data:;base64,"} + """ + file_input = document.get("file") + if file_input is None: + raise ValueError( + "document with type='file' must include a 'file' field containing " + "a file path (str), pathlib.Path, file-like object, or bytes" + ) + + file_bytes: bytes + mime_type: str = "application/octet-stream" + file_name: Optional[str] = None + + if isinstance(file_input, (str, Path)): + file_path = str(file_input) + if not os.path.isfile(file_path): + raise FileNotFoundError(f"File not found: {file_path}") + mime_type = get_mime_type(file_path) + file_name = os.path.basename(file_path) + with open(file_path, "rb") as f: + file_bytes = f.read() + elif isinstance(file_input, bytes): + file_bytes = file_input + elif isinstance(file_input, IOBase) or hasattr(file_input, "read"): + if hasattr(file_input, "name"): + file_name = getattr(file_input, "name", None) + if file_name: + mime_type = get_mime_type(file_name) + file_bytes = file_input.read() + if isinstance(file_bytes, str): + file_bytes = file_bytes.encode("utf-8") + else: + raise ValueError( + f"Unsupported file input type: {type(file_input)}. " + "Expected str (file path), pathlib.Path, bytes, or a file-like object." + ) + + if not file_bytes: + raise ValueError("File is empty or could not be read") + + if "mime_type" in document: + mime_type = document["mime_type"] + + if not _MIME_PATTERN.match(mime_type): + raise ValueError(f"Invalid MIME type: {mime_type}") + + base64_data = base64.b64encode(file_bytes).decode("utf-8") + data_uri = f"data:{mime_type};base64,{base64_data}" + + if mime_type.startswith("image/"): + verbose_logger.debug( + f"OCR file input: Converted file to image_url data URI " + f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" + ) + return {"type": "image_url", "image_url": data_uri} + else: + verbose_logger.debug( + f"OCR file input: Converted file to document_url data URI " + f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" + ) + return {"type": "document_url", "document_url": data_uri} diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json new file mode 100644 index 00000000000..fc79ba54759 --- /dev/null +++ b/litellm/provider_endpoints_support_backup.json @@ -0,0 +1,2748 @@ +{ + "_comment": "This file defines which endpoints are supported by each LiteLLM provider", + "_schema": { + "provider_slug": { + "display_name": "Display name shown in README (e.g., 'OpenAI (`openai`)')", + "url": "Link to provider documentation", + "endpoints": { + "chat_completions": "Supports /chat/completions endpoint", + "messages": "Supports /messages endpoint (Anthropic format)", + "responses": "Supports /responses endpoint (OpenAI/Anthropic unified)", + "embeddings": "Supports /embeddings endpoint", + "image_generations": "Supports /image/generations endpoint", + "audio_transcriptions": "Supports /audio/transcriptions endpoint", + "audio_speech": "Supports /audio/speech endpoint", + "moderations": "Supports /moderations endpoint", + "batches": "Supports /batches endpoint", + "rerank": "Supports /rerank endpoint", + "ocr": "Supports /ocr endpoint", + "search": "Supports /search endpoint", + "skills": "Supports /skills endpoint", + "interactions": "Supports /interactions endpoint (Google AI Interactions API)", + "a2a": "Supports /a2a/{agent}/message/send endpoint (A2A Protocol)", + "container": "Supports OpenAI's /containers endpoint", + "container_files": "Supports OpenAI's /containers/{id}/files endpoint", + "compact": "Supports /responses/compact endpoint", + "files": "Supports /files endpoint for file operations", + "image_edits": "Supports /images/edits endpoint for image editing", + "vector_stores_create": "Supports creating a new vector store via /vector_stores endpoint", + "vector_stores_search": "Supports searching a vector store via /vector_stores/{id}/search endpoint", + "video_generations": "Supports /videos/generations endpoint for video generation" + } + } + }, + "providers": { + "a2a": { + "display_name": "A2A (Agent-to-Agent) (`a2a`)", + "url": "https://docs.litellm.ai/docs/providers/a2a", + "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 + } + }, + "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", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "ai21": { + "display_name": "AI21 (`ai21`)", + "url": "https://docs.litellm.ai/docs/providers/ai21", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "ai21_chat": { + "display_name": "AI21 Chat (`ai21_chat`)", + "url": "https://docs.litellm.ai/docs/providers/ai21", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "amazon_nova": { + "display_name": "Amazon Nova (`amazon_nova`)", + "url": "https://docs.litellm.ai/docs/providers/amazon_nova", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "anthropic": { + "display_name": "Anthropic (`anthropic`)", + "url": "https://docs.litellm.ai/docs/providers/anthropic", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": true, + "rerank": false, + "skills": true, + "a2a": true, + "interactions": true, + "count_tokens": true + } + }, + "anthropic_text": { + "display_name": "Anthropic Text (`anthropic_text`)", + "url": "https://docs.litellm.ai/docs/providers/anthropic", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": true, + "rerank": false, + "skills": true, + "a2a": true, + "interactions": true + } + }, + "apertis": { + "display_name": "Apertis (`apertis`)", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, + "assemblyai": { + "display_name": "AssemblyAI (`assemblyai`)", + "url": "https://docs.litellm.ai/docs/pass_through/assembly_ai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": true, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "auto_router": { + "display_name": "Auto Router (`auto_router`)", + "url": "https://docs.litellm.ai/docs/proxy/auto_routing", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "bedrock": { + "display_name": "AWS - Bedrock (`bedrock`)", + "url": "https://docs.litellm.ai/docs/providers/bedrock", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true, + "a2a": true, + "interactions": true, + "bedrock_invoke": true, + "bedrock_converse": true, + "vector_stores_search": true, + "count_tokens": true, + "rag_ingest": true, + "rag_query": true + } + }, + "s3_vectors": { + "display_name": "AWS S3 Vectors (`s3_vectors`)", + "url": "https://docs.litellm.ai/docs/providers/s3_vectors", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false, + "vector_stores_create": true, + "vector_stores_search": true + } + }, + "sagemaker": { + "display_name": "AWS - Sagemaker (`sagemaker`)", + "url": "https://docs.litellm.ai/docs/providers/aws_sagemaker", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "aws_polly": { + "display_name": "AWS - Polly (`aws_polly`)", + "url": "https://docs.litellm.ai/docs/providers/aws_polly", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": true, + "moderations": false, + "batches": false, + "rerank": false + } + }, + "azure": { + "display_name": "Azure (`azure`)", + "url": "https://docs.litellm.ai/docs/providers/azure", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": true, + "rerank": false, + "a2a": true, + "interactions": true, + "vector_stores_search": true, + "assistants": true, + "fine_tuning": true, + "text_completion": true + } + }, + "azure_ai": { + "display_name": "Azure AI (`azure_ai`)", + "url": "https://docs.litellm.ai/docs/providers/azure_ai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "image_edits": true, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": true, + "rerank": false, + "ocr": true, + "a2a": true, + "interactions": true, + "vector_stores_create": true, + "vector_stores_search": true + } + }, + "azure_ai/doc-intelligence": { + "display_name": "Azure AI Document Intelligence (`azure_ai/doc-intelligence`)", + "url": "https://docs.litellm.ai/docs/providers/azure_document_intelligence", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "ocr": true + } + }, + "azure_ai/agents": { + "display_name": "Azure AI Foundry Agents (`azure_ai/agents`)", + "url": "https://docs.litellm.ai/docs/providers/azure_ai_agents", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "azure_text": { + "display_name": "Azure Text (`azure_text`)", + "url": "https://docs.litellm.ai/docs/providers/azure", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": true, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "baseten": { + "display_name": "Baseten (`baseten`)", + "url": "https://docs.litellm.ai/docs/providers/baseten", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "bytez": { + "display_name": "Bytez (`bytez`)", + "url": "https://docs.litellm.ai/docs/providers/bytez", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "cerebras": { + "display_name": "Cerebras (`cerebras`)", + "url": "https://docs.litellm.ai/docs/providers/cerebras", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "chutes": { + "display_name": "Chutes (`chutes`)", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, + "clarifai": { + "display_name": "Clarifai (`clarifai`)", + "url": "https://docs.litellm.ai/docs/providers/clarifai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "cloudflare": { + "display_name": "Cloudflare AI Workers (`cloudflare`)", + "url": "https://docs.litellm.ai/docs/providers/cloudflare_workers", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "codestral": { + "display_name": "Codestral (`codestral`)", + "url": "https://docs.litellm.ai/docs/providers/codestral", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "cohere": { + "display_name": "Cohere (`cohere`)", + "url": "https://docs.litellm.ai/docs/providers/cohere", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true, + "a2a": true, + "interactions": true + } + }, + "cohere_chat": { + "display_name": "Cohere Chat (`cohere_chat`)", + "url": "https://docs.litellm.ai/docs/providers/cohere", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "cometapi": { + "display_name": "CometAPI (`cometapi`)", + "url": "https://docs.litellm.ai/docs/providers/cometapi", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "compactifai": { + "display_name": "CompactifAI (`compactifai`)", + "url": "https://docs.litellm.ai/docs/providers/compactifai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "custom": { + "display_name": "Custom (`custom`)", + "url": "https://docs.litellm.ai/docs/providers/custom_llm_server", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "custom_openai": { + "display_name": "Custom OpenAI (`custom_openai`)", + "url": "https://docs.litellm.ai/docs/providers/openai_compatible", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": true, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "dashscope": { + "display_name": "Dashscope (`dashscope`)", + "url": "https://docs.litellm.ai/docs/providers/dashscope", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "databricks": { + "display_name": "Databricks (`databricks`)", + "url": "https://docs.litellm.ai/docs/providers/databricks", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "dataforseo": { + "display_name": "DataForSEO (`dataforseo`)", + "url": "https://docs.litellm.ai/docs/search/dataforseo", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, + "datarobot": { + "display_name": "DataRobot (`datarobot`)", + "url": "https://docs.litellm.ai/docs/providers/datarobot", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "deepgram": { + "display_name": "Deepgram (`deepgram`)", + "url": "https://docs.litellm.ai/docs/providers/deepgram", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": true, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "deepinfra": { + "display_name": "DeepInfra (`deepinfra`)", + "url": "https://docs.litellm.ai/docs/providers/deepinfra", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "deepseek": { + "display_name": "Deepseek (`deepseek`)", + "url": "https://docs.litellm.ai/docs/providers/deepseek", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "duckduckgo": { + "display_name": "DuckDuckGo (`duckduckgo`)", + "url": "https://docs.litellm.ai/docs/search/duckduckgo", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, + "elevenlabs": { + "display_name": "ElevenLabs (`elevenlabs`)", + "url": "https://docs.litellm.ai/docs/providers/elevenlabs", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": true, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "exa_ai": { + "display_name": "Exa AI (`exa_ai`)", + "url": "https://docs.litellm.ai/docs/search/exa_ai", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, + "brave": { + "display_name": "Brave Search (`brave`)", + "url": "https://docs.litellm.ai/docs/search/brave", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, + "empower": { + "display_name": "Empower (`empower`)", + "url": "https://docs.litellm.ai/docs/providers/empower", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "fal_ai": { + "display_name": "Fal AI (`fal_ai`)", + "url": "https://docs.litellm.ai/docs/providers/fal_ai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "featherless_ai": { + "display_name": "Featherless AI (`featherless_ai`)", + "url": "https://docs.litellm.ai/docs/providers/featherless_ai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "fireworks_ai": { + "display_name": "Fireworks AI (`fireworks_ai`)", + "url": "https://docs.litellm.ai/docs/providers/fireworks_ai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true, + "a2a": true, + "interactions": true + } + }, + "firecrawl": { + "display_name": "Firecrawl (`firecrawl`)", + "url": "https://docs.litellm.ai/docs/search/firecrawl", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, + "linkup": { + "display_name": "Linkup (`linkup`)", + "url": "https://docs.litellm.ai/docs/search/linkup", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, + "friendliai": { + "display_name": "FriendliAI (`friendliai`)", + "url": "https://docs.litellm.ai/docs/providers/friendliai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "galadriel": { + "display_name": "Galadriel (`galadriel`)", + "url": "https://docs.litellm.ai/docs/providers/galadriel", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "github_copilot": { + "display_name": "GitHub Copilot (`github_copilot`)", + "url": "https://docs.litellm.ai/docs/providers/github_copilot", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "chatgpt": { + "display_name": "ChatGPT Subscription (`chatgpt`)", + "url": "https://docs.litellm.ai/docs/providers/chatgpt", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, + "github": { + "display_name": "GitHub Models (`github`)", + "url": "https://docs.litellm.ai/docs/providers/github", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "gmi": { + "display_name": "GMI Cloud (`gmi`)", + "url": "https://docs.litellm.ai/docs/providers/gmi_cloud", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "vertex_ai": { + "display_name": "Google - Vertex AI (`vertex_ai`)", + "url": "https://docs.litellm.ai/docs/providers/vertex", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": true, + "moderations": false, + "batches": false, + "rerank": false, + "ocr": true, + "a2a": true, + "interactions": true, + "vector_stores_search": true, + "count_tokens": true, + "fine_tuning": true, + "rag_ingest": true, + "rag_query": true, + "generateContent": true, + "realtime": true + } + }, + "gemini": { + "display_name": "Google AI Studio - Gemini (`gemini`)", + "url": "https://docs.litellm.ai/docs/providers/gemini", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "interactions": true, + "a2a": true, + "vector_stores_search": true, + "count_tokens": true, + "rag_ingest": true, + "realtime": true, + "generateContent": true + } + }, + "gradient_ai": { + "display_name": "GradientAI (`gradient_ai`)", + "url": "https://docs.litellm.ai/docs/providers/gradient_ai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "groq": { + "display_name": "Groq AI (`groq`)", + "url": "https://docs.litellm.ai/docs/providers/groq", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "heroku": { + "display_name": "Heroku (`heroku`)", + "url": "https://docs.litellm.ai/docs/providers/heroku", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "hosted_vllm": { + "display_name": "Hosted VLLM (`hosted_vllm`)", + "url": "https://docs.litellm.ai/docs/providers/vllm", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": true, + "files": true, + "rerank": true, + "a2a": true, + "interactions": true + } + }, + "huggingface": { + "display_name": "Huggingface (`huggingface`)", + "url": "https://docs.litellm.ai/docs/providers/huggingface", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true, + "a2a": true, + "interactions": true + } + }, + "hyperbolic": { + "display_name": "Hyperbolic (`hyperbolic`)", + "url": "https://docs.litellm.ai/docs/providers/hyperbolic", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "watsonx": { + "display_name": "IBM - Watsonx.ai (`watsonx`)", + "url": "https://docs.litellm.ai/docs/providers/watsonx", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": true, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "infinity": { + "display_name": "Infinity (`infinity`)", + "url": "https://docs.litellm.ai/docs/providers/infinity", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false + } + }, + "jina_ai": { + "display_name": "Jina AI (`jina_ai`)", + "url": "https://docs.litellm.ai/docs/providers/jina_ai", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false + } + }, + "lambda_ai": { + "display_name": "Lambda AI (`lambda_ai`)", + "url": "https://docs.litellm.ai/docs/providers/lambda_ai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "lemonade": { + "display_name": "Lemonade (`lemonade`)", + "url": "https://docs.litellm.ai/docs/providers/lemonade", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "litellm_proxy": { + "display_name": "LiteLLM Proxy (`litellm_proxy`)", + "url": "https://docs.litellm.ai/docs/providers/litellm_proxy", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "llamafile": { + "display_name": "Llamafile (`llamafile`)", + "url": "https://docs.litellm.ai/docs/providers/llamafile", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "lm_studio": { + "display_name": "LM Studio (`lm_studio`)", + "url": "https://docs.litellm.ai/docs/providers/lm_studio", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "maritalk": { + "display_name": "Maritalk (`maritalk`)", + "url": "https://docs.litellm.ai/docs/providers/maritalk", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "meta_llama": { + "display_name": "Meta - Llama API (`meta_llama`)", + "url": "https://docs.litellm.ai/docs/providers/meta_llama", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "mistral": { + "display_name": "Mistral AI API (`mistral`)", + "url": "https://docs.litellm.ai/docs/providers/mistral", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "ocr": true, + "a2a": true, + "interactions": true + } + }, + "moonshot": { + "display_name": "Moonshot (`moonshot`)", + "url": "https://docs.litellm.ai/docs/providers/moonshot", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "docker_model_runner": { + "display_name": "Docker Model Runner (`docker_model_runner`)", + "url": "https://docs.litellm.ai/docs/providers/docker_model_runner", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "morph": { + "display_name": "Morph (`morph`)", + "url": "https://docs.litellm.ai/docs/providers/morph", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "nanogpt": { + "display_name": "NanoGPT (`nanogpt`)", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, + "nebius": { + "display_name": "Nebius AI Studio (`nebius`)", + "url": "https://docs.litellm.ai/docs/providers/nebius", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "nlp_cloud": { + "display_name": "NLP Cloud (`nlp_cloud`)", + "url": "https://docs.litellm.ai/docs/providers/nlp_cloud", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "novita": { + "display_name": "Novita AI (`novita`)", + "url": "https://novita.ai/models/llm?utm_source=github_litellm&utm_medium=github_readme&utm_campaign=github_link", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "nscale": { + "display_name": "Nscale (`nscale`)", + "url": "https://docs.litellm.ai/docs/providers/nscale", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "nvidia_nim": { + "display_name": "Nvidia NIM (`nvidia_nim`)", + "url": "https://docs.litellm.ai/docs/providers/nvidia_nim", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "oci": { + "display_name": "OCI (`oci`)", + "url": "https://docs.litellm.ai/docs/providers/oci", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "ollama": { + "display_name": "Ollama (`ollama`)", + "url": "https://docs.litellm.ai/docs/providers/ollama", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "ollama_chat": { + "display_name": "Ollama Chat (`ollama_chat`)", + "url": "https://docs.litellm.ai/docs/providers/ollama", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "oobabooga": { + "display_name": "Oobabooga (`oobabooga`)", + "url": "https://docs.litellm.ai/docs/providers/openai_compatible", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": true, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "openai": { + "display_name": "OpenAI (`openai`)", + "url": "https://docs.litellm.ai/docs/providers/openai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": true, + "rerank": false, + "container": true, + "compact": true, + "a2a": true, + "interactions": true, + "vector_store_files": true, + "vector_stores_create": true, + "vector_stores_search": true, + "assistants": true, + "container_files": true, + "fine_tuning": true, + "image_variations": true, + "rag_ingest": true, + "rag_query": true, + "realtime": true, + "text_completion": true + } + }, + "openai_like": { + "display_name": "OpenAI-like (`openai_like`)", + "url": "https://docs.litellm.ai/docs/providers/openai_compatible", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "assistants": true + } + }, + "openrouter": { + "display_name": "OpenRouter (`openrouter`)", + "url": "https://docs.litellm.ai/docs/providers/openrouter", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "ovhcloud": { + "display_name": "OVHCloud AI Endpoints (`ovhcloud`)", + "url": "https://docs.litellm.ai/docs/providers/ovhcloud", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": true, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "parallel_ai": { + "display_name": "Parallel AI (`parallel_ai`)", + "url": "https://docs.litellm.ai/docs/search/parallel_ai", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, + "perplexity": { + "display_name": "Perplexity AI (`perplexity`)", + "url": "https://docs.litellm.ai/docs/providers/perplexity", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true, + "a2a": true, + "interactions": true + } + }, + "petals": { + "display_name": "Petals (`petals`)", + "url": "https://docs.litellm.ai/docs/providers/petals", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "poe": { + "display_name": "Poe (`poe`)", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, + "publicai": { + "display_name": "PublicAI (`publicai`)", + "url": "https://docs.litellm.ai/docs/providers/publicai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "predibase": { + "display_name": "Predibase (`predibase`)", + "url": "https://docs.litellm.ai/docs/providers/predibase", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "recraft": { + "display_name": "Recraft (`recraft`)", + "url": "https://docs.litellm.ai/docs/providers/recraft", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false + } + }, + "replicate": { + "display_name": "Replicate (`replicate`)", + "url": "https://docs.litellm.ai/docs/providers/replicate", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "runwayml": { + "display_name": "RunwayML (`runwayml`)", + "url": "https://docs.litellm.ai/docs/providers/runwayml/videos", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": true, + "audio_transcriptions": false, + "audio_speech": true, + "moderations": false, + "batches": false, + "rerank": false, + "video_generations": true + } + }, + "sagemaker_chat": { + "display_name": "Sagemaker Chat (`sagemaker_chat`)", + "url": "https://docs.litellm.ai/docs/providers/aws_sagemaker", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "searxng": { + "display_name": "SearXNG (`searxng`)", + "url": "https://docs.litellm.ai/docs/search/searxng", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "search": true + } + }, + "sambanova": { + "display_name": "Sambanova (`sambanova`)", + "url": "https://docs.litellm.ai/docs/providers/sambanova", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "sap": { + "display_name": "SAP Generative AI Hub (`sap`)", + "url": "https://docs.litellm.ai/docs/providers/sap", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "scaleway": { + "display_name": "Scaleway (`scaleway`)", + "url": "https://docs.litellm.ai/docs/providers/scaleway", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "snowflake": { + "display_name": "Snowflake (`snowflake`)", + "url": "https://docs.litellm.ai/docs/providers/snowflake", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "synthetic": { + "display_name": "Synthetic (`synthetic`)", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, + "text-completion-codestral": { + "display_name": "Text Completion Codestral (`text-completion-codestral`)", + "url": "https://docs.litellm.ai/docs/providers/codestral", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "text-completion-openai": { + "display_name": "Text Completion OpenAI (`text-completion-openai`)", + "url": "https://docs.litellm.ai/docs/providers/text_completion_openai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": true, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "together_ai": { + "display_name": "Together AI (`together_ai`)", + "url": "https://docs.litellm.ai/docs/providers/togetherai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "topaz": { + "display_name": "Topaz (`topaz`)", + "url": "https://docs.litellm.ai/docs/providers/topaz", + "endpoints": { + "image_variations": true + } + }, + "tavily": { + "display_name": "Tavily (`tavily`)", + "url": "https://docs.litellm.ai/docs/search/tavily", + "endpoints": { + "search": true + } + }, + "triton": { + "display_name": "Triton (`triton`)", + "url": "https://docs.litellm.ai/docs/providers/triton-inference-server", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "v0": { + "display_name": "V0 (`v0`)", + "url": "https://docs.litellm.ai/docs/providers/v0", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "vercel_ai_gateway": { + "display_name": "Vercel AI Gateway (`vercel_ai_gateway`)", + "url": "https://docs.litellm.ai/docs/providers/vercel_ai_gateway", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "vllm": { + "display_name": "VLLM (`vllm`)", + "url": "https://docs.litellm.ai/docs/providers/vllm", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": true, + "files": true, + "rerank": true, + "a2a": true, + "interactions": true + } + }, + "volcengine": { + "display_name": "Volcengine (`volcengine`)", + "url": "https://docs.litellm.ai/docs/providers/volcano", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "voyage": { + "display_name": "Voyage AI (`voyage`)", + "url": "https://docs.litellm.ai/docs/providers/voyage", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": true + } + }, + "wandb": { + "display_name": "WandB Inference (`wandb`)", + "url": "https://docs.litellm.ai/docs/providers/wandb_inference", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "watsonx_text": { + "display_name": "Watsonx Text (`watsonx_text`)", + "url": "https://docs.litellm.ai/docs/providers/watsonx", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "xai": { + "display_name": "xAI (`xai`)", + "url": "https://docs.litellm.ai/docs/providers/xai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true, + "realtime": true + } + }, + "xinference": { + "display_name": "Xinference (`xinference`)", + "url": "https://docs.litellm.ai/docs/providers/xinference", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": true, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false + } + }, + "zai": { + "display_name": "Z.AI (Zhipu AI) (`zai`)", + "url": "https://docs.litellm.ai/docs/providers/zai", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "ragflow": { + "display_name": "RAGFlow (`ragflow`)", + "url": "https://docs.litellm.ai/docs/providers/ragflow", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "vector_stores_create": true, + "a2a": true, + "interactions": true + } + }, + "cursor": { + "display_name": "Cursor BYOK (`cursor`)", + "url": "https://docs.litellm.ai/docs/providers/cursor", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "langgraph": { + "display_name": "LangGraph (`langgraph`)", + "url": "https://docs.litellm.ai/docs/providers/langgraph", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "vertex_ai/agent_engine": { + "display_name": "Vertex AI Agent Engine (`vertex_ai/agent_engine`)", + "url": "https://docs.litellm.ai/docs/providers/vertex_ai_agent_engine", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "pydantic_ai_agents": { + "display_name": "Pydantic AI Agents (`pydantic_ai_agents`)", + "url": "https://docs.litellm.ai/docs/providers/pydantic_ai_agent", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true + } + }, + "stability": { + "display_name": "Stability AI (`stability`)", + "url": "https://docs.litellm.ai/docs/providers/stability", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": true, + "image_edits": true, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false + } + }, + "venice": { + "display_name": "Venice.ai (`venice`)", + "url": "https://docs.litellm.ai/docs/providers/venice", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "gigachat": { + "display_name": "GigaChat (`gigachat`)", + "url": "https://docs.litellm.ai/docs/providers/gigachat", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true + } + }, + "google_pse": { + "display_name": "Google PSE (`google_pse`)", + "url": "https://docs.litellm.ai/docs/search/google_pse", + "endpoints": { + "search": true + } + }, + "milvus": { + "display_name": "Milvus (`milvus`)", + "url": "https://docs.litellm.ai/docs/providers/milvus_vector_stores", + "endpoints": { + "vector_stores_search": true + } + }, + "minimax": { + "display_name": "Minimax (`minimax`)", + "url": "https://docs.litellm.ai/docs/providers/minimax", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true + } + }, + "pg_vector": { + "display_name": "PG Vector (`pg_vector`)", + "url": "https://docs.litellm.ai/docs/providers/pg_vector", + "endpoints": { + "vector_stores_search": true + } + }, + "helicone": { + "display_name": "Helicone (`helicone`)", + "url": "https://docs.litellm.ai/docs/providers/helicone", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true + } + }, + "llamagate": { + "display_name": "LlamaGate (`llamagate`)", + "url": "https://docs.litellm.ai/docs/providers/llamagate", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true + } + }, + "xiaomi_mimo": { + "display_name": "Xiaomi Mimo (`xiaomi_mimo`)", + "url": "https://docs.litellm.ai/docs/providers/xiaomi_mimo", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true + } + }, + "manus": { + "display_name": "Manus (`manus`)", + "url": "https://docs.litellm.ai/docs/providers/manus", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "files": true, + "rerank": false, + "a2a": true, + "interactions": true + } + }, + "sarvam": { + "display_name": "Sarvam (`sarvam`)", + "url": "https://docs.litellm.ai/docs/providers/sarvam", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true + } + } + }, + "endpoints": { + "a2a": { + "docs_label": "a2a", + "display_name": "A2A (Agent-to-Agent) protocol for agent communication", + "leftnav_label": "/a2a", + "provider_json_field": "a2a", + "url": "https://docs.litellm.ai/docs/a2a", + "bridges_to_chat_completion": true + }, + "messages": { + "docs_label": "anthropic_unified", + "display_name": "Anthropic Messages API", + "leftnav_label": "/messages", + "provider_json_field": "messages", + "url": "https://docs.litellm.ai/docs/anthropic_unified", + "bridges_to_chat_completion": true + }, + "anthropic_count_tokens": { + "docs_label": "anthropic_count_tokens", + "display_name": "Anthropic Count Tokens API", + "leftnav_label": "/count_tokens", + "provider_json_field": "count_tokens", + "url": "https://docs.litellm.ai/docs/anthropic_count_tokens" + }, + "apply_guardrail": { + "docs_label": "apply_guardrail", + "display_name": "Unified Apply Guardrail API", + "leftnav_label": "/guardrails/apply_guardrail", + "provider_json_field": "apply_guardrail", + "url": "https://docs.litellm.ai/docs/apply_guardrail" + }, + "assistants": { + "docs_label": "assistants", + "display_name": "OpenAI Assistants API", + "leftnav_label": "/assistants", + "provider_json_field": "assistants", + "url": "https://docs.litellm.ai/docs/assistants" + }, + "audio_transcription": { + "docs_label": "audio_transcription", + "display_name": "OpenAI Audio Transcription API", + "leftnav_label": "/audio/transcriptions", + "provider_json_field": "audio_transcriptions", + "url": "https://docs.litellm.ai/docs/audio_transcription" + }, + "batches": { + "docs_label": "batches", + "display_name": "OpenAI Batches API", + "leftnav_label": "/batches", + "provider_json_field": "batches", + "url": "https://docs.litellm.ai/docs/batches" + }, + "bedrock_invoke": { + "docs_label": "bedrock_invoke", + "display_name": "Bedrock Invoke API", + "leftnav_label": "/invoke", + "provider_json_field": "bedrock_invoke", + "url": "https://docs.litellm.ai/docs/bedrock_invoke" + }, + "bedrock_converse": { + "docs_label": "bedrock_converse", + "display_name": "Bedrock Converse API", + "leftnav_label": "/converse", + "provider_json_field": "bedrock_converse", + "url": "https://docs.litellm.ai/docs/bedrock_converse" + }, + "chat_completions": { + "docs_label": "chat_completions", + "display_name": "OpenAI Chat Completions API", + "leftnav_label": "/chat/completions", + "provider_json_field": "chat_completions", + "url": "https://docs.litellm.ai/docs/chat_completions" + }, + "container_files": { + "docs_label": "container_files", + "display_name": "OpenAI Container Files API", + "leftnav_label": "/create/container/files", + "provider_json_field": "container_files", + "url": "https://docs.litellm.ai/docs/container_files" + }, + "container": { + "docs_label": "containers", + "display_name": "OpenAI Containers API", + "leftnav_label": "/container", + "provider_json_field": "container", + "url": "https://docs.litellm.ai/docs/containers" + }, + "embeddings": { + "docs_label": "embedding/supported_embedding", + "display_name": "OpenAI Embeddings API", + "leftnav_label": "/embeddings", + "provider_json_field": "embeddings", + "url": "https://docs.litellm.ai/docs/embedding/supported_embedding" + }, + "files": { + "docs_label": "files", + "display_name": "OpenAI Files API", + "leftnav_label": "/files", + "provider_json_field": "files", + "url": "https://docs.litellm.ai/docs/proxy/litellm_managed_files" + }, + "fine_tuning": { + "docs_label": "fine_tuning", + "display_name": "OpenAI Fine-Tuning API", + "leftnav_label": "/fine_tuning", + "provider_json_field": "fine_tuning", + "url": "https://docs.litellm.ai/docs/proxy/managed_finetuning" + }, + "generateContent": { + "docs_label": "generateContent", + "display_name": "Google GenerateContent API", + "leftnav_label": "/generateContent", + "provider_json_field": "generateContent", + "url": "https://docs.litellm.ai/docs/generateContent", + "bridges_to_chat_completion": true + }, + "image_edits": { + "docs_label": "image_edits", + "display_name": "OpenAI Images Edits API", + "leftnav_label": "/images/edits", + "provider_json_field": "image_edits", + "url": "https://docs.litellm.ai/docs/image_edits" + }, + "image_generations": { + "docs_label": "image_generation", + "display_name": "OpenAI Images Generations API", + "leftnav_label": "/images/generations", + "provider_json_field": "image_generations", + "url": "https://docs.litellm.ai/docs/image_generation" + }, + "image_variations": { + "docs_label": "image_variations", + "display_name": "OpenAI Images Variations API", + "leftnav_label": "/images/variations", + "provider_json_field": "image_variations", + "url": "https://docs.litellm.ai/docs/image_variations" + }, + "interactions": { + "docs_label": "interactions", + "display_name": "Google Interactions API", + "leftnav_label": "/interactions", + "provider_json_field": "interactions", + "url": "https://docs.litellm.ai/docs/interactions", + "bridges_to_chat_completion": true + }, + "mcp": { + "docs_label": "mcp", + "display_name": "Model Context Protocol (MCP)", + "leftnav_label": "/mcp", + "provider_json_field": "mcp", + "url": "https://docs.litellm.ai/docs/mcp" + }, + "moderation": { + "docs_label": "moderation", + "display_name": "OpenAI Moderations API", + "leftnav_label": "/moderations", + "provider_json_field": "moderations", + "url": "https://docs.litellm.ai/docs/moderation" + }, + "ocr": { + "docs_label": "ocr", + "display_name": "Mistral OCR API", + "leftnav_label": "/ocr", + "provider_json_field": "ocr", + "url": "https://docs.litellm.ai/docs/ocr" + }, + "rag_ingest": { + "docs_label": "rag_ingest", + "display_name": "RAG Ingest API", + "leftnav_label": "/rag/ingest", + "provider_json_field": "rag_ingest", + "url": "https://docs.litellm.ai/docs/rag_ingest" + }, + "rag_query": { + "docs_label": "rag_query", + "display_name": "RAG Query API", + "leftnav_label": "/rag/query", + "provider_json_field": "rag_query", + "url": "https://docs.litellm.ai/docs/rag_query" + }, + "realtime": { + "docs_label": "realtime", + "display_name": "OpenAI Realtime API", + "leftnav_label": "/realtime", + "provider_json_field": "realtime", + "url": "https://docs.litellm.ai/docs/realtime" + }, + "rerank": { + "docs_label": "rerank", + "display_name": "Cohere Rerank API", + "leftnav_label": "/rerank", + "provider_json_field": "rerank", + "url": "https://docs.litellm.ai/docs/rerank" + }, + "responses": { + "docs_label": "response_api", + "display_name": "OpenAI Responses API", + "leftnav_label": "/responses", + "provider_json_field": "responses", + "url": "https://docs.litellm.ai/docs/response_api", + "bridges_to_chat_completion": true + }, + "response_api_compact": { + "docs_label": "response_api_compact", + "display_name": "OpenAI Responses API", + "leftnav_label": "/responses", + "provider_json_field": "compact", + "url": "https://docs.litellm.ai/docs/response_api" + }, + "search": { + "docs_label": "search", + "display_name": "Search API", + "leftnav_label": "/search", + "provider_json_field": "search", + "url": "https://docs.litellm.ai/docs/search" + }, + "skills": { + "docs_label": "skills", + "display_name": "Anthropic Skills API", + "leftnav_label": "/skills", + "provider_json_field": "skills", + "url": "https://docs.litellm.ai/docs/skills" + }, + "text_completion": { + "docs_label": "text_completion", + "display_name": "OpenAI Completions API", + "leftnav_label": "/completions", + "provider_json_field": "text_completion", + "url": "https://docs.litellm.ai/docs/text_completion", + "bridges_to_chat_completion": true + }, + "text_to_speech": { + "docs_label": "text_to_speech", + "display_name": "OpenAI Text-to-Speech API", + "leftnav_label": "/audio/speech", + "provider_json_field": "audio_speech", + "url": "https://docs.litellm.ai/docs/text_to_speech" + }, + "vector_store_files": { + "docs_label": "vector_store_files", + "display_name": "OpenAI Vector Store Files API", + "leftnav_label": "/vector_stores/files", + "provider_json_field": "vector_store_files", + "url": "https://docs.litellm.ai/docs/vector_store_files" + }, + "vector_stores_create": { + "docs_label": "vector_stores_create", + "display_name": "OpenAI Vector Stores Create API", + "leftnav_label": "/vector_stores/create", + "provider_json_field": "vector_stores_create", + "url": "https://docs.litellm.ai/docs/vector_stores/create" + }, + "vector_stores_search": { + "docs_label": "vector_stores_search", + "display_name": "OpenAI Vector Stores Search API", + "leftnav_label": "/vector_stores/search", + "provider_json_field": "vector_stores_search", + "url": "https://docs.litellm.ai/docs/vector_stores/search" + }, + "videos": { + "docs_label": "videos", + "display_name": "OpenAI Videos API", + "leftnav_label": "/videos", + "provider_json_field": "video_generations", + "url": "https://docs.litellm.ai/docs/videos" + } + } +} diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7484de33ce4..08213f40b43 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -331,7 +331,7 @@ class MCPServerManager: static_headers=server_config.get("static_headers", None), allow_all_keys=bool(server_config.get("allow_all_keys", False)), available_on_public_internet=bool( - server_config.get("available_on_public_internet", False) + server_config.get("available_on_public_internet", True) ), ) self.config_mcp_servers[server_id] = new_server @@ -634,7 +634,7 @@ class MCPServerManager: disallowed_tools=getattr(mcp_server, "disallowed_tools", None), allow_all_keys=mcp_server.allow_all_keys, available_on_public_internet=bool( - getattr(mcp_server, "available_on_public_internet", False) + getattr(mcp_server, "available_on_public_internet", True) ), updated_at=getattr(mcp_server, "updated_at", None), ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 48c837c1e4f..5b3d5bd60e2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -5,6 +5,7 @@ LiteLLM MCP Server Routes import asyncio import contextlib + import traceback import uuid from datetime import datetime @@ -84,6 +85,7 @@ except ImportError as e: _SESSION_MANAGERS_INITIALIZED = False _INITIALIZATION_LOCK = asyncio.Lock() + if MCP_AVAILABLE: from mcp.server import Server @@ -1919,65 +1921,86 @@ if MCP_AVAILABLE: mgr: "StreamableHTTPSessionManager", ) -> bool: """ - Handle stale MCP session IDs to prevent "Session not found" errors. - - When clients reconnect after a server restart or session cleanup, they may - send a session ID that no longer exists. This function handles two scenarios: - - 1. Non-DELETE requests: Strip the stale session ID header so the session - manager creates a fresh session transparently. - - 2. DELETE requests: Return success (200) immediately for idempotent behavior, - since the desired state (session doesn't exist) is already achieved. + Inspect the incoming ``mcp-session-id`` header **before** the + request reaches the MCP SDK. If the session is stale (not known + to this worker), strip the header so the SDK creates a fresh + stateless session instead of returning a 400. Returns: - True if the request was handled (DELETE on non-existent session) - False if the request should continue to the session manager + True if the request was fully handled (e.g. DELETE on + non-existent session). False if the request should continue + to the session manager. - Fixes https://github.com/BerriAI/litellm/issues/20292 + Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header = b"mcp-session-id" + _headers = scope.get("headers", []) + + def _normalize_header_name(header_name: Any) -> Optional[bytes]: + if isinstance(header_name, bytes): + return header_name.lower() + if isinstance(header_name, str): + return header_name.lower().encode("utf-8", errors="replace") + return None + _session_id: Optional[str] = None - for header_name, header_value in scope.get("headers", []): - if header_name == _mcp_session_header: - _session_id = header_value.decode("utf-8", errors="replace") + for header_name, header_value in _headers: + if _normalize_header_name(header_name) == _mcp_session_header: + if isinstance(header_value, bytes): + _session_id = header_value.decode("utf-8", errors="replace") + else: + _session_id = str(header_value) break if _session_id is None: return False + # Check in-memory session tracking known_sessions = getattr(mgr, "_server_instances", None) - if known_sessions is None or _session_id in known_sessions: - # Session exists or we can't check - let the session manager handle it + # If we cannot inspect known_sessions, let the manager handle it + if known_sessions is None: return False - # Session doesn't exist - handle based on request method + # If session exists in this worker's memory, let the manager handle it + try: + if _session_id in known_sessions: + return False + except Exception: + verbose_logger.debug( + "Unable to inspect active MCP sessions for '%s'. " + "Deferring to session manager.", + _session_id, + ) + return False + + # --- Session not in this worker's memory --- method = scope.get("method", "").upper() - + if method == "DELETE": - # Idempotent DELETE: session doesn't exist, return success verbose_logger.info( - f"DELETE request for non-existent MCP session '{_session_id}'. " - "Returning success (idempotent DELETE)." + "DELETE request for non-existent MCP session '%s'. " + "Returning success (idempotent DELETE).", + _session_id, ) success_response = JSONResponse( status_code=200, - content={"message": "Session terminated successfully"} + content={"message": "Session terminated successfully"}, ) await success_response(scope, receive, send) return True - else: - # Non-DELETE: strip stale session ID to allow new session creation - verbose_logger.warning( - "MCP session ID '%s' not found in active sessions. " - "Stripping stale header to force new session creation.", - _session_id, - ) - scope["headers"] = [ - (k, v) for k, v in scope["headers"] - if k != _mcp_session_header - ] - return False + + # Non-DELETE: strip stale session ID to allow new session creation + verbose_logger.warning( + "MCP session ID '%s' not found in this worker's memory. " + "Stripping stale header to force new session creation.", + _session_id, + ) + scope["headers"] = [ + (k, v) + for k, v in _headers + if _normalize_header_name(k) != _mcp_session_header + ] + return False async def handle_streamable_http_mcp( scope: Scope, receive: Receive, send: Send @@ -2055,7 +2078,9 @@ if MCP_AVAILABLE: # Handle stale session IDs - either strip them for reconnection # or return success for idempotent DELETE operations - handled = await _handle_stale_mcp_session(scope, receive, send, session_manager) + handled = await _handle_stale_mcp_session( + scope, receive, send, session_manager + ) if handled: # Request was fully handled (e.g., DELETE on non-existent session) return diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index de1609baf62..afedb6c8e72 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -431,6 +431,7 @@ class LiteLLMRoutes(enum.Enum): agent_routes = [ "/v1/agents", + "/v1/agents/{agent_id}", "/agents", "/a2a/{agent_id}", "/a2a/{agent_id}/message/send", @@ -1092,7 +1093,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): token_url: Optional[str] = None registration_url: Optional[str] = None allow_all_keys: bool = False - available_on_public_internet: bool = False + available_on_public_internet: bool = True @model_validator(mode="before") @classmethod @@ -1106,7 +1107,9 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): raise ValueError("args is required for stdio transport") elif transport in [MCPTransport.http, MCPTransport.sse]: if not values.get("url") and not values.get("spec_path"): - raise ValueError("url or spec_path is required for HTTP/SSE transport") + raise ValueError( + "url or spec_path is required for HTTP/SSE transport" + ) return values @model_validator(mode="before") @@ -1144,7 +1147,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): token_url: Optional[str] = None registration_url: Optional[str] = None allow_all_keys: bool = False - available_on_public_internet: bool = False + available_on_public_internet: bool = True @model_validator(mode="before") @classmethod @@ -1158,7 +1161,9 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): raise ValueError("args is required for stdio transport") elif transport in [MCPTransport.http, MCPTransport.sse]: if not values.get("url") and not values.get("spec_path"): - raise ValueError("url or spec_path is required for HTTP/SSE transport") + raise ValueError( + "url or spec_path is required for HTTP/SSE transport" + ) return values @@ -1199,7 +1204,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): token_url: Optional[str] = None registration_url: Optional[str] = None allow_all_keys: bool = False - available_on_public_internet: bool = False + available_on_public_internet: bool = True class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): @@ -1409,12 +1414,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @@ -1437,12 +1442,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -2268,6 +2273,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): end_user_tpm_limit: Optional[int] = None end_user_rpm_limit: Optional[int] = None end_user_max_budget: Optional[float] = None + end_user_model_max_budget: Optional[dict] = None # Organization Params organization_max_budget: Optional[float] = None @@ -2275,6 +2281,9 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): organization_rpm_limit: Optional[int] = None organization_metadata: Optional[dict] = None + # Project Params + project_metadata: Optional[dict] = None + # Time stamps last_refreshed_at: Optional[float] = None # last time joint view was pulled from db @@ -2408,6 +2417,7 @@ class UserAPIKeyAuth( key_alias=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, team_alias="system", user_id="system", + user_role=LitellmUserRoles.PROXY_ADMIN, ) @@ -2576,6 +2586,7 @@ class NewProjectRequest(LiteLLM_BudgetTable): team_id: str budget_id: Optional[str] = None metadata: Optional[dict] = None + tags: Optional[List[str]] = None models: List[str] = [] model_rpm_limit: Optional[dict] = None model_tpm_limit: Optional[dict] = None @@ -2585,6 +2596,11 @@ class NewProjectRequest(LiteLLM_BudgetTable): @model_validator(mode="before") @classmethod def set_model_info(cls, values): + if "tags" in values and values["tags"] is not None: + if not isinstance(values["tags"], list): + raise ValueError( + f"tags must be a list of strings, got {type(values['tags']).__name__}" + ) for field in LiteLLM_ManagementEndpoint_MetadataFields: if values.get(field) is not None: if values.get("metadata") is None: @@ -2602,6 +2618,7 @@ class UpdateProjectRequest(LiteLLM_BudgetTable): description: Optional[str] = None team_id: Optional[str] = None metadata: Optional[dict] = None + tags: Optional[List[str]] = None models: Optional[List[str]] = None model_rpm_limit: Optional[dict] = None model_tpm_limit: Optional[dict] = None @@ -2612,6 +2629,11 @@ class UpdateProjectRequest(LiteLLM_BudgetTable): @model_validator(mode="before") @classmethod def set_model_info(cls, values): + if "tags" in values and values["tags"] is not None: + if not isinstance(values["tags"], list): + raise ValueError( + f"tags must be a list of strings, got {type(values['tags']).__name__}" + ) for field in LiteLLM_ManagementEndpoint_MetadataFields: if values.get(field) is not None: if values.get("metadata") is None: @@ -2645,6 +2667,8 @@ class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): object_permission_id: Optional[str] = None created_by: str updated_by: str + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None litellm_budget_table: Optional[LiteLLM_BudgetTable] = None object_permission: Optional[LiteLLM_ObjectPermissionTable] = None @@ -3056,7 +3080,9 @@ class SpendLogsMetadata(TypedDict): str ] # S3/GCS object key for cold storage retrieval litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds - attempted_retries: Optional[int] # Number of retries attempted (0 = first attempt succeeded) + attempted_retries: Optional[ + int + ] # Number of retries attempted (0 = first attempt succeeded) max_retries: Optional[int] # Max retries configured for this request cost_breakdown: Optional[ CostBreakdown @@ -4117,10 +4143,10 @@ class SpendUpdateQueueItem(TypedDict, total=False): class ToolDiscoveryQueueItem(TypedDict, total=False): tool_name: str - origin: Optional[str] # MCP server name or "user_defined" + origin: Optional[str] # MCP server name or "user_defined" created_by: Optional[str] - key_hash: Optional[str] # hash of virtual key that triggered discovery - team_id: Optional[str] # team that triggered discovery + key_hash: Optional[str] # hash of virtual key that triggered discovery + team_id: Optional[str] # team that triggered discovery key_alias: Optional[str] # human-readable key alias @@ -4144,6 +4170,7 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): """Table for managing vector stores with target_model_names support.""" + unified_resource_id: str resource_object: Optional[Any] = None # VectorStoreCreateResponse model_mappings: Dict[str, str] diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index b411b81b434..65674d01be7 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -31,6 +31,23 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( router = APIRouter() +def _check_agent_management_permission(user_api_key_dict: UserAPIKeyAuth) -> None: + """ + Raises HTTP 403 if the caller does not have permission to create, update, + or delete agents. Only PROXY_ADMIN users are allowed to perform these + write operations. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={ + "error": "Only proxy admins can create, update, or delete agents. Your role={}".format( + user_api_key_dict.user_role + ) + }, + ) + + @router.get( "/v1/agents", tags=["[beta] A2A Agents"], @@ -164,6 +181,8 @@ async def create_agent( """ from litellm.proxy.proxy_server import prisma_client + _check_agent_management_permission(user_api_key_dict) + if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -302,6 +321,8 @@ async def update_agent( """ from litellm.proxy.proxy_server import prisma_client + _check_agent_management_permission(user_api_key_dict) + if prisma_client is None: raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value @@ -391,6 +412,8 @@ async def patch_agent( """ from litellm.proxy.proxy_server import prisma_client + _check_agent_management_permission(user_api_key_dict) + if prisma_client is None: raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value @@ -441,7 +464,10 @@ async def patch_agent( tags=["Agents"], dependencies=[Depends(user_api_key_auth)], ) -async def delete_agent(agent_id: str): +async def delete_agent( + agent_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Delete an agent @@ -460,6 +486,8 @@ async def delete_agent(agent_id: str): """ from litellm.proxy.proxy_server import prisma_client + _check_agent_management_permission(user_api_key_dict) + if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9921b74b561..553ba4d6c49 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -8,6 +8,7 @@ JWT token must have 'litellm_proxy_admin' in scope. import fnmatch import os +import re from typing import Any, List, Literal, Optional, Set, Tuple, cast from cryptography import x509 @@ -235,7 +236,17 @@ class JWTHandler: return self.litellm_jwtauth.team_id_default else: return default_value - # At this point, team_id is not the sentinel, so it should be a string + # AAD and other IdPs often send roles/groups as a list of strings. + # team_id_jwt_field is singular, so take the first element when a list + # is returned. This avoids "unhashable type: 'list'" errors downstream. + if isinstance(team_id, list): + if not team_id: + return default_value + verbose_proxy_logger.debug( + f"JWT Auth: team_id_jwt_field '{self.litellm_jwtauth.team_id_jwt_field}' " + f"returned a list {team_id}; using first element '{team_id[0]}' automatically." + ) + team_id = team_id[0] return team_id # type: ignore[return-value] elif self.litellm_jwtauth.team_id_default is not None: team_id = self.litellm_jwtauth.team_id_default @@ -453,6 +464,52 @@ class JWTHandler: scopes = [] return scopes + async def _resolve_jwks_url(self, url: str) -> str: + """ + If url points to an OIDC discovery document (*.well-known/openid-configuration), + fetch it and return the jwks_uri contained within. Otherwise return url unchanged. + This lets JWT_PUBLIC_KEY_URL be set to a well-known discovery endpoint instead of + requiring operators to manually find the JWKS URL. + """ + if ".well-known/openid-configuration" not in url: + return url + + cache_key = f"litellm_oidc_discovery_{url}" + cached_jwks_uri = await self.user_api_key_cache.async_get_cache(cache_key) + if cached_jwks_uri is not None: + return cached_jwks_uri + + verbose_proxy_logger.debug( + f"JWT Auth: Fetching OIDC discovery document from {url}" + ) + response = await self.http_handler.get(url) + if response.status_code != 200: + raise Exception( + f"JWT Auth: OIDC discovery endpoint {url} returned status {response.status_code}: {response.text}" + ) + try: + discovery = response.json() + except Exception as e: + raise Exception( + f"JWT Auth: Failed to parse OIDC discovery document at {url}: {e}" + ) + + jwks_uri = discovery.get("jwks_uri") + if not jwks_uri: + raise Exception( + f"JWT Auth: OIDC discovery document at {url} does not contain a 'jwks_uri' field." + ) + + verbose_proxy_logger.debug( + f"JWT Auth: Resolved OIDC discovery {url} -> jwks_uri={jwks_uri}" + ) + await self.user_api_key_cache.async_set_cache( + key=cache_key, + value=jwks_uri, + ttl=self.litellm_jwtauth.public_key_ttl, + ) + return jwks_uri + async def get_public_key(self, kid: Optional[str]) -> dict: keys_url = os.getenv("JWT_PUBLIC_KEY_URL") @@ -462,6 +519,7 @@ class JWTHandler: keys_url_list = [url.strip() for url in keys_url.split(",")] for key_url in keys_url_list: + key_url = await self._resolve_jwks_url(key_url) cache_key = f"litellm_jwt_auth_keys_{key_url}" cached_keys = await self.user_api_key_cache.async_get_cache(cache_key) @@ -913,8 +971,30 @@ class JWTAuthManager: if jwt_handler.is_required_team_id() is True: team_id_field = jwt_handler.litellm_jwtauth.team_id_jwt_field team_alias_field = jwt_handler.litellm_jwtauth.team_alias_jwt_field + hint = "" + if team_id_field: + # "roles.0" — dot-notation numeric indexing is not supported + if "." in team_id_field: + parts = team_id_field.rsplit(".", 1) + if parts[-1].isdigit(): + base_field = parts[0] + hint = ( + f" Hint: dot-notation array indexing (e.g. '{team_id_field}') is not " + f"supported. Use '{base_field}' instead — LiteLLM automatically " + f"uses the first element when the field value is a list." + ) + # "roles[0]" — bracket-notation indexing is also not supported in get_nested_value + elif "[" in team_id_field and team_id_field.endswith("]"): + m = re.match(r"^(\w+)\[(\d+)\]$", team_id_field) + if m: + base_field = m.group(1) + hint = ( + f" Hint: array indexing (e.g. '{team_id_field}') is not supported " + f"in team_id_jwt_field. Use '{base_field}' instead — LiteLLM " + f"automatically uses the first element when the field value is a list." + ) raise Exception( - f"No team found in token. Checked team_id field '{team_id_field}' and team_alias field '{team_alias_field}'" + f"No team found in token. Checked team_id field '{team_id_field}' and team_alias field '{team_alias_field}'.{hint}" ) return individual_team_id, team_object diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 133f8ec136d..8ad3b83c043 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -183,6 +183,9 @@ def _apply_budget_limits_to_end_user_params( if budget_info.max_budget is not None: end_user_params["end_user_max_budget"] = budget_info.max_budget + if budget_info.model_max_budget is not None: + end_user_params["end_user_model_max_budget"] = budget_info.model_max_budget + verbose_proxy_logger.debug(f"Applied budget limits to end user {end_user_id}") @@ -209,10 +212,12 @@ async def user_api_key_auth_websocket(websocket: WebSocket): api_key = websocket.headers.get("api-key") if not api_key: # Try extracting from WebSocket subprotocol (browser clients) - for protocol in websocket.headers.get("sec-websocket-protocol", "").split(","): + for protocol in websocket.headers.get("sec-websocket-protocol", "").split( + "," + ): protocol = protocol.strip() if protocol.startswith("openai-insecure-api-key."): - api_key = protocol[len("openai-insecure-api-key."):] + api_key = protocol[len("openai-insecure-api-key.") :] break if not api_key: await websocket.close(code=status.WS_1008_POLICY_VIOLATION) @@ -241,9 +246,20 @@ def update_valid_token_with_end_user_params( valid_token: UserAPIKeyAuth, end_user_params: dict ) -> UserAPIKeyAuth: valid_token.end_user_id = end_user_params.get("end_user_id") - valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") - valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") - valid_token.allowed_model_region = end_user_params.get("allowed_model_region") + # Only overwrite token fields when the DB-derived value is not None. + # This prevents DB lookups (where the budget table has no value set) + # from silently clearing values that a custom auth function may have + # already set on the token. + if end_user_params.get("end_user_tpm_limit") is not None: + valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] + if end_user_params.get("end_user_rpm_limit") is not None: + valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] + if end_user_params.get("allowed_model_region") is not None: + valid_token.allowed_model_region = end_user_params["allowed_model_region"] + if end_user_params.get("end_user_model_max_budget") is not None: + valid_token.end_user_model_max_budget = end_user_params[ + "end_user_model_max_budget" + ] return valid_token @@ -498,13 +514,29 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request=request, api_key=api_key, user_custom_auth=user_custom_auth ) if response is not None and isinstance(response, UserAPIKeyAuth): - return UserAPIKeyAuth.model_validate(response) + validated = UserAPIKeyAuth.model_validate(response) + validated = await _run_post_custom_auth_checks( + valid_token=validated, + request=request, + request_data=request_data, + route=route, + parent_otel_span=parent_otel_span, + ) + return validated elif response is not None and isinstance(response, str): api_key = response custom_auth_api_key = True elif user_custom_auth is not None: response = await user_custom_auth(request=request, api_key=api_key) # type: ignore - return UserAPIKeyAuth.model_validate(response) + validated = UserAPIKeyAuth.model_validate(response) + validated = await _run_post_custom_auth_checks( + valid_token=validated, + request=request, + request_data=request_data, + route=route, + parent_otel_span=parent_otel_span, + ) + return validated ### LITELLM-DEFINED AUTH FUNCTION ### #### IF JWT #### @@ -674,6 +706,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) + if _jwt_project_obj is not None: + valid_token.project_metadata = _jwt_project_obj.metadata # run through common checks _ = await common_checks( @@ -1210,6 +1244,21 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 model=current_model, ) + # Check 5b. End-user model max budget + end_user_mmb = valid_token.end_user_model_max_budget + if ( + end_user_mmb is not None + and isinstance(end_user_mmb, dict) + and len(end_user_mmb) > 0 + and current_model is not None + and valid_token.end_user_id is not None + ): + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=valid_token.end_user_id, + end_user_model_max_budget=end_user_mmb, + model=current_model, + ) + # Check 6: Additional Common Checks across jwt + key auth if valid_token.team_id is not None: try: @@ -1249,6 +1298,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) + if _project_obj is not None: + valid_token.project_metadata = _project_obj.metadata global_proxy_spend = None if ( @@ -1501,3 +1552,220 @@ def _update_key_budget_with_temp_budget_increase( temp_budget_increase = _get_temp_budget_increase(valid_token) or 0.0 valid_token.max_budget = valid_token.max_budget + temp_budget_increase return valid_token + + +async def _lookup_end_user_and_apply_budget( + valid_token: UserAPIKeyAuth, + route: str, + parent_otel_span: Optional[Span], + prisma_client, + user_api_key_cache, + proxy_logging_obj, +): + """Look up end_user from DB and apply budget limits to valid_token.""" + end_user_object = None + try: + end_user_object = await get_end_user_object( + end_user_id=valid_token.end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if end_user_object is not None: + end_user_params = { + "end_user_id": valid_token.end_user_id, + "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, + budget_info=end_user_object.litellm_budget_table, + end_user_id=valid_token.end_user_id, + ) + valid_token = update_valid_token_with_end_user_params( + valid_token=valid_token, end_user_params=end_user_params + ) + elif litellm.max_end_user_budget_id is not None: + from litellm.proxy.auth.auth_checks import get_default_end_user_budget + + default_budget = await get_default_end_user_budget( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) + if default_budget is not None: + end_user_params = {"end_user_id": valid_token.end_user_id} + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=default_budget, + end_user_id=valid_token.end_user_id, + ) + valid_token = update_valid_token_with_end_user_params( + valid_token=valid_token, end_user_params=end_user_params + ) + except Exception as e: + if isinstance(e, litellm.BudgetExceededError): + raise e + verbose_proxy_logger.debug(f"Unable to find user in db. Error - {str(e)}") + return valid_token, end_user_object + + +async def _run_post_custom_auth_checks( + valid_token: UserAPIKeyAuth, + request: Request, + request_data: dict, + route: str, + parent_otel_span: Optional[Span], +) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import ( + prisma_client, + user_api_key_cache, + proxy_logging_obj, + general_settings, + llm_router, + model_max_budget_limiter, + ) + + # 1. Look up end_user object from DB if end_user_id is set + end_user_object = None + if valid_token.end_user_id is not None: + valid_token, end_user_object = await _lookup_end_user_and_apply_budget( + valid_token=valid_token, + route=route, + parent_otel_span=parent_otel_span, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + # 2. Check token expiry + if valid_token.expires is not None: + current_time = datetime.now(timezone.utc) + if isinstance(valid_token.expires, datetime): + expiry_time = valid_token.expires + else: + expiry_time = datetime.fromisoformat(valid_token.expires) + if ( + expiry_time.tzinfo is None + or expiry_time.tzinfo.utcoffset(expiry_time) is None + ): + expiry_time = expiry_time.replace(tzinfo=timezone.utc) + if expiry_time < current_time: + raise ProxyException( + message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", + type=ProxyErrorTypes.expired_key, + code=400, + param=abbreviate_api_key(api_key=valid_token.token) + if valid_token.token + else "", + ) + + current_model = request_data.get("model", None) + + # 3. Check key-level model_max_budget + max_budget_per_model = valid_token.model_max_budget + if ( + max_budget_per_model is not None + and isinstance(max_budget_per_model, dict) + and len(max_budget_per_model) > 0 + and current_model is not None + and valid_token.token is not None + ): + await model_max_budget_limiter.is_key_within_model_budget( + user_api_key_dict=valid_token, + model=current_model, + ) + + # 4. Check end-user model_max_budget + end_user_mmb = valid_token.end_user_model_max_budget + if ( + end_user_mmb is not None + and isinstance(end_user_mmb, dict) + and len(end_user_mmb) > 0 + and current_model is not None + and valid_token.end_user_id is not None + ): + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=valid_token.end_user_id, + end_user_model_max_budget=end_user_mmb, + model=current_model, + ) + + # 5. Look up user object if user_id is set + user_object = None + if valid_token.user_id is not None: + try: + user_object = await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + # If user_role is PROXY_ADMIN on the token, create a synthetic user object + # so that admin route checks pass for custom auth + if valid_token.user_role == LitellmUserRoles.PROXY_ADMIN: + user_object = LiteLLM_UserTable( + user_id=valid_token.user_id, + user_role=LitellmUserRoles.PROXY_ADMIN, + spend=0.0, + ) + + # 6. Run common checks + if valid_token.team_id is not None: + try: + _team_obj = await get_team_object( + team_id=valid_token.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + _team_obj = LiteLLM_TeamTableCachedObj( + team_id=valid_token.team_id, + max_budget=valid_token.team_max_budget, + soft_budget=valid_token.team_soft_budget, + spend=valid_token.team_spend, + tpm_limit=valid_token.team_tpm_limit, + rpm_limit=valid_token.team_rpm_limit, + blocked=valid_token.team_blocked, + models=valid_token.team_models, + metadata=valid_token.team_metadata, + object_permission_id=valid_token.team_object_permission_id, + ) + else: + _team_obj = None + + _project_obj = None + if valid_token.project_id is not None: + _project_obj = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if _project_obj is not None: + valid_token.project_metadata = _project_obj.metadata + + _ = await common_checks( + request=request, + request_body=request_data, + team_object=_team_obj, + user_object=user_object, + end_user_object=end_user_object, + general_settings=general_settings, + global_proxy_spend=None, + route=route, + llm_router=llm_router, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + skip_budget_checks=False, + project_object=_project_obj, + ) + + return valid_token diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 40fae4e4a56..1269f58213a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -29,7 +29,7 @@ from litellm.constants import ( MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, STREAM_SSE_DATA_PREFIX, ) -from litellm.litellm_core_utils.dd_tracing import tracer +from litellm.litellm_core_utils.dd_tracing import set_active_span_tag, tracer from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, @@ -245,6 +245,26 @@ async def create_response( ) +def _add_dd_apm_tags_for_litellm_call_id(litellm_call_id: Optional[str]) -> None: + """ + Attach LiteLLM call id to the active Datadog APM span. + + This enables searching APM traces by LiteLLM call id returned in + `x-litellm-call-id`. + """ + if not litellm_call_id: + return + + try: + set_active_span_tag("litellm.call_id", str(litellm_call_id)) + except Exception: + # Tagging is best-effort and should never impact request processing. + verbose_proxy_logger.debug( + "Failed to tag active ddtrace span with litellm.call_id", + exc_info=True, + ) + + def _override_openai_response_model( *, response_obj: Any, @@ -642,6 +662,7 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_call_id"] = request.headers.get( "x-litellm-call-id", str(uuid.uuid4()) ) + _add_dd_apm_tags_for_litellm_call_id(self.data.get("litellm_call_id")) ### AUTO STREAM USAGE TRACKING ### # If always_include_stream_usage is enabled and this is a streaming request @@ -658,7 +679,6 @@ class ProxyBaseLLMRequestProcessing: and "include_usage" not in self.data["stream_options"] ): self.data["stream_options"]["include_usage"] = True - ### CALL HOOKS ### - modify/reject incoming data before calling the model ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 02fa84bae30..8c59c79ff0a 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -76,27 +76,29 @@ class SpendLogCleanup: "Max logs deleted - 1,00,000, rest of the logs will be deleted in next run" ) break - # Step 1: Find logs to delete - logs_to_delete = await prisma_client.db.litellm_spendlogs.find_many( - where={"startTime": {"lt": cutoff_date}}, - take=self.batch_size, + # Step 1: Find logs and delete them in one go without fetching to application + # Delete in batches, limited by self.batch_size + deleted_count = await prisma_client.db.execute_raw( + """ + DELETE FROM "LiteLLM_SpendLogs" + WHERE "request_id" IN ( + SELECT "request_id" FROM "LiteLLM_SpendLogs" + WHERE "startTime" < $1::timestamptz + LIMIT $2 + ) + """, + cutoff_date, + self.batch_size, ) - verbose_proxy_logger.info(f"Found {len(logs_to_delete)} logs in this batch") + verbose_proxy_logger.info(f"Deleted {deleted_count} logs in this batch") - if not logs_to_delete: + if deleted_count == 0: verbose_proxy_logger.info( f"No more logs to delete. Total deleted: {total_deleted}" ) break - request_ids = [log.request_id for log in logs_to_delete] - - # Step 2: Delete them in one go - await prisma_client.db.litellm_spendlogs.delete_many( - where={"request_id": {"in": request_ids}} - ) - - total_deleted += len(logs_to_delete) + total_deleted += deleted_count run_count += 1 # Add a small sleep to prevent overwhelming the database diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/results/block_insults_-_contentfilter_(denied_insults.yaml).json b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/results/insults_cf.json similarity index 100% rename from litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/results/block_insults_-_contentfilter_(denied_insults.yaml).json rename to litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/results/insults_cf.json diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/results/block_investment_-_contentfilter_(denied_financial_advice.yaml).json b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/results/investment_cf.json similarity index 100% rename from litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/results/block_investment_-_contentfilter_(denied_financial_advice.yaml).json rename to litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/results/investment_cf.json diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py index 01e820163fd..ca66b4da652 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/guardrail_benchmarks/test_eval.py @@ -18,6 +18,7 @@ Run a specific eval: import json import os +import re import time from datetime import datetime, timezone from typing import Any, Dict, List @@ -105,7 +106,25 @@ def _print_confusion_report(label: str, metrics: dict, wrong: list) -> None: def _save_confusion_results(label: str, metrics: dict, wrong: list, rows: list) -> dict: """Save confusion matrix results to a JSON file and return the result dict.""" os.makedirs(RESULTS_DIR, exist_ok=True) - safe_label = label.lower().replace(" ", "_").replace("—", "-") + # Build a short, filesystem-safe filename from the label. + # Full label is preserved inside the JSON; filename just needs to be + # unique and recognisable. Format: {topic}_{method_abbrev}.json + parts = label.split("\u2014") + topic = parts[0].strip().lower().replace("block ", "").replace(" ", "_") + method_full = parts[1].strip() if len(parts) > 1 else "" + method_name = re.sub(r"\s*\(.*?\)", "", method_full).strip().lower() + qualifier_match = re.search(r"\(([^)]+)\)", method_full) + qualifier = qualifier_match.group(1) if qualifier_match else "" + qualifier = re.sub(r"\.[a-z]+$", "", qualifier) # drop .yaml etc. + if method_name == "contentfilter": + safe_label = f"{topic}_cf" + elif qualifier: + safe_label = f"{topic}_{method_name}_{qualifier}" + else: + safe_label = f"{topic}_{method_name}" + safe_label = safe_label.replace(" ", "_") + safe_label = re.sub(r"[^a-z0-9_.\-]", "", safe_label) + safe_label = re.sub(r"_+", "_", safe_label).strip("_") result = { "label": label, "timestamp": datetime.now(timezone.utc).isoformat(), diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index d228bdb2129..a8d0e3e9af2 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -234,6 +234,14 @@ def _update_litellm_params_for_health_check( - for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID """ litellm_params["messages"] = _get_random_llm_message() + _health_check_max_tokens = model_info.get("health_check_max_tokens", None) + if _health_check_max_tokens is not None: + litellm_params["max_tokens"] = _health_check_max_tokens + elif "*" not in ( + model_info.get("health_check_model") or litellm_params.get("model") or "" + ): + litellm_params["max_tokens"] = 1 + _health_check_model = model_info.get("health_check_model", None) if _health_check_model is not None: litellm_params["model"] = _health_check_model @@ -321,7 +329,9 @@ async def perform_health_check( # Filter by model_id first so a single deployment is checked when id is specified if model_id is not None: - _by_id = [x for x in model_list if (x.get("model_info") or {}).get("id") == model_id] + _by_id = [ + x for x in model_list if (x.get("model_info") or {}).get("id") == model_id + ] if _by_id: model_list = _by_id elif model is not None: diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 4496ad92631..95b1836d8a9 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -33,6 +33,9 @@ from litellm.proxy.health_check import ( perform_health_check, run_with_timeout, ) +from litellm.proxy.middleware.in_flight_requests_middleware import ( + get_in_flight_requests, +) from litellm.secret_managers.main import get_secret #### Health ENDPOINTS #### @@ -1297,6 +1300,23 @@ async def health_readiness(): raise HTTPException(status_code=503, detail=f"Service Unhealthy ({str(e)})") +@router.get( + "/health/backlog", + tags=["health"], + dependencies=[Depends(user_api_key_auth)], +) +async def health_backlog(): + """ + Returns the number of HTTP requests currently in-flight on this uvicorn worker. + + Use this to measure per-pod queue depth. A high value means the worker is + processing many concurrent requests — requests arriving now will have to wait + for the event loop to get to them, adding latency before LiteLLM even starts + its own timer. + """ + return {"in_flight_requests": get_in_flight_requests()} + + @router.get( "/health/liveliness", # Historical LiteLLM name; doesn't match k8s terminology but kept for backwards compatibility tags=["health"], diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index b8c073dd061..5e48ef2879e 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -15,6 +15,7 @@ from litellm.types.utils import ( ) VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX = "virtual_key_spend" +END_USER_SPEND_CACHE_KEY_PREFIX = "end_user_model_spend" class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): @@ -83,6 +84,81 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return True + async def is_end_user_within_model_budget( + self, + end_user_id: str, + end_user_model_max_budget: dict, + model: str, + ) -> bool: + """ + Check if the end_user is within the model budget + + Raises: + BudgetExceededError: If the end_user has exceeded the model budget + """ + internal_model_max_budget: GenericBudgetConfigType = {} + + for _model, _budget_info in end_user_model_max_budget.items(): + internal_model_max_budget[_model] = BudgetConfig(**_budget_info) + + verbose_proxy_logger.debug( + "end_user internal_model_max_budget %s", + json.dumps(internal_model_max_budget, indent=4, default=str), + ) + + # check if current model is in internal_model_max_budget + _current_model_budget_info = self._get_request_model_budget_config( + model=model, internal_model_max_budget=internal_model_max_budget + ) + if _current_model_budget_info is None: + verbose_proxy_logger.debug( + f"Model {model} not found in end_user_model_max_budget" + ) + return True + + # check if current model is within budget + if ( + _current_model_budget_info.max_budget + and _current_model_budget_info.max_budget > 0 + ): + _current_spend = await self._get_end_user_spend_for_model( + end_user_id=end_user_id, + model=model, + key_budget_config=_current_model_budget_info, + ) + if ( + _current_spend is not None + and _current_model_budget_info.max_budget is not None + and _current_spend > _current_model_budget_info.max_budget + ): + raise litellm.BudgetExceededError( + message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", + current_cost=_current_spend, + max_budget=_current_model_budget_info.max_budget, + ) + + return True + + async def _get_end_user_spend_for_model( + self, + end_user_id: str, + model: str, + key_budget_config: BudgetConfig, + ) -> Optional[float]: + # 1. model: directly look up `model` + end_user_model_spend_cache_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" + _current_spend = await self.dual_cache.async_get_cache( + key=end_user_model_spend_cache_key, + ) + + if _current_spend is None: + # 2. If 1, does not exist, check if passed as {custom_llm_provider}/model + end_user_model_spend_cache_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{self._get_model_without_custom_llm_provider(model)}:{key_budget_config.budget_duration}" + _current_spend = await self.dual_cache.async_get_cache( + key=end_user_model_spend_cache_key, + ) + return _current_spend + async def _get_virtual_key_spend_for_model( self, user_api_key_hash: Optional[str], @@ -163,46 +239,77 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): user_api_key_model_max_budget: Optional[dict] = _metadata.get( "user_api_key_model_max_budget", None ) + user_api_key_end_user_model_max_budget: Optional[dict] = _metadata.get( + "user_api_key_end_user_model_max_budget", None + ) if ( user_api_key_model_max_budget is None or len(user_api_key_model_max_budget) == 0 + ) and ( + user_api_key_end_user_model_max_budget is None + or len(user_api_key_end_user_model_max_budget) == 0 ): verbose_proxy_logger.debug( - "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget is None or empty. `user_api_key_model_max_budget`=%s", - user_api_key_model_max_budget, + "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget and user_api_key_end_user_model_max_budget are None or empty." ) return + response_cost: float = standard_logging_payload.get("response_cost", 0) model = standard_logging_payload.get("model") virtual_key = standard_logging_payload.get("metadata", {}).get( "user_api_key_hash" ) + end_user_id = standard_logging_payload.get( + "end_user" + ) or standard_logging_payload.get("metadata", {}).get( + "user_api_key_end_user_id" + ) - if virtual_key is None or model is None: + if model is None: return - # Resolve per-model budget config (same logic as is_key_within_model_budget) - internal_model_max_budget: GenericBudgetConfigType = {} - for _model, _budget_info in user_api_key_model_max_budget.items(): - internal_model_max_budget[_model] = BudgetConfig(**_budget_info) - key_budget_config = self._get_request_model_budget_config( - model=model, internal_model_max_budget=internal_model_max_budget - ) - if key_budget_config is None or not key_budget_config.budget_duration: - verbose_proxy_logger.debug( - "Not incrementing model spend: no budget config or budget_duration for model=%s", - model, + if ( + virtual_key is not None + and user_api_key_model_max_budget is not None + and len(user_api_key_model_max_budget) > 0 + ): + internal_model_max_budget: GenericBudgetConfigType = {} + for _model, _budget_info in user_api_key_model_max_budget.items(): + internal_model_max_budget[_model] = BudgetConfig(**_budget_info) + key_budget_config = self._get_request_model_budget_config( + model=model, internal_model_max_budget=internal_model_max_budget ) - return + if key_budget_config is not None and key_budget_config.budget_duration: + virtual_spend_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}" + virtual_start_time_key = f"virtual_key_budget_start_time:{virtual_key}" + await self._increment_spend_for_key( + budget_config=key_budget_config, + spend_key=virtual_spend_key, + start_time_key=virtual_start_time_key, + response_cost=response_cost, + ) + + if ( + end_user_id is not None + and user_api_key_end_user_model_max_budget is not None + and len(user_api_key_end_user_model_max_budget) > 0 + ): + internal_model_max_budget: GenericBudgetConfigType = {} + for _model, _budget_info in user_api_key_end_user_model_max_budget.items(): + internal_model_max_budget[_model] = BudgetConfig(**_budget_info) + key_budget_config = self._get_request_model_budget_config( + model=model, internal_model_max_budget=internal_model_max_budget + ) + if key_budget_config is not None and key_budget_config.budget_duration: + end_user_spend_key = f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{key_budget_config.budget_duration}" + end_user_start_time_key = f"end_user_budget_start_time:{end_user_id}" + await self._increment_spend_for_key( + budget_config=key_budget_config, + spend_key=end_user_spend_key, + start_time_key=end_user_start_time_key, + response_cost=response_cost, + ) - virtual_spend_key = f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{key_budget_config.budget_duration}" - virtual_start_time_key = f"virtual_key_budget_start_time:{virtual_key}" - await self._increment_spend_for_key( - budget_config=key_budget_config, - spend_key=virtual_spend_key, - start_time_key=virtual_start_time_key, - response_cost=response_cost, - ) verbose_proxy_logger.debug( "current state of in memory cache %s", json.dumps( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d2312a00c3b..3168bcd812f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -248,13 +248,15 @@ def clean_headers( clean_headers = {} litellm_key_lower = ( litellm_key_header_name.lower() if litellm_key_header_name is not None else None - ) + ) for header, value in headers.items(): header_lower = header.lower() - + if header_lower == "authorization" and is_anthropic_oauth_key(value): clean_headers[header] = value - elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE: + elif ( + forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE + ): if litellm_key_lower and header_lower == litellm_key_lower: continue if header_lower == "authorization": @@ -840,11 +842,13 @@ async def add_litellm_data_to_request( # noqa: PLR0915 from litellm.types.proxy.litellm_pre_call_utils import SecretFields _raw_headers: Dict[str, str] = _safe_get_request_headers(request) - + forward_llm_auth = False if general_settings: - forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) - + forward_llm_auth = general_settings.get( + "forward_llm_provider_auth_headers", False + ) + _headers: Dict[str, str] = clean_headers( request.headers, litellm_key_header_name=( @@ -1019,6 +1023,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915 "spend_logs_metadata" ] + ## PROJECT-LEVEL TAGS + project_metadata = user_api_key_dict.project_metadata or {} + if "tags" in project_metadata and project_metadata["tags"] is not None: + data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=project_metadata["tags"], + ) + ## TEAM-LEVEL METADATA data = ( LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( @@ -1047,6 +1059,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data[_metadata_variable_name][ "user_api_key_model_max_budget" ] = user_api_key_dict.model_max_budget + data[_metadata_variable_name][ + "user_api_key_end_user_model_max_budget" + ] = user_api_key_dict.end_user_model_max_budget # User spend, budget - used by prometheus.py # Follow same pattern as team and API key budgets diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 12aa748bbc3..d58dca5aec0 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Set from fastapi import APIRouter, Depends, HTTPException, status @@ -94,6 +94,183 @@ async def _invalidate_cache_access_group(access_group_id: str) -> None: ) +# --------------------------------------------------------------------------- +# DB sync helpers (called inside a Prisma transaction) +# --------------------------------------------------------------------------- + + +async def _sync_add_access_group_to_teams( + tx, team_ids: List[str], access_group_id: str +) -> None: + """Add access_group_id to each team's access_group_ids (idempotent).""" + for team_id in team_ids: + team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) + if team is not None and access_group_id not in (team.access_group_ids or []): + await tx.litellm_teamtable.update( + where={"team_id": team_id}, + data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]}, + ) + + +async def _sync_remove_access_group_from_teams( + tx, team_ids: List[str], access_group_id: str +) -> None: + """Remove access_group_id from each team's access_group_ids (idempotent).""" + for team_id in team_ids: + team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) + if team is not None and access_group_id in (team.access_group_ids or []): + await tx.litellm_teamtable.update( + where={"team_id": team_id}, + data={"access_group_ids": [ag for ag in team.access_group_ids if ag != access_group_id]}, + ) + + +async def _sync_add_access_group_to_keys( + tx, key_tokens: List[str], access_group_id: str +) -> None: + """Add access_group_id to each key's access_group_ids (idempotent).""" + for token in key_tokens: + key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) + if key is not None and access_group_id not in (key.access_group_ids or []): + await tx.litellm_verificationtoken.update( + where={"token": token}, + data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]}, + ) + + +async def _sync_remove_access_group_from_keys( + tx, key_tokens: List[str], access_group_id: str +) -> None: + """Remove access_group_id from each key's access_group_ids (idempotent).""" + for token in key_tokens: + key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) + if key is not None and access_group_id in (key.access_group_ids or []): + await tx.litellm_verificationtoken.update( + where={"token": token}, + data={"access_group_ids": [ag for ag in key.access_group_ids if ag != access_group_id]}, + ) + + +# --------------------------------------------------------------------------- +# Cache patch helpers +# --------------------------------------------------------------------------- + + +async def _patch_team_caches_add_access_group( + team_ids: List[str], + access_group_id: str, + user_api_key_cache, + proxy_logging_obj, +) -> None: + """Patch cached team objects to include access_group_id.""" + for team_id in team_ids: + cached_team = await _get_team_object_from_cache( + key="team_id:{}".format(team_id), + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + ) + if cached_team is None: + continue + if cached_team.access_group_ids is None: + cached_team.access_group_ids = [access_group_id] + elif access_group_id not in cached_team.access_group_ids: + cached_team.access_group_ids = list(cached_team.access_group_ids) + [access_group_id] + else: + continue + await _cache_team_object( + team_id=team_id, + team_table=cached_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _patch_team_caches_remove_access_group( + team_ids: List[str], + access_group_id: str, + user_api_key_cache, + proxy_logging_obj, +) -> None: + """Patch cached team objects to remove access_group_id.""" + for team_id in team_ids: + cached_team = await _get_team_object_from_cache( + key="team_id:{}".format(team_id), + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + ) + if cached_team is not None and cached_team.access_group_ids: + cached_team.access_group_ids = [ + ag for ag in cached_team.access_group_ids if ag != access_group_id + ] + await _cache_team_object( + team_id=team_id, + team_table=cached_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _patch_key_caches_add_access_group( + key_tokens: List[str], + access_group_id: str, + user_api_key_cache, + proxy_logging_obj, +) -> None: + """Patch cached key objects to include access_group_id.""" + for token in key_tokens: + cached_key = await user_api_key_cache.async_get_cache(key=token) + if cached_key is None: + continue + if isinstance(cached_key, dict): + cached_key = UserAPIKeyAuth(**cached_key) + if not isinstance(cached_key, UserAPIKeyAuth): + continue + if cached_key.access_group_ids is None: + cached_key.access_group_ids = [access_group_id] + elif access_group_id not in cached_key.access_group_ids: + cached_key.access_group_ids = list(cached_key.access_group_ids) + [access_group_id] + else: + continue + await _cache_key_object( + hashed_token=token, + user_api_key_obj=cached_key, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _patch_key_caches_remove_access_group( + key_tokens: List[str], + access_group_id: str, + user_api_key_cache, + proxy_logging_obj, +) -> None: + """Patch cached key objects to remove access_group_id.""" + for token in key_tokens: + cached_key = await user_api_key_cache.async_get_cache(key=token) + if cached_key is None: + continue + if isinstance(cached_key, dict): + cached_key = UserAPIKeyAuth(**cached_key) + if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids: + cached_key.access_group_ids = [ + ag for ag in cached_key.access_group_ids if ag != access_group_id + ] + await _cache_key_object( + hashed_token=token, + user_api_key_obj=cached_key, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +# --------------------------------------------------------------------------- +# CRUD endpoints +# --------------------------------------------------------------------------- + + @router.post( "/v1/access_group", response_model=AccessGroupResponse, @@ -106,32 +283,42 @@ async def create_access_group( _require_proxy_admin(user_api_key_dict) prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - existing = await prisma_client.db.litellm_accessgrouptable.find_unique( - where={"access_group_name": data.access_group_name} - ) - if existing is not None: - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail=f"Access group '{data.access_group_name}' already exists", - ) - try: - record = await prisma_client.db.litellm_accessgrouptable.create( - data={ - "access_group_name": data.access_group_name, - "description": data.description, - "access_model_names": data.access_model_names or [], - "access_mcp_server_ids": data.access_mcp_server_ids or [], - "access_agent_ids": data.access_agent_ids or [], - "assigned_team_ids": data.assigned_team_ids or [], - "assigned_key_ids": data.assigned_key_ids or [], - "created_by": user_api_key_dict.user_id, - "updated_by": user_api_key_dict.user_id, - } - ) + async with prisma_client.db.tx() as tx: + existing = await tx.litellm_accessgrouptable.find_unique( + where={"access_group_name": data.access_group_name} + ) + if existing is not None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Access group '{data.access_group_name}' already exists", + ) + + record = await tx.litellm_accessgrouptable.create( + data={ + "access_group_name": data.access_group_name, + "description": data.description, + "access_model_names": data.access_model_names or [], + "access_mcp_server_ids": data.access_mcp_server_ids or [], + "access_agent_ids": data.access_agent_ids or [], + "assigned_team_ids": data.assigned_team_ids or [], + "assigned_key_ids": data.assigned_key_ids or [], + "created_by": user_api_key_dict.user_id, + "updated_by": user_api_key_dict.user_id, + } + ) + + # Sync team and key tables to reference the new access group + await _sync_add_access_group_to_teams( + tx, data.assigned_team_ids or [], record.access_group_id + ) + await _sync_add_access_group_to_keys( + tx, data.assigned_key_ids or [], record.access_group_id + ) + except HTTPException: + raise except Exception as e: # Race condition: another request created the same name between find_unique and create. - # Prisma raises UniqueViolationError (P2002) or similar for unique constraint. if "unique constraint" in str(e).lower() or "P2002" in str(e): raise HTTPException( status_code=status.HTTP_409_CONFLICT, @@ -139,8 +326,15 @@ async def create_access_group( ) raise - # Cache the newly created access group for read-heavy access patterns + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + await _cache_access_group_record(record) + await _patch_team_caches_add_access_group( + data.assigned_team_ids or [], record.access_group_id, user_api_key_cache, proxy_logging_obj + ) + await _patch_key_caches_add_access_group( + data.assigned_key_ids or [], record.access_group_id, user_api_key_cache, proxy_logging_obj + ) return _record_to_response(record) @@ -195,24 +389,54 @@ async def update_access_group( _require_proxy_admin(user_api_key_dict) prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - existing = await prisma_client.db.litellm_accessgrouptable.find_unique( - where={"access_group_id": access_group_id} - ) - if existing is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"Access group '{access_group_id}' not found", - ) - + update_fields = data.model_dump(exclude_unset=True) update_data: dict = {"updated_by": user_api_key_dict.user_id} - for field, value in data.model_dump(exclude_unset=True).items(): + for field, value in update_fields.items(): + if field in ("assigned_team_ids", "assigned_key_ids", "access_model_names", "access_mcp_server_ids", "access_agent_ids") and value is None: + value = [] update_data[field] = value + # Initialize delta lists before the try block so they remain accessible + # for cache updates after the transaction, even if an error path is added later. + teams_to_add: List[str] = [] + teams_to_remove: List[str] = [] + keys_to_add: List[str] = [] + keys_to_remove: List[str] = [] + try: - record = await prisma_client.db.litellm_accessgrouptable.update( - where={"access_group_id": access_group_id}, - data=update_data, - ) + async with prisma_client.db.tx() as tx: + # Read inside the transaction so delta computation is consistent with the write, + # avoiding a TOCTOU race where a concurrent update could make deltas stale. + existing = await tx.litellm_accessgrouptable.find_unique( + where={"access_group_id": access_group_id} + ) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Access group '{access_group_id}' not found", + ) + + old_team_ids: Set[str] = set(existing.assigned_team_ids or []) + old_key_ids: Set[str] = set(existing.assigned_key_ids or []) + new_team_ids: Set[str] = set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids + new_key_ids: Set[str] = set(update_fields["assigned_key_ids"] or []) if "assigned_key_ids" in update_fields else old_key_ids + + teams_to_add = list(new_team_ids - old_team_ids) + teams_to_remove = list(old_team_ids - new_team_ids) + keys_to_add = list(new_key_ids - old_key_ids) + keys_to_remove = list(old_key_ids - new_key_ids) + + record = await tx.litellm_accessgrouptable.update( + where={"access_group_id": access_group_id}, + data=update_data, + ) + + await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id) + await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id) + await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id) + await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id) + except HTTPException: + raise except Exception as e: # Unique constraint violation (e.g. access_group_name already exists). if "unique constraint" in str(e).lower() or "P2002" in str(e): @@ -222,8 +446,13 @@ async def update_access_group( ) raise - # Write the updated record into cache (same key, overwrites stale entry) + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + await _cache_access_group_record(record) + await _patch_team_caches_add_access_group(teams_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) + await _patch_team_caches_remove_access_group(teams_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj) + await _patch_key_caches_add_access_group(keys_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) + await _patch_key_caches_remove_access_group(keys_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj) return _record_to_response(record) @@ -240,9 +469,8 @@ async def delete_access_group( prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) try: - # Track affected team IDs and key tokens for cache invalidation - affected_team_ids: list = [] - affected_key_tokens: list = [] + affected_team_ids: List[str] = [] + affected_key_tokens: List[str] = [] async with prisma_client.db.tx() as tx: existing = await tx.litellm_accessgrouptable.find_unique( @@ -254,73 +482,61 @@ async def delete_access_group( detail=f"Access group '{access_group_id}' not found", ) - # Remove access_group_id from teams and keys that reference it + # Union of: teams that have this access_group_id in their own access_group_ids + # AND teams listed in assigned_team_ids (handles out-of-sync data from before this sync was added) teams_with_group = await tx.litellm_teamtable.find_many( where={"access_group_ids": {"hasSome": [access_group_id]}} ) - for team in teams_with_group: - affected_team_ids.append(team.team_id) - updated_ids = [tid for tid in (team.access_group_ids or []) if tid != access_group_id] - await tx.litellm_teamtable.update( - where={"team_id": team.team_id}, - data={"access_group_ids": updated_ids}, - ) + all_affected_team_ids: Set[str] = ( + {team.team_id for team in teams_with_group} + | set(existing.assigned_team_ids or []) + ) + affected_team_ids = list(all_affected_team_ids) + # Union of: keys that have this access_group_id in their own access_group_ids + # AND keys listed in assigned_key_ids (handles out-of-sync data) keys_with_group = await tx.litellm_verificationtoken.find_many( where={"access_group_ids": {"hasSome": [access_group_id]}} ) + all_affected_key_tokens: Set[str] = ( + {key.token for key in keys_with_group} + | set(existing.assigned_key_ids or []) + ) + affected_key_tokens = list(all_affected_key_tokens) + + # Update teams returned by find_many directly — we already have their data. + for team in teams_with_group: + await tx.litellm_teamtable.update( + where={"team_id": team.team_id}, + data={"access_group_ids": [ag for ag in (team.access_group_ids or []) if ag != access_group_id]}, + ) + # Use _sync_remove only for out-of-sync teams not found by the hasSome query. + out_of_sync_team_ids = set(existing.assigned_team_ids or []) - {t.team_id for t in teams_with_group} + await _sync_remove_access_group_from_teams(tx, list(out_of_sync_team_ids), access_group_id) + + # Update keys returned by find_many directly — we already have their data. for key in keys_with_group: - affected_key_tokens.append(key.token) - updated_ids = [kid for kid in (key.access_group_ids or []) if kid != access_group_id] await tx.litellm_verificationtoken.update( where={"token": key.token}, - data={"access_group_ids": updated_ids}, + data={"access_group_ids": [ag for ag in (key.access_group_ids or []) if ag != access_group_id]}, ) + # Use _sync_remove only for out-of-sync keys not found by the hasSome query. + out_of_sync_key_tokens = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group} + await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id) await tx.litellm_accessgrouptable.delete( where={"access_group_id": access_group_id} ) - # Invalidate the deleted access group from cache - await _invalidate_cache_access_group(access_group_id) - - # Patch cached team and key objects to remove the deleted access_group_id - # instead of fully invalidating them (keeps cache warm, avoids DB re-fetch) from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache - for team_id in affected_team_ids: - cached_team = await _get_team_object_from_cache( - key="team_id:{}".format(team_id), - proxy_logging_obj=proxy_logging_obj, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - ) - if cached_team is not None and cached_team.access_group_ids: - cached_team.access_group_ids = [ - ag_id for ag_id in cached_team.access_group_ids if ag_id != access_group_id - ] - await _cache_team_object( - team_id=team_id, - team_table=cached_team, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - for token in affected_key_tokens: - cached_key = await user_api_key_cache.async_get_cache(key=token) - if cached_key is not None: - if isinstance(cached_key, dict): - cached_key = UserAPIKeyAuth(**cached_key) - if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids: - cached_key.access_group_ids = [ - ag_id for ag_id in cached_key.access_group_ids if ag_id != access_group_id - ] - await _cache_key_object( - hashed_token=token, - user_api_key_obj=cached_key, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + await _invalidate_cache_access_group(access_group_id) + await _patch_team_caches_remove_access_group( + affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj + ) + await _patch_key_caches_remove_access_group( + affected_key_tokens, access_group_id, user_api_key_cache, proxy_logging_obj + ) except HTTPException: raise diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e535ccaaa46..5a0a05114a3 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -614,7 +614,7 @@ async def user_info( user_id is None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ): - return await _get_user_info_for_proxy_admin() + return await _get_user_info_for_proxy_admin(user_api_key_dict=user_api_key_dict) elif user_id is None: user_id = user_api_key_dict.user_id ## GET USER ROW ## @@ -714,7 +714,7 @@ async def user_info( raise handle_exception_on_proxy(e) -async def _get_user_info_for_proxy_admin(): +async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): """ Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying @@ -754,9 +754,23 @@ async def _get_user_info_for_proxy_admin(): _teams_in_db = [LiteLLM_TeamTable(**team) for team in _teams_in_db] _teams_in_db.sort(key=lambda x: (getattr(x, "team_alias", "") or "")) returned_keys = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) + + # Get admin's own user_id and user_info + admin_user_id = user_api_key_dict.user_id + admin_user_info = None + + if admin_user_id is not None: + admin_user_info = await prisma_client.get_data(user_id=admin_user_id) + if admin_user_info is not None: + admin_user_info = ( + admin_user_info.model_dump() + if isinstance(admin_user_info, BaseModel) + else admin_user_info + ) + return UserInfoResponse( - user_id=None, - user_info=None, + user_id=admin_user_id, + user_info=admin_user_info, keys=returned_keys, teams=_teams_in_db, ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b489369071f..9414ce6f686 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3983,6 +3983,10 @@ async def list_keys( status: Optional[str] = Query( None, description="Filter by status (e.g. 'deleted')" ), + project_id: Optional[str] = Query(None, description="Filter keys by project ID"), + access_group_id: Optional[str] = Query( + None, description="Filter keys by access group ID" + ), ) -> KeyListResponseObject: """ List all keys for a given user / team / organization. @@ -4076,6 +4080,8 @@ async def list_keys( sort_order=sort_order, expand=expand, status=status, + project_id=project_id, + access_group_id=access_group_id, ) verbose_proxy_logger.debug("Successfully prepared response") @@ -4109,13 +4115,23 @@ async def list_keys( dependencies=[Depends(user_api_key_auth)], ) @management_endpoint_wrapper -async def key_aliases() -> Dict[str, List[str]]: +async def key_aliases( + page: int = Query(1, ge=1, description="Page number"), + size: int = Query(50, ge=1, le=100, description="Page size"), + search: Optional[str] = Query( + None, description="Search key aliases (case-insensitive partial match)" + ), +) -> Dict[str, Any]: """ - Lists all key aliases + Lists key aliases with pagination and optional search. Returns: { - "aliases": List[str] + "aliases": List[str], + "total_count": int, + "current_page": int, + "total_pages": int, + "size": int, } """ try: @@ -4127,36 +4143,55 @@ async def key_aliases() -> Dict[str, List[str]]: verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") - where: Dict[str, Any] = {} - try: - where.update(_get_condition_to_filter_out_ui_session_tokens()) - except NameError: - # Helper may not exist in some builds; ignore if missing - pass + # Build a parameterized WHERE clause to avoid loading full rows into + # memory. Raw SQL is used because the Prisma client wrapper does not + # support column-level SELECT projection on find_many. + # + # $1 is always UI_SESSION_TOKEN_TEAM_ID (filters out UI session tokens). + query_params: List[Any] = [UI_SESSION_TOKEN_TEAM_ID] + where_parts = [ + "key_alias IS NOT NULL", + "key_alias != ''", + "(team_id IS NULL OR team_id != $1)", + ] + if search: + query_params.append(f"%{search}%") + where_parts.append(f"key_alias ILIKE ${len(query_params)}") - rows = await prisma_client.db.litellm_verificationtoken.find_many( - where=where, - order=[{"key_alias": "asc"}], + where_sql = " AND ".join(where_parts) + + count_sql = ( + f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' + ) + count_rows = await prisma_client.db.query_raw(count_sql, *query_params) + total_count = int(count_rows[0]["count"]) if count_rows else 0 + + aliases_params = query_params + [size, (page - 1) * size] + limit_idx = len(aliases_params) - 1 + offset_idx = len(aliases_params) + aliases_sql = ( + f"SELECT key_alias" + f' FROM "LiteLLM_VerificationToken"' + f" WHERE {where_sql}" + f" ORDER BY key_alias ASC" + f" LIMIT ${limit_idx} OFFSET ${offset_idx}" + ) + alias_rows = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + aliases: List[str] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] + + total_pages = -(-total_count // size) if total_count > 0 else 0 + verbose_proxy_logger.debug( + f"key_aliases: page={page}, size={size}, search={search!r}, " + f"total_count={total_count}, total_pages={total_pages}" ) - seen = set() - aliases: List[str] = [] - for row in rows: - alias = getattr(row, "key_alias", None) - if alias is None and isinstance(row, dict): - alias = row.get("key_alias") - - if not alias: - continue - - alias_str = str(alias).strip() - if alias_str and alias_str not in seen: - seen.add(alias_str) - aliases.append(alias_str) - - verbose_proxy_logger.debug(f"Returning {len(aliases)} key aliases") - - return {"aliases": aliases} + return { + "aliases": aliases, + "total_count": total_count, + "current_page": page, + "total_pages": total_pages, + "size": size, + } except Exception as e: verbose_proxy_logger.exception(f"Error in key_aliases: {e}") @@ -4223,6 +4258,8 @@ def _build_key_filter_conditions( admin_team_ids: Optional[List[str]], member_team_ids: Optional[List[str]] = None, include_created_by_keys: bool = False, + project_id: Optional[str] = None, + access_group_id: Optional[str] = None, ) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]: """Build filter conditions for key listing. @@ -4314,6 +4351,13 @@ def _build_key_filter_conditions( elif len(or_conditions) == 1: where.update(or_conditions[0]) + # Apply project_id and access_group_id as global AND filters so they + # narrow results across all visibility conditions (own keys, team keys, etc.) + if project_id: + where = {"AND": [where, {"project_id": project_id}]} + if access_group_id: + where = {"AND": [where, {"access_group_ids": {"hasSome": [access_group_id]}}]} + verbose_proxy_logger.debug(f"Filter conditions: {where}") return where @@ -4340,6 +4384,8 @@ async def _list_key_helper( sort_order: str = "desc", expand: Optional[List[str]] = None, status: Optional[str] = None, + project_id: Optional[str] = None, + access_group_id: Optional[str] = None, ) -> KeyListResponseObject: """ Helper function to list keys @@ -4373,6 +4419,8 @@ async def _list_key_helper( admin_team_ids=admin_team_ids, member_team_ids=member_team_ids, include_created_by_keys=include_created_by_keys, + project_id=project_id, + access_group_id=access_group_id, ) # Calculate skip for pagination diff --git a/litellm/proxy/management_endpoints/project_endpoints.py b/litellm/proxy/management_endpoints/project_endpoints.py index ba3238ebfd5..8f48f9def78 100644 --- a/litellm/proxy/management_endpoints/project_endpoints.py +++ b/litellm/proxy/management_endpoints/project_endpoints.py @@ -284,6 +284,7 @@ async def new_project( - model_tpm_limit: *Optional[dict]* - TPM limits per model. Example: {"gpt-4": 50000, "gpt-3.5-turbo": 100000} - budget_duration: *Optional[str]* - Frequency of reseting project budget - metadata: *Optional[dict]* - Metadata for project, store information for project. Example metadata - {"use_case_id": "SNOW-12345", "responsible_ai_id": "RAI-67890"} + - tags: *Optional[list]* - Tags for the project. Example: ["production", "api"] - blocked: *bool* - Flag indicating if the project is blocked or not - will stop all calls from keys with this project_id. - object_permission: Optional[LiteLLM_ObjectPermissionBase] - project-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission. @@ -339,6 +340,15 @@ async def new_project( ) try: + if getattr(data, "tags", None) is not None and not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Only premium users can add tags to projects. " + + CommonProxyErrors.not_premium_user.value + }, + ) + if not premium_user: raise HTTPException( status_code=403, @@ -348,6 +358,16 @@ async def new_project( }, ) + # ADD METADATA FIELDS + for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=data, + field_name=field, + value=getattr(data, field), + ) + delattr(data, field) + if prisma_client is None: raise HTTPException( status_code=500, @@ -463,7 +483,7 @@ async def new_project( response_model=LiteLLM_ProjectTable, ) @management_endpoint_wrapper -async def update_project( +async def update_project( # noqa: PLR0915 data: UpdateProjectRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -485,6 +505,7 @@ async def update_project( - model_rpm_limit: *Optional[dict]* - Updated RPM limits per model - model_tpm_limit: *Optional[dict]* - Updated TPM limits per model - budget_duration: *Optional[str]* - Updated budget duration + - tags: *Optional[list]* - Updated list of tags for the project - object_permission: Optional[LiteLLM_ObjectPermissionBase] - Updated object permission Example: @@ -514,6 +535,15 @@ async def update_project( ) try: + if getattr(data, "tags", None) is not None and not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Only premium users can add tags to projects. " + + CommonProxyErrors.not_premium_user.value + }, + ) + if not premium_user: raise HTTPException( status_code=403, @@ -523,6 +553,16 @@ async def update_project( }, ) + # ADD METADATA FIELDS + for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=data, + field_name=field, + value=getattr(data, field), + ) + delattr(data, field) + if prisma_client is None: raise HTTPException( status_code=500, diff --git a/litellm/proxy/middleware/in_flight_requests_middleware.py b/litellm/proxy/middleware/in_flight_requests_middleware.py new file mode 100644 index 00000000000..d615640d870 --- /dev/null +++ b/litellm/proxy/middleware/in_flight_requests_middleware.py @@ -0,0 +1,81 @@ +""" +Tracks the number of HTTP requests currently in-flight on this uvicorn worker. + +Used by /health/backlog to expose per-pod queue depth, and emitted as the +Prometheus gauge `litellm_in_flight_requests`. +""" + +import os +from typing import Optional + +from starlette.types import ASGIApp, Receive, Scope, Send + + +class InFlightRequestsMiddleware: + """ + ASGI middleware that increments a counter when a request arrives and + decrements it when the response is sent (or an error occurs). + + The counter is class-level and therefore scoped to a single uvicorn worker + process — exactly the per-pod granularity we want. + + Also updates the `litellm_in_flight_requests` Prometheus gauge if + prometheus_client is installed. The gauge is lazily initialised on the + first request so that PROMETHEUS_MULTIPROC_DIR is already set by the time + we register the metric. Initialisation is attempted only once — if + prometheus_client is absent the class remembers and never retries. + """ + + _in_flight: int = 0 + _gauge: Optional[object] = None + _gauge_init_attempted: bool = False + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + InFlightRequestsMiddleware._in_flight += 1 + gauge = InFlightRequestsMiddleware._get_gauge() + if gauge is not None: + gauge.inc() # type: ignore[union-attr] + try: + await self.app(scope, receive, send) + finally: + InFlightRequestsMiddleware._in_flight -= 1 + if gauge is not None: + gauge.dec() # type: ignore[union-attr] + + @staticmethod + def get_count() -> int: + """Return the number of HTTP requests currently in-flight.""" + return InFlightRequestsMiddleware._in_flight + + @staticmethod + def _get_gauge() -> Optional[object]: + if InFlightRequestsMiddleware._gauge_init_attempted: + return InFlightRequestsMiddleware._gauge + InFlightRequestsMiddleware._gauge_init_attempted = True + try: + from prometheus_client import Gauge + + kwargs = {} + if "PROMETHEUS_MULTIPROC_DIR" in os.environ: + # livesum aggregates across all worker processes in the scrape response + kwargs["multiprocess_mode"] = "livesum" + InFlightRequestsMiddleware._gauge = Gauge( + "litellm_in_flight_requests", + "Number of HTTP requests currently in-flight on this uvicorn worker", + **kwargs, + ) + except Exception: + InFlightRequestsMiddleware._gauge = None + return InFlightRequestsMiddleware._gauge + + +def get_in_flight_requests() -> int: + """Module-level convenience wrapper used by the /health/backlog endpoint.""" + return InFlightRequestsMiddleware.get_count() diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index c1092a06b48..4f31c762df1 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -1,9 +1,14 @@ #### OCR Endpoints ##### +import json +from typing import Any, Dict, Optional, cast + import orjson -from fastapi import APIRouter, Depends, Request, Response +from fastapi import APIRouter, Depends, Request, Response, UploadFile from fastapi.responses import ORJSONResponse +from litellm._logging import verbose_proxy_logger +from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -11,6 +16,171 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin router = APIRouter() +def _build_document_from_upload( + file_content: bytes, + filename: Optional[str], + content_type: Optional[str], +) -> Dict[str, str]: + """ + Convert uploaded file bytes into a Mistral-format document dict with base64 data URI. + + Delegates to convert_file_document_to_url_document after resolving MIME type + from the upload's content_type header or filename. + """ + mime_type = content_type.split(";")[0].strip() if content_type else None + if not mime_type or mime_type == "application/octet-stream": + if filename: + mime_type = get_mime_type(filename) + + return convert_file_document_to_url_document( + { + "type": "file", + "file": file_content, + "mime_type": mime_type or "application/octet-stream", + } + ) + + +async def _parse_multipart_form(request: Request) -> Dict[str, Any]: + """ + Extract OCR data from a multipart form request. + + Uses the cached form if already parsed by auth middleware, + otherwise parses the form from the request. + + Returns: + A dict with 'document', 'model', and any other OCR params. + """ + try: + form = await request.form() + except Exception as e: + raise ValueError( + f"Failed to parse multipart form data: {str(e)}. " + "When using curl with --form/-F, do NOT set the Content-Type header " + "manually — curl will set it automatically with the required boundary." + ) + + uploaded_file = form.get("file") + # request.form() may return either a FastAPI or Starlette UploadFile + # depending on middleware; check both via isinstance (FastAPI's UploadFile + # is a subclass of Starlette's) and fall back to duck-type check. + if uploaded_file is None or ( + not isinstance(uploaded_file, UploadFile) and not hasattr(uploaded_file, "read") + ): + raise ValueError( + "Multipart OCR request must include a 'file' field with the document to process" + ) + + uploaded_file = cast(UploadFile, uploaded_file) + + # Seek to start in case the file was already partially read by middleware + await uploaded_file.seek(0) + file_content = await uploaded_file.read() + if not file_content: + raise ValueError("Uploaded file is empty") + + document = _build_document_from_upload( + file_content=file_content, + filename=uploaded_file.filename, + content_type=uploaded_file.content_type, + ) + + data: Dict[str, Any] = {"document": document} + + for field_name, field_value in form.items(): + if field_name in ("file", "document"): + continue + # Try to parse JSON values (e.g. pages=[0,1,2]) + if isinstance(field_value, str): + try: + data[field_name] = json.loads(field_value) + except (json.JSONDecodeError, ValueError): + data[field_name] = field_value + else: + data[field_name] = field_value + + verbose_proxy_logger.debug( + f"OCR multipart form request parsed - model: {data.get('model')}, " + f"document_type: {document['type']}, " + f"filename: {uploaded_file.filename}" + ) + + return data + + +async def _parse_ocr_request(request: Request) -> Dict[str, Any]: + """ + Parse an OCR request, supporting both JSON and multipart form data. + + JSON body (existing behavior): + { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://..."} + } + + Multipart form data (new): + - file: the uploaded file + - model: model name (form field) + - Any other OCR params as form fields (pages, include_image_base64, etc.) + + Returns: + A dict suitable for passing to the OCR processing pipeline. + """ + content_type = request.headers.get("content-type", "") + + if "multipart/form-data" in content_type.lower(): + return await _parse_multipart_form(request) + + # --- JSON body (existing behavior) --- + try: + body = await request.body() + except RuntimeError: + # Body stream was consumed by auth middleware (e.g., form parsing). + body = b"" + + if not body: + # The body may be empty because the auth middleware already parsed + # it as form data (e.g., _read_request_body called request.form()). + # Check if form data is available. + if getattr(request, "_form", None) is not None: + verbose_proxy_logger.debug( + "OCR request body is empty but form data is available from middleware — " + "processing as multipart form." + ) + return await _parse_multipart_form(request) + + raise ValueError( + "Empty request body. For file uploads, use multipart/form-data content type " + "with a file field. When using curl with --form/-F, do NOT set the Content-Type " + "header manually." + ) + + try: + data = orjson.loads(body) + except orjson.JSONDecodeError as e: + raise ValueError( + f"Invalid JSON in request body: {e}. " + "Ensure the request body is valid JSON with Content-Type: application/json, " + "or use multipart/form-data for file uploads." + ) + + # Security: reject type="file" documents received via JSON. + # The "file" document type is designed for local SDK usage where the + # caller and the process share a filesystem. In the proxy context the + # caller is remote, so allowing a file-path string would let an + # authenticated user read arbitrary files from the server's filesystem. + # File uploads must go through multipart/form-data instead. + doc = data.get("document") if isinstance(data, dict) else None + if isinstance(doc, dict) and doc.get("type") == "file": + raise ValueError( + "document type 'file' is not supported through the JSON API. " + "To upload a local file, use multipart/form-data with a 'file' field. " + "For JSON requests, use 'document_url' or 'image_url' document types." + ) + + return data + + @router.post( "/v1/ocr", dependencies=[Depends(user_api_key_auth)], @@ -30,23 +200,30 @@ async def ocr( ): """ OCR endpoint for extracting text from documents and images. - - Follows the Mistral OCR API spec: - https://docs.mistral.ai/capabilities/vision/#optical-character-recognition-ocr - - Example: + + Supports two input modes: + + **1. JSON body** (Mistral OCR API compatible): ```bash curl -X POST "http://localhost:4000/v1/ocr" \ -H "Authorization: Bearer sk-1234" \ -H "Content-Type: application/json" \ -d '{ - "model": "mistral/mistral-ocr-latest", + "model": "mistral-ocr", "document": { "type": "document_url", "document_url": "https://arxiv.org/pdf/2201.04234" } }' ``` + + **2. Multipart form file upload**: + ```bash + curl -X POST "http://localhost:4000/v1/ocr" \ + -H "Authorization: Bearer sk-1234" \ + -F "model=mistral-ocr" \ + -F "file=@document.pdf" + ``` """ from litellm.proxy.proxy_server import ( general_settings, @@ -62,13 +239,14 @@ async def ocr( version, ) - # Read request body - body = await request.body() - data = orjson.loads(body) - - # Process request using ProxyBaseLLMRequestProcessing - processor = ProxyBaseLLMRequestProcessing(data=data) + data: dict = {} try: + # Parse request body (JSON or multipart form) + data = await _parse_ocr_request(request) + + # Process request using ProxyBaseLLMRequestProcessing + processor = ProxyBaseLLMRequestProcessing(data=data) + return await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, @@ -88,10 +266,10 @@ async def ocr( version=version, ) except Exception as e: + processor = ProxyBaseLLMRequestProcessing(data=data) raise await processor._handle_llm_api_exception( e=e, user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj, version=version, ) - diff --git a/litellm/proxy/prometheus_cleanup.py b/litellm/proxy/prometheus_cleanup.py new file mode 100644 index 00000000000..6353588532a --- /dev/null +++ b/litellm/proxy/prometheus_cleanup.py @@ -0,0 +1,47 @@ +""" +Prometheus multiprocess directory cleanup utilities. + +Wipes all .db files on startup so workers start with a clean slate. +""" + +from __future__ import annotations + +import glob +import os + +from litellm._logging import verbose_proxy_logger + + +def wipe_directory(directory: str) -> None: + """Delete all .db files in the directory. Called once before workers fork.""" + files = glob.glob(os.path.join(directory, "*.db")) + deleted = 0 + for filepath in files: + try: + os.remove(filepath) + deleted += 1 + except OSError as e: + verbose_proxy_logger.warning( + f"Failed to delete stale prometheus file {filepath}: {e}" + ) + if deleted: + verbose_proxy_logger.info( + f"Prometheus cleanup: wiped {deleted} stale .db files from {directory}" + ) + + +def mark_worker_exit(worker_pid: int) -> None: + """Remove prometheus .db files for a dead worker. Called by gunicorn child_exit hook.""" + if not os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + return + try: + from prometheus_client import multiprocess + + multiprocess.mark_process_dead(worker_pid) + verbose_proxy_logger.info( + f"Prometheus cleanup: marked worker {worker_pid} as dead" + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to mark prometheus worker {worker_pid} as dead: {e}" + ) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e91447af895..921d86c35c1 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -277,6 +277,15 @@ class ProxyInitializationHelpers: if max_requests_before_restart is not None: gunicorn_options["max_requests"] = max_requests_before_restart + # Clean up prometheus .db files when a worker exits (prevents ghost gauge values) + if os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + from litellm.proxy.prometheus_cleanup import mark_worker_exit + + def child_exit(server, worker): + mark_worker_exit(worker.pid) + + gunicorn_options["child_exit"] = child_exit + if ssl_certfile_path is not None and ssl_keyfile_path is not None: print( # noqa f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" # noqa @@ -314,6 +323,47 @@ class ProxyInitializationHelpers: return None # Let uvicorn choose the default loop on Windows return "uvloop" + @staticmethod + def _maybe_setup_prometheus_multiproc_dir( + num_workers: int, + litellm_settings: Optional[dict], + ) -> None: + """ + Auto-create PROMETHEUS_MULTIPROC_DIR when running with multiple workers + and prometheus is configured as a callback. + """ + import tempfile + + if num_workers <= 1 or litellm_settings is None: + return + + # Check if prometheus is in any callback list + callbacks = litellm_settings.get("callbacks") or [] + success_callbacks = litellm_settings.get("success_callback") or [] + failure_callbacks = litellm_settings.get("failure_callback") or [] + all_callbacks = callbacks + success_callbacks + failure_callbacks + if "prometheus" not in all_callbacks: + return + + from litellm.proxy.prometheus_cleanup import wipe_directory + + multiproc_dir = ( + os.environ.get("PROMETHEUS_MULTIPROC_DIR") + or os.environ.get("prometheus_multiproc_dir") + ) + + auto_created = not multiproc_dir + if not multiproc_dir: + multiproc_dir = os.path.join( + tempfile.gettempdir(), "litellm_prometheus_multiproc" + ) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir + + os.makedirs(multiproc_dir, exist_ok=True) + wipe_directory(multiproc_dir) + action = "Auto-created" if auto_created else "Using existing" + print(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}") # noqa + @click.command() @click.option( @@ -819,6 +869,12 @@ def run_server( # noqa: PLR0915 # DO NOT DELETE - enables global variables to work across files from litellm.proxy.proxy_server import app # noqa + # Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups + ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=num_workers, + litellm_settings=litellm_settings if config else None, + ) + # --- SEPARATE HEALTH APP LOGIC --- # To run the health app separately, use: # uvicorn litellm.proxy.health_app_factory:build_health_app --factory --host 0.0.0.0 --port=4001 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index be76c2ac5fb..bd5b5309e0f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -424,6 +424,9 @@ from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( router as user_agent_analytics_router, ) from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update +from litellm.proxy.middleware.in_flight_requests_middleware import ( + InFlightRequestsMiddleware, +) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_router @@ -1404,6 +1407,7 @@ app.add_middleware( ) app.add_middleware(PrometheusAuthMiddleware) +app.add_middleware(InFlightRequestsMiddleware) def mount_swagger_ui(): @@ -5298,13 +5302,15 @@ async def async_data_generator( ): verbose_proxy_logger.debug("inside generator") try: - # Use a list to accumulate response segments to avoid O(n^2) string concatenation - str_so_far_parts: list[str] = [] error_message: Optional[str] = None requested_model_from_client = _get_client_requested_model_for_streaming( request_data=request_data ) model_mismatch_logged = False + # Use a running string instead of list + join to avoid O(n^2) overhead. + # Previously "".join(str_so_far_parts) was called every chunk, re-joining + # the entire accumulated response. String += is O(n) amortized total. + _str_so_far: str = "" async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=response, @@ -5315,12 +5321,12 @@ async def async_data_generator( user_api_key_dict=user_api_key_dict, response=chunk, data=request_data, - str_so_far="".join(str_so_far_parts), + str_so_far=_str_so_far if _str_so_far else None, ) if isinstance(chunk, (ModelResponse, ModelResponseStream)): response_str = litellm.get_response_string(response_obj=chunk) - str_so_far_parts.append(response_str) + _str_so_far += response_str chunk, model_mismatch_logged = _restamp_streaming_chunk_model( chunk=chunk, diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 29c9cb571ca..ac5d9126145 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -1,6 +1,8 @@ import json import os -from typing import List +import re +from importlib.resources import files +from typing import Any, Dict, List, Optional import litellm from fastapi import APIRouter, Depends, HTTPException @@ -23,11 +25,107 @@ from litellm.types.proxy.public_endpoints.public_endpoints import ( AgentCreateInfo, ProviderCreateInfo, PublicModelHubInfo, + SupportedEndpointsResponse, ) from litellm.types.utils import LlmProviders router = APIRouter() +# --------------------------------------------------------------------------- +# /public/endpoints — helpers +# --------------------------------------------------------------------------- + +_ENDPOINT_METADATA: Dict[str, Dict[str, str]] = { + "chat_completions": {"label": "Chat Completions", "endpoint": "/chat/completions"}, + "messages": {"label": "Messages", "endpoint": "/messages"}, + "responses": {"label": "Responses", "endpoint": "/responses"}, + "embeddings": {"label": "Embeddings", "endpoint": "/embeddings"}, + "image_generations": {"label": "Image Generations", "endpoint": "/images/generations"}, + "audio_transcriptions": {"label": "Audio Transcriptions", "endpoint": "/audio/transcriptions"}, + "audio_speech": {"label": "Audio Speech", "endpoint": "/audio/speech"}, + "moderations": {"label": "Moderations", "endpoint": "/moderations"}, + "batches": {"label": "Batches", "endpoint": "/batches"}, + "rerank": {"label": "Rerank", "endpoint": "/rerank"}, + "ocr": {"label": "OCR", "endpoint": "/ocr"}, + "search": {"label": "Search", "endpoint": "/search"}, + "skills": {"label": "Skills", "endpoint": "/skills"}, + "interactions": {"label": "Interactions", "endpoint": "/interactions"}, + "a2a": {"label": "A2A (Agent Gateway)", "endpoint": "/a2a/{agent}/message/send"}, + "container": {"label": "Containers", "endpoint": "/containers"}, + "container_files": {"label": "Container Files", "endpoint": "/containers/{id}/files"}, + "compact": {"label": "Compact", "endpoint": "/responses/compact"}, + "files": {"label": "Files", "endpoint": "/files"}, + "image_edits": {"label": "Image Edits", "endpoint": "/images/edits"}, + "vector_stores_create": {"label": "Vector Stores (Create)", "endpoint": "/vector_stores"}, + "vector_stores_search": {"label": "Vector Stores (Search)", "endpoint": "/vector_stores/{id}/search"}, + "vector_store_files": {"label": "Vector Store Files", "endpoint": "/vector_stores/{id}/files"}, + "video_generations": {"label": "Video Generations", "endpoint": "/videos/generations"}, + "assistants": {"label": "Assistants", "endpoint": "/assistants"}, + "fine_tuning": {"label": "Fine Tuning", "endpoint": "/fine_tuning/jobs"}, + "text_completion": {"label": "Text Completion", "endpoint": "/completions"}, + "realtime": {"label": "Realtime", "endpoint": "/realtime"}, + "count_tokens": {"label": "Count Tokens", "endpoint": "/utils/token_counter"}, + "image_variations": {"label": "Image Variations", "endpoint": "/images/variations"}, + "generateContent": {"label": "Generate Content", "endpoint": "/generateContent"}, + "bedrock_invoke": {"label": "Bedrock Invoke", "endpoint": "/bedrock/invoke"}, + "bedrock_converse": {"label": "Bedrock Converse", "endpoint": "/bedrock/converse"}, + "rag_ingest": {"label": "RAG Ingest", "endpoint": "/rag/ingest"}, + "rag_query": {"label": "RAG Query", "endpoint": "/rag/query"}, +} + +_SLUG_SUFFIX_RE = re.compile(r"\s*\(`[^`]+`\)\s*$") + +# Loaded once on first request; never invalidated (local file, no TTL needed). +_cached_endpoints: Optional[List[Dict[str, Any]]] = None + + +def _clean_display_name(raw: str) -> str: + return _SLUG_SUFFIX_RE.sub("", raw).strip() + + +def _build_endpoints(raw: Dict[str, Any]) -> List[Dict[str, Any]]: + """Transform raw provider_endpoints_support_backup.json into the response shape.""" + providers: Dict[str, Any] = raw.get("providers", {}) + + # Collect endpoint keys in insertion order (union across all providers). + seen: set = set() + all_keys: List[str] = [] + for provider_data in providers.values(): + for key in provider_data.get("endpoints", {}): + if key not in seen: + seen.add(key) + all_keys.append(key) + + result: List[Dict[str, Any]] = [] + for key in all_keys: + meta = _ENDPOINT_METADATA.get(key) + label = meta["label"] if meta else key.replace("_", " ").title() + path = meta["endpoint"] if meta else "/" + key.replace("_", "/") + + supporting: List[Dict[str, str]] = [ + { + "slug": slug, + "display_name": _clean_display_name(pd.get("display_name", slug)), + } + for slug, pd in providers.items() + if pd.get("endpoints", {}).get(key) + ] + result.append({"key": key, "label": label, "endpoint": path, "providers": supporting}) + + return result + + +def _load_endpoints() -> List[Dict[str, Any]]: + raw = json.loads( + files("litellm") + .joinpath("provider_endpoints_support_backup.json") + .read_text(encoding="utf-8") + ) + return _build_endpoints(raw) + + +# --------------------------------------------------------------------------- + @router.get( "/public/model_hub", @@ -225,6 +323,24 @@ async def get_litellm_blog_posts(): return BlogPostsResponse(posts=posts) +@router.get( + "/public/endpoints", + tags=["public"], + response_model=SupportedEndpointsResponse, +) +async def get_supported_endpoints() -> SupportedEndpointsResponse: + """ + Return the list of LiteLLM proxy endpoints and which providers support each one. + + Reads from the bundled local backup file. Result is cached in-process for + the lifetime of the server process. + """ + global _cached_endpoints + if _cached_endpoints is None: + _cached_endpoints = SupportedEndpointsResponse(endpoints=_load_endpoints()) + return _cached_endpoints + + @router.get( "/public/agents/fields", tags=["public", "[beta] Agents"], @@ -233,7 +349,7 @@ async def get_litellm_blog_posts(): async def get_agent_fields() -> List[AgentCreateInfo]: """ Return agent type metadata required by the dashboard create-agent flow. - + If an agent has `inherit_credentials_from_provider`, the provider's credential fields are automatically appended to the agent's credential_fields. """ @@ -242,19 +358,19 @@ async def get_agent_fields() -> List[AgentCreateInfo]: "proxy", "public_endpoints", ) - + agent_create_fields_path = os.path.join(base_path, "agent_create_fields.json") provider_create_fields_path = os.path.join(base_path, "provider_create_fields.json") with open(agent_create_fields_path, "r") as f: agent_create_fields = json.load(f) - + with open(provider_create_fields_path, "r") as f: provider_create_fields = json.load(f) - + # Build a lookup map for providers by name provider_map = {p["provider"]: p for p in provider_create_fields} - + # Merge inherited credential fields for agent in agent_create_fields: inherit_from = agent.get("inherit_credentials_from_provider") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 440c9c1d829..f18556ac329 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -300,7 +300,7 @@ model LiteLLM_MCPServerTable { token_url String? registration_url String? allow_all_keys Boolean @default(false) - available_on_public_internet Boolean @default(false) + available_on_public_internet Boolean @default(true) } // Generate Tokens for Proxy @@ -390,6 +390,7 @@ model LiteLLM_DeletedVerificationToken { config Json @default("{}") user_id String? team_id String? + agent_id String? project_id String? permissions Json @default("{}") max_parallel_requests Int? @@ -504,6 +505,7 @@ model LiteLLM_SpendLogs { agent_id String? proxy_server_request Json? @default("{}") @@index([startTime]) + @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) } @@ -1094,4 +1096,19 @@ model LiteLLM_AccessGroupTable { created_by String? updated_at DateTime @default(now()) @updatedAt updated_by String? -} \ No newline at end of file +} +// Claude Code Plugin Marketplace table +model LiteLLM_ClaudeCodePluginTable { + id String @id @default(uuid()) + name String @unique + version String? + description String? + manifest_json String? + files_json String? @default("{}") + enabled Boolean @default(true) + created_at DateTime? @default(now()) + updated_at DateTime? @default(now()) @updatedAt + created_by String? + + @@map("LiteLLM_ClaudeCodePluginTable") +} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 320b58438a1..214a63dcbc9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2071,6 +2071,11 @@ class ProxyLogging: response_str = extract_text_from_a2a_response(response) if response_str is not None: + # Cache model-level guardrails check per-request to avoid repeated + # dict lookups + llm_router.get_deployment() per callback per chunk. + _cached_guardrail_data: Optional[dict] = None + _guardrail_data_computed = False + for callback in litellm.callbacks: try: _callback: Optional[CustomLogger] = None @@ -2078,14 +2083,16 @@ class ProxyLogging: # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks - ## CHECK FOR MODEL-LEVEL GUARDRAILS - modified_data = _check_and_merge_model_level_guardrails( - data=data, llm_router=llm_router - ) + ## CHECK FOR MODEL-LEVEL GUARDRAILS (cached per-request) + if not _guardrail_data_computed: + _cached_guardrail_data = _check_and_merge_model_level_guardrails( + data=data, llm_router=llm_router + ) + _guardrail_data_computed = True if ( callback.should_run_guardrail( - data=modified_data, + data=_cached_guardrail_data, event_type=GuardrailEventHooks.post_call, ) is not True @@ -4692,7 +4699,10 @@ async def update_spend_logs_job( from litellm.proxy.guardrails.usage_tracking import ( process_spend_logs_guardrail_usage, ) +<<<<<<< litellm_guardrail-filtering-dispatch +======= +>>>>>>> main await process_spend_logs_guardrail_usage( prisma_client=prisma_client, logs_to_process=logs_to_process, diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index ac597fc623d..3e64f61abdb 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -106,6 +106,8 @@ async def _arealtime( # noqa: PLR0915 client=client, timeout=timeout, headers=headers, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata(kwargs), ) elif _custom_llm_provider == "azure": api_base = ( @@ -277,6 +279,8 @@ async def _arealtime( # noqa: PLR0915 client=client, timeout=timeout, headers=headers, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata(kwargs), ) else: raise ValueError(f"Unsupported model: {model}") diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 276e0cb9a04..a627531e994 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -177,6 +177,19 @@ async def aresponses_api_with_mcp( "litellm_metadata", {} ).get("user_api_key_auth") + # Extract MCP auth headers from request (for dynamic auth when fetching tools) + mcp_auth_header: Optional[str] = None + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None + secret_fields = kwargs.get("secret_fields") + if secret_fields and isinstance(secret_fields, dict): + from litellm.responses.utils import ResponsesAPIRequestUtils + + mcp_auth_header, mcp_server_auth_headers, _, _ = ( + ResponsesAPIRequestUtils.extract_mcp_headers_from_request( + secret_fields=secret_fields, tools=tools + ) + ) + # Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods ( original_mcp_tools, @@ -185,6 +198,8 @@ async def aresponses_api_with_mcp( user_api_key_auth=user_api_key_auth, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, litellm_trace_id=kwargs.get("litellm_trace_id"), + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, ) openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( original_mcp_tools @@ -370,6 +385,8 @@ async def aresponses_api_with_mcp( ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( user_api_key_auth=user_api_key_auth, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, ) final_response = ( LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response( diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 377ce396457..bacc627cc84 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -120,7 +120,18 @@ async def acompletion_with_mcp( # noqa: PLR0915 (kwargs.get("metadata", {}) or {}).get("user_api_key_auth") ) - # Process MCP tools + # Extract MCP auth headers before fetching tools (needed for dynamic auth) + ( + mcp_auth_header, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( + secret_fields=kwargs.get("secret_fields"), + tools=tools, + ) + + # Process MCP tools (pass auth headers for dynamic auth) ( deduplicated_mcp_tools, tool_server_map, @@ -128,6 +139,8 @@ async def acompletion_with_mcp( # noqa: PLR0915 user_api_key_auth=user_api_key_auth, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, litellm_trace_id=kwargs.get("litellm_trace_id"), + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, ) openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( @@ -143,17 +156,6 @@ async def acompletion_with_mcp( # noqa: PLR0915 mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy ) - # Extract MCP auth headers - ( - mcp_auth_header, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( - secret_fields=kwargs.get("secret_fields"), - tools=tools, - ) - # Prepare call parameters # Remove keys that shouldn't be passed to acompletion clean_kwargs = {k: v for k, v in kwargs.items() if k not in ["acompletion"]} diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 805a1958552..5776ef95acb 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -99,6 +99,8 @@ class LiteLLM_Proxy_MCP_Handler: user_api_key_auth: Any, mcp_tools_with_litellm_proxy: Optional[Iterable[ToolParam]], litellm_trace_id: Optional[str] = None, + mcp_auth_header: Optional[str] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, ) -> tuple[List[MCPTool], List[str]]: """ Get available tools from the MCP server manager. @@ -106,6 +108,8 @@ class LiteLLM_Proxy_MCP_Handler: Args: user_api_key_auth: User authentication info for access control mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" + mcp_auth_header: Optional deprecated auth header for MCP servers + mcp_server_auth_headers: Optional server-specific auth headers (e.g. from x-mcp-{alias}-*) Returns: List of MCP tools @@ -133,13 +137,14 @@ class LiteLLM_Proxy_MCP_Handler: tools = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, - mcp_auth_header=None, + mcp_auth_header=mcp_auth_header, mcp_servers=mcp_servers, - mcp_server_auth_headers=None, + mcp_server_auth_headers=mcp_server_auth_headers, log_list_tools_to_spendlogs=True, list_tools_log_source="responses", litellm_trace_id=litellm_trace_id, ) + allowed_mcp_server_ids = ( await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth) ) @@ -278,6 +283,8 @@ class LiteLLM_Proxy_MCP_Handler: user_api_key_auth: Any, mcp_tools_with_litellm_proxy: List[ToolParam], litellm_trace_id: Optional[str] = None, + mcp_auth_header: Optional[str] = None, + mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, ) -> tuple[List[Any], dict[str, str]]: """ Process MCP tools through filtering and deduplication pipeline without OpenAI transformation. @@ -286,6 +293,8 @@ class LiteLLM_Proxy_MCP_Handler: Args: user_api_key_auth: User authentication info for access control mcp_tools_with_litellm_proxy: ToolParam objects with server_url starting with "litellm_proxy" + mcp_auth_header: Optional deprecated auth header for MCP servers + mcp_server_auth_headers: Optional server-specific auth headers (e.g. from x-mcp-{alias}-*) Returns: List of filtered and deduplicated MCP tools in their original format @@ -301,6 +310,8 @@ class LiteLLM_Proxy_MCP_Handler: user_api_key_auth=user_api_key_auth, mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy, litellm_trace_id=litellm_trace_id, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, ) # Step 2: Filter tools based on allowed_tools parameter diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index edcbb0d11b8..43ef4610b4b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1,5 +1,6 @@ import asyncio import json +import time import traceback from datetime import datetime from typing import Any, Dict, Optional @@ -7,7 +8,7 @@ from typing import Any, Dict, Optional import httpx import litellm -from litellm.constants import STREAM_SSE_DONE_STRING +from litellm.constants import LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -56,6 +57,7 @@ class BaseResponsesAPIStreamingIterator: self.completed_response: Optional[ResponsesAPIStreamingResponse] = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called + self._stream_created_time: float = time.time() # track request context for hooks self.litellm_metadata = litellm_metadata @@ -82,6 +84,18 @@ class BaseResponsesAPIStreamingIterator: self.response.headers or {} ) # GUARANTEE OPENAI HEADERS IN RESPONSE + def _check_max_streaming_duration(self) -> None: + """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" + if LITELLM_MAX_STREAMING_DURATION_SECONDS is None: + return + elapsed = time.time() - self._stream_created_time + if elapsed > LITELLM_MAX_STREAMING_DURATION_SECONDS: + raise litellm.Timeout( + message=f"Stream exceeded max streaming duration of {LITELLM_MAX_STREAMING_DURATION_SECONDS}s (elapsed {elapsed:.1f}s)", + model=self.model or "", + llm_provider=self.custom_llm_provider or "", + ) + def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]: """Process a single chunk of data from the stream""" if not chunk: @@ -357,6 +371,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): async def __anext__(self) -> ResponsesAPIStreamingResponse: try: + self._check_max_streaming_duration() while True: # Get the next chunk from the stream try: @@ -365,6 +380,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): self.finished = True raise StopAsyncIteration + self._check_max_streaming_duration() result = self._process_chunk(chunk) if self.finished: @@ -460,6 +476,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __next__(self): try: + self._check_max_streaming_duration() while True: # Get the next chunk from the stream try: @@ -468,6 +485,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): self.finished = True raise StopIteration + self._check_max_streaming_duration() result = self._process_chunk(chunk) if self.finished: diff --git a/litellm/router.py b/litellm/router.py index cbe5b414040..d89a5099b01 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6677,6 +6677,22 @@ class Router: # initialize client self._add_deployment(deployment=deployment) + # Register custom pricing in litellm.model_cost. + # Mirrors _create_deployment() logic to ensure dynamically-added deployments + # (e.g., loaded from DB) also have their custom pricing registered. + # Without this, _is_model_cost_zero() cannot detect explicitly-configured + # zero-cost models, causing budget checks to block free models. + _model_id = deployment.model_info.id + if _model_id is not None: + _model_info_dict: dict = deployment.model_info.model_dump( + exclude_none=True + ) + for field in CustomPricingLiteLLMParams.model_fields.keys(): + field_value = deployment.litellm_params.get(field) + if field_value is not None: + _model_info_dict[field] = field_value + litellm.register_model(model_cost={_model_id: _model_info_dict}) + # add to model names self._add_model_to_list_and_index_map( model=_deployment, model_id=deployment.model_info.id diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 2c75276d9ca..0856d8a6f9b 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -237,6 +237,7 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_remaining_api_key_tokens_for_model", "litellm_llm_api_failed_requests_metric", "litellm_callback_logging_failures_metric", + "litellm_in_flight_requests", ] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 15e8d1be930..c0aae9bc2de 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -699,7 +699,15 @@ class OpenAIChatCompletionAssistantMessage(TypedDict, total=False): role: Required[Literal["assistant"]] content: Optional[ Union[ - str, Iterable[Union[ChatCompletionTextObject, ChatCompletionThinkingBlock]] + str, + Iterable[ + Union[ + ChatCompletionTextObject, + ChatCompletionThinkingBlock, + ChatCompletionRedactedThinkingBlock, + ChatCompletionImageObject, + ] + ], ] ] name: Optional[str] @@ -786,17 +794,19 @@ ValidUserMessageContentTypes = [ "file", ] # used for validating user messages. Prevent users from accidentally sending anthropic messages. -# Assistant message content types (text, thinking, redacted_thinking) +# Assistant message content types (text, thinking, redacted_thinking, image_url) ValidAssistantMessageContentTypesLiteral = Literal[ "text", "thinking", "redacted_thinking", + "image_url", ] ValidAssistantMessageContentTypes = [ "text", "thinking", "redacted_thinking", + "image_url", ] # Combined valid content types for chat completion messages diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 7f99fd526c8..69b34a25a21 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -52,7 +52,7 @@ class MCPServer(BaseModel): env: Optional[Dict[str, str]] = None access_groups: Optional[List[str]] = None allow_all_keys: bool = False - available_on_public_internet: bool = False + available_on_public_internet: bool = True updated_at: Optional[datetime] = None model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index 57d68771c7f..caa9a978530 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -52,3 +52,19 @@ class AgentCreateInfo(BaseModel): credential_fields: List[AgentCredentialField] litellm_params_template: Optional[Dict[str, str]] = None model_template: Optional[str] = None + + +class EndpointProvider(BaseModel): + slug: str + display_name: str + + +class SupportedEndpoint(BaseModel): + key: str + label: str + endpoint: str + providers: List[EndpointProvider] + + +class SupportedEndpointsResponse(BaseModel): + endpoints: List[SupportedEndpoint] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 57563fc0bcc..f52288ea72a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14194,6 +14194,38 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-3.1-flash-image-preview": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true + }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -19178,6 +19210,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "gpt-audio-1.5": { + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "gpt-audio-2025-08-28": { "input_cost_per_audio_token": 3.2e-05, "input_cost_per_token": 2.5e-06, @@ -20895,6 +20960,38 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "gpt-realtime-1.5": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "openai", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, @@ -25060,6 +25157,25 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "openrouter/anthropic/claude-opus-4.6": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, "cache_creation_input_token_cost": 3.75e-06, @@ -26072,6 +26188,42 @@ "supports_prompt_caching": true, "supports_computer_use": false }, + "openrouter/openrouter/auto": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "max_input_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true + }, + "openrouter/openrouter/free": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "max_input_tokens": 200000, + "max_tokens": 200000, + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, + "openrouter/openrouter/bodybuilder": { + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "openrouter", + "max_input_tokens": 128000, + "max_tokens": 128000, + "mode": "chat" + }, "ovhcloud/DeepSeek-R1-Distill-Llama-70B": { "input_cost_per_token": 6.7e-07, "litellm_provider": "ovhcloud", @@ -26618,8 +26770,8 @@ "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", - "supports_function_calling": true, - "supports_tool_choice": true + "supports_function_calling": false, + "supports_tool_choice": false }, "publicai/swiss-ai/apertus-70b-instruct": { "input_cost_per_token": 0.0, @@ -26630,8 +26782,8 @@ "mode": "chat", "output_cost_per_token": 0.0, "source": "https://platform.publicai.co/docs", - "supports_function_calling": true, - "supports_tool_choice": true + "supports_function_calling": false, + "supports_tool_choice": false }, "publicai/aisingapore/Gemma-SEA-LION-v4-27B-IT": { "input_cost_per_token": 0.0, @@ -31545,6 +31697,19 @@ "output_cost_per_token_batches": 6e-06, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, + "vertex_ai/gemini-3.1-flash-image-preview": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" + }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -32946,6 +33111,7 @@ "supports_web_search": true }, "xai/grok-2-vision-1212": { + "deprecation_date": "2026-02-28", "input_cost_per_image": 2e-06, "input_cost_per_token": 2e-06, "litellm_provider": "xai", @@ -33050,6 +33216,7 @@ }, "xai/grok-3-mini": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2026-02-28", "input_cost_per_token": 3e-07, "litellm_provider": "xai", "max_input_tokens": 131072, @@ -33066,6 +33233,7 @@ }, "xai/grok-3-mini-beta": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2026-02-28", "input_cost_per_token": 3e-07, "litellm_provider": "xai", "max_input_tokens": 131072, diff --git a/poetry.lock b/poetry.lock index 34227a69ccb..0314a360542 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3222,15 +3222,15 @@ files = [ [[package]] name = "litellm-proxy-extras" -version = "0.4.48" +version = "0.4.49" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." optional = true python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" groups = ["main"] markers = "extra == \"proxy\"" files = [ - {file = "litellm_proxy_extras-0.4.48-py3-none-any.whl", hash = "sha256:097001fccec5dbf4cffd902114898a9cfeba62673202447d55d2d0286cf93126"}, - {file = "litellm_proxy_extras-0.4.48.tar.gz", hash = "sha256:5d5d8acf31b92d0cd6738555fb4a2411819755155438de9fb23c724c356400a2"}, + {file = "litellm_proxy_extras-0.4.49-py3-none-any.whl", hash = "sha256:aeb0e08b4705c19fdc5b75a43c608a82fc36032f6d83be509dbf37baea62f2cd"}, + {file = "litellm_proxy_extras-0.4.49.tar.gz", hash = "sha256:d9bdae54d1e3398f2e2025c9d8b98a19e226874337d540d5415922d7dbbc97bb"}, ] [[package]] @@ -7989,4 +7989,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "b9b1e47b3b84748c0053be6a544c2399bf2601746a4f88dcb1be7c5e4eeab359" +content-hash = "bbc7d43f5484af4c8877fe66e34f8283069528379af49d573036ba144cc2eb7a" diff --git a/pyproject.toml b/pyproject.toml index 8eb433fb064..07f3ea30fc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ boto3 = { version = "1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "0.4.48", optional = true} +litellm-proxy-extras = {version = "0.4.49", optional = true} rich = {version = "13.7.1", optional = true} litellm-enterprise = {version = "0.1.32", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index 6cdb2f63a33..67f390cf272 100644 --- a/requirements.txt +++ b/requirements.txt @@ -57,7 +57,7 @@ grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.48 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.49 # for proxy extras - e.g. prisma migrations llm-sandbox==0.3.31 # for skill execution in sandbox ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env diff --git a/schema.prisma b/schema.prisma index 440c9c1d829..a8cb297a3ed 100644 --- a/schema.prisma +++ b/schema.prisma @@ -300,7 +300,7 @@ model LiteLLM_MCPServerTable { token_url String? registration_url String? allow_all_keys Boolean @default(false) - available_on_public_internet Boolean @default(false) + available_on_public_internet Boolean @default(true) } // Generate Tokens for Proxy @@ -390,6 +390,7 @@ model LiteLLM_DeletedVerificationToken { config Json @default("{}") user_id String? team_id String? + agent_id String? project_id String? permissions Json @default("{}") max_parallel_requests Int? @@ -1094,4 +1095,19 @@ model LiteLLM_AccessGroupTable { created_by String? updated_at DateTime @default(now()) @updatedAt updated_by String? -} \ No newline at end of file +} +// Claude Code Plugin Marketplace table +model LiteLLM_ClaudeCodePluginTable { + id String @id @default(uuid()) + name String @unique + version String? + description String? + manifest_json String? + files_json String? @default("{}") + enabled Boolean @default(true) + created_at DateTime? @default(now()) + updated_at DateTime? @default(now()) @updatedAt + created_by String? + + @@map("LiteLLM_ClaudeCodePluginTable") +} diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 08b9351f9a3..61b1b1f8185 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -2316,5 +2316,3 @@ async def test_prometheus_token_metrics_with_prometheus_config(): raise AssertionError(f"Metric {metric_name} not found in registry") print("✓ All token metrics validated successfully!") - - # check final value of metrics in registry diff --git a/tests/litellm/llms/openai_like/test_assemblyai_provider.py b/tests/litellm/llms/openai_like/test_assemblyai_provider.py new file mode 100644 index 00000000000..7eee810b271 --- /dev/null +++ b/tests/litellm/llms/openai_like/test_assemblyai_provider.py @@ -0,0 +1,77 @@ +""" +Unit tests for the AssemblyAI LLM Gateway 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 + +ASSEMBLYAI_BASE_URL = "https://llm-gateway.assemblyai.com/v1" + + +def _get_config(): + provider = JSONProviderRegistry.get("assemblyai") + assert provider is not None + config_class = create_config_class(provider) + return config_class() + + +def test_assemblyai_provider_registered(): + provider = JSONProviderRegistry.get("assemblyai") + assert provider is not None + assert provider.base_url == ASSEMBLYAI_BASE_URL + assert provider.api_key_env == "ASSEMBLYAI_API_KEY" + + +def test_assemblyai_resolves_env_api_key(monkeypatch): + config = _get_config() + monkeypatch.setenv("ASSEMBLYAI_API_KEY", "test-key") + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == ASSEMBLYAI_BASE_URL + assert api_key == "test-key" + + +def test_assemblyai_complete_url_appends_endpoint(): + config = _get_config() + url = config.get_complete_url( + api_base=ASSEMBLYAI_BASE_URL, + api_key="test-key", + model="assemblyai/claude-sonnet-4-5-20250929", + optional_params={}, + litellm_params={}, + stream=False, + ) + assert url == f"{ASSEMBLYAI_BASE_URL}/chat/completions" + + +def test_assemblyai_provider_resolution(): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="assemblyai/claude-sonnet-4-5-20250929", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "claude-sonnet-4-5-20250929" + assert provider == "assemblyai" + assert api_base == ASSEMBLYAI_BASE_URL + + +def test_assemblyai_provider_config_manager(): + from litellm import LlmProviders + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model="claude-sonnet-4-5-20250929", provider=LlmProviders.ASSEMBLYAI + ) + + assert config is not None + assert config.custom_llm_provider == "assemblyai" diff --git a/tests/litellm/proxy/test_claude_code_marketplace.py b/tests/litellm/proxy/test_claude_code_marketplace.py new file mode 100644 index 00000000000..5376e81012b --- /dev/null +++ b/tests/litellm/proxy/test_claude_code_marketplace.py @@ -0,0 +1,18 @@ +import pytest + + +@pytest.mark.asyncio +async def test_claude_code_plugin_table_schema_exists(): + + with open("schema.prisma", "r") as f: + schema = f.read() + assert "LiteLLM_ClaudeCodePluginTable" in schema, ( + "LiteLLM_ClaudeCodePluginTable model missing from schema.prisma - " + "this causes AttributeError on all /claude-code/plugins endpoints" + ) + + with open("litellm/proxy/schema.prisma", "r") as f: + proxy_schema = f.read() + assert "LiteLLM_ClaudeCodePluginTable" in proxy_schema, ( + "LiteLLM_ClaudeCodePluginTable model missing from litellm/proxy/schema.prisma" + ) diff --git a/tests/llm_translation/realtime/test_realtime_guardrails_openai.py b/tests/llm_translation/realtime/test_realtime_guardrails_openai.py new file mode 100644 index 00000000000..a7913e6d761 --- /dev/null +++ b/tests/llm_translation/realtime/test_realtime_guardrails_openai.py @@ -0,0 +1,355 @@ +""" +Integration tests for RealTimeStreaming guardrails against a live OpenAI backend. + +These tests require OPENAI_API_KEY and are skipped if not set. + +They verify end-to-end that: + 1. A text message blocked by a guardrail -> error event sent to client, NO AI response. + 2. A voice transcript blocked by a guardrail -> error event sent, response.create NOT sent. + 3. A clean text message passes through and triggers a real OpenAI response. + +Run with: + poetry run pytest tests/llm_translation/realtime/test_realtime_guardrails_openai.py -v -s +""" + +import asyncio +import json +import os +from typing import List +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming +from litellm.types.guardrails import GuardrailEventHooks + +OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") +OPENAI_REALTIME_URL = ( + "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-12-17" +) + +pytestmark = pytest.mark.skipif( + not OPENAI_API_KEY, + reason="OPENAI_API_KEY not set - skipping OpenAI realtime integration tests", +) + +# A unique phrase guaranteed NOT to appear in normal assistant output. +BLOCKED_PHRASE = "XSECRETBLOCKTESTPHRASEX" + + +class PhraseBlockingGuardrail(CustomGuardrail): + """Blocks any message containing BLOCKED_PHRASE.""" + + async def apply_guardrail( + self, inputs, request_data, input_type, logging_obj=None + ): + for text in inputs.get("texts", []): + if BLOCKED_PHRASE in text: + raise ValueError( + "Content blocked: contains forbidden test phrase." + ) + return inputs + + +def _make_guardrail(event_hook=GuardrailEventHooks.pre_call): + return PhraseBlockingGuardrail( + guardrail_name="integration-test-guard", + event_hook=event_hook, + default_on=True, + ) + + +async def _wait_for_event( + client_events: List[dict], event_type: str, timeout: float = 15.0 +) -> dict: + """Poll client_events list until an event with matching type appears.""" + deadline = asyncio.get_event_loop().time() + timeout + while asyncio.get_event_loop().time() < deadline: + matching = [e for e in client_events if e.get("type") == event_type] + if matching: + return matching[0] + await asyncio.sleep(0.05) + raise TimeoutError( + f"Timed out waiting for '{event_type}'. Got so far: {[e.get('type') for e in client_events]}" + ) + + +async def _build_streaming(client_events: List[dict], backend_ws, request_data=None): + """Create a RealTimeStreaming with a mock client WebSocket that captures events.""" + client_ws = MagicMock() + input_queue: asyncio.Queue = asyncio.Queue() + + async def send_text(data: str): + client_events.append(json.loads(data)) + + client_ws.send_text = send_text + client_ws.receive_text = input_queue.get + + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + logging_obj.model_call_details = {} + + streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=backend_ws, + logging_obj=logging_obj, + request_data=request_data or {"guardrails": ["integration-test-guard"]}, + ) + return streaming, input_queue + + +@pytest.mark.asyncio +async def test_text_message_blocked_by_guardrail_no_ai_response(): + """ + Send a text message containing the blocked phrase. + Guardrail must: + - Send error event (guardrail_violation) to client. + - Send response.audio_transcript.delta with the block message to client. + - NOT forward response.create to OpenAI (no AI response). + """ + import websockets + + guardrail = _make_guardrail(GuardrailEventHooks.pre_call) + litellm.callbacks = [guardrail] + + client_events: List[dict] = [] + + try: + async with websockets.connect( + OPENAI_REALTIME_URL, + additional_headers={ + "Authorization": f"Bearer {OPENAI_API_KEY}", + "OpenAI-Beta": "realtime=v1", + }, + ) as backend_ws: + streaming, input_queue = await _build_streaming(client_events, backend_ws) + + # Start backend -> client forwarding + backend_task = asyncio.create_task( + streaming.backend_to_client_send_messages() + ) + # Start client -> backend forwarding (reads from input_queue) + client_task = asyncio.create_task(streaming.client_ack_messages()) + + try: + # Wait until session is ready + await _wait_for_event(client_events, "session.created", timeout=15) + + # Send the blocked message + response.create + blocked_item = json.dumps( + { + "type": "conversation.item.create", + "item": { + "role": "user", + "content": [ + { + "type": "input_text", + "text": f"Hello {BLOCKED_PHRASE}", + } + ], + }, + } + ) + await input_queue.put(blocked_item) + # Give guardrail time to process before the follow-up response.create + await asyncio.sleep(0.3) + await input_queue.put(json.dumps({"type": "response.create"})) + + # Allow time for guardrail round-trip + await asyncio.sleep(3.0) + + finally: + backend_task.cancel() + client_task.cancel() + await asyncio.gather(backend_task, client_task, return_exceptions=True) + + # --- Assertions --- + event_types = [e.get("type") for e in client_events] + + # 1. Must have received guardrail error + error_events = [e for e in client_events if e.get("type") == "error"] + assert len(error_events) >= 1, ( + f"Expected at least one error event but got: {event_types}" + ) + assert error_events[0]["error"]["type"] == "guardrail_violation", ( + f"Wrong error type: {error_events[0]}" + ) + + # 2. Must have the guardrail message surfaced as an AI transcript delta + transcript_deltas = [ + e + for e in client_events + if e.get("type") == "response.audio_transcript.delta" + ] + assert len(transcript_deltas) >= 1, ( + f"Expected guardrail message in transcript delta, got: {event_types}" + ) + + # 3. No real AI response should have been generated - response.done would only + # appear if we sent a response.create and OpenAI replied. We allow it in the + # synthetic form (empty output=[]) but NOT with actual AI content. + done_events = [e for e in client_events if e.get("type") == "response.done"] + for done in done_events: + output = done.get("response", {}).get("output", []) + ai_texts = [ + c.get("text", "") or c.get("transcript", "") + for item in output + for c in item.get("content", []) + ] + real_ai_text = " ".join(ai_texts).strip() + assert real_ai_text == "", ( + f"AI responded with real content even though message was blocked: {real_ai_text!r}" + ) + + finally: + litellm.callbacks = [] + + +@pytest.mark.asyncio +async def test_voice_transcript_blocked_by_guardrail(): + """ + Simulate a backend-side voice transcription event containing the blocked phrase. + Guardrail must block it - no response.create sent to OpenAI. + """ + from websockets.exceptions import ConnectionClosed + + guardrail = _make_guardrail(GuardrailEventHooks.realtime_input_transcription) + litellm.callbacks = [guardrail] + + client_events: List[dict] = [] + + # Build the transcript event that would come from the OpenAI backend + transcript_event = json.dumps( + { + "type": "conversation.item.input_audio_transcription.completed", + "transcript": f"This is {BLOCKED_PHRASE} in my voice message", + "item_id": "item_integ_test", + } + ).encode() + + # Mock backend that delivers the transcript then closes + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + transcript_event, + ConnectionClosed(None, None), + ] + ) + backend_ws.send = AsyncMock() + + try: + streaming, _ = await _build_streaming(client_events, backend_ws) + await streaming.backend_to_client_send_messages() + + event_types = [e.get("type") for e in client_events] + + # 1. Error event must be sent to client + error_events = [e for e in client_events if e.get("type") == "error"] + assert len(error_events) >= 1, ( + f"Expected guardrail error event, got: {event_types}" + ) + assert error_events[0]["error"]["type"] == "guardrail_violation" + + # 2. response.create must NOT have been sent to backend + sent_to_backend = [ + json.loads(c.args[0]) + for c in backend_ws.send.call_args_list + if c.args and isinstance(c.args[0], str) + ] + response_creates = [ + e for e in sent_to_backend if e.get("type") == "response.create" + ] + assert len(response_creates) == 0, ( + f"Guardrail should have stopped response.create, got: {sent_to_backend}" + ) + + # 3. Guardrail message surfaced as AI transcript delta + transcript_deltas = [ + e + for e in client_events + if e.get("type") == "response.audio_transcript.delta" + ] + assert len(transcript_deltas) >= 1, ( + f"Expected guardrail message in transcript delta, got: {event_types}" + ) + + finally: + litellm.callbacks = [] + + +@pytest.mark.asyncio +async def test_clean_text_message_passes_through_to_openai(): + """ + A clean message (no blocked phrase) must pass the guardrail and result in a real + AI response from OpenAI (response.done with non-empty output). + """ + import websockets + + guardrail = _make_guardrail(GuardrailEventHooks.pre_call) + litellm.callbacks = [guardrail] + + client_events: List[dict] = [] + + try: + async with websockets.connect( + OPENAI_REALTIME_URL, + additional_headers={ + "Authorization": f"Bearer {OPENAI_API_KEY}", + "OpenAI-Beta": "realtime=v1", + }, + ) as backend_ws: + streaming, input_queue = await _build_streaming(client_events, backend_ws) + + backend_task = asyncio.create_task( + streaming.backend_to_client_send_messages() + ) + client_task = asyncio.create_task(streaming.client_ack_messages()) + + try: + await _wait_for_event(client_events, "session.created", timeout=15) + + # Send a clean message + clean_item = json.dumps( + { + "type": "conversation.item.create", + "item": { + "role": "user", + "content": [ + {"type": "input_text", "text": "Reply with just: OK"} + ], + }, + } + ) + await input_queue.put(clean_item) + await asyncio.sleep(0.1) + await input_queue.put(json.dumps({"type": "response.create"})) + + # Wait for OpenAI to respond + await _wait_for_event(client_events, "response.done", timeout=30) + + finally: + backend_task.cancel() + client_task.cancel() + await asyncio.gather(backend_task, client_task, return_exceptions=True) + + # No guardrail error should have been sent + error_events = [e for e in client_events if e.get("type") == "error"] + guardrail_errors = [ + e for e in error_events if e.get("error", {}).get("type") == "guardrail_violation" + ] + assert len(guardrail_errors) == 0, ( + f"Clean message should not trigger guardrail, got: {guardrail_errors}" + ) + + # AI response must be present + done_events = [e for e in client_events if e.get("type") == "response.done"] + assert len(done_events) >= 1, ( + f"Expected response.done from OpenAI, got: {[e.get('type') for e in client_events]}" + ) + + finally: + litellm.callbacks = [] diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 88bca007740..9f902f2bd86 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -1214,6 +1214,181 @@ def test_anthropic_messages_pt_with_server_tool_use(): assert tool_use["id"] == "toolu_01XYZ789" +def test_convert_to_anthropic_tool_invoke_with_tool_results(): + """ + Test that non-web-search *_tool_result blocks (e.g. bash_code_execution_tool_result) + stored in provider_specific_fields["tool_results"] are paired with their server_tool_use + block when reconstructing assistant history. + + Regression for: server tool result blocks dropped on multi-turn replay + (bash_code_execution_tool_result, text_editor_code_execution_tool_result, etc.) + """ + tool_calls = [ + { + "id": "srvtoolu_01BASH", + "type": "function", + "function": { + "name": "bash_code_execution", + "arguments": '{"command": "python3 -c \\"print(2)\\""}', + }, + } + ] + + tool_results = [ + { + "type": "bash_code_execution_tool_result", + "tool_use_id": "srvtoolu_01BASH", + "content": { + "type": "bash_code_execution_result", + "stdout": "2\n", + "stderr": "", + "return_code": 0, + "content": [], + }, + } + ] + + result = convert_to_anthropic_tool_invoke(tool_calls, tool_results=tool_results) + + assert len(result) == 2 + # First: server_tool_use + assert result[0]["type"] == "server_tool_use" + assert result[0]["id"] == "srvtoolu_01BASH" + assert result[0]["name"] == "bash_code_execution" + # Second: bash_code_execution_tool_result paired correctly + assert result[1]["type"] == "bash_code_execution_tool_result" + assert result[1]["tool_use_id"] == "srvtoolu_01BASH" + + +def test_anthropic_messages_pt_raw_bash_tool_result_passthrough(): + """ + Test that raw assistant content lists containing bash_code_execution_tool_result + blocks are passed through intact to Anthropic. + + Regression: the raw-block passthrough only handled tool_search_tool_result; + bash_code_execution_tool_result and other *_tool_result types were silently dropped. + """ + messages = [ + {"role": "user", "content": "What is 1+1?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01BASH", + "name": "bash_code_execution", + "input": {"command": "python3 -c \"print(1+1)\""}, + }, + { + "type": "bash_code_execution_tool_result", + "tool_use_id": "srvtoolu_01BASH", + "content": { + "type": "bash_code_execution_result", + "stdout": "2\n", + "stderr": "", + "return_code": 0, + "content": [], + }, + }, + {"type": "text", "text": "The answer is 2."}, + ], + }, + {"role": "user", "content": "Thanks!"}, + ] + + result = anthropic_messages_pt( + messages, model="claude-sonnet-4-5", llm_provider="anthropic" + ) + + assistant_msg = next(m for m in result if m["role"] == "assistant") + content = assistant_msg["content"] + types = [c.get("type") for c in content] + + assert "server_tool_use" in types, "server_tool_use block must be preserved" + assert ( + "bash_code_execution_tool_result" in types + ), "bash_code_execution_tool_result block must not be dropped" + assert "text" in types + + # Result must immediately follow its server_tool_use + srv_idx = types.index("server_tool_use") + result_idx = types.index("bash_code_execution_tool_result") + assert result_idx == srv_idx + 1 + + bash_result = next( + c for c in content if c.get("type") == "bash_code_execution_tool_result" + ) + assert bash_result["tool_use_id"] == "srvtoolu_01BASH" + + +def test_anthropic_messages_pt_with_bash_tool_result_in_provider_specific_fields(): + """ + Test that anthropic_messages_pt correctly reconstructs bash_code_execution_tool_result + from provider_specific_fields["tool_results"] when replaying LiteLLM response objects. + + Regression: only web_search_results were read from provider_specific_fields; + tool_results (bash_code_execution_tool_result, etc.) were silently lost. + """ + messages = [ + {"role": "user", "content": "What is 1+1?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "srvtoolu_01BASH", + "type": "function", + "function": { + "name": "bash_code_execution", + "arguments": '{"command": "python3 -c \\"print(1+1)\\""}', + }, + } + ], + "provider_specific_fields": { + "tool_results": [ + { + "type": "bash_code_execution_tool_result", + "tool_use_id": "srvtoolu_01BASH", + "content": { + "type": "bash_code_execution_result", + "stdout": "2\n", + "stderr": "", + "return_code": 0, + "content": [], + }, + } + ] + }, + }, + {"role": "user", "content": "Thanks!"}, + ] + + result = anthropic_messages_pt( + messages, model="claude-sonnet-4-5", llm_provider="anthropic" + ) + + assistant_msg = next(m for m in result if m["role"] == "assistant") + content = assistant_msg["content"] + types = [c.get("type") for c in content] + + assert "server_tool_use" in types, "server_tool_use block must be reconstructed" + assert ( + "bash_code_execution_tool_result" in types + ), "bash_code_execution_tool_result must be paired from provider_specific_fields['tool_results']" + + # Result must immediately follow its server_tool_use + srv_idx = types.index("server_tool_use") + result_idx = types.index("bash_code_execution_tool_result") + assert result_idx == srv_idx + 1 + + srv = next(c for c in content if c.get("type") == "server_tool_use") + assert srv["id"] == "srvtoolu_01BASH" + bash_result = next( + c for c in content if c.get("type") == "bash_code_execution_tool_result" + ) + assert bash_result["tool_use_id"] == "srvtoolu_01BASH" + + # ============ parse_tool_call_arguments Tests ============ # Tests for the shared utility that parses tool call JSON arguments diff --git a/tests/llm_translation/test_skills_data/slack-gif-creator.zip b/tests/llm_translation/test_skills_data/slack-gif-creator.zip index 15c60e3667d..9827db9ac2d 100644 Binary files a/tests/llm_translation/test_skills_data/slack-gif-creator.zip and b/tests/llm_translation/test_skills_data/slack-gif-creator.zip differ diff --git a/tests/llm_translation/test_skills_data/slack-gif-creator/LICENSE.txt b/tests/llm_translation/test_skills_data/slack-gif-creator/LICENSE.txt deleted file mode 100644 index 7a4a3ea2424..00000000000 --- a/tests/llm_translation/test_skills_data/slack-gif-creator/LICENSE.txt +++ /dev/null @@ -1,202 +0,0 @@ - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. \ No newline at end of file diff --git a/tests/llm_translation/test_skills_data/slack-gif-creator/SKILL.md b/tests/llm_translation/test_skills_data/slack-gif-creator/SKILL.md index 16660d8ceb7..3cae971b731 100644 --- a/tests/llm_translation/test_skills_data/slack-gif-creator/SKILL.md +++ b/tests/llm_translation/test_skills_data/slack-gif-creator/SKILL.md @@ -1,7 +1,6 @@ --- name: slack-gif-creator description: Knowledge and utilities for creating animated GIFs optimized for Slack. Provides constraints, validation tools, and animation concepts. Use when users request animated GIFs for Slack like "make me a GIF of X doing Y for Slack." -license: Complete terms in LICENSE.txt --- # Slack GIF Creator diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index bae0b15dfec..c22c3537af8 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -3,6 +3,7 @@ import os import sys import pytest from typing import List, Any, cast +from unittest.mock import AsyncMock, patch sys.path.insert(0, os.path.abspath("../../..")) @@ -254,6 +255,76 @@ async def test_aresponses_api_with_mcp_mock_integration(): print(f"Other tools parsed: {len(other_parsed)}") +@pytest.mark.asyncio +async def test_aresponses_api_with_mcp_passes_mcp_server_auth_headers_to_process_tools(): + """ + Test that MCP auth headers from secret_fields (e.g. x-mcp-linear_config-authorization) + are passed to _process_mcp_tools_without_openai_transform when using the responses API. + """ + from litellm.responses.main import aresponses_api_with_mcp + + captured_process_kwargs = {} + + async def mock_process(**kwargs): + captured_process_kwargs.update(kwargs) + return ([], {}) + + mock_response = ResponsesAPIResponse( + **{ + "id": "resp_test", + "object": "response", + "created_at": 1234567890, + "status": "completed", + "error": None, + "incomplete_details": None, + "instructions": None, + "max_output_tokens": None, + "model": "gpt-4o", + "output": [{"type": "message", "id": "msg_1", "status": "completed", "role": "assistant", "content": []}], + "parallel_tool_calls": True, + "previous_response_id": None, + "reasoning": {"effort": None, "summary": None}, + "store": True, + "temperature": 1.0, + "text": {"format": {"type": "text"}}, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + "truncation": "disabled", + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + "user": None, + "metadata": {}, + } + ) + + mcp_tools = [{"type": "mcp", "server_url": "litellm_proxy"}] + secret_fields = { + "raw_headers": {"x-mcp-linear_config-authorization": "Bearer linear-token"}, + } + + with patch.object( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + mock_process, + ), patch( + "litellm.responses.main.aresponses", + new_callable=AsyncMock, + return_value=mock_response, + ): + await aresponses_api_with_mcp( + input=[{"role": "user", "type": "message", "content": "hi"}], + model="gpt-4o", + tools=mcp_tools, + secret_fields=secret_fields, + ) + + assert "mcp_server_auth_headers" in captured_process_kwargs + mcp_server_auth_headers = captured_process_kwargs["mcp_server_auth_headers"] + assert mcp_server_auth_headers is not None + assert "linear_config" in mcp_server_auth_headers + assert mcp_server_auth_headers["linear_config"]["Authorization"] == "Bearer linear-token" + + @pytest.mark.asyncio async def test_mcp_allowed_tools_filtering(): """ diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 930b0a03042..a81702d0db3 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -427,7 +427,8 @@ async def test_streamable_http_mcp_handler_mock(): # Call the handler await handle_streamable_http_mcp(mock_scope, mock_receive, mock_send) - # Verify session manager handle_request was called + # Verify session manager handle_request was called with correct args + # send is passed directly (no wrapper) mock_session_manager.handle_request.assert_called_once_with( mock_scope, mock_receive, mock_send ) @@ -658,9 +659,9 @@ async def test_list_tools_rest_api_server_not_found(): mock_manager.get_allowed_mcp_servers = AsyncMock( return_value=["non_existent_server_id"] ) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) # Return None when trying to get the server (server doesn't exist) mock_manager.get_mcp_server_by_id = MagicMock(return_value=None) @@ -731,9 +732,9 @@ async def test_list_tools_rest_api_success(): return_value=["test-server-123"] ) mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) # Mock the _get_tools_for_single_server function @@ -813,9 +814,9 @@ async def test_get_tools_from_mcp_servers(): ) mock_manager.get_mcp_server_by_id = lambda server_id: mock_server_1 if server_id == "server1_id" else mock_server_2 mock_manager._get_tools_from_server = AsyncMock(return_value=[mock_tool_1]) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) with patch( @@ -852,9 +853,9 @@ async def test_get_tools_from_mcp_servers(): mock_manager_2._get_tools_from_server = AsyncMock( side_effect=mock_get_tools_side_effect ) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager_2.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager_2.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) with patch( @@ -880,9 +881,9 @@ async def test_get_tools_from_mcp_servers(): ) mock_manager.get_mcp_server_by_id = lambda server_id: mock_server_1 if server_id == "server1_id" else (mock_server_2 if server_id == "server2_id" else mock_server_3) mock_manager._get_tools_from_server = AsyncMock(return_value=[mock_tool_1]) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) with patch( @@ -1816,9 +1817,9 @@ async def test_list_tool_rest_api_with_server_specific_auth(): mock_server.mcp_info = {"server_name": "zapier"} mock_manager.get_mcp_server_by_id.return_value = mock_server - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) mock_user_api_key_dict = UserAPIKeyAuth( @@ -1910,9 +1911,9 @@ async def test_list_tool_rest_api_with_default_auth(): mock_server.mcp_info = {"server_name": "unknown_server"} mock_manager.get_mcp_server_by_id.return_value = mock_server - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) mock_user_api_key_dict = UserAPIKeyAuth( @@ -2020,9 +2021,9 @@ async def test_list_tool_rest_api_all_servers_with_auth(): server_id ) ) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) mock_user_api_key_dict = UserAPIKeyAuth( @@ -2153,9 +2154,9 @@ async def test_filter_tools_by_allowed_tools_integration(): return_value=["test-server-123"] ) mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) # Mock the _get_tools_from_server method to return all tools @@ -2267,9 +2268,9 @@ async def test_filter_tools_by_disallowed_tools_integration(): return_value=["test-server-456"] ) mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) # Mock the _get_tools_from_server method to return all tools mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools) @@ -2367,9 +2368,9 @@ async def test_filter_tools_no_restrictions_integration(): return_value=["test-server-000"] ) mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server) - # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test) - mock_manager.filter_server_ids_by_ip = MagicMock( - side_effect=lambda server_ids, client_ip: server_ids + # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + side_effect=lambda server_ids, client_ip: (server_ids, 0) ) # Mock the _get_tools_from_server method to return all tools diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index c3f68762810..ed528f21e0d 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3668,9 +3668,10 @@ async def test_list_keys(prisma_client): async def test_key_aliases(prisma_client): """ Test the key_aliases function: - - Returns a list + - Returns a paginated response - Includes alias from a newly created key - - Aliases are unique and sorted + - Aliases are sorted + - Pagination and search params work correctly """ import asyncio import uuid @@ -3682,10 +3683,16 @@ async def test_key_aliases(prisma_client): setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") await litellm.proxy.proxy_server.prisma_client.connect() - # Basic call - response = await key_aliases() + # Basic call - check pagination response shape + response = await key_aliases(page=1, size=50) assert "aliases" in response assert isinstance(response["aliases"], list) + assert "total_count" in response + assert "current_page" in response + assert "total_pages" in response + assert "size" in response + assert response["current_page"] == 1 + assert response["size"] == 50 # Create a new user (and key) with a unique alias unique_id = str(uuid.uuid4()) @@ -3704,17 +3711,22 @@ async def test_key_aliases(prisma_client): # Allow async DB writes to settle await asyncio.sleep(2) - # Call again and validate - response_after = await key_aliases() + # Call again and validate alias is present + response_after = await key_aliases(page=1, size=50) aliases = response_after["aliases"] - - # Contains the new alias assert test_alias in aliases - - # Unique & sorted (endpoint dedupes and orders ascending) - assert len(aliases) == len(set(aliases)) assert aliases == sorted(aliases) + # Search by partial alias + partial = test_alias[:10] + search_response = await key_aliases(page=1, size=50, search=partial) + assert test_alias in search_response["aliases"] + + # Search with no match + no_match_response = await key_aliases(page=1, size=50, search="__no_match_xyz__") + assert len(no_match_response["aliases"]) == 0 + assert no_match_response["total_count"] == 0 + @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") @pytest.mark.asyncio diff --git a/tests/proxy_unit_tests/test_project_endpoints_prisma.py b/tests/proxy_unit_tests/test_project_endpoints_prisma.py index c98cb7efda0..77ed09a40f0 100644 --- a/tests/proxy_unit_tests/test_project_endpoints_prisma.py +++ b/tests/proxy_unit_tests/test_project_endpoints_prisma.py @@ -791,3 +791,78 @@ def test_litellm_entity_type_has_project(): assert hasattr(Litellm_EntityType, "PROJECT") assert Litellm_EntityType.PROJECT.value == "project" + + +@pytest.mark.asyncio +async def test_list_projects_returns_timestamps(): + """ + Test that /project/list returns created_at and updated_at for each project. + """ + from datetime import datetime, timezone + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.management_endpoints.project_endpoints import list_projects + from litellm.proxy._types import LiteLLM_ProjectTable + + now = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc) + + # Build a fake DB row that includes created_at and updated_at + fake_project = MagicMock() + fake_project.model_dump.return_value = { + "project_id": "proj-1", + "project_alias": "test-project", + "team_id": "team-1", + "created_by": "admin", + "updated_by": "admin", + "created_at": now, + "updated_at": now, + "models": [], + "spend": 0.0, + "blocked": False, + "budget_id": None, + "description": None, + "metadata": None, + "model_spend": None, + "model_rpm_limit": None, + "model_tpm_limit": None, + "object_permission_id": None, + "litellm_budget_table": None, + "object_permission": None, + } + # Make the fake row behave like a Pydantic model for FastAPI serialization + fake_project.project_id = "proj-1" + fake_project.created_at = now + fake_project.updated_at = now + + mock_prisma = MagicMock() + mock_prisma.db.litellm_projecttable.find_many = AsyncMock( + return_value=[fake_project] + ) + + with patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ): + response = await list_projects( + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + assert len(response) == 1 + project = response[0] + assert project.created_at == now + assert project.updated_at == now + + +def test_litellm_project_table_has_timestamp_fields(): + """ + Test that LiteLLM_ProjectTable model includes created_at and updated_at fields, + so the /project/list response_model exposes them. + """ + from litellm.proxy._types import LiteLLM_ProjectTable + + fields = LiteLLM_ProjectTable.model_fields + assert "created_at" in fields, "LiteLLM_ProjectTable must have created_at field" + assert "updated_at" in fields, "LiteLLM_ProjectTable must have updated_at field" diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index cc1cc278392..fd97a38b41e 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -2,7 +2,8 @@ import asyncio import json import os import sys -from typing import Any, Dict, List, Optional +from datetime import datetime +from typing import Any, Dict, List, Optional, Union from unittest.mock import Mock import pytest @@ -1486,12 +1487,46 @@ class MockPrismaClientDB: mock_key_data, ): self.db = MockDb(mock_team_data, mock_key_data) + + async def get_data( + self, + token: Optional[Union[str, list]] = None, + user_id: Optional[str] = None, + user_id_list: Optional[list] = None, + team_id: Optional[str] = None, + team_id_list: Optional[list] = None, + key_val: Optional[dict] = None, + table_name: Optional[str] = None, + query_type: str = "find_unique", + expires: Optional[datetime] = None, + reset_at: Optional[datetime] = None, + offset: Optional[int] = None, + limit: Optional[int] = None, + ): + """Mock get_data method to return user info for admin""" + from litellm.proxy._types import LiteLLM_UserTable + + # Return a proper LiteLLM_UserTable object when querying by user_id + if user_id: + return LiteLLM_UserTable( + user_id=user_id, + user_role="proxy_admin", + spend=0.0, + max_budget=None, + ) + return None @pytest.mark.asyncio async def test_get_user_info_for_proxy_admin(mock_team_data, mock_key_data): # Patch the prisma_client import - from litellm.proxy._types import UserInfoResponse + from litellm.proxy._types import UserAPIKeyAuth, UserInfoResponse + + # Create a mock user_api_key_dict for admin user + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user_123", + user_role="proxy_admin", + ) with patch( "litellm.proxy.proxy_server.prisma_client", @@ -1502,11 +1537,18 @@ async def test_get_user_info_for_proxy_admin(mock_team_data, mock_key_data): ) # Execute the function - result = await _get_user_info_for_proxy_admin() + result = await _get_user_info_for_proxy_admin( + user_api_key_dict=mock_user_api_key_dict + ) # Verify the result structure assert isinstance(result, UserInfoResponse) assert len(result.keys) == 2 + # Verify admin's user_id is populated + assert result.user_id == "admin_user_123" + # Verify admin's user_info is populated + assert result.user_info is not None + assert result.user_info["user_id"] == "admin_user_123" def test_custom_openid_response(): diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 352db384c88..030d452e55f 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -158,3 +158,109 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}" ) assert call_kwargs["response_cost"] == 0.05 + + +# Test is_end_user_within_model_budget +@pytest.mark.asyncio +async def test_is_end_user_within_model_budget(budget_limiter): + # Test when model is within budget + with patch.object( + budget_limiter, "_get_end_user_spend_for_model", return_value=50.0 + ): + assert ( + await budget_limiter.is_end_user_within_model_budget( + "test-user", + {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, + "gpt-4", + ) + is True + ) + + # Test when model exceeds budget + with patch.object( + budget_limiter, "_get_end_user_spend_for_model", return_value=150.0 + ): + with pytest.raises(litellm.BudgetExceededError): + await budget_limiter.is_end_user_within_model_budget( + "test-user", + {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, + "gpt-4", + ) + + # Test model not in budget config + assert ( + await budget_limiter.is_end_user_within_model_budget( + "test-user", + {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, + "non-existent", + ) + is True + ) + + +# Test _get_end_user_spend_for_model +@pytest.mark.asyncio +async def test_get_end_user_spend_for_model(budget_limiter): + budget_config = GenericBudgetInfo(budget_limit=100.0, time_period="1d") + + # Mock cache get + with patch.object(budget_limiter.dual_cache, "async_get_cache", return_value=50.0): + spend = await budget_limiter._get_end_user_spend_for_model( + end_user_id="test-user", model="gpt-4", key_budget_config=budget_config + ) + assert spend == 50.0 + + # Test with provider prefix + spend = await budget_limiter._get_end_user_spend_for_model( + end_user_id="test-user", + model="openai/gpt-4", + key_budget_config=budget_config, + ) + assert spend == 50.0 + + +@pytest.mark.asyncio +async def test_async_log_success_event_uses_end_user_model_budget_duration( + budget_limiter, +): + """ + async_log_success_event must use the per-model budget_duration for the end user cache key + """ + from litellm.proxy.hooks.model_max_budget_limiter import ( + END_USER_SPEND_CACHE_KEY_PREFIX, + ) + + end_user_id = "test-user" + model = "gpt-4" + budget_duration = "1d" + user_api_key_end_user_model_max_budget = { + model: {"budget_limit": 100.0, "time_period": budget_duration}, + } + kwargs = { + "standard_logging_object": { + "response_cost": 0.05, + "model": model, + "end_user": end_user_id, + "metadata": {"user_api_key_end_user_id": end_user_id}, + }, + "litellm_params": { + "metadata": { + "user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget + }, + }, + } + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + call_kwargs = mock_increment.call_args.kwargs + spend_key = call_kwargs["spend_key"] + assert spend_key == ( + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model}:{budget_duration}" + ) + assert call_kwargs["response_cost"] == 0.05 diff --git a/tests/test_litellm/caching/test_llm_caching_handler.py b/tests/test_litellm/caching/test_llm_caching_handler.py new file mode 100644 index 00000000000..0ac4ac5de79 --- /dev/null +++ b/tests/test_litellm/caching/test_llm_caching_handler.py @@ -0,0 +1,158 @@ +import asyncio +import os +import sys +import warnings + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.caching.llm_caching_handler import LLMClientCache + + +class MockAsyncClient: + """Mock async HTTP client with an async close method.""" + + def __init__(self): + self.closed = False + + async def close(self): + self.closed = True + + +class MockSyncClient: + """Mock sync HTTP client with a sync close method.""" + + def __init__(self): + self.closed = False + + def close(self): + self.closed = True + + +@pytest.mark.asyncio +async def test_remove_key_no_unawaited_coroutine_warning(): + """ + Test that evicting an async client from LLMClientCache does not produce + 'coroutine was never awaited' warnings. + + Regression test for https://github.com/BerriAI/litellm/issues/22128 + """ + cache = LLMClientCache(max_size_in_memory=2) + + mock_client = MockAsyncClient() + cache.cache_dict["test-key"] = mock_client + cache.ttl_dict["test-key"] = 0 # expired + + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter("always") + cache._remove_key("test-key") + # Let the event loop process the close task + await asyncio.sleep(0.1) + + coroutine_warnings = [ + w for w in caught_warnings if "coroutine" in str(w.message).lower() + ] + assert ( + len(coroutine_warnings) == 0 + ), f"Got unawaited coroutine warnings: {coroutine_warnings}" + + +@pytest.mark.asyncio +async def test_remove_key_closes_async_client(): + """ + Test that evicting an async client from the cache properly closes it. + """ + cache = LLMClientCache(max_size_in_memory=2) + + mock_client = MockAsyncClient() + cache.cache_dict["test-key"] = mock_client + cache.ttl_dict["test-key"] = 0 + + cache._remove_key("test-key") + # Let the event loop process the close task + await asyncio.sleep(0.1) + + assert mock_client.closed is True + assert "test-key" not in cache.cache_dict + assert "test-key" not in cache.ttl_dict + + +def test_remove_key_closes_sync_client(): + """ + Test that evicting a sync client from the cache properly closes it. + """ + cache = LLMClientCache(max_size_in_memory=2) + + mock_client = MockSyncClient() + cache.cache_dict["test-key"] = mock_client + cache.ttl_dict["test-key"] = 0 + + cache._remove_key("test-key") + + assert mock_client.closed is True + assert "test-key" not in cache.cache_dict + + +@pytest.mark.asyncio +async def test_eviction_closes_async_clients(): + """ + Test that cache eviction (when cache is full) properly closes async clients + without producing warnings. + """ + cache = LLMClientCache(max_size_in_memory=2, default_ttl=1) + + clients = [] + for i in range(2): + client = MockAsyncClient() + clients.append(client) + cache.set_cache(f"key-{i}", client) + + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter("always") + # This should trigger eviction of one of the existing entries + cache.set_cache("key-new", "new-value") + await asyncio.sleep(0.1) + + coroutine_warnings = [ + w for w in caught_warnings if "coroutine" in str(w.message).lower() + ] + assert ( + len(coroutine_warnings) == 0 + ), f"Got unawaited coroutine warnings: {coroutine_warnings}" + + +def test_remove_key_no_event_loop(): + """ + Test that _remove_key doesn't raise when there's no running event loop + (falls through to the RuntimeError except branch). + """ + cache = LLMClientCache(max_size_in_memory=2) + + mock_client = MockAsyncClient() + cache.cache_dict["test-key"] = mock_client + cache.ttl_dict["test-key"] = 0 + + # Should not raise even though there's no running event loop + cache._remove_key("test-key") + assert "test-key" not in cache.cache_dict + + +@pytest.mark.asyncio +async def test_background_tasks_cleaned_up_after_completion(): + """ + Test that completed close tasks are removed from the _background_tasks set. + """ + cache = LLMClientCache(max_size_in_memory=2) + + mock_client = MockAsyncClient() + cache.cache_dict["test-key"] = mock_client + cache.ttl_dict["test-key"] = 0 + + cache._remove_key("test-key") + # Let the task complete + await asyncio.sleep(0.1) + + assert len(cache._background_tasks) == 0 diff --git a/tests/test_litellm/caching/test_llm_client_cache_e2e.py b/tests/test_litellm/caching/test_llm_client_cache_e2e.py new file mode 100644 index 00000000000..a7d012d2269 --- /dev/null +++ b/tests/test_litellm/caching/test_llm_client_cache_e2e.py @@ -0,0 +1,47 @@ +"""e2e tests: httpx clients obtained via get_async_httpx_client must remain +usable after LLMClientCache evicts their cache entry.""" + +import pytest + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + + +@pytest.fixture(autouse=True) +def _tiny_client_cache(monkeypatch): + """Replace the global client cache with a size-1 cache so eviction + triggers on the second insert.""" + cache = LLMClientCache(max_size_in_memory=1, default_ttl=600) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", cache) + yield cache + + +@pytest.mark.asyncio +async def test_evicted_client_is_not_closed(): + """Get a client via get_async_httpx_client, evict it by caching a second + one, then verify the first client's transport is still open.""" + client_a = get_async_httpx_client(llm_provider="provider_a") + # This evicts client_a from cache (capacity=1) + client_b = get_async_httpx_client(llm_provider="provider_b") + + assert not client_a.client.is_closed + await client_a.client.aclose() + await client_b.client.aclose() + + +@pytest.mark.asyncio +async def test_expired_client_is_not_closed(): + """Get a client, expire it via TTL, then verify the client is still open.""" + cache = litellm.in_memory_llm_clients_cache + client = get_async_httpx_client(llm_provider="provider_ttl") + + # Force the entry to expire and trigger eviction + for key in list(cache.ttl_dict.keys()): + cache.ttl_dict[key] = 0 + # Also fix the heap entry so evict_cache finds it + cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap] + cache.evict_cache() + + assert not client.client.is_closed + await client.client.aclose() diff --git a/tests/test_litellm/caching/test_redis_connection_pool.py b/tests/test_litellm/caching/test_redis_connection_pool.py index b8922846e82..f6e429ceff9 100644 --- a/tests/test_litellm/caching/test_redis_connection_pool.py +++ b/tests/test_litellm/caching/test_redis_connection_pool.py @@ -4,14 +4,12 @@ Regression tests for Redis connection pool leak fixes (RC1-RC5). Tests are pure unit tests — no Redis server required. """ -import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest import redis.asyncio as async_redis from litellm._redis import get_redis_async_client, get_redis_connection_pool -from litellm.caching.llm_caching_handler import LLMClientCache def test_url_config_uses_passed_pool(): @@ -131,37 +129,3 @@ async def test_disconnect_idempotent(): await cache.disconnect() # should not raise -@pytest.mark.asyncio -async def test_eviction_calls_aclose(): - """When an async client is evicted from LLMClientCache, its aclose() - should be scheduled via create_task.""" - cache = LLMClientCache(max_size_in_memory=2, default_ttl=600) - - client = AsyncMock() - client.aclose = AsyncMock() - - cache.set_cache(key="client-0", value=client) - cache.set_cache(key="filler", value="x") - # Third insert triggers eviction of client-0 - cache.set_cache(key="trigger", value="y") - - # Let the scheduled task run - await asyncio.sleep(0.05) - - assert client.aclose.await_count > 0 - - -@pytest.mark.asyncio -async def test_eviction_non_closeable_safe(): - """Evicting plain values (strings, dicts, ints) should not crash.""" - cache = LLMClientCache(max_size_in_memory=2, default_ttl=600) - - cache.set_cache(key="str-val", value="hello") - cache.set_cache(key="dict-val", value={"foo": "bar"}) - # This evicts "str-val" — should not raise - cache.set_cache(key="int-val", value=42) - - await asyncio.sleep(0.05) - - # If we got here without exception, the test passes - assert cache.get_cache(key="int-val") == 42 diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index a840e2fe162..cd76ba1e863 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -1,7 +1,8 @@ """ Unit tests for Prometheus user and team count metrics """ -from unittest.mock import MagicMock +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch import pytest from prometheus_client import REGISTRY @@ -258,3 +259,267 @@ class TestPrometheusUserTeamCountMetrics: assert True except Exception as e: pytest.fail(f"Metrics should handle large values: {e}") + + +# --------------------------------------------------------------------------- +# Regression tests: team budget showing +Inf when user_api_key_team_max_budget +# is None in request metadata but the team has a real budget in the DB. +# --------------------------------------------------------------------------- + + +async def test_assemble_team_object_uses_db_max_budget_when_metadata_is_none( + prometheus_logger, +): + """ + When max_budget is None in request metadata (e.g. stale key cache), + _assemble_team_object must fall back to the value returned by get_team_object + so that _safe_get_remaining_budget does not return +Inf. + """ + db_team = MagicMock() + db_team.max_budget = 3000.0 + db_team.budget_reset_at = datetime(2026, 3, 1, tzinfo=timezone.utc) + + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = db_team + team_object = await prometheus_logger._assemble_team_object( + team_id="c5c33858-4379-4c90-8733-d9c58c312c10", + team_alias="ai-ml-local_dev", + spend=1617.02, + max_budget=None, # simulates None coming from request metadata + response_cost=0.5, + ) + + assert team_object.max_budget == 3000.0, ( + "max_budget should be populated from DB when metadata value is None" + ) + assert team_object.budget_reset_at == datetime(2026, 3, 1, tzinfo=timezone.utc) + + +async def test_assemble_team_object_does_not_override_metadata_max_budget( + prometheus_logger, +): + """ + When max_budget IS present in request metadata, it must not be overridden + by the DB value. + """ + db_team = MagicMock() + db_team.max_budget = 9999.0 + db_team.budget_reset_at = None + + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = db_team + team_object = await prometheus_logger._assemble_team_object( + team_id="team-1", + team_alias="my-team", + spend=50.0, + max_budget=100.0, # metadata has a real value + response_cost=1.0, + ) + + assert team_object.max_budget == 100.0, ( + "max_budget from metadata must not be replaced by the DB value" + ) + + +async def test_set_team_budget_metrics_after_api_request_no_inf_when_metadata_budget_none( + prometheus_logger, +): + """ + End-to-end: when user_api_key_team_max_budget is None in request metadata + but the team has a real budget in the DB, the metric must NOT be set to +Inf. + """ + prometheus_logger.litellm_remaining_team_budget_metric = MagicMock() + prometheus_logger.litellm_team_max_budget_metric = MagicMock() + prometheus_logger.litellm_team_budget_remaining_hours_metric = MagicMock() + + db_team = MagicMock() + db_team.max_budget = 3000.0 + db_team.budget_reset_at = datetime(2026, 3, 1, tzinfo=timezone.utc) + + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = db_team + await prometheus_logger._set_team_budget_metrics_after_api_request( + user_api_team="c5c33858-4379-4c90-8733-d9c58c312c10", + user_api_team_alias="ai-ml-local_dev", + team_spend=1617.02, + team_max_budget=None, # simulates stale key cache + response_cost=0.5, + ) + + set_call_args = ( + prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args + ) + assert set_call_args is not None, "remaining_team_budget_metric.labels().set was not called" + actual_value = set_call_args[0][0] + assert actual_value != float("inf"), ( + f"remaining_team_budget_metric must not be +Inf when team has a real budget; got {actual_value}" + ) + expected = 3000.0 - 1617.02 - 0.5 + assert abs(actual_value - expected) < 0.01, ( + f"Expected remaining budget ~{expected}, got {actual_value}" + ) + + +async def test_set_team_budget_metrics_after_api_request_inf_when_genuinely_no_budget( + prometheus_logger, +): + """ + When the team genuinely has no budget (max_budget=None in both metadata and + DB), +Inf is the correct value and must be preserved. + """ + prometheus_logger.litellm_remaining_team_budget_metric = MagicMock() + prometheus_logger.litellm_team_max_budget_metric = MagicMock() + prometheus_logger.litellm_team_budget_remaining_hours_metric = MagicMock() + + db_team = MagicMock() + db_team.max_budget = None + db_team.budget_reset_at = None + + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = db_team + await prometheus_logger._set_team_budget_metrics_after_api_request( + user_api_team="team-no-budget", + user_api_team_alias="no-budget-team", + team_spend=10.0, + team_max_budget=None, + response_cost=1.0, + ) + + set_call_args = ( + prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args + ) + assert set_call_args is not None + actual_value = set_call_args[0][0] + assert actual_value == float("inf"), ( + "remaining_team_budget_metric should be +Inf when team truly has no budget" + ) + + +# --------------------------------------------------------------------------- +# Regression tests: user budget showing +Inf when user_api_key_user_max_budget +# is None in request metadata but the user has a real budget in the DB. +# --------------------------------------------------------------------------- + + +async def test_assemble_user_object_uses_db_max_budget_when_metadata_is_none( + prometheus_logger, +): + """ + When max_budget is None in request metadata (e.g. stale key cache), + _assemble_user_object must fall back to the value returned by get_user_object + so that _safe_get_remaining_budget does not return +Inf. + """ + db_user = MagicMock() + db_user.max_budget = 500.0 + db_user.budget_reset_at = datetime(2026, 3, 1, tzinfo=timezone.utc) + + with patch("litellm.proxy.auth.auth_checks.get_user_object") as mock_get_user: + mock_get_user.return_value = db_user + user_object = await prometheus_logger._assemble_user_object( + user_id="user-abc-123", + spend=120.0, + max_budget=None, # simulates None coming from request metadata + response_cost=0.5, + ) + + assert user_object.max_budget == 500.0, ( + "max_budget should be populated from DB when metadata value is None" + ) + assert user_object.budget_reset_at == datetime(2026, 3, 1, tzinfo=timezone.utc) + + +async def test_assemble_user_object_does_not_override_metadata_max_budget( + prometheus_logger, +): + """ + When max_budget IS present in request metadata, it must not be overridden + by the DB value. + """ + db_user = MagicMock() + db_user.max_budget = 9999.0 + db_user.budget_reset_at = None + + with patch("litellm.proxy.auth.auth_checks.get_user_object") as mock_get_user: + mock_get_user.return_value = db_user + user_object = await prometheus_logger._assemble_user_object( + user_id="user-abc-123", + spend=50.0, + max_budget=100.0, # metadata has a real value + response_cost=1.0, + ) + + assert user_object.max_budget == 100.0, ( + "max_budget from metadata must not be replaced by the DB value" + ) + + +async def test_set_user_budget_metrics_after_api_request_no_inf_when_metadata_budget_none( + prometheus_logger, +): + """ + End-to-end: when user_max_budget is None in request metadata but the user + has a real budget in the DB, the metric must NOT be set to +Inf. + """ + prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() + prometheus_logger.litellm_user_max_budget_metric = MagicMock() + prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() + + db_user = MagicMock() + db_user.max_budget = 500.0 + db_user.budget_reset_at = datetime(2026, 3, 1, tzinfo=timezone.utc) + + with patch("litellm.proxy.auth.auth_checks.get_user_object") as mock_get_user: + mock_get_user.return_value = db_user + await prometheus_logger._set_user_budget_metrics_after_api_request( + user_id="user-abc-123", + user_spend=120.0, + user_max_budget=None, # simulates stale key cache + response_cost=0.5, + ) + + set_call_args = ( + prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args + ) + assert set_call_args is not None, "remaining_user_budget_metric.labels().set was not called" + actual_value = set_call_args[0][0] + assert actual_value != float("inf"), ( + f"remaining_user_budget_metric must not be +Inf when user has a real budget; got {actual_value}" + ) + expected = 500.0 - 120.0 - 0.5 + assert abs(actual_value - expected) < 0.01, ( + f"Expected remaining budget ~{expected}, got {actual_value}" + ) + + +async def test_set_user_budget_metrics_after_api_request_inf_when_genuinely_no_budget( + prometheus_logger, +): + """ + When the user genuinely has no budget (max_budget=None in both metadata and + DB), +Inf is the correct value and must be preserved. + """ + prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() + prometheus_logger.litellm_user_max_budget_metric = MagicMock() + prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() + + db_user = MagicMock() + db_user.max_budget = None + db_user.budget_reset_at = None + + with patch("litellm.proxy.auth.auth_checks.get_user_object") as mock_get_user: + mock_get_user.return_value = db_user + await prometheus_logger._set_user_budget_metrics_after_api_request( + user_id="user-no-budget", + user_spend=10.0, + user_max_budget=None, + response_cost=1.0, + ) + + set_call_args = ( + prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args + ) + assert set_call_args is not None + actual_value = set_call_args[0][0] + assert actual_value == float("inf"), ( + "remaining_user_budget_metric should be +Inf when user truly has no budget" + ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index b45cbbd99c0..7e8848be301 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -9,9 +9,19 @@ import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) +from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as gemini_image_generation_cost_calculator, +) +from litellm.llms.vertex_ai.image_generation.cost_calculator import ( + cost_calculator as vertex_image_generation_cost_calculator, +) from litellm.types.llms.openai import FileSearchTool, WebSearchOptions from litellm.types.utils import ( CompletionTokensDetailsWrapper, + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, ModelInfo, ModelResponse, PromptTokensDetailsWrapper, @@ -766,7 +776,14 @@ def test_service_tier_fallback_pricing(): assert abs(std_cost[1] - expected_standard_completion) < 1e-10, f"Standard completion cost mismatch: {std_cost[1]} vs {expected_standard_completion}" -def test_gemini_image_generation_cost_with_zero_text_tokens(): +@pytest.mark.parametrize( + "model", + [ + "gemini-3-pro-image-preview", + "gemini-3.1-flash-image-preview", + ], +) +def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): """ Test that image_tokens are correctly costed when text_tokens=0. @@ -779,7 +796,6 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") - model = "gemini-3-pro-image-preview" custom_llm_provider = "vertex_ai" # Usage from the issue: text_tokens=0, image_tokens=1120, reasoning_tokens=225 @@ -809,9 +825,9 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(): # Expected costs: # - text_tokens: 0 * output_cost_per_token = 0 - # - image_tokens: 1120 * output_cost_per_image_token = 1120 * 1.2e-04 = 0.1344 - # - reasoning_tokens: 225 * output_cost_per_token = 225 * 1.2e-05 = 0.0027 - # Total completion: ~0.1371 + # - image_tokens: 1120 * output_cost_per_image_token + # - reasoning_tokens: 225 * output_cost_per_token + # Total completion should include both image + reasoning costs. output_cost_per_image_token = model_cost_map.get("output_cost_per_image_token", 0) output_cost_per_token = model_cost_map.get("output_cost_per_token", 0) @@ -820,18 +836,151 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(): expected_reasoning_cost = 225 * output_cost_per_token # reasoning uses base token cost expected_completion_cost = expected_image_cost + expected_reasoning_cost - # The bug was: all 1345 tokens were treated as text = 1345 * 1.2e-05 = 0.01614 - # Fixed: image_tokens use image pricing = ~0.137 - - assert completion_cost > 0.10, ( - f"Completion cost should be > $0.10 (image tokens are expensive), got ${completion_cost:.6f}. " - f"Bug: tokens may be incorrectly treated as text tokens." + # The bug was: all completion tokens were treated as text tokens only. + bugged_text_only_cost = 1345 * output_cost_per_token + assert completion_cost > bugged_text_only_cost * 2, ( + f"Completion cost should be significantly larger than text-only bugged path. " + f"Expected > {bugged_text_only_cost * 2:.6f}, got {completion_cost:.6f}" ) assert round(completion_cost, 4) == round(expected_completion_cost, 4), ( f"Expected completion cost ${expected_completion_cost:.6f}, got ${completion_cost:.6f}" ) +def test_vertex_image_generation_cost_prefers_token_usage_metadata(): + """ + When usage metadata exists on image responses, Vertex image generation cost + should be calculated from token pricing, not flat output_cost_per_image. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gemini-3.1-flash-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") + + input_text_tokens = 50 + input_image_tokens = 1120 + output_image_tokens = 1120 + prompt_tokens = input_text_tokens + input_image_tokens + + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")], + usage=ImageUsage( + input_tokens=prompt_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=input_text_tokens, + image_tokens=input_image_tokens, + ), + output_tokens=output_image_tokens, + total_tokens=prompt_tokens + output_image_tokens, + ), + ) + + cost = vertex_image_generation_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_prompt_cost = prompt_tokens * model_info["input_cost_per_token"] + expected_completion_cost = output_image_tokens * model_info["output_cost_per_image_token"] + expected_total_cost = expected_prompt_cost + expected_completion_cost + + assert round(cost, 10) == round(expected_total_cost, 10) + # Ensure this is not falling back to flat per-image pricing. + assert cost != len(image_response.data) * model_info["output_cost_per_image"] + + +def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(): + """ + Without usage metadata, Vertex image generation cost should fall back to + output_cost_per_image * number_of_images. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gemini-3.1-flash-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") + + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] + ) + + cost = vertex_image_generation_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_cost = len(image_response.data) * model_info["output_cost_per_image"] + assert round(cost, 10) == round(expected_cost, 10) + + +def test_gemini_image_generation_cost_prefers_token_usage_metadata(): + """ + When usage metadata exists on image responses, Gemini image generation cost + should be calculated from token pricing, not flat output_cost_per_image. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gemini/gemini-3-pro-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + input_text_tokens = 20 + input_image_tokens = 1120 + output_image_tokens = 1120 + prompt_tokens = input_text_tokens + input_image_tokens + + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")], + usage=ImageUsage( + input_tokens=prompt_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=input_text_tokens, + image_tokens=input_image_tokens, + ), + output_tokens=output_image_tokens, + total_tokens=prompt_tokens + output_image_tokens, + ), + ) + + cost = gemini_image_generation_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_prompt_cost = prompt_tokens * model_info["input_cost_per_token"] + expected_completion_cost = output_image_tokens * model_info["output_cost_per_image_token"] + expected_total_cost = expected_prompt_cost + expected_completion_cost + + assert round(cost, 10) == round(expected_total_cost, 10) + # Ensure this is not falling back to flat per-image pricing. + assert cost != len(image_response.data) * model_info["output_cost_per_image"] + + +def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing(): + """ + Without usage metadata, Gemini image generation cost should fall back to + output_cost_per_image * number_of_images. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model = "gemini/gemini-3-pro-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] + ) + + cost = gemini_image_generation_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_cost = len(image_response.data) * model_info["output_cost_per_image"] + assert round(cost, 10) == round(expected_cost, 10) + + def test_bedrock_anthropic_prompt_caching(): """Test Bedrock Anthropic models with prompt caching return correct costs.""" model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" diff --git a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py new file mode 100644 index 00000000000..f09bdfae649 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py @@ -0,0 +1,126 @@ +""" +Tests for LITELLM_MAX_STREAMING_DURATION_SECONDS — the global cap on streaming response wall-clock time. + +Covers: + - CustomStreamWrapper (chat/completions) sync + async + - BaseResponsesAPIStreamingIterator (responses) sync + async +""" + +import os +import sys +import time +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_custom_stream_wrapper() -> CustomStreamWrapper: + """Build a minimal CustomStreamWrapper for testing.""" + return CustomStreamWrapper( + completion_stream=None, + model="test-model", + logging_obj=MagicMock(), + custom_llm_provider="openai", + ) + + +# --------------------------------------------------------------------------- +# CustomStreamWrapper (chat/completions) +# --------------------------------------------------------------------------- + + +class TestCustomStreamWrapperMaxDuration: + def test_should_not_raise_when_duration_is_none(self): + """No limit configured → never raises.""" + wrapper = _make_custom_stream_wrapper() + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", None): + wrapper._check_max_streaming_duration() # should not raise + + def test_should_not_raise_when_under_limit(self): + """Stream is under the limit → no error.""" + wrapper = _make_custom_stream_wrapper() + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 60.0): + wrapper._check_max_streaming_duration() # should not raise + + def test_should_raise_timeout_when_exceeded(self): + """Stream exceeded the limit → litellm.Timeout.""" + wrapper = _make_custom_stream_wrapper() + wrapper._stream_created_time = time.time() - 20 # simulate 20s elapsed + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): + with pytest.raises(litellm.Timeout, match="max streaming duration"): + wrapper._check_max_streaming_duration() + + def test_should_raise_on_sync_next_when_exceeded(self): + """__next__ should check the limit before iterating.""" + wrapper = _make_custom_stream_wrapper() + wrapper._stream_created_time = time.time() - 20 + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): + with pytest.raises(litellm.Timeout): + wrapper.__next__() + + @pytest.mark.asyncio + async def test_should_raise_on_async_anext_when_exceeded(self): + """__anext__ should check the limit before iterating.""" + wrapper = _make_custom_stream_wrapper() + wrapper._stream_created_time = time.time() - 20 + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): + with pytest.raises(litellm.Timeout): + await wrapper.__anext__() + + +# --------------------------------------------------------------------------- +# BaseResponsesAPIStreamingIterator (responses) +# --------------------------------------------------------------------------- + +class TestResponsesStreamingIteratorMaxDuration: + def _make_base_iterator(self): + """Build a minimal BaseResponsesAPIStreamingIterator for testing.""" + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + mock_response = MagicMock() + mock_response.headers = {} + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_logging_obj.start_time = time.time() + + mock_provider_config = MagicMock() + return BaseResponsesAPIStreamingIterator( + response=mock_response, + model="test-model", + responses_api_provider_config=mock_provider_config, + logging_obj=mock_logging_obj, + custom_llm_provider="openai", + ) + + def test_should_not_raise_when_duration_is_none(self): + it = self._make_base_iterator() + with patch( + "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", None + ): + it._check_max_streaming_duration() + + def test_should_not_raise_when_under_limit(self): + it = self._make_base_iterator() + with patch( + "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", 60.0 + ): + it._check_max_streaming_duration() + + def test_should_raise_timeout_when_exceeded(self): + it = self._make_base_iterator() + it._stream_created_time = time.time() - 20 + with patch( + "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0 + ): + with pytest.raises(litellm.Timeout, match="max streaming duration"): + it._check_max_streaming_duration() diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index bcda3c7bfac..7d38a5cc80a 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -379,7 +379,8 @@ async def test_realtime_guardrail_blocks_prompt_injection(): """ Test that when a transcription event containing prompt injection arrives from the backend, a registered guardrail blocks it — sending a warning to the client - and NOT sending response.create to the backend. + and voicing the guardrail violation message via response.cancel + + conversation.item.create + response.create. """ import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -430,19 +431,36 @@ async def test_realtime_guardrail_blocks_prompt_injection(): streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) await streaming.backend_to_client_send_messages() - # ASSERT 1: no response.create was sent to backend (injection blocked). + # ASSERT 1: the guardrail blocked the normal auto-response and instead + # injected a conversation.item.create + response.create to voice the + # violation message. There should be exactly ONE response.create (the + # guardrail-triggered one), preceded by a response.cancel and a + # conversation.item.create carrying the violation text. sent_to_backend = [ json.loads(c.args[0]) for c in backend_ws.send.call_args_list if c.args ] - response_creates = [ - e for e in sent_to_backend - if e.get("type") == "response.create" + response_cancels = [ + e for e in sent_to_backend if e.get("type") == "response.cancel" ] - assert len(response_creates) == 0, ( - f"Guardrail should prevent response.create for injected content, " - f"but got: {response_creates}" + assert len(response_cancels) == 1, ( + f"Guardrail should send response.cancel, got: {response_cancels}" + ) + guardrail_items = [ + e for e in sent_to_backend + if e.get("type") == "conversation.item.create" + ] + assert len(guardrail_items) == 1, ( + f"Guardrail should inject a conversation.item.create with violation message, " + f"got: {guardrail_items}" + ) + response_creates = [ + e for e in sent_to_backend if e.get("type") == "response.create" + ] + assert len(response_creates) == 1, ( + f"Guardrail should send exactly one response.create to voice the violation, " + f"got: {response_creates}" ) # ASSERT 2: error event was sent directly to the client WebSocket @@ -595,14 +613,26 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(): assert len(error_events) == 1, f"Expected one error event, got: {sent_texts}" assert error_events[0]["error"]["type"] == "guardrail_violation" - # ASSERT: blocked item was NOT forwarded to the backend + # ASSERT: the original blocked item was NOT forwarded to the backend. + # The guardrail handler injects its own conversation.item.create with + # the violation message — only that one should be present, not the + # original user message. sent_to_backend = [c.args[0] for c in backend_ws.send.call_args_list if c.args] forwarded_items = [ json.loads(m) for m in sent_to_backend if isinstance(m, str) and json.loads(m).get("type") == "conversation.item.create" ] - assert len(forwarded_items) == 0, ( - f"Blocked item should not be forwarded to backend, got: {forwarded_items}" + # Filter out guardrail-injected items (contain "Say exactly the following message") + original_items = [ + item for item in forwarded_items + if not any( + "Say exactly the following message" in c.get("text", "") + for c in item.get("item", {}).get("content", []) + if isinstance(c, dict) + ) + ] + assert len(original_items) == 0, ( + f"Blocked item should not be forwarded to backend, got: {original_items}" ) litellm.callbacks = [] # cleanup @@ -637,9 +667,10 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode(): assert streaming._has_realtime_guardrails() is True, ( "pre_call guardrail should be recognized as a realtime guardrail" ) - # pre_call guardrail should NOT trigger the audio/VAD session.update injection - assert streaming._has_audio_transcription_guardrails() is False, ( - "pre_call guardrail should not trigger audio transcription guardrail path" + # pre_call guardrail SHOULD trigger the audio/VAD session.update injection so + # that the LLM does not auto-respond before the guardrail can check the transcript. + assert streaming._has_audio_transcription_guardrails() is True, ( + "pre_call guardrail should trigger audio transcription guardrail path" ) litellm.callbacks = [] # cleanup @@ -711,10 +742,11 @@ async def test_realtime_session_created_injects_session_update_for_audio_guardra @pytest.mark.asyncio -async def test_realtime_session_created_no_injection_for_pre_call_only(): +async def test_realtime_session_created_injects_session_update_for_pre_call_guardrail(): """ - Test that when only a pre_call guardrail is configured (no audio transcription), - session.created does NOT trigger the session.update injection. + Test that when a pre_call guardrail is configured, session.created triggers the + session.update injection (create_response: false) so the LLM does not auto-respond + before the guardrail can check the voice transcript. """ import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -751,14 +783,15 @@ async def test_realtime_session_created_no_injection_for_pre_call_only(): streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) await streaming.backend_to_client_send_messages() - # No session.update should be injected + # session.update SHOULD be injected so the LLM waits for guardrail approval sent_to_backend = [ json.loads(c.args[0]) for c in backend_ws.send.call_args_list if c.args ] session_updates = [e for e in sent_to_backend if e.get("type") == "session.update"] - assert len(session_updates) == 0, ( - f"pre_call guardrail should NOT inject session.update, got: {sent_to_backend}" + assert len(session_updates) == 1, ( + f"pre_call guardrail should inject session.update to gate audio responses, got: {sent_to_backend}" ) + assert session_updates[0]["session"]["turn_detection"]["create_response"] is False litellm.callbacks = [] # cleanup diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 77c74a7847e..636e84fe796 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -16,13 +16,14 @@ from litellm.types.utils import Delta, ModelResponse, StreamingChoices def test_anthropic_experimental_pass_through_messages_handler(): """ - Test that api key is passed to litellm.completion + Test that api key is passed to litellm.responses for OpenAI models. + OpenAI and Azure models are routed directly to the Responses API. """ from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( anthropic_messages_handler, ) - with patch("litellm.completion", return_value="test-response") as mock_completion: + with patch("litellm.responses", return_value="test-response") as mock_responses: try: anthropic_messages_handler( max_tokens=100, @@ -32,19 +33,20 @@ def test_anthropic_experimental_pass_through_messages_handler(): ) except Exception as e: print(f"Error: {e}") - mock_completion.assert_called_once() - assert mock_completion.call_args.kwargs["api_key"] == "test-api-key" + mock_responses.assert_called_once() + assert mock_responses.call_args.kwargs["api_key"] == "test-api-key" def test_anthropic_experimental_pass_through_messages_handler_dynamic_api_key_and_api_base_and_custom_values(): """ - Test that api key is passed to litellm.completion + Test that api key, api base, and extra kwargs are forwarded to litellm.completion for Azure models. + Azure models are routed through chat/completions (not the Responses API). """ from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( anthropic_messages_handler, ) - with patch("litellm.completion", return_value="test-response") as mock_completion: + with patch("litellm.completion", return_value=MagicMock()) as mock_completion: try: anthropic_messages_handler( max_tokens=100, @@ -143,19 +145,19 @@ async def test_bedrock_converse_budget_tokens_preserved(): assert thinking_param.get("budget_tokens") == 1024, f"thinking.budget_tokens should be 1024, but got {thinking_param.get('budget_tokens')}" -def test_openai_model_with_thinking_converts_to_reasoning_effort(): +def test_openai_model_with_thinking_converts_to_reasoning(): """ - Test that when using a non-Anthropic model (like OpenAI gpt-5.2) with thinking parameter, - the thinking is converted to reasoning_effort and NOT passed as thinking. - - This ensures we don't regress on issue #16052 where non-Anthropic models would fail - with UnsupportedParamsError when thinking was passed directly. + Test that when using an OpenAI model with thinking parameter, the thinking is + converted to a Responses API `reasoning` param (NOT passed as thinking). + + OpenAI models are routed directly to the Responses API, so we verify that + litellm.responses() is called with `reasoning` properly set. """ from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( anthropic_messages_handler, ) - with patch("litellm.completion", return_value="test-response") as mock_completion: + with patch("litellm.responses", return_value="test-response") as mock_responses: try: anthropic_messages_handler( max_tokens=1024, @@ -170,20 +172,22 @@ def test_openai_model_with_thinking_converts_to_reasoning_effort(): except Exception as e: print(f"Error: {e}") - mock_completion.assert_called_once() - - call_kwargs = mock_completion.call_args.kwargs - - # Verify reasoning_effort is set (converted from thinking) - assert "reasoning_effort" in call_kwargs, "reasoning_effort should be passed to completion" + mock_responses.assert_called_once() - # reasoning_effort is transformed into a dict with effort and summary fields - expected_reasoning_effort = {"effort": "minimal", "summary": "detailed"} - assert call_kwargs["reasoning_effort"] == expected_reasoning_effort, \ - f"reasoning_effort should be {expected_reasoning_effort} for budget_tokens=1024, got {call_kwargs.get('reasoning_effort')}" + call_kwargs = mock_responses.call_args.kwargs - # Verify thinking is NOT passed (non-Claude model) - assert "thinking" not in call_kwargs, "thinking should NOT be passed for non-Claude models" + # Verify reasoning is set (converted from thinking) + assert "reasoning" in call_kwargs, "reasoning should be passed to litellm.responses" + + # budget_tokens=1024 -> effort="minimal" (< 2000 threshold) + expected_reasoning = {"effort": "minimal", "summary": "detailed"} + assert call_kwargs["reasoning"] == expected_reasoning, ( + f"reasoning should be {expected_reasoning} for budget_tokens=1024, " + f"got {call_kwargs.get('reasoning')}" + ) + + # Verify thinking is NOT passed directly to the Responses API + assert "thinking" not in call_kwargs, "thinking should NOT be passed directly to litellm.responses" class TestThinkingParameterTransformation: diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py new file mode 100644 index 00000000000..252ba230ff7 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -0,0 +1,987 @@ +""" +Tests for LiteLLMAnthropicToResponsesAPIAdapter +(litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py) +""" + +import json +import os +import sys +from typing import Any, Dict, List +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.abspath("../../../../../../..")) + +from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.types.llms.anthropic import AnthropicMessagesRequest + + +def _make_request(**overrides) -> AnthropicMessagesRequest: + base: dict = { + "model": "openai.gpt-5.1-codex", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 1024, + } + base.update(overrides) + return AnthropicMessagesRequest(**base) + + +_ADAPTER = LiteLLMAnthropicToResponsesAPIAdapter() + + +# --------------------------------------------------------------------------- +# context_management conversion +# --------------------------------------------------------------------------- + + +class TestContextManagementConversion: + """Anthropic dict -> OpenAI array conversion for context_management.""" + + def test_compact_edit_converted_to_array(self): + """compact_20260112 with trigger maps to OpenAI compaction entry.""" + cm = { + "edits": [ + { + "type": "compact_20260112", + "trigger": {"type": "input_tokens", "value": 150000}, + } + ] + } + result = _ADAPTER.translate_context_management_to_responses_api(cm) + assert result == [{"type": "compaction", "compact_threshold": 150000}] + + def test_compact_edit_without_trigger(self): + """compact_20260112 without a trigger still maps to a compaction entry.""" + cm = {"edits": [{"type": "compact_20260112"}]} + result = _ADAPTER.translate_context_management_to_responses_api(cm) + assert result == [{"type": "compaction"}] + + def test_unknown_edit_type_is_dropped(self): + """Anthropic-only edit types (e.g. clear_thinking) are silently dropped.""" + cm = {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]} + result = _ADAPTER.translate_context_management_to_responses_api(cm) + assert result is None + + def test_mixed_edits_only_known_types_kept(self): + """Only compact_20260112 is converted; unknown types are dropped.""" + cm = { + "edits": [ + {"type": "clear_thinking_20251015", "keep": "all"}, + { + "type": "compact_20260112", + "trigger": {"type": "input_tokens", "value": 200000}, + }, + ] + } + result = _ADAPTER.translate_context_management_to_responses_api(cm) + assert result == [{"type": "compaction", "compact_threshold": 200000}] + + def test_non_dict_returns_none(self): + result = _ADAPTER.translate_context_management_to_responses_api([]) # type: ignore + assert result is None + + def test_translate_request_includes_context_management(self): + """translate_request converts context_management and sets it on kwargs.""" + req = _make_request( + context_management={ + "edits": [ + { + "type": "compact_20260112", + "trigger": {"type": "input_tokens", "value": 100000}, + } + ] + } + ) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["context_management"] == [ + {"type": "compaction", "compact_threshold": 100000} + ] + + def test_translate_request_drops_anthropic_only_context_management(self): + """context_management with only unknown edit types is omitted from kwargs.""" + req = _make_request( + context_management={ + "edits": [{"type": "clear_thinking_20251015", "keep": "all"}] + } + ) + kwargs = _ADAPTER.translate_request(req) + assert "context_management" not in kwargs + + +# --------------------------------------------------------------------------- +# structured output via output_config +# --------------------------------------------------------------------------- + + +class TestOutputConfigStructuredOutput: + """output_config.format.json_schema -> OpenAI text.format conversion.""" + + _SCHEMA = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "email": {"type": "string"}, + }, + "required": ["name", "email"], + "additionalProperties": False, + } + + def test_output_config_format_json_schema_converted(self): + """output_config.format.json_schema is converted to OpenAI text.format.""" + req = _make_request( + output_config={"format": {"type": "json_schema", "schema": self._SCHEMA}} + ) + kwargs = _ADAPTER.translate_request(req) + assert "text" in kwargs + fmt = kwargs["text"]["format"] + assert fmt["type"] == "json_schema" + assert fmt["schema"] == self._SCHEMA + assert fmt["strict"] is True + assert fmt["name"] == "structured_output" + + def test_output_config_without_format_does_not_set_text(self): + """output_config with only non-format keys doesn't produce text.format.""" + req = _make_request(output_config={"effort": "high"}) + kwargs = _ADAPTER.translate_request(req) + assert "text" not in kwargs + + def test_output_format_still_works(self): + """The original output_format field still takes precedence when present.""" + req = _make_request( + output_format={"type": "json_schema", "schema": self._SCHEMA} + ) + kwargs = _ADAPTER.translate_request(req) + assert "text" in kwargs + assert kwargs["text"]["format"]["type"] == "json_schema" + + def test_output_format_takes_precedence_over_output_config(self): + """output_format takes precedence over output_config.format.""" + other_schema = {"type": "object", "properties": {"id": {"type": "integer"}}} + req = _make_request( + output_format={"type": "json_schema", "schema": self._SCHEMA}, + output_config={"format": {"type": "json_schema", "schema": other_schema}}, + ) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["text"]["format"]["schema"] == self._SCHEMA + + +# --------------------------------------------------------------------------- +# translate_messages_to_responses_input +# --------------------------------------------------------------------------- + +# Helper: cast plain dicts to the expected type so call sites stay clean. +def _translate_messages(messages: List[Any]) -> List[Dict[str, Any]]: + return _ADAPTER.translate_messages_to_responses_input(messages) # type: ignore[arg-type] + + +class TestTranslateMessagesToResponsesInput: + """Anthropic messages list -> OpenAI Responses API input items.""" + + def test_user_string_content(self): + """Plain string user message becomes a message with input_text.""" + messages = [{"role": "user", "content": "Hello world"}] + result = _translate_messages(messages) + assert result == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Hello world"}], + } + ] + + def test_user_list_text_block(self): + """User message with text content block maps to input_text.""" + messages = [ + { + "role": "user", + "content": [{"type": "text", "text": "What is 2+2?"}], + } + ] + result = _translate_messages(messages) + assert result == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "What is 2+2?"}], + } + ] + + def test_user_multiple_text_blocks(self): + """Multiple text blocks in a user message are all converted.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "First part."}, + {"type": "text", "text": "Second part."}, + ], + } + ] + result = _translate_messages(messages) + assert len(result) == 1 + assert result[0]["content"] == [ + {"type": "input_text", "text": "First part."}, + {"type": "input_text", "text": "Second part."}, + ] + + def test_user_base64_image(self): + """User message with base64 image source becomes input_image with data URL.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "abc123", + }, + } + ], + } + ] + result = _translate_messages(messages) + assert len(result) == 1 + assert result[0]["content"] == [ + {"type": "input_image", "image_url": "data:image/png;base64,abc123"} + ] + + def test_user_url_image(self): + """User message with URL image source becomes input_image with the URL.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": {"type": "url", "url": "https://example.com/img.jpg"}, + } + ], + } + ] + result = _translate_messages(messages) + assert result[0]["content"] == [ + {"type": "input_image", "image_url": "https://example.com/img.jpg"} + ] + + def test_user_base64_image_empty_data_skipped(self): + """Base64 image with empty data is skipped (no URL can be formed).""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": {"type": "base64", "media_type": "image/jpeg", "data": ""}, + } + ], + } + ] + result = _translate_messages(messages) + # No user_parts -> no message item appended + assert result == [] + + def test_user_tool_result_string_content(self): + """tool_result with string content becomes function_call_output.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_abc", + "content": "42 degrees", + } + ], + } + ] + result = _translate_messages(messages) + assert result == [ + { + "type": "function_call_output", + "call_id": "call_abc", + "output": "42 degrees", + } + ] + + def test_user_tool_result_list_content(self): + """tool_result with list of text blocks is joined into a single string.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_xyz", + "content": [ + {"type": "text", "text": "Line 1"}, + {"type": "text", "text": "Line 2"}, + ], + } + ], + } + ] + result = _translate_messages(messages) + assert result[0]["output"] == "Line 1\nLine 2" + + def test_user_tool_result_null_content(self): + """tool_result with null content becomes empty string output.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "call_null", "content": None} + ], + } + ] + result = _translate_messages(messages) + assert result[0]["output"] == "" + + def test_assistant_string_content(self): + """Plain string assistant message becomes a message with output_text.""" + messages = [{"role": "assistant", "content": "I can help with that."}] + result = _translate_messages(messages) + assert result == [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I can help with that."}], + } + ] + + def test_assistant_text_block(self): + """Assistant message with text block maps to output_text.""" + messages = [ + { + "role": "assistant", + "content": [{"type": "text", "text": "Here is the answer."}], + } + ] + result = _translate_messages(messages) + assert result[0]["content"] == [ + {"type": "output_text", "text": "Here is the answer."} + ] + + def test_assistant_tool_use_becomes_function_call(self): + """Assistant tool_use block becomes a top-level function_call item.""" + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "get_weather", + "input": {"location": "Boston"}, + } + ], + } + ] + result = _translate_messages(messages) + assert result == [ + { + "type": "function_call", + "call_id": "toolu_01", + "name": "get_weather", + "arguments": json.dumps({"location": "Boston"}), + } + ] + + def test_assistant_thinking_block_becomes_output_text(self): + """Assistant thinking block text is included as output_text.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Let me reason step by step."} + ], + } + ] + result = _translate_messages(messages) + assert result[0]["content"] == [ + {"type": "output_text", "text": "Let me reason step by step."} + ] + + def test_assistant_empty_thinking_block_skipped(self): + """Assistant thinking block with empty thinking text is skipped.""" + messages = [ + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": ""}], + } + ] + result = _translate_messages(messages) + assert result == [] + + def test_mixed_messages_ordering(self): + """Full multi-turn conversation is converted in order.""" + messages = [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_02", + "name": "get_weather", + "input": {"city": "NYC"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_02", + "content": "Sunny, 72F", + } + ], + }, + {"role": "assistant", "content": "It's sunny and 72°F in NYC."}, + ] + result = _translate_messages(messages) + types = [item["type"] for item in result] + assert types == ["message", "function_call", "function_call_output", "message"] + + def test_user_text_and_image_mixed(self): + """User message with both text and image produces both parts.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image:"}, + { + "type": "image", + "source": {"type": "url", "url": "https://example.com/cat.jpg"}, + }, + ], + } + ] + result = _translate_messages(messages) + assert len(result) == 1 + assert result[0]["content"][0] == {"type": "input_text", "text": "Describe this image:"} + assert result[0]["content"][1] == { + "type": "input_image", + "image_url": "https://example.com/cat.jpg", + } + + def test_unknown_image_source_type_skipped(self): + """Image block with unknown source type is silently skipped.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": {"type": "file_path", "path": "/tmp/img.png"}, + } + ], + } + ] + result = _translate_messages(messages) + assert result == [] + + +# --------------------------------------------------------------------------- +# translate_tools_to_responses_api +# --------------------------------------------------------------------------- + + +class TestTranslateToolsToResponsesAPI: + """Anthropic tool definitions -> Responses API function tools.""" + + def test_regular_tool_with_description_and_schema(self): + """Standard tool with description and input_schema is converted to function.""" + tools = [ + { + "name": "get_weather", + "description": "Get current weather for a city.", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ] + result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + assert result == [ + { + "type": "function", + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ] + + def test_tool_without_description(self): + """Tool without a description omits the description key.""" + tools = [{"name": "ping", "input_schema": {"type": "object", "properties": {}}}] + result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + assert result[0]["type"] == "function" + assert result[0]["name"] == "ping" + assert "description" not in result[0] + + def test_tool_without_input_schema(self): + """Tool without input_schema omits the parameters key.""" + tools = [{"name": "no_schema_tool", "description": "Does something."}] + result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + assert result[0]["type"] == "function" + assert "parameters" not in result[0] + + def test_web_search_tool_by_name(self): + """Tool named 'web_search' maps to web_search_preview.""" + tools = [{"name": "web_search", "type": "custom"}] + result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + assert result == [{"type": "web_search_preview"}] + + def test_web_search_tool_by_type_prefix(self): + """Tool with type starting with 'web_search' maps to web_search_preview.""" + tools = [{"name": "search", "type": "web_search_20250305"}] + result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + assert result == [{"type": "web_search_preview"}] + + def test_multiple_tools_order_preserved(self): + """Multiple tools are converted in order.""" + tools = [ + {"name": "tool_a", "description": "A"}, + {"name": "web_search", "type": "custom"}, + {"name": "tool_b", "description": "B"}, + ] + result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + assert len(result) == 3 + assert result[0]["name"] == "tool_a" + assert result[1] == {"type": "web_search_preview"} + assert result[2]["name"] == "tool_b" + + def test_empty_tools_list(self): + """Empty tools list returns empty list.""" + assert _ADAPTER.translate_tools_to_responses_api([]) == [] + + +# --------------------------------------------------------------------------- +# translate_tool_choice_to_responses_api +# --------------------------------------------------------------------------- + + +class TestTranslateToolChoiceToResponsesAPI: + """Anthropic tool_choice -> Responses API tool_choice.""" + + def test_auto_maps_to_auto(self): + assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "auto"}) == { + "type": "auto" + } + + def test_any_maps_to_required(self): + assert _ADAPTER.translate_tool_choice_to_responses_api({"type": "any"}) == { + "type": "required" + } + + def test_specific_tool_maps_to_function(self): + result = _ADAPTER.translate_tool_choice_to_responses_api( + {"type": "tool", "name": "get_weather"} + ) + assert result == {"type": "function", "name": "get_weather"} + + def test_unknown_type_defaults_to_auto(self): + result = _ADAPTER.translate_tool_choice_to_responses_api({"type": "none"}) + assert result == {"type": "auto"} + + +# --------------------------------------------------------------------------- +# translate_thinking_to_reasoning +# --------------------------------------------------------------------------- + + +class TestTranslateThinkingToReasoning: + """Anthropic thinking param -> Responses API reasoning param.""" + + def test_budget_high_effort(self): + result = _ADAPTER.translate_thinking_to_reasoning( + {"type": "enabled", "budget_tokens": 10000} + ) + assert result == {"effort": "high", "summary": "detailed"} + + def test_budget_above_threshold_high_effort(self): + result = _ADAPTER.translate_thinking_to_reasoning( + {"type": "enabled", "budget_tokens": 50000} + ) + assert result is not None + assert result["effort"] == "high" + + def test_budget_medium_effort(self): + result = _ADAPTER.translate_thinking_to_reasoning( + {"type": "enabled", "budget_tokens": 7500} + ) + assert result == {"effort": "medium", "summary": "detailed"} + + def test_budget_low_effort(self): + result = _ADAPTER.translate_thinking_to_reasoning( + {"type": "enabled", "budget_tokens": 3000} + ) + assert result == {"effort": "low", "summary": "detailed"} + + def test_budget_minimal_effort(self): + result = _ADAPTER.translate_thinking_to_reasoning( + {"type": "enabled", "budget_tokens": 500} + ) + assert result == {"effort": "minimal", "summary": "detailed"} + + def test_budget_at_exact_thresholds(self): + result_medium = _ADAPTER.translate_thinking_to_reasoning( + {"type": "enabled", "budget_tokens": 5000} + ) + assert result_medium is not None + assert result_medium["effort"] == "medium" + result_low = _ADAPTER.translate_thinking_to_reasoning( + {"type": "enabled", "budget_tokens": 2000} + ) + assert result_low is not None + assert result_low["effort"] == "low" + + def test_disabled_type_returns_none(self): + result = _ADAPTER.translate_thinking_to_reasoning({"type": "disabled"}) + assert result is None + + def test_non_dict_returns_none(self): + result = _ADAPTER.translate_thinking_to_reasoning("enabled") # type: ignore + assert result is None + + def test_missing_budget_defaults_to_minimal(self): + """Missing budget_tokens defaults to 0, which is < 2000 -> minimal.""" + result = _ADAPTER.translate_thinking_to_reasoning({"type": "enabled"}) + assert result == {"effort": "minimal", "summary": "detailed"} + + +# --------------------------------------------------------------------------- +# translate_request – broader coverage +# --------------------------------------------------------------------------- + + +class TestTranslateRequestBroaderCoverage: + """Full translate_request call: field-by-field mapping verification.""" + + def test_model_and_input_always_present(self): + req = _make_request() + kwargs = _ADAPTER.translate_request(req) + assert "model" in kwargs + assert "input" in kwargs + + def test_system_string_becomes_instructions(self): + req = _make_request(system="You are a helpful assistant.") + kwargs = _ADAPTER.translate_request(req) + assert kwargs["instructions"] == "You are a helpful assistant." + + def test_system_list_of_text_blocks_joined(self): + req = _make_request( + system=[ + {"type": "text", "text": "Be concise."}, + {"type": "text", "text": "Be helpful."}, + ] + ) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["instructions"] == "Be concise.\nBe helpful." + + def test_system_list_skips_non_text_blocks(self): + req = _make_request( + system=[ + {"type": "image", "source": {}}, + {"type": "text", "text": "Only text matters."}, + ] + ) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["instructions"] == "Only text matters." + + def test_max_tokens_mapped_to_max_output_tokens(self): + req = _make_request(max_tokens=512) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["max_output_tokens"] == 512 + + def test_temperature_passed_through(self): + req = _make_request(temperature=0.7) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["temperature"] == 0.7 + + def test_top_p_passed_through(self): + req = _make_request(top_p=0.9) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["top_p"] == 0.9 + + def test_tools_translated(self): + req = _make_request( + tools=[{"name": "calculator", "description": "Does math.", "input_schema": {}}] + ) + kwargs = _ADAPTER.translate_request(req) + assert len(kwargs["tools"]) == 1 + assert kwargs["tools"][0]["name"] == "calculator" + + def test_tool_choice_translated(self): + req = _make_request( + tools=[{"name": "do_thing"}], + tool_choice={"type": "tool", "name": "do_thing"}, + ) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["tool_choice"] == {"type": "function", "name": "do_thing"} + + def test_thinking_translated_to_reasoning(self): + req = _make_request(thinking={"type": "enabled", "budget_tokens": 12000}) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["reasoning"] == {"effort": "high", "summary": "detailed"} + + def test_disabled_thinking_not_included_in_kwargs(self): + req = _make_request(thinking={"type": "disabled"}) + kwargs = _ADAPTER.translate_request(req) + assert "reasoning" not in kwargs + + def test_metadata_user_id_mapped_to_user(self): + req = _make_request(metadata={"user_id": "user-42"}) + kwargs = _ADAPTER.translate_request(req) + assert kwargs["user"] == "user-42" + + def test_metadata_user_id_truncated_to_64_chars(self): + long_id = "x" * 100 + req = _make_request(metadata={"user_id": long_id}) + kwargs = _ADAPTER.translate_request(req) + assert len(kwargs["user"]) == 64 + + def test_no_optional_fields_does_not_add_spurious_keys(self): + req = _make_request() + kwargs = _ADAPTER.translate_request(req) + for key in ("instructions", "temperature", "top_p", "tools", "tool_choice", + "reasoning", "text", "context_management", "user"): + assert key not in kwargs, f"unexpected key: {key}" + + +# --------------------------------------------------------------------------- +# translate_response +# --------------------------------------------------------------------------- + + +def _make_mock_response( + output: list, + status: str = "completed", + response_id: str = "resp_001", + model: str = "gpt-4o", + input_tokens: int = 100, + output_tokens: int = 50, +) -> MagicMock: + """Build a minimal mock ResponsesAPIResponse.""" + usage = MagicMock() + usage.input_tokens = input_tokens + usage.output_tokens = output_tokens + + resp = MagicMock() + resp.id = response_id + resp.model = model + resp.status = status + resp.output = output + resp.usage = usage + return resp + + +def _make_output_message(texts: List[str]) -> MagicMock: + """Build a mock ResponseOutputMessage with output_text parts.""" + from openai.types.responses import ResponseOutputMessage # type: ignore[import] + + parts = [] + for t in texts: + part = MagicMock() + part.type = "output_text" + part.text = t + parts.append(part) + + msg = MagicMock(spec=ResponseOutputMessage) + msg.content = parts + return msg + + +def _make_function_call_item( + call_id: str, name: str, arguments: str +) -> MagicMock: + """Build a mock ResponseFunctionToolCall.""" + from openai.types.responses import ResponseFunctionToolCall # type: ignore[import] + + item = MagicMock(spec=ResponseFunctionToolCall) + item.call_id = call_id + item.id = call_id + item.name = name + item.arguments = arguments + return item + + +def _make_reasoning_item(summaries: List[str]) -> MagicMock: + """Build a mock ResponseReasoningItem.""" + from openai.types.responses import ResponseReasoningItem # type: ignore[import] + + summary_mocks = [] + for text in summaries: + s = MagicMock() + s.text = text + summary_mocks.append(s) + + item = MagicMock(spec=ResponseReasoningItem) + item.summary = summary_mocks + return item + + +class TestTranslateResponse: + """Responses API -> AnthropicMessagesResponse conversion.""" + + def test_output_text_message_becomes_text_block(self): + """ResponseOutputMessage with output_text parts -> Anthropic text content.""" + response = _make_mock_response(output=[_make_output_message(["Hello!"])]) + result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Hello!" + + def test_multiple_text_parts(self): + """Multiple output_text parts become multiple text content blocks.""" + response = _make_mock_response( + output=[_make_output_message(["Part 1", "Part 2"])] + ) + result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 2 + assert result["content"][0]["text"] == "Part 1" + assert result["content"][1]["text"] == "Part 2" + + def test_function_call_becomes_tool_use(self): + """ResponseFunctionToolCall -> Anthropic tool_use content block.""" + fc = _make_function_call_item("call_99", "get_weather", '{"city": "NYC"}') + response = _make_mock_response(output=[fc]) + result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 + block = result["content"][0] + assert block["type"] == "tool_use" + assert block["id"] == "call_99" + assert block["name"] == "get_weather" + assert block["input"] == {"city": "NYC"} + + def test_function_call_sets_stop_reason_tool_use(self): + """Presence of a function_call sets stop_reason to 'tool_use'.""" + fc = _make_function_call_item("call_1", "tool_a", "{}") + response = _make_mock_response(output=[fc]) + result: Any = _ADAPTER.translate_response(response) + assert result["stop_reason"] == "tool_use" + + def test_text_only_stop_reason_end_turn(self): + """Text-only response has stop_reason 'end_turn'.""" + response = _make_mock_response(output=[_make_output_message(["Hi"])]) + result: Any = _ADAPTER.translate_response(response) + assert result["stop_reason"] == "end_turn" + + def test_incomplete_status_sets_max_tokens(self): + """status='incomplete' overrides stop_reason to 'max_tokens'.""" + response = _make_mock_response( + output=[_make_output_message(["Truncated..."])], + status="incomplete", + ) + result: Any = _ADAPTER.translate_response(response) + assert result["stop_reason"] == "max_tokens" + + def test_reasoning_item_becomes_thinking_block(self): + """ResponseReasoningItem summaries -> Anthropic thinking content blocks.""" + reasoning = _make_reasoning_item(["Step 1: analyze. Step 2: conclude."]) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 + assert result["content"][0]["type"] == "thinking" + assert "Step 1" in result["content"][0]["thinking"] + + def test_empty_reasoning_summary_skipped(self): + """Reasoning item with empty text summary is not added to content.""" + reasoning = _make_reasoning_item([""]) + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [] + + def test_usage_mapped_correctly(self): + """Input/output tokens from ResponseAPIUsage are mapped to AnthropicUsage.""" + response = _make_mock_response( + output=[_make_output_message(["OK"])], + input_tokens=200, + output_tokens=75, + ) + result: Any = _ADAPTER.translate_response(response) + assert result["usage"]["input_tokens"] == 200 + assert result["usage"]["output_tokens"] == 75 + + def test_model_and_id_preserved(self): + """Model and response ID from the Responses API are forwarded.""" + response = _make_mock_response( + output=[_make_output_message(["Hi"])], + response_id="resp_xyz", + model="gpt-4-turbo", + ) + result: Any = _ADAPTER.translate_response(response) + assert result["id"] == "resp_xyz" + assert result["model"] == "gpt-4-turbo" + + def test_role_is_always_assistant(self): + response = _make_mock_response(output=[_make_output_message(["Hi"])]) + result: Any = _ADAPTER.translate_response(response) + assert result["role"] == "assistant" + + def test_type_is_always_message(self): + response = _make_mock_response(output=[_make_output_message(["Hi"])]) + result: Any = _ADAPTER.translate_response(response) + assert result["type"] == "message" + + def test_empty_output_list(self): + """Empty output list produces empty content with 'end_turn' stop reason.""" + response = _make_mock_response(output=[]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [] + assert result["stop_reason"] == "end_turn" + + def test_function_call_with_invalid_json_arguments(self): + """Invalid JSON in function_call arguments falls back to empty dict.""" + fc = _make_function_call_item("call_bad", "broken_tool", "not-valid-json") + response = _make_mock_response(output=[fc]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"][0]["input"] == {} + + def test_dict_output_message_item(self): + """Dict-shaped output message (type=message) is also handled.""" + output_item = { + "type": "message", + "content": [{"type": "output_text", "text": "Dict-based response"}], + } + response = _make_mock_response(output=[output_item]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Dict-based response" + + def test_dict_function_call_item(self): + """Dict-shaped function_call item is converted to tool_use block.""" + output_item = { + "type": "function_call", + "call_id": "call_dict_1", + "name": "search", + "arguments": '{"query": "cats"}', + } + response = _make_mock_response(output=[output_item]) + result: Any = _ADAPTER.translate_response(response) + assert result["content"][0]["type"] == "tool_use" + assert result["content"][0]["name"] == "search" + assert result["content"][0]["input"] == {"query": "cats"} + assert result["stop_reason"] == "tool_use" + + def test_mixed_reasoning_text_and_tool_use(self): + """Reasoning + text + tool_use in one response all convert correctly.""" + reasoning = _make_reasoning_item(["Thinking..."]) + text_msg = _make_output_message(["Here is my answer."]) + fc = _make_function_call_item("call_mix", "lookup", '{"id": 1}') + response = _make_mock_response(output=[reasoning, text_msg, fc]) + result: Any = _ADAPTER.translate_response(response) + types = [b["type"] for b in result["content"]] + assert "thinking" in types + assert "text" in types + assert "tool_use" in types + assert result["stop_reason"] == "tool_use" diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 26395597166..345f3ae7c5d 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2616,11 +2616,11 @@ def test_empty_assistant_message_handling(): empty or whitespace-only content with a placeholder to prevent AWS Bedrock Converse API 400 Bad Request errors. """ - # Import the litellm module that factory.py uses to ensure we patch the correct reference - import litellm.litellm_core_utils.prompt_templates.factory as factory_module from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, ) + # Import the litellm module that factory.py uses to ensure we patch the correct reference + import litellm.litellm_core_utils.prompt_templates.factory as factory_module # Test case 1: Empty string content - test with modify_params=True to prevent merging messages = [ @@ -3135,12 +3135,7 @@ def test_native_structured_output_no_fake_stream(): def test_transform_request_with_output_config(): """Test that outputConfig flows through _transform_request_helper into the final request.""" - from litellm.types.llms.bedrock import ( - JsonSchemaDefinition, - OutputConfigBlock, - OutputFormat, - OutputFormatStructure, - ) + from litellm.types.llms.bedrock import OutputConfigBlock, OutputFormat, OutputFormatStructure, JsonSchemaDefinition config = AmazonConverseConfig() @@ -3382,59 +3377,78 @@ def test_output_config_applies_additional_properties(): -def test_parallel_tool_calls_in_request_transformation(): - """Test that parallel_tool_calls is correctly placed in additionalModelRequestFields after full transformation""" - config = AmazonConverseConfig() - - messages = [ - {"role": "user", "content": "What's the weather in SF and NYC?"} - ] - - non_default_params = { - "parallel_tool_calls": False, - "tools": [ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the weather", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The location to get weather for" - } - }, - "required": ["location"] +_TOOL_PARAM = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The location to get weather for", } - } - } - ], - "max_tokens": 100, + }, + "required": ["location"], + }, + }, } - +] + + +def test_parallel_tool_calls_newer_model_adds_disable_flag(): + """Newer Claude models (4.5+) should get disable_parallel_tool_use in additionalModelRequestFields.""" + config = AmazonConverseConfig() + model = "anthropic.claude-sonnet-4-5-20250929-v1:0" + messages = [{"role": "user", "content": "What's the weather in SF and NYC?"}] + optional_params = config.map_openai_params( - non_default_params=non_default_params, + non_default_params={"parallel_tool_calls": False, "tools": _TOOL_PARAM}, optional_params={}, - model="anthropic.claude-sonnet-4-5-v2:0", + model=model, drop_params=False, ) - - # Transform the request + request_data = config.transform_request( - model="anthropic.claude-sonnet-4-5-v2:0", + model=model, messages=messages, optional_params=optional_params, litellm_params={}, headers={}, ) - - # Verify the structure + assert "additionalModelRequestFields" in request_data assert "tool_choice" in request_data["additionalModelRequestFields"] - assert "disable_parallel_tool_use" in request_data["additionalModelRequestFields"]["tool_choice"] assert request_data["additionalModelRequestFields"]["tool_choice"]["disable_parallel_tool_use"] is True + assert "parallel_tool_calls" not in request_data["additionalModelRequestFields"] + + +def test_parallel_tool_calls_older_model_drops_disable_flag(): + """Older Claude models (pre-4.5) must NOT receive disable_parallel_tool_use — Bedrock rejects it.""" + config = AmazonConverseConfig() + model = "anthropic.claude-3-5-sonnet-20241022-v2:0" + messages = [{"role": "user", "content": "What's the weather in SF and NYC?"}] + + optional_params = config.map_openai_params( + non_default_params={"parallel_tool_calls": False, "tools": _TOOL_PARAM}, + optional_params={}, + model=model, + drop_params=False, + ) + + request_data = config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + additional = request_data.get("additionalModelRequestFields", {}) + assert "tool_choice" not in additional + assert "parallel_tool_calls" not in additional class TestBedrockMinThinkingBudgetTokens: @@ -3479,3 +3493,262 @@ class TestBedrockMinThinkingBudgetTokens: drop_params=False, ) assert "thinking" not in result or result.get("thinking") is None + +def test_transform_response_with_both_json_tool_call_and_real_tool(): + """ + When Bedrock returns BOTH json_tool_call AND a real tool (get_weather), + only the real tool should remain in tool_calls. The json_tool_call should be filtered out. + Fixes https://github.com/BerriAI/litellm/issues/18381 + """ + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + response_json = { + "metrics": {"latencyMs": 200}, + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_json_001", + "name": "json_tool_call", + "input": { + "Current_Temperature": 62, + "Weather_Explanation": "Mild and cool.", + }, + } + }, + { + "toolUse": { + "toolUseId": "tooluse_weather_001", + "name": "get_weather", + "input": { + "location": "San Francisco, CA", + "unit": "fahrenheit", + }, + } + }, + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 100, + "outputTokens": 50, + "totalTokens": 150, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + + class MockResponse: + def json(self): + return response_json + + @property + def text(self): + return json.dumps(response_json) + + config = AmazonConverseConfig() + model_response = ModelResponse() + optional_params = {"json_mode": True} + + result = config._transform_response( + model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + response=MockResponse(), + model_response=model_response, + stream=False, + logging_obj=None, + optional_params=optional_params, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + # Only real tool should remain + assert result.choices[0].message.tool_calls is not None + assert len(result.choices[0].message.tool_calls) == 1 + assert result.choices[0].message.tool_calls[0].function.name == "get_weather" + assert ( + result.choices[0].message.tool_calls[0].function.arguments + == '{"location": "San Francisco, CA", "unit": "fahrenheit"}' + ) + + # json_tool_call content should be preserved as message text + content = result.choices[0].message.content + assert content is not None + parsed = json.loads(content) + assert parsed["Current_Temperature"] == 62 + assert parsed["Weather_Explanation"] == "Mild and cool." + + +def test_transform_response_does_not_mutate_optional_params(): + """ + Verify that optional_params still contains json_mode after _transform_response. + Previously, .pop() was used which mutated the caller's dict. + """ + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + response_json = { + "metrics": {"latencyMs": 50}, + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_001", + "name": "json_tool_call", + "input": {"result": "ok"}, + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 10, + "outputTokens": 5, + "totalTokens": 15, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + + class MockResponse: + def json(self): + return response_json + + @property + def text(self): + return json.dumps(response_json) + + config = AmazonConverseConfig() + model_response = ModelResponse() + optional_params = {"json_mode": True, "other_key": "value"} + + config._transform_response( + model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + response=MockResponse(), + model_response=model_response, + stream=False, + logging_obj=None, + optional_params=optional_params, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + # json_mode should still be in optional_params (not popped) + assert "json_mode" in optional_params + assert optional_params["json_mode"] is True + assert optional_params["other_key"] == "value" + + +def test_streaming_filters_json_tool_call_with_real_tools(): + """ + Simulate streaming chunks where both json_tool_call and a real tool arrive. + Verify json_tool_call chunks are converted to text content while real tool + chunks pass through normally. + """ + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + from litellm.types.llms.bedrock import ( + ContentBlockDeltaEvent, + ContentBlockStartEvent, + ) + + decoder = AWSEventStreamDecoder(model="test-model", json_mode=True) + + # Chunk 1: json_tool_call start + json_start = ContentBlockStartEvent( + toolUse={ + "toolUseId": "tooluse_json_001", + "name": "json_tool_call", + } + ) + tool_use_1, _, _ = decoder._handle_converse_start_event(json_start) + # json_tool_call start should be suppressed (return None tool_use) + assert tool_use_1 is None + # tool_calls_index should NOT have been incremented + assert decoder.tool_calls_index is None + + # Chunk 2: json_tool_call delta — should become text, not tool_use + json_delta = ContentBlockDeltaEvent(toolUse={"input": '{"temp": 62}'}) + text_2, tool_use_2, _, _, _ = decoder._handle_converse_delta_event( + json_delta, index=0 + ) + assert text_2 == '{"temp": 62}' + assert tool_use_2 is None + + # Chunk 3: json_tool_call stop + stop_tool = decoder._handle_converse_stop_event(index=0) + assert stop_tool is None + # _current_tool_name should be reset + assert decoder._current_tool_name is None + + # Chunk 4: real tool start + real_start = ContentBlockStartEvent( + toolUse={ + "toolUseId": "tooluse_weather_001", + "name": "get_weather", + } + ) + tool_use_4, _, _ = decoder._handle_converse_start_event(real_start) + assert tool_use_4 is not None + assert tool_use_4["function"]["name"] == "get_weather" + assert decoder.tool_calls_index == 0 + + # Chunk 5: real tool delta + real_delta = ContentBlockDeltaEvent( + toolUse={"input": '{"location": "SF"}'} + ) + text_5, tool_use_5, _, _, _ = decoder._handle_converse_delta_event( + real_delta, index=1 + ) + assert text_5 == "" + assert tool_use_5 is not None + assert tool_use_5["function"]["arguments"] == '{"location": "SF"}' + + +def test_streaming_without_json_mode_passes_all_tools(): + """ + Verify backward compatibility: when json_mode=False, all tools + (including json_tool_call if present) pass through unchanged. + """ + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + from litellm.types.llms.bedrock import ( + ContentBlockDeltaEvent, + ContentBlockStartEvent, + ) + + decoder = AWSEventStreamDecoder(model="test-model", json_mode=False) + + # json_tool_call start — should pass through when json_mode=False + json_start = ContentBlockStartEvent( + toolUse={ + "toolUseId": "tooluse_json_001", + "name": "json_tool_call", + } + ) + tool_use, _, _ = decoder._handle_converse_start_event(json_start) + assert tool_use is not None + assert tool_use["function"]["name"] == "json_tool_call" + assert decoder.tool_calls_index == 0 + + # json_tool_call delta — should be a tool_use, not text + json_delta = ContentBlockDeltaEvent(toolUse={"input": '{"data": 1}'}) + text, tool_use_delta, _, _, _ = decoder._handle_converse_delta_event( + json_delta, index=0 + ) + assert text == "" + assert tool_use_delta is not None + assert tool_use_delta["function"]["arguments"] == '{"data": 1}' + diff --git a/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py new file mode 100644 index 00000000000..33db9d33c1c --- /dev/null +++ b/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -0,0 +1,212 @@ +""" +Unit tests for extra_headers propagation in OpenAI image generation. + +Verifies that extra_headers passed to litellm.image_generation() / +litellm.aimage_generation() are forwarded to the OpenAI API client as +extra_headers in the images.generate() call. +""" + +import os +import sys +from unittest.mock import MagicMock, AsyncMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.openai.openai import OpenAIChatCompletion + + +@pytest.fixture +def openai_chat_completions(): + return OpenAIChatCompletion() + + +@pytest.fixture +def mock_logging_obj(): + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + logging_obj.post_call = MagicMock() + return logging_obj + + +class TestImageGenerationExtraHeaders: + """Test that extra_headers are properly injected into OpenAI image generation calls.""" + + def test_sync_image_generation_with_headers( + self, openai_chat_completions, mock_logging_obj + ): + """Sync image_generation should pass headers as extra_headers to images.generate().""" + mock_image_data = MagicMock() + mock_image_data.model_dump.return_value = { + "created": 1700000000, + "data": [{"url": "https://example.com/image.png"}], + } + + mock_openai_client = MagicMock() + mock_openai_client.images.generate.return_value = mock_image_data + mock_openai_client.api_key = "test-key" + mock_openai_client._base_url._uri_reference = "https://api.openai.com" + + test_headers = {"cf-aig-authorization": "Bearer custom-token"} + + openai_chat_completions.image_generation( + model="dall-e-3", + prompt="A white cat", + timeout=60.0, + optional_params={}, + logging_obj=mock_logging_obj, + api_key="test-key", + headers=test_headers, + client=mock_openai_client, + ) + + _, kwargs = mock_openai_client.images.generate.call_args + assert kwargs.get("extra_headers") == test_headers + + def test_sync_image_generation_without_headers( + self, openai_chat_completions, mock_logging_obj + ): + """Sync image_generation without headers should not inject extra_headers.""" + mock_image_data = MagicMock() + mock_image_data.model_dump.return_value = { + "created": 1700000000, + "data": [{"url": "https://example.com/image.png"}], + } + + mock_openai_client = MagicMock() + mock_openai_client.images.generate.return_value = mock_image_data + mock_openai_client.api_key = "test-key" + mock_openai_client._base_url._uri_reference = "https://api.openai.com" + + openai_chat_completions.image_generation( + model="dall-e-3", + prompt="A white cat", + timeout=60.0, + optional_params={}, + logging_obj=mock_logging_obj, + api_key="test-key", + client=mock_openai_client, + ) + + _, kwargs = mock_openai_client.images.generate.call_args + assert "extra_headers" not in kwargs + + @pytest.mark.asyncio + async def test_async_image_generation_with_headers( + self, openai_chat_completions, mock_logging_obj + ): + """Async aimage_generation should pass headers as extra_headers to images.generate().""" + mock_image_data = MagicMock() + mock_image_data.model_dump.return_value = { + "created": 1700000000, + "data": [{"url": "https://example.com/image.png"}], + } + + mock_openai_client = MagicMock() + mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.api_key = "test-key" + + test_headers = {"cf-aig-authorization": "Bearer custom-token"} + + await openai_chat_completions.aimage_generation( + prompt="A white cat", + data={"model": "dall-e-3", "prompt": "A white cat"}, + model_response=MagicMock(), + timeout=60.0, + logging_obj=mock_logging_obj, + api_key="test-key", + headers=test_headers, + client=mock_openai_client, + ) + + _, kwargs = mock_openai_client.images.generate.call_args + assert kwargs.get("extra_headers") == test_headers + + @pytest.mark.asyncio + async def test_async_image_generation_without_headers( + self, openai_chat_completions, mock_logging_obj + ): + """Async aimage_generation without headers should not inject extra_headers.""" + mock_image_data = MagicMock() + mock_image_data.model_dump.return_value = { + "created": 1700000000, + "data": [{"url": "https://example.com/image.png"}], + } + + mock_openai_client = MagicMock() + mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.api_key = "test-key" + + await openai_chat_completions.aimage_generation( + prompt="A white cat", + data={"model": "dall-e-3", "prompt": "A white cat"}, + model_response=MagicMock(), + timeout=60.0, + logging_obj=mock_logging_obj, + api_key="test-key", + client=mock_openai_client, + ) + + _, kwargs = mock_openai_client.images.generate.call_args + assert "extra_headers" not in kwargs + + def test_sync_image_generation_forwards_headers_to_async( + self, openai_chat_completions, mock_logging_obj + ): + """When aimg_generation=True, image_generation should forward headers to aimage_generation.""" + with patch.object( + openai_chat_completions, "aimage_generation" + ) as mock_aimage_gen: + mock_aimage_gen.return_value = MagicMock() + + test_headers = {"x-custom-header": "value"} + + openai_chat_completions.image_generation( + model="dall-e-3", + prompt="A white cat", + timeout=60.0, + optional_params={}, + logging_obj=mock_logging_obj, + api_key="test-key", + aimg_generation=True, + headers=test_headers, + ) + + mock_aimage_gen.assert_called_once() + call_kwargs = mock_aimage_gen.call_args[1] + assert call_kwargs["headers"] == test_headers + + +class TestImageGenerationEntryPointHeaders: + """Test that litellm.image_generation() passes headers through to the OpenAI provider.""" + + @pytest.mark.asyncio + async def test_extra_headers_reach_openai_provider(self): + """End-to-end: extra_headers from litellm.aimage_generation() reach OpenAI images.generate().""" + import litellm + + mock_image_data = MagicMock() + mock_image_data.model_dump.return_value = { + "created": 1700000000, + "data": [{"url": "https://example.com/image.png"}], + } + + mock_openai_client = MagicMock() + mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.api_key = "test-key" + mock_openai_client._base_url._uri_reference = "https://api.openai.com" + + test_headers = {"cf-aig-authorization": "Bearer my-secret"} + + await litellm.aimage_generation( + model="dall-e-3", + prompt="A white cat", + extra_headers=test_headers, + client=mock_openai_client, + api_key="test-key", + ) + + mock_openai_client.images.generate.assert_called_once() + _, kwargs = mock_openai_client.images.generate.call_args + assert kwargs.get("extra_headers") == test_headers diff --git a/tests/test_litellm/ocr/__init__.py b/tests/test_litellm/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py new file mode 100644 index 00000000000..492253e2f11 --- /dev/null +++ b/tests/test_litellm/ocr/test_ocr_file_input.py @@ -0,0 +1,464 @@ +""" +Tests for OCR file input support. + +Tests that: +1. The SDK document parameter with type="file" correctly converts file paths, + file objects, and raw bytes to base64 data URIs before sending to providers. +2. The proxy _build_document_from_upload helper correctly handles uploaded file bytes. +3. The proxy rejects type="file" documents received via JSON (security guard). +4. The proxy returns user-friendly errors for invalid JSON bodies. +""" +import base64 +import os +import tempfile +from io import BytesIO +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import orjson +import pytest + +from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type + + +class TestGetMimeType: + def test_should_detect_pdf_mime_type(self): + assert get_mime_type("document.pdf") == "application/pdf" + + def test_should_detect_png_mime_type(self): + assert get_mime_type("image.png") == "image/png" + + def test_should_detect_jpg_mime_type(self): + assert get_mime_type("photo.jpg") == "image/jpeg" + + def test_should_detect_jpeg_mime_type(self): + assert get_mime_type("photo.jpeg") == "image/jpeg" + + def test_should_detect_gif_mime_type(self): + assert get_mime_type("animation.gif") == "image/gif" + + def test_should_detect_webp_mime_type(self): + assert get_mime_type("image.webp") == "image/webp" + + def test_should_detect_tiff_mime_type(self): + assert get_mime_type("scan.tiff") == "image/tiff" + + def test_should_detect_tif_mime_type(self): + assert get_mime_type("scan.tif") == "image/tiff" + + def test_should_detect_bmp_mime_type(self): + assert get_mime_type("bitmap.bmp") == "image/bmp" + + def test_should_be_case_insensitive(self): + assert get_mime_type("DOCUMENT.PDF") == "application/pdf" + assert get_mime_type("IMAGE.PNG") == "image/png" + + def test_should_fallback_for_unknown_extension(self): + result = get_mime_type("file.xyz123") + assert isinstance(result, str) + + +class TestConvertFileDocumentToUrlDocument: + def test_should_convert_pdf_file_path_to_document_url(self): + """File path to a PDF should produce type=document_url with base64 data URI.""" + pdf_content = b"%PDF-1.4 test content" + + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: + f.write(pdf_content) + f.flush() + tmp_path = f.name + + try: + result = convert_file_document_to_url_document( + {"type": "file", "file": tmp_path} + ) + + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + b64_data = result["document_url"].split(";base64,")[1] + assert base64.b64decode(b64_data) == pdf_content + finally: + os.unlink(tmp_path) + + def test_should_convert_image_file_path_to_image_url(self): + """File path to a PNG image should produce type=image_url with base64 data URI.""" + png_content = b"\x89PNG\r\n\x1a\n fake png content" + + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: + f.write(png_content) + f.flush() + tmp_path = f.name + + try: + result = convert_file_document_to_url_document( + {"type": "file", "file": tmp_path} + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/png;base64,") + + b64_data = result["image_url"].split(";base64,")[1] + assert base64.b64decode(b64_data) == png_content + finally: + os.unlink(tmp_path) + + def test_should_convert_pathlib_path(self): + """pathlib.Path objects should work the same as string paths.""" + content = b"test pdf content" + + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: + f.write(content) + f.flush() + tmp_path = Path(f.name) + + try: + result = convert_file_document_to_url_document( + {"type": "file", "file": tmp_path} + ) + + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + finally: + os.unlink(str(tmp_path)) + + def test_should_convert_raw_bytes(self): + """Raw bytes should be converted using a fallback MIME type.""" + content = b"raw bytes content" + + result = convert_file_document_to_url_document( + {"type": "file", "file": content} + ) + + assert result["type"] == "document_url" + assert "base64," in result["document_url"] + + b64_data = result["document_url"].split(";base64,")[1] + assert base64.b64decode(b64_data) == content + + def test_should_convert_raw_bytes_with_explicit_mime_type(self): + """Raw bytes with explicit mime_type should use the specified MIME type.""" + content = b"raw pdf content" + + result = convert_file_document_to_url_document( + {"type": "file", "file": content, "mime_type": "application/pdf"} + ) + + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_convert_raw_bytes_with_image_mime_type(self): + """Raw bytes with an image MIME type should produce type=image_url.""" + content = b"raw image content" + + result = convert_file_document_to_url_document( + {"type": "file", "file": content, "mime_type": "image/jpeg"} + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/jpeg;base64,") + + def test_should_convert_file_like_object(self): + """BytesIO and other file-like objects should be supported.""" + content = b"file-like content" + file_obj = BytesIO(content) + + result = convert_file_document_to_url_document( + {"type": "file", "file": file_obj} + ) + + assert result["type"] == "document_url" + assert "base64," in result["document_url"] + + def test_should_convert_file_like_object_with_name(self): + """File-like objects with a .name attribute should detect MIME from the name.""" + content = b"file-like png content" + file_obj = BytesIO(content) + file_obj.name = "test_image.png" + + result = convert_file_document_to_url_document( + {"type": "file", "file": file_obj} + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/png;base64,") + + def test_should_raise_error_for_missing_file_field(self): + """Missing 'file' field should raise ValueError.""" + with pytest.raises(ValueError, match="must include a 'file' field"): + convert_file_document_to_url_document({"type": "file"}) + + def test_should_raise_error_for_nonexistent_file_path(self): + """Non-existent file path should raise FileNotFoundError.""" + with pytest.raises(FileNotFoundError, match="File not found"): + convert_file_document_to_url_document( + {"type": "file", "file": "/nonexistent/path/to/file.pdf"} + ) + + def test_should_raise_error_for_empty_file(self): + """Empty file should raise ValueError.""" + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: + tmp_path = f.name + + try: + with pytest.raises(ValueError, match="File is empty"): + convert_file_document_to_url_document( + {"type": "file", "file": tmp_path} + ) + finally: + os.unlink(tmp_path) + + def test_should_raise_error_for_unsupported_type(self): + """Unsupported file input types should raise ValueError.""" + with pytest.raises(ValueError, match="Unsupported file input type"): + convert_file_document_to_url_document({"type": "file", "file": 12345}) + + def test_should_raise_error_for_invalid_mime_type(self): + """MIME types with special characters should be rejected.""" + content = b"some content" + with pytest.raises(ValueError, match="Invalid MIME type"): + convert_file_document_to_url_document( + {"type": "file", "file": content, "mime_type": "text/html; charset=utf-8\nX-Injected: true"} + ) + + def test_should_override_mime_type_for_file_path(self): + """Explicit mime_type should override auto-detection from extension.""" + content = b"some content" + + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f: + f.write(content) + f.flush() + tmp_path = f.name + + try: + result = convert_file_document_to_url_document( + {"type": "file", "file": tmp_path, "mime_type": "image/png"} + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/png;base64,") + finally: + os.unlink(tmp_path) + + +class TestBuildDocumentFromUpload: + """Test the proxy endpoint's file upload to document conversion helper.""" + + @pytest.fixture(autouse=True) + def _import_helper(self): + """Import the proxy helper, skip if proxy deps aren't installed.""" + try: + from litellm.proxy.ocr_endpoints.endpoints import ( + _build_document_from_upload, + ) + + self._build = _build_document_from_upload + except ImportError: + pytest.skip("Proxy dependencies (fastapi/orjson) not installed") + + def test_should_build_document_url_for_pdf(self): + content = b"%PDF-1.4 test content" + + result = self._build( + file_content=content, + filename="document.pdf", + content_type="application/pdf", + ) + + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + b64_data = result["document_url"].split(";base64,")[1] + assert base64.b64decode(b64_data) == content + + def test_should_build_image_url_for_png(self): + content = b"\x89PNG fake png" + + result = self._build( + file_content=content, + filename="screenshot.png", + content_type="image/png", + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/png;base64,") + + def test_should_build_image_url_for_jpeg(self): + content = b"\xff\xd8\xff fake jpeg" + + result = self._build( + file_content=content, + filename="photo.jpg", + content_type="image/jpeg", + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/jpeg;base64,") + + def test_should_detect_mime_from_filename_when_content_type_is_octet_stream(self): + content = b"pdf content" + + result = self._build( + file_content=content, + filename="report.pdf", + content_type="application/octet-stream", + ) + + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_detect_mime_from_filename_when_content_type_is_none(self): + content = b"png content" + + result = self._build( + file_content=content, + filename="image.png", + content_type=None, + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/png;base64,") + + def test_should_fallback_to_octet_stream_for_unknown(self): + content = b"unknown content" + + result = self._build( + file_content=content, + filename=None, + content_type=None, + ) + + assert result["type"] == "document_url" + assert "application/octet-stream" in result["document_url"] + + def test_should_preserve_base64_content_correctly(self): + content = b"Hello, World! \x00\x01\x02\xff" + + result = self._build( + file_content=content, + filename="test.pdf", + content_type="application/pdf", + ) + + b64_data = result["document_url"].split(";base64,")[1] + assert base64.b64decode(b64_data) == content + + def test_should_strip_mime_parameters_from_content_type(self): + """Content-Type with parameters (e.g. charset) should be stripped to the base MIME type.""" + content = b"%PDF-1.4 test" + + result = self._build( + file_content=content, + filename="doc.pdf", + content_type="application/pdf; charset=utf-8", + ) + + assert result["type"] == "document_url" + assert result["document_url"].startswith("data:application/pdf;base64,") + + def test_should_strip_mime_parameters_with_multiple_params(self): + """Content-Type with multiple parameters should still be stripped correctly.""" + content = b"image data" + + result = self._build( + file_content=content, + filename="img.png", + content_type="image/png; charset=utf-8; boundary=something", + ) + + assert result["type"] == "image_url" + assert result["image_url"].startswith("data:image/png;base64,") + + +class TestProxySecurityGuard: + """Test that the proxy rejects type='file' documents in JSON requests + and that multipart form fields cannot override the constructed document.""" + + @pytest.fixture(autouse=True) + def _import_helpers(self): + """Import the proxy helpers, skip if proxy deps aren't installed.""" + try: + from litellm.proxy.ocr_endpoints.endpoints import ( + _parse_multipart_form, + _parse_ocr_request, + ) + + self._parse = _parse_ocr_request + self._parse_multipart = _parse_multipart_form + except ImportError: + pytest.skip("Proxy dependencies (fastapi/orjson) not installed") + + @pytest.mark.asyncio + async def test_should_reject_file_type_document_in_json_body(self): + """type='file' in a JSON body must be rejected to prevent server-side file reads.""" + body = orjson.dumps( + { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "file", "file": "/etc/passwd"}, + } + ) + + mock_request = MagicMock() + mock_request.headers = {"content-type": "application/json"} + mock_request.body = AsyncMock(return_value=body) + mock_request._form = None + + with pytest.raises(ValueError, match="not supported through the JSON API"): + await self._parse(mock_request) + + @pytest.mark.asyncio + async def test_should_accept_document_url_type_in_json_body(self): + """type='document_url' in a JSON body should pass through normally.""" + expected = { + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf", + }, + } + body = orjson.dumps(expected) + + mock_request = MagicMock() + mock_request.headers = {"content-type": "application/json"} + mock_request.body = AsyncMock(return_value=body) + mock_request._form = None + + result = await self._parse(mock_request) + assert result["document"]["type"] == "document_url" + + @pytest.mark.asyncio + async def test_should_raise_on_invalid_json_body(self): + """Invalid JSON should produce a user-friendly ValueError.""" + mock_request = MagicMock() + mock_request.headers = {"content-type": "application/json"} + mock_request.body = AsyncMock(return_value=b"not valid json{{{") + mock_request._form = None + + with pytest.raises(ValueError, match="Invalid JSON in request body"): + await self._parse(mock_request) + + @pytest.mark.asyncio + async def test_should_ignore_document_form_field_injection(self): + """A 'document' form field must not override the document built from the uploaded file.""" + from starlette.datastructures import UploadFile + + file_content = b"%PDF-1.4 legit content" + upload = UploadFile(filename="legit.pdf", file=BytesIO(file_content)) + + injected = '{"type": "file", "file": "/etc/passwd"}' + + mock_form = { + "file": upload, + "model": "mistral/mistral-ocr-latest", + "document": injected, + } + + mock_request = MagicMock() + mock_request.headers = {"content-type": "multipart/form-data; boundary=---"} + mock_request.form = AsyncMock(return_value=mock_form) + + result = await self._parse_multipart(mock_request) + + assert result["document"]["type"] == "document_url" + assert result["document"]["document_url"].startswith("data:application/pdf;base64,") + assert result["model"] == "mistral/mistral-ocr-latest" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 8cec3538077..fcf8f048190 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -7,6 +7,7 @@ from fastapi.testclient import TestClient from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints from litellm.proxy.agent_endpoints.endpoints import ( + _check_agent_management_permission, get_agent_daily_activity, router, user_api_key_auth, @@ -47,6 +48,16 @@ def _sample_agent_response( ) +def _make_app_with_role(role: LitellmUserRoles) -> TestClient: + """Create a TestClient where the auth dependency returns the given role.""" + test_app = FastAPI() + test_app.include_router(router) + test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user", user_role=role + ) + return TestClient(test_app) + + app = FastAPI() app.include_router(router) app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( @@ -258,3 +269,173 @@ async def test_get_agent_daily_activity_with_agent_names(monkeypatch): "agent-1": {"agent_name": "First Agent"}, "agent-2": {"agent_name": "Second Agent"}, } + + +# ---------- RBAC enforcement tests ---------- + + +class TestAgentRBACInternalUser: + """Internal users should be able to read agents but not create/update/delete.""" + + @pytest.fixture(autouse=True) + def _setup(self, monkeypatch): + self.internal_client = _make_app_with_role(LitellmUserRoles.INTERNAL_USER) + self.mock_registry = MagicMock() + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry) + + def test_should_allow_internal_user_to_list_agents(self, monkeypatch): + self.mock_registry.get_agent_list = MagicMock(return_value=[]) + resp = self.internal_client.get( + "/v1/agents", headers={"Authorization": "Bearer k"} + ) + assert resp.status_code == 200 + + def test_should_allow_internal_user_to_get_agent_by_id(self, monkeypatch): + self.mock_registry.get_agent_by_id = MagicMock( + return_value=_sample_agent_response() + ) + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + resp = self.internal_client.get( + "/v1/agents/agent-123", headers={"Authorization": "Bearer k"} + ) + assert resp.status_code == 200 + + def test_should_block_internal_user_from_creating_agent(self): + resp = self.internal_client.post( + "/v1/agents", + json=_sample_agent_config(), + headers={"Authorization": "Bearer k"}, + ) + assert resp.status_code == 403 + assert "Only proxy admins" in resp.json()["detail"]["error"] + + def test_should_block_internal_user_from_updating_agent(self): + resp = self.internal_client.put( + "/v1/agents/agent-123", + json=_sample_agent_config(), + headers={"Authorization": "Bearer k"}, + ) + assert resp.status_code == 403 + + def test_should_block_internal_user_from_patching_agent(self): + resp = self.internal_client.patch( + "/v1/agents/agent-123", + json={"agent_name": "new-name"}, + headers={"Authorization": "Bearer k"}, + ) + assert resp.status_code == 403 + + def test_should_block_internal_user_from_deleting_agent(self): + resp = self.internal_client.delete( + "/v1/agents/agent-123", headers={"Authorization": "Bearer k"} + ) + assert resp.status_code == 403 + + +class TestAgentRBACInternalUserViewOnly: + """View-only internal users should only be able to read agents.""" + + @pytest.fixture(autouse=True) + def _setup(self, monkeypatch): + self.viewer_client = _make_app_with_role( + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ) + self.mock_registry = MagicMock() + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry) + + def test_should_allow_view_only_user_to_list_agents(self): + self.mock_registry.get_agent_list = MagicMock(return_value=[]) + resp = self.viewer_client.get( + "/v1/agents", headers={"Authorization": "Bearer k"} + ) + assert resp.status_code == 200 + + def test_should_block_view_only_user_from_creating_agent(self): + resp = self.viewer_client.post( + "/v1/agents", + json=_sample_agent_config(), + headers={"Authorization": "Bearer k"}, + ) + assert resp.status_code == 403 + + def test_should_block_view_only_user_from_deleting_agent(self): + resp = self.viewer_client.delete( + "/v1/agents/agent-123", headers={"Authorization": "Bearer k"} + ) + assert resp.status_code == 403 + + +class TestAgentRBACProxyAdmin: + """Proxy admins should have full CRUD access to agents.""" + + @pytest.fixture(autouse=True) + def _setup(self, monkeypatch): + self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN) + self.mock_registry = MagicMock() + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry) + + def test_should_allow_admin_to_create_agent(self, monkeypatch): + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + self.mock_registry.get_agent_by_name = MagicMock(return_value=None) + self.mock_registry.add_agent_to_db = AsyncMock( + return_value=_sample_agent_response() + ) + self.mock_registry.register_agent = MagicMock() + resp = self.admin_client.post( + "/v1/agents", + json=_sample_agent_config(), + headers={"Authorization": "Bearer k"}, + ) + assert resp.status_code == 200 + + def test_should_allow_admin_to_delete_agent(self): + existing = { + "agent_id": "agent-123", + "agent_name": "Existing Agent", + "agent_card_params": _sample_agent_card_params(), + } + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=existing + ) + self.mock_registry.delete_agent_from_db = AsyncMock() + self.mock_registry.deregister_agent = MagicMock() + resp = self.admin_client.delete( + "/v1/agents/agent-123", headers={"Authorization": "Bearer k"} + ) + assert resp.status_code == 200 + + +class TestCheckAgentManagementPermission: + """Unit tests for the _check_agent_management_permission helper.""" + + def test_should_allow_proxy_admin(self): + auth = UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + _check_agent_management_permission(auth) + + @pytest.mark.parametrize( + "role", + [ + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ], + ) + def test_should_block_non_admin_roles(self, role): + from fastapi import HTTPException + + auth = UserAPIKeyAuth(user_id="user", user_role=role) + with pytest.raises(HTTPException) as exc_info: + _check_agent_management_permission(auth) + assert exc_info.value.status_code == 403 + + +class TestAgentRoutesIncludesAgentIdPattern: + """Verify that agent_routes includes the {agent_id} pattern for route access.""" + + def test_should_include_agent_id_pattern(self): + from litellm.proxy._types import LiteLLMRoutes + + assert "/v1/agents/{agent_id}" in LiteLLMRoutes.agent_routes.value diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py new file mode 100644 index 00000000000..73a97188424 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -0,0 +1,119 @@ +import pytest +from unittest.mock import AsyncMock, patch +import litellm +from litellm.proxy.auth.user_api_key_auth import ( + _run_post_custom_auth_checks, + update_valid_token_with_end_user_params, +) +from litellm.proxy._types import UserAPIKeyAuth + + +@pytest.mark.asyncio +async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): + # Test backwards compatibility + valid_token = UserAPIKeyAuth(token="test_token") + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock + ) as mock_common: + mock_common.return_value = True + result = await _run_post_custom_auth_checks( + valid_token=valid_token, + request=None, + request_data={}, + route="/v1/chat/completions", + parent_otel_span=None, + ) + assert result.token == "test_token" + assert getattr(result, "end_user_id", None) is None + mock_common.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_custom_auth_run_post_custom_auth_checks_with_end_user_budget_exceeded(): + valid_token = UserAPIKeyAuth( + token="test_token", + end_user_id="test_user", + end_user_model_max_budget={ + "gpt-4": {"budget_limit": 10.0, "time_period": "1d"} + }, + ) + request_data = {"model": "gpt-4"} + + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock + ): + with patch( + "litellm.proxy.proxy_server.model_max_budget_limiter.is_end_user_within_model_budget", + new_callable=AsyncMock, + ) as mock_budget_check: + mock_budget_check.side_effect = litellm.BudgetExceededError( + message="Exceeded budget", current_cost=20.0, max_budget=10.0 + ) + + with pytest.raises(litellm.BudgetExceededError): + await _run_post_custom_auth_checks( + valid_token=valid_token, + request=None, + request_data=request_data, + route="/v1/chat/completions", + parent_otel_span=None, + ) + mock_budget_check.assert_awaited_once() + + +def test_update_valid_token_does_not_override_custom_auth_values_with_none(): + """ + Greptile feedback: if custom auth sets end_user_model_max_budget on the token, + but the DB end_user has no model_max_budget in their budget table, the DB lookup + should NOT clear the custom-auth-provided value. + """ + custom_auth_budget = {"gpt-4": {"budget_limit": 5.0, "time_period": "1d"}} + valid_token = UserAPIKeyAuth( + token="test_token", + end_user_id="user_1", + end_user_tpm_limit=100, + end_user_rpm_limit=50, + end_user_model_max_budget=custom_auth_budget, + ) + + # Simulate DB lookup that found the end_user but budget table has no limits set + end_user_params = { + "end_user_id": "user_1", + "allowed_model_region": None, + # No tpm_limit, rpm_limit, or model_max_budget from DB + } + + result = update_valid_token_with_end_user_params(valid_token, end_user_params) + + # Custom-auth-provided values should be preserved, not cleared to None + assert result.end_user_tpm_limit == 100 + assert result.end_user_rpm_limit == 50 + assert result.end_user_model_max_budget == custom_auth_budget + assert result.end_user_id == "user_1" + + +def test_update_valid_token_db_values_override_custom_auth_when_set(): + """ + When the DB budget table has explicit values, they should override + whatever the custom auth function set (DB is source of truth). + """ + valid_token = UserAPIKeyAuth( + token="test_token", + end_user_id="user_1", + end_user_tpm_limit=100, + end_user_model_max_budget={"gpt-4": {"budget_limit": 5.0, "time_period": "1d"}}, + ) + + db_budget = {"gpt-4": {"budget_limit": 20.0, "time_period": "1d"}} + end_user_params = { + "end_user_id": "user_1", + "end_user_tpm_limit": 500, + "end_user_model_max_budget": db_budget, + } + + result = update_valid_token_with_end_user_params(valid_token, end_user_params) + + # DB values should win + assert result.end_user_tpm_limit == 500 + assert result.end_user_model_max_budget == db_budget diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index b56d13bb932..8418dde5e9c 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1485,4 +1485,273 @@ async def test_get_objects_resolves_org_by_name(): ) +# --------------------------------------------------------------------------- +# Fix 1: OIDC discovery URL resolution +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_resolve_jwks_url_passthrough_for_direct_jwks_url(): + """Non-discovery URLs are returned unchanged.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + url = "https://login.microsoftonline.com/common/discovery/keys" + result = await handler._resolve_jwks_url(url) + assert result == url + + +@pytest.mark.asyncio +async def test_resolve_jwks_url_resolves_oidc_discovery_document(): + """ + A .well-known/openid-configuration URL should be fetched and its + jwks_uri returned. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + cache = DualCache() + handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" + jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"jwks_uri": jwks_url, "issuer": "https://..."} + + with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response) as mock_get: + result = await handler._resolve_jwks_url(discovery_url) + + assert result == jwks_url + mock_get.assert_called_once_with(discovery_url) + + +@pytest.mark.asyncio +async def test_resolve_jwks_url_caches_resolved_jwks_uri(): + """Resolved jwks_uri is cached — second call does not hit the network.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + cache = DualCache() + handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" + jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" + + mock_response = MagicMock() + mock_response.json.return_value = {"jwks_uri": jwks_url} + + with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response) as mock_get: + first = await handler._resolve_jwks_url(discovery_url) + second = await handler._resolve_jwks_url(discovery_url) + + assert first == jwks_url + assert second == jwks_url + # Network should only be hit once + assert mock_get.call_count == 1 + + +@pytest.mark.asyncio +async def test_resolve_jwks_url_raises_if_no_jwks_uri_in_discovery_doc(): + """Raise a helpful error if the discovery document has no jwks_uri.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + discovery_url = "https://example.com/.well-known/openid-configuration" + mock_response = MagicMock() + mock_response.json.return_value = {"issuer": "https://example.com"} # no jwks_uri + + with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response): + with pytest.raises(Exception, match="jwks_uri"): + await handler._resolve_jwks_url(discovery_url) + + +# --------------------------------------------------------------------------- +# Fix 2: handle array values in team_id_jwt_field (e.g. AAD "roles" claim) +# --------------------------------------------------------------------------- + + +def _make_jwt_handler(team_id_jwt_field: str) -> JWTHandler: + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(team_id_jwt_field=team_id_jwt_field), + ) + return handler + + +def test_get_team_id_returns_first_element_when_roles_is_list(): + """ + AAD sends roles as a list. get_team_id() must return the first string + element rather than the raw list (which would later crash with + 'unhashable type: list'). + """ + handler = _make_jwt_handler("roles") + token = {"oid": "user-oid", "roles": ["team1"]} + result = handler.get_team_id(token=token, default_value=None) + assert result == "team1" + + +def test_get_team_id_returns_first_element_from_multi_value_roles_list(): + """When roles has multiple entries, the first one is used.""" + handler = _make_jwt_handler("roles") + token = {"roles": ["team2", "team1"]} + result = handler.get_team_id(token=token, default_value=None) + assert result == "team2" + + +def test_get_team_id_returns_default_when_roles_list_is_empty(): + """Empty list should fall back to default_value.""" + handler = _make_jwt_handler("roles") + token = {"roles": []} + result = handler.get_team_id(token=token, default_value="fallback") + assert result == "fallback" + + +def test_get_team_id_still_works_with_string_value(): + """String values (non-array) continue to work as before.""" + handler = _make_jwt_handler("appid") + token = {"appid": "my-team-id"} + result = handler.get_team_id(token=token, default_value=None) + assert result == "my-team-id" + + +def test_get_team_id_list_result_is_hashable(): + """ + The value returned by get_team_id() must be hashable so it can be + added to a set (the operation that previously crashed). + """ + handler = _make_jwt_handler("roles") + token = {"roles": ["team1"]} + result = handler.get_team_id(token=token, default_value=None) + # This must not raise TypeError + s: set = set() + s.add(result) + assert "team1" in s + + +# --------------------------------------------------------------------------- +# Fix 3: helpful error message for dot-notation array indexing (roles.0) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_hints_bracket_notation(): + """ + When team_id_jwt_field is set to 'roles.0' (unsupported dot-notation for + array indexing) and no team is found, the exception message should suggest + using 'roles' instead (and explain LiteLLM auto-unwraps list values). + """ + from unittest.mock import MagicMock + + from litellm.caching.dual_cache import DualCache + + handler = _make_jwt_handler("roles.0") + # token has roles as a list — dot-notation won't find anything + token = {"roles": ["team1"]} + + with pytest.raises(Exception) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + error_msg = str(exc_info.value) + # Should mention the bad field name and suggest the fix + assert "roles.0" in error_msg, f"Expected field name in: {error_msg}" + assert "roles" in error_msg and "list" in error_msg, ( + f"Expected hint about using 'roles' instead: {error_msg}" + ) + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_hints_bracket_index_notation(): + """ + When team_id_jwt_field is set to 'roles[0]' (bracket indexing, also unsupported + in get_nested_value) the error message should suggest using 'roles' instead. + """ + from unittest.mock import MagicMock + + from litellm.caching.dual_cache import DualCache + + handler = _make_jwt_handler("roles[0]") + token = {"roles": ["team1"]} + + with pytest.raises(Exception) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + error_msg = str(exc_info.value) + assert "roles[0]" in error_msg, f"Expected field name in: {error_msg}" + assert "roles" in error_msg and "list" in error_msg, ( + f"Expected hint about using 'roles' instead: {error_msg}" + ) + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_no_hint_for_valid_field(): + """ + When team_id_jwt_field is a normal field name (no dot-notation) the + error message should not contain a spurious bracket-notation hint. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.dual_cache import DualCache + + handler = _make_jwt_handler("appid") + token = {} # no appid — triggers the "no team found" path + + with pytest.raises(Exception) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + error_msg = str(exc_info.value) + assert "Hint" not in error_msg diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index 9b6e0631762..32fd0750de8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -99,10 +99,12 @@ def client_and_mocks(monkeypatch): mock_team_table = MagicMock() mock_team_table.find_many = AsyncMock(return_value=[]) + mock_team_table.find_unique = AsyncMock(return_value=None) mock_team_table.update = AsyncMock(return_value=None) mock_key_table = MagicMock() mock_key_table.find_many = AsyncMock(return_value=[]) + mock_key_table.find_unique = AsyncMock(return_value=None) mock_key_table.update = AsyncMock(return_value=None) @asynccontextmanager @@ -570,11 +572,13 @@ def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks): team_with_group.team_id = "team-1" team_with_group.access_group_ids = ["ag-to-delete", "ag-other"] mock_team_table.find_many = AsyncMock(return_value=[team_with_group]) + mock_team_table.find_unique = AsyncMock(return_value=team_with_group) key_with_group = MagicMock() key_with_group.token = "key-token-1" key_with_group.access_group_ids = ["ag-to-delete"] mock_key_table.find_many = AsyncMock(return_value=[key_with_group]) + mock_key_table.find_unique = AsyncMock(return_value=key_with_group) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 204 @@ -669,11 +673,13 @@ def test_delete_access_group_patches_cached_team_and_key( team_with_group.team_id = "team-1" team_with_group.access_group_ids = ["ag-to-delete", "ag-keep"] mock_team_table.find_many = AsyncMock(return_value=[team_with_group]) + mock_team_table.find_unique = AsyncMock(return_value=team_with_group) key_with_group = MagicMock() key_with_group.token = "hashed-key-1" key_with_group.access_group_ids = ["ag-to-delete"] mock_key_table.find_many = AsyncMock(return_value=[key_with_group]) + mock_key_table.find_unique = AsyncMock(return_value=key_with_group) # Build cached team object (returned from proxy_logging dual cache) if team_cache_group_ids is not None: @@ -762,6 +768,7 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): key_with_group.token = "hashed-key-dict" key_with_group.access_group_ids = ["ag-to-delete", "ag-other"] mock_key_table.find_many = AsyncMock(return_value=[key_with_group]) + mock_key_table.find_unique = AsyncMock(return_value=key_with_group) # No team in cache mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( @@ -882,3 +889,304 @@ def test_record_to_access_group_table(): assert result.access_group_name == "unit-test-group" assert result.access_model_names == ["gpt-4", "claude-3"] assert result.access_agent_ids == ["agent-1"] + + +# --------------------------------------------------------------------------- +# Sync tests: CREATE +# --------------------------------------------------------------------------- + + +def test_create_access_group_syncs_assigned_teams(client_and_mocks): + """Create adds access_group_id to each assigned team's access_group_ids in DB.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + team_record = MagicMock() + team_record.team_id = "team-1" + team_record.access_group_ids = [] + mock_team_table.find_unique = AsyncMock(return_value=team_record) + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "new-group", "assigned_team_ids": ["team-1"]}, + ) + assert resp.status_code == 201 + + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-1"}) + mock_team_table.update.assert_awaited_once() + call_kwargs = mock_team_table.update.call_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-1"} + # The newly created access group id ("ag-new") should be in the updated list + assert "ag-new" in call_kwargs["data"]["access_group_ids"] + + +def test_create_access_group_syncs_assigned_keys(client_and_mocks): + """Create adds access_group_id to each assigned key's access_group_ids in DB.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_key_table = mock_prisma.db.litellm_verificationtoken + + key_record = MagicMock() + key_record.token = "hashed-token-1" + key_record.access_group_ids = [] + mock_key_table.find_unique = AsyncMock(return_value=key_record) + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "new-group", "assigned_key_ids": ["hashed-token-1"]}, + ) + assert resp.status_code == 201 + + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "hashed-token-1"}) + mock_key_table.update.assert_awaited_once() + call_kwargs = mock_key_table.update.call_args.kwargs + assert call_kwargs["where"] == {"token": "hashed-token-1"} + assert "ag-new" in call_kwargs["data"]["access_group_ids"] + + +def test_create_access_group_skips_sync_for_nonexistent_team(client_and_mocks): + """Create skips updating a team that doesn't exist in DB.""" + client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + mock_team_table.find_unique = AsyncMock(return_value=None) + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "new-group", "assigned_team_ids": ["nonexistent-team"]}, + ) + assert resp.status_code == 201 + mock_team_table.update.assert_not_awaited() + + +def test_create_access_group_idempotent_team_sync(client_and_mocks): + """Create skips updating a team that already has the access_group_id.""" + client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + team_record = MagicMock() + team_record.team_id = "team-1" + team_record.access_group_ids = ["ag-new"] # already synced + mock_team_table.find_unique = AsyncMock(return_value=team_record) + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "new-group", "assigned_team_ids": ["team-1"]}, + ) + assert resp.status_code == 201 + mock_team_table.update.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Sync tests: UPDATE +# --------------------------------------------------------------------------- + + +def test_update_access_group_syncs_added_teams(client_and_mocks): + """Update adds access_group_id to newly assigned teams.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_team_ids=["team-existing"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + team_record = MagicMock() + team_record.team_id = "team-new" + team_record.access_group_ids = [] + mock_team_table.find_unique = AsyncMock(return_value=team_record) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_team_ids": ["team-existing", "team-new"]}, + ) + assert resp.status_code == 200 + + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-new"}) + mock_team_table.update.assert_awaited_once() + call_kwargs = mock_team_table.update.call_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-new"} + assert "ag-update" in call_kwargs["data"]["access_group_ids"] + + +def test_update_access_group_syncs_removed_teams(client_and_mocks): + """Update removes access_group_id from de-assigned teams.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_team_ids=["team-keep", "team-remove"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + team_to_remove = MagicMock() + team_to_remove.team_id = "team-remove" + team_to_remove.access_group_ids = ["ag-update"] + mock_team_table.find_unique = AsyncMock(return_value=team_to_remove) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_team_ids": ["team-keep"]}, + ) + assert resp.status_code == 200 + + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-remove"}) + mock_team_table.update.assert_awaited_once() + call_kwargs = mock_team_table.update.call_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-remove"} + assert "ag-update" not in call_kwargs["data"]["access_group_ids"] + + +def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks): + """Update does not sync teams when assigned_team_ids is absent from the payload.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_team_ids=["team-1"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"}) + assert resp.status_code == 200 + + mock_team_table.find_unique.assert_not_awaited() + mock_team_table.update.assert_not_awaited() + + +def test_update_access_group_syncs_added_keys(client_and_mocks): + """Update adds access_group_id to newly assigned keys.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_key_table = mock_prisma.db.litellm_verificationtoken + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_key_ids=["old-token"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + key_record = MagicMock() + key_record.token = "new-token" + key_record.access_group_ids = [] + mock_key_table.find_unique = AsyncMock(return_value=key_record) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_key_ids": ["old-token", "new-token"]}, + ) + assert resp.status_code == 200 + + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "new-token"}) + mock_key_table.update.assert_awaited_once() + call_kwargs = mock_key_table.update.call_args.kwargs + assert call_kwargs["where"] == {"token": "new-token"} + assert "ag-update" in call_kwargs["data"]["access_group_ids"] + + +def test_update_access_group_syncs_removed_keys(client_and_mocks): + """Update removes access_group_id from de-assigned keys.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_key_table = mock_prisma.db.litellm_verificationtoken + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_key_ids=["keep-token", "remove-token"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + key_to_remove = MagicMock() + key_to_remove.token = "remove-token" + key_to_remove.access_group_ids = ["ag-update"] + mock_key_table.find_unique = AsyncMock(return_value=key_to_remove) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_key_ids": ["keep-token"]}, + ) + assert resp.status_code == 200 + + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "remove-token"}) + mock_key_table.update.assert_awaited_once() + call_kwargs = mock_key_table.update.call_args.kwargs + assert call_kwargs["where"] == {"token": "remove-token"} + assert "ag-update" not in call_kwargs["data"]["access_group_ids"] + + +# --------------------------------------------------------------------------- +# Sync tests: DELETE (out-of-sync data handling) +# --------------------------------------------------------------------------- + + +def test_delete_access_group_handles_out_of_sync_assigned_teams(client_and_mocks): + """Delete includes teams from assigned_team_ids even when not found by hasSome query.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + # Access group has assigned_team_ids but the team's access_group_ids is not synced + existing = _make_access_group_record( + access_group_id="ag-to-delete", + assigned_team_ids=["team-out-of-sync"], + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + # hasSome query finds nothing (team's own access_group_ids is out of sync) + mock_team_table.find_many = AsyncMock(return_value=[]) + + out_of_sync_team = MagicMock() + out_of_sync_team.team_id = "team-out-of-sync" + out_of_sync_team.access_group_ids = [] # already clean, no update needed + mock_team_table.find_unique = AsyncMock(return_value=out_of_sync_team) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + # find_unique is called for the out-of-sync team (included via union with assigned_team_ids) + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-out-of-sync"}) + # No update needed since team's access_group_ids doesn't contain "ag-to-delete" + mock_team_table.update.assert_not_awaited() + + +def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks): + """Delete includes keys from assigned_key_ids even when not found by hasSome query.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_key_table = mock_prisma.db.litellm_verificationtoken + + existing = _make_access_group_record( + access_group_id="ag-to-delete", + assigned_key_ids=["token-out-of-sync"], + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + mock_key_table.find_many = AsyncMock(return_value=[]) + + out_of_sync_key = MagicMock() + out_of_sync_key.token = "token-out-of-sync" + out_of_sync_key.access_group_ids = [] + mock_key_table.find_unique = AsyncMock(return_value=out_of_sync_key) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "token-out-of-sync"}) + mock_key_table.update.assert_not_awaited() + + +def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks): + """Update with explicit null for assigned_*_ids clears the list and writes [] to DB.""" + client, _, mock_table, *_ = client_and_mocks + + existing = _make_access_group_record( + access_group_id="ag-update", + assigned_team_ids=["team-1"], + assigned_key_ids=["key-1"], + ) + mock_table.find_unique = AsyncMock(return_value=existing) + + # Sending null for assigned_team_ids and assigned_key_ids + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_team_ids": None, "assigned_key_ids": None}, + ) + assert resp.status_code == 200 + + # Verify the DB update was called with [] (not null) for list fields + update_call_kwargs = mock_table.update.call_args.kwargs + assert update_call_kwargs["data"]["assigned_team_ids"] == [] + assert update_call_kwargs["data"]["assigned_key_ids"] == [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index fb71adc1085..3a2e2f93949 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -45,6 +45,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( check_team_key_model_specific_limits, delete_verification_tokens, generate_key_helper_fn, + key_aliases, list_keys, prepare_key_update_data, reset_key_spend_fn, @@ -6069,6 +6070,121 @@ async def test_build_key_filter_admin_all_member_overlap(): ) +@pytest.mark.asyncio +async def test_build_key_filter_project_id(): + """ + Test that project_id is applied as a global AND condition, narrowing all results + to keys that belong to the specified project. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + user_id = "user-123" + project_id = "proj-abc" + + where = _build_key_filter_conditions( + user_id=user_id, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=None, + member_team_ids=None, + include_created_by_keys=False, + project_id=project_id, + ) + + # Should be wrapped in a top-level AND for the project_id filter + assert "AND" in where + and_parts = where["AND"] + assert len(and_parts) == 2 + + # Second part of AND should be the project_id filter + assert {"project_id": project_id} in and_parts + + +@pytest.mark.asyncio +async def test_build_key_filter_access_group_id(): + """ + Test that access_group_id is applied as a global AND condition using hasSome, + narrowing results to keys whose access_group_ids array contains the given ID. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + user_id = "user-123" + access_group_id = "ag-xyz" + + where = _build_key_filter_conditions( + user_id=user_id, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=None, + member_team_ids=None, + include_created_by_keys=False, + access_group_id=access_group_id, + ) + + # Should be wrapped in a top-level AND for the access_group_id filter + assert "AND" in where + and_parts = where["AND"] + assert len(and_parts) == 2 + + # Second part of AND should use hasSome for the array field + assert {"access_group_ids": {"hasSome": [access_group_id]}} in and_parts + + +@pytest.mark.asyncio +async def test_build_key_filter_project_id_and_access_group_id(): + """ + Test that project_id and access_group_id stack correctly when both are provided. + Both should be applied as AND conditions, narrowing results to keys that match both. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + user_id = "user-123" + project_id = "proj-abc" + access_group_id = "ag-xyz" + + where = _build_key_filter_conditions( + user_id=user_id, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=None, + member_team_ids=None, + include_created_by_keys=False, + project_id=project_id, + access_group_id=access_group_id, + ) + + # After project_id: {"AND": [visibility_where, {"project_id": ...}]} + # After access_group_id: {"AND": [above, {"access_group_ids": ...}]} + assert "AND" in where + outer_and = where["AND"] + assert len(outer_and) == 2 + + # The access_group_ids filter is the outermost AND + access_group_filter = outer_and[1] + assert access_group_filter == {"access_group_ids": {"hasSome": [access_group_id]}} + + # The project_id filter is nested one level in + inner = outer_and[0] + assert "AND" in inner + inner_and = inner["AND"] + assert {"project_id": project_id} in inner_and + + @pytest.mark.asyncio async def test_get_member_team_ids(): """ @@ -6210,3 +6326,97 @@ async def test_generate_key_helper_fn_agent_id(): assert key_data.get("agent_id") == "test-agent-456", ( f"Expected agent_id='test-agent-456' in key_data, got: {key_data.get('agent_id')}" ) + + +@pytest.mark.asyncio +async def test_key_aliases_response_shape(): + """Test that key_aliases returns the correct paginated response shape.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 2}], + [{"key_alias": "alias-alpha"}, {"key_alias": "alias-beta"}], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await key_aliases(page=1, size=50, search=None) + + assert result["aliases"] == ["alias-alpha", "alias-beta"] + assert result["total_count"] == 2 + assert result["current_page"] == 1 + assert result["total_pages"] == 1 + assert result["size"] == 50 + + # Both SQL calls must filter out null/empty aliases + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + aliases_sql = mock_prisma_client.db.query_raw.call_args_list[1].args[0] + assert "key_alias IS NOT NULL" in count_sql + assert "key_alias IS NOT NULL" in aliases_sql + + +@pytest.mark.asyncio +async def test_key_aliases_pagination_skip_take(): + """Test that LIMIT and OFFSET are correctly derived from page and size.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 120}], + [], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await key_aliases(page=3, size=25, search=None) + + assert result["current_page"] == 3 + assert result["size"] == 25 + assert result["total_count"] == 120 + assert result["total_pages"] == 5 # ceil(120 / 25) + + # aliases query params: [UI_SESSION_TOKEN_TEAM_ID, size=25, offset=50] + aliases_call_args = mock_prisma_client.db.query_raw.call_args_list[1].args + assert aliases_call_args[-2] == 25 # LIMIT = size + assert aliases_call_args[-1] == 50 # OFFSET = (3 - 1) * 25 + + +@pytest.mark.asyncio +async def test_key_aliases_search_filter(): + """Test that the search param adds a case-insensitive ILIKE condition.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 0}], + [], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await key_aliases(page=1, size=50, search="my-key") + + count_call = mock_prisma_client.db.query_raw.call_args_list[0] + count_sql = count_call.args[0] + count_params = count_call.args[1:] + + assert "ILIKE" in count_sql + assert "%my-key%" in count_params + + +@pytest.mark.asyncio +async def test_key_aliases_no_search_omits_ilike_filter(): + """Test that without a search term no ILIKE condition is added.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.db.query_raw = AsyncMock( + side_effect=[ + [{"count": 0}], + [], + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await key_aliases(page=1, size=50, search=None) + + count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0] + assert "ILIKE" not in count_sql + + diff --git a/tests/test_litellm/proxy/middleware/test_in_flight_requests_middleware.py b/tests/test_litellm/proxy/middleware/test_in_flight_requests_middleware.py new file mode 100644 index 00000000000..830bca49936 --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_in_flight_requests_middleware.py @@ -0,0 +1,98 @@ +""" +Tests for InFlightRequestsMiddleware. + +Verifies that in_flight_requests is incremented during a request and +decremented after it completes, including on errors. +""" +import asyncio + +import pytest +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.routing import Route +from starlette.testclient import TestClient + +from litellm.proxy.middleware.in_flight_requests_middleware import ( + InFlightRequestsMiddleware, + get_in_flight_requests, +) + + +@pytest.fixture(autouse=True) +def reset_state(): + """Reset class-level state between tests.""" + InFlightRequestsMiddleware._in_flight = 0 + yield + InFlightRequestsMiddleware._in_flight = 0 + + +def _make_app(handler): + from starlette.applications import Starlette + + app = Starlette(routes=[Route("/", handler)]) + app.add_middleware(InFlightRequestsMiddleware) + return app + + +# ── Structure ───────────────────────────────────────────────────────────────── + + +def test_is_not_base_http_middleware(): + """Must be pure ASGI — BaseHTTPMiddleware causes streaming degradation.""" + assert not issubclass(InFlightRequestsMiddleware, BaseHTTPMiddleware) + + +def test_has_asgi_call_protocol(): + assert "__call__" in InFlightRequestsMiddleware.__dict__ + + +# ── Counter behaviour ───────────────────────────────────────────────────────── + + +def test_counter_zero_at_start(): + assert get_in_flight_requests() == 0 + + +def test_counter_increments_inside_handler(): + captured = [] + + async def handler(request: Request) -> Response: + captured.append(InFlightRequestsMiddleware.get_count()) + return JSONResponse({}) + + TestClient(_make_app(handler)).get("/") + assert captured == [1] + + +def test_counter_returns_to_zero_after_request(): + async def handler(request: Request) -> Response: + return JSONResponse({}) + + TestClient(_make_app(handler)).get("/") + assert get_in_flight_requests() == 0 + + +def test_counter_decrements_after_error(): + """Counter must reach 0 even when the handler raises.""" + + async def handler(request: Request) -> Response: + return Response("boom", status_code=500) + + TestClient(_make_app(handler)).get("/") + assert get_in_flight_requests() == 0 + + +def test_non_http_scopes_not_counted(): + """Lifespan / websocket scopes must not touch the counter.""" + + class _InnerApp: + async def __call__(self, scope, receive, send): + pass + + mw = InFlightRequestsMiddleware(_InnerApp()) + + asyncio.get_event_loop().run_until_complete( + mw({"type": "lifespan"}, None, None) # type: ignore[arg-type] + ) + assert get_in_flight_requests() == 0 diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 5f5e2cf1ff8..53c98c8c400 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -357,3 +357,164 @@ def test_public_model_hub_mixed_health_statuses(): assert claude["health_checked_at"] is None app.dependency_overrides.clear() + +# --------------------------------------------------------------------------- +# /public/endpoints +# --------------------------------------------------------------------------- + +import litellm.proxy.public_endpoints.public_endpoints as _pe_module +from litellm.proxy.public_endpoints.public_endpoints import _build_endpoints, _clean_display_name + + +@pytest.fixture(autouse=False) +def reset_endpoints_cache(): + """Reset the module-level cache before and after each cache-related test.""" + original = _pe_module._cached_endpoints + _pe_module._cached_endpoints = None + yield + _pe_module._cached_endpoints = original + + +def _make_client(): + app = FastAPI() + app.include_router(router) + return TestClient(app) + + +def test_get_supported_endpoints_returns_200(reset_endpoints_cache): + response = _make_client().get("/public/endpoints") + assert response.status_code == 200 + + +def test_get_supported_endpoints_response_shape(reset_endpoints_cache): + data = _make_client().get("/public/endpoints").json() + assert "endpoints" in data + assert isinstance(data["endpoints"], list) + assert len(data["endpoints"]) > 0 + + +def test_get_supported_endpoints_item_fields(reset_endpoints_cache): + endpoints = _make_client().get("/public/endpoints").json()["endpoints"] + for item in endpoints: + assert "key" in item + assert "label" in item + assert "endpoint" in item + assert "providers" in item + assert isinstance(item["providers"], list) + + +def test_get_supported_endpoints_provider_fields(reset_endpoints_cache): + endpoints = _make_client().get("/public/endpoints").json()["endpoints"] + for item in endpoints: + for provider in item["providers"]: + assert "slug" in provider + assert "display_name" in provider + + +def test_get_supported_endpoints_paths_start_with_slash(reset_endpoints_cache): + endpoints = _make_client().get("/public/endpoints").json()["endpoints"] + for item in endpoints: + assert item["endpoint"].startswith("/"), f"Expected path starting with /, got: {item['endpoint']}" + + +def test_get_supported_endpoints_chat_completions_present(reset_endpoints_cache): + endpoints = _make_client().get("/public/endpoints").json()["endpoints"] + keys = [item["key"] for item in endpoints] + assert "chat_completions" in keys + + chat = next(item for item in endpoints if item["key"] == "chat_completions") + assert chat["endpoint"] == "/chat/completions" + assert chat["label"] == "Chat Completions" + assert len(chat["providers"]) > 0 + + +def test_get_supported_endpoints_display_names_have_no_slug_suffix(reset_endpoints_cache): + """Provider display_names must not contain the raw `` (`slug`) `` suffix.""" + import re + suffix_re = re.compile(r"\(`[^`]+`\)") + endpoints = _make_client().get("/public/endpoints").json()["endpoints"] + for item in endpoints: + for provider in item["providers"]: + assert not suffix_re.search(provider["display_name"]), ( + f"display_name still contains slug suffix: {provider['display_name']!r}" + ) + + +def test_get_supported_endpoints_is_cached(reset_endpoints_cache): + """`_load_endpoints` is called only once; subsequent requests use the cache.""" + client = _make_client() + with patch( + "litellm.proxy.public_endpoints.public_endpoints._load_endpoints", + wraps=_pe_module._load_endpoints, + ) as mock_load: + client.get("/public/endpoints") + client.get("/public/endpoints") + client.get("/public/endpoints") + + mock_load.assert_called_once() + + +# --------------------------------------------------------------------------- +# _build_endpoints unit tests (transformation logic) +# --------------------------------------------------------------------------- + +_MINIMAL_RAW = { + "providers": { + "openai": { + "display_name": "OpenAI (`openai`)", + "url": "https://example.com", + "endpoints": {"chat_completions": True, "embeddings": True, "images": False}, + }, + "anthropic": { + "display_name": "Anthropic (`anthropic`)", + "url": "https://example.com", + "endpoints": {"chat_completions": True, "embeddings": False, "images": False}, + }, + } +} + + +def test_build_endpoints_known_key_uses_metadata(): + result = _build_endpoints(_MINIMAL_RAW) + chat = next(e for e in result if e["key"] == "chat_completions") + assert chat["label"] == "Chat Completions" + assert chat["endpoint"] == "/chat/completions" + + +def test_build_endpoints_only_includes_supporting_providers(): + result = _build_endpoints(_MINIMAL_RAW) + embeddings = next(e for e in result if e["key"] == "embeddings") + slugs = [p["slug"] for p in embeddings["providers"]] + assert slugs == ["openai"] + + +def test_build_endpoints_unknown_key_derives_label_and_path(): + raw = { + "providers": { + "someprovider": { + "display_name": "Some Provider (`someprovider`)", + "endpoints": {"my_custom_endpoint": True}, + } + } + } + result = _build_endpoints(raw) + item = result[0] + assert item["key"] == "my_custom_endpoint" + assert item["label"] == "My Custom Endpoint" + assert item["endpoint"].startswith("/") + + +def test_build_endpoints_empty_providers_returns_empty(): + result = _build_endpoints({"providers": {}}) + assert result == [] + + +def test_clean_display_name_strips_suffix(): + assert _clean_display_name("OpenAI (`openai`)") == "OpenAI" + assert _clean_display_name("AI/ML API (`aiml`)") == "AI/ML API" + assert _clean_display_name("A2A (Agent-to-Agent) (`a2a`)") == "A2A (Agent-to-Agent)" + + +def test_clean_display_name_passthrough_when_no_suffix(): + assert _clean_display_name("OpenAI") == "OpenAI" + assert _clean_display_name("") == "" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index bf794478f10..9b905d24fd1 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -13,6 +13,7 @@ from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ProxyConfig, + _add_dd_apm_tags_for_litellm_call_id, _extract_error_from_sse_chunk, _get_cost_breakdown_from_logging_obj, _override_openai_response_model, @@ -79,6 +80,20 @@ class TestProxyBaseLLMRequestProcessing: pytest.fail("litellm_call_id is not a valid UUID") assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"] + def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch): + mock_set_active_span_tag = MagicMock(return_value=True) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "set_active_span_tag", + mock_set_active_span_tag, + ) + + _add_dd_apm_tags_for_litellm_call_id("test-call-id") + + mock_set_active_span_tag.assert_called_once_with( + "litellm.call_id", "test-call-id" + ) + @pytest.mark.asyncio async def test_should_apply_hierarchical_router_settings_as_override( self, monkeypatch diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py new file mode 100644 index 00000000000..e26f7fb9f20 --- /dev/null +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -0,0 +1,75 @@ +import pytest +from litellm.proxy.health_check import _update_litellm_params_for_health_check +from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers +from unittest.mock import AsyncMock, patch, MagicMock + + +@pytest.mark.asyncio +async def test_update_litellm_params_max_tokens_default(): + """ + Test that max_tokens defaults to 1 for non-wildcard models. + """ + model_info = {} + litellm_params = {"model": "gpt-4"} + + updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated_params["max_tokens"] == 1 + + +@pytest.mark.asyncio +async def test_update_litellm_params_max_tokens_custom(): + """ + Test that max_tokens respects health_check_max_tokens from model_info. + """ + model_info = {"health_check_max_tokens": 5} + litellm_params = {"model": "gpt-4"} + + updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated_params["max_tokens"] == 5 + + +@pytest.mark.asyncio +async def test_update_litellm_params_max_tokens_wildcard(): + """ + Test that max_tokens does NOT default to 1 for wildcard models. + """ + model_info = {} + litellm_params = {"model": "openai/*"} + + updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) + + # Should not be set to 1 + assert "max_tokens" not in updated_params or updated_params["max_tokens"] != 1 + + +@pytest.mark.asyncio +async def test_ahealth_check_wildcard_models_respects_max_tokens(): + """ + Test that ahealth_check_wildcard_models respects max_tokens if passed, + otherwise defaults to 10. + """ + with patch( + "litellm.litellm_core_utils.llm_request_utils.pick_cheapest_chat_models_from_llm_provider", + return_value=["gpt-4o-mini"], + ), patch("litellm.acompletion", new_callable=AsyncMock): + # Test Case 1: No max_tokens passed, should default to 10 + model_params = {} + await HealthCheckHelpers.ahealth_check_wildcard_models( + model="openai/*", + custom_llm_provider="openai", + model_params=model_params, + litellm_logging_obj=MagicMock(), + ) + assert model_params["max_tokens"] == 10 + + # Test Case 2: Custom health_check_max_tokens passed via model_params, should be respected + model_params = {"max_tokens": 3} + await HealthCheckHelpers.ahealth_check_wildcard_models( + model="openai/*", + custom_llm_provider="openai", + model_params=model_params, + litellm_logging_obj=MagicMock(), + ) + assert model_params["max_tokens"] == 3 diff --git a/tests/test_litellm/proxy/test_prometheus_cleanup.py b/tests/test_litellm/proxy/test_prometheus_cleanup.py new file mode 100644 index 00000000000..b3d785f1133 --- /dev/null +++ b/tests/test_litellm/proxy/test_prometheus_cleanup.py @@ -0,0 +1,113 @@ +""" +Tests for litellm.proxy.prometheus_cleanup.wipe_directory and +ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir. +""" + +from __future__ import annotations + +import os +from unittest.mock import patch + +import pytest + +from litellm.proxy.prometheus_cleanup import mark_worker_exit, wipe_directory +from litellm.proxy.proxy_cli import ProxyInitializationHelpers + + +class TestWipeDirectory: + def test_deletes_all_db_files(self, tmp_path): + (tmp_path / "counter_1234.db").touch() + (tmp_path / "histogram_5678.db").touch() + (tmp_path / "gauge_livesum_9999.db").touch() + wipe_directory(str(tmp_path)) + assert not list(tmp_path.glob("*.db")) + + +class TestMarkWorkerExit: + def test_calls_mark_process_dead_when_env_set(self, tmp_path): + with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}): + with patch( + "prometheus_client.multiprocess.mark_process_dead" + ) as mock_mark: + mark_worker_exit(12345) + mock_mark.assert_called_once_with(12345) + + def test_noop_when_env_not_set(self): + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + with patch( + "prometheus_client.multiprocess.mark_process_dead" + ) as mock_mark: + mark_worker_exit(12345) + mock_mark.assert_not_called() + + def test_exception_is_caught_and_logged(self, tmp_path): + with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}): + with patch( + "prometheus_client.multiprocess.mark_process_dead", + side_effect=FileNotFoundError("gone"), + ) as mock_mark: + # Should not raise + mark_worker_exit(99) + mock_mark.assert_called_once_with(99) + + +class TestMaybeSetupPrometheusMultiprocDir: + def test_respects_existing_env_var(self, tmp_path): + """When PROMETHEUS_MULTIPROC_DIR is already set, don't override it.""" + custom_dir = str(tmp_path / "custom_prom") + litellm_settings = {"callbacks": ["prometheus"]} + + with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": custom_dir}): + ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=4, + litellm_settings=litellm_settings, + ) + + assert os.environ["PROMETHEUS_MULTIPROC_DIR"] == custom_dir + assert os.path.isdir(custom_dir) + + @pytest.mark.parametrize( + "num_workers, litellm_settings", + [ + (1, {"callbacks": ["prometheus"]}), + (4, {"callbacks": ["langfuse"]}), + (4, None), + ], + ) + def test_noop_when_setup_not_needed(self, num_workers, litellm_settings): + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + os.environ.pop("prometheus_multiproc_dir", None) + + ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=num_workers, + litellm_settings=litellm_settings, + ) + + assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") is None + + @pytest.mark.parametrize( + "litellm_settings", + [ + {"callbacks": ["prometheus"]}, + {"success_callback": ["prometheus"]}, + ], + ) + def test_auto_creates_dir_when_prometheus_configured(self, litellm_settings): + """When multiple workers + prometheus callback, auto-creates temp dir.""" + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + os.environ.pop("prometheus_multiproc_dir", None) + + ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=4, + litellm_settings=litellm_settings, + ) + + result_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + assert result_dir is not None + assert os.path.isdir(result_dir) + + # Cleanup + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) diff --git a/tests/test_litellm/proxy/test_proxy_types.py b/tests/test_litellm/proxy/test_proxy_types.py index 0e47134478b..ae2b7bbf24c 100644 --- a/tests/test_litellm/proxy/test_proxy_types.py +++ b/tests/test_litellm/proxy/test_proxy_types.py @@ -45,3 +45,27 @@ def test_audit_log_masking(): json_before_value = json.loads(audit_log.before_value) assert json_before_value["token"] == "1q2132r222" assert json_before_value["key"] == "sk-1*****7890" + + +def test_internal_jobs_user_has_proxy_admin_role(): + """ + Test that the internal jobs system user has PROXY_ADMIN role. + + This is critical for key rotation to work properly. The system user needs + PROXY_ADMIN role to bypass team permission checks in + TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint() + + Regression test for: https://github.com/BerriAI/litellm/pull/21896 + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + # Get the system user used for internal jobs like key rotation + system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth() + + # Verify the system user has PROXY_ADMIN role + assert system_user.user_role == LitellmUserRoles.PROXY_ADMIN + + # Verify other expected properties + assert system_user.user_id == "system" + assert system_user.team_id == "system" + assert system_user.team_alias == "system" diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 1ffbb83caef..c1fa3ad0c43 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -151,28 +151,16 @@ async def test_should_delete_spend_logs(): @pytest.mark.asyncio async def test_cleanup_old_spend_logs_batch_deletion(): - from types import SimpleNamespace - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import AsyncMock, MagicMock # Setup Prisma client mock_prisma_client = MagicMock() mock_db = MagicMock() - # Mock spendlogs table - mock_spendlogs = MagicMock() - mock_spendlogs.find_many = AsyncMock() - mock_spendlogs.delete_many = AsyncMock() - - # Create 1500 mocked logs with .request_id - mock_logs = [SimpleNamespace(request_id=f"req_{i}") for i in range(1500)] - mock_spendlogs.find_many.side_effect = [ - mock_logs[:1000], # Batch 1 - mock_logs[1000:], # Batch 2 - [], # Done - ] + # Mock execute_raw to return deleted counts + mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0]) # Wire up mocks - mock_db.litellm_spendlogs = mock_spendlogs mock_prisma_client.db = mock_db # Mock Redis cache and pod_lock_manager @@ -189,15 +177,13 @@ async def test_cleanup_old_spend_logs_batch_deletion(): assert cleaner._should_delete_spend_logs() is True await cleaner.cleanup_old_spend_logs(mock_prisma_client) - # Validate batching and deletion - assert mock_spendlogs.find_many.call_count == 3 - assert mock_spendlogs.delete_many.call_count == 2 - mock_spendlogs.delete_many.assert_any_call( - where={"request_id": {"in": [f"req_{i}" for i in range(1000)]}} - ) - mock_spendlogs.delete_many.assert_any_call( - where={"request_id": {"in": [f"req_{i}" for i in range(1000, 1500)]}} - ) + # Validate batching and deletion via raw SQL + assert mock_db.execute_raw.call_count == 3 + + # Check the first call argument + call_args_sql = mock_db.execute_raw.call_args_list[0][0][0] + assert 'DELETE FROM "LiteLLM_SpendLogs"' in call_args_sql + assert 'WHERE "request_id" IN' in call_args_sql @pytest.mark.asyncio @@ -208,10 +194,7 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): # Setup Prisma client mock_prisma_client = MagicMock() mock_db = MagicMock() - mock_spendlogs = MagicMock() - mock_spendlogs.find_many = AsyncMock(return_value=[]) - mock_spendlogs.delete_many = AsyncMock() - mock_db.litellm_spendlogs = mock_spendlogs + mock_db.execute_raw = AsyncMock(return_value=0) mock_prisma_client.db = mock_db # Mock Redis cache and pod_lock_manager @@ -229,7 +212,7 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): await cleaner.cleanup_old_spend_logs(mock_prisma_client) # Verify the cutoff date is correct - cutoff_date = mock_spendlogs.find_many.call_args[1]["where"]["startTime"]["lt"] + cutoff_date = mock_db.execute_raw.call_args[0][1] expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400) assert ( abs((cutoff_date - expected_cutoff).total_seconds()) < 1 @@ -242,14 +225,12 @@ async def test_cleanup_old_spend_logs_no_retention_period(): Test that no logs are deleted when no retention period is set """ mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_spendlogs.find_many = AsyncMock() - mock_prisma_client.db.litellm_spendlogs.delete = AsyncMock() + mock_prisma_client.db.execute_raw = AsyncMock() cleaner = SpendLogCleanup(general_settings={}) # no retention await cleaner.cleanup_old_spend_logs(mock_prisma_client) - mock_prisma_client.db.litellm_spendlogs.find_many.assert_not_called() - mock_prisma_client.db.litellm_spendlogs.delete.assert_not_called() + mock_prisma_client.db.execute_raw.assert_not_called() def test_cleanup_batch_size_env_var(monkeypatch): diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index a238531d2e0..3ba41705733 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -90,6 +90,72 @@ async def test_acompletion_with_mcp_without_auto_execution_calls_model(monkeypat assert captured_secret_fields["value"] == {"api_key": "value"} +@pytest.mark.asyncio +async def test_acompletion_with_mcp_passes_mcp_server_auth_headers_to_process_tools( + monkeypatch, +): + """ + Test that MCP auth headers extracted from secret_fields (e.g. x-mcp-linear_config-authorization) + are passed to _process_mcp_tools_without_openai_transform for dynamic auth when fetching tools. + """ + tools = [{"type": "mcp", "server_url": "litellm_proxy"}] + mock_acompletion = AsyncMock(return_value="ok") + + captured_process_kwargs = {} + + async def mock_process(**kwargs): + captured_process_kwargs.update(kwargs) + return ([], {}) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_use_litellm_mcp_gateway", + staticmethod(lambda t: True), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_parse_mcp_tools", + staticmethod(lambda t: (t, [])), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + mock_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda *_, **__: ["openai-tool"]), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_auto_execute_tools", + staticmethod(lambda **_: False), + ) + + # secret_fields with raw_headers containing MCP auth - extract_mcp_headers_from_request + # will parse these and pass to _process_mcp_tools_without_openai_transform + secret_fields = { + "raw_headers": { + "x-mcp-linear_config-authorization": "Bearer linear-token", + }, + } + + with patch("litellm.acompletion", mock_acompletion): + await acompletion_with_mcp( + model="test-model", + messages=[], + tools=tools, + secret_fields=secret_fields, + ) + + assert "mcp_server_auth_headers" in captured_process_kwargs + mcp_server_auth_headers = captured_process_kwargs["mcp_server_auth_headers"] + assert mcp_server_auth_headers is not None + assert "linear_config" in mcp_server_auth_headers + assert mcp_server_auth_headers["linear_config"]["Authorization"] == "Bearer linear-token" + + @pytest.mark.asyncio async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch): from litellm.utils import CustomStreamWrapper diff --git a/tests/test_litellm/test_exception_header_preservation.py b/tests/test_litellm/test_exception_header_preservation.py index d3e33fa13b3..ec52d9fb746 100644 --- a/tests/test_litellm/test_exception_header_preservation.py +++ b/tests/test_litellm/test_exception_header_preservation.py @@ -16,6 +16,8 @@ from litellm.exceptions import ( ContentPolicyViolationError, ContextWindowExceededError, ImageFetchError, + MidStreamFallbackError, + RateLimitError, ) @@ -210,6 +212,46 @@ class TestExceptionAttributes: assert error.num_retries == 1 assert error.status_code == 400 + def test_midstream_fallback_error_status_code_propagation(self): + """ + MidStreamFallbackError should preserve the original status code and keep + message/request/response fields consistent after super().__init__(). + """ + original_req = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + original_resp = httpx.Response(status_code=429, request=original_req) + + rate_limit_error = RateLimitError( + message="Rate limit exceeded", + llm_provider="openai", + model="gpt-4o-mini", + response=original_resp, + ) + + midstream_error = MidStreamFallbackError( + message="stream broke", + model="gpt-4o-mini", + llm_provider="openai", + original_exception=rate_limit_error, + ) + + assert midstream_error.status_code == 429 + assert midstream_error.response.status_code == 429 + assert str(midstream_error.response.request.url) == "https://openai.com/v1/" + assert midstream_error.message == "litellm.MidStreamFallbackError: stream broke" + assert midstream_error.args == ("litellm.MidStreamFallbackError: stream broke",) + + # With no original exception, should default to 503. + midstream_fallback = MidStreamFallbackError( + message="stream broke without original", + model="gpt-4o-mini", + llm_provider="openai", + original_exception=None, + ) + + assert midstream_fallback.status_code == 503 + assert midstream_fallback.response.status_code == 503 + assert str(midstream_fallback.response.request.url) == "https://openai.com/v1/" + class TestProxyHeaderExtraction: """Test that proxy correctly extracts headers from exceptions.""" diff --git a/tests/test_litellm/test_project_tags_pydantic.py b/tests/test_litellm/test_project_tags_pydantic.py new file mode 100644 index 00000000000..b3f58df2325 --- /dev/null +++ b/tests/test_litellm/test_project_tags_pydantic.py @@ -0,0 +1,31 @@ +import pytest +from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest + + +def test_new_project_request_tags(): + # Test tags correctly stay top level initially + req = NewProjectRequest( + project_id="test_proj", team_id="team_1", tags=["tag1", "tag2"] + ) + + # tags should be top level initially + assert req.tags == ["tag1", "tag2"] + + +def test_update_project_request_tags(): + # Test tags correctly stay top level initially + req = UpdateProjectRequest(project_id="test_proj", tags=["new_tag"]) + + assert req.tags == ["new_tag"] + + +def test_new_project_request_invalid_tags_type(): + # tags must be a list — a string should raise a ValidationError + with pytest.raises(Exception): + NewProjectRequest(project_id="test_proj", team_id="team_1", tags="not-a-list") + + +def test_update_project_request_invalid_tags_type(): + # tags must be a list — a string should raise a ValidationError + with pytest.raises(Exception): + UpdateProjectRequest(project_id="test_proj", tags="not-a-list") diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py index 87cc9586665..054fe505764 100644 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ b/tests/test_litellm/types/llms/test_types_llms_openai.py @@ -169,3 +169,97 @@ class TestResponsesAPIResponseOutputText: ) assert response.output_text == "" + + +class TestAssistantMessageImageUrlContent: + """ + Regression tests for image_url blocks in assistant message content. + + Bug: ChatCompletionAssistantMessage.content did not include + ChatCompletionImageObject in its union, so Pydantic v2 silently dropped + image_url blocks (content → []) when serialising via AllMessageValues. + This affects users who store conversation history as JSON (e.g. in a DB) + and read it back typed as list[AllMessageValues]. + """ + + ASSISTANT_MESSAGE_WITH_IMAGE = { + "role": "assistant", + "content": [ + {"type": "text", "text": "Here is the image you requested:"}, + { + "type": "image_url", + "image_url": { + "url": ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA" + "DUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + ) + }, + }, + ], + } + + def test_assistant_message_image_url_preserved_single(self): + """ + TypeAdapter(ChatCompletionAssistantMessage): image_url block must survive + validate_python → dump_python without being dropped or raising an error. + """ + from typing import List + + from pydantic import TypeAdapter + + from litellm.types.llms.openai import ChatCompletionAssistantMessage + + adapter = TypeAdapter(ChatCompletionAssistantMessage) + validated = adapter.validate_python(self.ASSISTANT_MESSAGE_WITH_IMAGE) + dumped = adapter.dump_python(validated) + + raw_content = dumped.get("content") + # Pydantic may return a lazy SerializationIterator for Iterable fields; + # convert to list to consume it — this must not raise ValidationError. + content_blocks = list(raw_content) if raw_content is not None else [] + + assert len(content_blocks) == 2, ( + f"Expected 2 content blocks (text + image_url), got {len(content_blocks)}: {content_blocks}" + ) + types = [b.get("type") for b in content_blocks if isinstance(b, dict)] + assert "image_url" in types, f"image_url block was silently dropped; blocks: {content_blocks}" + + def test_assistant_message_image_url_preserved_in_all_message_values(self): + """ + TypeAdapter(List[AllMessageValues]) DB round-trip: image_url blocks in an + assistant message must not be silently dropped during dump_python(mode='json'). + + This is the primary failing path: conversation history stored as JSON in a + database and read back typed as list[AllMessageValues]. + """ + from typing import List + + from pydantic import TypeAdapter + + from litellm.types.llms.openai import AllMessageValues + + conversation = [ + { + "role": "user", + "content": "Generate an image of a banana wearing a LiteLLM costume", + }, + self.ASSISTANT_MESSAGE_WITH_IMAGE, + ] + + adapter = TypeAdapter(List[AllMessageValues]) + validated = adapter.validate_python(conversation) + dumped = adapter.dump_python(validated, mode="json") + + assistant = next((m for m in dumped if m.get("role") == "assistant"), None) + assert assistant is not None, "Assistant message missing after serialisation" + + content = assistant.get("content", []) + assert isinstance(content, list), f"content should be a list, got {type(content)}" + assert len(content) == 2, ( + f"Expected 2 content blocks (text + image_url), got {len(content)}: {content}" + ) + types = [b.get("type") for b in content if isinstance(b, dict)] + assert "image_url" in types, ( + f"image_url block was silently dropped during AllMessageValues serialisation; blocks: {content}" + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts new file mode 100644 index 00000000000..b382b1f2ad3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.test.ts @@ -0,0 +1,177 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useInfiniteKeyAliases } from "./useKeyAliases"; +import type { PaginatedKeyAliasResponse } from "@/components/networking"; + +// Mock networking module +vi.mock("@/components/networking", () => ({ + keyAliasesCall: vi.fn(), +})); + +// Mock useAuthorized hook +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock console methods to avoid noise +vi.spyOn(console, "log").mockImplementation(() => {}); +vi.spyOn(console, "error").mockImplementation(() => {}); + +import { keyAliasesCall } from "@/components/networking"; + +const mockKeyAliasesCall = vi.mocked(keyAliasesCall); + +const mockPage1: PaginatedKeyAliasResponse = { + aliases: ["alias-1", "alias-2"], + total_count: 3, + current_page: 1, + total_pages: 2, + size: 2, +}; + +const mockPage2: PaginatedKeyAliasResponse = { + aliases: ["alias-3"], + total_count: 3, + current_page: 2, + total_pages: 2, + size: 2, +}; + +const createWrapper = () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + }, + }); + return ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); +}; + +describe("useInfiniteKeyAliases", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" }); + mockKeyAliasesCall.mockResolvedValue(mockPage1); + }); + + it("should fetch the first page of key aliases", async () => { + const wrapper = createWrapper(); + const { result } = renderHook(() => useInfiniteKeyAliases(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, undefined); + expect(result.current.data?.pages[0]).toEqual(mockPage1); + }); + + it("should pass custom size parameter", async () => { + const wrapper = createWrapper(); + renderHook(() => useInfiniteKeyAliases(25), { wrapper }); + + await waitFor(() => { + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 25, undefined); + }); + }); + + it("should pass search parameter when provided", async () => { + const wrapper = createWrapper(); + renderHook(() => useInfiniteKeyAliases(50, "my-alias"), { wrapper }); + + await waitFor(() => { + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, "my-alias"); + }); + }); + + it("should not fetch when accessToken is not available", () => { + mockUseAuthorized.mockReturnValue({ accessToken: null }); + const wrapper = createWrapper(); + const { result } = renderHook(() => useInfiniteKeyAliases(), { wrapper }); + + expect(result.current.isFetching).toBe(false); + expect(mockKeyAliasesCall).not.toHaveBeenCalled(); + }); + + it("should expose hasNextPage when more pages are available", async () => { + const wrapper = createWrapper(); + const { result } = renderHook(() => useInfiniteKeyAliases(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.hasNextPage).toBe(true); + }); + + it("should return hasNextPage false when on last page", async () => { + const singlePage: PaginatedKeyAliasResponse = { + aliases: ["alias-1"], + total_count: 1, + current_page: 1, + total_pages: 1, + size: 50, + }; + mockKeyAliasesCall.mockResolvedValue(singlePage); + + const wrapper = createWrapper(); + const { result } = renderHook(() => useInfiniteKeyAliases(), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.hasNextPage).toBe(false); + }); + + it("should fetch the next page when fetchNextPage is called", async () => { + mockKeyAliasesCall + .mockResolvedValueOnce(mockPage1) + .mockResolvedValueOnce(mockPage2); + + const wrapper = createWrapper(); + const { result } = renderHook(() => useInfiniteKeyAliases(2), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + result.current.fetchNextPage(); + + await waitFor(() => { + expect(result.current.data?.pages).toHaveLength(2); + }); + + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 2, 2, undefined); + expect(result.current.data?.pages[1]).toEqual(mockPage2); + }); + + it("should include search in query key so search changes refetch from page 1", async () => { + const wrapper = createWrapper(); + const { result, rerender } = renderHook( + ({ search }: { search?: string }) => useInfiniteKeyAliases(50, search), + { wrapper, initialProps: { search: undefined } }, + ); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + mockKeyAliasesCall.mockResolvedValue({ + aliases: ["search-result"], + total_count: 1, + current_page: 1, + total_pages: 1, + size: 50, + }); + + rerender({ search: "search-result" }); + + await waitFor(() => { + expect(mockKeyAliasesCall).toHaveBeenCalledWith("test-token", 1, 50, "search-result"); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts new file mode 100644 index 00000000000..f67b15f3a9f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyAliases.ts @@ -0,0 +1,37 @@ +import { useInfiniteQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { keyAliasesCall, type PaginatedKeyAliasResponse } from "@/components/networking"; +import useAuthorized from "../useAuthorized"; + +const infiniteKeyAliasKeys = createQueryKeys("infiniteKeyAliases"); + +export const useInfiniteKeyAliases = ( + size: number = 50, + search?: string, +) => { + const { accessToken } = useAuthorized(); + return useInfiniteQuery({ + queryKey: infiniteKeyAliasKeys.list({ + filters: { + size, + ...(search && { search }), + }, + }), + queryFn: async ({ pageParam }) => { + return await keyAliasesCall( + accessToken!, + pageParam as number, + size, + search, + ); + }, + initialPageParam: 1, + getNextPageParam: (lastPage) => { + if (lastPage.current_page < lastPage.total_pages) { + return lastPage.current_page + 1; + } + return undefined; + }, + enabled: Boolean(accessToken), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts index 1643412d1e9..80cb69495da 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts @@ -50,6 +50,7 @@ const mockKeys: KeyResponse[] = [ config: {}, user_id: "user-1", team_id: null, + project_id: null, max_parallel_requests: 10, metadata: {}, tpm_limit: 1000, @@ -105,6 +106,7 @@ const mockKeys: KeyResponse[] = [ config: {}, user_id: "user-2", team_id: "team-1", + project_id: "project-1", max_parallel_requests: 5, metadata: {}, tpm_limit: 500, @@ -396,6 +398,76 @@ describe("useKeys", () => { }, ); }); + + it("should pass projectID filter to the API", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ + keys: [mockKeys[1]], + total_count: 1, + current_page: 1, + total_pages: 1, + }), + }); + + const { result } = renderHook( + () => useKeys(1, 10, { projectID: "project-1" }), + { wrapper }, + ); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + const callUrl = mockFetch.mock.calls[0][0]; + expect(callUrl).toContain("project_id=project-1"); + expect(result.current.data?.keys).toHaveLength(1); + expect(result.current.data?.keys[0].project_id).toBe("project-1"); + }); + + it("should pass both projectID and teamID filters to the API", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => ({ + keys: [mockKeys[1]], + total_count: 1, + current_page: 1, + total_pages: 1, + }), + }); + + const { result } = renderHook( + () => useKeys(1, 10, { projectID: "project-1", teamID: "team-1" }), + { wrapper }, + ); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + const callUrl = mockFetch.mock.calls[0][0]; + expect(callUrl).toContain("project_id=project-1"); + expect(callUrl).toContain("team_id=team-1"); + }); + + it("should not include project_id param when projectID is null", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockKeysResponse, + }); + + const { result } = renderHook( + () => useKeys(1, 10, { projectID: null }), + { wrapper }, + ); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + const callUrl = mockFetch.mock.calls[0][0]; + expect(callUrl).not.toContain("project_id"); + }); }); describe("useDeletedKeys", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index cf477a2e556..fbe5eccb75a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -33,6 +33,7 @@ export interface DeletedKeysResponse { export interface KeyListCallOptions { organizationID?: string | null; teamID?: string | null; + projectID?: string | null; selectedKeyAlias?: string | null; userID?: string | null; keyHash?: string | null; @@ -57,6 +58,7 @@ const keyListCall = async ( const params = new URLSearchParams( Object.entries({ team_id: options.teamID, + project_id: options.projectID, organization_id: options.organizationID, key_alias: options.selectedKeyAlias, key_hash: options.keyHash, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts new file mode 100644 index 00000000000..3943f23794e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts @@ -0,0 +1,70 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { + getProxyBaseUrl, + getGlobalLitellmHeaderName, + deriveErrorMessage, + handleError, +} from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { ProjectResponse, projectKeys } from "./useProjects"; + +// ── Types ──────────────────────────────────────────────────────────────────── + +export interface ProjectCreateParams { + project_alias?: string; + description?: string; + team_id: string; + models?: string[]; + max_budget?: number; + blocked?: boolean; + metadata?: Record; + model_rpm_limit?: Record; + model_tpm_limit?: Record; +} + +// ── Fetch function ─────────────────────────────────────────────────────────── + +const createProject = async ( + accessToken: string, + params: ProjectCreateParams, +): Promise => { + const baseUrl = getProxyBaseUrl(); + const url = `${baseUrl}/project/new`; + + const response = await fetch(url, { + method: "POST", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(params), + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + return response.json(); +}; + +// ── Hook ───────────────────────────────────────────────────────────────────── + +export const useCreateProject = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (params) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return createProject(accessToken, params); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: projectKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.ts new file mode 100644 index 00000000000..1d35ac1bf70 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjectDetails.ts @@ -0,0 +1,63 @@ +import { useQuery, useQueryClient } from "@tanstack/react-query"; +import { + getProxyBaseUrl, + getGlobalLitellmHeaderName, + deriveErrorMessage, + handleError, +} from "@/components/networking"; +import { all_admin_roles } from "@/utils/roles"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { ProjectResponse, projectKeys } from "./useProjects"; + +// ── Fetch function ─────────────────────────────────────────────────────────── + +const fetchProjectDetails = async ( + accessToken: string, + projectId: string, +): Promise => { + const baseUrl = getProxyBaseUrl(); + const url = `${baseUrl}/project/info?project_id=${encodeURIComponent(projectId)}`; + + const response = await fetch(url, { + method: "GET", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + return response.json(); +}; + +// ── Hook ───────────────────────────────────────────────────────────────────── + +export const useProjectDetails = (projectId?: string) => { + const { accessToken, userRole } = useAuthorized(); + const queryClient = useQueryClient(); + + return useQuery({ + queryKey: projectKeys.detail(projectId!), + queryFn: async () => fetchProjectDetails(accessToken!, projectId!), + enabled: + Boolean(accessToken && projectId) && + all_admin_roles.includes(userRole || ""), + + // Seed from the list cache when available + initialData: () => { + if (!projectId) return undefined; + + const projects = queryClient.getQueryData( + projectKeys.list({}), + ); + + return projects?.find((p) => p.project_id === projectId); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts new file mode 100644 index 00000000000..85c8b25645c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts @@ -0,0 +1,87 @@ +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { + getProxyBaseUrl, + getGlobalLitellmHeaderName, + deriveErrorMessage, + handleError, +} from "@/components/networking"; +import { all_admin_roles } from "@/utils/roles"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +// ── Types ──────────────────────────────────────────────────────────────────── + +export interface ProjectBudget { + budget_id: string; + max_budget: number | null; + soft_budget: number | null; + max_parallel_requests: number | null; + tpm_limit: number | null; + rpm_limit: number | null; + model_max_budget: Record | null; + budget_duration: string | null; +} + +export interface ProjectResponse { + project_id: string; + project_alias: string | null; + description: string | null; + team_id: string | null; + budget_id: string | null; + metadata: Record | null; + models: string[]; + spend: number; + model_spend: Record | null; + model_rpm_limit: Record | null; + model_tpm_limit: Record | null; + blocked: boolean; + object_permission_id: string | null; + created_at: string; + created_by: string; + updated_at: string; + updated_by: string; + litellm_budget_table: ProjectBudget | null; +} + +// ── Query keys (shared across project hooks) ───────────────────────────────── + +export const projectKeys = createQueryKeys("projects"); + +// ── Fetch function ─────────────────────────────────────────────────────────── + +const fetchProjects = async ( + accessToken: string, +): Promise => { + const baseUrl = getProxyBaseUrl(); + const url = `${baseUrl}/project/list`; + + const response = await fetch(url, { + method: "GET", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + return response.json(); +}; + +// ── Hook ───────────────────────────────────────────────────────────────────── + +export const useProjects = () => { + const { accessToken, userRole } = useAuthorized(); + + return useQuery({ + queryKey: projectKeys.list({}), + queryFn: async () => fetchProjects(accessToken!), + enabled: + Boolean(accessToken) && all_admin_roles.includes(userRole || ""), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts new file mode 100644 index 00000000000..e6cd3071f5f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts @@ -0,0 +1,75 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { + getProxyBaseUrl, + getGlobalLitellmHeaderName, + deriveErrorMessage, + handleError, +} from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { ProjectResponse, projectKeys } from "./useProjects"; + +// ── Types ──────────────────────────────────────────────────────────────────── + +export interface ProjectUpdateParams { + project_alias?: string; + description?: string; + team_id?: string; + models?: string[]; + max_budget?: number; + blocked?: boolean; + metadata?: Record; + model_rpm_limit?: Record; + model_tpm_limit?: Record; +} + +// ── Fetch function ─────────────────────────────────────────────────────────── + +const updateProject = async ( + accessToken: string, + projectId: string, + params: ProjectUpdateParams, +): Promise => { + const baseUrl = getProxyBaseUrl(); + const url = `${baseUrl}/project/update`; + + const response = await fetch(url, { + method: "POST", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ project_id: projectId, ...params }), + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + return response.json(); +}; + +// ── Hook ───────────────────────────────────────────────────────────────────── + +export const useUpdateProject = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation< + ProjectResponse, + Error, + { projectId: string; params: ProjectUpdateParams } + >({ + mutationFn: async ({ projectId, params }) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return updateProject(accessToken, projectId, params); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: projectKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 258c2ccb0e0..0b2f467e8f8 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -37,6 +37,7 @@ import UIThemeSettings from "@/components/ui_theme_settings"; import Usage from "@/components/usage"; import UserDashboard from "@/components/user_dashboard"; import { AccessGroupsPage } from "@/components/AccessGroups/AccessGroupsPage"; +import { ProjectsPage } from "@/components/Projects/ProjectsPage"; import VectorStoreManagement from "@/components/vector_store_management"; import ToolPolicies from "@/components/ToolPolicies"; import SpendLogsTable from "@/components/view_logs"; @@ -547,6 +548,8 @@ function CreateKeyPageContent() { ) : page == "access-groups" ? ( + ) : page == "projects" ? ( + ) : page == "vector-stores" ? ( ) : page == "tool-policies" ? ( diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.test.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.test.tsx new file mode 100644 index 00000000000..d61f0987ac2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.test.tsx @@ -0,0 +1,148 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../tests/test-utils"; +import AddMarginForm from "./add_margin_form"; +import { MarginConfig } from "./types"; + +vi.mock("../provider_info_helpers", () => ({ + Providers: { + OpenAI: "OpenAI", + Anthropic: "Anthropic", + }, + provider_map: { + OpenAI: "openai", + Anthropic: "anthropic", + }, + providerLogoMap: { + OpenAI: "https://example.com/openai.png", + Anthropic: "https://example.com/anthropic.png", + }, +})); + +vi.mock("./provider_display_helpers", () => ({ + handleImageError: vi.fn(), +})); + +const DEFAULT_PROPS = { + marginConfig: {} as MarginConfig, + selectedProvider: undefined, + marginType: "percentage" as const, + percentageValue: "", + fixedAmountValue: "", + onProviderChange: vi.fn(), + onMarginTypeChange: vi.fn(), + onPercentageChange: vi.fn(), + onFixedAmountChange: vi.fn(), + onAddProvider: vi.fn(), +}; + +describe("AddMarginForm", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render", () => { + renderWithProviders(); + expect(screen.getByRole("button", { name: /add provider margin/i })).toBeInTheDocument(); + }); + + it("should show the percentage input when marginType is percentage", () => { + renderWithProviders(); + expect(screen.getByPlaceholderText("10")).toBeInTheDocument(); + }); + + it("should show the fixed amount input when marginType is fixed", () => { + renderWithProviders(); + expect(screen.getByPlaceholderText("0.001")).toBeInTheDocument(); + }); + + it("should not show the fixed amount input when marginType is percentage", () => { + renderWithProviders(); + expect(screen.queryByPlaceholderText("0.001")).not.toBeInTheDocument(); + }); + + it("should not show the percentage input when marginType is fixed", () => { + renderWithProviders(); + expect(screen.queryByPlaceholderText("10")).not.toBeInTheDocument(); + }); + + it("should show the Percentage-based and Fixed Amount radio options", () => { + renderWithProviders(); + expect(screen.getByText("Percentage-based")).toBeInTheDocument(); + expect(screen.getByText("Fixed Amount")).toBeInTheDocument(); + }); + + it("should disable the submit button when no provider is selected (percentage mode)", () => { + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /add provider margin/i })).toBeDisabled(); + }); + + it("should disable the submit button when provider is selected but no percentage value (percentage mode)", () => { + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /add provider margin/i })).toBeDisabled(); + }); + + it("should enable the submit button when provider and percentage value are both provided", () => { + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /add provider margin/i })).not.toBeDisabled(); + }); + + it("should disable the submit button in fixed mode when no fixed amount is provided", () => { + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /add provider margin/i })).toBeDisabled(); + }); + + it("should enable the submit button in fixed mode when provider and fixed amount are provided", () => { + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /add provider margin/i })).not.toBeDisabled(); + }); + + it("should call onAddProvider when the enabled submit button is clicked", async () => { + const onAddProvider = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + + ); + + await user.click(screen.getByRole("button", { name: /add provider margin/i })); + expect(onAddProvider).toHaveBeenCalledTimes(1); + }); + + it("should call onMarginTypeChange when the Fixed Amount radio is clicked", async () => { + const onMarginTypeChange = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + + ); + + await user.click(screen.getByText("Fixed Amount")); + expect(onMarginTypeChange).toHaveBeenCalledWith("fixed"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.test.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.test.tsx new file mode 100644 index 00000000000..611c8609c36 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.test.tsx @@ -0,0 +1,98 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../tests/test-utils"; +import AddProviderForm from "./add_provider_form"; +import { DiscountConfig } from "./types"; + +vi.mock("../provider_info_helpers", () => ({ + Providers: { + OpenAI: "OpenAI", + Anthropic: "Anthropic", + }, + provider_map: { + OpenAI: "openai", + Anthropic: "anthropic", + }, + providerLogoMap: { + OpenAI: "https://example.com/openai.png", + Anthropic: "https://example.com/anthropic.png", + }, +})); + +vi.mock("./provider_display_helpers", () => ({ + handleImageError: vi.fn(), +})); + +const DEFAULT_PROPS = { + discountConfig: {} as DiscountConfig, + selectedProvider: undefined, + newDiscount: "", + onProviderChange: vi.fn(), + onDiscountChange: vi.fn(), + onAddProvider: vi.fn(), +}; + +describe("AddProviderForm", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render", () => { + renderWithProviders(); + expect(screen.getByRole("button", { name: /add provider discount/i })).toBeInTheDocument(); + }); + + it("should render the discount percentage input field", () => { + renderWithProviders(); + expect(screen.getByPlaceholderText("5")).toBeInTheDocument(); + }); + + it("should disable the submit button when no provider is selected and no discount is entered", () => { + renderWithProviders(); + expect(screen.getByRole("button", { name: /add provider discount/i })).toBeDisabled(); + }); + + it("should disable the submit button when a provider is selected but no discount is entered", () => { + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /add provider discount/i })).toBeDisabled(); + }); + + it("should disable the submit button when a discount is entered but no provider is selected", () => { + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /add provider discount/i })).toBeDisabled(); + }); + + it("should enable the submit button when both a provider and a discount value are provided", () => { + renderWithProviders( + + ); + expect(screen.getByRole("button", { name: /add provider discount/i })).not.toBeDisabled(); + }); + + it("should call onAddProvider when the enabled submit button is clicked", async () => { + const onAddProvider = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + + ); + + await user.click(screen.getByRole("button", { name: /add provider discount/i })); + expect(onAddProvider).toHaveBeenCalledTimes(1); + }); + + it("should show the percent sign next to the discount input", () => { + renderWithProviders(); + expect(screen.getByText("%")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx new file mode 100644 index 00000000000..db6899ba17f --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx @@ -0,0 +1,201 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../tests/test-utils"; +import CostTrackingSettings from "./cost_tracking_settings"; + +// Mock sub-hooks so we can control their state without network calls +const mockDiscountConfig = vi.fn(() => ({})); +const mockMarginConfig = vi.fn(() => ({})); + +vi.mock("./use_discount_config", () => ({ + useDiscountConfig: () => ({ + discountConfig: mockDiscountConfig(), + fetchDiscountConfig: vi.fn().mockResolvedValue(undefined), + handleAddProvider: vi.fn().mockResolvedValue(true), + handleRemoveProvider: vi.fn().mockResolvedValue(undefined), + handleDiscountChange: vi.fn().mockResolvedValue(undefined), + }), +})); + +vi.mock("./use_margin_config", () => ({ + useMarginConfig: () => ({ + marginConfig: mockMarginConfig(), + fetchMarginConfig: vi.fn().mockResolvedValue(undefined), + handleAddMargin: vi.fn().mockResolvedValue(true), + handleRemoveMargin: vi.fn().mockResolvedValue(undefined), + handleMarginChange: vi.fn().mockResolvedValue(undefined), + }), +})); + +vi.mock("./pricing_calculator/index", () => ({ + default: () =>
Pricing Calculator
, +})); + +vi.mock("../playground/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn().mockResolvedValue([]), +})); + +vi.mock("../HelpLink", () => ({ + DocsMenu: () => null, +})); + +vi.mock("./how_it_works", () => ({ + default: () =>
How It Works
, +})); + +vi.mock("../provider_info_helpers", () => ({ + Providers: { OpenAI: "OpenAI" }, + provider_map: { OpenAI: "openai" }, + providerLogoMap: {}, +})); + +vi.mock("./provider_display_helpers", () => ({ + getProviderDisplayInfo: vi.fn(() => ({ displayName: "OpenAI", logo: "", enumKey: "OpenAI" })), + handleImageError: vi.fn(), +})); + +const ADMIN_PROPS = { + userID: "user-1", + userRole: "proxy_admin", + accessToken: "test-token", +}; + +describe("CostTrackingSettings", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockDiscountConfig.mockReturnValue({}); + mockMarginConfig.mockReturnValue({}); + }); + + it("should return nothing when accessToken is null", () => { + const { container } = renderWithProviders( + + ); + expect(container.firstChild).toBeNull(); + }); + + it("should render the page title", () => { + renderWithProviders(); + expect(screen.getByText("Cost Tracking Settings")).toBeInTheDocument(); + }); + + it("should show the Provider Discounts accordion header for proxy_admin", () => { + renderWithProviders(); + expect(screen.getByText("Provider Discounts")).toBeInTheDocument(); + }); + + it("should show the Fee/Price Margin accordion header for proxy_admin", () => { + renderWithProviders(); + expect(screen.getByText("Fee/Price Margin")).toBeInTheDocument(); + }); + + it("should always show the Pricing Calculator section", () => { + renderWithProviders(); + // The accordion header text appears in the DOM; getAllByText tolerates duplicates + expect(screen.getAllByText("Pricing Calculator").length).toBeGreaterThan(0); + }); + + it("should show the pricing calculator component", async () => { + renderWithProviders(); + expect(await screen.findByTestId("pricing-calculator")).toBeInTheDocument(); + }); + + it("should not show Provider Discounts section for a non-admin role", () => { + renderWithProviders( + + ); + expect(screen.queryByText("Provider Discounts")).not.toBeInTheDocument(); + }); + + it("should not show Fee/Price Margin section for a non-admin role", () => { + renderWithProviders( + + ); + expect(screen.queryByText("Fee/Price Margin")).not.toBeInTheDocument(); + }); + + it("should show Provider Discounts for the 'Admin' role as well", () => { + renderWithProviders( + + ); + expect(screen.getByText("Provider Discounts")).toBeInTheDocument(); + }); + + it("should show the subtitle describing discount/margin configuration", () => { + renderWithProviders(); + expect( + screen.getByText(/configure cost discounts and margins/i) + ).toBeInTheDocument(); + }); + + describe("Add Provider Discount modal", () => { + it("should open the Add Provider Discount modal when the button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + // The button lives inside the Provider Discounts accordion — click the header to expand first + const accordionHeader = screen.getByText("Provider Discounts").closest("button"); + if (accordionHeader) { + await user.click(accordionHeader); + } + + const addButton = await screen.findByRole("button", { name: /add provider discount/i }); + await user.click(addButton); + + expect( + await screen.findByText("Add Provider Discount", { selector: "h2" }) + ).toBeInTheDocument(); + }); + }); + + describe("Add Provider Margin modal", () => { + it("should open the Add Provider Margin modal when the button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const accordionHeader = screen.getByText("Fee/Price Margin").closest("button"); + if (accordionHeader) { + await user.click(accordionHeader); + } + + const addButton = await screen.findByRole("button", { name: /add provider margin/i }); + await user.click(addButton); + + expect( + await screen.findByText("Add Provider Margin", { selector: "h2" }) + ).toBeInTheDocument(); + }); + }); + + describe("empty state messages", () => { + it("should show the empty state message when no discount config is loaded", async () => { + mockDiscountConfig.mockReturnValue({}); + renderWithProviders(); + + const accordionHeader = screen.getByText("Provider Discounts").closest("button"); + if (accordionHeader) { + await userEvent.setup().click(accordionHeader); + } + + expect( + await screen.findByText(/no provider discounts configured/i) + ).toBeInTheDocument(); + }); + + it("should show the empty state message when no margin config is loaded", async () => { + mockMarginConfig.mockReturnValue({}); + renderWithProviders(); + + const accordionHeader = screen.getByText("Fee/Price Margin").closest("button"); + if (accordionHeader) { + await userEvent.setup().click(accordionHeader); + } + + expect( + await screen.findByText(/no provider margins configured/i) + ).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/how_it_works.test.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/how_it_works.test.tsx new file mode 100644 index 00000000000..fa608f555ce --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/how_it_works.test.tsx @@ -0,0 +1,95 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../tests/test-utils"; +import HowItWorks from "./how_it_works"; + +vi.mock("@/app/(dashboard)/api-reference/components/CodeBlock", () => ({ + default: ({ code }: { code: string }) => 
{code}
, +})); + +describe("HowItWorks", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render", () => { + renderWithProviders(); + expect(screen.getByText("Cost Calculation")).toBeInTheDocument(); + }); + + it("should display the cost calculation formula", () => { + renderWithProviders(); + expect(screen.getByText(/final_cost = base_cost/i)).toBeInTheDocument(); + }); + + it("should display the valid range information", () => { + renderWithProviders(); + expect(screen.getByText(/0% and 100%/i)).toBeInTheDocument(); + }); + + it("should render the code block with a curl example", () => { + renderWithProviders(); + expect(screen.getByTestId("code-block")).toBeInTheDocument(); + expect(screen.getByTestId("code-block").textContent).toContain("curl"); + }); + + it("should show the response header names for discount verification", () => { + renderWithProviders(); + expect(screen.getByText("x-litellm-response-cost")).toBeInTheDocument(); + expect(screen.getByText("x-litellm-response-cost-original")).toBeInTheDocument(); + expect(screen.getByText("x-litellm-response-cost-discount-amount")).toBeInTheDocument(); + }); + + it("should not show calculated results initially when no input is provided", () => { + renderWithProviders(); + expect(screen.queryByText("Calculated Results")).not.toBeInTheDocument(); + }); + + it("should not show calculated results when only response cost is entered", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const responseCostInput = screen.getByPlaceholderText("0.0171938125"); + await user.type(responseCostInput, "0.01"); + + expect(screen.queryByText("Calculated Results")).not.toBeInTheDocument(); + }); + + it("should not show calculated results when only discount amount is entered", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const discountAmountInput = screen.getByPlaceholderText("0.0009049375"); + await user.type(discountAmountInput, "0.001"); + + expect(screen.queryByText("Calculated Results")).not.toBeInTheDocument(); + }); + + it("should show calculated results when both fields are filled", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const responseCostInput = screen.getByPlaceholderText("0.0171938125"); + const discountAmountInput = screen.getByPlaceholderText("0.0009049375"); + + await user.type(responseCostInput, "0.0171938125"); + await user.type(discountAmountInput, "0.0009049375"); + + expect(await screen.findByText("Calculated Results")).toBeInTheDocument(); + }); + + it("should show original cost, final cost, and discount amount in results", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.type(screen.getByPlaceholderText("0.0171938125"), "0.0171938125"); + await user.type(screen.getByPlaceholderText("0.0009049375"), "0.0009049375"); + + expect(await screen.findByText("Original Cost:")).toBeInTheDocument(); + expect(screen.getByText("Final Cost:")).toBeInTheDocument(); + expect(screen.getByText("Discount Amount:")).toBeInTheDocument(); + expect(screen.getByText("Discount Applied:")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.test.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.test.tsx new file mode 100644 index 00000000000..7697c6e7686 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.test.tsx @@ -0,0 +1,241 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { renderWithProviders } from "../../../tests/test-utils"; +import ProviderDiscountTable from "./provider_discount_table"; + +vi.mock("@heroicons/react/outline", () => ({ + TrashIcon: function TrashIcon() { return null; }, + PencilAltIcon: function PencilAltIcon() { return null; }, + CheckIcon: function CheckIcon() { return null; }, + XIcon: function XIcon() { return null; }, +})); + +vi.mock("@tremor/react", () => ({ + Table: ({ children }: any) => {children}
, + TableHead: ({ children }: any) => {children}, + TableRow: ({ children }: any) => {children}, + TableHeaderCell: ({ children }: any) => {children}, + TableBody: ({ children }: any) => {children}, + TableCell: ({ children }: any) => {children}, + Text: ({ children }: any) => {children}, + TextInput: ({ value, onValueChange, onKeyDown, placeholder, ...rest }: any) => ( + onValueChange?.(e.target.value)} + onKeyDown={onKeyDown} + placeholder={placeholder} + {...rest} + /> + ), + Icon: ({ icon: IconComponent, onClick }: any) => { + const name = IconComponent?.displayName ?? IconComponent?.name ?? "icon"; + return