Merge branch 'BerriAI:main' into bugfix-14404-image-gen-azure-managed-identity

README.md

@@ -273,7 +273,7 @@ echo 'LITELLM_SALT_KEY="sk-1234"' >> .env
 source .env
 
 # Start
-docker-compose up
+docker compose up
 ```
 
 

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

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)
 

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://<user>:<password>@<host>:<port>/<dbname>"
 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

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

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/)

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

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).
+
+</TabItem>
+</Tabs>

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

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
+
+<Tabs>
+<TabItem value="llama-1b" label="LLaMa 1B Model">
+
+```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)
+```
+
+</TabItem>
+<TabItem value="mistral-4b" label="Mistral 4B Model">
+
+```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)
+```
+
+</TabItem>
+</Tabs>
+
+**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:
+
+<Tabs>
+<TabItem value="env" label="Environment Variable">
+
+```bash
+export NVIDIA_NIM_API_KEY="nvapi-..."
+```
+
+</TabItem>
+<TabItem value="python" label="Python">
+
+```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-...",
+)
+```
+
+</TabItem>
+</Tabs>
+
+## 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/)
+

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
 ```
 
 

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
 ```
 
 </TabItem>

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
 

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
 
 <Tabs>
@@ -90,6 +90,51 @@ response = litellm.completion(
 ```
 
 </TabItem>
+<TabItem value="gitlab" label="GITLAB">
+
+**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?"}
+)
+```
+
+</TabItem>
+
 <TabItem value="proxy" label="PROXY">
 
 **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/<base_model>        # 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/<repo_name>",
+        "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
+    }
+)
 ```

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",

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

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()

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},

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/<repo_name>",
+    "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/<repo_name>",
+    "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/<base_model>",  # 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.

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",
+]

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()

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,
+        )

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))
 

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):

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,
     }

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(

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(

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
         )

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}"
+            )

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,
+        )

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:

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:

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,

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

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:

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():

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}
 

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,
+        )
+

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,

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:

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"

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 ##########
         #########################################################

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
+
+

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

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

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}")

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/<Model>/<endpoint>. Got: "
-                + endpoint,
+                "error": "Model missing from endpoint. Expected format: /model/<Model>/<endpoint>. 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
 

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

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)

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

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()

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

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

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

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
 

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]

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,

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):

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=())

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],

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,

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"

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"
     

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")

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`.

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

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

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

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/<repo_name>",
+            "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

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():

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

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"
+    )

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")

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)
 

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()

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"

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"]
+        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)

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)

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%)")

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"

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()

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",

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",

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<EntityUsageProps> = ({
   };
 
   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<EntityUsageProps> = ({
             },
             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<EntityUsageProps> = ({
       .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<EntityUsageProps> = ({
                     userRole={userRole}
                     teams={null}
                     premiumUser={premiumUser}
+                    showTags={entityType === "tag"}
                   />
                 </Card>
               </Col>

ui/litellm-dashboard/src/components/new_usage.tsx

@@ -269,6 +269,7 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({ 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<NewUsagePageProps> = ({ 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)

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<TopKeyViewProps> = ({ topKeys, accessToken, userID, userRole, teams, premiumUser }) => {
+const TopKeyView: React.FC<TopKeyViewProps> = ({ topKeys, accessToken, userID, userRole, teams, premiumUser, showTags = false }) => {
   const [isModalOpen, setIsModalOpen] = useState<boolean>(false)
   const [selectedKey, setSelectedKey] = useState<string | null>(null)
   const [keyData, setKeyData] = useState<any | undefined>(undefined)
   const [viewMode, setViewMode] = useState<"chart" | "table">("table")
+  const [expandedTags, setExpandedTags] = useState<Set<string>>(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<TopKeyViewProps> = ({ 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<TopKeyViewProps> = ({ 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 (
+        <div className="overflow-hidden">
+          <div className="flex flex-wrap items-center gap-1">
+            {displayTags.map((tag, index) => (
+              <Tooltip 
+                key={index} 
+                title={
+                  <div>
+                    <div><span className="text-gray-300">Tag Name:</span> {tag.tag}</div>
+                    <div><span className="text-gray-300">Spend:</span> {tag.usage > 0 && tag.usage < 0.01 ? '<$0.01' : `$${formatNumberWithCommas(tag.usage, 2)}`}</div>
+                  </div>
+                }
+              >
+                <span className="px-2 py-1 bg-gray-100 rounded-full text-xs">
+                  {tag.tag.slice(0, 7)}...
+                </span>
+              </Tooltip>
+            ))}
+            {hasMoreTags && (
+              <button
+                onClick={() => toggleTagsExpansion(apiKey)}
+                className="ml-1 p-1 hover:bg-gray-200 rounded-full transition-colors"
+                title={isExpanded ? "Show fewer tags" : "Show all tags"}
+              >
+                {isExpanded ? (
+                  <ChevronUpIcon className="h-3 w-3 text-gray-500" />
+                ) : (
+                  <ChevronDownIcon className="h-3 w-3 text-gray-500" />
+                )}
+              </button>
+            )}
+          </div>
+        </div>
+      );
+    }
+  }
+
+  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 || "-",

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
+}

ui/litellm-dashboard/src/components/view_logs/table.tsx

@@ -53,7 +53,7 @@ export function DataTable<TData, TValue>({
 
   return (
     <div className="rounded-lg custom-border overflow-x-auto w-full max-w-full box-border">
-      <Table className="[&_td]:py-0.5 [&_th]:py-1 table-fixed w-full box-border" style={{minWidth: '800px'}}>
+      <Table className="[&_td]:py-0.5 [&_th]:py-1 table-fixed w-full box-border" style={{minWidth: '400px'}}>
         <TableHead>
           {table.getHeaderGroups().map((headerGroup) => (
             <TableRow key={headerGroup.id}>

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(<TopKeyView {...mockProps} showTags={false} />);
+      expect(screen.queryByText('Tags')).not.toBeInTheDocument();
+    });
+
+    it('should show tags column when showTags is true', () => {
+      renderWithProviders(<TopKeyView {...mockProps} showTags={true} />);
+      expect(screen.getByText('Tags')).toBeInTheDocument();
+    });
+  });
+
+  describe('Tags Display and Sorting', () => {
+    beforeEach(() => {
+      renderWithProviders(<TopKeyView {...mockProps} topKeys={mockKeysWithTags} showTags={true} />);
+    });
+
+    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(<TopKeyView {...mockProps} topKeys={mockKeysWithTags} showTags={true} />);
+      
+      // 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(<TopKeyView {...mockProps} topKeys={mockKeysWithTags} showTags={true} />);
+      
+      // 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(<TopKeyView {...mockProps} topKeys={keysWithoutTags} showTags={true} />);
+      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(<TopKeyView {...mockProps} topKeys={keysWithUndefinedTags} showTags={true} />);
+      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(<TopKeyView {...mockProps} topKeys={keysWithNullTags} showTags={true} />);
+      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(<TopKeyView {...mockProps} topKeys={keysWithLongTags} showTags={true} />);
+      
+      // 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(<TopKeyView {...mockProps} topKeys={mixedSpendKeys} showTags={true} />);
+      
+      // 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(<TopKeyView {...mockProps} showTags={true} />);
+      
+      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(<TopKeyView {...mockProps} showTags={false} />);
+      
+      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(<TopKeyView {...mockProps} topKeys={simpleKeys} showTags={true} />);
+      
+      // Check that key alias is displayed
+      expect(screen.getByText('Test Key')).toBeInTheDocument();
+      
+      // Check that spend is formatted correctly
+      expect(screen.getByText('$25.50')).toBeInTheDocument();
+    });
+  });
+});

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',
+  },
 })