diff --git a/.circleci/config.yml b/.circleci/config.yml index faf43ff0b8b..7debc582915 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1477,6 +1477,7 @@ jobs: docker run -d \ -p 4000:4000 \ -e DATABASE_URL=$PROXY_DATABASE_URL \ + -e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \ -e DISABLE_SCHEMA_UPDATE="True" \ -v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/schema.prisma \ -v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/litellm/proxy/schema.prisma \ @@ -1912,6 +1913,7 @@ jobs: -e APORIA_API_BASE_1=$APORIA_API_BASE_1 \ -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \ -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \ + -e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \ -e USE_DDTRACE=True \ -e DD_API_KEY=$DD_API_KEY \ -e DD_SITE=$DD_SITE \ @@ -2962,6 +2964,7 @@ jobs: command: | docker run --name my-app \ -p 4000:4000 \ + -e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \ -e DATABASE_URL="postgresql://wrong:wrong@wrong:5432/wrong" \ myapp:latest \ --port 4000 > docker_output.log 2>&1 || true diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index 7e67aee8d73..0d3a9f2b5d4 100644 --- a/.github/workflows/test-litellm.yml +++ b/.github/workflows/test-litellm.yml @@ -31,6 +31,7 @@ jobs: poetry run pip install "pytest-retry==1.6.3" poetry run pip install pytest-xdist poetry run pip install "google-genai==1.22.0" + poetry run pip install "google-cloud-aiplatform>=1.38" poetry run pip install "fastapi-offline==1.7.3" - name: Setup litellm-enterprise as local package run: | diff --git a/cookbook/veo_video_generation.py b/cookbook/veo_video_generation.py new file mode 100644 index 00000000000..64a7207feb1 --- /dev/null +++ b/cookbook/veo_video_generation.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +""" +Complete example for Veo video generation through LiteLLM proxy. + +This script demonstrates how to: +1. Generate videos using Google's Veo model +2. Poll for completion status +3. Download the generated video file + +Requirements: +- LiteLLM proxy running with Google AI Studio pass-through configured +- Google AI Studio API key with Veo access +""" + +import json +import os +import time +import requests +from typing import Optional + + +class VeoVideoGenerator: + """Complete Veo video generation client using LiteLLM proxy.""" + + def __init__(self, base_url: str = "http://localhost:4000/gemini/v1beta", + api_key: str = "sk-1234"): + """ + Initialize the Veo video generator. + + Args: + base_url: Base URL for the LiteLLM proxy with Gemini pass-through + api_key: API key for LiteLLM proxy authentication + """ + self.base_url = base_url + self.api_key = api_key + self.headers = { + "x-goog-api-key": api_key, + "Content-Type": "application/json" + } + + def generate_video(self, prompt: str) -> Optional[str]: + """ + Initiate video generation with Veo. + + Args: + prompt: Text description of the video to generate + + Returns: + Operation name if successful, None otherwise + """ + print(f"šŸŽ¬ Generating video with prompt: '{prompt}'") + + url = f"{self.base_url}/models/veo-3.0-generate-preview:predictLongRunning" + payload = { + "instances": [{ + "prompt": prompt + }] + } + + try: + response = requests.post(url, headers=self.headers, json=payload) + response.raise_for_status() + + data = response.json() + operation_name = data.get("name") + + if operation_name: + print(f"āœ… Video generation started: {operation_name}") + return operation_name + else: + print("āŒ No operation name returned") + print(f"Response: {json.dumps(data, indent=2)}") + return None + + except requests.RequestException as e: + print(f"āŒ Failed to start video generation: {e}") + if hasattr(e, 'response') and e.response is not None: + try: + error_data = e.response.json() + print(f"Error details: {json.dumps(error_data, indent=2)}") + except: + print(f"Error response: {e.response.text}") + return None + + def wait_for_completion(self, operation_name: str, max_wait_time: int = 600) -> Optional[str]: + """ + Poll operation status until video generation is complete. + + Args: + operation_name: Name of the operation to monitor + max_wait_time: Maximum time to wait in seconds (default: 10 minutes) + + Returns: + Video URI if successful, None otherwise + """ + print("ā³ Waiting for video generation to complete...") + + operation_url = f"{self.base_url}/{operation_name}" + start_time = time.time() + poll_interval = 10 # Start with 10 seconds + + while time.time() - start_time < max_wait_time: + try: + print(f"šŸ” Polling status... ({int(time.time() - start_time)}s elapsed)") + + response = requests.get(operation_url, headers=self.headers) + response.raise_for_status() + + data = response.json() + + # Check for errors + if "error" in data: + print("āŒ Error in video generation:") + print(json.dumps(data["error"], indent=2)) + return None + + # Check if operation is complete + is_done = data.get("done", False) + + if is_done: + print("šŸŽ‰ Video generation complete!") + + try: + # Extract video URI from nested response + video_uri = data["response"]["generateVideoResponse"]["generatedSamples"][0]["video"]["uri"] + print(f"šŸ“¹ Video URI: {video_uri}") + return video_uri + except KeyError as e: + print(f"āŒ Could not extract video URI: {e}") + print("Full response:") + print(json.dumps(data, indent=2)) + return None + + # Wait before next poll, with exponential backoff + time.sleep(poll_interval) + poll_interval = min(poll_interval * 1.2, 30) # Cap at 30 seconds + + except requests.RequestException as e: + print(f"āŒ Error polling operation status: {e}") + time.sleep(poll_interval) + + print(f"ā° Timeout after {max_wait_time} seconds") + return None + + def download_video(self, video_uri: str, output_filename: str = "generated_video.mp4") -> bool: + """ + Download the generated video file. + + Args: + video_uri: URI of the video to download (from Google's response) + output_filename: Local filename to save the video + + Returns: + True if download successful, False otherwise + """ + print(f"ā¬‡ļø Downloading video...") + print(f"Original URI: {video_uri}") + + # Convert Google URI to LiteLLM proxy URI + # Example: files/abc123 -> /gemini/v1beta/files/abc123:download?alt=media + if video_uri.startswith("files/"): + download_path = f"{video_uri}:download?alt=media" + else: + download_path = video_uri + + litellm_download_url = f"{self.base_url}/{download_path}" + print(f"Download URL: {litellm_download_url}") + + try: + # Download with streaming and redirect handling + response = requests.get( + litellm_download_url, + headers=self.headers, + stream=True, + allow_redirects=True # Handle redirects automatically + ) + response.raise_for_status() + + # Save video file + with open(output_filename, 'wb') as f: + downloaded_size = 0 + for chunk in response.iter_content(chunk_size=8192): + if chunk: + f.write(chunk) + downloaded_size += len(chunk) + + # Progress indicator for large files + if downloaded_size % (1024 * 1024) == 0: # Every MB + print(f"šŸ“¦ Downloaded {downloaded_size / (1024*1024):.1f} MB...") + + # Verify file was created and has content + if os.path.exists(output_filename): + file_size = os.path.getsize(output_filename) + if file_size > 0: + print(f"āœ… Video downloaded successfully!") + print(f"šŸ“ Saved as: {output_filename}") + print(f"šŸ“ File size: {file_size / (1024*1024):.2f} MB") + return True + else: + print("āŒ Downloaded file is empty") + os.remove(output_filename) + return False + else: + print("āŒ File was not created") + return False + + except requests.RequestException as e: + print(f"āŒ Download failed: {e}") + if hasattr(e, 'response') and e.response is not None: + print(f"Status code: {e.response.status_code}") + print(f"Response headers: {dict(e.response.headers)}") + return False + + def generate_and_download(self, prompt: str, output_filename: str = None) -> bool: + """ + Complete workflow: generate video and download it. + + Args: + prompt: Text description for video generation + output_filename: Output filename (auto-generated if None) + + Returns: + True if successful, False otherwise + """ + # Auto-generate filename if not provided + if output_filename is None: + timestamp = int(time.time()) + safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip() + output_filename = f"veo_video_{safe_prompt.replace(' ', '_')}_{timestamp}.mp4" + + print("=" * 60) + print("šŸŽ¬ VEO VIDEO GENERATION WORKFLOW") + print("=" * 60) + + # Step 1: Generate video + operation_name = self.generate_video(prompt) + if not operation_name: + return False + + # Step 2: Wait for completion + video_uri = self.wait_for_completion(operation_name) + if not video_uri: + return False + + # Step 3: Download video + success = self.download_video(video_uri, output_filename) + + if success: + print("=" * 60) + print("šŸŽ‰ SUCCESS! Video generation complete!") + print(f"šŸ“ Video saved as: {output_filename}") + print("=" * 60) + else: + print("=" * 60) + print("āŒ FAILED! Video generation or download failed") + print("=" * 60) + + return success + + +def main(): + """ + Example usage of the VeoVideoGenerator. + + Configure these environment variables: + - LITELLM_BASE_URL: Your LiteLLM proxy URL (default: http://localhost:4000/gemini/v1beta) + - LITELLM_API_KEY: Your LiteLLM API key (default: sk-1234) + """ + + # Configuration from environment or defaults + base_url = os.getenv("LITELLM_BASE_URL", "http://localhost:4000/gemini/v1beta") + api_key = os.getenv("LITELLM_API_KEY", "sk-1234") + + print("šŸš€ Starting Veo Video Generation Example") + print(f"šŸ“” Using LiteLLM proxy at: {base_url}") + + # Initialize generator + generator = VeoVideoGenerator(base_url=base_url, api_key=api_key) + + # Example prompts - try different ones! + example_prompts = [ + "A cat playing with a ball of yarn in a sunny garden", + "Ocean waves crashing against rocky cliffs at sunset", + "A bustling city street with people walking and cars passing by", + "A peaceful forest with sunlight filtering through the trees" + ] + + # Use first example or get from user + prompt = example_prompts[0] + print(f"šŸŽ¬ Using prompt: '{prompt}'") + + # Generate and download video + success = generator.generate_and_download(prompt) + + if success: + print("\nāœ… Example completed successfully!") + print("šŸ’” Try modifying the prompt in the script for different videos!") + else: + print("\nāŒ Example failed!") + print("šŸ”§ Check your LiteLLM proxy configuration and Google AI Studio API key") + + # Troubleshooting tips + print("\nšŸ” Troubleshooting:") + print("1. Ensure LiteLLM proxy is running with Google AI Studio pass-through") + print("2. Verify your Google AI Studio API key has Veo access") + print("3. Check that your prompt meets Veo's content guidelines") + print("4. Review the LiteLLM proxy logs for detailed error information") + + +if __name__ == "__main__": + main() diff --git a/dist/litellm-1.57.6.tar.gz b/dist/litellm-1.57.6.tar.gz deleted file mode 100644 index 01a039cf6ee..00000000000 Binary files a/dist/litellm-1.57.6.tar.gz and /dev/null differ diff --git a/docs/my-website/docs/completion/input.md b/docs/my-website/docs/completion/input.md index 26629a0b8f8..9699d97b352 100644 --- a/docs/my-website/docs/completion/input.md +++ b/docs/my-website/docs/completion/input.md @@ -106,6 +106,7 @@ def completion( parallel_tool_calls: Optional[bool] = None, logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, + safety_identifier: Optional[str] = None, deployment_id=None, # soon to be deprecated params by OpenAI functions: Optional[List] = None, @@ -196,6 +197,8 @@ def completion( - `top_logprobs`: *int (optional)* - An integer between 0 and 5 specifying the number of most likely tokens to return at each token position, each with an associated log probability. `logprobs` must be set to true if this parameter is used. +- `safety_identifier`: *string (optional)* - A unique identifier for tracking and managing safety-related requests. This parameter helps with safety monitoring and compliance tracking. + - `headers`: *dict (optional)* - A dictionary of headers to be sent with the request. - `extra_headers`: *dict (optional)* - Alternative to `headers`, used to send extra headers in LLM API request. diff --git a/docs/my-website/docs/completion/web_search.md b/docs/my-website/docs/completion/web_search.md index fe49be852a7..262e3fc4f9c 100644 --- a/docs/my-website/docs/completion/web_search.md +++ b/docs/my-website/docs/completion/web_search.md @@ -8,10 +8,25 @@ Use web search with litellm | Feature | Details | |---------|---------| | Supported Endpoints | - `/chat/completions`
- `/responses` | -| Supported Providers | `openai`, `xai`, `vertex_ai`, `gemini`, `perplexity` | +| Supported Providers | `openai`, `xai`, `vertex_ai`, `anthropic`, `gemini`, `perplexity` | | LiteLLM Cost Tracking | āœ… Supported | | LiteLLM Version | `v1.71.0+` | +## Which Search Engine is Used? + +Each provider uses their own search backend: + +| Provider | Search Engine | Notes | +|----------|---------------|-------| +| **OpenAI** (`gpt-4o-search-preview`) | OpenAI's internal search | Real-time web data | +| **xAI** (`grok-3`) | xAI's search + X/Twitter | Real-time social media data | +| **Google AI/Vertex** (`gemini-2.0-flash`) | **Google Search** | Uses actual Google search results | +| **Anthropic** (`claude-3-5-sonnet`) | Anthropic's web search | Real-time web data | +| **Perplexity** | Perplexity's search engine | AI-powered search and reasoning | + +:::info +**Anthropic Web Search Models**: Claude models that support web search: `claude-3-5-sonnet-latest`, `claude-3-5-sonnet-20241022`, `claude-3-5-haiku-latest`, `claude-3-5-haiku-20241022`, `claude-3-7-sonnet-20250219` +::: ## `/chat/completions` (litellm.completion) @@ -56,6 +71,12 @@ model_list: model: xai/grok-3 api_key: os.environ/XAI_API_KEY + # Anthropic + - model_name: claude-3-5-sonnet-latest + litellm_params: + model: anthropic/claude-3-5-sonnet-latest + api_key: os.environ/ANTHROPIC_API_KEY + # VertexAI - model_name: gemini-2-flash litellm_params: @@ -143,6 +164,31 @@ response = completion( ) ``` +**Anthropic (using web_search_options)** +```python showLineNumbers +from litellm import completion + +# Customize search context size for Anthropic +response = completion( + model="anthropic/claude-3-5-sonnet-latest", + messages=[ + { + "role": "user", + "content": "What was a positive news story from today?", + } + ], + web_search_options={ + "search_context_size": "medium", # Options: "low", "medium" (default), "high" + "user_location": { + "type": "approximate", + "approximate": { + "city": "San Francisco", + }, + } + } +) +``` + **VertexAI/Gemini (using web_search_options)** ```python showLineNumbers from litellm import completion @@ -375,6 +421,9 @@ assert litellm.supports_web_search(model="openai/gpt-4o-search-preview") == True # Check xAI models assert litellm.supports_web_search(model="xai/grok-3") == True +# Check Anthropic models +assert litellm.supports_web_search(model="anthropic/claude-3-5-sonnet-latest") == True + # Check VertexAI models assert litellm.supports_web_search(model="gemini-2.0-flash") == True @@ -405,6 +454,14 @@ model_list: model_info: supports_web_search: True + # Anthropic + - model_name: claude-3-5-sonnet-latest + litellm_params: + model: anthropic/claude-3-5-sonnet-latest + api_key: os.environ/ANTHROPIC_API_KEY + model_info: + supports_web_search: True + # VertexAI - model_name: gemini-2-flash litellm_params: diff --git a/docs/my-website/docs/contributing.md b/docs/my-website/docs/contributing.md index 8fc64b8f287..8768e0b4c4d 100644 --- a/docs/my-website/docs/contributing.md +++ b/docs/my-website/docs/contributing.md @@ -14,6 +14,11 @@ git clone https://github.com/BerriAI/litellm.git Tell the proxy where the UI is located ```bash export PROXY_BASE_URL="http://localhost:3000/" + +### ALSO ### - set the basic env variables +DATABASE_URL = "postgresql://:@:/" +LITELLM_MASTER_KEY = "sk-1234" +STORE_MODEL_IN_DB = "True" ``` ```bash diff --git a/docs/my-website/docs/image_generation.md b/docs/my-website/docs/image_generation.md index 60a6356f012..7e7ff9922d6 100644 --- a/docs/my-website/docs/image_generation.md +++ b/docs/my-website/docs/image_generation.md @@ -124,8 +124,6 @@ Any non-openai params, will be treated as provider-specific params, and sent in - `size`: *string (optional)* The size of the generated images. Must be one of `1024x1024`, `1536x1024` (landscape), `1024x1536` (portrait), or `auto` (default value) for `gpt-image-1`, one of `256x256`, `512x512`, or `1024x1024` for `dall-e-2`, and one of `1024x1024`, `1792x1024`, or `1024x1792` for `dall-e-3`. -- `input_fidelity`: *string (optional)* Controls how closely the model follows the input prompt. Supported for `gpt-image-1` model. Higher fidelity may improve prompt adherence but could affect generation speed. - - `timeout`: *integer* - The maximum time, in seconds, to wait for the API to respond. Defaults to 600 seconds (10 minutes). - `user`: *string (optional)* A unique identifier representing your end-user, diff --git a/docs/my-website/docs/observability/callbacks.md b/docs/my-website/docs/observability/callbacks.md index 69cb0d053ee..040d83697d3 100644 --- a/docs/my-website/docs/observability/callbacks.md +++ b/docs/my-website/docs/observability/callbacks.md @@ -4,9 +4,14 @@ liteLLM provides `input_callbacks`, `success_callbacks` and `failure_callbacks`, making it easy for you to send data to a particular provider depending on the status of your responses. +:::tip +**New to LiteLLM Callbacks?** Check out our comprehensive [Callback Management Guide](./callback_management.md) to understand when to use different callback hooks like `async_log_success_event` vs `async_post_call_success_hook`. +::: + liteLLM supports: - [Custom Callback Functions](https://docs.litellm.ai/docs/observability/custom_callback) +- [Callback Management Guide](./callback_management.md) - **Comprehensive guide for choosing the right hooks** - [Lunary](https://lunary.ai/docs) - [Langfuse](https://langfuse.com/docs) - [LangSmith](https://www.langchain.com/langsmith) diff --git a/docs/my-website/docs/observability/cloudzero.md b/docs/my-website/docs/observability/cloudzero.md new file mode 100644 index 00000000000..f213ef64e13 --- /dev/null +++ b/docs/my-website/docs/observability/cloudzero.md @@ -0,0 +1,209 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# CloudZero Integration + +LiteLLM provides an integration with CloudZero's AnyCost API, allowing you to export your LLM usage data to CloudZero for cost tracking analysis. + +## Overview + +| Property | Details | +|----------|---------| +| Description | Export LiteLLM usage data to CloudZero AnyCost API for cost tracking and analysis | +| callback name | `cloudzero`| +| Supported Operations | • Automatic hourly data export
• Manual data export
• Dry run testing
• Cost and token usage tracking | +| Data Format | CloudZero Billing Format (CBF) with proper resource tagging | +| Export Frequency | Hourly (configurable via `CLOUDZERO_EXPORT_INTERVAL_MINUTES`) | + +## Environment Variables + +| Variable | Required | Description | Example | +|----------|----------|-------------|---------| +| `CLOUDZERO_API_KEY` | Yes | Your CloudZero API key | `cz_api_xxxxxxxxxx` | +| `CLOUDZERO_CONNECTION_ID` | Yes | CloudZero connection ID for data submission | `conn_xxxxxxxxxx` | +| `CLOUDZERO_TIMEZONE` | No | Timezone for date handling (default: UTC) | `America/New_York` | +| `CLOUDZERO_EXPORT_INTERVAL_MINUTES` | No | Export frequency in minutes (default: 60) | `60` | + +## Setup + +### End to End Video Walkthrough +This video walks through the entire process of setting up LiteLLM with CloudZero integration and viewing LiteLLM exported usage data in CloudZero. + + + +### Step 1: Configure Environment Variables + +Set your CloudZero credentials in your environment: + +```bash +export CLOUDZERO_API_KEY="cz_api_xxxxxxxxxx" +export CLOUDZERO_CONNECTION_ID="conn_xxxxxxxxxx" +export CLOUDZERO_TIMEZONE="UTC" # Optional, defaults to UTC +``` + +### Step 2: Enable CloudZero Integration + +Add the CloudZero callback to your LiteLLM configuration YAML file: + + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: sk-xxxxxxx + +litellm_settings: + callbacks: ["cloudzero"] # Enable CloudZero integration +``` + +### Step 3: Start LiteLLM Proxy + +Start your LiteLLM proxy with the configuration: + +```bash +litellm --config /path/to/config.yaml +``` + +## Testing Your Setup + +### Dry Run Export + +Call the dry run endpoint to test your CloudZero configuration without sending data to CloudZero. This endpoint will not send any data to CloudZero, but will return the data that would be exported. + +```bash +curl -X POST "http://localhost:4000/cloudzero/dry-run" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "limit": 10 + }' | jq +``` + +**Expected Response:** +```json +{ + "message": "CloudZero dry run export completed successfully.", + "status": "success", + "dry_run_data": { + "usage_data": [...], + "cbf_data": [...], + "summary": { + "total_cost": 0.05, + "total_tokens": 1250, + "total_records": 10 + } + } +} +``` + +### Manual Export + +Call the export endpoint to send data immediately to CloudZero. We suggest setting a small `limit` to test the export. This will only export the last 10 records to CloudZero. Note: Cloudzero can take up to 15 minutes to process the exported data. + +```bash +curl -X POST "http://localhost:4000/cloudzero/export" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "limit": 10 + }' | jq +``` + +**Expected Response:** +```json +{ + "message": "CloudZero export completed successfully", + "status": "success" +} +``` + +## Data Export Details + +### Automatic Export Schedule + +- **Frequency**: Every 60 minutes (configurable via `CLOUDZERO_EXPORT_INTERVAL_MINUTES`) +- **Data Processing**: LiteLLM automatically processes and exports usage data hourly +- **CloudZero Processing**: CloudZero typically takes 10-15 minutes to process data from LiteLLM + +### Data Format + +LiteLLM exports data in CloudZero Billing Format (CBF) with the following structure: + +```json +{ + "time/usage_start": "2024-01-15T14:00:00Z", + "cost/cost": 0.002, + "usage/amount": 150, + "usage/units": "tokens", + "resource/id": "czrn:litellm:openai:cross-region:team-123:llm-usage:gpt-4o", + "resource/service": "litellm", + "resource/account": "team-123", + "resource/region": "cross-region", + "resource/usage_family": "llm-usage", + "resource/tag:provider": "openai", + "resource/tag:model": "gpt-4o", + "resource/tag:prompt_tokens": "100", + "resource/tag:completion_tokens": "50" +} +``` + +### Resource Tagging + +LiteLLM automatically creates comprehensive resource tags for cost attribution: + +- **Provider Tags**: `openai`, `anthropic`, `azure`, etc. +- **Model Tags**: Specific model names like `gpt-4o`, `claude-3-sonnet` +- **Team/User Tags**: Team IDs and user IDs for cost allocation +- **Token Breakdown**: Separate tracking of prompt and completion tokens +- **Usage Metrics**: Total tokens consumed per request + +## Advanced Configuration + +### Custom Export Frequency + +Change the export frequency (not recommended to go below 60 minutes): + +```bash +export CLOUDZERO_EXPORT_INTERVAL_MINUTES=120 # Export every 2 hours +``` + +### Custom Time Range Export + +Export data for a specific time range: + +```bash +curl -X POST "http://localhost:4000/cloudzero/export" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "start_time_utc": "2024-01-15T00:00:00Z", + "end_time_utc": "2024-01-15T23:59:59Z", + "operation": "replace_hourly" + }' | jq +``` + +## Troubleshooting + +### Common Issues + +1. **Missing Credentials Error** + ``` + CloudZero configuration missing. Please set CLOUDZERO_API_KEY and CLOUDZERO_CONNECTION_ID environment variables. + ``` + **Solution**: Ensure both environment variables are set with valid values. + +2. **Connection Issues** + - Verify your CloudZero API key is valid + - Check that the connection ID exists in your CloudZero account + - Ensure your proxy has internet access to reach CloudZero's API + +3. **No Data in CloudZero** + - CloudZero can take 10-15 minutes to process data + - Check that your LiteLLM proxy is generating usage data + - Use the dry-run endpoint to verify data is being formatted correctly + +## Related Links + +- [CloudZero Documentation](https://docs.cloudzero.com/) +- [CloudZero AnyCost API](https://docs.cloudzero.com/reference/anycost-api) diff --git a/docs/my-website/docs/observability/custom_callback.md b/docs/my-website/docs/observability/custom_callback.md index cc586b2e5d9..c206c23d0f4 100644 --- a/docs/my-website/docs/observability/custom_callback.md +++ b/docs/my-website/docs/observability/custom_callback.md @@ -4,7 +4,6 @@ **For PROXY** [Go Here](../proxy/logging.md#custom-callback-class-async) ::: - ## Callback Class You can create a custom callback class to precisely log events as they occur in litellm. @@ -57,6 +56,17 @@ def async completion(): asyncio.run(completion()) ``` +## Common Hooks + +- `async_log_success_event` - Log successful API calls +- `async_log_failure_event` - Log failed API calls +- `log_pre_api_call` - Log before API call +- `log_post_api_call` - Log after API call + +**Proxy-only hooks** (only work with LiteLLM Proxy): +- `async_post_call_success_hook` - Access user data + modify responses +- `async_pre_call_hook` - Modify requests before sending + ## Callback Functions If you just want to log on a specific event (e.g. on input) - you can use callback functions. @@ -174,260 +184,87 @@ async def test_chat_openai(): asyncio.run(test_chat_openai()) ``` -:::info +## What's Available in kwargs? -We're actively trying to expand this to other event types. [Tell us if you need this!](https://github.com/BerriAI/litellm/issues/1007) -::: - -## What's in kwargs? - -Notice we pass in a kwargs argument to custom callback. -```python -def custom_callback( - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, end_time # start/end time -): - # Your custom code here - print("LITELLM: in custom callback function") - print("kwargs", kwargs) - print("completion_response", completion_response) - print("start_time", start_time) - print("end_time", end_time) -``` - -This is a dictionary containing all the model-call details (the params we receive, the values we send to the http endpoint, the response we receive, stacktrace in case of errors, etc.). - -This is all logged in the [model_call_details via our Logger](https://github.com/BerriAI/litellm/blob/fc757dc1b47d2eb9d0ea47d6ad224955b705059d/litellm/utils.py#L246). - -Here's exactly what you can expect in the kwargs dictionary: -```shell -### DEFAULT PARAMS ### -"model": self.model, -"messages": self.messages, -"optional_params": self.optional_params, # model-specific params passed in -"litellm_params": self.litellm_params, # litellm-specific params passed in (e.g. metadata passed to completion call) -"start_time": self.start_time, # datetime object of when call was started - -### PRE-API CALL PARAMS ### (check via kwargs["log_event_type"]="pre_api_call") -"input" = input # the exact prompt sent to the LLM API -"api_key" = api_key # the api key used for that LLM API -"additional_args" = additional_args # any additional details for that API call (e.g. contains optional params sent) - -### POST-API CALL PARAMS ### (check via kwargs["log_event_type"]="post_api_call") -"original_response" = original_response # the original http response received (saved via response.text) - -### ON-SUCCESS PARAMS ### (check via kwargs["log_event_type"]="successful_api_call") -"complete_streaming_response" = complete_streaming_response # the complete streamed response (only set if `completion(..stream=True)`) -"end_time" = end_time # datetime object of when call was completed - -### ON-FAILURE PARAMS ### (check via kwargs["log_event_type"]="failed_api_call") -"exception" = exception # the Exception raised -"traceback_exception" = traceback_exception # the traceback generated via `traceback.format_exc()` -"end_time" = end_time # datetime object of when call was completed -``` - - -### Cache hits - -Cache hits are logged in success events as `kwarg["cache_hit"]`. - -Here's an example of accessing it: - - ```python - import litellm -from litellm.integrations.custom_logger import CustomLogger -from litellm import completion, acompletion, Cache - -class MyCustomHandler(CustomLogger): - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - print(f"On Success") - print(f"Value of Cache hit: {kwargs['cache_hit']"}) - -async def test_async_completion_azure_caching(): - customHandler_caching = MyCustomHandler() - litellm.cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD']) - litellm.callbacks = [customHandler_caching] - unique_time = time.time() - response1 = await litellm.acompletion(model="azure/chatgpt-v-2", - messages=[{ - "role": "user", - "content": f"Hi šŸ‘‹ - i'm async azure {unique_time}" - }], - caching=True) - await asyncio.sleep(1) - print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}") - response2 = await litellm.acompletion(model="azure/chatgpt-v-2", - messages=[{ - "role": "user", - "content": f"Hi šŸ‘‹ - i'm async azure {unique_time}" - }], - caching=True) - await asyncio.sleep(1) # success callbacks are done in parallel - print(f"customHandler_caching.states post-cache hit: {customHandler_caching.states}") - assert len(customHandler_caching.errors) == 0 - assert len(customHandler_caching.states) == 4 # pre, post, success, success - ``` - -### Get complete streaming response - -LiteLLM will pass you the complete streaming response in the final streaming chunk as part of the kwargs for your custom callback function. +The kwargs dictionary contains all the details about your API call: ```python -# litellm.set_verbose = False - def custom_callback( - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, end_time # start/end time - ): - # print(f"streaming response: {completion_response}") - if "complete_streaming_response" in kwargs: - print(f"Complete Streaming Response: {kwargs['complete_streaming_response']}") - - # Assign the custom callback function - litellm.success_callback = [custom_callback] - - response = completion(model="claude-instant-1", messages=messages, stream=True) - for idx, chunk in enumerate(response): - pass -``` - - -### Log additional metadata - -LiteLLM accepts a metadata dictionary in the completion call. You can pass additional metadata into your completion call via `completion(..., metadata={"key": "value"})`. - -Since this is a [litellm-specific param](https://github.com/BerriAI/litellm/blob/b6a015404eed8a0fa701e98f4581604629300ee3/litellm/main.py#L235), it's accessible via kwargs["litellm_params"] - -```python -from litellm import completion -import os, litellm - -## set ENV variables -os.environ["OPENAI_API_KEY"] = "your-api-key" - -messages = [{ "content": "Hello, how are you?","role": "user"}] - -def custom_callback( - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, end_time # start/end time -): - print(kwargs["litellm_params"]["metadata"]) +def custom_callback(kwargs, completion_response, start_time, end_time): + # Access common data + model = kwargs.get("model") + messages = kwargs.get("messages", []) + cost = kwargs.get("response_cost", 0) + cache_hit = kwargs.get("cache_hit", False) - -# Assign the custom callback function -litellm.success_callback = [custom_callback] - -response = litellm.completion(model="gpt-3.5-turbo", messages=messages, metadata={"hello": "world"}) + # Access metadata you passed in + metadata = kwargs.get("litellm_params", {}).get("metadata", {}) ``` -## Examples +**Key fields in kwargs:** +- `model` - The model name +- `messages` - Input messages +- `response_cost` - Calculated cost +- `cache_hit` - Whether response was cached +- `litellm_params.metadata` - Your custom metadata -### Custom Callback to track costs for Streaming + Non-Streaming -By default, the response cost is accessible in the logging object via `kwargs["response_cost"]` on success (sync + async) +## Practical Examples + +### Track API Costs ```python +def track_cost_callback(kwargs, completion_response, start_time, end_time): + cost = kwargs["response_cost"] # litellm calculates this for you + print(f"Request cost: ${cost}") -# Step 1. Write your custom callback function -def track_cost_callback( - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, end_time # start/end time -): - try: - response_cost = kwargs["response_cost"] # litellm calculates response cost for you - print("regular response_cost", response_cost) - except: - pass - -# Step 2. Assign the custom callback function litellm.success_callback = [track_cost_callback] -# Step 3. Make litellm.completion call -response = completion( - model="gpt-3.5-turbo", - messages=[ - { - "role": "user", - "content": "Hi šŸ‘‹ - i'm openai" - } - ] -) - -print(response) +response = completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello"}]) ``` -### Custom Callback to log transformed Input to LLMs +### Log Inputs to LLMs ```python -def get_transformed_inputs( - kwargs, -): +def get_transformed_inputs(kwargs): params_to_model = kwargs["additional_args"]["complete_input_dict"] print("params to model", params_to_model) litellm.input_callback = [get_transformed_inputs] -def test_chat_openai(): - try: - response = completion(model="claude-2", - messages=[{ - "role": "user", - "content": "Hi šŸ‘‹ - i'm openai" - }]) - - print(response) - - except Exception as e: - print(e) - pass +response = completion(model="claude-2", messages=[{"role": "user", "content": "Hello"}]) ``` -#### Output -```shell -params to model {'model': 'claude-2', 'prompt': "\n\nHuman: Hi šŸ‘‹ - i'm openai\n\nAssistant: ", 'max_tokens_to_sample': 256} +### Send to External Service +```python +import requests + +def send_to_analytics(kwargs, completion_response, start_time, end_time): + data = { + "model": kwargs.get("model"), + "cost": kwargs.get("response_cost", 0), + "duration": (end_time - start_time).total_seconds() + } + requests.post("https://your-analytics.com/api", json=data) + +litellm.success_callback = [send_to_analytics] ``` -### Custom Callback to write to Mixpanel +## Common Issues + +### Callback Not Called +Make sure you: +1. Register callbacks correctly: `litellm.callbacks = [MyHandler()]` +2. Use the right hook names (check spelling) +3. Don't use proxy-only hooks in library mode + +### Performance Issues +- Use async hooks for I/O operations +- Don't block in callback functions +- Handle exceptions properly: ```python -import mixpanel -import litellm -from litellm import completion - -def custom_callback( - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, end_time # start/end time -): - # Your custom code here - mixpanel.track("LLM Response", {"llm_response": completion_response}) - - -# Assign the custom callback function -litellm.success_callback = [custom_callback] - -response = completion( - model="gpt-3.5-turbo", - messages=[ - { - "role": "user", - "content": "Hi šŸ‘‹ - i'm openai" - } - ] -) - -print(response) - +class SafeHandler(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + try: + await external_service(response_obj) + except Exception as e: + print(f"Callback error: {e}") # Log but don't break the flow ``` - - - - - - - - - - - diff --git a/docs/my-website/docs/pass_through/google_ai_studio.md b/docs/my-website/docs/pass_through/google_ai_studio.md index c3671f58d36..3de7c54aa7a 100644 --- a/docs/my-website/docs/pass_through/google_ai_studio.md +++ b/docs/my-website/docs/pass_through/google_ai_studio.md @@ -230,6 +230,13 @@ curl -X POST "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5 ``` +## **Example 4: Video Generation with Veo** + +Generate videos using Google's Veo model through LiteLLM pass-through routes. + +[**→ Complete Veo Video Generation Guide**](../proxy/veo_video_generation.md) + + ## Advanced Pre-requisites diff --git a/docs/my-website/docs/pass_through/intro.md b/docs/my-website/docs/pass_through/intro.md index 3d6286afcc5..38218224f11 100644 --- a/docs/my-website/docs/pass_through/intro.md +++ b/docs/my-website/docs/pass_through/intro.md @@ -11,3 +11,43 @@ These endpoints are useful for 2 scenarios: ## How is your request handled? The request is passed through to the provider's endpoint. The response is then passed back to the client. **No translation is done.** + +### Request Forwarding Process + +1. **Request Reception**: LiteLLM receives your request at `/provider/endpoint` +2. **Authentication**: Your LiteLLM API key is validated and mapped to the provider's API key +3. **Request Transformation**: Request is reformatted for the target provider's API +4. **Forwarding**: Request is sent to the actual provider endpoint +5. **Response Handling**: Provider response is returned directly to you + +### Authentication Flow + +```mermaid +graph LR + A[Client Request] --> B[LiteLLM Proxy] + B --> C[Validate LiteLLM API Key] + C --> D[Map to Provider API Key] + D --> E[Forward to Provider] + E --> F[Return Response] +``` + +**Key Points:** +- Use your **LiteLLM API key** in requests, not the provider's key +- LiteLLM handles the provider authentication internally +- Same authentication works across all passthrough endpoints + +### Error Handling + +**Provider Errors**: Forwarded directly to you with original error codes and messages + +**LiteLLM Errors**: +- `401`: Invalid LiteLLM API key +- `404`: Provider or endpoint not supported +- `500`: Internal routing/forwarding errors + +### Benefits + +- **Unified Authentication**: One API key for all providers +- **Centralized Logging**: All requests logged through LiteLLM +- **Cost Tracking**: Usage tracked across all endpoints +- **Access Control**: Same permissions apply to passthrough endpoints diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index 1356ec1744e..c191b742268 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -467,7 +467,7 @@ print(f"\nResponse: {resp}") ## Usage - 'thinking' / 'reasoning content' -This is currently only supported for Anthropic's Claude 3.7 Sonnet + Deepseek R1. +This is currently only supported for Anthropic's Claude 3.7 Sonnet + Deepseek R1 + GPT-OSS models. Works on v1.61.20+. diff --git a/docs/my-website/docs/providers/databricks.md b/docs/my-website/docs/providers/databricks.md index 8631cbfdad9..921b06a17b7 100644 --- a/docs/my-website/docs/providers/databricks.md +++ b/docs/my-website/docs/providers/databricks.md @@ -282,6 +282,11 @@ ModelResponse( ) ``` +### Citations + +Anthropic models served through Databricks can return citation metadata. LiteLLM +exposes these via `response.choices[0].message.provider_specific_fields["citations"]`. + ### Pass `thinking` to Anthropic models You can also pass the `thinking` parameter to Anthropic models. diff --git a/docs/my-website/docs/providers/vertex_partner.md b/docs/my-website/docs/providers/vertex_partner.md index cf780e35dbd..856f054b8e6 100644 --- a/docs/my-website/docs/providers/vertex_partner.md +++ b/docs/my-website/docs/providers/vertex_partner.md @@ -15,6 +15,7 @@ import TabItem from '@theme/TabItem'; | Mistral | `vertex_ai/mistral-*` | [Vertex AI - Mistral Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/mistral) | | AI21 (Jamba) | `vertex_ai/jamba-*` | [Vertex AI - AI21 Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/ai21) | | Qwen | `vertex_ai/qwen/*` | [Vertex AI - Qwen Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/qwen) | +| OpenAI (GPT-OSS) | `vertex_ai/openai/gpt-oss-*` | [Vertex AI - GPT-OSS Models](https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/) | | Model Garden | `vertex_ai/openai/{MODEL_ID}` or `vertex_ai/{MODEL_ID}` | [Vertex Model Garden](https://cloud.google.com/model-garden?hl=en) | ## Vertex AI - Anthropic (Claude) @@ -658,6 +659,141 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ +## VertexAI GPT-OSS Models + +| Property | Details | +|----------|---------| +| Provider Route | `vertex_ai/openai/{MODEL}` | +| Vertex Documentation | [Vertex AI - GPT-OSS Models](https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/) | + +**LiteLLM Supports all Vertex AI GPT-OSS Models.** Ensure you use the `vertex_ai/openai/` prefix for all Vertex AI GPT-OSS models. + +| Model Name | Usage | +|------------------|------------------------------| +| vertex_ai/openai/gpt-oss-20b-maas | `completion('vertex_ai/openai/gpt-oss-20b-maas', messages)` | + +#### Usage + + + + +```python +from litellm import completion +import os + +os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "" + +model = "openai/gpt-oss-20b-maas" + +vertex_ai_project = "your-vertex-project" # can also set this as os.environ["VERTEXAI_PROJECT"] +vertex_ai_location = "your-vertex-location" # can also set this as os.environ["VERTEXAI_LOCATION"] + +response = completion( + model="vertex_ai/" + model, + messages=[{"role": "user", "content": "hi"}], + vertex_ai_project=vertex_ai_project, + vertex_ai_location=vertex_ai_location, +) +print("\nModel Response", response) +``` + + + +**1. Add to config** + +```yaml +model_list: + - model_name: gpt-oss + litellm_params: + model: vertex_ai/openai/gpt-oss-20b-maas + vertex_ai_project: "my-test-project" + vertex_ai_location: "us-central1" +``` + +**2. Start proxy** + +```bash +litellm --config /path/to/config.yaml + +# RUNNING at http://0.0.0.0:4000 +``` + +**3. Test it!** + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "model": "gpt-oss", # šŸ‘ˆ the 'model_name' in config + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ], + }' +``` + + + + +#### Usage - `reasoning_effort` + +GPT-OSS models support the `reasoning_effort` parameter for enhanced reasoning capabilities. + + + + +```python +from litellm import completion + +response = completion( + model="vertex_ai/openai/gpt-oss-20b-maas", + messages=[{"role": "user", "content": "Solve this complex problem step by step"}], + reasoning_effort="low", # Options: "minimal", "low", "medium", "high" + vertex_ai_project="your-vertex-project", + vertex_ai_location="us-central1", +) +``` + + + + + +1. Setup config.yaml + +```yaml +model_list: +- model_name: gpt-oss + litellm_params: + model: vertex_ai/openai/gpt-oss-20b-maas + vertex_ai_project: "my-test-project" + vertex_ai_location: "us-central1" +``` + +2. Start proxy + +```bash +litellm --config /path/to/config.yaml +``` + +3. Test it! + +```bash +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{ + "model": "gpt-oss", + "messages": [{"role": "user", "content": "Solve this complex problem step by step"}], + "reasoning_effort": "low" + }' +``` + + + + ## Model Garden :::tip diff --git a/docs/my-website/docs/providers/volcano.md b/docs/my-website/docs/providers/volcano.md index 1742a43d819..efd1e02b60b 100644 --- a/docs/my-website/docs/providers/volcano.md +++ b/docs/my-website/docs/providers/volcano.md @@ -3,7 +3,7 @@ https://www.volcengine.com/docs/82379/1263482 :::tip -**We support ALL Volcengine NIM models, just set `model=volcengine/` as a prefix when sending litellm requests** +**We support ALL Volcengine models including Chat and Embeddings, just set `model=volcengine/` as a prefix when sending litellm requests** ::: @@ -11,6 +11,8 @@ https://www.volcengine.com/docs/82379/1263482 ```python # env variable os.environ['VOLCENGINE_API_KEY'] +# or +os.environ['ARK_API_KEY'] ``` ## Sample Usage @@ -64,9 +66,42 @@ for chunk in response: print(chunk) ``` +## Sample Usage - Embedding +```python +from litellm import embedding +import os -## Supported Models - šŸ’„ ALL Volcengine NIM Models Supported! -We support ALL `volcengine` models, just set `volcengine/` as a prefix when sending completion requests +os.environ['VOLCENGINE_API_KEY'] = "" +response = embedding( + model="volcengine/doubao-embedding-text-240715", + input=["hello world", "good morning"] +) +print(response) +``` + +### Supported Embedding Models +- `doubao-embedding-large` (2048 dimensions) +- `doubao-embedding-large-text-250515` (2048 dimensions) +- `doubao-embedding-large-text-240915` (4096 dimensions) +- `doubao-embedding` (2560 dimensions) +- `doubao-embedding-text-240715` (2560 dimensions) + +### Embedding Parameters +```python +from litellm import embedding + +response = embedding( + model="volcengine/doubao-embedding-text-240715", + input=["sample text"], + encoding_format="float", # optional: "float" (default), "base64" + user="user-123", # optional: user identifier for tracking +) +``` + +## Supported Models - šŸ’„ ALL Volcengine Models Supported! +We support ALL `volcengine` models for both chat completions and embeddings: +- **Chat Models**: Set `volcengine/` as a prefix when sending completion requests +- **Embedding Models**: Use the specific model names listed above (e.g., `volcengine/doubao-embedding-text-240715`) ## Sample Usage - LiteLLM Proxy @@ -74,14 +109,21 @@ We support ALL `volcengine` models, just set `volcengine/` as a ```yaml model_list: + # Chat model - model_name: volcengine-model litellm_params: model: volcengine/ api_key: os.environ/VOLCENGINE_API_KEY + # Embedding model + - model_name: volcengine-embedding + litellm_params: + model: volcengine/doubao-embedding-text-240715 + api_key: os.environ/VOLCENGINE_API_KEY ``` ### Send Request +#### Chat Completion ```shell curl --location 'http://localhost:4000/chat/completions' \ --header 'Authorization: Bearer sk-1234' \ @@ -95,4 +137,15 @@ curl --location 'http://localhost:4000/chat/completions' \ } ] }' +``` + +#### Embedding +```shell +curl --location 'http://localhost:4000/embeddings' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "model": "volcengine-embedding", + "input": ["hello world", "good morning"] +}' ``` \ No newline at end of file diff --git a/docs/my-website/docs/proxy/access_control.md b/docs/my-website/docs/proxy/access_control.md index 69b8a3ff6de..4ca3eb119d6 100644 --- a/docs/my-website/docs/proxy/access_control.md +++ b/docs/my-website/docs/proxy/access_control.md @@ -4,7 +4,7 @@ Role-based access control (RBAC) is based on Organizations, Teams and Internal U - `Organizations` are the top-level entities that contain Teams. - `Team` - A Team is a collection of multiple `Internal Users` -- `Internal Users` - users that can create keys, make LLM API calls, view usage on LiteLLM +- `Internal Users` - users that can create keys, make LLM API calls, view usage on LiteLLM. Users can be on multiple teams. - `Roles` define the permissions of an `Internal User` - `Virtual Keys` - Keys are used for authentication to the LiteLLM API. Keys are tied to a `Internal User` and `Team` diff --git a/docs/my-website/docs/proxy/admin_ui_sso.md b/docs/my-website/docs/proxy/admin_ui_sso.md index 86cb6b0bf8c..823301d4c38 100644 --- a/docs/my-website/docs/proxy/admin_ui_sso.md +++ b/docs/my-website/docs/proxy/admin_ui_sso.md @@ -235,6 +235,13 @@ Example setting a local image (on your container) ```shell UI_LOGO_PATH="ui_images/logo.jpg" ``` + +#### Or set your logo directly from Admin UI: +
+ + +
+ #### Set Custom Color Theme - Navigate to [/enterprise/enterprise_ui](https://github.com/BerriAI/litellm/blob/main/enterprise/enterprise_ui/_enterprise_colors.json) - Inside the `enterprise_ui` directory, rename `_enterprise_colors.json` to `enterprise_colors.json` diff --git a/docs/my-website/docs/proxy/call_hooks.md b/docs/my-website/docs/proxy/call_hooks.md index b4e22027d19..aef33f8c708 100644 --- a/docs/my-website/docs/proxy/call_hooks.md +++ b/docs/my-website/docs/proxy/call_hooks.md @@ -6,6 +6,10 @@ import Image from '@theme/IdealImage'; - Reject data before making llm api calls / before returning the response - Enforce 'user' param for all openai endpoint calls +:::tip +**Understanding Callback Hooks?** Check out our [Callback Management Guide](../observability/callback_management.md) to understand the differences between proxy-specific hooks like `async_pre_call_hook` and general logging hooks like `async_log_success_event`. +::: + See a complete example with our [parallel request rate limiter](https://github.com/BerriAI/litellm/blob/main/litellm/proxy/hooks/parallel_request_limiter.py) ## Quick Start diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 541dc6fb3c8..36f9bbc40a5 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -335,12 +335,15 @@ router_settings: | ANTHROPIC_API_KEY | API key for Anthropic service | ANTHROPIC_API_BASE | Base URL for Anthropic API. Default is https://api.anthropic.com | AWS_ACCESS_KEY_ID | Access Key ID for AWS services +| AWS_BATCH_ROLE_ARN | ARN of the AWS IAM role for batch operations | AWS_DEFAULT_REGION | Default AWS region for service interactions when AWS_REGION is not set | AWS_PROFILE_NAME | AWS CLI profile name to be used | AWS_REGION | AWS region for service interactions (takes precedence over AWS_DEFAULT_REGION) | AWS_REGION_NAME | Default AWS region for service interactions | AWS_ROLE_ARN | ARN of the AWS IAM role to assume for authentication | AWS_ROLE_NAME | Role name for AWS IAM usage +| AWS_S3_BUCKET_NAME | Name of the AWS S3 bucket for file operations +| AWS_S3_OUTPUT_BUCKET_NAME | Name of the AWS S3 output bucket for batch operations | AWS_SECRET_ACCESS_KEY | Secret Access Key for AWS services | AWS_SESSION_NAME | Name for AWS session | AWS_WEB_IDENTITY_TOKEN | Web identity token for AWS @@ -380,6 +383,8 @@ router_settings: | CIRCLE_OIDC_TOKEN_V2 | Version 2 of the OpenID Connect token for CircleCI | CLOUDZERO_API_KEY | CloudZero API key for authentication | CLOUDZERO_CONNECTION_ID | CloudZero connection ID for data submission +| CLOUDZERO_EXPORT_INTERVAL_MINUTES | Interval in minutes for CloudZero data export operations +| CLOUDZERO_MAX_FETCHED_DATA_RECORDS | Maximum number of data records to fetch from CloudZero | CLOUDZERO_TIMEZONE | Timezone for date handling (default: UTC) | CONFIG_FILE_PATH | File path for configuration file | CONFIDENT_API_KEY | API key for DeepEval integration @@ -412,6 +417,7 @@ router_settings: | DEFAULT_ALLOWED_FAILS | Maximum failures allowed before cooling down a model. Default is 3 | DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS | Default maximum tokens for Anthropic chat completions. Default is 4096 | DEFAULT_BATCH_SIZE | Default batch size for operations. Default is 512 +| DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS | Timeout in seconds for checking client disconnection. Default is 1 | DEFAULT_COOLDOWN_TIME_SECONDS | Duration in seconds to cooldown a model after failures. Default is 5 | DEFAULT_CRON_JOB_LOCK_TTL_SECONDS | Time-to-live for cron job locks in seconds. Default is 60 (1 minute) | DEFAULT_FAILURE_THRESHOLD_PERCENT | Threshold percentage of failures to cool down a deployment. Default is 0.5 (50%) @@ -438,6 +444,10 @@ router_settings: | DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET | Default high reasoning effort thinking budget. Default is 4096 | DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET | Default low reasoning effort thinking budget. Default is 1024 | DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET | Default medium reasoning effort thinking budget. Default is 2048 +| DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET | Default minimal reasoning effort thinking budget. Default is 512 +| DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH | Default minimal reasoning effort thinking budget for Gemini 2.5 Flash. Default is 512 +| DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE | Default minimal reasoning effort thinking budget for Gemini 2.5 Flash Lite. Default is 512 +| DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO | Default minimal reasoning effort thinking budget for Gemini 2.5 Pro. Default is 512 | DEFAULT_REDIS_SYNC_INTERVAL | Default Redis synchronization interval in seconds. Default is 1 | DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND | Default price per second for Replicate GPU. Default is 0.001400 | DEFAULT_REPLICATE_POLLING_DELAY_SECONDS | Default delay in seconds for Replicate polling. Default is 1 @@ -573,6 +583,10 @@ router_settings: | LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM | LITELLM_LOG | Enable detailed logging for LiteLLM | LITELLM_LOG_FILE | File path to write LiteLLM logs to. When set, logs will be written to both console and the specified file +| LITELLM_LOGGER_NAME | Name for OTEL logger +| LITELLM_METER_NAME | Name for OTEL Meter +| LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS | Optionally enable semantic logs for OTEL +| LITELLM_OTEL_INTEGRATION_ENABLE_METRICS | Optionally enable emantic metrics for OTEL | LITELLM_MASTER_KEY | Master key for proxy authentication | LITELLM_MODE | Operating mode for LiteLLM (e.g., production, development) | LITELLM_RATE_LIMIT_WINDOW_SIZE | Rate limit window size for LiteLLM. Default is 60 diff --git a/docs/my-website/docs/proxy/cost_tracking.md b/docs/my-website/docs/proxy/cost_tracking.md index 19e3344f21b..7f0a13f763f 100644 --- a/docs/my-website/docs/proxy/cost_tracking.md +++ b/docs/my-website/docs/proxy/cost_tracking.md @@ -505,11 +505,11 @@ litellm_settings: ### Disable user-agent tracking -You can disable user-agent tracking by setting `litellm_settings.disable_user_agent_tracking` to `true`. +You can disable user-agent tracking by setting `litellm_settings.disable_add_user_agent_to_request_tags` to `true`. ```yaml litellm_settings: - disable_user_agent_tracking: true + disable_add_user_agent_to_request_tags: true ``` ## ✨ (Enterprise) Generate Spend Reports diff --git a/docs/my-website/docs/proxy/enterprise.md b/docs/my-website/docs/proxy/enterprise.md index 468bcad2cf8..7d50aedb424 100644 --- a/docs/my-website/docs/proxy/enterprise.md +++ b/docs/my-website/docs/proxy/enterprise.md @@ -439,6 +439,33 @@ response = client.chat.completions.create( print(response) ``` + +**Using Headers:** + +```python +import openai +client = openai.OpenAI( + api_key="sk-1234", + base_url="http://0.0.0.0:4000" +) + +# Pass spend logs metadata via headers +response = client.chat.completions.create( + model="gpt-3.5-turbo", + messages = [ + { + "role": "user", + "content": "this is a test request, write a short poem" + } + ], + extra_headers={ + "x-litellm-spend-logs-metadata": '{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}' + } +) + +print(response) +``` + @@ -478,6 +505,43 @@ async function runOpenAI() { // Call the asynchronous function runOpenAI(); ``` + +**Using Headers:** + +```js +const openai = require('openai'); + +async function runOpenAI() { + const client = new openai.OpenAI({ + apiKey: 'sk-1234', + baseURL: 'http://0.0.0.0:4000' + }); + + try { + const response = await client.chat.completions.create({ + model: 'gpt-3.5-turbo', + messages: [ + { + role: 'user', + content: "this is a test request, write a short poem" + }, + ] + }, { + headers: { + 'x-litellm-spend-logs-metadata': '{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}' + } + }); + console.log(response); + } catch (error) { + console.log("got this exception from server"); + console.error(error); + } +} + +// Call the asynchronous function +runOpenAI(); +``` + @@ -502,6 +566,29 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ } }' ``` + + + + + +Pass `x-litellm-spend-logs-metadata` as a request header with JSON string + +```shell +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'x-litellm-spend-logs-metadata: {"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}' \ + --data '{ + "model": "gpt-3.5-turbo", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] +}' +``` + diff --git a/docs/my-website/docs/proxy/load_balancing.md b/docs/my-website/docs/proxy/load_balancing.md index fd95b57c1ba..bcbc4e93651 100644 --- a/docs/my-website/docs/proxy/load_balancing.md +++ b/docs/my-website/docs/proxy/load_balancing.md @@ -13,6 +13,23 @@ For more details on routing strategies / params, see [Routing](../routing.md) ::: +## How Load Balancing Works + +LiteLLM automatically distributes requests across multiple deployments of the same model using its built-in router. the proxy routes traffic to optimize performance and reliability. + +"simple-shuffle" routing strategy is used by default + +### Routing Strategies + +| Strategy | Description | When to Use | +|----------|-------------|-------------| +| **simple-shuffle** (recommended) | Randomly distributes requests | General purpose, good for even load distribution | +| **least-busy** | Routes to deployment with fewest active requests | High concurrency scenarios | +| **usage-based-routing** (bad for perf) | Routes to deployment with lowest current usage (RPM/TPM) | When you want to respect rate limits evenly | +| **latency-based-routing** | Routes to fastest responding deployment | Latency-critical applications | +| **cost-based-routing** | Routes to deployment with lowest cost | Cost-sensitive applications | + + ## Quick Start - Load Balancing #### Step 1 - Set deployments on config @@ -106,49 +123,14 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ ] }' ``` - - - -```python -from langchain.chat_models import ChatOpenAI -from langchain.prompts.chat import ( - ChatPromptTemplate, - HumanMessagePromptTemplate, - SystemMessagePromptTemplate, -) -from langchain.schema import HumanMessage, SystemMessage -import os - -os.environ["OPENAI_API_KEY"] = "anything" - -chat = ChatOpenAI( - openai_api_base="http://0.0.0.0:4000", - model="gpt-3.5-turbo", -) - -messages = [ - SystemMessage( - content="You are a helpful assistant that im using to make a test request to." - ), - HumanMessage( - content="test from litellm. tell me why it's amazing in 1 sentence" - ), -] -response = chat(messages) - -print(response) -``` - - - ### Test - Loadbalancing In this request, the following will occur: 1. A rate limit exception will be raised -2. LiteLLM proxy will retry the request on the model group (default is 3). +2. LiteLLM proxy will retry the request on the model group (default retries are 3). ```bash curl -X POST 'http://0.0.0.0:4000/chat/completions' \ @@ -256,4 +238,16 @@ model_group_alias: Optional[Dict[str, Union[str, RouterModelGroupAliasItem]]] = class RouterModelGroupAliasItem(TypedDict): model: str hidden: bool # if 'True', don't return on `/v1/models`, `/v1/model/info`, `/v1/model_group/info` -``` \ No newline at end of file +``` + +### When You'll See Load Balancing in Action + +**Immediate Effects:** + +- Different deployments serve subsequent requests (visible in logs) +- Better response times during high traffic + +**Observable Benefits:** +- **Higher throughput**: More requests handled simultaneously across deployments +- **Improved reliability**: If one deployment fails, traffic automatically routes to healthy ones +- **Better resource utilization**: Load spread evenly across all available deployments diff --git a/docs/my-website/docs/proxy/prometheus.md b/docs/my-website/docs/proxy/prometheus.md index dc7030949bd..8bbf737540d 100644 --- a/docs/my-website/docs/proxy/prometheus.md +++ b/docs/my-website/docs/proxy/prometheus.md @@ -63,7 +63,7 @@ Use this for for tracking per [user, key, team, etc.](virtual_keys) | Metric Name | Description | |----------------------|--------------------------------------| -| `litellm_spend_metric` | Total Spend, per `"user", "key", "model", "team", "end-user"` | +| `litellm_spend_metric` | Total Spend, per `"end_user", "hashed_api_key", "api_key_alias", "model", "team", "team_alias", "user"` | | `litellm_total_tokens_metric` | input + output tokens per `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model"` | | `litellm_input_tokens_metric` | input tokens per `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model"` | | `litellm_output_tokens_metric` | output tokens per `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model"` | @@ -73,9 +73,9 @@ Use this for for tracking per [user, key, team, etc.](virtual_keys) | Metric Name | Description | |----------------------|--------------------------------------| -| `litellm_team_max_budget_metric` | Max Budget for Team Labels: `"team_id", "team_alias"`| -| `litellm_remaining_team_budget_metric` | Remaining Budget for Team (A team created on LiteLLM) Labels: `"team_id", "team_alias"`| -| `litellm_team_budget_remaining_hours_metric` | Hours before the team budget is reset Labels: `"team_id", "team_alias"`| +| `litellm_team_max_budget_metric` | Max Budget for Team Labels: `"team", "team_alias"`| +| `litellm_remaining_team_budget_metric` | Remaining Budget for Team (A team created on LiteLLM) Labels: `"team", "team_alias"`| +| `litellm_team_budget_remaining_hours_metric` | Hours before the team budget is reset Labels: `"team", "team_alias"`| ### Virtual Key - Budget @@ -119,8 +119,8 @@ Use this to track overall LiteLLM Proxy usage. | Metric Name | Description | |----------------------|--------------------------------------| -| `litellm_proxy_failed_requests_metric` | Total number of failed responses from proxy - the client did not get a success response from litellm proxy. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "exception_status", "exception_class"` | -| `litellm_proxy_total_requests_metric` | Total number of requests made to the proxy server - track number of client side requests. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "status_code"` | +| `litellm_proxy_failed_requests_metric` | Total number of failed responses from proxy - the client did not get a success response from litellm proxy. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "exception_status", "exception_class", "route"` | +| `litellm_proxy_total_requests_metric` | Total number of requests made to the proxy server - track number of client side requests. Labels: `"end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "status_code", "user_email", "route"` | ## LLM Provider Metrics @@ -155,7 +155,7 @@ Use this for LLM API Error monitoring and tracking remaining rate limits and tok | Metric Name | Description | |----------------------|--------------------------------------| | `litellm_remaining_requests_metric` | Track `x-ratelimit-remaining-requests` returned from LLM API Deployment. Labels: `"model_group", "api_provider", "api_base", "litellm_model_name", "hashed_api_key", "api_key_alias"` | -| `litellm_remaining_tokens` | Track `x-ratelimit-remaining-tokens` return from LLM API Deployment. Labels: `"model_group", "api_provider", "api_base", "litellm_model_name", "hashed_api_key", "api_key_alias"` | +| `litellm_remaining_tokens_metric` | Track `x-ratelimit-remaining-tokens` return from LLM API Deployment. Labels: `"model_group", "api_provider", "api_base", "litellm_model_name", "hashed_api_key", "api_key_alias"` | ### Deployment State | Metric Name | Description | @@ -167,16 +167,22 @@ Use this for LLM API Error monitoring and tracking remaining rate limits and tok | Metric Name | Description | |----------------------|--------------------------------------| -| `litellm_deployment_cooled_down` | Number of times a deployment has been cooled down by LiteLLM load balancing logic. Labels: `"litellm_model_name", "model_id", "api_base", "api_provider", "exception_status"` | +| `litellm_deployment_cooled_down` | Number of times a deployment has been cooled down by LiteLLM load balancing logic. Labels: `"litellm_model_name", "model_id", "api_base", "api_provider"` | | `litellm_deployment_successful_fallbacks` | Number of successful fallback requests from primary model -> fallback model. Labels: `"requested_model", "fallback_model", "hashed_api_key", "api_key_alias", "team", "team_alias", "exception_status", "exception_class"` | | `litellm_deployment_failed_fallbacks` | Number of failed fallback requests from primary model -> fallback model. Labels: `"requested_model", "fallback_model", "hashed_api_key", "api_key_alias", "team", "team_alias", "exception_status", "exception_class"` | +## Request Counting Metrics + +| Metric Name | Description | +|----------------------|--------------------------------------| +| `litellm_requests_metric` | Total number of requests tracked per endpoint. Labels: `"end_user", "hashed_api_key", "api_key_alias", "model", "team", "team_alias", "user", "user_email"` | + ## Request Latency Metrics | Metric Name | Description | |----------------------|--------------------------------------| | `litellm_request_total_latency_metric` | Total latency (seconds) for a request to LiteLLM Proxy Server - tracked for labels "end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model" | -| `litellm_overhead_latency_metric` | Latency overhead (seconds) added by LiteLLM processing - tracked for labels "end_user", "hashed_api_key", "api_key_alias", "requested_model", "team", "team_alias", "user", "model" | +| `litellm_overhead_latency_metric` | Latency overhead (seconds) added by LiteLLM processing - tracked for labels "model_group", "api_provider", "api_base", "litellm_model_name", "hashed_api_key", "api_key_alias" | | `litellm_llm_api_latency_metric` | Latency (seconds) for just the LLM API call - tracked for labels "model", "hashed_api_key", "api_key_alias", "team", "team_alias", "requested_model", "end_user", "user" | | `litellm_llm_api_time_to_first_token_metric` | Time to first token for LLM API call - tracked for labels `model`, `hashed_api_key`, `api_key_alias`, `team`, `team_alias` [Note: only emitted for streaming requests] | @@ -486,7 +492,6 @@ Here is a screenshot of the metrics you can monitor with the LiteLLM Grafana Das | Metric Name | Description | |----------------------|--------------------------------------| | `litellm_llm_api_failed_requests_metric` | **deprecated** use `litellm_proxy_failed_requests_metric` | -| `litellm_requests_metric` | **deprecated** use `litellm_proxy_total_requests_metric` | diff --git a/docs/my-website/docs/proxy/request_headers.md b/docs/my-website/docs/proxy/request_headers.md index c250d42f7bb..eea66e5fa93 100644 --- a/docs/my-website/docs/proxy/request_headers.md +++ b/docs/my-website/docs/proxy/request_headers.md @@ -14,6 +14,8 @@ Special headers that are supported by LiteLLM. `x-litellm-num-retries`: Optional[int]: The number of retries for the request. +`x-litellm-spend-logs-metadata`: Optional[str]: JSON string containing custom metadata to include in spend logs. Example: `{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion"}`. [Learn More](../proxy/enterprise#tracking-spend-with-custom-metadata) + ## Anthropic Headers `anthropic-version` Optional[str]: The version of the Anthropic API to use. diff --git a/docs/my-website/docs/proxy/user_management_heirarchy.md b/docs/my-website/docs/proxy/user_management_heirarchy.md index 3565c9d257d..cb5cc0dd7a2 100644 --- a/docs/my-website/docs/proxy/user_management_heirarchy.md +++ b/docs/my-website/docs/proxy/user_management_heirarchy.md @@ -9,5 +9,5 @@ LiteLLM supports a hierarchy of users, teams, organizations, and budgets. - Organizations can have multiple teams. [API Reference](https://litellm-api.up.railway.app/#/organization%20management) - Teams can have multiple users. [API Reference](https://litellm-api.up.railway.app/#/team%20management) -- Users can have multiple keys. [API Reference](https://litellm-api.up.railway.app/#/budget%20management) +- Users can have multiple keys, and be on multiple teams. [API Reference](https://litellm-api.up.railway.app/#/budget%20management) - Keys can belong to either a team or a user. [API Reference](https://litellm-api.up.railway.app/#/end-user%20management) diff --git a/docs/my-website/docs/proxy/veo_video_generation.md b/docs/my-website/docs/proxy/veo_video_generation.md new file mode 100644 index 00000000000..14c263bf847 --- /dev/null +++ b/docs/my-website/docs/proxy/veo_video_generation.md @@ -0,0 +1,163 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Veo Video Generation with Google AI Studio + +Generate videos using Google's Veo model through LiteLLM's pass-through endpoints. + +## Quick Start + +LiteLLM allows you to use Google AI Studio's Veo video generation API through pass-through routes with zero configuration. + +### 1. Add Google AI Studio API Key to your environment + +```bash +export GEMINI_API_KEY="your_google_ai_studio_api_key" +``` + +### 2. Start LiteLLM Proxy + +```bash +litellm + +# RUNNING on http://0.0.0.0:4000 +``` + +### 3. Generate Video + + + + +```python +import requests +import time +import json + +# Configuration +BASE_URL = "http://localhost:4000/gemini/v1beta" +API_KEY = "anything" # Use "anything" as the key + +headers = { + "x-goog-api-key": API_KEY, + "Content-Type": "application/json" +} + +# Step 1: Initiate video generation +def generate_video(prompt): + url = f"{BASE_URL}/models/veo-3.0-generate-preview:predictLongRunning" + payload = { + "instances": [{ + "prompt": prompt + }] + } + + response = requests.post(url, headers=headers, json=payload) + response.raise_for_status() + + data = response.json() + return data.get("name") # Operation name + +# Step 2: Poll for completion +def wait_for_completion(operation_name): + operation_url = f"{BASE_URL}/{operation_name}" + + while True: + response = requests.get(operation_url, headers=headers) + response.raise_for_status() + + data = response.json() + + if data.get("done", False): + # Extract video URI + video_uri = data["response"]["generateVideoResponse"]["generatedSamples"][0]["video"]["uri"] + return video_uri + + time.sleep(10) # Wait 10 seconds before next poll + +# Step 3: Download video +def download_video(video_uri, filename="generated_video.mp4"): + # Replace Google URL with LiteLLM proxy URL + litellm_url = video_uri.replace( + "https://generativelanguage.googleapis.com/v1beta", + BASE_URL + ) + + response = requests.get(litellm_url, headers=headers, stream=True) + response.raise_for_status() + + with open(filename, 'wb') as f: + for chunk in response.iter_content(chunk_size=8192): + if chunk: + f.write(chunk) + + return filename + +# Complete workflow +prompt = "A cat playing with a ball of yarn in a sunny garden" + +print("Generating video...") +operation_name = generate_video(prompt) + +print("Waiting for completion...") +video_uri = wait_for_completion(operation_name) + +print("Downloading video...") +filename = download_video(video_uri) + +print(f"Video saved as: {filename}") +``` + + + + + +```bash +# Step 1: Initiate video generation +curl -X POST "http://localhost:4000/gemini/v1beta/models/veo-3.0-generate-preview:predictLongRunning" \ + -H "x-goog-api-key: anything" \ + -H "Content-Type: application/json" \ + -d '{ + "instances": [{ + "prompt": "A cat playing with a ball of yarn in a sunny garden" + }] + }' + +# Response will include operation name: +# {"name": "operations/generate_12345"} + +# Step 2: Poll for completion +curl -X GET "http://localhost:4000/gemini/v1beta/operations/generate_12345" \ + -H "x-goog-api-key: anything" + +# Step 3: Download video (when done=true) +curl -X GET "http://localhost:4000/gemini/v1beta/files/VIDEO_ID:download?alt=media" \ + -H "x-goog-api-key: anything" \ + --output generated_video.mp4 +``` + + + + +## Complete Example + +For a full working example with error handling and logging, see our [Veo Video Generation Cookbook](https://github.com/BerriAI/litellm/blob/main/cookbook/veo_video_generation.py). + +## How It Works + +1. **Video Generation Request**: Send a prompt to Veo's `predictLongRunning` endpoint +2. **Operation Polling**: Monitor the long-running operation until completion +3. **File Download**: Download the generated video through LiteLLM's pass-through with automatic redirect handling + +LiteLLM handles: +- āœ… Authentication with Google AI Studio +- āœ… Request routing and proxying +- āœ… Automatic redirect handling for file downloads + +## Configuration Options + +### Environment Variables + +```bash +export GEMINI_API_KEY="your_google_ai_studio_api_key" +``` + diff --git a/docs/my-website/docs/reasoning_content.md b/docs/my-website/docs/reasoning_content.md index f9cab01639d..12db17325d4 100644 --- a/docs/my-website/docs/reasoning_content.md +++ b/docs/my-website/docs/reasoning_content.md @@ -12,7 +12,7 @@ Requires LiteLLM v1.63.0+ Supported Providers: - Deepseek (`deepseek/`) - Anthropic API (`anthropic/`) -- Bedrock (Anthropic + Deepseek) (`bedrock/`) +- Bedrock (Anthropic + Deepseek + GPT-OSS) (`bedrock/`) - Vertex AI (Anthropic) (`vertexai/`) - OpenRouter (`openrouter/`) - XAI (`xai/`) @@ -20,6 +20,7 @@ Supported Providers: - Vertex AI (`vertex_ai/`) - Perplexity (`perplexity/`) - Mistral AI (Magistral models) (`mistral/`) +- Groq (`groq/`) LiteLLM will standardize the `reasoning_content` in the response and `thinking_blocks` in the assistant message. diff --git a/docs/my-website/docs/tutorials/claude_responses_api.md b/docs/my-website/docs/tutorials/claude_responses_api.md index 09b352a7663..5000161a520 100644 --- a/docs/my-website/docs/tutorials/claude_responses_api.md +++ b/docs/my-website/docs/tutorials/claude_responses_api.md @@ -12,6 +12,12 @@ This tutorial is based on [Anthropic's official LiteLLM configuration documentat ::: +
+ +### Video Walkthrough + + + ## Prerequisites - [Claude Code](https://docs.anthropic.com/en/docs/claude-code/overview) installed @@ -83,11 +89,17 @@ curl -X POST http://0.0.0.0:4000/v1/messages \ Configure Claude Code to use LiteLLM's unified endpoint: +Either a virtual key / master key can be used here + ```bash export ANTHROPIC_BASE_URL="http://0.0.0.0:4000" export ANTHROPIC_AUTH_TOKEN="$LITELLM_MASTER_KEY" ``` +:::tip +LITELLM_MASTER_KEY gives claude access to all proxy models, whereas a virtual key would be limited to the models set in UI +::: + #### Method 2: Provider-specific Pass-through Endpoint Alternatively, use the Anthropic pass-through endpoint: diff --git a/docs/my-website/img/admin_settings_ui_theme.png b/docs/my-website/img/admin_settings_ui_theme.png new file mode 100644 index 00000000000..81e6d761e17 Binary files /dev/null and b/docs/my-website/img/admin_settings_ui_theme.png differ diff --git a/docs/my-website/img/admin_settings_ui_theme_logo.png b/docs/my-website/img/admin_settings_ui_theme_logo.png new file mode 100644 index 00000000000..38f36e61602 Binary files /dev/null and b/docs/my-website/img/admin_settings_ui_theme_logo.png differ diff --git a/docs/my-website/release_notes/v1.74.7/index.md b/docs/my-website/release_notes/v1.74.7/index.md index e3a2ac0aa00..7d7a568e13f 100644 --- a/docs/my-website/release_notes/v1.74.7/index.md +++ b/docs/my-website/release_notes/v1.74.7/index.md @@ -148,7 +148,6 @@ Starting with this release, you can run health endpoints on an isolated process - New provider integration for v0.dev - [PR #12751](https://github.com/BerriAI/litellm/pull/12751), [Get Started](../../docs/providers/v0) - **[OpenAI](../../docs/providers/openai)** - Use OpenAI DeepResearch models with `litellm.completion` (`/chat/completions`) - [PR #12627](https://github.com/BerriAI/litellm/pull/12627) **DOC NEEDED** - - Add `input_fidelity` parameter for OpenAI image generation - [PR #12662](https://github.com/BerriAI/litellm/pull/12662), [Get Started](../../docs/image_generation) - **[Azure OpenAI](../../docs/providers/azure_openai)** - Use Azure OpenAI DeepResearch models with `litellm.completion` (`/chat/completions`) - [PR #12627](https://github.com/BerriAI/litellm/pull/12627) **DOC NEEDED** - Added `response_format` support for openai gpt-4.1 models - [PR #12745](https://github.com/BerriAI/litellm/pull/12745) diff --git a/docs/my-website/release_notes/v1.76.0-stable/index.md b/docs/my-website/release_notes/v1.76.0-stable/index.md new file mode 100644 index 00000000000..660c8cbcf02 --- /dev/null +++ b/docs/my-website/release_notes/v1.76.0-stable/index.md @@ -0,0 +1,189 @@ +--- +title: "[PRE-RELEASE]v1.76.0-stable - RPS Improvements" +slug: "v1-76-0" +date: 2025-08-23T10:00:00 +authors: + - name: Krrish Dholakia + title: CEO, LiteLLM + url: https://www.linkedin.com/in/krish-d/ + image_url: https://pbs.twimg.com/profile_images/1298587542745358340/DZv3Oj-h_400x400.jpg + - name: Ishaan Jaffer + title: CTO, LiteLLM + url: https://www.linkedin.com/in/reffajnaahsi/ + image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg + +hide_table_of_contents: false +--- + +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +:::info + +LiteLLM is hiring a **Founding Backend Engineer**, in San Francisco. + +[Apply here](https://www.ycombinator.com/companies/litellm/jobs/6uvoBp3-founding-backend-engineer) if you're interested! +::: + + + + + +## Deploy this version + +:::info + +This release is not live yet. +::: + + +--- + +## New Models / Updated Models + +#### Bugs +- **[OpenAI](../../docs/providers/openai)** + - Gpt-5 chat: clarify does not support function calling [PR #13612](https://github.com/BerriAI/litellm/pull/13612), s/o Ā @[superpoussin22](https://github.com/superpoussin22) +- **[VertexAI](../../docs/providers/vertex)** + - fix vertexai batch file format byĀ @[thiagosalvatore](https://github.com/thiagosalvatore)Ā inĀ [PR #13576](https://github.com/BerriAI/litellm/pull/13576) +- **[LiteLLM Proxy](../../docs/providers/litellm_proxy)** + - Add support for calling image_edits + image_generations via SDK to Proxy - [PR #13735](https://github.com/BerriAI/litellm/pull/13735) +- **[OpenRouter](../../docs/providers/openrouter)** + - Fix max_output_tokens value for anthropic Claude 4 - [PR #13526](https://github.com/BerriAI/litellm/pull/13526) +- **[Gemini](../../docs/providers/gemini)** + - Fix prompt caching cost calculation - [PR #13742](https://github.com/BerriAI/litellm/pull/13742) +- **[Azure](../../docs/providers/azure)** + - Support `../openai/v1/respones` api base - [PR #13526](https://github.com/BerriAI/litellm/pull/13526) + - Fix azure/gpt-5-chat max_input_tokens - [PR #13660](https://github.com/BerriAI/litellm/pull/13660) +- **[Groq](../../docs/providers/groq)** + - streaming ASCII encoding issue - [PR #13675](https://github.com/BerriAI/litellm/pull/13675) +- **[Baseten](../../docs/providers/baseten)** + - Refactored integration to use new openai-compatible endpoints - [PR #13783](https://github.com/BerriAI/litellm/pull/13783) +- **[Bedrock](../../docs/providers/bedrock)** + - fix application inference profile for pass-through endpoints for bedrock - [PR #13881](https://github.com/BerriAI/litellm/pull/13881) +- **[DataRobot](../../docs/providers/datarobot)** + - Updated URL handling for DataRobot provider URL - [PR #13880](https://github.com/BerriAI/litellm/pull/13880) + +#### Features +- **[Together AI](../../docs/providers/together)** + - Added Qwen3, Deepseek R1 0528 Throughput, GLM 4.5 and GPT-OSS models cost tracking - [PR #13637](https://github.com/BerriAI/litellm/pull/13637), s/o Ā @[Tasmay-Tibrewal](https://github.com/Tasmay-Tibrewal) +- **[Fireworks AI](../../docs/providers/fireworks_ai)** + - add fireworks_ai/accounts/fireworks/models/deepseek-v3-0324 - [PR #13821](https://github.com/BerriAI/litellm/pull/13821) +- **[VertexAI](../../docs/providers/vertex)** + - Add VertexAI qwen API Service - [PR #13828](https://github.com/BerriAI/litellm/pull/13828) + - Add new VertexAI image modelsĀ vertex_ai/imagen-4.0-generate-001,Ā vertex_ai/imagen-4.0-ultra-generate-001,Ā vertex_ai/imagen-4.0-fast-generate-001Ā  - [PR #13874](https://github.com/BerriAI/litellm/pull/13874) +- **[Anthropic](../../docs/providers/anthropic)** + - Add long context support w/ cost tracking - [PR #13759](https://github.com/BerriAI/litellm/pull/13759) +- **[DeepInfra](../../docs/providers/deepinfra)** + - Add rerank endpoint support for deepinfra - [PR #13820](https://github.com/BerriAI/litellm/pull/13820) + - Add new models for cost tracking - [PR #13883](https://github.com/BerriAI/litellm/pull/13883), s/o Ā @[Toy-97](https://github.com/Toy-97) +- **[Bedrock](../../docs/providers/bedrock)** + - Add tool prompt caching on async calls - [PR #13803](https://github.com/BerriAI/litellm/pull/13803), s/o Ā @[UlookEE](https://github.com/UlookEE) + - role chaining and session name with webauthentication for aws bedrock - [PR #13753](https://github.com/BerriAI/litellm/pull/13753), s/o @[RichardoC](https://github.com/RichardoC) +- **[Ollama](../../docs/providers/ollama)** + - Handle Ollama null response when using tool calling with non-tool trained models - [PR #13902](https://github.com/BerriAI/litellm/pull/13902) +- **[OpenRouter](../../docs/providers/openrouter)** + - Add deepseek/deepseek-chat-v3.1 support - [PR #13897](https://github.com/BerriAI/litellm/pull/13897) +- **[Mistral](../../docs/providers/mistral)** + - Add support for calling mistral files via chat completions - [PR #13866](https://github.com/BerriAI/litellm/pull/13866), s/o Ā @[jinskjoy](https://github.com/jinskjoy) + - Handle empty assistant content - [PR #13671](https://github.com/BerriAI/litellm/pull/13671) + - Support new ā€˜thinking’ response block - [PR #13671](https://github.com/BerriAI/litellm/pull/13671) +- **[Databricks](../../docs/providers/databricks)** + - remove deprecated dbrx models (dbrx-instruct, llama 3.1) - [PR #13843](https://github.com/BerriAI/litellm/pull/13843) +- **[AI/ML API](../../docs/providers/ai_ml_api)** + - Image gen api support - [PR #13893](https://github.com/BerriAI/litellm/pull/13893) + + +## LLM API Endpoints +#### Bugs +- **[Responses API](../../docs/response_api)** + - add default api version for openai responses api calls - [PR #13526](https://github.com/BerriAI/litellm/pull/13526) + - supportĀ allowed_openai_params - [PR #13671](https://github.com/BerriAI/litellm/pull/13671) + + +## MCP Gateway +#### Bugs +- fix StreamableHTTPSessionManager .run() error - [PR #13666](https://github.com/BerriAI/litellm/pull/13666) + +## Vector Stores +#### Bugs +- **[Bedrock](../../docs/providers/bedrock)** + - Using LiteLLM Managed Credentials for Query - [PR #13787](https://github.com/BerriAI/litellm/pull/13787) + +## Management Endpoints / UI +#### Bugs +- **[Passthrough](../../docs/pass_through/intro)** + - Fix query passthrough deletion - [PR #13622](https://github.com/BerriAI/litellm/pull/13622) + +#### Features +- **Models** + - Add Search Functionality for Public Model Names in Model Dashboard - [PR #13687](https://github.com/BerriAI/litellm/pull/13687) + - Auto-Add `azure/` to deployment Name in UI - [PR #13685](https://github.com/BerriAI/litellm/pull/13685) + - Models page row UI restructure - [PR #13771](https://github.com/BerriAI/litellm/pull/13771) +- **Notifications** + - Add new notifications toast UI everywhere - [PR #13813](https://github.com/BerriAI/litellm/pull/13813) +- **Keys** + - Fix key edit settings after regenerating a key - [PR #13815](https://github.com/BerriAI/litellm/pull/13815) + - Require team_id when creating service account keys - [PR #13873](https://github.com/BerriAI/litellm/pull/13873) + - Filter - show all options on filter option click - [PR #13858](https://github.com/BerriAI/litellm/pull/13858) +- **Usage** + - Fix ā€˜Cannot read properties of undefined’ exception on user agent activity tab - [PR #13892](https://github.com/BerriAI/litellm/pull/13892) +- **SSO** + - Free SSO usage for up to 5 users - [PR #13843](https://github.com/BerriAI/litellm/pull/13843) + +## Logging / Guardrail Integrations +#### Bugs +- **[Bedrock Guardrails](../../docs/proxy/guardrails/bedrock)** + - Add bedrock api key support - [PR #13835](https://github.com/BerriAI/litellm/pull/13835) +#### Features +- **[Datadog LLM Observability](../../docs/integrations/datadog)** + - Add support for Failure LoggingĀ [PR #13726](https://github.com/BerriAI/litellm/pull/13726) + - Add time to first token, litellm overhead, guardrail overhead latency metrics - [PR #13734](https://github.com/BerriAI/litellm/pull/13734) + - Add support for tracing guardrail input/output - [PR #13767](https://github.com/BerriAI/litellm/pull/13767) +- **[Langfuse OTEL](../../docs/integrations/langfuse)** + - Allow using Key/Team Based Logging - [PR #13791](https://github.com/BerriAI/litellm/pull/13791) +- **[AIM](../../docs/integrations/aim)** + - Migrate to new firewall API - [PR #13748](https://github.com/BerriAI/litellm/pull/13748) +- **[OTEL](../../docs/observability/opentelemetry_integration)** + - Add OTEL tracing for actual LLM API call - [PR #13836](https://github.com/BerriAI/litellm/pull/13836) +- **[MLFlow](../../docs/observability/mlflow_integration)** + - Include predicted output in MLflow tracing - [PR #13795](https://github.com/BerriAI/litellm/pull/13795), s/oĀ @TomeHirataĀ  + + +## Performance / Loadbalancing / Reliability improvements +#### Bugs +- **[Cooldowns](../../docs/routing#how-cooldowns-work)** + - don't return raw Azure Exceptions to client (can contain prompt leakage) - [PR #13529](https://github.com/BerriAI/litellm/pull/13529) +- **[Auto-router](../../docs/proxy/auto_routing)** + - Ensures the relevant dependencies for auto router existing on LiteLLM Docker - [PR #13788](https://github.com/BerriAI/litellm/pull/13788) +- **Model Alias** + - Fix calling key with access to model alias - [PR #13830](https://github.com/BerriAI/litellm/pull/13830) + +#### Features +- **[S3 Caching](../../docs/proxy/caching)** + - Use namespace as prefix for s3 cache - [PR #13704](https://github.com/BerriAI/litellm/pull/13704) + - Async S3 Caching support (4x RPS improvement) - [PR #13852](https://github.com/BerriAI/litellm/pull/13852), s/o @[michal-otmianowski](https://github.com/michal-otmianowski) +- **Model Group header forwarding** + - reuse same logic as global header forwarding - [PR #13741](https://github.com/BerriAI/litellm/pull/13741) + - add support for hosted_vllm on UI - [PR #13885](https://github.com/BerriAI/litellm/pull/13885) +- **Performance** + - Improve LiteLLM Python SDK RPS by +200 RPS (braintrust import + aiohttp transport fixes) - [PR #13839](https://github.com/BerriAI/litellm/pull/13839) + - Use O(1) Set lookups for model routing - [PR #13879](https://github.com/BerriAI/litellm/pull/13879) + - Reduce Significant CPU overhead from litellm_logging.py - [PR #13895](https://github.com/BerriAI/litellm/pull/13895) + - Improvements for Async Success Handler (Logging Callbacks) - Approx +130 RPS - [PR #13905](https://github.com/BerriAI/litellm/pull/13905) + + +## General Proxy Improvements +#### Bugs + +- **SDK** + - Fix litellm compatibility with newest release of openAI (>v1.100.0) - [PR #13728](https://github.com/BerriAI/litellm/pull/13728) +- **Helm** + - Add possibility to configure resources for migrations-job - [PR #13617](https://github.com/BerriAI/litellm/pull/13617) + - Ensure Helm chart auto generated master keys follow sk-xxxx format - [PR #13871](https://github.com/BerriAI/litellm/pull/13871) + - Enhance database configuration: add support for optional endpointKey - [PR #13763](https://github.com/BerriAI/litellm/pull/13763) +- **Rate Limits** + - fixing descriptor/response size mismatch on parallel_request_limiter_v3 - [PR #13863](https://github.com/BerriAI/litellm/pull/13863), s/o Ā @[luizrennocosta](https://github.com/luizrennocosta) +- **Non-root** + - fix permission access on prisma migrate in non-root image - [PR #13848](https://github.com/BerriAI/litellm/pull/13848), s/o @[Ithanil](https://github.com/Ithanil) \ No newline at end of file diff --git a/docs/my-website/release_notes/v1.76.3-stable/index.md b/docs/my-website/release_notes/v1.76.3-stable/index.md new file mode 100644 index 00000000000..0f997d6941a --- /dev/null +++ b/docs/my-website/release_notes/v1.76.3-stable/index.md @@ -0,0 +1,282 @@ +--- +title: "v1.76.3-stable - Performance, Video Generation & CloudZero Integration" +slug: "v1-76-3" +date: 2025-09-06T10:00:00 +authors: + - name: Krrish Dholakia + title: CEO, LiteLLM + url: https://www.linkedin.com/in/krish-d/ + image_url: https://pbs.twimg.com/profile_images/1298587542745358340/DZv3Oj-h_400x400.jpg + - name: Ishaan Jaffer + title: CTO, LiteLLM + url: https://www.linkedin.com/in/reffajnaahsi/ + image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg + +hide_table_of_contents: false +--- + +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +## Deploy this version + + + + +``` showLineNumbers title="docker run litellm" +docker run \ +-e STORE_MODEL_IN_DB=True \ +-p 4000:4000 \ +ghcr.io/berriai/litellm:v1.76.3 +``` + + + + +``` showLineNumbers title="pip install litellm" +pip install litellm==1.76.3 +``` + + + + +--- + +## Key Highlights + +- **Major Performance Improvements** +400 RPS when using correct amount of workers + CPU cores combination +- **Video Generation Support** - Added Google AI Studio and Vertex AI Veo Video Generation through LiteLLM Pass through routes +- **CloudZero Integration** - New cost tracking integration for exporting LiteLLM Usage and Spend data to CloudZero. + +## Major Changes +- **Performance Optimization**: LiteLLM Proxy now achieves +400 RPS when using correct amount of CPU cores - [PR #14153](https://github.com/BerriAI/litellm/pull/14153), [PR #14242](https://github.com/BerriAI/litellm/pull/14242) + + By default, LiteLLM will now use `num_workers = os.cpu_count()` to achieve optimal performance. + + **Override Options:** + + Set environment variable: + ```bash + DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 + ``` + + Or start LiteLLM Proxy with: + ```bash + litellm --num_workers 1 + ``` + +- **Security Fix**: Fixed memory_usage_in_mem_cache cache endpoint vulnerability - [PR #14229](https://github.com/BerriAI/litellm/pull/14229) + +--- + +## Performance Improvements + +This release includes significant performance optimizations. On our internal benchmarks we saw 1 instance get +400 RPS when using correct amount of workers + CPU cores combination. + +- **+400 RPS Performance Boost** - LiteLLM Proxy now uses correct amount of CPU cores for optimal performance - [PR #14153](https://github.com/BerriAI/litellm/pull/14153) +- **Default CPU Workers** - Changed DEFAULT_NUM_WORKERS_LITELLM_PROXY default to number of CPUs - [PR #14242](https://github.com/BerriAI/litellm/pull/14242) + + +--- + +## New Models / Updated Models + +#### New Model Support + +| Provider | Model | Context Window | Input ($/1M tokens) | Output ($/1M tokens) | Features | +| ----------- | -------------------------------------- | -------------- | ------------------- | -------------------- | -------- | +| OpenRouter | `openrouter/openai/gpt-4.1` | 1M | $2.00 | $8.00 | Chat completions with vision | +| OpenRouter | `openrouter/openai/gpt-4.1-mini` | 1M | $0.40 | $1.60 | Efficient chat completions | +| OpenRouter | `openrouter/openai/gpt-4.1-nano` | 1M | $0.10 | $0.40 | Ultra-efficient chat | +| Vertex AI | `vertex_ai/openai/gpt-oss-20b-maas` | 131K | $0.075 | $0.30 | Reasoning support | +| Vertex AI | `vertex_ai/openai/gpt-oss-120b-maas` | 131K | $0.15 | $0.60 | Advanced reasoning | +| Gemini | `gemini/veo-3.0-generate-preview` | 1K | - | $0.75/sec | Video generation | +| Gemini | `gemini/veo-3.0-fast-generate-preview` | 1K | - | $0.40/sec | Fast video generation | +| Gemini | `gemini/veo-2.0-generate-001` | 1K | - | $0.35/sec | Video generation | +| Volcengine | `doubao-embedding-large` | 4K | Free | Free | 2048-dim embeddings | +| Together AI | `together_ai/deepseek-ai/DeepSeek-V3.1` | 128K | $0.60 | $1.70 | Reasoning support | + +#### Features + +- **[Google Gemini](../../docs/providers/gemini)** + - Added 'thoughtSignature' support via 'thinking_blocks' - [PR #14122](https://github.com/BerriAI/litellm/pull/14122) + - Added support for reasoning_effort='minimal' for Gemini models - [PR #14262](https://github.com/BerriAI/litellm/pull/14262) +- **[OpenRouter](../../docs/providers/openrouter)** + - Added GPT-4.1 model family - [PR #14101](https://github.com/BerriAI/litellm/pull/14101) +- **[Groq](../../docs/providers/groq)** + - Added support for reasoning_effort parameter - [PR #14207](https://github.com/BerriAI/litellm/pull/14207) +- **[X.AI](../../docs/providers/xai)** + - Fixed XAI cost calculation - [PR #14127](https://github.com/BerriAI/litellm/pull/14127) +- **[Vertex AI](../../docs/providers/vertex)** + - Added support for GPT-OSS models on Vertex AI - [PR #14184](https://github.com/BerriAI/litellm/pull/14184) + - Added additionalProperties to Vertex AI Schema definition - [PR #14252](https://github.com/BerriAI/litellm/pull/14252) +- **[VLLM](../../docs/providers/vllm)** + - Handle output parsing responses API output - [PR #14121](https://github.com/BerriAI/litellm/pull/14121) +- **[Ollama](../../docs/providers/ollama)** + - Added unified 'thinking' param support via `reasoning_content` - [PR #14121](https://github.com/BerriAI/litellm/pull/14121) +- **[Anthropic](../../docs/providers/anthropic)** + - Added supported text field to anthropic citation response - [PR #14126](https://github.com/BerriAI/litellm/pull/14126) +- **[OCI Provider](../../docs/providers/oci)** + - Handle assistant messages with both content and tool_calls - [PR #14171](https://github.com/BerriAI/litellm/pull/14171) +- **[Bedrock](../../docs/providers/bedrock)** + - Fixed structure output - [PR #14130](https://github.com/BerriAI/litellm/pull/14130) + - Added initial support for Bedrock Batches API - [PR #14190](https://github.com/BerriAI/litellm/pull/14190) +- **[Databricks](../../docs/providers/databricks)** + - Added support for anthropic citation API in Databricks - [PR #14077](https://github.com/BerriAI/litellm/pull/14077) + +### Bug Fixes +- **[Google Gemini (Google AI Studio + Vertex AI)](../../docs/providers/gemini)** + - Fixed Gemini 2.5 Pro schema validation with OpenAI-style type arrays in tools - [PR #14154](https://github.com/BerriAI/litellm/pull/14154) + - Fixed Gemini Tool Calling empty enum property - [PR #14155](https://github.com/BerriAI/litellm/pull/14155) + +#### New Provider Support + +- **[Volcengine](../../docs/providers/volcengine)** + - Added Volcengine embedding module with handler and transformation logic - [PR #14028](https://github.com/BerriAI/litellm/pull/14028) + +--- + +## LLM API Endpoints + +#### Features + +- **[Images API](../../docs/image_generation)** + - Added pass through image generation and image editing on OpenAI - [PR #14292](https://github.com/BerriAI/litellm/pull/14292) + - Support extra_body parameter for image generation - [PR #14211](https://github.com/BerriAI/litellm/pull/14211) +- **[Responses API](../../docs/response_api)** + - Fixed response API for reasoning item in input for litellm proxy - [PR #14200](https://github.com/BerriAI/litellm/pull/14200) + - Added structured output for SDK - [PR #14206](https://github.com/BerriAI/litellm/pull/14206) +- **[Bedrock Passthrough](../../docs/pass_through/bedrock)** + - Support AWS_BEDROCK_RUNTIME_ENDPOINT on bedrock passthrough - [PR #14156](https://github.com/BerriAI/litellm/pull/14156) +- **[Google AI Studio Passthrough](../../docs/pass_through/google_ai_studio)** + - Allow using Veo Video Generation through LiteLLM Pass through routes - [PR #14228](https://github.com/BerriAI/litellm/pull/14228) +- **General** + - Added support for safety_identifier parameter in chat.completions.create - [PR #14174](https://github.com/BerriAI/litellm/pull/14174) + - Fixed misclassified 500 error on invalid image_url in /chat/completions request - [PR #14149](https://github.com/BerriAI/litellm/pull/14149) + - Fixed token count error for Gemini CLI - [PR #14133](https://github.com/BerriAI/litellm/pull/14133) + +#### Bugs + +- **General** + - Remove "/" or ":" from model name when being used as h11 header name - [PR #14191](https://github.com/BerriAI/litellm/pull/14191) + - Bug fix for openai.gpt-oss when using reasoning_effort parameter - [PR #14300](https://github.com/BerriAI/litellm/pull/14300) + +--- + +## Spend Tracking, Budgets and Rate Limiting + +### Features + - Added header support for spend_logs_metadata - [PR #14186](https://github.com/BerriAI/litellm/pull/14186) + - Litellm passthrough cost tracking for chat completion - [PR #14256](https://github.com/BerriAI/litellm/pull/14256) + +### Bug Fixes + - Fixed TPM Rate Limit Bug - [PR #14237](https://github.com/BerriAI/litellm/pull/14237) + - Fixed Key Budget not resets at expectable times - [PR #14241](https://github.com/BerriAI/litellm/pull/14241) + + + +## Management Endpoints / UI + +#### Features + +- **UI Improvements** + - Logs page screen size fixed - [PR #14135](https://github.com/BerriAI/litellm/pull/14135) + - Create Organization Tooltip added on Success - [PR #14132](https://github.com/BerriAI/litellm/pull/14132) + - Back to Keys should say Back to Logs - [PR #14134](https://github.com/BerriAI/litellm/pull/14134) + - Add client side pagination on All Models table - [PR #14136](https://github.com/BerriAI/litellm/pull/14136) + - Model Filters UI improvement - [PR #14131](https://github.com/BerriAI/litellm/pull/14131) + - Remove table filter on user info page - [PR #14169](https://github.com/BerriAI/litellm/pull/14169) + - Team name badge added on the User Details - [PR #14003](https://github.com/BerriAI/litellm/pull/14003) + - Fix: Log page parameter passing error - [PR #14193](https://github.com/BerriAI/litellm/pull/14193) +- **Authentication & Authorization** + - Support for ES256/ES384/ES512 and EdDSA JWT verification - [PR #14118](https://github.com/BerriAI/litellm/pull/14118) + - Ensure `team_id` is a required field for generating service account keys - [PR #14270](https://github.com/BerriAI/litellm/pull/14270) + +#### Bugs + +- **General** + - Validate store model in db setting - [PR #14269](https://github.com/BerriAI/litellm/pull/14269) + +--- + +## Logging / Guardrail Integrations + +#### Features + +- **[Datadog](../../docs/proxy/logging#datadog)** + - Ensure `apm_id` is set on DD LLM Observability traces - [PR #14272](https://github.com/BerriAI/litellm/pull/14272) +- **[Braintrust](../../docs/proxy/logging#braintrust)** + - Fix logging when OTEL is enabled - [PR #14122](https://github.com/BerriAI/litellm/pull/14122) +- **[OTEL](../../docs/proxy/logging#otel)** + - Optional Metrics and Logs following semantic conventions - [PR #14179](https://github.com/BerriAI/litellm/pull/14179) +- **[Slack Alerting](../../docs/proxy/alerting)** + - Added alert type to alert message to slack for easier handling - [PR #14176](https://github.com/BerriAI/litellm/pull/14176) + +#### Guardrails + - Added guardrail to the Anthropic API endpoint - [PR #14107](https://github.com/BerriAI/litellm/pull/14107) + +#### New Integration + +- **[CloudZero](../../docs/proxy/cost_tracking)** + - LiteLLM x CloudZero Integration for Cost Tracking - [PR #14296](https://github.com/BerriAI/litellm/pull/14296) + +--- + +## Performance / Loadbalancing / Reliability improvements + +#### Features + +- **Performance** + - LiteLLM Proxy: +400 RPS when using correct amount of CPU cores - [PR #14153](https://github.com/BerriAI/litellm/pull/14153) + - Allow using `x-litellm-stream-timeout` header for stream timeout in requests - [PR #14147](https://github.com/BerriAI/litellm/pull/14147) + - Change DEFAULT_NUM_WORKERS_LITELLM_PROXY default to number CPUs - [PR #14242](https://github.com/BerriAI/litellm/pull/14242) +- **Monitoring** + - Added Prometheus missing metrics - [PR #14139](https://github.com/BerriAI/litellm/pull/14139) +- **Timeout** + - **Stream Timeout Control** - Allow using `x-litellm-stream-timeout` header for stream timeout in requests - [PR #14147](https://github.com/BerriAI/litellm/pull/14147) +- **Routing** + - Fixed x-litellm-tags not routing with Responses API - [PR #14289](https://github.com/BerriAI/litellm/pull/14289) + +#### Bugs + +- **Security** + - Fixed memory_usage_in_mem_cache cache endpoint vulnerability - [PR #14229](https://github.com/BerriAI/litellm/pull/14229) + +--- + +## General Proxy Improvements + +#### Features + +- **SCIM Support** + - Added better SCIM debugging - [PR #14221](https://github.com/BerriAI/litellm/pull/14221) + - Bug fixes for handling SCIM Group Memberships - [PR #14226](https://github.com/BerriAI/litellm/pull/14226) +- **Kubernetes** + - Added optional PodDisruptionBudget for litellm proxy - [PR #14093](https://github.com/BerriAI/litellm/pull/14093) +- **Error Handling** + - Add model to azure error message - [PR #14294](https://github.com/BerriAI/litellm/pull/14294) + +--- + +## New Contributors +* @iabhi4 made their first contribution in [PR #14093](https://github.com/BerriAI/litellm/pull/14093) +* @zainhas made their first contribution in [PR #14087](https://github.com/BerriAI/litellm/pull/14087) +* @LifeDJIK made their first contribution in [PR #14146](https://github.com/BerriAI/litellm/pull/14146) +* @retanoj made their first contribution in [PR #14133](https://github.com/BerriAI/litellm/pull/14133) +* @zhxlp made their first contribution in [PR #14193](https://github.com/BerriAI/litellm/pull/14193) +* @kayoch1n made their first contribution in [PR #14191](https://github.com/BerriAI/litellm/pull/14191) +* @kutsushitaneko made their first contribution in [PR #14171](https://github.com/BerriAI/litellm/pull/14171) +* @mjmendo made their first contribution in [PR #14176](https://github.com/BerriAI/litellm/pull/14176) +* @HarshavardhanK made their first contribution in [PR #14213](https://github.com/BerriAI/litellm/pull/14213) +* @eycjur made their first contribution in [PR #14207](https://github.com/BerriAI/litellm/pull/14207) +* @22mSqRi made their first contribution in [PR #14241](https://github.com/BerriAI/litellm/pull/14241) +* @onlylhf made their first contribution in [PR #14028](https://github.com/BerriAI/litellm/pull/14028) +* @btpemercier made their first contribution in [PR #11319](https://github.com/BerriAI/litellm/pull/11319) +* @tremlin made their first contribution in [PR #14287](https://github.com/BerriAI/litellm/pull/14287) +* @TobiMayr made their first contribution in [PR #14262](https://github.com/BerriAI/litellm/pull/14262) +* @Eitan1112 made their first contribution in [PR #14252](https://github.com/BerriAI/litellm/pull/14252) + +--- + +## **[Full Changelog](https://github.com/BerriAI/litellm/compare/v1.76.1-nightly...v1.76.3-nightly)** diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index a2d781fa1c4..efee1a7783e 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -95,13 +95,14 @@ class PrometheusLogger(CustomLogger): self.litellm_llm_api_time_to_first_token_metric = self._histogram_factory( "litellm_llm_api_time_to_first_token_metric", "Time to first token for a models LLM API call", - labelnames=[ - "model", - "hashed_api_key", - "api_key_alias", - "team", - "team_alias", - ], + # labelnames=[ + # "model", + # "hashed_api_key", + # "api_key_alias", + # "team", + # "team_alias", + # ], + labelnames=self.get_labels_for_metric("litellm_llm_api_time_to_first_token_metric"), buckets=LATENCY_BUCKETS, ) @@ -109,15 +110,7 @@ class PrometheusLogger(CustomLogger): self.litellm_spend_metric = self._counter_factory( "litellm_spend_metric", "Total spend on LLM requests", - labelnames=[ - "end_user", - "hashed_api_key", - "api_key_alias", - "model", - "team", - "team_alias", - "user", - ], + labelnames=self.get_labels_for_metric("litellm_spend_metric"), ) # Counter for total_output_tokens @@ -243,25 +236,18 @@ class PrometheusLogger(CustomLogger): labelnames=["api_provider"], ) - # Get all keys - _logged_llm_labels = [ - UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, - UserAPIKeyLabelNames.MODEL_ID.value, - UserAPIKeyLabelNames.API_BASE.value, - UserAPIKeyLabelNames.API_PROVIDER.value, - ] - # Metric for deployment state self.litellm_deployment_state = self._gauge_factory( "litellm_deployment_state", "LLM Deployment Analytics - The state of the deployment: 0 = healthy, 1 = partial outage, 2 = complete outage", - labelnames=_logged_llm_labels, + labelnames=self.get_labels_for_metric("litellm_deployment_state") ) self.litellm_deployment_cooled_down = self._counter_factory( "litellm_deployment_cooled_down", "LLM Deployment Analytics - Number of times a deployment has been cooled down by LiteLLM load balancing logic. exception_status is the status of the exception that caused the deployment to be cooled down", - labelnames=_logged_llm_labels + [EXCEPTION_STATUS], + # labelnames=_logged_llm_labels + [EXCEPTION_STATUS], + labelnames=self.get_labels_for_metric("litellm_deployment_cooled_down") ) self.litellm_deployment_success_responses = self._counter_factory( @@ -327,6 +313,7 @@ class PrometheusLogger(CustomLogger): documentation="deprecated - use litellm_proxy_total_requests_metric. Total number of LLM calls to litellm - track total per API Key, team, user", labelnames=self.get_labels_for_metric("litellm_requests_metric"), ) + except Exception as e: print_verbose(f"Got exception on init prometheus client {str(e)}") raise e diff --git a/litellm/__init__.py b/litellm/__init__.py index 79865c83513..b89dde0add9 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -67,6 +67,7 @@ from litellm.constants import ( bedrock_embedding_models, known_tokenizer_config, BEDROCK_INVOKE_PROVIDERS_LITERAL, + BEDROCK_CONVERSE_MODELS, DEFAULT_MAX_TOKENS, DEFAULT_SOFT_BUDGET, DEFAULT_ALLOWED_FAILS, @@ -145,8 +146,11 @@ _custom_logger_compatible_callbacks_literal = Literal[ "aws_sqs", "vector_store_pre_call_hook", "dotprompt", + "cloudzero", ] -configured_cold_storage_logger: Optional[_custom_logger_compatible_callbacks_literal] = None +configured_cold_storage_logger: Optional[ + _custom_logger_compatible_callbacks_literal +] = None logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None _known_custom_logger_compatible_callbacks: List = list( get_args(_custom_logger_compatible_callbacks_literal) @@ -432,43 +436,10 @@ organization = None project = None config_path = None vertex_ai_safety_settings: Optional[dict] = None -BEDROCK_CONVERSE_MODELS = [ - "openai.gpt-oss-20b-1:0", - "openai.gpt-oss-120b-1:0", - "anthropic.claude-opus-4-1-20250805-v1:0", - "anthropic.claude-opus-4-20250514-v1:0", - "anthropic.claude-sonnet-4-20250514-v1:0", - "anthropic.claude-3-7-sonnet-20250219-v1:0", - "anthropic.claude-3-5-haiku-20241022-v1:0", - "anthropic.claude-3-5-sonnet-20241022-v2:0", - "anthropic.claude-3-5-sonnet-20240620-v1:0", - "anthropic.claude-3-opus-20240229-v1:0", - "anthropic.claude-3-sonnet-20240229-v1:0", - "anthropic.claude-3-haiku-20240307-v1:0", - "anthropic.claude-v2", - "anthropic.claude-v2:1", - "anthropic.claude-v1", - "anthropic.claude-instant-v1", - "ai21.jamba-instruct-v1:0", - "ai21.jamba-1-5-mini-v1:0", - "ai21.jamba-1-5-large-v1:0", - "meta.llama3-70b-instruct-v1:0", - "meta.llama3-8b-instruct-v1:0", - "meta.llama3-1-8b-instruct-v1:0", - "meta.llama3-1-70b-instruct-v1:0", - "meta.llama3-1-405b-instruct-v1:0", - "meta.llama3-70b-instruct-v1:0", - "mistral.mistral-large-2407-v1:0", - "mistral.mistral-large-2402-v1:0", - "mistral.mistral-small-2402-v1:0", - "meta.llama3-2-1b-instruct-v1:0", - "meta.llama3-2-3b-instruct-v1:0", - "meta.llama3-2-11b-instruct-v1:0", - "meta.llama3-2-90b-instruct-v1:0", -] ####### COMPLETION MODELS ################### -from typing import Set +from typing import Set + open_ai_chat_completion_models: Set = set() open_ai_text_completion_models: Set = set() cohere_models: Set = set() @@ -483,6 +454,7 @@ vertex_vision_models: Set = set() vertex_chat_models: Set = set() vertex_code_chat_models: Set = set() vertex_ai_image_models: Set = set() +vertex_ai_video_models: Set = set() vertex_text_models: Set = set() vertex_code_text_models: Set = set() vertex_embedding_models: Set = set() @@ -491,6 +463,7 @@ vertex_llama3_models: Set = set() vertex_deepseek_models: Set = set() vertex_ai_ai21_models: Set = set() vertex_mistral_models: Set = set() +vertex_openai_models: Set = set() ai21_models: Set = set() ai21_chat_models: Set = set() nlp_cloud_models: Set = set() @@ -544,6 +517,7 @@ recraft_models: Set = set() cometapi_models: Set = set() oci_models: Set = set() vercel_ai_gateway_models: Set = set() +volcengine_models: Set = set() def is_bedrock_pricing_only_model(key: str) -> bool: @@ -637,6 +611,12 @@ def add_known_models(): elif value.get("litellm_provider") == "vertex_ai-image-models": key = key.replace("vertex_ai/", "") vertex_ai_image_models.add(key) + elif value.get("litellm_provider") == "vertex_ai-video-models": + key = key.replace("vertex_ai/", "") + vertex_ai_video_models.add(key) + elif value.get("litellm_provider") == "vertex_ai-openai_models": + key = key.replace("vertex_ai/", "") + vertex_openai_models.add(key) elif value.get("litellm_provider") == "ai21": if value.get("mode") == "chat": ai21_chat_models.add(key) @@ -748,6 +728,8 @@ def add_known_models(): cometapi_models.add(key) elif value.get("litellm_provider") == "oci": oci_models.add(key) + elif value.get("litellm_provider") == "volcengine": + volcengine_models.add(key) add_known_models() @@ -840,6 +822,7 @@ model_list = list( | cometapi_models | oci_models | vercel_ai_gateway_models + | volcengine_models ) model_list_set = set(model_list) @@ -860,7 +843,12 @@ models_by_provider: dict = { "openrouter": openrouter_models, "vercel_ai_gateway": vercel_ai_gateway_models, "datarobot": datarobot_models, - "vertex_ai": vertex_chat_models | vertex_text_models | vertex_anthropic_models | vertex_vision_models | vertex_language_models | vertex_deepseek_models, + "vertex_ai": vertex_chat_models + | vertex_text_models + | vertex_anthropic_models + | vertex_vision_models + | vertex_language_models + | vertex_deepseek_models, "ai21": ai21_models, "bedrock": bedrock_models | bedrock_converse_models, "petals": petals_models, @@ -914,6 +902,7 @@ models_by_provider: dict = { "recraft": recraft_models, "cometapi": cometapi_models, "oci": oci_models, + "volcengine": volcengine_models, } # mapping for those models which have larger equivalents @@ -1157,7 +1146,9 @@ from .llms.topaz.image_variations.transformation import TopazImageVariationConfi from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig from .llms.groq.chat.transformation import GroqChatConfig from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig -from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig +from .llms.voyage.embedding.transformation_contextual import ( + VoyageContextualEmbeddingConfig, +) from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig from .llms.azure_ai.chat.transformation import AzureAIStudioConfig from .llms.mistral.chat.transformation import MistralConfig @@ -1221,7 +1212,9 @@ from .llms.jina_ai.embedding.transformation import JinaAIEmbeddingConfig from .llms.xai.chat.transformation import XAIChatConfig from .llms.xai.common_utils import XAIModelInfo from .llms.aiml.chat.transformation import AIMLChatConfig -from .llms.volcengine import VolcEngineConfig +from .llms.volcengine.chat.transformation import ( + VolcEngineChatConfig as VolcEngineConfig, +) from .llms.codestral.completion.transformation import CodestralTextCompletionConfig from .llms.azure.azure import ( AzureOpenAIError, diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 3ea0f95157f..0d250779da3 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -14,13 +14,15 @@ import asyncio import contextvars import os from functools import partial -from typing import Any, Coroutine, Dict, Literal, Optional, Union +from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast import httpx import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.batches.handler import AzureBatchesAPI +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.openai import OpenAIBatchesAPI from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction from litellm.secret_managers.main import get_secret_str @@ -31,13 +33,19 @@ from litellm.types.llms.openai import ( RetrieveBatchRequest, ) from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import LiteLLMBatch -from litellm.utils import client, get_litellm_params, supports_httpx_timeout +from litellm.types.utils import LiteLLMBatch, LlmProviders +from litellm.utils import ( + ProviderConfigManager, + client, + get_litellm_params, + supports_httpx_timeout, +) ####### ENVIRONMENT VARIABLES ################### openai_batches_instance = OpenAIBatchesAPI() azure_batches_instance = AzureBatchesAPI() vertex_ai_batches_instance = VertexAIBatchPrediction(gcs_bucket_name="") +base_llm_http_handler = BaseLLMHTTPHandler() ################################################# @@ -46,7 +54,7 @@ async def acreate_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -94,7 +102,7 @@ def create_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -111,8 +119,8 @@ def create_batch( proxy_server_request = kwargs.get("proxy_server_request", None) model_info = kwargs.get("model_info", None) _is_async = kwargs.pop("acreate_batch", False) is True - litellm_params = get_litellm_params(**kwargs) - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj", None) + litellm_params = dict(GenericLiteLLMParams(**kwargs)) + litellm_logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None)) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 litellm_logging_obj.update_environment_variables( @@ -142,6 +150,7 @@ def create_batch( timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 + _create_batch_request = CreateBatchRequest( completion_window=completion_window, @@ -151,6 +160,27 @@ def create_batch( extra_headers=extra_headers, extra_body=extra_body, ) + provider_config = ProviderConfigManager.get_provider_batches_config( + model="", + provider=LlmProviders(custom_llm_provider), + ) + if provider_config is not None: + response = base_llm_http_handler.create_batch( + provider_config=provider_config, + litellm_params=litellm_params, + create_batch_data=_create_batch_request, + headers=extra_headers or {}, + api_base=optional_params.api_base, + api_key=optional_params.api_key, + logging_obj=litellm_logging_obj, + _is_async=_is_async, + client=client + if client is not None + and isinstance(client, (HTTPHandler, AsyncHTTPHandler)) + else None, + timeout=timeout, + ) + return response api_base: Optional[str] = None if custom_llm_provider == "openai": # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -322,20 +352,21 @@ def retrieve_batch( """ try: optional_params = GenericLiteLLMParams(**kwargs) - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj", None) + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 litellm_params = get_litellm_params( custom_llm_provider=custom_llm_provider, **kwargs, ) - litellm_logging_obj.update_environment_variables( - model=None, - user=None, - optional_params=optional_params.model_dump(), - litellm_params=litellm_params, - custom_llm_provider=custom_llm_provider, - ) + if litellm_logging_obj is not None: + litellm_logging_obj.update_environment_variables( + model=None, + user=None, + optional_params=optional_params.model_dump(), + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, + ) if ( timeout is not None diff --git a/litellm/constants.py b/litellm/constants.py index 0655473301f..bbe2ae3c8d5 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -14,7 +14,9 @@ DEFAULT_S3_BATCH_SIZE = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) DEFAULT_SQS_FLUSH_INTERVAL_SECONDS = int( os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10) ) -DEFAULT_NUM_WORKERS_LITELLM_PROXY = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 4)) +DEFAULT_NUM_WORKERS_LITELLM_PROXY = int( + os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", os.cpu_count() or 4) +) DEFAULT_SQS_BATCH_SIZE = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) SQS_SEND_MESSAGE_ACTION = "SendMessage" SQS_API_VERSION = "2012-11-05" @@ -49,6 +51,23 @@ SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD = int( DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET = int( os.getenv("DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET", 0) ) + +# Gemini model-specific minimal thinking budget constants +DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH = int( + os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH", 1) +) +DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO = int( + os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO", 128) +) +DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int( + os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE", 512) +) + +# Generic fallback for unknown models +DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int( + os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) +) + DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int( os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024) ) @@ -395,6 +414,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES = { "reasoning_effort": None, "thinking": None, "web_search_options": None, + "safety_identifier": None, } openai_compatible_endpoints: List = [ @@ -745,6 +765,42 @@ BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[ "deepseek_r1", ] +BEDROCK_CONVERSE_MODELS = [ + "openai.gpt-oss-20b-1:0", + "openai.gpt-oss-120b-1:0", + "anthropic.claude-opus-4-1-20250805-v1:0", + "anthropic.claude-opus-4-20250514-v1:0", + "anthropic.claude-sonnet-4-20250514-v1:0", + "anthropic.claude-3-7-sonnet-20250219-v1:0", + "anthropic.claude-3-5-haiku-20241022-v1:0", + "anthropic.claude-3-5-sonnet-20241022-v2:0", + "anthropic.claude-3-5-sonnet-20240620-v1:0", + "anthropic.claude-3-opus-20240229-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", + "anthropic.claude-3-haiku-20240307-v1:0", + "anthropic.claude-v2", + "anthropic.claude-v2:1", + "anthropic.claude-v1", + "anthropic.claude-instant-v1", + "ai21.jamba-instruct-v1:0", + "ai21.jamba-1-5-mini-v1:0", + "ai21.jamba-1-5-large-v1:0", + "meta.llama3-70b-instruct-v1:0", + "meta.llama3-8b-instruct-v1:0", + "meta.llama3-1-8b-instruct-v1:0", + "meta.llama3-1-70b-instruct-v1:0", + "meta.llama3-1-405b-instruct-v1:0", + "meta.llama3-70b-instruct-v1:0", + "mistral.mistral-large-2407-v1:0", + "mistral.mistral-large-2402-v1:0", + "mistral.mistral-small-2402-v1:0", + "meta.llama3-2-1b-instruct-v1:0", + "meta.llama3-2-3b-instruct-v1:0", + "meta.llama3-2-11b-instruct-v1:0", + "meta.llama3-2-90b-instruct-v1:0", +] + + open_ai_embedding_models: set = set(["text-embedding-ada-002"]) cohere_embedding_models: set = set( [ @@ -834,6 +890,9 @@ AZURE_STORAGE_MSFT_VERSION = "2019-07-07" PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES = int( os.getenv("PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES", 5) ) +CLOUDZERO_EXPORT_INTERVAL_MINUTES = int( + os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60) +) MCP_TOOL_NAME_PREFIX = "mcp_tool" MAXIMUM_TRACEBACK_LINES_TO_LOG = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)) @@ -888,6 +947,8 @@ LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token" ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics" +CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data" +CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)) SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) diff --git a/litellm/files/main.py b/litellm/files/main.py index 5d0dc05771a..299e52895bf 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -50,7 +50,7 @@ vertex_ai_files_instance = VertexAIFilesHandler() async def acreate_file( file: FileTypes, purpose: Literal["assistants", "batch", "fine-tune"], - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -94,7 +94,7 @@ async def acreate_file( def create_file( file: FileTypes, purpose: Literal["assistants", "batch", "fine-tune"], - custom_llm_provider: Optional[Literal["openai", "azure", "vertex_ai"]] = None, + custom_llm_provider: Optional[Literal["openai", "azure", "vertex_ai", "bedrock"]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -109,7 +109,7 @@ def create_file( try: _is_async = kwargs.pop("acreate_file", False) is True optional_params = GenericLiteLLMParams(**kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + litellm_params_dict = dict(**kwargs) logging_obj = cast( Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj") ) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 1f575f27591..dcf707ebd51 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -37,6 +37,10 @@ class GenerateContentToCompletionHandler: completion_kwargs: Dict[str, Any] = dict(completion_request) + # feed metadata for custom callback + if extra_kwargs is not None and "metadata" in extra_kwargs: + completion_kwargs["metadata"] = extra_kwargs["metadata"] + if stream: completion_kwargs["stream"] = stream diff --git a/litellm/images/main.py b/litellm/images/main.py index 4993a48c724..2a8b62bce24 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -90,12 +90,12 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: response = init_response elif asyncio.iscoroutine(init_response): response = await init_response # type: ignore - + if response is None: raise ValueError( "Unable to get Image Response. Please pass a valid llm_provider." ) - + return response except Exception as e: custom_llm_provider = custom_llm_provider or "openai" @@ -108,6 +108,8 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: ) +# fmt: off + # Overload for when aimg_generation=True (returns Coroutine) @overload def image_generation( @@ -119,7 +121,6 @@ def image_generation( size: Optional[str] = None, style: Optional[str] = None, user: Optional[str] = None, - input_fidelity: Optional[str] = None, timeout=600, # default to 10 minutes api_key: Optional[str] = None, api_base: Optional[str] = None, @@ -128,10 +129,11 @@ def image_generation( *, aimg_generation: Literal[True], **kwargs, -) -> Coroutine[Any, Any, ImageResponse]: +) -> Coroutine[Any, Any, ImageResponse]: ... + # Overload for when aimg_generation=False or not specified (returns ImageResponse) @overload def image_generation( @@ -143,7 +145,6 @@ def image_generation( size: Optional[str] = None, style: Optional[str] = None, user: Optional[str] = None, - input_fidelity: Optional[str] = None, timeout=600, # default to 10 minutes api_key: Optional[str] = None, api_base: Optional[str] = None, @@ -152,9 +153,11 @@ def image_generation( *, aimg_generation: Literal[False] = False, **kwargs, -) -> ImageResponse: +) -> ImageResponse: ... +# fmt: on + @client def image_generation( # noqa: PLR0915 @@ -166,7 +169,6 @@ def image_generation( # noqa: PLR0915 size: Optional[str] = None, style: Optional[str] = None, user: Optional[str] = None, - input_fidelity: Optional[str] = None, timeout=600, # default to 10 minutes api_key: Optional[str] = None, api_base: Optional[str] = None, @@ -174,9 +176,9 @@ def image_generation( # noqa: PLR0915 custom_llm_provider=None, **kwargs, ) -> Union[ - ImageResponse, - Coroutine[Any, Any, ImageResponse], - ]: + ImageResponse, + Coroutine[Any, Any, ImageResponse], +]: """ Maps the https://api.openai.com/v1/images/generations endpoint. @@ -227,7 +229,6 @@ def image_generation( # noqa: PLR0915 "quality", "size", "style", - "input_fidelity", ] litellm_params = all_litellm_params default_params = openai_params + litellm_params @@ -255,7 +256,6 @@ def image_generation( # noqa: PLR0915 size=size, style=style, user=user, - input_fidelity=input_fidelity, custom_llm_provider=custom_llm_provider, provider_config=image_generation_config, **non_default_params, @@ -344,8 +344,10 @@ def image_generation( # noqa: PLR0915 litellm.LlmProviders.GEMINI, ): if image_generation_config is None: - raise ValueError(f"image generation config is not supported for {custom_llm_provider}") - + raise ValueError( + f"image generation config is not supported for {custom_llm_provider}" + ) + return llm_http_handler.image_generation_handler( api_key=api_key, model=model, @@ -360,6 +362,7 @@ def image_generation( # noqa: PLR0915 ) elif custom_llm_provider == "azure_ai": from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + api_base = AzureFoundryModelInfo.get_api_base(api_base) api_key = AzureFoundryModelInfo.get_api_key(api_key) if extra_headers is not None: @@ -420,7 +423,7 @@ def image_generation( # noqa: PLR0915 aimg_generation=aimg_generation, client=client, api_base=api_base, - api_key=api_key + api_key=api_key, ) elif custom_llm_provider == "vertex_ai": vertex_ai_project = ( @@ -705,7 +708,7 @@ def image_edit( litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("async_call", False) is True - #add images / or return a single image + # add images / or return a single image images = image if isinstance(image, list) else [image] # get llm provider logic @@ -716,11 +719,11 @@ def image_edit( ) # get provider config - image_edit_provider_config: Optional[ - BaseImageEditConfig - ] = ProviderConfigManager.get_provider_image_edit_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), + image_edit_provider_config: Optional[BaseImageEditConfig] = ( + ProviderConfigManager.get_provider_image_edit_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), + ) ) if image_edit_provider_config is None: diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 41db4a551bd..7da38e193b6 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -805,9 +805,9 @@ class SlackAlerting(CustomBatchLogger): ### UNIQUE CACHE KEY ### cache_key = provider + region_name - outage_value: Optional[ProviderRegionOutageModel] = ( - await self.internal_usage_cache.async_get_cache(key=cache_key) - ) + outage_value: Optional[ + ProviderRegionOutageModel + ] = await self.internal_usage_cache.async_get_cache(key=cache_key) if ( getattr(exception, "status_code", None) is None @@ -1367,12 +1367,13 @@ Model Info: # Get the current timestamp current_time = datetime.now().strftime("%H:%M:%S") _proxy_base_url = os.getenv("PROXY_BASE_URL", None) + # Use .name if it's an enum, otherwise use as is + alert_type_name = getattr(alert_type, 'name', alert_type) + alert_type_formatted = f"Alert type: `{alert_type_name}`" if alert_type == "daily_reports" or alert_type == "new_model_added": - formatted_message = message + formatted_message = alert_type_formatted + message else: - formatted_message = ( - f"Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" - ) + formatted_message = f"{alert_type_formatted}\nLevel: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" if kwargs: for key, value in kwargs.items(): @@ -1388,9 +1389,9 @@ Model Info: self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url ): - slack_webhook_url: Optional[Union[str, List[str]]] = ( - self.alert_to_webhook_url[alert_type] - ) + slack_webhook_url: Optional[ + Union[str, List[str]] + ] = self.alert_to_webhook_url[alert_type] elif self.default_webhook_url is not None: slack_webhook_url = self.default_webhook_url else: diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 531da933fcc..5bc6afb6dbc 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -1,13 +1,11 @@ # What is this? ## Log success + failure events to Braintrust -import copy import os from datetime import datetime from typing import Dict, Optional import httpx -from pydantic import BaseModel import litellm from litellm import verbose_logger @@ -24,7 +22,6 @@ API_BASE = "https://api.braintrustdata.com/v1" def get_utc_datetime(): import datetime as dt - from datetime import datetime if hasattr(dt, "UTC"): return datetime.now(dt.UTC) # type: ignore @@ -45,9 +42,9 @@ class BraintrustLogger(CustomLogger): "Authorization": "Bearer " + self.api_key, "Content-Type": "application/json", } - self._project_id_cache: Dict[ - str, str - ] = {} # Cache mapping project names to IDs + self._project_id_cache: Dict[str, str] = ( + {} + ) # Cache mapping project names to IDs self.global_braintrust_http_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) @@ -108,43 +105,6 @@ class BraintrustLogger(CustomLogger): except httpx.HTTPStatusError as e: raise Exception(f"Failed to register project: {e.response.text}") - @staticmethod - def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict: - """ - Adds metadata from proxy request headers to Braintrust logging if keys start with "braintrust_" - and overwrites litellm_params.metadata if already included. - - For example if you want to append your trace to an existing `trace_id` via header, send - `headers: { ..., langfuse_existing_trace_id: your-existing-trace-id }` via proxy request. - """ - if litellm_params is None: - return metadata - - if litellm_params.get("proxy_server_request") is None: - return metadata - - if metadata is None: - metadata = {} - - proxy_headers = ( - litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} - ) - - for metadata_param_key in proxy_headers: - if metadata_param_key.startswith("braintrust"): - trace_param_key = metadata_param_key.replace("braintrust", "", 1) - if trace_param_key in metadata: - verbose_logger.warning( - f"Overwriting Braintrust `{trace_param_key}` from request header" - ) - else: - verbose_logger.debug( - f"Found Braintrust `{trace_param_key}` in request header" - ) - metadata[trace_param_key] = proxy_headers.get(metadata_param_key) - - return metadata - async def create_default_project_and_experiment(self): project = await self.global_braintrust_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} @@ -169,7 +129,9 @@ class BraintrustLogger(CustomLogger): verbose_logger.debug("REACHES BRAINTRUST SUCCESS") try: litellm_call_id = kwargs.get("litellm_call_id") + standard_logging_object = kwargs.get("standard_logging_object", {}) prompt = {"messages": kwargs.get("messages")} + output = None choices = [] if response_obj is not None and ( @@ -192,33 +154,13 @@ class BraintrustLogger(CustomLogger): ): output = response_obj["data"] - litellm_params = kwargs.get("litellm_params", {}) - metadata = ( - litellm_params.get("metadata", {}) or {} - ) # if litellm_params['metadata'] == None - metadata = self.add_metadata_from_header(litellm_params, metadata) - clean_metadata = {} - try: - metadata = copy.deepcopy( - metadata - ) # Avoid modifying the original metadata - except Exception: - new_metadata = {} - for key, value in metadata.items(): - if ( - isinstance(value, list) - or isinstance(value, dict) - or isinstance(value, str) - or isinstance(value, int) - or isinstance(value, float) - ): - new_metadata[key] = copy.deepcopy(value) - metadata = new_metadata + litellm_params = kwargs.get("litellm_params", {}) or {} + dynamic_metadata = litellm_params.get("metadata", {}) or {} # Get project_id from metadata or create default if needed - project_id = metadata.get("project_id") + project_id = dynamic_metadata.get("project_id") if project_id is None: - project_name = metadata.get("project_name") + project_name = dynamic_metadata.get("project_name") project_id = ( self.get_project_id_sync(project_name) if project_name else None ) @@ -229,8 +171,9 @@ class BraintrustLogger(CustomLogger): project_id = self.default_project_id tags = [] - if isinstance(metadata, dict): - for key, value in metadata.items(): + + if isinstance(dynamic_metadata, dict): + for key, value in dynamic_metadata.items(): # generate langfuse tags - Default Tags sent to Langfuse from LiteLLM Proxy if ( litellm.langfuse_default_tags is not None @@ -239,25 +182,12 @@ class BraintrustLogger(CustomLogger): ): tags.append(f"{key}:{value}") - # clean litellm metadata before logging - if key in [ - "headers", - "endpoint", - "caching_groups", - "previous_models", - ]: - continue - else: - clean_metadata[key] = value + if ( + isinstance(value, str) and key not in standard_logging_object + ): # support logging dynamic metadata to braintrust + standard_logging_object[key] = value cost = kwargs.get("response_cost", None) - if cost is not None: - clean_metadata["litellm_response_cost"] = cost - - # metadata.model is required for braintrust to calculate the "Estimated cost" metric - litellm_model = kwargs.get("model", None) - if litellm_model is not None: - clean_metadata["model"] = litellm_model metrics: Optional[dict] = None usage_obj = getattr(response_obj, "usage", None) @@ -275,12 +205,12 @@ class BraintrustLogger(CustomLogger): } # Allow metadata override for span name - span_name = metadata.get("span_name", "Chat Completion") - + span_name = dynamic_metadata.get("span_name", "Chat Completion") + request_data = { "id": litellm_call_id, "input": prompt["messages"], - "metadata": clean_metadata, + "metadata": standard_logging_object, "tags": tags, "span_attributes": {"name": span_name, "type": "llm"}, } @@ -312,6 +242,7 @@ class BraintrustLogger(CustomLogger): verbose_logger.debug("REACHES BRAINTRUST SUCCESS") try: litellm_call_id = kwargs.get("litellm_call_id") + standard_logging_object = kwargs.get("standard_logging_object", {}) prompt = {"messages": kwargs.get("messages")} output = None choices = [] @@ -336,32 +267,12 @@ class BraintrustLogger(CustomLogger): output = response_obj["data"] litellm_params = kwargs.get("litellm_params", {}) - metadata = ( - litellm_params.get("metadata", {}) or {} - ) # if litellm_params['metadata'] == None - metadata = self.add_metadata_from_header(litellm_params, metadata) - clean_metadata = {} - new_metadata = {} - for key, value in metadata.items(): - if ( - isinstance(value, list) - or isinstance(value, str) - or isinstance(value, int) - or isinstance(value, float) - ): - new_metadata[key] = value - elif isinstance(value, BaseModel): - new_metadata[key] = value.model_dump_json() - elif isinstance(value, dict): - for k, v in value.items(): - if isinstance(v, datetime): - value[k] = v.isoformat() - new_metadata[key] = value + dynamic_metadata = litellm_params.get("metadata", {}) or {} # Get project_id from metadata or create default if needed - project_id = metadata.get("project_id") + project_id = dynamic_metadata.get("project_id") if project_id is None: - project_name = metadata.get("project_name") + project_name = dynamic_metadata.get("project_name") project_id = ( await self.get_project_id_async(project_name) if project_name @@ -374,8 +285,9 @@ class BraintrustLogger(CustomLogger): project_id = self.default_project_id tags = [] - if isinstance(metadata, dict): - for key, value in metadata.items(): + + if isinstance(dynamic_metadata, dict): + for key, value in dynamic_metadata.items(): # generate langfuse tags - Default Tags sent to Langfuse from LiteLLM Proxy if ( litellm.langfuse_default_tags is not None @@ -384,25 +296,12 @@ class BraintrustLogger(CustomLogger): ): tags.append(f"{key}:{value}") - # clean litellm metadata before logging - if key in [ - "headers", - "endpoint", - "caching_groups", - "previous_models", - ]: - continue - else: - clean_metadata[key] = value + if ( + isinstance(value, str) and key not in standard_logging_object + ): # support logging dynamic metadata to braintrust + standard_logging_object[key] = value cost = kwargs.get("response_cost", None) - if cost is not None: - clean_metadata["litellm_response_cost"] = cost - - # metadata.model is required for braintrust to calculate the "Estimated cost" metric - litellm_model = kwargs.get("model", None) - if litellm_model is not None: - clean_metadata["model"] = litellm_model metrics: Optional[dict] = None usage_obj = getattr(response_obj, "usage", None) @@ -430,13 +329,13 @@ class BraintrustLogger(CustomLogger): ) # Allow metadata override for span name - span_name = metadata.get("span_name", "Chat Completion") - + span_name = dynamic_metadata.get("span_name", "Chat Completion") + request_data = { "id": litellm_call_id, "input": prompt["messages"], "output": output, - "metadata": clean_metadata, + "metadata": standard_logging_object, "tags": tags, "span_attributes": {"name": span_name, "type": "llm"}, } diff --git a/litellm/integrations/cloudzero/cloudzero.py b/litellm/integrations/cloudzero/cloudzero.py index 85aa1679732..ab4ec234bf0 100644 --- a/litellm/integrations/cloudzero/cloudzero.py +++ b/litellm/integrations/cloudzero/cloudzero.py @@ -1,14 +1,15 @@ -import asyncio import os -from datetime import datetime, timedelta -from typing import Optional +from datetime import datetime +from typing import TYPE_CHECKING, Any, List, Optional, cast +import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger -from .cz_stream_api import CloudZeroStreamer -from .database import LiteLLMDatabase -from .transform import CBFTransformer +if TYPE_CHECKING: + from apscheduler.schedulers.asyncio import AsyncIOScheduler +else: + AsyncIOScheduler = Any class CloudZeroLogger(CustomLogger): @@ -29,20 +30,80 @@ class CloudZeroLogger(CustomLogger): self.api_key = api_key or os.getenv("CLOUDZERO_API_KEY") self.connection_id = connection_id or os.getenv("CLOUDZERO_CONNECTION_ID") self.timezone = timezone or os.getenv("CLOUDZERO_TIMEZONE", "UTC") + verbose_logger.debug(f"CloudZero Logger initialized with connection ID: {self.connection_id}, timezone: {self.timezone}") - async def export_usage_data(self, target_hour: datetime, limit: Optional[int] = 1000, operation: str = "replace_hourly"): + async def initialize_cloudzero_export_job(self): """ - Exports the usage data for a specific hour to CloudZero. + Handler for initializing CloudZero export job. - - Reads spend logs from the DB for the specified hour + Runs when CloudZero logger starts up. + + - If redis cache is available, we use the pod lock manager to acquire a lock and export the data. + - Ensures only one pod exports the data at a time. + - If redis cache is not available, we export the data directly. + """ + from litellm.constants import ( + CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME, + ) + from litellm.proxy.proxy_server import proxy_logging_obj + pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager + + # if using redis, ensure only one pod exports the data at a time + if pod_lock_manager and pod_lock_manager.redis_cache: + if await pod_lock_manager.acquire_lock( + cronjob_id=CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME + ): + try: + await self._hourly_usage_data_export() + finally: + await pod_lock_manager.release_lock( + cronjob_id=CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME + ) + else: + # if not using redis, export the data directly + await self._hourly_usage_data_export() + + async def _hourly_usage_data_export(self): + """ + Exports the hourly usage data to CloudZero. + + Start time: 1 hour ago + End time: current time + """ + from datetime import timedelta, timezone + + from litellm.constants import CLOUDZERO_MAX_FETCHED_DATA_RECORDS + current_time_utc = datetime.now(timezone.utc) + one_hour_ago_utc = current_time_utc - timedelta(hours=1) + await self.export_usage_data( + limit=CLOUDZERO_MAX_FETCHED_DATA_RECORDS, + operation="replace_hourly", + start_time_utc=one_hour_ago_utc, + end_time_utc=current_time_utc + ) + + + async def export_usage_data( + self, + limit: Optional[int] = None, + operation: str = "replace_hourly", + start_time_utc: Optional[datetime] = None, + end_time_utc: Optional[datetime] = None + ): + """ + Exports the usage data to CloudZero. + + - Reads data from the DB - Transforms the data to the CloudZero format - Sends the data to CloudZero Args: - target_hour: The specific hour to export data for - limit: Optional limit on number of records to export (default: 1000) + limit: Optional limit on number of records to export operation: CloudZero operation type ("replace_hourly" or "sum") """ + from litellm.integrations.cloudzero.cz_stream_api import CloudZeroStreamer + from litellm.integrations.cloudzero.database import LiteLLMDatabase + from litellm.integrations.cloudzero.transform import CBFTransformer try: verbose_logger.debug("CloudZero Logger: Starting usage data export") @@ -52,11 +113,27 @@ class CloudZeroLogger(CustomLogger): "CloudZero configuration missing. Please set CLOUDZERO_API_KEY and CLOUDZERO_CONNECTION_ID environment variables." ) - # Fetch and transform data using helper - cbf_data = await self._fetch_cbf_data_for_hour(target_hour, limit) + # Initialize database connection and load data + database = LiteLLMDatabase() + verbose_logger.debug("CloudZero Logger: Loading usage data from database") + data = await database.get_usage_data( + limit=limit, + start_time_utc=start_time_utc, + end_time_utc=end_time_utc + ) + + if data.is_empty(): + verbose_logger.info("CloudZero Logger: No usage data found to export") + return + + verbose_logger.debug(f"CloudZero Logger: Processing {len(data)} records") + + # Transform data to CloudZero CBF format + transformer = CBFTransformer() + cbf_data = transformer.transform(data) if cbf_data.is_empty(): - verbose_logger.info("CloudZero Logger: No usage data found to export") + verbose_logger.warning("CloudZero Logger: No valid data after transformation") return # Send data to CloudZero @@ -75,60 +152,86 @@ class CloudZeroLogger(CustomLogger): verbose_logger.error(f"CloudZero Logger: Error exporting usage data: {str(e)}") raise - async def _fetch_cbf_data_for_hour(self, target_hour: datetime, limit: Optional[int] = 1000): + async def dry_run_export_usage_data(self, limit: Optional[int] = 10000): """ - Helper method to fetch usage data for a specific hour and transform it to CloudZero CBF format. + Returns the data that would be exported to CloudZero without actually sending it. Args: - target_hour: The specific hour to fetch data for - limit: Optional limit on number of records to fetch (default: 1000) + limit: Limit number of records to display (default: 10000) Returns: - CBF formatted data ready for CloudZero ingestion - """ - # Initialize database connection and load data - database = LiteLLMDatabase() - verbose_logger.debug(f"CloudZero Logger: Loading spend logs for hour {target_hour}") - data = await database.get_usage_data_for_hour(target_hour=target_hour, limit=limit) - - if data.is_empty(): - verbose_logger.info("CloudZero Logger: No usage data found for the specified hour") - return data # Return empty data - - verbose_logger.debug(f"CloudZero Logger: Processing {len(data)} records") - - # Transform data to CloudZero CBF format - transformer = CBFTransformer() - cbf_data = transformer.transform(data) - - if cbf_data.is_empty(): - verbose_logger.warning("CloudZero Logger: No valid data after transformation") - - return cbf_data - - async def dry_run_export_usage_data(self, target_hour: datetime, limit: Optional[int] = 1000): - """ - Only prints the spend logs data for a specific hour that would be exported to CloudZero. - - Args: - target_hour: The specific hour to export data for - limit: Limit number of records to display (default: 1000) + dict: Contains usage_data, cbf_data, and summary statistics """ + from litellm.integrations.cloudzero.database import LiteLLMDatabase + from litellm.integrations.cloudzero.transform import CBFTransformer try: verbose_logger.debug("CloudZero Logger: Starting dry run export") - # Fetch and transform data using helper - cbf_data = await self._fetch_cbf_data_for_hour(target_hour, limit) + # Initialize database connection and load data + database = LiteLLMDatabase() + verbose_logger.debug("CloudZero Logger: Loading usage data for dry run") + data = await database.get_usage_data(limit=limit) + + if data.is_empty(): + verbose_logger.warning("CloudZero Dry Run: No usage data found") + return { + "usage_data": [], + "cbf_data": [], + "summary": { + "total_records": 0, + "total_cost": 0, + "total_tokens": 0, + "unique_accounts": 0, + "unique_services": 0 + } + } + + verbose_logger.debug(f"CloudZero Dry Run: Processing {len(data)} records...") + + # Convert usage data to dict format for response + usage_data_sample = data.head(50).to_dicts() # Return first 50 rows + + # Transform data to CloudZero CBF format + transformer = CBFTransformer() + cbf_data = transformer.transform(data) if cbf_data.is_empty(): - verbose_logger.warning("CloudZero Dry Run: No usage data found") - return + verbose_logger.warning("CloudZero Dry Run: No valid data after transformation") + return { + "usage_data": usage_data_sample, + "cbf_data": [], + "summary": { + "total_records": len(usage_data_sample), + "total_cost": sum(row.get('spend', 0) for row in usage_data_sample), + "total_tokens": sum(row.get('prompt_tokens', 0) + row.get('completion_tokens', 0) for row in usage_data_sample), + "unique_accounts": 0, + "unique_services": 0 + } + } - # Display the transformed data on screen - self._display_cbf_data_on_screen(cbf_data) + # Convert CBF data to dict format for response + cbf_data_dict = cbf_data.to_dicts() + + # Calculate summary statistics + total_cost = sum(record.get('cost/cost', 0) for record in cbf_data_dict) + unique_accounts = len(set(record.get('resource/account', '') for record in cbf_data_dict if record.get('resource/account'))) + unique_services = len(set(record.get('resource/service', '') for record in cbf_data_dict if record.get('resource/service'))) + total_tokens = sum(record.get('usage/amount', 0) for record in cbf_data_dict) verbose_logger.info(f"CloudZero Logger: Dry run completed for {len(cbf_data)} records") + return { + "usage_data": usage_data_sample, + "cbf_data": cbf_data_dict, + "summary": { + "total_records": len(cbf_data_dict), + "total_cost": total_cost, + "total_tokens": total_tokens, + "unique_accounts": unique_accounts, + "unique_services": unique_services + } + } + except Exception as e: verbose_logger.error(f"CloudZero Logger: Error in dry run export: {str(e)}") verbose_logger.error(f"CloudZero Dry Run Error: {str(e)}") @@ -155,6 +258,11 @@ class CloudZeroLogger(CustomLogger): cbf_table = Table(show_header=True, header_style="bold cyan", box=SIMPLE, padding=(0, 1)) cbf_table.add_column("time/usage_start", style="blue", no_wrap=False) cbf_table.add_column("cost/cost", style="green", justify="right", no_wrap=False) + cbf_table.add_column("entity_type", style="magenta", justify="right", no_wrap=False) + cbf_table.add_column("entity_id", style="magenta", justify="right", no_wrap=False) + cbf_table.add_column("team_id", style="cyan", no_wrap=False) + cbf_table.add_column("team_alias", style="cyan", no_wrap=False) + cbf_table.add_column("api_key_alias", style="yellow", no_wrap=False) cbf_table.add_column("usage/amount", style="yellow", justify="right", no_wrap=False) cbf_table.add_column("resource/id", style="magenta", no_wrap=False) cbf_table.add_column("resource/service", style="cyan", no_wrap=False) @@ -170,10 +278,20 @@ class CloudZeroLogger(CustomLogger): resource_service = str(record.get('resource/service', 'N/A')) resource_account = str(record.get('resource/account', 'N/A')) resource_region = str(record.get('resource/region', 'N/A')) + entity_type = str(record.get('entity_type', 'N/A')) + entity_id = str(record.get('entity_id', 'N/A')) + team_id = str(record.get('resource/tag:team_id', 'N/A')) + team_alias = str(record.get('resource/tag:team_alias', 'N/A')) + api_key_alias = str(record.get('resource/tag:api_key_alias', 'N/A')) cbf_table.add_row( time_usage_start, cost_cost, + entity_type, + entity_id, + team_id, + team_alias, + api_key_alias, usage_amount, resource_id, resource_service, @@ -199,55 +317,33 @@ class CloudZeroLogger(CustomLogger): console.print(f" Unique Services: {unique_services}") console.print("\n[dim]šŸ’” This is the CloudZero CBF format ready for AnyCost ingestion[/dim]") + + @staticmethod + async def init_cloudzero_background_job(scheduler: AsyncIOScheduler): + """ + Initialize the CloudZero background job. - async def init_background_job(self, redis_cache=None): + Starts the background job that exports the usage data to CloudZero every hour. """ - Initialize a background job that exports usage data every hour. - Uses PodLockManager to ensure only one instance runs the export at a time. + from litellm.constants import CLOUDZERO_EXPORT_INTERVAL_MINUTES + from litellm.integrations.custom_logger import CustomLogger - Args: - redis_cache: Redis cache instance for pod locking - """ - from litellm.proxy.db.db_transaction_queue.pod_lock_manager import ( - PodLockManager, + + prometheus_loggers: List[CustomLogger] = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=CloudZeroLogger + ) ) - - lock_manager = PodLockManager(redis_cache=redis_cache) - cronjob_id = "cloudzero_hourly_export" - - async def hourly_export_task(): - while True: - try: - # Calculate the previous completed hour - now = datetime.utcnow() - target_hour = now.replace(minute=0, second=0, microsecond=0) - # Export data for the previous hour to ensure all data is available - target_hour = target_hour - timedelta(hours=1) - - # Try to acquire lock - lock_acquired = await lock_manager.acquire_lock(cronjob_id) - - if lock_acquired: - try: - verbose_logger.info(f"CloudZero Background Job: Starting export for hour {target_hour}") - await self.export_usage_data(target_hour) - verbose_logger.info(f"CloudZero Background Job: Completed export for hour {target_hour}") - finally: - # Always release the lock - await lock_manager.release_lock(cronjob_id) - else: - verbose_logger.debug("CloudZero Background Job: Another instance is already running the export") - - # Wait until the next hour - next_hour = (datetime.utcnow() + timedelta(hours=1)).replace(minute=0, second=0, microsecond=0) - sleep_seconds = (next_hour - datetime.utcnow()).total_seconds() - await asyncio.sleep(sleep_seconds) - - except Exception as e: - verbose_logger.error(f"CloudZero Background Job: Error in hourly export task: {str(e)}") - # Sleep for 5 minutes before retrying on error - await asyncio.sleep(300) - - # Start the background task - asyncio.create_task(hourly_export_task()) - verbose_logger.debug("CloudZero Background Job: Initialized hourly export task") \ No newline at end of file + # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them + verbose_logger.debug("found %s cloudzero loggers", len(prometheus_loggers)) + if len(prometheus_loggers) > 0: + cloudzero_logger = cast(CloudZeroLogger, prometheus_loggers[0]) + verbose_logger.debug( + "Initializing remaining budget metrics as a cron job executing every %s minutes" + % CLOUDZERO_EXPORT_INTERVAL_MINUTES + ) + scheduler.add_job( + cloudzero_logger.initialize_cloudzero_export_job, + "interval", + minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES + ) \ No newline at end of file diff --git a/litellm/integrations/cloudzero/cz_resource_names.py b/litellm/integrations/cloudzero/cz_resource_names.py index 44147f9c210..f1098d20381 100644 --- a/litellm/integrations/cloudzero/cz_resource_names.py +++ b/litellm/integrations/cloudzero/cz_resource_names.py @@ -17,11 +17,16 @@ """CloudZero Resource Names (CZRN) generation and validation for LiteLLM resources.""" import re +from enum import Enum from typing import Any, cast import litellm +class CZEntityType(str, Enum): + TEAM = "team" + + class CZRNGenerator: """Generate CloudZero Resource Names (CZRNs) for LiteLLM resources.""" @@ -49,8 +54,8 @@ class CZRNGenerator: region = 'cross-region' # Use the actual entity_id (team_id or user_id) as the owner account - entity_id = row.get('entity_id', 'unknown') - owner_account_id = self._normalize_component(entity_id) + team_id = row.get('team_id', 'unknown') + owner_account_id = self._normalize_component(team_id) resource_type = 'llm-usage' diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 6d12c5cfbd9..71b4125ed75 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -12,14 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # -# CHANGELOG: 2025-07-23 - Added support for using LiteLLM_SpendLogs table for CBF mapping (ishaan-jaff) # CHANGELOG: 2025-01-19 - Refactored to use daily spend tables for proper CBF mapping (erik.peterson) # CHANGELOG: 2025-01-19 - Migrated from pandas to polars for database operations (erik.peterson) # CHANGELOG: 2025-01-19 - Initial database module for LiteLLM data extraction (erik.peterson) """Database connection and data extraction for LiteLLM.""" -from datetime import datetime, timedelta +from datetime import datetime from typing import Any, Dict, Optional import polars as pl @@ -37,61 +36,88 @@ class LiteLLMDatabase: ) return prisma_client - async def get_usage_data_for_hour(self, target_hour: datetime, limit: Optional[int] = 1000) -> pl.DataFrame: - """Retrieve spend logs for a specific hour from LiteLLM_SpendLogs table with batching.""" + async def get_usage_data( + self, + limit: Optional[int] = None, + start_time_utc: Optional[datetime] = None, + end_time_utc: Optional[datetime] = None + ) -> pl.DataFrame: + """Retrieve usage data from LiteLLM daily user spend table.""" client = self._ensure_prisma_client() - # Calculate hour range - hour_start = target_hour.replace(minute=0, second=0, microsecond=0) - hour_end = hour_start + timedelta(hours=1) + # Build WHERE clause for time filtering + where_conditions = [] + if start_time_utc: + where_conditions.append(f"dus.created_at >= '{start_time_utc.isoformat()}'") + if end_time_utc: + where_conditions.append(f"dus.created_at <= '{end_time_utc.isoformat()}'") - # Convert datetime objects to ISO format strings for PostgreSQL compatibility - hour_start_str = hour_start.isoformat() - hour_end_str = hour_end.isoformat() + where_clause = "" + if where_conditions: + where_clause = "WHERE " + " AND ".join(where_conditions) - # Query to get spend logs for the specific hour - query = """ - SELECT * - FROM "LiteLLM_SpendLogs" - WHERE "startTime" >= $1::timestamp - AND "startTime" < $2::timestamp - ORDER BY "startTime" ASC + # Query to get user spend data with team information + query = f""" + SELECT + dus.id, + dus.date, + dus.user_id, + dus.api_key, + dus.model, + dus.model_group, + dus.custom_llm_provider, + dus.prompt_tokens, + dus.completion_tokens, + dus.spend, + dus.api_requests, + dus.successful_requests, + dus.failed_requests, + dus.cache_creation_input_tokens, + dus.cache_read_input_tokens, + dus.created_at, + dus.updated_at, + vt.team_id, + vt.key_alias as api_key_alias, + tt.team_alias + FROM "LiteLLM_DailyUserSpend" dus + LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token + LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id + {where_clause} + ORDER BY dus.date DESC, dus.created_at DESC """ if limit: query += f" LIMIT {limit}" try: - db_response = await client.db.query_raw(query, hour_start_str, hour_end_str) - # Convert the response to polars DataFrame - return pl.DataFrame(db_response) if db_response else pl.DataFrame() + db_response = await client.db.query_raw(query) + # Convert the response to polars DataFrame with full schema inference + # This prevents schema mismatch errors when data types vary across rows + return pl.DataFrame(db_response, infer_schema_length=None) except Exception as e: - raise Exception(f"Error retrieving spend logs for hour {target_hour}: {str(e)}") - + raise Exception(f"Error retrieving usage data: {str(e)}") async def get_table_info(self) -> Dict[str, Any]: - """Get information about the LiteLLM_SpendLogs table.""" + """Get information about the daily user spend table.""" client = self._ensure_prisma_client() try: - # Get row count from SpendLogs table - spend_logs_count = await self._get_table_row_count('LiteLLM_SpendLogs') + # Get row count from user spend table + user_count = await self._get_table_row_count('LiteLLM_DailyUserSpend') - # Get column structure from spend logs table + # Get column structure from user spend table query = """ SELECT column_name, data_type, is_nullable FROM information_schema.columns - WHERE table_name = 'LiteLLM_SpendLogs' + WHERE table_name = 'LiteLLM_DailyUserSpend' ORDER BY ordinal_position; """ columns_response = await client.db.query_raw(query) return { 'columns': columns_response, - 'row_count': spend_logs_count, - 'table_breakdown': { - 'spend_logs': spend_logs_count - } + 'row_count': user_count, + 'table_name': 'LiteLLM_DailyUserSpend' } except Exception as e: raise Exception(f"Error getting table info: {str(e)}") diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index 7091ea26b95..e0263295388 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -# CHANGELOG: 2025-01-19 - Updated CBF transformation for LiteLLM_SpendLogs with hourly aggregation and team_id focus (ishaan-jaff) +# CHANGELOG: 2025-01-19 - Updated CBF transformation for daily spend tables and proper CloudZero mapping (erik.peterson) # CHANGELOG: 2025-01-19 - Migrated from pandas to polars for data transformation (erik.peterson) # CHANGELOG: 2025-01-19 - Initial CBF transformation module (erik.peterson) @@ -24,7 +24,7 @@ from typing import Any, Optional import polars as pl from ...types.integrations.cloudzero import CBFRecord -from .cz_resource_names import CZRNGenerator +from .cz_resource_names import CZEntityType, CZRNGenerator class CBFTransformer: @@ -35,160 +35,99 @@ class CBFTransformer: self.czrn_generator = CZRNGenerator() def transform(self, data: pl.DataFrame) -> pl.DataFrame: - """Transform LiteLLM SpendLogs data to hourly aggregated CBF format.""" + """Transform LiteLLM data to CBF format, dropping records with zero successful_requests or invalid CZRNs.""" if data.is_empty(): return pl.DataFrame() - # Filter out records with zero spend or invalid team_id + # Filter out records with zero successful_requests first original_count = len(data) - filtered_data = data.filter( - (pl.col('spend') > 0) & - (pl.col('team_id').is_not_null()) & - (pl.col('team_id') != "") - ) - filtered_count = len(filtered_data) - zero_spend_dropped = original_count - filtered_count + if 'successful_requests' in data.columns: + filtered_data = data.filter(pl.col('successful_requests') > 0) + zero_requests_dropped = original_count - len(filtered_data) + else: + filtered_data = data + zero_requests_dropped = 0 - if filtered_data.is_empty(): - from rich.console import Console - console = Console() - console.print(f"[yellow]āš ļø Dropped all {original_count:,} records due to zero spend or missing team_id[/yellow]") - return pl.DataFrame() - - # Aggregate data to hourly level - hourly_aggregated = self._aggregate_to_hourly(filtered_data) - - # Transform aggregated data to CBF format cbf_data = [] czrn_dropped_count = 0 - - for row in hourly_aggregated.iter_rows(named=True): + filtered_count = len(filtered_data) + + for row in filtered_data.iter_rows(named=True): try: cbf_record = self._create_cbf_record(row) + # Only include the record if CZRN generation was successful cbf_data.append(cbf_record) except Exception: # Skip records that fail CZRN generation czrn_dropped_count += 1 continue - # Print summary of transformations + # Print summary of dropped records if any from rich.console import Console console = Console() - if zero_spend_dropped > 0: - console.print(f"[yellow]āš ļø Dropped {zero_spend_dropped:,} of {original_count:,} records with zero spend or missing team_id[/yellow]") + if zero_requests_dropped > 0: + console.print(f"[yellow]āš ļø Dropped {zero_requests_dropped:,} of {original_count:,} records with zero successful_requests[/yellow]") if czrn_dropped_count > 0: - console.print(f"[yellow]āš ļø Dropped {czrn_dropped_count:,} of {len(hourly_aggregated):,} aggregated records due to invalid CZRNs[/yellow]") + console.print(f"[yellow]āš ļø Dropped {czrn_dropped_count:,} of {filtered_count:,} filtered records due to invalid CZRNs[/yellow]") if len(cbf_data) > 0: - console.print(f"[green]āœ“ Successfully transformed {len(cbf_data):,} hourly aggregated records[/green]") + console.print(f"[green]āœ“ Successfully transformed {len(cbf_data):,} records[/green]") return pl.DataFrame(cbf_data) - def _aggregate_to_hourly(self, data: pl.DataFrame) -> pl.DataFrame: - """Aggregate spend logs to hourly level by team_id, key_name, model, and tags.""" - - # Extract hour from startTime, skip tags and metadata for now - data_with_hour = data.with_columns([ - pl.col('startTime').str.to_datetime().dt.truncate('1h').alias('usage_hour'), - pl.lit([]).cast(pl.List(pl.String)).alias('parsed_tags'), # Empty tags list for now - pl.lit("").alias('key_name') # Empty key name for now - ]) - - # Skip tag explosion for now - just add a null tag column - all_data = data_with_hour.with_columns([ - pl.lit(None, dtype=pl.String).alias('tag') - ]) - - # Group by hour, team_id, key_name, model, provider, and tag - aggregated = all_data.group_by([ - 'usage_hour', - 'team_id', - 'key_name', - 'model', - 'model_group', - 'custom_llm_provider', - 'tag' - ]).agg([ - pl.col('spend').sum().alias('total_spend'), - pl.col('total_tokens').sum().alias('total_tokens'), - pl.col('prompt_tokens').sum().alias('total_prompt_tokens'), - pl.col('completion_tokens').sum().alias('total_completion_tokens'), - pl.col('request_id').count().alias('request_count'), - pl.col('api_key').first().alias('api_key_sample'), # Keep one for reference - pl.col('status').filter(pl.col('status') == 'success').count().alias('successful_requests'), - pl.col('status').filter(pl.col('status') != 'success').count().alias('failed_requests') - ]) - return aggregated - - def _create_cbf_record(self, row: dict[str, Any]) -> CBFRecord: - """Create a single CBF record from aggregated hourly spend data.""" + """Create a single CBF record from LiteLLM daily spend row.""" - # Helper function to extract scalar values from polars data - def extract_scalar(value): - if hasattr(value, 'item') and not isinstance(value, (str, int, float, bool)): - return value.item() if value is not None else None - return value + # Parse date (daily spend tables use date strings like '2025-04-19') + usage_date = self._parse_date(row.get('date')) - # Use the aggregated hour as usage time - usage_time = self._parse_datetime(extract_scalar(row.get('usage_hour'))) - - # Use team_id as the primary entity_id - entity_id = str(extract_scalar(row.get('team_id', ''))) - key_name = str(extract_scalar(row.get('key_name', ''))) - model = str(extract_scalar(row.get('model', ''))) - model_group = str(extract_scalar(row.get('model_group', ''))) - provider = str(extract_scalar(row.get('custom_llm_provider', ''))) - tag = extract_scalar(row.get('tag')) - - # Calculate aggregated metrics - total_spend = float(extract_scalar(row.get('total_spend', 0.0)) or 0.0) - total_tokens = int(extract_scalar(row.get('total_tokens', 0)) or 0) - total_prompt_tokens = int(extract_scalar(row.get('total_prompt_tokens', 0)) or 0) - total_completion_tokens = int(extract_scalar(row.get('total_completion_tokens', 0)) or 0) - request_count = int(extract_scalar(row.get('request_count', 0)) or 0) - successful_requests = int(extract_scalar(row.get('successful_requests', 0)) or 0) - failed_requests = int(extract_scalar(row.get('failed_requests', 0)) or 0) + # Calculate total tokens + prompt_tokens = int(row.get('prompt_tokens', 0)) + completion_tokens = int(row.get('completion_tokens', 0)) + total_tokens = prompt_tokens + completion_tokens # Create CloudZero Resource Name (CZRN) as resource_id - # Create a mock row for CZRN generation with team_id as entity_id - czrn_row = { - 'entity_id': entity_id, - 'entity_type': 'team', - 'model': model, - 'custom_llm_provider': provider, - 'api_key': str(extract_scalar(row.get('api_key_sample', ''))) - } - resource_id = self.czrn_generator.create_from_litellm_data(czrn_row) + resource_id = self.czrn_generator.create_from_litellm_data(row) - # Build dimensions for CloudZero tracking - dimensions = { - 'entity_type': 'team', - 'entity_id': entity_id, - 'key_name': key_name, - 'model': model, - 'model_group': model_group, - 'provider': provider, - 'request_count': str(request_count), - 'successful_requests': str(successful_requests), - 'failed_requests': str(failed_requests), - } + # Build dimensions for CloudZero + model = str(row.get('model', '')) + api_key_hash = str(row.get('api_key', ''))[:8] # First 8 chars for identification - # Add tag if present - if tag is not None and str(tag) not in ['', 'null', 'None']: - dimensions['tag'] = str(tag) + # Handle team information with fallbacks + team_id = row.get('team_id') + team_alias = row.get('team_alias') + + # Use team_alias if available, otherwise team_id, otherwise fallback to 'unknown' + entity_id = str(team_alias) if team_alias else (str(team_id) if team_id else 'unknown') + + dimensions = { + 'entity_type': CZEntityType.TEAM.value, + 'entity_id': entity_id, + 'team_id': str(team_id) if team_id else 'unknown', + 'team_alias': str(team_alias) if team_alias else 'unknown', + 'model': model, + 'model_group': str(row.get('model_group', '')), + 'provider': str(row.get('custom_llm_provider', '')), + 'api_key_prefix': api_key_hash, + 'api_key_alias': str(row.get('api_key_alias', '')), + 'api_requests': str(row.get('api_requests', 0)), + 'successful_requests': str(row.get('successful_requests', 0)), + 'failed_requests': str(row.get('failed_requests', 0)), + 'cache_creation_tokens': str(row.get('cache_creation_input_tokens', 0)), + 'cache_read_tokens': str(row.get('cache_read_input_tokens', 0)), + } # Extract CZRN components to populate corresponding CBF columns czrn_components = self.czrn_generator.extract_components(resource_id) - service_type, provider_czrn, region, owner_account_id, resource_type, cloud_local_id = czrn_components + service_type, provider, region, owner_account_id, resource_type, cloud_local_id = czrn_components # CloudZero CBF format with proper column names cbf_record = { # Required CBF fields - 'time/usage_start': usage_time.isoformat() if usage_time else None, # Required: ISO-formatted UTC datetime - 'cost/cost': total_spend, # Required: billed cost + 'time/usage_start': usage_date.isoformat() if usage_date else None, # Required: ISO-formatted UTC datetime + 'cost/cost': float(row.get('spend', 0.0)), # Required: billed cost 'resource/id': resource_id, # Required when resource tags are present # Usage metrics for token consumption @@ -206,41 +145,42 @@ class CBFTransformer: } # Add CZRN components that don't have direct CBF column mappings as resource tags - cbf_record['resource/tag:provider'] = provider_czrn # CZRN provider component + cbf_record['resource/tag:provider'] = provider # CZRN provider component cbf_record['resource/tag:model'] = cloud_local_id # CZRN cloud-local-id component (model) - + # Add resource tags for all dimensions (using resource/tag: format) for key, value in dimensions.items(): - # Ensure value is a scalar and not empty - if hasattr(value, 'item') and not isinstance(value, str): - value = value.item() if value is not None else None - if value is not None and str(value) not in ['', 'N/A', 'None', 'null']: # Only add non-empty tags + if value and value != 'N/A' and value != 'unknown': # Only add meaningful tags cbf_record[f'resource/tag:{key}'] = str(value) # Add token breakdown as resource tags for analysis - if total_prompt_tokens > 0: - cbf_record['resource/tag:prompt_tokens'] = str(total_prompt_tokens) - if total_completion_tokens > 0: - cbf_record['resource/tag:completion_tokens'] = str(total_completion_tokens) + if prompt_tokens > 0: + cbf_record['resource/tag:prompt_tokens'] = str(prompt_tokens) + if completion_tokens > 0: + cbf_record['resource/tag:completion_tokens'] = str(completion_tokens) if total_tokens > 0: cbf_record['resource/tag:total_tokens'] = str(total_tokens) return CBFRecord(cbf_record) - def _parse_datetime(self, datetime_obj) -> Optional[datetime]: - """Parse datetime object to ensure proper format.""" - if datetime_obj is None: + def _parse_date(self, date_str) -> Optional[datetime]: + """Parse date string from daily spend tables (e.g., '2025-04-19').""" + if date_str is None: return None - if isinstance(datetime_obj, datetime): - return datetime_obj + if isinstance(date_str, datetime): + return date_str - if isinstance(datetime_obj, str): + if isinstance(date_str, str): try: - # Try to parse ISO format - return pl.Series([datetime_obj]).str.to_datetime().item() + # Parse date string and set to midnight UTC for daily aggregation + return pl.Series([date_str]).str.to_datetime("%Y-%m-%d").item() except Exception: - return None + try: + # Fallback: try ISO format parsing + return pl.Series([date_str]).str.to_datetime().item() + except Exception: + return None return None diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 501185b207e..1ca45f907e1 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -119,11 +119,8 @@ class CustomGuardrail(CustomLogger): """ if "guardrails" in data: return data["guardrails"] - metadata = data.get("metadata") or {} - requested_guardrails = metadata.get("guardrails") or [] - if requested_guardrails: - return requested_guardrails - return requested_guardrails + metadata = data.get("litellm_metadata") or data.get("metadata", {}) + return metadata.get("guardrails") or [] def _guardrail_is_in_requested_guardrails( self, diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 4f9c6409770..200f2f283de 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -19,6 +19,7 @@ import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.datadog.datadog import DataDogLogger +from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_any_messages_to_chat_completion_str_messages_conversion, ) @@ -216,7 +217,7 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger): time_to_first_token=self._get_time_to_first_token_seconds(standard_logging_payload), ) - return LLMObsPayload( + payload: LLMObsPayload = LLMObsPayload( parent_id=metadata.get("parent_id", "undefined"), trace_id=standard_logging_payload.get("trace_id", str(uuid.uuid4())), span_id=metadata.get("span_id", str(uuid.uuid4())), @@ -230,6 +231,26 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger): self._get_datadog_tags(standard_logging_object=standard_logging_payload) ], ) + + apm_trace_id = self._get_apm_trace_id() + if apm_trace_id is not None: + payload["apm_id"] = apm_trace_id + + return payload + + def _get_apm_trace_id(self) -> Optional[str]: + """Retrieve the current APM trace ID if available.""" + 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: + trace_id = getattr(current_span, "trace_id", None) + if trace_id is not None: + return str(trace_id) + except Exception: + pass + return None def _assemble_error_info(self, standard_logging_payload: StandardLoggingPayload) -> Optional[DDLLMObsError]: """ diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 22ab3092901..e6f265ded58 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -15,6 +15,8 @@ from litellm.types.utils import ( StandardLoggingPayload, ) +# OpenTelemetry imports moved to individual functions to avoid import errors when not installed + if TYPE_CHECKING: from opentelemetry.sdk.trace.export import SpanExporter as _SpanExporter from opentelemetry.trace import Context as _Context @@ -41,6 +43,8 @@ else: Context = Any LITELLM_TRACER_NAME = os.getenv("OTEL_TRACER_NAME", "litellm") +LITELLM_METER_NAME = os.getenv("LITELLM_METER_NAME", "litellm") +LITELLM_LOGGER_NAME = os.getenv("LITELLM_LOGGER_NAME", "litellm") # Remove the hardcoded LITELLM_RESOURCE dictionary - we'll create it properly later RAW_REQUEST_SPAN_NAME = "raw_gen_ai_request" LITELLM_REQUEST_SPAN_NAME = "litellm_request" @@ -83,6 +87,8 @@ class OpenTelemetryConfig: exporter: Union[str, SpanExporter] = "console" endpoint: Optional[str] = None headers: Optional[str] = None + enable_metrics: bool = False + enable_events: bool = False @classmethod def from_env(cls): @@ -104,6 +110,14 @@ class OpenTelemetryConfig: headers = os.getenv( "OTEL_EXPORTER_OTLP_HEADERS", os.getenv("OTEL_HEADERS") ) # example: OTEL_HEADERS=x-honeycomb-team=B85YgLm96***" + enable_metrics: bool = ( + os.getenv("LITELLM_OTEL_INTEGRATION_ENABLE_METRICS", "false").lower() + == "true" + ) + enable_events: bool = ( + os.getenv("LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS", "false").lower() + == "true" + ) if exporter == "in_memory": return cls(exporter=InMemorySpanExporter()) @@ -111,6 +125,8 @@ class OpenTelemetryConfig: exporter=exporter, endpoint=endpoint, headers=headers, # example: OTEL_HEADERS=x-honeycomb-team=B85YgLm96***" + enable_metrics=enable_metrics, + enable_events=enable_events, ) @@ -119,27 +135,22 @@ class OpenTelemetry(CustomLogger): self, config: Optional[OpenTelemetryConfig] = None, callback_name: Optional[str] = None, + # injection points for testing + tracer_provider: Optional[Any] = None, + logger_provider: Optional[Any] = None, + meter_provider: Optional[Any] = None, **kwargs, ): - from opentelemetry import trace - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.trace import SpanKind if config is None: config = OpenTelemetryConfig.from_env() self.config = config + self.callback_name = callback_name self.OTEL_EXPORTER = self.config.exporter self.OTEL_ENDPOINT = self.config.endpoint self.OTEL_HEADERS = self.config.headers - provider = TracerProvider(resource=_get_litellm_resource()) - provider.add_span_processor(self._get_span_processor()) - self.callback_name = callback_name - - trace.set_tracer_provider(provider) - self.tracer = trace.get_tracer(LITELLM_TRACER_NAME) - - self.span_kind = SpanKind + self._init_tracing(tracer_provider) _debug_otel = str(os.getenv("DEBUG_OTEL", "False")).lower() @@ -156,6 +167,8 @@ class OpenTelemetry(CustomLogger): # init CustomLogger params super().__init__(**kwargs) + self._init_metrics(meter_provider) + self._init_logs(logger_provider) self._init_otel_logger_on_litellm_proxy() def _init_otel_logger_on_litellm_proxy(self): @@ -178,14 +191,109 @@ class OpenTelemetry(CustomLogger): litellm.service_callback.append("otel") setattr(proxy_server, "open_telemetry_logger", self) + def _init_tracing(self, tracer_provider): + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import SpanKind + + # use provided tracer or create a new one + if tracer_provider is None: + tracer_provider = TracerProvider(resource=_get_litellm_resource()) + # Only add OTLP span processor if we created the tracer provider ourselves + tracer_provider.add_span_processor(self._get_span_processor()) + + # register global provider and grab our tracer + trace.set_tracer_provider(tracer_provider) + self.tracer = trace.get_tracer(LITELLM_TRACER_NAME) + self.span_kind = SpanKind + + def _init_metrics(self, meter_provider): + if not self.config.enable_metrics: + self._operation_duration_histogram = None + self._token_usage_histogram = None + self._cost_histogram = None + return + + from opentelemetry import metrics + from opentelemetry.sdk.metrics import Histogram, MeterProvider + + # Only create OTLP infrastructure if no custom meter provider is provided + if meter_provider is None: + from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( + OTLPMetricExporter, + ) + from opentelemetry.sdk.metrics.export import ( + AggregationTemporality, + PeriodicExportingMetricReader, + ) + + _metric_exporter = OTLPMetricExporter( + endpoint=self.config.endpoint, + headers=OpenTelemetry._get_headers_dictionary(self.config.headers), + preferred_temporality={Histogram: AggregationTemporality.DELTA}, + ) + _metric_reader = PeriodicExportingMetricReader( + _metric_exporter, export_interval_millis=10000 + ) + + meter_provider = MeterProvider( + metric_readers=[_metric_reader], resource=_get_litellm_resource() + ) + meter = meter_provider.get_meter(__name__) + else: + # Use the provided meter provider as-is, without creating additional OTLP infrastructure + meter = meter_provider.get_meter(__name__) + + metrics.set_meter_provider(meter_provider) + + self._operation_duration_histogram = meter.create_histogram( + name="gen_ai.client.operation.duration", # Replace with semconv constant in otel 1.38 + description="GenAI operation duration", + unit="s", + ) + self._token_usage_histogram = meter.create_histogram( + name="gen_ai.client.token.usage", # Replace with semconv constant in otel 1.38 + description="GenAI token usage", + unit="{token}", + ) + self._cost_histogram = meter.create_histogram( + name="gen_ai.client.token.cost", + description="GenAI request cost", + unit="USD", + ) + + def _init_logs(self, logger_provider): + # nothing to do if events disabled + if not self.config.enable_events: + return + + from opentelemetry._logs import set_logger_provider + from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter + from opentelemetry.sdk._logs import LoggerProvider as OTLoggerProvider + from opentelemetry.sdk._logs.export import BatchLogRecordProcessor + + # set up log pipeline + if logger_provider is None: + logger_provider = OTLoggerProvider() + # Only add OTLP exporter if we created the logger provider ourselves + logger_provider.add_log_record_processor( + BatchLogRecordProcessor( + OTLPLogExporter( + endpoint=self.config.endpoint, + headers=self._get_headers_dictionary(self.config.headers), + ) + ) + ) + set_logger_provider(logger_provider) + def log_success_event(self, kwargs, response_obj, start_time, end_time): - self._handle_sucess(kwargs, response_obj, start_time, end_time) + self._handle_success(kwargs, response_obj, start_time, end_time) def log_failure_event(self, kwargs, response_obj, start_time, end_time): self._handle_failure(kwargs, response_obj, start_time, end_time) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self._handle_sucess(kwargs, response_obj, start_time, end_time) + self._handle_success(kwargs, response_obj, start_time, end_time) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): self._handle_failure(kwargs, response_obj, start_time, end_time) @@ -372,9 +480,9 @@ class OpenTelemetry(CustomLogger): def _get_dynamic_otel_headers_from_kwargs(self, kwargs) -> Optional[dict]: """Extract dynamic headers from kwargs if available.""" - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - kwargs.get("standard_callback_dynamic_params") - ) + standard_callback_dynamic_params: Optional[ + StandardCallbackDynamicParams + ] = kwargs.get("standard_callback_dynamic_params") if not standard_callback_dynamic_params: return None @@ -414,50 +522,185 @@ class OpenTelemetry(CustomLogger): # End of Team/Key Based Logging Control Flow ######################################################### - def _handle_sucess(self, kwargs, response_obj, start_time, end_time): - from opentelemetry import trace - from opentelemetry.trace import Status, StatusCode + def _handle_success(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug( "OpenTelemetry Logger: Logging kwargs: %s, OTEL config settings=%s", kwargs, self.config, ) + ctx, parent_span = self._get_span_context(kwargs) + + # 1. Primary span + span = self._start_primary_span(kwargs, response_obj, start_time, end_time, ctx) + + # 2. Raw‐request sub-span (if enabled) + self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) + + # 3. Guardrail span + self._create_guardrail_span(kwargs=kwargs, context=ctx) + + # 4. Metrics & cost recording + self._record_metrics(kwargs, response_obj, start_time, end_time) + + # 5. Semantic logs. + if self.config.enable_events: + self._emit_semantic_logs(kwargs, response_obj, span) + + # 6. End parent span + if parent_span is not None: + parent_span.end(end_time=self._to_ns(datetime.now())) + + def _start_primary_span(self, kwargs, response_obj, start_time, end_time, context): + from opentelemetry.trace import Status, StatusCode - _parent_context, parent_otel_span = self._get_span_context(kwargs) - # Span 1: Request sent to litellm SDK otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs) span = otel_tracer.start_span( name=self._get_span_name(kwargs), start_time=self._to_ns(start_time), - context=_parent_context, + context=context, ) span.set_status(Status(StatusCode.OK)) self.set_attributes(span, kwargs, response_obj) + span.end(end_time=self._to_ns(end_time)) + return span - if litellm.turn_off_message_logging is True: - pass - elif self.message_logging is not True: - pass - else: - # Span 2: Raw Request / Response to LLM - raw_request_span = otel_tracer.start_span( - name=RAW_REQUEST_SPAN_NAME, - start_time=self._to_ns(start_time), - context=trace.set_span_in_context(span), + def _maybe_log_raw_request( + self, kwargs, response_obj, start_time, end_time, parent_span + ): + from opentelemetry import trace + from opentelemetry.trace import Status, StatusCode + + # only log raw LLM request/response if message_logging is on and not globally turned off + if litellm.turn_off_message_logging or not self.message_logging: + return + + otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs) + raw_span = otel_tracer.start_span( + name=RAW_REQUEST_SPAN_NAME, + start_time=self._to_ns(start_time), + context=trace.set_span_in_context(parent_span), + ) + raw_span.set_status(Status(StatusCode.OK)) + self.set_raw_request_attributes(raw_span, kwargs, response_obj) + raw_span.end(end_time=self._to_ns(end_time)) + + def _record_metrics(self, kwargs, response_obj, start_time, end_time): + duration_s = (end_time - start_time).total_seconds() + params = kwargs.get("litellm_params") or {} + provider = params.get("custom_llm_provider", "Unknown") + + common_attrs = { + "gen_ai.operation.name": "chat", + "gen_ai.system": provider, + "gen_ai.request.model": kwargs.get("model"), + "gen_ai.framework": "litellm", + } + + std_log = kwargs.get("standard_logging_object") + md = getattr(std_log, "metadata", None) or (std_log or {}).get("metadata", {}) + for key in [ + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_team_id", + "user_api_key_org_id", + "user_api_key_user_id", + "user_api_key_team_alias", + "user_api_key_user_email", + "spend_logs_metadata", + "requester_ip_address", + "requester_metadata", + "user_api_key_end_user_id", + "prompt_management_metadata", + "applied_guardrails", + "mcp_tool_call_metadata", + "vector_store_request_metadata", + ]: + if md.get(key) is not None: + common_attrs[f"metadata.{key}"] = str(md[key]) + + if self._operation_duration_histogram: + self._operation_duration_histogram.record( + duration_s, attributes=common_attrs + ) + if ( + response_obj + and (usage := response_obj.get("usage")) + and self._token_usage_histogram + ): + in_attrs = {**common_attrs, "gen_ai.token.type": "input"} + out_attrs = {**common_attrs, "gen_ai.token.type": "completion"} + self._token_usage_histogram.record( + usage.get("prompt_tokens", 0), attributes=in_attrs + ) + self._token_usage_histogram.record( + usage.get("completion_tokens", 0), attributes=out_attrs + ) + + cost = kwargs.get("response_cost") + if self._cost_histogram and cost: + self._cost_histogram.record(cost, attributes=common_attrs) + + def _emit_semantic_logs(self, kwargs, response_obj, span: Span): + if not self.config.enable_events: + return + + from opentelemetry._logs import get_logger, LogRecord + otel_logger = get_logger(LITELLM_LOGGER_NAME) + + parent_ctx = span.get_span_context() + provider = (kwargs.get("litellm_params") or {}).get( + "custom_llm_provider", "Unknown" + ) + + # per-message events + for msg in kwargs.get("messages", []): + role = msg.get("role", "user") + attrs = {"event_name": "gen_ai.content.prompt", "gen_ai.system": provider} + if role == "tool" and msg.get("id"): + attrs["id"] = msg["id"] + if self.message_logging and msg.get("content"): + attrs["gen_ai.prompt"] = msg["content"] + + otel_logger.emit( + LogRecord( + attributes=attrs, + body=msg.copy(), + trace_id=parent_ctx.trace_id, + span_id=parent_ctx.span_id, + trace_flags=parent_ctx.trace_flags, + ) ) - raw_request_span.set_status(Status(StatusCode.OK)) - self.set_raw_request_attributes(raw_request_span, kwargs, response_obj) - raw_request_span.end(end_time=self._to_ns(end_time)) + # per-choice events + for idx, choice in enumerate(response_obj.get("choices", [])): + attrs = { + "event_name": "gen_ai.content.completion", + "gen_ai.system": provider, + "index": idx, + "finish_reason": choice.get("finish_reason"), + } + body_msg = choice.get("message", {}) + if self.message_logging and body_msg.get("content"): + attrs["message.content"] = body_msg["content"] + body = { + "index": idx, + "finish_reason": choice.get("finish_reason"), + "message": {"role": body_msg.get("role", "assistant")}, + } + if self.message_logging and body_msg.get("content"): + body["message"]["content"] = body_msg["content"] - span.end(end_time=self._to_ns(end_time)) + otel_logger.emit( + LogRecord( + attributes=attrs, + body=body, + trace_id=parent_ctx.trace_id, + span_id=parent_ctx.span_id, + trace_flags=parent_ctx.trace_flags, + ) + ) - # Create span for guardrail information - self._create_guardrail_span(kwargs=kwargs, context=_parent_context) - - if parent_otel_span is not None: - parent_otel_span.end(end_time=self._to_ns(datetime.now())) def _create_guardrail_span( self, kwargs: Optional[dict], context: Optional[Context] diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index fd82ecdf2b2..af51fe9ab79 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -38,6 +38,7 @@ try: from litellm_enterprise.integrations.prometheus import PrometheusLogger except Exception: PrometheusLogger = None +from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger from litellm.integrations.dotprompt import DotpromptManager from litellm.integrations.s3_v2 import S3Logger from litellm.integrations.sqs import SQSLogger @@ -86,6 +87,7 @@ class CustomLoggerRegistry: "dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler, "vector_store_pre_call_hook": VectorStorePreCallHook, "dotprompt": DotpromptManager, + "cloudzero": CloudZeroLogger, } try: diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 08f1d4c82d0..08e5323c30c 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -1,7 +1,7 @@ """ Helper utilities for parsing durations - 1s, 1d, 10d, 30d, 1mo, 2mo -duration_in_seconds is used in diff parts of the code base, example +duration_in_seconds is used in diff parts of the code base, example - Router - Provider budget routing - Proxy - Key, Team Generation """ @@ -192,6 +192,10 @@ def _handle_day_reset( current_time: datetime, base_midnight: datetime, value: int, timezone: timezone ) -> datetime: """Handle day-based reset times.""" + # Handle zero value - immediate expiration + if value == 0: + return current_time + if value == 1: # Daily reset at midnight return base_midnight + timedelta(days=1) elif value == 7: # Weekly reset on Monday at midnight @@ -234,6 +238,10 @@ def _handle_hour_reset( current_time: datetime, base_midnight: datetime, value: int ) -> datetime: """Handle hour-based reset times.""" + # Handle zero value - immediate expiration + if value == 0: + return current_time + current_hour = current_time.hour current_minute = current_time.minute current_second = current_time.second @@ -266,6 +274,10 @@ def _handle_minute_reset( current_time: datetime, base_midnight: datetime, value: int ) -> datetime: """Handle minute-based reset times.""" + # Handle zero value - immediate expiration + if value == 0: + return current_time + current_hour = current_time.hour current_minute = current_time.minute current_second = current_time.second @@ -306,6 +318,10 @@ def _handle_second_reset( current_time: datetime, base_midnight: datetime, value: int ) -> datetime: """Handle second-based reset times.""" + # Handle zero value - immediate expiration + if value == 0: + return current_time + current_hour = current_time.hour current_minute = current_time.minute current_second = current_time.second diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 2049480e264..c784568cc9f 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -320,6 +320,7 @@ def get_llm_provider( # noqa: PLR0915 or model in litellm.vertex_embedding_models or model in litellm.vertex_vision_models or model in litellm.vertex_ai_image_models + or model in litellm.vertex_ai_video_models ): custom_llm_provider = "vertex_ai" ## ai21 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7bc7702684d..0b152e0dda3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1164,7 +1164,6 @@ class Logging(LiteLLMLoggingBaseClass): used for consistent cost calculation across response headers + logging integrations. """ - if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"): hidden_params = getattr(result, "_hidden_params", {}) if ( @@ -3361,7 +3360,14 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 galileo_logger = GalileoObserve() _in_memory_loggers.append(galileo_logger) return galileo_logger # type: ignore - + elif logging_integration == "cloudzero": + from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger + for callback in _in_memory_loggers: + if isinstance(callback, CloudZeroLogger): + return callback # type: ignore + cloudzero_logger = CloudZeroLogger() + _in_memory_loggers.append(cloudzero_logger) + return cloudzero_logger # type: ignore elif logging_integration == "deepeval": for callback in _in_memory_loggers: if isinstance(callback, DeepEvalLogger): @@ -3581,6 +3587,11 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, GalileoObserve): return callback + elif logging_integration == "cloudzero": + from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger + for callback in _in_memory_loggers: + if isinstance(callback, CloudZeroLogger): + return callback elif logging_integration == "deepeval": for callback in _in_memory_loggers: if isinstance(callback, DeepEvalLogger): diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 7ad0038ecb2..c714e36b5f9 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -1,5 +1,6 @@ import json from typing import Any, Union + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH diff --git a/litellm/llms/base_llm/__init__.py b/litellm/llms/base_llm/__init__.py index 187c985fd67..665e242969c 100644 --- a/litellm/llms/base_llm/__init__.py +++ b/litellm/llms/base_llm/__init__.py @@ -1,5 +1,6 @@ from .anthropic_messages.transformation import BaseAnthropicMessagesConfig from .audio_transcription.transformation import BaseAudioTranscriptionConfig +from .batches.transformation import BaseBatchesConfig from .chat.transformation import BaseConfig from .embedding.transformation import BaseEmbeddingConfig from .image_edit.transformation import BaseImageEditConfig @@ -12,4 +13,5 @@ __all__ = [ "BaseAnthropicMessagesConfig", "BaseEmbeddingConfig", "BaseImageEditConfig", + "BaseBatchesConfig", ] diff --git a/litellm/llms/base_llm/batches/transformation.py b/litellm/llms/base_llm/batches/transformation.py new file mode 100644 index 00000000000..1d3e54fae67 --- /dev/null +++ b/litellm/llms/base_llm/batches/transformation.py @@ -0,0 +1,176 @@ +import types +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import httpx +from httpx import Headers + +from litellm.types.llms.openai import ( + AllMessageValues, + CreateBatchRequest, +) +from litellm.types.utils import LiteLLMBatch, LlmProviders + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + from ..chat.transformation import BaseLLMException as _BaseLLMException + + LiteLLMLoggingObj = _LiteLLMLoggingObj + BaseLLMException = _BaseLLMException +else: + LiteLLMLoggingObj = Any + BaseLLMException = Any + + +class BaseBatchesConfig(ABC): + """ + Abstract base class for batch processing configurations across different LLM providers. + + This class defines the interface that all provider-specific batch configurations + must implement to work with LiteLLM's unified batch processing system. + """ + + def __init__(self): + pass + + @property + @abstractmethod + def custom_llm_provider(self) -> LlmProviders: + """Return the LLM provider type for this configuration.""" + pass + + @classmethod + def get_config(cls): + """Get configuration dictionary for this class.""" + return { + k: v + for k, v in cls.__dict__.items() + if not k.startswith("__") + and not k.startswith("_abc") + and not isinstance( + v, + ( + types.FunctionType, + types.BuiltinFunctionType, + classmethod, + staticmethod, + ), + ) + and v is not None + } + + @abstractmethod + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate and prepare environment-specific headers and parameters. + + Args: + headers: HTTP headers dictionary + model: Model name + messages: List of messages + optional_params: Optional parameters + litellm_params: LiteLLM parameters + api_key: API key + api_base: API base URL + + Returns: + Updated headers dictionary + """ + pass + + @abstractmethod + def get_complete_batch_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: Dict, + litellm_params: Dict, + data: CreateBatchRequest, + ) -> str: + """ + Get the complete URL for batch creation request. + + Args: + api_base: Base API URL + api_key: API key + model: Model name + optional_params: Optional parameters + litellm_params: LiteLLM parameters + data: Batch creation request data + + Returns: + Complete URL for the batch request + """ + pass + + @abstractmethod + def transform_create_batch_request( + self, + model: str, + create_batch_data: CreateBatchRequest, + optional_params: dict, + litellm_params: dict, + ) -> Union[bytes, str, Dict[str, Any]]: + """ + Transform the batch creation request to provider-specific format. + + Args: + model: Model name + create_batch_data: Batch creation request data + optional_params: Optional parameters + litellm_params: LiteLLM parameters + + Returns: + Transformed request data + """ + pass + + @abstractmethod + def transform_create_batch_response( + self, + model: Optional[str], + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> LiteLLMBatch: + """ + Transform provider-specific batch response to LiteLLM format. + + Args: + model: Model name + raw_response: Raw HTTP response + logging_obj: Logging object + litellm_params: LiteLLM parameters + + Returns: + LiteLLM batch object + """ + pass + + @abstractmethod + def get_error_class( + self, error_message: str, status_code: int, headers: Union[Dict, Headers] + ) -> "BaseLLMException": + """ + Get the appropriate error class for this provider. + + Args: + error_message: Error message + status_code: HTTP status code + headers: Response headers + + Returns: + Provider-specific exception class + """ + pass diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 5c37a8b7547..35b76479cdc 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -35,6 +35,16 @@ class BaseFilesConfig(BaseConfig): def custom_llm_provider(self) -> LlmProviders: pass + @property + def file_upload_http_method(self) -> str: + """ + HTTP method to use for file uploads. + Override this in provider configs if they need different methods. + Default is POST (used by most providers like OpenAI, Anthropic). + S3-based providers like Bedrock should return "PUT". + """ + return "POST" + @abstractmethod def get_supported_openai_params( self, model: str diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 60d89c1610f..f925e6819dc 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -31,30 +31,26 @@ class BasePassthroughConfig(BaseLLMModelInfo): Args: endpoint: str - the endpoint to add to the url base_target_url: str - the base url to add the endpoint to - request_query_params: dict - the query params to add to the url + request_query_params: Optional[dict] - the query params to add to the url Returns: - str - the formatted url + httpx.URL - the formatted url """ from urllib.parse import urlencode import httpx - encoded_endpoint = httpx.URL(endpoint).path + base = base_target_url.rstrip('/') + endpoint = endpoint.lstrip('/') + full_url = f"{base}/{endpoint}" - # Ensure endpoint starts with '/' for proper URL construction - if not encoded_endpoint.startswith("/"): - encoded_endpoint = "/" + encoded_endpoint - - # Construct the full target URL using httpx - base_url = httpx.URL(base_target_url) - updated_url = base_url.copy_with(path=encoded_endpoint) + url = httpx.URL(full_url) if request_query_params: - # Create a new URL with the merged query params - updated_url = updated_url.copy_with( + url = url.copy_with( query=urlencode(request_query_params).encode("ascii") ) - return updated_url + + return url @abstractmethod def get_complete_url( diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py new file mode 100644 index 00000000000..ce580ebc624 --- /dev/null +++ b/litellm/llms/bedrock/batches/transformation.py @@ -0,0 +1,254 @@ +import os +import time +from typing import Any, Dict, List, Literal, Optional, Union, cast + +from httpx import Headers, Response + +from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.bedrock import ( + BedrockBatchJobStatus, + BedrockCreateBatchRequest, + BedrockCreateBatchResponse, + BedrockInputDataConfig, + BedrockOutputDataConfig, + BedrockS3InputDataConfig, + BedrockS3OutputDataConfig, +) +from litellm.types.llms.openai import ( + AllMessageValues, + CreateBatchRequest, +) +from litellm.types.utils import LiteLLMBatch, LlmProviders + +from ..base_aws_llm import BaseAWSLLM +from ..common_utils import CommonBatchFilesUtils + + +class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): + """ + Config for Bedrock Batches - handles batch job creation and management for Bedrock + """ + + def __init__(self): + super().__init__() + self.common_utils = CommonBatchFilesUtils() + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.BEDROCK + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate and prepare environment for Bedrock batch requests. + AWS credentials are handled by BaseAWSLLM. + """ + # Add any Bedrock-specific headers if needed + return headers + + def get_complete_batch_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: Dict, + litellm_params: Dict, + data: CreateBatchRequest, + ) -> str: + """ + Get the complete URL for Bedrock batch creation. + Bedrock batch jobs are created via the model invocation job API. + """ + aws_region_name = self._get_aws_region_name(optional_params, model) + + # Bedrock model invocation job endpoint + # Format: https://bedrock.{region}.amazonaws.com/model-invocation-job + bedrock_endpoint = f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job" + + return bedrock_endpoint + + + + + + + + def transform_create_batch_request( + self, + model: str, + create_batch_data: CreateBatchRequest, + optional_params: dict, + litellm_params: dict, + ) -> Dict[str, Any]: + """ + Transform the batch creation request to Bedrock format. + + Bedrock batch inference requires: + - modelId: The Bedrock model ID + - jobName: Unique name for the batch job + - inputDataConfig: Configuration for input data (S3 location) + - outputDataConfig: Configuration for output data (S3 location) + - roleArn: IAM role ARN for the batch job + """ + # Get required parameters + input_file_id = create_batch_data.get("input_file_id") + if not input_file_id: + raise ValueError("input_file_id is required for Bedrock batch creation") + + # Extract S3 information from file ID using common utility + input_bucket, input_key = self.common_utils.parse_s3_uri(input_file_id) + + # Get output S3 configuration + output_bucket = litellm_params.get("s3_output_bucket_name") or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") + if not output_bucket: + # Use same bucket as input if no output bucket specified + output_bucket = input_bucket + + # Get IAM role ARN + role_arn = ( + litellm_params.get("aws_batch_role_arn") + or optional_params.get("aws_batch_role_arn") + or os.getenv("AWS_BATCH_ROLE_ARN") + ) + if not role_arn: + raise ValueError( + "AWS IAM role ARN is required for Bedrock batch jobs. " + "Set 'aws_batch_role_arn' in litellm_params or AWS_BATCH_ROLE_ARN env var" + ) + + # Get the actual Bedrock model ID using common utility + bedrock_model_id = self.common_utils.extract_model_from_s3_file_path(input_file_id, optional_params) + + if not bedrock_model_id: + raise ValueError("Could not determine Bedrock model ID. Ensure the model is specified in the input file or passed as a parameter.") + + # Generate job name with the correct model ID using common utility + job_name = self.common_utils.generate_unique_job_name(bedrock_model_id, prefix="litellm") + output_key = f"litellm-batch-outputs/{job_name}/" + + # Build input data config + input_data_config: BedrockInputDataConfig = { + "s3InputDataConfig": BedrockS3InputDataConfig( + s3Uri=f"s3://{input_bucket}/{input_key}" + ) + } + + # Build output data config + output_data_config: BedrockOutputDataConfig = { + "s3OutputDataConfig": BedrockS3OutputDataConfig( + s3Uri=f"s3://{output_bucket}/{output_key}" + ) + } + + # Create Bedrock batch request with proper typing + bedrock_request: BedrockCreateBatchRequest = { + "modelId": bedrock_model_id, + "jobName": job_name, + "inputDataConfig": input_data_config, + "outputDataConfig": output_data_config, + "roleArn": role_arn + } + + # Add optional parameters if provided + completion_window = create_batch_data.get("completion_window") + if completion_window: + # Map OpenAI completion window to Bedrock timeout + # OpenAI uses "24h", Bedrock expects timeout in hours + if completion_window == "24h": + bedrock_request["timeoutDurationInHours"] = 24 + + # For Bedrock, we need to return a pre-signed request with AWS auth headers + # Use common utility for AWS signing + endpoint_url = f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job" + signed_headers, signed_data = self.common_utils.sign_aws_request( + service_name="bedrock", + data=bedrock_request, + endpoint_url=endpoint_url, + optional_params=optional_params, + method="POST" + ) + + # Return a pre-signed request format that the HTTP handler can use + return { + "method": "POST", + "url": endpoint_url, + "headers": signed_headers, + "data": signed_data.decode('utf-8') + } + + def transform_create_batch_response( + self, + model: Optional[str], + raw_response: Response, + logging_obj: Any, + litellm_params: dict, + ) -> LiteLLMBatch: + """ + Transform Bedrock batch creation response to LiteLLM format. + """ + try: + response_data: BedrockCreateBatchResponse = raw_response.json() + except Exception as e: + raise ValueError(f"Failed to parse Bedrock batch response: {e}") + + # Extract information from typed Bedrock response + job_arn = response_data.get("jobArn", "") + status: BedrockBatchJobStatus = response_data.get("status", "Submitted") + + # Map Bedrock status to OpenAI-compatible status + status_mapping: Dict[BedrockBatchJobStatus, str] = { + "Submitted": "validating", + "InProgress": "in_progress", + "Completed": "completed", + "Failed": "failed", + "Stopping": "cancelling", + "Stopped": "cancelled" + } + + openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status, "validating")) + + # Get original request data from litellm_params if available + original_request = litellm_params.get("original_batch_request", {}) + + # Create LiteLLM batch object + return LiteLLMBatch( + id=job_arn, # Use ARN as the batch ID + object="batch", + endpoint=original_request.get("endpoint", "/v1/chat/completions"), + errors=None, + input_file_id=original_request.get("input_file_id", ""), + completion_window=original_request.get("completion_window", "24h"), + status=openai_status, + output_file_id=None, # Will be populated when job completes + error_file_id=None, + created_at=int(time.time()), + in_progress_at=int(time.time()) if status == "InProgress" else None, + expires_at=None, + finalizing_at=None, + completed_at=None, + failed_at=None, + expired_at=None, + cancelling_at=None, + cancelled_at=None, + request_counts=None, + metadata=original_request.get("metadata", {}), + ) + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[Dict, Headers] + ) -> BaseLLMException: + """ + Get Bedrock-specific error class using common utility. + """ + return self.common_utils.get_error_class(error_message, status_code, headers) + + diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b93ca94bed4..fda9220ff7d 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -10,6 +10,8 @@ from typing import List, Literal, Optional, Tuple, Union, cast, overload import httpx import litellm +from litellm._logging import verbose_logger +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.llm_response_utils.convert_dict_to_response import ( @@ -48,14 +50,19 @@ from litellm.types.utils import ( ) from litellm.utils import add_dummy_tool, has_tool_call_blocks, supports_reasoning -from ..common_utils import BedrockError, BedrockModelInfo, get_bedrock_tool_name, get_anthropic_beta_from_headers +from ..common_utils import ( + BedrockError, + BedrockModelInfo, + get_anthropic_beta_from_headers, + get_bedrock_tool_name, +) # Computer use tool prefixes supported by Bedrock BEDROCK_COMPUTER_USE_TOOLS = [ "computer_use_preview", "computer_", "bash_", - "text_editor_" + "text_editor_", ] @@ -163,7 +170,9 @@ class AmazonConverseConfig(BaseConfig): # only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html supported_params.append("tool_choice") - if ( + if "gpt-oss" in model: + supported_params.append("reasoning_effort") + elif ( "claude-3-7" in model or "claude-sonnet-4" in model or "claude-opus-4" in model @@ -233,7 +242,7 @@ class AmazonConverseConfig(BaseConfig): """Check if computer use tools are being used in the request.""" if tools is None: return False - + for tool in tools: if "type" in tool: tool_type = tool["type"] @@ -247,17 +256,17 @@ class AmazonConverseConfig(BaseConfig): ) -> List[dict]: """Transform computer use tools to Bedrock format.""" transformed_tools: List[dict] = [] - + for tool in computer_use_tools: tool_type = tool.get("type", "") - + # Check if this is a computer use tool with the startswith method is_computer_use_tool = False for computer_use_prefix in BEDROCK_COMPUTER_USE_TOOLS: if tool_type.startswith(computer_use_prefix): is_computer_use_tool = True break - + transformed_tool: dict = {} if is_computer_use_tool: if tool_type.startswith("computer_") and "function" in tool: @@ -266,7 +275,7 @@ class AmazonConverseConfig(BaseConfig): transformed_tool = { "type": tool_type, "name": func.get("name", "computer"), - **func.get("parameters", {}) + **func.get("parameters", {}), } else: # Direct tools - just need to ensure name is present @@ -279,27 +288,29 @@ class AmazonConverseConfig(BaseConfig): else: # Pass through other tools as-is transformed_tool = dict(tool) - + transformed_tools.append(transformed_tool) - + return transformed_tools def _separate_computer_use_tools( self, tools: List[OpenAIChatCompletionToolParam], model: str - ) -> Tuple[List[OpenAIChatCompletionToolParam], List[OpenAIChatCompletionToolParam]]: + ) -> Tuple[ + List[OpenAIChatCompletionToolParam], List[OpenAIChatCompletionToolParam] + ]: """ Separate computer use tools from regular function tools. - + Args: tools: List of tools to separate model: The model name to check if it supports computer use - + Returns: Tuple of (computer_use_tools, regular_tools) """ computer_use_tools = [] regular_tools = [] - + for tool in tools: if "type" in tool: tool_type = tool["type"] @@ -314,15 +325,12 @@ class AmazonConverseConfig(BaseConfig): regular_tools.append(tool) else: regular_tools.append(tool) - + return computer_use_tools, regular_tools - - def _create_json_tool_call_for_response_format( self, json_schema: Optional[dict] = None, - schema_name: str = "json_tool_call", description: Optional[str] = None, ) -> ChatCompletionToolParam: """ @@ -344,10 +352,12 @@ class AmazonConverseConfig(BaseConfig): "properties": {}, } else: + # Use the schema as-is for Bedrock + # Bedrock requires the tool schema to be of type "object" and doesn't need unwrapping _input_schema = json_schema tool_param_function_chunk = ChatCompletionToolParamFunctionChunk( - name=schema_name, parameters=_input_schema + name=RESPONSE_FORMAT_TOOL_NAME, parameters=_input_schema ) if description: tool_param_function_chunk["description"] = description @@ -386,56 +396,9 @@ class AmazonConverseConfig(BaseConfig): for param, value in non_default_params.items(): if param == "response_format" and isinstance(value, dict): - ignore_response_format_types = ["text"] - if value["type"] in ignore_response_format_types: # value is a no-op - continue - - json_schema: Optional[dict] = None - schema_name: str = "" - description: Optional[str] = None - if "response_schema" in value: - json_schema = value["response_schema"] - schema_name = "json_tool_call" - elif "json_schema" in value: - json_schema = value["json_schema"]["schema"] - schema_name = value["json_schema"]["name"] - description = value["json_schema"].get("description") - - if "type" in value and value["type"] == "text": - continue - - """ - Follow similar approach to anthropic - translate to a single tool call. - - When using tools in this way: - https://docs.anthropic.com/en/docs/build-with-claude/tool-use#json-mode - - You usually want to provide a single tool - - You should set tool_choice (see Forcing tool use) to instruct the model to explicitly use that tool - - Remember that the model will pass the input to the tool, so the name of the tool and description should be from the model’s perspective. - """ - _tool = self._create_json_tool_call_for_response_format( - json_schema=json_schema, - schema_name=schema_name if schema_name != "" else "json_tool_call", - description=description, + optional_params = self._translate_response_format_param( + value=value, model=model, optional_params=optional_params, non_default_params=non_default_params, is_thinking_enabled=is_thinking_enabled ) - optional_params = self._add_tools_to_optional_params( - optional_params=optional_params, tools=[_tool] - ) - - if ( - litellm.utils.supports_tool_choice( - model=model, custom_llm_provider=self.custom_llm_provider - ) - and not is_thinking_enabled - ): - - optional_params["tool_choice"] = ToolChoiceValuesBlock( - tool=SpecificToolChoiceBlock( - name=schema_name if schema_name != "" else "json_tool_call" - ) - ) - optional_params["json_mode"] = True - if non_default_params.get("stream", False) is True: - optional_params["fake_stream"] = True if param == "max_tokens" or param == "max_completion_tokens": optional_params["maxTokens"] = value if param == "stream": @@ -466,14 +429,82 @@ class AmazonConverseConfig(BaseConfig): if param == "thinking": optional_params["thinking"] = value elif param == "reasoning_effort" and isinstance(value, str): - optional_params["thinking"] = AnthropicConfig._map_reasoning_effort( - value - ) + if "gpt-oss" in model: + # GPT-OSS models: keep reasoning_effort as-is + # It will be passed through to additionalModelRequestFields + optional_params["reasoning_effort"] = value + else: + # Anthropic and other models: convert to thinking parameter + optional_params["thinking"] = AnthropicConfig._map_reasoning_effort( + value + ) - self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=optional_params + # Only update thinking tokens for non-GPT-OSS models + if "gpt-oss" not in model: + self.update_optional_params_with_thinking_tokens( + non_default_params=non_default_params, optional_params=optional_params + ) + + return optional_params + + def _translate_response_format_param( + self, + value: dict, + model: str, + optional_params: dict, + non_default_params: dict, + is_thinking_enabled: bool, + ) -> dict: + """ + Handles translation of response_format parameter to Bedrock format. + + Returns `optional_params` with the translated response_format parameter. + """ + ignore_response_format_types = ["text"] + if value["type"] in ignore_response_format_types: # value is a no-op + return optional_params + + json_schema: Optional[dict] = None + description: Optional[str] = None + if "response_schema" in value: + json_schema = value["response_schema"] + elif "json_schema" in value: + json_schema = value["json_schema"]["schema"] + description = value["json_schema"].get("description") + + if "type" in value and value["type"] == "text": + return optional_params + + """ + Follow similar approach to anthropic - translate to a single tool call. + + When using tools in this way: - https://docs.anthropic.com/en/docs/build-with-claude/tool-use#json-mode + - You usually want to provide a single tool + - You should set tool_choice (see Forcing tool use) to instruct the model to explicitly use that tool + - Remember that the model will pass the input to the tool, so the name of the tool and description should be from the model’s perspective. + """ + _tool = self._create_json_tool_call_for_response_format( + json_schema=json_schema, + description=description, + ) + optional_params = self._add_tools_to_optional_params( + optional_params=optional_params, tools=[_tool] ) + if ( + litellm.utils.supports_tool_choice( + model=model, custom_llm_provider=self.custom_llm_provider + ) + and not is_thinking_enabled + ): + + optional_params["tool_choice"] = ToolChoiceValuesBlock( + tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME) + ) + optional_params["json_mode"] = True + if non_default_params.get("stream", False) is True: + optional_params["fake_stream"] = True + return optional_params def update_optional_params_with_thinking_tokens( @@ -597,7 +628,6 @@ class AmazonConverseConfig(BaseConfig): return {} - def _transform_request_helper( self, model: str, @@ -653,36 +683,38 @@ class AmazonConverseConfig(BaseConfig): ) original_tools = inference_params.pop("tools", []) - + # Initialize bedrock_tools bedrock_tools: List[ToolBlock] = [] - + # Collect anthropic_beta values from user headers anthropic_beta_list = [] if headers: user_betas = get_anthropic_beta_from_headers(headers) anthropic_beta_list.extend(user_betas) - + # Only separate tools if computer use tools are actually present if original_tools and self.is_computer_use_tool_used(original_tools, model): # Separate computer use tools from regular function tools computer_use_tools, regular_tools = self._separate_computer_use_tools( original_tools, model ) - + # Process regular function tools using existing logic bedrock_tools = _bedrock_tools_pt(regular_tools) - + # Add computer use tools and anthropic_beta if needed (only when computer use tools are present) if computer_use_tools: anthropic_beta_list.append("computer-use-2024-10-22") # Transform computer use tools to proper Bedrock format - transformed_computer_tools = self._transform_computer_use_tools(computer_use_tools) + transformed_computer_tools = self._transform_computer_use_tools( + computer_use_tools + ) additional_request_params["tools"] = transformed_computer_tools else: # No computer use tools, process all tools as regular tools bedrock_tools = _bedrock_tools_pt(original_tools) - + # Set anthropic_beta in additional_request_params if we have any beta features if anthropic_beta_list: # Remove duplicates while preserving order @@ -693,7 +725,7 @@ class AmazonConverseConfig(BaseConfig): unique_betas.append(beta) seen.add(beta) additional_request_params["anthropic_beta"] = unique_betas - + bedrock_tool_config: Optional[ToolConfigBlock] = None if len(bedrock_tools) > 0: tool_choice_values: ToolChoiceValuesBlock = inference_params.pop( @@ -1119,10 +1151,37 @@ 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: - # to support 'json_schema' logic on bedrock models + 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: + import json + + # 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 else: chat_completion_message["tool_calls"] = tools @@ -1182,7 +1241,6 @@ class AmazonConverseConfig(BaseConfig): if api_key: headers["Authorization"] = f"Bearer {api_key}" return headers - def should_fake_stream( self, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index c76fc0a80c3..831a6da93b3 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -6,6 +6,9 @@ import json import os from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union +if TYPE_CHECKING: + from litellm.types.llms.bedrock import BedrockCreateBatchRequest + import httpx import litellm @@ -608,3 +611,218 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: # Split comma-separated values and strip whitespace return [beta.strip() for beta in anthropic_beta_header.split(",")] + + +class CommonBatchFilesUtils: + """ + Common utilities for Bedrock batch and file operations. + Provides shared functionality to reduce code duplication between batches and files. + """ + + def __init__(self): + # Import here to avoid circular imports + from .base_aws_llm import BaseAWSLLM + self._base_aws = BaseAWSLLM() + + def get_bedrock_model_id_from_litellm_model(self, model: str) -> str: + """ + Extract the actual Bedrock model ID from LiteLLM model name. + + Args: + model: LiteLLM model name (e.g., "bedrock/anthropic.claude-3-sonnet-20240229-v1:0") + + Returns: + Bedrock model ID (e.g., "anthropic.claude-3-sonnet-20240229-v1:0") + """ + if model.startswith("bedrock/"): + return model[8:] # Remove "bedrock/" prefix + return model + + def parse_s3_uri(self, s3_uri: str) -> tuple: + """ + Parse S3 URI into bucket and key components. + + Args: + s3_uri: S3 URI (e.g., "s3://bucket/key/path") + + Returns: + Tuple of (bucket, key) + + Raises: + ValueError: If URI format is invalid + """ + if not s3_uri.startswith("s3://"): + raise ValueError(f"Invalid S3 URI format: {s3_uri}") + + s3_parts = s3_uri[5:].split("/", 1) # Remove "s3://" and split on first "/" + if len(s3_parts) != 2: + raise ValueError(f"Invalid S3 URI format: {s3_uri}") + + return s3_parts[0], s3_parts[1] # bucket, key + + def extract_model_from_s3_file_path(self, s3_uri: str, optional_params: dict) -> str: + """ + Extract model ID from S3 file path. + + The Bedrock file transformation creates S3 objects with the model name embedded: + Format: s3://bucket/litellm-bedrock-files-{model}-{uuid}.jsonl + """ + # Check if model is provided in optional_params first + if "model" in optional_params and optional_params["model"]: + return self.get_bedrock_model_id_from_litellm_model(optional_params["model"]) + + # Extract model from S3 URI path + # Expected format: s3://bucket/litellm-bedrock-files-{model}-{uuid}.jsonl + try: + bucket, object_key = self.parse_s3_uri(s3_uri) + + # Extract model from object key if it follows our naming pattern + if object_key.startswith("litellm-bedrock-files-"): + # Remove prefix and suffix to get model part + model_part = object_key[22:] # Remove "litellm-bedrock-files-" + # Find the last dash before the UUID + parts = model_part.split("-") + if len(parts) > 1: + # Reconstruct model name (everything except the last UUID part and .jsonl) + model_name = "-".join(parts[:-1]) + if model_name.endswith(".jsonl"): + model_name = model_name[:-6] # Remove .jsonl + return model_name + except Exception: + pass + + # Fallback to default model + return "anthropic.claude-3-5-sonnet-20240620-v1:0" + + def sign_aws_request( + self, + service_name: str, + data: Union[str, dict, "BedrockCreateBatchRequest"], + endpoint_url: str, + optional_params: dict, + method: str = "POST", + ) -> tuple: + """ + Sign AWS request using Signature Version 4. + + Args: + service_name: AWS service name ("bedrock" or "s3") + data: Request data (string or dict) + endpoint_url: Full endpoint URL + optional_params: Optional parameters containing AWS credentials + method: HTTP method (default: POST) + + Returns: + Tuple of (signed_headers, signed_data) + """ + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + # Get AWS credentials using existing methods + aws_region_name = self._base_aws._get_aws_region_name( + optional_params=optional_params, model="" + ) + credentials = self._base_aws.get_credentials( + aws_access_key_id=optional_params.get("aws_access_key_id"), + aws_secret_access_key=optional_params.get("aws_secret_access_key"), + aws_session_token=optional_params.get("aws_session_token"), + aws_region_name=aws_region_name, + aws_session_name=optional_params.get("aws_session_name"), + aws_profile_name=optional_params.get("aws_profile_name"), + aws_role_name=optional_params.get("aws_role_name"), + aws_web_identity_token=optional_params.get("aws_web_identity_token"), + aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + ) + + # Prepare the request data + if isinstance(data, dict): + import json + request_data = json.dumps(data) + else: + request_data = data + + # Prepare headers + headers = {"Content-Type": "application/json"} + + # Create AWS request and sign it + sigv4 = SigV4Auth(credentials, service_name, aws_region_name) + request = AWSRequest( + method=method.upper(), url=endpoint_url, data=request_data, headers=headers + ) + sigv4.add_auth(request) + prepped = request.prepare() + + return dict(prepped.headers), request_data.encode('utf-8') if isinstance(request_data, str) else request_data + + def generate_unique_job_name(self, model: str, prefix: str = "litellm") -> str: + """ + Generate a unique job name for AWS services. + AWS services often have length limits, so this creates a concise name. + + Args: + model: Model name to include in the job name + prefix: Prefix for the job name + + Returns: + Unique job name (≤ 63 characters for Bedrock compatibility) + """ + import fastuuid as uuid + unique_id = str(uuid.uuid4())[:8] + # Format: {prefix}-batch-{model}-{uuid} + # Example: litellm-batch-claude-266c398e + job_name = f"{prefix}-batch-{unique_id}" + + return job_name + + def get_s3_bucket_and_key_from_config( + self, + litellm_params: dict, + optional_params: dict, + bucket_env_var: str = "AWS_S3_BUCKET_NAME", + key_prefix: str = "litellm" + ) -> tuple: + """ + Get S3 bucket and generate a unique key from configuration. + + Args: + litellm_params: LiteLLM parameters + optional_params: Optional parameters + bucket_env_var: Environment variable name for bucket + key_prefix: Prefix for the S3 key + + Returns: + Tuple of (bucket_name, object_key) + """ + import time + import uuid + + # Get bucket name + bucket_name = ( + litellm_params.get("s3_bucket_name") + or optional_params.get("s3_bucket_name") + or os.getenv(bucket_env_var) + ) + if not bucket_name: + raise ValueError(f"S3 bucket name is required. Set 's3_bucket_name' parameter or {bucket_env_var} env var") + + # Generate unique object key + timestamp = int(time.time()) + unique_id = str(uuid.uuid4())[:8] + object_key = f"{key_prefix}-{timestamp}-{unique_id}" + + return bucket_name, object_key + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers] + ) -> BaseLLMException: + """ + Get Bedrock-specific error class. + """ + return BedrockError( + status_code=status_code, + message=error_message, + headers=headers + ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py new file mode 100644 index 00000000000..83bbad7e1e8 --- /dev/null +++ b/litellm/llms/bedrock/files/transformation.py @@ -0,0 +1,607 @@ +import json +import os +import time +import uuid +from typing import Any, Dict, List, Optional, Tuple, Union + +from httpx import Headers, Response + +from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.transformation import ( + BaseFilesConfig, + LiteLLMLoggingObj, +) +from litellm.types.llms.openai import ( + AllMessageValues, + CreateFileRequest, + FileTypes, + OpenAICreateFileRequestOptionalParams, + OpenAIFileObject, + PathLike, +) +from litellm.types.utils import ExtractedFileData, LlmProviders + +from ..base_aws_llm import BaseAWSLLM +from ..common_utils import BedrockError + + +class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): + """ + Config for Bedrock Files - handles S3 uploads for Bedrock batch processing + """ + + def __init__(self): + self.jsonl_transformation = BedrockJsonlFilesTransformation() + super().__init__() + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.BEDROCK + + @property + def file_upload_http_method(self) -> str: + """ + Bedrock files are uploaded to S3, which requires PUT requests + """ + return "PUT" + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + # No additional headers needed for S3 uploads - AWS credentials handled by BaseAWSLLM + return headers + + + + def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: + """ + Helper to extract content from various OpenAI file types and return as string. + + Handles: + - Direct content (str, bytes, IO[bytes]) + - Tuple formats: (filename, content, [content_type], [headers]) + - PathLike objects + """ + content: Union[str, bytes] = b"" + # Extract file content from tuple if necessary + if isinstance(openai_file_content, tuple): + # Take the second element which is always the file content + file_content = openai_file_content[1] + else: + file_content = openai_file_content + + # Handle different file content types + if isinstance(file_content, str): + # String content can be used directly + content = file_content + elif isinstance(file_content, bytes): + # Bytes content can be decoded + content = file_content + elif isinstance(file_content, PathLike): # PathLike + with open(str(file_content), "rb") as f: + content = f.read() + elif hasattr(file_content, "read"): # IO[bytes] + # File-like objects need to be read + content = file_content.read() + + # Ensure content is string + if isinstance(content, bytes): + content = content.decode("utf-8") + + return content + + def _get_s3_object_name_from_batch_jsonl( + self, + openai_jsonl_content: List[Dict[str, Any]], + ) -> str: + """ + Gets a unique S3 object name for the Bedrock batch processing job + + named as: litellm-bedrock-files/{model}/{uuid} + """ + _model = openai_jsonl_content[0].get("body", {}).get("model", "") + # Remove bedrock/ prefix if present + if _model.startswith("bedrock/"): + _model = _model[8:] + object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl" + return object_name + + def get_object_name( + self, extracted_file_data: ExtractedFileData, purpose: str + ) -> str: + """ + Get the object name for the request + """ + extracted_file_data_content = extracted_file_data.get("content") + + if extracted_file_data_content is None: + raise ValueError("file content is required") + + if purpose == "batch": + ## 1. If jsonl, check if there's a model name + file_content = self._get_content_from_openai_file( + extracted_file_data_content + ) + + # Split into lines and parse each line as JSON + openai_jsonl_content = [ + json.loads(line) for line in file_content.splitlines() if line.strip() + ] + if len(openai_jsonl_content) > 0: + return self._get_s3_object_name_from_batch_jsonl(openai_jsonl_content) + + ## 2. If not jsonl, return the filename + filename = extracted_file_data.get("filename") + if filename: + return filename + ## 3. If no file name, return timestamp + return str(int(time.time())) + + def get_complete_file_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: Dict, + litellm_params: Dict, + data: CreateFileRequest, + ) -> str: + """ + Get the complete S3 URL for the file upload request + """ + bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") + if not bucket_name: + raise ValueError("S3 bucket_name is required. Set 's3_bucket_name' in litellm_params or AWS_S3_BUCKET_NAME env var") + + aws_region_name = self._get_aws_region_name(optional_params, model) + + file_data = data.get("file") + purpose = data.get("purpose") + if file_data is None: + raise ValueError("file is required") + if purpose is None: + raise ValueError("purpose is required") + extracted_file_data = extract_file_data(file_data) + object_name = self.get_object_name(extracted_file_data, purpose) + + # S3 endpoint URL format + s3_endpoint_url = optional_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" + + return f"{s3_endpoint_url}/{bucket_name}/{object_name}" + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAICreateFileRequestOptionalParams]: + return [] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + return optional_params + + def _get_bedrock_provider_from_model(self, model: str) -> Optional[str]: + """ + Extract provider from Bedrock model name + """ + if model.startswith("anthropic."): + return "anthropic" + elif model.startswith("cohere."): + return "cohere" + elif model.startswith("meta.") or model.startswith("llama"): + return "meta" + elif model.startswith("mistral."): + return "mistral" + elif model.startswith("ai21."): + return "ai21" + elif model.startswith("amazon."): + return "amazon" + else: + return None + + def _map_openai_to_bedrock_params( + self, + openai_request_body: Dict[str, Any], + provider: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Transform OpenAI request body to Bedrock-compatible modelInput parameters using existing transformation logic + """ + _model = openai_request_body.get("model", "") + messages = openai_request_body.get("messages", []) + + # Use existing Anthropic transformation logic for Anthropic models + if provider == "anthropic": + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeConfig, + ) + + anthropic_config = AmazonAnthropicClaudeConfig() + + # Extract optional params (everything except model and messages) + optional_params = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]} + + # Transform using existing Anthropic logic + bedrock_params = anthropic_config.transform_request( + model=_model, + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={} + ) + + return bedrock_params + else: + # For other providers, use basic mapping + bedrock_params = { + "messages": messages, + **{k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]} + } + return bedrock_params + + def _transform_openai_jsonl_content_to_bedrock_jsonl_content( + self, openai_jsonl_content: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """ + Transforms OpenAI JSONL content to Bedrock batch format + + Bedrock batch format: { "recordId": "alphanumeric string", "modelInput": {JSON body} } + Example: + { + "recordId": "CALL0000001", + "modelInput": { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 1024, + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "Hello"}] + } + ] + } + } + """ + + bedrock_jsonl_content = [] + for idx, _openai_jsonl_content in enumerate(openai_jsonl_content): + # Extract the request body from OpenAI format + openai_body = _openai_jsonl_content.get("body", {}) + model = openai_body.get("model", "") + + # Determine provider from model name + provider = self._get_bedrock_provider_from_model(model) + + # Transform to Bedrock modelInput format + model_input = self._map_openai_to_bedrock_params( + openai_request_body=openai_body, + provider=provider + ) + + # Create Bedrock batch record + record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") + bedrock_record = { + "recordId": record_id, + "modelInput": model_input + } + + bedrock_jsonl_content.append(bedrock_record) + return bedrock_jsonl_content + + def transform_create_file_request( + self, + model: str, + create_file_data: CreateFileRequest, + optional_params: dict, + litellm_params: dict, + ) -> Union[bytes, str, dict]: + """ + Transform file request and return a pre-signed request for S3. + This keeps the HTTP handler clean by doing all the signing here. + """ + file_data = create_file_data.get("file") + if file_data is None: + raise ValueError("file is required") + extracted_file_data = extract_file_data(file_data) + extracted_file_data_content = extracted_file_data.get("content") + + # Get and transform the file content + if ( + create_file_data.get("purpose") == "batch" + and extracted_file_data.get("content_type") == "application/jsonl" + and extracted_file_data_content is not None + ): + ## Transform JSONL content to Bedrock format + original_file_content = self._get_content_from_openai_file( + extracted_file_data_content + ) + openai_jsonl_content = [ + json.loads(line) for line in original_file_content.splitlines() if line.strip() + ] + bedrock_jsonl_content = ( + self._transform_openai_jsonl_content_to_bedrock_jsonl_content( + openai_jsonl_content + ) + ) + file_content = "\n".join(json.dumps(item) for item in bedrock_jsonl_content) + elif isinstance(extracted_file_data_content, bytes): + file_content = extracted_file_data_content.decode('utf-8') + elif isinstance(extracted_file_data_content, str): + file_content = extracted_file_data_content + else: + raise ValueError("Unsupported file content type") + + # Get the S3 URL for upload + api_base = self.get_complete_file_url( + api_base=None, + api_key=None, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + data=create_file_data, + ) + + # Sign the request and return a pre-signed request object + signed_headers, signed_body = self._sign_s3_request( + content=file_content, + api_base=api_base, + optional_params=optional_params, + ) + + # Return a dict that tells the HTTP handler exactly what to do + return { + "method": "PUT", + "url": api_base, + "headers": signed_headers, + "data": signed_body or file_content, + } + + def _sign_s3_request( + self, + content: str, + api_base: str, + optional_params: dict, + ) -> Tuple[dict, str]: + """ + Sign S3 PUT request using the same proven logic as S3Logger. + Reuses the exact pattern from litellm/integrations/s3_v2.py + """ + try: + import hashlib + + import requests + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + # Get AWS credentials using existing methods + aws_region_name = self._get_aws_region_name( + optional_params=optional_params, model="" + ) + credentials = self.get_credentials( + aws_access_key_id=optional_params.get("aws_access_key_id"), + aws_secret_access_key=optional_params.get("aws_secret_access_key"), + aws_session_token=optional_params.get("aws_session_token"), + aws_region_name=aws_region_name, + aws_session_name=optional_params.get("aws_session_name"), + aws_profile_name=optional_params.get("aws_profile_name"), + aws_role_name=optional_params.get("aws_role_name"), + aws_web_identity_token=optional_params.get("aws_web_identity_token"), + aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + ) + + # Calculate SHA256 hash of the content (REQUIRED for S3) + content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() + + # Prepare headers with required S3 headers (same as s3_v2.py) + request_headers = { + "Content-Type": "application/json", # JSONL files are JSON content + "x-amz-content-sha256": content_hash, # REQUIRED by S3 + "Content-Language": "en", + "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + } + + # Use requests.Request to prepare the request (same pattern as s3_v2.py) + req = requests.Request("PUT", api_base, data=content, headers=request_headers) + prepped = req.prepare() + + # Sign the request with S3 service + aws_request = AWSRequest( + method=prepped.method, + url=prepped.url, + data=prepped.body, + headers=prepped.headers, + ) + + # Get region name for non-LLM API calls (same as s3_v2.py) + signing_region = self.get_aws_region_name_for_non_llm_api_calls( + aws_region_name=aws_region_name + ) + + SigV4Auth(credentials, "s3", signing_region).add_auth(aws_request) + + # Return signed headers and body + signed_body = aws_request.body + if isinstance(signed_body, bytes): + signed_body = signed_body.decode('utf-8') + elif signed_body is None: + signed_body = content # Fallback to original content + + return dict(aws_request.headers), signed_body + + def transform_create_file_response( + self, + model: Optional[str], + raw_response: Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> OpenAIFileObject: + """ + Transform S3 File upload response into OpenAI-style FileObject + """ + # For S3 uploads, we typically get an ETag and other metadata + response_headers = raw_response.headers + + # Extract S3 object information from the response + # S3 PUT object returns ETag and other metadata in headers + content_length = response_headers.get("Content-Length", "0") + + # Extract bucket and key from the request URL or litellm_params + bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") + + # Generate file ID in S3 format + object_key = getattr(logging_obj, 'object_key', None) or f"file-{int(time.time())}" + file_id = f"s3://{bucket_name}/{object_key}" + + # Extract filename from object key + filename = object_key.split("/")[-1] if "/" in object_key else object_key + + return OpenAIFileObject( + purpose="batch", # Default purpose for Bedrock files + id=file_id, + filename=filename, + created_at=int(time.time()), # Current timestamp + status="uploaded", + bytes=int(content_length) if content_length.isdigit() else 0, + object="file", + ) + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[Dict, Headers] + ) -> BaseLLMException: + return BedrockError( + status_code=status_code, message=error_message, headers=headers + ) + + +class BedrockJsonlFilesTransformation: + """ + Transforms OpenAI /v1/files/* requests to Bedrock S3 file uploads for batch processing + """ + + def transform_openai_file_content_to_bedrock_file_content( + self, openai_file_content: Optional[FileTypes] = None + ) -> Tuple[str, str]: + """ + Transforms OpenAI FileContentRequest to Bedrock S3 file format + """ + + if openai_file_content is None: + raise ValueError("contents of file are None") + # Read the content of the file + file_content = self._get_content_from_openai_file(openai_file_content) + + # Split into lines and parse each line as JSON + openai_jsonl_content = [ + json.loads(line) for line in file_content.splitlines() if line.strip() + ] + bedrock_jsonl_content = ( + self._transform_openai_jsonl_content_to_bedrock_jsonl_content( + openai_jsonl_content + ) + ) + bedrock_jsonl_string = "\n".join( + json.dumps(item) for item in bedrock_jsonl_content + ) + object_name = self._get_s3_object_name( + openai_jsonl_content=openai_jsonl_content + ) + return bedrock_jsonl_string, object_name + + def _transform_openai_jsonl_content_to_bedrock_jsonl_content( + self, openai_jsonl_content: List[Dict[str, Any]] + ): + """ + Delegate to the main BedrockFilesConfig transformation method + """ + config = BedrockFilesConfig() + return config._transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content) + + def _get_s3_object_name( + self, + openai_jsonl_content: List[Dict[str, Any]], + ) -> str: + """ + Gets a unique S3 object name for the Bedrock batch processing job + + named as: litellm-bedrock-files-{model}-{uuid} + """ + _model = openai_jsonl_content[0].get("body", {}).get("model", "") + # Remove bedrock/ prefix if present + if _model.startswith("bedrock/"): + _model = _model[8:] + object_name = f"litellm-bedrock-files-{_model}-{uuid.uuid4()}.jsonl" + return object_name + + + + def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str: + """ + Helper to extract content from various OpenAI file types and return as string. + + Handles: + - Direct content (str, bytes, IO[bytes]) + - Tuple formats: (filename, content, [content_type], [headers]) + - PathLike objects + """ + content: Union[str, bytes] = b"" + # Extract file content from tuple if necessary + if isinstance(openai_file_content, tuple): + # Take the second element which is always the file content + file_content = openai_file_content[1] + else: + file_content = openai_file_content + + # Handle different file content types + if isinstance(file_content, str): + # String content can be used directly + content = file_content + elif isinstance(file_content, bytes): + # Bytes content can be decoded + content = file_content + elif isinstance(file_content, PathLike): # PathLike + with open(str(file_content), "rb") as f: + content = f.read() + elif hasattr(file_content, "read"): # IO[bytes] + # File-like objects need to be read + content = file_content.read() + + # Ensure content is string + if isinstance(content, bytes): + content = content.decode("utf-8") + + return content + + def transform_s3_bucket_response_to_openai_file_object( + self, create_file_data: CreateFileRequest, s3_upload_response: Dict[str, Any] + ) -> OpenAIFileObject: + """ + Transforms S3 Bucket upload file response to OpenAI FileObject + """ + # S3 response typically contains ETag, key, etc. + object_key = s3_upload_response.get("Key", "") + bucket_name = s3_upload_response.get("Bucket", "") + + # Extract filename from object key + filename = object_key.split("/")[-1] if "/" in object_key else object_key + + return OpenAIFileObject( + purpose=create_file_data.get("purpose", "batch"), + id=f"s3://{bucket_name}/{object_key}", + filename=filename, + created_at=int(time.time()), # Current timestamp + status="uploaded", + bytes=s3_upload_response.get("ContentLength", 0), + object="file", + ) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index d7221ff4b7a..5791bfb8013 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -41,9 +41,15 @@ class BedrockPassthroughConfig( model_id=None, ) - api_base = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" + aws_bedrock_runtime_endpoint = optional_params.get("aws_bedrock_runtime_endpoint") + endpoint_url, _ = self.get_runtime_endpoint( + api_base=api_base, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, + aws_region_name=aws_region_name, + endpoint_type="runtime", + ) - return self.format_url(endpoint, api_base, request_query_params or {}), api_base + return self.format_url(endpoint, endpoint_url, request_query_params or {}), endpoint_url def sign_request( self, diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 4d8781fff2a..36b543086f5 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -212,6 +212,7 @@ class AsyncHTTPHandler: verify=ssl_config, cert=cert, headers=headers, + follow_redirects=True, ) async def close(self): @@ -687,6 +688,7 @@ class HTTPHandler: verify=ssl_config, cert=cert, headers=headers, + follow_redirects=True, ) else: self.client = client diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2faea53901c..13133a56aad 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -28,6 +28,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, ) from litellm.llms.base_llm.base_model_iterator import MockResponseIterator +from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.files.transformation import BaseFilesConfig @@ -58,6 +59,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) from litellm.types.llms.openai import ( + CreateBatchRequest, CreateFileRequest, OpenAIFileObject, ResponseInputParam, @@ -66,7 +68,12 @@ from litellm.types.llms.openai import ( from litellm.types.rerank import OptionalRerankParams, RerankResponse from litellm.types.responses.main import DeleteResponseResult from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import EmbeddingResponse, FileTypes, TranscriptionResponse +from litellm.types.utils import ( + EmbeddingResponse, + FileTypes, + LiteLLMBatch, + TranscriptionResponse, +) from litellm.types.vector_stores import ( VectorStoreCreateOptionalRequestParams, VectorStoreCreateResponse, @@ -2212,15 +2219,38 @@ class BaseLLMHTTPHandler: else: sync_httpx_client = client - if isinstance(transformed_request, str) or isinstance( - transformed_request, bytes - ): - upload_response = sync_httpx_client.post( - url=api_base, - headers=headers, - data=transformed_request, + if isinstance(transformed_request, dict) and "method" in transformed_request: + # Handle pre-signed requests (e.g., from Bedrock S3 uploads) + upload_response = getattr(sync_httpx_client, transformed_request["method"].lower())( + url=transformed_request["url"], + headers=transformed_request["headers"], + data=transformed_request["data"], timeout=timeout, ) + elif isinstance(transformed_request, str) or isinstance( + transformed_request, bytes + ): + # Handle traditional file uploads + # Ensure transformed_request is a string for httpx compatibility + if isinstance(transformed_request, bytes): + transformed_request = transformed_request.decode('utf-8') + + # Use the HTTP method specified by the provider config + http_method = provider_config.file_upload_http_method.upper() + if http_method == "PUT": + upload_response = sync_httpx_client.put( + url=api_base, + headers=headers, + data=transformed_request, + timeout=timeout, + ) + else: # Default to POST + upload_response = sync_httpx_client.post( + url=api_base, + headers=headers, + data=transformed_request, + timeout=timeout, + ) else: try: # Step 1: Initial request to get upload URL @@ -2280,16 +2310,52 @@ class BaseLLMHTTPHandler: ) else: async_httpx_client = client + + ######################################################### + # Debug Logging + ######################################################### + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": transformed_request, + "api_base": api_base, + "headers": headers, + }, + ) - if isinstance(transformed_request, str) or isinstance( - transformed_request, bytes - ): - upload_response = await async_httpx_client.post( - url=api_base, - headers=headers, - data=transformed_request, + if isinstance(transformed_request, dict) and "method" in transformed_request: + # Handle pre-signed requests (e.g., from Bedrock S3 uploads) + upload_response = await getattr(async_httpx_client, transformed_request["method"].lower())( + url=transformed_request["url"], + headers=transformed_request["headers"], + data=transformed_request["data"], timeout=timeout, ) + elif isinstance(transformed_request, str) or isinstance( + transformed_request, bytes + ): + # Handle traditional file uploads + # Ensure transformed_request is a string for httpx compatibility + if isinstance(transformed_request, bytes): + transformed_request = transformed_request.decode('utf-8') + + # Use the HTTP method specified by the provider config + http_method = provider_config.file_upload_http_method.upper() + if http_method == "PUT": + upload_response = await async_httpx_client.put( + url=api_base, + headers=headers, + data=transformed_request, + timeout=timeout, + ) + else: # Default to POST + upload_response = await async_httpx_client.post( + url=api_base, + headers=headers, + data=transformed_request, + timeout=timeout, + ) else: try: # Step 1: Initial request to get upload URL @@ -2330,6 +2396,188 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) + def create_batch( + self, + create_batch_data: "CreateBatchRequest", + litellm_params: dict, + provider_config: "BaseBatchesConfig", + headers: dict, + api_base: Optional[str], + api_key: Optional[str], + logging_obj: "LiteLLMLoggingObj", + _is_async: bool = False, + client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + ) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]: + """ + Creates a batch using provider-specific batch creation process + """ + # get config from model, custom llm provider + headers = provider_config.validate_environment( + api_key=api_key, + headers=headers, + model="", + messages=[], + optional_params={}, + litellm_params=litellm_params, + ) + + api_base = provider_config.get_complete_batch_url( + api_base=api_base, + api_key=api_key, + model="", + optional_params={}, + litellm_params=litellm_params, + data=create_batch_data, + ) + if api_base is None: + raise ValueError("api_base is required for create_batch") + + # Get the transformed request data + transformed_request = provider_config.transform_create_batch_request( + model="", + create_batch_data=create_batch_data, + litellm_params=litellm_params, + optional_params={}, + ) + + if _is_async: + return self.async_create_batch( + transformed_request=transformed_request, + litellm_params=litellm_params, + provider_config=provider_config, + headers=headers, + api_base=api_base, + logging_obj=logging_obj, + client=client, + timeout=timeout, + create_batch_data=create_batch_data, + ) + + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client() + else: + sync_httpx_client = client + + try: + if isinstance(transformed_request, dict) and "method" in transformed_request: + # Handle pre-signed requests (e.g., from Bedrock with AWS auth) + batch_response = getattr(sync_httpx_client, transformed_request["method"].lower())( + url=transformed_request["url"], + headers=transformed_request["headers"], + data=transformed_request["data"], + timeout=timeout, + ) + elif isinstance(transformed_request, dict): + # For other providers that use JSON requests + batch_response = sync_httpx_client.post( + url=api_base, + headers={**headers, "Content-Type": "application/json"}, + json=transformed_request, + timeout=timeout, + ) + else: + # Handle other request types if needed + batch_response = sync_httpx_client.post( + url=api_base, + headers=headers, + data=transformed_request, + timeout=timeout, + ) + except Exception as e: + verbose_logger.exception(f"Error creating batch: {e}") + raise self._handle_error( + e=e, + provider_config=provider_config, + ) + + # Store original request for response transformation + litellm_params_with_request = {**litellm_params, "original_batch_request": create_batch_data} + + return provider_config.transform_create_batch_response( + model=None, + raw_response=batch_response, + logging_obj=logging_obj, + litellm_params=litellm_params_with_request, + ) + + async def async_create_batch( + self, + transformed_request: Union[bytes, str, dict], + litellm_params: dict, + provider_config: "BaseBatchesConfig", + headers: dict, + api_base: str, + logging_obj: "LiteLLMLoggingObj", + client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + create_batch_data: Optional["CreateBatchRequest"] = None, + ): + """ + Async version of create_batch + """ + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=provider_config.custom_llm_provider + ) + else: + async_httpx_client = client + + ######################################################### + # Debug Logging + ######################################################### + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": transformed_request, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + if isinstance(transformed_request, dict) and "method" in transformed_request: + # Handle pre-signed requests (e.g., from Bedrock with AWS auth) + batch_response = await getattr(async_httpx_client, transformed_request["method"].lower())( + url=transformed_request["url"], + headers=transformed_request["headers"], + data=transformed_request["data"], + timeout=timeout, + ) + elif isinstance(transformed_request, dict): + # For other providers that use JSON requests + batch_response = await async_httpx_client.post( + url=api_base, + headers={**headers, "Content-Type": "application/json"}, + json=transformed_request, + timeout=timeout, + ) + else: + # Handle other request types if needed + batch_response = await async_httpx_client.post( + url=api_base, + headers=headers, + data=transformed_request, + timeout=timeout, + ) + except Exception as e: + verbose_logger.exception(f"Error creating batch: {e}") + raise self._handle_error( + e=e, + provider_config=provider_config, + ) + + # Store original request for response transformation (for async version) + litellm_params_with_request = {**litellm_params, "original_batch_request": create_batch_data or {}} + + return provider_config.transform_create_batch_response( + model=None, + raw_response=batch_response, + logging_obj=logging_obj, + litellm_params=litellm_params_with_request, + ) + def list_files(self): """ Lists all files @@ -2381,6 +2629,7 @@ class BaseLLMHTTPHandler: BaseVectorStoreConfig, BaseGoogleGenAIGenerateContentConfig, BaseAnthropicMessagesConfig, + BaseBatchesConfig, "BasePassthroughConfig", ], ): diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 908419f7193..d3df5bbf361 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -26,7 +26,6 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( - handle_messages_with_content_list_to_str_conversion, strip_name_from_messages, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -301,7 +300,6 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: """ Databricks does not support: - - content in list format. - 'name' in user message. """ new_messages = [] @@ -311,7 +309,6 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): else: _message = message new_messages.append(_message) - new_messages = handle_messages_with_content_list_to_str_conversion(new_messages) new_messages = strip_name_from_messages(new_messages) if is_async: @@ -379,6 +376,25 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): thinking_blocks.append(thinking_block) return reasoning_content, thinking_blocks + @staticmethod + def extract_citations( + content: Optional[AllDatabricksContentValues], + ) -> Optional[List[Any]]: + if content is None: + return None + citations = [] + if isinstance(content, list): + for item in content: + text = item.get("text", None) + if citations_item := item.get("citations"): + citations.append( + [ + {**citation, "supported_text": text} + for citation in citations_item + ] + ) + return citations or None + def _transform_dbrx_choices( self, choices: List[DatabricksChoice], json_mode: Optional[bool] = None ) -> List[Choices]: @@ -427,12 +443,19 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): choice["message"].get("content") ) + citations = DatabricksConfig.extract_citations( + choice["message"].get("content") + ) + translated_message = Message( role="assistant", content=content_str, reasoning_content=reasoning_content, thinking_blocks=thinking_blocks, tool_calls=choice["message"].get("tool_calls"), + provider_specific_fields={"citations": citations} + if citations is not None + else None, ) if finish_reason is None: @@ -561,6 +584,17 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator): for _tc in tool_calls: if _tc.get("function", {}).get("arguments") == "{}": _tc["function"]["arguments"] = "" # avoid invalid json + if isinstance(choice["delta"]["content"], list) and ( + content := choice["delta"]["content"] + ): + if citations := content[0].get("citations"): + # TODO: Databricks delta does not include supported text or chunk type. + # Add either here once Databricks supports it to enable citation linkage. + choice["delta"].setdefault("provider_specific_fields", {})[ + "citation" + ] = citations[ + 0 + ] # Databricks Content item always has citation as a list of list # extract the content str content_str = DatabricksConfig.extract_content_str( choice["delta"].get("content") diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index 86fa323f9e3..165301efb5c 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -6,6 +6,8 @@ from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast, import httpx from pydantic import BaseModel +import litellm +from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -55,6 +57,10 @@ class GroqChatConfig(OpenAILikeChatConfig): if key != "self" and value is not None: setattr(self.__class__, key, value) + @property + def custom_llm_provider(self) -> Optional[str]: + return "groq" + @classmethod def get_config(cls): return super().get_config() @@ -65,6 +71,15 @@ class GroqChatConfig(OpenAILikeChatConfig): base_params.remove("max_retries") except ValueError: pass + + try: + if litellm.supports_reasoning( + model=model, custom_llm_provider=self.custom_llm_provider + ): + base_params.append("reasoning_effort") + except Exception as e: + verbose_logger.debug(f"Error checking if model supports reasoning: {e}") + return base_params @overload diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 915d2029afe..3be373ca5e5 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -772,7 +772,14 @@ def adapt_messages_to_generic_oci_standard( tool_calls = message.get("tool_calls") tool_call_id = message.get("tool_call_id") - if role in ["system", "user", "assistant"] and content is not None: + if role == "assistant" and tool_calls is not None: + if not isinstance(tool_calls, list): + raise Exception("Prop `tool_calls` must be a list of tool calls") + new_messages.append( + adapt_messages_to_generic_oci_standard_tool_call(role, tool_calls) + ) + + elif role in ["system", "user", "assistant"] and content is not None: if not isinstance(content, (str, list)): raise Exception( "Prop `content` must be a string or a list of content items" @@ -781,13 +788,6 @@ def adapt_messages_to_generic_oci_standard( adapt_messages_to_generic_oci_standard_content_message(role, content) ) - elif role == "assistant" and tool_calls is not None: - if not isinstance(tool_calls, list): - raise Exception("Prop `tool_calls` must be a list of tool calls") - new_messages.append( - adapt_messages_to_generic_oci_standard_tool_call(role, tool_calls) - ) - elif role == "tool": if not isinstance(tool_call_id, str): raise Exception("Prop `tool_call_id` is required and must be a string") diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index d4ce4052a7e..c70fb97af74 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -137,6 +137,7 @@ class OllamaChatConfig(BaseConfig): "tool_choice", "functions", "response_format", + "reasoning_effort", ] def map_openai_params( @@ -175,6 +176,8 @@ class OllamaChatConfig(BaseConfig): if value.get("json_schema") and value["json_schema"].get("schema"): optional_params["format"] = value["json_schema"]["schema"] ### FUNCTION CALLING LOGIC ### + if param == "reasoning_effort" and value is not None: + optional_params["think"] = True if param == "tools": ## CHECK IF MODEL SUPPORTS TOOL CALLING ## try: @@ -212,9 +215,9 @@ class OllamaChatConfig(BaseConfig): litellm.add_function_to_prompt = ( True # so that main.py adds the function call to the prompt ) - optional_params[ - "functions_unsupported_model" - ] = non_default_params.get("functions") + optional_params["functions_unsupported_model"] = ( + non_default_params.get("functions") + ) non_default_params.pop("tool_choice", None) # causes ollama requests to hang non_default_params.pop("functions", None) # causes ollama requests to hang return optional_params @@ -346,11 +349,31 @@ class OllamaChatConfig(BaseConfig): ## RESPONSE OBJECT model_response.choices[0].finish_reason = "stop" + response_json_message = response_json.get("message") + if response_json_message is not None: + if "thinking" in response_json_message: + # remap 'thinking' to 'reasoning_content' + response_json_message["reasoning_content"] = response_json_message[ + "thinking" + ] + del response_json_message["thinking"] + elif response_json_message.get("content") is not None: + # parse reasoning content from content + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + _parse_content_for_reasoning, + ) + + reasoning_content, content = _parse_content_for_reasoning( + response_json_message["content"] + ) + response_json_message["reasoning_content"] = reasoning_content + response_json_message["content"] = content + if ( request_data.get("format", "") == "json" and litellm_params.get("function_name") is not None ): - function_call = json.loads(response_json["message"]["content"]) + function_call = json.loads(response_json_message["content"]) message = litellm.Message( content=None, tool_calls=[ @@ -367,11 +390,13 @@ class OllamaChatConfig(BaseConfig): "type": "function", } ], + reasoning_content=response_json_message.get("reasoning_content"), ) model_response.choices[0].message = message # type: ignore model_response.choices[0].finish_reason = "tool_calls" else: - _message = litellm.Message(**response_json["message"]) + + _message = litellm.Message(**response_json_message) model_response.choices[0].message = _message # type: ignore model_response.created = int(time.time()) model_response.model = "ollama_chat/" + model @@ -412,6 +437,9 @@ class OllamaChatConfig(BaseConfig): class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): + started_reasoning_content: bool = False + finished_reasoning_content: bool = False + def _is_function_call_complete(self, function_args: Union[str, dict]) -> bool: if isinstance(function_args, dict): return True @@ -465,8 +493,38 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): if is_function_call_complete: tool_call["id"] = str(uuid.uuid4()) + # PROCESS REASONING CONTENT + reasoning_content: Optional[str] = None + content: Optional[str] = None + if chunk["message"].get("thinking") is not None: + if self.started_reasoning_content is False: + reasoning_content = chunk["message"].get("thinking") + self.started_reasoning_content = True + elif self.finished_reasoning_content is False: + reasoning_content = chunk["message"].get("thinking") + self.finished_reasoning_content = True + elif chunk["message"].get("content") is not None: + message_content = chunk["message"].get("content") + if "" in message_content: + message_content = message_content.replace("", "") + + self.started_reasoning_content = True + + if "" in message_content and self.started_reasoning_content: + message_content = message_content.replace("", "") + self.finished_reasoning_content = True + + if ( + self.started_reasoning_content + and not self.finished_reasoning_content + ): + reasoning_content = message_content + else: + content = message_content + delta = Delta( - content=chunk["message"].get("content", ""), + content=content, + reasoning_content=reasoning_content, tool_calls=tool_calls, ) diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 4f7be507cc2..4a491c88963 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -19,13 +19,13 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMExcepti from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock from litellm.types.utils import ( + Delta, GenericStreamingChunk, ModelInfoBase, ModelResponse, ModelResponseStream, ProviderField, StreamingChoices, - Delta, ) from ..common_utils import OllamaError, _convert_image @@ -92,9 +92,9 @@ class OllamaConfig(BaseConfig): repeat_penalty: Optional[float] = None temperature: Optional[float] = None seed: Optional[int] = None - stop: Optional[ - list - ] = None # stop is a list based on this - https://github.com/ollama/ollama/pull/442 + stop: Optional[list] = ( + None # stop is a list based on this - https://github.com/ollama/ollama/pull/442 + ) tfs_z: Optional[float] = None num_predict: Optional[int] = None top_k: Optional[int] = None @@ -154,6 +154,7 @@ class OllamaConfig(BaseConfig): "stop", "response_format", "max_completion_tokens", + "reasoning_effort", ] def map_openai_params( @@ -166,19 +167,21 @@ class OllamaConfig(BaseConfig): for param, value in non_default_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["num_predict"] = value - if param == "stream": + elif param == "stream": optional_params["stream"] = value - if param == "temperature": + elif param == "temperature": optional_params["temperature"] = value - if param == "seed": + elif param == "seed": optional_params["seed"] = value - if param == "top_p": + elif param == "top_p": optional_params["top_p"] = value - if param == "frequency_penalty": + elif param == "frequency_penalty": optional_params["frequency_penalty"] = value - if param == "stop": + elif param == "stop": optional_params["stop"] = value - if param == "response_format" and isinstance(value, dict): + elif param == "reasoning_effort" and value is not None: + optional_params["think"] = True + elif param == "response_format" and isinstance(value, dict): if value["type"] == "json_object": optional_params["format"] = "json" elif value["type"] == "json_schema": @@ -258,12 +261,17 @@ class OllamaConfig(BaseConfig): api_key: Optional[str] = None, json_mode: Optional[bool] = None, ) -> ModelResponse: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + _parse_content_for_reasoning, + ) + response_json = raw_response.json() ## RESPONSE OBJECT model_response.choices[0].finish_reason = "stop" if request_data.get("format", "") == "json": # Check if response field exists and is not empty before parsing JSON response_text = response_json.get("response", "") + if not response_text or not response_text.strip(): # Handle empty response gracefully - set empty content message = litellm.Message(content="") @@ -288,7 +296,9 @@ class OllamaConfig(BaseConfig): "id": f"call_{str(uuid.uuid4())}", "function": { "name": function_call["name"], - "arguments": json.dumps(function_call["arguments"]), + "arguments": json.dumps( + function_call["arguments"] + ), }, "type": "function", } @@ -305,11 +315,28 @@ class OllamaConfig(BaseConfig): model_response.choices[0].finish_reason = "stop" except json.JSONDecodeError: # If JSON parsing fails, treat as regular text response - message = litellm.Message(content=response_text) + ## output parse reasoning content from response_text + reasoning_content: Optional[str] = None + content: Optional[str] = None + if response_text is not None: + reasoning_content, content = _parse_content_for_reasoning( + response_text + ) + message = litellm.Message( + content=content, reasoning_content=reasoning_content + ) model_response.choices[0].message = message # type: ignore model_response.choices[0].finish_reason = "stop" else: - model_response.choices[0].message.content = response_json["response"] # type: ignore + response_text = response_json.get("response", "") + content = None + reasoning_content = None + if response_text is not None and isinstance(response_text, str): + reasoning_content, content = _parse_content_for_reasoning(response_text) + else: + content = response_text # type: ignore + model_response.choices[0].message.content = content # type: ignore + model_response.choices[0].message.reasoning_content = reasoning_content # type: ignore model_response.created = int(time.time()) model_response.model = "ollama/" + model _prompt = request_data.get("prompt", "") @@ -434,12 +461,21 @@ class OllamaConfig(BaseConfig): class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): + def __init__( + self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False + ): + super().__init__(streaming_response, sync_stream, json_mode) + self.started_reasoning_content: bool = False + self.finished_reasoning_content: bool = False + def _handle_string_chunk( self, str_line: str ) -> Union[GenericStreamingChunk, ModelResponseStream]: return self.chunk_parser(json.loads(str_line)) - def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]: + def chunk_parser( + self, chunk: dict + ) -> Union[GenericStreamingChunk, ModelResponseStream]: try: if "error" in chunk: raise Exception(f"Ollama Error - {chunk}") @@ -469,12 +505,42 @@ class OllamaTextCompletionResponseIterator(BaseModelResponseIterator): ) elif chunk["response"]: text = chunk["response"] - return GenericStreamingChunk( - text=text, - is_finished=is_finished, - finish_reason="stop", + reasoning_content: Optional[str] = None + content: Optional[str] = None + if text is not None: + if "" in text: + text = text.replace("", "") + self.started_reasoning_content = True + elif "" in text: + text = text.replace("", "") + self.finished_reasoning_content = True + + if ( + self.started_reasoning_content + and not self.finished_reasoning_content + ): + reasoning_content = text + else: + content = text + + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + reasoning_content=reasoning_content, content=content + ), + ) + ], + finish_reason=finish_reason, usage=None, ) + # return GenericStreamingChunk( + # text=text, + # is_finished=is_finished, + # finish_reason="stop", + # usage=None, + # ) elif "thinking" in chunk and not chunk["response"]: # Return reasoning content as ModelResponseStream so UIs can render it thinking_content = chunk.get("thinking") or "" diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 9a8bb74d447..3902304a3b4 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -28,7 +28,18 @@ class OpenAIGPT5Config(OpenAIGPTConfig): base_gpt_series_params.extend(gpt_5_only_params) if not supports_tool_choice(model=model): base_gpt_series_params.remove("tool_choice") - return base_gpt_series_params + + non_supported_params = [ + "logprobs", + "top_p", + "presence_penalty", + "frequency_penalty", + "top_logprobs", + ] + + return [ + param for param in base_gpt_series_params if param not in non_supported_params + ] def map_openai_params( self, diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index be0ca3a7086..204916e3a48 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -158,6 +158,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "parallel_tool_calls", "audio", "web_search_options", + "safety_identifier", ] # works across all models model_specific_params = [] diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 150cffba21c..1cee13784e7 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -16,7 +16,6 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): ) -> List[OpenAIImageGenerationOptionalParams]: return [ "background", - "input_fidelity", "moderation", "n", "output_compression", diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index e70cadddaf7..392d47f9822 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,6 +1,15 @@ -from typing import TYPE_CHECKING, Any, Dict, Optional, Union, cast, get_type_hints +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Optional, + Union, + cast, + get_type_hints, +) import httpx +from openai.types.responses import ResponseReasoningItem from pydantic import BaseModel import litellm @@ -92,12 +101,67 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): # if it's pydantic, convert to dict if isinstance(item, BaseModel): validated_input.append(item.model_dump(exclude_none=True)) + elif isinstance(item, dict): + # Handle reasoning items specifically to filter out status=None + verbose_logger.debug(f"Handling reasoning item: {item}") + if item.get("type") == "reasoning": + # Type assertion since we know it's a dict at this point + dict_item = cast(Dict[str, Any], item) + filtered_item = self._handle_reasoning_item(dict_item) + else: + # For other dict items, just pass through + filtered_item = cast(Dict[str, Any], item) + validated_input.append(filtered_item) else: validated_input.append(item) - return validated_input + return validated_input # type: ignore # Input is expected to be either str or List, no single BaseModel expected return input + def _handle_reasoning_item(self, item: Dict[str, Any]) -> Dict[str, Any]: + """ + Handle reasoning items specifically to filter out status=None using OpenAI's model. + Issue: https://github.com/BerriAI/litellm/issues/13484 + OpenAI API does not accept ReasoningItem(status=None), so we need to: + 1. Check if the item is a reasoning type + 2. Create a ResponseReasoningItem object with the item data + 3. Convert it back to dict with exclude_none=True to filter None values + """ + verbose_logger.debug(f"Handling reasoning item: {item}") + if item.get("type") == "reasoning": + try: + # Ensure required fields are present for ResponseReasoningItem + item_data = dict(item) + if "id" not in item_data: + item_data["id"] = f"reasoning_{hash(str(item_data))}" + if "summary" not in item_data: + item_data["summary"] = ( + item_data.get("reasoning_content", "")[:100] + "..." + if len(item_data.get("reasoning_content", "")) > 100 + else item_data.get("reasoning_content", "") + ) + + # Create ResponseReasoningItem object from the item data + reasoning_item = ResponseReasoningItem(**item_data) + + # Convert back to dict with exclude_none=True to exclude None fields + dict_reasoning_item = reasoning_item.model_dump(exclude_none=True) + + return dict_reasoning_item + except Exception as e: + verbose_logger.debug( + f"Failed to create ResponseReasoningItem, falling back to manual filtering: {e}" + ) + # Fallback: manually filter out known None fields + filtered_item = { + k: v + for k, v in item.items() + if v is not None + or k not in {"status", "content", "encrypted_content"} + } + return filtered_item + return item + def transform_response_api_response( self, model: str, diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 8ab212e2558..267ca61ef5d 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -105,6 +105,64 @@ def _process_gemini_image(image_url: str, format: Optional[str] = None) -> PartT raise e +def _snake_to_camel(snake_str: str) -> str: + """Convert snake_case to camelCase""" + components = snake_str.split("_") + return components[0] + "".join(x.capitalize() for x in components[1:]) + + +def _camel_to_snake(camel_str: str) -> str: + """Convert camelCase to snake_case""" + import re + + return re.sub(r"(? Optional[str]: + """ + Get the equivalent key from available keys, checking both camelCase and snake_case variants + """ + if key in available_keys: + return key + + # Try camelCase version + camel_key = _snake_to_camel(key) + if camel_key in available_keys: + return camel_key + + # Try snake_case version + snake_key = _camel_to_snake(key) + if snake_key in available_keys: + return snake_key + + return None + + +def check_if_part_exists_in_parts( + parts: List[PartType], part: PartType, excluded_keys: List[str] = [] +) -> bool: + """ + Check if a part exists in a list of parts + Handles both camelCase and snake_case key variations (e.g., function_call vs functionCall) + """ + keys_to_compare = set(part.keys()) - set(excluded_keys) + for p in parts: + p_keys = set(p.keys()) + # Check if all keys in part have equivalent values in p + match_found = True + for key in keys_to_compare: + equivalent_key = _get_equivalent_key(key, p_keys) + if equivalent_key is None or p.get(equivalent_key, None) != part.get( + key, None + ): + match_found = False + break + + if match_found: + return True + return False + + def _gemini_convert_messages_with_history( # noqa: PLR0915 messages: List[AllMessageValues], ) -> List[ContentType]: @@ -236,10 +294,33 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 assistant_msg = ChatCompletionAssistantMessage(**msg_dict) # type: ignore _message_content = assistant_msg.get("content", None) reasoning_content = assistant_msg.get("reasoning_content", None) + thinking_blocks = assistant_msg.get("thinking_blocks") if reasoning_content is not None: assistant_content.append( PartType(thought=True, text=reasoning_content) ) + if thinking_blocks is not None: + for block in thinking_blocks: + block_thinking_str = block.get("thinking") + block_signature = block.get("signature") + if ( + block_thinking_str is not None + and block_signature is not None + ): + try: + assistant_content.append( + PartType( + thoughtSignature=block_signature, + **json.loads(block_thinking_str), + ) + ) + except Exception: + assistant_content.append( + PartType( + thoughtSignature=block_signature, + text=block_thinking_str, + ) + ) if _message_content is not None and isinstance(_message_content, list): _parts = [] for element in _message_content: @@ -262,9 +343,17 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 assistant_msg.get("tool_calls", []) is not None or assistant_msg.get("function_call") is not None ): # support assistant tool invoke conversion - assistant_content.extend( - convert_to_gemini_tool_call_invoke(assistant_msg) + gemini_tool_call_parts = convert_to_gemini_tool_call_invoke( + assistant_msg ) + ## check if gemini_tool_call already exists in assistant_content + for gemini_tool_call_part in gemini_tool_call_parts: + if not check_if_part_exists_in_parts( + assistant_content, + gemini_tool_call_part, + excluded_keys=["thoughtSignature"], + ): + assistant_content.append(gemini_tool_call_part) last_message_with_tool_calls = assistant_msg msg_i += 1 @@ -476,6 +565,7 @@ async def async_transform_request_body( optional_params=optional_params, ) + def _default_user_message_when_system_message_passed() -> ChatCompletionUserMessage: """ Returns a default user message when a "system" message is passed in gemini fails. @@ -484,6 +574,7 @@ def _default_user_message_when_system_message_passed() -> ChatCompletionUserMess """ return ChatCompletionUserMessage(content=".", role="user") + def _transform_system_message( supports_system_message: bool, messages: List[AllMessageValues] ) -> Tuple[Optional[SystemInstructions], List[AllMessageValues]]: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a376e0dd06e..9376b28cbec 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -30,6 +30,10 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET, + DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH, + DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO, + DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -43,6 +47,7 @@ from litellm.types.llms.gemini import BidiGenerateContentServerMessage from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionResponseMessage, + ChatCompletionThinkingBlock, ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ChatCompletionToolParamFunctionChunk, @@ -422,8 +427,25 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _map_reasoning_effort_to_thinking_budget( reasoning_effort: str, + model: Optional[str] = None, ) -> GeminiThinkingConfig: - if reasoning_effort == "low": + if reasoning_effort == "minimal": + # Use model-specific minimum thinking budget or fallback + # Check for exact matches first, then partial matches + if model and "gemini-2.5-flash-lite" in model.lower(): + budget = DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE + elif model and "gemini-2.5-pro" in model.lower(): + budget = DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO + elif model and "gemini-2.5-flash" in model.lower(): + budget = DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH + else: + budget = DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET + + return { + "thinkingBudget": budget, + "includeThoughts": True, + } + elif reasoning_effort == "low": return { "thinkingBudget": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, "includeThoughts": True, @@ -600,7 +622,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): optional_params["seed"] = value elif param == "reasoning_effort" and isinstance(value, str): optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(value) + VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( + value, model + ) ) elif param == "thinking": optional_params["thinkingConfig"] = ( @@ -794,6 +818,24 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return content_str, reasoning_content_str + def _extract_thinking_blocks_from_parts( + self, parts: List[HttpxPartType] + ) -> List[ChatCompletionThinkingBlock]: + """Extract thinking blocks from parts if present""" + thinking_blocks: List[ChatCompletionThinkingBlock] = [] + for part in parts: + if "thoughtSignature" in part: + part_copy = part.copy() + part_copy.pop("thoughtSignature") + thinking_blocks.append( + ChatCompletionThinkingBlock( + type="thinking", + thinking=json.dumps(part_copy), + signature=part["thoughtSignature"], + ) + ) + return thinking_blocks + def _extract_image_response_from_parts( self, parts: List[HttpxPartType] ) -> Optional[List[ImageURLListItem]]: @@ -1237,6 +1279,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tools: Optional[List[ChatCompletionToolCallChunk]] = [] functions: Optional[ChatCompletionToolCallFunctionChunk] = None cumulative_tool_call_index: int = 0 + thinking_blocks: Optional[List[ChatCompletionThinkingBlock]] = None for idx, candidate in enumerate(_candidates): if "content" not in candidate: @@ -1274,6 +1317,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) ) + thinking_blocks = ( + VertexGeminiConfig()._extract_thinking_blocks_from_parts( + parts=candidate["content"]["parts"] + ) + ) + if audio_response is not None: cast(Dict[str, Any], chat_completion_message)[ "audio" @@ -1310,6 +1359,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if functions is not None: chat_completion_message["function_call"] = functions + if thinking_blocks is not None: + chat_completion_message["thinking_blocks"] = thinking_blocks # type: ignore + if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( chat_completion_message=chat_completion_message, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py new file mode 100644 index 00000000000..86e36e802ed --- /dev/null +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/transformation.py @@ -0,0 +1,27 @@ +import litellm +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + + +class VertexAIGPTOSSTransformation(OpenAIGPTConfig): + """ + Transformation for GPT-OSS model from VertexAI + + https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas?hl=id + """ + def __init__(self): + super().__init__() + + def get_supported_openai_params(self, model: str) -> list: + base_gpt_series_params = super().get_supported_openai_params(model=model) + gpt_oss_only_params = ["reasoning_effort"] + base_gpt_series_params.extend(gpt_oss_only_params) + + ######################################################### + # VertexAI - GPT-OSS does not support tool calls + ######################################################### + if litellm.supports_function_calling(model=model) is False: + TOOL_CALLING_PARAMS_TO_REMOVE = ["tool", "tool_choice", "function_call", "functions"] + base_gpt_series_params = [param for param in base_gpt_series_params if param not in TOOL_CALLING_PARAMS_TO_REMOVE] + + return base_gpt_series_params + diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index f281cab3b58..ea29970f0aa 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -1,5 +1,6 @@ # What is this? ## API Handler for calling Vertex AI Partner Models +from enum import Enum from typing import Callable, Optional, Union import httpx # type: ignore @@ -27,6 +28,16 @@ class VertexAIError(Exception): self.message ) # Call the base class constructor with the parameters it needs +class PartnerModelPrefixes(str, Enum): + META_PREFIX = "meta/" + DEEPSEEK_PREFIX = "deepseek-ai" + MISTRAL_PREFIX = "mistral" + CODERESTAL_PREFIX = "codestral" + JAMBA_PREFIX = "jamba" + CLAUDE_PREFIX = "claude" + QWEN_PREFIX = "qwen" + GPT_OSS_PREFIX = "openai/gpt-oss-" + class VertexAIPartnerModels(VertexBase): def __init__(self) -> None: @@ -42,13 +53,14 @@ class VertexAIPartnerModels(VertexBase): bool: True if the model string is a Vertex AI Partner Model, False otherwise """ if ( - model.startswith("meta/") - or model.startswith("deepseek-ai") - or model.startswith("mistral") - or model.startswith("codestral") - or model.startswith("jamba") - or model.startswith("claude") - or model.startswith("qwen") + model.startswith(PartnerModelPrefixes.META_PREFIX) + or model.startswith(PartnerModelPrefixes.DEEPSEEK_PREFIX) + or model.startswith(PartnerModelPrefixes.MISTRAL_PREFIX) + or model.startswith(PartnerModelPrefixes.CODERESTAL_PREFIX) + or model.startswith(PartnerModelPrefixes.JAMBA_PREFIX) + or model.startswith(PartnerModelPrefixes.CLAUDE_PREFIX) + or model.startswith(PartnerModelPrefixes.QWEN_PREFIX) + or model.startswith(PartnerModelPrefixes.GPT_OSS_PREFIX) ): return True return False @@ -57,8 +69,9 @@ class VertexAIPartnerModels(VertexBase): def should_use_openai_handler(model: str): OPENAI_LIKE_VERTEX_PROVIDERS = [ "llama", - "deepseek-ai", - "qwen", + PartnerModelPrefixes.DEEPSEEK_PREFIX, + PartnerModelPrefixes.QWEN_PREFIX, + PartnerModelPrefixes.GPT_OSS_PREFIX, ] if any(provider in model for provider in OPENAI_LIKE_VERTEX_PROVIDERS): return True diff --git a/litellm/llms/volcengine/__init__.py b/litellm/llms/volcengine/__init__.py new file mode 100644 index 00000000000..0887937bed5 --- /dev/null +++ b/litellm/llms/volcengine/__init__.py @@ -0,0 +1,24 @@ +""" +Volcengine LLM Provider +Support for Volcengine (ByteDance) chat and embedding models +""" + +from .chat.transformation import VolcEngineChatConfig +from .common_utils import ( + VolcEngineError, + get_volcengine_base_url, + get_volcengine_headers, +) +from .embedding import VolcEngineEmbeddingConfig + +# For backward compatibility, keep the old class name +VolcEngineConfig = VolcEngineChatConfig + +__all__ = [ + "VolcEngineChatConfig", + "VolcEngineConfig", # backward compatibility + "VolcEngineEmbeddingConfig", + "VolcEngineError", + "get_volcengine_base_url", + "get_volcengine_headers", +] diff --git a/litellm/llms/volcengine.py b/litellm/llms/volcengine/chat/transformation.py similarity index 91% rename from litellm/llms/volcengine.py rename to litellm/llms/volcengine/chat/transformation.py index c878aaf933c..216570a1aba 100644 --- a/litellm/llms/volcengine.py +++ b/litellm/llms/volcengine/chat/transformation.py @@ -3,7 +3,7 @@ from typing import Optional, Union from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig -class VolcEngineConfig(OpenAILikeChatConfig): +class VolcEngineChatConfig(OpenAILikeChatConfig): frequency_penalty: Optional[int] = None function_call: Optional[Union[str, dict]] = None functions: Optional[list] = None @@ -82,17 +82,19 @@ class VolcEngineConfig(OpenAILikeChatConfig): if "thinking" in optional_params: thinking_value = optional_params.pop("thinking") - + # Handle disabled thinking case - don't add to extra_body if disabled if ( - thinking_value is not None - and isinstance(thinking_value, dict) + thinking_value is not None + and isinstance(thinking_value, dict) and thinking_value.get("type") == "disabled" ): # Skip adding thinking parameter when it's disabled pass else: # Add thinking parameter to extra_body for all other cases - optional_params.setdefault("extra_body", {})["thinking"] = thinking_value + optional_params.setdefault("extra_body", {})[ + "thinking" + ] = thinking_value return optional_params diff --git a/litellm/llms/volcengine/common_utils.py b/litellm/llms/volcengine/common_utils.py new file mode 100644 index 00000000000..0c8d3daebdc --- /dev/null +++ b/litellm/llms/volcengine/common_utils.py @@ -0,0 +1,62 @@ +""" +Common utilities for Volcengine LLM provider +""" + +from typing import Optional + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException + + +class VolcEngineError(BaseLLMException): + """ + Custom exception class for Volcengine provider errors. + """ + + def __init__( + self, status_code: int, message: str, headers: Optional[httpx.Headers] = None + ): + self.status_code = status_code + self.message = message + self.headers = headers or httpx.Headers() + super().__init__( + status_code=status_code, message=message, headers=dict(self.headers) + ) + + +def get_volcengine_base_url(api_base: Optional[str] = None) -> str: + """ + Get the base URL for Volcengine API calls. + + Args: + api_base: Optional custom API base URL + + Returns: + The base URL to use for API calls + """ + if api_base: + return api_base + return "https://ark.cn-beijing.volces.com" + + +def get_volcengine_headers(api_key: str, extra_headers: Optional[dict] = None) -> dict: + """ + Get headers for Volcengine API calls. + + Args: + api_key: The API key for authentication + extra_headers: Optional additional headers + + Returns: + Dictionary of headers + """ + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + + if extra_headers: + headers.update(extra_headers) + + return headers diff --git a/litellm/llms/volcengine/embedding/__init__.py b/litellm/llms/volcengine/embedding/__init__.py new file mode 100644 index 00000000000..7b3efc4f961 --- /dev/null +++ b/litellm/llms/volcengine/embedding/__init__.py @@ -0,0 +1,7 @@ +""" +Volcengine Embedding Module +""" + +from .transformation import VolcEngineEmbeddingConfig + +__all__ = ["VolcEngineEmbeddingConfig"] diff --git a/litellm/llms/volcengine/embedding/transformation.py b/litellm/llms/volcengine/embedding/transformation.py new file mode 100644 index 00000000000..20747b76725 --- /dev/null +++ b/litellm/llms/volcengine/embedding/transformation.py @@ -0,0 +1,211 @@ +""" +Volcengine Embedding Transformation +Transforms OpenAI embedding requests to Volcengine format +""" + +from typing import List, Optional, Union, Dict, Any +import httpx +from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues +from litellm.types.utils import EmbeddingResponse +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from ..common_utils import get_volcengine_base_url, get_volcengine_headers + + +class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): + """ + Configuration class for Volcengine embedding models. + Reference: https://ark.cn-beijing.volces.com/api/v3/embeddings + """ + + def __init__( + self, + encoding_format: Optional[str] = None, + ) -> None: + locals_ = locals().copy() + for key, value in locals_.items(): + if key != "self" and value is not None: + setattr(self.__class__, key, value) + + @classmethod + def get_config(cls): + return super().get_config() + + def get_supported_openai_params(self, model: str) -> List[str]: + """ + Get the list of OpenAI parameters supported by Volcengine embedding models. + + Args: + model: The model name + + Returns: + List of supported parameter names + """ + return [ + "encoding_format", + "user", + "extra_headers", + ] + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete URL for volcengine embedding API calls. + + Args: + api_base: Optional custom API base URL + api_key: API key (not used for URL construction) + model: Model name (not used for URL construction) + optional_params: Optional parameters (not used for URL construction) + litellm_params: LiteLLM parameters (not used for URL construction) + stream: Stream parameter (not used for URL construction) + + Returns: + Complete URL for the embedding API endpoint + """ + base_url = get_volcengine_base_url(api_base) + # Construct the complete URL with /embeddings endpoint + if base_url.endswith("/api/v3"): + return f"{base_url}/embeddings" + else: + return f"{base_url}/api/v3/embeddings" + + def map_openai_params( + self, + non_default_params: Dict[str, Any], + optional_params: Dict[str, Any], + model: str, + drop_params: bool, + ) -> Dict[str, Any]: + """ + Map OpenAI embedding parameters to Volcengine format. + + Args: + non_default_params: Parameters that are not default values + optional_params: Optional parameters dict to update + model: The model name + drop_params: Whether to drop unsupported parameters + + Returns: + Updated optional_params dict + """ + for param, value in non_default_params.items(): + if param == "encoding_format": + # Volcengine supports: float, base64, null + if value in ["float", "base64", None]: + optional_params["encoding_format"] = value + else: + if not drop_params: + raise ValueError( + f"Unsupported encoding_format: {value}. Volcengine supports: float, base64, null" + ) + elif param == "user": + # Keep user parameter as-is + optional_params["user"] = value + elif param in self.get_supported_openai_params(model): + optional_params[param] = value + elif not drop_params: + raise ValueError(f"Unsupported parameter for Volcengine: {param}") + + return optional_params + + + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict, + headers: dict, + ) -> dict: + """Transform embedding request to Volcengine format""" + # Prepare request data (only the JSON body, not the full request) + data = { + "model": model, + "input": input if isinstance(input, list) else [input], + } + + # Add optional parameters from optional_params + if "encoding_format" in optional_params: + encoding_format = optional_params["encoding_format"] + if encoding_format is not None: + data["encoding_format"] = encoding_format + + if "user" in optional_params: + user = optional_params["user"] + if user is not None: + data["user"] = user + + return data + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + request_data: dict, + optional_params: dict, + litellm_params: dict, + ) -> EmbeddingResponse: + """Transform Volcengine response to EmbeddingResponse""" + try: + response_json = raw_response.json() + except Exception as e: + raise ValueError(f"Failed to parse Volcengine response as JSON: {str(e)}") + + # Volcengine response format matches OpenAI format closely + # Just need to ensure all required fields are present + transformed_response = { + "object": "list", + "data": response_json.get("data", []), + "model": response_json.get("model", model), + "usage": response_json.get("usage", {}), + } + + # Add id if present + if "id" in response_json: + transformed_response["id"] = response_json["id"] + + # Create EmbeddingResponse from transformed data + return EmbeddingResponse(**transformed_response) + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """Validate environment and return headers""" + # Get Volcengine headers + if api_key is None: + raise ValueError("api_key is required for Volcengine authentication") + volcengine_headers = get_volcengine_headers(api_key) + return {**headers, **volcengine_headers} + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> BaseLLMException: + """Get error class for Volcengine errors""" + from ..common_utils import VolcEngineError + # Convert dict to httpx.Headers if needed + if isinstance(headers, dict): + headers = httpx.Headers(headers) + return VolcEngineError( + status_code=status_code, + message=error_message, + headers=headers, + ) diff --git a/litellm/main.py b/litellm/main.py index 786a0196e5e..decbebaf485 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -356,7 +356,8 @@ async def acompletion( logprobs: Optional[bool] = None, top_logprobs: Optional[int] = None, deployment_id=None, - reasoning_effort: Optional[Literal["minimal", "low", "medium", "high"]] = None, + reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "default"]] = None, + safety_identifier: Optional[str] = None, # set api_base, api_version, api_key base_url: Optional[str] = None, api_version: Optional[str] = None, @@ -493,6 +494,7 @@ async def acompletion( "api_key": api_key, "model_list": model_list, "reasoning_effort": reasoning_effort, + "safety_identifier": safety_identifier, "extra_headers": extra_headers, "acompletion": True, # assuming this is a required parameter "thinking": thinking, @@ -500,7 +502,7 @@ async def acompletion( } if custom_llm_provider is None: _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=completion_kwargs.get("base_url", None) + model=model, custom_llm_provider=custom_llm_provider, api_base=completion_kwargs.get("base_url", None) ) fallbacks = fallbacks or litellm.model_fallbacks @@ -895,7 +897,7 @@ def completion( # type: ignore # noqa: PLR0915 logit_bias: Optional[dict] = None, user: Optional[str] = None, # openai v1.0+ new params - reasoning_effort: Optional[Literal["minimal", "low", "medium", "high"]] = None, + reasoning_effort: Optional[Literal["none", "minimal", "low", "medium", "high", "default"]] = None, response_format: Optional[Union[dict, Type[BaseModel]]] = None, seed: Optional[int] = None, tools: Optional[List] = None, @@ -906,6 +908,7 @@ def completion( # type: ignore # noqa: PLR0915 web_search_options: Optional[OpenAIWebSearchOptions] = None, deployment_id=None, extra_headers: Optional[dict] = None, + safety_identifier: Optional[str] = None, # soon to be deprecated params by OpenAI functions: Optional[List] = None, function_call: Optional[str] = None, @@ -1243,6 +1246,7 @@ def completion( # type: ignore # noqa: PLR0915 "reasoning_effort": reasoning_effort, "thinking": thinking, "web_search_options": web_search_options, + "safety_identifier": safety_identifier, "allowed_openai_params": kwargs.get("allowed_openai_params"), } optional_params = get_optional_params( @@ -3667,7 +3671,7 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: model = args[0] if len(args) > 0 else kwargs["model"] ### PASS ARGS TO Embedding ### kwargs["aembedding"] = True - custom_llm_provider = None + custom_llm_provider = kwargs.get("custom_llm_provider", None) try: # Use a partial function to pass your keyword arguments func = partial(embedding, *args, **kwargs) @@ -3677,7 +3681,7 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: func_with_context = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) + model=model, custom_llm_provider=custom_llm_provider, api_base=kwargs.get("api_base", None) ) # Await normally @@ -4499,6 +4503,36 @@ def embedding( # noqa: PLR0915 client=client, aembedding=aembedding, ) + elif custom_llm_provider == "volcengine": + volcengine_key = ( + api_key + or litellm.api_key + or get_secret_str("ARK_API_KEY") + or get_secret_str("VOLCENGINE_API_KEY") + ) + if volcengine_key is None: + raise ValueError( + "Missing API key for Volcengine. Set ARK_API_KEY or VOLCENGINE_API_KEY environment variable or pass api_key parameter." + ) + if extra_headers is not None and isinstance(extra_headers, dict): + headers = extra_headers + else: + headers = {} + response = base_llm_http_handler.embedding( + model=model, + input=input, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + logging_obj=logging, + api_base=api_base, + optional_params=optional_params, + litellm_params={}, + model_response=EmbeddingResponse(), + api_key=volcengine_key, + client=client, + aembedding=aembedding, + headers=headers, + ) elif custom_llm_provider in litellm._custom_providers: custom_handler: Optional[CustomLLM] = None for item in litellm.custom_provider_map: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9abe3dc48fe..46eb48d2d42 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5817,16 +5817,6 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "groq/llama3-8b-8192": { - "max_tokens": 8192, - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "input_cost_per_token": 5e-08, - "output_cost_per_token": 8e-08, - "litellm_provider": "groq", - "mode": "chat", - "supports_tool_choice": true - }, "groq/llama-3.2-1b-preview": { "max_tokens": 8192, "max_input_tokens": 8192, @@ -5907,17 +5897,6 @@ "supports_tool_choice": true, "deprecation_date": "2025-04-14" }, - "groq/llama3-70b-8192": { - "max_tokens": 8192, - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "input_cost_per_token": 5.9e-07, - "output_cost_per_token": 7.9e-07, - "litellm_provider": "groq", - "mode": "chat", - "supports_response_schema": true, - "supports_tool_choice": true - }, "groq/llama-3.1-8b-instant": { "max_tokens": 8192, "max_input_tokens": 128000, @@ -6178,21 +6157,7 @@ "supports_tool_choice": true, "source": "https://inference-docs.cerebras.ai/support/pricing" }, - "cerebras/openai/gpt-oss-20b": { - "max_tokens": 32768, - "max_input_tokens": 131072, - "max_output_tokens": 32768, - "input_cost_per_token": 7e-08, - "output_cost_per_token": 3e-07, - "litellm_provider": "cerebras", - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_reasoning": true, - "supports_tool_choice": true, - "source": "https://inference-docs.cerebras.ai/support/pricing" - }, + "cerebras/openai/gpt-oss-120b": { "max_tokens": 32768, "max_input_tokens": 131072, @@ -8027,8 +7992,8 @@ "max_pdf_size_mb": 30, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, - "output_cost_per_token": 2.5e-06, - "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_reasoning_token": 3e-05, "output_cost_per_image": 0.039, "litellm_provider": "gemini", "mode": "chat", @@ -8391,8 +8356,8 @@ "max_pdf_size_mb": 30, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, - "output_cost_per_token": 2.5e-06, - "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_reasoning_token": 3e-05, "output_cost_per_image": 0.039, "litellm_provider": "vertex_ai-language-models", "mode": "chat", @@ -9519,6 +9484,48 @@ "source": "https://aistudio.google.com", "supports_tool_choice": true }, + "gemini/veo-3.0-generate-preview": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.75, + "litellm_provider": "gemini", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, + "gemini/veo-3.0-fast-generate-preview": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.40, + "litellm_provider": "gemini", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, + "gemini/veo-2.0-generate-001": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.35, + "litellm_provider": "gemini", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, "vertex_ai/claude-opus-4-1": { "max_tokens": 4096, "max_input_tokens": 200000, @@ -9905,6 +9912,28 @@ "supports_tool_choice": true, "supports_prompt_caching": true }, + "vertex_ai/openai/gpt-oss-20b-maas": { + "max_tokens": 32768, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "input_cost_per_token": 0.075e-06, + "output_cost_per_token": 0.30e-06, + "litellm_provider": "vertex_ai-openai_models", + "mode": "chat", + "supports_reasoning": true, + "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas" + }, + "vertex_ai/openai/gpt-oss-120b-maas": { + "max_tokens": 32768, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "input_cost_per_token": 0.15e-06, + "output_cost_per_token": 0.60e-06, + "litellm_provider": "vertex_ai-openai_models", + "mode": "chat", + "supports_reasoning": true, + "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas" + }, "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { "max_tokens": 32768, "max_input_tokens": 262144, @@ -10314,6 +10343,48 @@ "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "vertex_ai/veo-3.0-generate-preview": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.75, + "litellm_provider": "vertex_ai-video-models", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, + "vertex_ai/veo-3.0-fast-generate-preview": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.40, + "litellm_provider": "vertex_ai-video-models", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, + "vertex_ai/veo-2.0-generate-001": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.35, + "litellm_provider": "vertex_ai-video-models", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, "text-embedding-004": { "max_tokens": 2048, "max_input_tokens": 2048, @@ -20962,5 +21033,65 @@ "metadata": { "notes": "DALL-E 2 via AI/ML API - Reliable text-to-image generation" } + }, + "doubao-embedding-large": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 2048, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - large version with 2048 dimensions" + } + }, + "doubao-embedding-large-text-250515": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 2048, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - text-250515 version with 2048 dimensions" + } + }, + "doubao-embedding-large-text-240915": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 4096, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - text-240915 version with 4096 dimensions" + } + }, + "doubao-embedding": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 2560, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - standard version with 2560 dimensions" + } + }, + "doubao-embedding-text-240715": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 2560, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - text-240715 version with 2560 dimensions" + } } } \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 5c5f1cfe908..00000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index b3653c31435..b38272560d1 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,30 +1,27 @@ model_list: - - model_name: fake-openai-endpoint - litellm_params: - model: openai/fake - api_key: fake-key - api_base: https://exampleopenaiendpoint-production.up.railway.app/ - - model_name: gpt-5-mini - litellm_params: - model: azure/gpt-5-mini - api_base: os.environ/AZURE_GPT_5_MINI_API_BASE # runs os.getenv("AZURE_API_BASE") - api_key: os.environ/AZURE_GPT_5_MINI_API_KEY # runs os.getenv("AZURE_API_KEY") - stream_timeout: 60 - merge_reasoning_content_in_choices: true - model_info: - mode: chat + - model_name: fake-openai-endpoint + litellm_params: + model: openai/fake + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + - model_name: gpt-5-mini + litellm_params: + model: azure/gpt-5-mini + api_base: os.environ/AZURE_GPT_5_MINI_API_BASE # runs os.getenv("AZURE_API_BASE") + api_key: os.environ/AZURE_GPT_5_MINI_API_KEY # runs os.getenv("AZURE_API_KEY") + stream_timeout: 60 + merge_reasoning_content_in_choices: true + model_info: + mode: chat + - model_name: ollama-deepseek-r1 + litellm_params: + model: ollama/deepseek-r1:1.5b + model_info: + mode: chat router_settings: model_group_alias: {"my-fake-gpt-4": "fake-openai-endpoint"} - + litellm_settings: callbacks: ["otel"] - cache: true - cache_params: - type: redis - ttl: 600 - supported_call_types: ["acompletion", "completion"] - - model_group_settings: - forward_client_headers_to_llm_api: - - fake-openai-endpoint \ No newline at end of file + success_callback: ["braintrust"] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 87c5bd7a7c5..66bd5977551 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2,7 +2,16 @@ import enum import json import uuid from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + Union, +) import httpx from pydantic import ( @@ -388,7 +397,11 @@ class LiteLLMRoutes(enum.Enum): ] # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend - master_key_only_routes = ["/global/spend/reset"] + master_key_only_routes = [ + "/global/spend/reset", + "/memory-usage-in-mem-cache", + "/memory-usage-in-mem-cache-items", + ] key_management_routes = [ KeyManagementRoutes.KEY_GENERATE, @@ -774,7 +787,6 @@ class GenerateKeyRequest(KeyRequestBase): description="Type of key that determines default allowed routes.", ) - class GenerateKeyResponse(KeyRequestBase): key: str # type: ignore key_name: Optional[str] = None @@ -2908,6 +2920,12 @@ class LitellmDataForBackendLLMCall(TypedDict, total=False): user: Optional[str] num_retries: Optional[int] +class LitellmMetadataFromRequestHeaders(TypedDict, total=False): + """ + Headers a user can pass that will get added to litellm metadata for the request + """ + spend_logs_metadata: Optional[dict] + class JWTKeyItem(TypedDict, total=False): kid: str diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index a10a39a6a57..2de5ec1ee12 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -90,6 +90,17 @@ async def anthropic_response( # noqa: PLR0915 user_api_key_dict=user_api_key_dict, data=data, call_type="text_completion" ) + tasks = [] + tasks.append( + proxy_logging_obj.during_call_hook( + data=data, + user_api_key_dict=user_api_key_dict, + call_type=ProxyBaseLLMRequestProcessing._get_pre_call_type( + route_type="anthropic_messages" # type: ignore + ), + ) + ) + ### ROUTE THE REQUESTs ### router_model_names = llm_router.model_names if llm_router is not None else [] @@ -97,23 +108,21 @@ async def anthropic_response( # noqa: PLR0915 if ( llm_router is not None and data["model"] in router_model_names ): # model in router model list - llm_response = asyncio.create_task(llm_router.aanthropic_messages(**data)) + llm_coro = llm_router.aanthropic_messages(**data) elif ( llm_router is not None and llm_router.model_group_alias is not None and data["model"] in llm_router.model_group_alias ): # model set in model_group_alias - llm_response = asyncio.create_task(llm_router.aanthropic_messages(**data)) + llm_coro = llm_router.aanthropic_messages(**data) elif ( llm_router is not None and data["model"] in llm_router.deployment_names ): # model in router deployments, calling a specific deployment on the router - llm_response = asyncio.create_task( - llm_router.aanthropic_messages(**data, specific_deployment=True) - ) + llm_coro = llm_router.aanthropic_messages(**data, specific_deployment=True) elif ( llm_router is not None and data["model"] in llm_router.get_model_ids() ): # model in router model list - llm_response = asyncio.create_task(llm_router.aanthropic_messages(**data)) + llm_coro = llm_router.aanthropic_messages(**data) elif ( llm_router is not None and data["model"] not in router_model_names @@ -122,9 +131,9 @@ async def anthropic_response( # noqa: PLR0915 or len(llm_router.pattern_router.patterns) > 0 ) ): # model in router deployments, calling a specific deployment on the router - llm_response = asyncio.create_task(llm_router.aanthropic_messages(**data)) + llm_coro = llm_router.aanthropic_messages(**data) elif user_model is not None: # `litellm --model ` - llm_response = asyncio.create_task(litellm.anthropic_messages(**data)) + llm_coro = litellm.anthropic_messages(**data) else: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -134,8 +143,16 @@ async def anthropic_response( # noqa: PLR0915 }, ) - # Await the llm_response task - response = await llm_response + tasks.append(llm_coro) + + # wait for call to end + llm_responses = asyncio.gather( + *tasks + ) # run the moderation check in parallel to the actual llm api call + + responses = await llm_responses + + response = responses[1] hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = hidden_params.get("model_id", None) or "" @@ -183,6 +200,11 @@ async def anthropic_response( # noqa: PLR0915 headers=dict(fastapi_response.headers), ) + ### CALL HOOKS ### - modify outgoing data + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response # type: ignore + ) + verbose_proxy_logger.info("\nResponse from Litellm:\n{}".format(response)) return response except Exception as e: diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index f20c3debc23..5f78efbdf40 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -7,7 +7,6 @@ JWT token must have 'litellm_proxy_admin' in scope. """ import fnmatch -import json import os from typing import Any, List, Literal, Optional, Set, Tuple, cast @@ -484,7 +483,7 @@ class JWTHandler: # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret # the key in different ways (e.g. HS* and RS*)." - algorithms = ["RS256", "RS384", "RS512", "PS256", "PS384", "PS512"] + algorithms = ["RS256", "RS384", "RS512", "PS256", "PS384", "PS512", "ES256", "ES384", "ES512", "EdDSA"] audience = os.getenv("JWT_AUDIENCE") decode_options = None @@ -492,7 +491,7 @@ class JWTHandler: decode_options = {"verify_aud": False} import jwt - from jwt.algorithms import RSAAlgorithm + from jwt.api_jwk import PyJWK header = jwt.get_unverified_header(token) @@ -512,14 +511,21 @@ class JWTHandler: jwk["n"] = public_key["n"] if "e" in public_key: jwk["e"] = public_key["e"] + if "x" in public_key: + jwk["x"] = public_key["x"] + if "y" in public_key: + jwk["y"] = public_key["y"] + if "crv" in public_key: + jwk["crv"] = public_key["crv"] - public_key_rsa = RSAAlgorithm.from_jwk(json.dumps(jwk)) + # parse RSA/EC/OKP keys + public_key_obj = PyJWK.from_dict(jwk).key try: # decode the token using the public key payload = jwt.decode( token, - public_key_rsa, # type: ignore + public_key_obj, # type: ignore algorithms=algorithms, options=decode_options, audience=audience, @@ -534,9 +540,7 @@ class JWTHandler: raise Exception(f"Validation fails: {str(e)}") elif public_key is not None and isinstance(public_key, str): try: - cert = x509.load_pem_x509_certificate( - public_key.encode(), default_backend() - ) + cert = x509.load_pem_x509_certificate(public_key.encode(), default_backend()) # Extract public key key = cert.public_key().public_bytes( @@ -561,7 +565,7 @@ class JWTHandler: raise Exception(f"Validation fails: {str(e)}") raise Exception("Invalid JWT Submitted") - + async def close(self): await self.http_handler.close() @@ -1210,4 +1214,4 @@ class JWTAuthManager: end_user_object=end_user_object, token=api_key, team_membership=team_membership_object, - ) + ) \ No newline at end of file diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 68fa80c2b0f..e900975f1cc 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -109,7 +109,6 @@ async def create_streaming_response( final_status_code = default_status_code try: - # Handle coroutine that returns a generator if asyncio.iscoroutine(generator): generator = await generator @@ -118,7 +117,6 @@ async def create_streaming_response( first_chunk_value = await generator.__anext__() if first_chunk_value is not None: - try: error_code_from_chunk = await _parse_event_data_for_error( first_chunk_value @@ -132,7 +130,6 @@ async def create_streaming_response( verbose_proxy_logger.debug(f"Error parsing first chunk value: {e}") except StopAsyncIteration: - # Generator was empty. Default status async def empty_gen() -> AsyncGenerator[str, None]: if False: @@ -145,7 +142,6 @@ async def create_streaming_response( status_code=default_status_code, ) except Exception as e: - # Unexpected error consuming first chunk. verbose_proxy_logger.exception( f"Error consuming first chunk from generator: {e}" @@ -168,7 +164,6 @@ async def create_streaming_response( with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): yield first_chunk_value async for chunk in generator: - with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): yield chunk @@ -462,7 +457,6 @@ class ProxyBaseLLMRequestProcessing: ) or self._is_streaming_response( response ): # use generate_responses to stream responses - custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, call_id=logging_obj.litellm_call_id, @@ -480,7 +474,6 @@ class ProxyBaseLLMRequestProcessing: if route_type == "allm_passthrough_route": # Check if response is an async generator if self._is_streaming_response(response): - if asyncio.iscoroutine(response): generator = await response else: @@ -501,7 +494,6 @@ class ProxyBaseLLMRequestProcessing: headers=custom_headers, ) else: - selected_data_generator = select_data_generator( response=response, user_api_key_dict=user_api_key_dict, @@ -740,7 +732,11 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.debug("inside generator") try: str_so_far = "" - async for chunk in response: + async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=response, + request_data=request_data, + ): verbose_proxy_logger.debug( "async_data_generator: received streaming chunk - {}".format(chunk) ) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index e718255750a..d52592952bc 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -317,17 +317,26 @@ def get_remaining_tokens_and_requests_from_request_data(data: Dict) -> Dict[str, _metadata = data.get("metadata", None) or {} model_group = get_model_group_from_request_data(data) + # The h11 package considers "/" or ":" invalid and raise a LocalProtocolError + h11_model_group_name = ( + model_group.replace("/", "-").replace(":", "-") if model_group else None + ) + # Remaining Requests remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}" remaining_requests = _metadata.get(remaining_requests_variable_name, None) if remaining_requests: - headers[f"x-litellm-key-remaining-requests-{model_group}"] = remaining_requests + headers[f"x-litellm-key-remaining-requests-{h11_model_group_name}"] = ( + remaining_requests + ) # Remaining Tokens remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" remaining_tokens = _metadata.get(remaining_tokens_variable_name, None) if remaining_tokens: - headers[f"x-litellm-key-remaining-tokens-{model_group}"] = remaining_tokens + headers[f"x-litellm-key-remaining-tokens-{h11_model_group_name}"] = ( + remaining_tokens + ) return headers diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 8096d782dec..16ab2cc8058 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -5,10 +5,12 @@ import os import tracemalloc from collections import Counter -from fastapi import APIRouter +from fastapi import APIRouter, Depends from litellm import get_secret_str from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() @@ -84,7 +86,9 @@ if os.environ.get("LITELLM_PROFILE", "false").lower() == "true": @router.get("/memory-usage-in-mem-cache", include_in_schema=False) -async def memory_usage_in_mem_cache(): +async def memory_usage_in_mem_cache( + _: UserAPIKeyAuth = Depends(user_api_key_auth), +): # returns the size of all in-memory caches on the proxy server """ 1. user_api_key_cache @@ -121,7 +125,9 @@ async def memory_usage_in_mem_cache(): @router.get("/memory-usage-in-mem-cache-items", include_in_schema=False) -async def memory_usage_in_mem_cache_items(): +async def memory_usage_in_mem_cache_items( + _: UserAPIKeyAuth = Depends(user_api_key_auth), +): # returns the size of all in-memory caches on the proxy server """ 1. user_api_key_cache diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 373232e22d2..eb481b0a4f0 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -173,15 +173,24 @@ async def google_count_tokens(request: Request, model_name: str): """ from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.proxy_server import token_counter as internal_token_counter + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter data = await _read_request_body(request=request) contents = data.get("contents", []) #Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest + # Translate contents to openai format messages using the adapter + messages = ( + GoogleGenAIAdapter() + .translate_generate_content_to_completion(model_name, contents) + .get("messages", []) + ) + token_request = TokenCountRequest( model=model_name, - contents=contents + contents=contents, + messages=messages, # compatibility when use openai-like endpoint ) # Call the internal token counter function with direct request flag set to False @@ -192,11 +201,17 @@ async def google_count_tokens(request: Request, model_name: str): if token_response is not None: # cast the response to the well known format original_response: dict = token_response.original_response or {} - return TokenCountDetailsResponse( - totalTokens=original_response.get("totalTokens", 0), - promptTokensDetails=original_response.get("promptTokensDetails", []), - ) - + if original_response: + return TokenCountDetailsResponse( + totalTokens=original_response.get("totalTokens", 0), + promptTokensDetails=original_response.get("promptTokensDetails", []), + ) + else: + return TokenCountDetailsResponse( + totalTokens=token_response.total_tokens or 0, + promptTokensDetails=[], + ) + ######################################################### # Return the response in the well known format ######################################################### diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b04d14bcc8b..b3840761d2a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -68,6 +68,32 @@ end return results """ +TOKEN_INCREMENT_SCRIPT = """ +local results = {} + +-- Process each key/increment_value/ttl triplet +for i = 1, #KEYS do + local key = KEYS[i] + local increment_value = tonumber(ARGV[i * 2 - 1]) + local ttl_seconds = tonumber(ARGV[i * 2]) + + -- Increment the value + local new_value = redis.call('INCRBYFLOAT', key, increment_value) + + -- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL + -- ttl_seconds can be 0 (no TTL) or positive (set TTL) + if ttl_seconds and ttl_seconds > 0 then + local current_ttl = redis.call('TTL', key) + if current_ttl == -1 then + redis.call('EXPIRE', key, ttl_seconds) + end + end + + table.insert(results, new_value) +end + +return results +""" class RateLimitDescriptorRateLimitObject(TypedDict, total=False): requests_per_unit: Optional[int] @@ -109,8 +135,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): BATCH_RATE_LIMITER_SCRIPT ) ) + self.token_increment_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + TOKEN_INCREMENT_SCRIPT + ) + ) else: self.batch_rate_limiter_script = None + self.token_increment_script = None self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) @@ -567,6 +599,62 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return pipeline_operations + async def async_increment_tokens_with_ttl_preservation( + self, + pipeline_operations: List["RedisPipelineIncrementOperation"], + parent_otel_span: Optional[Span] = None, + ) -> None: + """ + Increment token counters using Lua script to preserve existing TTL. + This prevents TTL reset on every token increment. + """ + if not pipeline_operations: + return + + # Check if script is available + if self.token_increment_script is None: + verbose_proxy_logger.debug("TTL preservation script not available, using regular pipeline") + await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, + litellm_parent_otel_span=parent_otel_span, + ) + return + + try: + # Use Lua script for all operations + keys = [] + args = [] + + for op in pipeline_operations: + # Convert None TTL to 0 for Lua script + ttl_value = op["ttl"] if op["ttl"] is not None else 0 + + verbose_proxy_logger.debug( + f"Executing TTL-preserving increment for key={op['key']}, " + f"increment={op['increment_value']}, ttl={ttl_value}" + ) + keys.append(op["key"]) + args.extend([op["increment_value"], ttl_value]) + + await self.token_increment_script( + keys=keys, + args=args, + ) + + verbose_proxy_logger.debug( + f"Successfully executed TTL-preserving increment for {len(pipeline_operations)} keys" + ) + + except Exception as e: + verbose_proxy_logger.warning( + f"TTL preservation failed, falling back to regular pipeline: {str(e)}" + ) + # Fallback to regular pipeline on error + await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, + litellm_parent_otel_span=parent_otel_span, + ) + def get_rate_limit_type(self) -> Literal["output", "input", "total"]: from litellm.proxy.proxy_server import general_settings @@ -713,9 +801,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Execute all increments in a single pipeline if pipeline_operations: - await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - litellm_parent_otel_span=litellm_parent_otel_span, + await self.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations, + parent_otel_span=litellm_parent_otel_span, ) except Exception as e: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 32a283066ed..95f1eccffe4 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -291,6 +291,17 @@ class LiteLLMProxyRequestSetup: if num_retries_header is not None: return int(num_retries_header) return None + + @staticmethod + def _get_spend_logs_metadata_from_request_headers(headers: dict) -> Optional[dict]: + """ + Get the `spend_logs_metadata` from the request headers. + """ + from litellm.litellm_core_utils.safe_json_loads import safe_json_loads + spend_logs_metadata_header = headers.get("x-litellm-spend-logs-metadata", None) + if spend_logs_metadata_header is not None: + return safe_json_loads(spend_logs_metadata_header) + return None @staticmethod def _get_forwardable_headers( @@ -459,6 +470,30 @@ class LiteLLMProxyRequestSetup: data["num_retries"] = num_retries return data + + @staticmethod + def add_litellm_metadata_from_request_headers( + headers: dict, + data: dict, + _metadata_variable_name: str, + ) -> dict: + """ + Add litellm metadata from request headers + + Relevant issue: https://github.com/BerriAI/litellm/issues/14008 + """ + from litellm.proxy._types import LitellmMetadataFromRequestHeaders + metadata_from_headers = LitellmMetadataFromRequestHeaders() + spend_logs_metadata = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers) + if spend_logs_metadata is not None: + metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata + + ######################################################################################### + # Finally update the requests metadata with the `metadata_from_headers` + ######################################################################################### + if isinstance(data[_metadata_variable_name], dict): + data[_metadata_variable_name].update(metadata_from_headers) + return data @staticmethod def get_sanitized_user_information_from_key( @@ -642,7 +677,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915 from litellm.proxy.proxy_server import llm_router, premium_user from litellm.types.proxy.litellm_pre_call_utils import SecretFields - safe_add_api_version_from_query_params(data, request) _headers = clean_headers( request.headers, @@ -653,6 +687,24 @@ async def add_litellm_data_to_request( # noqa: PLR0915 ), ) + ########################################################## + # Init - Proxy Server Request + # we do this as soon as entering so we track the original request + ########################################################## + data["proxy_server_request"] = { + "url": str(request.url), + "method": request.method, + "headers": _headers, + "body": copy.copy(data), # use copy instead of deepcopy + } + + safe_add_api_version_from_query_params(data, request) + _metadata_variable_name = _get_metadata_variable_name(request) + if data.get(_metadata_variable_name, None) is None: + data[_metadata_variable_name] = {} + + + data.update( LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( headers=_headers, @@ -661,6 +713,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915 ) ) + data.update( + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=_headers, + data=data, + _metadata_variable_name=_metadata_variable_name, + ) + ) + # check for forwardable headers data = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group( data=data, headers=_headers, user_api_key_dict=user_api_key_dict @@ -674,13 +734,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915 if "user" not in data: data["user"] = user - # Include original request and headers in the data - data["proxy_server_request"] = { - "url": str(request.url), - "method": request.method, - "headers": _headers, - "body": copy.copy(data), # use copy instead of deepcopy - } data["secret_fields"] = SecretFields(raw_headers=dict(request.headers)) @@ -711,11 +764,6 @@ async def add_litellm_data_to_request( # noqa: PLR0915 verbose_proxy_logger.debug("receiving data: %s", data) - _metadata_variable_name = _get_metadata_variable_name(request) - - if data.get(_metadata_variable_name, None) is None: - data[_metadata_variable_name] = {} - # Parse metadata if it's a string (e.g., from multipart/form-data) if "metadata" in data and data["metadata"] is not None: if isinstance(data["metadata"], str): diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ddce7481ce7..8a3507e2398 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -346,6 +346,35 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: data_json["allowed_routes"] = ["info_routes"] return data_json +async def validate_team_id_used_in_service_account_request( + team_id: Optional[str], + prisma_client: Optional[PrismaClient], +): + """ + Validate team_id is used in the request body for generating a service account key + """ + if team_id is None: + raise HTTPException( + status_code=400, + detail="team_id is required for service account keys. Please specify `team_id` in the request body.", + ) + + if prisma_client is None: + raise HTTPException( + status_code=400, + detail="prisma_client is required for service account keys. Please specify `prisma_client` in the request body.", + ) + + # check if team_id exists in the database + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + ) + if team is None: + raise HTTPException( + status_code=400, + detail="team_id does not exist in the database. Please specify a valid `team_id` in the request body.", + ) + return True async def _common_key_generation_helper( # noqa: PLR0915 data: GenerateKeyRequest, @@ -372,9 +401,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 and data.metadata.get("service_account_id") is not None and data.team_id is None ): - raise HTTPException( - status_code=400, - detail="team_id is required for service account keys. Please specify `team_id` in the request body.", + await validate_team_id_used_in_service_account_request( + team_id=data.team_id, + prisma_client=prisma_client, ) # check if user set default key/generate params on config.yaml @@ -756,6 +785,11 @@ async def generate_service_account_key_fn( user_custom_key_generate, ) + await validate_team_id_used_in_service_account_request( + team_id=data.team_id, + prisma_client=prisma_client, + ) + verbose_proxy_logger.debug("entered /key/generate") if user_custom_key_generate is not None: @@ -1566,14 +1600,12 @@ async def generate_key_helper_fn( # noqa: PLR0915 if duration is None: # allow tokens that never expire expires = None else: - duration_s = duration_in_seconds(duration=duration) - expires = datetime.now(timezone.utc) + timedelta(seconds=duration_s) + expires = get_budget_reset_time(budget_duration=duration) if key_budget_duration is None: # one-time budget key_reset_at = None else: - duration_s = duration_in_seconds(duration=key_budget_duration) - key_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) + key_reset_at = get_budget_reset_time(budget_duration=key_budget_duration) if budget_duration is None: # one-time budget reset_at = None diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b8762899f1e..2e1a684e397 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -987,7 +987,7 @@ async def update_public_model_groups( try: # Update the public model groups import litellm - from litellm.proxy.proxy_server import proxy_config + from litellm.proxy.proxy_server import proxy_config, store_model_in_db # Check if user has admin permissions if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: @@ -1000,6 +1000,15 @@ async def update_public_model_groups( }, ) + # Check if STORE_MODEL_IN_DB is enabled + if store_model_in_db is not True: + raise HTTPException( + status_code=500, + detail={ + "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." + }, + ) + litellm.public_model_groups = request.model_groups # Load existing config diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py index bb07cdbd770..1bb59888405 100644 --- a/litellm/proxy/management_endpoints/scim/scim_transformations.py +++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py @@ -121,15 +121,16 @@ class ScimTransformations: if isinstance(team, dict): team = LiteLLM_TeamTable(**team) - # Get team members + # Get team members with proper display names scim_members: List[SCIMMember] = [] for member in team.members_with_roles or []: if isinstance(member, dict): member = Member(**member) + scim_members.append( SCIMMember( value=ScimTransformations._get_scim_member_value(member), - display=member.user_email, + display=ScimTransformations._get_scim_member_display(member), ) ) @@ -151,6 +152,24 @@ class ScimTransformations: @staticmethod def _get_scim_member_value(member: Member) -> str: - if member.user_email: + """ + Get the SCIM member value. Use user_email if available, otherwise use user_id. + SCIM member value should be the unique identifier for the user. + """ + if hasattr(member, "user_email") and member.user_email: return member.user_email + elif hasattr(member, "user_id"): + return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE + return ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE + + @staticmethod + def _get_scim_member_display(member: Member) -> str: + """ + Get the SCIM member display. Use user_email if available, otherwise use user_id. + SCIM member display should be the display name for the user. + """ + if hasattr(member, "user_email") and member.user_email: + return member.user_email + elif hasattr(member, "user_id"): + return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE return ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index e1d94ceaff2..b9929c3b433 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -22,6 +22,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( + LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, Member, @@ -237,6 +238,23 @@ async def _handle_team_membership_changes(user_id: str, existing_teams: List[str ) +async def _get_team_member_user_ids_from_team(team: LiteLLM_TeamTable) -> List[str]: + """ + Get the IDs of the members from a team. + + Use one source of truth for the member IDs: team.members_with_roles + + """ + member_user_ids: List[str] = [] + for member in team.members_with_roles or []: + if hasattr(member, "user_id") and member.user_id is not None: + member_user_ids.append(member.user_id) + elif isinstance(member, dict) and "user_id" in member: + user_id = member.get("user_id") + if user_id is not None: + member_user_ids.append(user_id) + return member_user_ids + # Dependency to set the correct SCIM Content-Type async def set_scim_content_type(response: Response): """Sets the Content-Type header to application/scim+json""" @@ -253,6 +271,12 @@ async def set_scim_content_type(response: Response): ) async def get_service_provider_config(request: Request): """Return SCIM Service Provider Configuration.""" + verbose_proxy_logger.debug( + "SCIM ServiceProviderConfig request: method=%s url=%s headers=%s", + request.method, + request.url, + dict(request.headers), + ) meta = { "resourceType": "ServiceProviderConfig", "location": str(request.url), @@ -275,6 +299,12 @@ async def get_users( """ Get a list of users according to SCIM v2 protocol """ + verbose_proxy_logger.debug( + "SCIM GET USERS request: startIndex=%s count=%s filter=%s", + startIndex, + count, + filter, + ) try: prisma_client = await _get_prisma_client_or_raise_exception() # Parse filter if provided (basic support) @@ -334,6 +364,7 @@ async def get_user( """ Get a single user by ID according to SCIM v2 protocol """ + verbose_proxy_logger.debug("SCIM GET USER request for user_id=%s", user_id) try: user = await _check_user_exists(user_id) @@ -357,7 +388,9 @@ async def create_user( Create a user according to SCIM v2 protocol """ try: - verbose_proxy_logger.debug("SCIM CREATE USER request: %s", user) + verbose_proxy_logger.debug( + "SCIM CREATE USER request: %s", user.model_dump() + ) prisma_client = await _get_prisma_client_or_raise_exception() # Extract data from SCIM user @@ -435,7 +468,11 @@ async def update_user( """ Update a user according to SCIM v2 protocol (full replacement) """ - verbose_proxy_logger.debug("SCIM PUT USER request: %s", user) + verbose_proxy_logger.debug( + "SCIM PUT USER request for user_id=%s: %s", + user_id, + user.model_dump(), + ) try: prisma_client = await _get_prisma_client_or_raise_exception() @@ -497,6 +534,9 @@ async def delete_user( """ Delete a user according to SCIM v2 protocol """ + verbose_proxy_logger.debug( + "SCIM DELETE USER request for user_id=%s", user_id + ) try: prisma_client = await _get_prisma_client_or_raise_exception() existing_user = await _check_user_exists(user_id) @@ -691,7 +731,11 @@ async def patch_user( """ Patch a user according to SCIM v2 protocol """ - verbose_proxy_logger.debug("SCIM PATCH USER request: %s", patch_ops) + verbose_proxy_logger.debug( + "SCIM PATCH USER request for user_id=%s: %s", + user_id, + patch_ops.model_dump(), + ) try: prisma_client = await _get_prisma_client_or_raise_exception() @@ -744,6 +788,12 @@ async def get_groups( """ Get a list of groups according to SCIM v2 protocol """ + verbose_proxy_logger.debug( + "SCIM GET GROUPS request: startIndex=%s count=%s filter=%s", + startIndex, + count, + filter, + ) try: prisma_client = await _get_prisma_client_or_raise_exception() # Parse filter if provided (basic support) @@ -814,6 +864,9 @@ async def get_group( """ Get a single group by ID according to SCIM v2 protocol """ + verbose_proxy_logger.debug( + "SCIM GET GROUP request for group_id=%s", group_id + ) try: team = await _check_team_exists(group_id) @@ -839,6 +892,10 @@ async def create_group( """ Create a group according to SCIM v2 protocol """ + verbose_proxy_logger.debug( + "SCIM CREATE GROUP request: %s", + group.model_dump(), + ) try: prisma_client = await _get_prisma_client_or_raise_exception() @@ -892,78 +949,51 @@ async def update_group( """ Update a group according to SCIM v2 protocol """ + verbose_proxy_logger.debug( + "SCIM PUT GROUP request for group_id=%s: %s", + group_id, + group.model_dump(), + ) try: prisma_client = await _get_prisma_client_or_raise_exception() existing_team = await _check_team_exists(group_id) # Extract valid member IDs member_ids = await _extract_group_member_ids(group) + verbose_proxy_logger.debug(f"SCIM PUT GROUP member_ids: {member_ids}") - # Update team in database + # Prepare update data existing_metadata = existing_team.metadata if existing_team.metadata else {} updated_metadata = {**existing_metadata, "scim_data": group.model_dump()} + update_data = { + "team_alias": group.displayName, + "metadata": safe_dumps(updated_metadata), + } + + # Update team in database updated_team = await prisma_client.db.litellm_teamtable.update( where={"team_id": group_id}, - data={ - "team_alias": group.displayName, - "members": member_ids, - "metadata": safe_dumps(updated_metadata), - }, + data=update_data, ) - # Handle user-team relationships - current_members = existing_team.members or [] - - # Add new members to team - for member_id in member_ids: - if member_id not in current_members: - user = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": member_id} - ) - if user: - current_user_teams = user.teams or [] - if group_id not in current_user_teams: - await prisma_client.db.litellm_usertable.update( - where={"user_id": member_id}, - data={"teams": {"push": group_id}}, - ) - - # Remove former members from team - for member_id in current_members: - if member_id not in member_ids: - user = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": member_id} - ) - if user: - current_user_teams = user.teams or [] - if group_id in current_user_teams: - new_teams = [t for t in current_user_teams if t != group_id] - await prisma_client.db.litellm_usertable.update( - where={"user_id": member_id}, data={"teams": new_teams} - ) - - # Get updated members for response - members = await _get_team_members_display(member_ids) - - team_created_at = ( - updated_team.created_at.isoformat() if updated_team.created_at else None - ) - team_updated_at = ( - updated_team.updated_at.isoformat() if updated_team.updated_at else None + # Handle user-team relationship changes using the same approach as patch_group + current_members = set(await _get_team_member_user_ids_from_team(existing_team)) + verbose_proxy_logger.debug(f"SCIM PUT GROUP current_members: {current_members}") + final_members = set(member_ids) + verbose_proxy_logger.debug(f"SCIM PUT GROUP final_members: {final_members}") + + await _handle_group_membership_changes( + group_id=group_id, + current_members=current_members, + final_members=final_members, ) - return SCIMGroup( - schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], - id=group_id, - displayName=updated_team.team_alias or group_id, - members=members, - meta={ - "resourceType": "Group", - "created": team_created_at, - "lastModified": team_updated_at, - }, + # Convert to SCIM format and return + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( + updated_team ) + return scim_group except Exception as e: raise handle_exception_on_proxy(e) @@ -980,6 +1010,9 @@ async def delete_group( """ Delete a group according to SCIM v2 protocol """ + verbose_proxy_logger.debug( + "SCIM DELETE GROUP request for group_id=%s", group_id + ) try: prisma_client = await _get_prisma_client_or_raise_exception() existing_team = await _check_team_exists(group_id) @@ -1135,7 +1168,11 @@ async def patch_group( """ Patch a group according to SCIM v2 protocol """ - verbose_proxy_logger.debug("SCIM PATCH GROUP request: %s", patch_ops) + verbose_proxy_logger.debug( + "SCIM PATCH GROUP request for group_id=%s: %s", + group_id, + patch_ops.model_dump(), + ) try: prisma_client = await _get_prisma_client_or_raise_exception() @@ -1147,7 +1184,7 @@ async def patch_group( ) # Track current members for comparison - current_members = set(existing_team.members or []) + current_members = set(await _get_team_member_user_ids_from_team(existing_team)) # Apply updates to the database updated_team = await _apply_group_patch_updates( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py new file mode 100644 index 00000000000..d6ab121096e --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -0,0 +1,545 @@ +""" +OpenAI Passthrough Logging Handler + +Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. +""" + +from datetime import datetime +from typing import List, Optional, Union +from urllib.parse import urlparse + +import httpx + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, +) +from litellm.llms.openai.openai import OpenAIConfig +from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( + BasePassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + PassthroughStandardLoggingPayload, +) +from litellm.types.utils import ImageResponse, LlmProviders, PassthroughCallTypes +from litellm.utils import ModelResponse, TextCompletionResponse + + +class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): + """ + OpenAI-specific passthrough logging handler that provides cost tracking for /chat/completions endpoints. + """ + + @property + def llm_provider_name(self) -> LlmProviders: + return LlmProviders.OPENAI + + def get_provider_config(self, model: str) -> OpenAIConfigType: + """Get OpenAI provider configuration for the given model.""" + return OpenAIConfig() + + @staticmethod + def is_openai_chat_completions_route(url_route: str) -> bool: + """Check if the URL route is an OpenAI chat completions endpoint.""" + if not url_route: + return False + parsed_url = urlparse(url_route) + return bool( + parsed_url.hostname + and ( + "api.openai.com" in parsed_url.hostname + or "openai.azure.com" in parsed_url.hostname + ) + and "/v1/chat/completions" in parsed_url.path + ) + + @staticmethod + def is_openai_image_generation_route(url_route: str) -> bool: + """Check if the URL route is an OpenAI image generation endpoint.""" + if not url_route: + return False + parsed_url = urlparse(url_route) + return bool( + parsed_url.hostname + and ( + "api.openai.com" in parsed_url.hostname + or "openai.azure.com" in parsed_url.hostname + ) + and "/v1/images/generations" in parsed_url.path + ) + + @staticmethod + def is_openai_image_editing_route(url_route: str) -> bool: + """Check if the URL route is an OpenAI image editing endpoint.""" + if not url_route: + return False + parsed_url = urlparse(url_route) + return bool( + parsed_url.hostname + and ( + "api.openai.com" in parsed_url.hostname + or "openai.azure.com" in parsed_url.hostname + ) + and "/v1/images/edits" in parsed_url.path + ) + + def _get_user_from_metadata( + self, + passthrough_logging_payload: PassthroughStandardLoggingPayload, + ) -> Optional[str]: + """Extract user information from passthrough logging payload.""" + request_body = passthrough_logging_payload.get("request_body") + if request_body: + return request_body.get("user") + return None + + @staticmethod + def _calculate_image_generation_cost( + model: str, + response_body: dict, + request_body: dict, + ) -> float: + """Calculate cost for OpenAI image generation.""" + try: + # Extract parameters from request + n = request_body.get("n", 1) + try: + n = int(n) + except Exception: + n = 1 + size = request_body.get("size", "1024x1024") + quality = request_body.get("quality", None) + + # Use LiteLLM's default image cost calculator + from litellm.cost_calculator import default_image_cost_calculator + + cost = default_image_cost_calculator( + model=model, + custom_llm_provider="openai", + quality=quality, + n=n, + size=size, + optional_params=request_body, + ) + + return cost + except Exception as e: + verbose_proxy_logger.warning( + f"Error calculating image generation cost: {str(e)}" + ) + return 0.0 + + @staticmethod + def _calculate_image_editing_cost( + model: str, + response_body: dict, + request_body: dict, + ) -> float: + """Calculate cost for OpenAI image editing.""" + try: + # Extract parameters from request + n = request_body.get("n", 1) + # Image edit typically uses multipart/form-data (because of files), so all fields arrive as strings (e.g., n = "1"). + try: + n = int(n) + except Exception: + n = 1 + size = request_body.get("size", "1024x1024") + + # Use LiteLLM's default image cost calculator + from litellm.cost_calculator import default_image_cost_calculator + + cost = default_image_cost_calculator( + model=model, + custom_llm_provider="openai", + quality=None, # Image editing doesn't have quality parameter + n=n, + size=size, + optional_params=request_body, + ) + + return cost + except Exception as e: + verbose_proxy_logger.warning( + f"Error calculating image editing cost: {str(e)}" + ) + return 0.0 + + @staticmethod + def openai_passthrough_handler( # noqa: PLR0915 + httpx_response: httpx.Response, + response_body: dict, + logging_obj: LiteLLMLoggingObj, + url_route: str, + result: str, + start_time: datetime, + end_time: datetime, + cache_hit: bool, + request_body: dict, + **kwargs, + ) -> PassThroughEndpointLoggingTypedDict: + """ + Handle OpenAI passthrough logging with cost tracking for chat completions, image generation, and image editing. + """ + # Check if this is a supported endpoint for cost tracking + is_chat_completions = ( + OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route) + ) + is_image_generation = ( + OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) + ) + is_image_editing = ( + OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) + ) + + if not (is_chat_completions or is_image_generation or is_image_editing): + # For unsupported endpoints, return None to let the system fall back to generic behavior + return { + "result": None, + "kwargs": kwargs, + } + + # Extract model from request or response + model = request_body.get("model", response_body.get("model", "")) + if not model: + verbose_proxy_logger.warning( + "No model found in request or response for OpenAI passthrough cost tracking" + ) + base_handler = OpenAIPassthroughLoggingHandler() + return base_handler.passthrough_chat_handler( + httpx_response=httpx_response, + response_body=response_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + + try: + response_cost = 0.0 + litellm_model_response: Optional[Union[ModelResponse, TextCompletionResponse, ImageResponse]] = None + handler_instance = OpenAIPassthroughLoggingHandler() + + if is_chat_completions: + # Handle chat completions with existing logic + provider_config = handler_instance.get_provider_config(model=model) + litellm_model_response = provider_config.transform_response( + raw_response=httpx_response, + model_response=litellm.ModelResponse(), + model=model, + messages=request_body.get("messages", []), + logging_obj=logging_obj, + optional_params=request_body.get("optional_params", {}), + api_key="", + request_data=request_body, + encoding=litellm.encoding, + json_mode=request_body.get("response_format", {}).get("type") + == "json_object", + litellm_params={}, + ) + + # Calculate cost using LiteLLM's cost calculator + response_cost = litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider="openai", + ) + elif is_image_generation: + # Handle image generation cost calculation + response_cost = ( + OpenAIPassthroughLoggingHandler._calculate_image_generation_cost( + model=model, + response_body=response_body, + request_body=request_body, + ) + ) + # Mark call type for downstream image-aware logic/metrics + try: + logging_obj.call_type = ( + PassthroughCallTypes.passthrough_image_generation.value + ) + except Exception: + pass + # Create a simple response object for logging + litellm_model_response = ImageResponse( + data=response_body.get("data", []), + model=model, + ) + # Set the calculated cost in _hidden_params to prevent recalculation + if not hasattr(litellm_model_response, "_hidden_params"): + litellm_model_response._hidden_params = {} + litellm_model_response._hidden_params["response_cost"] = response_cost + elif is_image_editing: + # Handle image editing cost calculation + response_cost = ( + OpenAIPassthroughLoggingHandler._calculate_image_editing_cost( + model=model, + response_body=response_body, + request_body=request_body, + ) + ) + # Mark call type for downstream image-aware logic/metrics + try: + logging_obj.call_type = ( + PassthroughCallTypes.passthrough_image_generation.value + ) + except Exception: + pass + # Create a simple response object for logging + litellm_model_response = ImageResponse( + data=response_body.get("data", []), + model=model, + ) + # Set the calculated cost in _hidden_params to prevent recalculation + if not hasattr(litellm_model_response, "_hidden_params"): + litellm_model_response._hidden_params = {} + litellm_model_response._hidden_params["response_cost"] = response_cost + + # Update kwargs with cost information + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = "openai" + + # Extract user information for tracking + passthrough_logging_payload: Optional[ + PassthroughStandardLoggingPayload + ] = kwargs.get("passthrough_logging_payload") + if passthrough_logging_payload: + user = handler_instance._get_user_from_metadata( + passthrough_logging_payload=passthrough_logging_payload, + ) + if user: + kwargs.setdefault("litellm_params", {}) + kwargs["litellm_params"].update( + {"proxy_server_request": {"body": {"user": user}}} + ) + + # Create standard logging object + if litellm_model_response is not None: + get_standard_logging_object_payload( + kwargs=kwargs, + init_response_obj=litellm_model_response, + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + status="success", + ) + + # Update logging object with cost information + logging_obj.model_call_details["model"] = model + logging_obj.model_call_details["custom_llm_provider"] = "openai" + logging_obj.model_call_details["response_cost"] = response_cost + + endpoint_type = ( + "chat_completions" + if is_chat_completions + else "image_generation" + if is_image_generation + else "image_editing" + ) + verbose_proxy_logger.debug( + f"OpenAI passthrough cost tracking - Endpoint: {endpoint_type}, Model: {model}, Cost: ${response_cost:.6f}" + ) + + return { + "result": litellm_model_response, + "kwargs": kwargs, + } + + except Exception as e: + verbose_proxy_logger.error( + f"Error in OpenAI passthrough cost tracking: {str(e)}" + ) + # Fall back to base handler without cost tracking + base_handler = OpenAIPassthroughLoggingHandler() + return base_handler.passthrough_chat_handler( + httpx_response=httpx_response, + response_body=response_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + + def _build_complete_streaming_response( + self, + all_chunks: list, + litellm_logging_obj: LiteLLMLoggingObj, + model: str, + ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: + """ + Builds complete response from raw chunks for OpenAI streaming responses. + + - Converts str chunks to generic chunks + - Converts generic chunks to litellm chunks (OpenAI format) + - Builds complete response from litellm chunks + """ + try: + # OpenAI's response iterator to parse chunks + from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator + + openai_iterator = OpenAIChatCompletionResponseIterator( + streaming_response=None, + sync_stream=False, + ) + + all_openai_chunks = [] + for chunk_str in all_chunks: + try: + # Parse the string chunk using the base iterator's string parser + from litellm.llms.base_llm.base_model_iterator import ( + BaseModelResponseIterator, + ) + + # Convert string chunk to dict + stripped_json_chunk = ( + BaseModelResponseIterator._string_to_dict_parser( + str_line=chunk_str + ) + ) + + if stripped_json_chunk: + # Parse the chunk using OpenAI's chunk parser + transformed_chunk = openai_iterator.chunk_parser( + chunk=stripped_json_chunk + ) + if transformed_chunk is not None: + all_openai_chunks.append(transformed_chunk) + + except (StopIteration, StopAsyncIteration, Exception) as e: + verbose_proxy_logger.debug(f"Error parsing streaming chunk: {e}") + continue + + if not all_openai_chunks: + verbose_proxy_logger.warning( + "No valid chunks found in streaming response" + ) + return None + + # Build complete response from chunks + complete_streaming_response = litellm.stream_chunk_builder( + chunks=all_openai_chunks + ) + + return complete_streaming_response + + except Exception as e: + verbose_proxy_logger.error( + f"Error building complete streaming response: {str(e)}" + ) + return None + + @staticmethod + def _handle_logging_openai_collected_chunks( + litellm_logging_obj: LiteLLMLoggingObj, + passthrough_success_handler_obj: PassThroughEndpointLogging, + url_route: str, + request_body: dict, + endpoint_type: EndpointType, + start_time: datetime, + all_chunks: List[str], + end_time: datetime, + ) -> PassThroughEndpointLoggingTypedDict: + """ + Handle logging for collected OpenAI streaming chunks with cost tracking. + """ + try: + # Extract model from request body + model = request_body.get("model", "gpt-4o") + + # Build complete response from chunks using our streaming handler + handler = OpenAIPassthroughLoggingHandler() + handler_instance = handler + complete_response = handler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + ) + + if complete_response is None: + verbose_proxy_logger.warning( + "Failed to build complete response from OpenAI streaming chunks" + ) + return { + "result": None, + "kwargs": {}, + } + + # Calculate cost using LiteLLM's cost calculator + response_cost = litellm.completion_cost( + completion_response=complete_response, + model=model, + custom_llm_provider="openai", + ) + + # Prepare kwargs for logging + kwargs = { + "response_cost": response_cost, + "model": model, + "custom_llm_provider": "openai", + } + + # Extract user information for tracking + passthrough_logging_payload: Optional[ + PassthroughStandardLoggingPayload + ] = litellm_logging_obj.model_call_details.get( + "passthrough_logging_payload" + ) + if passthrough_logging_payload: + user = handler_instance._get_user_from_metadata( + passthrough_logging_payload=passthrough_logging_payload, + ) + if user: + kwargs.setdefault("litellm_params", {}) + kwargs["litellm_params"].update( + {"proxy_server_request": {"body": {"user": user}}} + ) + + # Create standard logging object + get_standard_logging_object_payload( + kwargs=kwargs, + init_response_obj=complete_response, + start_time=start_time, + end_time=end_time, + logging_obj=litellm_logging_obj, + status="success", + ) + + # Update logging object with cost information + litellm_logging_obj.model_call_details["model"] = model + litellm_logging_obj.model_call_details["custom_llm_provider"] = "openai" + litellm_logging_obj.model_call_details["response_cost"] = response_cost + + verbose_proxy_logger.debug( + f"OpenAI streaming passthrough cost tracking - Model: {model}, Cost: ${response_cost:.6f}" + ) + + return { + "result": complete_response, + "kwargs": kwargs, + } + + except Exception as e: + verbose_proxy_logger.error( + f"Error in OpenAI streaming passthrough cost tracking: {str(e)}" + ) + return { + "result": None, + "kwargs": {}, + } diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index adedcaf781d..fccc65b8bd4 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -314,6 +314,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): return EndpointType.VERTEX_AI elif parsed_url.hostname == "api.anthropic.com": return EndpointType.ANTHROPIC + elif ( + parsed_url.hostname == "api.openai.com" + or parsed_url.hostname == "openai.azure.com" + or (parsed_url.hostname and "openai.com" in parsed_url.hostname) + ): + return EndpointType.OPENAI return EndpointType.GENERIC @staticmethod @@ -415,10 +421,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): for field_name, field_value in form_data.items(): if isinstance(field_value, (StarletteUploadFile, UploadFile)): - files[field_name] = ( - await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( - upload_file=field_value - ) + files[ + field_name + ] = await HttpPassThroughEndpointHelpers._build_request_files_from_upload_file( + upload_file=field_value ) else: form_data_dict[field_name] = field_value @@ -497,9 +503,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "passthrough_logging_payload": passthrough_logging_payload, } - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + logging_obj.model_call_details[ + "passthrough_logging_payload" + ] = passthrough_logging_payload return kwargs @@ -531,10 +537,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): subpath = subpath[1:] return base_target + subpath - + @staticmethod def _update_stream_param_based_on_request_body( - parsed_body: dict, + parsed_body: dict, stream: Optional[bool] = None, ) -> Optional[bool]: """ @@ -699,9 +705,11 @@ async def pass_through_request( # noqa: PLR0915 "headers": headers, }, ) - stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( - parsed_body=_parsed_body, - stream=stream, + stream = ( + HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( + parsed_body=_parsed_body, + stream=stream, + ) ) if stream: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 08b49bac383..2d5b0a686ce 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -14,6 +14,9 @@ from litellm.types.utils import StandardPassThroughResponseObject from .llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) +from .llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, +) from .llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) @@ -78,6 +81,7 @@ class PassThroughStreamingHandler: Supported endpoint types: - Anthropic - Vertex AI + - OpenAI """ all_chunks = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines( raw_bytes @@ -119,6 +123,23 @@ class PassThroughStreamingHandler: vertex_passthrough_logging_handler_result["result"] ) kwargs = vertex_passthrough_logging_handler_result["kwargs"] + elif endpoint_type == EndpointType.OPENAI: + openai_passthrough_logging_handler_result = ( + OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + ) + ) + standard_logging_response_object = ( + openai_passthrough_logging_handler_result["result"] + ) + kwargs = openai_passthrough_logging_handler_result["kwargs"] if standard_logging_response_object is None: standard_logging_response_object = StandardPassThroughResponseObject( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index ce576d5ac71..58fda370d93 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -162,9 +162,32 @@ class PassThroughEndpointLogging: cohere_passthrough_logging_handler_result["result"] ) kwargs = cohere_passthrough_logging_handler_result["kwargs"] - return_dict["standard_logging_response_object"] = ( - standard_logging_response_object - ) + elif self.is_openai_route(url_route) and self._is_supported_openai_endpoint(url_route): + from .llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, + ) + + openai_passthrough_logging_handler_result = ( + OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body or {}, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + ) + standard_logging_response_object = ( + openai_passthrough_logging_handler_result["result"] + ) + kwargs = openai_passthrough_logging_handler_result["kwargs"] + return_dict[ + "standard_logging_response_object" + ] = standard_logging_response_object return_dict["kwargs"] = kwargs return return_dict @@ -185,9 +208,9 @@ class PassThroughEndpointLogging: standard_logging_response_object: Optional[ PassThroughEndpointLoggingResultValues ] = None - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + logging_obj.model_call_details[ + "passthrough_logging_payload" + ] = passthrough_logging_payload if self.is_assemblyai_route(url_route): if ( AssemblyAIPassthroughLoggingHandler._should_log_request( @@ -286,6 +309,28 @@ class PassThroughEndpointLogging: return True return False + def is_openai_route(self, url_route: str): + """Check if the URL route is an OpenAI API route.""" + if not url_route: + return False + parsed_url = urlparse(url_route) + return parsed_url.hostname and ( + "api.openai.com" in parsed_url.hostname + or "openai.azure.com" in parsed_url.hostname + ) + + def _is_supported_openai_endpoint(self, url_route: str) -> bool: + """Check if the OpenAI endpoint is supported by the passthrough logging handler.""" + from .llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, + ) + + return ( + OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route) or + OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) or + OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) + ) + def _set_cost_per_request( self, logging_obj: LiteLLMLoggingObj, @@ -305,8 +350,8 @@ class PassThroughEndpointLogging: kwargs["response_cost"] = passthrough_logging_payload.get( "cost_per_request" ) - logging_obj.model_call_details["response_cost"] = ( - passthrough_logging_payload.get("cost_per_request") - ) + logging_obj.model_call_details[ + "response_cost" + ] = passthrough_logging_payload.get("cost_per_request") return kwargs diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index acff522196e..7ee09105254 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -3,4 +3,5 @@ model_list: litellm_params: model: openai/* api_base: https://exampleopenaiendpoint-production-0ee2.up.railway.app/ - mock_response: "hi" +litellm_settings: + callbacks: ["cloudzero"] \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 547aaf50788..9f1566b2e00 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -248,7 +248,9 @@ from litellm.proxy.management_endpoints.customer_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_update, +) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -295,7 +297,9 @@ from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMi from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config +from litellm.proxy.openai_files_endpoints.files_endpoints import ( + set_files_config, +) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -3807,13 +3811,13 @@ class ProxyStartupEvent: ######################################################## # CloudZero Background Job ######################################################## + from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger from litellm.proxy.spend_tracking.cloudzero_endpoints import ( - init_cloudzero_background_job, - is_cloudzero_setup_in_db, + is_cloudzero_setup, ) - if await is_cloudzero_setup_in_db(): - await init_cloudzero_background_job() + if await is_cloudzero_setup(): + await CloudZeroLogger.init_cloudzero_background_job(scheduler=scheduler) ######################################################## # Prometheus Background Job diff --git a/litellm/proxy/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index 67de202aa7a..502537cb70f 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -82,14 +82,8 @@ async def _get_cloudzero_settings(): cloudzero_config = await prisma_client.db.litellm_config.find_first( where={"param_name": "cloudzero_settings"} ) - - if not cloudzero_config or not cloudzero_config.param_value: - raise HTTPException( - status_code=400, - detail={ - "error": "CloudZero settings not configured. Please run /cloudzero/init first." - }, - ) + if cloudzero_config is None: + return {} settings = dict(cloudzero_config.param_value) @@ -257,62 +251,6 @@ async def update_cloudzero_settings( _cloudzero_background_job_initialized = False -async def init_cloudzero_background_job(): - """ - Initialize CloudZero background job if not already initialized. - This should be called from the proxy server startup. - """ - global _cloudzero_background_job_initialized - - if _cloudzero_background_job_initialized: - verbose_proxy_logger.debug( - "CloudZero background job already initialized, skipping" - ) - return - - try: - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - verbose_proxy_logger.warning( - "Prisma client not available, skipping CloudZero background job initialization" - ) - return - - # Get CloudZero settings from database - cloudzero_config = await prisma_client.db.litellm_config.find_first( - where={"param_name": "cloudzero_settings"} - ) - - if not cloudzero_config or not cloudzero_config.param_value: - verbose_proxy_logger.debug( - "CloudZero settings not configured, skipping background job initialization" - ) - return - - settings = dict(cloudzero_config.param_value) - - # Initialize CloudZero logger with credentials - from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger - - logger = CloudZeroLogger( - api_key=settings["api_key"], - connection_id=settings["connection_id"], - timezone=settings["timezone"], - ) - - # Initialize the background job - await logger.init_background_job() - - _cloudzero_background_job_initialized = True - verbose_proxy_logger.info("CloudZero background job initialized successfully") - - except Exception as e: - verbose_proxy_logger.error( - f"Error initializing CloudZero background job: {str(e)}" - ) - - async def is_cloudzero_setup_in_db() -> bool: """ Check if CloudZero is setup in the database. @@ -343,6 +281,47 @@ async def is_cloudzero_setup_in_db() -> bool: return False +def is_cloudzero_setup_in_config() -> bool: + """ + Check if CloudZero is setup in config.yaml or environment variables. + + CloudZero is considered setup in config if: + - "cloudzero" is in the callbacks list in config.yaml, OR + Returns: + bool: True if CloudZero is configured, False otherwise + """ + import litellm + return "cloudzero" in litellm.callbacks + + +async def is_cloudzero_setup() -> bool: + """ + Check if CloudZero is setup in either config.yaml/env vars OR database. + + CloudZero is considered setup if: + - CloudZero is configured in config.yaml callbacks, OR + - CloudZero environment variables are set, OR + - CloudZero settings exist in the database + + Returns: + bool: True if CloudZero is configured anywhere, False otherwise + """ + try: + # Check config.yaml/environment variables first + if is_cloudzero_setup_in_config(): + return True + + # Check database as fallback + if await is_cloudzero_setup_in_db(): + return True + + return False + + except Exception as e: + verbose_proxy_logger.error(f"Error checking CloudZero setup: {str(e)}") + return False + + @router.post( "/cloudzero/init", tags=["CloudZero"], @@ -383,9 +362,6 @@ async def init_cloudzero_settings( verbose_proxy_logger.info("CloudZero settings initialized successfully") - # Initialize background job after settings are saved - await init_cloudzero_background_job() - return CloudZeroInitResponse( message="CloudZero settings initialized successfully", status="success" ) @@ -412,15 +388,18 @@ async def cloudzero_dry_run_export( Perform a dry run export using the CloudZero logger. This endpoint uses the CloudZero logger to perform a dry run export, - which displays the data that would be exported without actually sending it to CloudZero. + which returns the data that would be exported without actually sending it to CloudZero. Parameters: - limit: Optional limit on number of records to process (default: 10000) + Returns: + - usage_data: Sample of the raw usage data (first 50 records) + - cbf_data: CloudZero CBF formatted data ready for export + - summary: Statistics including total cost, tokens, and record counts + Only admin users can perform CloudZero exports. """ - from datetime import datetime - # Validation if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( @@ -434,15 +413,17 @@ async def cloudzero_dry_run_export( # Initialize logger with credentials directly logger = CloudZeroLogger() - await logger.dry_run_export_usage_data( - target_hour=datetime.utcnow(), limit=request.limit + dry_run_result = await logger.dry_run_export_usage_data( + limit=request.limit ) verbose_proxy_logger.info("CloudZero dry run export completed successfully") return CloudZeroExportResponse( - message="CloudZero dry run export completed successfully. Check logs for output.", + message="CloudZero dry run export completed successfully.", status="success", + dry_run_data=dry_run_result, + summary=dry_run_result.get("summary") if dry_run_result else None, ) except Exception as e: @@ -477,7 +458,6 @@ async def cloudzero_export( Only admin users can perform CloudZero exports. """ - from datetime import datetime if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( @@ -494,20 +474,24 @@ async def cloudzero_export( # Initialize logger with credentials directly logger = CloudZeroLogger( - api_key=settings["api_key"], - connection_id=settings["connection_id"], - timezone=settings["timezone"], + api_key=settings.get("api_key"), + connection_id=settings.get("connection_id"), + timezone=settings.get("timezone"), ) await logger.export_usage_data( - target_hour=datetime.utcnow(), limit=request.limit, operation=request.operation, + start_time_utc=request.start_time_utc, + end_time_utc=request.end_time_utc, ) verbose_proxy_logger.info("CloudZero export completed successfully") return CloudZeroExportResponse( - message="CloudZero export completed successfully", status="success" + message="CloudZero export completed successfully", + status="success", + dry_run_data=None, + summary=None ) except Exception as e: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 9584baf7368..47ecbcf02c0 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1,9 +1,10 @@ import asyncio import contextvars from functools import partial -from typing import Any, Coroutine, Dict, Iterable, List, Literal, Optional, Union +from typing import Any, Coroutine, Dict, Iterable, List, Literal, Optional, Type, Union import httpx +from pydantic import BaseModel import litellm from litellm.constants import request_timeout @@ -135,9 +136,10 @@ async def aresponses_api_with_mcp( ) # Parse MCP tools and separate from other tools - mcp_tools_with_litellm_proxy, other_tools = ( - LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) - ) + ( + mcp_tools_with_litellm_proxy, + other_tools, + ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) # Get available tools from MCP manager if we have MCP tools openai_tools = [] @@ -254,6 +256,7 @@ async def aresponses( stream: Optional[bool] = None, temperature: Optional[float] = None, text: Optional["ResponseText"] = None, + text_format: Optional[Union[Type["BaseModel"], dict]] = None, tool_choice: Optional[ToolChoice] = None, tools: Optional[Iterable[ToolParam]] = None, top_p: Optional[float] = None, @@ -279,6 +282,14 @@ async def aresponses( loop = asyncio.get_event_loop() kwargs["aresponses"] = True + # Convert text_format to text parameter if provided + text = ResponsesAPIRequestUtils.convert_text_format_to_text_param( + text_format=text_format, text=text + ) + if text is not None: + # Update local_vars to include the converted text parameter + local_vars["text"] = text + # get custom llm provider so we can use this for mapping exceptions if custom_llm_provider is None: _, custom_llm_provider, _, _ = litellm.get_llm_provider( @@ -367,6 +378,7 @@ def responses( stream: Optional[bool] = None, temperature: Optional[float] = None, text: Optional["ResponseText"] = None, + text_format: Optional[Union[Type["BaseModel"], dict]] = None, tool_choice: Optional[ToolChoice] = None, tools: Optional[Iterable[ToolParam]] = None, top_p: Optional[float] = None, @@ -399,6 +411,14 @@ def responses( litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("aresponses", False) is True + # Convert text_format to text parameter if provided + text = ResponsesAPIRequestUtils.convert_text_format_to_text_param( + text_format=text_format, text=text + ) + if text is not None: + # Update local_vars to include the converted text parameter + local_vars["text"] = text + # get llm provider logic litellm_params = GenericLiteLLMParams(**kwargs) @@ -432,11 +452,11 @@ def responses( ) # get provider config - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) local_vars.update(kwargs) @@ -628,11 +648,11 @@ def delete_responses( raise ValueError("custom_llm_provider is required but passed as None") # get provider config - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if responses_api_provider_config is None: @@ -807,11 +827,11 @@ def get_responses( raise ValueError("custom_llm_provider is required but passed as None") # get provider config - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if responses_api_provider_config is None: @@ -963,11 +983,11 @@ def list_input_items( if custom_llm_provider is None: raise ValueError("custom_llm_provider is required but passed as None") - responses_api_provider_config: Optional[BaseResponsesAPIConfig] = ( - ProviderConfigManager.get_provider_responses_api_config( - model=None, - provider=litellm.LlmProviders(custom_llm_provider), - ) + responses_api_provider_config: Optional[ + BaseResponsesAPIConfig + ] = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=litellm.LlmProviders(custom_llm_provider), ) if responses_api_provider_config is None: @@ -1009,4 +1029,4 @@ def list_input_items( original_exception=e, completion_kwargs=local_vars, extra_kwargs=kwargs, - ) + ) \ No newline at end of file diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index ac59d28a50d..336f2b5a947 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,5 +1,17 @@ import base64 -from typing import Any, Dict, List, Optional, Union, cast, get_type_hints, overload +from typing import ( + Any, + Dict, + List, + Optional, + Type, + Union, + cast, + get_type_hints, + overload, +) + +from pydantic import BaseModel import litellm from litellm._logging import verbose_logger @@ -8,6 +20,7 @@ from litellm.types.llms.openai import ( ResponseAPIUsage, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse, + ResponseText, ) from litellm.types.responses.main import DecodedResponseId from litellm.types.utils import SpecialEnums, Usage @@ -24,7 +37,6 @@ class ResponsesAPIRequestUtils: custom_llm_provider: Optional[str], model: str, ): - if supported_params is None: return unsupported_params = {} @@ -302,6 +314,40 @@ class ResponsesAPIRequestUtils: ) return decoded_response_id.get("response_id", previous_response_id) + @staticmethod + def convert_text_format_to_text_param( + text_format: Optional[Union[Type["BaseModel"], dict]], + text: Optional["ResponseText"] = None, + ) -> Optional["ResponseText"]: + """ + Convert text_format parameter to text parameter for the responses API. + + Args: + text_format: Pydantic model class or dict to convert to response format + text: Existing text parameter (if provided, text_format is ignored) + + Returns: + ResponseText object with the converted format, or None if conversion fails + """ + if text_format is not None and text is None: + from litellm.llms.base_llm.base_utils import type_to_response_format_param + + # Convert Pydantic model to response format + response_format = type_to_response_format_param(text_format) + if response_format is not None: + # Create ResponseText object with the format + # The responses API expects the format to have name at the top level + text = { + "format": { + "type": response_format["type"], + "name": response_format["json_schema"]["name"], + "schema": response_format["json_schema"]["schema"], + "strict": response_format["json_schema"]["strict"], + } + } + return text + return text + class ResponseAPILoggingUtils: @staticmethod @@ -333,4 +379,4 @@ class ResponseAPILoggingUtils: prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, - ) + ) \ No newline at end of file diff --git a/litellm/router.py b/litellm/router.py index 190d19598c3..6255c2fdf92 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4562,6 +4562,20 @@ class Router: parent_otel_span=parent_otel_span, ttl=RoutingArgs.ttl.value, ) + + def _get_metadata_variable_name_from_kwargs(self, kwargs: dict) -> Literal["metadata", "litellm_metadata"]: + """ + Helper to return what the "metadata" field should be called in the request data + + - New endpoints return `litellm_metadata` + - Old endpoints return `metadata` + + Context: + - LiteLLM used `metadata` as an internal field for storing metadata + - OpenAI then started using this field for their metadata + - LiteLLM is now moving to using `litellm_metadata` for our metadata + """ + return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" def log_retry(self, kwargs: dict, e: Exception) -> dict: """ @@ -5451,7 +5465,7 @@ class Router: ## SET MODEL TO 'model=' - if base_model is None + not azure if custom_llm_provider == "azure" and base_model is None: verbose_router_logger.error( - "Could not identify azure model. Set azure 'base_model' for accurate max tokens, cost tracking, etc.- https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models" + f"Could not identify azure model '{_model}'. Set azure 'base_model' for accurate max tokens, cost tracking, etc.- https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models" ) elif custom_llm_provider != "azure": model = _model @@ -5658,6 +5672,11 @@ class Router: ) if supported_openai_params is None: supported_openai_params = [] + + # Get mode from database model_info if available, otherwise default to "chat" + db_model_info = model.get("model_info", {}) + mode = db_model_info.get("mode", "chat") + model_info = ModelMapInfo( key=model_group, max_tokens=None, @@ -5666,7 +5685,7 @@ class Router: input_cost_per_token=0, output_cost_per_token=0, litellm_provider=llm_provider, - mode="chat", + mode=mode, supported_openai_params=supported_openai_params, supports_system_messages=None, ) @@ -6783,6 +6802,7 @@ class Router: model=model, request_kwargs=request_kwargs, healthy_deployments=healthy_deployments, + metadata_variable_name=self._get_metadata_variable_name_from_kwargs(request_kwargs), ) if len(healthy_deployments) == 0: diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index 34261d83dcf..8094b5d86ac 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -6,7 +6,7 @@ Use this to route requests between Teams - If no default_deployments are set, return all deployments """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union from litellm._logging import verbose_logger from litellm.types.router import RouterErrors @@ -41,6 +41,7 @@ async def get_deployments_for_tag( model: str, # used to raise the correct error healthy_deployments: Union[List[Any], Dict[Any, Any]], request_kwargs: Optional[Dict[Any, Any]] = None, + metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata", ): """ Returns a list of deployments that match the requested model and tags in the request. @@ -63,9 +64,9 @@ async def get_deployments_for_tag( ) return healthy_deployments - verbose_logger.debug("request metadata: %s", request_kwargs.get("metadata")) - if "metadata" in request_kwargs: - metadata = request_kwargs["metadata"] + verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name)) + if metadata_variable_name in request_kwargs: + metadata = request_kwargs[metadata_variable_name] request_tags = metadata.get("tags") new_healthy_deployments = [] @@ -120,7 +121,8 @@ async def get_deployments_for_tag( def _get_tags_from_request_kwargs( - request_kwargs: Optional[Dict[Any, Any]] = None + request_kwargs: Optional[Dict[Any, Any]] = None, + metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata", ) -> List[str]: """ Helper to get tags from request kwargs @@ -133,11 +135,11 @@ def _get_tags_from_request_kwargs( """ if request_kwargs is None: return [] - if "metadata" in request_kwargs: - metadata = request_kwargs["metadata"] + if metadata_variable_name in request_kwargs: + metadata = request_kwargs[metadata_variable_name] return metadata.get("tags", []) elif "litellm_params" in request_kwargs: litellm_params = request_kwargs["litellm_params"] - _metadata = litellm_params.get("metadata", {}) + _metadata = litellm_params.get(metadata_variable_name, {}) return _metadata.get("tags", []) return [] diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py index 82fb4fe3887..75c55bcc93c 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -46,6 +46,7 @@ class LLMMetrics(TypedDict, total=False): class LLMObsPayload(TypedDict, total=False): parent_id: str trace_id: str + apm_id: str span_id: str name: str meta: Meta diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index e0ee950d260..955c1d888f2 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -154,6 +154,7 @@ class UserAPIKeyLabelNames(Enum): DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", "litellm_request_total_latency_metric", "litellm_overhead_latency_metric", "litellm_remaining_requests_metric", @@ -162,6 +163,7 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_proxy_failed_requests_metric", "litellm_deployment_latency_per_output_token", "litellm_requests_metric", + "litellm_spend_metric", "litellm_total_tokens_metric", "litellm_input_tokens_metric", "litellm_output_tokens_metric", @@ -173,9 +175,11 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_remaining_api_key_budget_metric", "litellm_api_key_max_budget_metric", "litellm_api_key_budget_remaining_hours_metric", + "litellm_deployment_state", "litellm_deployment_failure_responses", "litellm_deployment_total_requests", "litellm_deployment_success_responses", + "litellm_deployment_cooled_down", "litellm_pod_lock_manager_size", "litellm_in_memory_daily_spend_update_queue_size", "litellm_redis_daily_spend_update_queue_size", @@ -196,6 +200,14 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, ] + litellm_llm_api_time_to_first_token_metric = [ + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + ] + litellm_request_total_latency_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, @@ -282,6 +294,16 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER_EMAIL.value, ] + litellm_spend_metric = [ + UserAPIKeyLabelNames.END_USER.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.USER.value, + ] + litellm_input_tokens_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, @@ -315,6 +337,20 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.REQUESTED_MODEL.value, ] + litellm_deployment_state = [ + UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, + UserAPIKeyLabelNames.API_BASE.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + ] + + litellm_deployment_cooled_down = [ + UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.MODEL_ID.value, + UserAPIKeyLabelNames.API_BASE.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + ] + litellm_deployment_successful_fallbacks = [ UserAPIKeyLabelNames.REQUESTED_MODEL.value, UserAPIKeyLabelNames.FALLBACK_MODEL.value, diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 068c9db035b..5853f9de2aa 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -573,3 +573,84 @@ class AmazonDeepSeekR1StreamingResponse(TypedDict): generation_token_count: int stop_reason: Optional[str] prompt_token_count: int + + +################ Bedrock Batch Types ################# + + +class BedrockS3InputDataConfig(TypedDict): + """S3 input data configuration for Bedrock batch jobs.""" + s3Uri: str + + +class BedrockInputDataConfig(TypedDict): + """Input data configuration for Bedrock batch jobs.""" + s3InputDataConfig: BedrockS3InputDataConfig + + +class BedrockS3OutputDataConfig(TypedDict): + """S3 output data configuration for Bedrock batch jobs.""" + s3Uri: str + + +class BedrockOutputDataConfig(TypedDict): + """Output data configuration for Bedrock batch jobs.""" + s3OutputDataConfig: BedrockS3OutputDataConfig + + +class BedrockCreateBatchRequest(TypedDict, total=False): + """ + Request structure for creating a Bedrock batch inference job. + + Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_CreateModelInvocationJob.html + """ + jobName: str + roleArn: str + modelId: str + inputDataConfig: BedrockInputDataConfig + outputDataConfig: BedrockOutputDataConfig + timeoutDurationInHours: Optional[int] + clientRequestToken: Optional[str] + tags: Optional[List[dict]] + + +BedrockBatchJobStatus = Literal[ + "Submitted", + "InProgress", + "Completed", + "Failed", + "Stopping", + "Stopped" +] + + +class BedrockCreateBatchResponse(TypedDict): + """ + Response structure from creating a Bedrock batch inference job. + + Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_CreateModelInvocationJob.html + """ + jobArn: str + jobName: str + status: BedrockBatchJobStatus + + +class BedrockGetBatchResponse(TypedDict, total=False): + """ + Response structure from getting a Bedrock batch inference job. + + Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetModelInvocationJob.html + """ + jobArn: str + jobName: str + modelId: str + roleArn: str + status: BedrockBatchJobStatus + message: Optional[str] + submitTime: Optional[str] + lastModifiedTime: Optional[str] + endTime: Optional[str] + inputDataConfig: BedrockInputDataConfig + outputDataConfig: BedrockOutputDataConfig + timeoutDurationInHours: Optional[int] + clientRequestToken: Optional[str] diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index bb59b692ef7..112427c6b56 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -1,5 +1,5 @@ import json -from typing import Any, List, Literal, Optional, TypedDict, Union +from typing import Any, Dict, List, Literal, Optional, TypedDict, Union from pydantic import BaseModel from typing_extensions import ( @@ -24,9 +24,10 @@ class GenericStreamingChunk(TypedDict, total=False): usage: Optional[BaseModel] -class DatabricksTextContent(TypedDict): +class DatabricksTextContent(TypedDict, total=False): type: Literal["text"] text: Required[str] + citations: Optional[List[Dict[str, Any]]] class DatabricksReasoningSummary(TypedDict): @@ -35,9 +36,10 @@ class DatabricksReasoningSummary(TypedDict): signature: str -class DatabricksReasoningContent(TypedDict): +class DatabricksReasoningContent(TypedDict, total=False): type: Literal["reasoning"] - summary: List[DatabricksReasoningSummary] + summary: Required[List[DatabricksReasoningSummary]] + citations: Optional[List[Dict[str, Any]]] AllDatabricksContentListValues = Union[ diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 753cf884646..c6d126be681 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -794,6 +794,7 @@ class ChatCompletionRequest(TypedDict, total=False): response_format: dict seed: int service_tier: str + safety_identifier: str stop: Union[str, List[str]] stream_options: dict temperature: float @@ -919,7 +920,6 @@ OpenAIImageVariationOptionalParams = Literal["n", "size", "response_format", "us OpenAIImageGenerationOptionalParams = Literal[ "background", - "input_fidelity", "moderation", "n", "output_compression", @@ -1035,29 +1035,29 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): id: str created_at: int - error: Optional[dict] - incomplete_details: Optional[IncompleteDetails] - instructions: Optional[str] - metadata: Optional[Dict] - model: Optional[str] - object: Optional[str] + error: Optional[dict] = None + incomplete_details: Optional[IncompleteDetails] = None + instructions: Optional[str] = None + metadata: Optional[Dict] = None + model: Optional[str] = None + object: Optional[str] = None output: Union[ List[Union[ResponseOutputItem, Dict]], List[Union[GenericResponseOutputItem, OutputFunctionToolCall]], ] parallel_tool_calls: bool - temperature: Optional[float] + temperature: Optional[float] = None tool_choice: ToolChoice tools: Union[List[Tool], List[ResponseFunctionToolCall], List[Dict[str, Any]]] top_p: Optional[float] - max_output_tokens: Optional[int] - previous_response_id: Optional[str] - reasoning: Optional[Reasoning] - status: Optional[str] - text: Optional[Union["ResponseText", Dict[str, Any]]] - truncation: Optional[Literal["auto", "disabled"]] - usage: Optional[ResponseAPIUsage] - user: Optional[str] + max_output_tokens: Optional[int] = None + previous_response_id: Optional[str] = None + reasoning: Optional[Reasoning] = None + status: Optional[str] = None + text: Optional[Union["ResponseText", Dict[str, Any]]] = None + truncation: Optional[Literal["auto", "disabled"]] = None + usage: Optional[ResponseAPIUsage] = None + user: Optional[str] = None store: Optional[bool] = None # Define private attributes using PrivateAttr _hidden_params: dict = PrivateAttr(default_factory=dict) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 2931770cd6e..625a76b6789 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -41,6 +41,7 @@ class PartType(TypedDict, total=False): function_call: FunctionCall function_response: FunctionResponse thought: bool + thoughtSignature: str class HttpxFunctionCall(TypedDict): @@ -72,6 +73,7 @@ class HttpxPartType(TypedDict, total=False): executableCode: HttpxExecutableCode codeExecutionResult: HttpxCodeExecutionResult thought: bool + thoughtSignature: str class HttpxContentType(TypedDict, total=False): @@ -111,6 +113,7 @@ class Schema(TypedDict, total=False): pattern: str example: Any anyOf: List["Schema"] + additionalProperties: Any class FunctionDeclaration(TypedDict, total=False): @@ -245,10 +248,11 @@ class UsageMetadata(TypedDict, total=False): class TokenCountDetailsResponse(TypedDict): """ Response structure for token count details with modality breakdown. - + Example: {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} """ + totalTokens: int promptTokensDetails: List[PromptTokensDetails] diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 8a91dd14ebd..39facdd8e6f 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -5,6 +5,7 @@ from typing import Optional, TypedDict class EndpointType(str, Enum): VERTEX_AI = "vertex-ai" ANTHROPIC = "anthropic" + OPENAI = "openai" GENERIC = "generic" diff --git a/litellm/types/proxy/cloudzero_endpoints.py b/litellm/types/proxy/cloudzero_endpoints.py index f7f63233d4d..1d909bf7f8c 100644 --- a/litellm/types/proxy/cloudzero_endpoints.py +++ b/litellm/types/proxy/cloudzero_endpoints.py @@ -2,7 +2,8 @@ CloudZero endpoint types for LiteLLM Proxy """ -from typing import Optional +from datetime import datetime +from typing import Any, Dict, List, Optional from pydantic import BaseModel, Field @@ -27,6 +28,8 @@ class CloudZeroExportRequest(BaseModel): limit: Optional[int] = Field(None, description="Optional limit on number of records to export") operation: str = Field(default="replace_hourly", description="CloudZero operation type (replace_hourly or sum)") + start_time_utc: Optional[datetime] = Field(None, description="Start time for data export in UTC") + end_time_utc: Optional[datetime] = Field(None, description="End time for data export in UTC") class CloudZeroExportResponse(BaseModel): @@ -35,6 +38,8 @@ class CloudZeroExportResponse(BaseModel): message: str status: str records_exported: Optional[int] = None + dry_run_data: Optional[Dict[str, Any]] = Field(None, description="Dry run data including usage data and CBF transformed data") + summary: Optional[Dict[str, Any]] = Field(None, description="Summary statistics for dry run") class CloudZeroSettingsView(BaseModel): diff --git a/litellm/types/router.py b/litellm/types/router.py index 864fdbf79b8..8bed2bdccda 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -216,6 +216,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): auto_router_default_model: Optional[str] = None auto_router_embedding_model: Optional[str] = None + # Batch/File API Params + s3_bucket_name: Optional[str] = None + gcs_bucket_name: Optional[str] = None + def __init__( self, custom_llm_provider: Optional[str] = None, @@ -265,6 +269,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): auto_router_config: Optional[str] = None, auto_router_default_model: Optional[str] = None, auto_router_embedding_model: Optional[str] = None, + # Batch/File API Params + s3_bucket_name: Optional[str] = None, + gcs_bucket_name: Optional[str] = None, **params, ): args = locals() diff --git a/litellm/utils.py b/litellm/utils.py index 69f4603fea0..ccc4b475621 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -234,6 +234,7 @@ from litellm.llms.base_llm.base_utils import ( BaseLLMModelInfo, type_to_response_format_param, ) +from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.completion.transformation import BaseTextCompletionConfig from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig @@ -541,9 +542,9 @@ def function_setup( # noqa: PLR0915 function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = kwargs.pop("callbacks", None) + dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = ( + kwargs.pop("callbacks", None) + ) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) if len(all_callbacks) > 0: @@ -837,15 +838,13 @@ async def _client_async_logging_helper( # Async Logging Worker ################################################ from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine = logging_obj.async_success_handler( - result=result, - start_time=start_time, - end_time=end_time + async_coroutine=logging_obj.async_success_handler( + result=result, start_time=start_time, end_time=end_time ) ) - ################################################ # Sync Logging Worker ################################################ @@ -1299,9 +1298,9 @@ def client(original_function): # noqa: PLR0915 exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs[ - "retry_policy" - ] = reset_retry_policy() # prevent infinite loops + kwargs["retry_policy"] = ( + reset_retry_policy() + ) # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -2397,7 +2396,7 @@ def _should_drop_param(k, additional_drop_params) -> bool: def _get_non_default_params( - passed_params: dict, default_params: dict, additional_drop_params: Optional[bool] + passed_params: dict, default_params: dict, additional_drop_params: Optional[list] ) -> dict: non_default_params = {} for k, v in passed_params.items(): @@ -2509,9 +2508,8 @@ def get_optional_params_image_gen( size: Optional[str] = None, style: Optional[str] = None, user: Optional[str] = None, - input_fidelity: Optional[str] = None, custom_llm_provider: Optional[str] = None, - additional_drop_params: Optional[bool] = None, + additional_drop_params: Optional[list] = None, provider_config: Optional[BaseImageGenerationConfig] = None, drop_params: Optional[bool] = None, **kwargs, @@ -2546,7 +2544,6 @@ def get_optional_params_image_gen( "size": None, "style": None, "user": None, - "input_fidelity": None, } non_default_params = _get_non_default_params( @@ -2630,9 +2627,20 @@ def get_optional_params_image_gen( ) # Default to square if size not recognized optional_params["aspectRatio"] = aspect_ratio - for k in passed_params.keys(): - if k not in default_params.keys(): - optional_params[k] = passed_params[k] + openai_params: list[str] = list(default_params.keys()) + if provider_config is not None: + supported_params = provider_config.get_supported_openai_params( + model=model or "" + ) + openai_params = list(supported_params) + + optional_params = add_provider_specific_params_to_optional_params( + optional_params=optional_params, + passed_params=passed_params, + custom_llm_provider=custom_llm_provider or "", + openai_params=openai_params, + additional_drop_params=additional_drop_params, + ) return optional_params @@ -3109,10 +3117,10 @@ def pre_process_non_default_params( if "response_format" in non_default_params: if provider_config is not None: - non_default_params[ - "response_format" - ] = provider_config.get_json_schema_from_pydantic_object( - response_format=non_default_params["response_format"] + non_default_params["response_format"] = ( + provider_config.get_json_schema_from_pydantic_object( + response_format=non_default_params["response_format"] + ) ) else: non_default_params["response_format"] = type_to_response_format_param( @@ -3240,16 +3248,16 @@ def pre_process_optional_params( True # so that main.py adds the function call to the prompt ) if "tools" in non_default_params: - optional_params[ - "functions_unsupported_model" - ] = non_default_params.pop("tools") + optional_params["functions_unsupported_model"] = ( + non_default_params.pop("tools") + ) non_default_params.pop( "tool_choice", None ) # causes ollama requests to hang elif "functions" in non_default_params: - optional_params[ - "functions_unsupported_model" - ] = non_default_params.pop("functions") + optional_params["functions_unsupported_model"] = ( + non_default_params.pop("functions") + ) elif ( litellm.add_function_to_prompt ): # if user opts to add it to prompt instead @@ -3304,6 +3312,7 @@ def get_optional_params( # noqa: PLR0915 messages: Optional[List[AllMessageValues]] = None, thinking: Optional[AnthropicThinkingParam] = None, web_search_options: Optional[OpenAIWebSearchOptions] = None, + safety_identifier: Optional[str] = None, **kwargs, ): passed_params = locals().copy() @@ -3602,6 +3611,17 @@ def get_optional_params( # noqa: PLR0915 else False ), ) + elif provider_config is not None: + optional_params = provider_config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=( + drop_params + if drop_params is not None and isinstance(drop_params, bool) + else False + ), + ) else: # use generic openai-like param mapping optional_params = litellm.VertexAILlama3Config().map_openai_params( non_default_params=non_default_params, @@ -4065,7 +4085,7 @@ def add_provider_specific_params_to_optional_params( ): extra_body = passed_params.pop("extra_body", {}) for k in passed_params.keys(): - if k not in openai_params: + if k not in openai_params and passed_params[k] is not None: extra_body[k] = passed_params[k] optional_params.setdefault("extra_body", {}) initial_extra_body = { @@ -4087,7 +4107,7 @@ def add_provider_specific_params_to_optional_params( ) else: for k in passed_params.keys(): - if k not in openai_params: + if k not in openai_params and passed_params[k] is not None: optional_params[k] = passed_params[k] return optional_params @@ -4342,9 +4362,9 @@ def _count_characters(text: str) -> int: def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) -> str: - _choices: Union[ - List[Union[Choices, StreamingChoices]], List[StreamingChoices] - ] = response_obj.choices + _choices: Union[List[Union[Choices, StreamingChoices]], List[StreamingChoices]] = ( + response_obj.choices + ) response_str = "" for choice in _choices: @@ -6865,6 +6885,12 @@ class ProviderConfigManager: return litellm.VertexGeminiConfig() elif "claude" in model: return litellm.VertexAIAnthropicConfig() + elif "gpt-oss" in model: + from litellm.llms.vertex_ai.vertex_ai_partner_models.gpt_oss.transformation import ( + VertexAIGPTOSSTransformation, + ) + + return VertexAIGPTOSSTransformation() elif model in litellm.vertex_mistral_models: if "codestral" in model: return litellm.CodestralTextCompletionConfig() @@ -7091,6 +7117,12 @@ class ProviderConfigManager: ) return JinaAIEmbeddingConfig() + elif litellm.LlmProviders.VOLCENGINE == provider: + from litellm.llms.volcengine.embedding.transformation import ( + VolcEngineEmbeddingConfig, + ) + + return VolcEngineEmbeddingConfig() return None @staticmethod @@ -7272,6 +7304,21 @@ class ProviderConfigManager: from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig return VertexAIFilesConfig() + elif LlmProviders.BEDROCK == provider: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + return BedrockFilesConfig() + return None + + @staticmethod + def get_provider_batches_config( + model: str, + provider: LlmProviders, + ) -> Optional[BaseBatchesConfig]: + if LlmProviders.BEDROCK == provider: + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + return BedrockBatchesConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9abe3dc48fe..46eb48d2d42 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5817,16 +5817,6 @@ "supports_response_schema": true, "supports_tool_choice": true }, - "groq/llama3-8b-8192": { - "max_tokens": 8192, - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "input_cost_per_token": 5e-08, - "output_cost_per_token": 8e-08, - "litellm_provider": "groq", - "mode": "chat", - "supports_tool_choice": true - }, "groq/llama-3.2-1b-preview": { "max_tokens": 8192, "max_input_tokens": 8192, @@ -5907,17 +5897,6 @@ "supports_tool_choice": true, "deprecation_date": "2025-04-14" }, - "groq/llama3-70b-8192": { - "max_tokens": 8192, - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "input_cost_per_token": 5.9e-07, - "output_cost_per_token": 7.9e-07, - "litellm_provider": "groq", - "mode": "chat", - "supports_response_schema": true, - "supports_tool_choice": true - }, "groq/llama-3.1-8b-instant": { "max_tokens": 8192, "max_input_tokens": 128000, @@ -6178,21 +6157,7 @@ "supports_tool_choice": true, "source": "https://inference-docs.cerebras.ai/support/pricing" }, - "cerebras/openai/gpt-oss-20b": { - "max_tokens": 32768, - "max_input_tokens": 131072, - "max_output_tokens": 32768, - "input_cost_per_token": 7e-08, - "output_cost_per_token": 3e-07, - "litellm_provider": "cerebras", - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_reasoning": true, - "supports_tool_choice": true, - "source": "https://inference-docs.cerebras.ai/support/pricing" - }, + "cerebras/openai/gpt-oss-120b": { "max_tokens": 32768, "max_input_tokens": 131072, @@ -8027,8 +7992,8 @@ "max_pdf_size_mb": 30, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, - "output_cost_per_token": 2.5e-06, - "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_reasoning_token": 3e-05, "output_cost_per_image": 0.039, "litellm_provider": "gemini", "mode": "chat", @@ -8391,8 +8356,8 @@ "max_pdf_size_mb": 30, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, - "output_cost_per_token": 2.5e-06, - "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_reasoning_token": 3e-05, "output_cost_per_image": 0.039, "litellm_provider": "vertex_ai-language-models", "mode": "chat", @@ -9519,6 +9484,48 @@ "source": "https://aistudio.google.com", "supports_tool_choice": true }, + "gemini/veo-3.0-generate-preview": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.75, + "litellm_provider": "gemini", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, + "gemini/veo-3.0-fast-generate-preview": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.40, + "litellm_provider": "gemini", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, + "gemini/veo-2.0-generate-001": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.35, + "litellm_provider": "gemini", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, "vertex_ai/claude-opus-4-1": { "max_tokens": 4096, "max_input_tokens": 200000, @@ -9905,6 +9912,28 @@ "supports_tool_choice": true, "supports_prompt_caching": true }, + "vertex_ai/openai/gpt-oss-20b-maas": { + "max_tokens": 32768, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "input_cost_per_token": 0.075e-06, + "output_cost_per_token": 0.30e-06, + "litellm_provider": "vertex_ai-openai_models", + "mode": "chat", + "supports_reasoning": true, + "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas" + }, + "vertex_ai/openai/gpt-oss-120b-maas": { + "max_tokens": 32768, + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "input_cost_per_token": 0.15e-06, + "output_cost_per_token": 0.60e-06, + "litellm_provider": "vertex_ai-openai_models", + "mode": "chat", + "supports_reasoning": true, + "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas" + }, "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { "max_tokens": 32768, "max_input_tokens": 262144, @@ -10314,6 +10343,48 @@ "mode": "image_generation", "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "vertex_ai/veo-3.0-generate-preview": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.75, + "litellm_provider": "vertex_ai-video-models", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, + "vertex_ai/veo-3.0-fast-generate-preview": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.40, + "litellm_provider": "vertex_ai-video-models", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, + "vertex_ai/veo-2.0-generate-001": { + "max_tokens": 1024, + "max_input_tokens": 1024, + "output_cost_per_second": 0.35, + "litellm_provider": "vertex_ai-video-models", + "mode": "video_generation", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ], + "source": "https://ai.google.dev/gemini-api/docs/video" + }, "text-embedding-004": { "max_tokens": 2048, "max_input_tokens": 2048, @@ -20962,5 +21033,65 @@ "metadata": { "notes": "DALL-E 2 via AI/ML API - Reliable text-to-image generation" } + }, + "doubao-embedding-large": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 2048, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - large version with 2048 dimensions" + } + }, + "doubao-embedding-large-text-250515": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 2048, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - text-250515 version with 2048 dimensions" + } + }, + "doubao-embedding-large-text-240915": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 4096, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - text-240915 version with 4096 dimensions" + } + }, + "doubao-embedding": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 2560, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - standard version with 2560 dimensions" + } + }, + "doubao-embedding-text-240715": { + "max_tokens": 4096, + "max_input_tokens": 4096, + "output_vector_size": 2560, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "mode": "embedding", + "metadata": { + "notes": "Volcengine Doubao embedding model - text-240715 version with 2560 dimensions" + } } } \ No newline at end of file diff --git a/tests/batches_tests/bedrock_batch_completions.jsonl b/tests/batches_tests/bedrock_batch_completions.jsonl new file mode 100644 index 00000000000..3037b1031ea --- /dev/null +++ b/tests/batches_tests/bedrock_batch_completions.jsonl @@ -0,0 +1,3 @@ +{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}} +{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}} + diff --git a/tests/batches_tests/test_bedrock_files_and_batches.py b/tests/batches_tests/test_bedrock_files_and_batches.py new file mode 100644 index 00000000000..edf8c3c7745 --- /dev/null +++ b/tests/batches_tests/test_bedrock_files_and_batches.py @@ -0,0 +1,70 @@ + +# What is this? +## Unit Tests for OpenAI Batches API +import asyncio +import json +import os +import sys +import traceback +import tempfile +from dotenv import load_dotenv + +load_dotenv() +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system-path + + +import pytest +from typing import Optional +import litellm + + +@pytest.mark.asyncio() +async def test_async_create_file(): + """ + 1. Create File for Batch completion + 2. Create Batch Request + 3. Retrieve the specific batch + """ + litellm._turn_on_debug() + print("Testing async create batch") + + file_name = "bedrock_batch_completions.jsonl" + _current_dir = os.path.dirname(os.path.abspath(__file__)) + file_path = os.path.join(_current_dir, file_name) + file_obj = await litellm.acreate_file( + file=open(file_path, "rb"), + purpose="batch", + custom_llm_provider="bedrock", + s3_bucket_name="litellm-proxy", + ) + +@pytest.mark.asyncio() +async def test_async_file_and_batch(): + """ + Test file retrieval + """ + litellm._turn_on_debug() + file_name = "bedrock_batch_completions.jsonl" + _current_dir = os.path.dirname(os.path.abspath(__file__)) + file_path = os.path.join(_current_dir, file_name) + file_obj = await litellm.acreate_file( + file=open(file_path, "rb"), + purpose="batch", + custom_llm_provider="bedrock", + s3_bucket_name="litellm-proxy", + ) + print("CREATED FILE RESPONSE=", file_obj) + + # create batch + create_batch_response = await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id=file_obj.id, + metadata={"key1": "value1", "key2": "value2"}, + custom_llm_provider="bedrock", + aws_batch_role_arn="arn:aws:iam::888602223428:role/service-role/AmazonBedrockExecutionRoleForAgents_BB9HNW6V4CV" + ) + print("CREATED BATCH RESPONSE=", create_batch_response) + diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index c761da6d16c..7a803daf5d1 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -147,13 +147,13 @@ class TestBedrockNovaCanvasColorGuidedGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: litellm.in_memory_llm_clients_cache = InMemoryCache() return { - "model": "bedrock/amazon.nova-canvas-v1:0", - "n": 1, - "size": "320x320", - "imageGenerationConfig": {"cfgScale":6.5,"seed":12}, - "taskType": "COLOR_GUIDED_GENERATION", - "colorGuidedGenerationParams":{"colors":["#FFFFFF"]}, - "aws_region_name": "us-east-1", + "model": "bedrock/amazon.nova-canvas-v1:0", + "n": 1, + "size": "320x320", + "imageGenerationConfig": {"cfgScale": 6.5, "seed": 12}, + "taskType": "COLOR_GUIDED_GENERATION", + "colorGuidedGenerationParams": {"colors": ["#FFFFFF"]}, + "aws_region_name": "us-east-1", } @@ -161,22 +161,27 @@ class TestOpenAIDalle3(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "dall-e-3"} + class TestOpenAIGPTImage1(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "gpt-image-1"} + class TestRecraftImageGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "recraft/recraftv3"} + class TestAimlImageGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "aiml/flux-pro/v1.1"} + class TestGoogleImageGen(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "gemini/imagen-4.0-generate-001"} + class TestAzureOpenAIDalle3(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: litellm.set_verbose = True @@ -191,7 +196,6 @@ class TestAzureOpenAIDalle3(BaseImageGenTest): } }, } - class TestAzureFoundryFlux(BaseImageGenTest): @@ -268,117 +272,50 @@ async def test_aimage_generation_bedrock_with_optional_params(): pytest.fail(f"An exception occurred - {str(e)}") -@pytest.mark.asyncio -async def test_gpt_image_1_with_input_fidelity(): - """Test gpt-image-1 with input_fidelity parameter (mocked)""" - from unittest.mock import AsyncMock, patch - - # Mock OpenAI response - mock_openai_response = { - "created": 1703658209, - "data": [ - { - "url": "https://example.com/generated_image.png" - } - ] - } - - # Create a proper mock response object - class MockResponse: - def model_dump(self): - return mock_openai_response - - # Create a mock client with the images.generate method - mock_client = AsyncMock() - mock_client.images.generate = AsyncMock(return_value=MockResponse()) - - # Capture the actual arguments sent to OpenAI client - captured_args = None - captured_kwargs = None - - async def capture_generate_call(*args, **kwargs): - nonlocal captured_args, captured_kwargs - captured_args = args - captured_kwargs = kwargs - return MockResponse() - - mock_client.images.generate.side_effect = capture_generate_call - - # Mock the _get_openai_client method to return our mock client - with patch.object(litellm.main.openai_chat_completions, '_get_openai_client', return_value=mock_client): - response = await litellm.aimage_generation( - prompt="A cute baby sea otter", - model="gpt-image-1", - input_fidelity="high", - quality="medium", - size="1024x1024", - ) - - # Validate the response - assert response is not None - assert response.created == 1703658209 - assert response.data is not None - assert len(response.data) == 1 - assert response.data[0].url == "https://example.com/generated_image.png" - - # Validate that the OpenAI client was called with correct parameters - mock_client.images.generate.assert_called_once() - assert captured_kwargs is not None - assert captured_kwargs["model"] == "gpt-image-1" - assert captured_kwargs["prompt"] == "A cute baby sea otter" - assert captured_kwargs["input_fidelity"] == "high" - assert captured_kwargs["quality"] == "medium" - assert captured_kwargs["size"] == "1024x1024" - - @pytest.mark.asyncio async def test_aiml_image_generation_with_dynamic_api_key(): """ Test that when api_key is passed as a dynamic parameter to aimage_generation, it gets properly used for AIML provider authentication instead of falling back to environment variables. - + This test validates the fix for ensuring dynamic API keys are respected when making image generation requests to the AIML provider. """ from unittest.mock import AsyncMock, patch, MagicMock import httpx - + # Mock AIML response mock_aiml_response = { "created": 1703658209, - "data": [ - { - "url": "https://example.com/generated_image.png" - } - ] + "data": [{"url": "https://example.com/generated_image.png"}], } - + # Track captured arguments captured_headers = None captured_url = None captured_json_data = None - + def capture_post_call(*args, **kwargs): nonlocal captured_headers, captured_url, captured_json_data - captured_url = kwargs.get('url') or (args[0] if args else None) - captured_headers = kwargs.get('headers', {}) - captured_json_data = kwargs.get('json', {}) - + captured_url = kwargs.get("url") or (args[0] if args else None) + captured_headers = kwargs.get("headers", {}) + captured_json_data = kwargs.get("json", {}) + # Create a mock response mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = mock_aiml_response mock_response.text = json.dumps(mock_aiml_response) return mock_response - + # Mock the HTTP client that actually makes the request (sync version for image generation) - with patch('litellm.llms.custom_httpx.http_handler.HTTPHandler.post') as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.side_effect = capture_post_call - + # Test with dynamic api_key test_api_key = "test-dynamic-api-key-12345" - + response = await litellm.aimage_generation( prompt="A cute baby sea otter", model="aiml/flux-pro/v1.1", @@ -387,7 +324,7 @@ async def test_aiml_image_generation_with_dynamic_api_key(): # Validate the response (mocked response processing might not populate data correctly) assert response is not None - + # The most important validations: API key and endpoint usage # These prove that the dynamic API key was properly used assert captured_headers is not None @@ -398,10 +335,8 @@ async def test_aiml_image_generation_with_dynamic_api_key(): assert captured_url is not None assert "api.aimlapi.com" in captured_url assert "/v1/images/generations" in captured_url - + # Validate the request data assert captured_json_data is not None assert captured_json_data["prompt"] == "A cute baby sea otter" assert captured_json_data["model"] == "flux-pro/v1.1" - - diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 3228bd92189..0818d0b4b08 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1054,7 +1054,7 @@ def test_parse_content_for_reasoning(content, expected_reasoning, expected_conte ("gemini/gemini-1.5-pro", True), ("predibase/llama3-8b-instruct", True), ("gpt-3.5-turbo", False), - ("groq/llama3-70b-8192", True), + ("groq/llama-3.3-70b-versatile", True), ], ) def test_supports_response_schema(model, expected_bool): diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 5cb8295b1af..fc6983520fd 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -536,4 +536,57 @@ class BaseResponsesAPITest(ABC): # Validate final response structure validate_responses_api_response(final_response, final_chunk=True) assert final_response.output is not None - assert len(final_response.output) > 0 + + def test_openai_responses_api_dict_input_filtering(self): + """ + Test that regular dict inputs with status fields are properly filtered + to replicate exclude_unset=True behavior for non-Pydantic objects. + """ + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + # Test input with regular dict objects (like from JSON) + test_input = [ + { + "role": "user", + "content": "test" + }, + { + "id": "rs_123", + "summary": [{"text": "test", "type": "summary_text"}], + "type": "reasoning", + "content": None, # Should be filtered out + "encrypted_content": None, # Should be filtered out + "status": None # Should be filtered out + }, + { + "arguments": "{}", + "call_id": "call_123", + "name": "get_today", + "type": "function_call", + "id": "fc_123", + "status": "completed" # Should be preserved (not a default field) + } + ] + + config = OpenAIResponsesAPIConfig() + validated_input = config._validate_input_param(test_input) + + # Verify the results + assert len(validated_input) == 3 + + # Check reasoning item (index 1) + reasoning_item = validated_input[1] + assert reasoning_item["type"] == "reasoning" + assert "status" not in reasoning_item, "status field should be filtered out from reasoning item" + assert "content" not in reasoning_item, "content field should be filtered out from reasoning item" + assert "encrypted_content" not in reasoning_item, "encrypted_content field should be filtered out from reasoning item" + assert "id" in reasoning_item, "id field should be preserved" + assert "summary" in reasoning_item, "summary field should be preserved" + + # Check function call item (index 2) + function_call_item = validated_input[2] + assert function_call_item["type"] == "function_call" + assert "status" in function_call_item, "status field should be preserved in function call item" + assert function_call_item["status"] == "completed", "status value should be preserved" + + print("āœ… OpenAI Responses API dict input filtering test passed") diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index a1b05cbb4ae..59ca5cf61b8 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -633,14 +633,17 @@ def test_azure_openai_responses_bridge(): def test_azure_openai_gpt_5_responses_api(): - from litellm import responses + try: + from litellm import responses - litellm._turn_on_debug() + litellm._turn_on_debug() - response = responses( - model="azure/gpt-5", - input="Hello world", - api_key=os.getenv("AZURE_SWEDEN_API_KEY"), - api_base=os.getenv("AZURE_SWEDEN_API_BASE"), - ) - print(f"response: {response}") + response = responses( + model="azure/gpt-5", + input="Hello world", + api_key=os.getenv("AZURE_SWEDEN_API_KEY"), + api_base=os.getenv("AZURE_SWEDEN_API_BASE"), + ) + print(f"response: {response}") + except litellm.RateLimitError: + pytest.skip("Skipping test due to RateLimitError") diff --git a/tests/llm_translation/test_bedrock_gpt_oss.py b/tests/llm_translation/test_bedrock_gpt_oss.py index 61bce04e2d0..9487abfbc77 100644 --- a/tests/llm_translation/test_bedrock_gpt_oss.py +++ b/tests/llm_translation/test_bedrock_gpt_oss.py @@ -2,11 +2,13 @@ from base_llm_unit_tests import BaseLLMChatTest import pytest import sys import os +from unittest.mock import patch, MagicMock sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import litellm +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig class TestBedrockGPTOSS(BaseLLMChatTest): @@ -25,3 +27,27 @@ class TestBedrockGPTOSS(BaseLLMChatTest): Remove override once we have access to Bedrock prompt caching """ pass + + @pytest.mark.parametrize("model", [ + "bedrock/openai.gpt-oss-20b-1:0", + "bedrock/openai.gpt-oss-120b-1:0", + ]) + def test_reasoning_effort_transformation_gpt_oss(self, model): + """Test that reasoning_effort is handled correctly for GPT-OSS models.""" + config = AmazonConverseConfig() + + # Test GPT-OSS model - should keep reasoning_effort as-is + non_default_params = {"reasoning_effort": "low"} + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + # GPT-OSS should have reasoning_effort in result, not thinking + assert "reasoning_effort" in result + assert result["reasoning_effort"] == "low" + assert "thinking" not in result diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index 726371bd4e7..5bb4440d4d3 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -630,7 +630,6 @@ async def test_gemini_image_generation_async_stream(): and len(chunk.choices[0].delta.images) > 0 ): model_response_image = chunk.choices[0].delta.images[0]["image_url"] - print("MODEL_RESPONSE_IMAGE: ", model_response_image) assert model_response_image is not None assert model_response_image["url"].startswith("data:image/png;base64,") break @@ -660,3 +659,176 @@ def test_system_message_with_no_user_message(): assert response is not None assert response.choices[0].message.content is not None + +def get_current_weather(location, unit="fahrenheit"): + """Get the current weather in a given location""" + if "tokyo" in location.lower(): + return json.dumps({"location": "Tokyo", "temperature": "10", "unit": "celsius"}) + elif "san francisco" in location.lower(): + return json.dumps( + {"location": "San Francisco", "temperature": "72", "unit": "fahrenheit"} + ) + elif "paris" in location.lower(): + return json.dumps({"location": "Paris", "temperature": "22", "unit": "celsius"}) + else: + return json.dumps({"location": location, "temperature": "unknown"}) + + +def test_gemini_with_thinking(): + from litellm import completion + + litellm._turn_on_debug() + litellm.modify_params = True + model = "gemini/gemini-2.5-flash" + messages = [ + { + "role": "user", + "content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses", + } + ] + + tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state", + }, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + }, + }, + "required": ["location"], + }, + }, + } + ] + response = litellm.completion( + model=model, + messages=messages, + tools=tools, + tool_choice="auto", # auto is default, but we'll be explicit + reasoning_effort="low", + ) + print("Response\n", response) + response_message = response.choices[0].message + tool_calls = response_message.tool_calls + + print("Expecting there to be 3 tool calls") + assert len(tool_calls) > 0 # this has to call the function for SF, Tokyo and paris + + # Step 2: check if the model wanted to call a function + print(f"tool_calls: {tool_calls}") + if tool_calls: + # Step 3: call the function + # Note: the JSON response may not always be valid; be sure to handle errors + available_functions = { + "get_current_weather": get_current_weather, + } # only one function in this example, but you can have multiple + messages.append(response_message) # extend conversation with assistant's reply + print("Response message\n", response_message) + # Step 4: send the info for each function call and function response to the model + for tool_call in tool_calls: + function_name = tool_call.function.name + if function_name not in available_functions: + # the model called a function that does not exist in available_functions - don't try calling anything + return + function_to_call = available_functions[function_name] + function_args = json.loads(tool_call.function.arguments) + function_response = function_to_call( + location=function_args.get("location"), + unit=function_args.get("unit"), + ) + messages.append( + { + "tool_call_id": tool_call.id, + "role": "tool", + "name": function_name, + "content": function_response, + } + ) # extend conversation with function response + print(f"messages: {messages}") + second_response = litellm.completion( + model=model, + messages=messages, + seed=22, + reasoning_effort="low", + tools=tools, + drop_params=True, + ) # get a new response from the model where it can see the function response + print("second response\n", second_response) + + +def test_gemini_reasoning_effort_minimal(): + """ + Test that reasoning_effort='minimal' correctly maps to model-specific minimum thinking budgets + """ + from litellm.utils import return_raw_request + from litellm.types.utils import CallTypes + import json + + # Test with different Gemini models to verify model-specific mapping + test_cases = [ + ("gemini/gemini-2.5-flash", 1), # Flash: minimum 1 token + ("gemini/gemini-2.5-pro", 128), # Pro: minimum 128 tokens + ("gemini/gemini-2.5-flash-lite", 512), # Flash-Lite: minimum 512 tokens + ] + + for model, expected_min_budget in test_cases: + # Get the raw request to verify the thinking budget mapping + raw_request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": model, + "messages": [{"role": "user", "content": "Hello"}], + "reasoning_effort": "minimal", + }, + ) + + # Verify that the thinking config is set correctly + request_body = raw_request["raw_request_body"] + assert "generationConfig" in request_body, f"Model {model} should have generationConfig" + + generation_config = request_body["generationConfig"] + assert "thinkingConfig" in generation_config, f"Model {model} should have thinkingConfig" + + thinking_config = generation_config["thinkingConfig"] + assert "thinkingBudget" in thinking_config, f"Model {model} should have thinkingBudget" + + actual_budget = thinking_config["thinkingBudget"] + assert actual_budget == expected_min_budget, \ + f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}" + + # Verify that includeThoughts is True for minimal reasoning effort + assert thinking_config.get("includeThoughts", True), \ + f"Model {model} should have includeThoughts=True for minimal reasoning effort" + + # Test with unknown model (should use generic fallback) + try: + raw_request = return_raw_request( + endpoint=CallTypes.completion, + kwargs={ + "model": "gemini/unknown-model", + "messages": [{"role": "user", "content": "Hello"}], + "reasoning_effort": "minimal", + }, + ) + + request_body = raw_request["raw_request_body"] + generation_config = request_body["generationConfig"] + thinking_config = generation_config["thinkingConfig"] + # Should use generic fallback (128 tokens) + assert thinking_config["thinkingBudget"] == 128, \ + "Unknown model should use generic fallback of 128 tokens" + except Exception as e: + # If return_raw_request doesn't work for unknown models, that's okay + # The important part is that our known models work correctly + print(f"Note: Unknown model test skipped due to: {e}") + pass diff --git a/tests/llm_translation/test_groq.py b/tests/llm_translation/test_groq.py index dd1ae0aed6d..b9230d8eebb 100644 --- a/tests/llm_translation/test_groq.py +++ b/tests/llm_translation/test_groq.py @@ -1,6 +1,16 @@ -from base_llm_unit_tests import BaseLLMChatTest +import os +import sys +import pytest + +# sys.path.insert( +# 0, os.path.abspath("../..") +# ) # Adds the parent directory to the system path + +from base_llm_unit_tests import BaseLLMChatTest +from litellm.llms.groq.chat.transformation import GroqChatConfig + class TestGroq(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: return { @@ -10,3 +20,12 @@ class TestGroq(BaseLLMChatTest): def test_tool_call_no_arguments(self, tool_call_no_arguments): """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833""" pass + + def test_tool_call_with_empty_enum_property(self): + pass + + @pytest.mark.parametrize("model", ["groq/qwen/qwen3-32b", "groq/openai/gpt-oss-20b", "groq/openai/gpt-oss-120b"]) + def test_reasoning_effort_in_supported_params(self, model): + """Test that reasoning_effort is in the list of supported parameters for Groq""" + supported_params = GroqChatConfig().get_supported_openai_params(model=model) + assert "reasoning_effort" in supported_params diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index 0121eccaac3..285a406b3fb 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -664,3 +664,62 @@ async def test_openai_gpt5_reasoning(): ) print("response: ", response) assert response.choices[0].message.content is not None + +@pytest.mark.asyncio +async def test_openai_safety_identifier_parameter(): + """Test that safety_identifier parameter is correctly passed to the OpenAI API.""" + from openai import AsyncOpenAI + + litellm.set_verbose = True + client = AsyncOpenAI(api_key="fake-api-key") + + with patch.object( + client.chat.completions.with_raw_response, "create" + ) as mock_client: + try: + await litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello, how are you?"}], + safety_identifier="user_code_123456", + client=client, + ) + except Exception as e: + print(f"Error: {e}") + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + + # Verify the request contains the safety_identifier parameter + assert "safety_identifier" in request_body + # Verify safety_identifier is correctly sent to the API + assert request_body["safety_identifier"] == "user_code_123456" + + +def test_openai_safety_identifier_parameter_sync(): + """Test that safety_identifier parameter is correctly passed to the OpenAI API.""" + from openai import OpenAI + + litellm.set_verbose = True + client = OpenAI(api_key="fake-api-key") + + with patch.object( + client.chat.completions.with_raw_response, "create" + ) as mock_client: + try: + litellm.completion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello, how are you?"}], + safety_identifier="user_code_123456", + client=client, + ) + except Exception as e: + print(f"Error: {e}") + + mock_client.assert_called_once() + request_body = mock_client.call_args.kwargs + + # Verify the request contains the safety_identifier parameter + assert "safety_identifier" in request_body + # Verify safety_identifier is correctly sent to the API + assert request_body["safety_identifier"] == "user_code_123456" + diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index b908eabd0cf..a27fe738c7f 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -839,8 +839,8 @@ from test_completion import response_format_tests "model,region", [ ("vertex_ai/mistral-large-2411", "us-central1"), - ("vertex_ai/mistral-nemo@2407", "us-central1"), - ("vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", "us-south1") + ("vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", "us-south1"), + ("vertex_ai/openai/gpt-oss-20b-maas", "us-central1"), ], ) @pytest.mark.parametrize( @@ -911,6 +911,7 @@ async def test_partner_models_httpx(model, region, sync_mode): ("vertex_ai/meta/llama-4-scout-17b-16e-instruct-maas", "us-east5"), ("vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", "us-south1"), ("vertex_ai/mistral-large-2411", "us-central1"), # critical - we had this issue: https://github.com/BerriAI/litellm/issues/13888 + ("vertex_ai/openai/gpt-oss-20b-maas", "us-central1"), ], ) @pytest.mark.parametrize( diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py index bf482ca7527..39d4536d7aa 100644 --- a/tests/local_testing/test_completion_cost.py +++ b/tests/local_testing/test_completion_cost.py @@ -565,7 +565,7 @@ def test_groq_response_cost_tracking(is_streaming): response_cost = litellm.response_cost_calculator( response_object=response, - model="groq/llama3-70b-8192", + model="groq/llama-3.3-70b-versatile", custom_llm_provider="groq", call_type=CallTypes.acompletion.value, optional_params={}, diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index de1ba39e5bc..71249898f79 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -50,7 +50,6 @@ def get_current_weather(location, unit="fahrenheit"): "claude-3-haiku-20240307", "gemini/gemini-1.5-pro", "anthropic.claude-3-sonnet-20240229-v1:0", - "groq/llama-3.1-8b-instant", "cohere_chat/command-r", ], ) diff --git a/tests/local_testing/test_model_alias_map.py b/tests/local_testing/test_model_alias_map.py index e7d87a2946e..cf731d66283 100644 --- a/tests/local_testing/test_model_alias_map.py +++ b/tests/local_testing/test_model_alias_map.py @@ -35,7 +35,7 @@ def test_model_alias_map(caplog): for log in captured_logs: assert "ERROR" not in log - assert "llama3-8b-8192" in response.model + assert "llama-3.1-8b-instant" in response.model except litellm.ServiceUnavailableError: pass except Exception as e: diff --git a/tests/local_testing/test_router_cooldowns.py b/tests/local_testing/test_router_cooldown_handlers.py similarity index 100% rename from tests/local_testing/test_router_cooldowns.py rename to tests/local_testing/test_router_cooldown_handlers.py index cd178e2aaee..d3fd78063bf 100644 --- a/tests/local_testing/test_router_cooldowns.py +++ b/tests/local_testing/test_router_cooldown_handlers.py @@ -27,9 +27,9 @@ from litellm.router_utils.cooldown_handlers import ( _should_run_cooldown_logic, ) from litellm.types.router import ( + AllowedFailsPolicy, DeploymentTypedDict, LiteLLMParamsTypedDict, - AllowedFailsPolicy, ) diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index c0841d93c0f..ab6b8c5df5b 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -1422,8 +1422,6 @@ def test_bedrock_claude_3_streaming(): "claude-3-opus-20240229", "cohere.command-r-plus-v1:0", # bedrock "gpt-3.5-turbo", - # "databricks/databricks-dbrx-instruct", # databricks - "predibase/llama-3-8b-instruct", # predibase ], ) @pytest.mark.asyncio diff --git a/tests/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py index 386e9b299be..d0182a8d41d 100644 --- a/tests/otel_tests/test_prometheus.py +++ b/tests/otel_tests/test_prometheus.py @@ -516,9 +516,9 @@ async def test_key_budget_metrics(): ), "remaining budget should be less than 10.0 after first request" assert first_budget["total"] == 10.0, "Total budget metric is incorrect" print("first_budget['remaining_hours']", first_budget["remaining_hours"]) - # The budget reset time is now midnight, not exactly 7 days (168 hours) from creation - # So we'll check if it's within a reasonable range (5-7 days) - assert 120 <= first_budget["remaining_hours"] <= 168, "Budget remaining hours should be within a reasonable range (5-7 days)" + # The budget reset time is now standardized - for "7d" it resets on Monday at midnight + # So we'll check if it's within a reasonable range (0-7 days depending on current day of week) + assert 0 <= first_budget["remaining_hours"] <= 168, "Budget remaining hours should be within a reasonable range (0-7 days depending on day of week)" # Get key info and verify spend matches prometheus metrics key_info = await get_key_info(session, key) diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index d0403769425..57514817e88 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1375,3 +1375,142 @@ async def test_custom_validate_called(): pass # Assert custom_validate was called with the jwt token mock_custom_validate.assert_called_once_with({"sub": "test_user"}) + + +@pytest.mark.asyncio +async def test_auth_jwt_es256_jwk_path(monkeypatch): + import time, base64, jwt + from cryptography.hazmat.primitives.asymmetric import ec + from cryptography.hazmat.primitives import serialization + + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + + def b64url_uint(n: int, size: int) -> str: + return base64.urlsafe_b64encode(n.to_bytes(size, "big")).rstrip(b"=").decode() + + ec_key = ec.generate_private_key(ec.SECP256R1()) + ec_priv_pem = ec_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + + pub = ec_key.public_key().public_numbers() + ec_jwk = { + "kty": "EC", + "crv": "P-256", + "x": b64url_uint(pub.x, 32), + "y": b64url_uint(pub.y, 32), + "kid": "ec1", + "alg": "ES256", + "use": "sig", + } + + now = int(time.time()) + token = jwt.encode( + {"sub": "alice", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + ec_priv_pem, + algorithm="ES256", + headers={"kid": "ec1"}, + ) + + h = JWTHandler() + with patch.object(h, "get_public_key", new=AsyncMock(return_value=ec_jwk)): + claims = await h.auth_jwt(token) + assert claims["sub"] == "alice" + + +@pytest.mark.asyncio +async def test_auth_jwt_rs256_regression(monkeypatch): + """ + Regression: RSA path must still work (kty RSA, n/e) after EC support. + """ + import time, base64, jwt + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.hazmat.primitives import serialization + + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + + rsa_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + rsa_priv_pem = rsa_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + pub = rsa_key.public_key().public_numbers() + + def b64url(b: bytes) -> str: + return base64.urlsafe_b64encode(b).rstrip(b"=").decode() + + n = pub.n.to_bytes((pub.n.bit_length() + 7) // 8, "big") + e = pub.e.to_bytes((pub.e.bit_length() + 7) // 8, "big") + rsa_jwk = { + "kty": "RSA", + "n": b64url(n), + "e": b64url(e), + "kid": "rsa1", + "alg": "RS256", + "use": "sig", + } + + now = int(time.time()) + token = jwt.encode( + {"sub": "bob", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + rsa_priv_pem, + algorithm="RS256", + headers={"kid": "rsa1"}, + ) + + h = JWTHandler() + with patch.object(h, "get_public_key", new=AsyncMock(return_value=rsa_jwk)): + claims = await h.auth_jwt(token) + assert claims["sub"] == "bob" + + +@pytest.mark.asyncio +async def test_auth_jwt_mismatched_key_fails(monkeypatch): + """ + Negative: ES256 token must fail if JWKS returns an RSA key (mismatch). + """ + import time, base64, jwt + from cryptography.hazmat.primitives.asymmetric import ec, rsa + from cryptography.hazmat.primitives import serialization + + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + + # ES256 token + ec_key = ec.generate_private_key(ec.SECP256R1()) + ec_priv_pem = ec_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + now = int(time.time()) + token = jwt.encode( + {"sub": "mallory", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + ec_priv_pem, + algorithm="ES256", + headers={"kid": "ec1"}, + ) + + # RSA JWK (wrong key) + rsa_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + pub = rsa_key.public_key().public_numbers() + + def b64url(b: bytes) -> str: + return base64.urlsafe_b64encode(b).rstrip(b"=").decode() + + rsa_jwk = { + "kty": "RSA", + "n": b64url(pub.n.to_bytes((pub.n.bit_length() + 7) // 8, "big")), + "e": b64url(pub.e.to_bytes((pub.e.bit_length() + 7) // 8, "big")), + "kid": "rsa1", + "alg": "RS256", + "use": "sig", + } + + h = JWTHandler() + with patch.object(h, "get_public_key", new=AsyncMock(return_value=rsa_jwk)): + with pytest.raises(Exception) as exc: + await h.auth_jwt(token) + assert "Validation fails" in str(exc.value) \ No newline at end of file diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 48bb836dfd6..094df944bcc 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1690,3 +1690,38 @@ def test_handle_clientside_credential_with_responses_function(model_list): print( "āœ“ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata" ) + + +def test_get_metadata_variable_name_from_kwargs(model_list): + """ + Test _get_metadata_variable_name_from_kwargs method returns correct metadata variable name based on kwargs content. + """ + router = Router(model_list=model_list) + + # Test case 1: kwargs contains litellm_metadata - should return "litellm_metadata" + kwargs_with_litellm_metadata = { + "litellm_metadata": {"user": "test"}, + "metadata": {"other": "data"} + } + result = router._get_metadata_variable_name_from_kwargs(kwargs_with_litellm_metadata) + assert result == "litellm_metadata" + + # Test case 2: kwargs only contains metadata - should return "metadata" + kwargs_with_metadata_only = { + "metadata": {"user": "test"} + } + result = router._get_metadata_variable_name_from_kwargs(kwargs_with_metadata_only) + assert result == "metadata" + + # Test case 3: kwargs contains neither - should return "metadata" (default) + kwargs_empty = {} + result = router._get_metadata_variable_name_from_kwargs(kwargs_empty) + assert result == "metadata" + + # Test case 4: kwargs contains other keys but no metadata keys - should return "metadata" + kwargs_other = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}] + } + result = router._get_metadata_variable_name_from_kwargs(kwargs_other) + assert result == "metadata" diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index d389be79618..9cccdb51799 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -172,3 +172,27 @@ class TestSlackAlerting(unittest.TestCase): self.slack_alerting.update_values(alerting_args={"slack_alerting": "True"}) assert self.slack_alerting.periodic_started == True + + @patch("litellm.integrations.SlackAlerting.slack_alerting.datetime") + def test_alert_type_in_formatted_message(self, mock_datetime): + # Setup mocks + mock_datetime.now.return_value.strftime.return_value = "12:34:56" + + # Import required types + from litellm.types.integrations.slack_alerting import AlertType + + # Create a simple test message to check formatting + alert_type = AlertType.llm_exceptions + level = "Medium" + message = "Test alert message" + current_time = "12:34:56" + + # Test the specific formatting logic we're interested in + alert_type_formatted = f"Alert type: `{alert_type.name}`\n" + formatted_message = f"{alert_type_formatted}\n Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" + + # Verify alert_type is in the formatted message as expected + self.assertIn("Alert type: `llm_exceptions`", formatted_message) + self.assertIn("Level: `Medium`", formatted_message) + self.assertIn("Timestamp: `12:34:56`", formatted_message) + self.assertIn("Message: Test alert message", formatted_message) diff --git a/tests/test_litellm/integrations/cloudzero/test_dry_run_endpoint.py b/tests/test_litellm/integrations/cloudzero/test_dry_run_endpoint.py new file mode 100644 index 00000000000..5ba6457f376 --- /dev/null +++ b/tests/test_litellm/integrations/cloudzero/test_dry_run_endpoint.py @@ -0,0 +1,125 @@ +""" +Test the CloudZero dry run endpoint functionality +""" +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import polars as pl +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger + + +class TestCloudZeroDryRunEndpoint: + """Test suite for CloudZero dry run endpoint functionality.""" + + @pytest.mark.asyncio + async def test_dry_run_export_usage_data_returns_data(self): + """ + Test that dry_run_export_usage_data returns expected data structure + instead of just logging to console. + """ + logger = CloudZeroLogger() + + # Mock database data + mock_usage_data = pl.DataFrame({ + 'date': ['2025-01-19', '2025-01-20'], + 'model': ['gpt-4', 'gpt-3.5-turbo'], + 'custom_llm_provider': ['openai', 'openai'], + 'team_id': ['team1', 'team2'], + 'team_alias': ['Team One', 'Team Two'], + 'api_key_alias': ['key1', 'key2'], + 'prompt_tokens': [100, 200], + 'completion_tokens': [50, 100], + 'spend': [0.01, 0.02], + 'successful_requests': [1, 2] + }) + + # Mock CBF transformed data + mock_cbf_data = pl.DataFrame({ + 'time/usage_start': ['2025-01-19T00:00:00Z', '2025-01-20T00:00:00Z'], + 'cost/cost': [0.01, 0.02], + 'usage/amount': [150, 300], + 'resource/service': ['openai', 'openai'], + 'resource/account': ['litellm', 'litellm'], + 'resource/region': ['us-east-1', 'us-east-1'], + 'resource/id': ['gpt-4', 'gpt-3.5-turbo'], + 'entity_type': ['user', 'user'], + 'entity_id': ['team1', 'team2'], + 'resource/tag:team_id': ['team1', 'team2'], + 'resource/tag:team_alias': ['Team One', 'Team Two'], + 'resource/tag:api_key_alias': ['key1', 'key2'] + }) + + with patch('litellm.integrations.cloudzero.database.LiteLLMDatabase') as mock_db_class, \ + patch('litellm.integrations.cloudzero.transform.CBFTransformer') as mock_transformer_class: + + # Setup mocks + mock_db = AsyncMock() + mock_db.get_usage_data.return_value = mock_usage_data + mock_db_class.return_value = mock_db + + mock_transformer = MagicMock() + mock_transformer.transform.return_value = mock_cbf_data + mock_transformer_class.return_value = mock_transformer + + # Call the method + result = await logger.dry_run_export_usage_data(limit=1000) + + # Verify the result structure + assert isinstance(result, dict) + assert 'usage_data' in result + assert 'cbf_data' in result + assert 'summary' in result + + # Verify usage_data + assert isinstance(result['usage_data'], list) + assert len(result['usage_data']) == 2 + assert result['usage_data'][0]['model'] == 'gpt-4' + assert result['usage_data'][1]['model'] == 'gpt-3.5-turbo' + + # Verify cbf_data + assert isinstance(result['cbf_data'], list) + assert len(result['cbf_data']) == 2 + assert result['cbf_data'][0]['cost/cost'] == 0.01 + assert result['cbf_data'][1]['cost/cost'] == 0.02 + + # Verify summary + summary = result['summary'] + assert summary['total_records'] == 2 + assert summary['total_cost'] == 0.03 + assert summary['total_tokens'] == 450 # 150 + 300 + assert summary['unique_accounts'] == 1 + assert summary['unique_services'] == 1 + + @pytest.mark.asyncio + async def test_dry_run_export_usage_data_empty_data(self): + """ + Test that dry_run_export_usage_data handles empty data gracefully. + """ + logger = CloudZeroLogger() + + # Mock empty database data + mock_empty_data = pl.DataFrame() + + with patch('litellm.integrations.cloudzero.database.LiteLLMDatabase') as mock_db_class: + + # Setup mocks + mock_db = AsyncMock() + mock_db.get_usage_data.return_value = mock_empty_data + mock_db_class.return_value = mock_db + + # Call the method + result = await logger.dry_run_export_usage_data(limit=1000) + + # Verify the result structure for empty data + assert isinstance(result, dict) + assert result['usage_data'] == [] + assert result['cbf_data'] == [] + assert result['summary']['total_records'] == 0 + assert result['summary']['total_cost'] == 0 + assert result['summary']['total_tokens'] == 0 + diff --git a/tests/test_litellm/integrations/cloudzero/test_transform.py b/tests/test_litellm/integrations/cloudzero/test_transform.py new file mode 100644 index 00000000000..1f4db10cab8 --- /dev/null +++ b/tests/test_litellm/integrations/cloudzero/test_transform.py @@ -0,0 +1,183 @@ +import os +import sys +from datetime import datetime +from unittest.mock import MagicMock, patch + +import polars as pl +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.cloudzero.transform import CBFTransformer +from litellm.types.integrations.cloudzero import CBFRecord + + +class TestCBFTransformer: + """Test suite for CBFTransformer class.""" + + def test_init(self): + """Test CBFTransformer initialization.""" + transformer = CBFTransformer() + assert hasattr(transformer, 'czrn_generator') + assert transformer.czrn_generator is not None + + def test_transform_empty_dataframe(self): + """Test transform method with empty DataFrame.""" + transformer = CBFTransformer() + empty_df = pl.DataFrame() + + result = transformer.transform(empty_df) + + assert result.is_empty() + assert isinstance(result, pl.DataFrame) + + def test_transform_with_zero_successful_requests(self): + """Test transform method filters out records with zero successful_requests.""" + transformer = CBFTransformer() + data = pl.DataFrame({ + 'date': ['2025-01-19'], + 'successful_requests': [0], + 'spend': [10.0], + 'entity_id': ['test_entity'], + 'model': ['gpt-4'] + }) + + result = transformer.transform(data) + + assert result.is_empty() + + def test_transform_with_valid_data(self): + """Test transform method with valid data.""" + transformer = CBFTransformer() + with patch.object(transformer, '_create_cbf_record') as mock_create: + mock_create.return_value = CBFRecord({'test': 'data'}) + + data = pl.DataFrame({ + 'date': ['2025-01-19'], + 'successful_requests': [5], + 'spend': [10.0], + 'entity_id': ['test_entity'], + 'model': ['gpt-4'] + }) + + result = transformer.transform(data) + + assert len(result) == 1 + mock_create.assert_called_once() + + def test_transform_handles_czrn_generation_failures(self): + """Test transform method handles CZRN generation failures gracefully.""" + transformer = CBFTransformer() + with patch.object(transformer, '_create_cbf_record') as mock_create: + mock_create.side_effect = Exception("CZRN generation failed") + + data = pl.DataFrame({ + 'date': ['2025-01-19'], + 'successful_requests': [5], + 'spend': [10.0], + 'entity_id': ['test_entity'], + 'model': ['gpt-4'] + }) + + result = transformer.transform(data) + + assert result.is_empty() + + def test_create_cbf_record(self): + """Test _create_cbf_record method with valid row data.""" + transformer = CBFTransformer() + with patch.object(transformer.czrn_generator, 'create_from_litellm_data') as mock_czrn, \ + patch.object(transformer.czrn_generator, 'extract_components') as mock_extract: + + mock_czrn.return_value = 'test-czrn' + mock_extract.return_value = ('service', 'provider', 'region', 'account', 'resource', 'local_id') + + row = { + 'date': '2025-01-19', + 'spend': 10.5, + 'prompt_tokens': 100, + 'completion_tokens': 50, + 'entity_id': 'test_entity', + 'model': 'gpt-4', + 'entity_type': 'user', + 'model_group': 'openai', + 'custom_llm_provider': 'openai', + 'api_key': 'sk-test123', + 'api_requests': 5, + 'successful_requests': 5, + 'failed_requests': 0 + } + + result = transformer._create_cbf_record(row) + + assert isinstance(result, CBFRecord) + assert result['cost/cost'] == 10.5 + assert result['usage/amount'] == 150 # 100 + 50 + assert result['usage/units'] == 'tokens' + assert result['resource/id'] == 'test-czrn' + + def test_create_cbf_record_minimal_data(self): + """Test _create_cbf_record method with minimal row data.""" + transformer = CBFTransformer() + with patch.object(transformer.czrn_generator, 'create_from_litellm_data') as mock_czrn, \ + patch.object(transformer.czrn_generator, 'extract_components') as mock_extract: + + mock_czrn.return_value = 'test-czrn' + mock_extract.return_value = ('service', 'provider', 'region', 'account', 'resource', 'local_id') + + row = { + 'date': '2025-01-19', + 'spend': 0.0 + } + + result = transformer._create_cbf_record(row) + + assert isinstance(result, CBFRecord) + assert result['cost/cost'] == 0.0 + assert result['usage/amount'] == 0 # no tokens + assert result['usage/units'] == 'tokens' + + def test_parse_date_with_valid_string(self): + """Test _parse_date method with valid date string.""" + transformer = CBFTransformer() + + result = transformer._parse_date('2025-01-19') + + assert isinstance(result, datetime) + assert result.year == 2025 + assert result.month == 1 + assert result.day == 19 + + def test_parse_date_with_datetime_object(self): + """Test _parse_date method with datetime object.""" + transformer = CBFTransformer() + dt = datetime(2025, 1, 19) + + result = transformer._parse_date(dt) + + assert result == dt + + def test_parse_date_with_none(self): + """Test _parse_date method with None.""" + transformer = CBFTransformer() + + result = transformer._parse_date(None) + + assert result is None + + def test_parse_date_with_invalid_string(self): + """Test _parse_date method with invalid date string.""" + transformer = CBFTransformer() + + result = transformer._parse_date('invalid-date') + + assert result is None + + def test_parse_date_with_iso_format(self): + """Test _parse_date method with ISO format string.""" + transformer = CBFTransformer() + + result = transformer._parse_date('2025-01-19T10:30:00Z') + + assert isinstance(result, datetime) + assert result.year == 2025 \ No newline at end of file diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py index b4575a7ebdc..b1ce08de9e7 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py @@ -195,6 +195,32 @@ class TestDataDogLLMObsLogger: assert metadata["cache_hit"] == True assert metadata["cache_key"] == "test-cache-key-789" + def test_apm_id_included(self, mock_env_vars, mock_response_obj): + """Test that the current APM trace ID is attached to the payload""" + with patch('litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client'), \ + patch('asyncio.create_task'): + fake_tracer = MagicMock() + fake_span = MagicMock() + fake_span.trace_id = 987654321 + fake_tracer.current_span.return_value = fake_span + + with patch('litellm.integrations.datadog.datadog_llm_obs.tracer', fake_tracer): + logger = DataDogLLMObsLogger() + + standard_payload = create_standard_logging_payload_with_cache() + + kwargs = { + "standard_logging_object": standard_payload, + "litellm_params": {"metadata": {}} + } + + start_time = datetime.now() + end_time = datetime.now() + + payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) + + assert payload["apm_id"] == str(fake_span.trace_id) + def test_cache_metadata_fields(self, mock_env_vars, mock_response_obj): """Test that cache-related metadata fields are correctly tracked""" with patch('litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client'), \ diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json b/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json new file mode 100644 index 00000000000..913e3bfedae --- /dev/null +++ b/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json @@ -0,0 +1 @@ +{"litellm_trace_id": null, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "input": [{"role": "user", "content": "What is the capital of France?"}], "litellm_params": {"acompletion": true, "api_key": null, "force_timeout": 600, "logger_fn": null, "verbose": false, "custom_llm_provider": "bedrock", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse", "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "model_alias_map": {}, "completion_call_id": null, "aembedding": null, "metadata": {"requester_metadata": {}, "user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_user_id": null, "user_api_key_org_id": null, "user_api_key_team_alias": null, "user_api_key_end_user_id": null, "user_api_key_user_email": null, "user_api_key": "unused-for-aws-bedrock", "user_api_end_user_max_budget": null, "litellm_api_version": "1.72.3", "global_max_parallel_requests": null, "user_api_key_team_max_budget": null, "user_api_key_team_spend": null, "user_api_key_spend": 0.0, "user_api_key_max_budget": null, "user_api_key_model_max_budget": {}, "user_api_key_metadata": {}, "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "endpoint": "http://0.0.0.0:44444/chat/completions", "litellm_parent_otel_span": null, "requester_ip_address": "", "model_group": "claude-3-7-sonnet", "model_group_size": 1, "deployment": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "api_base": null, "caching_groups": null, "hidden_params": {"custom_llm_provider": "bedrock", "region_name": null, "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "api_base": null, "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "response_cost": 0.001047, "additional_headers": {"x-litellm-model-group": "claude-3-7-sonnet", "x-litellm-attempted-retries": 0, "x-litellm-attempted-fallbacks": 0}, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "litellm_overhead_time_ms": 231.156, "_response_ms": 236.798}}, "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "proxy_server_request": {"url": "http://0.0.0.0:44444/chat/completions", "method": "POST", "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "body": {"messages": [{"role": "user", "content": "What is the capital of France?"}], "model": "claude-3-7-sonnet", "stream": false}}, "preset_cache_key": null, "no-log": null, "stream_response": {}, "input_cost_per_token": null, "input_cost_per_second": null, "output_cost_per_token": null, "output_cost_per_second": null, "cooldown_time": null, "text_completion": null, "azure_ad_token_provider": null, "user_continue_message": null, "base_model": null, "litellm_trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "litellm_session_id": null, "hf_model_name": null, "custom_prompt_dict": {}, "litellm_metadata": null, "disable_add_transform_inline_image_block": null, "drop_params": null, "prompt_id": null, "prompt_variables": null, "async_call": null, "ssl_verify": null, "merge_reasoning_content_in_choices": false, "api_version": null, "azure_ad_token": null, "tenant_id": null, "client_id": null, "client_secret": null, "azure_username": null, "azure_password": null, "max_retries": 0, "timeout": 6000.0, "bucket_name": null, "vertex_credentials": null, "vertex_project": null, "use_litellm_proxy": false}, "applied_guardrails": [], "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "messages": [{"role": "user", "content": "What is the capital of France?"}], "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "start_time": "2025-06-22 10:59:08.159939", "stream": false, "user": null, "call_type": "acompletion", "completion_start_time": "2025-06-22 10:59:08.399523", "standard_callback_dynamic_params": {}, "stream_options": null, "max_retries": 0, "provider": "aws", "region": "us-west-2", "custom_llm_provider": "bedrock", "api_key": "", "additional_args": {"complete_input_dict": "{\"messages\": [{\"role\": \"user\", \"content\": [{\"text\": \"What is the capital of France?\"}]}], \"additionalModelRequestFields\": {\"provider\": \"aws\", \"region\": \"us-west-2\"}, \"system\": [], \"inferenceConfig\": {}}"}, "log_event_type": "post_api_call", "api_call_start_time": "2025-06-22 10:59:08.387641", "llm_api_duration_ms": 5.642, "original_response": "{\"metrics\":{\"latencyMs\":1513},\"output\":{\"message\":{\"content\":[{\"text\":\"The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.\"}],\"role\":\"assistant\"}},\"stopReason\":\"end_turn\",\"usage\":{\"cacheReadInputTokenCount\":0,\"cacheReadInputTokens\":0,\"cacheWriteInputTokenCount\":0,\"cacheWriteInputTokens\":0,\"inputTokens\":14,\"outputTokens\":67,\"totalTokens\":81}}", "end_time": "2025-06-22 10:59:08.399523", "cache_hit": null, "response_cost": 0.001047, "standard_logging_object": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "call_type": "acompletion", "cache_hit": null, "stream": true, "status": "success", "custom_llm_provider": "bedrock", "saved_cache_cost": 0.0, "startTime": 1750615148.162725, "endTime": 1750615148.399523, "completionStartTime": 1750615148.399523, "response_time": 0.23679804801940918, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "metadata": {"user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_org_id": null, "user_api_key_user_id": null, "user_api_key_team_alias": null, "user_api_key_user_email": null, "spend_logs_metadata": null, "requester_ip_address": "", "requester_metadata": {}, "user_api_key_end_user_id": null, "prompt_management_metadata": null, "applied_guardrails": [], "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "usage_object": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "requester_custom_headers": {"x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600"}}, "cache_key": null, "response_cost": 0.001047, "total_tokens": 81, "prompt_tokens": 14, "completion_tokens": 67, "request_tags": [], "end_user": "", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse", "model_group": "claude-3-7-sonnet", "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "requester_ip_address": "", "messages": [{"role": "user", "content": "What is the capital of France?"}], "response": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}, "model_parameters": {"stream": false}, "hidden_params": {"model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "cache_key": null, "api_base": null, "response_cost": 0.001047, "additional_headers": {}, "litellm_overhead_time_ms": 231.156, "batch_models": null, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "usage_object": null}, "model_map_information": {"model_map_key": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "model_map_value": {"key": "anthropic.claude-3-7-sonnet-20250219-v1:0", "max_tokens": 8192, "max_input_tokens": 200000, "max_output_tokens": 8192, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_reasoning_token": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "bedrock_converse", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": null, "supports_audio_output": null, "supports_pdf_input": true, "supports_embedding_image_input": null, "supports_native_streaming": null, "supports_web_search": null, "supports_url_context": null, "supports_reasoning": true, "supports_computer_use": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["max_tokens", "max_completion_tokens", "stream", "stream_options", "stop", "temperature", "top_p", "extra_headers", "response_format", "tools", "tool_choice", "thinking", "reasoning_effort"]}}, "error_str": null, "error_information": {"error_code": "", "error_class": "", "llm_provider": "", "traceback": "", "error_message": ""}, "response_cost_failure_debug_info": null, "guardrail_information": null, "standard_built_in_tools_params": {"web_search_options": null, "file_search": null}}, "async_complete_streaming_response": "ModelResponse(id='chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1', created=1750615148, model='arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0', object='chat.completion', system_fingerprint=None, choices=[Choices(finish_reason='stop', index=0, message=Message(content='The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None))], usage=Usage(completion_tokens=67, prompt_tokens=14, total_tokens=81, completion_tokens_details=None, prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None), cache_creation_input_tokens=0, cache_read_input_tokens=0))"} \ No newline at end of file diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_response.json b/tests/test_litellm/integrations/open_telemetry/data/captured_response.json new file mode 100644 index 00000000000..3cf77781cc2 --- /dev/null +++ b/tests/test_litellm/integrations/open_telemetry/data/captured_response.json @@ -0,0 +1 @@ +{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}} \ No newline at end of file diff --git a/tests/test_litellm/integrations/test_braintrust_span_name.py b/tests/test_litellm/integrations/test_braintrust_span_name.py index 10e512fc0ca..30381e99783 100644 --- a/tests/test_litellm/integrations/test_braintrust_span_name.py +++ b/tests/test_litellm/integrations/test_braintrust_span_name.py @@ -11,7 +11,7 @@ from litellm.integrations.braintrust_logging import BraintrustLogger class TestBraintrustSpanName(unittest.TestCase): """Test custom span_name functionality in Braintrust logging.""" - @patch('litellm.integrations.braintrust_logging.HTTPHandler') + @patch("litellm.integrations.braintrust_logging.HTTPHandler") def test_default_span_name(self, MockHTTPHandler): """Test that default span name is 'Chat Completion' when not provided.""" # Mock HTTP response @@ -22,39 +22,43 @@ class TestBraintrustSpanName(unittest.TestCase): # Setup logger = BraintrustLogger(api_key="test-key") logger.default_project_id = "test-project-id" - + # Create a properly structured mock response response_obj = litellm.ModelResponse( id="test-id", object="chat.completion", created=1234567890, model="gpt-3.5-turbo", - choices=[{ - "index": 0, - "message": {"role": "assistant", "content": "test response"}, - "finish_reason": "stop" - }], - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30} + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "test response"}, + "finish_reason": "stop", + } + ], + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, ) - + kwargs = { "litellm_call_id": "test-call-id", "messages": [{"role": "user", "content": "test"}], "litellm_params": {"metadata": {}}, "model": "gpt-3.5-turbo", - "response_cost": 0.001 + "response_cost": 0.001, } - + # Execute logger.log_success_event(kwargs, response_obj, datetime.now(), datetime.now()) - + # Verify call_args = mock_http_handler.post.call_args self.assertIsNotNone(call_args) - json_data = call_args.kwargs['json'] - self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Chat Completion') + json_data = call_args.kwargs["json"] + self.assertEqual( + json_data["events"][0]["span_attributes"]["name"], "Chat Completion" + ) - @patch('litellm.integrations.braintrust_logging.HTTPHandler') + @patch("litellm.integrations.braintrust_logging.HTTPHandler") def test_custom_span_name(self, MockHTTPHandler): """Test that custom span name is used when provided in metadata.""" # Mock HTTP response @@ -65,39 +69,43 @@ class TestBraintrustSpanName(unittest.TestCase): # Setup logger = BraintrustLogger(api_key="test-key") logger.default_project_id = "test-project-id" - + # Create a properly structured mock response response_obj = litellm.ModelResponse( id="test-id", object="chat.completion", created=1234567890, model="gpt-3.5-turbo", - choices=[{ - "index": 0, - "message": {"role": "assistant", "content": "test response"}, - "finish_reason": "stop" - }], - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30} + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "test response"}, + "finish_reason": "stop", + } + ], + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, ) - + kwargs = { "litellm_call_id": "test-call-id", "messages": [{"role": "user", "content": "test"}], "litellm_params": {"metadata": {"span_name": "Custom Operation"}}, "model": "gpt-3.5-turbo", - "response_cost": 0.001 + "response_cost": 0.001, } - + # Execute logger.log_success_event(kwargs, response_obj, datetime.now(), datetime.now()) - + # Verify call_args = mock_http_handler.post.call_args self.assertIsNotNone(call_args) - json_data = call_args.kwargs['json'] - self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Custom Operation') + json_data = call_args.kwargs["json"] + self.assertEqual( + json_data["events"][0]["span_attributes"]["name"], "Custom Operation" + ) - @patch('litellm.integrations.braintrust_logging.HTTPHandler') + @patch("litellm.integrations.braintrust_logging.HTTPHandler") def test_span_name_with_other_metadata(self, MockHTTPHandler): """Test that span_name works alongside other metadata fields.""" # Mock HTTP response @@ -108,21 +116,23 @@ class TestBraintrustSpanName(unittest.TestCase): # Setup logger = BraintrustLogger(api_key="test-key") logger.default_project_id = "test-project-id" - + # Create a properly structured mock response response_obj = litellm.ModelResponse( id="test-id", object="chat.completion", created=1234567890, model="gpt-3.5-turbo", - choices=[{ - "index": 0, - "message": {"role": "assistant", "content": "test response"}, - "finish_reason": "stop" - }], - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30} + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "test response"}, + "finish_reason": "stop", + } + ], + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, ) - + kwargs = { "litellm_call_id": "test-call-id", "messages": [{"role": "user", "content": "test"}], @@ -132,34 +142,40 @@ class TestBraintrustSpanName(unittest.TestCase): "project_id": "custom-project", "user_id": "user123", "session_id": "session456", - "environment": "production" + "environment": "production", } }, "model": "gpt-3.5-turbo", - "response_cost": 0.001 + "response_cost": 0.001, + "standard_logging_object": { + "user_id": "user123", + }, } - + # Execute logger.log_success_event(kwargs, response_obj, datetime.now(), datetime.now()) - + # Verify call_args = mock_http_handler.post.call_args self.assertIsNotNone(call_args) - json_data = call_args.kwargs['json'] - - # Check span name - self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Multi Metadata Test') - - # Check that other metadata is preserved (except for filtered keys) - event_metadata = json_data['events'][0]['metadata'] - self.assertEqual(event_metadata['user_id'], 'user123') - self.assertEqual(event_metadata['session_id'], 'session456') - self.assertEqual(event_metadata['environment'], 'production') - - # Span name should be in span_attributes, not in metadata - self.assertIn('span_name', event_metadata) # span_name is also kept in metadata + json_data = call_args.kwargs["json"] - @patch('litellm.integrations.braintrust_logging.get_async_httpx_client') + # Check span name + self.assertEqual( + json_data["events"][0]["span_attributes"]["name"], "Multi Metadata Test" + ) + + # Check that other metadata is preserved (except for filtered keys) + event_metadata = json_data["events"][0]["metadata"] + print(event_metadata) + self.assertEqual(event_metadata["user_id"], "user123") + self.assertEqual(event_metadata["session_id"], "session456") + self.assertEqual(event_metadata["environment"], "production") + + # Span name should be in span_attributes, not in metadata + self.assertIn("span_name", event_metadata) # span_name is also kept in metadata + + @patch("litellm.integrations.braintrust_logging.get_async_httpx_client") async def test_async_custom_span_name(self, mock_get_http_handler): """Test async logging with custom span name.""" # Mock async HTTP response @@ -170,38 +186,44 @@ class TestBraintrustSpanName(unittest.TestCase): # Setup logger = BraintrustLogger(api_key="test-key") logger.default_project_id = "test-project-id" - + # Create a properly structured mock response response_obj = litellm.ModelResponse( id="test-id", object="chat.completion", created=1234567890, model="gpt-3.5-turbo", - choices=[{ - "index": 0, - "message": {"role": "assistant", "content": "test response"}, - "finish_reason": "stop" - }], - usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30} + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "test response"}, + "finish_reason": "stop", + } + ], + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, ) - + kwargs = { "litellm_call_id": "test-call-id", "messages": [{"role": "user", "content": "test"}], "litellm_params": {"metadata": {"span_name": "Async Custom Operation"}}, "model": "gpt-3.5-turbo", - "response_cost": 0.001 + "response_cost": 0.001, } - + # Execute - await logger.async_log_success_event(kwargs, response_obj, datetime.now(), datetime.now()) - + await logger.async_log_success_event( + kwargs, response_obj, datetime.now(), datetime.now() + ) + # Verify call_args = mock_http_handler.post.call_args self.assertIsNotNone(call_args) - json_data = call_args.kwargs['json'] - self.assertEqual(json_data['events'][0]['span_attributes']['name'], 'Async Custom Operation') + json_data = call_args.kwargs["json"] + self.assertEqual( + json_data["events"][0]["span_attributes"]["name"], "Async Custom Operation" + ) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index e71b68ab934..182e0134928 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,4 +1,4 @@ -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest @@ -82,3 +82,101 @@ class TestCustomGuardrailDeploymentHook: # Verify messages were updated in result assert result["messages"] == mock_result["messages"] assert result["messages"] != original_messages + + +class TestCustomGuardrailShouldRunGuardrail: + + def test_should_run_guardrail_with_litellm_metadata(self): + """Test that should_run_guardrail works with litellm_metadata pattern""" + from litellm.types.guardrails import GuardrailEventHooks + + custom_guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.pre_call + ) + + # Test with guardrails in litellm_metadata + data = { + "model": "gpt-3.5-turbo", + "litellm_metadata": { + "guardrails": ["test_guardrail"] + } + } + + result = custom_guardrail.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + + assert result is True + + def test_should_run_guardrail_with_metadata(self): + """Test that should_run_guardrail works with metadata pattern""" + from litellm.types.guardrails import GuardrailEventHooks + + custom_guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.pre_call + ) + + # Test with guardrails in metadata + data = { + "model": "gpt-3.5-turbo", + "metadata": { + "guardrails": ["test_guardrail"] + } + } + + result = custom_guardrail.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + + assert result is True + + def test_should_run_guardrail_with_root_level_guardrails(self): + """Test that should_run_guardrail works with root level guardrails""" + from litellm.types.guardrails import GuardrailEventHooks + + custom_guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.pre_call + ) + + # Test with guardrails at root level + data = { + "model": "gpt-3.5-turbo", + "guardrails": ["test_guardrail"] + } + + result = custom_guardrail.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + + assert result is True + + + def test_should_run_guardrail_no_matching_guardrail(self): + """Test that should_run_guardrail returns False when guardrail name doesn't match""" + from litellm.types.guardrails import GuardrailEventHooks + + custom_guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.pre_call + ) + + # Test with different guardrail name + data = { + "model": "gpt-3.5-turbo", + "litellm_metadata": { + "guardrails": ["different_guardrail"] + } + } + + result = custom_guardrail.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + + assert result is False diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index e11895e30ea..7fb91f274d0 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -1,15 +1,142 @@ +import json import os import sys import unittest from unittest.mock import MagicMock, patch +from datetime import datetime, timedelta +import time # Adds the grandparent directory to sys.path to allow importing project modules sys.path.insert(0, os.path.abspath("../..")) from litellm.integrations.opentelemetry import OpenTelemetry from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk._logs import LoggerProvider as OTLoggerProvider +from opentelemetry.sdk._logs.export import SimpleLogRecordProcessor, InMemoryLogExporter +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + +class TestOpenTelemetryGuardrails(unittest.TestCase): + @patch("litellm.integrations.opentelemetry.datetime") + def test_create_guardrail_span_with_valid_info(self, mock_datetime): + # Setup + otel = OpenTelemetry() + otel.tracer = MagicMock() + mock_span = MagicMock() + otel.tracer.start_span.return_value = mock_span + + # Create guardrail information + guardrail_info = { + "guardrail_name": "test_guardrail", + "guardrail_mode": "input", + "masked_entity_count": {"CREDIT_CARD": 2}, + "guardrail_response": "filtered_content", + "start_time": 1609459200.0, + "end_time": 1609459201.0, + } + + # Create a kwargs dict with standard_logging_object containing guardrail information + kwargs = {"standard_logging_object": {"guardrail_information": guardrail_info}} + + # Call the method + otel._create_guardrail_span(kwargs=kwargs, context=None) + + # Assertions + otel.tracer.start_span.assert_called_once() + + # print all calls to mock_span.set_attribute + print("Calls to mock_span.set_attribute:") + for call in mock_span.set_attribute.call_args_list: + print(call) + + # Check that the span has the correct attributes set + mock_span.set_attribute.assert_any_call("guardrail_name", "test_guardrail") + mock_span.set_attribute.assert_any_call("guardrail_mode", "input") + mock_span.set_attribute.assert_any_call( + "guardrail_response", "filtered_content" + ) + mock_span.set_attribute.assert_any_call( + "masked_entity_count", safe_dumps({"CREDIT_CARD": 2}) + ) + + # Verify that the span was ended + mock_span.end.assert_called_once() + + def test_create_guardrail_span_with_no_info(self): + # Setup + otel = OpenTelemetry() + otel.tracer = MagicMock() + + # Test with no guardrail information + kwargs = {"standard_logging_object": {}} + otel._create_guardrail_span(kwargs=kwargs, context=None) + + # Verify that start_span was never called + otel.tracer.start_span.assert_not_called() + class TestOpenTelemetry(unittest.TestCase): + POLL_INTERVAL = 0.05 + POLL_TIMEOUT = 2.0 + MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + HERE = os.path.dirname(__file__) + + def wait_for_spans(self, exporter: InMemorySpanExporter, prefix: str): + """Poll until we see at least one span with an attribute key starting with `prefix`.""" + deadline = time.time() + self.POLL_TIMEOUT + while time.time() < deadline: + spans = exporter.get_finished_spans() + matches = [ + s + for s in spans + if s.attributes and any(str(k).startswith(prefix) for k in s.attributes) + ] + if matches: + return matches + time.sleep(self.POLL_INTERVAL) + return [] + + def wait_for_metric(self, reader: InMemoryMetricReader, name: str): + """Poll until we see a metric with the given name.""" + deadline = time.time() + self.POLL_TIMEOUT + while time.time() < deadline: + data = reader.get_metrics_data() + # guard against None or missing attribute + if not data or not hasattr(data, "resource_metrics"): + time.sleep(self.POLL_INTERVAL) + continue + + for rm in data.resource_metrics: + for sm in rm.scope_metrics: + for m in sm.metrics: + if m.name == name: + return m + + time.sleep(self.POLL_INTERVAL) + return None + + def wait_for_log(self, reader: InMemoryLogExporter, name: str): + """Poll until we see a log with the given name.""" + deadline = time.time() + self.POLL_TIMEOUT + while time.time() < deadline: + logs = reader.get_finished_logs() + if not logs: + time.sleep(self.POLL_INTERVAL) + continue + matches = [ + log + for log in logs + # if log.attributes and any(str(k).startswith(prefix) for k in log.attributes) + ] + if matches: + return matches + time.sleep(self.POLL_INTERVAL) + return [] + @patch("litellm.integrations.opentelemetry.datetime") def test_create_guardrail_span_with_valid_info(self, mock_datetime): # Setup @@ -79,7 +206,6 @@ class TestOpenTelemetry(unittest.TestCase): ) as mock_get_headers, patch.object( otel, "_get_tracer_with_dynamic_headers" ) as mock_get_tracer: - # Test case 1: With dynamic headers mock_get_headers.return_value = { "arize-space-id": "test-space", @@ -399,3 +525,229 @@ class TestOpenTelemetry(unittest.TestCase): self.assertEqual(attributes.get("service.name"), "litellm-service") # But other attributes from OTEL_RESOURCE_ATTRIBUTES should still be present self.assertEqual(attributes.get("extra.attr"), "extra-value") + + def test_handle_success_generates_spans_metrics_and_events(self): + # force both metrics & events on + os.environ["LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS"] = "true" + os.environ["LITELLM_OTEL_INTEGRATION_ENABLE_METRICS"] = "true" + + # ─── build in‐memory OTEL providers/exporters ───────────────────────────── + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + + log_exporter = InMemoryLogExporter() + logger_provider = OTLoggerProvider() + logger_provider.add_log_record_processor(SimpleLogRecordProcessor(log_exporter)) + + metric_reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[metric_reader]) + + # ─── instantiate our OpenTelemetry logger with test providers ─────────── + otel = OpenTelemetry( + tracer_provider=tracer_provider, + meter_provider=meter_provider, + logger_provider=logger_provider, + ) + + # OpenTelemetry attempts to set a global tracer provider, which can be set only once. + # so we hack here to set a local tracer deriver from the provider we created. + otel.tracer = tracer_provider.get_tracer(__name__) + + # ─── minimal input / output for a chat call ────────────────────────────── + start = datetime.utcnow() + end = start + timedelta(seconds=1) + + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_kwargs.json") + ) as f: + kwargs = json.load(f) + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_response.json") + ) as f: + response_obj = json.load(f) + + # ─── exercise the hook ─────────────────────────────────────────────────── + otel._handle_success(kwargs, response_obj, start, end) + + # ─── assert spans ──────────────────────────────────────────────────────── + spans = self.wait_for_spans(span_exporter, "gen_ai.") + self.assertTrue(spans, "Expected at least one gen_ai span") + + # verify our top‐level litellm_request span is present + names = [s.name for s in spans] + self.assertIn("litellm_request", names) + + # ─── assert metrics ────────────────────────────────────────────────────── + duration_metric = self.wait_for_metric( + metric_reader, "gen_ai.client.operation.duration" + ) + self.assertIsNotNone(duration_metric, "duration histogram was not recorded") + + # check that our model attribute made it onto at least one data point + found_dp = False + if ( + duration_metric + and hasattr(duration_metric, "data") + and hasattr(duration_metric.data, "data_points") + ): + found_dp = any( + dp.attributes.get("gen_ai.request.model") == self.MODEL + for dp in duration_metric.data.data_points + ) + self.assertTrue( + found_dp, "expected gen_ai.request.model attribute on a data point" + ) + + # ─── assert logs ─────────────────────────────────────────────────────── + logs = [] + logs = self.wait_for_log(log_exporter, "gen_ai.") + self.assertTrue(logs, "Expected at least one gen_ai log") + + user_logs = [log for log in logs if log.log_record.attributes.get("event_name") == "gen_ai.content.prompt"] + self.assertTrue(user_logs, "did not see a gen_ai.content.prompt log") + # check log bodies + user_prompt = user_logs[0].log_record.attributes.get("gen_ai.prompt") + self.assertEqual("What is the capital of France?", user_prompt, "did not see a prompt message") + + choice_logs = [log for log in logs if log.log_record.attributes.get("event_name") == "gen_ai.content.completion"] + self.assertTrue(choice_logs, "did not see a gen_ai.content.completion event") + + choice_response = choice_logs[0].log_record.body + self.assertIsNotNone(choice_response, "did not see a response message") + self.assertEqual("stop", choice_response.get("finish_reason"), "did not see expected finish reason") + + + def test_handle_success_spans_only(self): + # make sure neither events nor metrics is on + os.environ.pop("LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS", None) + os.environ.pop("LITELLM_OTEL_INTEGRATION_ENABLE_METRICS", None) + + # ─── build in‐memory OTEL providers/exporters ───────────────────────────── + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + + # no logs / no metrics + log_exporter = InMemoryLogExporter() + logger_provider = OTLoggerProvider() + logger_provider.add_log_record_processor(SimpleLogRecordProcessor(log_exporter)) + metric_reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[metric_reader]) + + # ─── instantiate our OpenTelemetry logger with test providers ─────────── + otel = OpenTelemetry( + tracer_provider=tracer_provider, + meter_provider=meter_provider, + logger_provider=logger_provider, # pass even if events disabled (safe) + ) + # bind our tracer to the test tracer provider (global registration is a no-op after the first time) + otel.tracer = tracer_provider.get_tracer(__name__) + + # ─── minimal input / output for a chat call ────────────────────────────── + start = datetime.utcnow() + end = start + timedelta(seconds=1) + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_kwargs.json") + ) as f: + kwargs = json.load(f) + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_response.json") + ) as f: + response_obj = json.load(f) + + # ─── exercise the hook ─────────────────────────────────────────────────── + otel._handle_success(kwargs, response_obj, start, end) + + # ─── assert spans only ─────────────────────────────────────────────────── + spans = span_exporter.get_finished_spans() + self.assertTrue(spans, "Expected at least one span") + # must have the top‐level litellm_request span + # self.assertIn( + # LITELLM_REQUEST_SPAN_NAME, + # [s.name for s in spans], + # "litellm_request span missing", + # ) + # model attribute should be on that span + found = any( + s.attributes + and s.attributes.get("gen_ai.request.model") == self.MODEL + for s in spans + ) + self.assertTrue(found, "expected gen_ai.request.model on span attributes") + + # no metrics recorded + self.assertIsNone( + self.wait_for_metric(metric_reader, "gen_ai.client.operation.duration"), + "Did not expect any metrics", + ) + # no logs emitted + logs = log_exporter.get_finished_logs() + self.assertFalse(logs, "Did not expect any logs") + + def test_handle_success_spans_and_metrics(self): + # only metrics on + os.environ.pop("LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS", None) + os.environ["LITELLM_OTEL_INTEGRATION_ENABLE_METRICS"] = "true" + + # ─── build in‐memory OTEL providers/exporters ───────────────────────────── + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + + log_exporter = InMemoryLogExporter() + logger_provider = OTLoggerProvider() + logger_provider.add_log_record_processor(SimpleLogRecordProcessor(log_exporter)) + metric_reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[metric_reader]) + + # ─── instantiate our OpenTelemetry logger with test providers ─────────── + otel = OpenTelemetry( + tracer_provider=tracer_provider, + meter_provider=meter_provider, + logger_provider=logger_provider, # needed if events were enabled + ) + otel.tracer = tracer_provider.get_tracer(__name__) + + # ─── minimal input / output for a chat call ────────────────────────────── + start = datetime.utcnow() + end = start + timedelta(seconds=1) + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_kwargs.json") + ) as f: + kwargs = json.load(f) + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_response.json") + ) as f: + response_obj = json.load(f) + + # ─── exercise the hook ─────────────────────────────────────────────────── + otel._handle_success(kwargs, response_obj, start, end) + + # ─── assert spans ──────────────────────────────────────────────────────── + spans = span_exporter.get_finished_spans() + self.assertTrue(spans, "Expected at least one span") + + # ─── assert metrics ────────────────────────────────────────────────────── + duration_metric = self.wait_for_metric( + metric_reader, "gen_ai.client.operation.duration" + ) + self.assertIsNotNone(duration_metric, "duration histogram was not recorded") + # model attribute should be present on a data point + found_dp = False + if ( + duration_metric + and hasattr(duration_metric, "data") + and hasattr(duration_metric.data, "data_points") + ): + found_dp = any( + dp.attributes.get("gen_ai.request.model") == self.MODEL + for dp in duration_metric.data.data_points + ) + self.assertTrue( + found_dp, "expected gen_ai.request.model attribute on a data point" + ) + + # ─── no events when only metrics enabled ───────────────────────────────── + logs = log_exporter.get_finished_logs() + self.assertFalse(logs, "Did not expect any logs") 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 1c91cc0fe8b..2fc710664e6 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -475,6 +475,239 @@ def test_transform_response_with_bash_tool(): assert args["command"] == "ls -la *.py" +def test_transform_response_with_structured_response_being_called(): + """Test response transformation with structured response.""" + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + # Simulate a Bedrock Converse response with a bash tool call + response_json = { + "additionalModelResponseFields": {}, + "metrics": {"latencyMs": 100.0}, + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_456", + "name": "json_tool_call", + "input": { + "Current_Temperature": 62, + "Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation."}, + } + } + ] + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 8, + "outputTokens": 3, + "totalTokens": 11, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + # Mock httpx.Response + 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, + "tools": [ + { + 'type': 'function', + 'function': { + 'name': 'get_weather', + 'description': 'Get the current weather in a given location', + 'parameters': { + 'type': 'object', + 'properties': { + 'location': { + 'type': 'string', + 'description': 'The city and state, e.g. San Francisco, CA' + }, + 'unit': { + 'type': 'string', + 'enum': ['celsius', 'fahrenheit'] + } + }, + 'required': ['location'] + } + } + }, + { + 'type': 'function', + 'function': { + 'name': 'json_tool_call', + 'parameters': { + '$schema': 'http://json-schema.org/draft-07/schema#', + 'type': 'object', + 'required': ['Weather_Explanation', 'Current_Temperature'], + 'properties': { + 'Weather_Explanation': { + 'type': ['string', 'null'], + 'description': '1-2 sentences explaining the weather in the location' + }, + 'Current_Temperature': { + 'type': ['number', 'null'], + 'description': 'Current temperature in the location' + } + }, + 'additionalProperties': False + } + } + } + ] + } + # Call the transformation logic + 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, + ) + # Check that the tool call is present in the returned message + assert result.choices[0].message.tool_calls is None + + assert result.choices[0].message.content is not None + assert result.choices[0].message.content == '{"Current_Temperature": 62, "Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation."}' + +def test_transform_response_with_structured_response_calling_tool(): + """Test response transformation with structured response.""" + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.types.utils import ModelResponse + + # Simulate a Bedrock Converse response with a bash tool call + response_json = { + "metrics": { + "latencyMs": 1148 + }, + "output": { + "message": + { + "content": [ + { + "text": "I\'ll check the current weather in San Francisco for you." + }, + { + "toolUse": { + "input": { + "location": "San Francisco, CA", + "unit": "celsius" + }, + "name": "get_weather", + "toolUseId": "tooluse_oKk__QrqSUmufMw3Q7vGaQ" + } + } + ], + "role": "assistant" + } + }, + "stopReason": "tool_use", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 534, + "outputTokens": 69, + "totalTokens": 603 + } + } + # Mock httpx.Response + 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, + "tools": [ + { + 'type': 'function', + 'function': { + 'name': 'get_weather', + 'description': 'Get the current weather in a given location', + 'parameters': { + 'type': 'object', + 'properties': { + 'location': { + 'type': 'string', + 'description': 'The city and state, e.g. San Francisco, CA' + }, + 'unit': { + 'type': 'string', + 'enum': ['celsius', 'fahrenheit'] + } + }, + 'required': ['location'] + } + } + }, + { + 'type': 'function', + 'function': { + 'name': 'json_tool_call', + 'parameters': { + '$schema': 'http://json-schema.org/draft-07/schema#', + 'type': 'object', + 'required': ['Weather_Explanation', 'Current_Temperature'], + 'properties': { + 'Weather_Explanation': { + 'type': ['string', 'null'], + 'description': '1-2 sentences explaining the weather in the location' + }, + 'Current_Temperature': { + 'type': ['number', 'null'], + 'description': 'Current temperature in the location' + } + }, + 'additionalProperties': False + } + } + } + ] + } + # Call the transformation logic + result = config._transform_response( + model="bedrock/eu.anthropic.claude-sonnet-4-20250514-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, + ) + # Check that the tool call is present in the returned message + 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": "celsius"}' + + @pytest.mark.asyncio async def test_bedrock_bash_tool_acompletion(): """Test Bedrock with bash tool for ls command using acompletion.""" @@ -938,6 +1171,68 @@ def test_transform_request_with_function_tool(): assert request_data["toolConfig"]["tools"][0]["toolSpec"]["name"] == "get_weather" +def test_map_openai_params_with_response_format(): + """Test map_openai_params with response_format.""" + config = AmazonConverseConfig() + + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + }, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["location"], + }, + } + } + ] + + json_schema = { + "type": "json_schema", + "json_schema": { + "name": "WeatherResult", + "schema": { + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "required": ["Weather_Explanation", "Current_Temperature"], + "properties": { + "Weather_Explanation": { + "type": ["string", "null"], + "description": "1-2 sentences explaining the weather in the location", + }, + "Current_Temperature": { + "type": ["number", "null"], + "description": "Current temperature in the location", + }, + }, + "additionalProperties": False, + }, + "strict": False, + }, + } + + optional_params = config.map_openai_params( + non_default_params={"response_format": json_schema}, + optional_params={"tools": tools}, + model="eu.anthropic.claude-sonnet-4-20250514-v1:0", + drop_params=False + ) + + assert "tools" in optional_params + assert len(optional_params["tools"]) == 2 + assert optional_params["tools"][1]["type"] == "function" + assert optional_params["tools"][1]["function"]["name"] == "json_tool_call" + + @pytest.mark.asyncio async def test_assistant_message_cache_control(): """Test that assistant messages with cache_control generate cachePoint blocks.""" diff --git a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py new file mode 100644 index 00000000000..7cb1ee2b54a --- /dev/null +++ b/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -0,0 +1,177 @@ +import os +import sys +from unittest.mock import patch + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig + + +def test_bedrock_passthrough_get_complete_url_default_endpoint(): + """Test get_complete_url with default AWS endpoint (no override)""" + config = BedrockPassthroughConfig() + + # Mock the methods following the pattern from test_base_aws_llm.py + with patch.object(config, '_get_aws_region_name', return_value="us-east-1"), \ + patch.object(config, 'get_runtime_endpoint', return_value=( + "https://bedrock-runtime.us-east-1.amazonaws.com", + "https://bedrock-runtime.us-east-1.amazonaws.com" + )) as mock_get_runtime: + + url, api_base = config.get_complete_url( + api_base=None, + api_key=None, + model="anthropic.claude-3-sonnet", + endpoint="/model/anthropic.claude-3-sonnet/invoke", + request_query_params=None, + litellm_params={} + ) + + # Verify get_runtime_endpoint was called with correct parameters + mock_get_runtime.assert_called_once_with( + api_base=None, + aws_bedrock_runtime_endpoint=None, + aws_region_name="us-east-1", + endpoint_type="runtime" + ) + + # Verify URL construction + assert str(url) == "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-sonnet/invoke" + assert api_base == "https://bedrock-runtime.us-east-1.amazonaws.com" + + +def test_bedrock_passthrough_get_complete_url_custom_endpoint_no_path(): + """Test get_complete_url with custom endpoint (no base path)""" + config = BedrockPassthroughConfig() + + with patch.object(config, '_get_aws_region_name', return_value="us-west-2"), \ + patch.object(config, 'get_runtime_endpoint', return_value=( + "http://proxy.com", + "http://proxy.com" + )) as mock_get_runtime: + + url, api_base = config.get_complete_url( + api_base="http://proxy.com", + api_key=None, + model="anthropic.claude-3-sonnet", + endpoint="/model/anthropic.claude-3-sonnet/invoke", + request_query_params=None, + litellm_params={} + ) + + # Verify get_runtime_endpoint was called with the api_base + mock_get_runtime.assert_called_once_with( + api_base="http://proxy.com", + aws_bedrock_runtime_endpoint=None, + aws_region_name="us-west-2", + endpoint_type="runtime" + ) + + # Verify URL construction + assert str(url) == "http://proxy.com/model/anthropic.claude-3-sonnet/invoke" + assert api_base == "http://proxy.com" + + +def test_bedrock_passthrough_get_complete_url_custom_endpoint_with_path(): + """Test get_complete_url with custom endpoint that has a base path""" + config = BedrockPassthroughConfig() + + with patch.object(config, '_get_aws_region_name', return_value="us-west-2"), \ + patch.object(config, 'get_runtime_endpoint', return_value=( + "http://proxy.com/bedrockproxy", + "http://proxy.com/bedrockproxy" + )) as mock_get_runtime: + + url, api_base = config.get_complete_url( + api_base="http://proxy.com/bedrockproxy", + api_key=None, + model="anthropic.claude-3-sonnet", + endpoint="/model/anthropic.claude-3-sonnet/invoke", + request_query_params=None, + litellm_params={ + "aws_bedrock_runtime_endpoint": "http://proxy.com/bedrockproxy" + } + ) + + # Verify get_runtime_endpoint was called with correct parameters + mock_get_runtime.assert_called_once_with( + api_base="http://proxy.com/bedrockproxy", + aws_bedrock_runtime_endpoint="http://proxy.com/bedrockproxy", + aws_region_name="us-west-2", + endpoint_type="runtime" + ) + + # Verify URL construction preserves the proxy path + assert str(url) == "http://proxy.com/bedrockproxy/model/anthropic.claude-3-sonnet/invoke" + assert api_base == "http://proxy.com/bedrockproxy" + + +def test_format_url_simple_joining(): + """Test format_url with simple URL joining""" + config = BedrockPassthroughConfig() + + result = config.format_url( + endpoint="model/test/invoke", + base_target_url="https://api.example.com", + request_query_params={} + ) + + assert str(result) == "https://api.example.com/model/test/invoke" + + +def test_format_url_preserves_proxy_paths(): + """Test format_url preserves proxy paths in base URL""" + config = BedrockPassthroughConfig() + + result = config.format_url( + endpoint="model/test/invoke", + base_target_url="http://proxy.com/bedrockproxy", + request_query_params={} + ) + + # This is the key test - proxy path should be preserved + assert str(result) == "http://proxy.com/bedrockproxy/model/test/invoke" + + +def test_format_url_with_query_parameters(): + """Test format_url properly handles query parameters""" + config = BedrockPassthroughConfig() + + result = config.format_url( + endpoint="model/test/invoke", + base_target_url="http://proxy.com/bedrockproxy", + request_query_params={"param1": "value1", "param2": "value2"} + ) + + # Should preserve proxy path and add query params + result_str = str(result) + assert "http://proxy.com/bedrockproxy/model/test/invoke" in result_str + assert "param1=value1" in result_str + assert "param2=value2" in result_str + + +def test_format_url_handles_trailing_slash_normalization(): + """Test format_url properly handles base URLs with and without trailing slashes""" + config = BedrockPassthroughConfig() + + # Test with trailing slash + result_with_slash = config.format_url( + endpoint="model/test/invoke", + base_target_url="http://proxy.com/bedrockproxy/", + request_query_params={} + ) + + # Test without trailing slash + result_without_slash = config.format_url( + endpoint="model/test/invoke", + base_target_url="http://proxy.com/bedrockproxy", + request_query_params={} + ) + + # Both should produce the same result + assert str(result_with_slash) == str(result_without_slash) + assert str(result_with_slash) == "http://proxy.com/bedrockproxy/model/test/invoke" + + diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index fc44d44aba9..51a2e971c09 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -10,7 +10,10 @@ sys.path.insert( ) # Adds the parent directory to the system path from unittest.mock import MagicMock, patch -from litellm.llms.databricks.chat.transformation import DatabricksConfig +from litellm.llms.databricks.chat.transformation import ( + DatabricksChatResponseIterator, + DatabricksConfig, +) def test_transform_choices(): @@ -85,8 +88,101 @@ def test_transform_choices_without_signature(): assert choices[0].message.reasoning_content == "i'm thinking without signature." assert choices[0].message.thinking_blocks is not None assert len(choices[0].message.thinking_blocks) == 1 - + # Verify the thinking block was created successfully without signature thinking_block = choices[0].message.thinking_blocks[0] assert thinking_block["type"] == "thinking" assert thinking_block["thinking"] == "i'm thinking without signature." + + +def test_transform_choices_with_citations(): + config = DatabricksConfig() + databricks_choices = [ + { + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Blue", + "citations": [ + { + "type": "char_location", + "cited_text": "The sky is blue.", + "document_index": 0, + "document_title": "My Document", + "start_char_index": 0, + "end_char_index": 50, + } + ], + } + ], + }, + "index": 0, + "finish_reason": "stop", + } + ] + + choices = config._transform_dbrx_choices(choices=databricks_choices) + + assert choices[0].message.provider_specific_fields == { + "citations": [ + [ + { + "type": "char_location", + "cited_text": "The sky is blue.", + "document_index": 0, + "document_title": "My Document", + "start_char_index": 0, + "end_char_index": 50, + "supported_text": "Blue", + } + ] + ] + } + + +def test_chunk_parser_with_citation(): + iterator = DatabricksChatResponseIterator(None, sync_stream=True) + chunk = { + "id": "1", + "object": "chat.completion.chunk", + "created": 0, + "model": "test", + "choices": [ + { + "delta": { + "content": [ + { + "type": "text", + "text": "", + "citations": [ + { + "type": "char_location", + "cited_text": "The sky is blue.", + "document_index": 0, + "document_title": "My Document", + "start_char_index": 0, + "end_char_index": 50, + } + ], + } + ], + }, + "index": 0, + "finish_reason": None, + } + ], + } + + parsed = iterator.chunk_parser(chunk) + assert parsed.choices[0].delta.provider_specific_fields == { + "citation": { + "type": "char_location", + "cited_text": "The sky is blue.", + "document_index": 0, + "document_title": "My Document", + "start_char_index": 0, + "end_char_index": 50, + } + } diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation_for_14158.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation_for_14158.py new file mode 100644 index 00000000000..950c1fcb4c9 --- /dev/null +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation_for_14158.py @@ -0,0 +1,229 @@ +import pytest +from litellm.llms.oci.chat.transformation import adapt_messages_to_generic_oci_standard + +def test_adapt_messages_with_empty_content_and_tool_calls(): + """Test that assistant messages with empty content and tool_calls are processed correctly.""" + # Arrange + messages_with_empty_content = [ + {"role": "user", "content": "Tell me the weather in Tokyo."}, + { + "role": "assistant", + "content": "", # Empty string + "tool_calls": [ + { + "id": "call_test_empty", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Tokyo"}' + } + } + ] + }, + { + "role": "tool", + "content": '{"weather": "Sunny", "temperature": "25°C"}', + "tool_call_id": "call_test_empty" + } + ] + + # Act + result = adapt_messages_to_generic_oci_standard(messages_with_empty_content) + + # Assert + assert len(result) == 3 + + # Check user message + assert result[0].role == "USER" + assert result[0].content[0].type == "TEXT" + assert result[0].content[0].text == "Tell me the weather in Tokyo." + + # Check assistant message with tool_calls (should prioritize tool_calls over empty content) + assert result[1].role == "ASSISTANT" + assert result[1].toolCalls is not None + assert len(result[1].toolCalls) == 1 + assert result[1].toolCalls[0].id == "call_test_empty" + assert result[1].toolCalls[0].name == "get_weather" + + # Check tool response message + assert result[2].role == "TOOL" # Tool responses have TOOL role, not USER + assert result[2].content[0].type == "TEXT" + assert "weather" in result[2].content[0].text + assert result[2].toolCallId == "call_test_empty" # Tool call ID is in separate field + +def test_adapt_messages_with_none_content_and_tool_calls(): + """Test that assistant messages with None content and tool_calls are processed correctly.""" + # Arrange + messages_with_none_content = [ + {"role": "user", "content": "Tell me the weather in Tokyo."}, + { + "role": "assistant", + "content": None, # None value + "tool_calls": [ + { + "id": "call_test_none", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Tokyo"}' + } + } + ] + }, + { + "role": "tool", + "content": '{"weather": "Sunny", "temperature": "25°C"}', + "tool_call_id": "call_test_none" + } + ] + + # Act + result = adapt_messages_to_generic_oci_standard(messages_with_none_content) + + # Assert + assert len(result) == 3 + + # Check assistant message prioritizes tool_calls over None content + assert result[1].role == "ASSISTANT" + assert result[1].toolCalls is not None + assert len(result[1].toolCalls) == 1 + assert result[1].toolCalls[0].id == "call_test_none" + +def test_adapt_messages_with_tool_calls_only(): + """Test that assistant messages with only tool_calls (no content field) are processed correctly.""" + # Arrange + messages_no_content = [ + {"role": "user", "content": "Tell me the weather in Tokyo."}, + { + "role": "assistant", + # No content field at all + "tool_calls": [ + { + "id": "call_test_no_content", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Tokyo"}' + } + } + ] + }, + { + "role": "tool", + "content": '{"weather": "Sunny", "temperature": "25°C"}', + "tool_call_id": "call_test_no_content" + } + ] + + # Act + result = adapt_messages_to_generic_oci_standard(messages_no_content) + + # Assert + assert len(result) == 3 + + # Check assistant message processes tool_calls correctly + assert result[1].role == "ASSISTANT" + assert result[1].toolCalls is not None + assert len(result[1].toolCalls) == 1 + assert result[1].toolCalls[0].id == "call_test_no_content" + +def test_adapt_messages_with_content_only(): + """Test that assistant messages with only content (no tool_calls) are processed correctly.""" + # Arrange + messages_content_only = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": "Hello! How can I help you today?" + } + ] + + # Act + result = adapt_messages_to_generic_oci_standard(messages_content_only) + + # Assert + assert len(result) == 2 + + # Check assistant message with content only + assert result[1].role == "ASSISTANT" + assert result[1].content[0].type == "TEXT" + assert result[1].content[0].text == "Hello! How can I help you today?" + assert result[1].toolCalls is None + +def test_adapt_messages_tool_id_tracking(): + """Test that tool call IDs are properly tracked for validation.""" + # Arrange + messages = [ + {"role": "user", "content": "Test"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "test_func", + "arguments": '{"param": "value"}' + } + } + ] + }, + { + "role": "tool", + "content": "Result", + "tool_call_id": "call_123" + } + ] + + # Act + result = adapt_messages_to_generic_oci_standard(messages) + + # Assert + # Tool call should be processed and ID should be available for validation + assert result[1].toolCalls[0].id == "call_123" + + # Tool response should reference the same ID + tool_response_text = result[2].content[0].text + # Tool response text is just the content, tool_call_id is separate + assert tool_response_text == "Result" # The actual content + assert result[2].toolCallId == "call_123" # Tool call ID is in separate field + +def test_adapt_messages_multiple_tool_calls(): + """Test that multiple tool calls in a single message are processed correctly.""" + # Arrange + messages = [ + {"role": "user", "content": "Test multiple tools"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "func1", + "arguments": '{"param": "value1"}' + } + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "func2", + "arguments": '{"param": "value2"}' + } + } + ] + } + ] + + # Act + result = adapt_messages_to_generic_oci_standard(messages) + + # Assert + assert len(result) == 2 + assert result[1].role == "ASSISTANT" + assert len(result[1].toolCalls) == 2 + assert result[1].toolCalls[0].id == "call_1" + assert result[1].toolCalls[1].id == "call_2" + diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py index 985d51f99da..452f5a94024 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -159,6 +159,261 @@ class TestOllamaConfig: assert result.choices[0]["finish_reason"] == "stop" # No usage assertions here as we don't need to test them in every case + def test_transform_response_with_thinking_tags(self): + """Test that responses with ... tags parse reasoning content correctly.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with thinking tags + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "I need to think about this problem step by stepHere is my answer", + "prompt_eval_count": 15, + "eval_count": 8, + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify reasoning content is extracted + assert ( + result.choices[0]["message"].reasoning_content + == "I need to think about this problem step by step" + ) + assert result.choices[0]["message"].content == "Here is my answer" + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_with_thinking_tags_alternative(self): + """Test that responses with ... tags parse reasoning content correctly.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with thinking tags (alternative format) + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Let me analyze this carefullyThe solution is X", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify reasoning content is extracted + assert ( + result.choices[0]["message"].reasoning_content + == "Let me analyze this carefully" + ) + assert result.choices[0]["message"].content == "The solution is X" + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_with_multiline_thinking_tags(self): + """Test that responses with multiline thinking content work correctly.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with multiline thinking content + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "\nThis is a complex problem.\nI need to break it down:\n1. First step\n2. Second step\nBased on my analysis, the answer is Y", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify multiline reasoning content is extracted + expected_reasoning = "\nThis is a complex problem.\nI need to break it down:\n1. First step\n2. Second step\n" + assert result.choices[0]["message"].reasoning_content == expected_reasoning + assert ( + result.choices[0]["message"].content + == "Based on my analysis, the answer is Y" + ) + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_thinking_only(self): + """Test response with only thinking content and no additional content.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with only thinking content + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Just internal thoughts, no response", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify reasoning content is extracted and content is empty + assert ( + result.choices[0]["message"].reasoning_content + == "Just internal thoughts, no response" + ) + assert result.choices[0]["message"].content == "" + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_json_mode_with_thinking_tags(self): + """Test JSON mode with thinking tags - should handle as text when JSON parsing fails.""" + # Initialize config + config = OllamaConfig() + + # Create mock response with thinking tags in JSON mode + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Planning my JSON responseThis is not valid JSON", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={"format": "json"}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify reasoning content is extracted even in JSON mode when JSON parsing fails + assert ( + result.choices[0]["message"].reasoning_content + == "Planning my JSON response" + ) + assert result.choices[0]["message"].content == "This is not valid JSON" + assert result.choices[0]["finish_reason"] == "stop" + + def test_transform_response_no_thinking_tags(self): + """Test that responses without thinking tags work normally.""" + # Initialize config + config = OllamaConfig() + + # Create mock response without thinking tags + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Regular response without any thinking tags", + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify no reasoning content is extracted + assert result.choices[0]["message"].reasoning_content is None + assert ( + result.choices[0]["message"].content + == "Regular response without any thinking tags" + ) + assert result.choices[0]["finish_reason"] == "stop" + class TestOllamaTextCompletionResponseIterator: def test_chunk_parser_with_thinking_field(self): @@ -199,10 +454,11 @@ class TestOllamaTextCompletionResponseIterator: result = iterator.chunk_parser(normal_chunk) - assert result["text"] == "Hello world" - assert result["is_finished"] is False - assert result["finish_reason"] == "stop" - assert result["usage"] is None + # Updated to handle ModelResponseStream return type + assert isinstance(result, ModelResponseStream) + assert result.choices and result.choices[0].delta is not None + assert result.choices[0].delta.content == "Hello world" + assert getattr(result.choices[0].delta, "reasoning_content", None) is None def test_chunk_parser_done_chunk(self): """Test that done chunks work correctly.""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index 6d46a40f6c6..a6a34518098 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -667,4 +667,4 @@ def test_get_supported_openai_params(): assert "temperature" in params assert "stream" in params assert "background" in params - assert "stream" in params + assert "stream" in params \ No newline at end of file diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index 3bdab355977..3e6a6a23468 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -41,3 +41,14 @@ def test_gpt5_temperature_error(config: OpenAIConfig): model="gpt-5", drop_params=False, ) + + +def test_gpt5_unsupported_params_drop(config: OpenAIConfig): + assert "top_p" not in config.get_supported_openai_params(model="gpt-5") + params = config.map_openai_params( + non_default_params={"top_p": 0.5}, + optional_params={}, + model="gpt-5", + drop_params=True, + ) + assert "top_p" not in params diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py new file mode 100644 index 00000000000..d6d33258576 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -0,0 +1,75 @@ +from litellm.llms.vertex_ai.gemini.transformation import check_if_part_exists_in_parts + + +def test_check_if_part_exists_in_parts(): + parts = [ + {"text": "Hello", "thought": True}, + {"text": "World", "thought": False}, + ] + part = {"text": "Hello", "thought": True} + new_part = {"text": "Hello World", "thought": True} + assert check_if_part_exists_in_parts(parts, part) + assert not check_if_part_exists_in_parts(parts, new_part, ["thought"]) + assert check_if_part_exists_in_parts(parts, new_part, ["text"]) + + +def test_check_if_part_exists_in_parts_camel_case_snake_case(): + """Test that function handles both camelCase and snake_case key variations""" + # Test snake_case to camelCase matching + parts_with_snake_case = [ + { + "function_call": { + "name": "get_current_weather", + "args": {"location": "San Francisco, CA"}, + } + }, + {"text": "Some other content"}, + ] + + part_with_camel_case = { + "functionCall": { + "name": "get_current_weather", + "args": {"location": "San Francisco, CA"}, + } + } + + # Should find match between function_call and functionCall + assert check_if_part_exists_in_parts(parts_with_snake_case, part_with_camel_case) + + # Test camelCase to snake_case matching + parts_with_camel_case = [ + {"functionCall": {"name": "calculate_sum", "args": {"a": 1, "b": 2}}} + ] + + part_with_snake_case = { + "function_call": {"name": "calculate_sum", "args": {"a": 1, "b": 2}} + } + + # Should find match between functionCall and function_call + assert check_if_part_exists_in_parts(parts_with_camel_case, part_with_snake_case) + + # Test no match when values differ + part_with_different_values = { + "function_call": {"name": "different_function", "args": {"x": 5}} + } + + assert not check_if_part_exists_in_parts( + parts_with_snake_case, part_with_different_values + ) + + # Test multiple keys with mixed casing + parts_mixed = [ + { + "function_call": {"name": "test"}, + "thoughtSignature": "reasoning", + "text": "content", + } + ] + + part_mixed_casing = { + "functionCall": {"name": "test"}, + "thought_signature": "reasoning", + "text": "content", + } + + assert check_if_part_exists_in_parts(parts_mixed, part_mixed_casing) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py new file mode 100644 index 00000000000..34046a00ee8 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py @@ -0,0 +1,238 @@ +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.llms.vertex_ai.vertex_ai_partner_models.gpt_oss.transformation import ( + VertexAIGPTOSSTransformation, +) + + +class TestVertexAIGPTOSSTransformation: + """Test class for VertexAI GPT-OSS transformation functionality.""" + + def test_supports_reasoning_effort(self): + """Test that reasoning_effort parameter is supported for GPT-OSS models.""" + config = VertexAIGPTOSSTransformation() + supported_params = config.get_supported_openai_params(model="openai/gpt-oss-20b-maas") + + assert "reasoning_effort" in supported_params + + def test_removes_tool_calling_params_when_not_supported(self): + """Test that tool calling parameters are removed when function calling is not supported.""" + config = VertexAIGPTOSSTransformation() + + # Mock litellm.supports_function_calling to return False + with patch('litellm.supports_function_calling', return_value=False): + supported_params = config.get_supported_openai_params(model="openai/gpt-oss-20b-maas") + + # Tool calling params should be removed + assert "tool" not in supported_params + assert "tool_choice" not in supported_params + assert "function_call" not in supported_params + assert "functions" not in supported_params + + # But reasoning_effort should still be there + assert "reasoning_effort" in supported_params + + +@pytest.mark.asyncio +async def test_vertex_ai_gpt_oss_simple_request(): + """ + Test that a simple request to vertex_ai/openai/gpt-oss-20b-maas lands at the correct URL + with the correct request body. + """ + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexLLM, + ) + + # Mock response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = { + "id": "chatcmpl-test123", + "object": "chat.completion", + "created": 1234567890, + "model": "openai/gpt-oss-20b-maas", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! I'm Litellm Bot, a helpful assistant. I don't have access to real-time weather information, but I'd be happy to help you with other questions or tasks!" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 42, + "completion_tokens": 28, + "total_tokens": 70 + } + } + + client = AsyncHTTPHandler() + + async def mock_post_func(*args, **kwargs): + return mock_response + + with patch.object(client, "post", side_effect=mock_post_func) as mock_post, \ + patch.object(VertexLLM, "_ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718")): + response = await litellm.acompletion( + model="vertex_ai/openai/gpt-oss-20b-maas", + messages=[ + { + "role": "system", + "content": "Your name is Litellm Bot, you are a helpful assistant" + }, + { + "role": "user", + "content": "Hello, what is your name and can you tell me the weather?" + } + ], + vertex_ai_location="us-central1", + vertex_ai_project="pathrise-convert-1606954137718", + client=client + ) + + # Verify the mock was called + mock_post.assert_called_once() + + # Get the call arguments + call_args = mock_post.call_args + # For side_effect, the URL is passed as kwargs['url'] + called_url = call_args.kwargs["url"] + request_body = json.loads(call_args.kwargs["data"]) + + # Verify the URL + expected_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/pathrise-convert-1606954137718/locations/us-central1/endpoints/openapi/chat/completions" + assert called_url == expected_url + + # Verify the request body + expected_request_body = { + 'model': 'openai/gpt-oss-20b-maas', + 'messages': [ + { + 'role': 'system', + 'content': 'Your name is Litellm Bot, you are a helpful assistant' + }, + { + 'role': 'user', + 'content': 'Hello, what is your name and can you tell me the weather?' + } + ], + 'stream': False + } + assert request_body == expected_request_body + + # Verify response structure + assert response.model == "openai/gpt-oss-20b-maas" + assert len(response.choices) == 1 + assert response.choices[0].message.role == "assistant" + + +@pytest.mark.asyncio +async def test_vertex_ai_gpt_oss_reasoning_effort(): + """ + Test that reasoning_effort parameter is correctly passed in the request body + for GPT-OSS models. + """ + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexLLM, + ) + + # Mock response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = { + "id": "chatcmpl-test456", + "object": "chat.completion", + "created": 1234567890, + "model": "openai/gpt-oss-20b-maas", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "I need to think about this carefully. The weather varies by location and time, so I would need to know your specific location to provide accurate weather information." + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 35, + "completion_tokens": 32, + "total_tokens": 67 + } + } + + client = AsyncHTTPHandler() + + async def mock_post_func(*args, **kwargs): + return mock_response + + with patch.object(client, "post", side_effect=mock_post_func) as mock_post, \ + patch.object(VertexLLM, "_ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718")): + response = await litellm.acompletion( + model="vertex_ai/openai/gpt-oss-20b-maas", + messages=[ + { + "role": "system", + "content": "Your name is Litellm Bot, you are a helpful assistant" + }, + { + "role": "user", + "content": "Hello, what is your name and can you tell me the weather?" + } + ], + reasoning_effort="low", + vertex_ai_location="us-central1", + vertex_ai_project="pathrise-convert-1606954137718", + client=client + ) + + # Verify the mock was called + mock_post.assert_called_once() + + # Get the call arguments + call_args = mock_post.call_args + request_body = json.loads(call_args.kwargs["data"]) + + # Verify reasoning_effort is in the request body + assert "reasoning_effort" in request_body + assert request_body["reasoning_effort"] == "low" + + # Verify other expected fields + expected_request_body = { + 'model': 'openai/gpt-oss-20b-maas', + 'messages': [ + { + 'role': 'system', + 'content': 'Your name is Litellm Bot, you are a helpful assistant' + }, + { + 'role': 'user', + 'content': 'Hello, what is your name and can you tell me the weather?' + } + ], + 'reasoning_effort': 'low', + 'stream': False + } + assert request_body == expected_request_body + + # Verify response structure + assert response.model == "openai/gpt-oss-20b-maas" + assert len(response.choices) == 1 + assert response.choices[0].message.role == "assistant" diff --git a/tests/test_litellm/llms/volcengine/__init__.py b/tests/test_litellm/llms/volcengine/__init__.py new file mode 100644 index 00000000000..6ac3aa6b71a --- /dev/null +++ b/tests/test_litellm/llms/volcengine/__init__.py @@ -0,0 +1 @@ +# Volcengine tests \ No newline at end of file diff --git a/tests/test_litellm/llms/volcengine/embedding/__init__.py b/tests/test_litellm/llms/volcengine/embedding/__init__.py new file mode 100644 index 00000000000..bb087ba3563 --- /dev/null +++ b/tests/test_litellm/llms/volcengine/embedding/__init__.py @@ -0,0 +1 @@ +# Volcengine embedding tests \ No newline at end of file diff --git a/tests/test_litellm/llms/test_volcengine.py b/tests/test_litellm/llms/volcengine/test_volcengine.py similarity index 97% rename from tests/test_litellm/llms/test_volcengine.py rename to tests/test_litellm/llms/volcengine/test_volcengine.py index 9db91217c28..59317914192 100644 --- a/tests/test_litellm/llms/test_volcengine.py +++ b/tests/test_litellm/llms/volcengine/test_volcengine.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch from pydantic import BaseModel -from litellm.llms.volcengine import VolcEngineConfig +from litellm.llms.volcengine.chat.transformation import VolcEngineChatConfig as VolcEngineConfig from litellm.utils import get_optional_params diff --git a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py new file mode 100644 index 00000000000..3be7f6ca8d4 --- /dev/null +++ b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py @@ -0,0 +1,262 @@ +""" +Integration tests for Volcengine embedding following LiteLLM testing patterns +Based on the BaseLLMEmbeddingTest framework +""" + +import os +import sys +from unittest.mock import MagicMock, patch +import pytest + +# Add parent directory to path for imports +sys.path.insert(0, os.path.abspath("../../../../..")) + +from tests.llm_translation.base_embedding_unit_tests import BaseLLMEmbeddingTest +import litellm +from litellm.types.utils import EmbeddingResponse + + +class TestVolcEngineEmbedding(BaseLLMEmbeddingTest): + """Test Volcengine embedding integration following LiteLLM patterns""" + + def get_custom_llm_provider(self) -> litellm.LlmProviders: + return litellm.LlmProviders.VOLCENGINE + + def get_base_embedding_call_args(self) -> dict: + return { + "model": "volcengine/doubao-embedding-text-240715", + } + + @pytest.mark.asyncio() + @pytest.mark.parametrize("sync_mode", [True, False]) + async def test_basic_embedding(self, sync_mode): + """Test basic embedding functionality with realistic response""" + litellm.set_verbose = True + embedding_call_args = self.get_base_embedding_call_args() + + # Mock the embedding functions to avoid actual API calls + with patch("litellm.embedding") as mock_embedding, patch("litellm.aembedding") as mock_aembedding: + # Create realistic Volcengine response + mock_response = MagicMock() + mock_response.model = "doubao-embedding-text-240715" + mock_response.object = "list" + mock_response.data = [ + { + "object": "embedding", + "embedding": [0.1, 0.2, 0.3] + [0.01 * i for i in range(1021)], # 1024-dim embedding + "index": 0 + }, + { + "object": "embedding", + "embedding": [0.4, 0.5, 0.6] + [0.02 * i for i in range(1021)], # 1024-dim embedding + "index": 1 + } + ] + mock_response.usage.prompt_tokens = 2 + mock_response.usage.total_tokens = 2 + + mock_embedding.return_value = mock_response + mock_aembedding.return_value = mock_response + + # Test sync mode + if sync_mode is True: + response = litellm.embedding( + **embedding_call_args, + input=["hello", "world"], + ) + + # Verify response structure matches Volcengine format + assert response.model == "doubao-embedding-text-240715" + assert response.object == "list" + assert len(response.data) == 2 + assert len(response.data[0]["embedding"]) == 1024 + assert response.usage.total_tokens > 0 + + # Test async mode + else: + response = await litellm.aembedding( + **embedding_call_args, + input=["hello", "world"], + ) + + # Verify response structure + assert response.model == "doubao-embedding-text-240715" + assert response.object == "list" + assert len(response.data) == 2 + assert len(response.data[0]["embedding"]) == 1024 + assert response.usage.total_tokens > 0 + + +def test_volcengine_embedding_with_encoding_formats(): + """Test Volcengine embedding with different encoding formats""" + + test_cases = [ + {"encoding_format": "float"}, + {"encoding_format": "base64"}, + {"encoding_format": None}, # Default + ] + + for params in test_cases: + with patch("litellm.embedding") as mock_embedding: + # Create mock response based on encoding format + mock_response = MagicMock() + mock_response.model = "doubao-embedding-text-240715" + mock_response.object = "list" + + if params["encoding_format"] == "base64": + # Simulate base64 encoded embeddings + mock_response.data = [ + { + "object": "embedding", + "embedding": "c29tZS1iYXNlNjQtZW5jb2RlZC1lbWJlZGRpbmc=", # base64 encoded + "index": 0 + } + ] + else: + # Float embeddings (default) + mock_response.data = [ + { + "object": "embedding", + "embedding": [0.1, 0.2, 0.3, -0.1] * 256, # 1024 dimensions + "index": 0 + } + ] + + mock_response.usage.prompt_tokens = 3 + mock_response.usage.total_tokens = 3 + mock_embedding.return_value = mock_response + + # Test the call + litellm.embedding( + model="volcengine/doubao-embedding-text-240715", + input=["test text"], + **params + ) + + # Verify the call was made with correct parameters + mock_embedding.assert_called_once() + call_args = mock_embedding.call_args + assert call_args[1]["model"] == "volcengine/doubao-embedding-text-240715" + assert call_args[1]["input"] == ["test text"] + + if params["encoding_format"] is not None: + assert call_args[1]["encoding_format"] == params["encoding_format"] + + +def test_volcengine_embedding_with_user_parameter(): + """Test Volcengine embedding with user parameter for tracking""" + + with patch("litellm.embedding") as mock_embedding: + mock_response = MagicMock() + mock_response.model = "doubao-embedding-text-240715" + mock_response.object = "list" + mock_response.data = [ + { + "object": "embedding", + "embedding": [0.1] * 1024, + "index": 0 + } + ] + mock_response.usage.prompt_tokens = 5 + mock_response.usage.total_tokens = 5 + mock_embedding.return_value = mock_response + + # Test with user parameter + litellm.embedding( + model="volcengine/doubao-embedding-text-240715", + input=["user tracking test"], + user="test-user-12345" + ) + + # Verify user parameter was passed + mock_embedding.assert_called_once() + call_args = mock_embedding.call_args + assert call_args[1]["user"] == "test-user-12345" + + +def test_volcengine_embedding_error_scenarios(): + """Test Volcengine embedding error handling in integration context""" + + error_scenarios = [ + # Invalid model name + { + "model": "volcengine/invalid-model-name", + "expected_error_pattern": "model" + }, + # Invalid encoding format + { + "model": "volcengine/doubao-embedding-text-240715", + "encoding_format": "invalid_format", + "expected_error_pattern": "encoding_format" + } + ] + + for scenario in error_scenarios: + with patch("litellm.embedding") as mock_embedding: + # Configure mock to raise appropriate errors + if "invalid-model" in scenario.get("model", ""): + mock_embedding.side_effect = Exception("Model not found") + elif scenario.get("encoding_format") == "invalid_format": + mock_embedding.side_effect = ValueError("Unsupported encoding_format") + + # Test that errors are properly raised + with pytest.raises(Exception) as exc_info: + test_params = {k: v for k, v in scenario.items() if k != "expected_error_pattern"} + litellm.embedding( + input=["test"], + **test_params + ) + + # Verify error message contains expected pattern + assert scenario["expected_error_pattern"].lower() in str(exc_info.value).lower() + + +def test_volcengine_embedding_with_multiple_inputs(): + """Test Volcengine embedding with various input lengths and types""" + + test_inputs = [ + # Single short text + ["hello"], + # Multiple short texts + ["hello", "world", "test"], + # Mixed length texts + ["short", "This is a much longer text that should be handled properly by the embedding service"], + # Unicode content + ["ęµ‹čÆ•äø­ę–‡ę–‡ęœ¬", "Test English text", "ę··åˆčÆ­čØ€ mixed language"], + # Many inputs (batch processing) + [f"Test sentence number {i}" for i in range(10)] + ] + + for test_input in test_inputs: + with patch("litellm.embedding") as mock_embedding: + # Create proportional mock response + mock_response = MagicMock() + mock_response.model = "doubao-embedding-text-240715" + mock_response.object = "list" + mock_response.data = [ + { + "object": "embedding", + "embedding": [0.1 * (i + 1)] * 1024, # Unique embedding per input + "index": i + } + for i in range(len(test_input)) + ] + mock_response.usage.prompt_tokens = len(test_input) * 5 # Realistic token estimate + mock_response.usage.total_tokens = len(test_input) * 5 + mock_embedding.return_value = mock_response + + # Test the call + response = litellm.embedding( + model="volcengine/doubao-embedding-text-240715", + input=test_input + ) + + # Verify response matches input count + assert len(response.data) == len(test_input) + for i, embedding_data in enumerate(response.data): + assert embedding_data["index"] == i + assert len(embedding_data["embedding"]) == 1024 + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py new file mode 100644 index 00000000000..b9ed4b9b508 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -0,0 +1,29 @@ +import sys +import os + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.proxy.common_utils.callback_utils import ( + get_remaining_tokens_and_requests_from_request_data, +) + + +def test_get_remaining_tokens_and_requests_from_request_data(): + model_group = "openrouter/google/gemini-2.0-flash-001" + casedata = { + "metadata": { + "model_group": model_group, + f"litellm-key-remaining-requests-{model_group}": 100, + f"litellm-key-remaining-tokens-{model_group}": 200, + } + } + + headers = get_remaining_tokens_and_requests_from_request_data(casedata) + + expected_name = "openrouter-google-gemini-2.0-flash-001" + assert headers == { + f"x-litellm-key-remaining-requests-{expected_name}": 100, + f"x-litellm-key-remaining-tokens-{expected_name}": 200, + } diff --git a/tests/test_litellm/proxy/google_endpoints/__init__.py b/tests/test_litellm/proxy/google_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py new file mode 100644 index 00000000000..2f2538bf9aa --- /dev/null +++ b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py @@ -0,0 +1,49 @@ +""" +Test for google_endpoints/endpoints.py +""" +import pytest +import sys, os +from dotenv import load_dotenv + + +from litellm.proxy.google_endpoints.endpoints import google_count_tokens +from litellm.types.llms.vertex_ai import TokenCountDetailsResponse +from starlette.requests import Request + +load_dotenv() + +sys.path.insert( + 0, os.path.abspath("../../../..") +) + +@pytest.mark.asyncio +async def test_proxy_gemini_to_openai_like_model_token_counting(): + """ + Test the token counting endpoint for proxing gemini to openai-like models. + """ + response: TokenCountDetailsResponse = await google_count_tokens( + request=Request( + scope={ + "type": "http", + "parsed_body": ( + [ + "contents" + ], + { + "contents": [ + { + "parts": [ + { + "text": "Hello, how are you?" + } + ] + } + ] + } + ) + } + ), + model_name="volcengine/foo", + ) + + assert response.get("totalTokens") > 0 \ No newline at end of file diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 694a49159c0..da4218a9547 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -934,3 +934,204 @@ async def test_team_member_rate_limits_v3(): assert team_member_descriptor["value"] == f"{_team_id}:{_user_id}", "Team member value should combine team_id and user_id" assert team_member_descriptor["rate_limit"]["requests_per_unit"] == 10, "Team member RPM limit should be set" assert team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000, "Team member TPM limit should be set" + + +@pytest.mark.asyncio +async def test_async_increment_tokens_with_ttl_preservation(): + """ + Test TTL preservation functionality for token increment operations. + + This test verifies that: + 1. Keys are created with proper TTL on first increment + 2. TTL is preserved on subsequent increments (not reset) + 3. Both TTL and non-TTL operations work correctly in the same call + + Environment variables required: + - REDIS_HOST: Redis server hostname + - REDIS_PORT: Redis server port + - REDIS_PASSWORD: Redis password (optional) + + Test scenario: + 1. First call: Create keys with TTL=60s and TTL=None + 2. Wait 2 seconds + 3. Second call: Increment same keys + 4. Verify TTL decreased but wasn't reset to 60s + """ + import os + import time + from litellm.caching.redis_cache import RedisCache + from litellm.types.caching import RedisPipelineIncrementOperation + + # Skip test if Redis environment variables are not set + redis_host = os.getenv("REDIS_HOST") + redis_port = os.getenv("REDIS_PORT") + redis_password = os.getenv("REDIS_PASSWORD") + + if not redis_host or not redis_port: + pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set") + + # Setup Redis cache + redis_cache = RedisCache( + host=redis_host, + port=int(redis_port), + password=redis_password, + ) + + local_cache = DualCache(redis_cache=redis_cache) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + # Verify Redis connection is working + try: + await redis_cache.ping() + except Exception as e: + pytest.skip(f"Redis connection failed: {str(e)}") + + # Test keys + test_key_with_ttl = "test_ttl_preservation:with_ttl" + test_key_without_ttl = "test_ttl_preservation:without_ttl" + + try: + # Clean up any existing test keys + try: + await redis_cache.async_delete_cache(test_key_with_ttl) + await redis_cache.async_delete_cache(test_key_without_ttl) + except Exception: + # Keys might not exist, ignore cleanup errors + pass + + # First increment: Create operations with mixed TTL scenarios + pipeline_operations_first = [ + RedisPipelineIncrementOperation( + key=test_key_with_ttl, + increment_value=10.0, + ttl=60 + ), + RedisPipelineIncrementOperation( + key=test_key_without_ttl, + increment_value=5.0, + ttl=None # No TTL + ) + ] + + # Execute first increment + await parallel_request_handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations_first + ) + + # Verify keys exist and check initial TTL + ttl_after_first = await redis_cache.async_get_ttl(test_key_with_ttl) + value_after_first_with_ttl = await redis_cache.async_get_cache(test_key_with_ttl) + value_after_first_without_ttl = await redis_cache.async_get_cache(test_key_without_ttl) + + assert value_after_first_with_ttl == 10.0, "First increment should set value to 10.0" + assert value_after_first_without_ttl == 5.0, "First increment should set value to 5.0" + assert ttl_after_first is not None and ttl_after_first > 0, "Key with TTL should have positive TTL after first increment" + assert ttl_after_first <= 60, "TTL should not exceed the set value" + + # Check TTL for key without TTL (should be None, meaning no expiry) + ttl_no_ttl_key = await redis_cache.async_get_ttl(test_key_without_ttl) + assert ttl_no_ttl_key is None, "Key without TTL should have no expiry (None from async_get_ttl)" + + # Wait a moment to ensure TTL decreases + await asyncio.sleep(2) + + # Second increment: Same operations to test TTL preservation + pipeline_operations_second = [ + RedisPipelineIncrementOperation( + key=test_key_with_ttl, + increment_value=15.0, + ttl=60 # Same TTL value + ), + RedisPipelineIncrementOperation( + key=test_key_without_ttl, + increment_value=7.0, + ttl=None # No TTL + ) + ] + + # Execute second increment + await parallel_request_handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations_second + ) + + # Verify TTL preservation and value updates + ttl_after_second = await redis_cache.async_get_ttl(test_key_with_ttl) + value_after_second_with_ttl = await redis_cache.async_get_cache(test_key_with_ttl) + value_after_second_without_ttl = await redis_cache.async_get_cache(test_key_without_ttl) + + assert value_after_second_with_ttl == 25.0, "Second increment should update value to 25.0" + assert value_after_second_without_ttl == 12.0, "Second increment should update value to 12.0" + + # Critical test: TTL should be preserved (not reset to 60) + assert ttl_after_second is not None, "TTL should still exist" + assert ttl_after_second < ttl_after_first, "TTL should have decreased (not been reset)" + assert ttl_after_second > 0, "TTL should still be positive" + + # TTL should not be close to the original 60 seconds (proving it wasn't reset) + assert ttl_after_second < 59, "TTL should be significantly less than original, proving preservation" + + # Key without TTL should still have no expiry + ttl_no_ttl_key_after_second = await redis_cache.async_get_ttl(test_key_without_ttl) + assert ttl_no_ttl_key_after_second is None, "Key without TTL should still have no expiry" + + finally: + # Clean up test keys + try: + await redis_cache.async_delete_cache(test_key_with_ttl) + await redis_cache.async_delete_cache(test_key_without_ttl) + except Exception: + # Ignore cleanup errors + pass + + # Properly close Redis connections to prevent warnings + try: + await redis_cache.disconnect() + except Exception: + # Ignore disconnect errors + pass + + +@pytest.mark.asyncio +async def test_async_increment_tokens_fallback_behavior(): + """ + Test fallback behavior when Lua script is not available. + """ + from litellm.types.caching import RedisPipelineIncrementOperation + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + # Mock the token_increment_script to None to simulate unavailable script + parallel_request_handler.token_increment_script = None + + # Mock the fallback method + fallback_called = False + original_method = parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline + + async def mock_fallback(*args, **kwargs): + nonlocal fallback_called + fallback_called = True + return await original_method(*args, **kwargs) + + parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = mock_fallback + + # Test operations + pipeline_operations = [ + RedisPipelineIncrementOperation( + key="test_fallback_key", + increment_value=10.0, + ttl=60 + ) + ] + + # Execute increment + await parallel_request_handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations + ) + + # Verify fallback was called + assert fallback_called, "Fallback method should be called when Lua script is not available" diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py index f2b374657e6..be24444afaf 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py @@ -224,10 +224,10 @@ class TestScimTransformations: result = ScimTransformations._get_scim_member_value(member_with_email) assert result == member_with_email.user_email - # Member without email + # Member without email should fall back to user_id member_without_email = Member(user_id="user-456", user_email=None, role="user") result = ScimTransformations._get_scim_member_value(member_without_email) - assert result == ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE + assert result == member_without_email.user_id class TestSCIMPatchOperations: diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 208c8774675..959275787c8 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -10,10 +10,13 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( create_user, get_service_provider_config, patch_user, + update_group, update_user, ) from litellm.types.proxy.management_endpoints.scim_v2 import ( SCIMFeature, + SCIMGroup, + SCIMMember, SCIMPatchOp, SCIMPatchOperation, SCIMServiceProviderConfig, @@ -678,4 +681,233 @@ async def test_update_group_metadata_serialization_issue(mocker): parsed_metadata = json.loads(metadata) assert "existing_key" in parsed_metadata assert "scim_data" in parsed_metadata - assert parsed_metadata["existing_key"] == "existing_value" \ No newline at end of file + + +@pytest.mark.asyncio +async def test_team_membership_management(mocker): + """ + Test that team membership changes work correctly: + - Adding members to team + - Removing members from team + - members_with_roles is used as source of truth + """ + from litellm.proxy._types import Member + from litellm.proxy.management_endpoints.scim.scim_v2 import ( + _get_team_member_user_ids_from_team, + _handle_group_membership_changes, + patch_team_membership, + ) + + # Mock team with members_with_roles as source of truth + mock_team = mocker.MagicMock() + mock_team.members_with_roles = [ + Member(user_id="user1", role="user"), + Member(user_id="user2", role="user") + ] + mock_team.members = ["user1", "user2", "user3"] # This should be ignored + + # Test that members_with_roles is source of truth + member_ids = await _get_team_member_user_ids_from_team(mock_team) + assert set(member_ids) == {"user1", "user2"} + assert "user3" not in member_ids # Should not be included even though in members + + # Mock patch_team_membership function + mock_patch_team_membership = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock() + ) + + # Test adding and removing members + group_id = "test-group-id" + current_members = {"user1", "user2"} + final_members = {"user2", "user3", "user4"} # Remove user1, add user3 and user4 + + await _handle_group_membership_changes( + group_id=group_id, + current_members=current_members, + final_members=final_members + ) + + # Verify patch_team_membership was called correctly + assert mock_patch_team_membership.call_count == 3 + + # Check calls for adding members + add_calls = [call for call in mock_patch_team_membership.call_args_list + if call[1]["teams_ids_to_add_user_to"] == [group_id]] + assert len(add_calls) == 2 # user3 and user4 + + add_user_ids = {call[1]["user_id"] for call in add_calls} + assert add_user_ids == {"user3", "user4"} + + # Check calls for removing members + remove_calls = [call for call in mock_patch_team_membership.call_args_list + if call[1]["teams_ids_to_remove_user_from"] == [group_id]] + assert len(remove_calls) == 1 # user1 + + remove_user_ids = {call[1]["user_id"] for call in remove_calls} + assert remove_user_ids == {"user1"} + + # Verify all calls have correct structure + for call in mock_patch_team_membership.call_args_list: + assert "user_id" in call[1] + assert "teams_ids_to_add_user_to" in call[1] + assert "teams_ids_to_remove_user_from" in call[1] + # Each call should either add OR remove, not both + add_teams = call[1]["teams_ids_to_add_user_to"] + remove_teams = call[1]["teams_ids_to_remove_user_from"] + assert (len(add_teams) > 0) != (len(remove_teams) > 0) # XOR - one should be empty + + +@pytest.mark.asyncio +async def test_update_group_e2e(mocker): + """ + End-to-end test for update_group endpoint: + - Updates group metadata (displayName) + - Handles complete member replacement (add/remove members) + - Verifies members_with_roles is updated as source of truth + - Tests the full flow from SCIM request to database updates + """ + from litellm.proxy._types import LiteLLM_TeamTable, Member + from litellm.proxy.management_endpoints.scim.scim_transformations import ( + ScimTransformations, + ) + from litellm.proxy.utils import safe_dumps + + # Setup test data + group_id = "test-team-123" + + # Mock existing team in database + existing_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Old Team Name", + members=["user1", "user2"], # This should be ignored + members_with_roles=[ + Member(user_id="user1", role="user"), + Member(user_id="user2", role="user") + ], + metadata={"existing_key": "existing_value"} + ) + + # Mock updated SCIM group request + scim_group_update = SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Updated Team Name", + members=[ + SCIMMember(value="user2", display="User Two"), # Keep user2 + SCIMMember(value="user3", display="User Three"), # Add user3 + SCIMMember(value="user4", display="User Four") # Add user4 + ] + ) + + # Mock prisma client + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_teamtable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + + # Mock database operations + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + + # Mock the updated team that gets returned from database + updated_team = LiteLLM_TeamTable( + team_id=group_id, + team_alias="Updated Team Name", + members=["user2", "user3", "user4"], + members_with_roles=[ + Member(user_id="user2", role="user"), + Member(user_id="user3", role="user"), + Member(user_id="user4", role="user") + ], + metadata={ + "existing_key": "existing_value", + "scim_data": scim_group_update.model_dump() + } + ) + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) + + # Mock user validation (all users exist) + mock_user = mocker.MagicMock() + mock_user.user_id = "test-user" + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) + + # Mock dependencies + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client) + ) + + # Mock patch_team_membership to track membership changes + mock_patch_team_membership = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock() + ) + + # Mock SCIM transformation + expected_scim_response = SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id=group_id, + displayName="Updated Team Name", + members=[ + SCIMMember(value="user2", display="user2"), + SCIMMember(value="user3", display="user3"), + SCIMMember(value="user4", display="user4") + ] + ) + mocker.patch.object( + ScimTransformations, + "transform_litellm_team_to_scim_group", + AsyncMock(return_value=expected_scim_response) + ) + + # Execute the update_group function + result = await update_group(group_id=group_id, group=scim_group_update) + + # Verify database update was called with correct data + mock_prisma_client.db.litellm_teamtable.update.assert_called_once() + update_call_args = mock_prisma_client.db.litellm_teamtable.update.call_args + + # Check the update parameters + assert update_call_args[1]["where"]["team_id"] == group_id + update_data = update_call_args[1]["data"] + assert update_data["team_alias"] == "Updated Team Name" + + # Verify metadata includes both existing data and SCIM data + metadata_str = update_data["metadata"] + import json + metadata = json.loads(metadata_str) + assert metadata["existing_key"] == "existing_value" + assert "scim_data" in metadata + assert metadata["scim_data"]["displayName"] == "Updated Team Name" + + # Verify team membership changes were handled correctly + assert mock_patch_team_membership.call_count == 3 # Remove user1, add user3, add user4 + + # Check membership changes + call_args_list = mock_patch_team_membership.call_args_list + + # Find remove operation (user1) + remove_calls = [call for call in call_args_list + if call[1]["teams_ids_to_remove_user_from"] == [group_id]] + assert len(remove_calls) == 1 + assert remove_calls[0][1]["user_id"] == "user1" + assert remove_calls[0][1]["teams_ids_to_add_user_to"] == [] + + # Find add operations (user3, user4) + add_calls = [call for call in call_args_list + if call[1]["teams_ids_to_add_user_to"] == [group_id]] + assert len(add_calls) == 2 + add_user_ids = {call[1]["user_id"] for call in add_calls} + assert add_user_ids == {"user3", "user4"} + + # Verify all add calls have empty remove lists + for call in add_calls: + assert call[1]["teams_ids_to_remove_user_from"] == [] + + # Verify the response + assert result.id == group_id + assert result.displayName == "Updated Team Name" + assert len(result.members) == 3 + + # Verify SCIM transformation was called with updated team + ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with(updated_team) \ No newline at end of file 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 e0102f8cd7a..3a597adef06 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 @@ -118,9 +118,9 @@ async def test_key_token_handling(monkeypatch): @pytest.mark.asyncio -async def test_budget_reset_at_first_of_month(monkeypatch): +async def test_budget_reset_and_expires_at_first_of_month(monkeypatch): """ - Test that when budget_duration is "1mo", budget_reset_at is set to first of next month + Test that when budget_duration, duration, and key_budget_duration are "1mo", budget_reset_at and expires are set to first of next month """ mock_prisma_client = AsyncMock() mock_insert_data = AsyncMock( @@ -152,10 +152,12 @@ async def test_budget_reset_at_first_of_month(monkeypatch): # Use monkeypatch to set the prisma_client monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - # Test key generation with budget_duration="1mo" + # Test key generation with budget_duration="1mo", duration="1mo", key_budget_duration="1mo" response = await generate_key_helper_fn( request_type="user", budget_duration="1mo", + duration="1mo", + key_budget_duration="1mo", user_id="test_user", ) @@ -171,17 +173,17 @@ async def test_budget_reset_at_first_of_month(monkeypatch): expected_month = now.month + 1 expected_year = now.year - # Parse the response date - response_date = response["budget_reset_at"] - - # Verify budget_reset_at is set to first of next month - assert ( - response_date.year == expected_year - ), f"Expected year {expected_year}, got {response_date.year}" - assert ( - response_date.month == expected_month - ), f"Expected month {expected_month}, got {response_date.month}" - assert response_date.day == 1, f"Expected day 1, got {response_date.day}" + # Verify budget_reset_at, expires is set to first of next month + for key in ["budget_reset_at", "expires"]: + response_date = response.get(key) + assert response_date is not None, f"{key} not found in response" + assert ( + response_date.year == expected_year + ), f"Expected year {expected_year}, got {response_date.year} for {key}" + assert ( + response_date.month == expected_month + ), f"Expected month {expected_month}, got {response_date.month} for {key}" + assert response_date.day == 1, f"Expected day 1, got {response_date.day} for {key}" @pytest.mark.asyncio @@ -574,3 +576,154 @@ async def test_update_service_account_works_with_team_id(): await prepare_key_update_data(data=data, existing_key_row=existing_key) + +@pytest.mark.asyncio +async def test_validate_team_id_used_in_service_account_request_requires_team_id(): + """ + Test that validate_team_id_used_in_service_account_request raises HTTPException + when team_id is None for service account key generation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_team_id_used_in_service_account_request, + ) + + mock_prisma_client = AsyncMock() + + # Test that HTTPException is raised when team_id is None + with pytest.raises(HTTPException) as exc_info: + await validate_team_id_used_in_service_account_request( + team_id=None, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "team_id is required for service account keys" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_validate_team_id_used_in_service_account_request_requires_prisma_client(): + """ + Test that validate_team_id_used_in_service_account_request raises HTTPException + when prisma_client is None for service account key generation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_team_id_used_in_service_account_request, + ) + + # Test that HTTPException is raised when prisma_client is None + with pytest.raises(HTTPException) as exc_info: + await validate_team_id_used_in_service_account_request( + team_id="test-team-id", + prisma_client=None, + ) + + assert exc_info.value.status_code == 400 + assert "prisma_client is required for service account keys" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_validate_team_id_used_in_service_account_request_checks_team_exists(): + """ + Test that validate_team_id_used_in_service_account_request validates that + the team_id exists in the database for service account key generation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_team_id_used_in_service_account_request, + ) + + mock_prisma_client = AsyncMock() + + # Mock the database query to return None (team doesn't exist) + mock_find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique + + # Test that HTTPException is raised when team doesn't exist in DB + with pytest.raises(HTTPException) as exc_info: + await validate_team_id_used_in_service_account_request( + team_id="non-existent-team-id", + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "team_id does not exist in the database" in str(exc_info.value.detail) + + # Verify the database was queried with the correct parameters + mock_find_unique.assert_called_once_with( + where={"team_id": "non-existent-team-id"} + ) + + +@pytest.mark.asyncio +async def test_validate_team_id_used_in_service_account_request_success(): + """ + Test that validate_team_id_used_in_service_account_request returns True + when team_id exists in the database for service account key generation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_team_id_used_in_service_account_request, + ) + + mock_prisma_client = AsyncMock() + + # Mock the database query to return a team object (team exists) + mock_team = {"team_id": "existing-team-id", "team_name": "Test Team"} + mock_find_unique = AsyncMock(return_value=mock_team) + mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique + + # Test that function returns True when team exists + result = await validate_team_id_used_in_service_account_request( + team_id="existing-team-id", + prisma_client=mock_prisma_client, + ) + + assert result is True + + # Verify the database was queried with the correct parameters + mock_find_unique.assert_called_once_with( + where={"team_id": "existing-team-id"} + ) + + +@pytest.mark.asyncio +async def test_generate_service_account_key_endpoint_validation(): + """ + Test that the /key/service-account/generate endpoint properly validates + team_id requirement and team existence in database. + """ + from unittest.mock import patch + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_service_account_key_fn, + ) + + # Test case 1: Missing team_id + with pytest.raises(HTTPException) as exc_info: + await generate_service_account_key_fn( + data=GenerateKeyRequest(team_id=None), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 400 + assert "team_id is required for service account keys" in str(exc_info.value.detail) + + # Test case 2: Team doesn't exist in database + with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma: + # Mock team not found + mock_find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.find_unique = mock_find_unique + + with pytest.raises(HTTPException) as exc_info: + await generate_service_account_key_fn( + data=GenerateKeyRequest(team_id="non-existent-team"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 400 + assert "team_id does not exist in the database" in str(exc_info.value.detail) + diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py new file mode 100644 index 00000000000..789b16f9515 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -0,0 +1,740 @@ +import json +import os +import sys +from datetime import datetime +from typing import Any, Dict, List +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, +) + + +class TestOpenAIPassthroughLoggingHandler: + """Test the OpenAI passthrough logging handler for cost tracking.""" + + def setup_method(self): + """Set up test fixtures""" + self.start_time = datetime.now() + self.end_time = datetime.now() + self.handler = OpenAIPassthroughLoggingHandler() + + # Mock OpenAI chat completions response + self.mock_openai_response = { + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o-2024-08-06", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! How can I help you today?" + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 20, + "completion_tokens": 15, + "total_tokens": 35 + } + } + + def _create_mock_logging_obj(self) -> LiteLLMLoggingObj: + """Create a mock logging object""" + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {} + return mock_logging_obj + + def _create_mock_httpx_response(self, response_data: dict = None) -> httpx.Response: + """Create a mock httpx response""" + if response_data is None: + response_data = self.mock_openai_response + + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.text = json.dumps(response_data) + mock_response.json.return_value = response_data + mock_response.headers = {"content-type": "application/json"} + return mock_response + + def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload: + """Create a mock passthrough logging payload""" + return PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/chat/completions", + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + request_method="POST", + ) + + def test_llm_provider_name(self): + """Test that the handler returns the correct provider name""" + assert self.handler.llm_provider_name == "openai" + + def test_get_provider_config(self): + """Test that the handler returns an OpenAI config""" + handler = OpenAIPassthroughLoggingHandler() + config = handler.get_provider_config(model="gpt-4o") + assert config is not None + # Verify it's an OpenAI config by checking if it has the expected methods + assert hasattr(config, 'transform_response') + + def test_is_openai_chat_completions_route(self): + """Test OpenAI chat completions route detection""" + # Positive cases + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/chat/completions") == True + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://openai.azure.com/v1/chat/completions") == True + + # Negative cases + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/models") == False + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("http://localhost:4000/openai/v1/chat/completions") == False + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.anthropic.com/v1/messages") == False + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("") == False + + def test_is_openai_image_generation_route(self): + """Test OpenAI image generation route detection""" + # Positive cases + assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/images/generations") == True + assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://openai.azure.com/v1/images/generations") == True + + # Negative cases + assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/chat/completions") == False + assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/images/edits") == False + assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("http://localhost:4000/openai/v1/images/generations") == False + assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("") == False + + def test_is_openai_image_editing_route(self): + """Test OpenAI image editing route detection""" + # Positive cases + assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/images/edits") == True + assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://openai.azure.com/v1/images/edits") == True + + # Negative cases + assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/chat/completions") == False + assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/images/generations") == False + assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("http://localhost:4000/openai/v1/images/edits") == False + assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("") == False + + @patch('litellm.completion_cost') + @patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload') + def test_openai_passthrough_handler_success(self, mock_get_standard_logging, mock_completion_cost): + """Test successful cost tracking for OpenAI chat completions""" + # Arrange + mock_completion_cost.return_value = 0.000045 + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + mock_httpx_response = self._create_mock_httpx_response() + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=self.mock_openai_response, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + **kwargs + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.000045 + assert result["kwargs"]["model"] == "gpt-4o" + assert result["kwargs"]["custom_llm_provider"] == "openai" + + # Verify cost calculation was called + mock_completion_cost.assert_called_once() + + # Verify logging object was updated + assert mock_logging_obj.model_call_details["response_cost"] == 0.000045 + assert mock_logging_obj.model_call_details["model"] == "gpt-4o" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" + + @patch('litellm.completion_cost') + def test_openai_passthrough_handler_non_chat_completions(self, mock_completion_cost): + """Test that non-chat-completions routes fall back to base handler""" + # Arrange + mock_httpx_response = self._create_mock_httpx_response() + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + } + + # Act - Use a non-chat-completions route + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body={"id": "file-123", "object": "file"}, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/files", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"purpose": "fine-tune"}, + **kwargs + ) + + # Assert - Should fall back to base handler for non-chat-completions + assert result is not None + assert "result" in result + assert "kwargs" in result + # Cost calculation may be called by the base handler fallback + # The important thing is that our specific OpenAI handler logic didn't run + + @patch('litellm.completion_cost') + @patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload') + def test_openai_passthrough_handler_with_user_tracking(self, mock_get_standard_logging, mock_completion_cost): + """Test cost tracking with user information""" + # Arrange + mock_completion_cost.return_value = 0.000123 + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + mock_httpx_response = self._create_mock_httpx_response() + mock_logging_obj = self._create_mock_logging_obj() + + # Create payload with user information + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/chat/completions", + request_body={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hello"}], + "user": "test_user_123" + }, + request_method="POST", + ) + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=self.mock_openai_response, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}], "user": "test_user_123"}, + **kwargs + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.000123 + + # Verify user information is included in litellm_params + assert "litellm_params" in result["kwargs"] + assert "proxy_server_request" in result["kwargs"]["litellm_params"] + assert "body" in result["kwargs"]["litellm_params"]["proxy_server_request"] + assert result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"] == "test_user_123" + + @patch('litellm.completion_cost') + def test_openai_passthrough_handler_cost_calculation_error(self, mock_completion_cost): + """Test error handling in cost calculation""" + # Arrange + mock_completion_cost.side_effect = Exception("Cost calculation failed") + + mock_httpx_response = self._create_mock_httpx_response() + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=self.mock_openai_response, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + **kwargs + ) + + # Assert - Should fall back to base handler when cost calculation fails + assert result is not None + assert "result" in result + assert "kwargs" in result + + def test_build_complete_streaming_response(self): + """Test the streaming response builder (placeholder implementation)""" + # This is a placeholder method that returns None for now + result = self.handler._build_complete_streaming_response( + all_chunks=["chunk1", "chunk2"], + litellm_logging_obj=self._create_mock_logging_obj(), + model="gpt-4o", + ) + + assert result is None # Placeholder implementation + + @patch('litellm.completion_cost') + @patch('litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload') + def test_different_models_cost_tracking(self, mock_get_standard_logging, mock_completion_cost): + """Test cost tracking for different OpenAI models""" + # Arrange + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + test_cases = [ + ("gpt-4o", 0.000045), + ("gpt-4o-mini", 0.000015), + ("gpt-3.5-turbo", 0.000002), + ] + + for model, expected_cost in test_cases: + mock_completion_cost.return_value = expected_cost + + mock_httpx_response = self._create_mock_httpx_response() + mock_httpx_response.json.return_value = { + **self.mock_openai_response, + "model": model + } + + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": model, + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body={**self.mock_openai_response, "model": model}, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": model, "messages": [{"role": "user", "content": "Hello"}]}, + **kwargs + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == expected_cost + assert result["kwargs"]["model"] == model + assert result["kwargs"]["custom_llm_provider"] == "openai" + + def test_static_methods(self): + """Test that static methods work correctly""" + # Test static method calls + assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/chat/completions") == True + # Test instance method + handler = OpenAIPassthroughLoggingHandler() + assert handler.get_provider_config("gpt-4o") is not None + + +class TestOpenAIPassthroughIntegration: + """Integration tests for OpenAI passthrough cost tracking""" + + def setup_method(self): + """Set up test fixtures""" + self.handler = PassThroughEndpointLogging() + self.start_time = datetime.now() + self.end_time = datetime.now() + + def _create_mock_logging_obj(self) -> LiteLLMLoggingObj: + """Create a mock logging object""" + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {} + return mock_logging_obj + + def _create_mock_httpx_response(self, response_data: dict = None) -> httpx.Response: + """Create a mock httpx response""" + if response_data is None: + response_data = {"id": "test", "choices": [{"message": {"content": "Hello"}}]} + + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.text = json.dumps(response_data) + mock_response.json.return_value = response_data + mock_response.headers = {"content-type": "application/json"} + return mock_response + + def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload: + """Create a mock passthrough logging payload""" + return PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/chat/completions", + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + request_method="POST", + ) + + def test_is_openai_route_detection(self): + """Test OpenAI route detection in the main success handler""" + # Positive cases + assert self.handler.is_openai_route("https://api.openai.com/v1/chat/completions") == True + assert self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions") == True + assert self.handler.is_openai_route("https://api.openai.com/v1/models") == True + + # Negative cases + assert self.handler.is_openai_route("http://localhost:4000/openai/v1/chat/completions") == False + assert self.handler.is_openai_route("https://api.anthropic.com/v1/messages") == False + assert self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") == False + assert self.handler.is_openai_route("") == False + + @patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler') + @pytest.mark.asyncio + async def test_success_handler_calls_openai_handler(self, mock_openai_handler): + """Test that the success handler calls our OpenAI handler for OpenAI routes""" + # Arrange + mock_openai_handler.return_value = { + "result": {"id": "chatcmpl-123"}, + "kwargs": { + "response_cost": 0.000045, + "model": "gpt-4o", + "custom_llm_provider": "openai" + } + } + + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.text = '{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}' + + mock_logging_obj = AsyncMock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.async_success_handler = AsyncMock() + + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/chat/completions", + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + request_method="POST", + ) + + # Act + result = await self.handler.pass_through_async_success_handler( + httpx_response=mock_httpx_response, + response_body={"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}, + passthrough_logging_payload=passthrough_payload, + ) + + # Assert + mock_openai_handler.assert_called_once() + # The success handler returns None on success, which is expected + assert result is None + + @pytest.mark.asyncio + async def test_success_handler_falls_back_for_non_openai_routes(self): + """Test that non-OpenAI routes don't call our handler""" + # Arrange + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.text = '{"status": "success"}' + mock_httpx_response.headers = {"content-type": "application/json"} + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {} + + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.anthropic.com/v1/messages", + request_body={"model": "claude-3-sonnet", "messages": [{"role": "user", "content": "Hello"}]}, + request_method="POST", + ) + + # Mock the _handle_logging method to capture calls + self.handler._handle_logging = AsyncMock() + + # Act + result = await self.handler.pass_through_async_success_handler( + httpx_response=mock_httpx_response, + response_body={"status": "success"}, + logging_obj=mock_logging_obj, + url_route="https://api.anthropic.com/v1/messages", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"model": "claude-3-sonnet", "messages": [{"role": "user", "content": "Hello"}]}, + passthrough_logging_payload=passthrough_payload, + ) + + # Assert - Should call the base handler, not our OpenAI handler + self.handler._handle_logging.assert_called_once() + + @patch('litellm.cost_calculator.default_image_cost_calculator') + def test_calculate_image_generation_cost(self, mock_image_cost_calculator): + """Test image generation cost calculation""" + # Arrange + mock_image_cost_calculator.return_value = 0.040 + model = "dall-e-3" + response_body = { + "data": [ + { + "url": "https://example.com/image1.png", + "revised_prompt": "A beautiful sunset over the ocean" + } + ] + } + request_body = { + "model": "dall-e-3", + "prompt": "A beautiful sunset over the ocean", + "n": 1, + "size": "1024x1024", + "quality": "standard" + } + + # Act + cost = OpenAIPassthroughLoggingHandler._calculate_image_generation_cost( + model=model, + response_body=response_body, + request_body=request_body, + ) + + # Assert + assert cost == 0.040 + mock_image_cost_calculator.assert_called_once_with( + model=model, + custom_llm_provider="openai", + quality="standard", + n=1, + size="1024x1024", + optional_params=request_body, + ) + + @patch('litellm.cost_calculator.default_image_cost_calculator') + def test_calculate_image_editing_cost(self, mock_image_cost_calculator): + """Test image editing cost calculation""" + # Arrange + mock_image_cost_calculator.return_value = 0.020 + model = "dall-e-2" + response_body = { + "data": [ + { + "url": "https://example.com/edited_image.png", + "revised_prompt": "A beautiful sunset over the ocean with added clouds" + } + ] + } + request_body = { + "model": "dall-e-2", + "prompt": "Add clouds to the sky", + "n": 1, + "size": "1024x1024" + } + + # Act + cost = OpenAIPassthroughLoggingHandler._calculate_image_editing_cost( + model=model, + response_body=response_body, + request_body=request_body, + ) + + # Assert + assert cost == 0.020 + mock_image_cost_calculator.assert_called_once_with( + model=model, + custom_llm_provider="openai", + quality=None, # Image editing doesn't have quality parameter + n=1, + size="1024x1024", + optional_params=request_body, + ) + + def test_cost_calculation_preservation(self): + """Test that manually calculated costs are preserved and not overridden.""" + # Create a logging object + logging_obj = LiteLLMLoggingObj( + model="dall-e-3", + messages=[{"role": "user", "content": "Generate an image"}], + stream=False, + call_type="pass_through_endpoint", + start_time=self.start_time, + litellm_call_id="test_123", + function_id="test_fn", + ) + + # Set a manually calculated cost in model_call_details + test_cost = 0.040000 + logging_obj.model_call_details["response_cost"] = test_cost + logging_obj.model_call_details["model"] = "dall-e-3" + logging_obj.model_call_details["custom_llm_provider"] = "openai" + + # Create an ImageResponse with cost in _hidden_params + from litellm.types.utils import ImageResponse + image_response = ImageResponse( + data=[{"url": "https://example.com/image.png"}], + model="dall-e-3", + ) + image_response._hidden_params = {"response_cost": test_cost} + + # Test the _response_cost_calculator method + calculated_cost = logging_obj._response_cost_calculator(result=image_response) + + assert calculated_cost == test_cost, f"Expected {test_cost}, got {calculated_cost}" + + @patch('litellm.cost_calculator.default_image_cost_calculator') + def test_openai_passthrough_handler_image_generation(self, mock_image_cost_calculator): + """Test successful cost tracking for OpenAI image generation""" + # Arrange + mock_image_cost_calculator.return_value = 0.040 + + mock_image_response = { + "data": [ + { + "url": "https://example.com/image1.png", + "revised_prompt": "A beautiful sunset over the ocean" + } + ] + } + + mock_httpx_response = self._create_mock_httpx_response(mock_image_response) + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "dall-e-3", + } + + request_body = { + "model": "dall-e-3", + "prompt": "A beautiful sunset over the ocean", + "n": 1, + "size": "1024x1024", + "quality": "standard" + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=mock_image_response, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/images/generations", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body=request_body, + **kwargs + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.040 + assert result["kwargs"]["model"] == "dall-e-3" + assert result["kwargs"]["custom_llm_provider"] == "openai" + + # Verify cost calculation was called + mock_image_cost_calculator.assert_called_once() + + # Verify logging object was updated + assert mock_logging_obj.model_call_details["response_cost"] == 0.040 + assert mock_logging_obj.model_call_details["model"] == "dall-e-3" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" + + @patch('litellm.cost_calculator.default_image_cost_calculator') + def test_openai_passthrough_handler_image_editing(self, mock_image_cost_calculator): + """Test successful cost tracking for OpenAI image editing""" + # Arrange + mock_image_cost_calculator.return_value = 0.020 + + mock_image_response = { + "data": [ + { + "url": "https://example.com/edited_image.png", + "revised_prompt": "A beautiful sunset over the ocean with added clouds" + } + ] + } + + mock_httpx_response = self._create_mock_httpx_response(mock_image_response) + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "dall-e-2", + } + + request_body = { + "model": "dall-e-2", + "prompt": "Add clouds to the sky", + "n": 1, + "size": "1024x1024" + } + + # Act + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=mock_image_response, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/images/edits", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body=request_body, + **kwargs + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.020 + assert result["kwargs"]["model"] == "dall-e-2" + assert result["kwargs"]["custom_llm_provider"] == "openai" + + # Verify cost calculation was called + mock_image_cost_calculator.assert_called_once() + + # Verify logging object was updated + assert mock_logging_obj.model_call_details["response_cost"] == 0.020 + assert mock_logging_obj.model_call_details["model"] == "dall-e-2" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index b8381201b0a..e296cb25f80 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1245,3 +1245,58 @@ async def test_delete_pass_through_endpoint_empty_list(): # Verify the exception assert exc_info.value.status_code == 400 assert "no pass-through endpoints setup" in str(exc_info.value.detail).lower() + + + +@pytest.mark.asyncio +async def test_pass_through_with_httpbin_redirect(): + """ + Integration test using httpbin.org redirect endpoint to test real redirect handling. + This tests the actual redirect handling capability end-to-end using the full pass_through_request function. + """ + from unittest.mock import MagicMock + + from fastapi import Request + from starlette.datastructures import Headers, QueryParams + + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + pass_through_request, + ) + + # Create mock request + mock_request = MagicMock(spec=Request) + mock_request.method = "GET" + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams("") + + # Mock the body method to return empty bytes for GET request + async def mock_body(): + return b"" + mock_request.body = mock_body + + # Mock user API key dict + mock_user_api_key_dict = MagicMock() + + try: + # Test with httpbin.org redirect endpoint + # This will redirect to httpbin.org/get + response = await pass_through_request( + request=mock_request, + target="https://httpbin.org/redirect/1", + custom_headers={}, + user_api_key_dict=mock_user_api_key_dict + ) + + # Should get the final response (200) from /get endpoint, not the redirect (302) + assert response.status_code == 200 + + # The response should be from the /get endpoint + response_content = response.body.decode('utf-8') + + # httpbin.org/get returns JSON with info about the request + assert '"url": "https://httpbin.org/get"' in response_content + print("GOT A Response from HTTPBIN=", response_content) + except Exception as e: + # If httpbin.org is not accessible, skip the test + import pytest + pytest.skip(f"Could not reach httpbin.org for integration test: {e}") diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 5104ffd80de..817f19d8d7d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch import pytest from fastapi import Request +import litellm from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import ( KeyAndTeamLoggingSettings, @@ -935,3 +936,126 @@ def test_add_headers_to_llm_call_by_model_group_existing_headers_in_data(): finally: # Restore original model_group_settings litellm.model_group_settings = original_model_group_settings + +import json +import time +from typing import Optional +from unittest.mock import AsyncMock + +from fastapi.responses import Response + +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import StandardLoggingPayload + + +class TestCustomLogger(CustomLogger): + def __init__(self): + self.standard_logging_object: Optional[StandardLoggingPayload] = None + super().__init__() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + print(f"SUCCESS CALLBACK CALLED! kwargs keys: {list(kwargs.keys())}") + self.standard_logging_object = kwargs.get("standard_logging_object") + print(f"Captured standard_logging_object: {self.standard_logging_object}") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + print(f"FAILURE CALLBACK CALLED! kwargs keys: {list(kwargs.keys())}") + +@pytest.mark.asyncio +async def test_add_litellm_metadata_from_request_headers(): + """ + Test that add_litellm_metadata_from_request_headers properly adds litellm metadata from request headers, + makes an LLM request using base_process_llm_request, sleeps for 3 seconds, and checks standard_logging_payload has spend_logs_metadata from headers + + Relevant issue: https://github.com/BerriAI/litellm/issues/14008 + """ + # Set up test logger + litellm._turn_on_debug() + test_logger = TestCustomLogger() + litellm.callbacks = [test_logger] + + # Prepare test data (ensure no streaming, add mock_response and api_key to route to litellm.acompletion) + headers = {"x-litellm-spend-logs-metadata": '{"user_id": "12345", "project_id": "proj_abc", "request_type": "chat_completion", "timestamp": "2025-09-02T10:30:00Z"}'} + data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "stream": False, "mock_response": "Hi", "api_key": "fake-key"} + + # Create mock request with headers + mock_request = MagicMock(spec=Request) + mock_request.headers = headers + mock_request.url.path = "/chat/completions" + + # Create mock response + mock_fastapi_response = MagicMock(spec=Response) + + # Create mock user API key dict + mock_user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + org_id="test-org" + ) + + # Create mock proxy logging object + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + + # Create async functions for the hooks + async def mock_during_call_hook(*args, **kwargs): + return None + + async def mock_pre_call_hook(*args, **kwargs): + return data + + async def mock_post_call_success_hook(*args, **kwargs): + # Return the response unchanged + return kwargs.get('response', args[2] if len(args) > 2 else None) + + mock_proxy_logging_obj.during_call_hook = mock_during_call_hook + mock_proxy_logging_obj.pre_call_hook = mock_pre_call_hook + mock_proxy_logging_obj.post_call_success_hook = mock_post_call_success_hook + + # Create mock proxy config + mock_proxy_config = MagicMock() + + # Create mock general settings + general_settings = {} + + # Create mock select_data_generator with correct signature + def mock_select_data_generator(response=None, user_api_key_dict=None, request_data=None): + async def mock_generator(): + yield "data: " + json.dumps({"choices": [{"delta": {"content": "Hello"}}]}) + "\n\n" + yield "data: [DONE]\n\n" + return mock_generator() + + # Create the processor + processor = ProxyBaseLLMRequestProcessing(data=data) + + # Call base_process_llm_request (it will use the mock_response="Hi" parameter) + result = await processor.base_process_llm_request( + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + route_type="acompletion", + proxy_logging_obj=mock_proxy_logging_obj, + general_settings=general_settings, + proxy_config=mock_proxy_config, + select_data_generator=mock_select_data_generator, + llm_router=None, + model="gpt-4", + is_streaming_request=False + ) + + # Sleep for 3 seconds to allow logging to complete + await asyncio.sleep(3) + + # Check if standard_logging_object was set + assert test_logger.standard_logging_object is not None, "standard_logging_object should be populated after LLM request" + + # Verify the logging object contains expected metadata + standard_logging_obj = test_logger.standard_logging_object + + print(f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}") + + SPEND_LOGS_METADATA = standard_logging_obj["metadata"]["spend_logs_metadata"] + assert SPEND_LOGS_METADATA == dict(json.loads(headers["x-litellm-spend-logs-metadata"])), "spend_logs_metadata should be the same as the headers" + + diff --git a/tests/test_litellm/responses/test_text_format_conversion.py b/tests/test_litellm/responses/test_text_format_conversion.py new file mode 100644 index 00000000000..20a87a4abbb --- /dev/null +++ b/tests/test_litellm/responses/test_text_format_conversion.py @@ -0,0 +1,161 @@ +import json +import os +import sys + +import pytest +from pydantic import BaseModel + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.types.llms.openai import ( + IncompleteDetails, + ResponseAPIUsage, + ResponsesAPIResponse, +) + + +class TestTextFormatConversion: + """Test text_format to text parameter conversion for responses API""" + + def get_base_completion_call_args(self): + """Get base arguments for completion call""" + return { + "model": "gpt-4o", + "api_key": "test-key", + "api_base": "https://api.openai.com/v1", + } + + @pytest.mark.asyncio + async def test_text_format_to_text_conversion(self): + """ + Test that when text_format parameter is passed to litellm.aresponses, + it gets converted to text parameter in the raw API call to OpenAI. + """ + from unittest.mock import AsyncMock, patch + + class TestResponse(BaseModel): + """Test Pydantic model for structured output""" + + answer: str + confidence: float + + class MockResponse: + """Mock response class for testing""" + + def __init__(self, json_data, status_code): + self._json_data = json_data + self.status_code = status_code + self.text = json.dumps(json_data) + + def json(self): + return self._json_data + + # Mock response from OpenAI + mock_response = { + "id": "resp_123", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "gpt-4o", + "output": [ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": '{"answer": "Paris", "confidence": 0.95}', + "annotations": [], + } + ], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 10, + "output_tokens": 20, + "total_tokens": 30, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "text": {"format": {"type": "json_object"}}, + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "temperature": 1.0, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": {"effort": None, "summary": None}, + "truncation": "disabled", + "user": None, + } + + base_completion_call_args = self.get_base_completion_call_args() + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + # Configure the mock to return our response + mock_post.return_value = MockResponse(mock_response, 200) + + litellm._turn_on_debug() + litellm.set_verbose = True + + # Call aresponses with text_format parameter + response = await litellm.aresponses( + input="What is the capital of France?", + text_format=TestResponse, + **base_completion_call_args, + ) + + # Verify the request was made correctly + mock_post.assert_called_once() + request_body = mock_post.call_args.kwargs["json"] + print("Request body:", json.dumps(request_body, indent=4)) + + # Validate that text_format was converted to text parameter + assert ( + "text" in request_body + ), "text parameter should be present in request body" + assert ( + "text_format" not in request_body + ), "text_format should not be in request body" + + # Validate the text parameter structure + text_param = request_body["text"] + assert "format" in text_param, "text parameter should have format field" + assert ( + text_param["format"]["type"] == "json_schema" + ), "format type should be json_schema" + assert "name" in text_param["format"], "format should have name field" + assert ( + text_param["format"]["name"] == "TestResponse" + ), "format name should match Pydantic model name" + assert "schema" in text_param["format"], "format should have schema field" + assert "strict" in text_param["format"], "format should have strict field" + + # Validate the schema structure + schema = text_param["format"]["schema"] + assert schema["type"] == "object", "schema type should be object" + assert "properties" in schema, "schema should have properties" + assert ( + "answer" in schema["properties"] + ), "schema should have answer property" + assert ( + "confidence" in schema["properties"] + ), "schema should have confidence property" + + # Validate other request parameters + assert request_body["input"] == "What is the capital of France?" + + # Validate the response + print("Response:", json.dumps(response, indent=4, default=str)) diff --git a/tests/local_testing/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py similarity index 81% rename from tests/local_testing/test_router_tag_routing.py rename to tests/test_litellm/router_strategy/test_router_tag_routing.py index 87cf2261a67..e78a16c6212 100644 --- a/tests/local_testing/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -63,6 +63,7 @@ async def test_router_free_paid_tier(): model="gpt-4", messages=[{"role": "user", "content": "Tell me a joke."}], metadata={"tags": ["free"]}, + mock_response="Tell me a joke.", ) print("Response: ", response) @@ -78,6 +79,7 @@ async def test_router_free_paid_tier(): model="gpt-4", messages=[{"role": "user", "content": "Tell me a joke."}], metadata={"tags": ["paid"]}, + mock_response="Tell me a joke.", ) print("Response: ", response) @@ -136,6 +138,7 @@ async def test_router_free_paid_tier_embeddings(): model="gpt-4", input="Tell me a joke.", metadata={"tags": ["free"]}, + mock_response=[1, 2, 3], ) print("Response: ", response) @@ -151,6 +154,7 @@ async def test_router_free_paid_tier_embeddings(): model="gpt-4", input="Tell me a joke.", metadata={"tags": ["paid"]}, + mock_response=[1, 2, 3], ) print("Response: ", response) @@ -205,6 +209,7 @@ async def test_default_tagged_deployments(): response = await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "Tell me a joke."}], + mock_response="Tell me a joke.", ) print("Response: ", response) @@ -220,6 +225,7 @@ async def test_default_tagged_deployments(): model="gpt-4", messages=[{"role": "user", "content": "Tell me a joke."}], metadata={"tags": ["default"]}, + mock_response="Tell me a joke.", ) print("Response: ", response) @@ -235,6 +241,7 @@ async def test_default_tagged_deployments(): model="gpt-4", messages=[{"role": "user", "content": "Tell me a joke."}], metadata={"tags": ["invalid-tag"]}, + mock_response="Tell me a joke.", ) print("Response: ", response) @@ -292,6 +299,7 @@ async def test_error_from_tag_routing(): model="gpt-4", messages=[{"role": "user", "content": "Tell me a joke."}], metadata={"tags": ["paid"]}, + mock_response="Tell me a joke.", ) pytest.fail("this should have failed - expected it to fail") @@ -315,3 +323,66 @@ def test_tag_routing_with_list_of_tags(): assert not is_valid_deployment_tag(["teamA", "teamB"], ["teamC"]) assert not is_valid_deployment_tag(["teamA", "teamB"], []) assert not is_valid_deployment_tag(["default"], ["teamA"]) + + +@pytest.mark.asyncio() +async def test_router_free_paid_tier_with_responses_api(): + """ + Pass list of orgs in 1 model definition, + expect a unique deployment for each to be created + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["free"], + }, + "model_info": {"id": "very-cheap-model"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["paid"], + }, + "model_info": {"id": "very-expensive-model"}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + # this should pick model with id == very-cheap-model + response = await router.aresponses( + model="gpt-4", + input="Tell me a joke.", + litellm_metadata={"tags": ["free"]}, + mock_response="Tell me a joke.", + ) + + print("Response: ", response) + + response_extra_info = response._hidden_params + print("response_extra_info: ", response_extra_info) + + assert response_extra_info["model_id"] == "very-cheap-model" + + for _ in range(5): + # this should pick model with id == very-cheap-model + response = await router.aresponses( + model="gpt-4", + input="Tell me a joke.", + litellm_metadata={"tags": ["paid"]}, + mock_response="Tell me a joke.", + ) + + print("Response: ", response) + + response_extra_info = response._hidden_params + print("response_extra_info: ", response_extra_info) + + assert response_extra_info["model_id"] == "very-expensive-model" \ No newline at end of file diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 49e2c6c1856..2bf2935901e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -170,7 +170,9 @@ def test_all_model_configs(): drop_params=False, ) == {"max_tokens": 10} - from litellm.llms.volcengine import VolcEngineConfig + from litellm.llms.volcengine.chat.transformation import ( + VolcEngineChatConfig as VolcEngineConfig, + ) assert "max_completion_tokens" in VolcEngineConfig().get_supported_openai_params( model="llama3" @@ -549,6 +551,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "completion", "embedding", "image_generation", + "video_generation", "moderation", "rerank", "responses", @@ -636,7 +639,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "type": "array", "items": { "type": "string", - "enum": ["text", "image", "audio", "code"], + "enum": ["text", "image", "audio", "code", "video"], }, }, "supports_native_streaming": {"type": "boolean"}, @@ -687,6 +690,7 @@ def test_get_model_info_gemini(): and not "gemma" in model and not "learnlm" in model and not "imagen" in model + and not "veo" in model ): assert info.get("tpm") is not None, f"{model} does not have tpm" assert info.get("rpm") is not None, f"{model} does not have rpm" @@ -979,10 +983,10 @@ class TestProxyFunctionCalling: # Groq models (mixed support) ("groq/gemma-7b-it", "litellm_proxy/groq/gemma-7b-it", True), ( - "groq/llama3-70b-8192", - "litellm_proxy/groq/llama3-70b-8192", - False, - ), # This model doesn't support function calling + "groq/llama-3.3-70b-versatile", + "litellm_proxy/groq/llama-3.3-70b-versatile", + True, + ), # Cohere models (generally don't support function calling) ("command-nightly", "litellm_proxy/command-nightly", False), ], diff --git a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx index 30f01d91db5..cca7bc08608 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx @@ -8,6 +8,7 @@ import { useReactTable, ColumnResizeMode, VisibilityState, + PaginationState, } from "@tanstack/react-table"; import React from "react"; import { @@ -18,7 +19,7 @@ import { TableRow, TableCell, } from "@tremor/react"; -import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, TableIcon } from "@heroicons/react/outline"; +import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon } from "@heroicons/react/outline"; // Extend the column meta type to include className declare module "@tanstack/react-table" { diff --git a/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx b/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx index 85f18ee89e1..2eff0bf3e36 100644 --- a/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx @@ -1,4 +1,4 @@ -import React, { useState, useEffect, useRef } from "react" +import React, { useState, useEffect, useRef, useMemo } from "react" import { Card, Title, @@ -61,7 +61,7 @@ import HealthCheckComponent from "../model_dashboard/HealthCheckComponent"; import PassThroughSettings from "../pass_through_settings"; import ModelGroupAliasSettings from "../model_group_alias_settings"; import { all_admin_roles } from "@/utils/roles"; -import { Table as TableInstance } from "@tanstack/react-table"; +import { Table as TableInstance, PaginationState } from "@tanstack/react-table"; import NotificationsManager from "../molecules/notifications_manager"; interface ModelDashboardProps { @@ -191,12 +191,21 @@ const ModelDashboard: React.FC = ({ const [currentTeam, setCurrentTeam] = useState("personal") // 'personal' or team_id const [modelViewMode, setModelViewMode] = useState<"current_team" | "all">("current_team") + // Add state for showing/hiding filters + const [showFilters, setShowFilters] = useState(false) + const [showColumnDropdown, setShowColumnDropdown] = useState(false) const [isDropdownOpen, setIsDropdownOpen] = useState(false) const [expandedRows, setExpandedRows] = useState>(new Set()) const dropdownRef = useRef(null) const tableRef = useRef>(null) + + // Pagination state + const [pagination, setPagination] = useState({ + pageIndex: 0, + pageSize: 50, + }) const [selectedTabIndex, setSelectedTabIndex] = useState(0) const handleCreateNewModelClick = () => { @@ -206,6 +215,63 @@ const ModelDashboard: React.FC = ({ setSelectedTabIndex(1) } + const resetFilters = () => { + setModelNameSearch("") + setSelectedModelGroup("all") + setSelectedModelAccessGroupFilter(null) + setCurrentTeam("personal") + setModelViewMode("current_team") + setPagination({ pageIndex: 0, pageSize: 50 }) + } + + // Memoize filtered data to prevent unnecessary re-calculations + const filteredData = useMemo(() => { + if (!modelData || !modelData.data || modelData.data.length === 0) { + return []; + } + + return modelData.data.filter((model: any) => { + const searchMatch = + modelNameSearch === "" || + model.model_name.toLowerCase().includes(modelNameSearch.toLowerCase()) + + const modelNameMatch = + selectedModelGroup === "all" || + model.model_name === selectedModelGroup || + !selectedModelGroup || + (selectedModelGroup === "wildcard" && model.model_name?.includes("*")) + + const accessGroupMatch = + selectedModelAccessGroupFilter === "all" || + model.model_info["access_groups"]?.includes(selectedModelAccessGroupFilter) || + !selectedModelAccessGroupFilter + + let teamAccessMatch = true + if (modelViewMode === "current_team") { + if (currentTeam === "personal") { + teamAccessMatch = model.model_info?.direct_access === true + } else { + teamAccessMatch = + model.model_info?.access_via_team_ids?.includes(currentTeam) === true + } + } + + return searchMatch && modelNameMatch && accessGroupMatch && teamAccessMatch + }); + }, [modelData, modelNameSearch, selectedModelGroup, selectedModelAccessGroupFilter, currentTeam, modelViewMode]); + + // Memoize paginated data + const paginatedData = useMemo(() => { + const startIndex = pagination.pageIndex * pagination.pageSize; + const endIndex = startIndex + pagination.pageSize; + return filteredData.slice(startIndex, endIndex); + }, [filteredData, pagination.pageIndex, pagination.pageSize]); + + // Reset pagination when filters change + useEffect(() => { + setPagination(prev => ({ ...prev, pageIndex: 0 })) + }, [modelNameSearch, selectedModelGroup, selectedModelAccessGroupFilter, currentTeam, modelViewMode]) + const setProviderModelsFn = (provider: Providers) => { const _providerModels = getProviderModels(provider, modelMap) setProviderModels(_providerModels) @@ -1006,119 +1072,166 @@ const ModelDashboard: React.FC = ({
-
-
- {/* Current Team Selector - Prominent */} -
-
-
- Current Team: - setCurrentTeam(value)} + > + +
+
+ Personal +
+
+ {teams + ?.filter((team) => team.team_id) + .map((team) => ( +
-
- Personal +
+ + {team.team_alias + ? `${team.team_alias.slice(0, 30)}...` + : `Team ${team.team_id.slice(0, 30)}...`} +
- {teams - ?.filter((team) => team.team_id) - .map((team) => ( - -
-
- - {team.team_alias - ? `${team.team_alias.slice(0, 30)}...` - : `Team ${team.team_id.slice(0, 30)}...`} - -
-
- ))} - -
- {modelViewMode === "current_team" && ( -
- -
- {currentTeam === "personal" ? ( - - To access these models: Create a Virtual Key without selecting a team on the{" "} - - Virtual Keys page - - - ) : ( - - To access these models: Create a Virtual Key and select Team as " - {currentTeam}" on the{" "} - - Virtual Keys page - - - )} -
-
- )} -
- - {/* Model View Mode Toggle - Also prominent */} -
- View: - -
+ ))} +
- {/* Other Filters */} -
-
- {/* Model Name Search */} -
- Search Public Model Name: - -
+
+ View: + +
+
+ + {modelViewMode === "current_team" && ( +
+ +
+ {currentTeam === "personal" ? ( + + To access these models: Create a Virtual Key without selecting a team on the{" "} + + Virtual Keys page + + + ) : ( + + To access these models: Create a Virtual Key and select Team as " + {currentTeam}" on the{" "} + + Virtual Keys page + + + )} +
+
+ )} +
+ {/* Search and Filter Controls */} +
+
+ {/* Search and Filter Controls */} +
+ {/* Model Name Search */} +
+ setModelNameSearch(e.target.value)} + /> + + + +
+ + {/* Filter Button */} + + + {/* Reset Filters Button */} + +
+ + {/* Additional Filters */} + {showFilters && ( +
{/* Model Name Filter */} -
- Filter by Public Model Name: +
-
- Filter by Model Access Group: + {/* Model Access Group Filter */} +
-
+ )} - {/* Results Count */} + {/* Results Count and Pagination Controls */}
- - Showing{" "} - {modelData && modelData.data.length > 0 - ? modelData.data.filter((model: any) => { - const searchMatch = - modelNameSearch === "" || - model.model_name.toLowerCase().includes(modelNameSearch.toLowerCase()) - - const modelNameMatch = - selectedModelGroup === "all" || - model.model_name === selectedModelGroup || - !selectedModelGroup - const accessGroupMatch = - selectedModelAccessGroupFilter === "all" || - model.model_info["access_groups"]?.includes(selectedModelAccessGroupFilter) || - !selectedModelAccessGroupFilter - let teamAccessMatch = true - if (modelViewMode === "current_team") { - if (currentTeam === "personal") { - teamAccessMatch = model.model_info?.direct_access === true - } else { - teamAccessMatch = - model.model_info?.access_via_team_ids?.includes(currentTeam) === true - } - } - - return searchMatch && modelNameMatch && accessGroupMatch && teamAccessMatch - }).length - : 0}{" "} - results - + + {filteredData.length > 0 ? ( + `Showing ${pagination.pageIndex * pagination.pageSize + 1} - ${Math.min( + (pagination.pageIndex + 1) * pagination.pageSize, + filteredData.length + )} of ${filteredData.length} results` + ) : ( + "Showing 0 results" + )} + + + {/* Pagination Controls */} + {filteredData.length > pagination.pageSize && ( +
+ + + +
+ )}
@@ -1202,38 +1320,7 @@ const ModelDashboard: React.FC = ({ expandedRows, setExpandedRows, )} - data={modelData.data.filter((model: any) => { - // Model name search filter - const searchMatch = - modelNameSearch === "" || - model.model_name.toLowerCase().includes(modelNameSearch.toLowerCase()) - - // Model name filter - const modelNameMatch = - selectedModelGroup === "all" || - model.model_name === selectedModelGroup || - !selectedModelGroup || - (selectedModelGroup === "wildcard" && model.model_name?.includes("*")) - // Model access group filter - const accessGroupMatch = - selectedModelAccessGroupFilter === "all" || - model.model_info["access_groups"]?.includes(selectedModelAccessGroupFilter) || - !selectedModelAccessGroupFilter - // Team access filter based on current team and view mode - let teamAccessMatch = true - if (modelViewMode === "current_team") { - if (currentTeam === "personal") { - // Show only models with direct access - teamAccessMatch = model.model_info?.direct_access === true - } else { - // Show only models accessible by the current team - teamAccessMatch = model.model_info?.access_via_team_ids?.includes(currentTeam) === true - } - } - // For 'all' mode, show all models (teamAccessMatch remains true) - - return searchMatch && modelNameMatch && accessGroupMatch && teamAccessMatch - })} + data={paginatedData} isLoading={false} table={tableRef} /> diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index c45116c0ddd..6c9ef66e95b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -80,6 +80,7 @@ export default function SpendLogsTable({ const [selectedKeyInfo, setSelectedKeyInfo] = useState(null) const [selectedKeyIdInfoView, setSelectedKeyIdInfoView] = useState(null) const [selectedStatus, setSelectedStatus] = useState("") + const [selectedEndUser, setSelectedEndUser] = useState("") const [filterByCurrentUser, setFilterByCurrentUser] = useState(userRole && internalUserRoles.includes(userRole)) const [activeTab, setActiveTab] = useState("request logs") @@ -193,6 +194,7 @@ export default function SpendLogsTable({ currentPage, pageSize, filterByCurrentUser ? userID : undefined, + selectedEndUser, selectedStatus, selectedModel, ) @@ -280,6 +282,7 @@ export default function SpendLogsTable({ } setSelectedStatus(filters["Status"] || "") setSelectedModel(filters["Model"] || "") + setSelectedEndUser(filters["End User"] || "") if (filters["Key Hash"]) { setSelectedKeyHash(filters["Key Hash"]) diff --git a/ui/litellm-dashboard/src/components/view_users.tsx b/ui/litellm-dashboard/src/components/view_users.tsx index 1d80e513e35..7c828f756b5 100644 --- a/ui/litellm-dashboard/src/components/view_users.tsx +++ b/ui/litellm-dashboard/src/components/view_users.tsx @@ -79,7 +79,6 @@ const ViewUserDashboard: React.FC = ({ accessToken, toke const [activeTab, setActiveTab] = useState("users") const [filters, setFilters] = useState(initialFilters) const [debouncedFilters, setDebouncedFilters, debouncer] = useDebouncedState(filters, { wait: 300 }) - const [showFilters, setShowFilters] = useState(false) const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false) const [invitationLinkData, setInvitationLinkData] = useState(null) const [baseUrl, setBaseUrl] = useState(null) @@ -330,209 +329,35 @@ const ViewUserDashboard: React.FC = ({ accessToken, toke -
-
-
- {/* Search and Filter Controls */} -
- {/* Email Search */} -
- updateFilters({ email: e.target.value })} - /> - - - -
- - {/* Filter Button */} - - - {/* Reset Filters Button */} - -
- - {/* Additional Filters */} - {showFilters && ( -
- {/* User ID Search */} -
- updateFilters({ user_id: e.target.value })} - /> - - - -
- - {/* Role Dropdown */} -
- -
- - {/* Team Dropdown */} -
- -
- - {/* SSO ID Search */} -
- updateFilters({ sso_user_id: e.target.value })} - /> -
-
- )} - - {/* Results Count and Pagination */} -
- - Showing{" "} - {userListResponse && userListResponse.users && userListResponse.users.length > 0 - ? (userListResponse.page - 1) * userListResponse.page_size + 1 - : 0}{" "} - -{" "} - {userListResponse && userListResponse.users - ? Math.min(userListResponse.page * userListResponse.page_size, userListResponse.total) - : 0}{" "} - of {userListResponse ? userListResponse.total : 0} results - - - {/* Pagination Buttons */} -
- - -
-
-
-
-
- { - setSelectedUser(user) - setEditModalVisible(true) - }} - handleDelete={handleDelete} - handleResetPassword={handleResetPassword} - enableSelection={selectionMode} - selectedUsers={selectedUsers} - onSelectionChange={handleSelectionChange} - /> -
- -
+ { + setSelectedUser(user) + setEditModalVisible(true) + }} + handleDelete={handleDelete} + handleResetPassword={handleResetPassword} + enableSelection={selectionMode} + selectedUsers={selectedUsers} + onSelectionChange={handleSelectionChange} + filters={filters} + updateFilters={updateFilters} + initialFilters={initialFilters} + teams={teams} + userListResponse={userListResponse} + currentPage={currentPage} + handlePageChange={handlePageChange} + />
diff --git a/ui/litellm-dashboard/src/components/view_users/table.tsx b/ui/litellm-dashboard/src/components/view_users/table.tsx index 0f1193c6e87..0b422ead7b3 100644 --- a/ui/litellm-dashboard/src/components/view_users/table.tsx +++ b/ui/litellm-dashboard/src/components/view_users/table.tsx @@ -15,12 +15,27 @@ import { TableBody, TableRow, TableCell, + Select, + SelectItem, } from "@tremor/react"; import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon } from "@heroicons/react/outline"; import { UserInfo } from "./types"; import UserInfoView from "./user_info_view"; import { columns as createColumns } from "./columns"; +interface FilterState { + email: string; + user_id: string; + user_role: string; + sso_user_id: string; + team: string; + model: string; + min_spend: number | null; + max_spend: number | null; + sort_by: string; + sort_order: "asc" | "desc"; +} + interface UserDataTableProps { data: UserInfo[]; columns: ColumnDef[]; @@ -39,6 +54,15 @@ interface UserDataTableProps { selectedUsers?: UserInfo[]; onSelectionChange?: (selectedUsers: UserInfo[]) => void; enableSelection?: boolean; + // Filter-related props + filters: FilterState; + updateFilters: (update: Partial) => void; + initialFilters: FilterState; + teams: any[] | null; + // Pagination props + userListResponse: any; + currentPage: number; + handlePageChange: (newPage: number) => void; } export function UserDataTable({ @@ -56,6 +80,13 @@ export function UserDataTable({ selectedUsers = [], onSelectionChange, enableSelection = false, + filters, + updateFilters, + initialFilters, + teams, + userListResponse, + currentPage, + handlePageChange, }: UserDataTableProps) { const [sorting, setSorting] = React.useState([ { @@ -65,6 +96,7 @@ export function UserDataTable({ ]); const [selectedUserId, setSelectedUserId] = React.useState(null); const [openInEditMode, setOpenInEditMode] = React.useState(false); + const [showFilters, setShowFilters] = React.useState(false); const handleUserClick = (userId: string, openInEditMode: boolean = false) => { setSelectedUserId(userId); @@ -171,9 +203,190 @@ export function UserDataTable({ } return ( -
-
- +
+ {/* Filter Section */} +
+
+ {/* Search and Filter Controls */} +
+ {/* Email Search */} +
+ updateFilters({ email: e.target.value })} + /> + + + +
+ + {/* Filter Button */} + + + {/* Reset Filters Button */} + +
+ + {/* Additional Filters */} + {showFilters && ( +
+ {/* User ID Search */} +
+ updateFilters({ user_id: e.target.value })} + /> + + + +
+ + {/* Role Dropdown */} +
+ +
+ + {/* Team Dropdown */} +
+ +
+ + {/* SSO ID Search */} +
+ updateFilters({ sso_user_id: e.target.value })} + /> +
+
+ )} + + {/* Results Count and Pagination */} +
+ + Showing{" "} + {userListResponse && userListResponse.users && userListResponse.users.length > 0 + ? (userListResponse.page - 1) * userListResponse.page_size + 1 + : 0}{" "} + -{" "} + {userListResponse && userListResponse.users + ? Math.min(userListResponse.page * userListResponse.page_size, userListResponse.total) + : 0}{" "} + of {userListResponse ? userListResponse.total : 0} results + + + {/* Pagination Buttons */} +
+ + +
+
+
+
+ + {/* Table Section */} +
+
+
+
{table.getHeaderGroups().map((headerGroup) => ( @@ -260,6 +473,8 @@ export function UserDataTable({ )}
+
+
); diff --git a/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx index d416df2d9b0..c36bde78a7d 100644 --- a/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx @@ -14,9 +14,9 @@ import { rolesWithWriteAccess } from "../../utils/roles" import { UserEditView } from "../user_edit_view" import OnboardingModal, { InvitationLink } from "../onboarding_link" import { formatNumberWithCommas, copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils" -import { CopyIcon, CheckIcon } from "lucide-react"; -import NotificationsManager from "../molecules/notifications_manager"; -import { getBudgetDurationLabel } from "../common_components/budget_duration_dropdown"; +import { CopyIcon, CheckIcon } from "lucide-react" +import NotificationsManager from "../molecules/notifications_manager" +import { getBudgetDurationLabel } from "../common_components/budget_duration_dropdown" interface UserInfoViewProps { userId: string @@ -57,16 +57,17 @@ export default function UserInfoView({ initialTab = 0, startInEditMode = false, }: UserInfoViewProps) { - const [userData, setUserData] = useState(null); - const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); - const [isLoading, setIsLoading] = useState(true); - const [isEditing, setIsEditing] = useState(startInEditMode); - const [userModels, setUserModels] = useState([]); - const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false); - const [invitationLinkData, setInvitationLinkData] = useState(null); - const [baseUrl, setBaseUrl] = useState(null); - const [activeTab, setActiveTab] = useState(initialTab); - const [copiedStates, setCopiedStates] = useState>({}); + const [userData, setUserData] = useState(null) + const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false) + const [isLoading, setIsLoading] = useState(true) + const [isEditing, setIsEditing] = useState(startInEditMode) + const [userModels, setUserModels] = useState([]) + const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false) + const [invitationLinkData, setInvitationLinkData] = useState(null) + const [baseUrl, setBaseUrl] = useState(null) + const [activeTab, setActiveTab] = useState(initialTab) + const [copiedStates, setCopiedStates] = useState>({}) + const [isTeamsExpanded, setIsTeamsExpanded] = useState(false) React.useEffect(() => { setBaseUrl(getProxyBaseUrl()) @@ -175,14 +176,14 @@ export default function UserInfoView({ } const copyToClipboard = async (text: string, key: string) => { - const success = await utilCopyToClipboard(text); + const success = await utilCopyToClipboard(text) if (success) { - setCopiedStates((prev) => ({ ...prev, [key]: true })); + setCopiedStates((prev) => ({ ...prev, [key]: true })) setTimeout(() => { - setCopiedStates((prev) => ({ ...prev, [key]: false })); - }, 2000); + setCopiedStates((prev) => ({ ...prev, [key]: false })) + }, 2000) } - }; + } return (
@@ -200,9 +201,9 @@ export default function UserInfoView({ icon={copiedStates["user-id"] ? : } onClick={() => copyToClipboard(userData.user_id, "user-id")} className={`left-2 z-10 transition-all duration-200 ${ - copiedStates["user-id"] - ? 'text-green-600 bg-green-50 border-green-200' - : 'text-gray-500 hover:text-gray-700 hover:bg-gray-100' + copiedStates["user-id"] + ? "text-green-600 bg-green-50 border-green-200" + : "text-gray-500 hover:text-gray-700 hover:bg-gray-100" }`} />
@@ -284,7 +285,35 @@ export default function UserInfoView({ Teams
- {userData.teams?.length || 0} teams + {userData.teams?.length && userData.teams?.length > 0 ? ( +
+ {userData.teams?.slice(0, isTeamsExpanded ? userData.teams.length : 20).map((team, index) => ( + + {team.team_alias} + + ))} + {!isTeamsExpanded && userData.teams?.length > 20 && ( + setIsTeamsExpanded(true)} + > + +{userData.teams.length - 20} more + + )} + {isTeamsExpanded && userData.teams?.length > 20 && ( + setIsTeamsExpanded(false)} + > + Show Less + + )} +
+ ) : ( + No teams + )}
@@ -344,9 +373,9 @@ export default function UserInfoView({ icon={copiedStates["user-id"] ? : } onClick={() => copyToClipboard(userData.user_id, "user-id")} className={`left-2 z-10 transition-all duration-200 ${ - copiedStates["user-id"] - ? 'text-green-600 bg-green-50 border-green-200' - : 'text-gray-500 hover:text-gray-700 hover:bg-gray-100' + copiedStates["user-id"] + ? "text-green-600 bg-green-50 border-green-200" + : "text-gray-500 hover:text-gray-700 hover:bg-gray-100" }`} />
@@ -384,11 +413,33 @@ export default function UserInfoView({ Teams
{userData.teams?.length && userData.teams?.length > 0 ? ( - userData.teams?.map((team, index) => ( - - {team.team_alias || team.team_id} - - )) + <> + {userData.teams?.slice(0, isTeamsExpanded ? userData.teams.length : 20).map((team, index) => ( + + {team.team_alias || team.team_id} + + ))} + {!isTeamsExpanded && userData.teams?.length > 20 && ( + setIsTeamsExpanded(true)} + > + +{userData.teams.length - 20} more + + )} + {isTeamsExpanded && userData.teams?.length > 20 && ( + setIsTeamsExpanded(false)} + > + Show Less + + )} + ) : ( No teams )}