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/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/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index 6a11d069fb0..854d781f546 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 ``` 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 902d0ffedba..6364b8c4444 100644 --- a/docs/my-website/docs/proxy/logging_spec.md +++ b/docs/my-website/docs/proxy/logging_spec.md @@ -163,17 +163,18 @@ A literal type with two possible values: ## StandardLoggingGuardrailInformation -| 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", "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 | +| 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 diff --git a/docs/my-website/docs/proxy/native_litellm_prompt.md b/docs/my-website/docs/proxy/native_litellm_prompt.md index ea326d00690..34edb66fc40 100644 --- a/docs/my-website/docs/proxy/native_litellm_prompt.md +++ b/docs/my-website/docs/proxy/native_litellm_prompt.md @@ -9,7 +9,7 @@ Store prompts as `.prompt` files in your repository and use them directly with L - **File System**: Store `.prompt` files locally - **BitBucket**: Store `.prompt` files in BitBucket repositories with team-based access control - +- **Gitlab**: Store `.prompt` files in Gitlab repositories with team-based access control ## Quick Start @@ -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/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/__init__.py b/litellm/__init__.py index 328b5a6d89b..43ae1750cd1 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", ] @@ -1060,6 +1061,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 @@ -1160,6 +1162,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 @@ -1358,3 +1361,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/constants.py b/litellm/constants.py index 3ff9a4b6fb0..318b23c72ce 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -374,7 +374,7 @@ OPENAI_TRANSCRIPTION_PARAMS = [ "timestamp_granularities", ] -OPENAI_EMBEDDING_PARAMS = ["dimensions", "encoding_format", "user"] +OPENAI_EMBEDDING_PARAMS = ["dimensions", "encoding_format", "user", "input_type"] DEFAULT_EMBEDDING_PARAM_VALUES = { **{k: None for k in OPENAI_EMBEDDING_PARAMS}, 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..a091633d7e2 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", {}) + 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", {}) + 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 265e1eccb4d..b5ab5aeefe3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3669,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( 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..89b4e1b0866 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -445,22 +445,25 @@ class BedrockModelInfo(BaseLLMModelInfo): @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..0d25440cd72 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -1,33 +1,46 @@ """ -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 from litellm.types.llms.bedrock import ( + 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 +54,140 @@ 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} - ) + 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) + """ + if inference_params.get("inputType"): + input_type = inference_params["inputType"] else: - # Text input - transformed_request = TwelveLabsMarengoEmbeddingRequest( - inputType="text", inputText=input + raise ValueError("input_type is required") + + # 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", "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 +246,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/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 5371b9a4b61..6586a83f06d 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, @@ -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/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/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index d8569b01b83..3bd5a4847f3 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} 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/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 877df1bb780..521251f50cd 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, @@ -12830,9 +12896,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": [ @@ -18347,6 +18413,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", @@ -19662,6 +19742,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, @@ -21028,6 +21112,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, @@ -21050,6 +21138,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/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_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index fb7ada8ab10..60d4e32ebbd 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -289,8 +289,8 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 def get_model_group_from_litellm_kwargs(kwargs: dict) -> Optional[str]: _litellm_params = kwargs.get("litellm_params", None) or {} - _metadata = _litellm_params.get(get_metadata_variable_name_from_kwargs(kwargs)) or {} - _model_group = _metadata.get("model_group", None) + _metadata = _litellm_params.get(get_metadata_variable_name_from_litellm_params(_litellm_params)) or {} + _model_group = _metadata.get("model_group", None) or kwargs.get("model", None) if _model_group is not None: return _model_group @@ -367,8 +367,8 @@ def add_guardrail_to_applied_guardrails_header( _metadata["applied_guardrails"] = [guardrail_name] -def get_metadata_variable_name_from_kwargs( - kwargs: dict +def get_metadata_variable_name_from_litellm_params( + litellm_params: dict ) -> Literal["metadata", "litellm_metadata"]: """ Helper to return what the "metadata" field should be called in the request data @@ -381,4 +381,4 @@ def get_metadata_variable_name_from_kwargs( - OpenAI then started using this field for their metadata - LiteLLM is now moving to using `litellm_metadata` for our metadata """ - return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" + return "litellm_metadata" if "litellm_metadata" in litellm_params else "metadata" diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index f498d647a5e..c88ebe16d99 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -727,6 +727,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) return + outputs: List[BedrockGuardrailOutput] = ( + response.get("outputs", []) or [] + ) + if not any(output.get("text") for output in outputs): + 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/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..5d0157f4361 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,18 @@ 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. Reuses v3 limiter's Redis-based tracking (works across multiple instances) + 2. Only enforces priority limits when model is saturated (>80% usage) + 3. When under capacity, allows all requests (generous behavior) + 4. When saturated, enforces strict priority-based limits (fairness) + + How it works: + - Uses v3 limiter's counter keys to check model-wide saturation + - Saturation check reads existing counters without incrementing + - Priority enforcement reuses v3 limiter's atomic Lua scripts """ def __init__(self, internal_usage_cache: DualCache): self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache) @@ -57,6 +63,107 @@ 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 + + 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 +171,10 @@ 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. + Only called when system is saturated. """ descriptors: List[RateLimitDescriptor] = [] @@ -79,8 +185,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if model_group_info is None: return descriptors - # Get priority weight - priority_weight = self._get_priority_weight(priority) + # Get normalized priority weight (handles over-allocation) + normalized_weights = self._normalize_priority_weights() + priority_weight = normalized_weights.get(priority, None) if priority else None + if priority_weight is None: + # Fallback to non-normalized weight + priority_weight = self._get_priority_weight(priority) + # Create priority-specific rate limits # Use model:priority as the key to separate different priority levels @@ -88,16 +199,17 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): rate_limit_config: RateLimitDescriptorRateLimitObject = {} - # Apply priority weight to model limits + # Apply normalized priority weight to model limits if model_group_info.tpm is not None: - # Reserve portion of TPM based on priority + # Reserve portion of TPM based on normalized 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 + # Reserve portion of RPM based on normalized priority reserved_rpm = int(model_group_info.rpm * priority_weight) rate_limit_config["requests_per_unit"] = reserved_rpm + if rate_limit_config: rate_limit_config["window_size"] = self.v3_limiter.window_size @@ -112,6 +224,171 @@ 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 _handle_generous_mode( + self, + model: str, + model_group_info: ModelGroupInfo, + user_api_key_dict: UserAPIKeyAuth, + key_priority: Optional[str], + ) -> None: + """ + Handle rate limiting in generous mode (under saturation threshold). + + In this mode, we enforce model-wide capacity but NOT priority-specific limits. + This allows lower-priority users to borrow unused capacity from higher-priority users. + + Args: + model: Model name + model_group_info: Model configuration + user_api_key_dict: User authentication info + key_priority: User's priority level + + Raises: + HTTPException: If model capacity is reached + """ + descriptor = self._create_model_tracking_descriptor( + model=model, + model_group_info=model_group_info, + high_limit_multiplier=1, # Enforce actual limits in generous mode + ) + + response = await self.v3_limiter.should_rate_limit( + descriptors=[descriptor], + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + + if response["overall_code"] == "OVER_LIMIT": + for status in response["statuses"]: + if status["code"] == "OVER_LIMIT": + 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", + }, + ) + + async def _handle_strict_mode( + self, + model: str, + model_group_info: ModelGroupInfo, + user_api_key_dict: UserAPIKeyAuth, + key_priority: Optional[str], + saturation: float, + data: dict, + ) -> None: + """ + Handle rate limiting in strict mode (above saturation threshold). + + In this mode, we enforce priority-specific limits using normalized weights. + + 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 priority-specific limit is exceeded + """ + # Create priority-based descriptors + descriptors = self._create_priority_based_descriptors( + model=model, + user_api_key_dict=user_api_key_dict, + priority=key_priority, + ) + + if not descriptors: + verbose_proxy_logger.debug("No rate limit descriptors created, allowing request") + return + + # Track model-wide usage for future saturation checks + # Why tracking_multiplier: v3_limiter.should_rate_limit() both increments AND checks limits. + # We need the increment (for saturation detection) but NOT the limit check (priority limits handle enforcement). + # Setting limit to 10x capacity ensures tracking never blocks while keeping accurate counters. + tracking_multiplier = litellm.priority_reservation_settings.tracking_multiplier + tracking_descriptor = self._create_model_tracking_descriptor( + model=model, + model_group_info=model_group_info, + high_limit_multiplier=tracking_multiplier, + ) + + await self.v3_limiter.should_rate_limit( + descriptors=[tracking_descriptor], + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + + # Enforce priority-specific limits + response = await self.v3_limiter.should_rate_limit( + descriptors=descriptors, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + + if response["overall_code"] == "OVER_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']}, " + 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%}", + }, + ) + else: + # Store response for post-call hook + data["litellm_proxy_rate_limit_response"] = response + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -130,60 +407,73 @@ 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. + + This hook implements a two-mode rate limiting strategy: + - Generous mode (< 80% saturation): Enforces model capacity, allows priority borrowing + - Strict mode (>= 80% saturation): Enforces normalized priority-based limits + + 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 + # Check current saturation level 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, + 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", - }, - ) + + data["litellm_model_saturation"] = saturation + + # Route to appropriate mode based on saturation + if saturation < saturation_threshold: + await self._handle_generous_mode( + model=model, + model_group_info=model_group_info, + user_api_key_dict=user_api_key_dict, + key_priority=key_priority, + ) else: - # Store response for post-call hook - data["litellm_proxy_rate_limit_response"] = response - + await self._handle_strict_mode( + 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 0a49d7f6759..2e3057cdba3 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -20,6 +20,8 @@ from typing import ( cast, ) +from fastapi import HTTPException + from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger @@ -843,7 +845,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): _get_parent_otel_span_from_kwargs, ) from litellm.proxy.common_utils.callback_utils import ( - get_metadata_variable_name_from_kwargs, + get_metadata_variable_name_from_litellm_params, get_model_group_from_litellm_kwargs, ) from litellm.types.caching import RedisPipelineIncrementOperation @@ -861,7 +863,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Get metadata from kwargs litellm_metadata = kwargs["litellm_params"].get( - get_metadata_variable_name_from_kwargs(kwargs), {} + get_metadata_variable_name_from_litellm_params(kwargs["litellm_params"]), {} ) if litellm_metadata is None: return 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 f25ed8a7bfd..8aa3b90d954 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -57,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 """ @@ -95,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 @@ -183,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 @@ -226,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, @@ -310,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( @@ -327,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, @@ -509,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} @@ -550,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( @@ -604,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, }, ) @@ -669,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, @@ -701,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() @@ -764,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("/"): @@ -829,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, @@ -864,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 @@ -876,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 @@ -888,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) @@ -956,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: @@ -1133,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, @@ -1181,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, @@ -1200,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 @@ -1219,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. """ @@ -1237,14 +1182,9 @@ 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 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..8c96c2ab96a --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -0,0 +1,204 @@ +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 _build_complete_streaming_response( + all_chunks: List[str], + litellm_logging_obj: LiteLLMLoggingObj, + model: str, + url_route: str, + ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: + parsed_chunks = [] + if "generateContent" in url_route or "streamGenerateContent" in url_route: + gemini_iterator: Any = GeminiModelResponseIterator( + streaming_response=None, + sync_stream=False, + logging_obj=litellm_logging_obj, + ) + chunk_parsing_logic: Any = gemini_iterator._common_chunk_parsing_logic + parsed_chunks = [chunk_parsing_logic(chunk) for chunk in all_chunks] + else: + return None + + if len(parsed_chunks) == 0: + return None + + 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) + + 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="gemini", + ) + + 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/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3f28ba92d36..7a0343e2e96 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -96,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") @@ -111,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:] @@ -206,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 @@ -237,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 @@ -254,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) @@ -278,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), @@ -301,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: @@ -432,10 +411,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 @@ -485,9 +462,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 ), ) ) @@ -521,16 +496,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. @@ -581,6 +552,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 @@ -632,9 +604,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 @@ -701,9 +671,7 @@ async def pass_through_request( # noqa: PLR0915 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: @@ -721,11 +689,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: @@ -742,9 +708,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( @@ -766,20 +730,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) @@ -787,9 +747,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( @@ -811,9 +769,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) @@ -835,6 +791,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, ) ) @@ -865,9 +822,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)) ) ######################################################### @@ -930,6 +885,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 @@ -965,16 +921,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 @@ -988,6 +940,7 @@ 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 @@ -1301,7 +1254,7 @@ async def websocket_passthrough_request( # noqa: PLR0915 logging_obj.model_call_details[ "custom_llm_provider" ] = "vertex_ai_language_models" - verbose_proxy_logger.info( + verbose_proxy_logger.debug( f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from server setup response" ) else: @@ -1465,7 +1418,7 @@ async def websocket_passthrough_request( # noqa: PLR0915 if websocket.client_state != WebSocketState.DISCONNECTED: await websocket.close( - code=exc.status_code if hasattr(exc, "status_code") else 1011, + code=getattr(exc, "status_code", 1011), reason="Upstream connection rejected", ) except Exception as e: @@ -1644,15 +1597,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( @@ -1689,9 +1669,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) @@ -1710,9 +1688,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, @@ -1739,9 +1715,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( @@ -1845,11 +1819,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 @@ -1857,9 +1827,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 @@ -1890,13 +1858,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( @@ -1923,9 +1887,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() @@ -1943,9 +1905,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) @@ -1978,9 +1938,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 @@ -1996,21 +1954,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 @@ -2018,9 +1968,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 @@ -2036,9 +1984,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]) @@ -2076,6 +2022,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/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 94517235a0c..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,14 @@ 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"] @@ -81,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, @@ -103,6 +103,7 @@ class PassThroughEndpointLogging: start_time: datetime, end_time: datetime, cache_hit: bool, + custom_llm_provider: Optional[str] = None, **kwargs, ): return_dict = { @@ -110,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 = ( @@ -142,28 +155,22 @@ 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 @@ -172,24 +179,21 @@ class PassThroughEndpointLogging: OpenAIPassthroughLoggingHandler, ) - openai_passthrough_logging_handler_result = ( - OpenAIPassthroughLoggingHandler.openai_passthrough_handler( - httpx_response=httpx_response, - response_body=response_body or {}, - logging_obj=logging_obj, - url_route=url_route, - result=result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, - request_body=request_body, - **kwargs, - ) - ) - standard_logging_response_object = ( - openai_passthrough_logging_handler_result["result"] + 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, @@ -216,6 +220,7 @@ class PassThroughEndpointLogging: return_dict[ "standard_logging_response_object" ] = standard_logging_response_object + return_dict["kwargs"] = kwargs return return_dict @@ -231,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, @@ -263,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, @@ -352,10 +344,16 @@ class PassThroughEndpointLogging: 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 ( @@ -386,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_server.py b/litellm/proxy/proxy_server.py index b1a269b31cd..55e4a96b5b0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -155,7 +155,6 @@ 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, ) - from litellm.proxy._experimental.mcp_server.rest_endpoints import ( router as mcp_rest_endpoints_router, ) @@ -254,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, @@ -301,7 +302,9 @@ from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMi from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config +from litellm.proxy.openai_files_endpoints.files_endpoints import ( + set_files_config, +) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -1874,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, @@ -2608,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 @@ -2615,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 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 9ae87316476..6f0eb51959c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4755,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 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/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 bcf0fa13746..4e93e167530 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -123,12 +123,18 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): max_output_tokens: Required[Optional[int]] input_cost_per_token: Required[float] input_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing - input_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing + input_cost_per_token_priority: Optional[ + float + ] # OpenAI priority service tier pricing cache_creation_input_token_cost: Optional[float] cache_creation_input_token_cost_above_1hr: Optional[float] cache_read_input_token_cost: Optional[float] - cache_read_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing - cache_read_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing + cache_read_input_token_cost_flex: Optional[ + float + ] # OpenAI flex service tier pricing + cache_read_input_token_cost_priority: Optional[ + float + ] # OpenAI priority service tier pricing input_cost_per_character: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_token_above_128k_tokens: Optional[float] # only for vertex ai models @@ -147,7 +153,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_batches: Optional[float] output_cost_per_token: Required[float] output_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing - output_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing + output_cost_per_token_priority: Optional[ + float + ] # OpenAI priority service tier pricing output_cost_per_character: Optional[float] # only for vertex ai models output_cost_per_audio_token: Optional[float] output_cost_per_token_above_128k_tokens: Optional[ @@ -1417,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) @@ -2638,6 +2649,7 @@ class SpecialEnums(Enum): class ServiceTier(Enum): """Enum for service tier types used in cost calculations.""" + FLEX = "flex" PRIORITY = "priority" @@ -2684,13 +2696,24 @@ CostResponseTypes = Union[ class PriorityReservationSettings(BaseModel): """ Settings for priority-based rate limiting reservation. - + Defines what priority to assign to keys without explicit priority metadata. The priority_reservation mapping is configured separately via litellm.priority_reservation. """ + default_priority: float = Field( default=0.5, - description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation." + 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.80, + 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..f963e8c9443 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, ) @@ -2801,6 +2802,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 +7208,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 +7587,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 877df1bb780..521251f50cd 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, @@ -12830,9 +12896,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": [ @@ -18347,6 +18413,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", @@ -19662,6 +19742,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, @@ -21028,6 +21112,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, @@ -21050,6 +21138,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/pyproject.toml b/pyproject.toml index 3a6262d8833..df6b0074911 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" 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/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index b98a1b16bef..c4d1655594b 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1366,3 +1366,59 @@ 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" + ) + + # Mock Bedrock API with no output text + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", + "name": "top_song", + "input": { + "sign": "WZPZ" + } + } + } + ] + } + }, + "stopReason": "tool_use" + } + + data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Hello"}, + ], + } + mock_user_api_key_dict = UserAPIKeyAuth() + + return await guardrail.async_post_call_success_hook( + data=data, + response=mock_bedrock_response, + user_api_key_dict=mock_user_api_key_dict, + ) + # If no error is raised, then the test passes + print("✅ No output text in response test passed") \ 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..903fd310262 100644 --- a/tests/llm_translation/test_bedrock_embedding.py +++ b/tests/llm_translation/test_bedrock_embedding.py @@ -170,6 +170,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/abc123", "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_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/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/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/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/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..b43ec226842 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,156 @@ 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 throws an error for TwelveLabs models but not others""" + litellm.set_verbose = True + client = HTTPHandler() + test_api_key = "test-bearer-token-12345" + + # Test TwelveLabs model - should throw error + 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 throws an error for TwelveLabs + with pytest.raises(Exception) as exc_info: + 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 throw an error + ) + + # Verify the error message contains the expected text + assert "input_type is required" in str(exc_info.value) + + # 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/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/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 05e0d4287a3..2b7c080e65d 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,10 +401,39 @@ async def test_concurrent_pre_call_hooks_stress(): successful_requests = [] rate_limited_requests = [] + # 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 + async def mock_should_rate_limit(descriptors, parent_otel_span=None): - """Mock rate limiter that allows premium users, limits some standard users.""" + """Mock rate limiter that handles saturation-aware descriptors.""" descriptor = descriptors[0] - priority = descriptor["value"].split(":")[-1] + 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, + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 10000, + } + ], + } + + # Handle priority-specific enforcement in strict mode + if descriptor_key == "priority_model": + # Extract priority from value like "pre-call-stress-model:premium" + priority = descriptor_value.split(":")[-1] if priority == "premium": # Allow all premium requests @@ -410,7 +442,7 @@ async def test_concurrent_pre_call_hooks_stress(): "statuses": [ { "code": "OK", - "descriptor_key": descriptor["value"], + "descriptor_key": descriptor_value, "rate_limit_type": "tokens_per_unit", "limit_remaining": 1000, } @@ -426,7 +458,7 @@ async def test_concurrent_pre_call_hooks_stress(): "statuses": [ { "code": "OVER_LIMIT", - "descriptor_key": descriptor["value"], + "descriptor_key": descriptor_value, "rate_limit_type": "tokens_per_unit", "limit_remaining": 0, } @@ -438,9 +470,22 @@ async def test_concurrent_pre_call_hooks_stress(): "statuses": [ { "code": "OK", - "descriptor_key": descriptor["value"], + "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, } ], } @@ -466,6 +511,8 @@ async def test_concurrent_pre_call_hooks_stress(): 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 ): try: result = await handler.async_pre_call_hook( @@ -534,7 +581,7 @@ 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%}" + ), f"Standard success rate should be >= 50% (with 30% random limiting, allows for variance), got {standard_success_rate:.2%}" assert ( premium_success_rate > standard_success_rate ), "Premium should have higher success rate than standard" @@ -550,3 +597,608 @@ 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: No Rate Limiting When At Capacity + + System: 100 RPM capacity + Key A: priority_reservation=0.75 (75 RPM reserved) + Key B: priority_reservation=0.25 (25 RPM reserved) + Traffic A: 50 RPM + Traffic B: 50 RPM + Expected A: 50 RPM (no limiting, under reserved capacity) + Expected B: 50 RPM (no limiting, under reserved capacity) + + When traffic is under individual reservations, no rate limiting should occur. + """ + 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 50 requests from each priority (within capacity) + tasks = [] + + for i in range(50): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(50): + 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 - No Rate Limiting When At Capacity:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A: {successful_requests['key_a']}/50 successful (reserved 75 RPM)") + print(f" - Key B: {successful_requests['key_b']}/50 successful (reserved 25 RPM)") + print(f" - Total successful: {total_successful}/100") + print(f" - Total rate limited: {total_rate_limited}/100") + + # Both keys should get all their requests since they're under capacity + assert successful_requests["key_a"] >= 45, f"Key A should get ≥45 requests, got {successful_requests['key_a']}" + assert successful_requests["key_b"] >= 45, f"Key B should get ≥45 requests, got {successful_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%)") 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..6f87d8f6ab5 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py @@ -0,0 +1,287 @@ +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" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 1cbe6420f6e..63436899fcc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1971,3 +1971,62 @@ async def test_model_info_v1_oci_secrets_not_leaked(): 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/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/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/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/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/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', + }, })