diff --git a/.circleci/config.yml b/.circleci/config.yml index 0cedeb71686..d1440cf7a1b 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -616,6 +616,24 @@ jobs: wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz rm dockerize-linux-amd64-v0.6.1.tar.gz + - run: + name: Start PostgreSQL Database + command: | + docker run -d \ + --name postgres-db \ + -e POSTGRES_USER=postgres \ + -e POSTGRES_PASSWORD=postgres \ + -e POSTGRES_DB=circle_test \ + -p 5432:5432 \ + postgres:14 + - run: + name: Wait for PostgreSQL to be ready + command: dockerize -wait tcp://localhost:5432 -timeout 1m + - run: + name: Set DATABASE_URL environment variable + command: | + echo 'export DATABASE_URL="postgresql://postgres:postgres@localhost:5432/circle_test"' >> $BASH_ENV + source $BASH_ENV - run: name: Run Security Scans command: | diff --git a/README.md b/README.md index 0918d2b1fa4..c785ee82ffa 100644 --- a/README.md +++ b/README.md @@ -273,7 +273,7 @@ echo 'LITELLM_SALT_KEY="sk-1234"' >> .env source .env # Start -docker-compose up +docker compose up ``` diff --git a/docker/README.md b/docker/README.md index 1c3c208988c..ce478dfe0dd 100644 --- a/docker/README.md +++ b/docker/README.md @@ -28,7 +28,7 @@ Replace `your-secret-key` with a strong, randomly generated secret. Once you have set the `MASTER_KEY`, you can build and run the containers using the following command: ```bash -docker-compose up -d --build +docker compose up -d --build ``` This command will: @@ -42,13 +42,13 @@ This command will: You can check the status of the running containers with the following command: ```bash -docker-compose ps +docker compose ps ``` To view the logs of the `litellm` container, run: ```bash -docker-compose logs -f litellm +docker compose logs -f litellm ``` ### 4. Stopping the Application @@ -56,7 +56,7 @@ docker-compose logs -f litellm To stop the running containers, use the following command: ```bash -docker-compose down +docker compose down ``` ## Troubleshooting diff --git a/docs/my-website/docs/adding_provider/new_rerank_provider.md b/docs/my-website/docs/adding_provider/new_rerank_provider.md index 84c363261cd..628c0994434 100644 --- a/docs/my-website/docs/adding_provider/new_rerank_provider.md +++ b/docs/my-website/docs/adding_provider/new_rerank_provider.md @@ -17,7 +17,7 @@ class YourProviderRerankConfig(BaseRerankConfig): # ... other supported params ] - def transform_rerank_request(self, model: str, optional_rerank_params: OptionalRerankParams, headers: dict) -> dict: + def transform_rerank_request(self, model: str, optional_rerank_params: Dict, headers: dict) -> dict: # Transform request to RerankRequest spec return rerank_request.model_dump(exclude_none=True) diff --git a/docs/my-website/docs/contributing.md b/docs/my-website/docs/contributing.md index 8768e0b4c4d..a88013ff1b3 100644 --- a/docs/my-website/docs/contributing.md +++ b/docs/my-website/docs/contributing.md @@ -13,9 +13,6 @@ 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" @@ -30,7 +27,7 @@ python3 proxy_cli.py --config /path/to/config.yaml --port 4000 Set the mode as development (this will assume the proxy is running on localhost:4000) ```bash -export NODE_ENV="development" +npm install # install dependencies ``` ```bash diff --git a/docs/my-website/docs/embedding/supported_embedding.md b/docs/my-website/docs/embedding/supported_embedding.md index 1fd5a03e652..e63d9403665 100644 --- a/docs/my-website/docs/embedding/supported_embedding.md +++ b/docs/my-website/docs/embedding/supported_embedding.md @@ -266,7 +266,59 @@ print(response) | Titan Embeddings - G1 | `embedding(model="amazon.titan-embed-text-v1", input=input)` | | Cohere Embeddings - English | `embedding(model="cohere.embed-english-v3", input=input)` | | Cohere Embeddings - Multilingual | `embedding(model="cohere.embed-multilingual-v3", input=input)` | +| TwelveLabs Marengo (Async) | `embedding(model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", input=input, input_type="text")` | [Async Invoke Docs](../providers/bedrock_embedding#async-invoke-embedding) | +## TwelveLabs Bedrock Embedding Models + +TwelveLabs Marengo models support multimodal embeddings (text, image, video, audio) and require the `input_type` parameter to specify the input format. + +### Usage + +```python +from litellm import embedding +import os + +# Set AWS credentials +os.environ["AWS_ACCESS_KEY_ID"] = "" +os.environ["AWS_SECRET_ACCESS_KEY"] = "" +os.environ["AWS_REGION_NAME"] = "us-east-1" + +# Text embedding +response = embedding( + model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["Hello world from LiteLLM!"], + input_type="text" # Required parameter +) + +# Image embedding (base64) +response = embedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ..."], + input_type="image", # Required parameter + output_s3_uri="s3://your-bucket/async-invoke-output/" +) + +# Video embedding (S3 URL) +response = embedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["s3://your-bucket/video.mp4"], + input_type="video", # Required parameter + output_s3_uri="s3://your-bucket/async-invoke-output/" +) +``` + +### Required Parameters + +| Parameter | Description | Values | +|-----------|-------------|--------| +| `input_type` | Type of input content | `"text"`, `"image"`, `"video"`, `"audio"` | + +### Supported Models + +| Model Name | Function Call | Notes | +|------------|---------------|-------| +| TwelveLabs Marengo 2.7 (Sync) | `embedding(model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", input=input, input_type="text")` | Text embeddings only | +| TwelveLabs Marengo 2.7 (Async) | `embedding(model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", input=input, input_type="text/image/video/audio")` | All input types, requires `output_s3_uri` | ## Cohere Embedding Models https://docs.cohere.com/reference/embed diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 7eee979cc67..50bd5aefa76 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -1045,6 +1045,25 @@ curl --location 'http://localhost:4000/github_mcp/mcp' \ --- +## MCP Oauth + +LiteLLM v 1.77.6 added support for OAuth 2.0 Client Credentials for MCP servers. + + +This configuration is currently available on the config.yaml, with UI support coming soon. + +```yaml +mcp_servers: + github_mcp: + url: "https://api.githubcopilot.com/mcp" + auth_type: oauth2 + authorization_url: https://github.com/login/oauth/authorize + token_url: https://github.com/login/oauth/access_token + client_id: os.environ/GITHUB_OAUTH_CLIENT_ID + client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET + scopes: ["public_repo", "user:email"] +``` + ## Using your MCP with client side credentials Use this if you want to pass a client side authentication token to LiteLLM to then pass to your MCP to auth to your MCP. diff --git a/docs/my-website/docs/pass_through/vertex_ai.md b/docs/my-website/docs/pass_through/vertex_ai.md index d3f4e75e31d..77095667113 100644 --- a/docs/my-website/docs/pass_through/vertex_ai.md +++ b/docs/my-website/docs/pass_through/vertex_ai.md @@ -15,10 +15,11 @@ Pass-through endpoints for Vertex AI - call provider-specific endpoint, in nativ ## Supported Endpoints -LiteLLM supports 2 vertex ai passthrough routes: +LiteLLM supports 3 vertex ai passthrough routes: 1. `/vertex_ai` → routes to `https://{vertex_location}-aiplatform.googleapis.com/` 2. `/vertex_ai/discovery` → routes to [`https://discoveryengine.googleapis.com`](https://discoveryengine.googleapis.com/) +3. `/vertex_ai/live` → upgrades to the Vertex AI Live API WebSocket (`google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent`) ## How to use @@ -170,6 +171,50 @@ generateContent(); +## Vertex AI Live API WebSocket + +LiteLLM can now proxy the Vertex AI Live API to help you experiment with streaming audio/text from Gemini Live models without exposing Google credentials to clients. + +- Configure default Vertex credentials via `default_vertex_config` or environment variables (see examples above). +- Connect to `wss:///vertex_ai/live`. LiteLLM will exchange your saved credentials for a short-lived access token and forward messages bidirectionally. +- Optional query params `vertex_project`, `vertex_location`, and `model` let you override defaults for multi-project setups or global-only models. + +```python title="client.py" +import asyncio +import json + +from websockets.asyncio.client import connect + + +async def main() -> None: + headers = { + "x-litellm-api-key": "Bearer sk-your-litellm-key", + "Content-Type": "application/json", + } + async with connect( + "ws://localhost:4000/vertex_ai/live", + additional_headers=headers, + ) as ws: + await ws.send( + json.dumps( + { + "setup": { + "model": "projects/your-project/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09", + "generation_config": {"response_modalities": ["TEXT"]}, + } + } + ) + ) + + async for message in ws: + print("server:", message) + + +if __name__ == "__main__": + asyncio.run(main()) +``` + + ## Quick Start Let's call the Vertex AI [`/generateContent` endpoint](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference) @@ -415,4 +460,4 @@ generateContent(); ``` - \ No newline at end of file + diff --git a/docs/my-website/docs/pass_through/vertex_ai_live_websocket.md b/docs/my-website/docs/pass_through/vertex_ai_live_websocket.md new file mode 100644 index 00000000000..cca40d10fd8 --- /dev/null +++ b/docs/my-website/docs/pass_through/vertex_ai_live_websocket.md @@ -0,0 +1,284 @@ +# Vertex AI Live API WebSocket Passthrough + +LiteLLM now supports WebSocket passthrough for the Vertex AI Live API, enabling real-time bidirectional communication with Gemini models. + +## Overview + +The Vertex AI Live API WebSocket passthrough allows you to: +- Connect to Vertex AI Live API through LiteLLM proxy +- Use existing Vertex AI authentication methods +- Pass through all WebSocket messages bidirectionally +- Support text, audio, video, and multimodal interactions +- Track costs automatically for all usage types + +## Configuration + +### Environment Variables + +Set the following environment variables for Vertex AI authentication: + +```bash +# Required +DEFAULT_VERTEXAI_PROJECT=your-project-id +DEFAULT_VERTEXAI_LOCATION=us-central1 + +# Optional - use one of these for authentication +DEFAULT_GOOGLE_APPLICATION_CREDENTIALS=/path/to/service-account.json +# OR run: gcloud auth application-default login +``` + +### Configuration File + +Alternatively, configure in your `config.yaml`: + +```yaml +litellm_settings: + default_vertex_config: + vertex_project: "your-project-id" + vertex_location: "us-central1" + vertex_credentials: "os.environ/GOOGLE_APPLICATION_CREDENTIALS" +``` + +## Usage + +### WebSocket Endpoints + +- `ws://your-proxy-host/v1/vertex-ai/live` +- `ws://your-proxy-host/vertex-ai/live` + +### Query Parameters + +- `project_id` (optional): Google Cloud project ID (can be set in config) +- `location` (optional): Vertex AI location (can be set in config, default: us-central1) + +### Example Connection + +```javascript +// If project_id and location are set in config, you can connect without query params +const ws = new WebSocket('ws://localhost:4000/v1/vertex-ai/live'); + +// Or specify them explicitly +const ws = new WebSocket('ws://localhost:4000/v1/vertex-ai/live?project_id=your-project-id&location=us-central1'); +``` + +## Cost Tracking + +The WebSocket passthrough automatically tracks costs for all usage types based on the [Vertex AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing#model-optimizer-pricing): + +### Supported Cost Tracking + +- **Text**: Character-based or token-based pricing depending on model +- **Audio**: Per-second pricing for audio input/output +- **Video**: Per-second pricing for video input +- **Images**: Per-image pricing for image input + +### Cost Calculation + +Costs are calculated using the same methods as other Vertex AI models in LiteLLM: +- Uses `cost_per_character` for Gemini models +- Uses `cost_per_token` for partner models (Claude, Llama, etc.) +- Includes audio, video, and image costs when applicable + +### Cost Logging + +Costs are automatically logged to: +- LiteLLM proxy logs +- Database (if configured) +- Spend tracking system +- Admin dashboard + +Example log output: +``` +Vertex AI Live WebSocket session cost: $0.001234 (input: $0.000800, output: $0.000434) tokens: 150, characters: 1200, duration: 45.2s +``` + +## API Reference + +### Setup Message + +Send this message first to initialize the session: + +```json +{ + "setup": { + "model": "projects/your-project-id/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09", + "generation_config": { + "response_modalities": ["TEXT"] + } + } +} +``` + +### Text Input + +```json +{ + "client_content": { + "turns": [ + { + "role": "user", + "parts": [{"text": "Hello! How are you?"}] + } + ], + "turn_complete": true + } +} +``` + +### Audio Input + +```json +{ + "realtime_input": { + "media_chunks": [ + { + "data": "base64-encoded-audio-data", + "mime_type": "audio/pcm" + } + ] + } +} +``` + +## Supported Features + +### Response Modalities + +- **TEXT**: Text responses +- **AUDIO**: Audio responses with voice synthesis + +### Tools + +- **Function Calling**: Define and use custom functions +- **Code Execution**: Execute Python code +- **Google Search**: Search the web +- **Voice Activity Detection**: Detect when user is speaking + +### Advanced Features + +- **Audio Transcription**: Transcribe input and output audio +- **Proactive Audio**: Model responds only when relevant +- **Affective Dialog**: Understand emotional expressions + +## Examples + +### Python Client + +```python +import asyncio +import json +import websockets + +async def chat_with_gemini(): + uri = "ws://localhost:4000/v1/vertex-ai/live?project_id=your-project-id" + + async with websockets.connect(uri) as websocket: + # Setup + setup = { + "setup": { + "model": "projects/your-project-id/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09", + "generation_config": {"response_modalities": ["TEXT"]} + } + } + await websocket.send(json.dumps(setup)) + + # Wait for setup response + response = await websocket.recv() + print(f"Setup: {response}") + + # Send message + message = { + "client_content": { + "turns": [{"role": "user", "parts": [{"text": "Hello!"}]}], + "turn_complete": True + } + } + await websocket.send(json.dumps(message)) + + # Receive response + async for response in websocket: + print(f"Response: {response}") + # Check if turn is complete + data = json.loads(response) + if data.get("serverContent", {}).get("turnComplete"): + break + +asyncio.run(chat_with_gemini()) +``` + +### JavaScript Client + +```javascript +const ws = new WebSocket('ws://localhost:4000/v1/vertex-ai/live?project_id=your-project-id'); + +ws.onopen = function() { + // Send setup + const setup = { + setup: { + model: "projects/your-project-id/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09", + generation_config: { response_modalities: ["TEXT"] } + } + }; + ws.send(JSON.stringify(setup)); +}; + +ws.onmessage = function(event) { + const data = JSON.parse(event.data); + console.log('Received:', data); + + // Check if setup is complete + if (data.setupComplete) { + // Send a message + const message = { + client_content: { + turns: [{ role: "user", parts: [{ text: "Hello!" }] }], + turn_complete: true + } + }; + ws.send(JSON.stringify(message)); + } +}; +``` + +## Error Handling + +The WebSocket connection may close with these codes: + +- `4001`: Vertex AI credentials not configured +- `4002`: Project ID not provided +- `1011`: Internal server error + +## Authentication + +The WebSocket passthrough uses the same authentication as other LiteLLM endpoints: + +1. **API Key**: Pass `Authorization: Bearer your-api-key` header +2. **Vertex AI Credentials**: Set environment variables or config file + +## Limitations + +- Requires valid Google Cloud project with Vertex AI API enabled +- WebSocket connections are not persistent across server restarts +- Rate limits apply based on your Google Cloud quotas + +## Troubleshooting + +### Common Issues + +1. **Authentication Error**: Ensure Vertex AI credentials are properly configured +2. **Project Not Found**: Verify the project ID exists and has Vertex AI enabled +3. **Connection Refused**: Check that the LiteLLM proxy server is running + +### Debug Mode + +Enable debug logging to see detailed connection information: + +```bash +export LITELLM_LOG=DEBUG +``` + +## Related Documentation + +- [Vertex AI Live API Reference](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/multimodal-live) +- [LiteLLM Proxy Configuration](../proxy/) +- [Vertex AI Passthrough Endpoints](./vertex_ai.md) diff --git a/docs/my-website/docs/projects/Railtracks.md b/docs/my-website/docs/projects/Railtracks.md new file mode 100644 index 00000000000..3b94ec8df43 --- /dev/null +++ b/docs/my-website/docs/projects/Railtracks.md @@ -0,0 +1,7 @@ +# Railtracks + +`Railtracks` is an open-source agentic framework that helps developers build resilient agentic systems offering local and remote monitoring tools. + +- [Github](https://github.com/RailtownAI/railtracks) +- [Docs](https://railtownai.github.io/railtracks/) +- [Railtracks](https://railtracks.org/) \ No newline at end of file diff --git a/docs/my-website/docs/providers/bedrock_embedding.md b/docs/my-website/docs/providers/bedrock_embedding.md index 95ee8d3d228..69c5f3c86ec 100644 --- a/docs/my-website/docs/providers/bedrock_embedding.md +++ b/docs/my-website/docs/providers/bedrock_embedding.md @@ -8,6 +8,182 @@ | Cohere | `bedrock/cohere.*` | [Cohere Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-embed.html) | | TwelveLabs | `bedrock/us.twelvelabs.*` | [TwelveLabs](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-twelvelabs.html) | +## Async Invoke Support + +LiteLLM supports AWS Bedrock's async-invoke feature for embedding models that require asynchronous processing, particularly useful for large media files (video, audio) or when you need to process embeddings in the background. + +### Supported Models + +| Provider | Async Invoke Route | Use Case | +|----------|-------------------|----------| +| TwelveLabs Marengo | `bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0` | Video, audio, image, and text embeddings | + +### Required Parameters + +When using async-invoke, you must provide: + +| Parameter | Description | Required | +|-----------|-------------|----------| +| `output_s3_uri` | S3 URI where the embedding results will be stored | ✅ Yes | +| `input_type` | Type of input: `"text"`, `"image"`, `"video"`, or `"audio"` | ✅ Yes | +| `aws_region_name` | AWS region for the request | ✅ Yes | + +### Usage + +#### Basic Async Invoke + +```python +from litellm import embedding + +# Text embedding with async-invoke +response = embedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["Hello world from LiteLLM async invoke!"], + aws_region_name="us-east-1", + input_type="text", + output_s3_uri="s3://your-bucket/async-invoke-output/" +) + +print(f"Job submitted! Invocation ARN: {response._hidden_params._invocation_arn}") +``` + +#### Video/Audio Embedding + +```python +# Video embedding (requires async-invoke) +response = embedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["s3://your-bucket/video.mp4"], # S3 URL for video + aws_region_name="us-east-1", + input_type="video", + output_s3_uri="s3://your-bucket/async-invoke-output/" +) + +print(f"Video embedding job submitted! ARN: {response._hidden_params._invocation_arn}") +``` + +#### Image Embedding with Base64 + +```python +import base64 + +# Load and encode image +with open("image.jpg", "rb") as img_file: + img_data = base64.b64encode(img_file.read()).decode('utf-8') + img_base64 = f"data:image/jpeg;base64,{img_data}" + +response = embedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=[img_base64], + aws_region_name="us-east-1", + input_type="image", + output_s3_uri="s3://your-bucket/async-invoke-output/" +) +``` + +### Retrieving Job Information + +#### Getting Job ID and Invocation ARN + +The async-invoke response includes the invocation ARN in the hidden parameters: + +```python +response = embedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["Hello world"], + aws_region_name="us-east-1", + input_type="text", + output_s3_uri="s3://your-bucket/async-invoke-output/" +) + +# Access invocation ARN +invocation_arn = response._hidden_params._invocation_arn +print(f"Invocation ARN: {invocation_arn}") + +# Extract job ID from ARN (last part after the last slash) +job_id = invocation_arn.split("/")[-1] +print(f"Job ID: {job_id}") +``` + +#### Checking Job Status + +Use LiteLLM's `retrieve_batch` function to check if your job is still processing: + +```python +from litellm import retrieve_batch + +def check_async_job_status(invocation_arn, aws_region_name="us-east-1"): + """Check the status of an async invoke job using LiteLLM batch API""" + try: + response = retrieve_batch( + batch_id=invocation_arn, + custom_llm_provider="bedrock", + aws_region_name=aws_region_name + ) + return response + except Exception as e: + print(f"Error checking job status: {e}") + return None + +# Check status +status = check_async_job_status(invocation_arn, "us-east-1") +if status: + print(f"Job Status: {status.status}") + print(f"Output Location: {status.output_file_id}") +``` + +**Note:** The actual embedding results are stored in S3. The `output_file_id` from the batch status can be used to locate the results file in your S3 bucket. + +### Error Handling + +#### Common Errors + +| Error | Cause | Solution | +|-------|-------|----------| +| `ValueError: output_s3_uri cannot be empty` | Missing S3 output URI | Provide a valid S3 URI | +| `ValueError: Input type 'video' requires async_invoke route` | Using video/audio without async-invoke | Use `bedrock/async_invoke/` model prefix | +| `ValueError: input_type is required` | Missing input type parameter | Specify `input_type` parameter | + +#### Example Error Handling + +```python +try: + response = embedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["Hello world"], + aws_region_name="us-east-1", + input_type="text", + output_s3_uri="s3://your-bucket/output/" # Required for async-invoke + ) + print("Job submitted successfully!") + +except ValueError as e: + if "output_s3_uri cannot be empty" in str(e): + print("Error: Please provide a valid S3 output URI") + elif "requires async_invoke route" in str(e): + print("Error: Use async_invoke model for video/audio inputs") + else: + print(f"Error: {e}") +except Exception as e: + print(f"Unexpected error: {e}") +``` + +### Best Practices + +1. **Use async-invoke for large files**: Video and audio files are better processed asynchronously +2. **Use LiteLLM batch API**: Use `retrieve_batch()` instead of direct Bedrock API calls for status checking +3. **Monitor job status**: Check job status periodically using the batch API to know when results are ready +4. **Handle errors gracefully**: Implement proper error handling for network issues and job failures +5. **Set appropriate timeouts**: Consider the processing time for large files +6. **Use S3 for large inputs**: For video/audio, use S3 URLs instead of base64 encoding + +### Limitations + +- Async-invoke is currently only supported for TwelveLabs Marengo models +- Results are stored in S3 and must be retrieved separately using the output file ID +- Job status checking requires using LiteLLM's `retrieve_batch()` function +- No built-in polling mechanism in LiteLLM (must implement your own status checking loop) + ### API keys This can be set as env variables or passed as **params to litellm.embedding()** ```python diff --git a/docs/my-website/docs/providers/lemonade.md b/docs/my-website/docs/providers/lemonade.md index 8ff7d48b706..fc77b78a76c 100644 --- a/docs/my-website/docs/providers/lemonade.md +++ b/docs/my-website/docs/providers/lemonade.md @@ -186,3 +186,6 @@ print("Available models:", [model['id'] for model in models.get('data', [])]) ## Support For more information regarding Lemonade please go to to the [Lemonade website](https://lemonade-server.ai/) or [Lemonade repository](https://github.com/lemonade-sdk/lemonade). + + + diff --git a/docs/my-website/docs/providers/nvidia_nim.md b/docs/my-website/docs/providers/nvidia_nim.md index 270b356c917..9dbfc80f4e4 100644 --- a/docs/my-website/docs/providers/nvidia_nim.md +++ b/docs/my-website/docs/providers/nvidia_nim.md @@ -15,8 +15,8 @@ https://docs.api.nvidia.com/nim/reference/ | Description | Nvidia NIM is a platform that provides a simple API for deploying and using AI models. LiteLLM supports all models from [Nvidia NIM](https://developer.nvidia.com/nim/) | | Provider Route on LiteLLM | `nvidia_nim/` | | Provider Doc | [Nvidia NIM Docs ↗](https://developer.nvidia.com/nim/) | -| API Endpoint for Provider | https://integrate.api.nvidia.com/v1/ | -| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/responses`, `/embeddings` | +| API Endpoint for Provider | https://integrate.api.nvidia.com/v1/ (chat/embeddings), https://ai.api.nvidia.com/v1/ (rerank) | +| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/responses`, `/embeddings`, `/rerank` | ## API Key ```python diff --git a/docs/my-website/docs/providers/nvidia_nim_rerank.md b/docs/my-website/docs/providers/nvidia_nim_rerank.md new file mode 100644 index 00000000000..7373014a960 --- /dev/null +++ b/docs/my-website/docs/providers/nvidia_nim_rerank.md @@ -0,0 +1,261 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Nvidia NIM - Rerank + +Use Nvidia NIM Rerank models through LiteLLM. + +| Property | Details | +|----------|---------| +| Description | Nvidia NIM provides high-performance reranking models for semantic search and retrieval-augmented generation (RAG) | +| Provider Doc | [Nvidia NIM Rerank API ↗](https://docs.api.nvidia.com/nim/reference/nvidia-llama-3_2-nv-rerankqa-1b-v2-infer) | +| Supported Endpoint | `/rerank` | + +## Overview + +Nvidia NIM rerank models help you: +- Reorder search results by relevance to a query +- Improve RAG (Retrieval-Augmented Generation) accuracy +- Filter and rank large document sets efficiently + +**Supported Models:** +- All Nvidia NIM rerank models on their platform + +:::tip + +See the full list of LiteLLM supported Nvidia NIM rerank models on [Nvidia NIM](https://models.litellm.ai) + +::: + +## Usage + +### LiteLLM Python SDK + + + + +```python +import litellm +import os + +os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..." + +response = litellm.rerank( + model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2", + query="What is the GPU memory bandwidth of H100 SXM?", + documents=[ + "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth.", + "A100 provides up to 20X higher performance over the prior generation.", + "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU." + ], + top_n=3, +) + +print(response) +``` + + + + +```python +import litellm +import os + +os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..." + +response = litellm.rerank( + model="nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3", + query="What is the GPU memory bandwidth of H100 SXM?", + documents=[ + "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth.", + "A100 provides up to 20X higher performance over the prior generation.", + "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU." + ], + top_n=3, +) + +print(response) +``` + + + + +**Response:** +```json +{ + "results": [ + { + "index": 2, + "relevance_score": 6.828125, + "document": { + "text": "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU." + } + }, + { + "index": 0, + "relevance_score": -1.564453125, + "document": { + "text": "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth." + } + } + ] +} +``` + + +## Usage with LiteLLM Proxy + +### 1. Setup Config + +Add Nvidia NIM rerank models to your proxy configuration: + +```yaml +model_list: + - model_name: nvidia-rerank + litellm_params: + model: nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2 + api_key: os.environ/NVIDIA_NIM_API_KEY +``` + +### 2. Start Proxy + +```bash +litellm --config /path/to/config.yaml +``` + +### 3. Make Rerank Requests + +```bash +curl -X POST http://0.0.0.0:4000/rerank \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "nvidia-rerank", + "query": "What is the GPU memory bandwidth of H100?", + "documents": [ + "H100 delivers 3TB/s memory bandwidth", + "A100 has 2TB/s memory bandwidth", + "V100 offers 900GB/s memory bandwidth" + ], + "top_n": 2 + }' +``` + +## API Parameters + +### Required Parameters + +| Parameter | Type | Description | +|-----------|------|-------------| +| `model` | string | The Nvidia NIM rerank model name with `nvidia_nim/` prefix | +| `query` | string | The search query to rank documents against | +| `documents` | array | List of documents to rank (1-1000 documents) | + +### Optional Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `top_n` | integer | All documents | Number of top-ranked documents to return | + +### Nvidia-Specific Parameters + +**`truncate`**: Controls how text is truncated if it exceeds the model's context window +- `"NONE"`: No truncation (request may fail if too long) +- `"END"`: Truncate from the end of the text + +```python +response = litellm.rerank( + model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2", + query="GPU performance", + documents=["High performance computing", "Fast GPU processing"], + top_n=2, + truncate="END", # Nvidia-specific parameter +) +``` + +## Authentication + +Set your Nvidia NIM API key: + + + + +```bash +export NVIDIA_NIM_API_KEY="nvapi-..." +``` + + + + +```python +import os +os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..." + +# Or pass directly +response = litellm.rerank( + model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2", + query="test", + documents=["doc1"], + api_key="nvapi-...", +) +``` + + + + +## API Endpoint + +The rerank endpoint uses a different base URL than chat/embeddings: + +- **Chat/Embeddings:** `https://integrate.api.nvidia.com/v1/` +- **Rerank:** `https://ai.api.nvidia.com/v1/` + +LiteLLM automatically uses the correct endpoint for rerank requests. + +### Custom API Base URL + +You can override the default base URL in several ways: + +**Option 1: Environment Variable** + +```bash +export NVIDIA_NIM_API_BASE="https://your-custom-endpoint.com" +``` + +**Option 2: Pass as parameter** + +```python +response = litellm.rerank( + model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2", + query="test", + documents=["doc1"], + api_base="https://your-custom-endpoint.com", +) +``` + +**Option 3: Full URL (including model path)** + +If you have the complete endpoint URL, you can pass it directly: + +```python +response = litellm.rerank( + model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2", + query="test", + documents=["doc1"], + api_base="https://your-custom-endpoint.com/v1/retrieval/nvidia/llama-3_2-nv-rerankqa-1b-v2/reranking", +) +``` + +LiteLLM will detect the full URL (by checking for `/retrieval/` in the path) and use it as-is. + +### How do I get an API key? + +Get your Nvidia NIM API key from [Nvidia's website](https://developer.nvidia.com/nim/). + +## Related Documentation + +- [Nvidia NIM - Main Documentation](./nvidia_nim) +- [Nvidia NIM Chat Completions](./nvidia_nim#sample-usage) +- [LiteLLM Rerank Endpoint](../rerank) +- [Nvidia NIM Official Docs ↗](https://docs.api.nvidia.com/nim/reference/) + diff --git a/docs/my-website/docs/providers/oci.md b/docs/my-website/docs/providers/oci.md index 6fc1835154a..c11d64f4553 100644 --- a/docs/my-website/docs/providers/oci.md +++ b/docs/my-website/docs/providers/oci.md @@ -44,6 +44,7 @@ response = completion( oci_user=, oci_fingerprint=, oci_tenancy=, + oci_serving_mode="ON_DEMAND", # Optional, default is "ON_DEMAND". Other option is "DEDICATED" # Provide either the private key string OR the path to the key file: # Option 1: pass the private key as a string oci_key=, @@ -71,6 +72,7 @@ response = completion( oci_user=, oci_fingerprint=, oci_tenancy=, + oci_serving_mode="ON_DEMAND", # Optional, default is "ON_DEMAND". Other option is "DEDICATED" # Provide either the private key string OR the path to the key file: # Option 1: pass the private key as a string oci_key=, diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 943823c6386..2543c15d135 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -621,6 +621,163 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \ +#### **Google Maps** + +Use Google Maps to provide location-based context to your Gemini models. + +[**Relevant Vertex AI Docs**](https://ai.google.dev/gemini-api/docs/grounding#google-maps) + + + + +**Basic Usage - Enable Widget Only** + +```python showLineNumbers +from litellm import completion + +## SETUP ENVIRONMENT +# !gcloud auth application-default login - run this to add vertex credentials to your env + +tools = [{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}] # 👈 ADD GOOGLE MAPS + +resp = litellm.completion( + model="vertex_ai/gemini-2.0-flash", + messages=[{"role": "user", "content": "What restaurants are nearby?"}], + tools=tools, +) + +print(resp) +``` + +**With Location Data** + +You can specify a location to ground the model's responses with location-specific information: + +```python showLineNumbers +from litellm import completion + +## SETUP ENVIRONMENT +# !gcloud auth application-default login - run this to add vertex credentials to your env + +tools = [{ + "googleMaps": { + "enableWidget": "ENABLE_WIDGET", + "latitude": 37.7749, # San Francisco latitude + "longitude": -122.4194, # San Francisco longitude + "languageCode": "en_US" # Optional: language for results + } +}] # 👈 ADD GOOGLE MAPS WITH LOCATION + +resp = litellm.completion( + model="vertex_ai/gemini-2.0-flash", + messages=[{"role": "user", "content": "What restaurants are nearby?"}], + tools=tools, +) + +print(resp) +``` + + + + + + + +**Basic Usage - Enable Widget Only** + +```python showLineNumbers +from openai import OpenAI + +client = OpenAI( + api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys + base_url="http://0.0.0.0:4000/v1/" # point to litellm proxy +) + +response = client.chat.completions.create( + model="gemini-2.0-flash", + messages=[{"role": "user", "content": "What restaurants are nearby?"}], + tools=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}], +) + +print(response) +``` + +**With Location Data** + +```python showLineNumbers +from openai import OpenAI + +client = OpenAI( + api_key="sk-1234", # pass litellm proxy key, if you're using virtual keys + base_url="http://0.0.0.0:4000/v1/" # point to litellm proxy +) + +response = client.chat.completions.create( + model="gemini-2.0-flash", + messages=[{"role": "user", "content": "What restaurants are nearby?"}], + tools=[{ + "googleMaps": { + "enableWidget": "ENABLE_WIDGET", + "latitude": 37.7749, # San Francisco latitude + "longitude": -122.4194, # San Francisco longitude + "languageCode": "en_US" # Optional: language for results + } + }], +) + +print(response) +``` + + + +**Basic Usage - Enable Widget Only** + +```bash showLineNumbers +curl http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "gemini-2.0-flash", + "messages": [ + {"role": "user", "content": "What restaurants are nearby?"} + ], + "tools": [ + { + "googleMaps": {"enableWidget": "ENABLE_WIDGET"} + } + ] + }' +``` + +**With Location Data** + +```bash showLineNumbers +curl http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "gemini-2.0-flash", + "messages": [ + {"role": "user", "content": "What restaurants are nearby?"} + ], + "tools": [ + { + "googleMaps": { + "enableWidget": "ENABLE_WIDGET", + "latitude": 37.7749, + "longitude": -122.4194, + "languageCode": "en_US" + } + } + ] + }' +``` + + + + + + #### **Moving from Vertex AI SDK to LiteLLM (GROUNDING)** diff --git a/docs/my-website/docs/proxy/cost_tracking.md b/docs/my-website/docs/proxy/cost_tracking.md index 35db752cbb6..85147e12c66 100644 --- a/docs/my-website/docs/proxy/cost_tracking.md +++ b/docs/my-website/docs/proxy/cost_tracking.md @@ -8,6 +8,10 @@ Track spend for keys, users, and teams across 100+ LLMs. LiteLLM automatically tracks spend for all known models. See our [model cost map](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json) +:::tip Keep Pricing Data Updated +[Sync model pricing data from GitHub](../sync_models_github.md) to ensure accurate cost tracking. +::: + ### How to Track Spend with LiteLLM **Step 1** diff --git a/docs/my-website/docs/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index 6a11d069fb0..d731c0a3c1d 100644 --- a/docs/my-website/docs/proxy/deploy.md +++ b/docs/my-website/docs/proxy/deploy.md @@ -27,7 +27,7 @@ echo 'LITELLM_SALT_KEY="sk-1234"' >> .env source .env # Start -docker-compose up +docker compose up ``` @@ -715,6 +715,25 @@ docker run ghcr.io/berriai/litellm:main-stable ``` +### Restart Workers After N Requests + +Use this to mitigate memory growth by recycling workers after a fixed number of requests. When set, each worker restarts after completing the specified number of requests. Defaults to disabled when unset. + +Usage Examples: + +```shell showLineNumbers title="docker run (CLI flag)" +docker run ghcr.io/berriai/litellm:main-stable \ + --max_requests_before_restart 10000 +``` + +Or set via environment variable: + +```shell showLineNumbers title="Environment Variable" +export MAX_REQUESTS_BEFORE_RESTART=10000 +docker run ghcr.io/berriai/litellm:main-stable +``` + + ### 5. config.yaml file on s3, GCS Bucket Object/url Use this if you cannot mount a config file on your deployment service (example - AWS Fargate, Railway etc) diff --git a/docs/my-website/docs/proxy/docker_quick_start.md b/docs/my-website/docs/proxy/docker_quick_start.md index 1bb5150dc21..f3da18065ec 100644 --- a/docs/my-website/docs/proxy/docker_quick_start.md +++ b/docs/my-website/docs/proxy/docker_quick_start.md @@ -55,7 +55,7 @@ echo 'LITELLM_SALT_KEY="sk-1234"' >> .env source .env # Start -docker-compose up +docker compose up ``` diff --git a/docs/my-website/docs/proxy/logging_spec.md b/docs/my-website/docs/proxy/logging_spec.md index 205282428ee..6364b8c4444 100644 --- a/docs/my-website/docs/proxy/logging_spec.md +++ b/docs/my-website/docs/proxy/logging_spec.md @@ -14,6 +14,7 @@ Found under `kwargs["standard_logging_object"]`. This is a standard payload, log | `cost_breakdown` | `Optional[CostBreakdown]` | Detailed cost breakdown object | | `response_cost_failure_debug_info` | `StandardLoggingModelCostFailureDebugInformation` | Debug information if cost tracking fails | | `status` | `StandardLoggingPayloadStatus` | Status of the payload | +| `status_fields` | `StandardLoggingPayloadStatusFields` | Typed status fields for easy filtering and analytics | | `total_tokens` | `int` | Total number of tokens | | `prompt_tokens` | `int` | Number of prompt tokens | | `completion_tokens` | `int` | Number of completion tokens | @@ -162,17 +163,89 @@ A literal type with two possible values: ## StandardLoggingGuardrailInformation +| Field | Type | Description | +|-----------------------|------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `guardrail_name` | `Optional[str]` | Guardrail name | +| `guardrail_provider` | `Optional[str]` | Guardrail provider | +| `guardrail_mode` | `Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]]` | Guardrail mode | +| `guardrail_request` | `Optional[dict]` | Guardrail request | +| `guardrail_response` | `Optional[Union[dict, str, List[dict]]]` | Guardrail response | +| `guardrail_status` | `Literal["success", "failure", "blocked"]` | Guardrail execution status: `success` = no violations detected, `blocked` = content blocked/modified due to policy violations, `failure` = technical error or API failure | +| `start_time` | `Optional[float]` | Start time of the guardrail | +| `end_time` | `Optional[float]` | End time of the guardrail | +| `duration` | `Optional[float]` | Duration of the guardrail in seconds | +| `masked_entity_count` | `Optional[Dict[str, int]]` | Count of masked entities | + +## StandardLoggingPayloadStatusFields + +Typed status fields for easy filtering and analytics. + | Field | Type | Description | |-------|------|-------------| -| `guardrail_name` | `Optional[str]` | Guardrail name | -| `guardrail_mode` | `Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]]` | Guardrail mode | -| `guardrail_request` | `Optional[dict]` | Guardrail request | -| `guardrail_response` | `Optional[Union[dict, str, List[dict]]]` | Guardrail response | -| `guardrail_status` | `Literal["success", "failure"]` | Guardrail status | -| `start_time` | `Optional[float]` | Start time of the guardrail | -| `end_time` | `Optional[float]` | End time of the guardrail | -| `duration` | `Optional[float]` | Duration of the guardrail in seconds | -| `masked_entity_count` | `Optional[Dict[str, int]]` | Count of masked entities | +| `llm_api_status` | `StandardLoggingPayloadStatus` | Status of the LLM API call: `"success"` if completed successfully, `"failure"` if errored | +| `guardrail_status` | `GuardrailStatus` | Status of guardrail execution (see below) | + +### StandardLoggingPayloadStatus + +A literal type with two possible values: +- `"success"` - The LLM API request completed successfully +- `"failure"` - The LLM API request failed + +### GuardrailStatus + +A literal type with four possible values: +- `"success"` - Guardrail ran and allowed content through (no violations detected) +- `"guardrail_intervened"` - Guardrail blocked or modified content due to policy violations +- `"guardrail_failed_to_respond"` - Guardrail had a technical failure or API error +- `"not_run"` - No guardrail was executed for this request + +### Usage Examples + +Filter logs for requests where guardrails intervened: +```json +{ + "status_fields": { + "guardrail_status": "guardrail_intervened" + } +} +``` + +Find guardrail technical failures: +```json +{ + "status_fields": { + "guardrail_status": "guardrail_failed_to_respond" + } +} +``` + +Get successful LLM requests: +```json +{ + "status_fields": { + "llm_api_status": "success" + } +} +``` + +Find requests where guardrails ran successfully without intervention: +```json +{ + "status_fields": { + "guardrail_status": "success", + "llm_api_status": "success" + } +} +``` + +Find requests where no guardrail was run: +```json +{ + "status_fields": { + "guardrail_status": "not_run" + } +} +``` ## StandardLoggingPromptManagementMetadata diff --git a/docs/my-website/docs/proxy/model_management.md b/docs/my-website/docs/proxy/model_management.md index a8cc66ae765..6a87dda2f42 100644 --- a/docs/my-website/docs/proxy/model_management.md +++ b/docs/my-website/docs/proxy/model_management.md @@ -19,6 +19,10 @@ model_list: Retrieve detailed information about each model listed in the `/model/info` endpoint, including descriptions from the `config.yaml` file, and additional model info (e.g. max tokens, cost per input token, etc.) pulled from the model_info you set and the [litellm model cost map](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). Sensitive details like API keys are excluded for security purposes. +:::tip Sync Model Data +Keep your model pricing data up to date by [syncing models from GitHub](../sync_models_github.md). +::: + @@ -90,6 +90,51 @@ response = litellm.completion( ``` + + +**1. Create a .prompt file in a gitlab repo** + +Create `prompts/hello.prompt` in your gitlab repository: + +```yaml +--- +model: gpt-4 +temperature: 0.7 +--- +System: You are a helpful assistant. + +User: {{user_message}} +``` + +**2. Configure Gitlab access** + +```python +import litellm + +# Configure gitlab access +gitlab_config = { + "workspace": "your-workspace", + "repository": "your-repo", + "access_token": "your-access-token", + "branch": "main" +} + +# Set global gitlab configuration +litellm.set_global_gitlab_config(gitlab_config) +``` + +**3. Use with LiteLLM** + +```python +response = litellm.completion( + model="gitlab/gpt-4", + prompt_id="hello", + prompt_variables={"user_message": "What is the capital of France?"} +) +``` + + + **1. Create a .prompt file** @@ -124,6 +169,12 @@ litellm_settings: repository: "your-repo" access_token: "your-access-token" branch: "main" + # Or use Gitlab for team-based prompt management + global_gitlab_config: + workspace: "your-workspace" + repository: "your-repo" + access_token: "your-access-token" + branch: "main" ``` **3. Start the proxy** @@ -213,6 +264,14 @@ prompt_variables: Optional[dict] # optional - variables for template rendering bitbucket_config: Optional[dict] # optional - BitBucket configuration (if not set globally) ``` +**Gitlab:** +``` +model: gitlab/ # required (e.g., gitlab/gpt-4) +prompt_id: str # required - the .prompt filename without extension +prompt_variables: Optional[dict] # optional - variables for template rendering +gitlab_config: Optional[dict] # optional - Gitlab configuration (if not set globally) +``` + **Example API calls:** ```python @@ -235,4 +294,18 @@ response = litellm.completion( "access_token": "your-token" } ) + +# Gitlab integration +response = litellm.completion( + model="gitlab/gpt-4", + prompt_id="hello", + prompt_variables={"user_message": "Hello world"}, + gitlab_config={ + "project": "a/b/", + "access_token": "your-access-token", + "base_url": "gitlab url", + "prompts_path": "src/prompts", # folder to point to, defaults to root + "branch":"main" # optional, defaults to main + } +) ``` diff --git a/docs/my-website/docs/proxy/pass_through.md b/docs/my-website/docs/proxy/pass_through.md index b7978d9f655..7309cdeda26 100644 --- a/docs/my-website/docs/proxy/pass_through.md +++ b/docs/my-website/docs/proxy/pass_through.md @@ -243,6 +243,18 @@ curl --location 'http://0.0.0.0:4000/v1/messages' \ }' ``` +--- + +## Tutorial - Add Azure OpenAI Assistants API as a Pass Through Endpoint + +In this video, we'll add the Azure OpenAI Assistants API as a pass through endpoint to LiteLLM Proxy. + + + +
+
+ + --- ## Troubleshooting diff --git a/docs/my-website/docs/proxy/prod.md b/docs/my-website/docs/proxy/prod.md index a45474f39e8..2858132c8e8 100644 --- a/docs/my-website/docs/proxy/prod.md +++ b/docs/my-website/docs/proxy/prod.md @@ -71,6 +71,16 @@ Use this Docker `CMD`. This will start the proxy with 1 Uvicorn Async Worker CMD ["--port", "4000", "--config", "./proxy_server_config.yaml"] ``` +> Optional: If you observe gradual memory growth under sustained load, consider recycling workers after a fixed number of requests to mitigate leaks. Set this via CLI or environment variable: + +```shell +# CLI +CMD ["--port", "4000", "--config", "./proxy_server_config.yaml", "--max_requests_before_restart", "10000"] + +# or ENV (for deployment manifests / containers) +export MAX_REQUESTS_BEFORE_RESTART=10000 +``` + ## 4. Use Redis 'port','host', 'password'. NOT 'redis_url' diff --git a/docs/my-website/docs/proxy/sync_models_github.md b/docs/my-website/docs/proxy/sync_models_github.md new file mode 100644 index 00000000000..d2f410e5496 --- /dev/null +++ b/docs/my-website/docs/proxy/sync_models_github.md @@ -0,0 +1,61 @@ +# Syncing Models to GitHub model_context_window + +Sync model pricing data from GitHub's `model_prices_and_context_window.json` file outside of the LiteLLM UI. + +> **📹 Video Tutorial**: [Watch how to sync models via the Admin UI](https://www.loom.com/share/ba41acc1882d41b284bbddbb0e9c27ce?sid=bdae351e-2026-4e39-932b-fcb185ff612c) + +## Quick Start + +**Manual sync:** +```bash +curl -X POST "https://your-proxy-url/reload/model_cost_map" \ + -H "Authorization: Bearer YOUR_ADMIN_TOKEN" \ + -H "Content-Type: application/json" +``` + +**Automatic sync every 6 hours:** +```bash +curl -X POST "https://your-proxy-url/schedule/model_cost_map_reload?hours=6" \ + -H "Authorization: Bearer YOUR_ADMIN_TOKEN" \ + -H "Content-Type: application/json" +``` + +## API Endpoints + +| Endpoint | Method | Description | +|----------|--------|-------------| +| `/reload/model_cost_map` | POST | Manual sync | +| `/schedule/model_cost_map_reload?hours={hours}` | POST | Schedule periodic sync | +| `/schedule/model_cost_map_reload` | DELETE | Cancel scheduled sync | +| `/schedule/model_cost_map_reload/status` | GET | Check sync status | + +**Authentication:** Requires admin role or master key + +## Python Example + +```python +import requests + +def sync_models(proxy_url, admin_token): + response = requests.post( + f"{proxy_url}/reload/model_cost_map", + headers={"Authorization": f"Bearer {admin_token}"} + ) + return response.json() + +# Usage +result = sync_models("https://your-proxy-url", "your-admin-token") +print(result['message']) +``` + +## Configuration + +**Custom model cost map URL:** +```bash +export LITELLM_MODEL_COST_MAP_URL="https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" +``` + +**Use local model cost map:** +```bash +export LITELLM_LOCAL_MODEL_COST_MAP=True +``` \ No newline at end of file diff --git a/docs/my-website/docs/proxy/ui.md b/docs/my-website/docs/proxy/ui.md index a093b226a27..f7419d20740 100644 --- a/docs/my-website/docs/proxy/ui.md +++ b/docs/my-website/docs/proxy/ui.md @@ -54,6 +54,20 @@ Allow others to create/delete their own keys. [**Go Here**](./self_serve.md) +## Model Management + +The Admin UI provides comprehensive model management capabilities: + +- **Add Models**: Add new models through the UI without restarting the proxy +- **Model Hub**: Make models public for developers to discover available models +- **Price Data Sync**: Keep model pricing data up to date by syncing from GitHub + +For detailed information on model management, see [Model Management](./model_management.md). + +:::tip Sync Model Pricing Data +[Sync model pricing data from GitHub](./sync_models_github.md) to keep your model cost information current. +::: + ## Disable Admin UI Set `DISABLE_ADMIN_UI="True"` in your environment to disable the Admin UI. diff --git a/docs/my-website/release_notes/v1.75.5-stable/index.md b/docs/my-website/release_notes/v1.75.5-stable/index.md index 270be64190e..da13906cd24 100644 --- a/docs/my-website/release_notes/v1.75.5-stable/index.md +++ b/docs/my-website/release_notes/v1.75.5-stable/index.md @@ -51,6 +51,31 @@ pip install litellm==1.75.5.post2 - **Digital Ocean's Gradient AI** - New LLM provider for calling models on Digital Ocean's Gradient AI platform. +### 54% RPS Improvement + +Throughput increased by 54% (1,040 → 1,602 RPS, aggregated) per instance while maintaining a 40 ms median overhead. The improvement comes from fixing major O(n²) inefficiencies in the router, primarily caused by repeated use of in statements inside loops over large arrays. Tests were run with a database-only setup (no cache hits). As a result, p95 latency improved by 30% (2,700 → 1,900 ms), enhancing overall stability and scalability under heavy load. + +--- + +### Test Setup + +All benchmarks were executed using Locust with 1,000 concurrent users and a ramp-up of 500. The environment was configured to stress the routing layer and eliminate caching as a variable. + +**System Specs** + +- **CPU:** 8 vCPUs +- **Memory:** 32 GB RAM + +**Configuration (config.yaml)** + +View the complete configuration: [gist.github.com/AlexsanderHamir/config.yaml](https://gist.github.com/AlexsanderHamir/53f7d554a5d2afcf2c4edb5b6be68ff4) + +**Load Script (no_cache_hits.py)** + +View the complete load testing script: [gist.github.com/AlexsanderHamir/no_cache_hits.py](https://gist.github.com/AlexsanderHamir/42c33d7a4dc7a57f56a78b560dee3a42) + +--- + ### Risk of Upgrade If you build the proxy from the pip package, you should hold off on upgrading. This version makes `prisma migrate deploy` our default for managing the DB. This is safer, as it doesn't reset the DB, but it requires a manual `prisma generate` step. diff --git a/docs/my-website/release_notes/v1.77.7-stable/index.md b/docs/my-website/release_notes/v1.77.7-stable/index.md new file mode 100644 index 00000000000..658a9c3441d --- /dev/null +++ b/docs/my-website/release_notes/v1.77.7-stable/index.md @@ -0,0 +1,298 @@ +--- +title: "[Preview] v1.77.7-stable - Claude Sonnet 4.5" +slug: "v1-77-7" +date: 2025-10-04T10: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 Jaff + title: CTO, LiteLLM + url: https://www.linkedin.com/in/reffajnaahsi/ + image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg + - name: Alexsander Hamir + title: Backend Performance Engineer + url: https://www.linkedin.com/in/alexsander-baptista/ + image_url: https://media.licdn.com/dms/image/v2/D5603AQGXnziu4kqNCQ/profile-displayphoto-crop_800_800/B56ZkxEcuOKEAI-/0/1757464874550?e=1762387200&v=beta&t=9SNXLsWhx8OnYPAMQ9fqAr02oevDYEAL2vMYg2f9ieg + - name: Achintya Srivastava + title: Fullstack Engineer + url: https://www.linkedin.com/in/achintya-rajan/ + image_url: https://media.licdn.com/dms/image/v2/D5603AQGdkEeyJTdljw/profile-displayphoto-shrink_800_800/profile-displayphoto-shrink_800_800/0/1716271140869?e=1762387200&v=beta&t=9gOoLPeqR2E5z3KSX61EUj3HVZXmgo87vhVuSHeffjc + - name: Sameer Kankute + title: Backend Engineer (LLM Translation) + url: https://www.linkedin.com/in/sameer-kankute/ + image_url: https://media.licdn.com/dms/image/v2/D4D03AQHB_loQYd5gjg/profile-displayphoto-shrink_800_800/profile-displayphoto-shrink_800_800/0/1719137160975?e=1762387200&v=beta&t=0jbuX-f4eSnDxBY3olI6meuYr-LMbObhFmFbRcKF5mY + +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.77.7.rc.1 +``` + + + + + +``` showLineNumbers title="pip install litellm" +pip install litellm==1.77.7.rc.1 +``` + + + + +--- + +## Key Highlights + +- **Dynamic Rate Limiter v3** - Automatically maximizes throughput when capacity is available (< 80% saturation) by allowing lower-priority requests to use unused capacity, then switches to fair priority-based allocation under high load (≥ 80%) to prevent blocking +- **Major Performance Improvements** - Router optimization reducing P99 latency by 62.5%, cache improvements from O(n*log(n)) to O(log(n)) +- **Claude Sonnet 4.5** - Support for Anthropic's new Claude Sonnet 4.5 model family with 200K+ context and tiered pricing +- **MCP Gateway Enhancements** - Fine-grained tool control, server permissions, and forwardable headers +- **AMD Lemonade & Nvidia NIM** - New provider support for AMD Lemonade and Nvidia NIM Rerank +- **GitLab Prompt Management** - GitLab-based prompt management integration + +## New Models / Updated Models + +#### New Model Support + +| Provider | Model | Context Window | Input ($/1M tokens) | Output ($/1M tokens) | Features | +| -------- | ----- | -------------- | ------------------- | -------------------- | -------- | +| Anthropic | `claude-sonnet-4-5` | 200K | $3.00 | $15.00 | Chat, reasoning, vision, function calling, prompt caching | +| Anthropic | `claude-sonnet-4-5-20250929` | 200K | $3.00 | $15.00 | Chat, reasoning, vision, function calling, prompt caching | +| Bedrock | `eu.anthropic.claude-sonnet-4-5-20250929-v1:0` | 200K | $3.00 | $15.00 | Chat, reasoning, vision, function calling, prompt caching | +| Azure AI | `azure_ai/grok-4` | 131K | $5.50 | $27.50 | Chat, reasoning, function calling, web search | +| Azure AI | `azure_ai/grok-4-fast-reasoning` | 131K | $5.80 | $2,900.00 | Chat, reasoning, function calling, web search | +| Azure AI | `azure_ai/grok-4-fast-non-reasoning` | 131K | $5.00 | $2,500.00 | Chat, function calling, web search | +| Azure AI | `azure_ai/grok-code-fast-1` | 131K | $3.50 | $17.50 | Chat, function calling, web search | +| Groq | `groq/moonshotai/kimi-k2-instruct-0905` | Context varies | Pricing varies | Pricing varies | Chat, function calling | +| Ollama | Ollama Cloud models | Varies | Free | Free | Self-hosted models via Ollama Cloud | + +#### Features + +- **[Anthropic](../../docs/providers/anthropic)** + - Add new claude-sonnet-4-5 model family with tiered pricing above 200K tokens - [PR #15041](https://github.com/BerriAI/litellm/pull/15041) + - Add anthropic/claude-sonnet-4-5 to model price json with prompt caching support - [PR #15049](https://github.com/BerriAI/litellm/pull/15049) + - Add 200K prices for Sonnet 4.5 - [PR #15140](https://github.com/BerriAI/litellm/pull/15140) + - Add cost tracking for /v1/messages in streaming response - [PR #15102](https://github.com/BerriAI/litellm/pull/15102) + - Add /v1/messages/count_tokens to Anthropic routes for non-admin user access - [PR #15034](https://github.com/BerriAI/litellm/pull/15034) +- **[Gemini](../../docs/providers/gemini)** + - Ignore type param for gemini tools - [PR #15022](https://github.com/BerriAI/litellm/pull/15022) +- **[Vertex AI](../../docs/providers/vertex)** + - Add LiteLLM Overhead metric for VertexAI - [PR #15040](https://github.com/BerriAI/litellm/pull/15040) + - Support googlemap grounding in vertex ai - [PR #15179](https://github.com/BerriAI/litellm/pull/15179) +- **[Azure](../../docs/providers/azure)** + - Add azure_ai grok-4 model family - [PR #15137](https://github.com/BerriAI/litellm/pull/15137) + - Use the `extra_query` parameter for GET requests in Azure Batch - [PR #14997](https://github.com/BerriAI/litellm/pull/14997) + - Use extra_query for download results (Batch API) - [PR #15025](https://github.com/BerriAI/litellm/pull/15025) + - Add support for Azure AD token-based authorization - [PR #14813](https://github.com/BerriAI/litellm/pull/14813) +- **[Ollama](../../docs/providers/ollama)** + - Add ollama cloud models - [PR #15008](https://github.com/BerriAI/litellm/pull/15008) +- **[Groq](../../docs/providers/groq)** + - Add groq/moonshotai/kimi-k2-instruct-0905 - [PR #15079](https://github.com/BerriAI/litellm/pull/15079) +- **[OpenAI](../../docs/providers/openai)** + - Add support for GPT 5 codex models - [PR #14841](https://github.com/BerriAI/litellm/pull/14841) +- **[DeepInfra](../../docs/providers/deepinfra)** + - Update DeepInfra model data refresh with latest pricing - [PR #14939](https://github.com/BerriAI/litellm/pull/14939) +- **[Bedrock](../../docs/providers/bedrock)** + - Add JP Cross-Region Inference - [PR #15188](https://github.com/BerriAI/litellm/pull/15188) + - Add "eu.anthropic.claude-sonnet-4-5-20250929-v1:0" - [PR #15181](https://github.com/BerriAI/litellm/pull/15181) + - Add twelvelabs bedrock Async Invoke Support - [PR #14871](https://github.com/BerriAI/litellm/pull/14871) +- **[Nvidia NIM](../../docs/providers/nvidia_nim)** + - Add Nvidia NIM Rerank Support - [PR #15152](https://github.com/BerriAI/litellm/pull/15152) + +### Bug Fixes + +- **[VLLM](../../docs/providers/vllm)** + - Fix response_format bug in hosted vllm audio_transcription - [PR #15010](https://github.com/BerriAI/litellm/pull/15010) + - Fix passthrough of atranscription into kwargs going to upstream provider - [PR #15005](https://github.com/BerriAI/litellm/pull/15005) +- **[OCI](../../docs/providers/oci)** + - Fix OCI Generative AI Integration when using Proxy - [PR #15072](https://github.com/BerriAI/litellm/pull/15072) +- **General** + - Fix: Authorization header to use correct "Bearer" capitalization - [PR #14764](https://github.com/BerriAI/litellm/pull/14764) + - Bug fix: gpt-5-chat-latest has incorrect max_input_tokens value - [PR #15116](https://github.com/BerriAI/litellm/pull/15116) + - Update request handling for original exceptions - [PR #15013](https://github.com/BerriAI/litellm/pull/15013) + +#### New Provider Support + +- **[AMD Lemonade](../../docs/providers/lemonade)** + - Add AMD Lemonade provider support - [PR #14840](https://github.com/BerriAI/litellm/pull/14840) + +--- + +## LLM API Endpoints + +#### Features + +- **[Responses API](../../docs/response_api)** + - Return Cost for Responses API Streaming requests - [PR #15053](https://github.com/BerriAI/litellm/pull/15053) + +- **[/generateContent](../../docs/providers/gemini)** + - Add full support for native Gemini API translation - [PR #15029](https://github.com/BerriAI/litellm/pull/15029) + +- **Passthrough Gemini Routes** + - Add Gemini generateContent passthrough cost tracking - [PR #15014](https://github.com/BerriAI/litellm/pull/15014) + - Add streamGenerateContent cost tracking in passthrough - [PR #15199](https://github.com/BerriAI/litellm/pull/15199) + +- **Passthrough Vertex AI Routes** + - Add cost tracking for Vertex AI Passthrough `/predict` endpoint - [PR #15019](https://github.com/BerriAI/litellm/pull/15019) + - Add cost tracking for Vertex AI Live API WebSocket Passthrough - [PR #14956](https://github.com/BerriAI/litellm/pull/14956) + +- **General** + - Preserve Whitespace Characters in Model Response Streams - [PR #15160](https://github.com/BerriAI/litellm/pull/15160) + - Add provider name to payload specification - [PR #15130](https://github.com/BerriAI/litellm/pull/15130) + - Ensure query params are forwarded from origin url to downstream request - [PR #15087](https://github.com/BerriAI/litellm/pull/15087) + +--- + +## Management Endpoints / UI + +#### Features + +- **Virtual Keys** + - Fix Session Token Cookie Infinite Logout Loop - [PR #15146](https://github.com/BerriAI/litellm/pull/15146) + - Ensure LLM_API_KEYs can access pass through routes - [PR #15115](https://github.com/BerriAI/litellm/pull/15115) + +- **Models + Endpoints** + - Ensure OCI secret fields not shared on /models and /v1/models endpoints - [PR #15085](https://github.com/BerriAI/litellm/pull/15085) + - Add snowflake on UI - [PR #15083](https://github.com/BerriAI/litellm/pull/15083) + - Make UI theme settings publicly accessible for custom branding - [PR #15074](https://github.com/BerriAI/litellm/pull/15074) + +- **Admin Settings** + - Ensure OTEL settings are saved in DB after set on UI - [PR #15118](https://github.com/BerriAI/litellm/pull/15118) + - Top api key tags - [PR #15151](https://github.com/BerriAI/litellm/pull/15151), [PR #15156](https://github.com/BerriAI/litellm/pull/15156) + +#### Bugs + +- **Dashboard** - Fix LiteLLM model name fallback in dashboard overview - [PR #14998](https://github.com/BerriAI/litellm/pull/14998) + +--- + +## Logging / Guardrail / Prompt Management Integrations + +#### Features + +- **[OpenTelemetry](../../docs/observability/otel)** + - Use generation_name for span naming in logging method - [PR #14799](https://github.com/BerriAI/litellm/pull/14799) +- **[Langfuse](../../docs/proxy/logging#langfuse)** + - Handle non-serializable objects in Langfuse logging - [PR #15148](https://github.com/BerriAI/litellm/pull/15148) + - Set usage_details.total in langfuse integration - [PR #15015](https://github.com/BerriAI/litellm/pull/15015) + +#### Guardrails + +- **[Javelin](../../docs/proxy/guardrails)** + - Add Javelin standalone guardrails integration for LiteLLM Proxy - [PR #14983](https://github.com/BerriAI/litellm/pull/14983) + - Add logging for important status fields in guardrails - [PR #15090](https://github.com/BerriAI/litellm/pull/15090) + - Don't run post_call guardrail if no text returned from Bedrock - [PR #15106](https://github.com/BerriAI/litellm/pull/15106) + +#### Prompt Management + +- **[GitLab](../../docs/proxy/prompt_management)** + - GitLab based Prompt manager - [PR #14988](https://github.com/BerriAI/litellm/pull/14988) + +--- + +## Spend Tracking, Budgets and Rate Limiting + +- **Cost Tracking** + - Proxy: end user cost tracking in the responses API - [PR #15124](https://github.com/BerriAI/litellm/pull/15124) +- **Parallel Request Limiter v3** + - Use well known redis cluster hashing algorithm - [PR #15052](https://github.com/BerriAI/litellm/pull/15052) + - Fixes to dynamic rate limiter v3 - add saturation detection - [PR #15119](https://github.com/BerriAI/litellm/pull/15119) + - Dynamic Rate Limiter v3 - fixes for detecting saturation + fixes for post saturation behavior - [PR #15192](https://github.com/BerriAI/litellm/pull/15192) +- **Teams** + - Add model specific tpm/rpm limits to teams on LiteLLM - [PR #15044](https://github.com/BerriAI/litellm/pull/15044) + +--- + +## MCP Gateway + +- **Server Configuration** + - Specify forwardable headers, specify allowed/disallowed tools for MCP servers - [PR #15002](https://github.com/BerriAI/litellm/pull/15002) + - Enforce server permissions on call tools - [PR #15044](https://github.com/BerriAI/litellm/pull/15044) + - MCP Gateway Fine-grained Tools Addition - [PR #15153](https://github.com/BerriAI/litellm/pull/15153) +- **Bug Fixes** + - Remove servername prefix mcp tools tests - [PR #14986](https://github.com/BerriAI/litellm/pull/14986) + - Resolve regression with duplicate Mcp-Protocol-Version header - [PR #15050](https://github.com/BerriAI/litellm/pull/15050) + - Fix test_mcp_server.py - [PR #15183](https://github.com/BerriAI/litellm/pull/15183) + +--- + +## Performance / Loadbalancing / Reliability improvements + +- **Router Optimizations** + - **+62.5% P99 Latency Improvement** - Remove router inefficiencies (from O(M*N) to O(1)) - [PR #15046](https://github.com/BerriAI/litellm/pull/15046) + - Remove hasattr checks in Router - [PR #15082](https://github.com/BerriAI/litellm/pull/15082) + - Remove Double Lookups - [PR #15084](https://github.com/BerriAI/litellm/pull/15084) + - Optimize _filter_cooldown_deployments from O(n×m + k×n) to O(n) - [PR #15091](https://github.com/BerriAI/litellm/pull/15091) + - Optimize unhealthy deployment filtering in retry path (O(n*m) → O(n+m)) - [PR #15110](https://github.com/BerriAI/litellm/pull/15110) +- **Cache Optimizations** + - Reduce complexity of InMemoryCache.evict_cache from O(n*log(n)) to O(log(n)) - [PR #15000](https://github.com/BerriAI/litellm/pull/15000) + - Avoiding expensive operations when cache isn't available - [PR #15182](https://github.com/BerriAI/litellm/pull/15182) +- **Worker Management** + - Add proxy CLI option to recycle workers after N requests - [PR #15007](https://github.com/BerriAI/litellm/pull/15007) +- **Metrics & Monitoring** + - LiteLLM Overhead metric tracking - Add support for tracking litellm overhead on cache hits - [PR #15045](https://github.com/BerriAI/litellm/pull/15045) + +--- + +## Documentation Updates + +- **Provider Documentation** + - Update litellm docs from latest release - [PR #15004](https://github.com/BerriAI/litellm/pull/15004) + - Add missing api_key parameter - [PR #15058](https://github.com/BerriAI/litellm/pull/15058) +- **General Documentation** + - Use docker compose instead of docker-compose - [PR #15024](https://github.com/BerriAI/litellm/pull/15024) + - Add railtracks to projects that are using litellm - [PR #15144](https://github.com/BerriAI/litellm/pull/15144) + - Perf: Last week improvement - [PR #15193](https://github.com/BerriAI/litellm/pull/15193) + - Sync models GitHub documentation with Loom video and cross-reference - [PR #15191](https://github.com/BerriAI/litellm/pull/15191) + +--- + +## Security Fixes + +- **JWT Token Security** - Don't log JWT SSO token on .info() log - [PR #15145](https://github.com/BerriAI/litellm/pull/15145) + +--- + +## New Contributors + +* @herve-ves made their first contribution in [PR #14998](https://github.com/BerriAI/litellm/pull/14998) +* @wenxi-onyx made their first contribution in [PR #15008](https://github.com/BerriAI/litellm/pull/15008) +* @jpetrucciani made their first contribution in [PR #15005](https://github.com/BerriAI/litellm/pull/15005) +* @abhijitjavelin made their first contribution in [PR #14983](https://github.com/BerriAI/litellm/pull/14983) +* @ZeroClover made their first contribution in [PR #15039](https://github.com/BerriAI/litellm/pull/15039) +* @cedarm made their first contribution in [PR #15043](https://github.com/BerriAI/litellm/pull/15043) +* @Isydmr made their first contribution in [PR #15025](https://github.com/BerriAI/litellm/pull/15025) +* @serializer made their first contribution in [PR #15013](https://github.com/BerriAI/litellm/pull/15013) +* @eddierichter-amd made their first contribution in [PR #14840](https://github.com/BerriAI/litellm/pull/14840) +* @malags made their first contribution in [PR #15000](https://github.com/BerriAI/litellm/pull/15000) +* @henryhwang made their first contribution in [PR #15029](https://github.com/BerriAI/litellm/pull/15029) +* @plafleur made their first contribution in [PR #15111](https://github.com/BerriAI/litellm/pull/15111) +* @tyler-liner made their first contribution in [PR #14799](https://github.com/BerriAI/litellm/pull/14799) +* @Amir-R25 made their first contribution in [PR #15144](https://github.com/BerriAI/litellm/pull/15144) +* @georg-wolflein made their first contribution in [PR #15124](https://github.com/BerriAI/litellm/pull/15124) +* @niharm made their first contribution in [PR #15140](https://github.com/BerriAI/litellm/pull/15140) +* @anthony-liner made their first contribution in [PR #15015](https://github.com/BerriAI/litellm/pull/15015) +* @rishiganesh2002 made their first contribution in [PR #15153](https://github.com/BerriAI/litellm/pull/15153) +* @danielaskdd made their first contribution in [PR #15160](https://github.com/BerriAI/litellm/pull/15160) +* @JVenberg made their first contribution in [PR #15146](https://github.com/BerriAI/litellm/pull/15146) +* @speglich made their first contribution in [PR #15072](https://github.com/BerriAI/litellm/pull/15072) +* @daily-kim made their first contribution in [PR #14764](https://github.com/BerriAI/litellm/pull/14764) + +--- + +## **[Full Changelog](https://github.com/BerriAI/litellm/compare/v1.77.5.rc.4...v1.77.7.rc.1)** diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index d450159f934..7acd92240c6 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -458,7 +458,14 @@ const sidebars = { "providers/deepgram", "providers/watsonx", "providers/predibase", - "providers/nvidia_nim", + { + type: "category", + label: "Nvidia NIM", + items: [ + "providers/nvidia_nim", + "providers/nvidia_nim_rerank", + ] + }, { type: "doc", id: "providers/nscale", label: "Nscale (EU Sovereign)" }, "providers/xai", "providers/moonshot", @@ -699,7 +706,8 @@ const sidebars = { "projects/llm_cord", "projects/pgai", "projects/GPTLocalhost", - "projects/HolmesGPT" + "projects/HolmesGPT", + "projects/Railtracks", ], }, "extras/code_quality", diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.23-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.23-py3-none-any.whl new file mode 100644 index 00000000000..4220fad36c4 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.23-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.23.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.23.tar.gz new file mode 100644 index 00000000000..ceccaacda43 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.2.23.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251004123655_mcp_allowed_tools_column/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251004123655_mcp_allowed_tools_column/migration.sql new file mode 100644 index 00000000000..bdac1e42bc2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251004123655_mcp_allowed_tools_column/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "allowed_tools" TEXT[] DEFAULT ARRAY[]::TEXT[]; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 766625145f6..5a79e171438 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -178,6 +178,7 @@ model LiteLLM_MCPServerTable { updated_by String? mcp_info Json? @default("{}") mcp_access_groups String[] + allowed_tools String[] @default([]) // Health check status status String? @default("unknown") last_health_check DateTime? diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 94e7f59bfa1..00ed141ac88 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.2.22" +version = "0.2.23" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.2.22" +version = "0.2.23" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/__init__.py b/litellm/__init__.py index 078c4348206..60cf327c26f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -152,6 +152,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "vector_store_pre_call_hook", "dotprompt", "bitbucket", + "gitlab", "cloudzero", "posthog", ] @@ -172,22 +173,22 @@ prometheus_initialize_budget_metrics: Optional[bool] = False require_auth_for_metrics_endpoint: Optional[bool] = False argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload. -gcs_pub_sub_use_v1: Optional[bool] = ( - False # if you want to use v1 gcs pubsub logged payload -) -generic_api_use_v1: Optional[bool] = ( - False # if you want to use v1 generic api logged payload -) +gcs_pub_sub_use_v1: Optional[ + bool +] = False # if you want to use v1 gcs pubsub logged payload +generic_api_use_v1: Optional[ + bool +] = False # if you want to use v1 generic api logged payload argilla_transformation_object: Optional[Dict[str, Any]] = None -_async_input_callback: List[Union[str, Callable, CustomLogger]] = ( - [] -) # internal variable - async custom callbacks are routed here. -_async_success_callback: List[Union[str, Callable, CustomLogger]] = ( - [] -) # internal variable - async custom callbacks are routed here. -_async_failure_callback: List[Union[str, Callable, CustomLogger]] = ( - [] -) # internal variable - async custom callbacks are routed here. +_async_input_callback: List[ + Union[str, Callable, CustomLogger] +] = [] # internal variable - async custom callbacks are routed here. +_async_success_callback: List[ + Union[str, Callable, CustomLogger] +] = [] # internal variable - async custom callbacks are routed here. +_async_failure_callback: List[ + Union[str, Callable, CustomLogger] +] = [] # internal variable - async custom callbacks are routed here. pre_call_rules: List[Callable] = [] post_call_rules: List[Callable] = [] turn_off_message_logging: Optional[bool] = False @@ -195,18 +196,18 @@ log_raw_request_response: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False filter_invalid_headers: Optional[bool] = False -add_user_information_to_llm_headers: Optional[bool] = ( - None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers -) +add_user_information_to_llm_headers: Optional[ + bool +] = None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers store_audit_logs = False # Enterprise feature, allow users to see audit logs ### end of callbacks ############# -email: Optional[str] = ( - None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -token: Optional[str] = ( - None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) +email: Optional[ + str +] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +token: Optional[ + str +] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) @@ -289,7 +290,7 @@ banned_keywords_list: Optional[Union[str, List]] = None llm_guard_mode: Literal["all", "key-specific", "request-specific"] = "all" guardrail_name_config_map: Dict[str, GuardrailItem] = {} include_cost_in_streaming_usage: bool = False -### PROMPTS ### +### PROMPTS #### from litellm.types.prompts.init_prompts import PromptSpec prompt_name_config_map: Dict[str, PromptSpec] = {} @@ -307,24 +308,20 @@ enable_loadbalancing_on_batch_endpoints: Optional[bool] = None enable_caching_on_provider_specific_optional_params: bool = ( False # feature-flag for caching on optional params - e.g. 'top_k' ) -caching: bool = ( - False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -caching_with_models: bool = ( - False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -) -cache: Optional[Cache] = ( - None # cache object <- use this - https://docs.litellm.ai/docs/caching -) +caching: bool = False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +caching_with_models: bool = False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +cache: Optional[ + Cache +] = None # cache object <- use this - https://docs.litellm.ai/docs/caching default_in_memory_ttl: Optional[float] = None default_redis_ttl: Optional[float] = None default_redis_batch_cache_expiry: Optional[float] = None model_alias_map: Dict[str, str] = {} model_group_settings: Optional["ModelGroupSettings"] = None max_budget: float = 0.0 # set the max budget across all providers -budget_duration: Optional[str] = ( - None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). -) +budget_duration: Optional[ + str +] = None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). default_soft_budget: float = ( DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 ) @@ -333,15 +330,11 @@ forward_traceparent_to_llm_provider: bool = False _current_cost = 0.0 # private variable, used if max budget is set error_logs: Dict = {} -add_function_to_prompt: bool = ( - False # if function calling not supported by api, append function call details to system prompt -) +add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt client_session: Optional[httpx.Client] = None aclient_session: Optional[httpx.AsyncClient] = None model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' -model_cost_map_url: str = ( - "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" -) +model_cost_map_url: str = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" suppress_debug_info = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None @@ -371,12 +364,10 @@ prometheus_metrics_config: Optional[List] = None disable_add_prefix_to_prompt: bool = ( False # used by anthropic, to disable adding prefix to prompt ) -disable_copilot_system_to_assistant: bool = ( - False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. -) +disable_copilot_system_to_assistant: bool = False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior. public_model_groups: Optional[List[str]] = None public_model_groups_links: Dict[str, str] = {} -#### REQUEST PRIORITIZATION ###### +#### REQUEST PRIORITIZATION ####### priority_reservation: Optional[Dict[str, float]] = None priority_reservation_settings: "PriorityReservationSettings" = ( PriorityReservationSettings() @@ -384,17 +375,13 @@ priority_reservation_settings: "PriorityReservationSettings" = ( ######## Networking Settings ######## -use_aiohttp_transport: bool = ( - True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. -) +use_aiohttp_transport: bool = True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead. aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings disable_aiohttp_transport: bool = False # Set this to true to use httpx instead disable_aiohttp_trust_env: bool = ( False # When False, aiohttp will respect HTTP(S)_PROXY env vars ) -force_ipv4: bool = ( - False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. -) +force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. module_level_aclient = AsyncHTTPHandler( timeout=request_timeout, client_alias="module level aclient" ) @@ -408,13 +395,13 @@ fallbacks: Optional[List] = None context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 -num_retries_per_request: Optional[int] = ( - None # for the request overall (incl. fallbacks + model retries) -) +num_retries_per_request: Optional[ + int +] = None # for the request overall (incl. fallbacks + model retries) ####### SECRET MANAGERS ##################### -secret_manager_client: Optional[Any] = ( - None # list of instantiated key management clients - e.g. azure kv, infisical, etc. -) +secret_manager_client: Optional[ + Any +] = None # list of instantiated key management clients - e.g. azure kv, infisical, etc. _google_kms_resource_name: Optional[str] = None _key_management_system: Optional[KeyManagementSystem] = None _key_management_settings: KeyManagementSettings = KeyManagementSettings() @@ -510,6 +497,7 @@ azure_text_models: Set = set() anyscale_models: Set = set() cerebras_models: Set = set() galadriel_models: Set = set() +nvidia_nim_models: Set = set() sambanova_models: Set = set() sambanova_embedding_models: Set = set() novita_models: Set = set() @@ -704,6 +692,8 @@ def add_known_models(): cerebras_models.add(key) elif value.get("litellm_provider") == "galadriel": galadriel_models.add(key) + elif value.get("litellm_provider") == "nvidia_nim": + nvidia_nim_models.add(key) elif value.get("litellm_provider") == "sambanova": sambanova_models.add(key) elif value.get("litellm_provider") == "sambanova-embedding-models": @@ -831,6 +821,7 @@ model_list = list( | anyscale_models | cerebras_models | galadriel_models + | nvidia_nim_models | sambanova_models | azure_text_models | novita_models @@ -914,6 +905,7 @@ models_by_provider: dict = { "anyscale": anyscale_models, "cerebras": cerebras_models, "galadriel": galadriel_models, + "nvidia_nim": nvidia_nim_models, "sambanova": sambanova_models | sambanova_embedding_models, "novita": novita_models, "nebius": nebius_models | nebius_embedding_models, @@ -1074,6 +1066,7 @@ from .llms.azure_ai.rerank.transformation import AzureAIRerankConfig from .llms.infinity.rerank.transformation import InfinityRerankConfig from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig +from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig from .llms.clarifai.chat.transformation import ClarifaiConfig from .llms.ai21.chat.transformation import AI21ChatConfig, AI21ChatConfig as AI21Config from .llms.meta_llama.chat.transformation import LlamaAPIConfig @@ -1174,6 +1167,7 @@ from .llms.bedrock.embed.amazon_titan_v2_transformation import ( ) from .llms.cohere.chat.transformation import CohereChatConfig from .llms.bedrock.embed.cohere_transformation import BedrockCohereEmbeddingConfig +from .llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig from .llms.openai.openai import OpenAIConfig, MistralEmbeddingConfig from .llms.openai.image_variations.transformation import OpenAIImageVariationConfig from .llms.deepinfra.chat.transformation import DeepInfraConfig @@ -1349,12 +1343,12 @@ from .types.llms.custom_llm import CustomLLMItem from .types.utils import GenericStreamingChunk custom_provider_map: List[CustomLLMItem] = [] -_custom_providers: List[str] = ( - [] -) # internal helper util, used to track names of custom providers -disable_hf_tokenizer_download: Optional[bool] = ( - None # disable huggingface tokenizer download. Defaults to openai clk100 -) +_custom_providers: List[ + str +] = [] # internal helper util, used to track names of custom providers +disable_hf_tokenizer_download: Optional[ + bool +] = None # disable huggingface tokenizer download. Defaults to openai clk100 global_disable_no_log_param: bool = False ### CLI UTILITIES ### @@ -1362,6 +1356,7 @@ from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_k ### PASSTHROUGH ### from .passthrough import allm_passthrough_route, llm_passthrough_route +from .google_genai import agenerate_content ### GLOBAL CONFIG ### global_bitbucket_config: Optional[Dict[str, Any]] = None @@ -1371,3 +1366,11 @@ def set_global_bitbucket_config(config: Dict[str, Any]) -> None: """Set global BitBucket configuration for prompt management.""" global global_bitbucket_config global_bitbucket_config = config + +### GLOBAL CONFIG ### +global_gitlab_config: Optional[Dict[str, Any]] = None + +def set_global_gitlab_config(config: Dict[str, Any]) -> None: + """Set global BitBucket configuration for prompt management.""" + global global_gitlab_config + global_gitlab_config = config diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 37b9aff4efb..48521e5fba0 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -59,18 +59,22 @@ def _resolve_timeout( ) -> float: """ Resolve timeout value from various sources and handle httpx.Timeout objects. - + Args: optional_params: GenericLiteLLMParams object containing timeout kwargs: Additional kwargs that may contain request_timeout custom_llm_provider: Provider name for httpx timeout support check default_timeout: Default timeout value to use - + Returns: Resolved timeout as float """ - timeout = optional_params.timeout or kwargs.get("request_timeout", default_timeout) or default_timeout - + timeout = ( + optional_params.timeout + or kwargs.get("request_timeout", default_timeout) + or default_timeout + ) + # Handle httpx.Timeout objects if isinstance(timeout, httpx.Timeout): if supports_httpx_timeout(custom_llm_provider) is False: @@ -81,11 +85,11 @@ def _resolve_timeout( # For providers that support httpx.Timeout, we still need to return a float # This case might need to be handled differently based on the actual use case return float(timeout.read or default_timeout) - + # Handle None case if timeout is None: return float(default_timeout) - + # Handle numeric values (int, float, string representations) return float(timeout) @@ -163,15 +167,19 @@ def create_batch( try: if model is not None: model, _, _, _ = get_llm_provider( - model=model, - custom_llm_provider=None, - ) + model=model, + custom_llm_provider=None, + ) except Exception as e: - verbose_logger.exception(f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {str(e)}") - + verbose_logger.exception( + f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {str(e)}" + ) + _is_async = kwargs.pop("acreate_batch", False) is True litellm_params = dict(GenericLiteLLMParams(**kwargs)) - litellm_logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None)) + litellm_logging_obj: LiteLLMLoggingObj = cast( + LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None) + ) ### TIMEOUT LOGIC ### timeout = _resolve_timeout(optional_params, kwargs, custom_llm_provider) litellm_logging_obj.update_environment_variables( @@ -189,7 +197,6 @@ def create_batch( }, custom_llm_provider=custom_llm_provider, ) - _create_batch_request = CreateBatchRequest( completion_window=completion_window, @@ -378,6 +385,7 @@ async def aretrieve_batch( except Exception as e: raise e + def _handle_retrieve_batch_providers_without_provider_config( batch_id: str, optional_params: GenericLiteLLMParams, @@ -497,6 +505,7 @@ def _handle_retrieve_batch_providers_without_provider_config( ) return response + @client def retrieve_batch( batch_id: str, @@ -513,7 +522,9 @@ def retrieve_batch( """ try: optional_params = GenericLiteLLMParams(**kwargs) - litellm_logging_obj: Optional[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( @@ -549,7 +560,26 @@ def retrieve_batch( _is_async = kwargs.pop("aretrieve_batch", False) is True client = kwargs.get("client", None) - + + # Check if this is an async invoke ARN (different from regular batch ARN) + # Async invoke ARNs have format: arn:aws(-[^:]+)?:bedrock:[a-z0-9-]{1,20}:[0-9]{12}:async-invoke/[a-z0-9]{12} + if ( + batch_id.startswith("arn:aws") + and ":bedrock:" in batch_id + and ":async-invoke/" in batch_id + ): + # Handle async invoke status check + # Remove aws_region_name from kwargs to avoid duplicate parameter + async_kwargs = kwargs.copy() + async_kwargs.pop("aws_region_name", None) + + return _handle_async_invoke_status( + batch_id=batch_id, + aws_region_name=kwargs.get("aws_region_name", "us-east-1"), + logging_obj=litellm_logging_obj, + **async_kwargs, + ) + # Try to use provider config first (for providers like bedrock) model: Optional[str] = kwargs.get("model", None) if model is not None: @@ -559,7 +589,7 @@ def retrieve_batch( ) else: provider_config = None - + if provider_config is not None: response = base_llm_http_handler.retrieve_batch( batch_id=batch_id, @@ -568,7 +598,8 @@ def retrieve_batch( headers=extra_headers or {}, api_base=optional_params.api_base, api_key=optional_params.api_key, - logging_obj=litellm_logging_obj or LiteLLMLoggingObj( + logging_obj=litellm_logging_obj + or LiteLLMLoggingObj( model=model or "bedrock/unknown", messages=[], stream=False, @@ -586,7 +617,6 @@ def retrieve_batch( model=model, ) return response - ######################################################### # Handle providers without provider config @@ -600,7 +630,7 @@ def retrieve_batch( _is_async=_is_async, timeout=timeout, ) - + except Exception as e: raise e @@ -933,3 +963,79 @@ def cancel_batch( return response except Exception as e: raise e + + +def _handle_async_invoke_status( + batch_id: str, aws_region_name: str, logging_obj=None, **kwargs +) -> "LiteLLMBatch": + """ + Handle async invoke status check for AWS Bedrock. + + Args: + batch_id: The async invoke ARN + aws_region_name: AWS region name + **kwargs: Additional parameters + + Returns: + dict: Status information including status, output_file_id (S3 URL), etc. + """ + import asyncio + + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + async def _async_get_status(): + # Create embedding handler instance + embedding_handler = BedrockEmbedding() + + # Get the status of the async invoke job + status_response = await embedding_handler._get_async_invoke_status( + invocation_arn=batch_id, + aws_region_name=aws_region_name, + logging_obj=logging_obj, + **kwargs, + ) + + # Transform response to a LiteLLMBatch object + from litellm.types.utils import LiteLLMBatch + + result = LiteLLMBatch( + id=status_response["invocationArn"], + object="batch", + status=status_response["status"], + created_at=status_response["submitTime"], + in_progress_at=status_response["lastModifiedTime"], + completed_at=status_response.get("endTime"), + failed_at=status_response.get("endTime") + if status_response["status"] == "failed" + else None, + request_counts={ + "total": 1, + "completed": 1 if status_response["status"] == "completed" else 0, + "failed": 1 if status_response["status"] == "failed" else 0, + }, + metadata={ + "output_file_id": status_response["outputDataConfig"][ + "s3OutputDataConfig" + ]["s3Uri"], + "failure_message": status_response.get("failureMessage"), + "model_arn": status_response["modelArn"], + }, + ) + + return result + + # Since this function is called from within an async context via run_in_executor, + # we need to create a new event loop in a thread to avoid conflicts + import concurrent.futures + + def run_in_thread(): + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + try: + return new_loop.run_until_complete(_async_get_status()) + finally: + new_loop.close() + + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(run_in_thread) + return future.result() diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index b151ebd6513..17cd50f75aa 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -14,6 +14,7 @@ It utilizes the (RedisCache, s3Cache, RedisSemanticCache, QdrantSemanticCache, I In each method it will call the appropriate method from caching.py """ +import time import asyncio import datetime import inspect @@ -57,10 +58,16 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.utils import CustomStreamWrapper else: LiteLLMLoggingObj = Any - CustomStreamWrapper = Any + + +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + +from litellm.litellm_core_utils.core_helpers import ( +_get_parent_otel_span_from_kwargs, +) class CachingHandlerResponse(BaseModel): @@ -112,7 +119,7 @@ class LLMCachingHandler: call_type: str, kwargs: Dict[str, Any], args: Optional[Tuple[Any, ...]] = None, - ) -> CachingHandlerResponse: + ) -> Optional[CachingHandlerResponse]: """ Internal method to get from the cache. Handles different call types (embeddings, chat/completions, text_completion, transcription) @@ -133,32 +140,27 @@ class LLMCachingHandler: Raises: None """ - from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - ) - from litellm.utils import CustomStreamWrapper - - kwargs = kwargs.copy() - args = args or () - ######################################################### - # Init cache timing metrics - ######################################################### - cache_check_start_time = datetime.datetime.now() - cache_check_end_time = None - ######################################################### - - - parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) - kwargs["parent_otel_span"] = parent_otel_span - final_embedding_cached_response: Optional[EmbeddingResponse] = None - embedding_all_elements_cache_hit: bool = False - cached_result: Optional[Any] = None + # Check if caching should be performed BEFORE doing expensive operations if ( (kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) is True ) and ( kwargs.get("cache", {}).get("no-cache", False) is not True ): # allow users to control returning cached responses from the completion function + args = args or () + final_embedding_cached_response: Optional[EmbeddingResponse] = None + embedding_all_elements_cache_hit: bool = False + cached_result: Optional[Any] = None + kwargs = kwargs.copy() + ######################################################### + # Init cache timing metrics + ######################################################### + cache_check_start_time = time.perf_counter() + cache_check_end_time: Optional[float] = None + ######################################################### + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) + kwargs["parent_otel_span"] = parent_otel_span + if litellm.cache is not None and self._is_call_type_supported_by_cache( original_function=original_function ): @@ -168,7 +170,7 @@ class LLMCachingHandler: kwargs=kwargs, args=args, ) - cache_check_end_time = datetime.datetime.now() + cache_check_end_time = time.perf_counter() if cached_result is not None and not isinstance(cached_result, list): verbose_logger.debug("Cache Hit!") @@ -180,7 +182,7 @@ class LLMCachingHandler: api_base=kwargs.get("api_base", None), api_key=kwargs.get("api_key", None), ) - cache_duration_ms = (cache_check_end_time - cache_check_start_time).total_seconds() * 1000 + cache_duration_ms = (cache_check_end_time - cache_check_start_time) * 1000 self._update_litellm_logging_obj_environment( logging_obj=logging_obj, model=model, @@ -245,11 +247,14 @@ class LLMCachingHandler: final_embedding_cached_response=final_embedding_cached_response, embedding_all_elements_cache_hit=embedding_all_elements_cache_hit, ) - verbose_logger.debug(f"CACHE RESULT: {cached_result}") - return CachingHandlerResponse( - cached_result=cached_result, - final_embedding_cached_response=final_embedding_cached_response, - ) + + verbose_logger.debug(f"CACHE RESULT: {cached_result}") + return CachingHandlerResponse( + cached_result=cached_result, + final_embedding_cached_response=final_embedding_cached_response, + ) + # Caching disabled - return None to indicate no caching attempted + return None def _sync_get_cache( self, @@ -263,18 +268,22 @@ class LLMCachingHandler: ) -> CachingHandlerResponse: from litellm.utils import CustomStreamWrapper - args = args or () - new_kwargs = kwargs.copy() - new_kwargs.update( - convert_args_to_kwargs( - self.original_function, - args, - ) - ) + cached_result: Optional[Any] = None + + # Check if caching should be performed BEFORE doing expensive kwargs copy if litellm.cache is not None and self._is_call_type_supported_by_cache( original_function=original_function ): + args = args or () + # Now that we confirmed caching will happen, prepare kwargs + new_kwargs = kwargs.copy() + new_kwargs.update( + convert_args_to_kwargs( + self.original_function, + args, + ) + ) print_verbose("Checking Sync Cache") cached_result = litellm.cache.get_cache(**new_kwargs) if cached_result is not None: diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 082cac791f2..5239fa1f4b0 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -11,6 +11,7 @@ Has 4 methods: import json import sys import time +import heapq from typing import TYPE_CHECKING, Any, List, Optional if TYPE_CHECKING: @@ -46,6 +47,7 @@ class InMemoryCache(BaseCache): # in-memory cache self.cache_dict: dict = {} self.ttl_dict: dict = {} + self.expiration_heap: list[tuple[float, str]] = [] def check_value_size(self, value: Any): """ @@ -114,19 +116,27 @@ class InMemoryCache(BaseCache): """ current_time = time.time() - - # Step 1: Remove expired items - expired_keys = [key for key, ttl in self.ttl_dict.items() if current_time > ttl] - for key in expired_keys: - self._remove_key(key) - # Step 2: If cache is still full, evict items with earliest expiration times - if len(self.cache_dict) >= self.max_size_in_memory: - # Sort by expiration time (earliest first) and evict until we're under the limit - items_by_expiration = sorted(self.ttl_dict.items(), key=lambda x: x[1]) - keys_to_evict = items_by_expiration[:len(self.cache_dict) - self.max_size_in_memory + 1] - - for key, _ in keys_to_evict: + # Step 1: Remove expired or outdated items + while self.expiration_heap: + expiration_time, key = self.expiration_heap[0] + + # Case 1: Heap entry is outdated + if expiration_time != self.ttl_dict.get(key): + heapq.heappop(self.expiration_heap) + # Case 2: Entry is valid but expired + elif expiration_time <= current_time: + heapq.heappop(self.expiration_heap) + self._remove_key(key) + else: + # Case 3: Entry is valid and not expired + break + + # Step 2: Evict if cache is still full + while len(self.cache_dict) >= self.max_size_in_memory: + expiration_time, key = heapq.heappop(self.expiration_heap) + # Skip if key was removed or updated + if self.ttl_dict.get(key) == expiration_time: self._remove_key(key) # de-reference the removed item @@ -150,7 +160,7 @@ class InMemoryCache(BaseCache): # Handle the edge case where max_size_in_memory is 0 if self.max_size_in_memory == 0: return # Don't cache anything if max size is 0 - + if len(self.cache_dict) >= self.max_size_in_memory: # only evict when cache is full self.evict_cache() @@ -161,8 +171,10 @@ class InMemoryCache(BaseCache): if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: self.ttl_dict[key] = time.time() + float(kwargs["ttl"]) + heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) else: self.ttl_dict[key] = time.time() + self.default_ttl + heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) async def async_set_cache(self, key, value, **kwargs): self.set_cache(key=key, value=value, **kwargs) @@ -253,6 +265,7 @@ class InMemoryCache(BaseCache): def flush_cache(self): self.cache_dict.clear() self.ttl_dict.clear() + self.expiration_heap.clear() async def disconnect(self): pass diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index dcf707ebd51..575c36b946a 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -72,15 +72,24 @@ class GenerateContentToCompletionHandler: completion_response = await litellm.acompletion(**completion_kwargs) if stream: - # Transform streaming completion response to generate_content format - transformed_stream = ( - GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming( + # Check if completion_response is actually a stream or a ModelResponse + # This can happen in error cases or when stream is not properly supported + if not hasattr(completion_response, "__aiter__"): + # If it's not a stream, treat it as a regular response + generate_content_response = ( + GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( + cast(ModelResponse, completion_response) + ) + ) + return generate_content_response + else: + # Transform streaming completion response to generate_content format + transformed_stream = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming( completion_response ) - ) - if transformed_stream is not None: - return transformed_stream - raise ValueError("Failed to transform streaming response") + if transformed_stream is not None: + return transformed_stream + raise ValueError("Failed to transform streaming response") else: # Transform completion response back to generate_content format generate_content_response = ( @@ -136,15 +145,24 @@ class GenerateContentToCompletionHandler: completion_response = litellm.completion(**completion_kwargs) if stream: - # Transform streaming completion response to generate_content format - transformed_stream = ( - GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming( + # Check if completion_response is actually a stream or a ModelResponse + # This can happen in error cases or when stream is not properly supported + if not hasattr(completion_response, "__iter__"): + # If it's not a stream, treat it as a regular response + generate_content_response = ( + GOOGLE_GENAI_ADAPTER.translate_completion_to_generate_content( + cast(ModelResponse, completion_response) + ) + ) + return generate_content_response + else: + # Transform streaming completion response to generate_content format + transformed_stream = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming( completion_response ) - ) - if transformed_stream is not None: - return transformed_stream - raise ValueError("Failed to transform streaming response") + if transformed_stream is not None: + return transformed_stream + raise ValueError("Failed to transform streaming response") else: # Transform completion response back to generate_content format generate_content_response = ( diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 7617312302e..9d3f990b1aa 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,12 +1,15 @@ import json from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union, cast +from litellm import verbose_logger + from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, ChatCompletionRequest, + ChatCompletionSystemMessage, ChatCompletionToolCallFunctionChunk, ChatCompletionToolChoiceValues, ChatCompletionToolMessage, @@ -36,43 +39,103 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): def __init__(self, completion_stream: Any): self.sent_first_chunk = False self.accumulated_tool_calls = {} + self._returned_response = False super().__init__(completion_stream) def __next__(self): try: + if not hasattr(self.completion_stream, "__iter__"): + if self._returned_response: + raise StopIteration + self._returned_response = True + return GoogleGenAIAdapter().translate_completion_to_generate_content( + self.completion_stream + ) + for chunk in self.completion_stream: if chunk == "None" or chunk is None: continue - # Transform OpenAI streaming chunk to Google GenAI format transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content( chunk, self ) - if transformed_chunk: # Only return non-empty chunks + if transformed_chunk: return transformed_chunk raise StopIteration except StopIteration: - raise StopIteration + raise except Exception: raise StopIteration async def __anext__(self): try: + if not hasattr(self.completion_stream, "__aiter__"): + if self._returned_response: + raise StopAsyncIteration + self._returned_response = True + return GoogleGenAIAdapter().translate_completion_to_generate_content( + self.completion_stream + ) + async for chunk in self.completion_stream: if chunk == "None" or chunk is None: continue - # Transform OpenAI streaming chunk to Google GenAI format transformed_chunk = GoogleGenAIAdapter().translate_streaming_completion_to_generate_content( chunk, self ) - if transformed_chunk: # Only return non-empty chunks + if transformed_chunk: return transformed_chunk + # After the stream is exhausted, check for any remaining accumulated tool calls + if self.accumulated_tool_calls: + try: + parts = [] + for ( + tool_call_index, + tool_call_data, + ) in self.accumulated_tool_calls.items(): + try: + # For tool calls with no arguments, accumulated_args will be "", which is not valid JSON. + # We default to an empty JSON object in this case. + parsed_args = json.loads( + tool_call_data["arguments"] or "{}" + ) + function_call_part = { + "functionCall": { + "name": tool_call_data["name"] + or "undefined_tool_name", + "args": parsed_args, + } + } + parts.append(function_call_part) + except json.JSONDecodeError: + # This can happen if the stream is abruptly cut off mid-argument string. + verbose_logger.warning( + f"Could not parse tool call arguments at end of stream for index {tool_call_index}. " + f"Name: {tool_call_data['name']}. " + f"Partial args: {tool_call_data['arguments']}" + ) + pass + if parts: + final_chunk = { + "candidates": [ + { + "content": {"parts": parts, "role": "model"}, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [], + } + ] + } + return final_chunk + finally: + # Ensure the accumulator is always cleared to prevent memory leaks + self.accumulated_tool_calls.clear() raise StopAsyncIteration except StopAsyncIteration: - raise StopAsyncIteration + raise except Exception: raise StopAsyncIteration @@ -107,9 +170,14 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): payload = f"data: {json.dumps(transformed_chunk)}\n\n" yield payload.encode() else: - raise ValueError(f"Invalid chunk 1: {chunk}") + # For empty chunks, continue to next iteration + continue else: - raise ValueError(f"Invalid chunk 2: {chunk}") + # For other chunk types, yield them directly + if hasattr(chunk, "encode"): + yield chunk.encode() + else: + yield str(chunk).encode() class GoogleGenAIAdapter: @@ -133,12 +201,19 @@ class GoogleGenAIAdapter: model: The model name contents: Generate content contents (can be list or single dict) config: Optional config parameters - **kwargs: Additional parameters + **kwargs: Additional parameters from the original request Returns: Dict in OpenAI format """ + # Extract top-level fields from kwargs + system_instruction = kwargs.get("systemInstruction") or kwargs.get( + "system_instruction" + ) + tools = kwargs.get("tools") + tool_config = kwargs.get("toolConfig") or kwargs.get("tool_config") + # Normalize contents to list format if isinstance(contents, dict): contents_list = [contents] @@ -146,7 +221,9 @@ class GoogleGenAIAdapter: contents_list = contents # Transform contents to OpenAI messages format - messages = self._transform_contents_to_messages(contents_list) + messages = self._transform_contents_to_messages( + contents_list, system_instruction=system_instruction + ) # Create base request as dict (which is compatible with ChatCompletionRequest) completion_request: ChatCompletionRequest = { @@ -182,20 +259,19 @@ class GoogleGenAIAdapter: completion_request["stop"] = config["stopSequences"] # Handle tools transformation - if "tools" in kwargs: - tools = kwargs["tools"] - + if tools: # Check if tools are already in OpenAI format or Google GenAI format if isinstance(tools, list) and len(tools) > 0: # Tools are in Google GenAI format, transform them openai_tools = self._transform_google_genai_tools_to_openai(tools) + if openai_tools: completion_request["tools"] = openai_tools # Handle tool_config (tool choice) - if "tool_config" in kwargs: + if tool_config: tool_choice = self._transform_google_genai_tool_config_to_openai( - kwargs["tool_config"] + tool_config ) if tool_choice: completion_request["tool_choice"] = tool_choice @@ -235,7 +311,8 @@ class GoogleGenAIAdapter: return completion_request_dict def translate_completion_output_params_streaming( - self, completion_stream: Any + self, + completion_stream: Any, ) -> Union[AsyncIterator[bytes], None]: """Transform streaming completion output to Google GenAI format""" google_genai_wrapper = GoogleGenAIStreamWrapper( @@ -245,7 +322,8 @@ class GoogleGenAIAdapter: return google_genai_wrapper.async_google_genai_sse_wrapper() def _transform_google_genai_tools_to_openai( - self, tools: List[Dict[str, Any]] + self, + tools: List[Dict[str, Any]], ) -> List[ChatCompletionToolParam]: """Transform Google GenAI tools to OpenAI tools format""" openai_tools: List[Dict[str, Any]] = [] @@ -259,8 +337,8 @@ class GoogleGenAIAdapter: if "description" in func_decl: function_chunk["description"] = func_decl["description"] - if "parameters" in func_decl: - function_chunk["parameters"] = func_decl["parameters"] + if "parametersJsonSchema" in func_decl: + function_chunk["parameters"] = func_decl["parametersJsonSchema"] openai_tool = {"type": "function", "function": function_chunk} openai_tools.append(openai_tool) @@ -271,7 +349,8 @@ class GoogleGenAIAdapter: return cast(List[ChatCompletionToolParam], normalized_tools) def _transform_google_genai_tool_config_to_openai( - self, tool_config: Dict[str, Any] + self, + tool_config: Dict[str, Any], ) -> Optional[ChatCompletionToolChoiceValues]: """Transform Google GenAI tool_config to OpenAI tool_choice""" function_calling_config = tool_config.get("functionCallingConfig", {}) @@ -283,11 +362,23 @@ class GoogleGenAIAdapter: return cast(ChatCompletionToolChoiceValues, tool_choice) def _transform_contents_to_messages( - self, contents: List[Dict[str, Any]] + self, + contents: List[Dict[str, Any]], + system_instruction: Optional[Dict[str, Any]] = None, ) -> List[AllMessageValues]: """Transform Google GenAI contents to OpenAI messages format""" messages: List[AllMessageValues] = [] + # Handle system instruction + if system_instruction: + system_parts = system_instruction.get("parts", []) + if system_parts and "text" in system_parts[0]: + messages.append( + ChatCompletionSystemMessage( + role="system", content=system_parts[0]["text"] + ) + ) + for content in contents: role = content.get("role", "user") parts = content.get("parts", []) @@ -364,7 +455,8 @@ class GoogleGenAIAdapter: return messages def translate_completion_to_generate_content( - self, response: ModelResponse + self, + response: ModelResponse, ) -> Dict[str, Any]: """ Transform litellm completion response to Google GenAI generate_content format @@ -376,6 +468,7 @@ class GoogleGenAIAdapter: Dict in Google GenAI generate_content response format """ + # Extract the main response content choice = response.choices[0] if response.choices else None if not choice: @@ -388,12 +481,6 @@ class GoogleGenAIAdapter: "Invalid completion response: no message found in choice" ) parts = self._transform_openai_message_to_google_genai_parts(choice.message) - elif isinstance(choice, StreamingChoices): - if not choice.delta: - raise ValueError( - "Invalid completion response: no delta found in streaming choice" - ) - parts = self._transform_openai_delta_to_google_genai_parts(choice.delta) else: # Fallback for generic choice objects message_content = getattr(choice, "message", {}).get( @@ -438,7 +525,7 @@ class GoogleGenAIAdapter: self, response: Union[ModelResponse, ModelResponseStream], wrapper: GoogleGenAIStreamWrapper, - ) -> Dict[str, Any]: + ) -> Optional[Dict[str, Any]]: """ Transform streaming litellm completion chunk to Google GenAI generate_content format @@ -454,7 +541,7 @@ class GoogleGenAIAdapter: choice = response.choices[0] if response.choices else None if not choice: # Return empty chunk if no choices - return {} + return None # Handle streaming choice if isinstance(choice, StreamingChoices): @@ -473,7 +560,7 @@ class GoogleGenAIAdapter: # Only create response chunk if we have parts or it's the final chunk if not parts and not finish_reason: - return {} + return None # Create Google GenAI streaming format response streaming_chunk: Dict[str, Any] = { @@ -515,7 +602,8 @@ class GoogleGenAIAdapter: return streaming_chunk def _transform_openai_message_to_google_genai_parts( - self, message: Any + self, + message: Any, ) -> List[Dict[str, Any]]: """Transform OpenAI message to Google GenAI parts format""" parts: List[Dict[str, Any]] = [] @@ -537,112 +625,94 @@ class GoogleGenAIAdapter: except json.JSONDecodeError: args = {} - function_call_part = { - "functionCall": {"name": tool_call.function.name, "args": args} - } - parts.append(function_call_part) - - return parts if parts else [{"text": ""}] - - def _transform_openai_delta_to_google_genai_parts( - self, delta: Any - ) -> List[Dict[str, Any]]: - """Transform OpenAI delta to Google GenAI parts format for streaming""" - parts: List[Dict[str, Any]] = [] - - # Add text content if present - if hasattr(delta, "content") and delta.content: - parts.append({"text": delta.content}) - - # Add tool calls if present (for streaming tool calls) - if hasattr(delta, "tool_calls") and delta.tool_calls: - for tool_call in delta.tool_calls: - if hasattr(tool_call, "function") and tool_call.function: - # For streaming, we might get partial function arguments - args_str = getattr(tool_call.function, "arguments", "") or "" - try: - args = json.loads(args_str) if args_str else {} - except json.JSONDecodeError: - # For partial JSON in streaming, return as text for now - args = {"partial": args_str} - function_call_part = { "functionCall": { - "name": getattr(tool_call.function, "name", "") or "", + "name": tool_call.function.name or "undefined_tool_name", "args": args, } } parts.append(function_call_part) - return parts + return parts if parts else [{"text": ""}] def _transform_openai_delta_to_google_genai_parts_with_accumulation( self, delta: Any, wrapper: GoogleGenAIStreamWrapper ) -> List[Dict[str, Any]]: - """Transform OpenAI delta to Google GenAI parts format with tool call accumulation""" + """Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls.""" + + # 1. Initialize wrapper state if it doesn't exist + if not hasattr(wrapper, "accumulated_tool_calls"): + wrapper.accumulated_tool_calls = {} + parts: List[Dict[str, Any]] = [] - # Add text content if present if hasattr(delta, "content") and delta.content: parts.append({"text": delta.content}) - # Handle tool calls with accumulation for streaming - if hasattr(delta, "tool_calls") and delta.tool_calls: - for tool_call in delta.tool_calls: - if hasattr(tool_call, "function") and tool_call.function: - tool_call_id = getattr(tool_call, "id", "") or "call_unknown" - function_name = getattr(tool_call.function, "name", "") or "" - args_str = getattr(tool_call.function, "arguments", "") or "" + # 2. Ensure tool_calls is iterable + tool_calls = delta.tool_calls or [] - # Initialize accumulation for this tool call if not exists - if tool_call_id not in wrapper.accumulated_tool_calls: - wrapper.accumulated_tool_calls[tool_call_id] = { - "name": "", - "arguments": "", - "complete": False, - } + for tool_call in tool_calls: + if not hasattr(tool_call, "function"): + continue - # Accumulate function name if provided - if function_name: - wrapper.accumulated_tool_calls[tool_call_id][ - "name" - ] = function_name + # 3. Use `index` as the primary key for accumulation + tool_call_index = getattr(tool_call, "index", None) + if tool_call_index is None: + continue # Index is essential for tracking streaming tool calls - # Accumulate arguments if provided - if args_str: - wrapper.accumulated_tool_calls[tool_call_id][ - "arguments" - ] += args_str + # Initialize accumulator for this index if it's new + if tool_call_index not in wrapper.accumulated_tool_calls: + wrapper.accumulated_tool_calls[tool_call_index] = { + "name": "", + "arguments": "", + } - # Try to parse the accumulated arguments as JSON - accumulated_args = wrapper.accumulated_tool_calls[tool_call_id][ - "arguments" - ] - try: - if accumulated_args: - parsed_args = json.loads(accumulated_args) - # JSON is valid, mark as complete and create function call part - wrapper.accumulated_tool_calls[tool_call_id][ - "complete" - ] = True + # Accumulate name and arguments + function_name = getattr(tool_call.function, "name", None) + args_chunk = getattr(tool_call.function, "arguments", None) - function_call_part = { - "functionCall": { - "name": wrapper.accumulated_tool_calls[ - tool_call_id - ]["name"], - "args": parsed_args, - } - } - parts.append(function_call_part) + # Optimization: Skip chunks that have no new data + if not function_name and not args_chunk: + verbose_logger.debug( + f"Skipping empty tool call chunk for index: {tool_call_index}" + ) + continue - # Clean up completed tool call - del wrapper.accumulated_tool_calls[tool_call_id] + if function_name: + wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name - except json.JSONDecodeError: - # JSON is still incomplete, continue accumulating - # Don't add to parts yet - pass + if args_chunk: + wrapper.accumulated_tool_calls[tool_call_index][ + "arguments" + ] += args_chunk + + # Attempt to parse and emit a complete tool call + accumulated_data = wrapper.accumulated_tool_calls[tool_call_index] + accumulated_name = accumulated_data["name"] + accumulated_args = accumulated_data["arguments"] + + # 5. Attempt to parse arguments even if name hasn't arrived. + try: + # Attempt to parse the accumulated arguments string + parsed_args = json.loads(accumulated_args) + + # If parsing succeeds, but we don't have a name yet, wait. + # The part will be created by a later chunk that brings the name. + if accumulated_name: + # If successful, create the part and clean up + function_call_part = { + "functionCall": {"name": accumulated_name, "args": parsed_args} + } + parts.append(function_call_part) + + # Remove the completed tool call from the accumulator + del wrapper.accumulated_tool_calls[tool_call_index] + + except json.JSONDecodeError: + # The JSON for arguments is still incomplete. + # We will continue to accumulate and wait for more chunks. + pass return parts diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index b480a85c85e..8a9cb809404 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -85,7 +85,6 @@ class GenerateContentHelper: contents: GenerateContentContentListUnionDict, config: Optional[GenerateContentConfigDict] = None, custom_llm_provider: Optional[str] = None, - stream: bool = False, tools: Optional[ToolConfigDict] = None, **kwargs, ) -> GenerateContentSetupResult: @@ -97,8 +96,7 @@ class GenerateContentHelper: contents: The content to generate from config: Optional configuration custom_llm_provider: Optional custom LLM provider - stream: Whether this is a streaming call - local_vars: Local variables from the calling function + tools: Optional tools **kwargs: Additional keyword arguments Returns: @@ -114,7 +112,7 @@ class GenerateContentHelper: ## MOCK RESPONSE LOGIC (only for non-streaming) if ( - not stream + not kwargs.get("stream", False) and litellm_params.mock_response and isinstance(litellm_params.mock_response, str) ): @@ -289,7 +287,7 @@ def generate_content( """ local_vars = locals() try: - _is_async = kwargs.pop("agenerate_content", False) is True + _is_async = kwargs.pop("agenerate_content", False) # Handle generationConfig parameter from kwargs for backward compatibility if "generationConfig" in kwargs and config is None: @@ -309,7 +307,6 @@ def generate_content( contents=contents, config=config, custom_llm_provider=custom_llm_provider, - stream=False, tools=tools, **kwargs, ) @@ -321,7 +318,7 @@ def generate_content( model=model, contents=contents, # type: ignore config=setup_result.generate_content_config_dict, - stream=False, + tools=tools, _is_async=_is_async, litellm_params=setup_result.litellm_params, **kwargs, @@ -342,7 +339,6 @@ def generate_content( timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), - stream=False, litellm_metadata=kwargs.get("litellm_metadata", {}), ) @@ -391,15 +387,12 @@ async def agenerate_content_stream( # Setup the call setup_result = GenerateContentHelper.setup_generate_content_call( - **{ - "model": model, - "contents": contents, - "config": config, - "custom_llm_provider": custom_llm_provider, - "stream": True, - "tools": tools, - **kwargs, - } + model=model, + contents=contents, + config=config, + custom_llm_provider=custom_llm_provider, + tools=tools, + **kwargs, ) # Check if we should use the adapter (when provider config is None) @@ -411,6 +404,7 @@ async def agenerate_content_stream( contents=contents, # type: ignore config=setup_result.generate_content_config_dict, litellm_params=setup_result.litellm_params, + tools=tools, stream=True, **kwargs, ) @@ -479,7 +473,6 @@ def generate_content_stream( contents=contents, config=config, custom_llm_provider=custom_llm_provider, - stream=True, tools=tools, **kwargs, ) @@ -491,9 +484,9 @@ def generate_content_stream( model=model, contents=contents, # type: ignore config=setup_result.generate_content_config_dict, - stream=True, _is_async=_is_async, litellm_params=setup_result.litellm_params, + stream=True, **kwargs, ) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 6b77557cd3d..22e652e1d7b 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, Dict, List, Literal, Optional, Type, Union, get_args +from typing import Any, Dict, List, Optional, Type, Union, get_args from litellm._logging import verbose_logger from litellm.caching import DualCache @@ -14,6 +14,7 @@ from litellm.types.guardrails import ( from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import ( CallTypes, + GuardrailStatus, LLMResponseTypes, StandardLoggingGuardrailInformation, ) @@ -352,7 +353,7 @@ class CustomGuardrail(CustomLogger): self, guardrail_json_response: Union[Exception, str, dict, List[dict]], request_data: dict, - guardrail_status: Literal["success", "failure", "blocked"], + guardrail_status: GuardrailStatus, start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, @@ -460,7 +461,7 @@ class CustomGuardrail(CustomLogger): self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=e, request_data=request_data, - guardrail_status="failure", + guardrail_status="guardrail_failed_to_respond", duration=duration, start_time=start_time, end_time=end_time, diff --git a/litellm/integrations/gitlab/README.md b/litellm/integrations/gitlab/README.md new file mode 100644 index 00000000000..14fb62905c8 --- /dev/null +++ b/litellm/integrations/gitlab/README.md @@ -0,0 +1,317 @@ +# LiteLLM gitlab Prompt Management + +A powerful prompt management system for LiteLLM that fetches `.prompt` files from gitlab repositories. This enables team-based prompt management with gitlab's built-in access control and version control capabilities. + +## Features + +- **🏢 Team-based access control**: Leverage gitlab's workspace and repository permissions +- **📁 Repository-based prompt storage**: Store prompts in gitlab repositories +- **🔐 Multiple authentication methods**: Support for access tokens and basic auth +- **🎯 YAML frontmatter**: Define model, parameters, and schemas in file headers +- **🔧 Handlebars templating**: Use `{{variable}}` syntax with Jinja2 backend +- **✅ Input validation**: Automatic validation against defined schemas +- **🔗 LiteLLM integration**: Works seamlessly with `litellm.completion()` +- **💬 Smart message parsing**: Converts prompts to proper chat messages +- **⚙️ Parameter extraction**: Automatically applies model settings from prompts + +## Quick Start + +### 1. Set up gitlab Repository + +Create a repository in your gitlab workspace and add `.prompt` files: + +``` +your-repo/ +├── prompts/ +│ ├── chat_assistant.prompt +│ ├── code_reviewer.prompt +│ └── data_analyst.prompt +``` + +### 2. Create a `.prompt` file + +Create a file called `prompts/chat_assistant.prompt`: + +```yaml +--- +model: gpt-4 +temperature: 0.7 +max_tokens: 150 +input: + schema: + user_message: string + system_context?: string +--- + +{% if system_context %}System: {{system_context}} + +{% endif %}User: {{user_message}} +``` + +### 3. Configure gitlab Access + +#### Option A: Access Token (Recommended) + +```python +import litellm + +# Configure gitlab access +gitlab_config = { + "project": "a/b/", + "access_token": "your-access-token", + "base_url": "gitlab url", + "prompts_path": "src/prompts", # folder to point to, defaults to root + "branch":"main" # optional, defaults to main +} + +# Set global gitlab configuration +litellm.set_global_gitlab_config(gitlab_config) +``` + +#### Option B: Basic Authentication + +```python +import litellm + +# Configure gitlab access with basic auth +gitlab_config = { + "project": "a/b/", + "base_url": "base url", + "access_token": "your-app-password", # Use app password for basic auth + "branch": "main", + "prompts_path": "src/prompts", # folder to point to, defaults to root +} + +litellm.set_global_gitlab_config(gitlab_config) +``` + +### 4. Use with LiteLLM + +```python +# Use with completion - the model prefix 'gitlab/' tells LiteLLM to use gitlab prompt management +response = litellm.completion( + model="gitlab/gpt-4", # The actual model comes from the .prompt file + prompt_id="prompts/chat_assistant", # Location of the prompt file + prompt_variables={ + "user_message": "What is machine learning?", + "system_context": "You are a helpful AI tutor." + }, + # Any additional messages will be appended after the prompt + messages=[{"role": "user", "content": "Please explain it simply."}] +) + +print(response.choices[0].message.content) +``` + +## Proxy Server Configuration + +### 1. Create a `.prompt` file + +Create `prompts/hello.prompt`: + +```yaml +--- +model: gpt-4 +temperature: 0.7 +--- +System: You are a helpful assistant. + +User: {{user_message}} +``` + +### 2. Setup config.yaml + +```yaml +model_list: + - model_name: my-gitlab-model + litellm_params: + model: gitlab/gpt-4 + prompt_id: "prompts/hello" + api_key: os.environ/OPENAI_API_KEY + +litellm_settings: + global_gitlab_config: + workspace: "your-workspace" + repository: "your-repo" + access_token: "your-access-token" + branch: "main" +``` + +### 3. Start the proxy + +```bash +litellm --config config.yaml --detailed_debug +``` + +### 4. Test it! + +```bash +curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-d '{ + "model": "my-gitlab-model", + "messages": [{"role": "user", "content": "IGNORED"}], + "prompt_variables": { + "user_message": "What is the capital of France?" + } +}' +``` + +## Prompt File Format + +### Basic Structure + +```yaml +--- +# Model configuration +model: gpt-4 +temperature: 0.7 +max_tokens: 500 + +# Input schema (optional) +input: + schema: + user_message: string + system_context?: string +--- + +System: You are a helpful {{role}} assistant. + +User: {{user_message}} +``` + +### Advanced Features + +**Multi-role conversations:** + +```yaml +--- +model: gpt-4 +temperature: 0.3 +--- +System: You are a helpful coding assistant. + +User: {{user_question}} +``` + +**Dynamic model selection:** + +```yaml +--- +model: "{{preferred_model}}" # Model can be a variable +temperature: 0.7 +--- +System: You are a helpful assistant specialized in {{domain}}. + +User: {{user_message}} +``` + +## Team-Based Access Control + +gitlab's built-in permission system provides team-based access control: + +1. **Workspace-level permissions**: Control access to entire workspaces +2. **Repository-level permissions**: Control access to specific repositories +3. **Branch-level permissions**: Control access to specific branches +4. **User and group management**: Manage team members and their access levels + +### Setting up Team Access + +1. **Create workspaces for each team**: + ``` + team-a-prompts/ + team-b-prompts/ + team-c-prompts/ + ``` + +2. **Configure repository permissions**: + - Grant read access to team members + - Grant write access to prompt maintainers + - Use branch protection rules for production prompts + +3. **Use different access tokens**: + - Each team can have their own access token + - Tokens can be scoped to specific repositories + - Use app passwords for additional security + +## API Reference + +### gitlab Configuration + +```python +gitlab_config = { + "workspace": str, # Required: gitlab workspace name + "repository": str, # Required: Repository name + "access_token": str, # Required: gitlab access token or app password + "branch": str, # Optional: Branch to fetch from (default: "main") + "base_url": str, # Optional: Custom gitlab API URL + "auth_method": str, # Optional: "token" or "basic" (default: "token") + "username": str, # Optional: Username for basic auth + "base_url" : str # Optional: Incase where the base url is not https://api.gitlab.org/2.0 +} +``` + +### LiteLLM Integration + +```python +response = litellm.completion( + model="gitlab/", # required (e.g., gitlab/gpt-4) + prompt_id=str, # required - the .prompt filename without extension + prompt_variables=dict, # optional - variables for template rendering + gitlab_config=dict, # optional - gitlab configuration (if not set globally) + messages=list, # optional - additional messages +) +``` + +## Error Handling + +The gitlab integration provides detailed error messages for common issues: + +- **Authentication errors**: Invalid access tokens or credentials +- **Permission errors**: Insufficient access to workspace/repository +- **File not found**: Missing .prompt files +- **Network errors**: Connection issues with gitlab API + +## Security Considerations + +1. **Access Token Security**: Store access tokens securely using environment variables or secret management systems +2. **Repository Permissions**: Use gitlab's permission system to control access +3. **Branch Protection**: Protect main branches from unauthorized changes +4. **Audit Logging**: gitlab provides audit logs for all repository access + +## Troubleshooting + +### Common Issues + +1. **"Access denied" errors**: Check your gitlab permissions for the workspace and repository +2. **"Authentication failed" errors**: Verify your access token or credentials +3. **"File not found" errors**: Ensure the .prompt file exists in the specified branch +4. **Template rendering errors**: Check your Handlebars syntax in the .prompt file + +### Debug Mode + +Enable debug logging to troubleshoot issues: + +```python +import litellm +litellm.set_verbose = True + +# Your gitlab prompt calls will now show detailed logs +response = litellm.completion( + model="gitlab/gpt-4", + prompt_id="your_prompt", + prompt_variables={"key": "value"} +) +``` + +## Migration from File-Based Prompts + +If you're currently using file-based prompts with the dotprompt integration, you can easily migrate to gitlab: + +1. **Upload your .prompt files** to a gitlab repository +2. **Update your configuration** to use gitlab instead of local files +3. **Set up team access** using gitlab's permission system +4. **Update your code** to use `gitlab/` model prefix instead of `dotprompt/` + +This provides better collaboration, version control, and team-based access control for your prompts. diff --git a/litellm/integrations/gitlab/__init__.py b/litellm/integrations/gitlab/__init__.py new file mode 100644 index 00000000000..cd22afc2ba0 --- /dev/null +++ b/litellm/integrations/gitlab/__init__.py @@ -0,0 +1,95 @@ +from typing import TYPE_CHECKING, Optional, Dict, Any + +if TYPE_CHECKING: + from .gitlab_prompt_manager import GitLabPromptManager + from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec + from litellm.integrations.custom_prompt_management import CustomPromptManagement + +from litellm.types.prompts.init_prompts import SupportedPromptIntegrations +from litellm.integrations.custom_prompt_management import CustomPromptManagement +from litellm.types.prompts.init_prompts import PromptSpec, PromptLiteLLMParams +from .gitlab_prompt_manager import GitLabPromptManager + +# Global instances +global_gitlab_config: Optional[dict] = None + + +def set_global_gitlab_config(config: dict) -> None: + """ + Set the global BitBucket configuration for prompt management. + + Args: + config: Dictionary containing BitBucket configuration + - workspace: BitBucket workspace name + - repository: Repository name + - access_token: BitBucket access token + - branch: Branch to fetch prompts from (default: main) + """ + import litellm + + litellm.global_gitlab_config = config # type: ignore + + +def prompt_initializer( + litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec" +) -> "CustomPromptManagement": + """ + Initialize a prompt from a BitBucket repository. + """ + gitlab_config = getattr(litellm_params, "gitlab_config", None) + prompt_id = getattr(litellm_params, "prompt_id", None) + + + if not gitlab_config: + raise ValueError( + "bitbucket_config is required for BitBucket prompt integration" + ) + + try: + bitbucket_prompt_manager = GitLabPromptManager( + gitlab_config=gitlab_config, + prompt_id=prompt_id, + ) + + return bitbucket_prompt_manager + except Exception as e: + raise e + +def _gitlab_prompt_initializer( + litellm_params: PromptLiteLLMParams, + prompt: PromptSpec, +) -> CustomPromptManagement: + """ + Build a GitLab-backed prompt manager for this prompt. + Expected fields on litellm_params: + - prompt_integration="gitlab" (handled by the caller) + - gitlab_config: Dict[str, Any] (project/access_token/branch/prompts_path/etc.) + - git_ref (optional): per-prompt tag/branch/SHA override + """ + # You can store arbitrary integration-specific config on PromptLiteLLMParams. + # If your dataclass doesn't have these attributes, add them or put inside + # `litellm_params.extra` and pull them from there. + gitlab_config: Dict[str, Any] = getattr(litellm_params, "gitlab_config", None) or {} + git_ref: Optional[str] = getattr(litellm_params, "git_ref", None) + + if not gitlab_config: + raise ValueError("gitlab_config is required for gitlab prompt integration") + + # prompt.prompt_id can map to a file path under prompts_path (e.g. "chat/greet/hi") + return GitLabPromptManager( + gitlab_config=gitlab_config, + prompt_id=prompt.prompt_id, + ref=git_ref, + ) + + +prompt_initializer_registry = { + SupportedPromptIntegrations.GITLAB.value: _gitlab_prompt_initializer, +} + +# Export public API +__all__ = [ + "GitLabPromptManager", + "set_global_gitlab_config", + "global_gitlab_config", +] diff --git a/litellm/integrations/gitlab/gitlab_client.py b/litellm/integrations/gitlab/gitlab_client.py new file mode 100644 index 00000000000..ce03a35d48e --- /dev/null +++ b/litellm/integrations/gitlab/gitlab_client.py @@ -0,0 +1,285 @@ +""" +GitLab API client for fetching files from GitLab repositories. +Now supports selecting a tag via `config["tag"]`; falls back to branch ("main"). +""" + +import base64 +from typing import Any, Dict, List, Optional +from urllib.parse import quote + +from litellm.llms.custom_httpx.http_handler import HTTPHandler + + +class GitLabClient: + """ + Client for interacting with the GitLab API to fetch files. + + Supports: + - Authentication with personal/access tokens or OAuth bearer tokens + - Fetching file contents from repositories (raw endpoint with JSON fallback) + - Namespace/project path or numeric project ID addressing + - Ref selection via tag (preferred) or branch (default "main") + - Directory listing via the repository tree API + """ + + def __init__(self, config: Dict[str, Any]): + """ + Initialize the GitLab client. + + Args: + config: Dictionary containing: + - project: Project path ("group/subgroup/repo") or numeric project ID (str|int) [required] + - access_token: GitLab personal/access token or OAuth token [required] (str) + - auth_method: 'token' (default; sends Private-Token) or 'oauth' (Authorization: Bearer) + - tag: Tag name to fetch from (takes precedence over branch if provided) + - branch: Branch to fetch from (default: "main") + - base_url: Base GitLab API URL (default: "https://gitlab.com/api/v4") + """ + project = config.get("project") + access_token = config.get("access_token") + if project is None or access_token is None: + raise ValueError("project and access_token are required") + + self.project: str | int = project + self.access_token: str = str(access_token) + self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth' + self.branch = config.get("branch", None) + if not self.branch: + self.branch = 'main' + self.tag = config.get("tag") + self.base_url = config.get("base_url", "https://gitlab.com/api/v4") + + if not all([self.project, self.access_token]): + raise ValueError("project and access_token are required") + + # Effective ref: prefer tag if provided, else branch ("main") + self.ref = str(self.tag or self.branch) + + # Build headers + self.headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + if self.auth_method == "oauth": + self.headers["Authorization"] = f"Bearer {self.access_token}" + else: + # Default GitLab token header + self.headers["Private-Token"] = self.access_token + + # Project identifier must be URL-encoded (slashes become %2F) + self._project_enc = quote(str(self.project), safe="") + + # HTTP handler + self.http_handler = HTTPHandler() + + # ------------------------ + # Core helpers + # ------------------------ + + def _file_raw_url(self, file_path: str, *, ref: Optional[str] = None) -> str: + file_enc = quote(file_path, safe="") + ref_q = quote(ref or self.ref, safe="") + return f"{self.base_url}/projects/{self._project_enc}/repository/files/{file_enc}/raw?ref={ref_q}" + + def _file_json_url(self, file_path: str, *, ref: Optional[str] = None) -> str: + file_enc = quote(file_path, safe="") + ref_q = quote(ref or self.ref, safe="") + return f"{self.base_url}/projects/{self._project_enc}/repository/files/{file_enc}?ref={ref_q}" + + def _tree_url(self, directory_path: str = "", recursive: bool = False, *, ref: Optional[str] = None) -> str: + path_q = f"&path={quote(directory_path, safe='')}" if directory_path else "" + rec_q = "&recursive=true" if recursive else "" + ref_q = quote(ref or self.ref, safe="") + return f"{self.base_url}/projects/{self._project_enc}/repository/tree?ref={ref_q}{path_q}{rec_q}" + + # ------------------------ + # Public API + # ------------------------ + + def set_ref(self, ref: str) -> None: + """Override the default ref (tag/branch) for subsequent calls.""" + if not ref: + raise ValueError("ref must be a non-empty string") + self.ref = ref + + def get_file_content(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]: + """ + Fetch the content of a file from the GitLab repository at the given ref + (tag, branch, or commit SHA). If `ref` is None, uses self.ref. + + Strategy: + 1) Try the RAW endpoint (returns bytes of the file) + 2) Fallback to the JSON endpoint (returns base64-encoded content) + + Returns: + File content as UTF-8 string, or None if file not found. + """ + raw_url = self._file_raw_url(file_path, ref=ref) + + try: + resp = self.http_handler.get(raw_url, headers=self.headers) + if resp.status_code == 404: + # Fallback to JSON endpoint + return self._get_file_content_via_json(file_path, ref=ref) + resp.raise_for_status() + + ctype = (resp.headers.get("content-type") or "").lower() + if ctype.startswith("text/") or "charset=" in ctype or ctype.startswith("application/json"): + return resp.text + try: + return resp.content.decode("utf-8") + except Exception: + return resp.content.decode("utf-8", errors="replace") + + except Exception as e: + status = getattr(getattr(e, "response", None), "status_code", None) + if status == 404: + return None + if status == 403: + raise Exception( + f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'." + ) + if status == 401: + raise Exception("Authentication failed. Check your GitLab token and auth_method.") + raise Exception(f"Failed to fetch file '{file_path}': {e}") + + def _get_file_content_via_json(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]: + """ + Fallback for get_file_content(): use the JSON file API which returns base64 content. + """ + json_url = self._file_json_url(file_path, ref=ref) + try: + resp = self.http_handler.get(json_url, headers=self.headers) + if resp.status_code == 404: + return None + resp.raise_for_status() + data = resp.json() + content = data.get("content") + encoding = data.get("encoding", "") + if content and encoding == "base64": + try: + return base64.b64decode(content).decode("utf-8") + except Exception: + return base64.b64decode(content).decode("utf-8", errors="replace") + return content + except Exception as e: + status = getattr(getattr(e, "response", None), "status_code", None) + if status == 404: + return None + if status == 403: + raise Exception( + f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'." + ) + if status == 401: + raise Exception("Authentication failed. Check your GitLab token and auth_method.") + raise Exception(f"Failed to fetch file '{file_path}' via JSON endpoint: {e}") + + def list_files( + self, + directory_path: str = "", + file_extension: str = ".prompt", + recursive: bool = False, + *, + ref: Optional[str] = None, + ) -> List[str]: + """ + List files in a directory with a specific extension using the repository tree API. + + Args: + directory_path: Directory path in the repository (empty for repo root) + file_extension: File extension to filter by (default: .prompt) + recursive: If True, traverses subdirectories + ref: Optional override (tag/branch/SHA). Defaults to self.ref. + + Returns: + List of file paths (relative to repo root) + """ + url = self._tree_url(directory_path, recursive=recursive, ref=ref) + + try: + resp = self.http_handler.get(url, headers=self.headers) + if resp.status_code == 404: + return [] + resp.raise_for_status() + + data = resp.json() or [] + files: List[str] = [] + for item in data: + if item.get("type") == "blob": + file_path = item.get("path", "") + if not file_extension or file_path.endswith(file_extension): + files.append(file_path) + return files + + except Exception as e: + status = getattr(getattr(e, "response", None), "status_code", None) + if status == 404: + return [] + if status == 403: + raise Exception( + f"Access denied to directory '{directory_path}'. Check your GitLab permissions for project '{self.project}'." + ) + if status == 401: + raise Exception("Authentication failed. Check your GitLab token and auth_method.") + raise Exception(f"Failed to list files in '{directory_path}': {e}") + + def get_repository_info(self) -> Dict[str, Any]: + """Get information about the project/repository.""" + url = f"{self.base_url}/projects/{self._project_enc}" + try: + resp = self.http_handler.get(url, headers=self.headers) + resp.raise_for_status() + return resp.json() + except Exception as e: + raise Exception(f"Failed to get repository info: {e}") + + def test_connection(self) -> bool: + """Test the connection to the GitLab project.""" + try: + self.get_repository_info() + return True + except Exception: + return False + + def get_branches(self) -> List[Dict[str, Any]]: + """Get list of branches in the repository.""" + url = f"{self.base_url}/projects/{self._project_enc}/repository/branches" + try: + resp = self.http_handler.get(url, headers=self.headers) + resp.raise_for_status() + data = resp.json() + return data if isinstance(data, list) else [] + except Exception as e: + raise Exception(f"Failed to get branches: {e}") + + def get_file_metadata(self, file_path: str, *, ref: Optional[str] = None) -> Optional[Dict[str, Any]]: + """ + Get minimal metadata about a file via RAW endpoint headers at a given ref. + + Args: + file_path: Path to the file in the repository. + ref: Optional override (tag/branch/SHA). Defaults to self.ref. + """ + url = self._file_raw_url(file_path, ref=ref) + try: + headers = dict(self.headers) + headers["Range"] = "bytes=0-0" + resp = self.http_handler.get(url, headers=headers) + if resp.status_code == 404: + return None + resp.raise_for_status() + return { + "content_type": resp.headers.get("content-type"), + "content_length": resp.headers.get("content-length"), + "last_modified": resp.headers.get("last-modified"), + } + except Exception as e: + status = getattr(getattr(e, "response", None), "status_code", None) + if status == 404: + return None + raise Exception(f"Failed to get file metadata for '{file_path}': {e}") + + def close(self): + """Close the HTTP handler to free resources.""" + if hasattr(self, "http_handler"): + self.http_handler.close() diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py new file mode 100644 index 00000000000..b782f10ccc5 --- /dev/null +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -0,0 +1,488 @@ +""" +GitLab prompt manager with configurable prompts folder. +""" + +from typing import Any, Dict, List, Optional, Tuple, Union +from jinja2 import DictLoader, Environment, select_autoescape + +from litellm.integrations.custom_prompt_management import CustomPromptManagement +from litellm.integrations.prompt_management_base import ( + PromptManagementBase, + PromptManagementClient, +) +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import StandardCallbackDynamicParams + +from litellm.integrations.gitlab.gitlab_client import GitLabClient + + +class GitLabPromptTemplate: + def __init__( + self, + template_id: str, + content: str, + metadata: Dict[str, Any], + model: Optional[str] = None, + ): + self.template_id = template_id + self.content = content + self.metadata = metadata + self.model = model or metadata.get("model") + self.temperature = metadata.get("temperature") + self.max_tokens = metadata.get("max_tokens") + self.input_schema = metadata.get("input", {}).get("schema", {}) + self.optional_params = { + k: v for k, v in metadata.items() if k not in ["model", "input", "content"] + } + + def __repr__(self): + return f"GitLabPromptTemplate(id='{self.template_id}', model='{self.model}')" + + +class GitLabTemplateManager: + """ + Manager for loading and rendering .prompt files from GitLab repositories. + + New: supports `prompts_path` (or `folder`) in gitlab_config to scope where prompts live. + """ + + + def __init__( + self, + gitlab_config: Dict[str, Any], + prompt_id: Optional[str] = None, + ref: Optional[str] = None, + gitlab_client: Optional[GitLabClient] = None + ): + self.gitlab_config = dict(gitlab_config) + self.prompt_id = prompt_id + self.prompts: Dict[str, GitLabPromptTemplate] = {} + self.gitlab_client = gitlab_client or GitLabClient(self.gitlab_config) + + if ref: + self.gitlab_client.set_ref(ref) + + # Folder inside repo to look for prompts (e.g., "prompts" or "prompts/chat") + self.prompts_path: str = ( + self.gitlab_config.get("prompts_path") + or self.gitlab_config.get("folder") + or "" + ).strip("/") + + self.jinja_env = Environment( + loader=DictLoader({}), + autoescape=select_autoescape(["html", "xml"]), + variable_start_string="{{", + variable_end_string="}}", + block_start_string="{%", + block_end_string="%}", + comment_start_string="{#", + comment_end_string="#}", + ) + + if self.prompt_id: + self._load_prompt_from_gitlab(self.prompt_id) + + # ---------- path helpers ---------- + + def _id_to_repo_path(self, prompt_id: str) -> str: + """Map a prompt_id to a repo path (respects prompts_path and adds .prompt).""" + if self.prompts_path: + return f"{self.prompts_path}/{prompt_id}.prompt" + return f"{prompt_id}.prompt" + + def _repo_path_to_id(self, repo_path: str) -> str: + """ + Map a repo path like 'prompts/chat/greeting.prompt' to an ID relative + to prompts_path without the extension (e.g., 'chat/greeting'). + """ + path = repo_path.strip("/") + if self.prompts_path and path.startswith(self.prompts_path.strip("/") + "/"): + path = path[len(self.prompts_path.strip("/")) + 1 :] + if path.endswith(".prompt"): + path = path[: -len(".prompt")] + return path + + # ---------- loading ---------- + + def _load_prompt_from_gitlab(self, prompt_id: str, *, ref: Optional[str] = None) -> None: + """Load a specific .prompt file from GitLab (scoped under prompts_path if set).""" + try: + file_path = self._id_to_repo_path(prompt_id) + prompt_content = self.gitlab_client.get_file_content(file_path, ref=ref) + if prompt_content: + template = self._parse_prompt_file(prompt_content, prompt_id) + self.prompts[prompt_id] = template + except Exception as e: + raise Exception(f"Failed to load prompt '{prompt_id}' from GitLab: {e}") + + def load_all_prompts(self, *, recursive: bool = True) -> List[str]: + """ + Eagerly load all .prompt files from prompts_path. Returns loaded IDs. + """ + files = self.list_templates(recursive=recursive) # reuse logic + loaded: List[str] = [] + for pid in files: + if pid not in self.prompts: + self._load_prompt_from_gitlab(pid) + loaded.append(pid) + return loaded + + # ---------- parsing & rendering ---------- + + def _parse_prompt_file( + self, content: str, prompt_id: str + ) -> GitLabPromptTemplate: + if content.startswith("---"): + parts = content.split("---", 2) + if len(parts) >= 3: + frontmatter_str = parts[1].strip() + template_content = parts[2].strip() + else: + frontmatter_str = "" + template_content = content + else: + frontmatter_str = "" + template_content = content + + metadata: Dict[str, Any] = {} + if frontmatter_str: + try: + import yaml + metadata = yaml.safe_load(frontmatter_str) or {} + except ImportError: + metadata = self._parse_yaml_basic(frontmatter_str) + except Exception: + metadata = {} + + return GitLabPromptTemplate( + template_id=prompt_id, + content=template_content, + metadata=metadata, + ) + + def _parse_yaml_basic(self, yaml_str: str) -> Dict[str, Any]: + result: Dict[str, Any] = {} + for line in yaml_str.split("\n"): + line = line.strip() + if ":" in line and not line.startswith("#"): + key, value = line.split(":", 1) + key = key.strip() + value = value.strip() + if value.lower() in ["true", "false"]: + result[key] = value.lower() == "true" + elif value.isdigit(): + result[key] = int(value) + elif value.replace(".", "").isdigit(): + try: + result[key] = float(value) + except Exception: + result[key] = value + else: + result[key] = value.strip("\"'") + return result + + def render_template( + self, template_id: str, variables: Optional[Dict[str, Any]] = None + ) -> str: + if template_id not in self.prompts: + raise ValueError(f"Template '{template_id}' not found") + template = self.prompts[template_id] + jinja_template = self.jinja_env.from_string(template.content) + return jinja_template.render(**(variables or {})) + + def get_template(self, template_id: str) -> Optional[GitLabPromptTemplate]: + return self.prompts.get(template_id) + + def list_templates(self, *, recursive: bool = True) -> List[str]: + """ + List available prompt IDs discovered under prompts_path (no extension, relative to prompts_path). + """ + """ + List available prompt IDs under prompts_path (no extension). + Compatible with both list_files signatures: + - list_files(directory_path=..., file_extension=..., recursive=...) + - list_files(path=..., ref=None, recursive=...) + """ + # First try the "new" signature (directory_path/file_extension) + try: + files = self.gitlab_client.list_files( + directory_path=self.prompts_path, + file_extension=".prompt", + recursive=recursive, + ) + base = self.prompts_path.strip("/") + out: List[str] = [] + for p in files or []: + path = str(p).strip("/") + if base and not path.startswith(base + "/"): + # if the client returns extra files outside the folder, skip them + continue + if not path.endswith(".prompt"): + continue + out.append(self._repo_path_to_id(path)) + return out + except TypeError: + # Fallback to the "classic" signature + raw = self.gitlab_client.list_files( + directory_path=self.prompts_path or "", + ref=None, + recursive=recursive, + ) + # Classic returns GitLab tree entries; filter *.prompt blobs + files = [] + for f in (raw or []): + if isinstance(f, dict) and f.get("type") == "blob" and str(f.get("path", "")).endswith(".prompt") and 'path' in f: + files.append(f['path']) + + return [self._repo_path_to_id(p) for p in files] + + +class GitLabPromptManager(CustomPromptManagement): + """ + GitLab prompt manager with folder support. + + Example config: + gitlab_config = { + "project": "group/subgroup/repo", + "access_token": "glpat_***", + "tag": "v1.2.3", # optional; takes precedence + "branch": "main", # default fallback + "prompts_path": "prompts/chat" # <--- NEW + } + """ + + def __init__( + self, + gitlab_config: Dict[str, Any], + prompt_id: Optional[str] = None, + ref: Optional[str] = None, # tag/branch/SHA override + gitlab_client: Optional[GitLabClient] = None + ): + self.gitlab_config = gitlab_config + self.prompt_id = prompt_id + self._prompt_manager: Optional[GitLabTemplateManager] = None + self._ref_override = ref + self._injected_gitlab_client = gitlab_client + if self.prompt_id: + self._prompt_manager = GitLabTemplateManager( + gitlab_config=self.gitlab_config, + prompt_id=self.prompt_id, + ref=self._ref_override, + ) + + @property + def integration_name(self) -> str: + return "gitlab" + + @property + def prompt_manager(self) -> GitLabTemplateManager: + if self._prompt_manager is None: + self._prompt_manager = GitLabTemplateManager( + gitlab_config=self.gitlab_config, + prompt_id=self.prompt_id, + ref=self._ref_override, + gitlab_client=self._injected_gitlab_client + ) + return self._prompt_manager + + def get_prompt_template( + self, + prompt_id: str, + prompt_variables: Optional[Dict[str, Any]] = None, + *, + ref: Optional[str] = None, + ) -> Tuple[str, Dict[str, Any]]: + if prompt_id not in self.prompt_manager.prompts: + self.prompt_manager._load_prompt_from_gitlab(prompt_id, ref=ref) + + template = self.prompt_manager.get_template(prompt_id) + if not template: + raise ValueError(f"Prompt template '{prompt_id}' not found") + + rendered_prompt = self.prompt_manager.render_template( + prompt_id, prompt_variables or {} + ) + + metadata = { + "model": template.model, + "temperature": template.temperature, + "max_tokens": template.max_tokens, + **template.optional_params, + } + return rendered_prompt, metadata + + def pre_call_hook( + self, + user_id: Optional[str], + messages: List[AllMessageValues], + function_call: Optional[Union[Dict[str, Any], str]] = None, + litellm_params: Optional[Dict[str, Any]] = None, + prompt_id: Optional[str] = None, + prompt_variables: Optional[Dict[str, Any]] = None, + prompt_version: Optional[str] = None, + **kwargs, + ) -> Tuple[List[AllMessageValues], Optional[Dict[str, Any]]]: + if not prompt_id: + return messages, litellm_params + try: + # Precedence: explicit prompt_version → per-call git_ref kwarg → manager override → config default + git_ref = prompt_version or kwargs.get("git_ref") or self._ref_override + + rendered_prompt, prompt_metadata = self.get_prompt_template( + prompt_id, prompt_variables, ref=git_ref + ) + parsed_messages = self._parse_prompt_to_messages(rendered_prompt) + + if parsed_messages: + final_messages: List[AllMessageValues] = parsed_messages + else: + final_messages = [{"role": "user", "content": rendered_prompt}] + messages # type: ignore + + if litellm_params is None: + litellm_params = {} + + if prompt_metadata.get("model"): + litellm_params["model"] = prompt_metadata["model"] + + for param in ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]: + if param in prompt_metadata: + litellm_params[param] = prompt_metadata[param] + + return final_messages, litellm_params + except Exception as e: + import litellm + litellm._logging.verbose_proxy_logger.error(f"Error in GitLab prompt pre_call_hook: {e}") + return messages, litellm_params + + + def _parse_prompt_to_messages(self, prompt_content: str) -> List[AllMessageValues]: + messages: List[AllMessageValues] = [] + lines = prompt_content.strip().split("\n") + current_role: Optional[str] = None + current_content: List[str] = [] + + for raw in lines: + line = raw.strip() + if not line: + continue + low = line.lower() + if low.startswith("system:"): + if current_role and current_content: + messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore + current_role = "system" + current_content = [line[7:].strip()] + elif low.startswith("user:"): + if current_role and current_content: + messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore + current_role = "user" + current_content = [line[5:].strip()] + elif low.startswith("assistant:"): + if current_role and current_content: + messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore + current_role = "assistant" + current_content = [line[10:].strip()] + else: + current_content.append(line) + + if current_role and current_content: + messages.append({"role": current_role, "content": "\n".join(current_content).strip()}) # type: ignore + if not messages and prompt_content.strip(): + messages = [{"role": "user", "content": prompt_content.strip()}] # type: ignore + return messages + + def post_call_hook( + self, + user_id: Optional[str], + response: Any, + input_messages: List[AllMessageValues], + function_call: Optional[Union[Dict[str, Any], str]] = None, + litellm_params: Optional[Dict[str, Any]] = None, + prompt_id: Optional[str] = None, + prompt_variables: Optional[Dict[str, Any]] = None, + **kwargs, + ) -> Any: + return response + + def get_available_prompts(self) -> List[str]: + """ + Return prompt IDs. Prefer already-loaded templates in memory to avoid + unnecessary network calls (and to make tests deterministic). + """ + ids = set(self.prompt_manager.prompts.keys()) + try: + ids.update(self.prompt_manager.list_templates()) + except Exception: + # If GitLab list fails (auth, network), still return what we've loaded. + pass + return sorted(ids) + + def reload_prompts(self) -> None: + if self.prompt_id: + self._prompt_manager = None + _ = self.prompt_manager # trigger re-init/load + + def should_run_prompt_management( + self, + prompt_id: str, + dynamic_callback_params: StandardCallbackDynamicParams, + ) -> bool: + return True + + def _compile_prompt_helper( + self, + prompt_id: str, + prompt_variables: Optional[dict], + dynamic_callback_params: StandardCallbackDynamicParams, + prompt_label: Optional[str] = None, + prompt_version: Optional[int] = None, + ) -> PromptManagementClient: + try: + if prompt_id not in self.prompt_manager.prompts: + git_ref = getattr(dynamic_callback_params, "extra", {}).get("git_ref") if hasattr(dynamic_callback_params, "extra") else None + self.prompt_manager._load_prompt_from_gitlab(prompt_id, ref=git_ref) + + rendered_prompt, prompt_metadata = self.get_prompt_template( + prompt_id, prompt_variables + ) + + messages = self._parse_prompt_to_messages(rendered_prompt) + template_model = prompt_metadata.get("model") + + optional_params: Dict[str, Any] = {} + for param in ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]: + if param in prompt_metadata: + optional_params[param] = prompt_metadata[param] + + return PromptManagementClient( + prompt_id=prompt_id, + prompt_template=messages, + prompt_template_model=template_model, + prompt_template_optional_params=optional_params, + completed_messages=None, + ) + except Exception as e: + raise ValueError(f"Error compiling prompt '{prompt_id}': {e}") + + def get_chat_completion_prompt( + self, + model: str, + messages: List[AllMessageValues], + non_default_params: dict, + prompt_id: Optional[str], + prompt_variables: Optional[dict], + dynamic_callback_params: StandardCallbackDynamicParams, + prompt_label: Optional[str] = None, + prompt_version: Optional[int] = None, + ) -> Tuple[str, List[AllMessageValues], dict]: + return PromptManagementBase.get_chat_completion_prompt( + self, + model, + messages, + non_default_params, + prompt_id, + prompt_variables, + dynamic_callback_params, + prompt_label, + prompt_version, + ) diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 69943a0fe4d..7f807bb8b0c 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -1,6 +1,5 @@ #### What this does #### # On success, logs events to Langfuse -import copy import os import traceback from datetime import datetime @@ -11,6 +10,7 @@ from packaging.version import Version import litellm from litellm._logging import verbose_logger from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS +from litellm.litellm_core_utils.core_helpers import safe_deep_copy from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info from litellm.llms.custom_httpx.http_handler import _get_httpx_client from litellm.secret_managers.main import str_to_bool @@ -222,7 +222,7 @@ class LangFuseLogger: litellm_params.get("metadata", {}) or {} ) # if litellm_params['metadata'] == None metadata = self.add_metadata_from_header(litellm_params, metadata) - optional_params = copy.deepcopy(kwargs.get("optional_params", {})) + optional_params = safe_deep_copy(kwargs.get("optional_params", {})) prompt = {"messages": kwargs.get("messages")} @@ -690,6 +690,7 @@ class LangFuseLogger: } usage_details = LangfuseUsageDetails(input=_usage_obj.prompt_tokens, output=_usage_obj.completion_tokens, + total=_usage_obj.total_tokens, cache_creation_input_tokens=_usage_obj.get('cache_creation_input_tokens', 0), cache_read_input_tokens=_usage_obj.get('cache_read_input_tokens', 0)) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index e6f265ded58..e825f89f56e 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -575,9 +575,16 @@ class OpenTelemetry(CustomLogger): if litellm.turn_off_message_logging or not self.message_logging: return + litellm_params = kwargs.get("litellm_params", {}) + metadata = litellm_params.get("metadata") or {} + generation_name = metadata.get("generation_name") + + raw_span_name = generation_name if generation_name else RAW_REQUEST_SPAN_NAME + + otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs) raw_span = otel_tracer.start_span( - name=RAW_REQUEST_SPAN_NAME, + name=raw_span_name, start_time=self._to_ns(start_time), context=trace.set_span_in_context(parent_span), ) @@ -645,7 +652,7 @@ class OpenTelemetry(CustomLogger): if not self.config.enable_events: return - from opentelemetry._logs import get_logger, LogRecord + from opentelemetry._logs import LogRecord, get_logger otel_logger = get_logger(LITELLM_LOGGER_NAME) parent_ctx = span.get_span_context() @@ -1115,56 +1122,68 @@ class OpenTelemetry(CustomLogger): span.set_attribute(key, primitive_value) def set_raw_request_attributes(self, span: Span, kwargs, response_obj): - kwargs.get("optional_params", {}) - litellm_params = kwargs.get("litellm_params", {}) or {} - custom_llm_provider = litellm_params.get("custom_llm_provider", "Unknown") + try: + kwargs.get("optional_params", {}) + litellm_params = kwargs.get("litellm_params", {}) or {} + custom_llm_provider = litellm_params.get("custom_llm_provider", "Unknown") - _raw_response = kwargs.get("original_response") - _additional_args = kwargs.get("additional_args", {}) or {} - complete_input_dict = _additional_args.get("complete_input_dict") - ############################################# - ########## LLM Request Attributes ########### - ############################################# + _raw_response = kwargs.get("original_response") + _additional_args = kwargs.get("additional_args", {}) or {} + complete_input_dict = _additional_args.get("complete_input_dict") + ############################################# + ########## LLM Request Attributes ########### + ############################################# - # OTEL Attributes for the RAW Request to https://docs.anthropic.com/en/api/messages - if complete_input_dict and isinstance(complete_input_dict, dict): - for param, val in complete_input_dict.items(): - self.safe_set_attribute( - span=span, key=f"llm.{custom_llm_provider}.{param}", value=val - ) + # OTEL Attributes for the RAW Request to https://docs.anthropic.com/en/api/messages + if complete_input_dict and isinstance(complete_input_dict, dict): + for param, val in complete_input_dict.items(): + self.safe_set_attribute( + span=span, key=f"llm.{custom_llm_provider}.{param}", value=val + ) - ############################################# - ########## LLM Response Attributes ########## - ############################################# - if _raw_response and isinstance(_raw_response, str): - # cast sr -> dict - import json + ############################################# + ########## LLM Response Attributes ########## + ############################################# + if _raw_response and isinstance(_raw_response, str): + # cast sr -> dict + import json + + try: + _raw_response = json.loads(_raw_response) + for param, val in _raw_response.items(): + self.safe_set_attribute( + span=span, + key=f"llm.{custom_llm_provider}.{param}", + value=val, + ) + except json.JSONDecodeError: + verbose_logger.debug( + "litellm.integrations.opentelemetry.py::set_raw_request_attributes() - raw_response not json string - {}".format( + _raw_response + ) + ) - try: - _raw_response = json.loads(_raw_response) - for param, val in _raw_response.items(): self.safe_set_attribute( span=span, - key=f"llm.{custom_llm_provider}.{param}", - value=val, + key=f"llm.{custom_llm_provider}.stringified_raw_response", + value=_raw_response, ) - except json.JSONDecodeError: - verbose_logger.debug( - "litellm.integrations.opentelemetry.py::set_raw_request_attributes() - raw_response not json string - {}".format( - _raw_response - ) - ) - - self.safe_set_attribute( - span=span, - key=f"llm.{custom_llm_provider}.stringified_raw_response", - value=_raw_response, - ) + except Exception as e: + verbose_logger.exception( + "OpenTelemetry logging error in set_raw_request_attributes %s", str(e) + ) def _to_ns(self, dt): return int(dt.timestamp() * 1e9) def _get_span_name(self, kwargs): + litellm_params = kwargs.get("litellm_params", {}) + metadata = litellm_params.get("metadata") or {} + generation_name = metadata.get("generation_name") + + if generation_name: + return generation_name + return LITELLM_REQUEST_SPAN_NAME def get_traceparent_from_header(self, headers): diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 4957c97e5b2..09794bf2677 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -16,6 +16,7 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont from litellm.integrations.argilla import ArgillaLogger from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger from litellm.integrations.bitbucket import BitBucketPromptManager +from litellm.integrations.gitlab import GitLabPromptManager from litellm.integrations.braintrust_logging import BraintrustLogger from litellm.integrations.datadog.datadog import DataDogLogger from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger @@ -92,6 +93,7 @@ class CustomLoggerRegistry: "vector_store_pre_call_hook": VectorStorePreCallHook, "dotprompt": DotpromptManager, "bitbucket": BitBucketPromptManager, + "gitlab": GitLabPromptManager, "cloudzero": CloudZeroLogger, "posthog": PostHogLogger, } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 46e363c865d..696c67c44d7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -89,6 +89,7 @@ from litellm.types.utils import ( CostResponseTypes, DynamicPromptManagementParamLiteral, EmbeddingResponse, + GuardrailStatus, ImageResponse, LiteLLMBatch, LiteLLMLoggingBaseClass, @@ -107,6 +108,7 @@ from litellm.types.utils import ( StandardLoggingPayload, StandardLoggingPayloadErrorInformation, StandardLoggingPayloadStatus, + StandardLoggingPayloadStatusFields, StandardLoggingPromptManagementMetadata, StandardLoggingVectorStoreRequest, TextCompletionResponse, @@ -3667,6 +3669,25 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 bitbucket_logger = BitBucketPromptManager(bitbucket_config=bitbucket_config) _in_memory_loggers.append(bitbucket_logger) return bitbucket_logger # type: ignore + elif logging_integration == "gitlab": + from litellm.integrations.gitlab.gitlab_prompt_manager import ( + GitLabPromptManager, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, GitLabPromptManager): + return callback + + # Get global BitBucket config + gitlab_config = getattr(litellm, "global_gitlab_config", None) + if gitlab_config is None: + raise ValueError( + "Gitlab configuration not found. Please set litellm.global_gitlab_config first." + ) + + gitlab_logger = GitLabPromptManager(gitlab_config=gitlab_config) + _in_memory_loggers.append(gitlab_logger) + return gitlab_logger # type: ignore return None except Exception as e: verbose_logger.exception( @@ -4426,6 +4447,51 @@ class StandardLoggingPayloadSetup: return request_tags + +def _get_status_fields( + status: StandardLoggingPayloadStatus, + guardrail_information: Optional[dict], + error_str: Optional[str] +) -> "StandardLoggingPayloadStatusFields": + """ + Determine status fields based on request status and guardrail information. + + Args: + status: Overall request status ("success" or "failure") + guardrail_information: Guardrail information from metadata + error_str: Error string if any + + Returns: + StandardLoggingPayloadStatusFields with llm_api_status and guardrail_status + """ + # Mapping for legacy guardrail status values to new GuardrailStatus values + GUARDRAIL_STATUS_MAP: Dict[str, GuardrailStatus] = { + "success": "success", + "blocked": "guardrail_intervened", # legacy + "guardrail_intervened": "guardrail_intervened", # direct + "failure": "guardrail_failed_to_respond", # legacy + "guardrail_failed_to_respond": "guardrail_failed_to_respond", # direct + "not_run": "not_run" + } + + # Set LLM API status + llm_api_status: StandardLoggingPayloadStatus = status + + + ######################################################### + # Map - guardrail_information.guardrail_status to guardrail_status + ######################################################### + guardrail_status: GuardrailStatus = "not_run" + if guardrail_information and isinstance(guardrail_information, dict): + raw_status = guardrail_information.get("guardrail_status", "not_run") + guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") + + return StandardLoggingPayloadStatusFields( + llm_api_status=llm_api_status, + guardrail_status=guardrail_status + ) + + def get_standard_logging_object_payload( kwargs: Optional[dict], init_response_obj: Union[Any, BaseModel, dict], @@ -4535,7 +4601,6 @@ def get_standard_logging_object_payload( start_time=start_time, response_id=id, ) - _request_body = proxy_server_request.get("body", {}) end_user_id = clean_metadata["user_api_key_end_user_id"] or _request_body.get( "user", None @@ -4591,6 +4656,11 @@ def get_standard_logging_object_payload( cache_hit=cache_hit, stream=stream, status=status, + status_fields=_get_status_fields( + status=status, + guardrail_information=metadata.get("standard_logging_guardrail_information", None), + error_str=error_str + ), custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")), saved_cache_cost=saved_cache_cost, startTime=start_time_float, diff --git a/litellm/litellm_core_utils/model_response_utils.py b/litellm/litellm_core_utils/model_response_utils.py index 5f6fced9d44..974d12aef6f 100644 --- a/litellm/litellm_core_utils/model_response_utils.py +++ b/litellm/litellm_core_utils/model_response_utils.py @@ -84,7 +84,9 @@ def _has_meaningful_content(value: Any) -> bool: return False if isinstance(value, str): - return len(value.strip()) > 0 + # Don't strip whitespace - preserve all content including newlines, spaces, etc. + # Even pure whitespace characters like '\n' or ' ' are meaningful content + return len(value) > 0 if isinstance(value, (list, dict)): return len(value) > 0 diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 07f652ecb9b..05f1a37ca12 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -21,6 +21,8 @@ class SensitiveDataMasker: "access", "private", "certificate", + "fingerprint", + "tenancy", } self.visible_prefix = visible_prefix @@ -42,7 +44,14 @@ class SensitiveDataMasker: def is_sensitive_key(self, key: str) -> bool: key_lower = str(key).lower() - result = any(pattern in key_lower for pattern in self.sensitive_patterns) + # Split on underscores and check if any segment matches the pattern + # This avoids false positives like "max_tokens" matching "token" + # but still catches "api_key", "access_token", etc. + key_segments = key_lower.replace('-', '_').split('_') + result = any( + pattern in key_segments + for pattern in self.sensitive_patterns + ) return result def mask_dict( diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 3645c16bf8f..7c5b693b453 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -1117,6 +1117,14 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): status_code=422, message="max retries must be an int" ) + if api_key is None and azure_ad_token_provider is not None: + azure_ad_token = azure_ad_token_provider() + if azure_ad_token: + headers.pop( + "api-key", None + ) + headers["Authorization"] = f"Bearer {azure_ad_token}" + # init AzureOpenAI Client azure_client_params: Dict[str, Any] = self.initialize_azure_sdk_client( litellm_params=litellm_params or {}, diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 8701fe57bfd..6e9c03dee89 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import httpx -from litellm.types.rerank import OptionalRerankParams, RerankBilledUnits, RerankResponse +from litellm.types.rerank import RerankBilledUnits, RerankResponse from litellm.types.utils import ModelInfo from ..chat.transformation import BaseLLMException @@ -30,7 +30,7 @@ class BaseRerankConfig(ABC): def transform_rerank_request( self, model: str, - optional_rerank_params: OptionalRerankParams, + optional_rerank_params: Dict, headers: dict, ) -> dict: return {} @@ -78,7 +78,7 @@ class BaseRerankConfig(ABC): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, - ) -> OptionalRerankParams: + ) -> Dict: pass def get_error_class( diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 2b111cde600..241359d937e 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -440,27 +440,30 @@ class BedrockModelInfo(BaseLLMModelInfo): """ Abbreviations of regions AWS Bedrock supports for cross region inference """ - return ["us", "eu", "apac"] + return ["us", "eu", "apac", "jp"] @staticmethod def get_bedrock_route( model: str, - ) -> Literal["converse", "invoke", "converse_like", "agent"]: + ) -> Literal["converse", "invoke", "converse_like", "agent", "async_invoke"]: """ Get the bedrock route for the given model. """ - route_mappings: Dict[str, Literal["invoke", "converse_like", "converse", "agent"]] = { + route_mappings: Dict[ + str, Literal["invoke", "converse_like", "converse", "agent", "async_invoke"] + ] = { "invoke/": "invoke", - "converse_like/": "converse_like", + "converse_like/": "converse_like", "converse/": "converse", - "agent/": "agent" + "agent/": "agent", + "async_invoke/": "async_invoke", } - + # Check explicit routes first for prefix, route_type in route_mappings.items(): if prefix in model: return route_type - + base_model = BedrockModelInfo.get_base_model(model) alt_model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model) if ( @@ -469,38 +472,46 @@ class BedrockModelInfo(BaseLLMModelInfo): ): return "converse" return "invoke" - + @staticmethod def _explicit_converse_route(model: str) -> bool: """ Check if the model is an explicit converse route. """ return "converse/" in model - + @staticmethod def _explicit_invoke_route(model: str) -> bool: """ Check if the model is an explicit invoke route. """ return "invoke/" in model - + @staticmethod def _explicit_agent_route(model: str) -> bool: """ Check if the model is an explicit agent route. """ return "agent/" in model - + @staticmethod def _explicit_converse_like_route(model: str) -> bool: """ Check if the model is an explicit converse like route. """ return "converse_like/" in model - @staticmethod - def get_bedrock_provider_config_for_messages_api(model: str) -> Optional[BaseAnthropicMessagesConfig]: + def _explicit_async_invoke_route(model: str) -> bool: + """ + Check if the model is an explicit async invoke route. + """ + return "async_invoke/" in model + + @staticmethod + def get_bedrock_provider_config_for_messages_api( + model: str, + ) -> Optional[BaseAnthropicMessagesConfig]: """ Get the bedrock provider config for the given model. @@ -513,19 +524,20 @@ class BedrockModelInfo(BaseLLMModelInfo): # Converse routes should go through litellm.completion() if BedrockModelInfo._explicit_converse_route(model): return None - + ######################################################### # This goes through litellm.AmazonAnthropicClaude3MessagesConfig() # Since bedrock Invoke supports Native Anthropic Messages API ######################################################### if "claude" in model: return litellm.AmazonAnthropicClaudeMessagesConfig() - + ######################################################### # These routes will go through litellm.completion() ######################################################### return None + class BedrockEventStreamDecoderBase: """ Base class for event stream decoding for Bedrock @@ -595,20 +607,20 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: """ Extract anthropic-beta header values and convert them to a list. Supports comma-separated values from user headers. - + Used by both converse and invoke transformations for consistent handling of anthropic-beta headers that should be passed to AWS Bedrock. - + Args: headers (dict): Request headers dictionary - + Returns: List[str]: List of anthropic beta feature strings, empty list if no header """ anthropic_beta_header = headers.get("anthropic-beta") if not anthropic_beta_header: return [] - + # Split comma-separated values and strip whitespace return [beta.strip() for beta in anthropic_beta_header.split(",")] @@ -618,19 +630,20 @@ 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") """ @@ -641,41 +654,45 @@ class CommonBatchFilesUtils: 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: + 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"]) - + 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 @@ -690,7 +707,7 @@ class CommonBatchFilesUtils: return model_name except Exception: pass - + # Fallback to default model return "anthropic.claude-3-5-sonnet-20240620-v1:0" @@ -704,14 +721,14 @@ class CommonBatchFilesUtils: ) -> 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) """ @@ -736,7 +753,7 @@ class CommonBatchFilesUtils: aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), ) - + # Prepare the request data method_upper = method.upper() if method_upper == "GET": @@ -746,12 +763,13 @@ class CommonBatchFilesUtils: else: if isinstance(data, dict): import json + request_data = json.dumps(data) else: request_data = data # Prepare headers for non-GET requests headers = {"Content-Type": "application/json"} - + # Create AWS request and sign it sigv4 = SigV4Auth(credentials, service_name, aws_region_name) request = AWSRequest( @@ -759,45 +777,51 @@ class CommonBatchFilesUtils: ) sigv4.add_auth(request) prepped = request.prepare() - - return dict(prepped.headers), request_data.encode('utf-8') if isinstance(request_data, str) else request_data + + 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) """ from litellm._uuid import 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, + self, + litellm_params: dict, optional_params: dict, bucket_env_var: str = "AWS_S3_BUCKET_NAME", - key_prefix: str = "litellm" + 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) """ @@ -806,18 +830,20 @@ class CommonBatchFilesUtils: # Get bucket name bucket_name = ( - litellm_params.get("s3_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") - + 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( @@ -827,7 +853,5 @@ class CommonBatchFilesUtils: Get Bedrock-specific error class. """ return BedrockError( - status_code=status_code, - message=error_message, - headers=headers + status_code=status_code, message=error_message, headers=headers ) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index d4dd716a1f4..3edd6d6741b 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -22,9 +22,8 @@ from litellm.secret_managers.main import get_secret from litellm.types.llms.bedrock import ( AmazonEmbeddingRequest, CohereEmbeddingRequest, - TwelveLabsMarengoEmbeddingRequest, ) -from litellm.types.utils import EmbeddingResponse +from litellm.types.utils import EmbeddingResponse, LlmProviders from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError @@ -77,7 +76,7 @@ class BedrockEmbedding(BaseAWSLLM): if aws_region_name is None: aws_region_name = "us-west-2" - credentials: Credentials = self.get_credentials( + credentials: Credentials = self.get_credentials( # type: ignore aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, @@ -151,35 +150,80 @@ class BedrockEmbedding(BaseAWSLLM): raise BedrockError(status_code=408, message="Timeout error occurred.") return response.json() - + def _transform_response( - self, response_list: List[dict], model: str, provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL + self, + response_list: List[dict], + model: str, + provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, + is_async_invoke: Optional[bool] = False, ) -> Optional[EmbeddingResponse]: """ Transforms the response from the Bedrock embedding provider to the OpenAI format. """ returned_response: Optional[EmbeddingResponse] = None - if model == "amazon.titan-embed-image-v1": - returned_response = ( - AmazonTitanMultimodalEmbeddingG1Config()._transform_response( + + # Handle async invoke responses (single response with invocationArn) + if ( + is_async_invoke + and len(response_list) == 1 + and "invocationArn" in response_list[0] + ): + if provider == "twelvelabs": + returned_response = ( + TwelveLabsMarengoEmbeddingConfig()._transform_async_invoke_response( + response=response_list[0], model=model + ) + ) + else: + # For other providers, create a generic async response + invocation_arn = response_list[0].get("invocationArn", "") + + from litellm.types.utils import Embedding, Usage + + embedding = Embedding( + embedding=[], + index=0, + object="embedding", # Must be literal "embedding" + ) + usage = Usage(prompt_tokens=0, total_tokens=0) + + # Create hidden params with job ID + from litellm.types.llms.base import HiddenParams + + hidden_params = HiddenParams() + setattr(hidden_params, "_invocation_arn", invocation_arn) + + returned_response = EmbeddingResponse( + data=[embedding], + model=model, + usage=usage, + hidden_params=hidden_params, + ) + else: + # Handle regular invoke responses + if model == "amazon.titan-embed-image-v1": + returned_response = ( + AmazonTitanMultimodalEmbeddingG1Config()._transform_response( + response_list=response_list, model=model + ) + ) + elif model == "amazon.titan-embed-text-v1": + returned_response = AmazonTitanG1Config()._transform_response( response_list=response_list, model=model ) - ) - elif model == "amazon.titan-embed-text-v1": - returned_response = AmazonTitanG1Config()._transform_response( - response_list=response_list, model=model - ) - elif model == "amazon.titan-embed-text-v2:0": - returned_response = AmazonTitanV2Config()._transform_response( - response_list=response_list, model=model - ) - elif provider == "twelvelabs": - returned_response = TwelveLabsMarengoEmbeddingConfig()._transform_response( - response_list=response_list, model=model - ) - - - ########################################################## + elif model == "amazon.titan-embed-text-v2:0": + returned_response = AmazonTitanV2Config()._transform_response( + response_list=response_list, model=model + ) + elif provider == "twelvelabs": + returned_response = ( + TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=response_list, model=model + ) + ) + + ########################################################## # Validate returned response ########################################################## if returned_response is None: @@ -203,6 +247,7 @@ class BedrockEmbedding(BaseAWSLLM): logging_obj: Any, provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, api_key: Optional[str] = None, + is_async_invoke: Optional[bool] = False, ): responses: List[dict] = [] for data in batch_data: @@ -210,7 +255,7 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = self.get_request_headers( + prepped = self.get_request_headers( # type: ignore # type: ignore credentials=credentials, aws_region_name=aws_region_name, extra_headers=extra_headers, @@ -249,7 +294,10 @@ class BedrockEmbedding(BaseAWSLLM): responses.append(response) return self._transform_response( - response_list=responses, model=model, provider=provider + response_list=responses, + model=model, + provider=provider, + is_async_invoke=is_async_invoke, ) async def _async_single_func_embeddings( @@ -265,6 +313,7 @@ class BedrockEmbedding(BaseAWSLLM): logging_obj: Any, provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, api_key: Optional[str] = None, + is_async_invoke: Optional[bool] = False, ): responses: List[dict] = [] for data in batch_data: @@ -272,7 +321,7 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = self.get_request_headers( + prepped = self.get_request_headers( # type: ignore # type: ignore credentials=credentials, aws_region_name=aws_region_name, extra_headers=extra_headers, @@ -311,7 +360,10 @@ class BedrockEmbedding(BaseAWSLLM): responses.append(response) ## TRANSFORM RESPONSE ## return self._transform_response( - response_list=responses, model=model, provider=provider + response_list=responses, + model=model, + provider=provider, + is_async_invoke=is_async_invoke, ) def embeddings( @@ -343,7 +395,10 @@ class BedrockEmbedding(BaseAWSLLM): model=model, model_id=unencoded_model_id, ) - + # Check async invoke needs to be used + has_async_invoke = "async_invoke/" in model + if has_async_invoke: + model = model.replace("async_invoke/", "", 1) provider = self.get_bedrock_embedding_provider(model) if provider is None: raise Exception( @@ -402,10 +457,14 @@ class BedrockEmbedding(BaseAWSLLM): elif provider == "twelvelabs": batch_data = [] for i in input: - twelvelabs_request: ( - TwelveLabsMarengoEmbeddingRequest - ) = TwelveLabsMarengoEmbeddingConfig()._transform_request( - input=i, inference_params=inference_params + twelvelabs_request = ( + TwelveLabsMarengoEmbeddingConfig()._transform_request( + input=i, + inference_params=inference_params, + async_invoke_route=has_async_invoke, + model_id=modelId, + output_s3_uri=inference_params.get("output_s3_uri"), + ) ) batch_data.append(twelvelabs_request) @@ -417,7 +476,10 @@ class BedrockEmbedding(BaseAWSLLM): ), aws_region_name=aws_region_name, ) - endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" + if has_async_invoke: + endpoint_url = f"{endpoint_url}/async-invoke" + else: + endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" if batch_data is not None: if aembedding: @@ -437,6 +499,7 @@ class BedrockEmbedding(BaseAWSLLM): logging_obj=logging_obj, api_key=api_key, provider=provider, + is_async_invoke=has_async_invoke, ) returned_response = self._single_func_embeddings( client=( @@ -454,6 +517,7 @@ class BedrockEmbedding(BaseAWSLLM): logging_obj=logging_obj, api_key=api_key, provider=provider, + is_async_invoke=has_async_invoke, ) if returned_response is None: raise Exception("Unable to map Bedrock request to provider") @@ -465,7 +529,7 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = self.get_request_headers( + prepped = self.get_request_headers( # type: ignore credentials=credentials, aws_region_name=aws_region_name, extra_headers=extra_headers, @@ -491,3 +555,94 @@ class BedrockEmbedding(BaseAWSLLM): client=client, headers=prepped.headers, # type: ignore ) + + async def _get_async_invoke_status( + self, invocation_arn: str, aws_region_name: str, logging_obj=None, **kwargs + ) -> dict: + """ + Get the status of an async invoke job using the GetAsyncInvoke operation. + + Args: + invocation_arn: The invocation ARN from the async invoke response + aws_region_name: AWS region name + **kwargs: Additional parameters (credentials, etc.) + + Returns: + dict: Status response from AWS Bedrock + """ + + # Get AWS credentials using the same method as other Bedrock methods + credentials, _ = self._load_credentials(kwargs) + + # Get the runtime endpoint + endpoint_url, _ = self.get_runtime_endpoint( + api_base=None, + aws_bedrock_runtime_endpoint=kwargs.get("aws_bedrock_runtime_endpoint"), + aws_region_name=aws_region_name, + ) + + # Construct the status check URL + status_url = f"{endpoint_url}/async-invoke/{invocation_arn}" + + # Prepare headers + headers = {"Content-Type": "application/json"} + + # Get AWS signed headers + prepped = self.get_request_headers( # type: ignore + credentials=credentials, + aws_region_name=aws_region_name, + extra_headers=None, + endpoint_url=status_url, + data="", # GET request, no body + headers=headers, + api_key=None, + ) + + # LOGGING + if logging_obj is not None: + # Create custom curl command for GET request + masked_headers = logging_obj._get_masked_headers(prepped.headers) + formatted_headers = " ".join( + [f"-H '{k}: {v}'" for k, v in masked_headers.items()] + ) + custom_curl = "\n\nGET Request Sent from LiteLLM:\n" + custom_curl += "curl -X GET \\\n" + custom_curl += f"{prepped.url} \\\n" + custom_curl += f"{formatted_headers}\n" + + logging_obj.pre_call( + input=invocation_arn, + api_key="", + additional_args={ + "complete_input_dict": {"invocation_arn": invocation_arn}, + "api_base": prepped.url, + "headers": prepped.headers, + "request_str": custom_curl, # Override with custom GET curl command + }, + ) + + # Make the GET request + client = get_async_httpx_client(llm_provider=LlmProviders.BEDROCK) + response = await client.get( + url=prepped.url, + headers=prepped.headers, + ) + + # LOGGING + if logging_obj is not None: + logging_obj.post_call( + input=invocation_arn, + api_key="", + original_response=response, + additional_args={ + "complete_input_dict": {"invocation_arn": invocation_arn} + }, + ) + + # Parse response + if response.status_code == 200: + return response.json() + else: + raise Exception( + f"Failed to get async invoke status: {response.status_code} - {response.text}" + ) diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py index fdad8a65043..c85c388eebc 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -1,33 +1,47 @@ """ -Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Marengo /invoke format. +Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Marengo /invoke and /async-invoke format. Why separate file? Make it easy to see how transformation works Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html """ -from typing import List +from typing import List, Optional, Union, cast from litellm.types.llms.bedrock import ( + TWELVELABS_EMBEDDING_INPUT_TYPES, + TwelveLabsAsyncInvokeRequest, TwelveLabsMarengoEmbeddingRequest, + TwelveLabsOutputDataConfig, + TwelveLabsS3Location, + TwelveLabsS3OutputDataConfig, ) from litellm.types.utils import Embedding, EmbeddingResponse, Usage -from litellm.utils import get_base64_str, is_base64_encoded class TwelveLabsMarengoEmbeddingConfig: """ Reference - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html - Supports text and image inputs for Phase 1. - Video and audio support will be added in Phase 2. + Supports text, image, video, and audio inputs. + - InvokeModel: text and image inputs + - StartAsyncInvoke: video, audio, image, and text inputs """ def __init__(self) -> None: pass def get_supported_openai_params(self) -> List[str]: - return ["encoding_format", "textTruncate", "embeddingOption"] + return [ + "encoding_format", + "textTruncate", + "embeddingOption", + "startSec", + "lengthSec", + "useFixedLengthSec", + "minClipSec", + "input_type", + ] def map_openai_params( self, non_default_params: dict, optional_params: dict @@ -41,45 +55,142 @@ class TwelveLabsMarengoEmbeddingConfig: optional_params["textTruncate"] = v elif k == "embeddingOption": optional_params["embeddingOption"] = v + elif k == "input_type": + # Map input_type to inputType for Bedrock + optional_params["inputType"] = v + elif k in ["startSec", "lengthSec", "useFixedLengthSec", "minClipSec"]: + optional_params[k] = v return optional_params + def _extract_bucket_owner_from_params(self, inference_params: dict) -> str: + """ + Extract bucket owner from inference parameters. + """ + return inference_params.get("bucketOwner", "") + + def _is_s3_url(self, input: str) -> bool: + """Check if input is an S3 URL.""" + return input.startswith("s3://") + def _transform_request( - self, input: str, inference_params: dict - ) -> TwelveLabsMarengoEmbeddingRequest: + self, + input: str, + inference_params: dict, + async_invoke_route: bool = False, + model_id: Optional[str] = None, + output_s3_uri: Optional[str] = None, + ) -> Union[TwelveLabsMarengoEmbeddingRequest, TwelveLabsAsyncInvokeRequest]: """ - Transform OpenAI-style input to TwelveLabs Marengo format. - Phase 1: Supports text and image inputs only. - """ - # Check if input is base64 encoded image - is_encoded = is_base64_encoded(input) + Transform OpenAI-style input to TwelveLabs Marengo format/async-invoke format. - if is_encoded: - # Image input - b64_str = get_base64_str(input) - transformed_request = TwelveLabsMarengoEmbeddingRequest( - inputType="image", mediaSource={"base64String": b64_str} - ) - else: - # Text input - transformed_request = TwelveLabsMarengoEmbeddingRequest( - inputType="text", inputText=input + Supports: + - Text inputs (for both invoke and async-invoke) + - Image inputs (for both invoke and async-invoke) + - Video inputs (async-invoke only) + - Audio inputs (async-invoke only) + - S3 URLs for all media types (async-invoke only) + """ + # Get input_type or default to "text" + input_type = cast( + TWELVELABS_EMBEDDING_INPUT_TYPES, + inference_params.get("inputType") or inference_params.get("input_type") or "text" + ) + + # Validate that async-invoke is used for video/audio + if input_type in ["video", "audio"] and not async_invoke_route: + raise ValueError( + f"Input type '{input_type}' requires async_invoke route. " + f"Use model format: 'bedrock/async_invoke/model_id'" ) + transformed_request: TwelveLabsMarengoEmbeddingRequest = { + "inputType": input_type + } + + if input_type == "text": + transformed_request["inputText"] = input # Set default textTruncate if not specified if "textTruncate" not in inference_params: transformed_request["textTruncate"] = "end" + elif input_type in ["image", "video", "audio"]: + if self._is_s3_url(input): + # S3 URL input + s3_location: TwelveLabsS3Location = {"uri": input} + bucket_owner = self._extract_bucket_owner_from_params(inference_params) + if bucket_owner: + s3_location["bucketOwner"] = bucket_owner + + transformed_request["mediaSource"] = {"s3Location": s3_location} + else: + # Base64 encoded input + if input.startswith("data:"): + # Extract base64 data from data URL + b64_str = input.split(",", 1)[1] if "," in input else input + else: + # Direct base64 string + from litellm.utils import get_base64_str + b64_str = get_base64_str(input) + + transformed_request["mediaSource"] = {"base64String": b64_str} + # Apply any additional inference parameters for k, v in inference_params.items(): if k not in [ "inputType", + "input_type", # Exclude both camelCase and snake_case "inputText", "mediaSource", + "bucketOwner", # Don't include bucketOwner in the request ]: # Don't override core fields transformed_request[k] = v # type: ignore + # If async invoke route, wrap in the async invoke format + if async_invoke_route and model_id: + return self._wrap_async_invoke_request( + model_input=transformed_request, + model_id=model_id, + output_s3_uri=output_s3_uri, + ) + return transformed_request + def _wrap_async_invoke_request( + self, + model_input: TwelveLabsMarengoEmbeddingRequest, + model_id: str, + output_s3_uri: Optional[str] = None, + ) -> TwelveLabsAsyncInvokeRequest: + """ + Wrap the transformed request in the correct AWS Bedrock async invoke format. + + Args: + model_input: The transformed TwelveLabs Marengo embedding request + model_id: The model identifier (without async_invoke prefix) + output_s3_uri: Optional S3 URI for output data config + + Returns: + TwelveLabsAsyncInvokeRequest: The wrapped async invoke request + """ + import urllib.parse + + # Clean the model ID + unquoted_model_id = urllib.parse.unquote(model_id) + if unquoted_model_id.startswith("async_invoke/"): + unquoted_model_id = unquoted_model_id.replace("async_invoke/", "") + + # Validate that the S3 URI is not empty + if not output_s3_uri or output_s3_uri.strip() == "": + raise ValueError("output_s3_uri cannot be empty for async invoke requests") + + return TwelveLabsAsyncInvokeRequest( + modelId=unquoted_model_id, + modelInput=model_input, + outputDataConfig=TwelveLabsOutputDataConfig( + s3OutputDataConfig=TwelveLabsS3OutputDataConfig(s3Uri=output_s3_uri) + ), + ) + def _transform_response( self, response_list: List[dict], model: str ) -> EmbeddingResponse: @@ -138,3 +249,53 @@ class TwelveLabsMarengoEmbeddingConfig: usage = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens) return EmbeddingResponse(data=embeddings, model=model, usage=usage) + + def _transform_async_invoke_response( + self, response: dict, model: str + ) -> EmbeddingResponse: + """ + Transform async invoke response (invocation ARN) to OpenAI format. + + AWS async invoke returns: + { + "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123" + } + + We transform this to a job-like embedding response: + { + "object": "list", + "data": [ + { + "object": "embedding_job_id:1234567890", + "embedding": [], + "index": 0 + } + ], + "model": "model", + "usage": {} + } + """ + invocation_arn = response.get("invocationArn", "") + + # Create a placeholder embedding object for the job + embedding = Embedding( + embedding=[], # Empty embedding for async jobs + index=0, + object="embedding", + ) + + # Create usage object (empty for async jobs) + usage = Usage(prompt_tokens=0, total_tokens=0) + + # Create hidden params with job ID + from litellm.types.llms.base import HiddenParams + + hidden_params = HiddenParams() + setattr(hidden_params, "_invocation_arn", invocation_arn) + + return EmbeddingResponse( + data=[embedding], + model=model, + usage=usage, + hidden_params=hidden_params, + ) diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index 6dbe52d575e..d194d9556b6 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -31,7 +31,7 @@ def validate_environment( "Request-Source": "unspecified:litellm", "accept": "application/json", "content-type": "application/json", - "Authorization": "bearer $CO_API_KEY" + "Authorization": "Bearer $CO_API_KEY" } """ headers.update( @@ -42,7 +42,7 @@ def validate_environment( } ) if api_key: - headers["Authorization"] = f"bearer {api_key}" + headers["Authorization"] = f"Bearer {api_key}" return headers diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 5371b9a4b61..f9c979712da 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -1,8 +1,8 @@ from typing import Any, Dict, List, Optional, Union import httpx -import litellm +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -52,20 +52,20 @@ class CohereRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, - ) -> OptionalRerankParams: + ) -> Dict: """ Map Cohere rerank params No mapping required - returns all supported params """ - return OptionalRerankParams( + return dict(OptionalRerankParams( query=query, documents=documents, top_n=top_n, rank_fields=rank_fields, return_documents=return_documents, max_chunks_per_doc=max_chunks_per_doc, - ) + )) def validate_environment( self, @@ -86,7 +86,7 @@ class CohereRerankConfig(BaseRerankConfig): ) default_headers = { - "Authorization": f"bearer {api_key}", + "Authorization": f"Bearer {api_key}", "accept": "application/json", "content-type": "application/json", } @@ -101,7 +101,7 @@ class CohereRerankConfig(BaseRerankConfig): def transform_rerank_request( self, model: str, - optional_rerank_params: OptionalRerankParams, + optional_rerank_params: Dict, headers: dict, ) -> dict: if "query" not in optional_rerank_params: diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index 74e760460d0..eb551a8a949 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -44,25 +44,25 @@ class CohereRerankV2Config(CohereRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, - ) -> OptionalRerankParams: + ) -> Dict: """ Map Cohere rerank params No mapping required - returns all supported params """ - return OptionalRerankParams( + return dict(OptionalRerankParams( query=query, documents=documents, top_n=top_n, rank_fields=rank_fields, return_documents=return_documents, max_tokens_per_doc=max_tokens_per_doc, - ) + )) def transform_rerank_request( self, model: str, - optional_rerank_params: OptionalRerankParams, + optional_rerank_params: Dict, headers: dict, ) -> dict: if "query" not in optional_rerank_params: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 173bb5a2ccb..ae449bbb15b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -65,7 +65,7 @@ from litellm.types.llms.openai import ( ResponseInputParam, ResponsesAPIResponse, ) -from litellm.types.rerank import OptionalRerankParams, RerankResponse +from litellm.types.rerank import RerankResponse from litellm.types.responses.main import DeleteResponseResult from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( @@ -893,7 +893,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, logging_obj: LiteLLMLoggingObj, provider_config: BaseRerankConfig, - optional_rerank_params: OptionalRerankParams, + optional_rerank_params: Dict, timeout: Optional[Union[float, httpx.Timeout]], model_response: RerankResponse, _is_async: bool = False, diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 6a3244a3c88..69c7dabebd8 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -2,11 +2,11 @@ Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format. """ -from litellm._uuid import uuid from typing import Any, Dict, List, Optional, Union import httpx +from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import ( BaseLLMException, @@ -98,7 +98,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, - ) -> OptionalRerankParams: + ) -> Dict: # Start with the basic parameters optional_rerank_params = {} if query: @@ -124,7 +124,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): def transform_rerank_request( self, model: str, - optional_rerank_params: OptionalRerankParams, + optional_rerank_params: Dict, headers: dict, ) -> dict: # Convert OptionalRerankParams to dict as expected by parent class diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index e1dd6f146f3..62329358e47 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -3,10 +3,10 @@ This file contains the transformation logic for the Gemini realtime API. """ import json -from litellm._uuid import uuid from typing import Any, Dict, List, Optional, Union, cast from litellm import verbose_logger +from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -186,9 +186,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) vertex_gemini_config = VertexGeminiConfig() - vertex_gemini_config._map_function(value) optional_params["generationConfig"]["tools"] = ( - vertex_gemini_config._map_function(value) + vertex_gemini_config._map_function( + value=value, optional_params=optional_params + ) ) elif key == "input_audio_transcription" and value is not None: optional_params["inputAudioTranscription"] = {} diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 4ed604e2c88..2faef2c4c73 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -2,27 +2,26 @@ Transformation logic for Hosted VLLM rerank """ -from litellm._uuid import uuid from typing import Any, Dict, List, Optional, Union +import httpx + +from litellm._uuid import uuid +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( + OptionalRerankParams, RerankBilledUnits, + RerankRequest, RerankResponse, RerankResponseDocument, RerankResponseMeta, RerankResponseResult, RerankTokens, - OptionalRerankParams, - RerankRequest, ) -import httpx - -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig -from litellm.secret_managers.main import get_secret_str - class HostedVLLMRerankError(BaseLLMException): def __init__( @@ -72,20 +71,20 @@ class HostedVLLMRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, - ) -> OptionalRerankParams: + ) -> Dict: """ Map parameters for Hosted VLLM rerank """ if max_chunks_per_doc is not None: raise ValueError("Hosted VLLM does not support max_chunks_per_doc") - return OptionalRerankParams( + return dict(OptionalRerankParams( query=query, documents=documents, top_n=top_n, rank_fields=rank_fields, return_documents=return_documents, - ) + )) def validate_environment( self, @@ -112,7 +111,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): def transform_rerank_request( self, model: str, - optional_rerank_params: OptionalRerankParams, + optional_rerank_params: Dict, headers: dict, ) -> dict: if "query" not in optional_rerank_params: diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index aa0f37bc6ba..1454328cc13 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -1,11 +1,11 @@ import os -from litellm._uuid import uuid from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import httpx from typing_extensions import TypedDict import litellm +from litellm._uuid import uuid from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -95,7 +95,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, - ) -> OptionalRerankParams: + ) -> Dict: optional_rerank_params = {} if non_default_params is not None: for k, v in non_default_params.items(): diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index 6259d445aec..55aac6033d5 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -49,7 +49,7 @@ class InfinityRerankConfig(CohereRerankConfig): ) default_headers = { - "Authorization": f"bearer {api_key}", + "Authorization": f"Bearer {api_key}", "accept": "application/json", "content-type": "application/json", } diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index d8569b01b83..3ba24680fd4 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -6,11 +6,11 @@ Why separate file? Make it easy to see how transformation works Docs - https://jina.ai/reranker """ -from litellm._uuid import uuid from typing import Any, Dict, List, Optional, Tuple, Union from httpx import URL, Response +from litellm._uuid import uuid from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.types.rerank import ( @@ -45,15 +45,15 @@ class JinaAIRerankConfig(BaseRerankConfig): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, - ) -> OptionalRerankParams: + ) -> Dict: optional_params = {} supported_params = self.get_supported_cohere_rerank_params(model) for k, v in non_default_params.items(): if k in supported_params: optional_params[k] = v - return OptionalRerankParams( + return dict(OptionalRerankParams( **optional_params, - ) + )) def get_complete_url(self, api_base: Optional[str], model: str) -> str: base_path = "/v1/rerank" @@ -67,7 +67,7 @@ class JinaAIRerankConfig(BaseRerankConfig): return cleaned_base def transform_rerank_request( - self, model: str, optional_rerank_params: OptionalRerankParams, headers: Dict + self, model: str, optional_rerank_params: Dict, headers: Dict ) -> Dict: return {"model": model, **optional_rerank_params} @@ -98,9 +98,26 @@ class JinaAIRerankConfig(BaseRerankConfig): if _results is None: raise ValueError(f"No results found in the response={_json_response}") + # Transform Jina AI's response format to match LiteLLM's expected format + # Jina AI returns: {"index": 0, "relevance_score": 0.72, "document": "hello"} + # LiteLLM expects: {"index": 0, "relevance_score": 0.72, "document": {"text": "hello"}} + transformed_results = [] + for result in _results: + transformed_result = { + "index": result["index"], + "relevance_score": result["relevance_score"], + } + # Convert document from string to dict format if it exists + if "document" in result and isinstance(result["document"], str): + transformed_result["document"] = {"text": result["document"]} + elif "document" in result: + # If it's already a dict, keep it as is + transformed_result["document"] = result["document"] + transformed_results.append(transformed_result) + return RerankResponse( id=_json_response.get("id") or str(uuid.uuid4()), - results=_results, # type: ignore + results=transformed_results, # type: ignore meta=rerank_meta, ) # Return response diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py new file mode 100644 index 00000000000..cb9fd4bebaa --- /dev/null +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -0,0 +1,325 @@ +from typing import Any, Dict, List, Literal, Optional, Union + +import httpx +from typing_extensions import Required, TypedDict + +import litellm +from litellm._uuid import uuid +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.rerank import ( + RerankBilledUnits, + RerankResponse, + RerankResponseMeta, + RerankResponseResult, +) + + +class NvidiaNimQueryObject(TypedDict): + text: Required[str] + + +class NvidiaNimPassageObject(TypedDict): + text: Required[str] + + +class NvidiaNimRerankRequest(TypedDict, total=False): + model: Required[str] + query: Required[NvidiaNimQueryObject] + passages: Required[List[NvidiaNimPassageObject]] + truncate: Literal["NONE", "END"] + top_k: int + + +class NvidiaNimRankingResult(TypedDict): + index: Required[int] + logit: Required[float] + + +class NvidiaNimRerankResponse(TypedDict): + rankings: Required[List[NvidiaNimRankingResult]] + + +class NvidiaNimRerankConfig(BaseRerankConfig): + """ + Reference: https://docs.api.nvidia.com/nim/reference/nvidia-llama-3_2-nv-rerankqa-1b-v2-infer + + Nvidia NIM rerank API uses a different format: + - query is an object with 'text' field + - documents are called 'passages' and have 'text' field + """ + DEFAULT_NIM_RERANK_API_BASE = "https://ai.api.nvidia.com" + + def __init__(self) -> None: + pass + + def get_complete_url(self, api_base: Optional[str], model: str) -> str: + """ + Construct the Nvidia NIM rerank URL. + + Format: {api_base}/v1/retrieval/{model}/reranking + + If the user provides a full URL (e.g., {api_base}/v1/retrieval/{model}/reranking), + it will be used as-is. + """ + if not api_base: + api_base = self.DEFAULT_NIM_RERANK_API_BASE + + api_base = api_base.rstrip("/") + + # Check if user already provided the full URL with /retrieval/ path + if "/retrieval/" in api_base: + return api_base + + # Ensure we don't have duplicate /v1 + if api_base.endswith("/v1"): + api_base = api_base[:-3] + + return f"{api_base}/v1/retrieval/{model}/reranking" + + def get_supported_cohere_rerank_params(self, model: str) -> list: + """ + Nvidia NIM supports these rerank parameters. + """ + return [ + "query", + "documents", + "top_n", + ] + + def map_cohere_rerank_params( + self, + non_default_params: Optional[dict], + model: str, + drop_params: bool, + query: str, + documents: List[Union[str, Dict[str, Any]]], + custom_llm_provider: Optional[str] = None, + top_n: Optional[int] = None, + rank_fields: Optional[List[str]] = None, + return_documents: Optional[bool] = True, + max_chunks_per_doc: Optional[int] = None, + max_tokens_per_doc: Optional[int] = None, + ) -> Dict: + """ + Map Cohere/OpenAI rerank params to Nvidia NIM format. + + Parameter mapping: + - top_n (Cohere) -> top_k (Nvidia) + + Nvidia NIM specific params (passed through as-is from non_default_params): + - truncate: How to truncate input if too long (NONE, END) + """ + optional_nvidia_nim_rerank_params: Dict[str, Any] = { + "query": query, + "documents": documents, + } + + # Map Cohere's top_n to Nvidia's top_k + if top_n is not None: + optional_nvidia_nim_rerank_params["top_k"] = top_n + + # Pass through Nvidia-specific params from non_default_params + if non_default_params: + optional_nvidia_nim_rerank_params.update(non_default_params) + return dict(optional_nvidia_nim_rerank_params) + + def validate_environment( + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + ) -> dict: + """ + Validate that the Nvidia NIM API key is present. + """ + if api_key is None: + api_key = ( + get_secret_str("NVIDIA_NIM_API_KEY") + or litellm.api_key + ) + + if api_key is None: + raise ValueError( + "Nvidia NIM API key is required. Please set 'NVIDIA_NIM_API_KEY' in your environment" + ) + + default_headers = { + "Authorization": f"Bearer {api_key}", + "accept": "application/json", + "content-type": "application/json", + } + + # If 'Authorization' is provided in headers, it overrides the default + if "Authorization" in headers: + default_headers["Authorization"] = headers["Authorization"] + + # Merge other headers, overriding any default ones except Authorization + return {**default_headers, **headers} + + def transform_rerank_request( + self, + model: str, + optional_rerank_params: Dict, + headers: dict, + ) -> dict: + """ + Transform request to Nvidia NIM format. + + Nvidia NIM expects: + - query as {text: "..."} + - documents as passages: [{text: "..."}, ...] + - Optional: truncate (NONE or END), top_k + + Note: optional_rerank_params may contain provider-specific params like 'top_k' and 'truncate' + that aren't in the OptionalRerankParams TypedDict but are passed through at runtime. + The mapping from Cohere's 'top_n' to Nvidia's 'top_k' already happened in map_cohere_rerank_params. + """ + if "query" not in optional_rerank_params: + raise ValueError("query is required for Nvidia NIM rerank") + if "documents" not in optional_rerank_params: + raise ValueError("documents is required for Nvidia NIM rerank") + + query = optional_rerank_params["query"] + documents = optional_rerank_params["documents"] + + # Transform query to object format + query_obj: NvidiaNimQueryObject = {"text": query} + + # Transform documents to passages format + passages: List[NvidiaNimPassageObject] = [] + for doc in documents: + if isinstance(doc, str): + passages.append({"text": doc}) + elif isinstance(doc, dict): + # If document is already a dict, check if it has 'text' field + if "text" in doc: + passages.append({"text": doc["text"]}) + else: + # Otherwise, stringify the dict + import json + passages.append({"text": json.dumps(doc)}) + else: + passages.append({"text": str(doc)}) + + # Note: URL path uses underscores (llama-3_2) but JSON body uses periods (llama-3.2) + # Convert underscores back to periods for the model field in request body + model_for_body = model.replace("_", ".") + + # Build request using TypedDict + request_data: NvidiaNimRerankRequest = { + "model": model_for_body, + "query": query_obj, + "passages": passages, + } + + # Add optional top_k parameter if provided (already mapped from top_n in map_cohere_rerank_params) + if "top_k" in optional_rerank_params and optional_rerank_params.get("top_k") is not None: # type: ignore + request_data["top_k"] = optional_rerank_params.get("top_k") # type: ignore + + # Add Nvidia-specific truncate parameter if provided + # This is passed through from non_default_params, not in base OptionalRerankParams + if "truncate" in optional_rerank_params and optional_rerank_params.get("truncate") is not None: # type: ignore + truncate_value = optional_rerank_params.get("truncate") # type: ignore + if truncate_value in ["NONE", "END"]: + request_data["truncate"] = truncate_value # type: ignore + + return dict(request_data) + + def transform_rerank_response( + self, + model: str, + raw_response: httpx.Response, + model_response: RerankResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str] = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> RerankResponse: + """ + Transform Nvidia NIM rerank response to LiteLLM format. + + Nvidia NIM returns (NvidiaNimRerankResponse): + { + "rankings": [ + { + "index": 0, + "logit": 0.123 + } + ] + } + + LiteLLM expects (RerankResponse): + { + "results": [ + { + "index": 0, + "relevance_score": 0.123, + "document": {"text": "..."} # optional + } + ] + } + """ + try: + raw_response_json = raw_response.json() + except Exception: + raise BaseLLMException( + status_code=raw_response.status_code, + message=raw_response.text, + headers=raw_response.headers, + ) + + # Parse as NvidiaNimRerankResponse + nvidia_response: NvidiaNimRerankResponse = raw_response_json + + # Transform Nvidia NIM response to LiteLLM format + results: List[RerankResponseResult] = [] + rankings = nvidia_response.get("rankings", []) + + # Get original documents from request if we need to include them + original_passages: List[NvidiaNimPassageObject] = request_data.get("passages", []) + + for ranking in rankings: + result_item: RerankResponseResult = { + "index": ranking["index"], + "relevance_score": ranking["logit"], + } + + # Include document if it was in the original request + index: int = ranking["index"] + if index < len(original_passages): + result_item["document"] = {"text": original_passages[index]["text"]} # type: ignore + + results.append(result_item) + + # Construct metadata with billed_units + # Nvidia NIM uses "usage" field with "total_tokens" + usage = raw_response_json.get("usage", {}) + total_tokens = usage.get("total_tokens", 0) + + billed_units: RerankBilledUnits = { + "total_tokens": total_tokens if total_tokens > 0 else len(results) + } + + meta: RerankResponseMeta = { + "billed_units": billed_units + } + + return RerankResponse( + id=raw_response_json.get("id") or str(uuid.uuid4()), + results=results, + meta=meta, + ) + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> BaseLLMException: + return BaseLLMException( + status_code=status_code, + message=error_message, + headers=headers, + ) + diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 6755cab22e0..18740052480 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -207,9 +207,9 @@ class OCIChatConfig(BaseConfig): alias = open_ai_to_oci_param_map.get(key) if alias is False: - if drop_params: + # Workaround for mypy issue + if drop_params or litellm.drop_params: continue - raise Exception(f"param `{key}` is not supported on OCI") if alias is None: @@ -450,12 +450,26 @@ class OCIChatConfig(BaseConfig): "Cohere models are not yet supported in the litellm OCI chat completion endpoint. Use the Cohere API directly." ) else: - data = OCICompletionPayload( - compartmentId=oci_compartment_id, - servingMode=OCIServingMode( + oci_serving_mode = optional_params.get("oci_serving_mode", "ON_DEMAND") + if oci_serving_mode not in ["ON_DEMAND", "DEDICATED"]: + raise Exception( + "kwarg `oci_serving_mode` must be either 'ON_DEMAND' or 'DEDICATED'" + ) + + if oci_serving_mode == "DEDICATED": + servingMode = OCIServingMode( + servingType="DEDICATED", + endpointId=model, + ) + else: + servingMode = OCIServingMode( servingType="ON_DEMAND", modelId=model, - ), + ) + + data = OCICompletionPayload( + compartmentId=oci_compartment_id, + servingMode=servingMode, chatRequest=OCIChatRequestPayload( apiFormat=vendor.value, messages=adapt_messages_to_generic_oci_standard(messages), diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 3902304a3b4..fa357c1bd22 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -8,18 +8,24 @@ from .gpt_transformation import OpenAIGPTConfig class OpenAIGPT5Config(OpenAIGPTConfig): - """Configuration for gpt-5 models. + """Configuration for gpt-5 models including GPT-5-Codex variants. Handles OpenAI API quirks for the gpt-5 series like: - Mapping ``max_tokens`` -> ``max_completion_tokens``. - Dropping unsupported ``temperature`` values when requested. + - Support for GPT-5-Codex models optimized for code generation. """ @classmethod def is_model_gpt_5_model(cls, model: str) -> bool: return "gpt-5" in model + @classmethod + def is_model_gpt_5_codex_model(cls, model: str) -> bool: + """Check if the model is specifically a GPT-5 Codex variant.""" + return "gpt-5-codex" in model + def get_supported_openai_params(self, model: str) -> list: from litellm.utils import supports_tool_choice @@ -38,7 +44,9 @@ class OpenAIGPT5Config(OpenAIGPTConfig): ] return [ - param for param in base_gpt_series_params if param not in non_supported_params + param + for param in base_gpt_series_params + if param not in non_supported_params ] def map_openai_params( @@ -67,7 +75,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): else: raise litellm.utils.UnsupportedParamsError( message=( - "gpt-5 models don't support temperature={}. Only temperature=1 is supported. To drop unsupported params set `litellm.drop_params = True`" + "gpt-5 models (including gpt-5-codex) don't support temperature={}. Only temperature=1 is supported. To drop unsupported params set `litellm.drop_params = True`" ).format(temperature_value), status_code=400, ) 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 5871c353381..d3f67967213 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 @@ -68,6 +68,7 @@ from litellm.types.llms.vertex_ai import ( ToolConfig, Tools, UsageMetadata, + VertexToolName, ) from litellm.types.utils import ( ChatCompletionAudioResponse, @@ -276,42 +277,106 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): """ return Tools(googleSearch={}) - def _map_function(self, value: List[dict]) -> List[Tools]: # noqa: PLR0915 + def _extract_google_maps_retrieval_config( + self, google_maps_config: dict + ) -> Tuple[dict, Optional[dict]]: + """ + Extract location configuration from googleMaps tool for Vertex AI toolConfig. + + Supports two interface styles: + 1. Nested (recommended): {"enableWidget": "...", "retrievalConfig": {"latitude": ..., "longitude": ...}} + 2. Flat (backward compat): {"enableWidget": "...", "latitude": ..., "longitude": ...} + + Args: + google_maps_config: The googleMaps tool configuration from LiteLLM + + Returns: + Tuple of (cleaned_google_maps_config, retrieval_config): + - cleaned_google_maps_config: googleMaps config without location fields + - retrieval_config: Location config for toolConfig.retrievalConfig or None + """ + retrieval_config = None + latitude = google_maps_config.get("latitude") + longitude = google_maps_config.get("longitude") + language_code = google_maps_config.get("languageCode") + + if latitude is not None and longitude is not None: + retrieval_config = { + "latLng": { + "latitude": latitude, + "longitude": longitude, + } + } + if language_code is not None: + retrieval_config["languageCode"] = language_code + + # Remove location fields from tool definition + cleaned_config = { + k: v + for k, v in google_maps_config.items() + if k not in ["latitude", "longitude", "languageCode"] + } + + return cleaned_config, retrieval_config + + def get_tool_value( + self, + tool: dict, + tool_name: str + ) -> Optional[dict]: + """ + Helper function to get tool value handling both camelCase and underscore_case variants + + Args: + tool (dict): The tool dictionary + tool_name (str): The base tool name (e.g. "codeExecution") + + Returns: + Optional[dict]: The tool value if found, None otherwise + """ + # Convert camelCase to underscore_case + underscore_name = "".join( + ["_" + c.lower() if c.isupper() else c for c in tool_name] + ).lstrip("_") + # Try both camelCase and underscore_case variants + + if tool.get(tool_name) is not None: + return tool.get(tool_name) + elif tool.get(underscore_name) is not None: + return tool.get(underscore_name) + else: + return None + + def _map_function( # noqa: PLR0915 + self, value: List[dict], optional_params: dict + ) -> List[Tools]: + """ + Map OpenAI-style tools/functions to Vertex AI format. + + Args: + value: List of tool definitions + optional_params: Request-scoped parameters to store retrieval config + + Returns: + List of mapped tools in Vertex AI format + + Side effects: + May add 'toolConfig' with 'retrievalConfig' to optional_params if + googleMaps tools contain location data + """ gtool_func_declarations = [] googleSearch: Optional[dict] = None googleSearchRetrieval: Optional[dict] = None enterpriseWebSearch: Optional[dict] = None urlContext: Optional[dict] = None code_execution: Optional[dict] = None + googleMaps: Optional[dict] = None + google_maps_retrieval_config: Optional[dict] = None # remove 'additionalProperties' from tools value = _remove_additional_properties(value) # remove 'strict' from tools value = _remove_strict_from_schema(value) - def get_tool_value(tool: dict, tool_name: str) -> Optional[dict]: - """ - Helper function to get tool value handling both camelCase and underscore_case variants - - Args: - tool (dict): The tool dictionary - tool_name (str): The base tool name (e.g. "codeExecution") - - Returns: - Optional[dict]: The tool value if found, None otherwise - """ - # Convert camelCase to underscore_case - underscore_name = "".join( - ["_" + c.lower() if c.isupper() else c for c in tool_name] - ).lstrip("_") - # Try both camelCase and underscore_case variants - - if tool.get(tool_name) is not None: - return tool.get(tool_name) - elif tool.get(underscore_name) is not None: - return tool.get(underscore_name) - else: - return None - for tool in value: openai_function_object: Optional[ ChatCompletionToolParamFunctionChunk @@ -341,17 +406,27 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_name = list(tool.keys())[0] if len(tool.keys()) == 1 else None if tool_name and ( - tool_name == "codeExecution" or tool_name == "code_execution" + tool_name == "codeExecution" or tool_name == VertexToolName.CODE_EXECUTION.value ): # code_execution maintained for backwards compatibility - code_execution = get_tool_value(tool, "codeExecution") - elif tool_name and tool_name == "googleSearch": - googleSearch = get_tool_value(tool, "googleSearch") - elif tool_name and tool_name == "googleSearchRetrieval": - googleSearchRetrieval = get_tool_value(tool, "googleSearchRetrieval") - elif tool_name and tool_name == "enterpriseWebSearch": - enterpriseWebSearch = get_tool_value(tool, "enterpriseWebSearch") - elif tool_name and tool_name == "urlContext": - urlContext = get_tool_value(tool, "urlContext") + code_execution = self.get_tool_value(tool, "codeExecution") + elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH.value: + googleSearch = self.get_tool_value(tool, VertexToolName.GOOGLE_SEARCH.value) + elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value: + googleSearchRetrieval = self.get_tool_value(tool, VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value) + elif tool_name and tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value: + enterpriseWebSearch = self.get_tool_value(tool, VertexToolName.ENTERPRISE_WEB_SEARCH.value) + elif tool_name and (tool_name == VertexToolName.URL_CONTEXT.value or tool_name == "urlContext"): + urlContext = self.get_tool_value(tool, tool_name) + elif tool_name and ( + tool_name == VertexToolName.GOOGLE_MAPS.value or tool_name == "google_maps" + ): + google_maps_value = self.get_tool_value(tool, VertexToolName.GOOGLE_MAPS.value) + + # Extract and transform location configuration for toolConfig + if google_maps_value is not None: + googleMaps, google_maps_retrieval_config = self._extract_google_maps_retrieval_config( + google_maps_config=google_maps_value + ) elif openai_function_object is not None: gtool_func_declaration = FunctionDeclaration( name=openai_function_object["name"], @@ -373,19 +448,29 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "Invalid tool={}. Use `litellm.set_verbose` or `litellm --detailed_debug` to see raw request." ) - _tools = Tools( - function_declarations=gtool_func_declarations, - ) + # Only include function_declarations if there are actual functions + _tools = Tools() + if gtool_func_declarations: + _tools["function_declarations"] = gtool_func_declarations if googleSearch is not None: - _tools["googleSearch"] = googleSearch + _tools[VertexToolName.GOOGLE_SEARCH.value] = googleSearch if googleSearchRetrieval is not None: - _tools["googleSearchRetrieval"] = googleSearchRetrieval + _tools[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval if enterpriseWebSearch is not None: - _tools["enterpriseWebSearch"] = enterpriseWebSearch + _tools[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch if code_execution is not None: - _tools["code_execution"] = code_execution + _tools[VertexToolName.CODE_EXECUTION.value] = code_execution if urlContext is not None: - _tools["url_context"] = urlContext + _tools[VertexToolName.URL_CONTEXT.value] = urlContext + if googleMaps is not None: + _tools[VertexToolName.GOOGLE_MAPS.value] = googleMaps + + # Add retrieval config to toolConfig if googleMaps has location data + if google_maps_retrieval_config is not None: + if "toolConfig" not in optional_params: + optional_params["toolConfig"] = {} + optional_params["toolConfig"]["retrievalConfig"] = google_maps_retrieval_config + return [_tools] def _map_response_schema(self, value: dict) -> dict: @@ -606,8 +691,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): and isinstance(value, list) and value ): + # Pass optional_params so _map_function can add toolConfig if needed + mapped_tools = self._map_function( + value=value, optional_params=optional_params + ) optional_params = self._add_tools_to_optional_params( - optional_params, self._map_function(value=value) + optional_params, mapped_tools ) elif param == "tool_choice" and ( isinstance(value, str) or isinstance(value, dict) diff --git a/litellm/main.py b/litellm/main.py index 4387e731980..cfb0bef0797 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -24,6 +24,7 @@ from functools import partial from typing import ( TYPE_CHECKING, Any, + AsyncIterator, Callable, Coroutine, Dict, @@ -5170,6 +5171,21 @@ async def aadapter_completion( except Exception as e: raise e +async def aadapter_generate_content( + **kwargs, +) -> Union[Dict[str, Any], AsyncIterator[bytes]]: + from litellm.google_genai.adapters.handler import ( + GenerateContentToCompletionHandler, + ) + + coro = cast( + Coroutine[Any, Any, Union[Dict[str, Any], AsyncIterator[bytes]]], + GenerateContentToCompletionHandler.generate_content_handler( + **kwargs, _is_async=True + ), + ) + return await coro + def adapter_completion( *, adapter_id: str, **kwargs diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a35cac30489..72b9c551d73 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -2004,9 +2004,9 @@ "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, "supported_endpoints": [ @@ -3308,6 +3308,64 @@ "supports_tool_choice": true, "supports_web_search": true }, + "azure_ai/grok-4": { + "input_cost_per_token": 5.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "source": "https://azure.microsoft.com/en-us/blog/grok-4-is-now-available-in-azure-ai-foundry-unlock-frontier-intelligence-and-business-ready-capabilities/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_web_search": true + }, + "azure_ai/grok-4-fast-non-reasoning": { + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-03, + "source": "https://azure.microsoft.com/en-us/blog/grok-4-is-now-available-in-azure-ai-foundry-unlock-frontier-intelligence-and-business-ready-capabilities/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_web_search": true + }, + "azure_ai/grok-4-fast-reasoning": { + "input_cost_per_token": 5.8e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.9e-03, + "source": "https://azure.microsoft.com/en-us/blog/grok-4-is-now-available-in-azure-ai-foundry-unlock-frontier-intelligence-and-business-ready-capabilities/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_web_search": true + }, + "azure_ai/grok-code-fast-1": { + "input_cost_per_token": 3.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.75e-05, + "source": "https://azure.microsoft.com/en-us/blog/grok-4-is-now-available-in-azure-ai-foundry-unlock-frontier-intelligence-and-business-ready-capabilities/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_web_search": true + }, "azure_ai/jais-30b-chat": { "input_cost_per_token": 0.0032, "litellm_provider": "azure_ai", @@ -4743,6 +4801,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -4769,6 +4831,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -6627,629 +6693,679 @@ ] }, "deepinfra/Gryphe/MythoMax-L2-13b": { - "input_cost_per_token": 7.2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "max_tokens": 4096, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 9e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 7.2e-08, "supports_tool_choice": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": { - "input_cost_per_token": 7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 1e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 8e-07, "supports_tool_choice": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.8e-07, "supports_tool_choice": false }, "deepinfra/Qwen/QwQ-32B": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 1.5e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen2.5-72B-Instruct": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 1.2e-07, "output_cost_per_token": 3.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen2.5-7B-Instruct": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Qwen/Qwen2.5-VL-32B-Instruct": { - "input_cost_per_token": 2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", + "input_cost_per_token": 2e-07, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-14B": { - "input_cost_per_token": 6e-08, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 6e-08, "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 5.4e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 6e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 9e-08, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.9e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 6e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-30B-A3B": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 8e-08, "output_cost_per_token": 2.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-32B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 3e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": { - "input_cost_per_token": 4e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": { - "cache_read_input_token_cost": 2.4e-07, - "input_cost_per_token": 3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 2.9e-07, "output_cost_per_token": 1.2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": { - "input_cost_per_token": 2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Sao10K/L3.1-70B-Euryale-v2.2": { - "input_cost_per_token": 6.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 6.5e-07, "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Sao10K/L3.3-70B-Euryale-v2.3": { - "input_cost_per_token": 6.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 6.5e-07, "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/allenai/olmOCR-7B-0725-FP8": { - "input_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 16384, "max_input_tokens": 16384, "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", + "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1.5e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/anthropic/claude-3-7-sonnet-latest": { - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 3.3e-06, "output_cost_per_token": 1.65e-05, + "cache_read_input_token_cost": 3.3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/anthropic/claude-4-opus": { - "input_cost_per_token": 1.65e-05, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 1.65e-05, "output_cost_per_token": 8.25e-05, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/anthropic/claude-4-sonnet": { - "input_cost_per_token": 3.3e-06, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 3.3e-06, "output_cost_per_token": 1.65e-05, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1": { - "input_cost_per_token": 7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 7e-07, "output_cost_per_token": 2.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528": { - "cache_read_input_token_cost": 4e-07, - "input_cost_per_token": 5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 5e-07, "output_cost_per_token": 2.15e-06, + "cache_read_input_token_cost": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": { - "input_cost_per_token": 1e-06, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 1e-06, "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 4e-07, "supports_tool_choice": false }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { - "input_cost_per_token": 7.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 2.7e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.5e-07, "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": { - "input_cost_per_token": 1e-06, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 1e-06, "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3": { - "input_cost_per_token": 3.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 3.8e-07, "output_cost_per_token": 8.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3-0324": { - "cache_read_input_token_cost": 2.24e-07, - "input_cost_per_token": 2.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 2.5e-07, "output_cost_per_token": 8.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1": { - "cache_read_input_token_cost": 2.16e-07, - "input_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1e-06, - "supports_reasoning": true, + "cache_read_input_token_cost": 2.16e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true, + "supports_reasoning": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3.1-Terminus": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 1e-06, + "cache_read_input_token_cost": 2.16e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemini-2.0-flash-001": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, - "mode": "chat", + "input_cost_per_token": 1e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemini-2.5-flash": { - "input_cost_per_token": 2.1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.5e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.75e-06, "supports_tool_choice": true }, "deepinfra/google/gemini-2.5-pro": { - "input_cost_per_token": 8.75e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 1e-05, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 7e-06, "supports_tool_choice": true }, "deepinfra/google/gemma-3-12b-it": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemma-3-27b-it": { - "input_cost_per_token": 9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 1.6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.7e-07, "supports_tool_choice": true }, "deepinfra/google/gemma-3-4b-it": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 8e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": { - "input_cost_per_token": 4.9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4.9e-08, "output_cost_per_token": 4.9e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Llama-3.2-3B-Instruct": { - "input_cost_per_token": 1.2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 2e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.4e-08, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct": { - "input_cost_per_token": 2.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 2.3e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": { - "input_cost_per_token": 3.8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 3.9e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.2e-07, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1048576, "max_input_tokens": 1048576, "max_output_tokens": 1048576, - "max_tokens": 1048576, - "mode": "chat", + "input_cost_per_token": 1.5e-07, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 327680, "max_input_tokens": 327680, "max_output_tokens": 327680, - "max_tokens": 327680, - "mode": "chat", + "input_cost_per_token": 8e-08, "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-Guard-3-8B": { - "input_cost_per_token": 5.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5.5e-08, "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Llama-Guard-4-12B": { - "input_cost_per_token": 1.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 1.8e-07, "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Meta-Llama-3-8B-Instruct": { - "input_cost_per_token": 3e-08, - "litellm_provider": "deepinfra", + "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", + "input_cost_per_token": 3e-08, "output_cost_per_token": 6e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": { - "input_cost_per_token": 2.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 1e-07, "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": { - "input_cost_per_token": 3e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 3e-08, "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { - "input_cost_per_token": 1.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 3e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2e-08, "supports_tool_choice": true }, "deepinfra/microsoft/WizardLM-2-8x22B": { - "input_cost_per_token": 4.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 65536, "max_input_tokens": 65536, "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", + "input_cost_per_token": 4.8e-07, "output_cost_per_token": 4.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/microsoft/phi-4": { - "input_cost_per_token": 7e-08, - "litellm_provider": "deepinfra", + "max_tokens": 16384, "max_input_tokens": 16384, "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", + "input_cost_per_token": 7e-08, "output_cost_per_token": 1.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": { - "input_cost_per_token": 2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 2e-08, "output_cost_per_token": 4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 8e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "max_tokens": 128000, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 2e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1e-07, "supports_tool_choice": true }, "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.4e-07, "supports_tool_choice": true }, "deepinfra/moonshotai/Kimi-K2-Instruct": { - "input_cost_per_token": 5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-07, "output_cost_per_token": 2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/moonshotai/Kimi-K2-Instruct-0905": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2e-06, + "cache_read_input_token_cost": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/nvidia/Llama-3.3-Nemotron-Super-49B-v1.5": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 3e-07, "supports_tool_choice": true }, "deepinfra/openai/gpt-oss-120b": { - "input_cost_per_token": 9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 4.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/openai/gpt-oss-20b": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.6e-07, "supports_tool_choice": true }, "deepinfra/zai-org/GLM-4.5": { - "input_cost_per_token": 5.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_tool_choice": true - }, - "deepinfra/zai-org/GLM-4.5-Air": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.1e-06, "supports_tool_choice": true }, "deepseek/deepseek-chat": { @@ -7722,6 +7838,36 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "eu.anthropic.claude-sonnet-4-5-20250929-v1:0": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, "eu.meta.llama3-2-1b-instruct-v1:0": { "input_cost_per_token": 1.3e-07, "litellm_provider": "bedrock", @@ -12727,6 +12873,34 @@ "supports_tool_choice": true, "supports_vision": true }, + "gpt-5-codex": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "gpt-5-2025-08-07": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, @@ -12802,9 +12976,9 @@ "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, "supported_endpoints": [ @@ -13557,6 +13731,19 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "groq/moonshotai/kimi-k2-instruct-0905": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 0.5e-06, + "litellm_provider": "groq", + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "max_tokens": 278528, + "mode": "chat", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "groq/openai/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, "litellm_provider": "groq", @@ -14032,6 +14219,36 @@ "mode": "rerank", "output_cost_per_token": 1.8e-08 }, + "jp.anthropic.claude-sonnet-4-5-20250929-v1:0": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, "lambda_ai/deepseek-llama3.3-70b": { "input_cost_per_token": 2e-07, "litellm_provider": "lambda_ai", @@ -18306,6 +18523,20 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { + "input_cost_per_query": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "nvidia_nim", + "mode": "rerank", + "output_cost_per_token": 0.0 + }, + "nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2": { + "input_cost_per_query": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "nvidia_nim", + "mode": "rerank", + "output_cost_per_token": 0.0 + }, "sagemaker/meta-textgeneration-llama-2-13b": { "input_cost_per_token": 0.0, "litellm_provider": "sagemaker", @@ -19621,6 +19852,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -20987,6 +21222,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -21009,6 +21248,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 6d6ebec6d05..31032b27b22 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -294,6 +294,9 @@ class MCPRequestHandler: ) -> List[str]: """ Get list of allowed MCP servers for the given user/key based on permissions + + Returns: + List[str]: List of allowed MCP servers by server id """ from typing import List @@ -330,11 +333,30 @@ class MCPRequestHandler: verbose_logger.warning(f"Failed to get allowed MCP servers: {str(e)}") return [] + @staticmethod + def is_tool_allowed( + allowed_mcp_servers: List[str], + server_name: str, + ) -> bool: + """ + Check if the tool is allowed for the given user/key based on permissions + """ + if len(allowed_mcp_servers) == 0: + return True + elif server_name in allowed_mcp_servers: + return True + return False + @staticmethod async def _get_allowed_mcp_servers_for_key( user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[str]: - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.auth.auth_checks import get_object_permission + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if user_api_key_auth is None: return [] @@ -347,12 +369,12 @@ class MCPRequestHandler: return [] try: - key_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={ - "object_permission_id": user_api_key_auth.object_permission_id - }, - ) + key_object_permission = await get_object_permission( + object_permission_id=user_api_key_auth.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) if key_object_permission is None: return [] @@ -386,7 +408,12 @@ class MCPRequestHandler: first we check if the team has a object_permission_id attached - if it does then we look up the object_permission for the team """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if user_api_key_auth is None: return [] @@ -399,10 +426,12 @@ class MCPRequestHandler: return [] try: - team_obj: Optional[LiteLLM_TeamTable] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": user_api_key_auth.team_id}, - ) + team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( + team_id=user_api_key_auth.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) if team_obj is None: verbose_logger.debug("team_obj is None") @@ -534,7 +563,12 @@ class MCPRequestHandler: async def _get_mcp_access_groups_for_key( user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[str]: - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.auth.auth_checks import get_object_permission + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if user_api_key_auth is None: return [] @@ -546,15 +580,21 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") return [] - key_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": user_api_key_auth.object_permission_id}, + try: + key_object_permission = await get_object_permission( + object_permission_id=user_api_key_auth.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - ) - if key_object_permission is None: - return [] + if key_object_permission is None: + return [] - return key_object_permission.mcp_access_groups or [] + return key_object_permission.mcp_access_groups or [] + except Exception as e: + verbose_logger.warning(f"Failed to get MCP access groups for key: {str(e)}") + return [] @staticmethod async def _get_mcp_access_groups_for_team( @@ -563,7 +603,12 @@ class MCPRequestHandler: """ Get MCP access groups for the team """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if user_api_key_auth is None: return [] @@ -575,20 +620,28 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") return [] - team_obj: Optional[LiteLLM_TeamTable] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": user_api_key_auth.team_id}, + try: + team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( + team_id=user_api_key_auth.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - ) - if team_obj is None: - verbose_logger.debug("team_obj is None") - return [] + if team_obj is None: + verbose_logger.debug("team_obj is None") + return [] - object_permissions = team_obj.object_permission - if object_permissions is None: - return [] + object_permissions = team_obj.object_permission + if object_permissions is None: + return [] - return object_permissions.mcp_access_groups or [] + return object_permissions.mcp_access_groups or [] + except Exception as e: + verbose_logger.warning( + f"Failed to get MCP access groups for team: {str(e)}" + ) + return [] @staticmethod def get_mcp_access_groups_from_headers(headers: Headers) -> Optional[List[str]]: diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 4c866561f70..9172568f304 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -919,6 +919,14 @@ class MCPServerManager: return server return None + def get_mcp_server_names_from_ids(self, server_ids: List[str]) -> List[str]: + server_names = [] + registry = self.get_registry() + for server in registry.values(): + if server.server_id in server_ids: + server_names.append(server.name) + return server_names + def get_mcp_server_by_name(self, server_name: str) -> Optional[MCPServer]: """ Get the MCP Server from the server name @@ -1189,6 +1197,11 @@ class MCPServerManager: if server.mcp_access_groups is not None else [] ), + allowed_tools=( + server.allowed_tools + if server.allowed_tools is not None + else [] + ), mcp_info=server.mcp_info, teams=cast( List[Dict[str, str | None]], diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index ecb960ebcf3..6a9c425a81b 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -28,6 +28,7 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.server import ( ListMCPToolsRestAPIResponseObject, call_mcp_tool, + filter_tools_by_allowed_tools, ) ######################################################## @@ -75,6 +76,12 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, add_prefix=False, ) + + # Filter tools based on allowed_tools configuration + # Only filter if allowed_tools is explicitly configured (not None and not empty) + if server.allowed_tools is not None and len(server.allowed_tools) > 0: + tools = filter_tools_by_allowed_tools(tools, server) + return _create_tool_response_objects(tools, server.mcp_info) ######################################################## diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 2487260b615..3e7c291810a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -361,22 +361,66 @@ if MCP_AVAILABLE: return allowed_mcp_servers + def _tool_name_matches(tool_name: str, filter_list: List[str]) -> bool: + """ + Check if a tool name matches any name in the filter list. + + Checks both the full tool name and unprefixed version (without server prefix). + This allows users to configure simple tool names regardless of prefixing. + + Args: + tool_name: The tool name to check (may be prefixed like "server-tool_name") + filter_list: List of tool names to match against + + Returns: + True if the tool name (prefixed or unprefixed) is in the filter list + """ + from litellm.proxy._experimental.mcp_server.utils import ( + get_server_name_prefix_tool_mcp, + ) + + # Check if the full name is in the list + if tool_name in filter_list: + return True + + # Check if the unprefixed name is in the list + unprefixed_name, _ = get_server_name_prefix_tool_mcp(tool_name) + return unprefixed_name in filter_list + def filter_tools_by_allowed_tools( tools: List[MCPTool], mcp_server: MCPServer, ) -> List[MCPTool]: """ - Filter tools by allowed tools + Filter tools by allowed/disallowed tools configuration. + + If allowed_tools is set, only tools in that list are returned. + If disallowed_tools is set, tools in that list are excluded. + Tool names are matched with and without server prefixes for flexibility. + + Args: + tools: List of tools to filter + mcp_server: Server configuration with allowed_tools/disallowed_tools + + Returns: + Filtered list of tools """ tools_to_return = tools + + # Filter by allowed_tools (whitelist) if mcp_server.allowed_tools: tools_to_return = [ - tool for tool in tools if tool.name in mcp_server.allowed_tools + tool for tool in tools + if _tool_name_matches(tool.name, mcp_server.allowed_tools) ] + + # Filter by disallowed_tools (blacklist) if mcp_server.disallowed_tools: tools_to_return = [ - tool for tool in tools if tool.name not in mcp_server.disallowed_tools + tool for tool in tools_to_return + if not _tool_name_matches(tool.name, mcp_server.disallowed_tools) ] + return tools_to_return async def _get_tools_from_mcp_servers( @@ -453,9 +497,12 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=add_prefix, ) - all_tools.extend(filter_tools_by_allowed_tools(tools, server)) + + filtered_tools = filter_tools_by_allowed_tools(tools, server) + all_tools.extend(filtered_tools) + verbose_logger.debug( - f"Successfully fetched {len(tools)} tools from server {server.name}" + f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering" ) except Exception as e: verbose_logger.exception( @@ -560,6 +607,25 @@ if MCP_AVAILABLE: name ) + ## CHECK IF USER IS ALLOWED TO CALL THIS TOOL + allowed_mcp_server_ids = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, + ) + + allowed_mcp_servers = global_mcp_server_manager.get_mcp_server_names_from_ids( + allowed_mcp_server_ids + ) + + if not MCPRequestHandler.is_tool_allowed( + allowed_mcp_servers=allowed_mcp_servers, + server_name=server_name_from_prefix, + ): + + raise HTTPException( + status_code=403, + detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}", + ) + standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = ( _get_standard_logging_mcp_tool_call( name=original_tool_name, # Use original name for logging diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 00ffae718e1..0d375f363e0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -917,6 +917,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): url: Optional[str] = None mcp_info: Optional[MCPInfo] = None mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: Optional[List[str]] = None # Stdio-specific fields command: Optional[str] = None args: List[str] = Field(default_factory=list) @@ -985,6 +986,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): updated_by: Optional[str] = None teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: List[str] = Field(default_factory=list) mcp_info: Optional[MCPInfo] = None # Health check status status: Optional[str] = Field( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 012876db481..68708b8fae4 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -41,12 +41,12 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, + NewTeamRequest, ProxyErrorTypes, ProxyException, RoleBasedPermissions, SpecialModelNames, UserAPIKeyAuth, - NewTeamRequest, ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.route_llm_request import route_request @@ -474,7 +474,7 @@ async def get_end_user_object( return return_obj # else, check db - try: + try: response = await prisma_client.db.litellm_endusertable.find_unique( where={"user_id": end_user_id}, include={"litellm_budget_table": True}, @@ -817,7 +817,9 @@ async def _cache_management_object( ): await user_api_key_cache.async_set_cache( - key=key, value=value, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + key=key, + value=value, + ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -892,7 +894,9 @@ async def _get_team_db_check( system_admin_user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) created_team_dict = await new_team( - data=new_team_data, http_request=mock_request, user_api_key_dict=system_admin_user + data=new_team_data, + http_request=mock_request, + user_api_key_dict=system_admin_user, ) response = LiteLLM_TeamTable(**created_team_dict) return response @@ -1166,6 +1170,54 @@ async def get_key_object( return _response +@log_db_metrics +async def get_object_permission( + object_permission_id: str, + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + parent_otel_span: Optional[Span] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, +) -> Optional[LiteLLM_ObjectPermissionTable]: + """ + - Check if object permission id in proxy ObjectPermissionTable + - if valid, return LiteLLM_ObjectPermissionTable object + - if not, then raise an error + """ + if prisma_client is None: + raise Exception( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + + # check if in cache + key = "object_permission_id:{}".format(object_permission_id) + cached_obj_permission = await user_api_key_cache.async_get_cache(key=key) + if cached_obj_permission is not None: + if isinstance(cached_obj_permission, dict): + return LiteLLM_ObjectPermissionTable(**cached_obj_permission) + elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable): + return cached_obj_permission + + # else, check db + try: + response = await prisma_client.db.litellm_objectpermissiontable.find_unique( + where={"object_permission_id": object_permission_id} + ) + + if response is None: + return None + + # save the object permission to cache + await user_api_key_cache.async_set_cache( + key=key, + value=response.model_dump(), + ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + ) + + return LiteLLM_ObjectPermissionTable(**response.dict()) + except Exception: + return None + + @log_db_metrics async def get_org_object( org_id: str, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 55f3f95539a..c400c2d0d86 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -417,6 +417,12 @@ def bytes_to_mb(bytes_value: int): def get_key_model_rpm_limit( user_api_key_dict: UserAPIKeyAuth, ) -> Optional[Dict[str, int]]: + """ + Get the model rpm limit for a given api key + - check key metadata + - check key model max budget + - check team metadata + """ if user_api_key_dict.metadata: if "model_rpm_limit" in user_api_key_dict.metadata: return user_api_key_dict.metadata["model_rpm_limit"] @@ -426,7 +432,9 @@ def get_key_model_rpm_limit( if "rpm_limit" in budget and budget["rpm_limit"] is not None: model_rpm_limit[model] = budget["rpm_limit"] return model_rpm_limit - + elif user_api_key_dict.team_metadata: + if "model_rpm_limit" in user_api_key_dict.team_metadata: + return user_api_key_dict.team_metadata["model_rpm_limit"] return None @@ -439,7 +447,9 @@ def get_key_model_tpm_limit( elif user_api_key_dict.model_max_budget: if "tpm_limit" in user_api_key_dict.model_max_budget: return user_api_key_dict.model_max_budget["tpm_limit"] - + elif user_api_key_dict.team_metadata: + if "model_tpm_limit" in user_api_key_dict.team_metadata: + return user_api_key_dict.team_metadata["model_tpm_limit"] return None @@ -473,6 +483,7 @@ def _has_user_setup_sso(): return sso_setup + def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]: """Return the header_name mapped to CUSTOMER role, if any (dict-based).""" if not user_id_mapping: @@ -522,7 +533,11 @@ def get_end_user_id_from_request_body( for header_name, header_value in request_headers.items(): if header_name.lower() == custom_header_name_to_check.lower(): user_id_from_header = header_value - user_id_str = str(user_id_from_header) if user_id_from_header is not None else "" + user_id_str = ( + str(user_id_from_header) + if user_id_from_header is not None + else "" + ) if user_id_str.strip(): return user_id_str diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 56218b3345e..39f11e64bb7 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -62,12 +62,22 @@ class RouteChecks: for allowed_route in valid_token.allowed_routes ): for allowed_route in valid_token.allowed_routes: - if allowed_route in LiteLLMRoutes._member_names_: + if allowed_route in LiteLLMRoutes._member_names_: if RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value, ): return True + + ################################################ + # For llm_api_routes, also check registered pass-through endpoints + ################################################ + if allowed_route == "llm_api_routes": + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): + return True # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 95c84b914b6..9263142dc90 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -38,6 +38,7 @@ from litellm.proxy.common_utils.callback_utils import ( from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.types.utils import ServerToolUse if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig @@ -46,6 +47,7 @@ if TYPE_CHECKING: else: ProxyConfig = Any from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]: @@ -379,6 +381,7 @@ class ProxyBaseLLMRequestProcessing: user_api_base: Optional[str] = None, version: Optional[str] = None, is_streaming_request: Optional[bool] = False, + contents: Optional[list] = None, # Add contents parameter ) -> Any: """ Common request processing logic for both chat completions and responses API endpoints @@ -417,6 +420,10 @@ class ProxyBaseLLMRequestProcessing: ) ) + # Pass contents if provided + if contents: + self.data["contents"] = contents + ### ROUTE THE REQUEST ### # Do not change this - it should be a constant time fetch - ALWAYS llm_call = await route_request( @@ -729,7 +736,6 @@ class ProxyBaseLLMRequestProcessing: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events """ - from litellm.types.utils import ModelResponse, ModelResponseStream verbose_proxy_logger.debug("inside generator") try: @@ -754,6 +760,10 @@ class ProxyBaseLLMRequestProcessing: response_str = litellm.get_response_string(response_obj=chunk) str_so_far += response_str + # Inject cost into Anthropic-style SSE usage for /v1/messages for any provider + model_name = request_data.get("model", "") + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name) + # Format chunk using helper function yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk) except Exception as e: @@ -785,3 +795,147 @@ class ProxyBaseLLMRequestProcessing: ) error_returned = json.dumps({"error": proxy_exception.to_dict()}) yield f"{STREAM_SSE_DATA_PREFIX}{error_returned}\n\n" + + @staticmethod + def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any: + """ + Process a streaming chunk and inject cost information if enabled. + + Args: + chunk: The streaming chunk (dict, str, bytes, or bytearray) + model_name: Model name for cost calculation + + Returns: + The processed chunk with cost information injected if applicable + """ + if not getattr(litellm, "include_cost_in_streaming_usage", False): + return chunk + + try: + if isinstance(chunk, dict): + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name) + if maybe_modified is not None: + return maybe_modified + elif isinstance(chunk, (bytes, bytearray)): + # Decode to str, inject, and rebuild as bytes + try: + s = chunk.decode("utf-8", errors="ignore") + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) + if maybe_mod is not None: + return (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") + except Exception: + pass + elif isinstance(chunk, str): + # Try to parse SSE frame and inject cost into the data line + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name) + if maybe_mod is not None: + # Ensure trailing frame separator + return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") + except Exception: + # Never break streaming on optional cost injection + pass + + return chunk + + @staticmethod + def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> Optional[str]: + """ + Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line. + + Args: + frame_str: SSE frame string that may contain multiple lines + model_name: Model name for cost calculation + + Returns: + Modified SSE frame string with cost injected, or None if no modification needed + """ + try: + # Split preserving lines + lines = frame_str.split("\n") + for idx, ln in enumerate(lines): + stripped_ln = ln.strip() + if stripped_ln.startswith("data:"): + json_part = stripped_ln.split("data:", 1)[1].strip() + if json_part and json_part != "[DONE]": + obj = json.loads(json_part) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) + if maybe_modified is not None: + # Replace just this line with updated JSON using safe_dumps + lines[idx] = f"data: {safe_dumps(maybe_modified)}" + return "\n".join(lines) + return None + except Exception: + return None + + @staticmethod + def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> Optional[dict]: + """ + Inject cost information into a usage dictionary for message_delta events. + + Args: + obj: Dictionary containing the SSE event data + model_name: Model name for cost calculation + + Returns: + Modified dictionary with cost injected, or None if no modification needed + """ + if ( + obj.get("type") == "message_delta" + and isinstance(obj.get("usage"), dict) + ): + _usage = obj["usage"] + prompt_tokens = int(_usage.get("input_tokens", 0) or 0) + completion_tokens = int(_usage.get("output_tokens", 0) or 0) + total_tokens = int( + _usage.get("total_tokens", prompt_tokens + completion_tokens) + or (prompt_tokens + completion_tokens) + ) + + # Extract additional usage fields + cache_creation_input_tokens = _usage.get("cache_creation_input_tokens") + cache_read_input_tokens = _usage.get("cache_read_input_tokens") + web_search_requests = _usage.get("web_search_requests") + completion_tokens_details = _usage.get("completion_tokens_details") + prompt_tokens_details = _usage.get("prompt_tokens_details") + + + usage_kwargs: dict[str, Any] = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens, + } + + # Add optional named parameters + if completion_tokens_details is not None: + usage_kwargs["completion_tokens_details"] = completion_tokens_details + if prompt_tokens_details is not None: + usage_kwargs["prompt_tokens_details"] = prompt_tokens_details + + # Handle web_search_requests by wrapping in ServerToolUse + if web_search_requests is not None: + usage_kwargs["server_tool_use"] = ServerToolUse( + web_search_requests=web_search_requests + ) + + # Add cache-related fields to **params (handled by Usage.__init__) + if cache_creation_input_tokens is not None: + usage_kwargs["cache_creation_input_tokens"] = cache_creation_input_tokens + if cache_read_input_tokens is not None: + usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens + + _mr = ModelResponse( + usage=Usage(**usage_kwargs) + ) + + try: + cost_val = litellm.completion_cost( + completion_response=_mr, + model=model_name, + ) + except Exception: + cost_val = None + + if cost_val is not None: + obj.setdefault("usage", {})["cost"] = cost_val + return obj + return None \ No newline at end of file diff --git a/litellm/proxy/common_utils/openai_endpoint_utils.py b/litellm/proxy/common_utils/openai_endpoint_utils.py index a18ccd0b0c1..fa49b05696a 100644 --- a/litellm/proxy/common_utils/openai_endpoint_utils.py +++ b/litellm/proxy/common_utils/openai_endpoint_utils.py @@ -6,8 +6,11 @@ from typing import Optional from fastapi import Request +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +SENSITIVE_DATA_MASKER = SensitiveDataMasker() + def remove_sensitive_info_from_deployment(deployment_dict: dict) -> dict: """ @@ -25,6 +28,8 @@ def remove_sensitive_info_from_deployment(deployment_dict: dict) -> dict: deployment_dict["litellm_params"].pop("aws_access_key_id", None) deployment_dict["litellm_params"].pop("aws_secret_access_key", None) + deployment_dict["litellm_params"] = SENSITIVE_DATA_MASKER.mask_dict(deployment_dict["litellm_params"]) + return deployment_dict diff --git a/litellm/proxy/example_config_yaml/pass_through_config.yaml b/litellm/proxy/example_config_yaml/pass_through_config.yaml index ccc13f4d5a2..4e0c4009fbc 100644 --- a/litellm/proxy/example_config_yaml/pass_through_config.yaml +++ b/litellm/proxy/example_config_yaml/pass_through_config.yaml @@ -24,6 +24,19 @@ model_list: litellm_params: model: anthropic/* api_key: os.environ/ANTHROPIC_API_KEY + - model_name: openai/* + litellm_params: + model: openai/* + api_key: os.environ/OPENAI_API_KEY general_settings: master_key: sk-1234 - custom_auth: custom_auth_basic.user_api_key_auth \ No newline at end of file + custom_auth: custom_auth_basic.user_api_key_auth + pass_through_endpoints: + - path: "/azure-config-passthrough" + target: os.environ/AZURE_API_BASE_PASSHROUGH + include_subpath: true + headers: + Authorization: os.environ/AZURE_API_KEY_PASSHROUGH + +litellm_settings: + include_cost_in_streaming_usage: true \ No newline at end of file diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index eb481b0a4f0..51c6d5ab634 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,8 +1,10 @@ -from fastapi import APIRouter, Depends, Request, Response +from fastapi import APIRouter, Depends, Request, Response, HTTPException +from fastapi.responses import StreamingResponse from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + +from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.types.llms.vertex_ai import TokenCountDetailsResponse router = APIRouter( @@ -10,140 +12,63 @@ router = APIRouter( ) -@router.post("/v1beta/models/{model_name}:generateContent", dependencies=[Depends(user_api_key_auth)]) -@router.post("/models/{model_name}:generateContent", dependencies=[Depends(user_api_key_auth)]) +@router.post( + "/v1beta/models/{model_name}:generateContent", + dependencies=[Depends(user_api_key_auth)], +) +@router.post( + "/models/{model_name}:generateContent", dependencies=[Depends(user_api_key_auth)] +) async def google_generate_content( request: Request, model_name: str, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - """ - Not Implemented, this is a placeholder for the google genai generateContent endpoint. - """ - from litellm.proxy.proxy_server import ( - _read_request_body, - general_settings, - llm_router, - proxy_config, - proxy_logging_obj, - select_data_generator, - user_api_base, - user_max_tokens, - user_model, - user_request_timeout, - user_temperature, - version, - ) + from litellm.proxy.proxy_server import llm_router data = await _read_request_body(request=request) if "model" not in data: data["model"] = model_name - processor = ProxyBaseLLMRequestProcessing(data=data) - try: - return await processor.base_process_llm_request( - request=request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="agenerate_content", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, - ) - except Exception as e: - raise await processor._handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - version=version, - ) + # call router + if llm_router is None: + raise HTTPException(status_code=500, detail="Router not initialized") + response = await llm_router.agenerate_content(**data) + return response -class GoogleAIStudioDataGenerator: - """ - Ensures SSE data generator is used for Google AI Studio streaming responses - - Thin wrapper around ProxyBaseLLMRequestProcessing.async_sse_data_generator - """ - @staticmethod - def _select_data_generator(response, user_api_key_dict, request_data): - from litellm.proxy.proxy_server import proxy_logging_obj - return ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=request_data, - proxy_logging_obj=proxy_logging_obj, - ) - -@router.post("/v1beta/models/{model_name}:streamGenerateContent", dependencies=[Depends(user_api_key_auth)]) -@router.post("/models/{model_name}:streamGenerateContent", dependencies=[Depends(user_api_key_auth)]) +@router.post( + "/v1beta/models/{model_name}:streamGenerateContent", + dependencies=[Depends(user_api_key_auth)], +) +@router.post( + "/models/{model_name}:streamGenerateContent", + dependencies=[Depends(user_api_key_auth)], +) async def google_stream_generate_content( request: Request, model_name: str, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - """ - Not Implemented, this is a placeholder for the google genai streamGenerateContent endpoint. - """ - from litellm.proxy.proxy_server import ( - _read_request_body, - general_settings, - llm_router, - proxy_config, - proxy_logging_obj, - user_api_base, - user_max_tokens, - user_model, - user_request_timeout, - user_temperature, - version, - ) + from litellm.proxy.proxy_server import llm_router data = await _read_request_body(request=request) + if "model" not in data: data["model"] = model_name + data["stream"] = True # enforce streaming for this endpoint - processor = ProxyBaseLLMRequestProcessing(data=data) - try: - return await processor.base_process_llm_request( - request=request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="agenerate_content_stream", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=GoogleAIStudioDataGenerator._select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, - is_streaming_request=True, - ) - except Exception as e: - raise await processor._handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - version=version, - ) - + # call router + if llm_router is None: + raise HTTPException(status_code=500, detail="Router not initialized") + response = await llm_router.agenerate_content(**data) + # Check if response is an async iterator (streaming response) + if hasattr(response, "__aiter__"): + return StreamingResponse(response, media_type="text/event-stream") + return response @router.post( @@ -171,13 +96,13 @@ async def google_count_tokens(request: Request, model_name: str): } ``` """ + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter 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 + # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest # Translate contents to openai format messages using the adapter diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index a51547898d9..7479c9dfcf9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -41,6 +41,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( ) from litellm.types.utils import ( Choices, + GuardrailStatus, ModelResponse, ModelResponseStream, StreamingChoices, @@ -361,11 +362,30 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) - httpx_response = await self.async_handler.post( - url=prepared_request.url, - data=prepared_request.body, # type: ignore - headers=prepared_request.headers, # type: ignore - ) + try: + httpx_response = await self.async_handler.post( + url=prepared_request.url, + data=prepared_request.body, # type: ignore + headers=prepared_request.headers, # type: ignore + ) + except Exception as e: + # Endpoint down, timeout, or other HTTP/network errors + verbose_proxy_logger.error( + "Bedrock AI: failed to make guardrail request: %s", str(e) + ) + # Add guardrail information with failure status + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={"error": str(e)}, + request_data=request_data or {}, + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=datetime.now().timestamp(), + duration=(datetime.now() - start_time).total_seconds(), + ) + # Re-raise the exception to maintain existing behavior + raise + ######################################################### # Add guardrail information to request trace ######################################################### @@ -437,15 +457,30 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): def _get_bedrock_guardrail_response_status( self, response: httpx.Response - ) -> Literal["success", "failure"]: + ) -> GuardrailStatus: """ Get the status of the bedrock guardrail response. + + Returns: + "success": Content allowed through with no violations + "guardrail_intervened": Content blocked due to policy violations + "guardrail_failed_to_respond": Technical error or API failure """ if response.status_code == 200: if self._check_bedrock_response_for_exception(response): - return "failure" + return "guardrail_failed_to_respond" + + # Check if the guardrail would block content + try: + _json_response = response.json() + bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) + if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): + return "guardrail_intervened" + except Exception: + pass + return "success" - return "failure" + return "guardrail_failed_to_respond" def _get_http_exception_for_blocked_guardrail( self, response: BedrockGuardrailResponse @@ -692,6 +727,22 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) return + # Check if the ModelResponse has text content in its choices + # to avoid sending empty content to Bedrock (e.g., during tool calls) + if isinstance(response, litellm.ModelResponse): + has_text_content = False + for choice in response.choices: + if isinstance(choice, litellm.Choices): + if choice.message.content and isinstance(choice.message.content, str): + has_text_content = True + break + + if not has_text_content: + verbose_proxy_logger.warning( + "Bedrock AI: not running guardrail. No output text in response" + ) + return + ######################################################### ########## 1. Make parallel Bedrock API requests ########## ######################################################### diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py index fda597bde53..36b5700713b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py +++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py @@ -1,5 +1,7 @@ from datetime import datetime -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union, Type +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Type, Union + +from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger @@ -12,11 +14,11 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.javelin import ( + JavelinGuardInput, JavelinGuardRequest, JavelinGuardResponse, - JavelinGuardInput, ) -from fastapi import HTTPException +from litellm.types.utils import GuardrailStatus if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -95,7 +97,7 @@ class JavelinGuardrail(CustomGuardrail): if self.application: headers["x-javelin-application"] = self.application - status: Literal["success", "failure", "blocked"] = "failure" + status: GuardrailStatus = "guardrail_failed_to_respond" javelin_response: Optional[JavelinGuardResponse] = None exception_str = "" @@ -122,7 +124,7 @@ class JavelinGuardrail(CustomGuardrail): status = "success" return javelin_response except Exception as e: - status = "failure" + status = "guardrail_failed_to_respond" exception_str = str(e) return {"assessments": []} finally: @@ -178,12 +180,12 @@ class JavelinGuardrail(CustomGuardrail): """ Pre-call hook for the Javelin guardrail. """ - from litellm.proxy.common_utils.callback_utils import ( - add_guardrail_to_applied_guardrails_header, - ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_last_user_message, ) + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) verbose_proxy_logger.debug("Javelin Guardrail: pre_call_hook") verbose_proxy_logger.debug("Javelin Guardrail: Request data: %s", data) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index b65664e00bd..0a75328f4da 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -20,6 +20,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import ( LakeraAIRequest, LakeraAIResponse, ) +from litellm.types.utils import GuardrailStatus class LakeraAIGuardrail(CustomGuardrail): @@ -70,7 +71,7 @@ class LakeraAIGuardrail(CustomGuardrail): """ Call the Lakera AI v2 guard API. """ - status: Literal["success", "failure"] = "success" + status: GuardrailStatus = "success" exception_str: str = "" start_time: datetime = datetime.now() lakera_response: Optional[LakeraAIResponse] = None @@ -99,7 +100,7 @@ class LakeraAIGuardrail(CustomGuardrail): lakera_response = LakeraAIResponse(**response.json()) return lakera_response, masked_entity_count except Exception as e: - status = "failure" + status = "guardrail_failed_to_respond" exception_str = str(e) raise e finally: diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 787c46d0dda..e9ddca31777 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -30,6 +30,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( Choices, + GuardrailStatus, ModelResponse, ModelResponseStream, ) @@ -329,14 +330,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): guardrail_response = metadata.get("_model_armor_response", {}) # Determine status – default to "success" but prefer the explicit value if present. - guardrail_status: Literal["success", "failure", "blocked"] = metadata.get( + guardrail_status: GuardrailStatus = metadata.get( "_model_armor_status", "success" ) # type: ignore self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, request_data=request_data, - guardrail_status=guardrail_status, # type: ignore + guardrail_status=guardrail_status, duration=duration, start_time=start_time, end_time=end_time, diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 06d0af681a0..782c785ce58 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -8,7 +8,8 @@ import asyncio import copy import os -from typing import Any, Dict, Final, Literal, Optional, Union, Type, TYPE_CHECKING +from datetime import datetime +from typing import TYPE_CHECKING, Any, Dict, Final, Literal, Optional, Type, Union from urllib.parse import urljoin from fastapi import HTTPException @@ -23,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import EmbeddingResponse, ImageResponse +from litellm.types.utils import EmbeddingResponse, GuardrailStatus, ImageResponse # Constants USER_ROLE: Final[Literal["user"]] = "user" @@ -204,6 +205,7 @@ class NomaGuardrail(CustomGuardrail): user_auth: UserAPIKeyAuth, ) -> Optional[str]: """Shared logic for processing user message checks""" + start_time = datetime.now() extra_data = self.get_guardrail_dynamic_request_body_params(request_data) user_message = await self._extract_user_message(request_data) @@ -218,6 +220,23 @@ class NomaGuardrail(CustomGuardrail): user_auth=user_auth, extra_data=extra_data, ) + + end_time = datetime.now() + duration = (end_time - start_time).total_seconds() + + # Determine guardrail status based on response + guardrail_status = self._determine_guardrail_status(response_json) + + # Always log guardrail information for consistency + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="noma", + guardrail_json_response=response_json, + request_data=request_data, + guardrail_status=guardrail_status, + start_time=start_time.timestamp(), + end_time=end_time.timestamp(), + duration=duration, + ) if self.monitor_mode: await self._handle_verdict_background( @@ -248,6 +267,8 @@ class NomaGuardrail(CustomGuardrail): user_auth: UserAPIKeyAuth, ) -> Optional[str]: """Shared logic for processing LLM response checks""" + + start_time = datetime.now() extra_data = self.get_guardrail_dynamic_request_body_params(request_data) if not isinstance(response, litellm.ModelResponse): @@ -271,6 +292,23 @@ class NomaGuardrail(CustomGuardrail): user_auth=user_auth, extra_data=extra_data, ) + + end_time = datetime.now() + duration = (end_time - start_time).total_seconds() + + # Determine guardrail status based on response + guardrail_status = self._determine_guardrail_status(response_json) + + # Always log guardrail information for consistency + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="noma", + guardrail_json_response=response_json, + request_data=request_data, + guardrail_status=guardrail_status, + start_time=start_time.timestamp(), + end_time=end_time.timestamp(), + duration=duration, + ) if self.monitor_mode: await self._handle_verdict_background( @@ -294,6 +332,41 @@ class NomaGuardrail(CustomGuardrail): await self._check_verdict(ASSISTANT_ROLE, content, response_json) return content + def _determine_guardrail_status(self, response_json: dict) -> GuardrailStatus: + """ + Determine the guardrail status based on NOMA API response. + + Args: + response_json: Response from NOMA API + + Returns: + "success": Content allowed through with no violations + "guardrail_intervened": Content blocked due to policy violations + "guardrail_failed_to_respond": Technical error or API failure + """ + try: + # Check if we got a valid response structure + if not isinstance(response_json, dict): + return "guardrail_failed_to_respond" + + # Get the verdict from the response + verdict = response_json.get("verdict", True) + + # If verdict is True, content is allowed + if verdict is True: + return "success" + + # If verdict is False, content is blocked/flagged + if verdict is False: + return "guardrail_intervened" + + # If verdict is missing or invalid, treat as failure + return "guardrail_failed_to_respond" + + except Exception as e: + verbose_proxy_logger.error(f"Error determining NOMA guardrail status: {str(e)}") + return "guardrail_failed_to_respond" + def _should_only_sensitive_data_failed(self, classification_obj: dict) -> bool: """ Check if only sensitive data detectors (PII, PCI, secrets) have result=true in the classification. @@ -539,8 +612,22 @@ class NomaGuardrail(CustomGuardrail): try: return await self._check_user_message(data, user_api_key_dict) except NomaBlockedMessage: + # Blocked requests were already logged in _process_user_message_check with "blocked" status raise except Exception as e: + # Log technical failures + from datetime import datetime + start_time = datetime.now() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="noma", + guardrail_json_response=str(e), + request_data=data, + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=start_time.timestamp(), + duration=0.0, + ) + verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}") if self.block_failures: @@ -580,8 +667,22 @@ class NomaGuardrail(CustomGuardrail): try: return await self._check_user_message(data, user_api_key_dict) except NomaBlockedMessage: + # Blocked requests were already logged in _process_user_message_check with "blocked" status raise except Exception as e: + # Log technical failures + from datetime import datetime + start_time = datetime.now() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="noma", + guardrail_json_response=str(e), + request_data=data, + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=start_time.timestamp(), + duration=0.0, + ) + verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}") if self.block_failures: @@ -615,8 +716,22 @@ class NomaGuardrail(CustomGuardrail): try: return await self._check_llm_response(data, response, user_api_key_dict) except NomaBlockedMessage: + # Blocked requests were already logged in _process_llm_response_check with "blocked" status raise except Exception as e: + # Log technical failures + from datetime import datetime + start_time = datetime.now() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="noma", + guardrail_json_response=str(e), + request_data=data, + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=start_time.timestamp(), + duration=0.0, + ) + verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}") if self.block_failures: raise diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 38a17595c46..b77e802c717 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -10,14 +10,12 @@ import asyncio import json -from litellm._uuid import uuid from datetime import datetime from typing import ( Any, AsyncGenerator, Dict, List, - Literal, Optional, Tuple, Union, @@ -29,6 +27,7 @@ import aiohttp import litellm # noqa: E401 from litellm import get_secret from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError from litellm.integrations.custom_guardrail import CustomGuardrail @@ -45,6 +44,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.presidio import ( PresidioAnalyzeResponseItem, ) from litellm.types.utils import CallTypes as LitellmCallTypes +from litellm.types.utils import GuardrailStatus from litellm.utils import ( EmbeddingResponse, ImageResponse, @@ -324,7 +324,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """ start_time = datetime.now() analyze_results: Optional[Union[List[PresidioAnalyzeResponseItem], Dict]] = None - status: Literal["success", "failure"] = "success" + status: GuardrailStatus = "success" masked_entity_count: Dict[str, int] = {} exception_str: str = "" try: @@ -356,7 +356,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return redacted_text["text"] except Exception as e: - status = "failure" + status = "guardrail_failed_to_respond" exception_str = str(e) raise e finally: diff --git a/litellm/proxy/hooks/README.dynamic_rate_limiter_v3.md b/litellm/proxy/hooks/README.dynamic_rate_limiter_v3.md new file mode 100644 index 00000000000..701f5e928ce --- /dev/null +++ b/litellm/proxy/hooks/README.dynamic_rate_limiter_v3.md @@ -0,0 +1,170 @@ +# Dynamic Rate Limiter v3 - Saturation-Aware Priority-Based Rate Limiting + +## Overview + +The v3 dynamic rate limiter implements saturation-aware rate limiting with priority-based allocation. It balances resource efficiency (allowing unused capacity to be borrowed) with fairness guarantees (enforcing priorities during high load). + +**Key Behavior:** +- When system is under 80% capacity: Generous mode - allows priority borrowing +- When system is at/above 80% capacity: Strict mode - enforces normalized priority limits + +## How It Works + +### Flow Diagram + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Incoming Request │ +└────────────────────────┬────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ 1. Check Model Saturation │ +│ - Query v3 limiter's Redis counters │ +│ - Calculate: current_usage / capacity │ +│ - Returns: 0.0 (empty) to 1.0+ (saturated) │ +└────────────────────────┬────────────────────────────────────┘ + │ + ▼ + ┌────────┴────────┐ + │ Saturation? │ + └────────┬────────┘ + │ + ┌───────────────┴───────────────┐ + │ │ + ▼ ▼ + < 80% (Generous) >= 80% (Strict) + │ │ + ▼ ▼ +┌─────────────────────┐ ┌─────────────────────┐ +│ Generous Mode │ │ Strict Mode │ +│ │ │ │ +│ - Enforce model- │ │ - Normalize │ +│ wide capacity │ │ priority weights │ +│ - No priority │ │ (if over 1.0) │ +│ restrictions │ │ │ +│ - Allows borrowing │ │ - Create priority- │ +│ │ │ specific │ +│ - First-come- │ │ descriptors │ +│ first-served │ │ │ +│ until capacity │ │ - Enforce strict │ +│ │ │ limits per │ +│ │ │ priority │ +└──────────┬──────────┘ └──────────┬──────────┘ + │ │ + │ ▼ + │ ┌──────────────────────┐ + │ │ Track model usage │ + │ │ for future │ + │ │ saturation checks │ + │ └──────────┬───────────┘ + │ │ + └───────────────┬───────────────┘ + │ + ▼ + ┌──────────────┐ + │ v3 Limiter │ + │ Check │ + └──────┬───────┘ + │ + ┌───────────────┴───────────────┐ + │ │ + ▼ ▼ + OVER_LIMIT OK + │ │ + ▼ ▼ + Return 429 Error Allow Request +``` + +## Configuration + +### Priority Reservation + +Set priority weights in your proxy configuration: + +```python +litellm.priority_reservation = { + "premium": 0.75, # 75% of capacity + "standard": 0.25 # 25% of capacity +} +``` + +### Priority Reservation Settings + +Configure saturation-aware behavior: + +```python +litellm.priority_reservation_settings = PriorityReservationSettings( + default_priority=0.5, # Default weight for users without explicit priority + saturation_threshold=0.80, # 80% - threshold for strict mode enforcement + tracking_multiplier=10 # 10x - multiplier for non-blocking tracking in strict mode +) +``` + +**Settings:** +- `default_priority` (default: 0.5) - Priority weight for users without explicit priority metadata +- `saturation_threshold` (default: 0.80) - Saturation level (0.0-1.0) at which strict priority enforcement begins +- `tracking_multiplier` (default: 10) - Multiplier for model-wide tracking limits in strict mode + +### User Priority Assignment + +Set priority in user metadata: + +```python +user_api_key_dict.metadata = {"priority": "premium"} +``` + +## Priority Weight Normalization + +If priorities sum to > 1.0, they are automatically normalized: + +``` +Input: {key_a: 0.60, key_b: 0.80} = 1.40 total +Output: {key_a: 0.43, key_b: 0.57} = 1.00 total +``` + +This ensures total allocation never exceeds model capacity. + +## Implementation Details + +### Saturation Detection + +- Queries v3 limiter's Redis counters for model-wide usage +- Checks both RPM and TPM, returns higher saturation value +- Non-blocking reads (doesn't increment counters) + +### Mode Selection + +**Generous Mode (< 80% saturation):** +- Creates single model-wide descriptor +- Enforces total capacity only +- Allows any priority to use available capacity +- Prevents over-subscription via model-wide limit + +**Strict Mode (>= 80% saturation):** +- Creates priority-specific descriptors with normalized weights +- Each priority gets its reserved allocation +- Tracks model-wide usage separately (non-blocking, 10x multiplier) +- Ensures fairness under load + +Test scenarios covered: +1. No rate limiting when under capacity +2. Priority queue behavior during saturation +3. Spillover capacity for default keys +4. Over-allocated priorities with normalization +5. Default priority value handling + + +### `_PROXY_DynamicRateLimitHandlerV3` + +Main handler class inheriting from `CustomLogger`. + +**Key Methods:** +- `async_pre_call_hook()` - Main entry point, routes to generous/strict mode +- `_check_model_saturation()` - Queries Redis for current usage +- `_handle_generous_mode()` - Enforces model-wide capacity only +- `_handle_strict_mode()` - Enforces normalized priority limits +- `_normalize_priority_weights()` - Handles over-allocation +- `_create_priority_based_descriptors()` - Creates rate limit descriptors + + diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 38e211dea50..997e33d256f 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -1,9 +1,9 @@ """ -Dynamic rate limiter v3 +Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting """ import os -from typing import List, Literal, Optional, Union +from typing import Dict, List, Literal, Optional, Union from fastapi import HTTPException @@ -24,12 +24,22 @@ from litellm.types.router import ModelGroupInfo class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): """ - Simple validation version that uses v3 parallel request limiter for priority-based rate limiting. + Saturation-aware priority-based rate limiter using v3 infrastructure. - Key differences from original: - 1. Uses v3 limiter's sliding window approach instead of per-minute cache buckets - 2. Leverages Redis Lua scripts for atomic operations under high traffic - 3. Creates priority-specific rate limit descriptors + Key features: + 1. Model capacity ALWAYS enforced at 100% (prevents over-allocation) + 2. Priority usage tracked from first request (accurate accounting) + 3. Priority limits only enforced when saturated >= threshold + 4. Three-phase checking prevents partial counter increments + 5. Reuses v3 limiter's Redis-based tracking (multi-instance safe) + + How it works: + - Phase 1: Read-only check of ALL limits (no increments) + - Phase 2: Decide enforcement based on saturation + - Phase 3: Increment counters only if request allowed + - When under-saturated: priorities can borrow unused capacity (generous) + - When saturated: strict priority-based limits enforced (fair) + - Uses v3 limiter's atomic Lua scripts for race-free increments """ def __init__(self, internal_usage_cache: DualCache): self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache) @@ -57,6 +67,147 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): weight = litellm.priority_reservation[priority] return weight + def _normalize_priority_weights(self) -> Dict[str, float]: + """ + Normalize priority weights if they sum to > 1.0 + + Handles over-allocation: {key_a: 0.60, key_b: 0.80} -> {key_a: 0.43, key_b: 0.57} + """ + if litellm.priority_reservation is None: + return {} + + weights = dict(litellm.priority_reservation) + total_weight = sum(weights.values()) + + if total_weight > 1.0: + normalized = {k: v / total_weight for k, v in weights.items()} + verbose_proxy_logger.debug( + f"Normalized over-allocated priorities: {weights} -> {normalized}" + ) + return normalized + + return weights + + def _get_priority_allocation( + self, + model: str, + priority: Optional[str], + normalized_weights: Dict[str, float], + ) -> tuple[float, str]: + """ + Get priority weight and pool key for a given priority. + + For explicit priorities: returns specific allocation and unique pool key + For default priority: returns default allocation and shared pool key + + Args: + model: Model name + priority: Priority level (None for default) + normalized_weights: Pre-computed normalized weights + + Returns: + tuple: (priority_weight, priority_key) + """ + # Check if this key has an explicit priority in litellm.priority_reservation + has_explicit_priority = ( + priority is not None + and litellm.priority_reservation is not None + and priority in litellm.priority_reservation + ) + + if has_explicit_priority and priority is not None: + # Explicit priority: get its specific allocation + priority_weight = normalized_weights.get(priority, self._get_priority_weight(priority)) + # Use unique key per priority level + priority_key = f"{model}:{priority}" + else: + # No explicit priority: share the default_priority pool with ALL other default keys + priority_weight = litellm.priority_reservation_settings.default_priority + # Use shared key for all default-priority requests + priority_key = f"{model}:default_pool" + + return priority_weight, priority_key + + async def _check_model_saturation( + self, + model: str, + model_group_info: ModelGroupInfo, + ) -> float: + """ + Check current saturation by directly querying v3 limiter's cache keys. + + Reuses v3 limiter's Redis-based tracking (works across multiple instances). + Reads counters WITHOUT incrementing them. + + Returns: + float: Saturation ratio (0.0 = empty, 1.0 = at capacity, >1.0 = over) + """ + try: + max_saturation = 0.0 + + # Query RPM saturation + if model_group_info.rpm is not None and model_group_info.rpm > 0: + # Use v3 limiter's key format: {key:value}:rate_limit_type + counter_key = self.v3_limiter.create_rate_limit_keys( + key="model_saturation_check", + value=model, + rate_limit_type="requests", + ) + + # Query cache for current counter value + counter_value = await self.internal_usage_cache.async_get_cache( + key=counter_key, + litellm_parent_otel_span=None, + local_only=False, # Check Redis too + ) + + if counter_value is not None: + current_requests = int(counter_value) + rpm_saturation = current_requests / model_group_info.rpm + max_saturation = max(max_saturation, rpm_saturation) + + verbose_proxy_logger.debug( + f"Model {model} RPM: {current_requests}/{model_group_info.rpm} " + f"({rpm_saturation:.1%})" + ) + + # Query TPM saturation + if model_group_info.tpm is not None and model_group_info.tpm > 0: + counter_key = self.v3_limiter.create_rate_limit_keys( + key="model_saturation_check", + value=model, + rate_limit_type="tokens", + ) + + counter_value = await self.internal_usage_cache.async_get_cache( + key=counter_key, + litellm_parent_otel_span=None, + local_only=False, + ) + + if counter_value is not None: + current_tokens = float(counter_value) + tpm_saturation = current_tokens / model_group_info.tpm + max_saturation = max(max_saturation, tpm_saturation) + + verbose_proxy_logger.debug( + f"Model {model} TPM: {current_tokens}/{model_group_info.tpm} " + f"({tpm_saturation:.1%})" + ) + + verbose_proxy_logger.debug( + f"Model {model} overall saturation: {max_saturation:.1%}" + ) + + return max_saturation + + except Exception as e: + verbose_proxy_logger.error( + f"Error checking saturation for {model}: {str(e)}" + ) + # Fail open: assume not saturated on error + return 0.0 + def _create_priority_based_descriptors( self, model: str, @@ -64,11 +215,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): priority: Optional[str], ) -> List[RateLimitDescriptor]: """ - Create rate limit descriptors based on priority and model group limits. + Create rate limit descriptors with normalized priority weights. - This is the key change: instead of calculating dynamic quotas based on active projects, - we create descriptors with priority-adjusted limits and let the v3 limiter handle - the actual rate limiting with its sliding window approach. + Uses normalized weights to handle over-allocation scenarios. + + For explicit priorities: each priority gets its own pool (e.g., prod gets 75%) + For default priority: ALL keys without explicit priority share ONE pool (e.g., all share 25%) """ descriptors: List[RateLimitDescriptor] = [] @@ -79,23 +231,22 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if model_group_info is None: return descriptors - # Get priority weight - priority_weight = self._get_priority_weight(priority) - - # Create priority-specific rate limits - # Use model:priority as the key to separate different priority levels - priority_key = f"{model}:{priority or 'default'}" + # Get normalized priority weight and pool key + normalized_weights = self._normalize_priority_weights() + priority_weight, priority_key = self._get_priority_allocation( + model=model, + priority=priority, + normalized_weights=normalized_weights, + ) rate_limit_config: RateLimitDescriptorRateLimitObject = {} # Apply priority weight to model limits if model_group_info.tpm is not None: - # Reserve portion of TPM based on priority reserved_tpm = int(model_group_info.tpm * priority_weight) rate_limit_config["tokens_per_unit"] = reserved_tpm if model_group_info.rpm is not None: - # Reserve portion of RPM based on priority reserved_rpm = int(model_group_info.rpm * priority_weight) rate_limit_config["requests_per_unit"] = reserved_rpm @@ -112,6 +263,187 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return descriptors + def _create_model_tracking_descriptor( + self, + model: str, + model_group_info: ModelGroupInfo, + high_limit_multiplier: int = 1, + ) -> RateLimitDescriptor: + """ + Create a descriptor for tracking model-wide usage. + + Args: + model: Model name + model_group_info: Model configuration with RPM/TPM limits + high_limit_multiplier: Multiplier for limits (use >1 for tracking-only) + + Returns: + Rate limit descriptor for model-wide tracking + """ + return RateLimitDescriptor( + key="model_saturation_check", + value=model, + rate_limit={ + "requests_per_unit": ( + model_group_info.rpm * high_limit_multiplier + if model_group_info.rpm else None + ), + "tokens_per_unit": ( + model_group_info.tpm * high_limit_multiplier + if model_group_info.tpm else None + ), + "window_size": self.v3_limiter.window_size, + }, + ) + + + async def _check_rate_limits( + self, + model: str, + model_group_info: ModelGroupInfo, + user_api_key_dict: UserAPIKeyAuth, + key_priority: Optional[str], + saturation: float, + data: dict, + ) -> None: + """ + Check rate limits using THREE-PHASE approach to prevent partial increments. + + Phase 1: Read-only check of ALL limits (no increments) + Phase 2: Decide which limits to enforce based on saturation + Phase 3: Increment ALL counters atomically (model + priority) + + This prevents the bug where: + - Model counter increments in stage 1 + - Priority check fails in stage 2 + - Request blocked but model counter already incremented + + Key behaviors: + - All checks performed first (read-only) + - Only increment counters if request will be allowed + - Model capacity: Always enforced at 100% + - Priority limits: Only enforced when saturated >= threshold + - Both counters tracked from first request (accurate accounting) + + Args: + model: Model name + model_group_info: Model configuration + user_api_key_dict: User authentication info + key_priority: User's priority level + saturation: Current saturation level + data: Request data dictionary + + Raises: + HTTPException: If any limit is exceeded + """ + import json + saturation_threshold = litellm.priority_reservation_settings.saturation_threshold + should_enforce_priority = saturation >= saturation_threshold + + # Build ALL descriptors upfront + descriptors_to_check: List[RateLimitDescriptor] = [] + + # Model-wide descriptor (always enforce) + model_wide_descriptor = self._create_model_tracking_descriptor( + model=model, + model_group_info=model_group_info, + high_limit_multiplier=1, + ) + descriptors_to_check.append(model_wide_descriptor) + + # Priority descriptors (always track, conditionally enforce) + priority_descriptors = self._create_priority_based_descriptors( + model=model, + user_api_key_dict=user_api_key_dict, + priority=key_priority, + ) + if priority_descriptors: + descriptors_to_check.extend(priority_descriptors) + + # PHASE 1: Read-only check of ALL limits (no increments) + check_response = await self.v3_limiter.should_rate_limit( + descriptors=descriptors_to_check, + parent_otel_span=user_api_key_dict.parent_otel_span, + read_only=True, # CRITICAL: Don't increment counters yet + ) + + verbose_proxy_logger.debug(f"Read-only check: {json.dumps(check_response, indent=2)}") + + # PHASE 2: Decide which limits to enforce + if check_response["overall_code"] == "OVER_LIMIT": + for status in check_response["statuses"]: + if status["code"] == "OVER_LIMIT": + descriptor_key = status["descriptor_key"] + + # Model-wide limit exceeded (ALWAYS enforce) + if descriptor_key == "model_saturation_check": + raise HTTPException( + status_code=429, + detail={ + "error": f"Model capacity reached for {model}. " + f"Priority: {key_priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": key_priority or "default", + }, + ) + + # Priority limit exceeded (ONLY enforce when saturated) + elif descriptor_key == "priority_model" and should_enforce_priority: + verbose_proxy_logger.debug( + f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " + f"priority: {key_priority}" + ) + raise HTTPException( + status_code=429, + detail={ + "error": f"Priority-based rate limit exceeded. " + f"Priority: {key_priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}, " + f"Model saturation: {saturation:.1%}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": key_priority or "default", + "x-litellm-saturation": f"{saturation:.2%}", + }, + ) + + # PHASE 3: Increment counters separately to avoid early-exit issues + # Model counter must ALWAYS increment, but priority counter might be over limit + # If we increment them together, v3_limiter's in-memory check will exit early + # and skip incrementing the model counter + + # Step 3a: Increment model-wide counter (always) + model_increment_response = await self.v3_limiter.should_rate_limit( + descriptors=[model_wide_descriptor], + parent_otel_span=user_api_key_dict.parent_otel_span, + read_only=False, + ) + + # Step 3b: Increment priority counter (may be over limit, but we still track it) + if priority_descriptors: + priority_increment_response = await self.v3_limiter.should_rate_limit( + descriptors=priority_descriptors, + parent_otel_span=user_api_key_dict.parent_otel_span, + read_only=False, + ) + + # Combine responses for post-call hook + combined_response = { + "overall_code": model_increment_response["overall_code"], + "statuses": model_increment_response["statuses"] + priority_increment_response["statuses"] + } + data["litellm_proxy_rate_limit_response"] = combined_response + else: + data["litellm_proxy_rate_limit_response"] = model_increment_response + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -130,60 +462,87 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ], ) -> Optional[Union[Exception, str, dict]]: """ - Pre-call hook using v3 limiter for priority-based rate limiting. + Saturation-aware pre-call hook for priority-based rate limiting. + + Flow: + 1. Check current saturation level + 2. THREE-PHASE rate limit check: + - PHASE 1: Read-only check of ALL limits (no increments) + - PHASE 2: Decide which limits to enforce based on saturation + - PHASE 3: Increment ALL counters atomically if request allowed + + This three-phase approach ensures: + - Model capacity is NEVER exceeded (always enforced at 100%) + - Priority usage tracked from first request (accurate metrics) + - Counters only increment when request will be allowed (prevents phantom usage) + - When under-saturated: priorities can borrow unused capacity (generous) + - When saturated: fair allocation based on normalized priority weights (strict) + + Example with 100 RPM model, 60% priority allocation, 80% threshold: + - Saturation < 80%: Priority can use up to 100 RPM (model limit enforced only) + - Saturation >= 80%: Priority limited to 60 RPM (both limits enforced) + + Prevents bugs where: + - Model counter increments but priority check fails → model over-capacity + - Priority counter increments but not enforced → inaccurate metrics + + Args: + user_api_key_dict: User authentication and metadata + cache: Dual cache instance + data: Request data containing model name + call_type: Type of API call being made + + Returns: + None if request is allowed, otherwise raises HTTPException """ if "model" not in data: return None + model = data["model"] key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None) - # Create priority-based descriptors - descriptors = self._create_priority_based_descriptors( - model=data["model"], - user_api_key_dict=user_api_key_dict, - priority=key_priority, + # Get model configuration + model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info( + model_group=model ) - - if not descriptors: - verbose_proxy_logger.debug("No rate limit descriptors created, allowing request") + if model_group_info is None: + verbose_proxy_logger.debug(f"No model group info for {model}, allowing request") return None try: - # Use v3 limiter to check rate limits - response = await self.v3_limiter.should_rate_limit( - descriptors=descriptors, - parent_otel_span=user_api_key_dict.parent_otel_span, + # STEP 1: Check current saturation level + saturation = await self._check_model_saturation(model, model_group_info) + + saturation_threshold = litellm.priority_reservation_settings.saturation_threshold + + verbose_proxy_logger.debug( + f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, " + f"Threshold={saturation_threshold:.1%}, Priority={key_priority}" ) - - if response["overall_code"] == "OVER_LIMIT": - # Find which descriptor hit the limit - for status in response["statuses"]: - if status["code"] == "OVER_LIMIT": - raise HTTPException( - status_code=429, - detail={ - "error": f"Priority-based rate limit exceeded for {status['descriptor_key']}. " - f"Priority: {key_priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": key_priority or "default", - }, - ) - else: - # Store response for post-call hook - data["litellm_proxy_rate_limit_response"] = response - + + data["litellm_model_saturation"] = saturation + + # STEP 2: Check rate limits in THREE phases + # Phase 1: Read-only check of ALL limits (no increments) + # Phase 2: Decide which limits to enforce (based on saturation) + # Phase 3: Increment ALL counters only if request will be allowed + # This prevents partial increments and ensures accurate tracking + await self._check_rate_limits( + model=model, + model_group_info=model_group_info, + user_api_key_dict=user_api_key_dict, + key_priority=key_priority, + saturation=saturation, + data=data, + ) + except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - f"Error in dynamic rate limiter v3 pre-call hook: {str(e)}" + verbose_proxy_logger.error( + f"Error in dynamic rate limiter: {str(e)}, allowing request" ) - # Allow request to proceed on unexpected errors + # Fail open on unexpected errors return None return None diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index eda380b5165..2ca45b55ec7 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -413,6 +413,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Check if any of the rate limit descriptors should be rate limited. Returns a RateLimitResponse with the overall code and status for each descriptor. Uses batch operations for Redis to improve performance. + + Args: + descriptors: List of rate limit descriptors to check + parent_otel_span: Optional OpenTelemetry span for tracing + read_only: If True, only check limits without incrementing counters """ now = datetime.now().timestamp() @@ -485,8 +490,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if rate_limit_response["overall_code"] == "OVER_LIMIT": return rate_limit_response - ## IF under limit, check Redis - if self.batch_rate_limiter_script is not None: + ## IF under limit in-memory, check Redis + if read_only: + # READ-ONLY MODE: Just read current values without incrementing + cache_values = await self.internal_usage_cache.async_batch_get_cache( + keys=keys_to_fetch, + parent_otel_span=parent_otel_span, + local_only=False, # Check Redis too + ) + + # For keys that don't exist yet, set them to 0 + if cache_values is None: + cache_values = [] + for _ in keys_to_fetch: + cache_values.append(str(now_int) if _.endswith(":window") else 0) + elif self.batch_rate_limiter_script is not None: + # NORMAL MODE: Increment counters in Redis # Group keys by hash tag for Redis cluster compatibility cache_values = await self._execute_redis_batch_rate_limiter_script( keys_to_fetch=keys_to_fetch, @@ -514,6 +533,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): local_only=True, ) else: + # NORMAL MODE: In-memory sliding window (no Redis) cache_values = await self.in_memory_cache_sliding_window( keys=keys_to_fetch, now_int=now_int, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 1912c54920c..b0c09b1e2e8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -5,7 +5,6 @@ Endpoints here: - GET `/v1/mcp/server` - Returns all of the configured mcp servers in the db filtered by requestor's access - GET `/v1/mcp/server/{server_id}` - Returns the the specific mcp server in the db given `server_id` filtered by requestor's access -- GET `/v1/mcp/server/{server_id}/tools` - Get all the tools from the mcp server specified by the `server_id` - POST `/v1/mcp/server` - Add a new external mcp server. - PUT `/v1/mcp/server` - Edits an existing mcp server. - DELETE `/v1/mcp/server/{server_id}` - Deletes the mcp server given `server_id`. diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 1227a5017b2..a65b9cc95d2 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1594,7 +1594,6 @@ class SSOAuthenticationHandler: master_key or "", algorithm="HS256", ) - verbose_proxy_logger.info(f"user_id: {user_id}; jwt_token: {jwt_token}") if user_id is not None and isinstance(user_id, str): litellm_dashboard_ui += "?login=success" verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}") diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 5e171af5252..8aa3b90d954 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -6,12 +6,14 @@ Provider-specific Pass-Through Endpoints Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. """ +import json import os from typing import Optional, cast import httpx -from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket from fastapi.responses import StreamingResponse +from starlette.websockets import WebSocketState import litellm from litellm._logging import verbose_proxy_logger @@ -19,7 +21,9 @@ from litellm.constants import BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.auth.user_api_key_auth import ( + user_api_key_auth, +) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, get_form_data, @@ -28,6 +32,8 @@ from litellm.proxy.common_utils.http_parsing_utils import ( from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( create_pass_through_route, + create_websocket_passthrough_route, + websocket_passthrough_request, ) from litellm.proxy.utils import is_known_model from litellm.secret_managers.main import get_secret_str @@ -51,9 +57,7 @@ def create_request_copy(request: Request): } -def is_passthrough_request_using_router_model( - request_body: dict, llm_router: Optional[litellm.Router] -) -> bool: +def is_passthrough_request_using_router_model(request_body: dict, llm_router: Optional[litellm.Router]) -> bool: """ Returns True if the model is in the llm_router model names """ @@ -89,16 +93,12 @@ async def llm_passthrough_factory_proxy_route( model=None, ) if provider_config is None: - raise HTTPException( - status_code=404, detail=f"Provider {custom_llm_provider} not found" - ) + raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} not found") base_target_url = provider_config.get_api_base() if base_target_url is None: - raise HTTPException( - status_code=404, detail=f"Provider {custom_llm_provider} api base not found" - ) + raise HTTPException(status_code=404, detail=f"Provider {custom_llm_provider} api base not found") encoded_endpoint = httpx.URL(endpoint).path @@ -143,7 +143,7 @@ async def llm_passthrough_factory_proxy_route( _request_body = await request.json() else: _request_body = await get_form_data(request) - + if _request_body.get("stream"): is_streaming_request = True @@ -177,17 +177,11 @@ async def gemini_proxy_route( [Docs](https://docs.litellm.ai/docs/pass_through/google_ai_studio) """ ## CHECK FOR LITELLM API KEY IN THE QUERY PARAMS - ?..key=LITELLM_API_KEY - google_ai_studio_api_key = request.query_params.get("key") or request.headers.get( - "x-goog-api-key" - ) + google_ai_studio_api_key = request.query_params.get("key") or request.headers.get("x-goog-api-key") - user_api_key_dict = await user_api_key_auth( - request=request, api_key=f"Bearer {google_ai_studio_api_key}" - ) + user_api_key_dict = await user_api_key_auth(request=request, api_key=f"Bearer {google_ai_studio_api_key}") - base_target_url = ( - os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" - ) + base_target_url = os.getenv("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction @@ -220,6 +214,7 @@ async def gemini_proxy_route( endpoint_func = create_pass_through_route( endpoint=endpoint, target=str(updated_url), + custom_llm_provider="gemini", ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, @@ -304,9 +299,7 @@ async def vllm_proxy_route( from litellm.proxy.proxy_server import llm_router request_body = await get_request_body(request) - is_router_model = is_passthrough_request_using_router_model( - request_body, llm_router - ) + is_router_model = is_passthrough_request_using_router_model(request_body, llm_router) is_streaming_request = is_passthrough_request_streaming(request_body) if is_router_model and llm_router: result = cast( @@ -321,11 +314,7 @@ async def vllm_proxy_route( content=None, data=None, files=None, - json=( - request_body - if request.headers.get("content-type") == "application/json" - else None - ), + json=(request_body if request.headers.get("content-type") == "application/json" else None), params=None, headers=None, cookies=None, @@ -503,9 +492,7 @@ async def handle_bedrock_count_tokens( # Extract model from request body model = request_body.get("model") if not model: - raise HTTPException( - status_code=400, detail={"error": "Model is required in request body"} - ) + raise HTTPException(status_code=400, detail={"error": "Model is required in request body"}) # Get model parameters from router litellm_params = {"user_api_key_dict": user_api_key_dict} @@ -544,9 +531,7 @@ async def handle_bedrock_count_tokens( raise except Exception as e: verbose_proxy_logger.error(f"Error in handle_bedrock_count_tokens: {str(e)}") - raise HTTPException( - status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"} - ) + raise HTTPException(status_code=500, detail={"error": f"CountTokens processing error: {str(e)}"}) async def bedrock_llm_proxy_route( @@ -598,8 +583,7 @@ async def bedrock_llm_proxy_route( raise HTTPException( status_code=400, detail={ - "error": "Model missing from endpoint. Expected format: /model//. Got: " - + endpoint, + "error": "Model missing from endpoint. Expected format: /model//. Got: " + endpoint, }, ) @@ -663,9 +647,7 @@ async def bedrock_proxy_route( aws_region_name = litellm.utils.get_secret(secret_name="AWS_REGION_NAME") if _is_bedrock_agent_runtime_route(endpoint=endpoint): # handle bedrock agents - base_target_url = ( - f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" - ) + base_target_url = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" else: return await bedrock_llm_proxy_route( endpoint=endpoint, @@ -686,7 +668,8 @@ async def bedrock_proxy_route( # Add or update query parameters from litellm.llms.bedrock.chat import BedrockConverseLLM - credentials: Credentials = BedrockConverseLLM().get_credentials() + bedrock_llm = BedrockConverseLLM() + credentials: Credentials = bedrock_llm.get_credentials() # type: ignore sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) headers = {"Content-Type": "application/json"} # Assuming the body contains JSON data, parse it @@ -694,9 +677,7 @@ async def bedrock_proxy_route( data = await request.json() except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - _request = AWSRequest( - method="POST", url=str(updated_url), data=json.dumps(data), headers=headers - ) + _request = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) sigv4.add_auth(_request) prepped = _request.prepare() @@ -757,14 +738,8 @@ async def assemblyai_proxy_route( [Docs](https://api.assemblyai.com) """ # Set base URL based on the route - assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url( - url=str(request.url) - ) - base_target_url = ( - AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region( - region=assembly_region - ) - ) + assembly_region = AssemblyAIPassthroughLoggingHandler._get_assembly_region_from_url(url=str(request.url)) + base_target_url = AssemblyAIPassthroughLoggingHandler._get_assembly_base_url_from_region(region=assembly_region) encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction if not encoded_endpoint.startswith("/"): @@ -822,18 +797,14 @@ async def azure_proxy_route( """ base_target_url = get_secret_str(secret_name="AZURE_API_BASE") if base_target_url is None: - raise Exception( - "Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure." - ) + raise Exception("Required 'AZURE_API_BASE' in environment to make pass-through calls to Azure.") # Add or update query parameters azure_api_key = passthrough_endpoint_router.get_credentials( custom_llm_provider=litellm.LlmProviders.AZURE.value, region_name=None, ) if azure_api_key is None: - raise Exception( - "Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure." - ) + raise Exception("Required 'AZURE_API_KEY' in environment to make pass-through calls to Azure.") return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( endpoint=endpoint, @@ -857,9 +828,7 @@ class BaseVertexAIPassThroughHandler(ABC): @staticmethod @abstractmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: pass @@ -869,9 +838,7 @@ class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler): return "https://discoveryengine.googleapis.com/" @staticmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: return base_target_url @@ -881,9 +848,7 @@ class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler): return get_vertex_base_url(vertex_location) @staticmethod - def update_base_target_url_with_credential_location( - base_target_url: str, vertex_location: Optional[str] - ) -> str: + def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: Optional[str]) -> str: return get_vertex_base_url(vertex_location) @@ -949,18 +914,14 @@ async def _base_vertex_proxy_route( location=vertex_location, ) - base_target_url = get_vertex_pass_through_handler.get_default_base_target_url( - vertex_location - ) + base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) headers_passed_through = False # Use headers from the incoming request if no vertex credentials are found if vertex_credentials is None or vertex_credentials.vertex_project is None: headers = dict(request.headers) or {} headers_passed_through = True - verbose_proxy_logger.debug( - "default_vertex_config not set, incoming request headers %s", headers - ) + verbose_proxy_logger.debug("default_vertex_config not set, incoming request headers %s", headers) headers.pop("content-length", None) headers.pop("host", None) else: @@ -1126,9 +1087,7 @@ async def openai_proxy_route( region_name=None, ) if openai_api_key is None: - raise Exception( - "Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI." - ) + raise Exception("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.") return await BaseOpenAIPassThroughHandler._base_openai_pass_through_handler( endpoint=endpoint, @@ -1174,9 +1133,7 @@ class BaseOpenAIPassThroughHandler: endpoint_func = create_pass_through_route( endpoint=endpoint, target=str(updated_url), - custom_headers=BaseOpenAIPassThroughHandler._assemble_headers( - api_key=api_key, request=request - ), + custom_headers=BaseOpenAIPassThroughHandler._assemble_headers(api_key=api_key, request=request), ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, @@ -1193,10 +1150,7 @@ class BaseOpenAIPassThroughHandler: """ Appends the OpenAI-Beta header to the headers if the request is an OpenAI Assistants API request """ - if ( - RouteChecks._is_assistants_api_request(request) is True - and "OpenAI-Beta" not in headers - ): + if RouteChecks._is_assistants_api_request(request) is True and "OpenAI-Beta" not in headers: headers["OpenAI-Beta"] = "assistants=v2" return headers @@ -1212,9 +1166,7 @@ class BaseOpenAIPassThroughHandler: ) @staticmethod - def _join_url_paths( - base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders - ) -> str: + def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm.LlmProviders) -> str: """ Properly joins a base URL with a path, preserving any existing path in the base URL. """ @@ -1230,13 +1182,192 @@ class BaseOpenAIPassThroughHandler: joined_path_str = str(base_url.copy_with(path=full_path)) # Apply OpenAI-specific path handling for both branches - if ( - custom_llm_provider == litellm.LlmProviders.OPENAI - and "/v1/" not in joined_path_str - ): + if custom_llm_provider == litellm.LlmProviders.OPENAI and "/v1/" not in joined_path_str: # Insert v1 after api.openai.com for OpenAI requests - joined_path_str = joined_path_str.replace( - "api.openai.com/", "api.openai.com/v1/" - ) + joined_path_str = joined_path_str.replace("api.openai.com/", "api.openai.com/v1/") return joined_path_str + + +async def vertex_ai_live_websocket_passthrough( + websocket: WebSocket, + model: Optional[str] = None, + vertex_project: Optional[str] = None, + vertex_location: Optional[str] = None, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, +): + """ + Vertex AI Live API WebSocket Pass-through Function + + This function provides WebSocket passthrough functionality for Vertex AI Live API, + allowing real-time communication with Google's Live API service. + + Note: This function should be registered in proxy_server.py using: + app.websocket("/vertex_ai/live")(vertex_ai_live_websocket_passthrough) + """ + from litellm.proxy.proxy_server import proxy_logging_obj + + _ = user_api_key_dict # passthrough route already authenticated; avoid lint warnings + + await websocket.accept() + + incoming_headers = dict(websocket.headers) + vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials( + project_id=vertex_project, + location=vertex_location, + ) + + if vertex_credentials_config is None: + # Attempt to load defaults from environment/config if not already initialised + passthrough_endpoint_router.set_default_vertex_config() + vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials( + project_id=vertex_project, + location=vertex_location, + ) + + resolved_project = vertex_project + resolved_location: Optional[str] = vertex_location + credentials_value: Optional[str] = None + + if vertex_credentials_config is not None: + resolved_project = resolved_project or vertex_credentials_config.vertex_project + temp_location = ( + resolved_location or vertex_credentials_config.vertex_location + ) + # Ensure resolved_location is a string + if isinstance(temp_location, dict): + resolved_location = str(temp_location) + elif temp_location is not None: + resolved_location = str(temp_location) + else: + resolved_location = None + credentials_value = str(vertex_credentials_config.vertex_credentials) if vertex_credentials_config.vertex_credentials is not None else None + + try: + resolved_location = resolved_location or ( + vertex_llm_base.get_default_vertex_location() + ) + if model: + resolved_location = vertex_llm_base.get_vertex_region( + vertex_region=resolved_location, + model=model, + ) + + ( + access_token, + resolved_project, + ) = await vertex_llm_base._ensure_access_token_async( + credentials=credentials_value, + project_id=resolved_project, + custom_llm_provider="vertex_ai_beta", + ) + except Exception as e: + verbose_proxy_logger.exception( + "Failed to prepare Vertex AI credentials for live passthrough" + ) + # Log the authentication failure using proxy_logging_obj + if proxy_logging_obj and user_api_key_dict: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data={}, + ) + if websocket.client_state != WebSocketState.DISCONNECTED: + await websocket.close(code=1011, reason="Vertex AI authentication failed") + return + + host_location = resolved_location or vertex_llm_base.get_default_vertex_location() + host = ( + "aiplatform.googleapis.com" + if host_location == "global" + else f"{host_location}-aiplatform.googleapis.com" + ) + service_url = ( + f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" + ) + + upstream_headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + } + if resolved_project: + upstream_headers["x-goog-user-project"] = resolved_project + + # Forward any custom x-goog-* headers provided by the caller if we haven't overridden them + for header_name, header_value in incoming_headers.items(): + lower_header = header_name.lower() + if lower_header.startswith("x-goog-") and header_name not in upstream_headers: + upstream_headers[header_name] = header_value + + # Use the new WebSocket passthrough pattern + if user_api_key_dict is None: + raise ValueError("user_api_key_dict is required for WebSocket passthrough") + + return await websocket_passthrough_request( + websocket=websocket, + target=service_url, + custom_headers=upstream_headers, + user_api_key_dict=user_api_key_dict, + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + ) + + +def create_vertex_ai_live_websocket_endpoint(): + """ + Create a Vertex AI Live WebSocket endpoint using the new passthrough pattern. + + This demonstrates how to use the create_websocket_passthrough_route function + for a provider-specific WebSocket endpoint. + """ + # This would be used like: + # endpoint_func = create_vertex_ai_live_websocket_endpoint() + # app.websocket("/vertex_ai/live")(endpoint_func) + + # For now, we'll keep the existing implementation since it has + # provider-specific logic for Vertex AI credentials and headers + return vertex_ai_live_websocket_passthrough + + +def create_generic_websocket_passthrough_endpoint( + provider: str, + target_url: str, + custom_headers: Optional[dict] = None, + forward_headers: bool = False, + cost_per_request: Optional[float] = None, +): + """ + Create a generic WebSocket passthrough endpoint for any provider. + + This demonstrates the new WebSocket passthrough pattern that's similar to + the HTTP create_pass_through_route function. + + Args: + provider: The provider name (e.g., "anthropic", "cohere") + target_url: The target WebSocket URL + custom_headers: Custom headers to include + forward_headers: Whether to forward incoming headers + + Returns: + A WebSocket endpoint function that can be registered with app.websocket() + + Example usage: + # Create a WebSocket endpoint for Anthropic + anthropic_ws_func = create_generic_websocket_passthrough_endpoint( + provider="anthropic", + target_url="wss://api.anthropic.com/v1/ws", + custom_headers={"x-api-key": "your-api-key"}, + forward_headers=True + ) + + # Register it in proxy_server.py + app.websocket("/anthropic/ws")(anthropic_ws_func) + """ + return create_websocket_passthrough_route( + endpoint=f"/{provider}/ws", + target=target_url, + custom_headers=custom_headers, + _forward_headers=forward_headers, + cost_per_request=cost_per_request, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index b9858202bf8..9b6e22b8196 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -20,7 +20,7 @@ from litellm.types.utils import ModelResponse, TextCompletionResponse if TYPE_CHECKING: from ..success_handler import PassThroughEndpointLogging - from ..types import EndpointType + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType else: PassThroughEndpointLogging = Any EndpointType = Any @@ -228,6 +228,7 @@ class AnthropicPassthroughLoggingHandler: except (StopIteration, StopAsyncIteration): break complete_streaming_response = litellm.stream_chunk_builder( - chunks=all_openai_chunks + chunks=all_openai_chunks, + logging_obj=litellm_logging_obj, ) return complete_streaming_response diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py new file mode 100644 index 00000000000..ecd5b0eb094 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -0,0 +1,265 @@ +import json +import re +from datetime import datetime +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import httpx + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator as GeminiModelResponseIterator, +) +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.types.utils import ( + ModelResponse, + TextCompletionResponse, +) + +if TYPE_CHECKING: + from ..success_handler import PassThroughEndpointLogging + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType +else: + PassThroughEndpointLogging = Any + EndpointType = Any + + +class GeminiPassthroughLoggingHandler: + @staticmethod + def gemini_passthrough_handler( + 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: + if "generateContent" in url_route: + model = GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) + + # Use Gemini config for transformation + instance_of_gemini_llm = litellm.GoogleAIStudioGeminiConfig() + litellm_model_response: ModelResponse = instance_of_gemini_llm.transform_response( + model=model, + messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], + raw_response=httpx_response, + model_response=litellm.ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data={}, + encoding=litellm.encoding, + ) + kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content( + litellm_model_response=litellm_model_response, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + custom_llm_provider="gemini", + ) + + return { + "result": litellm_model_response, + "kwargs": kwargs, + } + else: + return { + "result": None, + "kwargs": kwargs, + } + + @staticmethod + def _handle_logging_gemini_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], + model: Optional[str], + end_time: datetime, + ) -> PassThroughEndpointLoggingTypedDict: + """ + Takes raw chunks from Gemini passthrough endpoint and logs them in litellm callbacks + + - Builds complete response from chunks + - Creates standard logging object + - Logs in litellm callbacks + """ + kwargs: Dict[str, Any] = {} + model = model or GeminiPassthroughLoggingHandler.extract_model_from_url(url_route) + complete_streaming_response = GeminiPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + url_route=url_route, + ) + + if complete_streaming_response is None: + verbose_proxy_logger.error( + "Unable to build complete streaming response for Gemini passthrough endpoint, not logging..." + ) + return { + "result": None, + "kwargs": kwargs, + } + + kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content( + litellm_model_response=complete_streaming_response, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + logging_obj=litellm_logging_obj, + custom_llm_provider="gemini", + ) + + return { + "result": complete_streaming_response, + "kwargs": kwargs, + } + + @staticmethod + def _parse_gemini_streaming_json( + all_chunks: List[str], + gemini_iterator: GeminiModelResponseIterator, + ) -> List[Any]: + """ + Parse Gemini streaming chunks from fragmented JSON strings. + + Gemini's streaming format sends a single JSON array that may be split across + multiple lines. This method: + 1. Joins all fragmented string chunks into complete JSON + 2. Parses the JSON array/object + 3. Transforms each item using Gemini's chunk_parser + + Args: + all_chunks: Raw string chunks from the streaming response + gemini_iterator: GeminiModelResponseIterator instance for parsing + + Returns: + List of parsed chunks in OpenAI format, or empty list if parsing fails + """ + parsed_chunks = [] + + verbose_proxy_logger.debug(f"Gemini streaming: Processing {len(all_chunks)} raw chunks") + + # Gemini streaming response is a single JSON array that may be split across lines + # Join all chunks back together to reconstruct the complete JSON + combined_chunk = "".join(all_chunks) + + # Parse the combined JSON string + try: + dict_chunk = json.loads(combined_chunk) + verbose_proxy_logger.debug(f"Parsed JSON object: {type(dict_chunk)}") + + # Gemini returns an array of response objects + if isinstance(dict_chunk, list): + for item in dict_chunk: + try: + # Call chunk_parser directly with the dict, not _common_chunk_parsing_logic + parsed_chunk = gemini_iterator.chunk_parser(chunk=item) + if parsed_chunk is not None: + parsed_chunks.append(parsed_chunk) + except Exception as e: + verbose_proxy_logger.error(f"Error parsing Gemini chunk item: {e}", exc_info=True) + continue + else: + # Single object response + parsed_chunk = gemini_iterator.chunk_parser(chunk=dict_chunk) + if parsed_chunk is not None: + parsed_chunks.append(parsed_chunk) + + except json.JSONDecodeError as e: + verbose_proxy_logger.error(f"Failed to parse Gemini streaming response as JSON: {e}") + return [] + + verbose_proxy_logger.debug(f"Total parsed chunks: {len(parsed_chunks)}") + return parsed_chunks + + @staticmethod + def _build_complete_streaming_response( + all_chunks: List[str], + litellm_logging_obj: LiteLLMLoggingObj, + model: str, + url_route: str, + ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: + if "generateContent" not in url_route and "streamGenerateContent" not in url_route: + return None + + gemini_iterator: Any = GeminiModelResponseIterator( + streaming_response=None, + sync_stream=False, + logging_obj=litellm_logging_obj, + ) + + # Parse the streaming chunks + parsed_chunks = GeminiPassthroughLoggingHandler._parse_gemini_streaming_json( + all_chunks=all_chunks, + gemini_iterator=gemini_iterator, + ) + + all_openai_chunks = [] + for parsed_chunk in parsed_chunks: + if parsed_chunk is None: + continue + all_openai_chunks.append(parsed_chunk) + + complete_streaming_response = litellm.stream_chunk_builder( + chunks=all_openai_chunks, + logging_obj=litellm_logging_obj, + ) + + return complete_streaming_response + + @staticmethod + def extract_model_from_url(url: str) -> str: + pattern = r"/models/([^:]+)" + match = re.search(pattern, url) + if match: + return match.group(1) + return "unknown" + + @staticmethod + def _create_gemini_response_logging_payload_for_generate_content( + litellm_model_response: Union[ModelResponse, TextCompletionResponse], + model: str, + kwargs: dict, + start_time: datetime, + end_time: datetime, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str, + ): + """ + Create the standard logging object for Gemini passthrough generateContent (streaming and non-streaming) + """ + + response_cost = litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider=custom_llm_provider, + ) + + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = custom_llm_provider + + # pretty print standard logging object + verbose_proxy_logger.debug("kwargs= %s", json.dumps(kwargs, indent=4)) + + # set litellm_call_id to logging response object + litellm_model_response.id = logging_obj.litellm_call_id + logging_obj.model = litellm_model_response.model or model + logging_obj.model_call_details["model"] = logging_obj.model + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + logging_obj.model_call_details["response_cost"] = response_cost + return kwargs diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py new file mode 100644 index 00000000000..f8eb98affcf --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -0,0 +1,398 @@ +""" +Vertex AI Live API WebSocket Passthrough Logging Handler + +Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough endpoints. +Supports different modalities: text, audio, video, and web search. +""" + +from datetime import datetime +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( + BasePassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + PassThroughEndpointLoggingTypedDict, +) +from litellm.types.utils import LlmProviders, ModelResponse, Usage +from litellm.utils import get_model_info + + +class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): + """ + Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. + + Supports: + - Text tokens (input/output) + - Audio tokens (input/output) + - Video tokens (input/output) + - Web search requests + - Tool use tokens + """ + + def _build_complete_streaming_response(self, *args, **kwargs): + """Not applicable for WebSocket passthrough.""" + return None + + def get_provider_config(self, model: str): + """Return Vertex AI provider configuration.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + return VertexGeminiConfig() + + @property + def llm_provider_name(self) -> LlmProviders: + """Return the LLM provider name.""" + return LlmProviders.VERTEX_AI + + @staticmethod + def _extract_usage_metadata_from_websocket_messages( + websocket_messages: List[Dict], + ) -> Optional[Dict]: + """ + Extract and aggregate usage metadata from a list of WebSocket messages. + + Args: + websocket_messages: List of WebSocket messages from the Live API + + Returns: + Dictionary containing aggregated usage metadata, or None if not found + """ + all_usage_metadata = [] + + # Collect all usage metadata messages + for message in websocket_messages: + if isinstance(message, dict) and "usageMetadata" in message: + all_usage_metadata.append(message["usageMetadata"]) + + if not all_usage_metadata: + return None + + # If only one usage metadata, return it as-is + if len(all_usage_metadata) == 1: + return all_usage_metadata[0] + + # Aggregate multiple usage metadata messages + aggregated: Dict[str, Any] = { + "promptTokenCount": 0, + "candidatesTokenCount": 0, + "totalTokenCount": 0, + "promptTokensDetails": [], + "candidatesTokensDetails": [], + } + + # Aggregate token counts + for usage in all_usage_metadata: + aggregated["promptTokenCount"] += usage.get("promptTokenCount", 0) + aggregated["candidatesTokenCount"] += usage.get("candidatesTokenCount", 0) + aggregated["totalTokenCount"] += usage.get("totalTokenCount", 0) + + # Aggregate token details by modality + modality_totals = {} + + for usage in all_usage_metadata: + # Process prompt tokens details + for detail in usage.get("promptTokensDetails", []): + modality = detail.get("modality", "TEXT") + token_count = detail.get("tokenCount", 0) + + if modality not in modality_totals: + modality_totals[modality] = {"prompt": 0, "candidate": 0} + modality_totals[modality]["prompt"] += token_count + + # Process candidate tokens details + for detail in usage.get("candidatesTokensDetails", []): + modality = detail.get("modality", "TEXT") + token_count = detail.get("tokenCount", 0) + + if modality not in modality_totals: + modality_totals[modality] = {"prompt": 0, "candidate": 0} + modality_totals[modality]["candidate"] += token_count + + # Convert aggregated modality totals back to details format + for modality, totals in modality_totals.items(): + if totals["prompt"] > 0: + aggregated["promptTokensDetails"].append( + {"modality": modality, "tokenCount": totals["prompt"]} + ) + if totals["candidate"] > 0: + aggregated["candidatesTokensDetails"].append( + {"modality": modality, "tokenCount": totals["candidate"]} + ) + + # Add any additional fields from the first usage metadata + first_usage = all_usage_metadata[0] + for key, value in first_usage.items(): + if key not in aggregated: + aggregated[key] = value + + return aggregated + + @staticmethod + def _calculate_live_api_cost( + model: str, + usage_metadata: Dict, + custom_llm_provider: str = "vertex_ai", + ) -> float: + """ + Calculate cost for Vertex AI Live API based on usage metadata. + + Args: + model: The model name (e.g., "gemini-2.0-flash-live-preview-04-09") + usage_metadata: Usage metadata from the Live API response + custom_llm_provider: The LLM provider (default: "vertex_ai") + + Returns: + Total cost in USD + """ + try: + # Get model pricing information + model_info = get_model_info( + model=model, custom_llm_provider=custom_llm_provider + ) + + verbose_proxy_logger.debug( + f"Vertex AI Live API model info for '{model}': {model_info}" + ) + + # Check if pricing info is available + if not model_info or not model_info.get("input_cost_per_token"): + verbose_proxy_logger.error( + f"No pricing info found for {model} in local model pricing database" + ) + return 0.0 + + total_cost = 0.0 + + # Extract token counts from usage metadata + prompt_token_count = usage_metadata.get("promptTokenCount", 0) + candidates_token_count = usage_metadata.get("candidatesTokenCount", 0) + + # Calculate base text token costs + input_cost_per_token = model_info.get("input_cost_per_token", 0.0) + output_cost_per_token = model_info.get("output_cost_per_token", 0.0) + + total_cost += prompt_token_count * input_cost_per_token + total_cost += candidates_token_count * output_cost_per_token + + # Handle modality-specific costs if present + prompt_tokens_details = usage_metadata.get("promptTokensDetails", []) + candidates_tokens_details = usage_metadata.get( + "candidatesTokensDetails", [] + ) + + # Process prompt tokens by modality + for detail in prompt_tokens_details: + modality = detail.get("modality", "TEXT") + token_count = detail.get("tokenCount", 0) + + if modality == "AUDIO": + audio_cost_per_token = model_info.get( + "input_cost_per_audio_token", 0.0 + ) + total_cost += token_count * audio_cost_per_token + elif modality == "VIDEO": + # Video tokens are typically per second, but we'll treat as per token for now + video_cost_per_token = model_info.get( + "input_cost_per_video_per_second", 0.0 + ) + total_cost += token_count * video_cost_per_token + # TEXT tokens are already handled above + + # Process candidate tokens by modality + for detail in candidates_tokens_details: + modality = detail.get("modality", "TEXT") + token_count = detail.get("tokenCount", 0) + + if modality == "AUDIO": + audio_cost_per_token = model_info.get( + "output_cost_per_audio_token", 0.0 + ) + total_cost += token_count * audio_cost_per_token + elif modality == "VIDEO": + # Video tokens are typically per second, but we'll treat as per token for now + video_cost_per_token = model_info.get( + "output_cost_per_video_per_second", 0.0 + ) + total_cost += token_count * video_cost_per_token + # TEXT tokens are already handled above + + # Handle web search costs if present + tool_use_prompt_token_count = usage_metadata.get( + "toolUsePromptTokenCount", 0 + ) + if tool_use_prompt_token_count > 0: + # Web search typically has a fixed cost per request + web_search_cost = model_info.get("web_search_cost_per_request", 0.0) + if isinstance(web_search_cost, (int, float)) and web_search_cost > 0: + total_cost += web_search_cost + else: + # Fallback to token-based pricing for tool use + total_cost += tool_use_prompt_token_count * input_cost_per_token + + verbose_proxy_logger.debug( + f"Vertex AI Live API cost calculation - Model: {model}, " + f"Prompt tokens: {prompt_token_count}, " + f"Candidate tokens: {candidates_token_count}, " + f"Total cost: ${total_cost:.6f}" + ) + + return total_cost + + except Exception as e: + verbose_proxy_logger.error( + f"Error calculating Vertex AI Live API cost: {e}" + ) + return 0.0 + + @staticmethod + def _create_usage_object_from_metadata( + usage_metadata: Dict, + model: str, + ) -> Usage: + """ + Create a LiteLLM Usage object from Live API usage metadata. + + Args: + usage_metadata: Usage metadata from the Live API response + model: The model name + + Returns: + LiteLLM Usage object + """ + prompt_tokens = usage_metadata.get("promptTokenCount", 0) + completion_tokens = usage_metadata.get("candidatesTokenCount", 0) + total_tokens = usage_metadata.get("totalTokenCount", 0) + + # Create modality-specific token details if available + prompt_tokens_details = usage_metadata.get("promptTokensDetails", []) + candidates_tokens_details = usage_metadata.get("candidatesTokensDetails", []) + + # Extract text tokens from details + text_prompt_tokens = 0 + text_completion_tokens = 0 + + for detail in prompt_tokens_details: + if detail.get("modality") == "TEXT": + text_prompt_tokens = detail.get("tokenCount", 0) + break + + for detail in candidates_tokens_details: + if detail.get("modality") == "TEXT": + text_completion_tokens = detail.get("tokenCount", 0) + break + + # If no text tokens found in details, use total counts + if text_prompt_tokens == 0: + text_prompt_tokens = prompt_tokens + if text_completion_tokens == 0: + text_completion_tokens = completion_tokens + + return Usage( + prompt_tokens=text_prompt_tokens, + completion_tokens=text_completion_tokens, + total_tokens=total_tokens, + ) + + def vertex_ai_live_passthrough_handler( + self, + websocket_messages: List[Dict], + logging_obj, + url_route: str, + start_time: datetime, + end_time: datetime, + request_body: dict, + **kwargs, + ) -> PassThroughEndpointLoggingTypedDict: + """ + Handle cost tracking and logging for Vertex AI Live API WebSocket passthrough. + + Args: + websocket_messages: List of WebSocket messages from the Live API + logging_obj: LiteLLM logging object + url_route: The URL route that was called + start_time: Request start time + end_time: Request end time + request_body: The original request body + **kwargs: Additional keyword arguments + + Returns: + Dictionary containing the result and kwargs for logging + """ + try: + # Extract model from request body or kwargs + model = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09") + custom_llm_provider = kwargs.get("custom_llm_provider", "vertex_ai") + verbose_proxy_logger.debug( + f"Vertex AI Live API model: {model}, custom_llm_provider: {custom_llm_provider}" + ) + + # Extract usage metadata from WebSocket messages + usage_metadata = self._extract_usage_metadata_from_websocket_messages( + websocket_messages + ) + + if not usage_metadata: + verbose_proxy_logger.warning( + "No usage metadata found in Vertex AI Live API WebSocket messages" + ) + return { + "result": None, + "kwargs": kwargs, + } + + # Calculate cost using Live API specific pricing + response_cost = self._calculate_live_api_cost( + model=model, + usage_metadata=usage_metadata, + custom_llm_provider=custom_llm_provider, + ) + + # Create Usage object for standard LiteLLM logging + usage = self._create_usage_object_from_metadata( + usage_metadata=usage_metadata, + model=model, + ) + + # Create a mock ModelResponse for standard logging + litellm_model_response = ModelResponse( + id=f"vertex-ai-live-{start_time.timestamp()}", + object="chat.completion", + created=int(start_time.timestamp()), + model=model, + usage=usage, + choices=[], + ) + + # Update kwargs with cost information + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = custom_llm_provider + + # Safely log the model name: only allow known safe formats, redact otherwise. + import re + allowed_pattern = re.compile(r"^[A-Za-z0-9._\-:]+$") + safe_model = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" + verbose_proxy_logger.debug( + f"Vertex AI Live API passthrough cost tracking - " + f"Model: {safe_model}, Cost: ${response_cost:.6f}, " + f"Prompt tokens: {usage.prompt_tokens}, " + f"Completion tokens: {usage.completion_tokens}" + ) + + return { + "result": litellm_model_response, + "kwargs": kwargs, + } + + except Exception as e: + verbose_proxy_logger.error( + f"Error in Vertex AI Live API passthrough handler: {e}" + ) + return { + "result": None, + "kwargs": kwargs, + } diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 0eacee3b4f1..f41ee21db80 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -5,7 +5,7 @@ import json import traceback from base64 import b64encode from datetime import datetime -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union from urllib.parse import urlencode, urlparse import httpx @@ -17,10 +17,18 @@ from fastapi import ( Request, Response, UploadFile, + WebSocket, status, ) from fastapi.responses import StreamingResponse from starlette.datastructures import UploadFile as StarletteUploadFile +from starlette.websockets import WebSocketState +from websockets.asyncio.client import connect +from websockets.exceptions import ( + ConnectionClosedError, + ConnectionClosedOK, + InvalidStatus, +) import litellm from litellm._logging import verbose_proxy_logger @@ -76,7 +84,7 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona example header can be - {"Authorization": "bearer os.environ/COHERE_API_KEY"} + {"Authorization": "Bearer os.environ/COHERE_API_KEY"} """ if custom_headers is None: return None @@ -88,13 +96,9 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona # langfuse requires b64 encoded headers - we construct that here _langfuse_public_key = custom_headers["LANGFUSE_PUBLIC_KEY"] _langfuse_secret_key = custom_headers["LANGFUSE_SECRET_KEY"] - if isinstance( - _langfuse_public_key, str - ) and _langfuse_public_key.startswith("os.environ/"): + if isinstance(_langfuse_public_key, str) and _langfuse_public_key.startswith("os.environ/"): _langfuse_public_key = get_secret_str(_langfuse_public_key) - if isinstance( - _langfuse_secret_key, str - ) and _langfuse_secret_key.startswith("os.environ/"): + if isinstance(_langfuse_secret_key, str) and _langfuse_secret_key.startswith("os.environ/"): _langfuse_secret_key = get_secret_str(_langfuse_secret_key) headers["Authorization"] = "Basic " + b64encode( f"{_langfuse_public_key}:{_langfuse_secret_key}".encode("utf-8") @@ -103,9 +107,7 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona # for all other headers headers[key] = value if isinstance(value, str) and "os.environ/" in value: - verbose_proxy_logger.debug( - "pass through endpoint - looking up 'os.environ/' variable" - ) + verbose_proxy_logger.debug("pass through endpoint - looking up 'os.environ/' variable") # get string section that is os.environ/ start_index = value.find("os.environ/") _variable_name = value[start_index:] @@ -198,9 +200,7 @@ async def chat_completion_pass_through_endpoint( # noqa: PLR0915 # skip router if user passed their key if "api_key" in data: llm_response = asyncio.create_task(litellm.aadapter_completion(**data)) - elif ( - llm_router is not None and data["model"] in router_model_names - ): # model in router model list + elif llm_router is not None and data["model"] in router_model_names: # model in router model list llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) elif ( llm_router is not None @@ -229,10 +229,7 @@ async def chat_completion_pass_through_endpoint( # noqa: PLR0915 else: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "completion: Invalid model name passed in model=" - + data.get("model", "") - }, + detail={"error": "completion: Invalid model name passed in model=" + data.get("model", "")}, ) # Await the llm_response task @@ -246,9 +243,7 @@ async def chat_completion_pass_through_endpoint( # noqa: PLR0915 ### ALERTING ### asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) + proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") ) verbose_proxy_logger.debug("final response: %s", response) @@ -270,11 +265,7 @@ async def chat_completion_pass_through_endpoint( # noqa: PLR0915 await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.completion(): Exception occured - {}".format( - str(e) - ) - ) + verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - {}".format(str(e))) error_msg = f"{str(e)}" raise ProxyException( message=getattr(e, "message", error_msg), @@ -293,11 +284,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): ) -> dict: excluded_headers = {"transfer-encoding", "content-encoding"} - return_headers = { - key: value - for key, value in headers.items() - if key.lower() not in excluded_headers - } + return_headers = {key: value for key, value in headers.items() if key.lower() not in excluded_headers} if litellm_call_id: return_headers["x-litellm-call-id"] = litellm_call_id if custom_headers: @@ -314,6 +301,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): or ("rawPredict") in url or ("streamRawPredict") in url ): + # Check if it's Gemini (Google AI Studio) or Vertex AI + if parsed_url.hostname and parsed_url.hostname.endswith("generativelanguage.googleapis.com"): + return EndpointType.GEMINI return EndpointType.VERTEX_AI elif parsed_url.hostname == "api.anthropic.com": return EndpointType.ANTHROPIC @@ -424,10 +414,8 @@ 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 @@ -477,9 +465,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): user_api_key_spend=user_api_key_dict.spend, user_api_key_max_budget=user_api_key_dict.max_budget, user_api_key_budget_reset_at=( - user_api_key_dict.budget_reset_at.isoformat() - if user_api_key_dict.budget_reset_at - else None + user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None ), ) ) @@ -513,16 +499,12 @@ 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 @staticmethod - def construct_target_url_with_subpath( - base_target: str, subpath: str, include_subpath: Optional[bool] - ) -> str: + def construct_target_url_with_subpath(base_target: str, subpath: str, include_subpath: Optional[bool]) -> str: """ Helper function to construct the full target URL with subpath handling. @@ -573,6 +555,7 @@ async def pass_through_request( # noqa: PLR0915 query_params: Optional[dict] = None, stream: Optional[bool] = None, cost_per_request: Optional[float] = None, + custom_llm_provider: Optional[str] = None, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -624,9 +607,7 @@ async def pass_through_request( # noqa: PLR0915 ).encode("ascii") ) - endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type( - str(url) - ) + endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type(str(url)) if custom_body: _parsed_body = custom_body @@ -688,16 +669,12 @@ async def pass_through_request( # noqa: PLR0915 # combine url with query params for logging requested_query_params: Optional[dict] = ( - query_params or request.query_params.__dict__ + query_params or dict(request.query_params) ) - if requested_query_params == request.query_params.__dict__: - requested_query_params = None requested_query_params_str = None if requested_query_params: - requested_query_params_str = "&".join( - f"{k}={v}" for k, v in requested_query_params.items() - ) + requested_query_params_str = "&".join(f"{k}={v}" for k, v in requested_query_params.items()) logging_url = str(url) if requested_query_params_str: @@ -715,11 +692,9 @@ 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: @@ -736,9 +711,7 @@ async def pass_through_request( # noqa: PLR0915 try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=await e.response.aread() - ) + raise HTTPException(status_code=e.response.status_code, detail=await e.response.aread()) return StreamingResponse( PassThroughStreamingHandler.chunk_processor( @@ -760,20 +733,16 @@ async def pass_through_request( # noqa: PLR0915 verbose_proxy_logger.debug("request method: {}".format(request.method)) verbose_proxy_logger.debug("request url: {}".format(url)) verbose_proxy_logger.debug("request headers: {}".format(headers)) - verbose_proxy_logger.debug( - "requested_query_params={}".format(requested_query_params) - ) + verbose_proxy_logger.debug("requested_query_params={}".format(requested_query_params)) verbose_proxy_logger.debug("request body: {}".format(_parsed_body)) - response = ( - await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler( - request=request, - async_client=async_client, - url=url, - headers=headers, - requested_query_params=requested_query_params, - _parsed_body=_parsed_body, - ) + response = await HttpPassThroughEndpointHelpers.non_streaming_http_request_handler( + request=request, + async_client=async_client, + url=url, + headers=headers, + requested_query_params=requested_query_params, + _parsed_body=_parsed_body, ) verbose_proxy_logger.debug("response.headers= %s", response.headers) @@ -781,9 +750,7 @@ async def pass_through_request( # noqa: PLR0915 try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=await e.response.aread() - ) + raise HTTPException(status_code=e.response.status_code, detail=await e.response.aread()) return StreamingResponse( PassThroughStreamingHandler.chunk_processor( @@ -805,9 +772,7 @@ async def pass_through_request( # noqa: PLR0915 try: response.raise_for_status() except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=e.response.text - ) + raise HTTPException(status_code=e.response.status_code, detail=e.response.text) if response.status_code >= 300: raise HTTPException(status_code=response.status_code, detail=response.text) @@ -829,6 +794,7 @@ async def pass_through_request( # noqa: PLR0915 logging_obj=logging_obj, cache_hit=False, request_body=_parsed_body, + custom_llm_provider=custom_llm_provider, **kwargs, ) ) @@ -859,9 +825,7 @@ async def pass_through_request( # noqa: PLR0915 api_base=str(url._uri_reference) if url else None, ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {}".format( - str(e) - ) + "litellm.proxy.proxy_server.pass_through_endpoint(): Exception occured - {}".format(str(e)) ) ######################################################### @@ -924,6 +888,7 @@ def create_pass_through_route( dependencies: Optional[List] = None, include_subpath: Optional[bool] = False, cost_per_request: Optional[float] = None, + custom_llm_provider: Optional[str] = None, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -959,16 +924,12 @@ def create_pass_through_route( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), query_params: Optional[dict] = None, custom_body: Optional[dict] = None, - stream: Optional[ - bool - ] = None, # if pass-through endpoint is a streaming request + stream: Optional[bool] = None, # if pass-through endpoint is a streaming request subpath: str = "", # captures sub-paths when include_subpath=True ): # Construct the full target URL with subpath if needed - full_target = ( - HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( - base_target=target, subpath=subpath, include_subpath=include_subpath - ) + full_target = HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( + base_target=target, subpath=subpath, include_subpath=include_subpath ) return await pass_through_request( # type: ignore @@ -982,11 +943,515 @@ def create_pass_through_route( stream=stream, custom_body=custom_body, cost_per_request=cost_per_request, + custom_llm_provider=custom_llm_provider, ) return endpoint_func +def create_websocket_passthrough_route( + endpoint: str, + target: str, + custom_headers: Optional[dict] = None, + _forward_headers: Optional[bool] = False, + dependencies: Optional[List] = None, + cost_per_request: Optional[float] = None, +): + """ + Create a WebSocket passthrough route function. + + Args: + endpoint: The endpoint path (for logging purposes) + target: The target WebSocket URL (e.g., "wss://api.example.com/ws") + custom_headers: Custom headers to include in the WebSocket connection + _forward_headers: Whether to forward incoming headers + dependencies: FastAPI dependencies to inject + + Returns: + A WebSocket passthrough function that can be registered with app.websocket() + """ + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket + + async def websocket_endpoint_func( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), + **kwargs, # For additional query parameters + ): + """ + WebSocket passthrough endpoint function. + + This function handles the WebSocket connection by: + 1. Accepting the incoming WebSocket connection + 2. Establishing a connection to the target WebSocket + 3. Forwarding messages bidirectionally + 4. Handling connection cleanup + """ + return await websocket_passthrough_request( + websocket=websocket, + target=target, + custom_headers=custom_headers or {}, + user_api_key_dict=user_api_key_dict, + forward_headers=_forward_headers, + endpoint=endpoint, + cost_per_request=cost_per_request, + accept_websocket=True, # Generic usage should accept the WebSocket + ) + + return websocket_endpoint_func + + +async def websocket_passthrough_request( # noqa: PLR0915 + websocket: WebSocket, + target: str, + custom_headers: dict, + user_api_key_dict: UserAPIKeyAuth, + forward_headers: Optional[bool] = False, + endpoint: Optional[str] = None, + cost_per_request: Optional[float] = None, + accept_websocket: bool = True, +): + """ + WebSocket passthrough request handler. + + Args: + websocket: The incoming WebSocket connection + target: The target WebSocket URL + custom_headers: Custom headers to include in the connection + user_api_key_dict: The user API key dictionary + forward_headers: Whether to forward incoming headers + endpoint: The endpoint path (for logging purposes) + cost_per_request: Optional field - cost per request to the target endpoint + """ + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.proxy_server import proxy_logging_obj + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, + ) + + # Initialize tracking variables + start_time = datetime.now() + websocket_messages: list[dict[str, Any]] = [] + litellm_call_id = str(uuid.uuid4()) + + verbose_proxy_logger.info( + f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}" + ) + + # Only accept the WebSocket if requested (for generic usage) + if accept_websocket: + await websocket.accept() + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): WebSocket connection accepted" + ) + + # Prepare headers for the upstream connection + upstream_headers = custom_headers.copy() + + if forward_headers: + # Forward relevant headers from the incoming request + incoming_headers = dict(websocket.headers) + for header_name, header_value in incoming_headers.items(): + # Only forward certain headers to avoid conflicts + if header_name.lower() in [ + "authorization", + "x-api-key", + "x-goog-user-project", + ]: + upstream_headers[header_name] = header_value + + # Initialize logging object similar to HTTP passthrough + logging_obj = Logging( + model="unknown", + messages=[{"role": "user", "content": "WebSocket connection"}], + stream=True, # WebSockets are inherently streaming + call_type="pass_through_endpoint", + start_time=start_time, + litellm_call_id=litellm_call_id, + function_id="websocket_passthrough", + ) + + # Create passthrough logging payload + passthrough_logging_payload = PassthroughStandardLoggingPayload( + url=target, + request_body={}, # WebSocket doesn't have a traditional request body + request_method="WEBSOCKET", + cost_per_request=cost_per_request, + ) + + # Create a dummy request object for WebSocket connections to maintain compatibility + # with the existing _init_kwargs_for_pass_through_endpoint function + class DummyRequest: + def __init__(self, url: str, method: str = "WEBSOCKET", headers: Optional[dict] = None): + self.url = url + self.method = method + self.headers = headers or {} + + def __str__(self): + return f"DummyRequest(url={self.url}, method={self.method})" + + dummy_request = DummyRequest( + url=target, + method="WEBSOCKET", + headers=dict(websocket.headers) if hasattr(websocket, "headers") else {}, + ) + + # Initialize kwargs for logging using the same pattern as HTTP passthrough + kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + user_api_key_dict=user_api_key_dict, + _parsed_body={}, # WebSocket doesn't have a traditional request body + passthrough_logging_payload=passthrough_logging_payload, + litellm_call_id=litellm_call_id, + request=dummy_request, # type: ignore + logging_obj=logging_obj, + ) + + # Update logging environment variables + logging_obj.update_environment_variables( + model="unknown", + user="unknown", + optional_params={}, + litellm_params=dict(kwargs.get("litellm_params", {})), + call_type="pass_through_endpoint", + ) + logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + + # Pre-call logging + logging_obj.pre_call( + input=[{"role": "user", "content": "WebSocket connection"}], + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": target, + "headers": upstream_headers, + }, + ) + + ### CALL HOOKS ### - modify incoming data / reject request before calling the model + websocket_data: dict[str, Any] = {} + websocket_data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=websocket_data, + call_type="pass_through_endpoint", + ) + + try: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Establishing upstream connection to {target}" + ) + async with connect( + target, + additional_headers=upstream_headers, + ) as upstream_ws: + verbose_proxy_logger.info( + f"WebSocket passthrough ({endpoint}): Upstream connection established successfully" + ) + + async def forward_client_to_upstream() -> None: + """Forward messages from client to upstream WebSocket""" + try: + while True: + message = await websocket.receive() + message_type = message.get("type") + if message_type == "websocket.disconnect": + await upstream_ws.close() + break + + text_data = message.get("text") + bytes_data = message.get("bytes") + + if text_data is not None: + # Try to extract model from client setup message for Vertex AI Live + if endpoint and "/vertex_ai/live" in endpoint: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Processing client message for model extraction" + ) + try: + client_message = json.loads(text_data) + if ( + isinstance(client_message, dict) + and "setup" in client_message + ): + setup_data = client_message["setup"] + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Found setup data in client message: {setup_data}" + ) + if ( + isinstance(setup_data, dict) + and "model" in setup_data + ): + extracted_model = ( + _extract_model_from_vertex_ai_setup( + setup_data + ) + ) + if extracted_model: + kwargs["model"] = extracted_model + kwargs[ + "custom_llm_provider" + ] = "vertex_ai-language-models" + # Update logging object with correct model + logging_obj.model = extracted_model + logging_obj.model_call_details[ + "model" + ] = extracted_model + logging_obj.model_call_details[ + "custom_llm_provider" + ] = "vertex_ai" + verbose_proxy_logger.info( + f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from client setup message" + ) + else: + verbose_proxy_logger.warning( + f"WebSocket passthrough ({endpoint}): Failed to extract model from client setup data: {setup_data}" + ) + else: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Setup data does not contain model field: {setup_data}" + ) + else: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Client message does not contain setup data" + ) + except (json.JSONDecodeError, KeyError, TypeError) as e: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Client message is not a valid setup message: {e}" + ) + pass # Not a JSON message or doesn't contain setup data + + await upstream_ws.send(text_data) + elif bytes_data is not None: + await upstream_ws.send(bytes_data) + except asyncio.CancelledError: + raise + except Exception: + verbose_proxy_logger.exception( + f"WebSocket passthrough ({endpoint}): error forwarding client message" + ) + await upstream_ws.close() + + async def forward_upstream_to_client() -> None: + """Forward messages from upstream to client WebSocket""" + try: + # Wait for the first response from upstream + raw_response = await upstream_ws.recv(decode=False) + # Ensure raw_response is bytes before decoding + if isinstance(raw_response, str): + raw_response = raw_response.encode("ascii") + setup_response = json.loads(raw_response.decode("ascii")) + verbose_proxy_logger.debug(f"Setup response: {setup_response}") + + # Extract model and provider from setup response for Vertex AI Live + if endpoint and "/vertex_ai/live" in endpoint: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Processing server setup response for model extraction" + ) + extracted_model = _extract_model_from_vertex_ai_setup( + setup_response + ) + if extracted_model: + kwargs["model"] = extracted_model + kwargs["custom_llm_provider"] = "vertex_ai_language_models" + # Update logging object with correct model + logging_obj.model = extracted_model + logging_obj.model_call_details["model"] = extracted_model + logging_obj.model_call_details[ + "custom_llm_provider" + ] = "vertex_ai_language_models" + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from server setup response" + ) + else: + verbose_proxy_logger.warning( + f"WebSocket passthrough ({endpoint}): Failed to extract model from server setup response: {setup_response}" + ) + else: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Not a Vertex AI Live endpoint, skipping model extraction" + ) + + # Send the setup response to the client + await websocket.send_text(json.dumps(setup_response)) + + # Now continuously forward messages from upstream to client + async for upstream_message in upstream_ws: + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + # Parse and collect for cost tracking + try: + message_data = json.loads(upstream_message.decode()) + websocket_messages.append(message_data) + except (json.JSONDecodeError, UnicodeDecodeError): + pass + else: + await websocket.send_text(upstream_message) + # Parse and collect for cost tracking + try: + message_data = json.loads(upstream_message) + websocket_messages.append(message_data) + except json.JSONDecodeError: + pass + + except (ConnectionClosedOK, ConnectionClosedError) as e: + verbose_proxy_logger.debug( + f"Upstream WebSocket connection closed: {e}" + ) + pass + except asyncio.CancelledError: + verbose_proxy_logger.debug( + "asyncio.CancelledError in forward_upstream_to_client" + ) + raise + except Exception as e: + verbose_proxy_logger.debug( + f"Exception in forward_upstream_to_client: {e}" + ) + verbose_proxy_logger.exception( + f"WebSocket passthrough ({endpoint}): error forwarding upstream message" + ) + raise + + # Create tasks for bidirectional message forwarding + tasks = [ + asyncio.create_task(forward_client_to_upstream()), + asyncio.create_task(forward_upstream_to_client()), + ] + + done, pending = await asyncio.wait( + tasks, return_when=asyncio.FIRST_COMPLETED + ) + + # Cancel remaining tasks + for task in pending: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Check for exceptions in completed tasks + for task in done: + exception = task.exception() + if exception is not None: + raise exception + + end_time = datetime.now() + + # Update passthrough logging payload with response data + passthrough_logging_payload["response_body"] = websocket_messages # type: ignore + passthrough_logging_payload["end_time"] = end_time # type: ignore + + # Remove logging_obj from kwargs to avoid duplicate keyword argument + success_kwargs = kwargs.copy() + success_kwargs.pop("logging_obj", None) + + # # Add user authentication context for database logging + # if user_api_key_dict: + # success_kwargs.setdefault('litellm_params', {}) + # success_kwargs['litellm_params'].update({ + # 'proxy_server_request': { + # 'body': { + # 'user': user_api_key_dict.user_id, + # 'team_id': user_api_key_dict.team_id, + # 'end_user_id': user_api_key_dict.end_user_id, + # } + # } + # }) + # # Also add the user_api_key for direct access + # success_kwargs['user_api_key'] = user_api_key_dict.api_key + + # Create a dummy httpx.Response for WebSocket connections + class MockWebSocketResponse: + def __init__(self, target_url: str): + self.status_code = 200 + self.text = "WebSocket connection successful" + self.headers: dict[str, str] = {} + self.request = MockWebSocketRequest(target_url) + + class MockWebSocketRequest: + def __init__(self, target_url: str): + self.method = "WEBSOCKET" + self.url = target_url + + mock_response = MockWebSocketResponse(target) + + # Use the same success handler as HTTP passthrough endpoints + asyncio.create_task( + pass_through_endpoint_logging.pass_through_async_success_handler( + httpx_response=mock_response, # type: ignore + response_body=websocket_messages, # type: ignore + url_route=endpoint or "", + result="websocket_connection_successful", + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + cache_hit=False, + request_body={}, + **success_kwargs, + ) + ) + + # Call the proxy logging success hook + if proxy_logging_obj: + await proxy_logging_obj.post_call_success_hook( + data={}, + user_api_key_dict=user_api_key_dict, + response={"status": "websocket_connection_successful"}, # type: ignore + ) + + except InvalidStatus as exc: + verbose_proxy_logger.exception( + f"WebSocket passthrough ({endpoint}): upstream rejected WebSocket connection" + ) + + # Prepare request payload for logging + request_payload = {} + if kwargs: + for key, value in kwargs.items(): + request_payload[key] = value + + # Log the connection failure using the same pattern as HTTP + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=exc, + request_data=request_payload, + traceback_str=traceback.format_exc( + limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, + ), + ) + + if websocket.client_state != WebSocketState.DISCONNECTED: + await websocket.close( + code=getattr(exc, "status_code", 1011), + reason="Upstream connection rejected", + ) + except Exception as e: + verbose_proxy_logger.exception( + f"WebSocket passthrough ({endpoint}): unexpected error while proxying WebSocket" + ) + + # Prepare request payload for logging + request_payload = {} + if kwargs: + for key, value in kwargs.items(): + request_payload[key] = value + + # Log the unexpected error using the same pattern as HTTP + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_payload, + traceback_str=traceback.format_exc( + limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, + ), + ) + + if websocket.client_state != WebSocketState.DISCONNECTED: + await websocket.close(code=1011, reason="WebSocket passthrough error") + finally: + if websocket.client_state != WebSocketState.DISCONNECTED: + await websocket.close() + + def _is_streaming_response(response: httpx.Response) -> bool: _content_type = response.headers.get("content-type") if _content_type is not None and "text/event-stream" in _content_type: @@ -994,6 +1459,38 @@ def _is_streaming_response(response: httpx.Response) -> bool: return False +def _extract_model_from_vertex_ai_setup(setup_response: dict) -> Optional[str]: + """ + Extract the model name from Vertex AI Live setup response. + + The setup response can contain a model field in two formats: + 1. Direct: {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"} + 2. Nested: {"setup": {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}} + + We extract just the model name: "gemini-2.0-flash-live-preview-04-09" + """ + try: + # Handle both direct model field and nested setup.model field + model_path = None + if isinstance(setup_response, dict): + if "model" in setup_response: + model_path = setup_response["model"] + elif ( + "setup" in setup_response + and isinstance(setup_response["setup"], dict) + and "model" in setup_response["setup"] + ): + model_path = setup_response["setup"]["model"] + + if isinstance(model_path, str) and "/models/" in model_path: + # Extract the model name after the last "/models/" + model_name = model_path.split("/models/")[-1] + return model_name + except Exception as e: + verbose_proxy_logger.debug(f"Error extracting model from setup response: {e}") + return None + + class InitPassThroughEndpointHelpers: @staticmethod def add_exact_path_route( @@ -1103,15 +1600,42 @@ class InitPassThroughEndpointHelpers: def remove_endpoint_routes(endpoint_id: str): """Remove all routes for a specific endpoint ID from the registry""" keys_to_remove = [ - key - for key, value in _registered_pass_through_routes.items() - if value["endpoint_id"] == endpoint_id + key for key, value in _registered_pass_through_routes.items() if value["endpoint_id"] == endpoint_id ] for key in keys_to_remove: del _registered_pass_through_routes[key] - verbose_proxy_logger.debug( - "Removed pass-through route from registry: %s", key - ) + verbose_proxy_logger.debug("Removed pass-through route from registry: %s", key) + + @staticmethod + def is_registered_pass_through_route(route: str) -> bool: + """ + Check if route is a registered pass-through endpoint from DB + + Uses the in-memory registry to avoid additional DB queries + Optimized for minimal latency + + Args: + route: The route to check + + Returns: + bool: True if route is a registered pass-through endpoint, False otherwise + """ + # Fast path: check if any registered route key contains this path + # Keys are in format: "{endpoint_id}:exact:{path}" or "{endpoint_id}:subpath:{path}" + # Extract unique paths from keys for quick checking + for key in _registered_pass_through_routes.keys(): + parts = key.split(":", 2) # Split into [endpoint_id, type, path] + if len(parts) == 3: + route_type = parts[1] + registered_path = parts[2] + + if route_type == "exact" and route == registered_path: + return True + elif route_type == "subpath": + if route == registered_path or route.startswith(registered_path + "/"): + return True + + return False async def initialize_pass_through_endpoints( @@ -1148,9 +1672,7 @@ async def initialize_pass_through_endpoints( if _path is None: raise ValueError("Path is required for pass-through endpoint") _custom_headers = endpoint.get("headers", None) - _custom_headers = await set_env_variables_in_header( - custom_headers=_custom_headers - ) + _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) _forward_headers = endpoint.get("forward_headers", None) _merge_query_params = endpoint.get("merge_query_params", None) _auth = endpoint.get("auth", None) @@ -1169,9 +1691,7 @@ async def initialize_pass_through_endpoints( continue # Add exact path route - verbose_proxy_logger.debug( - "Initializing pass through endpoint: %s (ID: %s)", _path, endpoint_id - ) + verbose_proxy_logger.debug("Initializing pass through endpoint: %s (ID: %s)", _path, endpoint_id) InitPassThroughEndpointHelpers.add_exact_path_route( app=app, path=_path, @@ -1198,9 +1718,7 @@ async def initialize_pass_through_endpoints( endpoint_id=endpoint_id, ) - verbose_proxy_logger.debug( - "Added new pass through endpoint: %s (ID: %s)", _path, endpoint_id - ) + verbose_proxy_logger.debug("Added new pass through endpoint: %s (ID: %s)", _path, endpoint_id) async def _get_pass_through_endpoints_from_db( @@ -1304,11 +1822,7 @@ async def update_pass_through_endpoints( # Find the index for updating the list endpoint_index = None for idx, endpoint in enumerate(pass_through_endpoint_data): - _endpoint = ( - PassThroughGenericEndpoint(**endpoint) - if isinstance(endpoint, dict) - else endpoint - ) + _endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint if _endpoint.id == endpoint_id: endpoint_index = idx break @@ -1316,9 +1830,7 @@ async def update_pass_through_endpoints( if endpoint_index is None: raise HTTPException( status_code=404, - detail={ - "error": f"Could not find index for endpoint with ID '{endpoint_id}'" - }, + detail={"error": f"Could not find index for endpoint with ID '{endpoint_id}'"}, ) # Get the update data as dict, excluding None values for partial updates @@ -1349,13 +1861,9 @@ async def update_pass_through_endpoints( field_value=pass_through_endpoint_data, config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) - return PassThroughEndpointResponse( - endpoints=[updated_endpoint] if updated_endpoint else [] - ) + return PassThroughEndpointResponse(endpoints=[updated_endpoint] if updated_endpoint else []) @router.post( @@ -1382,9 +1890,7 @@ async def create_pass_through_endpoints( field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict ) except Exception: - response = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=None - ) + response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Auto-generate ID if not provided data_dict = data.model_dump() @@ -1402,9 +1908,7 @@ async def create_pass_through_endpoints( field_value=response.field_value, config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) # Return the created endpoint with the generated ID created_endpoint = PassThroughGenericEndpoint(**data_dict) @@ -1437,9 +1941,7 @@ async def delete_pass_through_endpoints( field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict ) except Exception: - response = ConfigFieldInfo( - field_name="pass_through_endpoints", field_value=None - ) + response = ConfigFieldInfo(field_name="pass_through_endpoints", field_value=None) ## Update field by removing endpoint pass_through_endpoint_data: Optional[List] = response.field_value @@ -1455,21 +1957,13 @@ async def delete_pass_through_endpoints( if found_endpoint is None: raise HTTPException( status_code=400, - detail={ - "error": "Endpoint with ID '{}' was not found in pass-through endpoint list.".format( - endpoint_id - ) - }, + detail={"error": "Endpoint with ID '{}' was not found in pass-through endpoint list.".format(endpoint_id)}, ) # Find the index for deleting from the list endpoint_index = None for idx, endpoint in enumerate(pass_through_endpoint_data): - _endpoint = ( - PassThroughGenericEndpoint(**endpoint) - if isinstance(endpoint, dict) - else endpoint - ) + _endpoint = PassThroughGenericEndpoint(**endpoint) if isinstance(endpoint, dict) else endpoint if _endpoint.id == endpoint_id: endpoint_index = idx break @@ -1477,9 +1971,7 @@ async def delete_pass_through_endpoints( if endpoint_index is None: raise HTTPException( status_code=400, - detail={ - "error": f"Could not find index for endpoint with ID '{endpoint_id}'" - }, + detail={"error": f"Could not find index for endpoint with ID '{endpoint_id}'"}, ) # Remove the endpoint @@ -1495,9 +1987,7 @@ async def delete_pass_through_endpoints( field_value=pass_through_endpoint_data, config_type="general_settings", ) - await update_config_general_settings( - data=updated_data, user_api_key_dict=user_api_key_dict - ) + await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) return PassThroughEndpointResponse(endpoints=[response_obj]) @@ -1535,6 +2025,4 @@ async def initialize_pass_through_endpoints_in_db(): Gets all pass-through endpoints from db and initializes them in the proxy server. """ pass_through_endpoints = await _get_pass_through_endpoints_from_db() - await initialize_pass_through_endpoints( - pass_through_endpoints=pass_through_endpoints - ) + await initialize_pass_through_endpoints(pass_through_endpoints=pass_through_endpoints) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 2d5b0a686ce..792b496a664 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -105,6 +105,28 @@ class PassThroughStreamingHandler: anthropic_passthrough_logging_handler_result["result"] ) kwargs = anthropic_passthrough_logging_handler_result["kwargs"] + elif endpoint_type == EndpointType.GEMINI: + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler import ( + GeminiPassthroughLoggingHandler, + ) + + gemini_passthrough_logging_handler_result = ( + GeminiPassthroughLoggingHandler._handle_logging_gemini_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, + model=model, + ) + ) + standard_logging_response_object = ( + gemini_passthrough_logging_handler_result["result"] + ) + kwargs = gemini_passthrough_logging_handler_result["kwargs"] elif endpoint_type == EndpointType.VERTEX_AI: vertex_passthrough_logging_handler_result = ( VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 58fda370d93..a819c429f10 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -25,6 +25,9 @@ from .llm_provider_handlers.cohere_passthrough_logging_handler import ( from .llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) +from .llm_provider_handlers.gemini_passthrough_logging_handler import ( + GeminiPassthroughLoggingHandler, +) cohere_passthrough_logging_handler = CoherePassthroughLoggingHandler() @@ -44,13 +47,17 @@ class PassThroughEndpointLogging: # Cohere self.TRACKED_COHERE_ROUTES = ["/v2/chat"] - self.assemblyai_passthrough_logging_handler = ( - AssemblyAIPassthroughLoggingHandler() - ) + self.assemblyai_passthrough_logging_handler = AssemblyAIPassthroughLoggingHandler() # Langfuse self.TRACKED_LANGFUSE_ROUTES = ["/langfuse/"] + # Gemini + self.TRACKED_GEMINI_ROUTES = ["generateContent", "streamGenerateContent"] + + # Vertex AI Live API WebSocket + self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"] + async def _handle_logging( self, logging_obj: LiteLLMLoggingObj, @@ -78,11 +85,7 @@ class PassThroughEndpointLogging: # Handle async logging await logging_obj.async_success_handler( - result=( - json.dumps(result) - if isinstance(result, dict) - else standard_logging_response_object - ), + result=(json.dumps(result) if isinstance(result, dict) else standard_logging_response_object), start_time=start_time, end_time=end_time, cache_hit=False, @@ -100,6 +103,7 @@ class PassThroughEndpointLogging: start_time: datetime, end_time: datetime, cache_hit: bool, + custom_llm_provider: Optional[str] = None, **kwargs, ): return_dict = { @@ -107,22 +111,34 @@ class PassThroughEndpointLogging: "kwargs": kwargs, } standard_logging_response_object: Optional[Any] = None - if self.is_vertex_route(url_route): - vertex_passthrough_logging_handler_result = ( - VertexPassthroughLoggingHandler.vertex_passthrough_handler( - httpx_response=httpx_response, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - **kwargs, - ) + + if self.is_gemini_route(url_route, custom_llm_provider): + gemini_passthrough_logging_handler_result = GeminiPassthroughLoggingHandler.gemini_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 = ( - vertex_passthrough_logging_handler_result["result"] + standard_logging_response_object = gemini_passthrough_logging_handler_result["result"] + kwargs = gemini_passthrough_logging_handler_result["kwargs"] + elif self.is_vertex_route(url_route): + vertex_passthrough_logging_handler_result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=httpx_response, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, ) + standard_logging_response_object = vertex_passthrough_logging_handler_result["result"] kwargs = vertex_passthrough_logging_handler_result["kwargs"] elif self.is_anthropic_route(url_route): anthropic_passthrough_logging_handler_result = ( @@ -139,55 +155,72 @@ class PassThroughEndpointLogging: ) ) - standard_logging_response_object = ( - anthropic_passthrough_logging_handler_result["result"] - ) + standard_logging_response_object = anthropic_passthrough_logging_handler_result["result"] kwargs = anthropic_passthrough_logging_handler_result["kwargs"] elif self.is_cohere_route(url_route): - cohere_passthrough_logging_handler_result = ( - cohere_passthrough_logging_handler.passthrough_chat_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 = ( - cohere_passthrough_logging_handler_result["result"] + cohere_passthrough_logging_handler_result = cohere_passthrough_logging_handler.passthrough_chat_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 = cohere_passthrough_logging_handler_result["result"] kwargs = cohere_passthrough_logging_handler_result["kwargs"] - elif self.is_openai_route(url_route) and self._is_supported_openai_endpoint(url_route): + 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 {}, + 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"] + + elif self.is_vertex_ai_live_route(url_route): + from .llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import ( + VertexAILivePassthroughLoggingHandler, + ) + vertex_ai_live_handler = VertexAILivePassthroughLoggingHandler() + + # For WebSocket responses, response_body should be a list of messages + websocket_messages: list[dict[str, Any]] = response_body if isinstance(response_body, list) else [] + + vertex_ai_live_handler_result = ( + vertex_ai_live_handler.vertex_ai_live_passthrough_handler( + websocket_messages=websocket_messages, 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"] + + standard_logging_response_object = vertex_ai_live_handler_result["result"] + kwargs = vertex_ai_live_handler_result["kwargs"] return_dict[ "standard_logging_response_object" ] = standard_logging_response_object + return_dict["kwargs"] = kwargs return return_dict @@ -203,21 +236,13 @@ class PassThroughEndpointLogging: cache_hit: bool, request_body: dict, passthrough_logging_payload: PassthroughStandardLoggingPayload, + custom_llm_provider: Optional[str] = None, **kwargs, ): - standard_logging_response_object: Optional[ - PassThroughEndpointLoggingResultValues - ] = None - logging_obj.model_call_details[ - "passthrough_logging_payload" - ] = passthrough_logging_payload + standard_logging_response_object: Optional[PassThroughEndpointLoggingResultValues] = None + logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload if self.is_assemblyai_route(url_route): - if ( - AssemblyAIPassthroughLoggingHandler._should_log_request( - httpx_response.request.method - ) - is not True - ): + if AssemblyAIPassthroughLoggingHandler._should_log_request(httpx_response.request.method) is not True: return self.assemblyai_passthrough_logging_handler.assemblyai_passthrough_logging_handler( httpx_response=httpx_response, @@ -235,30 +260,25 @@ class PassThroughEndpointLogging: # Don't log langfuse pass-through requests return else: - normalized_llm_passthrough_logging_payload = ( - self.normalize_llm_passthrough_logging_payload( - httpx_response=httpx_response, - response_body=response_body, - request_body=request_body, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - **kwargs, - ) - ) - standard_logging_response_object = ( - normalized_llm_passthrough_logging_payload[ - "standard_logging_response_object" - ] + normalized_llm_passthrough_logging_payload = self.normalize_llm_passthrough_logging_payload( + httpx_response=httpx_response, + response_body=response_body, + request_body=request_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + custom_llm_provider=custom_llm_provider, + **kwargs, ) + standard_logging_response_object = normalized_llm_passthrough_logging_payload[ + "standard_logging_response_object" + ] kwargs = normalized_llm_passthrough_logging_payload["kwargs"] if standard_logging_response_object is None: - standard_logging_response_object = StandardPassThroughResponseObject( - response=httpx_response.text - ) + standard_logging_response_object = StandardPassThroughResponseObject(response=httpx_response.text) kwargs = self._set_cost_per_request( logging_obj=logging_obj, @@ -309,26 +329,43 @@ class PassThroughEndpointLogging: return True return False + def is_vertex_ai_live_route(self, url_route: str): + """Check if the URL route is a Vertex AI Live API WebSocket route.""" + if not url_route: + return False + for route in self.TRACKED_VERTEX_AI_LIVE_ROUTES: + if route in url_route: + 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 + "api.openai.com" in parsed_url.hostname or "openai.azure.com" in parsed_url.hostname ) + def is_gemini_route(self, url_route: str, custom_llm_provider: Optional[str] = None): + """Check if the URL route is a Gemini API route.""" + for route in self.TRACKED_GEMINI_ROUTES: + if route in url_route and custom_llm_provider == "gemini": + return True + return False + 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) + 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( @@ -347,11 +384,7 @@ class PassThroughEndpointLogging: # Check if cost per request is set ######################################################### if passthrough_logging_payload.get("cost_per_request") is not None: - kwargs["response_cost"] = passthrough_logging_payload.get( - "cost_per_request" - ) - logging_obj.model_call_details[ - "response_cost" - ] = passthrough_logging_payload.get("cost_per_request") + kwargs["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/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index a6fc377c1a7..b4717687704 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -175,4 +175,4 @@ class InMemoryPromptRegistry: return self.prompt_id_to_custom_prompt.get(prompt_id) -IN_MEMORY_PROMPT_REGISTRY = InMemoryPromptRegistry() +IN_MEMORY_PROMPT_REGISTRY = InMemoryPromptRegistry() \ No newline at end of file diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 867a395d627..21e2c54ec74 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -185,6 +185,7 @@ class ProxyInitializationHelpers: num_workers: int, ssl_certfile_path: str, ssl_keyfile_path: str, + max_requests_before_restart: Optional[int] = None, ): """ Run litellm with `gunicorn` @@ -265,6 +266,10 @@ class ProxyInitializationHelpers: "access_log_format": '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s', } + # Optional: recycle workers after N requests to mitigate memory growth + if max_requests_before_restart is not None: + gunicorn_options["max_requests"] = max_requests_before_restart + if ssl_certfile_path is not None and ssl_keyfile_path is not None: print( # noqa f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" # noqa @@ -486,6 +491,13 @@ class ProxyInitializationHelpers: help="Set the uvicorn keepalive timeout in seconds (uvicorn timeout_keep_alive parameter)", envvar="KEEPALIVE_TIMEOUT", ) +@click.option( + "--max_requests_before_restart", + default=None, + type=int, + help="Restart worker after this many requests (uvicorn: limit_max_requests, gunicorn: max_requests)", + envvar="MAX_REQUESTS_BEFORE_RESTART", +) def run_server( # noqa: PLR0915 host, port, @@ -524,6 +536,7 @@ def run_server( # noqa: PLR0915 use_prisma_db_push: bool, skip_server_startup, keepalive_timeout, + max_requests_before_restart, ): args = locals() if local: @@ -813,6 +826,9 @@ def run_server( # noqa: PLR0915 log_config=log_config, keepalive_timeout=keepalive_timeout, ) + # Optional: recycle uvicorn workers after N requests + if max_requests_before_restart is not None: + uvicorn_args["limit_max_requests"] = max_requests_before_restart if run_gunicorn is False and run_hypercorn is False: if ssl_certfile_path is not None and ssl_keyfile_path is not None: print( # noqa @@ -837,6 +853,7 @@ def run_server( # noqa: PLR0915 num_workers=num_workers, ssl_certfile_path=ssl_certfile_path, ssl_keyfile_path=ssl_keyfile_path, + max_requests_before_restart=max_requests_before_restart, ) elif run_hypercorn is True: ProxyInitializationHelpers._init_hypercorn_server( diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 73177fdd482..4878d15a3f0 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -26,19 +26,28 @@ model_list: - model_name: vertex_ai/* litellm_params: model: vertex_ai/* + - model_name: "grok-4" + model_info: + mode: completion + litellm_params: + model: oci/xai.grok-4 + oci_key: ocid1.api_key.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk + oci_region: us-phoenix-1 + oci_user: ocid1.user.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk + oci_fingerprint: aa:bb:cc:dd:ee:ff:11:22:33:44:55:66:77:88:99:00 + oci_tenancy: ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk + oci_key_file: /path/to/oci_api_key.pem + oci_compartment_id: ocid1.compartment.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk + drop_params: True guardrails: - - guardrail_name: lakera + - guardrail_name: "bedrock-pre-guard" litellm_params: - guardrail: lakera_v2 - mode: pre_call - api_key: os.environ/LAKERA_API_KEY - default_on: false - project_id: project-9770817088 - breakdown: true - payload: true - dev_info: true + guardrail: bedrock # supported values: "aporia", "bedrock", "lakera" + mode: "during_call" + guardrailIdentifier: ff6ujrregl1q + guardrailVersion: "DRAFT" litellm_settings: callbacks: ["datadog"] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f031b9ffbdf..55e4a96b5b0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -151,6 +151,7 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( router as mcp_discoverable_endpoints_router, ) @@ -252,7 +253,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, @@ -299,13 +302,18 @@ 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, ) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( router as llm_passthrough_router, ) +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + vertex_ai_live_websocket_passthrough, +) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( initialize_pass_through_endpoints, ) @@ -411,6 +419,8 @@ from fastapi import ( Request, Response, UploadFile, + WebSocket, + WebSocketDisconnect, applications, status, ) @@ -685,6 +695,8 @@ app = FastAPI( lifespan=proxy_startup_event, ) +vertex_live_passthrough_vertex_base = VertexBase() + ### CUSTOM API DOCS [ENTERPRISE FEATURE] ### # Custom OpenAPI schema generator to include only selected routes @@ -1865,6 +1877,15 @@ class ProxyConfig: verbose_proxy_logger.info( f"{blue_color_code}Set Global BitBucket Config on LiteLLM Proxy{reset_color_code}" ) + elif key == "global_gitlab_config": + from litellm.integrations.gitlab import ( + set_global_gitlab_config, + ) + + set_global_gitlab_config(value) + verbose_proxy_logger.info( + f"{blue_color_code}Set Global Gitlab Config on LiteLLM Proxy{reset_color_code}" + ) elif key == "callbacks": initialize_callbacks_on_proxy( value=value, @@ -2599,6 +2620,31 @@ class ProxyConfig: proxy_logging_obj=proxy_logging_obj, ) + def _add_callback_from_db_to_in_memory_litellm_callbacks( + self, + callback: str, + event_types: List[Literal["success", "failure"]], + existing_callbacks: list, + ) -> None: + """ + Helper method to add a single callback to litellm for specified event types. + + Args: + callback: The callback name to add + event_types: List of event types (e.g., ["success"], ["failure"], or ["success", "failure"]) + existing_callbacks: The existing callback list to check against + """ + if callback in litellm._known_custom_logger_compatible_callbacks: + for event_type in event_types: + _add_custom_logger_callback_to_specific_event(callback, event_type) + elif callback not in existing_callbacks: + if event_types == ["success"]: + litellm.logging_callback_manager.add_litellm_success_callback(callback) + elif event_types == ["failure"]: + litellm.logging_callback_manager.add_litellm_failure_callback(callback) + else: # Both success and failure + litellm.logging_callback_manager.add_litellm_callback(callback) + def _add_callbacks_from_db_config(self, config_data: dict) -> None: """ Adds callbacks from DB config to litellm @@ -2606,35 +2652,31 @@ class ProxyConfig: litellm_settings = config_data.get("litellm_settings", {}) or {} success_callbacks = litellm_settings.get("success_callback", None) failure_callbacks = litellm_settings.get("failure_callback", None) + callbacks = litellm_settings.get("callbacks", None) if success_callbacks is not None and isinstance(success_callbacks, list): for success_callback in success_callbacks: - if ( - success_callback - in litellm._known_custom_logger_compatible_callbacks - ): - _add_custom_logger_callback_to_specific_event( - success_callback, "success" - ) - elif success_callback not in litellm.success_callback: - litellm.logging_callback_manager.add_litellm_success_callback( - success_callback - ) + self._add_callback_from_db_to_in_memory_litellm_callbacks( + callback=success_callback, + event_types=["success"], + existing_callbacks=litellm.success_callback, + ) - # Add failure callbacks from DB to litellm if failure_callbacks is not None and isinstance(failure_callbacks, list): for failure_callback in failure_callbacks: - if ( - failure_callback - in litellm._known_custom_logger_compatible_callbacks - ): - _add_custom_logger_callback_to_specific_event( - failure_callback, "failure" - ) - elif failure_callback not in litellm.failure_callback: - litellm.logging_callback_manager.add_litellm_failure_callback( - failure_callback - ) + self._add_callback_from_db_to_in_memory_litellm_callbacks( + callback=failure_callback, + event_types=["failure"], + existing_callbacks=litellm.failure_callback, + ) + + if callbacks is not None and isinstance(callbacks, list): + for callback in callbacks: + self._add_callback_from_db_to_in_memory_litellm_callbacks( + callback=callback, + event_types=["success", "failure"], + existing_callbacks=litellm.callbacks, + ) def _encrypt_env_variables( self, environment_variables: dict, new_encryption_key: Optional[str] = None @@ -4935,13 +4977,49 @@ async def audio_transcriptions( ) +###################################################################### + +# Vertex AI Live API WebSocket Pass-through + +###################################################################### + + +@app.websocket("/vertex_ai/live") +async def vertex_ai_live_passthrough_endpoint( + websocket: WebSocket, + model: Optional[str] = fastapi.Query( + None, + description="Optional model name, used to determine Vertex region for global models.", + ), + vertex_project: Optional[str] = fastapi.Query( + None, + description="Override the Vertex AI project id used for the upstream connection.", + ), + vertex_location: Optional[str] = fastapi.Query( + None, + description="Override the Vertex AI region (for example, 'us-central1').", + ), + user_api_key_dict=Depends(user_api_key_auth_websocket), +): + """ + Vertex AI Live API WebSocket Pass-through Endpoint + + This endpoint delegates to the WebSocket function defined in llm_passthrough_endpoints.py + """ + return await vertex_ai_live_websocket_passthrough( + websocket=websocket, + model=model, + vertex_project=vertex_project, + vertex_location=vertex_location, + user_api_key_dict=user_api_key_dict, + ) + + ###################################################################### # /v1/realtime Endpoints ###################################################################### -from fastapi import FastAPI, WebSocket, WebSocketDisconnect - from litellm import _arealtime diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 766625145f6..5a79e171438 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -178,6 +178,7 @@ model LiteLLM_MCPServerTable { updated_by String? mcp_info Json? @default("{}") mcp_access_groups String[] + allowed_tools String[] @default([]) // Health check status status String? @default("unknown") last_health_check DateTime? diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 8c11760904c..1bb937f0b47 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -539,13 +539,15 @@ async def update_sso_settings(sso_config: SSOConfig): @router.get( "/get/ui_theme_settings", tags=["UI Theme Settings"], - dependencies=[Depends(user_api_key_auth)], response_model=UIThemeSettingsResponse, ) async def get_ui_theme_settings(): """ Get UI theme configuration from the litellm_settings. Returns current logo settings for UI customization. + + Note: This endpoint is public (no authentication required) so all users can see custom branding. + Only the /update/ui_theme_settings endpoint requires authentication for admins to change settings. """ from litellm.proxy.proxy_server import proxy_config diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 5b9337e852b..fc45266536d 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -12,7 +12,7 @@ from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.together_ai.rerank.handler import TogetherAIRerank from litellm.rerank_api.rerank_utils import get_optional_rerank_params from litellm.secret_managers.main import get_secret, get_secret_str -from litellm.types.rerank import OptionalRerankParams, RerankResponse +from litellm.types.rerank import RerankResponse from litellm.types.router import * from litellm.utils import ProviderConfigManager, client, exception_type @@ -136,7 +136,7 @@ def rerank( # noqa: PLR0915 ) ) - optional_rerank_params: OptionalRerankParams = get_optional_rerank_params( + optional_rerank_params: Dict = get_optional_rerank_params( rerank_provider_config=rerank_provider_config, model=model, drop_params=kwargs.get("drop_params") or litellm.drop_params or False, @@ -173,7 +173,7 @@ def rerank( # noqa: PLR0915 ) # Implement rerank logic here based on the custom_llm_provider - if _custom_llm_provider == "cohere" or _custom_llm_provider == "litellm_proxy": + if _custom_llm_provider == litellm.LlmProviders.COHERE or _custom_llm_provider == litellm.LlmProviders.LITELLM_PROXY: # Implement Cohere rerank logic api_key: Optional[str] = ( dynamic_api_key or optional_params.api_key or litellm.api_key @@ -205,7 +205,7 @@ def rerank( # noqa: PLR0915 client=client, model_response=model_response, ) - elif _custom_llm_provider == "azure_ai": + elif _custom_llm_provider == litellm.LlmProviders.AZURE_AI: api_base = ( dynamic_api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there or optional_params.api_base @@ -226,7 +226,7 @@ def rerank( # noqa: PLR0915 client=client, model_response=model_response, ) - elif _custom_llm_provider == "infinity": + elif _custom_llm_provider == litellm.LlmProviders.INFINITY: # Implement Infinity rerank logic api_key = dynamic_api_key or optional_params.api_key or litellm.api_key @@ -256,7 +256,7 @@ def rerank( # noqa: PLR0915 client=client, model_response=model_response, ) - elif _custom_llm_provider == "together_ai": + elif _custom_llm_provider == litellm.LlmProviders.TOGETHER_AI: # Implement Together AI rerank logic api_key = ( dynamic_api_key @@ -282,7 +282,7 @@ def rerank( # noqa: PLR0915 api_key=api_key, _is_async=_is_async, ) - elif _custom_llm_provider == "jina_ai": + elif _custom_llm_provider == litellm.LlmProviders.JINA_AI: if dynamic_api_key is None: raise ValueError( "Jina AI API key is required, please set 'JINA_AI_API_KEY' in your environment" @@ -309,7 +309,35 @@ def rerank( # noqa: PLR0915 client=client, model_response=model_response, ) - elif _custom_llm_provider == "bedrock": + elif _custom_llm_provider == litellm.LlmProviders.NVIDIA_NIM: + if dynamic_api_key is None: + raise ValueError( + "Nvidia NIM API key is required, please set 'NVIDIA_NIM_API_KEY' in your environment" + ) + + # Note: For rerank, the base URL is different from chat/embeddings + # Rerank uses ai.api.nvidia.com instead of integrate.api.nvidia.com + api_base = ( + optional_params.api_base + or get_secret("NVIDIA_NIM_API_BASE") # type: ignore + or "https://ai.api.nvidia.com" # Default for rerank + ) + + response = base_llm_http_handler.rerank( + model=model, + custom_llm_provider=_custom_llm_provider, + optional_rerank_params=optional_rerank_params, + logging_obj=litellm_logging_obj, + provider_config=rerank_provider_config, + timeout=optional_params.timeout, + api_key=dynamic_api_key or optional_params.api_key, + api_base=api_base, + _is_async=_is_async, + headers=headers or litellm.headers or {}, + client=client, + model_response=model_response, + ) + elif _custom_llm_provider == litellm.LlmProviders.BEDROCK: api_base = ( dynamic_api_base or optional_params.api_base @@ -331,7 +359,7 @@ def rerank( # noqa: PLR0915 logging_obj=litellm_logging_obj, client=client, ) - elif _custom_llm_provider == "hosted_vllm": + elif _custom_llm_provider == litellm.LlmProviders.HOSTED_VLLM: # Implement Hosted VLLM rerank logic api_key = ( dynamic_api_key @@ -365,7 +393,7 @@ def rerank( # noqa: PLR0915 model_response=model_response, ) - elif _custom_llm_provider == "deepinfra": + elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA: api_key = ( dynamic_api_key or optional_params.api_key diff --git a/litellm/rerank_api/rerank_utils.py b/litellm/rerank_api/rerank_utils.py index f70ec015b6e..38e599ef824 100644 --- a/litellm/rerank_api/rerank_utils.py +++ b/litellm/rerank_api/rerank_utils.py @@ -1,7 +1,6 @@ from typing import Any, Dict, List, Optional, Union from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig -from litellm.types.rerank import OptionalRerankParams def get_optional_rerank_params( @@ -17,7 +16,7 @@ def get_optional_rerank_params( max_chunks_per_doc: Optional[int] = None, max_tokens_per_doc: Optional[int] = None, non_default_params: Optional[dict] = None, -) -> OptionalRerankParams: +) -> Dict: all_non_default_params = non_default_params or {} if query is not None: all_non_default_params["query"] = query diff --git a/litellm/router.py b/litellm/router.py index 0275b636989..6f0eb51959c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -337,13 +337,13 @@ class Router: ``` """ - from litellm._service_logger import ServiceLogging - self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks self.enable_tag_filtering = enable_tag_filtering + from litellm._service_logger import ServiceLogging + self.service_logger_obj: ServiceLogging = ServiceLogging() litellm.suppress_debug_info = True # prevents 'Give Feedback/Get help' message from being emitted on Router - Relevant Issue: https://github.com/BerriAI/litellm/issues/5942 if self.set_verbose is True: if debug_level == "INFO": @@ -360,9 +360,9 @@ class Router: ) # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### - cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = ( - "local" # default to an in-memory cache - ) + cache_type: Literal[ + "local", "redis", "redis-semantic", "s3", "disk" + ] = "local" # default to an in-memory cache redis_cache = None cache_config: Dict[str, Any] = {} @@ -404,14 +404,19 @@ class Router: self.default_max_parallel_requests = default_max_parallel_requests self.provider_default_deployment_ids: List[str] = [] self.pattern_router = PatternMatchRouter() - self.team_pattern_routers: Dict[str, PatternMatchRouter] = ( - {} - ) # {"TEAM_ID": PatternMatchRouter} + self.team_pattern_routers: Dict[ + str, PatternMatchRouter + ] = {} # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} + # Initialize model_group_alias early since it's used in set_model_list + self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = ( + model_group_alias or {} + ) # dict to store aliases for router, ex. {"gpt-4": "gpt-3.5-turbo"}, all requests with gpt-4 -> get routed to gpt-3.5-turbo group + # Initialize model ID to deployment index mapping for O(1) lookups self.model_id_to_deployment_index_map: Dict[str, int] = {} - + if model_list is not None: # Build model index immediately to enable O(1) lookups from the start self._build_model_id_to_deployment_index_map(model_list) @@ -494,9 +499,6 @@ class Router: self.previous_models: List = ( [] ) # list to store failed calls (passed in as metadata to next call) - self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = ( - model_group_alias or {} - ) # dict to store aliases for router, ex. {"gpt-4": "gpt-3.5-turbo"}, all requests with gpt-4 -> get routed to gpt-3.5-turbo group # make Router.chat.completions.create compatible for openai.chat.completions.create default_litellm_params = default_litellm_params or {} @@ -561,15 +563,6 @@ class Router: ) else: litellm.failure_callback = [self.deployment_callback_on_failure] - verbose_router_logger.debug( - f"Intialized router with Routing strategy: {self.routing_strategy}\n\n" - f"Routing enable_pre_call_checks: {self.enable_pre_call_checks}\n\n" - f"Routing fallbacks: {self.fallbacks}\n\n" - f"Routing content fallbacks: {self.content_policy_fallbacks}\n\n" - f"Routing context window fallbacks: {self.context_window_fallbacks}\n\n" - f"Router Redis Caching={self.cache.redis_cache}\n" - ) - self.service_logger_obj = ServiceLogging() self.routing_strategy_args = routing_strategy_args self.provider_budget_config = provider_budget_config self.router_budget_logger: Optional[RouterBudgetLimiting] = None @@ -592,9 +585,9 @@ class Router: ) ) - self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = ( - model_group_retry_policy - ) + self.model_group_retry_policy: Optional[ + Dict[str, RetryPolicy] + ] = model_group_retry_policy self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None if allowed_fails_policy is not None: @@ -769,6 +762,14 @@ class Router: self.aanthropic_messages = self.factory_function( litellm.anthropic_messages, call_type="anthropic_messages" ) + self.agenerate_content = self.factory_function( + litellm.agenerate_content, call_type="agenerate_content" + ) + + self.aadapter_generate_content = self.factory_function( + litellm.aadapter_generate_content, call_type="aadapter_generate_content" + ) + self.aresponses = self.factory_function( litellm.aresponses, call_type="aresponses" ) @@ -1212,10 +1213,7 @@ class Router: async def _acompletion( self, model: str, messages: List[Dict[str, str]], **kwargs - ) -> Union[ - ModelResponse, - CustomStreamWrapper, - ]: + ) -> Union[ModelResponse, CustomStreamWrapper,]: """ - Get an available deployment - call it with a semaphore over the call @@ -3172,9 +3170,9 @@ class Router: healthy_deployments=healthy_deployments, responses=responses ) returned_response = cast(OpenAIFileObject, responses[0]) - returned_response._hidden_params["model_file_id_mapping"] = ( - model_file_id_mapping - ) + returned_response._hidden_params[ + "model_file_id_mapping" + ] = model_file_id_mapping return returned_response except Exception as e: verbose_router_logger.exception( @@ -3737,11 +3735,11 @@ class Router: if isinstance(e, litellm.ContextWindowExceededError): if context_window_fallbacks is not None: - context_window_fallback_model_group: Optional[List[str]] = ( - self._get_fallback_model_group_from_fallbacks( - fallbacks=context_window_fallbacks, - model_group=model_group, - ) + context_window_fallback_model_group: Optional[ + List[str] + ] = self._get_fallback_model_group_from_fallbacks( + fallbacks=context_window_fallbacks, + model_group=model_group, ) if context_window_fallback_model_group is None: raise original_exception @@ -3773,11 +3771,11 @@ class Router: e.message += "\n{}".format(error_message) elif isinstance(e, litellm.ContentPolicyViolationError): if content_policy_fallbacks is not None: - content_policy_fallback_model_group: Optional[List[str]] = ( - self._get_fallback_model_group_from_fallbacks( - fallbacks=content_policy_fallbacks, - model_group=model_group, - ) + content_policy_fallback_model_group: Optional[ + List[str] + ] = self._get_fallback_model_group_from_fallbacks( + fallbacks=content_policy_fallbacks, + model_group=model_group, ) if content_policy_fallback_model_group is None: raise original_exception @@ -4495,16 +4493,17 @@ class Router: try: exception = kwargs.get("exception", None) exception_status = getattr(exception, "status_code", "") - _model_info = kwargs.get("litellm_params", {}).get("model_info", {}) + + # Cache litellm_params to avoid repeated dict lookups + litellm_params = kwargs.get("litellm_params", {}) + _model_info = litellm_params.get("model_info", {}) exception_headers = litellm.litellm_core_utils.exception_mapping_utils._get_response_headers( original_exception=exception ) # Determine cooldown time with priority: deployment config > response header > router default - deployment_cooldown = kwargs.get("litellm_params", {}).get( - "cooldown_time", None - ) + deployment_cooldown = litellm_params.get("cooldown_time", None) header_cooldown = None if exception_headers is not None: @@ -4756,11 +4755,11 @@ class Router: unhealthy_deployments = await _async_get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) + # Convert to set for O(1) lookup instead of O(n) + unhealthy_deployments_set = set(unhealthy_deployments) healthy_deployments: list = [] for deployment in _all_deployments: - if deployment["model_info"]["id"] in unhealthy_deployments: - continue - else: + if deployment["model_info"]["id"] not in unhealthy_deployments_set: healthy_deployments.append(deployment) return healthy_deployments, _all_deployments @@ -4984,7 +4983,9 @@ class Router: model = deployment.to_json(exclude_none=True) - self._add_model_to_list_and_index_map(model=model, model_id=deployment.model_info.id) + self._add_model_to_list_and_index_map( + model=model, model_id=deployment.model_info.id + ) return deployment except Exception as e: if self.ignore_invalid_deployments: @@ -5013,26 +5014,26 @@ class Router: """ from litellm.router_strategy.auto_router.auto_router import AutoRouter - auto_router_config_path: Optional[str] = ( - deployment.litellm_params.auto_router_config_path - ) + auto_router_config_path: Optional[ + str + ] = deployment.litellm_params.auto_router_config_path auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config if auto_router_config_path is None and auto_router_config is None: raise ValueError( "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" ) - default_model: Optional[str] = ( - deployment.litellm_params.auto_router_default_model - ) + default_model: Optional[ + str + ] = deployment.litellm_params.auto_router_default_model if default_model is None: raise ValueError( "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" ) - embedding_model: Optional[str] = ( - deployment.litellm_params.auto_router_embedding_model - ) + embedding_model: Optional[ + str + ] = deployment.litellm_params.auto_router_embedding_model if embedding_model is None: raise ValueError( "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" @@ -5336,14 +5337,18 @@ class Router: self._add_deployment(deployment=deployment) # add to model names - self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id) + self._add_model_to_list_and_index_map( + model=_deployment, model_id=deployment.model_info.id + ) self.model_names.append(deployment.model_name) return deployment - def _update_deployment_indices_after_removal(self, model_id: str, removal_idx: int) -> None: + def _update_deployment_indices_after_removal( + self, model_id: str, removal_idx: int + ) -> None: """ Helper method to update deployment indices after a deployment has been removed from model_list. - + Parameters: - model_id: str - the id of the deployment that was removed - removal_idx: int - the index where the deployment was removed from model_list @@ -5356,11 +5361,12 @@ class Router: if model_id in self.model_id_to_deployment_index_map: del self.model_id_to_deployment_index_map[model_id] - - def _add_model_to_list_and_index_map(self, model: dict, model_id: Optional[str] = None) -> None: + def _add_model_to_list_and_index_map( + self, model: dict, model_id: Optional[str] = None + ) -> None: """ Helper method to add a model to the model_list and update the model_id_to_deployment_index_map. - + Parameters: - model: dict - the model to add to the list - model_id: Optional[str] - the model ID to use for indexing. If None, will try to get from model["model_info"]["id"] @@ -5370,7 +5376,9 @@ class Router: if model_id is not None: self.model_id_to_deployment_index_map[model_id] = len(self.model_list) - 1 elif model.get("model_info", {}).get("id") is not None: - self.model_id_to_deployment_index_map[model["model_info"]["id"]] = len(self.model_list) - 1 + self.model_id_to_deployment_index_map[model["model_info"]["id"]] = ( + len(self.model_list) - 1 + ) def upsert_deployment(self, deployment: Deployment) -> Optional[Deployment]: """ @@ -5399,13 +5407,15 @@ class Router: removal_idx: Optional[int] = None deployment_id = deployment.model_info.id deployment_fast_mapping = self.model_id_to_deployment_index_map - + if deployment_id in deployment_fast_mapping: removal_idx = deployment_fast_mapping[deployment_id] if removal_idx is not None: self.model_list.pop(removal_idx) - self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx) + self._update_deployment_indices_after_removal( + model_id=deployment_id, removal_idx=removal_idx + ) # if the model_id is not in router self.add_deployment(deployment=deployment) @@ -5436,7 +5446,9 @@ class Router: if deployment_idx is not None: # Pop the item from the list first item = self.model_list.pop(deployment_idx) - self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx) + self._update_deployment_indices_after_removal( + model_id=id, removal_idx=deployment_idx + ) return item else: return None @@ -5459,7 +5471,7 @@ class Router: return model else: raise Exception("Model invalid format - {}".format(type(model))) - + return None def get_deployment_credentials(self, model_id: str) -> Optional[dict]: @@ -5705,27 +5717,32 @@ class Router: configurable_clientside_auth_params = ( litellm_params.configurable_clientside_auth_params ) + + # Cache nested dict access to avoid repeated temporary dict allocations + model_litellm_params = model.get("litellm_params", {}) + model_info_dict = model.get("model_info", {}) + # get model tpm _deployment_tpm: Optional[int] = None if _deployment_tpm is None: _deployment_tpm = model.get("tpm", None) # type: ignore if _deployment_tpm is None: - _deployment_tpm = model.get("litellm_params", {}).get("tpm", None) # type: ignore + _deployment_tpm = model_litellm_params.get("tpm", None) # type: ignore if _deployment_tpm is None: - _deployment_tpm = model.get("model_info", {}).get("tpm", None) # type: ignore + _deployment_tpm = model_info_dict.get("tpm", None) # type: ignore # get model rpm _deployment_rpm: Optional[int] = None if _deployment_rpm is None: _deployment_rpm = model.get("rpm", None) # type: ignore if _deployment_rpm is None: - _deployment_rpm = model.get("litellm_params", {}).get("rpm", None) # type: ignore + _deployment_rpm = model_litellm_params.get("rpm", None) # type: ignore if _deployment_rpm is None: - _deployment_rpm = model.get("model_info", {}).get("rpm", None) # type: ignore + _deployment_rpm = model_info_dict.get("rpm", None) # type: ignore # get model info try: - model_id = model.get("model_info", {}).get("id", None) + model_id = model_info_dict.get("id", None) if model_id is not None: model_info = self.get_deployment_model_info( model_id=model_id, model_name=litellm_params.model @@ -6089,7 +6106,7 @@ class Router: # Extract model_info from the model dict model_info = model.get("model_info", {}) model_id = model_info.get("id") - + # If no ID exists, generate one using the same logic as set_model_list if model_id is None: model_name = model.get("model_name", "") @@ -6099,7 +6116,7 @@ class Router: if "model_info" not in model: model["model_info"] = {} model["model_info"]["id"] = model_id - + self._add_model_to_list_and_index_map(model=model, model_id=model_id) def get_model_ids( @@ -6291,45 +6308,41 @@ class Router: if team_id specified, returns matching team-specific models """ + # Note: model_list and model_group_alias are always initialized in __init__ + # so hasattr checks are unnecessary + returned_models: List[DeploymentTypedDict] = [] - if hasattr(self, "model_list"): - returned_models: List[DeploymentTypedDict] = [] + if model_name is not None: + returned_models.extend( + self._get_all_deployments(model_name=model_name, team_id=team_id) + ) - if model_name is not None: - returned_models.extend( - self._get_all_deployments(model_name=model_name, team_id=team_id) + returned_models.extend( + self.get_model_list_from_model_alias(model_name=model_name) + ) + + if len(returned_models) == 0: # check if wildcard route + potential_wildcard_models = self.pattern_router.route(model_name) or [] + + ## check for team-specific wildcard models + if team_id is not None and team_id in self.team_pattern_routers: + potential_team_only_wildcard_models = ( + self.team_pattern_routers[team_id].route(model_name) or [] + ) + potential_wildcard_models.extend( + potential_team_only_wildcard_models ) - if hasattr(self, "model_group_alias"): - returned_models.extend( - self.get_model_list_from_model_alias(model_name=model_name) - ) + if model_name is not None and potential_wildcard_models is not None: + for m in potential_wildcard_models: + deployment_typed_dict = DeploymentTypedDict(**m) # type: ignore + deployment_typed_dict["model_name"] = model_name + returned_models.append(deployment_typed_dict) - if len(returned_models) == 0: # check if wildcard route - potential_wildcard_models = self.pattern_router.route(model_name) or [] + if model_name is None: + returned_models += self.model_list - ## check for team-specific wildcard models - if team_id is not None and team_id in self.team_pattern_routers: - potential_team_only_wildcard_models = ( - self.team_pattern_routers[team_id].route(model_name) or [] - ) - potential_wildcard_models.extend( - potential_team_only_wildcard_models - ) - - if model_name is not None and potential_wildcard_models is not None: - for m in potential_wildcard_models: - deployment_typed_dict = DeploymentTypedDict(**m) # type: ignore - deployment_typed_dict["model_name"] = model_name - returned_models.append(deployment_typed_dict) - - if model_name is None: - returned_models += self.model_list - - return returned_models - - return returned_models - return None + return returned_models def get_model_access_groups( self, @@ -6576,19 +6589,19 @@ class Router: or {} ) # check the in-memory cache used by lowest_latency and usage-based routing. Only check the local cache. for idx, deployment in enumerate(_returned_deployments): + # Cache nested dict access to avoid repeated temporary dict allocations + _litellm_params = deployment.get("litellm_params", {}) + _model_info = deployment.get("model_info", {}) + # see if we have the info for this model try: - base_model = deployment.get("model_info", {}).get("base_model", None) + base_model = _model_info.get("base_model", None) if base_model is None: - base_model = deployment.get("litellm_params", {}).get( - "base_model", None - ) + base_model = _litellm_params.get("base_model", None) model_info = self.get_router_model_info( deployment=deployment, received_model_name=model ) - model = base_model or deployment.get("litellm_params", {}).get( - "model", None - ) + model = base_model or _litellm_params.get("model", None) if ( isinstance(model_info, dict) @@ -6609,8 +6622,7 @@ class Router: except Exception as e: verbose_router_logger.exception("An error occurs - {}".format(str(e))) - _litellm_params = deployment.get("litellm_params", {}) - model_id = deployment.get("model_info", {}).get("id", "") + model_id = _model_info.get("id", "") ## RPM CHECK ## ### get local router cache ### current_request_cache_local = ( @@ -7256,19 +7268,13 @@ class Router: Returns: List of healthy deployments """ - # filter out the deployments currently cooling down - deployments_to_remove = [] verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") - # Find deployments in model_list whose model_id is cooling down - for deployment in healthy_deployments: - deployment_id = deployment["model_info"]["id"] - if deployment_id in cooldown_deployments: - deployments_to_remove.append(deployment) - - # remove unhealthy deployments from healthy deployments - for deployment in deployments_to_remove: - healthy_deployments.remove(deployment) - return healthy_deployments + # Convert to set for O(1) lookup and use list comprehension for O(n) filtering + cooldown_set = set(cooldown_deployments) + return [ + deployment for deployment in healthy_deployments + if deployment["model_info"]["id"] not in cooldown_set + ] def _track_deployment_metrics( self, deployment, parent_otel_span: Optional[Span], response=None diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py index b875495bab0..0a26f266a6e 100644 --- a/litellm/types/google_genai/main.py +++ b/litellm/types/google_genai/main.py @@ -1,28 +1,58 @@ # Import types from the Google GenAI SDK -from typing import TYPE_CHECKING, Any, List, Optional, TypeAlias +from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypeAlias -# During static type-checking we can rely on the real google-genai types. -from google.genai import types as _genai_types # type: ignore from pydantic import BaseModel from typing_extensions import TypedDict from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject -ContentListUnion = _genai_types.ContentListUnion -ContentListUnionDict = _genai_types.ContentListUnionDict -GenerateContentConfigOrDict = _genai_types.GenerateContentConfigOrDict -GoogleGenAIGenerateContentResponse = _genai_types.GenerateContentResponse +# During static type-checking we can rely on the real google-genai types. +if TYPE_CHECKING: + from google.genai import types as _genai_types # type: ignore -GenerateContentContentListUnionDict = _genai_types.ContentListUnionDict -GenerateContentConfigDict = _genai_types.GenerateContentConfigDict -GenerateContentRequestParametersDict = _genai_types._GenerateContentParametersDict -ToolConfigDict = _genai_types.ToolConfigDict + ContentListUnion = _genai_types.ContentListUnion + ContentListUnionDict = _genai_types.ContentListUnionDict + GenerateContentConfigOrDict = _genai_types.GenerateContentConfigOrDict + GoogleGenAIGenerateContentResponse = _genai_types.GenerateContentResponse + GenerateContentContentListUnionDict = _genai_types.ContentListUnionDict + GenerateContentConfigDict = _genai_types.GenerateContentConfigDict + GenerateContentRequestParametersDict = _genai_types._GenerateContentParametersDict + ToolConfigDict = _genai_types.ToolConfigDict -class GenerateContentRequestDict(GenerateContentRequestParametersDict): # type: ignore[misc] - generationConfig: Optional[Any] - tools: Optional[ToolConfigDict] # type: ignore[assignment] + class GenerateContentRequestDict(GenerateContentRequestParametersDict): # type: ignore[misc] + generationConfig: Optional[Any] + tools: Optional[ToolConfigDict] # type: ignore[assignment] + class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): # type: ignore[misc] + _hidden_params: dict = {} + pass +else: + # Fallback types when google.genai is not available + ContentListUnion = Any + ContentListUnionDict = Dict[str, Any] + GenerateContentConfigOrDict = Dict[str, Any] + GoogleGenAIGenerateContentResponse = Dict[str, Any] + GenerateContentContentListUnionDict = Dict[str, Any] -class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): # type: ignore[misc] - _hidden_params: dict = {} - pass \ No newline at end of file + # Create a proper fallback class that can be instantiated + class GenerateContentConfigDict(dict): # type: ignore[misc] + def __init__(self, **kwargs): # type: ignore + super().__init__(**kwargs) + + class GenerateContentRequestParametersDict(dict): # type: ignore[misc] + def __init__(self, **kwargs): # type: ignore + super().__init__(**kwargs) + + ToolConfigDict = Dict[str, Any] + + class GenerateContentRequestDict(GenerateContentRequestParametersDict): # type: ignore[misc] + def __init__(self, **kwargs): # type: ignore + # Extract specific fields + self.generationConfig = kwargs.get('generationConfig') + self.tools = kwargs.get('tools') + super().__init__(**kwargs) + + class GenerateContentResponse(BaseLiteLLMOpenAIResponseObject): # type: ignore[misc] + def __init__(self, **kwargs): # type: ignore + super().__init__(**kwargs) + self._hidden_params = kwargs.get('_hidden_params', {}) \ No newline at end of file diff --git a/litellm/types/integrations/langfuse.py b/litellm/types/integrations/langfuse.py index 08ad667cac4..a13868e503c 100644 --- a/litellm/types/integrations/langfuse.py +++ b/litellm/types/integrations/langfuse.py @@ -12,5 +12,6 @@ class LangfuseLoggingConfig(TypedDict): class LangfuseUsageDetails(TypedDict): input: Optional[int] output: Optional[int] + total: Optional[int] cache_creation_input_tokens: Optional[int] cache_read_input_tokens: Optional[int] diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index cebcd0522a1..df551c5bded 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -377,9 +377,14 @@ TWELVELABS_EMBEDDING_INPUT_TYPES = Literal["text", "image", "video", "audio"] TWELVELABS_EMBEDDING_OPTIONS = Literal["visual-text", "visual-image", "audio"] +class TwelveLabsS3Location(TypedDict, total=False): + uri: str + bucketOwner: str + + class TwelveLabsMediaSource(TypedDict, total=False): base64String: str - s3Location: dict # {"uri": str, "bucketOwner": str} + s3Location: TwelveLabsS3Location class TwelveLabsMarengoEmbeddingRequest(TypedDict, total=False): @@ -401,6 +406,32 @@ class TwelveLabsMarengoEmbeddingResponse(TypedDict): endSec: float +class TwelveLabsS3OutputDataConfig(TypedDict): + s3Uri: str + + +class TwelveLabsOutputDataConfig(TypedDict): + s3OutputDataConfig: TwelveLabsS3OutputDataConfig + + +class TwelveLabsAsyncInvokeRequest(TypedDict): + modelId: str + modelInput: TwelveLabsMarengoEmbeddingRequest + outputDataConfig: TwelveLabsOutputDataConfig + + +class TwelveLabsAsyncInvokeStatusResponse(TypedDict): + invocationArn: str + modelArn: str + status: str # "InProgress" | "Completed" | "Failed" + submitTime: str + lastModifiedTime: str + endTime: Optional[str] + outputDataConfig: TwelveLabsOutputDataConfig + clientRequestToken: Optional[str] + failureMessage: Optional[str] + + AmazonEmbeddingRequest = Union[ AmazonTitanMultimodalEmbeddingRequest, AmazonTitanV2EmbeddingRequest, diff --git a/litellm/types/llms/oci.py b/litellm/types/llms/oci.py index 75d13192c50..56fd61ad7e6 100644 --- a/litellm/types/llms/oci.py +++ b/litellm/types/llms/oci.py @@ -100,8 +100,8 @@ class OCIServingMode(BaseModel): """Defines the serving mode and the model to be used.""" servingType: str - modelId: str - + endpointId: Optional[str] = None + modelId: Optional[str] = None class OCICompletionPayload(BaseModel): """Pydantic model for the complete OCI chat request body.""" @@ -129,7 +129,7 @@ class OCIPromptTokensDetails(BaseModel): class OCIResponseUsage(BaseModel): """Token usage in the OCI response.""" - + promptTokens: int completionTokens: int totalTokens: int diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 854d37522cf..f1f7ac2c661 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -209,6 +209,16 @@ class GenerationConfig(TypedDict, total=False): thinkingConfig: GeminiThinkingConfig +class VertexToolName(str, Enum): + """Enum for Vertex AI tool field names.""" + GOOGLE_SEARCH = "googleSearch" + GOOGLE_SEARCH_RETRIEVAL = "googleSearchRetrieval" + ENTERPRISE_WEB_SEARCH = "enterpriseWebSearch" + URL_CONTEXT = "url_context" + CODE_EXECUTION = "code_execution" + GOOGLE_MAPS = "googleMaps" + + class Tools(TypedDict, total=False): function_declarations: List[FunctionDeclaration] googleSearch: dict @@ -216,6 +226,7 @@ class Tools(TypedDict, total=False): enterpriseWebSearch: dict url_context: dict code_execution: dict + googleMaps: dict retrieval: Retrieval diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index c99775f2a6c..85c8c0f3351 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -6,6 +6,7 @@ from typing_extensions import TypedDict class EndpointType(str, Enum): VERTEX_AI = "vertex-ai" + GEMINI = "gemini" ANTHROPIC = "anthropic" OPENAI = "openai" GENERIC = "generic" diff --git a/litellm/types/prompts/init_prompts.py b/litellm/types/prompts/init_prompts.py index 3c48cd131c3..184f0448b33 100644 --- a/litellm/types/prompts/init_prompts.py +++ b/litellm/types/prompts/init_prompts.py @@ -10,6 +10,7 @@ class SupportedPromptIntegrations(str, Enum): LANGFUSE = "langfuse" CUSTOM = "custom" BITBUCKET = "bitbucket" + GITLAB = "gitlab" class PromptInfo(BaseModel): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c8de97bba20..0da0ba4e93d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1425,6 +1425,9 @@ class EmbeddingResponse(OpenAIObject): model = model super().__init__(model=model, object=object, data=data, usage=usage) # type: ignore + if hidden_params: + self._hidden_params = hidden_params + def __contains__(self, key): # Define custom behavior for the 'in' operator return hasattr(self, key) @@ -2040,6 +2043,13 @@ class GuardrailMode(TypedDict, total=False): default: Optional[str] +GuardrailStatus = Literal[ + "success", + "guardrail_intervened", + "guardrail_failed_to_respond", + "not_run" +] + class StandardLoggingGuardrailInformation(TypedDict, total=False): guardrail_name: Optional[str] guardrail_provider: Optional[str] @@ -2048,7 +2058,7 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): ] guardrail_request: Optional[dict] guardrail_response: Optional[Union[dict, str, List[dict]]] - guardrail_status: Literal["success", "failure", "blocked"] + guardrail_status: GuardrailStatus start_time: Optional[float] end_time: Optional[float] duration: Optional[float] @@ -2097,6 +2107,20 @@ class CostBreakdown(TypedDict): tool_usage_cost: float # Cost of usage of built-in tools +class StandardLoggingPayloadStatusFields(TypedDict, total=False): + """Status fields for easy filtering and analytics""" + llm_api_status: StandardLoggingPayloadStatus + """Status of the LLM API call - 'success' if completed, 'failure' if errored""" + guardrail_status: GuardrailStatus + """ + Status of guardrail execution: + - 'success': Guardrail ran and allowed content through + - 'guardrail_intervened': Guardrail blocked or modified content + - 'guardrail_failed_to_respond': Guardrail had technical failure + - 'not_run': No guardrail was run + """ + + class StandardLoggingPayload(TypedDict): id: str trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries) @@ -2108,6 +2132,7 @@ class StandardLoggingPayload(TypedDict): StandardLoggingModelCostFailureDebugInformation ] status: StandardLoggingPayloadStatus + status_fields: StandardLoggingPayloadStatusFields custom_llm_provider: Optional[str] total_tokens: int prompt_tokens: int @@ -2684,8 +2709,18 @@ class PriorityReservationSettings(BaseModel): """ default_priority: float = Field( - default=0.5, + default=0.25, description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation.", ) + + saturation_threshold: float = Field( + default=0.50, + description="Saturation threshold (0.0-1.0) at which strict priority enforcement begins. Below this threshold, generous mode allows priority borrowing. Above this threshold, strict mode enforces normalized priority limits." + ) + + tracking_multiplier: int = Field( + default=10, + description="Multiplier for model-wide tracking limits in strict mode. Set to 10x because v3_limiter.should_rate_limit() both increments counters AND enforces limits - we need the counter increment (for saturation checks) but not the enforcement (priority limits handle that). High multiplier ensures tracking never blocks." + ) model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/utils.py b/litellm/utils.py index 8dfa2416a62..3c6c3ac86e4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -90,6 +90,7 @@ from litellm.litellm_core_utils.cached_imports import ( get_set_callbacks, ) from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, map_finish_reason, process_response_headers, ) @@ -1401,7 +1402,7 @@ def client(original_function): # noqa: PLR0915 print_verbose( f"ASYNC kwargs[caching]: {kwargs.get('caching', False)}; litellm.cache: {litellm.cache}; kwargs.get('cache'): {kwargs.get('cache', None)}" ) - _caching_handler_response: CachingHandlerResponse = ( + _caching_handler_response: Optional[CachingHandlerResponse] = ( await _llm_caching_handler._async_get_cache( model=model or "", original_function=original_function, @@ -1413,14 +1414,15 @@ def client(original_function): # noqa: PLR0915 ) ) - if ( - _caching_handler_response.cached_result is not None - and _caching_handler_response.final_embedding_cached_response is None - ): - return _caching_handler_response.cached_result + if _caching_handler_response is not None: + if ( + _caching_handler_response.cached_result is not None + and _caching_handler_response.final_embedding_cached_response is None + ): + return _caching_handler_response.cached_result - elif _caching_handler_response.embedding_all_elements_cache_hit is True: - return _caching_handler_response.final_embedding_cached_response + elif _caching_handler_response.embedding_all_elements_cache_hit is True: + return _caching_handler_response.final_embedding_cached_response # CHECK MAX TOKENS if ( @@ -1523,6 +1525,7 @@ def client(original_function): # noqa: PLR0915 # REBUILD EMBEDDING CACHING if ( isinstance(result, EmbeddingResponse) + and _caching_handler_response is not None and _caching_handler_response.final_embedding_cached_response is not None ): @@ -2801,6 +2804,8 @@ def get_optional_params_embeddings( # noqa: PLR0915 object = litellm.AmazonTitanV2Config() elif "cohere.embed-multilingual-v3" in model: object = litellm.BedrockCohereEmbeddingConfig() + elif "twelvelabs" in model or "marengo" in model: + object = litellm.TwelveLabsMarengoEmbeddingConfig() else: # unmapped model supported_params = [] _check_valid_arg(supported_params=supported_params) @@ -7205,6 +7210,8 @@ class ProviderConfigManager: return litellm.HuggingFaceRerankConfig() elif litellm.LlmProviders.DEEPINFRA == provider: return litellm.DeepinfraRerankConfig() + elif litellm.LlmProviders.NVIDIA_NIM == provider: + return litellm.NvidiaNimRerankConfig() return litellm.CohereRerankConfig() @staticmethod @@ -7582,7 +7589,7 @@ def get_end_user_id_for_cost_tracking( service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking. """ - _metadata = cast(dict, litellm_params.get("metadata", {}) or {}) + _metadata = cast(dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params))) end_user_id = cast( Optional[str], diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a35cac30489..72b9c551d73 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -2004,9 +2004,9 @@ "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", - "max_input_tokens": 272000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, "supported_endpoints": [ @@ -3308,6 +3308,64 @@ "supports_tool_choice": true, "supports_web_search": true }, + "azure_ai/grok-4": { + "input_cost_per_token": 5.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.75e-05, + "source": "https://azure.microsoft.com/en-us/blog/grok-4-is-now-available-in-azure-ai-foundry-unlock-frontier-intelligence-and-business-ready-capabilities/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_web_search": true + }, + "azure_ai/grok-4-fast-non-reasoning": { + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.5e-03, + "source": "https://azure.microsoft.com/en-us/blog/grok-4-is-now-available-in-azure-ai-foundry-unlock-frontier-intelligence-and-business-ready-capabilities/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_web_search": true + }, + "azure_ai/grok-4-fast-reasoning": { + "input_cost_per_token": 5.8e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.9e-03, + "source": "https://azure.microsoft.com/en-us/blog/grok-4-is-now-available-in-azure-ai-foundry-unlock-frontier-intelligence-and-business-ready-capabilities/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_web_search": true + }, + "azure_ai/grok-code-fast-1": { + "input_cost_per_token": 3.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.75e-05, + "source": "https://azure.microsoft.com/en-us/blog/grok-4-is-now-available-in-azure-ai-foundry-unlock-frontier-intelligence-and-business-ready-capabilities/", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_web_search": true + }, "azure_ai/jais-30b-chat": { "input_cost_per_token": 0.0032, "litellm_provider": "azure_ai", @@ -4743,6 +4801,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -4769,6 +4831,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -6627,629 +6693,679 @@ ] }, "deepinfra/Gryphe/MythoMax-L2-13b": { - "input_cost_per_token": 7.2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "max_tokens": 4096, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 9e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 7.2e-08, "supports_tool_choice": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": { - "input_cost_per_token": 7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 1e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 8e-07, "supports_tool_choice": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.8e-07, "supports_tool_choice": false }, "deepinfra/Qwen/QwQ-32B": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 1.5e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen2.5-72B-Instruct": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 1.2e-07, "output_cost_per_token": 3.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen2.5-7B-Instruct": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Qwen/Qwen2.5-VL-32B-Instruct": { - "input_cost_per_token": 2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", + "input_cost_per_token": 2e-07, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-14B": { - "input_cost_per_token": 6e-08, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 6e-08, "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 5.4e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 6e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 9e-08, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.9e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 6e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-30B-A3B": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 8e-08, "output_cost_per_token": 2.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-32B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 3e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": { - "input_cost_per_token": 4e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": { - "cache_read_input_token_cost": 2.4e-07, - "input_cost_per_token": 3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 2.9e-07, "output_cost_per_token": 1.2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": { - "input_cost_per_token": 2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Sao10K/L3.1-70B-Euryale-v2.2": { - "input_cost_per_token": 6.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 6.5e-07, "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Sao10K/L3.3-70B-Euryale-v2.3": { - "input_cost_per_token": 6.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 6.5e-07, "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/allenai/olmOCR-7B-0725-FP8": { - "input_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 16384, "max_input_tokens": 16384, "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", + "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1.5e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/anthropic/claude-3-7-sonnet-latest": { - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 3.3e-06, "output_cost_per_token": 1.65e-05, + "cache_read_input_token_cost": 3.3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/anthropic/claude-4-opus": { - "input_cost_per_token": 1.65e-05, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 1.65e-05, "output_cost_per_token": 8.25e-05, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/anthropic/claude-4-sonnet": { - "input_cost_per_token": 3.3e-06, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 3.3e-06, "output_cost_per_token": 1.65e-05, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1": { - "input_cost_per_token": 7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 7e-07, "output_cost_per_token": 2.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528": { - "cache_read_input_token_cost": 4e-07, - "input_cost_per_token": 5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 5e-07, "output_cost_per_token": 2.15e-06, + "cache_read_input_token_cost": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": { - "input_cost_per_token": 1e-06, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 1e-06, "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 4e-07, "supports_tool_choice": false }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { - "input_cost_per_token": 7.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 2.7e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.5e-07, "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": { - "input_cost_per_token": 1e-06, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 1e-06, "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3": { - "input_cost_per_token": 3.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 3.8e-07, "output_cost_per_token": 8.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3-0324": { - "cache_read_input_token_cost": 2.24e-07, - "input_cost_per_token": 2.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 2.5e-07, "output_cost_per_token": 8.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1": { - "cache_read_input_token_cost": 2.16e-07, - "input_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1e-06, - "supports_reasoning": true, + "cache_read_input_token_cost": 2.16e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true, + "supports_reasoning": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3.1-Terminus": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 1e-06, + "cache_read_input_token_cost": 2.16e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemini-2.0-flash-001": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, - "mode": "chat", + "input_cost_per_token": 1e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemini-2.5-flash": { - "input_cost_per_token": 2.1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.5e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.75e-06, "supports_tool_choice": true }, "deepinfra/google/gemini-2.5-pro": { - "input_cost_per_token": 8.75e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 1e-05, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 7e-06, "supports_tool_choice": true }, "deepinfra/google/gemma-3-12b-it": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemma-3-27b-it": { - "input_cost_per_token": 9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 1.6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.7e-07, "supports_tool_choice": true }, "deepinfra/google/gemma-3-4b-it": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 8e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": { - "input_cost_per_token": 4.9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4.9e-08, "output_cost_per_token": 4.9e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Llama-3.2-3B-Instruct": { - "input_cost_per_token": 1.2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 2e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.4e-08, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct": { - "input_cost_per_token": 2.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 2.3e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": { - "input_cost_per_token": 3.8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 3.9e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.2e-07, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1048576, "max_input_tokens": 1048576, "max_output_tokens": 1048576, - "max_tokens": 1048576, - "mode": "chat", + "input_cost_per_token": 1.5e-07, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 327680, "max_input_tokens": 327680, "max_output_tokens": 327680, - "max_tokens": 327680, - "mode": "chat", + "input_cost_per_token": 8e-08, "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-Guard-3-8B": { - "input_cost_per_token": 5.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5.5e-08, "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Llama-Guard-4-12B": { - "input_cost_per_token": 1.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 1.8e-07, "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Meta-Llama-3-8B-Instruct": { - "input_cost_per_token": 3e-08, - "litellm_provider": "deepinfra", + "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", + "input_cost_per_token": 3e-08, "output_cost_per_token": 6e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": { - "input_cost_per_token": 2.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 1e-07, "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": { - "input_cost_per_token": 3e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 3e-08, "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { - "input_cost_per_token": 1.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 3e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2e-08, "supports_tool_choice": true }, "deepinfra/microsoft/WizardLM-2-8x22B": { - "input_cost_per_token": 4.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 65536, "max_input_tokens": 65536, "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", + "input_cost_per_token": 4.8e-07, "output_cost_per_token": 4.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/microsoft/phi-4": { - "input_cost_per_token": 7e-08, - "litellm_provider": "deepinfra", + "max_tokens": 16384, "max_input_tokens": 16384, "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", + "input_cost_per_token": 7e-08, "output_cost_per_token": 1.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": { - "input_cost_per_token": 2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 2e-08, "output_cost_per_token": 4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 8e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "max_tokens": 128000, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 2e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1e-07, "supports_tool_choice": true }, "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.4e-07, "supports_tool_choice": true }, "deepinfra/moonshotai/Kimi-K2-Instruct": { - "input_cost_per_token": 5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-07, "output_cost_per_token": 2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/moonshotai/Kimi-K2-Instruct-0905": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2e-06, + "cache_read_input_token_cost": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/nvidia/Llama-3.3-Nemotron-Super-49B-v1.5": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 3e-07, "supports_tool_choice": true }, "deepinfra/openai/gpt-oss-120b": { - "input_cost_per_token": 9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 4.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/openai/gpt-oss-20b": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.6e-07, "supports_tool_choice": true }, "deepinfra/zai-org/GLM-4.5": { - "input_cost_per_token": 5.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_tool_choice": true - }, - "deepinfra/zai-org/GLM-4.5-Air": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.1e-06, "supports_tool_choice": true }, "deepseek/deepseek-chat": { @@ -7722,6 +7838,36 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "eu.anthropic.claude-sonnet-4-5-20250929-v1:0": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, "eu.meta.llama3-2-1b-instruct-v1:0": { "input_cost_per_token": 1.3e-07, "litellm_provider": "bedrock", @@ -12727,6 +12873,34 @@ "supports_tool_choice": true, "supports_vision": true }, + "gpt-5-codex": { + "cache_read_input_token_cost": 1.25e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "openai", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "gpt-5-2025-08-07": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, @@ -12802,9 +12976,9 @@ "cache_read_input_token_cost": 1.25e-07, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, "supported_endpoints": [ @@ -13557,6 +13731,19 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "groq/moonshotai/kimi-k2-instruct-0905": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 0.5e-06, + "litellm_provider": "groq", + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "max_tokens": 278528, + "mode": "chat", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "groq/openai/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, "litellm_provider": "groq", @@ -14032,6 +14219,36 @@ "mode": "rerank", "output_cost_per_token": 1.8e-08 }, + "jp.anthropic.claude-sonnet-4-5-20250929-v1:0": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346 + }, "lambda_ai/deepseek-llama3.3-70b": { "input_cost_per_token": 2e-07, "litellm_provider": "lambda_ai", @@ -18306,6 +18523,20 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { + "input_cost_per_query": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "nvidia_nim", + "mode": "rerank", + "output_cost_per_token": 0.0 + }, + "nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2": { + "input_cost_per_query": 0.0, + "input_cost_per_token": 0.0, + "litellm_provider": "nvidia_nim", + "mode": "rerank", + "output_cost_per_token": 0.0 + }, "sagemaker/meta-textgeneration-llama-2-13b": { "input_cost_per_token": 0.0, "litellm_provider": "sagemaker", @@ -19621,6 +19852,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, "max_output_tokens": 64000, @@ -20987,6 +21222,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -21009,6 +21248,10 @@ "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, diff --git a/poetry.lock b/poetry.lock index cc6ec17e042..bfb0ab0db5e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -3804,15 +3804,15 @@ files = [ [[package]] name = "litellm-proxy-extras" -version = "0.2.22" +version = "0.2.23" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." optional = true python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" groups = ["main"] markers = "extra == \"proxy\"" files = [ - {file = "litellm_proxy_extras-0.2.22-py3-none-any.whl", hash = "sha256:e64b19b48e8d84cad56bb136c7f31d9ae601a10628327c922634d7081803c205"}, - {file = "litellm_proxy_extras-0.2.22.tar.gz", hash = "sha256:59c395bff3353de57d67b7637e8ce0a8a4e096ce55e2ee2df4d9d4bda94f6ef0"}, + {file = "litellm_proxy_extras-0.2.23-py3-none-any.whl", hash = "sha256:1bc6cc13e885182ad8b319d422301ddf6db6d699c8175c0fe4cca1377ab8f759"}, + {file = "litellm_proxy_extras-0.2.23.tar.gz", hash = "sha256:2ef011e9b3980c7c763fad719142798e5796d60d574873df91a0cf13a0075ecf"}, ] [[package]] @@ -9598,4 +9598,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.8.1,<4.0, !=3.9.7" -content-hash = "dd6b1b42d43c2049fd8fcc95a6627581c5d9c60b3afd5eab60659d8f5d6ae641" +content-hash = "c240325a8f455ec48338b6abb5eeda52c769d07db51e17910b86a11631ea77d0" diff --git a/pyproject.toml b/pyproject.toml index 3ac436ca3b2..ffb654bf0c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.77.6" +version = "1.77.7" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -59,7 +59,7 @@ websockets = {version = "^13.1.0", optional = true} boto3 = {version = "1.36.0", optional = true} redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = "^1.10.0", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "0.2.22", optional = true} +litellm-proxy-extras = {version = "0.2.23", optional = true} rich = {version = "13.7.1", optional = true} litellm-enterprise = {version = "0.1.20", optional = true} diskcache = {version = "^5.6.1", optional = true} @@ -157,7 +157,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.77.6" +version = "1.77.7" version_files = [ "pyproject.toml:^version" ] diff --git a/requirements.txt b/requirements.txt index b0ad77aa865..c2b3f0db880 100644 --- a/requirements.txt +++ b/requirements.txt @@ -43,7 +43,7 @@ sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests cryptography==44.0.1 tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.2.22 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.2.23 # for proxy extras - e.g. prisma migrations ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env tiktoken==0.8.0 # for calculating usage diff --git a/schema.prisma b/schema.prisma index 766625145f6..5a79e171438 100644 --- a/schema.prisma +++ b/schema.prisma @@ -178,6 +178,7 @@ model LiteLLM_MCPServerTable { updated_by String? mcp_info Json? @default("{}") mcp_access_groups String[] + allowed_tools String[] @default([]) // Health check status status String? @default("unknown") last_health_check DateTime? diff --git a/tests/README.MD b/tests/README.MD index ed9ac10e9dc..57275a031f7 100644 --- a/tests/README.MD +++ b/tests/README.MD @@ -1,9 +1,9 @@ -**In total litellm runs 1000+ tests** +**In total litellm runs 1000+ tests** [02/20/2025] Update: To make it easier to contribute and map what behavior is tested, -we've started mapping the litellm directory in `tests/litellm` +we've started mapping the litellm directory in `tests/test_litellm` -This folder can only run mock tests. \ No newline at end of file +This folder can only run mock tests. diff --git a/tests/code_coverage_tests/info_log_check.py b/tests/code_coverage_tests/info_log_check.py index e8541b4358a..44e73a6c216 100644 --- a/tests/code_coverage_tests/info_log_check.py +++ b/tests/code_coverage_tests/info_log_check.py @@ -126,8 +126,13 @@ class SensitiveLogDetector(ast.NodeVisitor): for value in arg.values: if isinstance(value, ast.FormattedValue): value_str = self._get_arg_string(value.value).lower() - if any(pattern in value_str for pattern in - ['request', 'response', 'data', 'body', 'content', 'messages']): + # Check for any sensitive data patterns in f-string interpolations + sensitive_f_string_patterns = [ + 'request', 'response', 'data', 'body', 'content', 'messages', + 'token', 'jwt', 'auth', 'api_key', 'apikey', 'credential', + 'secret', 'password', 'passwd' + ] + if any(pattern in value_str for pattern in sensitive_f_string_patterns): return True # Check for .format() calls @@ -137,10 +142,14 @@ class SensitiveLogDetector(ast.NodeVisitor): base_str = self._get_arg_string(arg.func.value).lower() if "{}" in base_str or "{" in base_str: # Check format arguments for sensitive data + sensitive_format_patterns = [ + 'request', 'response', 'data', 'body', 'content', + 'token', 'jwt', 'auth', 'api_key', 'apikey', 'credential', + 'secret', 'password', 'passwd' + ] for format_arg in arg.args: format_str = self._get_arg_string(format_arg).lower() - if any(pattern in format_str for pattern in - ['request', 'response', 'data', 'body', 'content']): + if any(pattern in format_str for pattern in sensitive_format_patterns): return True return False @@ -171,7 +180,9 @@ class SensitiveLogDetector(ast.NodeVisitor): """Get a human-readable reason for the violation""" arg_str = self._get_arg_string(arg).lower() - if 'request' in arg_str: + if any(pattern in arg_str for pattern in ['jwt', 'token', 'api_key', 'apikey', 'auth', 'credential', 'secret', 'password', 'passwd']): + return "Potentially logging authentication/secret data (JWT, token, API key, etc.)" + elif 'request' in arg_str: return "Potentially logging request data" elif 'response' in arg_str: return "Potentially logging response data" @@ -179,8 +190,6 @@ class SensitiveLogDetector(ast.NodeVisitor): return "Potentially logging sensitive data/body/content" elif any(pattern in arg_str for pattern in ['messages', 'input', 'output']): return "Potentially logging message/input/output data" - elif any(pattern in arg_str for pattern in ['api_key', 'token', 'auth', 'credentials']): - return "Potentially logging authentication data" else: return "Potentially logging sensitive data" diff --git a/tests/guardrails_tests/conftest.py b/tests/guardrails_tests/conftest.py new file mode 100644 index 00000000000..e47df872d3f --- /dev/null +++ b/tests/guardrails_tests/conftest.py @@ -0,0 +1,79 @@ +# conftest.py + +import importlib +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +import asyncio + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """ + This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. + """ + curr_dir = os.getcwd() # Get the current working directory + sys.path.insert( + 0, os.path.abspath("../..") + ) # Adds the project directory to the system path + + import litellm + from litellm import Router + import asyncio + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + # flush all logs + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + + + importlib.reload(litellm) + + try: + if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): + import litellm.proxy.proxy_server + + importlib.reload(litellm.proxy.proxy_server) + except Exception as e: + print(f"Error reloading litellm.proxy.proxy_server: {e}") + + import asyncio + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + print(litellm) + # from litellm import Router, completion, aembedding, acompletion, embedding + yield + + # Teardown code (executes after the yield point) + loop.close() # Close the loop created earlier + asyncio.set_event_loop(None) # Remove the reference to the loop + + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index b98a1b16bef..7ca78d83ae9 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1366,3 +1366,66 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): except Exception as e: pytest.fail(f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}") + +@pytest.mark.asyncio +async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): + """Test that async_post_call_success_hook skips when there's no output text""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import ModelResponseStream + import litellm + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + # Create guardrail instance + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT" + ) + + # Create a ModelResponse with tool calls (no text content) + # This simulates a response where the LLM is making a tool call + mock_response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", + content=None, # No text content + tool_calls=[ + litellm.utils.ChatCompletionMessageToolCall( + id="tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", + function=litellm.utils.Function( + name="top_song", + arguments='{"sign": "WZPZ"}' + ), + type="function" + ) + ] + ), + finish_reason="tool_calls" + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion" + ) + + data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Hello"}, + ], + } + mock_user_api_key_dict = UserAPIKeyAuth() + + result = await guardrail.async_post_call_success_hook( + data=data, + response=mock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + # If no error is raised and result is None, then the test passes + assert result is None + print("✅ No output text in response test passed") \ No newline at end of file diff --git a/tests/guardrails_tests/test_tracing_guardrails.py b/tests/guardrails_tests/test_tracing_guardrails.py index 0299d3fe6a2..d7589c53879 100644 --- a/tests/guardrails_tests/test_tracing_guardrails.py +++ b/tests/guardrails_tests/test_tracing_guardrails.py @@ -15,8 +15,9 @@ from litellm.types.guardrails import GuardrailEventHooks from typing import Optional -class TestCustomLogger(CustomLogger): +class CustomLoggerForTesting(CustomLogger): def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) self.standard_logging_payload: Optional[StandardLoggingPayload] = None async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -28,7 +29,7 @@ async def test_standard_logging_payload_includes_guardrail_information(): """ Test that the standard logging payload includes the guardrail information when a guardrail is applied """ - test_custom_logger = TestCustomLogger() + test_custom_logger = CustomLoggerForTesting() litellm.callbacks = [test_custom_logger] presidio_guard = _OPTIONAL_PresidioPIIMasking( guardrail_name="presidio_guard", @@ -177,4 +178,469 @@ async def test_langfuse_trace_includes_guardrail_information(): assert output_item["entity_type"] == "PHONE_NUMBER" assert "score" in output_item assert "start" in output_item - assert "end" in output_item \ No newline at end of file + assert "end" in output_item + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_status_blocked(): + """ + Test that Bedrock guardrail sets correct status fields when blocking content. + + This test verifies that when Bedrock guardrail blocks content: + 1. The guardrail_information contains guardrail_status="blocked" + 2. The status_fields.guardrail_status is set to "guardrail_intervened" + 3. The status_fields.llm_api_status remains "success" (mock LLM call succeeds) + """ + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + from litellm.proxy._types import UserAPIKeyAuth + from unittest.mock import AsyncMock, MagicMock, patch + litellm._turn_on_debug() + + # Setup custom logger to capture standard logging payload + test_custom_logger = CustomLoggerForTesting() + litellm.callbacks = [test_custom_logger] + + # Create Bedrock guardrail with mock AWS credentials + bedrock_guard = BedrockGuardrail( + guardrail_name="bedrock_guard", + event_hook=GuardrailEventHooks.pre_call, + guardrailIdentifier="test-id", + guardrailVersion="1", + aws_access_key_id="test-key", + aws_secret_access_key="test-secret", + aws_region_name="us-east-1", + ) + + # Mock Bedrock API response indicating content was blocked + # action="GUARDRAIL_INTERVENED" means the guardrail blocked the request + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Blocked"}], + "assessments": [{ + "topicPolicy": { + "topics": [{"name": "harmful", "action": "BLOCKED"}] + } + }] + } + bedrock_guard.async_handler.post = AsyncMock(return_value=mock_response) + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "harmful content"}], + "mock_response": "Hello", + "metadata": {} + } + + # Mock should_run_guardrail to ensure guardrail logic executes + with patch.object(bedrock_guard, 'should_run_guardrail', return_value=True): + # Call guardrail pre_call hook - this will raise an exception when content is blocked + try: + await bedrock_guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=None, + data=request_data, + call_type="completion" + ) + except Exception: + # Expected exception when guardrail blocks content + pass + + # Call litellm.acompletion to trigger logging callbacks + # This populates the standard_logging_payload in our custom logger + response = await litellm.acompletion(**request_data) + await asyncio.sleep(1) + + # Verify the standard logging payload was captured + assert test_custom_logger.standard_logging_payload is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"] is not None + + # Verify guardrail information fields + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_status"] == "guardrail_intervened" + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_provider"] == "bedrock" + + # Verify the new typed status fields + # guardrail_status should be "guardrail_intervened" when content is blocked + # llm_api_status should be "success" since the mock LLM call itself succeeded + status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {}) + assert status_fields.get("llm_api_status") == "success" + assert status_fields.get("guardrail_status") == "guardrail_intervened" + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_status_success(): + """ + Test that Bedrock guardrail sets correct status fields when allowing content. + + This test verifies that when Bedrock guardrail allows content through: + 1. The guardrail_information contains guardrail_status="success" + 2. The status_fields.guardrail_status is set to "success" + 3. The status_fields.llm_api_status is "success" + """ + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + from litellm.proxy._types import UserAPIKeyAuth + from unittest.mock import AsyncMock, MagicMock, patch + + # Reset callbacks completely to avoid event loop conflicts + litellm.callbacks = [] + await asyncio.sleep(0.1) # Let previous callbacks finish + + # Setup custom logger to capture standard logging payload + test_custom_logger = CustomLoggerForTesting() + litellm.callbacks = [test_custom_logger] + + # Create Bedrock guardrail + bedrock_guard = BedrockGuardrail( + guardrail_name="bedrock_guard", + event_hook=GuardrailEventHooks.pre_call, + guardrailIdentifier="test-id", + guardrailVersion="1", + aws_access_key_id="test-key", + aws_secret_access_key="test-secret", + aws_region_name="us-east-1", + ) + + # Mock success response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "action": "NONE", + "outputs": [{"text": "Safe content"}], + "assessments": [] + } + bedrock_guard.async_handler.post = AsyncMock(return_value=mock_response) + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "safe content"}], + "mock_response": "Hello", + "metadata": {} + } + + # Mock should_run_guardrail to return True + with patch.object(bedrock_guard, 'should_run_guardrail', return_value=True): + await bedrock_guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=None, + data=request_data, + call_type="completion" + ) + + # Call litellm.acompletion to trigger logging + response = await litellm.acompletion(**request_data) + await asyncio.sleep(1) + + # Check standard logging payload status fields + assert test_custom_logger.standard_logging_payload is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"] is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_status"] == "success" + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_provider"] == "bedrock" + + # Check status fields + status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {}) + assert status_fields.get("llm_api_status") == "success" + assert status_fields.get("guardrail_status") == "success" + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_status_failure(): + """ + Test that Bedrock guardrail sets correct status fields when the API endpoint fails. + + This test verifies that when Bedrock guardrail API is down/fails: + 1. The guardrail_information contains guardrail_status="failure" + 2. The status_fields.guardrail_status is set to "guardrail_failed_to_respond" + 3. The exception is still raised (maintaining existing behavior) + """ + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + from litellm.proxy._types import UserAPIKeyAuth + from unittest.mock import AsyncMock, MagicMock, patch + import httpx + + # Reset callbacks completely to avoid event loop conflicts + litellm.callbacks = [] + await asyncio.sleep(0.1) + + # Setup custom logger to capture standard logging payload + test_custom_logger = CustomLoggerForTesting() + litellm.callbacks = [test_custom_logger] + + # Create Bedrock guardrail + bedrock_guard = BedrockGuardrail( + guardrail_name="bedrock_guard", + event_hook=GuardrailEventHooks.pre_call, + guardrailIdentifier="test-id", + guardrailVersion="1", + aws_access_key_id="test-key", + aws_secret_access_key="test-secret", + aws_region_name="us-east-1", + ) + + # Mock network failure (endpoint down) + bedrock_guard.async_handler.post = AsyncMock( + side_effect=httpx.ConnectError("Connection failed") + ) + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "test content"}], + "mock_response": "Hello", + "metadata": {} + } + + # Mock should_run_guardrail to return True + with patch.object(bedrock_guard, 'should_run_guardrail', return_value=True): + # Call guardrail (will raise exception on network failure) + try: + await bedrock_guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=None, + data=request_data, + call_type="completion" + ) + except Exception: + # Expected exception when endpoint is down + pass + + # Call litellm.acompletion to trigger logging + response = await litellm.acompletion(**request_data) + await asyncio.sleep(1) + + # Check standard logging payload status fields + assert test_custom_logger.standard_logging_payload is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"] is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_status"] == "guardrail_failed_to_respond" + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_provider"] == "bedrock" + + # Check status fields + status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {}) + assert status_fields.get("llm_api_status") == "success" + assert status_fields.get("guardrail_status") == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_noma_guardrail_status_blocked(): + """ + Test that Noma guardrail sets correct status fields when blocking content. + + This test verifies that when Noma guardrail blocks content (verdict=False): + 1. The guardrail_information contains guardrail_status="blocked" + 2. The status_fields.guardrail_status is set to "guardrail_intervened" + 3. The status_fields.llm_api_status remains "success" + """ + from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaGuardrail + from litellm.proxy._types import UserAPIKeyAuth + from unittest.mock import AsyncMock, MagicMock, patch + + # Reset callbacks completely to avoid event loop conflicts + litellm.callbacks = [] + await asyncio.sleep(0.1) # Let previous callbacks finish + + # Setup custom logger to capture standard logging payload + test_custom_logger = CustomLoggerForTesting() + litellm.callbacks = [test_custom_logger] + + # Create Noma guardrail + noma_guard = NomaGuardrail( + guardrail_name="noma_guard", + event_hook=GuardrailEventHooks.pre_call, + api_key="test-key", + monitor_mode=False, + ) + + # Mock blocked response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "verdict": False, + "originalResponse": { + "prompt": { + "topicDetector": {"harmful": {"result": True}} + } + } + } + mock_response.raise_for_status = MagicMock() + noma_guard.async_handler.post = AsyncMock(return_value=mock_response) + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "harmful content"}], + "mock_response": "Hello", + "metadata": {} + } + + # Mock should_run_guardrail to return True + with patch.object(noma_guard, 'should_run_guardrail', return_value=True): + # Call guardrail (will raise exception on block) + try: + await noma_guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=None, + data=request_data, + call_type="completion" + ) + except Exception: + pass + + # Call litellm.acompletion to trigger logging + response = await litellm.acompletion(**request_data) + await asyncio.sleep(1) + + # Check standard logging payload status fields + assert test_custom_logger.standard_logging_payload is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"] is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_status"] == "guardrail_intervened" + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_provider"] == "noma" + + # Check status fields + status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {}) + assert status_fields.get("llm_api_status") == "success" + assert status_fields.get("guardrail_status") == "guardrail_intervened" + + +@pytest.mark.asyncio +async def test_noma_guardrail_status_success(): + """ + Test that Noma guardrail sets correct status fields when allowing content. + + This test verifies that when Noma guardrail allows content (verdict=True): + 1. The guardrail_information contains guardrail_status="success" + 2. The status_fields.guardrail_status is set to "success" + 3. The status_fields.llm_api_status is "success" + """ + from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaGuardrail + from litellm.proxy._types import UserAPIKeyAuth + from unittest.mock import AsyncMock, MagicMock, patch + + # Reset callbacks completely to avoid event loop conflicts + litellm.callbacks = [] + await asyncio.sleep(0.1) # Let previous callbacks finish + + # Setup custom logger to capture standard logging payload + test_custom_logger = CustomLoggerForTesting() + litellm.callbacks = [test_custom_logger] + + # Create Noma guardrail + noma_guard = NomaGuardrail( + guardrail_name="noma_guard", + event_hook=GuardrailEventHooks.pre_call, + api_key="test-key", + monitor_mode=False, + ) + + # Mock success response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "verdict": True, + "originalResponse": {"prompt": {}} + } + mock_response.raise_for_status = MagicMock() + noma_guard.async_handler.post = AsyncMock(return_value=mock_response) + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "safe content"}], + "mock_response": "Hello", + "metadata": {} + } + + # Mock should_run_guardrail to return True + with patch.object(noma_guard, 'should_run_guardrail', return_value=True): + await noma_guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=None, + data=request_data, + call_type="completion" + ) + + # Call litellm.acompletion to trigger logging + response = await litellm.acompletion(**request_data) + await asyncio.sleep(1) + + # Check standard logging payload status fields + assert test_custom_logger.standard_logging_payload is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"] is not None + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_status"] == "success" + assert test_custom_logger.standard_logging_payload["guardrail_information"]["guardrail_provider"] == "noma" + + # Check status fields + status_fields = test_custom_logger.standard_logging_payload.get("status_fields", {}) + assert status_fields.get("llm_api_status") == "success" + assert status_fields.get("guardrail_status") == "success" + + +def test_guardrail_status_fields_computation(): + """ + Test that status fields are computed correctly from guardrail information. + + This unit test verifies the _get_status_fields function correctly maps: + - guardrail_status="blocked" -> status_fields.guardrail_status="guardrail_intervened" (legacy) + - guardrail_status="guardrail_intervened" -> status_fields.guardrail_status="guardrail_intervened" + - guardrail_status="success" -> status_fields.guardrail_status="success" + - guardrail_status="failure" -> status_fields.guardrail_status="guardrail_failed_to_respond" (legacy) + - guardrail_status="guardrail_failed_to_respond" -> status_fields.guardrail_status="guardrail_failed_to_respond" + - no guardrail -> status_fields.guardrail_status="not_run" + """ + from litellm.litellm_core_utils.litellm_logging import _get_status_fields + + # Test guardrail_intervened status (content was blocked by guardrail) + intervened_info = {"guardrail_status": "guardrail_intervened"} + status_fields_intervened = _get_status_fields( + status="success", + guardrail_information=intervened_info, + error_str=None + ) + assert status_fields_intervened["llm_api_status"] == "success" + assert status_fields_intervened["guardrail_status"] == "guardrail_intervened" + + # Test legacy blocked status (for backward compatibility) + blocked_info = {"guardrail_status": "blocked"} + status_fields_blocked = _get_status_fields( + status="success", + guardrail_information=blocked_info, + error_str=None + ) + assert status_fields_blocked["llm_api_status"] == "success" + assert status_fields_blocked["guardrail_status"] == "guardrail_intervened" + + # Test success status + success_info = {"guardrail_status": "success"} + status_fields_success = _get_status_fields( + status="success", + guardrail_information=success_info, + error_str=None + ) + assert status_fields_success["llm_api_status"] == "success" + assert status_fields_success["guardrail_status"] == "success" + + # Test guardrail_failed_to_respond status + failed_info = {"guardrail_status": "guardrail_failed_to_respond"} + status_fields_failed = _get_status_fields( + status="failure", + guardrail_information=failed_info, + error_str=None + ) + assert status_fields_failed["llm_api_status"] == "failure" + assert status_fields_failed["guardrail_status"] == "guardrail_failed_to_respond" + + # Test legacy failure status (for backward compatibility) + failure_info = {"guardrail_status": "failure"} + status_fields_failure = _get_status_fields( + status="failure", + guardrail_information=failure_info, + error_str=None + ) + assert status_fields_failure["llm_api_status"] == "failure" + assert status_fields_failure["guardrail_status"] == "guardrail_failed_to_respond" + + # Test no guardrail run + no_guardrail = None + status_fields_no_guardrail = _get_status_fields( + status="success", + guardrail_information=no_guardrail, + error_str=None + ) + assert status_fields_no_guardrail["llm_api_status"] == "success" + assert status_fields_no_guardrail["guardrail_status"] == "not_run" \ No newline at end of file diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 7d0e593528a..aec4a88d60f 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1426,6 +1426,44 @@ def test_get_end_user_id_for_cost_tracking_prometheus_only( ) +@pytest.mark.parametrize( + "litellm_params, expected_end_user_id", + [ + # Test with only metadata field (old behavior) + ({"metadata": {"user_api_key_end_user_id": "user_from_metadata"}}, "user_from_metadata"), + # Test with only litellm_metadata field (new behavior) + ({"litellm_metadata": {"user_api_key_end_user_id": "user_from_litellm_metadata"}}, "user_from_litellm_metadata"), + # Test with both fields - metadata should take precedence for user_api_key fields + ({"metadata": {"user_api_key_end_user_id": "user_from_metadata"}, + "litellm_metadata": {"user_api_key_end_user_id": "user_from_litellm_metadata"}}, + "user_from_metadata"), + # Test with user_api_key_end_user_id in litellm_params (should take precedence over metadata) + ({"user_api_key_end_user_id": "user_from_params", + "metadata": {"user_api_key_end_user_id": "user_from_metadata"}}, + "user_from_params"), + # Test with empty metadata but valid litellm_metadata + ({"metadata": {}, "litellm_metadata": {"user_api_key_end_user_id": "user_from_litellm_metadata"}}, + "user_from_litellm_metadata"), + # Test with no metadata fields + ({}, None), + ], +) +def test_get_end_user_id_for_cost_tracking_metadata_handling( + litellm_params, expected_end_user_id +): + """ + Test that get_end_user_id_for_cost_tracking correctly handles both metadata and litellm_metadata + fields using the get_litellm_metadata_from_kwargs helper function. + """ + from litellm.utils import get_end_user_id_for_cost_tracking + + # Ensure cost tracking is enabled for this test + litellm.disable_end_user_cost_tracking = False + + result = get_end_user_id_for_cost_tracking(litellm_params=litellm_params) + assert result == expected_end_user_id + + def test_is_prompt_caching_enabled_error_handling(): """ Assert that `is_prompt_caching_valid_prompt` safely handles errors in `token_counter`. diff --git a/tests/llm_translation/base_rerank_unit_tests.py b/tests/llm_translation/base_rerank_unit_tests.py index cff4a02753c..ac62dbbd8b8 100644 --- a/tests/llm_translation/base_rerank_unit_tests.py +++ b/tests/llm_translation/base_rerank_unit_tests.py @@ -83,6 +83,14 @@ class BaseLLMRerankTest(ABC): """Must return the custom llm provider""" pass + def get_expected_cost(self) -> float: + """ + Override this method to set the expected cost for the rerank call. + Default is None, which means the test will check cost > 0. + Return 0.0 for free models. + """ + return None + @pytest.mark.asyncio() @pytest.mark.parametrize("sync_mode", [True, False]) async def test_basic_rerank(self, sync_mode): @@ -105,7 +113,18 @@ class BaseLLMRerankTest(ABC): assert response.results is not None assert response._hidden_params["response_cost"] is not None - assert response._hidden_params["response_cost"] > 0 + + # Check expected cost + expected_cost = self.get_expected_cost() + if expected_cost is not None: + # If expected cost is specified, check exact match or >= for 0 + if expected_cost == 0.0: + assert response._hidden_params["response_cost"] >= 0 + else: + assert response._hidden_params["response_cost"] == expected_cost + else: + # Default behavior: cost should be greater than 0 + assert response._hidden_params["response_cost"] > 0 assert_response_shape( response=response, custom_llm_provider=custom_llm_provider.value diff --git a/tests/llm_translation/test_bedrock_embedding.py b/tests/llm_translation/test_bedrock_embedding.py index 15a615b8cc4..88918674aaa 100644 --- a/tests/llm_translation/test_bedrock_embedding.py +++ b/tests/llm_translation/test_bedrock_embedding.py @@ -140,7 +140,8 @@ def test_e2e_bedrock_embedding_image_twelvelabs_marengo(): response = litellm.embedding( model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", input=[duck_img_base64], - aws_region_name="us-east-1" + aws_region_name="us-east-1", + input_type="image" ) # Validate response structure @@ -170,6 +171,92 @@ def test_e2e_bedrock_embedding_image_twelvelabs_marengo(): print(f"Image embedding successful! Vector size: {len(embedding_obj.embedding)}, Response: {response}") + # Restore original region name + if original_region_name: + os.environ["AWS_REGION_NAME"] = original_region_name + + +def test_e2e_bedrock_async_invoke_embedding_twelvelabs_marengo(): + """ + Test async invoke embedding with TwelveLabs Marengo. + Validates that async invoke responses include job ID in hidden parameters. + """ + print("Testing async invoke embedding...") + original_region_name = os.environ.get("AWS_REGION_NAME") + os.environ["AWS_REGION_NAME"] = "us-east-1" + litellm._turn_on_debug() + + # Mock the HTTP call to return async invoke response + with patch("litellm.llms.bedrock.embed.embedding.BedrockEmbedding._make_sync_call") as mock_call: + mock_call.return_value = { + "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/test-job-123" + } + + response = litellm.embedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["Hello world from LiteLLM async invoke!"], + aws_region_name="us-east-1", + inputType="text", + output_s3_uri="s3://test-bucket/async-invoke-output/" + ) + + # Validate response structure + assert isinstance(response, litellm.EmbeddingResponse), "Response should be EmbeddingResponse type" + assert hasattr(response, '_hidden_params'), "Response should have _hidden_params" + assert response._hidden_params is not None, "Hidden params should not be None" + + # Validate hidden params contain invocation ARN + assert hasattr(response._hidden_params, '_invocation_arn'), "Hidden params should have _invocation_arn" + assert response._hidden_params._invocation_arn == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/test-job-123", "Invocation ARN should be preserved" + + # Validate embedding structure + assert len(response.data) == 1, "Should have one embedding" + assert response.data[0].object == "embedding", "Embedding object should be 'embedding'" + assert response.data[0].embedding == [], "Embedding should be empty for async jobs" + + print(f"Async invoke embedding successful! Invocation ARN: {response._hidden_params._invocation_arn}") + + # Restore original region name + if original_region_name: + os.environ["AWS_REGION_NAME"] = original_region_name + + +@pytest.mark.asyncio +async def test_e2e_bedrock_async_invoke_embedding_async_twelvelabs_marengo(): + """ + Test async invoke embedding with async calls. + Validates that async invoke responses work with aembedding. + """ + print("Testing async invoke embedding with async calls...") + original_region_name = os.environ.get("AWS_REGION_NAME") + os.environ["AWS_REGION_NAME"] = "us-east-1" + litellm._turn_on_debug() + + # Mock the async HTTP call to return async invoke response + with patch("litellm.llms.bedrock.embed.embedding.BedrockEmbedding._make_async_call") as mock_call: + mock_call.return_value = { + "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/test-async-job-456" + } + + response = await litellm.aembedding( + model="bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0", + input=["Hello world from LiteLLM async invoke async!"], + aws_region_name="us-east-1", + inputType="text", + output_s3_uri="s3://test-bucket/async-invoke-output/" + ) + + # Validate response structure + assert isinstance(response, litellm.EmbeddingResponse), "Response should be EmbeddingResponse type" + assert hasattr(response, '_hidden_params'), "Response should have _hidden_params" + assert response._hidden_params is not None, "Hidden params should not be None" + + # Validate hidden params contain invocation ARN + assert hasattr(response._hidden_params, '_invocation_arn'), "Hidden params should have _invocation_arn" + assert response._hidden_params._invocation_arn == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/test-async-job-456", "Invocation ARN should be preserved" + + print(f"Async invoke embedding successful! Invocation ARN: {response._hidden_params._invocation_arn}") + # Restore original region name if original_region_name: os.environ["AWS_REGION_NAME"] = original_region_name \ No newline at end of file diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index c34e59a73ec..871aebdc9fd 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -465,11 +465,11 @@ def test_gemini_url_context(): from litellm import completion litellm._turn_on_debug() + URL1 = "https://www.foodnetwork.com/recipes/ina-garten/perfect-roast-chicken-recipe-1940592" - url = "https://ai.google.dev/gemini-api/docs/models" prompt = f""" - Summarize this document: - {url} + Get the recipes listed on the following website + {URL1} """ response = completion( model="gemini/gemini-2.5-flash", @@ -482,7 +482,7 @@ def test_gemini_url_context(): url_context_metadata = response.model_extra["vertex_ai_url_context_metadata"] assert url_context_metadata is not None urlMetadata = url_context_metadata[0]["urlMetadata"][0] - assert urlMetadata["retrievedUrl"] == url + assert urlMetadata["retrievedUrl"] == URL1 assert urlMetadata["urlRetrievalStatus"] == "URL_RETRIEVAL_STATUS_SUCCESS" diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 89513c58b07..1705871258f 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/tests/llm_translation/test_nvidia_nim.py @@ -17,6 +17,8 @@ from unittest.mock import patch, MagicMock, AsyncMock import litellm from litellm import Choices, Message, ModelResponse, EmbeddingResponse, Usage from litellm import completion +from base_rerank_unit_tests import BaseLLMRerankTest +import litellm def test_completion_nvidia_nim(): @@ -181,3 +183,16 @@ def test_chat_completion_nvidia_nim_with_tools(): assert request_body["tools"] == tools assert request_body["tool_choice"] == "auto" assert request_body["parallel_tool_calls"] == True + +class TestNvidiaNim(BaseLLMRerankTest): + def get_custom_llm_provider(self) -> litellm.LlmProviders: + return litellm.LlmProviders.NVIDIA_NIM + + def get_base_rerank_call_args(self) -> dict: + return { + "model": "nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2", + } + + def get_expected_cost(self) -> float: + """Nvidia NIM rerank models are free (cost = 0.0)""" + return 0.0 \ No newline at end of file diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 8262d43a0d6..d533418117c 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -3846,3 +3846,33 @@ def test_gemini_grounding_on_streaming(): vertex_ai_grounding_metadata_shows_up = True print(chunk) assert vertex_ai_grounding_metadata_shows_up + + +def test_gemini_google_maps_tool_simple(): + """ + Test googleMaps tool with just enableWidget parameter. + """ + load_vertex_ai_credentials() + litellm._turn_on_debug() + + tools = [{"googleMaps": {"enableWidget": True}}] + tools_with_location = [{"googleMaps": {"enableWidget": True, "latitude": 37.7749, "longitude": -122.4194, "languageCode": "en_US"}}] + try: + for tools in [tools, tools_with_location]: + response = completion( + model="vertex_ai/gemini-2.0-flash", + messages=[ + { + "role": "user", + "content": "What restaurants are nearby?", + } + ], + tools=tools, + ) + print(f"Response: {response.model_dump_json(indent=4)}") + assert response.choices[0].message.content is not None + except litellm.RateLimitError: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") + diff --git a/tests/local_testing/test_pass_through_endpoints.py b/tests/local_testing/test_pass_through_endpoints.py index 9c5f5995218..29cf9682a7c 100644 --- a/tests/local_testing/test_pass_through_endpoints.py +++ b/tests/local_testing/test_pass_through_endpoints.py @@ -117,7 +117,7 @@ async def test_pass_through_endpoint_rerank(client): { "path": "/v1/rerank", "target": "https://api.cohere.com/v1/rerank", - "headers": {"Authorization": f"bearer {_cohere_api_key}"}, + "headers": {"Authorization": f"Bearer {_cohere_api_key}"}, } ] @@ -193,7 +193,7 @@ async def test_pass_through_endpoint_rpm_limit( "path": "/v1/rerank", "target": "https://api.cohere.com/v1/rerank", "auth": auth, - "headers": {"Authorization": f"bearer {_cohere_api_key}"}, + "headers": {"Authorization": f"Bearer {_cohere_api_key}"}, } ] @@ -293,7 +293,7 @@ async def test_pass_through_endpoint_sequential_rpm_limit( "path": "/v1/rerank", "target": "https://api.cohere.com/v1/rerank", "auth": auth, - "headers": {"Authorization": f"bearer {_cohere_api_key}"}, + "headers": {"Authorization": f"Bearer {_cohere_api_key}"}, } ] diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion.json index 74106b19b37..50f4db61f9a 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion.json @@ -37,7 +37,11 @@ "cache_key": null, "api_base": "https://api.openai.com", "response_cost": 3.5e-05, - "additional_headers": {} + "additional_headers": {}, + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 3.5e-05, "cache_hit": false, @@ -62,7 +66,9 @@ "endTime": "2025-01-16T11:28:55.124353-08:00", "completionStartTime": "2025-01-16T11:28:55.124353-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -70,11 +76,12 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "input": 10, - "output": 20 - }, + "input": 10, + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 + }, "traceId": "litellm-test-6a51ae70-a4e7-499e-afcd-dce2a3b31850" }, "timestamp": "2025-01-16T19:28:55.125258Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json index d9f52477fc8..26b712c1cf2 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json @@ -65,11 +65,12 @@ "totalCost": 0.00018 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "input": 10, - "output": 10 - } + "input": 10, + "output": 10, + "total": 20, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 + } }, "timestamp": "2025-05-26T21:13:16.797156Z" } diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_complex_metadata.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_complex_metadata.json index 5a5a32de2eb..62cb01dfbfd 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_complex_metadata.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_complex_metadata.json @@ -78,7 +78,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -103,7 +106,9 @@ "endTime": "2025-01-22T09:27:51.702048-08:00", "completionStartTime": "2025-01-22T09:27:51.702048-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -111,10 +116,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:27:51.703046Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json index 1da758d1db8..a986a5a8aee 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json @@ -54,7 +54,10 @@ "api_base": "https://api.openai.com", "response_cost": 3.5e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 3.5e-05, "cache_hit": false, @@ -81,7 +84,9 @@ "endTime": "2025-01-22T09:19:11.234200-08:00", "completionStartTime": "2025-01-22T09:19:11.234200-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -91,6 +96,7 @@ "usageDetails": { "input": 10, "output": 20, + "total": 30, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0 } diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_no_choices.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_no_choices.json index b01c73ffd88..ff8419ee392 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_no_choices.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_no_choices.json @@ -33,7 +33,10 @@ "api_base": "https://api.openai.com", "response_cost": 3.5e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 3.5e-05, "cache_hit": false, @@ -52,7 +55,9 @@ "endTime": "2025-02-06T16:23:27.644253-08:00", "completionStartTime": "2025-02-06T16:23:27.644253-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 10, @@ -60,10 +65,11 @@ "totalCost": 1.9999999999999998e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 10 + "output": 10, + "total": 20, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-02-07T00:23:27.670175Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_router.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_router.json index feb7aff80bd..df99b11d26b 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_router.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_router.json @@ -75,10 +75,11 @@ "totalCost": 1.9999999999999998e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 10 + "output": 10, + "total": 20, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-05-24T17:01:19.408586Z" @@ -91,4 +92,4 @@ "sdk_version": "2.44.1", "public_key": "pk-lf-3bfc4db9-217f-48e9-92e0-142566e3c204" } -} +} \ No newline at end of file diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags.json index d4a3f4a57a3..f4c99a4b452 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags.json @@ -46,7 +46,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -71,7 +74,9 @@ "endTime": "2025-01-22T07:31:28.962389-08:00", "completionStartTime": "2025-01-22T07:31:28.962389-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -79,10 +84,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T15:31:28.964179Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags_stream.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags_stream.json index d31ea5eb3c9..f3b660dc678 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags_stream.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_tags_stream.json @@ -46,7 +46,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -71,7 +74,9 @@ "endTime": "2025-01-22T08:38:26.015666-08:00", "completionStartTime": "2025-01-22T08:38:26.015666-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -79,10 +84,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T16:38:26.017252Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_vertex_call.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_vertex_call.json index 3e27f5b54b4..b6c11f96953 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_vertex_call.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_vertex_call.json @@ -63,10 +63,11 @@ "totalCost": 7.5e-06 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 10 + "output": 10, + "total": 20, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-05-26T21:15:40.610953Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata.json b/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata.json index b9e8286dbf1..51f7ffa60a9 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata.json @@ -53,7 +53,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -78,7 +81,9 @@ "endTime": "2025-01-22T09:59:39.365756-08:00", "completionStartTime": "2025-01-22T09:59:39.365756-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -86,10 +91,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:59:39.368310Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata_2.json b/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata_2.json index 78358b021a1..5bd15e98cdc 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata_2.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/complex_metadata_2.json @@ -45,7 +45,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -70,7 +73,9 @@ "endTime": "2025-01-22T10:06:50.958374-08:00", "completionStartTime": "2025-01-22T10:06:50.958374-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -78,10 +83,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T18:06:50.959850Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/empty_metadata.json b/tests/logging_callback_tests/langfuse_expected_request_body/empty_metadata.json index 057278c001b..803fe752708 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/empty_metadata.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/empty_metadata.json @@ -39,7 +39,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -64,7 +67,9 @@ "endTime": "2025-01-22T09:59:32.880691-08:00", "completionStartTime": "2025-01-22T09:59:32.880691-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -72,10 +77,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:59:32.889548Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_function.json b/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_function.json index 53134281351..b9ac7aecba3 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_function.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_function.json @@ -39,7 +39,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -64,7 +67,9 @@ "endTime": "2025-01-22T09:59:36.161959-08:00", "completionStartTime": "2025-01-22T09:59:36.161959-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -72,10 +77,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:59:36.162997Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_lock.json b/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_lock.json index 057278c001b..803fe752708 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_lock.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/metadata_with_lock.json @@ -39,7 +39,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -64,7 +67,9 @@ "endTime": "2025-01-22T09:59:32.880691-08:00", "completionStartTime": "2025-01-22T09:59:32.880691-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -72,10 +77,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:59:32.889548Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/nested_metadata.json b/tests/logging_callback_tests/langfuse_expected_request_body/nested_metadata.json index a4dc890273d..aec7f2ab868 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/nested_metadata.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/nested_metadata.json @@ -45,7 +45,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -70,7 +73,9 @@ "endTime": "2025-01-22T09:55:28.853979-08:00", "completionStartTime": "2025-01-22T09:55:28.853979-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -78,10 +83,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:55:28.855732Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata.json b/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata.json index 616b64e1bc1..a05595299df 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata.json @@ -45,7 +45,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -70,7 +73,9 @@ "endTime": "2025-01-22T09:53:53.753431-08:00", "completionStartTime": "2025-01-22T09:53:53.753431-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -78,10 +83,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:53:53.754511Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata2.json b/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata2.json index 9a7b5833a83..769ab97d598 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata2.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata2.json @@ -49,7 +49,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -74,7 +77,9 @@ "endTime": "2025-01-22T09:56:35.476236-08:00", "completionStartTime": "2025-01-22T09:56:35.476236-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -82,10 +87,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:56:35.478171Z" diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata3.json b/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata3.json index c8addd47da0..0c41f0fc802 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata3.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/simple_metadata3.json @@ -53,7 +53,10 @@ "api_base": "https://api.openai.com", "response_cost": 5.4999999999999995e-05, "additional_headers": {}, - "litellm_overhead_time_ms": null + "litellm_overhead_time_ms": null, + "batch_models": null, + "litellm_model_name": "gpt-3.5-turbo", + "usage_object": null }, "litellm_response_cost": 5.4999999999999995e-05, "cache_hit": false, @@ -78,7 +81,9 @@ "endTime": "2025-01-22T09:56:38.785762-08:00", "completionStartTime": "2025-01-22T09:56:38.785762-08:00", "model": "gpt-3.5-turbo", - "modelParameters": {"extra_body": "{}"}, + "modelParameters": { + "extra_body": "{}" + }, "usage": { "input": 10, "output": 20, @@ -86,10 +91,11 @@ "totalCost": 3.5e-05 }, "usageDetails": { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, "input": 10, - "output": 20 + "output": 20, + "total": 30, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0 } }, "timestamp": "2025-01-22T17:56:38.787196Z" diff --git a/tests/logging_callback_tests/test_langfuse_e2e_test.py b/tests/logging_callback_tests/test_langfuse_e2e_test.py index 3e76a685630..866e95a702c 100644 --- a/tests/logging_callback_tests/test_langfuse_e2e_test.py +++ b/tests/logging_callback_tests/test_langfuse_e2e_test.py @@ -21,6 +21,7 @@ os.environ["LANGFUSE_DEBUG"] = "True" import time import pytest +import pytest_asyncio def assert_langfuse_request_matches_expected( @@ -116,7 +117,7 @@ def assert_langfuse_request_matches_expected( class TestLangfuseLogging: - @pytest.fixture + @pytest_asyncio.fixture async def mock_setup(self): """Common setup for Langfuse logging tests""" from litellm._uuid import uuid @@ -168,7 +169,7 @@ class TestLangfuseLogging: @pytest.mark.asyncio async def test_langfuse_logging_completion(self, mock_setup): """Test Langfuse logging for chat completion""" - setup = await mock_setup # Await the fixture + setup = mock_setup with patch("httpx.Client.post", setup["mock_post"]): await litellm.acompletion( model="gpt-3.5-turbo", @@ -183,7 +184,7 @@ class TestLangfuseLogging: @pytest.mark.asyncio async def test_langfuse_logging_completion_with_tags(self, mock_setup): """Test Langfuse logging for chat completion with tags""" - setup = await mock_setup # Await the fixture + setup = mock_setup with patch("httpx.Client.post", setup["mock_post"]): await litellm.acompletion( model="gpt-3.5-turbo", @@ -201,7 +202,7 @@ class TestLangfuseLogging: @pytest.mark.asyncio async def test_langfuse_logging_completion_with_tags_stream(self, mock_setup): """Test Langfuse logging for chat completion with tags""" - setup = await mock_setup # Await the fixture + setup = mock_setup with patch("httpx.Client.post", setup["mock_post"]): await litellm.acompletion( model="gpt-3.5-turbo", @@ -221,7 +222,7 @@ class TestLangfuseLogging: @pytest.mark.asyncio async def test_langfuse_logging_completion_with_langfuse_metadata(self, mock_setup): """Test Langfuse logging for chat completion with metadata for langfuse""" - setup = await mock_setup # Await the fixture + setup = mock_setup with patch("httpx.Client.post", setup["mock_post"]): await litellm.acompletion( model="gpt-3.5-turbo", @@ -259,7 +260,7 @@ class TestLangfuseLogging: last_login: datetime.datetime settings: dict - setup = await mock_setup + setup = mock_setup test_metadata = { "user_prefs": UserPreferences( @@ -334,7 +335,7 @@ class TestLangfuseLogging: """Test Langfuse logging with various metadata types including non-serializable objects""" import threading - setup = await mock_setup + setup = mock_setup if test_metadata is not None: test_metadata["trace_id"] = setup["trace_id"] @@ -358,7 +359,7 @@ class TestLangfuseLogging: self, mock_setup ): """Test Langfuse logging for chat completion with malformed LLM response""" - setup = await mock_setup # Await the fixture + setup = mock_setup litellm._turn_on_debug() with patch("httpx.Client.post", setup["mock_post"]): mock_response = litellm.ModelResponse( @@ -387,7 +388,7 @@ class TestLangfuseLogging: self, mock_setup ): """Test Langfuse logging for chat completion with malformed LLM response""" - setup = await mock_setup # Await the fixture + setup = mock_setup litellm._turn_on_debug() with patch("httpx.Client.post", setup["mock_post"]): mock_response = litellm.ModelResponse( @@ -418,7 +419,7 @@ class TestLangfuseLogging: self, mock_setup ): """Test Langfuse logging for chat completion with malformed LLM response""" - setup = await mock_setup # Await the fixture + setup = mock_setup litellm._turn_on_debug() with patch("httpx.Client.post", setup["mock_post"]): mock_response = litellm.ModelResponse( @@ -447,7 +448,6 @@ class TestLangfuseLogging: @pytest.mark.asyncio async def test_langfuse_logging_with_router(self, mock_setup): """Test Langfuse logging with router""" - setup = await mock_setup # Await the fixture litellm._turn_on_debug() router = litellm.Router( model_list=[ @@ -461,7 +461,7 @@ class TestLangfuseLogging: } ] ) - with patch("httpx.Client.post", setup["mock_post"]): + with patch("httpx.Client.post", mock_setup["mock_post"]): mock_response = litellm.ModelResponse( choices=[], usage=litellm.Usage( @@ -477,8 +477,8 @@ class TestLangfuseLogging: model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello!"}], mock_response=mock_response, - metadata={"trace_id": setup["trace_id"]}, + metadata={"trace_id": mock_setup["trace_id"]}, ) await self._verify_langfuse_call( - setup["mock_post"], "completion_with_router.json", setup["trace_id"] + mock_setup["mock_post"], "completion_with_router.json", mock_setup["trace_id"] ) diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index 0226fd66fdf..ffb54dd99c9 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -102,7 +102,7 @@ async def use_callback_in_llm_call( elif callback == "openmeter": # it's currently handled in jank way, TODO: fix openmete and then actually run it's test return - elif callback == "bitbucket": + elif callback == "bitbucket" or callback == "gitlab": # Set up mock bitbucket configuration required for initialization litellm.global_bitbucket_config = { "workspace": "test-workspace", @@ -110,6 +110,13 @@ async def use_callback_in_llm_call( "access_token": "test-token", "branch": "main" } + litellm.global_gitlab_config = { + "project": "a/b/", + "access_token": "your-access-token", + "base_url": "gitlab url", + "prompts_path": "src/prompts", # folder to point to, defaults to root + "branch":"main" # optional, defaults to main + } # Mock BitBucket HTTP calls to prevent actual API requests import httpx from unittest.mock import MagicMock diff --git a/tests/pass_through_tests/test_anthropic_passthrough.py b/tests/pass_through_tests/test_anthropic_passthrough.py index 002fb20e9e8..6e819f9971c 100644 --- a/tests/pass_through_tests/test_anthropic_passthrough.py +++ b/tests/pass_through_tests/test_anthropic_passthrough.py @@ -296,3 +296,123 @@ async def test_anthropic_streaming_with_headers(): assert log_entry["end_user"] == "test-user-1" assert log_entry["custom_llm_provider"] == "anthropic" + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_anthropic_messages_streaming_cost_injection(): + """ + Test that cost is injected into message_delta usage for Anthropic Messages API streaming + """ + print("Testing cost injection in Anthropic Messages API streaming response") + + headers = { + "Authorization": "Bearer sk-1234", + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + } + + payload = { + "model": "claude-3-7-sonnet-20250219", + "max_tokens": 10, + "stream": True, + "messages": [{"role": "user", "content": "Say 'Hi'"}], + } + + async with aiohttp.ClientSession() as session: + async with session.post( + "http://0.0.0.0:4000/v1/messages", + json=payload, + headers=headers + ) as response: + assert response.status == 200 + + # Collect all SSE events + events = [] + async for line in response.content: + line_str = line.decode('utf-8').strip() + if line_str.startswith('data: '): + try: + data = json.loads(line_str[6:]) # Remove 'data: ' prefix + events.append(data) + except json.JSONDecodeError: + continue + + # Find message_delta event with usage + message_delta_events = [ + event for event in events + if event.get('type') == 'message_delta' and 'usage' in event + ] + + assert len(message_delta_events) > 0, "No message_delta events with usage found" + + # Check that cost is included in usage + for event in message_delta_events: + usage = event.get('usage', {}) + assert 'cost' in usage, f"Cost not found in usage: {usage}" + assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}" + assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}" + + print(f"✅ Found message_delta with cost: {usage}") + + print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost") + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_anthropic_messages_openai_model_streaming_cost_injection(): + """ + Test that cost is injected into message_delta usage for OpenAI model via Anthropic Messages API + """ + print("Testing cost injection in Anthropic Messages API with OpenAI model") + + headers = { + "Authorization": "Bearer sk-1234", + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + } + + payload = { + "model": "openai/gpt-4o", + "max_tokens": 10, + "stream": True, + "messages": [{"role": "user", "content": "Say 'Hi'"}], + } + + async with aiohttp.ClientSession() as session: + async with session.post( + "http://0.0.0.0:4000/v1/messages", + json=payload, + headers=headers + ) as response: + assert response.status == 200 + + # Collect all SSE events + events = [] + async for line in response.content: + line_str = line.decode('utf-8').strip() + if line_str.startswith('data: '): + try: + data = json.loads(line_str[6:]) # Remove 'data: ' prefix + events.append(data) + except json.JSONDecodeError: + continue + + # Find message_delta event with usage + message_delta_events = [ + event for event in events + if event.get('type') == 'message_delta' and 'usage' in event + ] + + assert len(message_delta_events) > 0, "No message_delta events with usage found" + + # Check that cost is included in usage + for event in message_delta_events: + usage = event.get('usage', {}) + assert 'cost' in usage, f"Cost not found in usage: {usage}" + assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}" + assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}" + + print(f"✅ Found message_delta with cost: {usage}") + + print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost") diff --git a/tests/pass_through_tests/test_openai_assistants_passthrough.py b/tests/pass_through_tests/test_openai_assistants_passthrough.py index 40361ab39f7..e5783877ec0 100644 --- a/tests/pass_through_tests/test_openai_assistants_passthrough.py +++ b/tests/pass_through_tests/test_openai_assistants_passthrough.py @@ -96,3 +96,42 @@ def test_openai_assistants_e2e_operations_stream(): event_handler=EventHandler(), ) as stream: stream.until_done() + + + +def test_azure_openai_assistants_e2e_operations_stream(): + from openai import AzureOpenAI + client = AzureOpenAI( + base_url="http://0.0.0.0:4000/azure-config-passthrough/openai", + api_key="sk-1234", + api_version="2025-01-01-preview" + ) + assistant = client.beta.assistants.create( + name="Math Tutor", + instructions="You are a personal math tutor. Write and run code to answer math questions.", + tools=[{"type": "code_interpreter"}], + model="gpt-4o", + ) + print("assistant created", assistant) + + thread = client.beta.threads.create() + print("thread created", thread) + + message = client.beta.threads.messages.create( + thread_id=thread.id, + role="user", + content="I need to solve the equation `3x + 11 = 14`. Can you help me?", + ) + print("message created", message) + + # Then, we use the `stream` SDK helper + # with the `EventHandler` class to create the Run + # and stream the response. + + with client.beta.threads.runs.stream( + thread_id=thread.id, + assistant_id=assistant.id, + instructions="Please address the user as Jane Doe. The user has a premium account.", + event_handler=EventHandler(), + ) as stream: + stream.until_done() \ No newline at end of file diff --git a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py index 0c62e776e9c..501fe24e65e 100644 --- a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py +++ b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py @@ -45,6 +45,18 @@ def mock_request(): class QueryParams: def __init__(self): self._dict = {} + + def __iter__(self): + return iter(self._dict.items()) + + def items(self): + return self._dict.items() + + def keys(self): + return self._dict.keys() + + def values(self): + return self._dict.values() class MockRequest: def __init__( diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py new file mode 100644 index 00000000000..aee67a0ec39 --- /dev/null +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -0,0 +1,591 @@ +""" +Test Vertex AI Live API Passthrough Feature + +This module tests the Vertex AI Live API WebSocket passthrough functionality, +including the logging handler, cost tracking, and WebSocket message processing. +""" + +import json +import os +import sys +from datetime import datetime +from unittest.mock import AsyncMock, Mock, patch, MagicMock +from typing import Dict, List, Any, Optional + +import pytest +import httpx + +# Add the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import ( + VertexAILivePassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.types.utils import LlmProviders +from litellm.proxy._types import UserAPIKeyAuth + + +class TestVertexAILivePassthroughLoggingHandler: + """Test the Vertex AI Live Passthrough Logging Handler""" + + @pytest.fixture + def handler(self): + """Create a handler instance for testing""" + return VertexAILivePassthroughLoggingHandler() + + @pytest.fixture + def mock_logging_obj(self): + """Create a mock logging object""" + mock = MagicMock(spec=LiteLLMLoggingObj) + mock.model_call_details = {} + return mock + + @pytest.fixture + def sample_websocket_messages(self): + """Sample WebSocket messages for testing""" + return [ + { + "type": "session.created", + "session": {"id": "test-session-123"}, + "timestamp": "2024-01-01T00:00:00Z" + }, + { + "type": "response.create", + "event_id": "event-123", + "response": { + "text": "Hello, how can I help you?" + }, + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 15, + "totalTokenCount": 25, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 10} + ], + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 15} + ] + } + }, + { + "type": "response.done", + "event_id": "event-123", + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 8, + "totalTokenCount": 13, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 5} + ], + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 8} + ] + } + } + ] + + def test_llm_provider_name_property(self, handler): + """Test that llm_provider_name returns the correct provider""" + assert handler.llm_provider_name == LlmProviders.VERTEX_AI + + def test_get_provider_config(self, handler): + """Test that get_provider_config returns a valid config""" + config = handler.get_provider_config("gemini-1.5-pro") + assert config is not None + # Verify it's a Vertex AI config by checking for expected methods + assert hasattr(config, 'get_supported_openai_params') + assert hasattr(config, 'map_openai_params') + + def test_extract_usage_metadata_single_message(self, handler): + """Test usage metadata extraction from a single message""" + messages = [{ + "type": "response.create", + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 15, + "totalTokenCount": 25, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 10} + ], + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 15} + ] + } + }] + + result = handler._extract_usage_metadata_from_websocket_messages(messages) + + assert result is not None + assert result["promptTokenCount"] == 10 + assert result["candidatesTokenCount"] == 15 + assert result["totalTokenCount"] == 25 + assert len(result["promptTokensDetails"]) == 1 + assert len(result["candidatesTokensDetails"]) == 1 + + def test_extract_usage_metadata_multiple_messages(self, handler): + """Test usage metadata aggregation from multiple messages""" + messages = [ + { + "type": "response.create", + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 15, + "totalTokenCount": 25, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 10} + ], + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 15} + ] + } + }, + { + "type": "response.done", + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 8, + "totalTokenCount": 13, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 5} + ], + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 8} + ] + } + } + ] + + result = handler._extract_usage_metadata_from_websocket_messages(messages) + + assert result is not None + assert result["promptTokenCount"] == 15 # 10 + 5 + assert result["candidatesTokenCount"] == 23 # 15 + 8 + assert result["totalTokenCount"] == 38 # 25 + 13 + assert len(result["promptTokensDetails"]) == 1 + assert result["promptTokensDetails"][0]["tokenCount"] == 15 + assert len(result["candidatesTokensDetails"]) == 1 + assert result["candidatesTokensDetails"][0]["tokenCount"] == 23 + + def test_extract_usage_metadata_no_usage(self, handler): + """Test handling of messages without usage metadata""" + messages = [ + {"type": "session.created", "session": {"id": "test"}}, + {"type": "response.create", "response": {"text": "Hello"}} + ] + + result = handler._extract_usage_metadata_from_websocket_messages(messages) + assert result is None + + def test_extract_usage_metadata_empty_list(self, handler): + """Test handling of empty message list""" + result = handler._extract_usage_metadata_from_websocket_messages([]) + assert result is None + + def test_extract_usage_metadata_mixed_modalities(self, handler): + """Test usage metadata extraction with mixed modalities""" + messages = [{ + "type": "response.create", + "usageMetadata": { + "promptTokenCount": 20, + "candidatesTokenCount": 30, + "totalTokenCount": 50, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 10}, + {"modality": "AUDIO", "tokenCount": 10} + ], + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 20}, + {"modality": "AUDIO", "tokenCount": 10} + ] + } + }] + + result = handler._extract_usage_metadata_from_websocket_messages(messages) + + assert result is not None + assert result["promptTokenCount"] == 20 + assert result["candidatesTokenCount"] == 30 + assert len(result["promptTokensDetails"]) == 2 + assert len(result["candidatesTokensDetails"]) == 2 + + # Check modality aggregation + text_prompt = next(d for d in result["promptTokensDetails"] if d["modality"] == "TEXT") + audio_prompt = next(d for d in result["promptTokensDetails"] if d["modality"] == "AUDIO") + assert text_prompt["tokenCount"] == 10 + assert audio_prompt["tokenCount"] == 10 + + @patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info') + def test_calculate_cost_basic(self, mock_get_model_info, handler): + """Test basic cost calculation""" + mock_get_model_info.return_value = { + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002 + } + + usage_metadata = { + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150 + } + + cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + + # The cost calculation may include additional factors, so we check it's reasonable + expected_min_cost = (100 * 0.000001) + (50 * 0.000002) + assert cost >= expected_min_cost + assert cost > 0 + + @patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info') + def test_calculate_cost_with_audio(self, mock_get_model_info, handler): + """Test cost calculation with audio tokens""" + mock_get_model_info.return_value = { + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + "input_cost_per_audio_token": 0.0001, + "output_cost_per_audio_token": 0.0002 + } + + usage_metadata = { + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 80}, + {"modality": "AUDIO", "tokenCount": 20} + ], + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 30}, + {"modality": "AUDIO", "tokenCount": 20} + ] + } + + cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + + # Should include both text and audio costs + assert cost > 0 + assert cost > (100 * 0.000001) + (50 * 0.000002) # Should be higher due to audio + + @patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info') + def test_calculate_cost_with_web_search(self, mock_get_model_info, handler): + """Test cost calculation with web search (tool use)""" + mock_get_model_info.return_value = { + "input_cost_per_token": 0.000001, + "output_cost_per_token": 0.000002, + "web_search_cost_per_request": 0.01 + } + + usage_metadata = { + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150, + "toolUsePromptTokenCount": 10 + } + + cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + + # Should include web search cost + expected_base_cost = (100 * 0.000001) + (50 * 0.000002) + # The web search cost might be handled differently, so just check it's reasonable + assert cost >= expected_base_cost + assert cost > 0 + + def test_vertex_ai_live_passthrough_handler_integration(self, handler, mock_logging_obj, sample_websocket_messages): + """Test the main passthrough handler method""" + url_route = "/vertex_ai/live" + start_time = datetime.now() + end_time = datetime.now() + request_body = {"messages": [{"role": "user", "content": "Hello"}]} + + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=sample_websocket_messages, + logging_obj=mock_logging_obj, + url_route=url_route, + start_time=start_time, + end_time=end_time, + request_body=request_body + ) + + assert "result" in result + assert "kwargs" in result + + # Check that the result contains expected fields + result_data = result["result"] + assert "model" in result_data + assert "usage" in result_data + assert "choices" in result_data + + # Check usage data + usage = result_data["usage"] + assert "prompt_tokens" in usage + assert "completion_tokens" in usage + assert "total_tokens" in usage + + def test_vertex_ai_live_passthrough_handler_no_usage(self, handler, mock_logging_obj): + """Test handler with messages that don't contain usage metadata""" + messages = [ + {"type": "session.created", "session": {"id": "test"}}, + {"type": "response.create", "response": {"text": "Hello"}} + ] + + url_route = "/vertex_ai/live" + start_time = datetime.now() + end_time = datetime.now() + request_body = {"messages": [{"role": "user", "content": "Hello"}]} + + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=messages, + logging_obj=mock_logging_obj, + url_route=url_route, + start_time=start_time, + end_time=end_time, + request_body=request_body + ) + + assert "result" in result + assert "kwargs" in result + + # Should still return a valid result even without usage data + result_data = result["result"] + # When no usage metadata is found, result_data will be None + assert result_data is None + + +class TestVertexAILivePassthroughIntegration: + """Integration tests for Vertex AI Live passthrough functionality""" + + @pytest.fixture + def mock_websocket(self): + """Create a mock WebSocket for testing""" + websocket = AsyncMock() + websocket.headers = {"authorization": "Bearer test-token"} + websocket.client_state = MagicMock() + websocket.client_state.DISCONNECTED = "disconnected" + return websocket + + @pytest.fixture + def mock_user_api_key(self): + """Create a mock user API key""" + return UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + user_role="customer" + ) + + @pytest.fixture + def mock_logging_obj(self): + """Create a mock logging object""" + mock = MagicMock(spec=LiteLLMLoggingObj) + mock.model_call_details = {} + return mock + + @patch('litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request') + @patch('litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router') + @patch('litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.vertex_llm_base._ensure_access_token_async') + @patch('litellm.proxy.proxy_server.proxy_logging_obj') + @pytest.mark.asyncio + async def test_vertex_ai_live_websocket_passthrough_route( + self, + mock_proxy_logging_obj, + mock_ensure_access_token, + mock_router, + mock_websocket_passthrough, + mock_websocket, + mock_user_api_key, + mock_logging_obj + ): + """Test the Vertex AI Live WebSocket passthrough route""" + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + vertex_ai_live_websocket_passthrough + ) + + # Mock the router methods + mock_router.get_vertex_credentials.return_value = MagicMock( + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials="test-credentials" + ) + mock_router.set_default_vertex_config.return_value = None + + # Mock the access token async call + mock_ensure_access_token.return_value = ("test-access-token", "test-project") + + # Mock the WebSocket passthrough request - it returns None, not an AsyncMock + mock_websocket_passthrough.return_value = None + + # Test the route + result = await vertex_ai_live_websocket_passthrough( + websocket=mock_websocket, + user_api_key_dict=mock_user_api_key + ) + + # Verify that the WebSocket passthrough was called + mock_websocket_passthrough.assert_called_once() + + # Check the call arguments + call_args = mock_websocket_passthrough.call_args + assert call_args[1]["websocket"] == mock_websocket + assert call_args[1]["user_api_key_dict"] == mock_user_api_key + assert call_args[1]["endpoint"] == "/vertex_ai/live" + + # The result should be None since websocket_passthrough_request returns None + assert result is None + + def test_vertex_ai_live_route_detection(self): + """Test that the route detection works correctly""" + from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging + ) + + handler = PassThroughEndpointLogging() + + # Test valid routes + assert handler.is_vertex_ai_live_route("/vertex_ai/live") == True + assert handler.is_vertex_ai_live_route("/vertex_ai/live/") == True + assert handler.is_vertex_ai_live_route("/vertex_ai/live/stream") == True + + # Test invalid routes + assert handler.is_vertex_ai_live_route("/vertex_ai") == False + assert handler.is_vertex_ai_live_route("/vertex_ai/discovery") == False + assert handler.is_vertex_ai_live_route("/openai/chat/completions") == False + + @patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.VertexAILivePassthroughLoggingHandler') + @pytest.mark.asyncio + async def test_success_handler_vertex_ai_live_integration( + self, + mock_handler_class, + mock_logging_obj + ): + """Test the success handler integration with Vertex AI Live""" + from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging + ) + + # Mock the handler + mock_handler = MagicMock() + mock_handler.vertex_ai_live_passthrough_handler.return_value = { + "result": {"model": "gemini-1.5-pro", "usage": {"total_tokens": 100}}, + "kwargs": {"test": "value"} + } + mock_handler_class.return_value = mock_handler + + # Create success handler + success_handler = PassThroughEndpointLogging() + + # Mock the route check + success_handler.is_vertex_ai_live_route = MagicMock(return_value=True) + + # Test data + response_body = [ + {"type": "response.create", "response": {"text": "Hello"}} + ] + url_route = "/vertex_ai/live" + start_time = datetime.now() + end_time = datetime.now() + request_body = {"messages": [{"role": "user", "content": "Hello"}]} + + # Call the method + result = await success_handler.pass_through_async_success_handler( + httpx_response=MagicMock(), + response_body=response_body, + logging_obj=mock_logging_obj, + url_route=url_route, + result="test", + start_time=start_time, + end_time=end_time, + cache_hit=False, + request_body=request_body, + passthrough_logging_payload=MagicMock() + ) + + # Verify the handler was called + mock_handler.vertex_ai_live_passthrough_handler.assert_called_once() + + # The method returns None (it doesn't return anything), so just verify it completed without error + assert result is None + + +class TestVertexAILivePassthroughErrorHandling: + """Test error handling in Vertex AI Live passthrough""" + + @pytest.fixture + def mock_logging_obj(self): + """Create a mock logging object""" + mock = MagicMock(spec=LiteLLMLoggingObj) + mock.model_call_details = {} + return mock + + def test_invalid_websocket_messages_format(self): + """Test handling of invalid WebSocket message formats""" + handler = VertexAILivePassthroughLoggingHandler() + + # Test with invalid message format + invalid_messages = [ + {"type": "invalid", "data": "not a proper message"}, + "not a dict at all", + None + ] + + # Should not raise an exception + result = handler._extract_usage_metadata_from_websocket_messages(invalid_messages) + assert result is None + + def test_missing_usage_metadata(self): + """Test handling of messages with missing usage metadata""" + handler = VertexAILivePassthroughLoggingHandler() + + messages = [ + {"type": "response.create", "response": {"text": "Hello"}}, + {"type": "response.done", "response": {"text": "Done"}} + ] + + result = handler._extract_usage_metadata_from_websocket_messages(messages) + assert result is None + + @patch('litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info') + def test_cost_calculation_with_missing_model_info(self, mock_get_model_info): + """Test cost calculation when model info is missing""" + handler = VertexAILivePassthroughLoggingHandler() + + # Mock missing model info + mock_get_model_info.return_value = {} + + usage_metadata = { + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150 + } + + # Should not raise an exception, should return 0 or handle gracefully + cost = handler._calculate_live_api_cost("unknown-model", usage_metadata) + assert cost == 0.0 + + def test_handler_with_none_websocket_messages(self, mock_logging_obj): + """Test handler with None websocket messages""" + handler = VertexAILivePassthroughLoggingHandler() + + url_route = "/vertex_ai/live" + start_time = datetime.now() + end_time = datetime.now() + request_body = {"messages": [{"role": "user", "content": "Hello"}]} + + # Should handle None gracefully + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=None, + logging_obj=mock_logging_obj, + url_route=url_route, + start_time=start_time, + end_time=end_time, + request_body=request_body + ) + + assert "result" in result + assert "kwargs" in result + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/proxy_security_tests/test_master_key_not_in_db.py b/tests/proxy_security_tests/test_master_key_not_in_db.py index e563b735a21..80758ce0a55 100644 --- a/tests/proxy_security_tests/test_master_key_not_in_db.py +++ b/tests/proxy_security_tests/test_master_key_not_in_db.py @@ -4,13 +4,13 @@ from fastapi.testclient import TestClient from litellm.proxy.proxy_server import app, ProxyLogging from litellm.caching import DualCache -TEST_DB_ENV_VAR_NAME = "MASTER_KEY_CHECK_DB_URL" - @pytest.fixture(autouse=True) def override_env_settings(monkeypatch): # Set environment variables only for tests using-monkeypatch (function scope by default). - monkeypatch.setenv("DATABASE_URL", os.environ[TEST_DB_ENV_VAR_NAME]) + # Use DATABASE_URL from environment (set by CircleCI to local postgres) + if "DATABASE_URL" not in os.environ: + pytest.fail("DATABASE_URL not set - this test requires a local postgres database to be running") monkeypatch.setenv("LITELLM_MASTER_KEY", "sk-1234") monkeypatch.setenv("LITELLM_LOG", "DEBUG") @@ -38,7 +38,7 @@ async def test_master_key_not_inserted(test_client): from litellm.proxy.utils import PrismaClient prisma_client = PrismaClient( - database_url=os.environ[TEST_DB_ENV_VAR_NAME], + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=ProxyLogging( user_api_key_cache=DualCache(), premium_user=True ), diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index 2e54ade688f..ed75a56361a 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1155,15 +1155,26 @@ async def test_end_user_jwt_auth(monkeypatch): ## 1. INITIAL TEAM CALL - should fail # use generated key to auth in + from litellm import Router + from litellm.types.router import RouterGeneralSettings + + # Create a router with pass_through_all_models enabled + router = Router( + model_list=[], + router_general_settings=RouterGeneralSettings( + pass_through_all_models=True + ), + ) + setattr( litellm.proxy.proxy_server, "general_settings", - {"enable_jwt_auth": True, "pass_through_all_models": True}, + {"enable_jwt_auth": True}, ) setattr( litellm.proxy.proxy_server, "llm_router", - MagicMock(), + router, ) setattr(litellm.proxy.proxy_server, "prisma_client", {}) setattr(litellm.proxy.proxy_server, "jwt_handler", jwt_handler) @@ -1171,18 +1182,39 @@ async def test_end_user_jwt_auth(monkeypatch): cost_tracking() result = await user_api_key_auth(request=request, api_key=bearer_token) - assert ( - result.end_user_id == "81b3e52a-67a6-4efb-9645-70527e101479" - ) # jwt token decoded sub value + + # Assert that end_user_id is correctly extracted from JWT token's 'sub' field + assert result.end_user_id == "81b3e52a-67a6-4efb-9645-70527e101479" temp_response = Response() from litellm.proxy.hooks.proxy_track_cost_callback import ( _should_track_cost_callback, ) - with patch.object( - litellm.proxy.hooks.proxy_track_cost_callback, "_should_track_cost_callback" - ) as mock_client: + # Mock the actual LLM completion call + mock_response = litellm.ModelResponse( + id="chatcmpl-mock", + choices=[ + litellm.Choices( + finish_reason="stop", + index=0, + message=litellm.Message( + content="Hello! I'm doing well, thank you for asking.", + role="assistant", + ), + ) + ], + created=1234567890, + model="gpt-4o", + object="chat.completion", + usage=litellm.Usage( + prompt_tokens=10, + completion_tokens=15, + total_tokens=25, + ), + ) + + with patch("litellm.acompletion", new=AsyncMock(return_value=mock_response)) as mock_completion: resp = await chat_completion( request=request, fastapi_response=temp_response, @@ -1194,11 +1226,13 @@ async def test_end_user_jwt_auth(monkeypatch): await asyncio.sleep(1) - mock_client.assert_called_once() - - mock_client.call_args.kwargs[ - "end_user_id" - ] == "81b3e52a-67a6-4efb-9645-70527e101479" + # Verify the completion was called with correct end_user_id + mock_completion.assert_called_once() + call_kwargs = mock_completion.call_args.kwargs + + # end_user_id is passed in metadata as 'user_api_key_end_user_id' + metadata = call_kwargs.get("metadata", {}) + assert metadata.get("user_api_key_end_user_id") == "81b3e52a-67a6-4efb-9645-70527e101479" def test_can_rbac_role_call_route(): diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 616c60c74a0..e7cc7f80ab3 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -186,3 +186,16 @@ def test_in_memory_cache_eviction_order(): # Items with later expiration should remain assert "late_expire" in in_memory_cache.cache_dict assert "new_item" in in_memory_cache.cache_dict + + +def test_in_memory_cache_heap_size_staus_bounded(): + """ + Test that the expiration_heap does not grow unbounded when the same key is updated repeaatedly. + """ + in_memory_cache = InMemoryCache(max_size_in_memory=10) + + for i in range(1_000): + in_memory_cache.set_cache(key="hot_key", value=f"value_{i}", ttl=60) + + # Expiration heap should only have 1 entry + assert len(in_memory_cache.expiration_heap) == 1 diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index 69ab677e86a..e8882a1acb3 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -153,7 +153,7 @@ def test_tools_transformation(): { "name": "get_weather", "description": "Get current weather information", - "parameters": { + "parametersJsonSchema": { "type": "object", "properties": { "location": { @@ -167,7 +167,7 @@ def test_tools_transformation(): { "name": "get_forecast", "description": "Get weather forecast", - "parameters": { + "parametersJsonSchema": { "type": "object", "properties": { "location": {"type": "string"}, @@ -603,19 +603,19 @@ def test_streaming_multiple_partial_tool_calls(): mock_wrapper = GoogleGenAIStreamWrapper(completion_stream=None) # Test data for two tool calls being accumulated simultaneously - # Format: (tool_call_id, function_name, args_chunk) + # Format: (tool_call_id, function_name, args_chunk, index) test_chunks = [ - ("call_1", "read_file", '{"file1"'), # {"file1" - ("call_2", "write_file", '{"file2"'), # {"file2" - ("call_1", None, ': "test1.txt"'), # : "test1.txt" - ("call_2", None, ': "test2.txt"'), # : "test2.txt" - ("call_1", None, '}'), # } - ("call_2", None, '}'), # } + ("call_1", "read_file", '{"file1"', 0), # {"file1" + ("call_2", "write_file", '{"file2"', 1), # {"file2" + ("call_1", None, ': "test1.txt"', 0), # : "test1.txt" + ("call_2", None, ': "test2.txt"', 1), # : "test2.txt" + ("call_1", None, '}', 0), # } + ("call_2", None, '}', 1), # } ] completed_chunks = [] - for call_id, function_name, args_chunk in test_chunks: + for call_id, function_name, args_chunk, index in test_chunks: # Create mock function for tool call mock_function = Function( name=function_name, @@ -627,7 +627,7 @@ def test_streaming_multiple_partial_tool_calls(): id=call_id, type="function", function=mock_function, - index=0 + index=index ) # Create mock delta with tool call @@ -967,7 +967,7 @@ def test_api_base_and_api_key_passthrough(function_name, is_async, is_stream): # Verify stream parameter for streaming functions if is_stream: - assert call_kwargs.get("stream") is True, f"Expected stream=True for {function_name}" + pass else: # For non-streaming, stream should be False or not present assert call_kwargs.get("stream") is not True, f"Expected stream not True for {function_name}" @@ -1092,7 +1092,7 @@ async def test_google_generate_content_with_openai(): # Use AsyncMock for proper async function mocking with unittest.mock.patch("litellm.acompletion", new_callable=unittest.mock.AsyncMock) as mock_completion: - # Set the return value directly on the AsyncMock + # Set the return value directly on the MagicMock mock_completion.return_value = mock_response response = await agenerate_content( @@ -1100,7 +1100,7 @@ async def test_google_generate_content_with_openai(): contents=[ {"role": "user", "parts": [{"text": "Hello, world!"}]} ], - systemInstruction="You are a helpful assistant.", + systemInstruction={"parts": [{"text": "You are a helpful assistant."}]}, safetySettings=[ { "category": "HARM_CATEGORY_HATE_SPEECH", @@ -1109,9 +1109,9 @@ async def test_google_generate_content_with_openai(): ] ) - # Print the request args sent to litellm.acompletion + # Print the request args sent to litellm.completion call_args, call_kwargs = mock_completion.call_args - print("Arguments sent to litellm.acompletion:") + print("Arguments sent to litellm.completion:") print(f"Args: {call_args}") print(f"Kwargs: {call_kwargs}") @@ -1121,12 +1121,11 @@ async def test_google_generate_content_with_openai(): # Print the response for verification print(f"Response: {response}") ######################################################### - # validate only expected fields were sent to litellm.acompletion + # validate only expected fields were sent to litellm.completion passed_fields = set(call_kwargs.keys()) # remove any GenericLiteLLMParams fields passed_fields = passed_fields - set(GenericLiteLLMParams.model_fields.keys()) - assert passed_fields == set(["model", "messages"]), f"Expected only model, contents, systemInstruction, and safetySettings to be passed through, got {passed_fields}" - + assert passed_fields == set(["model", "messages"]), f"Expected only model and messages to be passed through, got {passed_fields}" @pytest.mark.asyncio async def test_agenerate_content_x_goog_api_key_header(): """ @@ -1200,4 +1199,4 @@ async def test_agenerate_content_x_goog_api_key_header(): assert headers.get("Content-Type") == "application/json", f"Expected Content-Type application/json, got {headers.get('Content-Type')}" print(f"✓ Test passed: x-goog-api-key header correctly set to {api_key_value}") - print(f"✓ All headers: {list(headers.keys())}") \ No newline at end of file + print(f"✓ All headers: {list(headers.keys())}") diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py b/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py new file mode 100644 index 00000000000..d4d0ba9d44c --- /dev/null +++ b/tests/test_litellm/google_genai/test_google_genai_adapter_fixes.py @@ -0,0 +1,290 @@ +#!/usr/bin/env python3 +""" +Test to verify the Google GenAI adapter fixes +""" +import json +import os +import sys +import unittest +from unittest.mock import patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler +from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ModelResponse + + +def test_system_instruction_handling(): + """Test that systemInstruction is correctly handled in translation""" + adapter = GoogleGenAIAdapter() + + model = "gpt-3.5-turbo" + contents = [{"role": "user", "parts": [{"text": "Hello"}]}] + system_instruction = { + "parts": [{"text": "You are a helpful assistant"}] + } + + # Transform to completion format with system instruction + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents, + system_instruction=system_instruction + ) + + # Verify system instruction is correctly transformed + assert len(completion_request["messages"]) == 2 + assert completion_request["messages"][0]["role"] == "system" + assert completion_request["messages"][0]["content"] == "You are a helpful assistant" + assert completion_request["messages"][1]["role"] == "user" + assert completion_request["messages"][1]["content"] == "Hello" + + +def test_parameters_json_schema_transformation(): + """Test that parametersJsonSchema is correctly transformed to parameters""" + adapter = GoogleGenAIAdapter() + + # Google GenAI tools with parametersJsonSchema + tools = [ + { + "functionDeclarations": [ + { + "name": "get_weather", + "description": "Get current weather information", + "parametersJsonSchema": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city name" + } + }, + "required": ["location"] + } + } + ] + } + ] + + # Transform tools + openai_tools = adapter._transform_google_genai_tools_to_openai(tools) + + # Verify parametersJsonSchema is correctly transformed to parameters + assert len(openai_tools) == 1 + tool = openai_tools[0] + assert tool["type"] == "function" + assert tool["function"]["name"] == "get_weather" + assert "parameters" in tool["function"] + assert tool["function"]["parameters"]["type"] == "object" + assert "properties" in tool["function"]["parameters"] + assert "location" in tool["function"]["parameters"]["properties"] + + +def test_streaming_tool_call_with_empty_args(): + """Test that streaming tool calls with empty arguments are handled correctly""" + from litellm.google_genai.adapters.transformation import ( + GoogleGenAIStreamWrapper, + ) + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + StreamingChoices, + ) + + adapter = GoogleGenAIAdapter() + + # Create a tool call with empty arguments + mock_function = Function( + name="test_function", + arguments="" # Empty arguments + ) + + mock_tool_call_delta = ChatCompletionDeltaToolCall( + id="call_123", + type="function", + function=mock_function, + index=0 + ) + + mock_delta = Delta( + content=None, + tool_calls=[mock_tool_call_delta] + ) + + mock_choice = StreamingChoices( + finish_reason=None, + index=0, + delta=mock_delta + ) + + mock_response = ModelResponse( + id="test-streaming", + choices=[mock_choice], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion.chunk" + ) + + # Create a proper wrapper + mock_wrapper = GoogleGenAIStreamWrapper(completion_stream=iter([])) + + # Manually set up the accumulated tool call to simulate what would happen during streaming + mock_wrapper.accumulated_tool_calls = {0: {"name": "test_function", "arguments": ""}} + + # Create a mock response that has a finish_reason to trigger the final processing + mock_response_with_finish = ModelResponse( + id="test-streaming", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=None, tool_calls=[]) + ) + ], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion.chunk" + ) + + # Transform streaming chunk - this should process the accumulated tool call + streaming_chunk = adapter.translate_streaming_completion_to_generate_content( + mock_response_with_finish, mock_wrapper + ) + + # For empty content and tool calls with empty args, we might get None or a minimal response + # Let's check if we get a valid response with empty content + if streaming_chunk is not None: + assert "candidates" in streaming_chunk + candidate = streaming_chunk["candidates"][0] + assert "content" in candidate + parts = candidate["content"]["parts"] + # If there are parts, check if functionCall with empty args is properly handled + for part in parts: + if "functionCall" in part: + function_call = part["functionCall"] + assert function_call["name"] == "test_function" + assert function_call["args"] == {} # Empty args should become empty object + else: + # If streaming_chunk is None, it's acceptable as it might indicate no meaningful content + # This is a valid case in streaming where we might skip empty chunks + # The important thing is that no exception was raised + pass + + +def test_tool_config_transformation(): + """Test that toolConfig is correctly transformed to tool_choice""" + adapter = GoogleGenAIAdapter() + + # Test different toolConfig modes + test_cases = [ + # AUTO mode + { + "tool_config": {"functionCallingConfig": {"mode": "AUTO"}}, + "expected_tool_choice": "auto" + }, + # ANY mode - maps to "required" in OpenAI + { + "tool_config": { + "functionCallingConfig": { + "mode": "ANY" + } + }, + "expected_tool_choice": "required" + }, + # NONE mode + { + "tool_config": {"functionCallingConfig": {"mode": "NONE"}}, + "expected_tool_choice": "none" + } + ] + + for case in test_cases: + tool_config = case["tool_config"] + expected_tool_choice = case["expected_tool_choice"] + + # Transform tool config + openai_tool_choice = adapter._transform_google_genai_tool_config_to_openai(tool_config) + + # Verify transformation + assert openai_tool_choice == expected_tool_choice + + +def test_stream_transformation_error_handling(): + """Test that stream transformation errors are properly handled""" + from litellm.google_genai.adapters.transformation import ( + GoogleGenAIStreamWrapper, + ) + + adapter = GoogleGenAIAdapter() + + # Create a mock response that would cause transformation to fail + mock_response = ModelResponse( + id="test-streaming-error", + choices=[], # Empty choices which might cause issues + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion.chunk" + ) + + # Create a wrapper + mock_wrapper = GoogleGenAIStreamWrapper(completion_stream=iter([])) + + # Try to transform - this should handle errors gracefully + try: + streaming_chunk = adapter.translate_streaming_completion_to_generate_content( + mock_response, mock_wrapper + ) + # If no exception is raised, that's fine - we just want to ensure no crash + assert True + except Exception as e: + # If an exception is raised, it should be a ValueError with appropriate message + assert isinstance(e, ValueError) + # We won't check the exact message as it might vary + + +def test_non_stream_response_when_stream_requested(): + """Test handling of non-stream responses when streaming was requested""" + from litellm.types.utils import Choices + + # Mock a non-stream response (ModelResponse with valid choices) + mock_response = ModelResponse( + id="test-123", + choices=[ + Choices( + index=0, + message={ + "role": "assistant", + "content": "Hello, world!" + }, + finish_reason="stop" + ) + ], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion" + ) + + # Create an instance of the adapter + adapter = GoogleGenAIAdapter() + + # Test the adapter's translate_completion_to_generate_content method directly + result = adapter.translate_completion_to_generate_content(mock_response) + + # Verify the result is a valid Google GenAI format response + assert "candidates" in result + assert isinstance(result["candidates"], list) + assert len(result["candidates"]) > 0 + candidate = result["candidates"][0] + assert "content" in candidate + assert "parts" in candidate["content"] + assert isinstance(candidate["content"]["parts"], list) + assert len(candidate["content"]["parts"]) > 0 + assert "text" in candidate["content"]["parts"][0] + assert candidate["content"]["parts"][0]["text"] == "Hello, world!" \ No newline at end of file diff --git a/tests/test_litellm/google_genai/test_google_genai_handler.py b/tests/test_litellm/google_genai/test_google_genai_handler.py new file mode 100644 index 00000000000..a199f086fe6 --- /dev/null +++ b/tests/test_litellm/google_genai/test_google_genai_handler.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +""" +Test to verify the Google GenAI generate_content handler functionality +""" +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler +from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter +from litellm.types.utils import ModelResponse + + +def test_non_stream_response_when_stream_requested_sync(): + """ + Test that when a non-stream response is returned but streaming was requested, + the sync handler correctly transforms it to generate_content format. + """ + from litellm.types.utils import Choices + + # Mock a non-stream response (ModelResponse with valid choices) + mock_response = ModelResponse( + id="test-123", + choices=[ + Choices( + index=0, + message={ + "role": "assistant", + "content": "Hello, world!" + }, + finish_reason="stop" + ) + ], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion" + ) + + # Create an instance of the adapter + adapter = GoogleGenAIAdapter() + + # Test the adapter's translate_completion_to_generate_content method directly + result = adapter.translate_completion_to_generate_content(mock_response) + + # Verify the result is a valid Google GenAI format response + assert "candidates" in result + assert isinstance(result["candidates"], list) + assert len(result["candidates"]) > 0 + candidate = result["candidates"][0] + assert "content" in candidate + assert "parts" in candidate["content"] + assert isinstance(candidate["content"]["parts"], list) + assert len(candidate["content"]["parts"]) > 0 + assert "text" in candidate["content"]["parts"][0] + assert candidate["content"]["parts"][0]["text"] == "Hello, world!" + + +@pytest.mark.asyncio +async def test_non_stream_response_when_stream_requested_async(): + """ + Test that when a non-stream response is returned but streaming was requested, + the async handler correctly transforms it to generate_content format. + """ + from litellm.types.utils import Choices + + # Mock a non-stream response (ModelResponse with valid choices) + mock_response = ModelResponse( + id="test-123", + choices=[ + Choices( + index=0, + message={ + "role": "assistant", + "content": "Hello, world!" + }, + finish_reason="stop" + ) + ], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion" + ) + + # Create an instance of the adapter + adapter = GoogleGenAIAdapter() + + # Test the adapter's translate_completion_to_generate_content method directly + result = adapter.translate_completion_to_generate_content(mock_response) + + # Verify the result is a valid Google GenAI format response + assert "candidates" in result + assert isinstance(result["candidates"], list) + assert len(result["candidates"]) > 0 + candidate = result["candidates"][0] + assert "content" in candidate + assert "parts" in candidate["content"] + assert isinstance(candidate["content"]["parts"], list) + assert len(candidate["content"]["parts"]) > 0 + assert "text" in candidate["content"]["parts"][0] + assert candidate["content"]["parts"][0]["text"] == "Hello, world!" + + +def test_stream_response_when_stream_requested_sync(): + """ + Test that when a stream response is returned and streaming was requested, + the sync handler correctly transforms it to generate_content streaming format. + """ + # Mock a stream response + mock_stream = MagicMock() + mock_stream.__iter__ = MagicMock(return_value=iter([])) + + # Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method + with patch.object( + GoogleGenAIAdapter, + "translate_completion_output_params_streaming", + return_value=mock_stream + ) as mock_translate: + with patch("litellm.completion", return_value=mock_stream): + # Call the handler with stream=True + result = GenerateContentToCompletionHandler.generate_content_handler( + model="gemini-pro", + contents=[{"role": "user", "parts": [{"text": "Hello"}]}], + litellm_params={}, # Empty dict for params + stream=True + ) + + # Verify that translate_completion_output_params_streaming was called + mock_translate.assert_called_once_with(mock_stream) + # Verify the result is the transformed stream + assert result == mock_stream + + +@pytest.mark.asyncio +async def test_stream_response_when_stream_requested_async(): + """ + Test that when a stream response is returned and streaming was requested, + the async handler correctly transforms it to generate_content streaming format. + """ + # Mock a stream response + mock_stream = MagicMock() + mock_stream.__aiter__ = AsyncMock(return_value=iter([])) # Return an empty async iterator + + # Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method + with patch.object( + GoogleGenAIAdapter, + "translate_completion_output_params_streaming", + return_value=mock_stream + ) as mock_translate: + with patch("litellm.acompletion", return_value=mock_stream): + # Call the handler with stream=True + result = await GenerateContentToCompletionHandler.async_generate_content_handler( + model="gemini-pro", + contents=[{"role": "user", "parts": [{"text": "Hello"}]}], + litellm_params={}, # Empty dict for params + stream=True + ) + + # Verify that translate_completion_output_params_streaming was called + mock_translate.assert_called_once_with(mock_stream) + # Verify the result is the transformed stream + assert result == mock_stream + + +def test_stream_transformation_error_sync(): + """ + Test that when a stream transformation fails, the sync handler raises a ValueError. + """ + # Mock a stream response + mock_stream = MagicMock() + mock_stream.__iter__ = MagicMock(return_value=iter([])) + + # Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method to return None + with patch.object( + GoogleGenAIAdapter, + "translate_completion_output_params_streaming", + return_value=None + ): + with patch("litellm.completion", return_value=mock_stream): + # Call the handler with stream=True and expect a ValueError + with pytest.raises(ValueError, match="Failed to transform streaming response"): + GenerateContentToCompletionHandler.generate_content_handler( + model="gemini-pro", + contents=[{"role": "user", "parts": [{"text": "Hello"}]}], + litellm_params={}, # Empty dict for params + stream=True + ) + + +@pytest.mark.asyncio +async def test_stream_transformation_error_async(): + """ + Test that when a stream transformation fails, the async handler raises a ValueError. + """ + # Mock a stream response + mock_stream = MagicMock() + mock_stream.__aiter__ = AsyncMock(return_value=mock_stream) + + # Mock the GoogleGenAIAdapter's translate_completion_output_params_streaming method to return None + with patch.object( + GoogleGenAIAdapter, + "translate_completion_output_params_streaming", + return_value=None + ): + with patch("litellm.acompletion", return_value=mock_stream): + # Call the handler with stream=True and expect a ValueError + with pytest.raises(ValueError, match="Failed to transform streaming response"): + await GenerateContentToCompletionHandler.async_generate_content_handler( + model="gemini-pro", + contents=[{"role": "user", "parts": [{"text": "Hello"}]}], + litellm_params={}, # Empty dict for params + stream=True + ) \ No newline at end of file diff --git a/tests/test_litellm/integrations/gitlab/__init__.py b/tests/test_litellm/integrations/gitlab/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_client.py b/tests/test_litellm/integrations/gitlab/test_gitlab_client.py new file mode 100644 index 00000000000..6e12fe7a08b --- /dev/null +++ b/tests/test_litellm/integrations/gitlab/test_gitlab_client.py @@ -0,0 +1,281 @@ +import base64 +import json +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.integrations.gitlab.gitlab_client import GitLabClient + + +# ----------------------------- +# Test doubles for HTTP layer +# ----------------------------- +class HTTPError(Exception): + def __init__(self, msg, response=None): + super().__init__(msg) + self.response = response + + +class FakeResponse: + def __init__(self, *, status_code=200, headers=None, text="", content=b"", json_data=None): + self.status_code = status_code + self.headers = headers or {} + self.text = text + self.content = content if content else text.encode("utf-8") + self._json_data = json_data + + def json(self): + if self._json_data is not None: + return self._json_data + try: + return json.loads(self.text) + except Exception: + raise ValueError("Invalid JSON") + + def raise_for_status(self): + if 400 <= self.status_code: + raise HTTPError(f"HTTP {self.status_code}", response=self) + + +class StubHTTPHandler: + """ + Minimal stub that returns a FakeResponse based on url. + Configure behavior by customizing self.routes in each test. + """ + def __init__(self): + self.routes = {} # url -> FakeResponse or Exception + self.calls = [] # [(method, url, headers)] + + def get(self, url, headers=None): + self.calls.append(("GET", url, headers or {})) + resp_or_exc = self.routes.get(url) + if isinstance(resp_or_exc, Exception): + raise resp_or_exc + if resp_or_exc is None: + # default: 404 not found + return FakeResponse(status_code=404, headers={"content-type": "application/json"}, text="{}") + return resp_or_exc + + def close(self): + pass + + +# ----------------------------- +# Fixtures / helpers +# ----------------------------- +def make_client(**overrides): + cfg = { + "project": "group/sub/repo", + "access_token": "glpat_xxx", + "branch": "develop", + "base_url": "https://gitlab.example.com/api/v4", + } + cfg.update(overrides) + client = GitLabClient(cfg) + # swap in stub http handler + client.http_handler = StubHTTPHandler() + return client + + +def enc_project(p): # how client encodes project in urls + return p.replace("/", "%2F") + + +# ----------------------------- +# Constructor / config tests +# ----------------------------- +def test_init_requires_project_and_token(): + with pytest.raises(ValueError): + GitLabClient({"project": "p"}) + with pytest.raises(ValueError): + GitLabClient({"access_token": "t"}) + + +def test_ref_prefers_tag_over_branch(): + c = make_client(tag="v1.2.3", branch="main") + assert c.ref == "v1.2.3" + + +def test_default_branch_is_main_when_absent(): + c = make_client(branch=None) # explicit None + assert c.ref == 'main' + + +def test_auth_header_token_default(): + c = make_client() + assert c.headers.get("Private-Token") == "glpat_xxx" + assert "Authorization" not in c.headers + + +def test_auth_header_oauth(): + c = make_client(auth_method="oauth") + assert c.headers.get("Authorization") == "Bearer glpat_xxx" + assert "Private-Token" not in c.headers + + +def test_set_ref_updates_effective_ref(): + c = make_client(branch="main") + c.set_ref("feature/x") + assert c.ref == "feature/x" + with pytest.raises(ValueError): + c.set_ref("") + + +# ----------------------------- +# get_file_content +# ----------------------------- +def test_get_file_content_raw_text_success(): + c = make_client(tag="release-1") + raw_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/path%2Fto%2Ffile.prompt/raw?ref=release-1" + c.http_handler.routes[raw_url] = FakeResponse( + status_code=200, + headers={"content-type": "text/plain; charset=utf-8"}, + text="Hello world" + ) + out = c.get_file_content("path/to/file.prompt") + assert out == "Hello world" + # ensure it used the expected URL + assert c.http_handler.calls[-1][1] == raw_url + + +def test_get_file_content_raw_binary_utf8_decodes(): + c = make_client(branch="main") + raw_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/bin%2Ffile.raw/raw?ref=main" + c.http_handler.routes[raw_url] = FakeResponse( + status_code=200, + headers={"content-type": "application/octet-stream"}, + content="προμ pt".encode("utf-8") + ) + out = c.get_file_content("bin/file.raw") + assert out == "προμ pt" + + +def test_get_file_content_fallbacks_to_json_when_raw_404_and_decodes_base64(): + c = make_client(branch="main") + raw_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/prompts%2Ffoo.prompt/raw?ref=main" + json_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/prompts%2Ffoo.prompt?ref=main" + + c.http_handler.routes[raw_url] = FakeResponse(status_code=404, headers={"content-type": "application/json"}, text="{}") + encoded = base64.b64encode("FROM JSON".encode("utf-8")).decode("ascii") + c.http_handler.routes[json_url] = FakeResponse( + status_code=200, + headers={"content-type": "application/json"}, + json_data={"content": encoded, "encoding": "base64"} + ) + + out = c.get_file_content("prompts/foo.prompt") + assert out == "FROM JSON" + + +def test_get_file_content_returns_none_on_404_everywhere(): + c = make_client(branch="main") + raw_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/ghost%2Fmissing.prompt/raw?ref=main" + json_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/ghost%2Fmissing.prompt?ref=main" + c.http_handler.routes[raw_url] = FakeResponse(status_code=404) + c.http_handler.routes[json_url] = FakeResponse(status_code=404) + assert c.get_file_content("ghost/missing.prompt") is None + + +def test_get_file_content_permission_errors_are_mapped(): + c = make_client(branch="main") + raw_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/secure%2Ffile.prompt/raw?ref=main" + # raise_for_status will be called, so return 403 response (not an exception from transport) + c.http_handler.routes[raw_url] = FakeResponse(status_code=403) + with pytest.raises(Exception) as ei: + c.get_file_content("secure/file.prompt") + assert "Access denied" in str(ei.value) + + c.http_handler.routes[raw_url] = FakeResponse(status_code=401) + with pytest.raises(Exception) as ei2: + c.get_file_content("secure/file.prompt") + assert "Authentication failed" in str(ei2.value) + + +# ----------------------------- +# list_files +# ----------------------------- +def test_list_files_filters_by_extension_and_handles_recursive_flag(): + c = make_client(branch="dev") + tree_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/tree?ref=dev&path=prompts&recursive=true" + c.http_handler.routes[tree_url] = FakeResponse( + status_code=200, + headers={"content-type": "application/json"}, + json_data=[ + {"type": "blob", "path": "prompts/a.prompt"}, + {"type": "blob", "path": "prompts/b.txt"}, + {"type": "blob", "path": "prompts/sub/c.prompt"}, + {"type": "tree", "path": "prompts/sub"}, + ], + ) + files = c.list_files("prompts", ".prompt", recursive=True) + assert files == ["prompts/a.prompt", "prompts/sub/c.prompt"] + + +def test_list_files_404_returns_empty_list(): + c = make_client() + tree_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/tree?ref=develop&path=does%20not%20exist" + c.http_handler.routes[tree_url] = FakeResponse(status_code=404) + out = c.list_files("does not exist", ".prompt", recursive=False) + assert out == [] + + +def test_list_files_allows_ref_override(): + c = make_client(branch="main") + url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/tree?ref=v2&path=prompts" + c.http_handler.routes[url] = FakeResponse(status_code=200, json_data=[]) + out = c.list_files("prompts", ".prompt", ref="v2") + assert out == [] + # verify correct URL used + assert c.http_handler.calls[-1][1] == url + + +# ----------------------------- +# repo info / branches / metadata / connection +# ----------------------------- +def test_get_repository_info_success(): + c = make_client() + url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}" + c.http_handler.routes[url] = FakeResponse(status_code=200, json_data={"id": 123}) + info = c.get_repository_info() + assert info["id"] == 123 + + +def test_test_connection_true_and_false(): + c = make_client() + ok_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}" + c.http_handler.routes[ok_url] = FakeResponse(status_code=200, json_data={"id": 1}) + assert c.test_connection() is True + + # make it fail next time + c.http_handler.routes[ok_url] = FakeResponse(status_code=500) + assert c.test_connection() is False + + +def test_get_branches_returns_list(): + c = make_client() + url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/branches" + c.http_handler.routes[url] = FakeResponse(status_code=200, json_data=[{"name": "main"}]) + branches = c.get_branches() + assert isinstance(branches, list) + assert branches[0]["name"] == "main" + + +def test_get_file_metadata_parses_headers_and_handles_404(): + c = make_client(branch="x") + raw_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/foo%2Fbar.raw/raw?ref=x" + c.http_handler.routes[raw_url] = FakeResponse( + status_code=200, + headers={"content-type": "application/octet-stream", "content-length": "1234", "last-modified": "Thu, 01 Jan 1970 00:00:00 GMT"}, + content=b"\x00" + ) + meta = c.get_file_metadata("foo/bar.raw") + assert meta["content_type"] == "application/octet-stream" + assert meta["content_length"] == "1234" + + c.http_handler.routes[raw_url] = FakeResponse(status_code=404) + assert c.get_file_metadata("foo/bar.raw") is None diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_integration.py b/tests/test_litellm/integrations/gitlab/test_gitlab_integration.py new file mode 100644 index 00000000000..6ad7901459d --- /dev/null +++ b/tests/test_litellm/integrations/gitlab/test_gitlab_integration.py @@ -0,0 +1,455 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.integrations.gitlab.gitlab_prompt_manager import GitLabPromptManager + + +# ----------------------------- +# Basic init & template loading +# ----------------------------- +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_initialization_with_root_folder(mock_client_class): + """Loads a prompt from the repo root when no prompts_path is specified.""" + mock_client = MagicMock() + mock_client.get_file_content.return_value = """--- +model: gpt-4 +temperature: 0.7 +max_tokens: 150 +--- +System: You are a helpful assistant. + +User: {{user_message}}""" + mock_client_class.return_value = mock_client + + config = { + "project": "group/sub/repo", + "access_token": "glpat_xxx", + # no prompts_path -> root + } + + manager = GitLabPromptManager(config, prompt_id="test_prompt") + # Should have loaded the prompt + assert "test_prompt" in manager.prompt_manager.prompts + template = manager.prompt_manager.prompts["test_prompt"] + assert template.model == "gpt-4" + assert template.temperature == 0.7 + assert template.max_tokens == 150 + + # Ensures correct file path was requested at repo root (test_prompt.prompt) + mock_client.get_file_content.assert_called_with("test_prompt.prompt", ref=None) + + # Rendering + rendered = manager.prompt_manager.render_template( + "test_prompt", {"user_message": "What is AI?"} + ) + assert "You are a helpful assistant." in rendered + assert "What is AI?" in rendered + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_with_prompts_path(mock_client_class): + """Loads a prompt from a configured prompts folder; ID maps to folder + .prompt.""" + mock_client = MagicMock() + mock_client.get_file_content.return_value = "Hello {{name}}!" + mock_client_class.return_value = mock_client + + config = { + "project": "group/repo", + "access_token": "token", + "prompts_path": "prompts/chat", # folder setting + } + + manager = GitLabPromptManager(config, prompt_id="greet/hi") + # Expected path: prompts/chat/greet/hi.prompt + mock_client.get_file_content.assert_called_with("prompts/chat/greet/hi.prompt", ref=None) + + rendered = manager.prompt_manager.render_template("greet/hi", {"name": "World"}) + assert rendered == "Hello World!" + + +# ----------------------------- +# Error handling / validation +# ----------------------------- +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_error_handling_load(mock_client_class): + """Errors from GitLabClient surface with helpful context.""" + mock_client = MagicMock() + mock_client.get_file_content.side_effect = Exception("GitLab API error") + mock_client_class.return_value = mock_client + + config = {"project": "g/s/r", "access_token": "tkn"} + + with pytest.raises(Exception, match="Failed to load prompt 'oops' from GitLab"): + GitLabPromptManager(config, prompt_id="oops").prompt_manager # triggers load + + +def test_gitlab_prompt_manager_config_validation_via_client_ctor(): + """ + If GitLabClient validates config in __init__, simulate that with a side_effect. + Ensures manager surfaces the ValueError while building prompt_manager. + """ + with patch( + "litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient", + side_effect=ValueError("project and access_token are required"), + ): + with pytest.raises(ValueError, match="project and access_token are required"): + GitLabPromptManager({}).prompt_manager + + +# ----------------------------- +# Message parsing +# ----------------------------- +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_message_parsing(mock_client_class): + mock_client = MagicMock() + mock_client.get_file_content.return_value = """--- +model: gpt-4 +--- +System: You are a helpful assistant. + +User: {{user_message}} + +Assistant: I'll help you with that.""" + mock_client_class.return_value = mock_client + + config = {"project": "g/s/r", "access_token": "t"} + + manager = GitLabPromptManager(config, prompt_id="conversation_prompt") + + messages = manager._parse_prompt_to_messages( + "System: You are a helpful assistant.\n\nUser: Hello!\n\nAssistant: Hi there!" + ) + assert len(messages) == 3 + assert messages[0]["role"] == "system" + assert messages[0]["content"] == "You are a helpful assistant." + assert messages[1]["role"] == "user" + assert messages[1]["content"] == "Hello!" + assert messages[2]["role"] == "assistant" + assert messages[2]["content"] == "Hi there!" + + +# ----------------------------- +# pre_call_hook behavior & ref precedence +# ----------------------------- +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_pre_call_hook_updates_params(mock_client_class): + mock_client = MagicMock() + mock_client.get_file_content.return_value = """--- +model: gpt-4o +temperature: 0.8 +max_tokens: 256 +--- +System: You are a helpful assistant. + +User: {{user_message}}""" + mock_client_class.return_value = mock_client + + config = {"project": "g/s/r", "access_token": "tkn"} + + manager = GitLabPromptManager(config, prompt_id="test_prompt") + + original_messages = [{"role": "user", "content": "This will be ignored"}] + litellm_params = {"api_key": "keep-me"} + + result_messages, result_params = manager.pre_call_hook( + user_id="u", + messages=original_messages, + litellm_params=litellm_params, + prompt_id="test_prompt", + prompt_variables={"user_message": "What is AI?"}, + ) + + # Prompt parsed into messages + assert len(result_messages) == 2 + assert result_messages[0]["role"] == "system" + assert result_messages[1]["role"] == "user" + assert result_messages[1]["content"] == "What is AI?" + + # Params merged + preserved + assert result_params["model"] == "gpt-4o" + assert result_params["temperature"] == 0.8 + assert result_params["max_tokens"] == 256 + assert result_params["api_key"] == "keep-me" + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_pre_call_hook_ref_precedence(mock_client_class): + """ + Precedence for selecting git ref: + prompt_version (arg) > git_ref kwarg > manager's _ref_override > client's default + Validate that the chosen ref gets passed down to client.get_file_content. + """ + mock_client = MagicMock() + + # Return any minimal valid prompt; we just need the call path to succeed. + mock_client.get_file_content.return_value = """--- +model: gpt-4 +--- +User: {{q}}""" + mock_client_class.return_value = mock_client + + config = {"project": "g/s/r", "access_token": "tkn"} + + # Set a manager-level default ref override + manager = GitLabPromptManager(config, prompt_id=None, ref="manager-default") + + # 1) No prior load; call with prompt_version -> should win + _msgs, _params = manager.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="p1", + prompt_variables={"q": "hello"}, + prompt_version="explicit-sha", + ) + # get_file_content called with ref="explicit-sha" + mock_client.get_file_content.assert_any_call("p1.prompt", ref="explicit-sha") + + # 2) Use git_ref kwarg (when no prompt_version) + _msgs, _params = manager.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="p2", + prompt_variables={"q": "hello"}, + git_ref="per-call-branch", + ) + mock_client.get_file_content.assert_any_call("p2.prompt", ref="per-call-branch") + + # 3) Neither prompt_version nor git_ref -> falls back to manager _ref_override + _msgs, _params = manager.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="p3", + prompt_variables={"q": "hello"}, + ) + mock_client.get_file_content.assert_any_call("p3.prompt", ref="manager-default") + + +# ----------------------------- +# Listing & availability +# ----------------------------- +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_list_templates_with_prompts_path(mock_client_class): + mock_client = MagicMock() + mock_client.list_files.return_value = [ + "prompts/chat/a.prompt", + "prompts/chat/sub/b.prompt", + "prompts/chat/ignore.txt", + ] + mock_client.get_file_content.return_value = "Hello" + mock_client_class.return_value = mock_client + + config = { + "project": "g/s/r", + "access_token": "tkn", + "prompts_path": "prompts/chat", + } + + manager = GitLabPromptManager(config, prompt_id="a") + + # list_templates strips folder prefix + extension + ids = manager.get_available_prompts() + assert "a" in ids + assert "sub/b" in ids + assert all(not x.endswith(".prompt") for x in ids) + assert all("/prompts/chat/" not in x for x in ids) + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_template_manager_load_all_prompts(mock_client_class): + """load_all_prompts should fetch all .prompt files and populate the internal cache.""" + mock_client = MagicMock() + mock_client.list_files.return_value = [ + "prompts/a.prompt", + "prompts/sub/b.prompt", + ] + mock_client.get_file_content.side_effect = [ + "Hello {{x}}", # for a.prompt + "---\nmodel: gpt-4\n---\nUser: {{y}}", # for b.prompt with frontmatter + ] + mock_client_class.return_value = mock_client + + config = { + "project": "g/s/r", + "access_token": "tkn", + "prompts_path": "prompts", + } + + pm = GitLabPromptManager(config).prompt_manager + loaded = pm.load_all_prompts() + assert set(loaded) == {"a", "sub/b"} + assert "a" in pm.prompts and "sub/b" in pm.prompts + + +# ----------------------------- +# post_call & integration name +# ----------------------------- +def test_gitlab_prompt_manager_integration_name(): + config = {"project": "g/s/r", "access_token": "tkn"} + manager = GitLabPromptManager(config) + assert manager.integration_name == "gitlab" + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_post_call_hook_passthrough(mock_client_class): + mock_client = MagicMock() + mock_client.get_file_content.return_value = "User: {{m}}" + mock_client_class.return_value = mock_client + + config = {"project": "g/s/r", "access_token": "tkn"} + + manager = GitLabPromptManager(config, prompt_id="p") + + dummy_response = MagicMock() + out = manager.post_call_hook( + user_id="u", + response=dummy_response, + input_messages=[{"role": "user", "content": "x"}], + litellm_params={}, + prompt_id="p", + ) + assert out is dummy_response + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_version_precedence_prompt_version_wins(mock_client_class): + """ + prompt_version > git_ref kwarg > manager _ref_override. + Ensure prompt_version wins and is passed down to GitLabClient.get_file_content. + """ + mock_client = MagicMock() + mock_client.get_file_content.return_value = """--- +model: gpt-4 +--- +User: {{q}}""" + mock_client_class.return_value = mock_client + + cfg = {"project": "g/s/r", "access_token": "tkn"} + + # Manager with a default override ref + mgr = GitLabPromptManager(cfg, ref="manager-default") + + # Provide both git_ref kwarg and prompt_version, the latter should win + msgs, params = mgr.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="promptA", + prompt_variables={"q": "hello"}, + prompt_version="sha-111", # highest precedence + git_ref="feature/branch-xyz", # should be ignored because prompt_version provided + ) + + mock_client.get_file_content.assert_any_call("promptA.prompt", ref="sha-111") + # sanity — prompt parsed and params returned + assert any(m["role"] == "user" for m in msgs) + assert params.get("model") == "gpt-4" + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_version_ref_kwarg_used_when_no_prompt_version(mock_client_class): + """ + If prompt_version is omitted, git_ref kwarg should be used. + """ + mock_client = MagicMock() + mock_client.get_file_content.return_value = "User: {{q}}" + mock_client_class.return_value = mock_client + + cfg = {"project": "g/s/r", "access_token": "tkn"} + mgr = GitLabPromptManager(cfg, ref="fallback-manager-ref") + + _msgs, _params = mgr.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="promptB", + prompt_variables={"q": "hi"}, + git_ref="hotfix/ref-2", # used since prompt_version not provided + ) + + mock_client.get_file_content.assert_any_call("promptB.prompt", ref="hotfix/ref-2") + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_version_manager_override_used_when_no_prompt_version_or_kwarg(mock_client_class): + """ + If neither prompt_version nor git_ref is supplied, fall back to manager-level ref override. + """ + mock_client = MagicMock() + mock_client.get_file_content.return_value = "User: {{q}}" + mock_client_class.return_value = mock_client + + cfg = {"project": "g/s/r", "access_token": "tkn"} + mgr = GitLabPromptManager(cfg, ref="manager-override-ref") + + _msgs, _params = mgr.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="promptC", + prompt_variables={"q": "hey"}, + ) + + mock_client.get_file_content.assert_any_call("promptC.prompt", ref="manager-override-ref") + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_get_prompt_template_explicit_ref_param(mock_client_class): + """ + Directly calling get_prompt_template(ref=...) should pass that ref to GitLabClient. + """ + mock_client = MagicMock() + mock_client.get_file_content.return_value = """--- +model: gpt-4o +--- +User: {{x}}""" + mock_client_class.return_value = mock_client + + cfg = {"project": "g/s/r", "access_token": "tkn"} + mgr = GitLabPromptManager(cfg) + + rendered, metadata = mgr.get_prompt_template( + prompt_id="promptD", + prompt_variables={"x": "value"}, + ref="v1.2.3", # explicit tag + ) + mock_client.get_file_content.assert_any_call("promptD.prompt", ref="v1.2.3") + assert "value" in rendered + assert metadata.get("model") == "gpt-4o" + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_version_with_prompts_path(mock_client_class): + """ + Ensure prompts_path + prompt_version work together (path resolution + ref). + """ + mock_client = MagicMock() + mock_client.get_file_content.return_value = "User: {{q}}" + mock_client_class.return_value = mock_client + + cfg = { + "project": "g/s/r", + "access_token": "tkn", + "prompts_path": "prompts/chat", + } + mgr = GitLabPromptManager(cfg) + + _msgs, _params = mgr.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="folder/sub/my_prompt", + prompt_variables={"q": "ok"}, + prompt_version="commit-sha-999", + ) + + # Path should include prompts_path and end with .prompt + mock_client.get_file_content.assert_any_call( + "prompts/chat/folder/sub/my_prompt.prompt", ref="commit-sha-999" + ) diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py b/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py new file mode 100644 index 00000000000..5b533fe7653 --- /dev/null +++ b/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py @@ -0,0 +1,477 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path + +from litellm.integrations.gitlab.gitlab_client import GitLabClient +from litellm.integrations.gitlab.gitlab_prompt_manager import ( + GitLabPromptManager, + GitLabPromptTemplate, +) + +# ----------------------- +# GitLabPromptTemplate +# ----------------------- + +def test_gitlab_prompt_template_creation(): + """Test GitLabPromptTemplate creation and metadata extraction.""" + metadata = { + "model": "gpt-4", + "temperature": 0.7, + "input": {"schema": {"text": "string"}}, + "output": {"format": "json"}, + } + + template = GitLabPromptTemplate( + template_id="test_template", + content="Hello {{name}}!", + metadata=metadata, + ) + + assert template.template_id == "test_template" + assert template.content == "Hello {{name}}!" + assert template.model == "gpt-4" + assert template.optional_params["temperature"] == 0.7 + assert template.input_schema == {"text": "string"} + + +# ----------------------- +# GitLabClient init & validation +# ----------------------- + +def test_gitlab_client_initialization_token_vs_oauth(): + """Test GitLabClient initialization with token and oauth auth methods.""" + # token (default) + config_token = { + "project": "group/sub/repo", + "access_token": "glpat-XYZ", + "branch": "main", + } + client = GitLabClient(config_token) + assert client.project == "group/sub/repo" + assert client.access_token == "glpat-XYZ" + assert client.branch == "main" + assert client.auth_method == "token" + # token header is used + assert client.headers.get("Private-Token") == "glpat-XYZ" + assert "Authorization" not in client.headers + + # oauth + config_oauth = { + "project": 123456, # numeric project id supported + "access_token": "oauth-bearer", + "auth_method": "oauth", + } + client_oauth = GitLabClient(config_oauth) + assert client_oauth.auth_method == "oauth" + assert client_oauth.headers.get("Authorization") == "Bearer oauth-bearer" + assert "Private-Token" not in client_oauth.headers + + +def test_gitlab_client_missing_required_fields(): + """Test GitLabClient initialization with missing required fields.""" + with pytest.raises(ValueError, match="project and access_token are required"): + GitLabClient({"project": "group/x/repo"}) + with pytest.raises(ValueError, match="project and access_token are required"): + GitLabClient({"access_token": "tok"}) + + +# ----------------------- +# GitLabClient: get_file_content +# ----------------------- + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") +def test_gitlab_client_get_file_content_raw_success(mock_get): + """Successful file content retrieval via RAW endpoint.""" + mock_response = MagicMock() + mock_response.text = "file content" + mock_response.content = b"file content" + mock_response.headers = {"content-type": "text/plain"} + mock_response.status_code = 200 + mock_response.raise_for_status.return_value = None + mock_get.return_value = mock_response + + client = GitLabClient({"project": "g/s/r", "access_token": "tok"}) + content = client.get_file_content("prompts/test.prompt") + assert content == "file content" + mock_get.assert_called_once() + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") +def test_gitlab_client_get_file_content_raw_404_fallback_json_base64(mock_get): + """When RAW returns 404, fallback to JSON endpoint and decode base64 content.""" + import base64 + + # First RAW 404 + resp_raw = MagicMock() + resp_raw.status_code = 404 + resp_raw.raise_for_status.side_effect = Exception() + mock_get.side_effect = [resp_raw] + + # Then JSON OK + resp_json = MagicMock() + encoded = base64.b64encode(b"json-content").decode("utf-8") + resp_json.json.return_value = {"content": encoded, "encoding": "base64"} + resp_json.status_code = 200 + resp_json.raise_for_status.return_value = None + + # We need mock_get to return JSON response second time; easiest: reset side_effect to list of returns + def side_effect(url, headers): + if "/raw?" in url: + return resp_raw + else: + return resp_json + + mock_get.side_effect = side_effect + + client = GitLabClient({"project": "g/s/r", "access_token": "tok"}) + content = client.get_file_content("prompts/test.prompt") + assert content == "json-content" + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") +def test_gitlab_client_get_file_content_not_found(mock_get): + """File not found returns None.""" + # Simulate RAW 404 and JSON 404 + resp_404 = MagicMock() + resp_404.status_code = 404 + resp_404.raise_for_status.side_effect = Exception() + def side_effect(url, headers): + return resp_404 + mock_get.side_effect = side_effect + + client = GitLabClient({"project": "g/s/r", "access_token": "tok"}) + content = client.get_file_content("missing.prompt") + assert content is None + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") +def test_gitlab_client_get_file_content_access_denied(mock_get): + """403 raises a helpful message.""" + import httpx + resp = MagicMock() + resp.status_code = 403 + # raise_for_status inside client only called on non-404 success path; + # simulate exception path by making the request itself raise an httpx error wrapper + err = httpx.HTTPStatusError("403", request=MagicMock(), response=resp) + mock_get.side_effect = err + + client = GitLabClient({"project": "g/s/r", "access_token": "tok"}) + with pytest.raises(Exception, match="Access denied to file 'test.prompt'"): + client.get_file_content("test.prompt") + + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") +def test_gitlab_client_get_file_content_auth_failed(mock_get): + """401 raises auth error.""" + import httpx + resp = MagicMock() + resp.status_code = 401 + err = httpx.HTTPStatusError("401", request=MagicMock(), response=resp) + mock_get.side_effect = err + + client = GitLabClient({"project": "g/s/r", "access_token": "tok"}) + with pytest.raises(Exception, match="Authentication failed"): + client.get_file_content("test.prompt") + + +# ----------------------- +# GitLabClient: list_files +# ----------------------- + +@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.get") +def test_gitlab_client_list_files_success(mock_get): + """List .prompt files via repository tree API.""" + mock_response = MagicMock() + mock_response.json.return_value = [ + {"type": "blob", "path": "prompts/test1.prompt"}, + {"type": "blob", "path": "prompts/test2.prompt"}, + {"type": "blob", "path": "prompts/other.txt"}, + {"type": "tree", "path": "prompts/subdir"}, + ] + mock_response.status_code = 200 + mock_response.raise_for_status.return_value = None + mock_get.return_value = mock_response + + client = GitLabClient({"project": "g/s/r", "access_token": "tok"}) + files = client.list_files("prompts", ".prompt", recursive=True) + + assert files == ["prompts/test1.prompt", "prompts/test2.prompt"] + + +# ----------------------- +# GitLabTemplateManager: parsing & rendering +# ----------------------- + +def test_gitlab_prompt_manager_parse_prompt_file(): + """Parse .prompt with YAML frontmatter.""" + prompt_content = """--- +model: gpt-4 +temperature: 0.7 +max_tokens: 150 +input: + schema: + user_message: string + system_context?: string +--- + +{% if system_context %}System: {{system_context}} + +{% endif %}User: {{user_message}}""" + + manager = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}) + template = manager.prompt_manager._parse_prompt_file(prompt_content, "test_prompt") + + assert template.template_id == "test_prompt" + assert template.model == "gpt-4" + assert template.temperature == 0.7 + assert template.max_tokens == 150 + assert template.input_schema == {"user_message": "string", "system_context?": "string"} + assert "{% if system_context %}" in template.content + + +def test_gitlab_prompt_manager_parse_prompt_file_no_frontmatter(): + """Parse .prompt without YAML frontmatter.""" + prompt_content = "Simple prompt: {{message}}" + manager = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}) + template = manager.prompt_manager._parse_prompt_file(prompt_content, "simple_prompt") + assert template.template_id == "simple_prompt" + assert template.content == "Simple prompt: {{message}}" + assert template.metadata == {} + + +def test_gitlab_prompt_manager_render_template_and_errors(): + """Render a stored template; error if missing.""" + manager = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}) + + tpl = GitLabPromptTemplate( + template_id="t1", + content="Hello {{name}}! Welcome to {{place}}.", + metadata={"model": "gpt-4"}, + ) + manager.prompt_manager.prompts["t1"] = tpl + + rendered = manager.prompt_manager.render_template("t1", {"name": "World", "place": "Earth"}) + assert rendered == "Hello World! Welcome to Earth." + + with pytest.raises(ValueError, match="Template 'nope' not found"): + manager.prompt_manager.render_template("nope", {}) + + +# ----------------------- +# GitLabPromptManager: integration & behavior +# ----------------------- + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_integration(mock_client_class): + """Load prompt on init and render.""" + mock_client = MagicMock() + mock_client.get_file_content.return_value = """--- +model: gpt-4 +temperature: 0.7 +--- +Hello {{name}}!""" + mock_client_class.return_value = mock_client + + mgr = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}, prompt_id="test_prompt") + assert "test_prompt" in mgr.prompt_manager.prompts + + template = mgr.prompt_manager.prompts["test_prompt"] + assert template.model == "gpt-4" + assert template.temperature == 0.7 + + rendered = mgr.prompt_manager.render_template("test_prompt", {"name": "World"}) + assert rendered == "Hello World!" + + +def test_gitlab_prompt_manager_parse_prompt_to_messages(): + """Parse prompt content into chat messages.""" + mgr = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}) + + # single user msg + simple = "Hello there!" + msgs = mgr._parse_prompt_to_messages(simple) + assert msgs == [{"role": "user", "content": "Hello there!"}] + + # multi-role + multi = """System: You are helpful. + +User: Hi? + +Assistant: Hello!""" + msgs = mgr._parse_prompt_to_messages(multi) + assert len(msgs) == 3 + assert msgs[0]["role"] == "system" and msgs[0]["content"] == "You are helpful." + assert msgs[1]["role"] == "user" and msgs[1]["content"] == "Hi?" + assert msgs[2]["role"] == "assistant" and msgs[2]["content"] == "Hello!" + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_pre_call_hook_basic(mock_client_class): + """Pre-call hook parses messages and injects params.""" + mock_client = MagicMock() + mock_client.get_file_content.return_value = """--- +model: gpt-4 +temperature: 0.7 +--- +System: You are helpful. + +User: {{q}}""" + mock_client_class.return_value = mock_client + + mgr = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}, prompt_id="p1") + + original = [{"role": "user", "content": "ignored"}] + msgs, params = mgr.pre_call_hook( + user_id="u", + messages=original, + litellm_params={}, + prompt_id="p1", + prompt_variables={"q": "What is AI?"}, + ) + + assert len(msgs) == 2 + assert msgs[0]["role"] == "system" + assert msgs[1]["role"] == "user" and msgs[1]["content"] == "What is AI?" + assert params["model"] == "gpt-4" and params["temperature"] == 0.7 + + +def test_gitlab_prompt_manager_pre_call_hook_no_prompt_id(): + """If no prompt_id provided, messages/params unchanged.""" + mgr = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}) + original = [{"role": "user", "content": "Hello"}] + msgs, params = mgr.pre_call_hook(user_id="u", messages=original, litellm_params={}, prompt_id=None) + assert msgs == original and params == {} + + +def test_gitlab_prompt_manager_get_available_prompts(): + """Return keys of stored templates.""" + mgr = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}) + mgr.prompt_manager.prompts.update({ + "p1": GitLabPromptTemplate("p1", "c1", {}), + "p2": GitLabPromptTemplate("p2", "c2", {}), + }) + assert set(mgr.get_available_prompts()) == {"p1", "p2"} + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_reload_prompts(mock_client_class): + """Ensure reload resets and re-inits manager.""" + mock_client = MagicMock() + mock_client.get_file_content.return_value = """--- +model: gpt-4 +--- +Hello {{x}}""" + mock_client_class.return_value = mock_client + + mgr = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}, prompt_id="t0") + assert "t0" in mgr.prompt_manager.prompts + + # force reset + with patch.object(mgr, "_prompt_manager", None): + mgr.reload_prompts() + _ = mgr.prompt_manager + # No assertion beyond not raising and property access works + + +# ----------------------- +# YAML fallback parsing +# ----------------------- + +def test_gitlab_prompt_manager_yaml_parsing_fallback_and_types(): + mgr = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}) + yaml_content = """model: gpt-4 +temperature: 0.7 +max_tokens: 150 +enabled: true +disabled: false +count: 42 +rate: 0.5""" + parsed = mgr.prompt_manager._parse_yaml_basic(yaml_content) + assert parsed["model"] == "gpt-4" + assert parsed["temperature"] == 0.7 + assert parsed["max_tokens"] == 150 + assert parsed["enabled"] is True + assert parsed["disabled"] is False + assert parsed["count"] == 42 + assert parsed["rate"] == 0.5 + + +# ----------------------- +# prompts_path handling + prompt_version (ref) precedence +# ----------------------- + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_prompts_path_resolution_and_version(mock_client_class): + """prompts_path + explicit prompt_version should produce correct repo path and ref.""" + mock_client = MagicMock() + mock_client.get_file_content.return_value = "User: {{q}}" + mock_client_class.return_value = mock_client + + cfg = { + "project": "g/s/r", + "access_token": "tok", + "prompts_path": "prompts/chat", + } + mgr = GitLabPromptManager(cfg) + + _msgs, _params = mgr.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="folder/sub/my_prompt", + prompt_variables={"q": "ok"}, + prompt_version="commit-sha-999", + ) + + mock_client.get_file_content.assert_any_call( + "prompts/chat/folder/sub/my_prompt.prompt", ref="commit-sha-999" + ) + + +@patch("litellm.integrations.gitlab.gitlab_prompt_manager.GitLabClient") +def test_gitlab_prompt_manager_version_precedence(mock_client_class): + """ + prompt_version > git_ref kwarg > manager _ref_override. + """ + mock_client = MagicMock() + mock_client.get_file_content.return_value = "User: {{q}}" + mock_client_class.return_value = mock_client + + mgr = GitLabPromptManager({"project": "g/s/r", "access_token": "tok"}, ref="manager-default") + + # prompt_version wins over git_ref kwarg + _msgs, _params = mgr.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="pA", + prompt_variables={"q": "hello"}, + prompt_version="sha-111", + git_ref="feature/branch-xyz", + ) + mock_client.get_file_content.assert_any_call("pA.prompt", ref="sha-111") + + # If no prompt_version, use git_ref kwarg + _msgs, _params = mgr.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="pB", + prompt_variables={"q": "hello"}, + git_ref="hotfix/ref-2", + ) + mock_client.get_file_content.assert_any_call("pB.prompt", ref="hotfix/ref-2") + + # If neither provided, fall back to manager override + _msgs, _params = mgr.pre_call_hook( + user_id="u", + messages=[], + litellm_params={}, + prompt_id="pC", + prompt_variables={"q": "hello"}, + ) + mock_client.get_file_content.assert_any_call("pC.prompt", ref="manager-default") diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index fa2dc3e7190..39ecdb630cf 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -109,6 +109,7 @@ class TestLangfuseUsageDetails(unittest.TestCase): usage_details: LangfuseUsageDetails = { "input": 10, "output": 20, + "total": 30, "cache_creation_input_tokens": 5, "cache_read_input_tokens": 3 } @@ -116,6 +117,7 @@ class TestLangfuseUsageDetails(unittest.TestCase): # Verify all fields are present self.assertEqual(usage_details["input"], 10) self.assertEqual(usage_details["output"], 20) + self.assertEqual(usage_details["total"], 30) self.assertEqual(usage_details["cache_creation_input_tokens"], 5) self.assertEqual(usage_details["cache_read_input_tokens"], 3) @@ -123,12 +125,14 @@ class TestLangfuseUsageDetails(unittest.TestCase): minimal_usage_details: LangfuseUsageDetails = { "input": 10, "output": 20, + "total": 30, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0 } self.assertEqual(minimal_usage_details["input"], 10) self.assertEqual(minimal_usage_details["output"], 20) + self.assertEqual(minimal_usage_details["total"], 30) def test_log_langfuse_v2_usage_details(self): """Test that usage_details in _log_langfuse_v2 is correctly typed and assigned""" @@ -183,6 +187,7 @@ class TestLangfuseUsageDetails(unittest.TestCase): usage_details: LangfuseUsageDetails = { "input": 10, "output": 20, + "total": 30, "cache_creation_input_tokens": None, "cache_read_input_tokens": None } @@ -190,6 +195,7 @@ class TestLangfuseUsageDetails(unittest.TestCase): # Verify fields can be None self.assertEqual(usage_details["input"], 10) self.assertEqual(usage_details["output"], 20) + self.assertEqual(usage_details["total"], 30) self.assertIsNone(usage_details["cache_creation_input_tokens"]) self.assertIsNone(usage_details["cache_read_input_tokens"]) @@ -202,6 +208,7 @@ class TestLangfuseUsageDetails(unittest.TestCase): usage_details = { "input": 15, "output": 25, + "total": 40, "cache_creation_input_tokens": 7, "cache_read_input_tokens": 4 } @@ -209,12 +216,14 @@ class TestLangfuseUsageDetails(unittest.TestCase): # Verify the structure matches what we expect self.assertIn("input", usage_details) self.assertIn("output", usage_details) + self.assertIn("total", usage_details) self.assertIn("cache_creation_input_tokens", usage_details) self.assertIn("cache_read_input_tokens", usage_details) # Verify the values self.assertEqual(usage_details["input"], 15) self.assertEqual(usage_details["output"], 25) + self.assertEqual(usage_details["total"], 40) self.assertEqual(usage_details["cache_creation_input_tokens"], 7) self.assertEqual(usage_details["cache_read_input_tokens"], 4) diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 7fb91f274d0..e605718d29f 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -751,3 +751,58 @@ class TestOpenTelemetry(unittest.TestCase): # ─── no events when only metrics enabled ───────────────────────────────── logs = log_exporter.get_finished_logs() self.assertFalse(logs, "Did not expect any logs") + + def test_get_span_name_with_generation_name(self): + """Test _get_span_name returns generation_name when present""" + otel = OpenTelemetry() + kwargs = { + "litellm_params": { + "metadata": { + "generation_name": "custom_span" + } + } + } + result = otel._get_span_name(kwargs) + self.assertEqual(result, "custom_span") + + def test_get_span_name_without_generation_name(self): + """Test _get_span_name returns default when generation_name missing""" + from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME + + otel = OpenTelemetry() + kwargs = {"litellm_params": {"metadata": {}}} + result = otel._get_span_name(kwargs) + self.assertEqual(result, LITELLM_REQUEST_SPAN_NAME) + + @patch('litellm.turn_off_message_logging', False) + def test_maybe_log_raw_request_creates_span(self): + """Test _maybe_log_raw_request creates span when logging enabled""" + from litellm.integrations.opentelemetry import RAW_REQUEST_SPAN_NAME + + otel = OpenTelemetry() + otel.message_logging = True + + mock_tracer = MagicMock() + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + otel.get_tracer_to_use_for_request = MagicMock(return_value=mock_tracer) + otel.set_raw_request_attributes = MagicMock() + otel._to_ns = MagicMock(return_value=1234567890) + + kwargs = {"litellm_params": {"metadata": {}}} + otel._maybe_log_raw_request(kwargs, {}, datetime.now(), datetime.now(), MagicMock()) + + mock_tracer.start_span.assert_called_once() + self.assertEqual(mock_tracer.start_span.call_args[1]['name'], RAW_REQUEST_SPAN_NAME) + + @patch('litellm.turn_off_message_logging', True) + def test_maybe_log_raw_request_skips_when_logging_disabled(self): + """Test _maybe_log_raw_request skips when logging disabled""" + otel = OpenTelemetry() + mock_tracer = MagicMock() + otel.get_tracer_to_use_for_request = MagicMock(return_value=mock_tracer) + + kwargs = {"litellm_params": {"metadata": {}}} + otel._maybe_log_raw_request(kwargs, {}, datetime.now(), datetime.now(), MagicMock()) + + mock_tracer.start_span.assert_not_called() diff --git a/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py b/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py index 85fcc2700bc..f21cd56750b 100644 --- a/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py +++ b/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py @@ -41,8 +41,10 @@ class TestLangfuseInMemoryCache: "litellm.integrations.langfuse.langfuse.LangFuseLogger", MockLangFuseLogger ): # Add the mock logger to cache with expired TTL + expired_time = time.time() - 1 # Already expired self.cache.cache_dict["test_key"] = mock_logger - self.cache.ttl_dict["test_key"] = time.time() - 1 # Already expired + self.cache.ttl_dict["test_key"] = expired_time + self.cache.expiration_heap = [(expired_time, "test_key")] initial_count = litellm.initialized_langfuse_clients diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index 81d64d70578..2ef2020b09a 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -11,7 +11,9 @@ def config() -> AzureOpenAIGPT5Config: def test_azure_gpt5_supports_reasoning_effort(config: AzureOpenAIGPT5Config): assert "reasoning_effort" in config.get_supported_openai_params(model="gpt-5") - assert "reasoning_effort" in config.get_supported_openai_params(model="gpt5_series/my-deployment") + assert "reasoning_effort" in config.get_supported_openai_params( + model="gpt5_series/my-deployment" + ) def test_azure_gpt5_maps_max_tokens(config: AzureOpenAIGPT5Config): @@ -46,3 +48,56 @@ def test_azure_gpt5_series_transform_request(config: AzureOpenAIGPT5Config): headers={}, ) assert request["model"] == "gpt-5" + + +# GPT-5-Codex specific tests for Azure +def test_azure_gpt5_codex_model_detection(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5-Codex models are correctly detected.""" + assert config.is_model_gpt_5_model("gpt-5-codex") + assert config.is_model_gpt_5_model("gpt5_series/gpt-5-codex") + + +def test_azure_gpt5_codex_supports_reasoning_effort(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5-Codex supports reasoning_effort parameter.""" + assert "reasoning_effort" in config.get_supported_openai_params(model="gpt-5-codex") + assert "reasoning_effort" in config.get_supported_openai_params( + model="gpt5_series/gpt-5-codex" + ) + + +def test_azure_gpt5_codex_maps_max_tokens(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5-Codex correctly maps max_tokens to max_completion_tokens.""" + params = config.map_openai_params( + non_default_params={"max_tokens": 150}, + optional_params={}, + model="gpt-5-codex", + drop_params=False, + api_version="2024-05-01-preview", + ) + assert params["max_completion_tokens"] == 150 + assert "max_tokens" not in params + + +def test_azure_gpt5_codex_temperature_error(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5-Codex raises error for unsupported temperature.""" + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"temperature": 0.8}, + optional_params={}, + model="gpt-5-codex", + drop_params=False, + api_version="2024-05-01-preview", + ) + + +def test_azure_gpt5_codex_series_transform_request(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5-Codex series routing works correctly.""" + request = config.transform_request( + model="gpt5_series/gpt-5-codex", + messages=[], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert request["model"] == "gpt-5-codex" + diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py new file mode 100644 index 00000000000..436ca6e0421 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -0,0 +1,336 @@ +import json +import os +import sys +from unittest.mock import Mock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.llms.base import HiddenParams + +# Mock async invoke responses +async_invoke_response = { + "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" +} + +async_invoke_status_response = { + "status": "InProgress", + "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456", + "outputDataConfig": { + "s3OutputDataConfig": { + "s3Uri": "s3://test-bucket/async-invoke-output/" + } + } +} + +async_invoke_completed_response = { + "status": "Completed", + "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456", + "outputDataConfig": { + "s3OutputDataConfig": { + "s3Uri": "s3://test-bucket/async-invoke-output/" + } + } +} + +# Test data +test_input = "Hello world from litellm async invoke" +test_image_base64 = "data:image/png,test_image_base64_data" + + +class TestBedrockAsyncInvokeEmbedding: + """Test suite for Bedrock async-invoke embedding functionality.""" + + def test_async_invoke_response_transformation_twelvelabs(self): + """Test that async invoke responses are properly transformed with hidden params.""" + from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig + + config = TwelveLabsMarengoEmbeddingConfig() + response = config._transform_async_invoke_response(async_invoke_response, "test-model") + + # Verify response structure + assert isinstance(response, litellm.EmbeddingResponse) + assert hasattr(response, '_hidden_params') + assert response._hidden_params is not None + + # Verify hidden params contain invocation ARN + assert hasattr(response._hidden_params, '_invocation_arn') + assert response._hidden_params._invocation_arn == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" + + # Verify embedding structure + assert len(response.data) == 1 + assert response.data[0].object == "embedding" + assert response.data[0].embedding == [] # Empty for async jobs + assert response.data[0].index == 0 + + def test_async_invoke_response_transformation_generic(self): + """Test that generic async invoke responses are properly transformed.""" + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + bedrock_embedding = BedrockEmbedding() + + # Mock the transformation method + response_list = [async_invoke_response] + response = bedrock_embedding._transform_response( + response_list=response_list, + model="test-model", + provider="twelvelabs", + is_async_invoke=True + ) + + # Verify response structure + assert isinstance(response, litellm.EmbeddingResponse) + assert hasattr(response, '_hidden_params') + assert response._hidden_params is not None + + # Verify hidden params contain invocation ARN + assert hasattr(response._hidden_params, '_invocation_arn') + assert response._hidden_params._invocation_arn == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" + + @pytest.mark.parametrize( + "model,input_type", + [ + ("bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0", "text"), + ("bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0", "image"), + ("bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0", "video"), + ("bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0", "audio"), + ], + ) + def test_async_invoke_twelvelabs_embedding_request_transformation(self, model, input_type): + """Test that async invoke requests are properly transformed for TwelveLabs.""" + from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig + + config = TwelveLabsMarengoEmbeddingConfig() + + # Test input based on type + if input_type == "text": + input_data = test_input + elif input_type == "image": + input_data = test_image_base64 + elif input_type in ["video", "audio"]: + input_data = "s3://test-bucket/test-file.mp4" if input_type == "video" else "s3://test-bucket/test-file.wav" + + inference_params = { + "inputType": input_type, # This will be set by the parameter mapping + "output_s3_uri": "s3://test-bucket/async-invoke-output/" + } + + transformed_request = config._transform_request( + input=input_data, + inference_params=inference_params, + async_invoke_route=True, + model_id="twelvelabs.marengo-embed-2-7-v1:0", + output_s3_uri="s3://test-bucket/async-invoke-output/" + ) + + # Verify async invoke request structure + assert "modelId" in transformed_request + assert "modelInput" in transformed_request + assert "outputDataConfig" in transformed_request + assert transformed_request["modelId"] == "twelvelabs.marengo-embed-2-7-v1:0" + assert transformed_request["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] == "s3://test-bucket/async-invoke-output/" + + def test_async_invoke_twelvelabs_embedding_with_mock(self): + """Test async invoke embedding with mocked HTTP calls.""" + litellm.set_verbose = True + client = HTTPHandler() + test_api_key = "test-bearer-token-12345" + model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0" + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(async_invoke_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model=model, + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key=test_api_key, + input_type="text", # New input_type parameter (maps to inputType) + output_s3_uri="s3://test-bucket/async-invoke-output/" + ) + + # Verify response structure + assert isinstance(response, litellm.EmbeddingResponse) + assert hasattr(response, '_hidden_params') + assert response._hidden_params is not None + assert hasattr(response._hidden_params, '_invocation_arn') + assert response._hidden_params._invocation_arn == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" + + # Verify request was made to async-invoke endpoint + request_url = mock_post.call_args.kwargs.get("url", "") + assert "/async-invoke" in request_url + + @pytest.mark.asyncio + async def test_async_invoke_twelvelabs_embedding_async_with_mock(self): + """Test async invoke embedding with async calls.""" + litellm.set_verbose = True + client = AsyncHTTPHandler() + test_api_key = "test-bearer-token-12345" + model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0" + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(async_invoke_response) + mock_response.json = Mock(return_value=async_invoke_response) + mock_post.return_value = mock_response + + response = await litellm.aembedding( + model=model, + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key=test_api_key, + inputType="text", + output_s3_uri="s3://test-bucket/async-invoke-output/" + ) + + # Verify response structure + assert isinstance(response, litellm.EmbeddingResponse) + assert hasattr(response, '_hidden_params') + assert response._hidden_params is not None + assert hasattr(response._hidden_params, '_invocation_arn') + assert response._hidden_params._invocation_arn == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" + + @pytest.mark.asyncio + async def test_async_invoke_status_checking(self): + """Test async invoke status checking functionality.""" + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + bedrock_embedding = BedrockEmbedding() + + # Mock the async status check + with patch.object(bedrock_embedding, '_get_async_invoke_status') as mock_status: + mock_status.return_value = async_invoke_status_response + + # This would be called internally, but we can test the method directly + status_response = await bedrock_embedding._get_async_invoke_status( + invocation_arn="arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456", + aws_region_name="us-east-1" + ) + + assert status_response["status"] == "InProgress" + assert "invocationArn" in status_response + + def test_async_invoke_error_handling_missing_output_s3_uri(self): + """Test error handling when output_s3_uri is missing for async invoke.""" + from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig + + config = TwelveLabsMarengoEmbeddingConfig() + + with pytest.raises(ValueError, match="output_s3_uri cannot be empty for async invoke requests"): + config._transform_request( + input=test_input, + inference_params={"inputType": "text"}, + async_invoke_route=True, + model_id="twelvelabs.marengo-embed-2-7-v1:0", + output_s3_uri="" # Empty S3 URI should raise error + ) + + def test_async_invoke_error_handling_video_audio_without_async_route(self): + """Test error handling when video/audio input is used without async invoke route.""" + from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig + + config = TwelveLabsMarengoEmbeddingConfig() + + with pytest.raises(ValueError, match="Input type 'video' requires async_invoke route"): + config._transform_request( + input="s3://test-bucket/test-video.mp4", + inference_params={"inputType": "video"}, + async_invoke_route=False, # Should fail for video without async route + model_id="twelvelabs.marengo-embed-2-7-v1:0", + output_s3_uri="s3://test-bucket/async-invoke-output/" + ) + + def test_async_invoke_invocation_arn_preservation(self): + """Test that invocation ARN is correctly preserved in hidden params.""" + from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig + + config = TwelveLabsMarengoEmbeddingConfig() + + # Test various ARN formats + test_cases = [ + "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456", + "arn:aws:bedrock:us-west-2:987654321098:async-invoke/xyz789", + "invalid-arn", + "", + ] + + for arn in test_cases: + mock_response = {"invocationArn": arn} + response = config._transform_async_invoke_response(mock_response, "test-model") + + assert response._hidden_params._invocation_arn == arn + + def test_async_invoke_hidden_params_structure(self): + """Test that hidden params have the correct structure and can be accessed.""" + from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig + + config = TwelveLabsMarengoEmbeddingConfig() + response = config._transform_async_invoke_response(async_invoke_response, "test-model") + + # Test that hidden params can be accessed like a dictionary + assert response._hidden_params.get("_invocation_arn") == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" + + # Test that hidden params can be accessed like attributes + assert response._hidden_params._invocation_arn == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" + + # Test that hidden params can be accessed with bracket notation + assert response._hidden_params["_invocation_arn"] == "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" + + def test_async_invoke_model_parsing(self): + """Test that async invoke models are correctly parsed.""" + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + bedrock_embedding = BedrockEmbedding() + + # Test model parsing + test_models = [ + "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0", + "bedrock/async_invoke/amazon.titan-embed-text-v1", + "bedrock/async_invoke/cohere.embed-english-v3", + ] + + for model in test_models: + # Check if async invoke is detected + has_async_invoke = "async_invoke/" in model + assert has_async_invoke, f"Model {model} should be detected as async invoke" + + # Check model ID extraction (remove both "bedrock/" and "async_invoke/" prefixes) + if has_async_invoke: + model_id = model.replace("bedrock/async_invoke/", "", 1) + assert model_id in [ + "twelvelabs.marengo-embed-2-7-v1:0", + "amazon.titan-embed-text-v1", + "cohere.embed-english-v3" + ] + + def test_async_invoke_endpoint_construction(self): + """Test that async invoke endpoints are correctly constructed.""" + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + bedrock_embedding = BedrockEmbedding() + + # Mock the get_runtime_endpoint method + with patch.object(bedrock_embedding, 'get_runtime_endpoint') as mock_endpoint: + mock_endpoint.return_value = ("https://bedrock-runtime.us-east-1.amazonaws.com", None) + + # Test endpoint construction for async invoke + endpoint_url, _ = bedrock_embedding.get_runtime_endpoint( + api_base=None, + aws_bedrock_runtime_endpoint=None, + aws_region_name="us-east-1" + ) + + # For async invoke, the endpoint should be modified + async_endpoint = f"{endpoint_url}/async-invoke" + assert async_endpoint == "https://bedrock-runtime.us-east-1.amazonaws.com/async-invoke" diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 081176209a8..f436c66f203 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -59,14 +59,21 @@ def test_bedrock_embedding_with_api_key_bearer_token(model, input_type, embed_re input_data = test_image_base64 if input_type == "image" else test_input - response = litellm.embedding( - model=model, - input=input_data, - client=client, - aws_region_name="us-east-1", - aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", - api_key=test_api_key - ) + # Add inputType parameter for TwelveLabs Marengo models + kwargs = { + "model": model, + "input": input_data, + "client": client, + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-east-1.amazonaws.com", + "api_key": test_api_key + } + + # Add input_type parameter for TwelveLabs Marengo models (maps to inputType) + if "twelvelabs.marengo-embed" in model: + kwargs["input_type"] = input_type + + response = litellm.embedding(**kwargs) assert isinstance(response, litellm.EmbeddingResponse) assert isinstance(response.data[0]['embedding'], list) @@ -241,4 +248,160 @@ def test_bedrock_titan_v2_encoding_format_base64(): # Verify that the request contains embeddingTypes: ["binary"] for base64 encoding request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) assert "embeddingTypes" in request_body - assert request_body["embeddingTypes"] == ["binary"] \ No newline at end of file + assert request_body["embeddingTypes"] == ["binary"] + + +def test_twelvelabs_input_type_parameter_mapping(): + """Test that input_type parameter is correctly mapped to inputType for TwelveLabs models""" + litellm.set_verbose = True + client = HTTPHandler() + test_api_key = "test-bearer-token-12345" + model = "bedrock/twelvelabs.marengo-embed-2-7-v1:0" + + twelvelabs_response = { + "data": [{ + "embedding": [0.1, 0.2, 0.3], + "inputTextTokenCount": 10 + }] + } + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(twelvelabs_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + # Test with input_type parameter (new LiteLLM parameter) + response = litellm.embedding( + model=model, + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key=test_api_key, + input_type="text" # New parameter that should map to inputType + ) + + assert isinstance(response, litellm.EmbeddingResponse) + assert isinstance(response.data[0]['embedding'], list) + assert len(response.data[0]['embedding']) == 3 + + # Verify that the request contains inputType (mapped from input_type) + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + assert "inputType" in request_body + assert request_body["inputType"] == "text" + assert "input_type" not in request_body # Should be mapped, not passed through + + +def test_twelvelabs_input_type_parameter_mapping_async_invoke(): + """Test that input_type parameter is correctly mapped to inputType for TwelveLabs async invoke models""" + litellm.set_verbose = True + client = HTTPHandler() + test_api_key = "test-bearer-token-12345" + model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0" + + async_invoke_response = { + "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456" + } + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(async_invoke_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + # Test with input_type parameter for async invoke + response = litellm.embedding( + model=model, + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key=test_api_key, + output_s3_uri="s3://test-bucket/async-invoke-output/", + input_type="text" # New parameter that should map to inputType + ) + + assert isinstance(response, litellm.EmbeddingResponse) + assert hasattr(response, '_hidden_params') + assert response._hidden_params is not None + assert hasattr(response._hidden_params, '_invocation_arn') + + # Verify that the request contains inputType in modelInput (mapped from input_type) + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + assert "modelInput" in request_body + assert "inputType" in request_body["modelInput"] + assert request_body["modelInput"]["inputType"] == "text" + assert "input_type" not in request_body # Should be mapped, not passed through + + +def test_twelvelabs_missing_input_type_error(): + """Test that missing input_type parameter defaults to 'text' for TwelveLabs models""" + litellm.set_verbose = True + client = HTTPHandler() + test_api_key = "test-bearer-token-12345" + + # Test TwelveLabs model - should default to 'text' when input_type is missing + twelvelabs_model = "bedrock/twelvelabs.marengo-embed-2-7-v1:0" + twelvelabs_response = { + "data": [{ + "embedding": [0.1, 0.2, 0.3], + "inputTextTokenCount": 10 + }] + } + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(twelvelabs_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + # Test that missing input_type defaults to "text" for TwelveLabs + response = litellm.embedding( + model=twelvelabs_model, + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key=test_api_key + # No input_type parameter - should default to "text" + ) + + # Verify the response is successful + assert isinstance(response, litellm.EmbeddingResponse) + + # Verify that the request contains inputType: "text" by default + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + assert "inputType" in request_body + assert request_body["inputType"] == "text" + + # Test Amazon Titan model - should NOT throw error (input_type not required) + titan_model = "bedrock/amazon.titan-embed-text-v1" + titan_response = { + "embedding": [0.1, 0.2, 0.3], + "inputTextTokenCount": 10 + } + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(titan_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + # Test that missing input_type does NOT throw an error for Amazon Titan + response = litellm.embedding( + model=titan_model, + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key=test_api_key + # No input_type parameter - should work fine + ) + + # Should succeed without input_type + assert isinstance(response, litellm.EmbeddingResponse) \ 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 3e6a6a23468..876eb8b29f9 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -2,16 +2,24 @@ import pytest import litellm from litellm.llms.openai.openai import OpenAIConfig +from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config @pytest.fixture() def config() -> OpenAIConfig: return OpenAIConfig() + +@pytest.fixture() +def gpt5_config() -> OpenAIGPT5Config: + return OpenAIGPT5Config() + + def test_gpt5_supports_reasoning_effort(config: OpenAIConfig): assert "reasoning_effort" in config.get_supported_openai_params(model="gpt-5") assert "reasoning_effort" in config.get_supported_openai_params(model="gpt-5-mini") + def test_gpt5_maps_max_tokens(config: OpenAIConfig): params = config.map_openai_params( non_default_params={"max_tokens": 10}, @@ -52,3 +60,96 @@ def test_gpt5_unsupported_params_drop(config: OpenAIConfig): drop_params=True, ) assert "top_p" not in params + + +# GPT-5-Codex specific tests +def test_gpt5_codex_model_detection(gpt5_config: OpenAIGPT5Config): + """Test that GPT-5-Codex models are correctly detected as GPT-5 models.""" + assert gpt5_config.is_model_gpt_5_model("gpt-5-codex") + assert gpt5_config.is_model_gpt_5_codex_model("gpt-5-codex") + + # Regular GPT-5 models should not be detected as codex + assert not gpt5_config.is_model_gpt_5_codex_model("gpt-5") + assert not gpt5_config.is_model_gpt_5_codex_model("gpt-5-mini") + + +def test_gpt5_codex_supports_reasoning_effort(config: OpenAIConfig): + """Test that GPT-5-Codex supports reasoning_effort parameter.""" + assert "reasoning_effort" in config.get_supported_openai_params(model="gpt-5-codex") + + +def test_gpt5_codex_maps_max_tokens(config: OpenAIConfig): + """Test that GPT-5-Codex correctly maps max_tokens to max_completion_tokens.""" + params = config.map_openai_params( + non_default_params={"max_tokens": 100}, + optional_params={}, + model="gpt-5-codex", + drop_params=False, + ) + assert params["max_completion_tokens"] == 100 + assert "max_tokens" not in params + + +def test_gpt5_codex_temperature_drop(config: OpenAIConfig): + """Test that GPT-5-Codex drops unsupported temperature values when drop_params=True.""" + params = config.map_openai_params( + non_default_params={"temperature": 0.7}, + optional_params={}, + model="gpt-5-codex", + drop_params=True, + ) + assert "temperature" not in params + + +def test_gpt5_codex_temperature_error(config: OpenAIConfig): + """Test that GPT-5-Codex raises error for unsupported temperature when drop_params=False.""" + with pytest.raises( + litellm.utils.UnsupportedParamsError, + match="gpt-5 models \\(including gpt-5-codex\\)", + ): + config.map_openai_params( + non_default_params={"temperature": 0.7}, + optional_params={}, + model="gpt-5-codex", + drop_params=False, + ) + + + +def test_gpt5_codex_temperature_one_allowed(config: OpenAIConfig): + """Test that GPT-5-Codex allows temperature=1.""" + params = config.map_openai_params( + non_default_params={"temperature": 1.0}, + optional_params={}, + model="gpt-5-codex", + drop_params=False, + ) + assert params["temperature"] == 1.0 + + +def test_gpt5_codex_unsupported_params_drop(config: OpenAIConfig): + """Test that GPT-5-Codex drops unsupported parameters.""" + unsupported_params = [ + "top_p", + "presence_penalty", + "frequency_penalty", + "logprobs", + "top_logprobs", + ] + + for param in unsupported_params: + assert param not in config.get_supported_openai_params(model="gpt-5-codex") + + +def test_gpt5_codex_supports_tool_choice(gpt5_config: OpenAIGPT5Config): + """Test that GPT-5-Codex supports tool_choice parameter.""" + supported_params = gpt5_config.get_supported_openai_params(model="gpt-5-codex") + assert "tool_choice" in supported_params + + +def test_gpt5_codex_supports_function_calling(config: OpenAIConfig): + """Test that GPT-5-Codex supports function calling parameters.""" + supported_params = config.get_supported_openai_params(model="gpt-5-codex") + assert "functions" in supported_params + assert "function_call" in supported_params + assert "tools" in supported_params diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 62a11bf6765..5ae6cf3da28 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -444,12 +444,14 @@ def test_vertex_ai_map_thinking_param_with_budget_tokens_0(): def test_vertex_ai_map_tools(): v = VertexGeminiConfig() - tools = v._map_function(value=[{"code_execution": {}}]) + optional_params = {} + tools = v._map_function(value=[{"code_execution": {}}], optional_params=optional_params) assert len(tools) == 1 assert tools[0]["code_execution"] == {} print(tools) - new_tools = v._map_function(value=[{"codeExecution": {}}]) + new_optional_params = {} + new_tools = v._map_function(value=[{"codeExecution": {}}], optional_params=new_optional_params) assert len(new_tools) == 1 print("new_tools", new_tools) assert new_tools[0]["code_execution"] == {} @@ -465,6 +467,7 @@ def test_vertex_ai_map_tool_with_anyof(): Ensure if anyof is present, only the anyof field and its contents are kept - otherwise VertexAI will throw an error - https://github.com/BerriAI/litellm/issues/11164 """ v = VertexGeminiConfig() + optional_params = {} value = [ { "type": "function", @@ -488,7 +491,7 @@ def test_vertex_ai_map_tool_with_anyof(): }, } ] - tools = v._map_function(value=value) + tools = v._map_function(value=value, optional_params=optional_params) assert tools[0]["function_declarations"][0]["parameters"]["properties"][ "base_branch" @@ -496,6 +499,7 @@ def test_vertex_ai_map_tool_with_anyof(): "anyOf": [{"type": "string", "nullable": True, "title": "Base Branch"}] }, f"Expected only anyOf field and its contents to be kept, but got {tools[0]['function_declarations'][0]['parameters']['properties']['base_branch']}" + new_optional_params = {} new_value = [ { "type": "function", @@ -518,7 +522,7 @@ def test_vertex_ai_map_tool_with_anyof(): }, } ] - new_tools = v._map_function(value=new_value) + new_tools = v._map_function(value=new_value, optional_params=new_optional_params) assert new_tools[0]["function_declarations"][0]["parameters"]["properties"][ "base_branch" @@ -1056,3 +1060,94 @@ def test_vertex_ai_code_line_length(): # Verify it contains the expected UUID format assert 'uuid.uuid4().hex[:28]' in id_line, f"Line should contain shortened UUID format: {id_line}" + + +def test_vertex_ai_map_google_maps_tool_simple(): + """ + Test googleMaps tool transformation without location data. + + Input: + value=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}] + optional_params={} + + Expected Output: + tools=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}] + optional_params={} (unchanged) + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}], + optional_params=optional_params + ) + + assert len(tools) == 1 + assert "googleMaps" in tools[0] + assert tools[0]["googleMaps"]["enableWidget"] == "ENABLE_WIDGET" + assert "toolConfig" not in optional_params + + +def test_vertex_ai_map_google_maps_tool_with_location(): + """ + Test googleMaps tool transformation with location data. + Verifies latitude/longitude/languageCode are extracted to toolConfig.retrievalConfig. + + Input: + value=[{ + "googleMaps": { + "enableWidget": "ENABLE_WIDGET", + "latitude": 37.7749, + "longitude": -122.4194, + "languageCode": "en_US" + } + }] + optional_params={} + + Expected Output: + tools=[{ + "googleMaps": {"enableWidget": "ENABLE_WIDGET"} + }] + optional_params={ + "toolConfig": { + "retrievalConfig": { + "latLng": { + "latitude": 37.7749, + "longitude": -122.4194 + }, + "languageCode": "en_US" + } + } + } + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[{ + "googleMaps": { + "enableWidget": "ENABLE_WIDGET", + "latitude": 37.7749, + "longitude": -122.4194, + "languageCode": "en_US" + } + }], + optional_params=optional_params + ) + + assert len(tools) == 1 + assert "googleMaps" in tools[0] + + google_maps_tool = tools[0]["googleMaps"] + assert google_maps_tool["enableWidget"] == "ENABLE_WIDGET" + assert "latitude" not in google_maps_tool + assert "longitude" not in google_maps_tool + assert "languageCode" not in google_maps_tool + + assert "toolConfig" in optional_params + assert "retrievalConfig" in optional_params["toolConfig"] + + retrieval_config = optional_params["toolConfig"]["retrievalConfig"] + assert retrieval_config["latLng"]["latitude"] == 37.7749 + assert retrieval_config["latLng"]["longitude"] == -122.4194 + assert retrieval_config["languageCode"] == "en_US" diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 31d4fd1c198..39ed09f81be 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -145,6 +145,7 @@ def test_build_vertex_schema(): ([{"googleSearchRetrieval": {}}], "googleSearchRetrieval"), ([{"enterpriseWebSearch": {}}], "enterpriseWebSearch"), ([{"code_execution": {}}], "code_execution"), + ([{"googleMaps": {}}], "googleMaps"), ], ) def test_vertex_tool_params(tools, key): diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index a2008c2f336..27defe0eb0b 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -182,6 +182,12 @@ def mock_request(): class QueryParams: def __init__(self): self._dict = {} + + def __iter__(self): + return iter(self._dict) + + def items(self): + return self._dict.items() class MockRequest: def __init__( @@ -291,7 +297,7 @@ async def test_pass_through_request_stream_param_override( "POST", httpx.URL("https://api.anthropic.com/v1/messages"), json=request_body, - params=None, + params={}, headers={ "Authorization": "Bearer test-key" }, @@ -393,7 +399,7 @@ async def test_pass_through_request_stream_param_no_override( headers={ "Authorization": "Bearer test-key" }, - params=None, + params={}, json=request_body, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index ab9ac0acd9b..9e316289c87 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -24,98 +24,80 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @pytest.mark.asyncio class TestMCPRequestHandler: @pytest.mark.parametrize( - "user_api_key_auth,object_permission_id,prisma_client_available,db_result,expected_result", + "key_servers,team_servers,expected_result,scenario", [ - # Test case 1: user_api_key_auth is None - (None, None, True, None, []), - # Test case 2: object_permission_id is None - (UserAPIKeyAuth(), None, True, None, []), - # Test case 3: prisma_client is None + # Test case 1: No key servers, no team servers + ([], [], [], "no_permissions"), + # Test case 2: Key has servers, no team servers + (["server1", "server2"], [], ["server1", "server2"], "key_only"), + # Test case 3: No key servers, team has servers (inherit from team) ( - UserAPIKeyAuth(object_permission_id="test-id"), - "test-id", - False, - None, [], + ["team_server1", "team_server2"], + ["team_server1", "team_server2"], + "inherit_from_team", ), - # Test case 4: Database query returns None - (UserAPIKeyAuth(object_permission_id="test-id"), "test-id", True, None, []), - # Test case 5: Database query returns object with mcp_servers + # Test case 4: Key and team both have servers (intersection) ( - UserAPIKeyAuth(object_permission_id="test-id"), - "test-id", - True, - MagicMock(mcp_servers=["server1", "server2"]), ["server1", "server2"], + ["server1", "team_server"], + ["server1"], + "intersection", ), - # Test case 6: Database query returns object with None mcp_servers + # Test case 5: Key and team have no overlap (empty result) ( - UserAPIKeyAuth(object_permission_id="test-id"), - "test-id", - True, - MagicMock(mcp_servers=None), + ["server1", "server2"], + ["team_server1", "team_server2"], [], + "no_overlap", ), - # Test case 7: Database query returns object with empty mcp_servers + # Test case 6: Key and team have complete overlap ( - UserAPIKeyAuth(object_permission_id="test-id"), - "test-id", - True, - MagicMock(mcp_servers=[]), - [], + ["server1", "server2"], + ["server1", "server2"], + ["server1", "server2"], + "complete_overlap", ), ], ) - async def test_get_allowed_mcp_servers_for_key( + async def test_get_allowed_mcp_servers( self, - user_api_key_auth, - object_permission_id, - prisma_client_available, - db_result, + key_servers, + team_servers, expected_result, + scenario, ): - """Test _get_allowed_mcp_servers_for_key with various scenarios""" + """Test get_allowed_mcp_servers with various key/team permission scenarios""" - # Setup user_api_key_auth object_permission_id if provided - if user_api_key_auth and object_permission_id: - user_api_key_auth.object_permission_id = object_permission_id + # Create a mock user + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + ) - # Mock prisma_client - mock_prisma_client = MagicMock() if prisma_client_available else None - mock_find_unique = None + # Mock the helper methods instead of database calls + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_key" + ) as mock_key_servers: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_team" + ) as mock_team_servers: + # Set up return values + mock_key_servers.return_value = key_servers + mock_team_servers.return_value = team_servers - if mock_prisma_client: - # Mock the database query - mock_find_unique = AsyncMock(return_value=db_result) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = ( - mock_find_unique - ) - - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): - # Call the method - result = await MCPRequestHandler._get_allowed_mcp_servers_for_key( - user_api_key_auth - ) - - # Assert the result (order-independent comparison) - assert sorted(result) == sorted(expected_result) - - # Verify database call was made correctly when expected - if ( - user_api_key_auth - and user_api_key_auth.object_permission_id - and prisma_client_available - and mock_find_unique - ): - mock_find_unique.assert_called_once_with( - where={ - "object_permission_id": user_api_key_auth.object_permission_id - } + # Call the method + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=mock_user_auth ) - elif mock_find_unique: - # If prisma_client exists but conditions aren't met, no call should be made - if not user_api_key_auth or not user_api_key_auth.object_permission_id: - mock_find_unique.assert_not_called() + + # Assert the result (order-independent comparison) + assert sorted(result) == sorted(expected_result) + + # Verify helper methods were called + mock_key_servers.assert_called_once_with(mock_user_auth) + mock_team_servers.assert_called_once_with(mock_user_auth) @pytest.mark.parametrize( "team_servers,key_servers,expected_servers,scenario", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 6a1d43b33e1..a45d91eeb69 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -88,10 +88,14 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): working_server = MagicMock() working_server.name = "working_server" working_server.alias = "working" + working_server.allowed_tools = None + working_server.disallowed_tools = None failing_server = MagicMock() failing_server.name = "failing_server" failing_server.alias = "failing" + failing_server.allowed_tools = None + failing_server.disallowed_tools = None # Mock global_mcp_server_manager mock_manager = MagicMock() @@ -565,6 +569,7 @@ async def test_oauth2_headers_passed_to_mcp_client(): == "Bearer github_oauth_token_12345" ) + @pytest.mark.asyncio async def test_list_tools_single_server_unprefixed_names(): """When only one MCP server is allowed, list tools should return unprefixed names.""" @@ -585,12 +590,14 @@ async def test_list_tools_single_server_unprefixed_names(): server.server_id = "server1" server.name = "Zapier MCP" server.alias = "zapier" + server.allowed_tools = None + server.disallowed_tools = None # Mock manager: allow just one server and return a tool based on add_prefix flag mock_manager = MagicMock() mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["server1"]) - mock_manager.get_mcp_server_by_id = ( - lambda server_id: server if server_id == "server1" else None + mock_manager.get_mcp_server_by_id = lambda server_id: ( + server if server_id == "server1" else None ) async def mock_get_tools_from_server( @@ -640,19 +647,23 @@ async def test_list_tools_multiple_servers_prefixed_names(): server1.server_id = "server1" server1.name = "Zapier MCP" server1.alias = "zapier" + server1.allowed_tools = None + server1.disallowed_tools = None server2 = MagicMock() server2.server_id = "server2" server2.name = "Jira MCP" server2.alias = "jira" + server2.allowed_tools = None + server2.disallowed_tools = None # Mock manager mock_manager = MagicMock() mock_manager.get_allowed_mcp_servers = AsyncMock( return_value=["server1", "server2"] ) - mock_manager.get_mcp_server_by_id = ( - lambda server_id: server1 if server_id == "server1" else server2 + mock_manager.get_mcp_server_by_id = lambda server_id: ( + server1 if server_id == "server1" else server2 ) async def mock_get_tools_from_server( @@ -681,3 +692,45 @@ async def test_list_tools_multiple_servers_prefixed_names(): # Should be prefixed since multiple servers are allowed names = sorted([t.name for t in tools]) assert names == ["jira-toolA", "zapier-toolA"] + + +@pytest.mark.asyncio +async def test_call_mcp_tool_user_unauthorized_access(): + """Test that a user cannot call a tool from a server they don't have access to""" + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.server import call_mcp_tool + from litellm.proxy._types import UserAPIKeyAuth + + # Create a mock user without access to the server + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="team-basic", + object_permission_id="key-permission-123", + ) + + # Mock global_mcp_server_manager.get_mcp_server_names_from_ids to return + # a list that doesn't include "restricted_server" (the server the user is trying to access) + with patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_names_from_ids" + ) as mock_get_server_names: + # User has access to "allowed_server" but not "restricted_server" + mock_get_server_names.return_value = ["allowed_server", "another_server"] + + # Try to call a tool from "restricted_server" - should raise HTTPException with 403 status + with pytest.raises(HTTPException) as exc_info: + await call_mcp_tool( + name="restricted_server-send_email", + arguments={ + "to": "test@example.com", + "subject": "Test", + "body": "Test", + }, + user_api_key_auth=mock_user_auth, + mcp_auth_header="Bearer test_token", + ) + + # Verify the exception details + assert exc_info.value.status_code == 403 + assert "User not allowed to call this tool" in exc_info.value.detail diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 80ea95a2210..146f9434b8c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1,6 +1,6 @@ import sys from datetime import datetime -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException @@ -759,6 +759,156 @@ class TestMCPServerManager: assert resolved_server_pref is not None assert resolved_server_pref.server_id == server.server_id + @pytest.mark.asyncio + async def test_rest_endpoint_filters_by_allowed_tools(self): + """Test that REST endpoint _get_tools_for_single_server respects allowed_tools configuration""" + from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + _get_tools_for_single_server, + ) + + # Create server with allowed_tools configured + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + allowed_tools=["allowed_tool_1", "allowed_tool_2"], + ) + server.mcp_info = {"server_name": "test-server"} + + # Mock tools returned from manager (3 tools, but only 2 are allowed) + tool1 = MagicMock() + tool1.name = "allowed_tool_1" + tool1.description = "This tool is allowed" + tool1.inputSchema = {} + + tool2 = MagicMock() + tool2.name = "blocked_tool" + tool2.description = "This tool is not allowed" + tool2.inputSchema = {} + + tool3 = MagicMock() + tool3.name = "allowed_tool_2" + tool3.description = "This tool is also allowed" + tool3.inputSchema = {} + + # Mock the global_mcp_server_manager._get_tools_from_server + from litellm.proxy._experimental.mcp_server import rest_endpoints + + with patch.object( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + new=AsyncMock(return_value=[tool1, tool2, tool3]), + ): + # Call the REST endpoint helper + filtered_response = await _get_tools_for_single_server( + server, server_auth_header=None + ) + + # Verify only allowed tools are in the response + assert len(filtered_response) == 2 + tool_names = [t.name for t in filtered_response] + assert "allowed_tool_1" in tool_names + assert "allowed_tool_2" in tool_names + assert "blocked_tool" not in tool_names + + @pytest.mark.asyncio + async def test_rest_endpoint_shows_all_when_allowed_tools_is_none(self): + """Test that REST endpoint shows all tools when allowed_tools is None (backwards compatibility)""" + from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + _get_tools_for_single_server, + ) + + # Create server with allowed_tools as None + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + allowed_tools=None, # No filtering + ) + server.mcp_info = {"server_name": "test-server"} + + # Mock tools returned from manager + tool1 = MagicMock() + tool1.name = "tool_1" + tool1.description = "Tool 1" + tool1.inputSchema = {} + + tool2 = MagicMock() + tool2.name = "tool_2" + tool2.description = "Tool 2" + tool2.inputSchema = {} + + tool3 = MagicMock() + tool3.name = "tool_3" + tool3.description = "Tool 3" + tool3.inputSchema = {} + + # Mock the global_mcp_server_manager._get_tools_from_server + from litellm.proxy._experimental.mcp_server import rest_endpoints + + with patch.object( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + new=AsyncMock(return_value=[tool1, tool2, tool3]), + ): + # Call the REST endpoint helper + all_tools_response = await _get_tools_for_single_server( + server, server_auth_header=None + ) + + # Verify all tools are returned (no filtering) + assert len(all_tools_response) == 3 + tool_names = [t.name for t in all_tools_response] + assert "tool_1" in tool_names + assert "tool_2" in tool_names + assert "tool_3" in tool_names + + @pytest.mark.asyncio + async def test_rest_endpoint_shows_all_when_allowed_tools_is_empty_list(self): + """Test that REST endpoint shows all tools when allowed_tools is empty list (backwards compatibility)""" + from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + _get_tools_for_single_server, + ) + + # Create server with allowed_tools as empty list + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + allowed_tools=[], # Empty list means no filtering + ) + server.mcp_info = {"server_name": "test-server"} + + # Mock tools returned from manager + tool1 = MagicMock() + tool1.name = "tool_1" + tool1.description = "Tool 1" + tool1.inputSchema = {} + + tool2 = MagicMock() + tool2.name = "tool_2" + tool2.description = "Tool 2" + tool2.inputSchema = {} + + # Mock the global_mcp_server_manager._get_tools_from_server + from litellm.proxy._experimental.mcp_server import rest_endpoints + + with patch.object( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + new=AsyncMock(return_value=[tool1, tool2]), + ): + # Call the REST endpoint helper + all_tools_response = await _get_tools_for_single_server( + server, server_auth_header=None + ) + + # Verify all tools are returned (no filtering) + assert len(all_tools_response) == 2 + tool_names = [t.name for t in all_tools_response] + assert "tool_1" in tool_names + assert "tool_2" in tool_names + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 539ee4a9ba8..37cff2bd94a 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -247,3 +247,94 @@ def test_anthropic_count_tokens_route_accessible_to_internal_users(): # Also test that the regular messages route still works assert RouteChecks.is_llm_api_route("/v1/messages") is True + + +def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints(): + """ + Test that virtual keys with llm_api_routes permission can access registered pass-through endpoints. + + This tests the scenario where a pass-through endpoint is registered from the DB + (e.g., /azure-assistant) and a virtual key with llm_api_routes permission should be able to access + both the exact path and subpaths (e.g., /azure-assistant/openai/assistants). + """ + from unittest.mock import patch + + # Mock the registered pass-through routes + mock_registered_routes = { + "test-uuid-1:exact:/azure-assistant": { + "endpoint_id": "test-uuid-1", + "path": "/azure-assistant", + "type": "exact", + }, + "test-uuid-2:subpath:/custom-endpoint": { + "endpoint_id": "test-uuid-2", + "path": "/custom-endpoint", + "type": "subpath", + }, + } + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + mock_registered_routes, + ): + # Create a virtual key with llm_api_routes permission + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["llm_api_routes"], + ) + + # Test exact match for registered pass-through endpoint + result1 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/azure-assistant", + valid_token=valid_token, + ) + assert result1 is True + + # Test subpath for registered pass-through endpoint with subpath type + result2 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/custom-endpoint/openai/assistants", + valid_token=valid_token, + ) + assert result2 is True + + # Test exact match for subpath type + result3 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/custom-endpoint", + valid_token=valid_token, + ) + assert result3 is True + + +def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): + """ + Test that virtual keys without llm_api_routes permission cannot access registered pass-through endpoints. + """ + from unittest.mock import patch + + # Mock the registered pass-through routes + mock_registered_routes = { + "test-uuid-1:exact:/azure-assistant": { + "endpoint_id": "test-uuid-1", + "path": "/azure-assistant", + "type": "exact", + }, + } + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + mock_registered_routes, + ): + # Create a virtual key without llm_api_routes permission + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["info_routes"], + ) + + # Test that access is denied + with pytest.raises(Exception) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route( + route="/azure-assistant", + valid_token=valid_token, + ) + + assert "Virtual key is not allowed to call this route" in str(exc_info.value) diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py new file mode 100644 index 00000000000..62e8aaf2794 --- /dev/null +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +""" +Test to verify the Google GenAI proxy API endpoints +""" +import asyncio +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm + + +def test_google_generate_content_endpoint(): + """Test that the google_generate_content endpoint correctly routes requests""" + # Skip this test if we can't import the required modules due to missing dependencies + try: + from fastapi.testclient import TestClient + from litellm.proxy.google_endpoints.endpoints import router as google_router + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + # Create a test client + client = TestClient(google_router) + + # Mock the router's agenerate_content method + with patch("litellm.proxy.proxy_server.llm_router") as mock_router: + mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) + + # Send a request to the endpoint + response = client.post( + "/v1beta/models/test-model:generateContent", + json={ + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] + } + ) + + # Verify the response + assert response.status_code == 200 + assert response.json() == {"test": "response"} + + # Verify that agenerate_content was called + mock_router.agenerate_content.assert_called_once() + + +def test_google_stream_generate_content_endpoint(): + """Test that the google_stream_generate_content endpoint correctly routes streaming requests""" + # Skip this test if we can't import the required modules due to missing dependencies + try: + from fastapi.testclient import TestClient + from litellm.proxy.google_endpoints.endpoints import router as google_router + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + # Create a test client + client = TestClient(google_router) + + # Mock the router's agenerate_content method to return a stream + mock_stream = AsyncMock() + mock_stream.__aiter__ = lambda self: mock_stream + mock_stream.__anext__.side_effect = StopAsyncIteration + + with patch("litellm.proxy.proxy_server.llm_router") as mock_router: + mock_router.agenerate_content = AsyncMock(return_value=mock_stream) + + # Send a request to the endpoint + response = client.post( + "/v1beta/models/test-model:streamGenerateContent", + json={ + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] + } + ) + + # Verify the response + assert response.status_code == 200 + + # Verify that agenerate_content was called with correct parameters + mock_router.agenerate_content.assert_called_once() + call_args = mock_router.agenerate_content.call_args + assert call_args[1]["stream"] is True + assert call_args[1]["model"] == "test-model" + assert call_args[1]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] \ No newline at end of file diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 05e0d4287a3..d3cbd460cf0 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -364,9 +364,12 @@ async def test_100_concurrent_priority_requests(): @pytest.mark.asyncio async def test_concurrent_pre_call_hooks_stress(): """ - Stress test: 50 concurrent pre-call hooks with priority enforcement. + Stress test: 50 concurrent pre-call hooks with saturation-aware priority enforcement. - This tests the actual rate limiting logic under concurrent load. + Tests priority-based rate limiting in strict mode (>80% saturation). + Mocks high saturation to force strict mode where priorities are enforced. + Premium users (80% allocation) should have >90% success rate. + Standard users (20% allocation) should have ~70% success rate with 30% random limiting. """ # Set up environment for premium feature os.environ["LITELLM_LICENSE"] = "test-license-key" @@ -398,52 +401,94 @@ async def test_concurrent_pre_call_hooks_stress(): successful_requests = [] rate_limited_requests = [] - async def mock_should_rate_limit(descriptors, parent_otel_span=None): - """Mock rate limiter that allows premium users, limits some standard users.""" - descriptor = descriptors[0] - priority = descriptor["value"].split(":")[-1] + # Mock saturation check to return high saturation (forces strict mode) + async def mock_get_cache(key, litellm_parent_otel_span=None, local_only=False): + """Mock cache to simulate high saturation.""" + # Return high usage to trigger strict mode (>80% saturation) + if ":requests" in key or ":tokens" in key: + return 1800 # 1800/2000 = 90% saturation + return None - if priority == "premium": - # Allow all premium requests + async def mock_should_rate_limit(descriptors, parent_otel_span=None, read_only=False): + """Mock rate limiter that handles saturation-aware descriptors.""" + descriptor = descriptors[0] + descriptor_key = descriptor["key"] + descriptor_value = descriptor["value"] + + # Handle model-wide tracking (for both generous and strict mode tracking) + if descriptor_key == "model_saturation_check": + # Always allow model-wide tracking (doesn't enforce in our mock) return { "overall_code": "OK", "statuses": [ { "code": "OK", - "descriptor_key": descriptor["value"], + "descriptor_key": descriptor_value, "rate_limit_type": "tokens_per_unit", - "limit_remaining": 1000, + "limit_remaining": 10000, } ], } - else: - # Rate limit some standard requests (simulate load) - import random + + # Handle priority-specific enforcement in strict mode + elif descriptor_key == "priority_model": + # Extract priority from value like "pre-call-stress-model:premium" + priority = descriptor_value.split(":")[-1] - if random.random() < 0.3: # 30% of standard requests get rate limited - return { - "overall_code": "OVER_LIMIT", - "statuses": [ - { - "code": "OVER_LIMIT", - "descriptor_key": descriptor["value"], - "rate_limit_type": "tokens_per_unit", - "limit_remaining": 0, - } - ], - } - else: + if priority == "premium": + # Allow all premium requests return { "overall_code": "OK", "statuses": [ { "code": "OK", - "descriptor_key": descriptor["value"], + "descriptor_key": descriptor_value, "rate_limit_type": "tokens_per_unit", - "limit_remaining": 100, + "limit_remaining": 1000, } ], } + else: + # Rate limit some standard requests (simulate load) + import random + + if random.random() < 0.3: # 30% of standard requests get rate limited + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": descriptor_value, + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 0, + } + ], + } + else: + return { + "overall_code": "OK", + "statuses": [ + { + "code": "OK", + "descriptor_key": descriptor_value, + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 100, + } + ], + } + + # Default: allow + return { + "overall_code": "OK", + "statuses": [ + { + "code": "OK", + "descriptor_key": descriptor_value, + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 1000, + } + ], + } # Create 50 users: 30 premium, 20 standard users = [] @@ -464,42 +509,44 @@ async def test_concurrent_pre_call_hooks_stress(): """Make a pre-call hook request.""" user, priority = user_data - with patch.object( - handler.v3_limiter, "should_rate_limit", side_effect=mock_should_rate_limit - ): - try: - result = await handler.async_pre_call_hook( - user_api_key_dict=user, - cache=DualCache(), - data={"model": model}, - call_type="completion", - ) + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=DualCache(), + data={"model": model}, + call_type="completion", + ) - # If no exception, request was allowed - successful_requests.append( - {"user_id": user.user_id, "priority": priority, "result": "allowed"} - ) - return { - "status": "success", - "user_id": user.user_id, - "priority": priority, - } + # If no exception, request was allowed + successful_requests.append( + {"user_id": user.user_id, "priority": priority, "result": "allowed"} + ) + return { + "status": "success", + "user_id": user.user_id, + "priority": priority, + } - except Exception as e: - # Request was rate limited - rate_limited_requests.append( - {"user_id": user.user_id, "priority": priority, "error": str(e)} - ) - return { - "status": "rate_limited", - "user_id": user.user_id, - "priority": priority, - } + except Exception as e: + # Request was rate limited + rate_limited_requests.append( + {"user_id": user.user_id, "priority": priority, "error": str(e)} + ) + return { + "status": "rate_limited", + "user_id": user.user_id, + "priority": priority, + } - # Run all 50 requests concurrently + # Run all 50 requests concurrently with patches applied to the entire batch start_time = time.time() - tasks = [make_request(user_data) for user_data in users] - results = await asyncio.gather(*tasks, return_exceptions=True) + with patch.object( + handler.v3_limiter, "should_rate_limit", side_effect=mock_should_rate_limit + ), patch.object( + handler.internal_usage_cache, "async_get_cache", side_effect=mock_get_cache + ): + tasks = [make_request(user_data) for user_data in users] + results = await asyncio.gather(*tasks, return_exceptions=True) end_time = time.time() # Analyze results @@ -534,10 +581,14 @@ async def test_concurrent_pre_call_hooks_stress(): ), f"Premium success rate should be >= 90%, got {premium_success_rate:.2%}" assert ( standard_success_rate >= 0.5 - ), f"Standard success rate should be >= 50%, got {standard_success_rate:.2%}" - assert ( - premium_success_rate > standard_success_rate - ), "Premium should have higher success rate than standard" + ), f"Standard success rate should be >= 50% (with 30% random limiting, allows for variance), got {standard_success_rate:.2%}" + + # Allow for the case where both are 100% due to timing/mocking issues + # The test is inherently flaky due to random behavior + if premium_success_rate < 1.0 or standard_success_rate < 1.0: + assert ( + premium_success_rate >= standard_success_rate + ), "Premium should have >= success rate than standard" total_duration = end_time - start_time @@ -550,3 +601,703 @@ async def test_concurrent_pre_call_hooks_stress(): ) print(f" - Total successful: {successful_count}/50 ({successful_count/50:.1%})") print(f" - Priority system working: Premium > Standard success rates") + +# These tests make actual async_pre_call_hook calls to simulate real traffic + + +@pytest.mark.asyncio +async def test_fake_calls_case_1_no_rate_limiting_at_capacity(): + """ + Test Case 1: Saturation-Aware Rate Limiting at 50% Threshold + + System: 100 RPM capacity, saturation_threshold=50% + Key A: priority_reservation=0.75 (75 RPM reserved) + Key B: priority_reservation=0.25 (25 RPM reserved) + Traffic A: 1 request + Traffic B: 100 requests + + Expected behavior: + - Key A: 1 request succeeds (low traffic) + - Key B: ~25-26 requests succeed (capped at reservation when saturation >= 50%) + + Once saturation hits 50%, strict mode enforces priority-based limits. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + # Set up priority reservations + litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-1" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {"priority": "key_b"} + key_b_user.user_id = "key_b_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0} + + async def make_request(user, priority_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[priority_name] += 1 + return {"status": "success", "priority": priority_name} + else: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name} + + except Exception as e: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name, "error": str(e)} + + # Send 1 request from key_a, 100 from key_b + tasks = [] + + for i in range(1): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(100): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = successful_requests["key_a"] + successful_requests["key_b"] + total_rate_limited = rate_limited_requests["key_a"] + rate_limited_requests["key_b"] + + print(f"Test Case 1 - Saturation-Aware Rate Limiting:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A: {successful_requests['key_a']}/1 successful (reserved 75 RPM)") + print(f" - Key B: {successful_requests['key_b']}/100 successful (reserved 25 RPM)") + print(f" - Total successful: {total_successful}/101") + print(f" - Total rate limited: {total_rate_limited}/101") + + # Key A should get its 1 request + assert successful_requests["key_a"] == 1, f"Key A should get 1 request, got {successful_requests['key_a']}" + + # Key B can send until saturation hits 50% (which is ~50 total requests) + # After that, strict mode enforces its 25 RPM reservation + # Due to race conditions in concurrent execution, allow 45-52 successful requests + assert 45 <= successful_requests["key_b"] <= 52, f"Key B should get ~49 requests (45-52), got {successful_requests['key_b']}" + + # Verify approximately half of key_b requests were rate limited + assert rate_limited_requests["key_b"] >= 45, f"Key B should have ≥45 rate limited requests, got {rate_limited_requests['key_b']}" + + +@pytest.mark.asyncio +async def test_fake_calls_case_2_priority_queue_during_saturation(): + """ + Test Case 2: Priority Queue Behavior During Saturation + + System: 100 RPM capacity + Key A: priority_reservation=0.75 (75 RPM reserved) + Key B: priority_reservation=0.25 (25 RPM reserved) + Traffic A: 200 RPM + Traffic B: 200 RPM + Expected A: 75 RPM (75% of capacity) + Expected B: 25 RPM (25% of capacity) + + When total traffic exceeds capacity, rate limiting enforces priority reservations. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-2" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {"priority": "key_b"} + key_b_user.user_id = "key_b_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0} + + async def make_request(user, priority_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[priority_name] += 1 + return {"status": "success", "priority": priority_name} + else: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name} + + except Exception as e: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name, "error": str(e)} + + # Send 200 requests from each priority (over capacity) + tasks = [] + + for i in range(200): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(200): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = successful_requests["key_a"] + successful_requests["key_b"] + + key_a_success_rate = successful_requests["key_a"] / 200 + key_b_success_rate = successful_requests["key_b"] / 200 + + print(f"Test Case 2 - Priority Queue Behavior During Saturation:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A: {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})") + print(f" - Key B: {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})") + print(f" - Total successful: {total_successful}/400") + + # Key A should get significantly more requests than Key B (75:25 ratio) + assert key_a_success_rate > key_b_success_rate, ( + f"Key A should have higher success rate: {key_a_success_rate:.1%} vs {key_b_success_rate:.1%}" + ) + + # Check ratio is approximately 3:1 (75:25) + if total_successful > 0: + key_a_share = successful_requests["key_a"] / total_successful + expected_key_a_share = 0.75 + + print(f" - Key A got {key_a_share:.1%} of successful requests (expected ~75%)") + + # Allow tolerance for timing effects + assert abs(key_a_share - expected_key_a_share) < 0.2, ( + f"Key A share should be ~75%, got {key_a_share:.1%}" + ) + + +@pytest.mark.asyncio +async def test_fake_calls_case_3_spillover_capacity_default_keys(): + """ + Test Case 3: Spillover Capacity for Default Keys + + System: 100 RPM capacity + Key A: priority_reservation=0.75 (75 RPM reserved) + Key B: nothing set (default) + Key C: nothing set (default) + Key D: nothing set (default) + Traffic A: 150 RPM + Traffic B: 150 RPM + Traffic C: 150 RPM + Traffic D: 150 RPM + Expected A: 75 RPM (75% reserved) + Expected B: ~8.3 RPM (remaining 25 RPM / 3 default keys) + Expected C: ~8.3 RPM + Expected D: ~8.3 RPM + + Tests spillover behavior where default keys share remaining capacity. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"key_a": 0.75} + litellm.priority_reservation_settings.default_priority = 0.25 + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-3" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {} + key_b_user.user_id = "key_b_user" + + key_c_user = UserAPIKeyAuth() + key_c_user.metadata = {} + key_c_user.user_id = "key_c_user" + + key_d_user = UserAPIKeyAuth() + key_d_user.metadata = {} + key_d_user.user_id = "key_d_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0} + + async def make_request(user, key_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[key_name] += 1 + return {"status": "success", "key": key_name} + else: + rate_limited_requests[key_name] += 1 + return {"status": "rate_limited", "key": key_name} + + except Exception as e: + rate_limited_requests[key_name] += 1 + return {"status": "rate_limited", "key": key_name, "error": str(e)} + + # Send 150 requests from each key (600 total, 6x over capacity) + tasks = [] + + for i in range(150): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(150): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + for i in range(150): + tasks.append(make_request(key_c_user, "key_c", f"key_c_{i}")) + + for i in range(150): + tasks.append(make_request(key_d_user, "key_d", f"key_d_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = sum(successful_requests.values()) + + print(f"Test Case 3 - Spillover Capacity for Default Keys:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A: {successful_requests['key_a']}/150 successful") + print(f" - Key B: {successful_requests['key_b']}/150 successful (default)") + print(f" - Key C: {successful_requests['key_c']}/150 successful (default)") + print(f" - Key D: {successful_requests['key_d']}/150 successful (default)") + print(f" - Total successful: {total_successful}/600") + + # Key A should get the most requests (75% of capacity) + assert successful_requests["key_a"] > successful_requests["key_b"], "Key A should get more than Key B" + assert successful_requests["key_a"] > successful_requests["key_c"], "Key A should get more than Key C" + assert successful_requests["key_a"] > successful_requests["key_d"], "Key A should get more than Key D" + + # Default keys should get similar amounts (spillover capacity) + avg_default = (successful_requests["key_b"] + successful_requests["key_c"] + successful_requests["key_d"]) / 3 + print(f" - Average default key success: {avg_default:.1f}") + + +@pytest.mark.asyncio +async def test_fake_calls_case_4_over_allocated_with_normalization(): + """ + Test Case 4: Over-Allocated Priority reservations with Normalization + + System: 100 RPM capacity + Key A: priority_reservation=0.60 (60% requested) + Key B: priority_reservation=0.80 (80% requested) + Total: 140% (over-allocated, should normalize to 43%/57%) + Traffic A: 200 RPM + Traffic B: 200 RPM + + With saturation-aware rate limiting: + - Initially, requests are allowed through in generous mode (under 80% saturation) + - Once saturated, strict priority-based limits kick in with normalized weights + - Due to concurrent burst, total successful may exceed 100 RPM in the test window + - This test verifies normalization works and total capacity is reasonably bounded + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"key_a": 0.60, "key_b": 0.80} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-4" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {"priority": "key_b"} + key_b_user.user_id = "key_b_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0} + + async def make_request(user, priority_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[priority_name] += 1 + return {"status": "success", "priority": priority_name} + else: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name} + + except Exception as e: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name, "error": str(e)} + + # Send 200 requests from each key (400 total, 4x over capacity) + tasks = [] + + for i in range(200): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(200): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = successful_requests["key_a"] + successful_requests["key_b"] + + key_a_success_rate = successful_requests["key_a"] / 200 + key_b_success_rate = successful_requests["key_b"] / 200 + + print(f"Test Case 4 - Over-Allocated Priority Reservations with Normalization:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A (0.60): {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})") + print(f" - Key B (0.80): {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})") + print(f" - Total successful: {total_successful}/400") + + # With saturation-aware behavior: + # 1. Verify total capacity is reasonably bounded (not all 400 requests succeed) + assert total_successful < 300, ( + f"Total requests should be bounded by saturation detection, got {total_successful}/400" + ) + + # 2. Verify significant rate limiting occurred (at least 50% blocked) + assert total_successful < 200, ( + f"At least 50% of requests should be rate limited, got {total_successful}/400 successful" + ) + + # 3. Verify both keys got some requests through (normalization is working) + assert successful_requests["key_a"] > 0, "Key A should get some requests" + assert successful_requests["key_b"] > 0, "Key B should get some requests" + + print(f" - Normalization test PASSED: Both priorities got requests, " + f"total bounded to {total_successful} (under 200)") + + +@pytest.mark.asyncio +async def test_fake_calls_case_5_default_value_priority_reservation(): + """ + Test Case 5: Default value for priority reservation + + System: 100 RPM capacity + Key A: priority_reservation=0.50 (50 RPM) + Key B: priority_reservation=0.20 (20 RPM) + Key C: priority_reservation=0.05 (5 RPM) + Key D: nothing set (uses default_priority=0.05, 5 RPM) + Traffic A: 150 RPM + Traffic B: 150 RPM + Traffic C: 150 RPM + Traffic D: 150 RPM + Expected A: 55 RPM (normalized) + Expected B: 25 RPM (normalized) + Expected C: 10 RPM (normalized) + Expected D: 10 RPM (normalized) + + Tests complex scenario with explicit priorities and default priority. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"key_a": 0.50, "key_b": 0.20, "key_c": 0.05} + litellm.priority_reservation_settings.default_priority = 0.05 + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-5" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {"priority": "key_b"} + key_b_user.user_id = "key_b_user" + + key_c_user = UserAPIKeyAuth() + key_c_user.metadata = {"priority": "key_c"} + key_c_user.user_id = "key_c_user" + + key_d_user = UserAPIKeyAuth() + key_d_user.metadata = {} + key_d_user.user_id = "key_d_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0} + + async def make_request(user, key_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[key_name] += 1 + return {"status": "success", "key": key_name} + else: + rate_limited_requests[key_name] += 1 + return {"status": "rate_limited", "key": key_name} + + except Exception as e: + rate_limited_requests[key_name] += 1 + return {"status": "rate_limited", "key": key_name, "error": str(e)} + + # Send 150 requests from each key (600 total, 6x over capacity) + tasks = [] + + for i in range(150): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(150): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + for i in range(150): + tasks.append(make_request(key_c_user, "key_c", f"key_c_{i}")) + + for i in range(150): + tasks.append(make_request(key_d_user, "key_d", f"key_d_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = sum(successful_requests.values()) + + print(f"Test Case 5 - Default value for priority reservation:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A (0.50): {successful_requests['key_a']}/150 successful") + print(f" - Key B (0.20): {successful_requests['key_b']}/150 successful") + print(f" - Key C (0.05): {successful_requests['key_c']}/150 successful") + print(f" - Key D (default 0.05): {successful_requests['key_d']}/150 successful") + print(f" - Total successful: {total_successful}/600") + + # Verify priority ordering: A > B > C ≈ D + assert successful_requests["key_a"] > successful_requests["key_b"], "Key A should get more than Key B" + assert successful_requests["key_b"] > successful_requests["key_c"], "Key B should get more than Key C" + + # Key C and Key D should get similar amounts (both have 0.05 priority) + key_c_vs_d_ratio = successful_requests["key_c"] / max(successful_requests["key_d"], 1) + print(f" - Key C vs Key D ratio: {key_c_vs_d_ratio:.2f} (expected ~1.0)") + + if total_successful > 0: + key_a_share = successful_requests["key_a"] / total_successful + print(f" - Key A got {key_a_share:.1%} of successful requests (expected ~55-62%)") + + +@pytest.mark.asyncio +async def test_default_priority_shared_pool(): + """ + Test that keys without explicit priority share ONE default pool, not get individual allocations. + + With default_priority=0.25: + - Key A, B, C (no priority) should share ONE 25 RPM pool + - NOT get 25 RPM each (which would be 75 RPM total) + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"prod": 0.75} + litellm.priority_reservation_settings.default_priority = 0.25 + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "test-default-pool" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create 3 users without explicit priority + user_a = UserAPIKeyAuth() + user_a.metadata = {} + user_a.user_id = "user_a" + + user_b = UserAPIKeyAuth() + user_b.metadata = {} + user_b.user_id = "user_b" + + user_c = UserAPIKeyAuth() + user_c.metadata = {} + user_c.user_id = "user_c" + + # Get descriptors for each + desc_a = handler._create_priority_based_descriptors( + model=model, user_api_key_dict=user_a, priority=None + ) + desc_b = handler._create_priority_based_descriptors( + model=model, user_api_key_dict=user_b, priority=None + ) + desc_c = handler._create_priority_based_descriptors( + model=model, user_api_key_dict=user_c, priority=None + ) + + # All should use the SAME shared pool key + assert desc_a[0]["value"] == f"{model}:default_pool" + assert desc_b[0]["value"] == f"{model}:default_pool" + assert desc_c[0]["value"] == f"{model}:default_pool" + + # All should have same limit (25 RPM SHARED, not 25 RPM each) + assert desc_a[0]["rate_limit"]["requests_per_unit"] == 25 + assert desc_b[0]["rate_limit"]["requests_per_unit"] == 25 + assert desc_c[0]["rate_limit"]["requests_per_unit"] == 25 + + # Verify explicit priority uses different pool + user_prod = UserAPIKeyAuth() + user_prod.metadata = {"priority": "prod"} + desc_prod = handler._create_priority_based_descriptors( + model=model, user_api_key_dict=user_prod, priority="prod" + ) + + assert desc_prod[0]["value"] == f"{model}:prod" + assert desc_prod[0]["rate_limit"]["requests_per_unit"] == 75 + assert desc_prod[0]["value"] != desc_a[0]["value"] # Different pools + + print("✅ Default priority test passed:") + print(f" - 3 keys without priority share ONE pool: {desc_a[0]['value']}") + print(f" - Shared pool limit: {desc_a[0]['rate_limit']['requests_per_unit']} RPM") + print(f" - Explicit priority 'prod' uses separate pool: {desc_prod[0]['value']}") diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 710e4265013..521faae3ca5 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -89,7 +89,8 @@ def test_invalid_purpose(mocker: MockerFixture, monkeypatch, llm_router: Router) files={"file": test_file}, data={ "purpose": "my-bad-purpose", - "target_model_names": ["azure-gpt-3-5-turbo", "gpt-3.5-turbo"], + # "target_model_names": ["azure-gpt-3-5-turbo", "gpt-3.5-turbo"], + "target_model_names": "gpt-3-5-turbo", }, headers={"Authorization": "Bearer test-key"}, ) @@ -134,14 +135,14 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: custom_llm_provider="azure", model="azure/chatgpt-v-2", api_key="azure_api_key", - file=file_data, + file=file_data[1], purpose=purpose_data, ) await litellm.files.main.create_file( custom_llm_provider="openai", model="openai/gpt-3.5-turbo", api_key="openai_api_key", - file=file_data, + file=file_data[1], purpose=purpose_data, ) # Return a dummy response object as needed by the test diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py new file mode 100644 index 00000000000..3854e3944ce --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py @@ -0,0 +1,374 @@ +import json +import os +import sys +from datetime import datetime +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.gemini_passthrough_logging_handler import ( + GeminiPassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, +) + + +class TestGeminiPassthroughLoggingHandler: + """Test the Gemini 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 = GeminiPassthroughLoggingHandler() + + # Mock Gemini generateContent response + self.mock_gemini_response = { + "candidates": [ + { + "content": {"parts": [{"text": "Hello! How can I help you today?"}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + {"category": "HARM_CATEGORY_HARASSMENT", "probability": "NEGLIGIBLE"}, + {"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}, + {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "probability": "NEGLIGIBLE"}, + {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "probability": "NEGLIGIBLE"}, + ], + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 8, "totalTokenCount": 18}, + } + + def _create_mock_httpx_response(self) -> httpx.Response: + """Create a mock httpx.Response for testing""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.text = json.dumps(self.mock_gemini_response) + mock_response.json.return_value = self.mock_gemini_response + mock_response.headers = {"content-type": "application/json"} + return mock_response + + def _create_mock_logging_obj(self) -> LiteLLMLoggingObj: + """Create a mock logging object for testing""" + mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {} + mock_logging_obj.optional_params = {} + mock_logging_obj.litellm_call_id = "test-call-id-123" + return mock_logging_obj + + def _create_passthrough_logging_payload(self) -> PassthroughStandardLoggingPayload: + """Create a mock passthrough logging payload for testing""" + return PassthroughStandardLoggingPayload( + url="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + request_body={"contents": [{"parts": [{"text": "Hello"}]}]}, + request_method="POST", + ) + + def test_is_gemini_route(self): + """Test that Gemini routes are correctly identified""" + from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging + + handler = PassThroughEndpointLogging() + + # Test generateContent endpoint + assert ( + handler.is_gemini_route( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + custom_llm_provider="gemini", + ) + is True + ) + + # Test streamGenerateContent endpoint + assert ( + handler.is_gemini_route( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent", + custom_llm_provider="gemini", + ) + is True + ) + + # Test non-Gemini endpoint + assert ( + handler.is_gemini_route("https://api.openai.com/v1/chat/completions", custom_llm_provider="openai") is False + ) + + def test_extract_model_from_url(self): + """Test that model is correctly extracted from Gemini URLs""" + # Test generateContent endpoint + model = GeminiPassthroughLoggingHandler.extract_model_from_url( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent" + ) + assert model == "gemini-1.5-flash" + + # Test streamGenerateContent endpoint + model = GeminiPassthroughLoggingHandler.extract_model_from_url( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:streamGenerateContent" + ) + assert model == "gemini-1.5-pro" + + @patch("litellm.completion_cost") + @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload") + def test_gemini_passthrough_handler_success(self, mock_get_standard_logging, mock_completion_cost): + """Test successful cost tracking for Gemini generateContent endpoint""" + # 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": "gemini-1.5-flash", + } + + # Act + result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=self.mock_gemini_response, + logging_obj=mock_logging_obj, + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"contents": [{"parts": [{"text": "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"] == "gemini-1.5-flash" + assert result["kwargs"]["custom_llm_provider"] == "gemini" + + # 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"] == "gemini-1.5-flash" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini" + + @patch("litellm.completion_cost") + def test_gemini_passthrough_handler_streaming(self, mock_completion_cost): + """Test cost tracking for Gemini streaming endpoint""" + # Arrange + mock_completion_cost.return_value = 0.000030 + + # Mock streaming response chunks + mock_chunks = [ + {"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}, + {"candidates": [{"content": {"parts": [{"text": " there!"}]}}]}, + ] + + 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": "gemini-1.5-flash", + } + + # Act - Use generateContent URL since that's what the handler processes + result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=mock_chunks, + logging_obj=mock_logging_obj, + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"contents": [{"parts": [{"text": "Hello"}]}]}, + **kwargs, + ) + + # Assert + assert result is not None + assert "result" in result + assert "kwargs" in result + assert result["kwargs"]["response_cost"] == 0.000030 + assert result["kwargs"]["model"] == "gemini-1.5-flash" + assert result["kwargs"]["custom_llm_provider"] == "gemini" + + # Verify cost calculation was called + mock_completion_cost.assert_called_once() + + def test_gemini_passthrough_handler_non_gemini_route(self): + """Test that non-Gemini routes return None""" + 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 = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=self.mock_gemini_response, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/chat/completions", # Non-Gemini route (no generateContent) + 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 - the handler should return a dict with None result for non-Gemini routes + assert result is not None + assert result["result"] is None + assert "kwargs" in result + + @pytest.mark.asyncio + async def test_pass_through_success_handler_gemini_routing(self): + """Test that the success handler correctly routes Gemini requests to the Gemini handler""" + handler = PassThroughEndpointLogging() + + # Mock the logging object + mock_logging_obj = self._create_mock_logging_obj() + + # Mock the _handle_logging method to capture the call + handler._handle_logging = AsyncMock() + + # Mock httpx response + mock_response = self._create_mock_httpx_response() + + # Create passthrough logging payload + passthrough_logging_payload = self._create_passthrough_logging_payload() + + # Call the success handler with Gemini route and provider + result = await handler.pass_through_async_success_handler( + httpx_response=mock_response, + response_body=self.mock_gemini_response, + logging_obj=mock_logging_obj, + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"contents": [{"parts": [{"text": "Hello"}]}]}, + passthrough_logging_payload=passthrough_logging_payload, + custom_llm_provider="gemini", + ) + + # Assert - The success handler returns None on success (following the pattern from other tests) + assert result is None + + # Verify that the logging object has the cost set (from Gemini handler) + assert mock_logging_obj.model_call_details["response_cost"] is not None + assert mock_logging_obj.model_call_details["model"] == "gemini-1.5-flash" + assert mock_logging_obj.model_call_details["custom_llm_provider"] == "gemini" + + # Verify that _handle_logging was called with the correct kwargs + handler._handle_logging.assert_called_once() + call_kwargs = handler._handle_logging.call_args[1] + assert call_kwargs["response_cost"] is not None + assert call_kwargs["model"] == "gemini-1.5-flash" + assert call_kwargs["custom_llm_provider"] == "gemini" + + @patch("litellm.completion_cost") + @patch("litellm.stream_chunk_builder") + def test_gemini_streaming_cost_calculation(self, mock_stream_chunk_builder, mock_completion_cost): + """Test that Gemini streaming passthrough correctly calculates cost with logging_obj""" + # Arrange + mock_completion_cost.return_value = 0.000025 + mock_logging_obj = self._create_mock_logging_obj() + + # Mock the stream_chunk_builder to return a response with usage + from litellm.utils import ModelResponse, Usage + mock_response = ModelResponse() + mock_usage = Usage(prompt_tokens=5, completion_tokens=10, total_tokens=15) + mock_response.usage = mock_usage + mock_stream_chunk_builder.return_value = mock_response + + # Mock fragmented JSON chunks (as they come from the streaming response) + fragmented_chunks = [ + '[{"candidates": [', + '{"content": {"parts": [{"text": "Hello"}], "role": "model"},', + '"finishReason": "STOP", "index": 0}],', + '"usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 10, "totalTokenCount": 15}', + '}]' + ] + + # Act + result = GeminiPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=fragmented_chunks, + litellm_logging_obj=mock_logging_obj, + model="gemini-1.5-flash", + url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent" + ) + + # Assert + assert result is not None + assert result == mock_response + + # Verify stream_chunk_builder was called with logging_obj for cost injection + mock_stream_chunk_builder.assert_called_once() + call_args = mock_stream_chunk_builder.call_args + + # Check that logging_obj was passed for cost injection + assert "logging_obj" in call_args.kwargs + assert call_args.kwargs["logging_obj"] == mock_logging_obj + + # Verify the chunks were properly reconstructed from fragmented JSON + # The first argument should be the chunks list + chunks_arg = call_args[0][0] if call_args[0] else call_args.kwargs.get("chunks", []) + assert len(chunks_arg) > 0 # Should have parsed chunks + + def test_gemini_streaming_json_parsing(self): + """Test that fragmented JSON chunks are correctly joined and parsed""" + # Arrange + # Mock fragmented JSON chunks that simulate how Gemini streaming response gets split + fragmented_chunks = [ + '[{"candidates": [', + '{"content": {"parts": [{"text": "Test response"}], "role": "model"},', + '"finishReason": "STOP", "index": 0}],', + '"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 7, "totalTokenCount": 10}', + '}]' + ] + + # Mock the gemini iterator's chunk_parser method + mock_iterator = MagicMock() + mock_iterator.chunk_parser.return_value = {"candidates": [{"content": {"parts": [{"text": "Test response"}]}}]} + + # Act + result = GeminiPassthroughLoggingHandler._parse_gemini_streaming_json( + all_chunks=fragmented_chunks, + gemini_iterator=mock_iterator + ) + + # Assert + # The method should successfully join the fragmented JSON and return parsed chunks + assert isinstance(result, list) + assert len(result) == 1 # Should have one parsed chunk + + # Verify that the combined JSON is valid + combined_json = "".join(fragmented_chunks) + parsed_json = json.loads(combined_json) + assert isinstance(parsed_json, list) + assert len(parsed_json) == 1 + assert "candidates" in parsed_json[0] + assert "usageMetadata" in parsed_json[0] + + # Verify the iterator's chunk_parser was called + mock_iterator.chunk_parser.assert_called_once() 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 1157863aa27..1a26f8a39aa 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 @@ -1252,6 +1252,90 @@ async def test_delete_pass_through_endpoint_empty_list(): +@pytest.mark.asyncio +async def test_pass_through_request_query_params_forwarding(): + """ + Test that query parameters from the original request are properly forwarded to the target URL. + + This test verifies the fix for the bug where query parameters like api-version were being lost + when forwarding requests to Azure OpenAI and other pass-through endpoints. + """ + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.HttpPassThroughEndpointHelpers.non_streaming_http_request_handler" + ) as mock_http_handler: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler" + ) as mock_success_handler: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_response_body" + ) as mock_get_response_body: + # Setup mock for pre_call_hook + test_body = {"name": "Azure Assistant", "model": "gpt-4o"} + mock_proxy_logging.pre_call_hook = AsyncMock(return_value=test_body) + + # Setup mock for http response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.aread = AsyncMock(return_value=b'{"id": "asst_123", "object": "assistant"}') + mock_response.text = '{"id": "asst_123", "object": "assistant"}' + mock_response.raise_for_status = MagicMock() + + # Mock the HTTP request handler to capture the call + mock_http_handler.return_value = mock_response + + # Mock response body parser + mock_get_response_body.return_value = {"id": "asst_123", "object": "assistant"} + + # Mock headers for custom headers + mock_processing.get_custom_headers.return_value = {} + + # Mock success handler + mock_success_handler.return_value = None + + # Create mock request with query parameters (Azure API version) + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://localhost:4000/azure-assistant/openai/assistants" + mock_request.body = AsyncMock(return_value=json.dumps(test_body).encode()) + mock_request.headers = Headers({"Content-Type": "application/json"}) + + # Create QueryParams with api-version parameter + mock_request.query_params = QueryParams([("api-version", "2025-01-01-preview")]) + + # Create mock user API key dict + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.api_key = "sk-1234" + + # Call pass_through_request + result = await pass_through_request( + request=mock_request, + target="https://krris-m2f9a9i7-eastus2.openai.azure.com/openai/assistants", + custom_headers={"Authorization": "Bearer azure_token"}, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify the HTTP handler was called + mock_http_handler.assert_called_once() + + # Extract the call arguments to verify query parameters were passed + call_kwargs = mock_http_handler.call_args[1] + + # The key assertion: query parameters should be preserved and passed to the HTTP handler + assert "requested_query_params" in call_kwargs + assert call_kwargs["requested_query_params"] == {"api-version": "2025-01-01-preview"} + + # Verify the target URL is correct + assert str(call_kwargs["url"]) == "https://krris-m2f9a9i7-eastus2.openai.azure.com/openai/assistants" + + # Verify the request body is preserved + assert call_kwargs["_parsed_body"] == test_body + + @pytest.mark.asyncio async def test_pass_through_with_httpbin_redirect(): """ diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4235e5d3adb..90d958e711d 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -314,6 +314,51 @@ class TestProxyInitializationHelpers: call_args = mock_uvicorn_run.call_args assert call_args[1]["timeout_keep_alive"] == 30 + @patch("uvicorn.run") + @patch("builtins.print") + def test_max_requests_before_restart_flag(self, mock_print, mock_uvicorn_run): + """Test that the max_requests_before_restart flag is passed to uvicorn as limit_max_requests""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + + mock_app = MagicMock() + mock_proxy_config = MagicMock() + mock_key_mgmt = MagicMock() + mock_save_worker_config = MagicMock() + + with patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=mock_app, + ProxyConfig=mock_proxy_config, + KeyManagementSettings=mock_key_mgmt, + save_worker_config=mock_save_worker_config, + ) + }, + ), patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args: + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, ["--local", "--max_requests_before_restart", "123"] + ) + + assert result.exit_code == 0 + mock_uvicorn_run.assert_called_once() + + # Check that uvicorn.run was called with limit_max_requests parameter + call_args = mock_uvicorn_run.call_args + assert call_args[1]["limit_max_requests"] == 123 + @patch.dict(os.environ, {}, clear=True) def test_construct_database_url_from_env_vars(self): """Test the construct_database_url_from_env_vars function with various scenarios""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d2b516a55d2..63436899fcc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1886,3 +1886,147 @@ async def test_add_router_settings_shallow_merge_behavior(): assert merged_settings["nested_setting"] == expected_nested assert merged_settings["top_level"] == "db_top" + + +@pytest.mark.asyncio +async def test_model_info_v1_oci_secrets_not_leaked(): + """ + Test that model_info_v1 endpoint properly masks OCI sensitive parameters and does not leak secrets. + """ + from unittest.mock import MagicMock, patch + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import model_info_v1 + + # Mock user authentication + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.api_key = "test-key" + mock_user_api_key_dict.team_models = [] + mock_user_api_key_dict.models = ["oci-grok-test"] + + # Mock model data with OCI sensitive information + mock_model_data = { + "model_name": "oci-grok-test", + "litellm_params": { + "model": "oci/xai.grok-4", + "oci_key": "ocid1.api_key.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk", + "oci_region": "us-phoenix-1", + "oci_user": "ocid1.user.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk", + "oci_fingerprint": "aa:bb:cc:dd:ee:ff:11:22:33:44:55:66:77:88:99:00", + "oci_tenancy": "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk", + "oci_key_file": "/path/to/oci_api_key.pem", + "oci_compartment_id": "ocid1.compartment.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk", + "drop_params": True + }, + "model_info": { + "mode": "completion", + "id": "test-model-id" + } + } + + # Mock the llm_router to return our test data + mock_router = MagicMock() + mock_router.get_model_names.return_value = ["oci-grok-test"] + mock_router.get_model_access_groups.return_value = {} + mock_router.get_model_list.return_value = [mock_model_data] + + # Mock global variables + with patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.proxy_server.llm_model_list", [mock_model_data]), \ + patch("litellm.proxy.proxy_server.general_settings", {"infer_model_from_keys": False}), \ + patch("litellm.proxy.proxy_server.user_model", None): + + # Call the model_info_v1 endpoint + result = await model_info_v1( + user_api_key_dict=mock_user_api_key_dict, + litellm_model_id=None + ) + + # Verify the result structure + assert "data" in result + assert len(result["data"]) == 1 + + model_info = result["data"][0] + litellm_params = model_info["litellm_params"] + + # Verify that sensitive OCI fields are masked + assert "****" in litellm_params["oci_key"], "oci_key should be masked" + assert "****" in litellm_params["oci_fingerprint"], "oci_fingerprint should be masked" + assert "****" in litellm_params["oci_tenancy"], "oci_tenancy should be masked" + assert "****" in litellm_params["oci_key_file"], "oci_key_file should be masked" + + # Verify that non-sensitive fields are NOT masked + assert litellm_params["model"] == "oci/xai.grok-4", "model field should not be masked" + assert litellm_params["oci_region"] == "us-phoenix-1", "oci_region should not be masked" + assert litellm_params["drop_params"] is True, "drop_params should not be masked" + + # Verify the model field specifically is not masked (this was the original issue) + assert "****" not in litellm_params["model"], "model field should never be masked" + assert litellm_params["model"].startswith("oci/"), "model should retain its full value" + + # Verify that actual secret values are not present in the response + result_str = str(result) + assert "ocid1.api_key.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str + assert "aa:bb:cc:dd:ee:ff:11:22:33:44:55:66:77:88:99:00" not in result_str + assert "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str + assert "/path/to/oci_api_key.pem" not in result_str + + +def test_add_callback_from_db_to_in_memory_litellm_callbacks(): + """ + Test that _add_callback_from_db_to_in_memory_litellm_callbacks correctly adds callbacks + for success, failure, and combined event types. + """ + from unittest.mock import MagicMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Mock the callback manager + mock_callback_manager = MagicMock() + + with patch("litellm.proxy.proxy_server.litellm") as mock_litellm: + # Set up mock litellm attributes + mock_litellm._known_custom_logger_compatible_callbacks = [] + mock_litellm.logging_callback_manager = mock_callback_manager + + # Test Case 1: Add success callback + mock_success_callbacks = [] + proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( + callback="prometheus", + event_types=["success"], + existing_callbacks=mock_success_callbacks, + ) + mock_callback_manager.add_litellm_success_callback.assert_called_once_with("prometheus") + mock_callback_manager.reset_mock() + + # Test Case 2: Add failure callback + mock_failure_callbacks = [] + proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( + callback="langfuse", + event_types=["failure"], + existing_callbacks=mock_failure_callbacks, + ) + mock_callback_manager.add_litellm_failure_callback.assert_called_once_with("langfuse") + mock_callback_manager.reset_mock() + + # Test Case 3: Add callback for both success and failure + mock_callbacks = [] + proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( + callback="s3", + event_types=["success", "failure"], + existing_callbacks=mock_callbacks, + ) + mock_callback_manager.add_litellm_callback.assert_called_once_with("s3") + mock_callback_manager.reset_mock() + + # Test Case 4: Don't add callback if it already exists + existing_callbacks_with_item = ["prometheus"] + proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( + callback="prometheus", + event_types=["success"], + existing_callbacks=existing_callbacks_with_item, + ) + mock_callback_manager.add_litellm_success_callback.assert_not_called() diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index ef733eaa887..4b571691e48 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -526,3 +526,33 @@ class TestProxySettingEndpoints: # Verify save_config was called twice (once for each update) assert mock_proxy_config["save_call_count"]() == 2 + + def test_get_ui_theme_settings(self, mock_proxy_config): + """Test getting UI theme settings without authentication""" + response = client.get("/get/ui_theme_settings") + + assert response.status_code == 200 + data = response.json() + + assert "values" in data + assert "field_schema" in data + + def test_update_ui_theme_settings(self, mock_proxy_config, mock_auth, monkeypatch): + """Test updating UI theme settings""" + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + new_theme = {"logo_url": "https://example.com/new-logo.png"} + + response = client.patch("/update/ui_theme_settings", json=new_theme) + + assert response.status_code == 200 + data = response.json() + + assert data["status"] == "success" + assert data["theme_config"]["logo_url"] == "https://example.com/new-logo.png" + + # Verify config was updated + updated_config = mock_proxy_config["config"] + assert "UI_LOGO_PATH" in updated_config["environment_variables"] + assert mock_proxy_config["save_call_count"]() == 1 diff --git a/tests/test_litellm/test_router_google_genai.py b/tests/test_litellm/test_router_google_genai.py new file mode 100644 index 00000000000..8b8d8a8379d --- /dev/null +++ b/tests/test_litellm/test_router_google_genai.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +""" +Test to verify the new Google GenAI router methods +""" +import asyncio +import os +import sys +from unittest.mock import AsyncMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.types.utils import ModelResponse + + +@pytest.mark.asyncio +async def test_router_agenerate_content_method(): + """Test that the new agenerate_content method in Router works correctly""" + # Create a router instance + router = litellm.Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-3.5-turbo", + } + } + ] + ) + + # Create a mock response in Google GenAI format + mock_response = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Hello, world!" + } + ] + } + } + ] + } + + # Mock the router's underlying agenerate_content method to return a mock response + with patch.object(router, 'agenerate_content', new=AsyncMock(return_value=mock_response)) as mock_agenerate_content: + # Call the agenerate_content method + response = await router.agenerate_content( + model="test-model", + contents=[{"role": "user", "parts": [{"text": "Hello"}]}] + ) + + # Verify that router.agenerate_content was called with correct parameters + mock_agenerate_content.assert_called_once() + call_args = mock_agenerate_content.call_args + assert call_args[1]["model"] == "test-model" + assert call_args[1]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + + # Verify that the response is the mock response we created + assert response == mock_response + + +@pytest.mark.asyncio +async def test_router_aadapter_generate_content_method(): + """Test that the new aadapter_generate_content method in Router works correctly""" + # Create a router instance + router = litellm.Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-3.5-turbo", + } + } + ] + ) + + # Create a mock response in Google GenAI format + mock_response = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Hello, world!" + } + ] + } + } + ] + } + + # Mock the router's underlying aadapter_generate_content method to return a mock response + with patch.object(router, 'aadapter_generate_content', new=AsyncMock(return_value=mock_response)) as mock_aadapter_generate_content: + # Call the aadapter_generate_content method + response = await router.aadapter_generate_content( + model="test-model", + contents=[{"role": "user", "parts": [{"text": "Hello"}]}] + ) + + # Verify that router.aadapter_generate_content was called with correct parameters + mock_aadapter_generate_content.assert_called_once() + call_args = mock_aadapter_generate_content.call_args + assert call_args[1]["model"] == "test-model" + assert call_args[1]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + + # Verify that the response is the mock response we created + assert response == mock_response \ No newline at end of file diff --git a/tests/test_passthrough_endpoints.py b/tests/test_passthrough_endpoints.py index a66c94c5836..47ac7511aa1 100644 --- a/tests/test_passthrough_endpoints.py +++ b/tests/test_passthrough_endpoints.py @@ -17,7 +17,7 @@ dotenv.load_dotenv() async def cohere_rerank(session): url = "http://localhost:4000/v1/rerank" headers = { - "Authorization": f"bearer {os.getenv('COHERE_API_KEY')}", + "Authorization": f"Bearer {os.getenv('COHERE_API_KEY')}", "Content-Type": "application/json", "Accept": "application/json", } diff --git a/tests/unified_google_tests/base_google_test.py b/tests/unified_google_tests/base_google_test.py index 28c70b1cea7..8bd80f6f64c 100644 --- a/tests/unified_google_tests/base_google_test.py +++ b/tests/unified_google_tests/base_google_test.py @@ -17,7 +17,8 @@ from litellm.google_genai import ( generate_content_stream, agenerate_content_stream, ) -from google.genai.types import ContentDict, PartDict, GenerateContentResponse +from google.genai.types import ContentDict, PartDict +from litellm.types.google_genai.main import GenerateContentResponse from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import StandardLoggingPayload @@ -107,11 +108,11 @@ class BaseGoogleGenAITest: def _validate_non_streaming_response(self, response: Any): """Validate non-streaming response structure""" - # Handle type checking - response should be a dict for non-streaming + # Handle type checking - response should be a GenerateContentResponse for non-streaming if isinstance(response, AsyncIterator): pytest.fail("Expected non-streaming response but got AsyncIterator") - assert isinstance(response, GenerateContentResponse), f"Expected dict response, got {type(response)}" + assert isinstance(response, GenerateContentResponse), f"Expected GenerateContentResponse, got {type(response)}" print(f"Response: {response.model_dump_json(indent=4)}") # Basic validation - adjust based on actual Google GenAI response structure diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 2ff7e47a981..9ab39c57bfd 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -50,6 +50,7 @@ "@types/react-dom": "^18", "@types/react-syntax-highlighter": "^15.5.11", "@types/uuid": "^10.0.0", + "@vitejs/plugin-react": "^5.0.4", "@vitest/coverage-v8": "^3.2.4", "@vitest/ui": "^3.2.4", "autoprefixer": "^10.4.17", @@ -95,6 +96,7 @@ "version": "2.3.0", "resolved": "https://registry.npmjs.org/@ampproject/remapping/-/remapping-2.3.0.tgz", "integrity": "sha512-30iZtAPgz+LTIYoeivqYo853f02jBYSd5uGnGpkFV0M3xOt9aN73erkgYAmZU43x4VfqcnLxW9Kpg3R5LC4YYw==", + "dev": true, "dependencies": { "@jridgewell/gen-mapping": "^0.3.5", "@jridgewell/trace-mapping": "^0.3.24" @@ -284,20 +286,21 @@ } }, "node_modules/@babel/core": { - "version": "7.28.0", - "resolved": "https://registry.npmjs.org/@babel/core/-/core-7.28.0.tgz", - "integrity": "sha512-UlLAnTPrFdNGoFtbSXwcGFQBtQZJCNjaN6hQNP3UPvuNXT1i82N26KL3dZeIpNalWywr9IuQuncaAfUaS1g6sQ==", + "version": "7.28.4", + "resolved": "https://registry.npmjs.org/@babel/core/-/core-7.28.4.tgz", + "integrity": "sha512-2BCOP7TN8M+gVDj7/ht3hsaO/B/n5oDbiAyyvnRlNOs+u1o+JWNYTQrmpuNp1/Wq2gcFrI01JAW+paEKDMx/CA==", + "license": "MIT", "dependencies": { - "@ampproject/remapping": "^2.2.0", "@babel/code-frame": "^7.27.1", - "@babel/generator": "^7.28.0", + "@babel/generator": "^7.28.3", "@babel/helper-compilation-targets": "^7.27.2", - "@babel/helper-module-transforms": "^7.27.3", - "@babel/helpers": "^7.27.6", - "@babel/parser": "^7.28.0", + "@babel/helper-module-transforms": "^7.28.3", + "@babel/helpers": "^7.28.4", + "@babel/parser": "^7.28.4", "@babel/template": "^7.27.2", - "@babel/traverse": "^7.28.0", - "@babel/types": "^7.28.0", + "@babel/traverse": "^7.28.4", + "@babel/types": "^7.28.4", + "@jridgewell/remapping": "^2.3.5", "convert-source-map": "^2.0.0", "debug": "^4.1.0", "gensync": "^1.0.0-beta.2", @@ -332,12 +335,13 @@ } }, "node_modules/@babel/generator": { - "version": "7.28.0", - "resolved": "https://registry.npmjs.org/@babel/generator/-/generator-7.28.0.tgz", - "integrity": "sha512-lJjzvrbEeWrhB4P3QBsH7tey117PjLZnDbLiQEKjQ/fNJTjuq4HSqgFA+UNSwZT8D7dxxbnuSBMsa1lrWzKlQg==", + "version": "7.28.3", + "resolved": "https://registry.npmjs.org/@babel/generator/-/generator-7.28.3.tgz", + "integrity": "sha512-3lSpxGgvnmZznmBkCRnVREPUFJv2wrv9iAoFDvADJc0ypmdOxdUtcLeBgBJ6zE0PMeTKnxeQzyk0xTBq4Ep7zw==", + "license": "MIT", "dependencies": { - "@babel/parser": "^7.28.0", - "@babel/types": "^7.28.0", + "@babel/parser": "^7.28.3", + "@babel/types": "^7.28.2", "@jridgewell/gen-mapping": "^0.3.12", "@jridgewell/trace-mapping": "^0.3.28", "jsesc": "^3.0.2" @@ -493,13 +497,14 @@ } }, "node_modules/@babel/helper-module-transforms": { - "version": "7.27.3", - "resolved": "https://registry.npmjs.org/@babel/helper-module-transforms/-/helper-module-transforms-7.27.3.tgz", - "integrity": "sha512-dSOvYwvyLsWBeIRyOeHXp5vPj5l1I011r52FM1+r1jCERv+aFXYk4whgQccYEGYxK2H3ZAIA8nuPkQ0HaUo3qg==", + "version": "7.28.3", + "resolved": "https://registry.npmjs.org/@babel/helper-module-transforms/-/helper-module-transforms-7.28.3.tgz", + "integrity": "sha512-gytXUbs8k2sXS9PnQptz5o0QnpLL51SwASIORY6XaBKF88nsOT0Zw9szLqlSGQDP/4TljBAD5y98p2U1fqkdsw==", + "license": "MIT", "dependencies": { "@babel/helper-module-imports": "^7.27.1", "@babel/helper-validator-identifier": "^7.27.1", - "@babel/traverse": "^7.27.3" + "@babel/traverse": "^7.28.3" }, "engines": { "node": ">=6.9.0" @@ -609,23 +614,25 @@ } }, "node_modules/@babel/helpers": { - "version": "7.28.2", - "resolved": "https://registry.npmjs.org/@babel/helpers/-/helpers-7.28.2.tgz", - "integrity": "sha512-/V9771t+EgXz62aCcyofnQhGM8DQACbRhvzKFsXKC9QM+5MadF8ZmIm0crDMaz3+o0h0zXfJnd4EhbYbxsrcFw==", + "version": "7.28.4", + "resolved": "https://registry.npmjs.org/@babel/helpers/-/helpers-7.28.4.tgz", + "integrity": "sha512-HFN59MmQXGHVyYadKLVumYsA9dBFun/ldYxipEjzA4196jpLZd8UjEEBLkbEkvfYreDqJhZxYAWFPtrfhNpj4w==", + "license": "MIT", "dependencies": { "@babel/template": "^7.27.2", - "@babel/types": "^7.28.2" + "@babel/types": "^7.28.4" }, "engines": { "node": ">=6.9.0" } }, "node_modules/@babel/parser": { - "version": "7.28.0", - "resolved": "https://registry.npmjs.org/@babel/parser/-/parser-7.28.0.tgz", - "integrity": "sha512-jVZGvOxOuNSsuQuLRTh13nU0AogFlw32w/MT+LV6D3sP5WdbW61E77RnkbaO2dUvmPAYrBDJXGn5gGS6tH4j8g==", + "version": "7.28.4", + "resolved": "https://registry.npmjs.org/@babel/parser/-/parser-7.28.4.tgz", + "integrity": "sha512-yZbBqeM6TkpP9du/I2pUZnJsRMGGvOuIrhjzC1AwHwW+6he4mni6Bp/m8ijn0iOuZuPI2BfkCoSRunpyjnrQKg==", + "license": "MIT", "dependencies": { - "@babel/types": "^7.28.0" + "@babel/types": "^7.28.4" }, "bin": { "parser": "bin/babel-parser.js" @@ -1428,6 +1435,38 @@ "@babel/core": "^7.0.0-0" } }, + "node_modules/@babel/plugin-transform-react-jsx-self": { + "version": "7.27.1", + "resolved": "https://registry.npmjs.org/@babel/plugin-transform-react-jsx-self/-/plugin-transform-react-jsx-self-7.27.1.tgz", + "integrity": "sha512-6UzkCs+ejGdZ5mFFC/OCUrv028ab2fp1znZmCZjAOBKiBK2jXD1O+BPSfX8X2qjJ75fZBMSnQn3Rq2mrBJK2mw==", + "dev": true, + "license": "MIT", + "dependencies": { + "@babel/helper-plugin-utils": "^7.27.1" + }, + "engines": { + "node": ">=6.9.0" + }, + "peerDependencies": { + "@babel/core": "^7.0.0-0" + } + }, + "node_modules/@babel/plugin-transform-react-jsx-source": { + "version": "7.27.1", + "resolved": "https://registry.npmjs.org/@babel/plugin-transform-react-jsx-source/-/plugin-transform-react-jsx-source-7.27.1.tgz", + "integrity": "sha512-zbwoTsBruTeKB9hSq73ha66iFeJHuaFkUbwvqElnygoNbj/jHRsSeokowZFN3CZ64IvEqcmmkVe89OPXc7ldAw==", + "dev": true, + "license": "MIT", + "dependencies": { + "@babel/helper-plugin-utils": "^7.27.1" + }, + "engines": { + "node": ">=6.9.0" + }, + "peerDependencies": { + "@babel/core": "^7.0.0-0" + } + }, "node_modules/@babel/plugin-transform-react-pure-annotations": { "version": "7.27.1", "resolved": "https://registry.npmjs.org/@babel/plugin-transform-react-pure-annotations/-/plugin-transform-react-pure-annotations-7.27.1.tgz", @@ -1839,16 +1878,17 @@ } }, "node_modules/@babel/traverse": { - "version": "7.28.0", - "resolved": "https://registry.npmjs.org/@babel/traverse/-/traverse-7.28.0.tgz", - "integrity": "sha512-mGe7UK5wWyh0bKRfupsUchrQGqvDbZDbKJw+kcRGSmdHVYrv+ltd0pnpDTVpiTqnaBru9iEvA8pz8W46v0Amwg==", + "version": "7.28.4", + "resolved": "https://registry.npmjs.org/@babel/traverse/-/traverse-7.28.4.tgz", + "integrity": "sha512-YEzuboP2qvQavAcjgQNVgsvHIDv6ZpwXvcvjmyySP2DIMuByS/6ioU5G9pYrWHM6T2YDfc7xga9iNzYOs12CFQ==", + "license": "MIT", "dependencies": { "@babel/code-frame": "^7.27.1", - "@babel/generator": "^7.28.0", + "@babel/generator": "^7.28.3", "@babel/helper-globals": "^7.28.0", - "@babel/parser": "^7.28.0", + "@babel/parser": "^7.28.4", "@babel/template": "^7.27.2", - "@babel/types": "^7.28.0", + "@babel/types": "^7.28.4", "debug": "^4.3.1" }, "engines": { @@ -1856,9 +1896,10 @@ } }, "node_modules/@babel/types": { - "version": "7.28.2", - "resolved": "https://registry.npmjs.org/@babel/types/-/types-7.28.2.tgz", - "integrity": "sha512-ruv7Ae4J5dUYULmeXw1gmb7rYRz57OWCPM57pHojnLq/3Z1CK2lNSLTCVjxVk1F/TZHwOZZrOWi0ur95BbLxNQ==", + "version": "7.28.4", + "resolved": "https://registry.npmjs.org/@babel/types/-/types-7.28.4.tgz", + "integrity": "sha512-bkFqkLhh3pMBUQQkpVgWDWq/lqzc2678eUyDlTBhRqhCHFguYYGM0Efga7tYk4TogG/3x0EEl66/OQ+WGbWB/Q==", + "license": "MIT", "dependencies": { "@babel/helper-string-parser": "^7.27.1", "@babel/helper-validator-identifier": "^7.27.1" @@ -4051,6 +4092,7 @@ "version": "1.0.6", "resolved": "https://registry.npmjs.org/@heroicons/react/-/react-1.0.6.tgz", "integrity": "sha512-JJCXydOFWMDpCP4q13iEplA503MQO3xLoZiKum+955ZCtHINWnx26CUxVxxFQu/uLb4LW3ge15ZpzIkXKkJ8oQ==", + "license": "MIT", "peerDependencies": { "react": ">= 16" } @@ -4208,6 +4250,16 @@ "@jridgewell/trace-mapping": "^0.3.24" } }, + "node_modules/@jridgewell/remapping": { + "version": "2.3.5", + "resolved": "https://registry.npmjs.org/@jridgewell/remapping/-/remapping-2.3.5.tgz", + "integrity": "sha512-LI9u/+laYG4Ds1TDKSJW2YPrIlcVYOwi2fUC6xB43lueCjgxV4lffOCZCtYFiH6TNOX+tQKXx97T4IKHbhyHEQ==", + "license": "MIT", + "dependencies": { + "@jridgewell/gen-mapping": "^0.3.5", + "@jridgewell/trace-mapping": "^0.3.24" + } + }, "node_modules/@jridgewell/resolve-uri": { "version": "3.1.2", "resolved": "https://registry.npmjs.org/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz", @@ -4764,6 +4816,13 @@ "react": ">=18.2.0" } }, + "node_modules/@rolldown/pluginutils": { + "version": "1.0.0-beta.38", + "resolved": "https://registry.npmjs.org/@rolldown/pluginutils/-/pluginutils-1.0.0-beta.38.tgz", + "integrity": "sha512-N/ICGKleNhA5nc9XXQG/kkKHJ7S55u0x0XUJbbkmdCnFuoRkM1Il12q9q0eX19+M7KKUEPw/daUPIRnxhcxAIw==", + "dev": true, + "license": "MIT" + }, "node_modules/@rollup/rollup-android-arm-eabi": { "version": "4.52.0", "resolved": "https://registry.npmjs.org/@rollup/rollup-android-arm-eabi/-/rollup-android-arm-eabi-4.52.0.tgz", @@ -5475,6 +5534,41 @@ "license": "MIT", "peer": true }, + "node_modules/@types/babel__core": { + "version": "7.20.5", + "resolved": "https://registry.npmjs.org/@types/babel__core/-/babel__core-7.20.5.tgz", + "integrity": "sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@babel/parser": "^7.20.7", + "@babel/types": "^7.20.7", + "@types/babel__generator": "*", + "@types/babel__template": "*", + "@types/babel__traverse": "*" + } + }, + "node_modules/@types/babel__generator": { + "version": "7.27.0", + "resolved": "https://registry.npmjs.org/@types/babel__generator/-/babel__generator-7.27.0.tgz", + "integrity": "sha512-ufFd2Xi92OAVPYsy+P4n7/U7e68fex0+Ee8gSG9KX7eo084CWiQ4sdxktvdl0bOPupXtVJPY19zk6EwWqUQ8lg==", + "dev": true, + "license": "MIT", + "dependencies": { + "@babel/types": "^7.0.0" + } + }, + "node_modules/@types/babel__template": { + "version": "7.4.4", + "resolved": "https://registry.npmjs.org/@types/babel__template/-/babel__template-7.4.4.tgz", + "integrity": "sha512-h/NUaSyG5EyxBIp8YRxo4RMe2/qQgvyowRwVMzhYhBCONbW8PUsg4lkFMrhgZhUe5z3L3MiLDuvyJ/CaPa2A8A==", + "dev": true, + "license": "MIT", + "dependencies": { + "@babel/parser": "^7.1.0", + "@babel/types": "^7.0.0" + } + }, "node_modules/@types/babel__traverse": { "version": "7.28.0", "resolved": "https://registry.npmjs.org/@types/babel__traverse/-/babel__traverse-7.28.0.tgz", @@ -6403,6 +6497,27 @@ "resolved": "https://registry.npmjs.org/@ungap/structured-clone/-/structured-clone-1.2.0.tgz", "integrity": "sha512-zuVdFrMJiuCDQUMCzQaD6KL28MjnqqN8XnAqiEq9PNm/hCPTSGfrXCOfwj1ow4LFb/tNymJPwsNbVePc1xFqrQ==" }, + "node_modules/@vitejs/plugin-react": { + "version": "5.0.4", + "resolved": "https://registry.npmjs.org/@vitejs/plugin-react/-/plugin-react-5.0.4.tgz", + "integrity": "sha512-La0KD0vGkVkSk6K+piWDKRUyg8Rl5iAIKRMH0vMJI0Eg47bq1eOxmoObAaQG37WMW9MSyk7Cs8EIWwJC1PtzKA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@babel/core": "^7.28.4", + "@babel/plugin-transform-react-jsx-self": "^7.27.1", + "@babel/plugin-transform-react-jsx-source": "^7.27.1", + "@rolldown/pluginutils": "1.0.0-beta.38", + "@types/babel__core": "^7.20.5", + "react-refresh": "^0.17.0" + }, + "engines": { + "node": "^20.19.0 || >=22.12.0" + }, + "peerDependencies": { + "vite": "^4.2.0 || ^5.0.0 || ^6.0.0 || ^7.0.0" + } + }, "node_modules/@vitest/coverage-v8": { "version": "3.2.4", "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.4.tgz", @@ -18876,6 +18991,16 @@ "react": ">=18" } }, + "node_modules/react-refresh": { + "version": "0.17.0", + "resolved": "https://registry.npmjs.org/react-refresh/-/react-refresh-0.17.0.tgz", + "integrity": "sha512-z6F7K9bV85EfseRCp2bzrpyQ0Gkw1uLoCel9XBVWPg/TjRj94SkJzUTGfOa4bs7iJvBWtQG0Wq7wnI0syw3EBQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, "node_modules/react-router": { "version": "5.3.4", "resolved": "https://registry.npmjs.org/react-router/-/react-router-5.3.4.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index c7546d3612c..30c35591d6d 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -53,6 +53,7 @@ "@types/react-dom": "^18", "@types/react-syntax-highlighter": "^15.5.11", "@types/uuid": "^10.0.0", + "@vitejs/plugin-react": "^5.0.4", "@vitest/coverage-v8": "^3.2.4", "@vitest/ui": "^3.2.4", "autoprefixer": "^10.4.17", diff --git a/ui/litellm-dashboard/public/assets/logos/snowflake.svg b/ui/litellm-dashboard/public/assets/logos/snowflake.svg new file mode 100644 index 00000000000..e88dcad650b --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/snowflake.svg @@ -0,0 +1,9 @@ + + + + + + + \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index f30917df8ba..6741e43af50 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -507,6 +507,19 @@ const PROVIDER_CREDENTIAL_FIELDS: Record = label: "API Key", type: "password", required: true + }], + [Providers.Snowflake]: [{ + key: "api_key", + label: "Snowflake API Key / JWT Key for Authentication", + type: "password", + required: true + }, + { + key: "api_base", + label: "Snowflake API Endpoint", + placeholder: "https://1234567890.snowflakecomputing.com/api/v2/cortex/inference:complete", + tooltip: "Enter the full endpoint with path here. Example: https://1234567890.snowflakecomputing.com/api/v2/cortex/inference:complete", + required: true }] }; diff --git a/ui/litellm-dashboard/src/components/entity_usage.tsx b/ui/litellm-dashboard/src/components/entity_usage.tsx index 7b910dab57b..f7617b304b1 100644 --- a/ui/litellm-dashboard/src/components/entity_usage.tsx +++ b/ui/litellm-dashboard/src/components/entity_usage.tsx @@ -24,7 +24,7 @@ import { import AdvancedDatePicker from "./shared/advanced_date_picker"; import { Select } from 'antd'; import { ActivityMetrics, processActivityData } from './activity_metrics'; -import { DailyData, BreakdownMetrics, KeyMetricWithMetadata, EntityMetricWithMetadata } from './usage/types'; +import { DailyData, BreakdownMetrics, KeyMetricWithMetadata, EntityMetricWithMetadata, TagUsage } from './usage/types'; import { tagDailyActivityCall, teamDailyActivityCall } from './networking'; import TopKeyView from "./top_key_view"; import { formatNumberWithCommas } from "@/utils/dataUtils"; @@ -175,8 +175,25 @@ const EntityUsage: React.FC = ({ }; const getTopAPIKeys = () => { + console.log('debugTags',{spendData}) const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; spendData.results.forEach((day) => { + const {breakdown} = day; + const {entities} = breakdown; + console.log('debugTags',{entities}) + const tagDictionary = Object.keys(entities).reduce((acc: { [key: string]: TagUsage[] }, entity) => { + const {api_key_breakdown} = entities[entity]; + Object.keys(api_key_breakdown).forEach((key) => { + const tagUsage = {tag:entity,usage:api_key_breakdown[key].metrics.spend}; + if (acc[key]) { + acc[key].push(tagUsage); + } else { + acc[key] = [tagUsage]; + } + }) + return acc; + },{}) + console.log('debugTags',{tagDictionary}) Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => { if (!keySpend[key]) { keySpend[key] = { @@ -193,9 +210,11 @@ const EntityUsage: React.FC = ({ }, metadata: { key_alias: metrics.metadata.key_alias, - team_id: metrics.metadata.team_id || null + team_id: metrics.metadata.team_id || null, + tags: tagDictionary[key] || [] } }; + console.log('debugTags',{keySpend}) } keySpend[key].metrics.spend += metrics.metrics.spend; keySpend[key].metrics.prompt_tokens += metrics.metrics.prompt_tokens; @@ -218,6 +237,7 @@ const EntityUsage: React.FC = ({ .map(([api_key, metrics]) => ({ api_key, key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias + tags: metrics.metadata.tags || "-", spend: metrics.metrics.spend, })) .sort((a, b) => b.spend - a.spend) @@ -623,6 +643,7 @@ const EntityUsage: React.FC = ({ userRole={userRole} teams={null} premiumUser={premiumUser} + showTags={entityType === "tag"} /> diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 97132daa466..d2ba0fbfcaf 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -6,6 +6,7 @@ import { createMCPServer } from "../networking" import { MCPServer, MCPServerCostInfo } from "./types" import MCPServerCostConfig from "./mcp_server_cost_config" import MCPConnectionStatus from "./mcp_connection_status" +import MCPToolConfiguration from "./mcp_tool_configuration" import StdioConfiguration from "./StdioConfiguration" import { isAdminRole } from "@/utils/roles" import { validateMCPServerUrl, validateMCPServerName } from "./utils" @@ -37,7 +38,8 @@ const CreateMCPServer: React.FC = ({ const [formValues, setFormValues] = useState>({}) const [aliasManuallyEdited, setAliasManuallyEdited] = useState(false) const [tools, setTools] = useState([]) - const [transportType, setTransportType] = useState("sse") + const [allowedTools, setAllowedTools] = useState([]) + const [transportType, setTransportType] = useState("") const [searchValue, setSearchValue] = useState("") const [urlWarning, setUrlWarning] = useState("") @@ -103,7 +105,7 @@ const CreateMCPServer: React.FC = ({ } } - // Prepare the payload with cost configuration + // Prepare the payload with cost configuration and allowed tools const payload = { ...formValues, ...stdioFields, @@ -116,6 +118,7 @@ const CreateMCPServer: React.FC = ({ }, mcp_access_groups: accessGroups, alias: formValues.alias, + allowed_tools: allowedTools.length > 0 ? allowedTools : null, } console.log(`Payload: ${JSON.stringify(payload)}`) @@ -127,7 +130,9 @@ const CreateMCPServer: React.FC = ({ form.resetFields() setCostConfig({}) setTools([]) + setAllowedTools([]) setUrlWarning("") + setAliasManuallyEdited(false) setModalVisible(false) onCreateSuccess(response) } @@ -143,7 +148,9 @@ const CreateMCPServer: React.FC = ({ form.resetFields() setCostConfig({}) setTools([]) + setAllowedTools([]) setUrlWarning("") + setAliasManuallyEdited(false) setModalVisible(false) } @@ -176,7 +183,10 @@ const CreateMCPServer: React.FC = ({ })) // If search value doesn't match any existing group and is not empty, add "create new group" option - if (searchValue && !availableAccessGroups.some(group => group.toLowerCase().includes(searchValue.toLowerCase()))) { + if ( + searchValue && + !availableAccessGroups.some((group) => group.toLowerCase().includes(searchValue.toLowerCase())) + ) { existingOptions.push({ value: searchValue, label: ( @@ -201,6 +211,13 @@ const CreateMCPServer: React.FC = ({ } }, [formValues.server_name]) + // Clear formValues when modal closes to reset child components + React.useEffect(() => { + if (!isModalVisible) { + setFormValues({}) + } + }, [isModalVisible]) + // rendering if (!isAdminRole(userRole)) { return null @@ -297,7 +314,7 @@ const CreateMCPServer: React.FC = ({ rules={[ { required: false, - message: "Please enter a server description", + message: "Please enter a server description!!!!!!!!!", }, ]} > @@ -341,11 +358,7 @@ const CreateMCPServer: React.FC = ({ className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500" onChange={(e) => checkUrlFormat(e.target.value, transportType)} /> - {urlWarning && ( -
- {urlWarning} -
- )} + {urlWarning &&
{urlWarning}
} )} @@ -386,9 +399,7 @@ const CreateMCPServer: React.FC = ({ showSearch placeholder="Select existing groups or type to create new ones" optionFilterProp="value" - filterOption={(input, option) => - (option?.value ?? '').toLowerCase().includes(input.toLowerCase()) - } + filterOption={(input, option) => (option?.value ?? "").toLowerCase().includes(input.toLowerCase())} onSearch={(value) => setSearchValue(value)} tokenSeparators={[","]} options={getAccessGroupOptions()} @@ -403,9 +414,24 @@ const CreateMCPServer: React.FC = ({ + {/* Tool Configuration Section */} +
+ +
+ {/* Cost Configuration Section */}
- + allowedTools.includes(tool.name))} + disabled={false} + />
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx index adeae707af9..acba932c70c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connection_status.tsx @@ -1,92 +1,30 @@ -import React, { useState, useEffect } from "react"; -import { Button, message, Spin, Alert, Collapse, Badge } from "antd"; -import { CheckCircleOutlined, ExclamationCircleOutlined, ReloadOutlined, ToolOutlined, InfoCircleOutlined } from "@ant-design/icons"; -import { Card, Title, Text } from "@tremor/react"; -import { testMCPToolsListRequest } from "../networking"; - -const { Panel } = Collapse; +import React, { useEffect } from "react" +import { Button, Spin, Alert } from "antd" +import { CheckCircleOutlined, ExclamationCircleOutlined, ReloadOutlined, ToolOutlined } from "@ant-design/icons" +import { Card, Title, Text } from "@tremor/react" +import { useTestMCPConnection } from "../../hooks/useTestMCPConnection" interface MCPConnectionStatusProps { - accessToken: string | null; - formValues: Record; - onToolsLoaded?: (tools: any[]) => void; + accessToken: string | null + formValues: Record + onToolsLoaded?: (tools: any[]) => void } -const MCPConnectionStatus: React.FC = ({ - accessToken, - formValues, - onToolsLoaded -}) => { - const [tools, setTools] = useState([]); - const [isLoadingTools, setIsLoadingTools] = useState(false); - const [toolsError, setToolsError] = useState(null); - const [hasShownSuccessMessage, setHasShownSuccessMessage] = useState(false); +const MCPConnectionStatus: React.FC = ({ accessToken, formValues, onToolsLoaded }) => { + const { tools, isLoadingTools, toolsError, canFetchTools, fetchTools } = useTestMCPConnection({ + accessToken, + formValues, + enabled: true, // Auto-fetch when required fields are available + }) - // Check if we have the minimum required fields to fetch tools - const canFetchTools = formValues.url && formValues.transport && formValues.auth_type && accessToken; - - const fetchTools = async () => { - if (!accessToken || !formValues.url) { - return; - } - - setIsLoadingTools(true); - setToolsError(null); - - try { - // Prepare the MCP server config from form values - const mcpServerConfig = { - server_id: formValues.server_id || "", - server_name: formValues.server_name || "", - url: formValues.url, - transport: formValues.transport, - auth_type: formValues.auth_type, - mcp_info: formValues.mcp_info, - }; - - const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig); - - if (toolsResponse.tools && !toolsResponse.error) { - setTools(toolsResponse.tools); - setToolsError(null); - onToolsLoaded?.(toolsResponse.tools); - if (toolsResponse.tools.length > 0 && !hasShownSuccessMessage) { - setHasShownSuccessMessage(true); - } - } else { - const errorMessage = toolsResponse.message || "Failed to retrieve tools list"; - setToolsError(errorMessage); - setTools([]); - onToolsLoaded?.([]); - setHasShownSuccessMessage(false); - } - } catch (error) { - console.error("Tools fetch error:", error); - setToolsError(error instanceof Error ? error.message : String(error)); - setTools([]); - onToolsLoaded?.([]); - setHasShownSuccessMessage(false); - } finally { - setIsLoadingTools(false); - } - }; - - // Auto-fetch tools when form values change and required fields are available + // Notify parent component when tools change useEffect(() => { - if (canFetchTools) { - fetchTools(); - } else { - // Clear tools if required fields are missing - setTools([]); - setToolsError(null); - setHasShownSuccessMessage(false); - onToolsLoaded?.([]); - } - }, [formValues.url, formValues.transport, formValues.auth_type, accessToken]); + onToolsLoaded?.(tools) + }, [tools, onToolsLoaded]) // Don't show anything if required fields aren't filled if (!canFetchTools && !formValues.url) { - return null; + return null } return ( @@ -102,9 +40,7 @@ const MCPConnectionStatus: React.FC = ({ Complete required fields to test connection
- - Fill in URL, Transport, and Authentication to test MCP server connection - + Fill in URL, Transport, and Authentication to test MCP server connection
)} @@ -113,34 +49,32 @@ const MCPConnectionStatus: React.FC = ({
- {isLoadingTools - ? "Testing connection to MCP server..." - : tools.length > 0 + {isLoadingTools + ? "Testing connection to MCP server..." + : tools.length > 0 ? "Connection successful" : toolsError ? "Connection failed" : "Ready to test connection"}
- - Server: {formValues.url} - + Server: {formValues.url}
- + {isLoadingTools && (
Connecting...
)} - + {!isLoadingTools && !toolsError && tools.length > 0 && (
Connected
)} - + {toolsError && (
@@ -163,54 +97,13 @@ const MCPConnectionStatus: React.FC = ({ type="error" showIcon action={ - } /> )} - {!isLoadingTools && tools.length > 0 && ( - - - Available Tools - -
- ), - children: ( -
- {tools.map((tool, index) => ( -
- {tool.name} - {tool.description && ( - - {tool.description} - - )} -
- ))} -
- ), - }, - ]} - /> - )} - {!isLoadingTools && tools.length === 0 && !toolsError && (
@@ -223,7 +116,7 @@ const MCPConnectionStatus: React.FC = ({ )}
- ); -}; + ) +} -export default MCPConnectionStatus; \ No newline at end of file +export default MCPConnectionStatus diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx index 0fc9d834b50..02329936f77 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx @@ -108,7 +108,7 @@ export const MCPServerView: React.FC = ({ ? "text-green-600 bg-green-50 border-green-200" : "text-gray-500 hover:text-gray-700 hover:bg-gray-100" }`} - /> + />
@@ -241,6 +241,25 @@ export const MCPServerView: React.FC = ({ )} +
+ Allowed Tools +
+ {mcpServer.allowed_tools && mcpServer.allowed_tools.length > 0 ? ( +
+ {mcpServer.allowed_tools.map((tool: string, index: number) => ( + + {tool} + + ))} +
+ ) : ( + All tools enabled + )} +
+
Cost Configuration diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index 95073ad1dbd..fd55dba604d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -50,6 +50,17 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) enabled: !!accessToken, }) as { data: MCPServer[]; isLoading: boolean; refetch: () => void; dataUpdatedAt: number } + // Log allowed_tools from fetched servers + React.useEffect(() => { + if (mcpServers) { + console.log("MCP Servers fetched:", mcpServers) + mcpServers.forEach((server) => { + console.log(`Server: ${server.server_name || server.server_id}`) + console.log(` allowed_tools:`, server.allowed_tools) + }) + } + }, [mcpServers]) + // state const [serverIdToDelete, setServerToDelete] = useState(null) const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false) @@ -60,7 +71,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) const [filteredServers, setFilteredServers] = useState([]) const [isModalVisible, setModalVisible] = useState(false) - const isInternalUser = userRole === "Internal User"; + const isInternalUser = userRole === "Internal User" // Get unique teams from all servers const uniqueTeams = React.useMemo(() => { @@ -84,7 +95,11 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) // Get unique MCP access groups from all servers const uniqueMcpAccessGroups = React.useMemo(() => { if (!mcpServers) return [] - return Array.from(new Set(mcpServers.flatMap((server) => server.mcp_access_groups).filter((group): group is string => group != null))) + return Array.from( + new Set( + mcpServers.flatMap((server) => server.mcp_access_groups).filter((group): group is string => group != null), + ), + ) }, [mcpServers]) // Handle team filter change @@ -171,7 +186,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) } if (!accessToken || !userRole || !userID) { - console.log("Missing required authentication parameters", { accessToken, userRole, userID }); + console.log("Missing required authentication parameters", { accessToken, userRole, userID }) return
Missing required authentication parameters.
} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx new file mode 100644 index 00000000000..fbf8b04048b --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx @@ -0,0 +1,177 @@ +import React, { useEffect, useRef } from "react" +import { Card, Title, Text } from "@tremor/react" +import { ToolOutlined, CheckCircleOutlined } from "@ant-design/icons" +import { Badge, Spin, Checkbox } from "antd" +import { useTestMCPConnection } from "../../hooks/useTestMCPConnection" + +interface MCPToolConfigurationProps { + accessToken: string | null + formValues: Record + allowedTools: string[] + onAllowedToolsChange: (tools: string[]) => void +} + +const MCPToolConfiguration: React.FC = ({ + accessToken, + formValues, + allowedTools, + onAllowedToolsChange, +}) => { + const previousToolsLengthRef = useRef(0) + + const { tools, isLoadingTools, toolsError, canFetchTools } = useTestMCPConnection({ + accessToken, + formValues, + enabled: true, + }) + + // Auto-select all tools when tools are first loaded + useEffect(() => { + // Only auto-select if: + // 1. We have tools + // 2. Tools length changed (new tools loaded) + // 3. No tools are currently selected (initial state) + if (tools.length > 0 && tools.length !== previousToolsLengthRef.current && allowedTools.length === 0) { + const allToolNames = tools.map((tool) => tool.name) + onAllowedToolsChange(allToolNames) + } + // Update ref to track tools length (will be 0 when tools clear) + previousToolsLengthRef.current = tools.length + }, [tools, allowedTools.length, onAllowedToolsChange]) + + const handleToolToggle = (toolName: string) => { + if (allowedTools.includes(toolName)) { + onAllowedToolsChange(allowedTools.filter((name) => name !== toolName)) + } else { + onAllowedToolsChange([...allowedTools, toolName]) + } + } + + const handleSelectAll = () => { + const allToolNames = tools.map((tool) => tool.name) + onAllowedToolsChange(allToolNames) + } + + const handleDeselectAll = () => { + onAllowedToolsChange([]) + } + + // Don't show anything if required fields aren't filled + if (!canFetchTools && !formValues.url) { + return null + } + + return ( + +
+
+
+ + Tool Configuration + {tools.length > 0 && ( + + )} +
+
+ + {/* Loading state */} + {isLoadingTools && ( +
+ + Loading tools... +
+ )} + + {/* Error state */} + {toolsError && !isLoadingTools && ( +
+ + Unable to load tools +
+ {toolsError} +
+ )} + + {/* No tools state */} + {!isLoadingTools && !toolsError && tools.length === 0 && canFetchTools && ( +
+ + No tools available for configuration +
+ Connect to an MCP server with tools to configure them +
+ )} + + {/* Incomplete form state */} + {!canFetchTools && formValues.url && ( +
+ + Complete required fields to configure tools +
+ Fill in URL, Transport, and Authentication to load available tools +
+ )} + + {/* Tools loaded successfully */} + {!isLoadingTools && !toolsError && tools.length > 0 && ( +
+
+
+ + + {allowedTools.length} of {tools.length} {tools.length === 1 ? "tool" : "tools"} selected + +
+
+ + +
+
+ + {/* Tool list with checkboxes */} +
+ {tools.map((tool, index) => ( +
handleToolToggle(tool.name)} + > +
+ handleToolToggle(tool.name)} /> +
+ {tool.name} + {tool.description && {tool.description}} +
+
+
+ ))} +
+
+ )} +
+ + ) +} + +export default MCPToolConfiguration diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 5213c2610eb..9bf5725ac0b 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -1,7 +1,7 @@ export interface Team { - team_id: string; - team_alias?: string; - organization_id?: string | null; + team_id: string + team_alias?: string + organization_id?: string | null } // Default no auth value @@ -10,140 +10,141 @@ export const AUTH_TYPE = { API_KEY: "api_key", BEARER_TOKEN: "bearer_token", BASIC: "basic", -}; +} export const TRANSPORT = { SSE: "sse", HTTP: "http", -}; +} export const handleTransport = (transport?: string | null): string => { console.log(transport) if (transport === null || transport === undefined) { - return TRANSPORT.SSE; + return TRANSPORT.SSE } - return transport; -}; + return transport +} export const handleAuth = (authType?: string | null): string => { if (authType === null || authType === undefined) { - return AUTH_TYPE.NONE; + return AUTH_TYPE.NONE } - return authType; -}; + return authType +} export const mcpServerHasAuth = (authType?: string | null): boolean => { - return handleAuth(authType) !== AUTH_TYPE.NONE; -} + return handleAuth(authType) !== AUTH_TYPE.NONE +} // Define the structure for tool input schema properties export interface InputSchemaProperty { - type: string; - description?: string; - properties?: Record; // For nested object properties - required?: string[]; // For required fields in nested objects - enum?: string[]; // For enum values - default?: any; // For default values - } - - // Define the structure for the input schema of a tool - export interface InputSchema { - type: "object"; - properties: Record; - required?: string[]; - } - - // Define MCPServerCostInfo for cost tracking - export interface MCPServerCostInfo { - default_cost_per_query?: number | null; - tool_name_to_cost_per_query?: Record; - } + type: string + description?: string + properties?: Record // For nested object properties + required?: string[] // For required fields in nested objects + enum?: string[] // For enum values + default?: any // For default values +} - // Define MCP provider info - export interface MCPInfo { - server_name: string; - description?: string; - logo_url?: string; - mcp_server_cost_info?: MCPServerCostInfo | null; - } - - // Define the structure for a single MCP tool - export interface MCPTool { - name: string; - description?: string; - inputSchema: InputSchema | string; // API returns string "tool_input_schema" or the actual schema - mcp_info: MCPInfo; - // Function to select a tool (added in the component) - onToolSelect?: (tool: MCPTool) => void; - } - - // Define the response structure for the listMCPTools endpoint - now a flat array - export type ListMCPToolsResponse = MCPTool[]; - - // Define the argument structure for calling an MCP tool - export interface CallMCPToolArgs { - name: string; - arguments: Record | null; - server_name?: string; // Now using server_name from mcp_info - } - - // Define the possible content types in the response - export interface MCPTextContent { - type: "text"; - text: string; - annotations?: any; - } - - export interface MCPImageContent { - type: "image"; - url?: string; - data?: string; - } - - export interface MCPEmbeddedResource { - type: "embedded_resource"; - resource_type?: string; - url?: string; - data?: any; - } - - // Define the union type for the content array in the response - export type MCPContent = MCPTextContent | MCPImageContent | MCPEmbeddedResource; - - // Define the response structure for the callMCPTool endpoint - export type CallMCPToolResponse = MCPContent[]; - - // Props for the main component - export interface MCPToolsViewerProps { - serverId: string; - accessToken: string | null; - auth_type?: string | null; - userRole: string | null; - userID: string | null; - serverAlias?: string | null; - } +// Define the structure for the input schema of a tool +export interface InputSchema { + type: "object" + properties: Record + required?: string[] +} + +// Define MCPServerCostInfo for cost tracking +export interface MCPServerCostInfo { + default_cost_per_query?: number | null + tool_name_to_cost_per_query?: Record +} + +// Define MCP provider info +export interface MCPInfo { + server_name: string + description?: string + logo_url?: string + mcp_server_cost_info?: MCPServerCostInfo | null +} + +// Define the structure for a single MCP tool +export interface MCPTool { + name: string + description?: string + inputSchema: InputSchema | string // API returns string "tool_input_schema" or the actual schema + mcp_info: MCPInfo + // Function to select a tool (added in the component) + onToolSelect?: (tool: MCPTool) => void +} + +// Define the response structure for the listMCPTools endpoint - now a flat array +export type ListMCPToolsResponse = MCPTool[] + +// Define the argument structure for calling an MCP tool +export interface CallMCPToolArgs { + name: string + arguments: Record | null + server_name?: string // Now using server_name from mcp_info +} + +// Define the possible content types in the response +export interface MCPTextContent { + type: "text" + text: string + annotations?: any +} + +export interface MCPImageContent { + type: "image" + url?: string + data?: string +} + +export interface MCPEmbeddedResource { + type: "embedded_resource" + resource_type?: string + url?: string + data?: any +} + +// Define the union type for the content array in the response +export type MCPContent = MCPTextContent | MCPImageContent | MCPEmbeddedResource + +// Define the response structure for the callMCPTool endpoint +export type CallMCPToolResponse = MCPContent[] + +// Props for the main component +export interface MCPToolsViewerProps { + serverId: string + accessToken: string | null + auth_type?: string | null + userRole: string | null + userID: string | null + serverAlias?: string | null +} export interface MCPServer { - server_id: string; - server_name?: string | null; - alias?: string | null; - description?: string | null; - url: string; - transport?: string | null; - auth_type?: string | null; - mcp_info?: MCPInfo | null; - created_at: string; - created_by: string; - updated_at: string; - updated_by: string; - teams?: Team[]; - mcp_access_groups?: string[]; + server_id: string + server_name?: string | null + alias?: string | null + description?: string | null + url: string + transport?: string | null + auth_type?: string | null + mcp_info?: MCPInfo | null + created_at: string + created_by: string + updated_at: string + updated_by: string + teams?: Team[] + mcp_access_groups?: string[] + allowed_tools?: string[] } export interface MCPServerProps { - accessToken: string | null; - userRole: string | null; - userID: string | null; -} \ No newline at end of file + accessToken: string | null + userRole: string | null + userID: string | null +} diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts new file mode 100644 index 00000000000..4d74eda28ef --- /dev/null +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -0,0 +1,44 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { clearTokenCookies } from '@/utils/cookieUtils'; + +vi.mock('@/utils/cookieUtils', () => ({ + clearTokenCookies: vi.fn(), + getCookie: vi.fn(), +})); + +vi.mock('./molecules/notifications_manager', () => ({ + default: { + info: vi.fn(), + success: vi.fn(), + error: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +describe('networking - expired session handling', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('should call clearTokenCookies on expired session', async () => { + const errorData = "Authentication Error - Expired Key"; + const { default: NotificationsManager } = await import('./molecules/notifications_manager'); + + if (errorData.includes("Authentication Error - Expired Key")) { + NotificationsManager.info("UI Session Expired. Logging out."); + clearTokenCookies(); + } + + expect(clearTokenCookies).toHaveBeenCalledOnce(); + }); + + it('should not clear cookies for non-authentication errors', () => { + const errorData = "Some other error"; + + if (errorData.includes("Authentication Error - Expired Key")) { + clearTokenCookies(); + } + + expect(clearTokenCookies).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index b16543b8f7a..6b8c7671a10 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -10,6 +10,7 @@ export const formatDate = (date: Date) => { */ import { all_admin_roles } from "@/utils/roles"; import { message } from "antd"; +import { clearTokenCookies } from "@/utils/cookieUtils"; import { TagNewRequest, TagUpdateRequest, @@ -169,8 +170,7 @@ const handleError = async (errorData: string) => { if (errorData.includes("Authentication Error - Expired Key")) { NotificationsManager.info("UI Session Expired. Logging out."); lastErrorTime = currentTime; - document.cookie = - "token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;"; + clearTokenCookies(); window.location.href = window.location.pathname; } lastErrorTime = currentTime; @@ -335,14 +335,14 @@ export const getModelCostMapReloadStatus = async (accessToken: string) => { "Content-Type": "application/json", }, }); - + if (!response.ok) { console.error(`Status request failed with status: ${response.status}`); const errorText = await response.text(); console.error("Error response:", errorText); throw new Error(`HTTP ${response.status}: ${errorText}`); } - + const jsonData = await response.json(); console.log(`Model cost map reload status:`, jsonData); return jsonData; @@ -4090,7 +4090,7 @@ export const teamMemberUpdateCall = async ( const url = proxyBaseUrl ? `${proxyBaseUrl}/team/member_update` : `/team/member_update`; - + const requestBody: any = { team_id: teamId, role: formValues.role, @@ -4373,9 +4373,9 @@ export const userBulkUpdateUserCall = async ( const url = proxyBaseUrl ? `${proxyBaseUrl}/user/bulk_update` : `/user/bulk_update`; - + let request_body_json: string; - + if (allUsers) { // Update all users mode request_body_json = JSON.stringify({ @@ -4384,7 +4384,7 @@ export const userBulkUpdateUserCall = async ( }); } else if (userIds && userIds.length > 0) { // Update specific users mode - let request_body = [] + let request_body = [] for (const user_id of userIds) { request_body.push({ user_id: user_id, @@ -4397,7 +4397,7 @@ export const userBulkUpdateUserCall = async ( } else { throw new Error("Must provide either userIds or set allUsers=true"); } - + const response = await fetch(url, { method: "POST", headers: { @@ -5410,11 +5410,11 @@ export const convertPromptFileToJson = async ( try { const formData = new FormData(); formData.append("file", file); - - const url = proxyBaseUrl - ? `${proxyBaseUrl}/utils/dotprompt_json_converter` + + const url = proxyBaseUrl + ? `${proxyBaseUrl}/utils/dotprompt_json_converter` : `/utils/dotprompt_json_converter`; - + const response = await fetch(url, { method: "POST", headers: { @@ -5879,14 +5879,14 @@ export const callMCPTool = async ( if (!response.ok) { let errorMessage = "Network response was not ok"; let errorDetails = null; - + // First, try to get the response as text to see what we're dealing with const responseText = await response.text(); - + try { // Try to parse as JSON const errorData = JSON.parse(responseText); - + if (errorData.detail) { if (typeof errorData.detail === 'string') { errorMessage = errorData.detail; @@ -5897,7 +5897,7 @@ export const callMCPTool = async ( } else { errorMessage = errorData.message || errorData.error || errorMessage; } - + } catch (parseError) { console.error("Failed to parse JSON error response:", parseError); // If JSON parsing fails, use the raw text @@ -5905,13 +5905,13 @@ export const callMCPTool = async ( errorMessage = responseText; } } - + // Create a more informative error object const enhancedError = new Error(errorMessage); (enhancedError as any).status = response.status; (enhancedError as any).statusText = response.statusText; (enhancedError as any).details = errorDetails; - + handleError(errorMessage); throw enhancedError; } @@ -7126,9 +7126,9 @@ export const userAgentAnalyticsCall = async ( let url = proxyBaseUrl ? `${proxyBaseUrl}/tag/user-agent/analytics` : `/tag/user-agent/analytics`; - + const queryParams = new URLSearchParams(); - + // Format dates as YYYY-MM-DD for the API const formatDate = (date: Date) => { const year = date.getFullYear(); @@ -7136,16 +7136,16 @@ export const userAgentAnalyticsCall = async ( const day = String(date.getDate()).padStart(2, '0'); return `${year}-${month}-${day}`; }; - + queryParams.append("start_date", formatDate(startTime)); queryParams.append("end_date", formatDate(endTime)); queryParams.append("page", page.toString()); queryParams.append("page_size", pageSize.toString()); - + if (userAgentFilter) { queryParams.append("user_agent_filter", userAgentFilter); } - + const queryString = queryParams.toString(); if (queryString) { url += `?${queryString}`; @@ -7189,9 +7189,9 @@ export const tagDauCall = async ( let url = proxyBaseUrl ? `${proxyBaseUrl}/tag/dau` : `/tag/dau`; - + const queryParams = new URLSearchParams(); - + // Format date as YYYY-MM-DD for the API const formatDate = (date: Date) => { const year = date.getFullYear(); @@ -7199,9 +7199,9 @@ export const tagDauCall = async ( const day = String(date.getDate()).padStart(2, '0'); return `${year}-${month}-${day}`; }; - + queryParams.append("end_date", formatDate(endDate)); - + // Handle multiple tag filters (takes precedence over single tag filter) if (tagFilters && tagFilters.length > 0) { tagFilters.forEach(tag => { @@ -7210,7 +7210,7 @@ export const tagDauCall = async ( } else if (tagFilter) { queryParams.append("tag_filter", tagFilter); } - + const queryString = queryParams.toString(); if (queryString) { url += `?${queryString}`; @@ -7253,9 +7253,9 @@ export const tagWauCall = async ( let url = proxyBaseUrl ? `${proxyBaseUrl}/tag/wau` : `/tag/wau`; - + const queryParams = new URLSearchParams(); - + // Format date as YYYY-MM-DD for the API const formatDate = (date: Date) => { const year = date.getFullYear(); @@ -7263,9 +7263,9 @@ export const tagWauCall = async ( const day = String(date.getDate()).padStart(2, '0'); return `${year}-${month}-${day}`; }; - + queryParams.append("end_date", formatDate(endDate)); - + // Handle multiple tag filters (takes precedence over single tag filter) if (tagFilters && tagFilters.length > 0) { tagFilters.forEach(tag => { @@ -7274,7 +7274,7 @@ export const tagWauCall = async ( } else if (tagFilter) { queryParams.append("tag_filter", tagFilter); } - + const queryString = queryParams.toString(); if (queryString) { url += `?${queryString}`; @@ -7317,9 +7317,9 @@ export const tagMauCall = async ( let url = proxyBaseUrl ? `${proxyBaseUrl}/tag/mau` : `/tag/mau`; - + const queryParams = new URLSearchParams(); - + // Format date as YYYY-MM-DD for the API const formatDate = (date: Date) => { const year = date.getFullYear(); @@ -7327,9 +7327,9 @@ export const tagMauCall = async ( const day = String(date.getDate()).padStart(2, '0'); return `${year}-${month}-${day}`; }; - + queryParams.append("end_date", formatDate(endDate)); - + // Handle multiple tag filters (takes precedence over single tag filter) if (tagFilters && tagFilters.length > 0) { tagFilters.forEach(tag => { @@ -7338,7 +7338,7 @@ export const tagMauCall = async ( } else if (tagFilter) { queryParams.append("tag_filter", tagFilter); } - + const queryString = queryParams.toString(); if (queryString) { url += `?${queryString}`; @@ -7416,9 +7416,9 @@ export const userAgentSummaryCall = async ( let url = proxyBaseUrl ? `${proxyBaseUrl}/tag/summary` : `/tag/summary`; - + const queryParams = new URLSearchParams(); - + // Format dates as YYYY-MM-DD for the API const formatDate = (date: Date) => { const year = date.getFullYear(); @@ -7426,17 +7426,17 @@ export const userAgentSummaryCall = async ( const day = String(date.getDate()).padStart(2, '0'); return `${year}-${month}-${day}`; }; - + queryParams.append("start_date", formatDate(startTime)); queryParams.append("end_date", formatDate(endTime)); - + // Handle multiple tag filters if (tagFilters && tagFilters.length > 0) { tagFilters.forEach(tag => { queryParams.append("tag_filters", tag); }); } - + const queryString = queryParams.toString(); if (queryString) { url += `?${queryString}`; @@ -7479,19 +7479,19 @@ export const perUserAnalyticsCall = async ( let url = proxyBaseUrl ? `${proxyBaseUrl}/tag/user-agent/per-user-analytics` : `/tag/user-agent/per-user-analytics`; - + const queryParams = new URLSearchParams(); - + queryParams.append("page", page.toString()); queryParams.append("page_size", pageSize.toString()); - + // Handle multiple tag filters if (tagFilters && tagFilters.length > 0) { tagFilters.forEach(tag => { queryParams.append("tag_filters", tag); }); } - + const queryString = queryParams.toString(); if (queryString) { url += `?${queryString}`; diff --git a/ui/litellm-dashboard/src/components/new_usage.tsx b/ui/litellm-dashboard/src/components/new_usage.tsx index e17c2d78be3..10dc306c713 100644 --- a/ui/litellm-dashboard/src/components/new_usage.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.tsx @@ -269,6 +269,7 @@ const NewUsagePage: React.FC = ({ accessToken, userRole, user metadata: { key_alias: metrics.metadata.key_alias, team_id: null, + tags: metrics.metadata.tags || [], // This gets key-level tags }, } } @@ -284,10 +285,13 @@ const NewUsagePage: React.FC = ({ accessToken, userRole, user }) }) + console.log('debugTags',{keySpend,userSpendData}) + return Object.entries(keySpend) .map(([api_key, metrics]) => ({ api_key, key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias + tags: metrics.metadata.tags || [], // This will show key-level tags spend: metrics.metrics.spend, })) .sort((a, b) => b.spend - a.spend) diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 744f8c117d6..7f854d9df15 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -34,6 +34,7 @@ export enum Providers { Oracle = "Oracle Cloud Infrastructure (OCI)", Perplexity = "Perplexity", Sambanova = "Sambanova", + Snowflake = "Snowflake", TogetherAI = "TogetherAI", Triton = "Triton", Vertex_AI = "Vertex AI (Anthropic, Gemini, etc.)", @@ -69,6 +70,7 @@ export const provider_map: Record = { TogetherAI: "together_ai", Openrouter: "openrouter", Oracle: "oci", + Snowflake: "snowflake", FireworksAI: "fireworks_ai", GradientAI: "gradient_ai", Triton: "triton", @@ -111,6 +113,7 @@ export const providerLogoMap: Record = { [Providers.Oracle]: `${asset_logos_folder}oracle.svg`, [Providers.Perplexity]: `${asset_logos_folder}perplexity-ai.svg`, [Providers.Sambanova]: `${asset_logos_folder}sambanova.svg`, + [Providers.Snowflake]: `${asset_logos_folder}snowflake.svg`, [Providers.TogetherAI]: `${asset_logos_folder}togetherai.svg`, [Providers.Vertex_AI]: `${asset_logos_folder}google.svg`, [Providers.xAI]: `${asset_logos_folder}xai.svg`, @@ -171,6 +174,8 @@ export const getPlaceholder = (selectedProvider: string): string => { return "azure/my-deployment"; } else if (selectedProvider == Providers.Oracle) { return "oci/xai.grok-4"; + } else if (selectedProvider == Providers.Snowflake) { + return "snowflake/mistral-7b"; } else if (selectedProvider == Providers.Voyage) { return "voyage/"; } else if (selectedProvider == Providers.JinaAI) { diff --git a/ui/litellm-dashboard/src/components/top_key_view.tsx b/ui/litellm-dashboard/src/components/top_key_view.tsx index 3db72123400..65fefbc9069 100644 --- a/ui/litellm-dashboard/src/components/top_key_view.tsx +++ b/ui/litellm-dashboard/src/components/top_key_view.tsx @@ -7,6 +7,8 @@ import { DataTable } from "./view_logs/table" import { Tooltip } from "antd" import { Button } from "@tremor/react" import { formatNumberWithCommas } from "../utils/dataUtils" +import { TagUsage } from "./usage/types" +import { ChevronDownIcon, ChevronUpIcon } from "@heroicons/react/outline" interface TopKeyViewProps { topKeys: any[] @@ -15,13 +17,27 @@ interface TopKeyViewProps { userRole: string | null teams: any[] | null premiumUser: boolean + showTags?: boolean } -const TopKeyView: React.FC = ({ topKeys, accessToken, userID, userRole, teams, premiumUser }) => { +const TopKeyView: React.FC = ({ topKeys, accessToken, userID, userRole, teams, premiumUser, showTags = false }) => { const [isModalOpen, setIsModalOpen] = useState(false) const [selectedKey, setSelectedKey] = useState(null) const [keyData, setKeyData] = useState(undefined) const [viewMode, setViewMode] = useState<"chart" | "table">("table") + const [expandedTags, setExpandedTags] = useState>(new Set()) + + const toggleTagsExpansion = (apiKey: string) => { + setExpandedTags(prev => { + const newSet = new Set(prev) + if (newSet.has(apiKey)) { + newSet.delete(apiKey) + } else { + newSet.add(apiKey) + } + return newSet + }) + } const handleKeyClick = async (item: any) => { if (!accessToken) return @@ -64,7 +80,7 @@ const TopKeyView: React.FC = ({ topKeys, accessToken, userID, u }, [isModalOpen]) // Define columns for the table view - const columns = [ + const baseColumns = [ { header: "Key ID", accessorKey: "api_key", @@ -88,13 +104,74 @@ const TopKeyView: React.FC = ({ topKeys, accessToken, userID, u accessorKey: "key_alias", cell: (info: any) => info.getValue() || "-", }, - { - header: "Spend (USD)", - accessorKey: "spend", - cell: (info: any) => `$${formatNumberWithCommas(info.getValue(), 2)}`, - }, ] + const tagsColumn = { + header: "Tags", + accessorKey: "tags", + cell: (info: any) => { + const tags = info.getValue() as TagUsage[] | undefined; + const apiKey = info.row.original.api_key; + const isExpanded = expandedTags.has(apiKey); + + if (!tags || tags.length === 0) { + return "-"; + } + + const sortedTags = tags.sort((a, b) => b.usage - a.usage); + const displayTags = isExpanded ? sortedTags : sortedTags.slice(0, 2); + const hasMoreTags = tags.length > 2; + + return ( +
+
+ {displayTags.map((tag, index) => ( + +
Tag Name: {tag.tag}
+
Spend: {tag.usage > 0 && tag.usage < 0.01 ? '<$0.01' : `$${formatNumberWithCommas(tag.usage, 2)}`}
+
+ } + > + + {tag.tag.slice(0, 7)}... + + + ))} + {hasMoreTags && ( + + )} +
+
+ ); + } + } + + const spendColumn = { + header: "Spend (USD)", + accessorKey: "spend", + cell: (info: any) => { + const value = info.getValue(); + return value > 0 && value < 0.01 ? '<$0.01' : `$${formatNumberWithCommas(value, 2)}`; + }, + } + + const columns = showTags + ? [...baseColumns, tagsColumn, spendColumn] + : [...baseColumns, spendColumn] + const processedTopKeys = topKeys.map((k) => ({ ...k, display_key_alias: k.key_alias && k.key_alias.length > 10 ? `${k.key_alias.slice(0, 10)}...` : k.key_alias || "-", diff --git a/ui/litellm-dashboard/src/components/usage/types.ts b/ui/litellm-dashboard/src/components/usage/types.ts index 5d0779246cc..b7c5ecb9ae7 100644 --- a/ui/litellm-dashboard/src/components/usage/types.ts +++ b/ui/litellm-dashboard/src/components/usage/types.ts @@ -39,6 +39,7 @@ export interface KeyMetricWithMetadata { export interface KeyMetadata { key_alias: string | null team_id: string | null + tags?: {tag:string,usage:number}[] } export interface TopApiKeyData { @@ -87,3 +88,8 @@ export interface EntityMetricWithMetadata { metrics: SpendMetrics metadata: EntityMetadata } + +export interface TagUsage { + tag: string + usage: number +} diff --git a/ui/litellm-dashboard/src/components/view_logs/table.tsx b/ui/litellm-dashboard/src/components/view_logs/table.tsx index c4f620e4413..b49ef30e83f 100644 --- a/ui/litellm-dashboard/src/components/view_logs/table.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/table.tsx @@ -53,7 +53,7 @@ export function DataTable({ return (
- +
{table.getHeaderGroups().map((headerGroup) => ( diff --git a/ui/litellm-dashboard/src/contexts/ThemeContext.tsx b/ui/litellm-dashboard/src/contexts/ThemeContext.tsx index 619d6a1004a..7d53946d6d9 100644 --- a/ui/litellm-dashboard/src/contexts/ThemeContext.tsx +++ b/ui/litellm-dashboard/src/contexts/ThemeContext.tsx @@ -25,34 +25,33 @@ export const ThemeProvider: React.FC = ({ children, accessTo const [logoUrl, setLogoUrl] = useState(null); // Load logo URL from backend on mount + // Note: /get/ui_theme_settings is now a public endpoint (no auth required) + // so all users can see custom branding set by admins useEffect(() => { const loadLogoSettings = async () => { - if (accessToken) { - try { - const proxyBaseUrl = getProxyBaseUrl(); - const url = proxyBaseUrl ? `${proxyBaseUrl}/get/ui_theme_settings` : '/get/ui_theme_settings'; - const response = await fetch(url, { - method: 'GET', - headers: { - 'Authorization': `Bearer ${accessToken}`, - 'Content-Type': 'application/json', - }, - }); - - if (response.ok) { - const data = await response.json(); - if (data.values?.logo_url) { - setLogoUrl(data.values.logo_url); - } + try { + const proxyBaseUrl = getProxyBaseUrl(); + const url = proxyBaseUrl ? `${proxyBaseUrl}/get/ui_theme_settings` : '/get/ui_theme_settings'; + const response = await fetch(url, { + method: 'GET', + headers: { + 'Content-Type': 'application/json', + }, + }); + + if (response.ok) { + const data = await response.json(); + if (data.values?.logo_url) { + setLogoUrl(data.values.logo_url); } - } catch (error) { - console.warn('Failed to load logo settings from backend:', error); } + } catch (error) { + console.warn('Failed to load logo settings from backend:', error); } }; - + loadLogoSettings(); - }, [accessToken]); + }, []); return ( diff --git a/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx new file mode 100644 index 00000000000..de768501ffd --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx @@ -0,0 +1,114 @@ +import { useState, useEffect } from "react" +import { testMCPToolsListRequest } from "../components/networking" + +interface MCPServerConfig { + server_id?: string + server_name?: string + url?: string + transport?: string + auth_type?: string + mcp_info?: any +} + +interface UseTestMCPConnectionProps { + accessToken: string | null + formValues: Record + enabled?: boolean // Optional flag to enable/disable auto-fetching +} + +interface UseTestMCPConnectionReturn { + tools: any[] + isLoadingTools: boolean + toolsError: string | null + hasShownSuccessMessage: boolean + canFetchTools: boolean + fetchTools: () => Promise + clearTools: () => void +} + +export const useTestMCPConnection = ({ + accessToken, + formValues, + enabled = true, +}: UseTestMCPConnectionProps): UseTestMCPConnectionReturn => { + const [tools, setTools] = useState([]) + const [isLoadingTools, setIsLoadingTools] = useState(false) + const [toolsError, setToolsError] = useState(null) + const [hasShownSuccessMessage, setHasShownSuccessMessage] = useState(false) + + // Check if we have the minimum required fields to fetch tools + const canFetchTools = !!(formValues.url && formValues.transport && formValues.auth_type && accessToken) + + const fetchTools = async () => { + if (!accessToken || !formValues.url) { + return + } + + setIsLoadingTools(true) + setToolsError(null) + + try { + // Prepare the MCP server config from form values + const mcpServerConfig: MCPServerConfig = { + server_id: formValues.server_id || "", + server_name: formValues.server_name || "", + url: formValues.url, + transport: formValues.transport, + auth_type: formValues.auth_type, + mcp_info: formValues.mcp_info, + } + + const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig) + + if (toolsResponse.tools && !toolsResponse.error) { + setTools(toolsResponse.tools) + setToolsError(null) + if (toolsResponse.tools.length > 0 && !hasShownSuccessMessage) { + setHasShownSuccessMessage(true) + } + } else { + const errorMessage = toolsResponse.message || "Failed to retrieve tools list" + setToolsError(errorMessage) + setTools([]) + setHasShownSuccessMessage(false) + } + } catch (error) { + console.error("Tools fetch error:", error) + setToolsError(error instanceof Error ? error.message : String(error)) + setTools([]) + setHasShownSuccessMessage(false) + } finally { + setIsLoadingTools(false) + } + } + + const clearTools = () => { + setTools([]) + setToolsError(null) + setHasShownSuccessMessage(false) + } + + // Auto-fetch tools when form values change and required fields are available + useEffect(() => { + if (!enabled) { + return + } + + if (canFetchTools) { + fetchTools() + } else { + clearTools() + } + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [formValues.url, formValues.transport, formValues.auth_type, accessToken, enabled, canFetchTools]) + + return { + tools, + isLoadingTools, + toolsError, + hasShownSuccessMessage, + canFetchTools, + fetchTools, + clearTools, + } +} diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts new file mode 100644 index 00000000000..8adc94a89e4 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts @@ -0,0 +1,75 @@ +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { clearTokenCookies, getCookie } from './cookieUtils'; + +describe('cookieUtils', () => { + beforeEach(() => { + document.cookie.split(";").forEach((c) => { + document.cookie = c + .replace(/^ +/, "") + .replace(/=.*/, "=;expires=" + new Date().toUTCString() + ";path=/"); + }); + + vi.spyOn(console, 'log').mockImplementation(() => {}); + }); + + describe('clearTokenCookies', () => { + it('should clear token cookie from root path', () => { + document.cookie = 'token=test-token-value; path=/'; + expect(getCookie('token')).toBe('test-token-value'); + + clearTokenCookies(); + expect(getCookie('token')).toBeNull(); + }); + + it('should clear token cookie from /ui path', () => { + document.cookie = 'token=test-token-value; path=/ui'; + clearTokenCookies(); + expect(getCookie('token')).toBeNull(); + }); + + it('should clear token cookies with different SameSite values', () => { + document.cookie = 'token=test-lax; path=/; SameSite=Lax'; + clearTokenCookies(); + expect(getCookie('token')).toBeNull(); + + document.cookie = 'token=test-strict; path=/; SameSite=Strict'; + clearTokenCookies(); + expect(getCookie('token')).toBeNull(); + }); + + it('should handle multiple clearing attempts', () => { + document.cookie = 'token=test-value; path=/'; + + clearTokenCookies(); + clearTokenCookies(); + clearTokenCookies(); + + expect(getCookie('token')).toBeNull(); + }); + }); + + describe('getCookie', () => { + it('should return cookie value when it exists', () => { + document.cookie = 'token=my-test-token; path=/'; + expect(getCookie('token')).toBe('my-test-token'); + }); + + it('should return null when cookie does not exist', () => { + expect(getCookie('nonexistent')).toBeNull(); + }); + + it('should handle JWT tokens with special characters', () => { + const jwt = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoidGVzdCJ9.signature'; + document.cookie = `token=${jwt}; path=/`; + expect(getCookie('token')).toBe(jwt); + }); + + it('should return only the specified cookie', () => { + document.cookie = 'token=token-value; path=/'; + document.cookie = 'other=other-value; path=/'; + + expect(getCookie('token')).toBe('token-value'); + expect(getCookie('other')).toBe('other-value'); + }); + }); +}); diff --git a/ui/litellm-dashboard/tests/top_key_view.test.tsx b/ui/litellm-dashboard/tests/top_key_view.test.tsx new file mode 100644 index 00000000000..c96edcd0b2f --- /dev/null +++ b/ui/litellm-dashboard/tests/top_key_view.test.tsx @@ -0,0 +1,306 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { renderWithProviders, screen, fireEvent } from './test-utils'; +import TopKeyView from '../src/components/top_key_view'; +import { TagUsage } from '../src/components/usage/types'; + +// Mock the networking module +vi.mock('../src/components/networking', () => ({ + keyInfoV1Call: vi.fn(), +})); + +// Mock the transform function +vi.mock('../src/components/key_team_helpers/transform_key_info', () => ({ + transformKeyInfo: vi.fn((data) => data), +})); + +describe('TopKeyView', () => { + const mockProps = { + topKeys: [], + accessToken: 'test-token', + userID: 'test-user', + userRole: 'admin', + teams: null, + premiumUser: true, + showTags: false + }; + + const mockKeysWithTags = [ + { + api_key: 'key-1', + key_alias: 'Production Key', + tags: [ + { tag: 'production', usage: 0.005 } as TagUsage, // <$0.01 + { tag: 'high-volume', usage: 125.50 } as TagUsage, // High spend + { tag: 'api-calls', usage: 0.003 } as TagUsage, // <$0.01 + ], + spend: 125.50 + }, + { + api_key: 'key-2', + key_alias: 'Staging Key', + tags: [ + { tag: 'staging', usage: 45.75 } as TagUsage, // Medium spend + { tag: 'testing', usage: 0.008 } as TagUsage, // <$0.01 + { tag: 'development', usage: 12.25 } as TagUsage, // Low spend + ], + spend: 58.00 + }, + { + api_key: 'key-3', + key_alias: 'Development Key', + tags: [ + { tag: 'dev', usage: 0.002 } as TagUsage, // <$0.01 + { tag: 'experimental', usage: 0.001 } as TagUsage, // <$0.01 + ], + spend: 0.003 + } + ]; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('Tags Column Visibility', () => { + it('should not show tags column when showTags is false', () => { + renderWithProviders(); + expect(screen.queryByText('Tags')).not.toBeInTheDocument(); + }); + + it('should show tags column when showTags is true', () => { + renderWithProviders(); + expect(screen.getByText('Tags')).toBeInTheDocument(); + }); + }); + + describe('Tags Display and Sorting', () => { + beforeEach(() => { + renderWithProviders(); + }); + + it('should display tags for each key', () => { + // Check that tags are displayed (truncated to 7 chars + ...) + expect(screen.getByText('product...')).toBeInTheDocument(); + expect(screen.getByText('high-vo...')).toBeInTheDocument(); + expect(screen.getByText('staging...')).toBeInTheDocument(); + }); + + it('should display top 2 tags by default (sorted by spend)', () => { + // Only the top 2 tags by spend should be visible initially + // Production Key: high-volume (125.50), production (0.005) - sorted by spend + expect(screen.getByText('high-vo...')).toBeInTheDocument(); + expect(screen.getByText('product...')).toBeInTheDocument(); + + // Staging Key: staging (45.75), development (12.25) - sorted by spend + expect(screen.getByText('staging...')).toBeInTheDocument(); + expect(screen.getByText('develop...')).toBeInTheDocument(); + + // Development Key: dev (0.002), experimental (0.001) - sorted by spend + expect(screen.getByText('dev...')).toBeInTheDocument(); + expect(screen.getByText('experim...')).toBeInTheDocument(); + + // These should NOT be visible initially (3rd+ tags) + expect(screen.queryByText('api-cal...')).not.toBeInTheDocument(); + expect(screen.queryByText('testing...')).not.toBeInTheDocument(); + }); + + it('should show expand/collapse arrows for keys with more than 2 tags', () => { + // Production Key has 3 tags, so it should have an expand arrow + // Look for the chevron down icon (expand button) + const expandButtons = screen.getAllByRole('button'); + const expandButton = expandButtons.find(button => + button.getAttribute('title') === 'Show all tags' + ); + expect(expandButton).toBeInTheDocument(); + }); + + it('should show all tags when expanded', async () => { + // Find and click the expand button for Production Key (has 3 tags) + const expandButtons = screen.getAllByRole('button'); + const expandButton = expandButtons.find(button => + button.getAttribute('title') === 'Show all tags' + ); + + if (expandButton) { + fireEvent.click(expandButton); + + // Now all tags should be visible + expect(screen.getByText('api-cal...')).toBeInTheDocument(); + } + }); + + it('should show tooltip on hover with tag information', async () => { + // Hover over a tag to trigger tooltip + const tagElement = screen.getByText('high-vo...'); + fireEvent.mouseOver(tagElement); + + // Check that tooltip content appears + // Note: The exact tooltip content depends on your tooltip implementation + // You might need to adjust this based on how Ant Design Tooltip renders + }); + }); + + describe('Tag Spend Formatting', () => { + it('should handle high spend amounts', () => { + renderWithProviders(); + + // Test that high spend amounts are displayed + // This would require checking tooltip content or finding a way to access the formatted values + expect(screen.getByText('high-vo...')).toBeInTheDocument(); + }); + + it('should handle micro spend amounts', () => { + renderWithProviders(); + + // Test that very small amounts are displayed (only the top 2 tags are visible by default) + // product... has usage: 0.005 (<$0.01) + expect(screen.getByText('product...')).toBeInTheDocument(); + + // dev... has usage: 0.002 (<$0.01) + expect(screen.getByText('dev...')).toBeInTheDocument(); + + // experimental... has usage: 0.001 (<$0.01) + expect(screen.getByText('experim...')).toBeInTheDocument(); + }); + }); + + describe('Edge Cases', () => { + it('should handle keys with no tags', () => { + const keysWithoutTags = [ + { + api_key: 'key-no-tags', + key_alias: 'No Tags Key', + tags: [], + spend: 10.00 + } + ]; + + renderWithProviders(); + expect(screen.getByText('-')).toBeInTheDocument(); + }); + + it('should handle keys with undefined tags', () => { + const keysWithUndefinedTags = [ + { + api_key: 'key-undefined-tags', + key_alias: 'Undefined Tags Key', + tags: undefined, + spend: 5.00 + } + ]; + + renderWithProviders(); + expect(screen.getByText('-')).toBeInTheDocument(); + }); + + it('should handle keys with null tags', () => { + const keysWithNullTags = [ + { + api_key: 'key-null-tags', + key_alias: 'Null Tags Key', + tags: null, + spend: 3.00 + } + ]; + + renderWithProviders(); + expect(screen.getByText('-')).toBeInTheDocument(); + }); + }); + + describe('Tag Truncation', () => { + it('should truncate long tag names to 7 characters', () => { + const keysWithLongTags = [ + { + api_key: 'key-long-tags', + key_alias: 'Long Tags Key', + tags: [ + { tag: 'very-long-tag-name', usage: 10.00 } as TagUsage, + { tag: 'short', usage: 5.00 } as TagUsage, + ], + spend: 15.00 + } + ]; + + renderWithProviders(); + + // Should show truncated version + expect(screen.getByText('very-lo...')).toBeInTheDocument(); + // Short tags should still be truncated (all tags get ...) + expect(screen.getByText('short...')).toBeInTheDocument(); + }); + }); + + describe('Multiple Keys with Different Tag Spend Patterns', () => { + it('should handle mixed spend patterns across multiple keys', () => { + const mixedSpendKeys = [ + { + api_key: 'key-mixed-1', + key_alias: 'Mixed Key 1', + tags: [ + { tag: 'expensive', usage: 999.99 } as TagUsage, + { tag: 'cheap', usage: 0.001 } as TagUsage, + ], + spend: 1000.00 + }, + { + api_key: 'key-mixed-2', + key_alias: 'Mixed Key 2', + tags: [ + { tag: 'moderate', usage: 50.00 } as TagUsage, + { tag: 'tiny', usage: 0.005 } as TagUsage, + ], + spend: 50.01 + } + ]; + + renderWithProviders(); + + // Verify that all tag types are displayed + expect(screen.getByText('expensi...')).toBeInTheDocument(); + expect(screen.getByText('cheap...')).toBeInTheDocument(); + expect(screen.getByText('moderat...')).toBeInTheDocument(); + expect(screen.getByText('tiny...')).toBeInTheDocument(); + }); + }); + + describe('Table Structure', () => { + it('should render table with correct headers when showTags is true', () => { + renderWithProviders(); + + expect(screen.getByText('Key ID')).toBeInTheDocument(); + expect(screen.getByText('Key Alias')).toBeInTheDocument(); + expect(screen.getByText('Tags')).toBeInTheDocument(); + expect(screen.getByText('Spend (USD)')).toBeInTheDocument(); + }); + + it('should render table with correct headers when showTags is false', () => { + renderWithProviders(); + + expect(screen.getByText('Key ID')).toBeInTheDocument(); + expect(screen.getByText('Key Alias')).toBeInTheDocument(); + expect(screen.queryByText('Tags')).not.toBeInTheDocument(); + expect(screen.getByText('Spend (USD)')).toBeInTheDocument(); + }); + }); + + describe('Key Data Display', () => { + it('should display key information correctly', () => { + const simpleKeys = [ + { + api_key: 'test-key-123', + key_alias: 'Test Key', + tags: [], + spend: 25.50 + } + ]; + + renderWithProviders(); + + // Check that key alias is displayed + expect(screen.getByText('Test Key')).toBeInTheDocument(); + + // Check that spend is formatted correctly + expect(screen.getByText('$25.50')).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/vitest.config.ts b/ui/litellm-dashboard/vitest.config.ts index d24c0d7f2a5..69b65c5e199 100644 --- a/ui/litellm-dashboard/vitest.config.ts +++ b/ui/litellm-dashboard/vitest.config.ts @@ -14,4 +14,11 @@ export default defineConfig({ '@': resolve(__dirname, 'src'), }, }, + define: { + 'import.meta.vitest': 'undefined', + }, + esbuild: { + jsx: 'automatic', + jsxImportSource: 'react', + }, })