diff --git a/docs/my-website/blog/gemin_3.1/index.md b/docs/my-website/blog/gemin_3.1/index.md new file mode 100644 index 00000000000..b81595e4bd5 --- /dev/null +++ b/docs/my-website/blog/gemin_3.1/index.md @@ -0,0 +1,150 @@ +--- +slug: gemini_3_1_pro +title: "DAY 0 Support: Gemini 3.1 Pro on LiteLLM" +date: 2026-02-19T10:00:00 +authors: + - name: Sameer Kankute + title: SWE @ LiteLLM (LLM Translation) + url: https://www.linkedin.com/in/sameer-kankute/ + image_url: https://pbs.twimg.com/profile_images/2001352686994907136/ONgNuSk5_400x400.jpg + - name: Krrish Dholakia + title: "CEO, LiteLLM" + url: https://www.linkedin.com/in/krish-d/ + image_url: https://pbs.twimg.com/profile_images/1298587542745358340/DZv3Oj-h_400x400.jpg + - name: Ishaan Jaff + title: "CTO, LiteLLM" + url: https://www.linkedin.com/in/reffajnaahsi/ + image_url: https://pbs.twimg.com/profile_images/1613813310264340481/lz54oEiB_400x400.jpg +description: "Guide to using Gemini 3.1 Pro on LiteLLM Proxy and SDK with day 0 support." +tags: [gemini, day 0 support, llms] +hide_table_of_contents: false +--- + + +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Gemini 3.1 Pro Day 0 Support + +LiteLLM now supports `gemini-3.1-pro-preview` and all the new API changes along with it. + +## Deploy this version + + + + +``` showLineNumbers title="docker run litellm" +docker run \ +-e STORE_MODEL_IN_DB=True \ +-p 4000:4000 \ +ghcr.io/berriai/litellm:main-v1.81.9-stable.gemini.3.1-pro +``` + + + + + +``` showLineNumbers title="pip install litellm" +pip install litellm==v1.81.9-stable.gemini.3.1-pro +``` + + + + +## What's New + +### 1. New Thinking Levels: `thinkingLevel` with MINIMAL & MEDIUM + +Gemini 3.1 Pro introduces support for **medium** thinking level + +LiteLLM automatically maps the OpenAI `reasoning_effort` parameter to Gemini's `thinkingLevel`, so you can use familiar `reasoning_effort` values (`minimal`, `low`, `medium`, `high`) without changing your code! + +--- +## Supported Endpoints + +LiteLLM provides **full end-to-end support** for Gemini 3.1 Pro on: + +- ✅ `/v1/chat/completions` - OpenAI-compatible chat completions endpoint +- ✅ `/v1/responses` - OpenAI Responses API endpoint (streaming and non-streaming) +- ✅ [`/v1/messages`](../../docs/anthropic_unified) - Anthropic-compatible messages endpoint +- ✅ `/v1/generateContent` – [Google Gemini API](../../docs/generateContent.md) compatible endpoint + +All endpoints support: +- Streaming and non-streaming responses +- Function calling with thought signatures +- Multi-turn conversations +- All Gemini 3-specific features +- Conversion of provider specific thinking related param to thinkingLevel + +## Quick Start + + + + +**Basic Usage with MEDIUM thinking (NEW)** + +```python +from litellm import completion + +# No need to make any changes to your code as we map openai reasoning param to thinkingLevel +response = completion( + model="gemini/gemini-3.1-pro-preview", + messages=[{"role": "user", "content": "Solve this complex math problem: 25 * 4 + 10"}], + reasoning_effort="medium", # NEW: MEDIUM thinking level +) + +print(response.choices[0].message.content) +``` + + + + + +**1. Setup config.yaml** + +```yaml +model_list: + - model_name: gemini-3.1-pro-preview + litellm_params: + model: gemini/gemini-3.1-pro-preview + api_key: os.environ/GEMINI_API_KEY + - model_name: vertex-gemini-3.1-pro-preview + litellm_params: + model: vertex_ai/gemini-3.1-pro-preview +``` + +**2. Start proxy** + +```bash +litellm --config /path/to/config.yaml +``` + +**3. Call with MEDIUM thinking** + +```bash +curl -X POST http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer " \ + -d '{ + "model": "gemini-3.1-pro-preview", + "messages": [{"role": "user", "content": "Complex reasoning task"}], + "reasoning_effort": "medium" + }' +``` + + + + +--- + +## `reasoning_effort` Mapping for Gemini 3+ + +| reasoning_effort | thinking_level | +|------------------|----------------| +| `minimal` | `minimal` | +| `low` | `low` | +| `medium` | `medium` | +| `high` | `high` | +| `disable` | `minimal` | +| `none` | `minimal` | + diff --git a/docs/my-website/docs/completion/usage.md b/docs/my-website/docs/completion/usage.md index c388e5bfee1..d610afeae55 100644 --- a/docs/my-website/docs/completion/usage.md +++ b/docs/my-website/docs/completion/usage.md @@ -50,3 +50,51 @@ for chunk in completion: print(chunk.choices[0].delta) ``` + +### Proxy: Always Include Streaming Usage + +When using the LiteLLM Proxy, you can configure it to automatically include usage information in all streaming responses, even if the client doesn't send `stream_options={"include_usage": True}`. + +#### Configuration + +Add the following to your config.yaml: + +```yaml +general_settings: + always_include_stream_usage: true +``` + +Alternatively, configure it through the UI: + +1. Navigate to the LiteLLM Proxy UI +2. Go to `Settings` > `Router Settings` > `General` +3. Find the `always_include_stream_usage` setting +4. Toggle it to `true` +5. Click `Update` to save + +#### How it works + +When `always_include_stream_usage` is enabled: +- All streaming requests will automatically have `stream_options={"include_usage": True}` added +- Clients will receive usage information in the final chunk, even if they didn't explicitly request it +- If a client already provides `stream_options`, `include_usage: True` will be added without overwriting other options +- Non-streaming requests are not affected + +#### Example + +With this setting enabled, a simple streaming request like: + +```bash +curl -X POST http://localhost:4000/v1/chat/completions \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hello!"}], + "stream": true + }' +``` + +Will automatically receive usage information in the response, without needing to explicitly include `stream_options`. + +``` diff --git a/docs/my-website/docs/proxy/pass_through.md b/docs/my-website/docs/proxy/pass_through.md index cf8168764b8..f47d7064140 100644 --- a/docs/my-website/docs/proxy/pass_through.md +++ b/docs/my-website/docs/proxy/pass_through.md @@ -58,6 +58,17 @@ Configure the required authentication and pricing: - The Bria API requires an `api_token` header - Enter your Bria API key as the value for the `api_token` header +**Default Query Parameters (Optional):** +- Add query parameters that will be automatically sent with every request +- Perfect for API versioning, format specifications, or default configurations +- Clients can override these parameters by providing their own values +- Example: `version=v1`, `format=json`, `timeout=30` + + + **Pricing Configuration:** - Set a cost per request (e.g., $12.00 in this example) - This enables cost tracking and billing for your users @@ -112,6 +123,9 @@ general_settings: content-type: application/json accept: application/json forward_headers: true # Forward all incoming headers + default_query_params: # Optional: Default query parameters + version: "v1" # Always send version=v1 + format: "json" # Default format (can be overridden) ``` ### Start and Test @@ -166,6 +180,9 @@ general_settings: auth: boolean # Enable LiteLLM authentication (Enterprise) forward_headers: boolean # Forward all incoming headers include_subpath: boolean # If true, forwards requests to sub-paths (default: false) + methods: list[string] # Optional: HTTP methods (e.g., ["GET", "POST"]). If not specified, all methods are supported. + default_query_params: # Optional: Default query parameters sent with every request + : string # Key-value pairs (e.g., version: "v1", format: "json") headers: # Custom headers to add Authorization: string # Auth header for target API content-type: string # Request content type @@ -177,11 +194,17 @@ general_settings: ### Header Options - **Authorization**: Authentication for the target API -- **content-type**: Request body format specification +- **content-type**: Request body format specification - **accept**: Expected response format - **LANGFUSE_PUBLIC_KEY/SECRET_KEY**: For Langfuse integration - **Custom headers**: Any additional key-value pairs +### Default Query Parameters +- **Parameter precedence**: Client params > URL params > default params +- **Use cases**: API versioning, authentication tokens, format control, feature flags +- **Override capability**: Clients can override any default parameter +- **Examples**: `version: "v1"`, `format: "json"`, `timeout: "30"` + ### Sub-path Routing By default, pass-through endpoints only match the **exact path** specified. To forward requests to sub-paths, set `include_subpath: true`: @@ -201,6 +224,92 @@ general_settings: --- +### Default Query Parameters + +Pass-through endpoints support default query parameters that are automatically added to every request. This is useful for API versioning, format specifications, authentication tokens, or any default configuration. + +#### How It Works + +**Parameter Precedence (highest to lowest priority):** +1. **Client-provided parameters** (in the request URL) +2. **URL parameters** (from the target URL) +3. **Default parameters** (from configuration) + +#### Example Configuration + +```yaml +general_settings: + pass_through_endpoints: + - path: "/api/v1" + target: "https://external-api.com/service?timeout=60" # URL has timeout=60 + default_query_params: + version: "v1" # Always add version=v1 + format: "json" # Default format=json (can be overridden) + auth_level: "basic" # Always add auth_level=basic +``` + +#### Request Examples + +**Client Request:** `GET /api/v1/users` +**Actual Backend Call:** `https://external-api.com/service?version=v1&format=json&auth_level=basic&timeout=60` + +**Client Request:** `GET /api/v1/users?format=xml&custom=value` +**Actual Backend Call:** `https://external-api.com/service?version=v1&auth_level=basic&timeout=60&format=xml&custom=value` +- Client `format=xml` overrides default `format=json` +- Default `version=v1` and `auth_level=basic` are preserved +- URL `timeout=60` is preserved +- Client `custom=value` is added + +#### Use Cases + +- **API Versioning**: Always send `version=v2` to maintain compatibility +- **Authentication**: Add authentication tokens like `api_key=default_key` +- **Format Control**: Default to `format=json` but allow client override +- **Rate Limiting**: Set `rate_limit=standard` as default +- **Feature Flags**: Enable `experimental=false` by default + +--- + +You can configure different target URLs for the same path using different HTTP methods. This is useful when different backends handle different operations: + + + +```yaml +general_settings: + pass_through_endpoints: + # GET requests to /azure/kb go to read API + - path: "/azure/kb" + target: "https://read-api.example.com/knowledge-base" + methods: ["GET"] + headers: + Authorization: "bearer os.environ/READ_API_KEY" + + # POST requests to /azure/kb go to write API + - path: "/azure/kb" + target: "https://write-api.example.com/knowledge-base" + methods: ["POST"] + headers: + Authorization: "bearer os.environ/WRITE_API_KEY" + + # PUT requests to /azure/kb go to update API + - path: "/azure/kb" + target: "https://update-api.example.com/knowledge-base" + methods: ["PUT"] + headers: + Authorization: "bearer os.environ/UPDATE_API_KEY" +``` + +**Key Points:** +- If `methods` is not specified, the endpoint supports all HTTP methods (GET, POST, PUT, DELETE, PATCH) +- Multiple endpoints can share the same path as long as they have different methods +- You can specify multiple methods for a single endpoint: `methods: ["GET", "POST"]` +- This allows you to route to different backends based on the operation type + +--- + ## Advanced: Custom Adapters For complex integrations (like Anthropic/Bedrock clients), you can create custom adapters that translate between different API schemas. diff --git a/docs/my-website/docs/proxy/project_management.md b/docs/my-website/docs/proxy/project_management.md new file mode 100644 index 00000000000..06ed5b4a0d5 --- /dev/null +++ b/docs/my-website/docs/proxy/project_management.md @@ -0,0 +1,318 @@ +# [Beta] Project Management + +Projects in LiteLLM sit between teams and keys in the organizational hierarchy, enabling fine-grained access control and budget management for specific use cases or applications. + +```mermaid +graph TD + A[Organization] --> B[Team 1] + A --> C[Team 2] + B --> D[Project A] + B --> E[Project B] + C --> F[Project C] + D --> G[API Key 1] + D --> H[API Key 2] + E --> I[API Key 3] + F --> J[API Key 4] + + style A fill:#e1f5ff + style B fill:#fff4e6 + style C fill:#fff4e6 + style D fill:#f3e5f5 + style E fill:#f3e5f5 + style F fill:#f3e5f5 + style G fill:#e8f5e9 + style H fill:#e8f5e9 + style I fill:#e8f5e9 + style J fill:#e8f5e9 +``` + +**Hierarchy**: `Organizations > Teams > Projects > Keys` + +## Quick Start + +This walkthrough shows how to create a project, generate an API key, make requests, and view project-level spend tracking in the UI. + +### Step 1: Create a Project + +```bash showLineNumbers +curl --location 'http://0.0.0.0:4000/project/new' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{ + "project_alias": "flight-search-assistant", + "team_id": "ad898803-c8a3-4f4a-976a-a3c372cffa45", + "models": ["gpt-4", "gpt-3.5-turbo"], + "max_budget": 100, + "metadata": { + "use_case_id": "SNOW-12345", + "responsible_ai_id": "RAI-67890" + } +}' | jq +``` + +**Response:** +```json +{ + "project_id": "e402a141-725a-4437-bff5-d47459189716", + "project_alias": "flight-search-assistant", + "team_id": "ad898803-c8a3-4f4a-976a-a3c372cffa45", + "models": ["gpt-4", "gpt-3.5-turbo"], + "max_budget": 100, + ... +} +``` + +### Step 2: Generate API Key for Project + +```bash showLineNumbers +curl 'http://0.0.0.0:4000/key/generate' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data-raw '{ + "models": ["gpt-3.5-turbo", "gpt-4"], + "metadata": {"user": "ishaan@berri.ai"}, + "project_id": "e402a141-725a-4437-bff5-d47459189716" +}' | jq +``` + +**Response:** +```json +{ + "key": "sk-W8VbscpfuyvHm5TkxRYiXA", + "key_name": "sk-...YiXA", + "project_id": "e402a141-725a-4437-bff5-d47459189716", + ... +} +``` + +### Step 3: Use API Key in Chat Completions + +```bash showLineNumbers +curl http://localhost:4000/v1/chat/completions \ +--header 'Content-Type: application/json' \ +--header 'Authorization: Bearer sk-W8VbscpfuyvHm5TkxRYiXA' \ +--data '{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "What is litellm?"}] +}' | jq +``` + +### Step 4: View Project Spend in UI + +Navigate to the **Logs** page in the LiteLLM Admin UI. You'll see the `user_api_key_project_id` tracked in the request metadata: + +![Project Spend Tracking](/img/project_spend.png) + +As shown above, the spend logs metadata includes: +- `"user_api_key_project_id": "e402a141-725a-4437-bff5-d47459189716"` - Links the request to your project +- All costs and token usage are automatically attributed to the project +- You can query and filter logs by project ID for detailed reporting + +## API Endpoints + +### POST /project/new + +Create a new project. + +**Who can call**: Admins or Team Admins + +**Parameters**: +- `project_alias` (string, optional): Human-readable name for the project +- `team_id` (string, required): The team this project belongs to +- `models` (array, optional): List of models the project can access +- `max_budget` (float, optional): Maximum spend budget for the project +- `tpm_limit` (int, optional): Tokens per minute limit +- `rpm_limit` (int, optional): Requests per minute limit +- `budget_duration` (string, optional): Budget reset period (e.g., "30d", "1mo") +- `metadata` (object, optional): Custom metadata for the project +- `blocked` (boolean, optional): Block all API calls for this project + +**Example**: + +```bash +curl --location 'http://0.0.0.0:4000/project/new' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{ + "project_alias": "hotel-recommendations", + "team_id": "team-123", + "models": ["claude-3-sonnet"], + "max_budget": 200, + "tpm_limit": 100000, + "metadata": { + "use_case_id": "SNOW-12346", + "cost_center": "travel-products" + } +}' +``` + +**Response**: + +```json +{ + "project_id": "project-def", + "project_alias": "hotel-recommendations", + "team_id": "team-123", + "models": ["claude-3-sonnet"], + "spend": 0.0, + "budget_id": "budget-xyz", + "metadata": { + "use_case_id": "SNOW-12346", + "cost_center": "travel-products" + }, + "created_at": "2025-01-15T10:00:00Z", + "updated_at": "2025-01-15T10:00:00Z" +} +``` + +### POST /project/update + +Update an existing project. + +**Who can call**: Admins or Team Admins + +**Parameters**: +- `project_id` (string, required): The project to update +- `project_alias` (string, optional): Updated project name +- `team_id` (string, optional): Move project to different team +- `models` (array, optional): Updated list of allowed models +- `max_budget` (float, optional): Updated budget +- `tpm_limit` (int, optional): Updated TPM limit +- `rpm_limit` (int, optional): Updated RPM limit +- `metadata` (object, optional): Updated metadata +- `blocked` (boolean, optional): Updated blocked status + +**Example**: + +```bash +curl --location 'http://0.0.0.0:4000/project/update' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{ + "project_id": "project-abc", + "max_budget": 200, + "tpm_limit": 200000, + "metadata": { + "status": "production" + } +}' +``` + +### GET /project/info + +Get information about a specific project. + +**Parameters**: +- `project_id` (string, required): Query parameter + +**Example**: + +```bash +curl --location 'http://0.0.0.0:4000/project/info?project_id=project-abc' \ +--header 'Authorization: Bearer sk-1234' +``` + +**Response**: + +```json +{ + "project_id": "project-abc", + "project_alias": "flight-search-assistant", + "team_id": "team-123", + "models": ["gpt-4", "gpt-3.5-turbo"], + "spend": 45.67, + "model_spend": { + "gpt-4": 42.30, + "gpt-3.5-turbo": 3.37 + }, + "litellm_budget_table": { + "budget_id": "budget-xyz", + "max_budget": 100.0, + "tpm_limit": 100000, + "rpm_limit": 100 + }, + "metadata": { + "use_case_id": "SNOW-12345" + } +} +``` + +### GET /project/list + +List all projects the user has access to. + +**Example**: + +```bash +curl --location 'http://0.0.0.0:4000/project/list' \ +--header 'Authorization: Bearer sk-1234' +``` + +**Response**: + +```json +[ + { + "project_id": "project-abc", + "project_alias": "flight-search-assistant", + "team_id": "team-123", + "spend": 45.67 + }, + { + "project_id": "project-def", + "project_alias": "hotel-recommendations", + "team_id": "team-123", + "spend": 23.45 + } +] +``` + +### DELETE /project/delete + +Delete one or more projects. + +**Who can call**: Admins only + +**Parameters**: +- `project_ids` (array, required): List of project IDs to delete + +**Example**: + +```bash +curl --location --request DELETE 'http://0.0.0.0:4000/project/delete' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{ + "project_ids": ["project-abc", "project-def"] +}' +``` + +**Note**: Projects with associated API keys cannot be deleted. Delete or reassign the keys first. + +## Model-Specific Quotas + +You can set different quotas for different models within a project: + +```bash +curl --location 'http://0.0.0.0:4000/project/new' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{ + "project_alias": "multi-model-project", + "team_id": "team-123", + "models": ["gpt-4", "gpt-3.5-turbo", "claude-3-sonnet"], + "max_budget": 500, + "metadata": { + "model_tpm_limit": { + "gpt-4": 50000, + "gpt-3.5-turbo": 200000, + "claude-3-sonnet": 100000 + }, + "model_rpm_limit": { + "gpt-4": 50, + "gpt-3.5-turbo": 500, + "claude-3-sonnet": 100 + } + } +}' +``` diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index 65b7ad7773a..90b1beefa0f 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -1047,6 +1047,8 @@ For long-running conversations, you can enable **server-side compaction** so tha Supported on the OpenAI Responses API when using the `openai` or `azure` provider. Pass `context_management` with a compaction entry and `compact_threshold` (token count; minimum 1000). When the context crosses the threshold, the server compacts in-stream and continues. Chain turns with `previous_response_id` or by appending output items to your next input array. See [OpenAI Compaction guide](https://developers.openai.com/api/docs/guides/compaction) for details. +> **Note:** You can use openai `context_management` format with Anthropic models via LiteLLM via responses API. LiteLLM will automatically translate this format for Anthropic and handle context management for you. + For explicit control over when compaction runs, use the standalone compact endpoint (`POST /v1/responses/compact`) instead. ### Python SDK diff --git a/docs/my-website/img/passthrough_method_setup.png b/docs/my-website/img/passthrough_method_setup.png new file mode 100644 index 00000000000..584e3b966c6 Binary files /dev/null and b/docs/my-website/img/passthrough_method_setup.png differ diff --git a/docs/my-website/img/passthrough_query_default.png b/docs/my-website/img/passthrough_query_default.png new file mode 100644 index 00000000000..fb97e69001e Binary files /dev/null and b/docs/my-website/img/passthrough_query_default.png differ diff --git a/docs/my-website/img/project_spend.png b/docs/my-website/img/project_spend.png new file mode 100644 index 00000000000..955d1786ba1 Binary files /dev/null and b/docs/my-website/img/project_spend.png differ diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 3acfa3937a9..805ce5d8517 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -410,6 +410,7 @@ const sidebars = { items: [ "proxy/users", "proxy/team_budgets", + "proxy/project_management", "proxy/ui_team_soft_budget_alerts", "proxy/tag_budgets", "proxy/customers", @@ -781,13 +782,13 @@ const sidebars = { "providers/bedrock_batches", "providers/bedrock_realtime_with_audio", "providers/aws_polly", - "providers/bedrock_vector_store", - ] - }, - "providers/litellm_proxy", - "providers/abliteration", - "providers/ai21", - "providers/aiml", + "providers/bedrock_vector_store", + ] + }, + "providers/litellm_proxy", + "providers/abliteration", + "providers/ai21", + "providers/aiml", "providers/aleph_alpha", "providers/amazon_nova", "providers/anyscale", diff --git a/docs/my-website/static/img/project_spend.png b/docs/my-website/static/img/project_spend.png new file mode 100644 index 00000000000..955d1786ba1 Binary files /dev/null and b/docs/my-website/static/img/project_spend.png differ diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index e481cdc995c..b6c9104b232 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -1,309 +1,311 @@ -""" -PagerDuty Alerting Integration - -Handles two types of alerts: -- High LLM API Failure Rate. Configure X fails in Y seconds to trigger an alert. -- High Number of Hanging LLM Requests. Configure X hangs in Y seconds to trigger an alert. - -Note: This is a Free feature on the regular litellm docker image. - -However, this is under the enterprise license -""" - -import asyncio -import os -from datetime import datetime, timedelta, timezone -from typing import List, Literal, Optional, Union - -from litellm._logging import verbose_logger -from litellm.caching import DualCache -from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting -from litellm.llms.custom_httpx.http_handler import ( - AsyncHTTPHandler, - get_async_httpx_client, - httpxSpecialProvider, -) -from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.integrations.pagerduty import ( - AlertingConfig, - PagerDutyInternalEvent, - PagerDutyPayload, - PagerDutyRequestBody, -) -from litellm.types.utils import ( - CallTypesLiteral, - StandardLoggingPayload, - StandardLoggingPayloadErrorInformation, -) - -PAGERDUTY_DEFAULT_FAILURE_THRESHOLD = 60 -PAGERDUTY_DEFAULT_FAILURE_THRESHOLD_WINDOW_SECONDS = 60 -PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS = 60 -PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS = 600 - - -class PagerDutyAlerting(SlackAlerting): - """ - Tracks failed requests and hanging requests separately. - If threshold is crossed for either type, triggers a PagerDuty alert. - """ - - def __init__( - self, alerting_args: Optional[Union[AlertingConfig, dict]] = None, **kwargs - ): - super().__init__() - _api_key = os.getenv("PAGERDUTY_API_KEY") - if not _api_key: - raise ValueError("PAGERDUTY_API_KEY is not set") - - self.api_key: str = _api_key - alerting_args = alerting_args or {} - self.pagerduty_alerting_args: AlertingConfig = AlertingConfig( - failure_threshold=alerting_args.get( - "failure_threshold", PAGERDUTY_DEFAULT_FAILURE_THRESHOLD - ), - failure_threshold_window_seconds=alerting_args.get( - "failure_threshold_window_seconds", - PAGERDUTY_DEFAULT_FAILURE_THRESHOLD_WINDOW_SECONDS, - ), - hanging_threshold_seconds=alerting_args.get( - "hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS - ), - hanging_threshold_window_seconds=alerting_args.get( - "hanging_threshold_window_seconds", - PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS, - ), - ) - - # Separate storage for failures vs. hangs - self._failure_events: List[PagerDutyInternalEvent] = [] - self._hanging_events: List[PagerDutyInternalEvent] = [] - - # ------------------ MAIN LOGIC ------------------ # - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - """ - Record a failure event. Only send an alert to PagerDuty if the - configured *failure* threshold is exceeded in the specified window. - """ - now = datetime.now(timezone.utc) - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object" - ) - if not standard_logging_payload: - raise ValueError( - "standard_logging_object is required for PagerDutyAlerting" - ) - - # Extract error details - error_info: Optional[StandardLoggingPayloadErrorInformation] = ( - standard_logging_payload.get("error_information") or {} - ) - _meta = standard_logging_payload.get("metadata") or {} - - self._failure_events.append( - PagerDutyInternalEvent( - failure_event_type="failed_response", - timestamp=now, - error_class=error_info.get("error_class"), - error_code=error_info.get("error_code"), - error_llm_provider=error_info.get("llm_provider"), - user_api_key_hash=_meta.get("user_api_key_hash"), - user_api_key_alias=_meta.get("user_api_key_alias"), - user_api_key_spend=_meta.get("user_api_key_spend"), - user_api_key_max_budget=_meta.get("user_api_key_max_budget"), - user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"), - user_api_key_org_id=_meta.get("user_api_key_org_id"), - user_api_key_team_id=_meta.get("user_api_key_team_id"), - user_api_key_user_id=_meta.get("user_api_key_user_id"), - user_api_key_team_alias=_meta.get("user_api_key_team_alias"), - user_api_key_end_user_id=_meta.get("user_api_key_end_user_id"), - user_api_key_user_email=_meta.get("user_api_key_user_email"), - user_api_key_request_route=_meta.get("user_api_key_request_route"), - user_api_key_auth_metadata=_meta.get("user_api_key_auth_metadata"), - ) - ) - - # Prune + Possibly alert - window_seconds = self.pagerduty_alerting_args.get( - "failure_threshold_window_seconds", 60 - ) - threshold = self.pagerduty_alerting_args.get("failure_threshold", 1) - - # If threshold is crossed, send PD alert for failures - await self._send_alert_if_thresholds_crossed( - events=self._failure_events, - window_seconds=window_seconds, - threshold=threshold, - alert_prefix="High LLM API Failure Rate", - ) - - async def async_pre_call_hook( - self, - user_api_key_dict: UserAPIKeyAuth, - cache: DualCache, - data: dict, - call_type: CallTypesLiteral, - ) -> Optional[Union[Exception, str, dict]]: - """ - Example of detecting hanging requests by waiting a given threshold. - If the request didn't finish by then, we treat it as 'hanging'. - """ - verbose_logger.info("Inside Proxy Logging Pre-call hook!") - asyncio.create_task( - self.hanging_response_handler( - request_data=data, user_api_key_dict=user_api_key_dict - ) - ) - return None - - async def hanging_response_handler( - self, request_data: Optional[dict], user_api_key_dict: UserAPIKeyAuth - ): - """ - Checks if request completed by the time 'hanging_threshold_seconds' elapses. - If not, we classify it as a hanging request. - """ - verbose_logger.debug( - f"Inside Hanging Response Handler!..sleeping for {self.pagerduty_alerting_args.get('hanging_threshold_seconds', PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS)} seconds" - ) - await asyncio.sleep( - self.pagerduty_alerting_args.get( - "hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS - ) - ) - - if await self._request_is_completed(request_data=request_data): - return # It's not hanging if completed - - # Otherwise, record it as hanging - self._hanging_events.append( - PagerDutyInternalEvent( - failure_event_type="hanging_response", - timestamp=datetime.now(timezone.utc), - error_class="HangingRequest", - error_code="HangingRequest", - error_llm_provider="HangingRequest", - user_api_key_hash=user_api_key_dict.api_key, - user_api_key_alias=user_api_key_dict.key_alias, - 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_org_id=user_api_key_dict.org_id, - user_api_key_team_id=user_api_key_dict.team_id, - user_api_key_user_id=user_api_key_dict.user_id, - user_api_key_team_alias=user_api_key_dict.team_alias, - user_api_key_end_user_id=user_api_key_dict.end_user_id, - user_api_key_user_email=user_api_key_dict.user_email, - user_api_key_request_route=user_api_key_dict.request_route, - user_api_key_auth_metadata=user_api_key_dict.metadata, - ) - ) - - # Prune + Possibly alert - window_seconds = self.pagerduty_alerting_args.get( - "hanging_threshold_window_seconds", - PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS, - ) - threshold: int = self.pagerduty_alerting_args.get( - "hanging_threshold_fails", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS - ) - - # If threshold is crossed, send PD alert for hangs - await self._send_alert_if_thresholds_crossed( - events=self._hanging_events, - window_seconds=window_seconds, - threshold=threshold, - alert_prefix="High Number of Hanging LLM Requests", - ) - - # ------------------ HELPERS ------------------ # - - async def _send_alert_if_thresholds_crossed( - self, - events: List[PagerDutyInternalEvent], - window_seconds: int, - threshold: int, - alert_prefix: str, - ): - """ - 1. Prune old events - 2. If threshold is reached, build alert, send to PagerDuty - 3. Clear those events - """ - cutoff = datetime.now(timezone.utc) - timedelta(seconds=window_seconds) - pruned = [e for e in events if e.get("timestamp", datetime.min) > cutoff] - - # Update the reference list - events.clear() - events.extend(pruned) - - # Check threshold - verbose_logger.debug( - f"Have {len(events)} events in the last {window_seconds} seconds. Threshold is {threshold}" - ) - if len(events) >= threshold: - # Build short summary of last N events - error_summaries = self._build_error_summaries(events, max_errors=5) - alert_message = ( - f"{alert_prefix}: {len(events)} in the last {window_seconds} seconds." - ) - custom_details = {"recent_errors": error_summaries} - - await self.send_alert_to_pagerduty( - alert_message=alert_message, - custom_details=custom_details, - ) - - # Clear them after sending an alert, so we don't spam - events.clear() - - def _build_error_summaries( - self, events: List[PagerDutyInternalEvent], max_errors: int = 5 - ) -> List[PagerDutyInternalEvent]: - """ - Build short text summaries for the last `max_errors`. - Example: "ValueError (code: 500, provider: openai)" - """ - recent = events[-max_errors:] - summaries = [] - for fe in recent: - # If any of these is None, show "N/A" to avoid messing up the summary string - fe.pop("timestamp") - summaries.append(fe) - return summaries - - async def send_alert_to_pagerduty(self, alert_message: str, custom_details: dict): - """ - Send [critical] Alert to PagerDuty - - https://developer.pagerduty.com/api-reference/YXBpOjI3NDgyNjU-pager-duty-v2-events-api - """ - try: - verbose_logger.debug(f"Sending alert to PagerDuty: {alert_message}") - async_client: AsyncHTTPHandler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.LoggingCallback - ) - payload: PagerDutyRequestBody = PagerDutyRequestBody( - payload=PagerDutyPayload( - summary=alert_message, - severity="critical", - source="LiteLLM Alert", - component="LiteLLM", - custom_details=custom_details, - ), - routing_key=self.api_key, - event_action="trigger", - ) - - return await async_client.post( - url="https://events.pagerduty.com/v2/enqueue", - json=dict(payload), - headers={"Content-Type": "application/json"}, - ) - except Exception as e: - verbose_logger.exception(f"Error sending alert to PagerDuty: {e}") +""" +PagerDuty Alerting Integration + +Handles two types of alerts: +- High LLM API Failure Rate. Configure X fails in Y seconds to trigger an alert. +- High Number of Hanging LLM Requests. Configure X hangs in Y seconds to trigger an alert. + +Note: This is a Free feature on the regular litellm docker image. + +However, this is under the enterprise license +""" + +import asyncio +import os +from datetime import datetime, timedelta, timezone +from typing import List, Optional, Union + +from litellm._logging import verbose_logger +from litellm.caching import DualCache +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.integrations.pagerduty import ( + AlertingConfig, + PagerDutyInternalEvent, + PagerDutyPayload, + PagerDutyRequestBody, +) +from litellm.types.utils import ( + CallTypesLiteral, + StandardLoggingPayload, + StandardLoggingPayloadErrorInformation, +) + +PAGERDUTY_DEFAULT_FAILURE_THRESHOLD = 60 +PAGERDUTY_DEFAULT_FAILURE_THRESHOLD_WINDOW_SECONDS = 60 +PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS = 60 +PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS = 600 + + +class PagerDutyAlerting(SlackAlerting): + """ + Tracks failed requests and hanging requests separately. + If threshold is crossed for either type, triggers a PagerDuty alert. + """ + + def __init__( + self, alerting_args: Optional[Union[AlertingConfig, dict]] = None, **kwargs + ): + super().__init__() + _api_key = os.getenv("PAGERDUTY_API_KEY") + if not _api_key: + raise ValueError("PAGERDUTY_API_KEY is not set") + + self.api_key: str = _api_key + alerting_args = alerting_args or {} + self.pagerduty_alerting_args: AlertingConfig = AlertingConfig( + failure_threshold=alerting_args.get( + "failure_threshold", PAGERDUTY_DEFAULT_FAILURE_THRESHOLD + ), + failure_threshold_window_seconds=alerting_args.get( + "failure_threshold_window_seconds", + PAGERDUTY_DEFAULT_FAILURE_THRESHOLD_WINDOW_SECONDS, + ), + hanging_threshold_seconds=alerting_args.get( + "hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS + ), + hanging_threshold_window_seconds=alerting_args.get( + "hanging_threshold_window_seconds", + PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS, + ), + ) + + # Separate storage for failures vs. hangs + self._failure_events: List[PagerDutyInternalEvent] = [] + self._hanging_events: List[PagerDutyInternalEvent] = [] + + # ------------------ MAIN LOGIC ------------------ # + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + """ + Record a failure event. Only send an alert to PagerDuty if the + configured *failure* threshold is exceeded in the specified window. + """ + now = datetime.now(timezone.utc) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + if not standard_logging_payload: + raise ValueError( + "standard_logging_object is required for PagerDutyAlerting" + ) + + # Extract error details + error_info: Optional[StandardLoggingPayloadErrorInformation] = ( + standard_logging_payload.get("error_information") or {} + ) + _meta = standard_logging_payload.get("metadata") or {} + + self._failure_events.append( + PagerDutyInternalEvent( + failure_event_type="failed_response", + timestamp=now, + error_class=error_info.get("error_class"), + error_code=error_info.get("error_code"), + error_llm_provider=error_info.get("llm_provider"), + user_api_key_hash=_meta.get("user_api_key_hash"), + user_api_key_alias=_meta.get("user_api_key_alias"), + user_api_key_spend=_meta.get("user_api_key_spend"), + user_api_key_max_budget=_meta.get("user_api_key_max_budget"), + user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"), + user_api_key_org_id=_meta.get("user_api_key_org_id"), + user_api_key_team_id=_meta.get("user_api_key_team_id"), + user_api_key_project_id=_meta.get("user_api_key_project_id"), + user_api_key_user_id=_meta.get("user_api_key_user_id"), + user_api_key_team_alias=_meta.get("user_api_key_team_alias"), + user_api_key_end_user_id=_meta.get("user_api_key_end_user_id"), + user_api_key_user_email=_meta.get("user_api_key_user_email"), + user_api_key_request_route=_meta.get("user_api_key_request_route"), + user_api_key_auth_metadata=_meta.get("user_api_key_auth_metadata"), + ) + ) + + # Prune + Possibly alert + window_seconds = self.pagerduty_alerting_args.get( + "failure_threshold_window_seconds", 60 + ) + threshold = self.pagerduty_alerting_args.get("failure_threshold", 1) + + # If threshold is crossed, send PD alert for failures + await self._send_alert_if_thresholds_crossed( + events=self._failure_events, + window_seconds=window_seconds, + threshold=threshold, + alert_prefix="High LLM API Failure Rate", + ) + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> Optional[Union[Exception, str, dict]]: + """ + Example of detecting hanging requests by waiting a given threshold. + If the request didn't finish by then, we treat it as 'hanging'. + """ + verbose_logger.info("Inside Proxy Logging Pre-call hook!") + asyncio.create_task( + self.hanging_response_handler( + request_data=data, user_api_key_dict=user_api_key_dict + ) + ) + return None + + async def hanging_response_handler( + self, request_data: Optional[dict], user_api_key_dict: UserAPIKeyAuth + ): + """ + Checks if request completed by the time 'hanging_threshold_seconds' elapses. + If not, we classify it as a hanging request. + """ + verbose_logger.debug( + f"Inside Hanging Response Handler!..sleeping for {self.pagerduty_alerting_args.get('hanging_threshold_seconds', PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS)} seconds" + ) + await asyncio.sleep( + self.pagerduty_alerting_args.get( + "hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS + ) + ) + + if await self._request_is_completed(request_data=request_data): + return # It's not hanging if completed + + # Otherwise, record it as hanging + self._hanging_events.append( + PagerDutyInternalEvent( + failure_event_type="hanging_response", + timestamp=datetime.now(timezone.utc), + error_class="HangingRequest", + error_code="HangingRequest", + error_llm_provider="HangingRequest", + user_api_key_hash=user_api_key_dict.api_key, + user_api_key_alias=user_api_key_dict.key_alias, + 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_org_id=user_api_key_dict.org_id, + user_api_key_team_id=user_api_key_dict.team_id, + user_api_key_project_id=user_api_key_dict.project_id, + user_api_key_user_id=user_api_key_dict.user_id, + user_api_key_team_alias=user_api_key_dict.team_alias, + user_api_key_end_user_id=user_api_key_dict.end_user_id, + user_api_key_user_email=user_api_key_dict.user_email, + user_api_key_request_route=user_api_key_dict.request_route, + user_api_key_auth_metadata=user_api_key_dict.metadata, + ) + ) + + # Prune + Possibly alert + window_seconds = self.pagerduty_alerting_args.get( + "hanging_threshold_window_seconds", + PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS, + ) + threshold: int = self.pagerduty_alerting_args.get( + "hanging_threshold_fails", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS + ) + + # If threshold is crossed, send PD alert for hangs + await self._send_alert_if_thresholds_crossed( + events=self._hanging_events, + window_seconds=window_seconds, + threshold=threshold, + alert_prefix="High Number of Hanging LLM Requests", + ) + + # ------------------ HELPERS ------------------ # + + async def _send_alert_if_thresholds_crossed( + self, + events: List[PagerDutyInternalEvent], + window_seconds: int, + threshold: int, + alert_prefix: str, + ): + """ + 1. Prune old events + 2. If threshold is reached, build alert, send to PagerDuty + 3. Clear those events + """ + cutoff = datetime.now(timezone.utc) - timedelta(seconds=window_seconds) + pruned = [e for e in events if e.get("timestamp", datetime.min) > cutoff] + + # Update the reference list + events.clear() + events.extend(pruned) + + # Check threshold + verbose_logger.debug( + f"Have {len(events)} events in the last {window_seconds} seconds. Threshold is {threshold}" + ) + if len(events) >= threshold: + # Build short summary of last N events + error_summaries = self._build_error_summaries(events, max_errors=5) + alert_message = ( + f"{alert_prefix}: {len(events)} in the last {window_seconds} seconds." + ) + custom_details = {"recent_errors": error_summaries} + + await self.send_alert_to_pagerduty( + alert_message=alert_message, + custom_details=custom_details, + ) + + # Clear them after sending an alert, so we don't spam + events.clear() + + def _build_error_summaries( + self, events: List[PagerDutyInternalEvent], max_errors: int = 5 + ) -> List[PagerDutyInternalEvent]: + """ + Build short text summaries for the last `max_errors`. + Example: "ValueError (code: 500, provider: openai)" + """ + recent = events[-max_errors:] + summaries = [] + for fe in recent: + # If any of these is None, show "N/A" to avoid messing up the summary string + fe.pop("timestamp") + summaries.append(fe) + return summaries + + async def send_alert_to_pagerduty(self, alert_message: str, custom_details: dict): + """ + Send [critical] Alert to PagerDuty + + https://developer.pagerduty.com/api-reference/YXBpOjI3NDgyNjU-pager-duty-v2-events-api + """ + try: + verbose_logger.debug(f"Sending alert to PagerDuty: {alert_message}") + async_client: AsyncHTTPHandler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + payload: PagerDutyRequestBody = PagerDutyRequestBody( + payload=PagerDutyPayload( + summary=alert_message, + severity="critical", + source="LiteLLM Alert", + component="LiteLLM", + custom_details=custom_details, + ), + routing_key=self.api_key, + event_action="trigger", + ) + + return await async_client.post( + url="https://events.pagerduty.com/v2/enqueue", + json=dict(payload), + headers={"Content-Type": "application/json"}, + ) + except Exception as e: + verbose_logger.exception(f"Error sending alert to PagerDuty: {e}") diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.41-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.41-py3-none-any.whl new file mode 100644 index 00000000000..9d7fdb78f72 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.41-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.41.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.41.tar.gz new file mode 100644 index 00000000000..a478356f886 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.41.tar.gz differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.42-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.42-py3-none-any.whl new file mode 100644 index 00000000000..c2eedc2a258 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.42-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.42.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.42.tar.gz new file mode 100644 index 00000000000..fc9ff018078 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.42.tar.gz differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.43-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.43-py3-none-any.whl new file mode 100644 index 00000000000..ee821fed313 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.43-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.43.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.43.tar.gz new file mode 100644 index 00000000000..d0304bd9825 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.43.tar.gz differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.44-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.44-py3-none-any.whl new file mode 100644 index 00000000000..29eb20f0d97 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.44-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.44.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.44.tar.gz new file mode 100644 index 00000000000..7b3070f71a2 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.44.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000000_add_project_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000000_add_project_table/migration.sql new file mode 100644 index 00000000000..f1d3129bb36 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000000_add_project_table/migration.sql @@ -0,0 +1,35 @@ +-- CreateTable +CREATE TABLE "LiteLLM_ProjectTable" ( + "project_id" TEXT NOT NULL, + "project_alias" TEXT, + "team_id" TEXT, + "budget_id" TEXT, + "metadata" JSONB NOT NULL DEFAULT '{}', + "models" TEXT[], + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "model_spend" JSONB NOT NULL DEFAULT '{}', + "blocked" BOOLEAN NOT NULL DEFAULT false, + "object_permission_id" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_ProjectTable_pkey" PRIMARY KEY ("project_id") +); + +-- AddForeignKey +ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_team_id_fkey" FOREIGN KEY ("team_id") REFERENCES "LiteLLM_TeamTable"("team_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_ProjectTable" ADD CONSTRAINT "LiteLLM_ProjectTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AlterTable: Add project_id to LiteLLM_VerificationToken +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "project_id" TEXT; + +-- AddForeignKey +ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_project_id_fkey" FOREIGN KEY ("project_id") REFERENCES "LiteLLM_ProjectTable"("project_id") ON DELETE SET NULL ON UPDATE CASCADE; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000001_add_project_fields/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000001_add_project_fields/migration.sql new file mode 100644 index 00000000000..48328b4d6a2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000001_add_project_fields/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable: Add new fields to LiteLLM_ProjectTable +ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN "description" TEXT; +ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN "model_rpm_limit" JSONB NOT NULL DEFAULT '{}'; +ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN "model_tpm_limit" JSONB NOT NULL DEFAULT '{}'; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260218231534_add_last_active_to_key_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260218231534_add_last_active_to_key_table/migration.sql new file mode 100644 index 00000000000..ded1856059b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260218231534_add_last_active_to_key_table/migration.sql @@ -0,0 +1,6 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "last_active" TIMESTAMP(3); + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "last_active" TIMESTAMP(3); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219105005_add_project_id_to_deleted_keys/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219105005_add_project_id_to_deleted_keys/migration.sql new file mode 100644 index 00000000000..59bdc86adbb --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219105005_add_project_id_to_deleted_keys/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "project_id" TEXT; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8bd46672ae6..45cd90f3413 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -24,6 +24,7 @@ model LiteLLM_BudgetTable { updated_at DateTime @default(now()) @updatedAt @map("updated_at") updated_by String organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget + projects LiteLLM_ProjectTable[] // multiple projects can have the same budget keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget @@ -135,6 +136,81 @@ model LiteLLM_TeamTable { litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + projects LiteLLM_ProjectTable[] +} + +// Projects sit between teams and keys for use-case management +model LiteLLM_ProjectTable { + project_id String @id @default(uuid()) + project_alias String? + description String? + team_id String? + budget_id String? + metadata Json @default("{}") + models String[] + spend Float @default(0.0) + model_spend Json @default("{}") + model_rpm_limit Json @default("{}") + model_tpm_limit Json @default("{}") + blocked Boolean @default(false) + object_permission_id String? + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + + // Relations + litellm_team_table LiteLLM_TeamTable? @relation(fields: [team_id], references: [team_id]) + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + keys LiteLLM_VerificationToken[] + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) +} + +// Audit table for deleted teams - preserves spend and team information for historical tracking +model LiteLLM_DeletedTeamTable { + id String @id @default(uuid()) + team_id String // Original team_id + team_alias String? + organization_id String? + object_permission_id String? + admins String[] + members String[] + members_with_roles Json @default("{}") + metadata Json @default("{}") + max_budget Float? + soft_budget Float? + spend Float @default(0.0) + models String[] + max_parallel_requests Int? + tpm_limit BigInt? + rpm_limit BigInt? + budget_duration String? + budget_reset_at DateTime? + blocked Boolean @default(false) + model_spend Json @default("{}") + model_max_budget Json @default("{}") + router_settings Json? @default("{}") + team_member_permissions String[] @default([]) + access_group_ids String[] @default([]) + policies String[] @default([]) + model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases + allow_team_guardrail_config Boolean @default(false) + + // Original timestamps from team creation/updates + created_at DateTime? @map("created_at") + updated_at DateTime? @map("updated_at") + + // Deletion metadata + deleted_at DateTime @default(now()) @map("deleted_at") + deleted_by String? @map("deleted_by") // User who deleted the team + deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion + litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided + + @@index([team_id]) + @@index([deleted_at]) + @@index([organization_id]) + @@index([team_alias]) + @@index([created_at]) } // Audit table for deleted teams - preserves spend and team information for historical tracking @@ -230,6 +306,7 @@ model LiteLLM_ObjectPermissionTable { agents String[] @default([]) agent_access_groups String[] @default([]) teams LiteLLM_TeamTable[] + projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] @@ -284,6 +361,7 @@ model LiteLLM_VerificationToken { router_settings Json? @default("{}") user_id String? team_id String? + project_id String? permissions Json @default("{}") max_parallel_requests Int? metadata Json @default("{}") @@ -306,6 +384,7 @@ model LiteLLM_VerificationToken { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + last_active DateTime? // When this key was last used rotation_count Int? @default(0) // Number of times key has been rotated auto_rotate Boolean? @default(false) // Whether this key should be auto-rotated rotation_interval String? // How often to rotate (e.g., "30d", "90d") @@ -313,6 +392,7 @@ model LiteLLM_VerificationToken { key_rotation_at DateTime? // When this key should next be rotated litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) + litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" @@ -353,6 +433,7 @@ model LiteLLM_DeletedVerificationToken { config Json @default("{}") user_id String? team_id String? + project_id String? permissions Json @default("{}") max_parallel_requests Int? metadata Json @default("{}") @@ -376,6 +457,7 @@ model LiteLLM_DeletedVerificationToken { created_by String? // Original creator updated_at DateTime? // Last update timestamp before deletion updated_by String? // Last user who updated before deletion + last_active DateTime? // When this key was last used before deletion rotation_count Int? @default(0) auto_rotate Boolean? @default(false) rotation_interval String? @@ -448,7 +530,7 @@ model LiteLLM_SpendLogs { custom_llm_provider String? @default("") // litellm used custom_llm_provider api_base String? @default("") user String? @default("") - metadata Json? @default("{}") + metadata Json? @default("{}") // project_id stored here cache_hit String? @default("") cache_key String? @default("") request_tags Json? @default("[]") diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 7ef0409b6b8..9ceee1e343f 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.40" +version = "0.4.44" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.40" +version = "0.4.44" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index a06c7173ea3..9b79a38214c 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -67,7 +67,7 @@ "compact-2026-01-12": null, "computer-use-2025-01-24": "computer-use-2025-01-24", "computer-use-2025-11-24": "computer-use-2025-11-24", - "context-1m-2025-08-07": null, + "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", "effort-2025-11-24": null, "fast-mode-2026-02-01": null, diff --git a/litellm/constants.py b/litellm/constants.py index 17ad742e419..7a2107c8c82 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -578,7 +578,6 @@ OPENAI_CHAT_COMPLETION_PARAMS = [ "thinking", "web_search_options", "service_tier", - "store", "prompt_cache_key", "prompt_cache_retention", "safety_identifier", @@ -644,6 +643,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES = { "prompt_cache_retention": None, "store": None, "metadata": None, + "context_management": None, } openai_compatible_endpoints: List = [ @@ -1474,3 +1474,8 @@ MICROSOFT_USER_FIRST_NAME_ATTRIBUTE = str( MICROSOFT_USER_LAST_NAME_ATTRIBUTE = str( os.getenv("MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "surname") ) + +# Policy template enrichment +MAX_COMPETITOR_NAMES = int(os.getenv("MAX_COMPETITOR_NAMES", 100)) +COMPETITOR_LLM_TEMPERATURE = float(os.getenv("COMPETITOR_LLM_TEMPERATURE", 0.3)) +DEFAULT_COMPETITOR_DISCOVERY_MODEL = "gpt-4o-mini" diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py index 205c5c89e35..ea80b258540 100644 --- a/litellm/integrations/SlackAlerting/budget_alert_types.py +++ b/litellm/integrations/SlackAlerting/budget_alert_types.py @@ -74,6 +74,14 @@ class ProjectedLimitExceededAlert(BaseBudgetAlertType): return user_info.token or "default_id" +class ProjectBudgetAlert(BaseBudgetAlertType): + def get_event_message(self) -> str: + return "Project Budget: " + + def get_id(self, user_info: CallInfo) -> str: + return user_info.token or "default_id" + + def get_budget_alert_type( type: Literal[ "token_budget", @@ -84,6 +92,7 @@ def get_budget_alert_type( "organization_budget", "proxy_budget", "projected_limit_exceeded", + "project_budget", ], ) -> BaseBudgetAlertType: """Factory function to get the appropriate budget alert type class""" @@ -97,6 +106,7 @@ def get_budget_alert_type( "organization_budget": OrganizationBudgetAlert(), "token_budget": TokenBudgetAlert(), "projected_limit_exceeded": ProjectedLimitExceededAlert(), + "project_budget": ProjectBudgetAlert(), } if type in alert_types: diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 8fb3e132ded..a525856db82 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -538,6 +538,7 @@ class SlackAlerting(CustomBatchLogger): "organization_budget", "proxy_budget", "projected_limit_exceeded", + "project_budget", ], user_info: CallInfo, ): @@ -1378,9 +1379,13 @@ Model Info: """ if self.alerting is None: return - + # Start periodic flush if not already started - if not self.periodic_started and self.alerting is not None and len(self.alerting) > 0: + if ( + not self.periodic_started + and self.alerting is not None + and len(self.alerting) > 0 + ): asyncio.create_task(self.periodic_flush()) self.periodic_started = True diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 1277cac51d7..d7858d71eb3 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -16,6 +16,7 @@ from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.tools import ( get_litellm_web_search_tool, + get_litellm_web_search_tool_openai, is_web_search_tool, is_web_search_tool_chat_completion, ) @@ -77,7 +78,13 @@ class WebSearchInterceptionLogger(CustomLogger): that we can intercept and execute ourselves. """ # Check if this is for an enabled provider - custom_llm_provider = kwargs.get("litellm_params", {}).get("custom_llm_provider", "") + # Try top-level kwargs first, then nested litellm_params, then derive from model name + custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get("custom_llm_provider", "") + if not custom_llm_provider: + try: + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", "")) + except Exception: + custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: return None @@ -101,7 +108,7 @@ class WebSearchInterceptionLogger(CustomLogger): for tool in tools: if is_web_search_tool(tool): # Convert to LiteLLM standard web search tool - converted_tool = get_litellm_web_search_tool() + converted_tool = get_litellm_web_search_tool_openai() converted_tools.append(converted_tool) verbose_logger.debug( f"WebSearchInterception: Converted {tool.get('name', 'unknown')} " @@ -111,8 +118,9 @@ class WebSearchInterceptionLogger(CustomLogger): # Keep other tools as-is converted_tools.append(tool) - # Return modified kwargs with converted tools - return {"tools": converted_tools} + # Update tools in-place and return full kwargs + kwargs["tools"] = converted_tools + return kwargs @classmethod def from_config_yaml( diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index c39d150fb19..7ef2b35004d 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -49,6 +49,39 @@ def get_litellm_web_search_tool() -> Dict[str, Any]: } +def get_litellm_web_search_tool_openai() -> Dict[str, Any]: + """ + Get the standard LiteLLM web search tool definition in OpenAI format. + + Used by async_pre_call_deployment_hook which runs in the chat completions + path where tools must be in OpenAI format (type: "function" with + function.parameters). + + Returns: + Dict containing the OpenAI-style tool definition. + """ + return { + "type": "function", + "function": { + "name": LITELLM_WEB_SEARCH_TOOL_NAME, + "description": ( + "Search the web for information. Use this when you need current " + "information or answers to questions that require up-to-date data." + ), + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to execute" + } + }, + "required": ["query"] + } + } + } + + def is_web_search_tool_chat_completion(tool: Dict[str, Any]) -> bool: """ Check if a tool is a web search tool for Chat Completions API (strict check). diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6a14e42c485..5f555d83cf1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1,5444 +1,5446 @@ -# What is this? -## Common Utility file for Logging handler -# Logging function -> log the exact model details + what's being sent | Non-Blocking -import copy -import datetime -import json -import os -import re -import subprocess -import sys -import time -import traceback -from datetime import datetime as dt_object -from functools import lru_cache -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - List, - Literal, - Optional, - Tuple, - Type, - Union, - cast, -) - -from httpx import Response -from pydantic import BaseModel - -import litellm -from litellm import ( - _custom_logger_compatible_callbacks_literal, - json_logs, - log_raw_request_response, - turn_off_message_logging, -) -from litellm._logging import _is_debugging_on, verbose_logger -from litellm._uuid import uuid -from litellm.batches.batch_utils import _handle_completed_batch -from litellm.caching.caching import DualCache, InMemoryCache -from litellm.caching.caching_handler import LLMCachingHandler -from litellm.constants import ( - DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, - DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, - SENTRY_DENYLIST, - SENTRY_PII_DENYLIST, -) -from litellm.cost_calculator import ( - RealtimeAPITokenUsageProcessor, - _select_model_name_for_cost_calc, -) -from litellm.integrations.agentops import AgentOps -from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook -from litellm.integrations.arize.arize import ArizeLogger -from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.deepeval.deepeval import DeepEvalLogger -from litellm.integrations.mlflow import MlflowLogger -from litellm.integrations.sqs import SQSLogger -from litellm.litellm_core_utils.core_helpers import reconstruct_model_name -from litellm.litellm_core_utils.get_litellm_params import get_litellm_params -from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( - StandardBuiltInToolCostTracking, -) -from litellm.litellm_core_utils.model_param_helper import ModelParamHelper -from litellm.litellm_core_utils.redact_messages import ( - redact_message_input_output_from_custom_logger, - redact_message_input_output_from_logging, -) -from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.llms.base_llm.search.transformation import SearchResponse -from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.types.agents import LiteLLMSendMessageResponse -from litellm.types.containers.main import ContainerObject -from litellm.types.llms.openai import ( - AllMessageValues, - Batch, - FineTuningJob, - HttpxBinaryResponseContent, - OpenAIFileObject, - OpenAIModerationResponse, - ResponseAPIUsage, - ResponseCompletedEvent, - ResponsesAPIResponse, -) -from litellm.types.mcp import MCPPostCallResponseObject -from litellm.types.prompts.init_prompts import PromptSpec -from litellm.types.rerank import RerankResponse -from litellm.types.utils import ( - CachingDetails, - CallTypes, - CostBreakdown, - CostResponseTypes, - CustomPricingLiteLLMParams, - DynamicPromptManagementParamLiteral, - EmbeddingResponse, - GuardrailStatus, - ImageResponse, - LiteLLMBatch, - LiteLLMLoggingBaseClass, - LiteLLMRealtimeStreamLoggingObject, - ModelResponse, - ModelResponseStream, - RawRequestTypedDict, - StandardBuiltInToolsParams, - StandardCallbackDynamicParams, - StandardLoggingAdditionalHeaders, - StandardLoggingHiddenParams, - StandardLoggingMCPToolCall, - StandardLoggingMetadata, - StandardLoggingModelCostFailureDebugInformation, - StandardLoggingModelInformation, - StandardLoggingPayload, - StandardLoggingPayloadErrorInformation, - StandardLoggingPayloadStatus, - StandardLoggingPayloadStatusFields, - StandardLoggingPromptManagementMetadata, - StandardLoggingVectorStoreRequest, - TextCompletionResponse, - TranscriptionResponse, - Usage, -) -from litellm.types.videos.main import VideoObject -from litellm.utils import _get_base_model_from_metadata, executor, print_verbose - -from ..integrations.argilla import ArgillaLogger -from ..integrations.arize.arize_phoenix import ArizePhoenixLogger -from ..integrations.athina import AthinaLogger -from ..integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger -from ..integrations.azure_storage.azure_storage import AzureBlobStorageLogger -from ..integrations.custom_prompt_management import CustomPromptManagement -from ..integrations.datadog.datadog import DataDogLogger -from ..integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger -from ..integrations.dotprompt import DotpromptManager -from ..integrations.dynamodb import DyanmoDBLogger -from ..integrations.galileo import GalileoObserve -from ..integrations.gcs_bucket.gcs_bucket import GCSBucketLogger -from ..integrations.gcs_pubsub.pub_sub import GcsPubSubLogger -from ..integrations.greenscale import GreenscaleLogger -from ..integrations.helicone import HeliconeLogger -from ..integrations.humanloop import HumanloopLogger -from ..integrations.lago import LagoLogger -from ..integrations.langfuse.langfuse import LangFuseLogger -from ..integrations.langfuse.langfuse_handler import LangFuseHandler -from ..integrations.langfuse.langfuse_prompt_management import LangfusePromptManagement -from ..integrations.langsmith import LangsmithLogger -from ..integrations.literal_ai import LiteralAILogger -from ..integrations.logfire_logger import LogfireLevel, LogfireLogger -from ..integrations.lunary import LunaryLogger -from ..integrations.openmeter import OpenMeterLogger -from ..integrations.opik.opik import OpikLogger -from ..integrations.posthog import PostHogLogger -from ..integrations.prompt_layer import PromptLayerLogger -from ..integrations.s3 import S3Logger -from ..integrations.s3_v2 import S3Logger as S3V2Logger -from ..integrations.supabase import Supabase -from ..integrations.traceloop import TraceloopLogger -from .exception_mapping_utils import _get_response_headers -from .initialize_dynamic_callback_params import ( - initialize_standard_callback_dynamic_params as _initialize_standard_callback_dynamic_params, -) -from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache - -if TYPE_CHECKING: - from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig -try: - from litellm_enterprise.enterprise_callbacks.callback_controls import ( - EnterpriseCallbackControls, - ) - from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import ( - PagerDutyAlerting, - ) - from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ( - ResendEmailLogger, - ) - from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import ( - SendGridEmailLogger, - ) - from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import ( - SMTPEmailLogger, - ) - from litellm_enterprise.litellm_core_utils.litellm_logging import ( - StandardLoggingPayloadSetup as EnterpriseStandardLoggingPayloadSetup, - ) - - from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger - - EnterpriseStandardLoggingPayloadSetupVAR: Optional[ - Type[EnterpriseStandardLoggingPayloadSetup] - ] = EnterpriseStandardLoggingPayloadSetup -except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Unable to import GenericAPILogger - LiteLLM Enterprise Feature - {str(e)}" - ) - GenericAPILogger = CustomLogger # type: ignore - ResendEmailLogger = CustomLogger # type: ignore - SendGridEmailLogger = CustomLogger # type: ignore - SMTPEmailLogger = CustomLogger # type: ignore - PagerDutyAlerting = CustomLogger # type: ignore - EnterpriseCallbackControls = None # type: ignore - EnterpriseStandardLoggingPayloadSetupVAR = None -_in_memory_loggers: List[Any] = [] - -_STANDARD_LOGGING_METADATA_KEYS: frozenset = frozenset( - StandardLoggingMetadata.__annotations__.keys() -) - -### GLOBAL VARIABLES ### - -# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys -_CUSTOM_PRICING_KEYS: frozenset = frozenset( - CustomPricingLiteLLMParams.model_fields.keys() -) - -sentry_sdk_instance = None -capture_exception = None -add_breadcrumb = None -slack_app = None -alerts_channel = None -heliconeLogger = None -athinaLogger = None -promptLayerLogger = None -logfireLogger = None -weightsBiasesLogger = None -customLogger = None -langFuseLogger = None -openMeterLogger = None -lagoLogger = None -dataDogLogger = None -prometheusLogger = None -dynamoLogger = None -s3Logger = None -greenscaleLogger = None -lunaryLogger = None -supabaseClient = None -deepevalLogger = None -callback_list: Optional[List[str]] = [] -user_logger_fn = None -additional_details: Optional[Dict[str, str]] = {} -local_cache: Optional[Dict[str, str]] = {} -last_fetched_at = None -last_fetched_at_keys = None - - -#### -class ServiceTraceIDCache: - def __init__(self) -> None: - self.cache = InMemoryCache() - - def get_cache(self, litellm_call_id: str, service_name: str) -> Optional[str]: - key_name = "{}:{}".format(service_name, litellm_call_id) - response = self.cache.get_cache(key=key_name) - return response - - def set_cache(self, litellm_call_id: str, service_name: str, trace_id: str) -> None: - key_name = "{}:{}".format(service_name, litellm_call_id) - self.cache.set_cache(key=key_name, value=trace_id) - return None - - -in_memory_trace_id_cache = ServiceTraceIDCache() -in_memory_dynamic_logger_cache = DynamicLoggingCache() - -# Cached lazy import for PrometheusLogger -# Module-level cache to avoid repeated imports while preserving memory benefits -_PrometheusLogger = None - - -def _get_cached_prometheus_logger(): - """ - Get cached PrometheusLogger class. - Lazy imports on first call to avoid loading prometheus.py and utils.py at import time (60MB saved). - Subsequent calls use cached class for better performance. - """ - global _PrometheusLogger - if _PrometheusLogger is None: - from litellm.integrations.prometheus import PrometheusLogger - - _PrometheusLogger = PrometheusLogger - return _PrometheusLogger - - -class Logging(LiteLLMLoggingBaseClass): - global supabaseClient, promptLayerLogger, weightsBiasesLogger, logfireLogger, capture_exception, add_breadcrumb, lunaryLogger, logfireLogger, prometheusLogger, slack_app - custom_pricing: bool = False - stream_options = None - litellm_request_debug: bool = False - - def __init__( - self, - model: str, - messages, - stream, - call_type, - start_time, - litellm_call_id: str, - function_id: str, - litellm_trace_id: Optional[str] = None, - dynamic_input_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - dynamic_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - dynamic_async_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - dynamic_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - dynamic_async_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = None, - applied_guardrails: Optional[List[str]] = None, - kwargs: Optional[Dict] = None, - log_raw_request_response: bool = False, - ): - _input: Optional[str] = messages # save original value of messages - if messages is not None: - if isinstance(messages, str): - messages = [ - {"role": "user", "content": messages} - ] # convert text completion input to the chat completion format - elif ( - isinstance(messages, list) - and len(messages) > 0 - and isinstance(messages[0], str) - ): - new_messages = [] - for m in messages: - new_messages.append({"role": "user", "content": m}) - messages = new_messages - - self.model = model - self.messages = copy.deepcopy(messages) if messages is not None else None - self.stream = stream - self.start_time = start_time # log the call start time - self.call_type = call_type - self.litellm_call_id = litellm_call_id - self.litellm_trace_id: str = ( - litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) - ) - self.function_id = function_id - self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[ - Any - ] = [] # for generating complete stream response - self.log_raw_request_response = log_raw_request_response - - # Initialize dynamic callbacks - self.dynamic_input_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_input_callbacks - self.dynamic_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_success_callbacks - self.dynamic_async_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_async_success_callbacks - self.dynamic_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_failure_callbacks - self.dynamic_async_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = dynamic_async_failure_callbacks - - # Process dynamic callbacks - self.process_dynamic_callbacks() - - ## DYNAMIC LANGFUSE / GCS / logging callback KEYS ## - self.standard_callback_dynamic_params: StandardCallbackDynamicParams = ( - self.initialize_standard_callback_dynamic_params(kwargs) - ) - self.standard_built_in_tools_params: StandardBuiltInToolsParams = ( - self.initialize_standard_built_in_tools_params(kwargs) - ) - ## TIME TO FIRST TOKEN LOGGING ## - self.completion_start_time: Optional[datetime.datetime] = None - self._llm_caching_handler: Optional[LLMCachingHandler] = None - - # INITIAL LITELLM_PARAMS - litellm_params = {} - if kwargs is not None: - litellm_params = get_litellm_params(**kwargs) - litellm_params = scrub_sensitive_keys_in_metadata(litellm_params) - - self.litellm_params = litellm_params - - # Initialize cost breakdown field - self.cost_breakdown: Optional[CostBreakdown] = None - - # Init Caching related details - self.caching_details: Optional[CachingDetails] = None - - # Passthrough endpoint guardrails config for field targeting - self.passthrough_guardrails_config: Optional[Dict[str, Any]] = None - - self.model_call_details: Dict[str, Any] = { - "litellm_trace_id": litellm_trace_id, - "litellm_call_id": litellm_call_id, - "input": _input, - "litellm_params": litellm_params, - "applied_guardrails": applied_guardrails, - "model": model, - } - - def process_dynamic_callbacks(self): - """ - Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks - - If a callback is in litellm._known_custom_logger_compatible_callbacks, it needs to be intialized and added to the respective dynamic_* callback list. - """ - # Process input callbacks - self.dynamic_input_callbacks = self._process_dynamic_callback_list( - self.dynamic_input_callbacks, dynamic_callbacks_type="input" - ) - - # Process failure callbacks - self.dynamic_failure_callbacks = self._process_dynamic_callback_list( - self.dynamic_failure_callbacks, dynamic_callbacks_type="failure" - ) - - # Process async failure callbacks - self.dynamic_async_failure_callbacks = self._process_dynamic_callback_list( - self.dynamic_async_failure_callbacks, dynamic_callbacks_type="async_failure" - ) - - # Process success callbacks - self.dynamic_success_callbacks = self._process_dynamic_callback_list( - self.dynamic_success_callbacks, dynamic_callbacks_type="success" - ) - - # Process async success callbacks - self.dynamic_async_success_callbacks = self._process_dynamic_callback_list( - self.dynamic_async_success_callbacks, dynamic_callbacks_type="async_success" - ) - - def _process_dynamic_callback_list( - self, - callback_list: Optional[List[Union[str, Callable, CustomLogger]]], - dynamic_callbacks_type: Literal[ - "input", "success", "failure", "async_success", "async_failure" - ], - ) -> Optional[List[Union[str, Callable, CustomLogger]]]: - """ - Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks - - - If a callback is in litellm._known_custom_logger_compatible_callbacks, - replace the string with the initialized callback class. - - If dynamic callback is a "success" callback that is a known_custom_logger_compatible_callbacks then add it to dynamic_async_success_callbacks - - If dynamic callback is a "failure" callback that is a known_custom_logger_compatible_callbacks then add it to dynamic_failure_callbacks - """ - if callback_list is None: - return None - - processed_list: List[Union[str, Callable, CustomLogger]] = [] - for callback in callback_list: - if ( - isinstance(callback, str) - and callback in litellm._known_custom_logger_compatible_callbacks - ): - callback_class = _init_custom_logger_compatible_class( - callback, internal_usage_cache=None, llm_router=None # type: ignore - ) - if callback_class is not None: - processed_list.append(callback_class) - - # If processing dynamic_success_callbacks, add to dynamic_async_success_callbacks - if dynamic_callbacks_type == "success": - if self.dynamic_async_success_callbacks is None: - self.dynamic_async_success_callbacks = [] - self.dynamic_async_success_callbacks.append(callback_class) - elif dynamic_callbacks_type == "failure": - if self.dynamic_async_failure_callbacks is None: - self.dynamic_async_failure_callbacks = [] - self.dynamic_async_failure_callbacks.append(callback_class) - else: - processed_list.append(callback) - return processed_list - - def initialize_standard_callback_dynamic_params( - self, kwargs: Optional[Dict] = None - ) -> StandardCallbackDynamicParams: - """ - Initialize the standard callback dynamic params from the kwargs - - checks if langfuse_secret_key, gcs_bucket_name in kwargs and sets the corresponding attributes in StandardCallbackDynamicParams - """ - - return _initialize_standard_callback_dynamic_params(kwargs) - - def initialize_standard_built_in_tools_params( - self, kwargs: Optional[Dict] = None - ) -> StandardBuiltInToolsParams: - """ - Initialize the standard built-in tools params from the kwargs - - checks if web_search_options in kwargs or tools and sets the corresponding attribute in StandardBuiltInToolsParams - """ - return StandardBuiltInToolsParams( - web_search_options=StandardBuiltInToolCostTracking._get_web_search_options( - kwargs or {} - ), - file_search=StandardBuiltInToolCostTracking._get_file_search_tool_call( - kwargs or {} - ), - ) - - def update_environment_variables( - self, - litellm_params: Dict, - optional_params: Dict, - model: Optional[str] = None, - user: Optional[str] = None, - **additional_params, - ): - self.optional_params = optional_params - if model is not None: - self.model = model - self.user = user - self.litellm_params = { - **self.litellm_params, - **scrub_sensitive_keys_in_metadata(litellm_params), - } - self.litellm_request_debug = litellm_params.get("litellm_request_debug", False) - self.logger_fn = litellm_params.get("logger_fn", None) - if _is_debugging_on() or self.litellm_request_debug: - verbose_logger.debug(f"self.optional_params: {self.optional_params}") - - self.model_call_details.update( - { - "model": self.model, - "messages": self.messages, - "optional_params": self.optional_params, - "litellm_params": self.litellm_params, - "start_time": self.start_time, - "stream": self.stream, - "user": user, - "call_type": str(self.call_type), - "litellm_call_id": self.litellm_call_id, - "completion_start_time": self.completion_start_time, - "standard_callback_dynamic_params": self.standard_callback_dynamic_params, - **self.optional_params, - **additional_params, - } - ) - - ## check if stream options is set ## - used by CustomStreamWrapper for easy instrumentation - if "stream_options" in additional_params: - self.stream_options = additional_params["stream_options"] - ## check if custom pricing set ## - if any( - litellm_params.get(key) is not None - for key in _CUSTOM_PRICING_KEYS & litellm_params.keys() - ): - self.custom_pricing = True - - if "custom_llm_provider" in self.model_call_details: - self.custom_llm_provider = self.model_call_details["custom_llm_provider"] - - def update_messages(self, messages: List[AllMessageValues]): - """ - Update the logged value of the messages in the model_call_details - - Allows pre-call hooks to update the messages before the call is made - """ - self.messages = messages - self.model_call_details["messages"] = messages - - def should_run_prompt_management_hooks( - self, - non_default_params: Dict, - prompt_id: Optional[str] = None, - tools: Optional[List[Dict]] = None, - ) -> bool: - """ - Return True if prompt management hooks should be run - """ - if prompt_id: - return True - - if self._should_run_prompt_management_hooks_without_prompt_id( - non_default_params=non_default_params, - tools=tools, - ): - return True - - return False - - def _should_run_prompt_management_hooks_without_prompt_id( - self, - non_default_params: Dict, - tools: Optional[List[Dict]] = None, - ) -> bool: - """ - Certain prompt management hooks don't need a `prompt_id` to be passed in, they are triggered by dynamic params - - eg. AnthropicCacheControlHook and BedrockKnowledgeBaseHook both don't require a `prompt_id` to be passed in, they are triggered by dynamic params - """ - for param in non_default_params: - if param in DynamicPromptManagementParamLiteral.list_all_params(): - return True - - ############################################################################# - # Check if Vector Store / Knowledge Base hooks should be applied to the prompt - ############################################################################# - if litellm.vector_store_registry is not None: - if litellm.vector_store_registry.get_vector_store_to_run( - non_default_params=non_default_params, tools=tools - ): - return True - return False - - def get_chat_completion_prompt( - self, - model: str, - messages: List[AllMessageValues], - non_default_params: Dict, - prompt_variables: Optional[dict], - prompt_id: Optional[str] = None, - prompt_spec: Optional[PromptSpec] = None, - prompt_management_logger: Optional[CustomLogger] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ) -> Tuple[str, List[AllMessageValues], dict]: - custom_logger = ( - prompt_management_logger - or self.get_custom_logger_for_prompt_management( - model=model, - non_default_params=non_default_params, - prompt_id=prompt_id, - prompt_spec=prompt_spec, - dynamic_callback_params=self.standard_callback_dynamic_params, - ) - ) - - if custom_logger: - ( - model, - messages, - non_default_params, - ) = custom_logger.get_chat_completion_prompt( - model=model, - messages=messages, - non_default_params=non_default_params or {}, - prompt_id=prompt_id, - prompt_spec=prompt_spec, - prompt_variables=prompt_variables, - dynamic_callback_params=self.standard_callback_dynamic_params, - prompt_label=prompt_label, - prompt_version=prompt_version, - ) - self.messages = messages - return model, messages, non_default_params - - async def async_get_chat_completion_prompt( - self, - model: str, - messages: List[AllMessageValues], - non_default_params: Dict, - prompt_variables: Optional[dict], - prompt_id: Optional[str] = None, - prompt_spec: Optional[PromptSpec] = None, - prompt_management_logger: Optional[CustomLogger] = None, - tools: Optional[List[Dict]] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ) -> Tuple[str, List[AllMessageValues], dict]: - custom_logger = ( - prompt_management_logger - or self.get_custom_logger_for_prompt_management( - model=model, - tools=tools, - non_default_params=non_default_params, - prompt_id=prompt_id, - prompt_spec=prompt_spec, - dynamic_callback_params=self.standard_callback_dynamic_params, - ) - ) - - if custom_logger: - ( - model, - messages, - non_default_params, - ) = await custom_logger.async_get_chat_completion_prompt( - model=model, - messages=messages, - non_default_params=non_default_params or {}, - prompt_id=prompt_id, - prompt_spec=prompt_spec, - prompt_variables=prompt_variables, - dynamic_callback_params=self.standard_callback_dynamic_params, - litellm_logging_obj=self, - tools=tools, - prompt_label=prompt_label, - prompt_version=prompt_version, - ) - self.messages = messages - return model, messages, non_default_params - - def _auto_detect_prompt_management_logger( - self, - prompt_id: str, - prompt_spec: Optional[PromptSpec], - dynamic_callback_params: StandardCallbackDynamicParams, - ) -> Optional[CustomLogger]: - """ - Auto-detect which prompt management system owns the given prompt_id. - - This allows a user to just pass prompt_id in the completion call and it will be auto-detected which system owns this prompt. - - Args: - prompt_id: The prompt ID to check - dynamic_callback_params: Dynamic callback parameters for should_run_prompt_management checks - - Returns: - A CustomLogger instance if a matching prompt management system is found, None otherwise - """ - prompt_management_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=CustomPromptManagement - ) - ) - - for logger in prompt_management_loggers: - if isinstance(logger, CustomPromptManagement): - try: - if logger.should_run_prompt_management( - prompt_id=prompt_id, - prompt_spec=prompt_spec, - dynamic_callback_params=dynamic_callback_params, - ): - self.model_call_details[ - "prompt_integration" - ] = logger.__class__.__name__ - return logger - except Exception: - # If check fails, continue to next logger - continue - - return None - - def get_custom_logger_for_prompt_management( - self, - model: str, - non_default_params: Dict, - tools: Optional[List[Dict]] = None, - prompt_id: Optional[str] = None, - prompt_spec: Optional[PromptSpec] = None, - dynamic_callback_params: Optional[StandardCallbackDynamicParams] = None, - ) -> Optional[CustomLogger]: - """ - Get a custom logger for prompt management based on model name or available callbacks. - - Args: - model: The model name to check for prompt management integration - non_default_params: Non-default parameters passed to the completion call - tools: Optional tools passed to the completion call - prompt_id: Optional prompt ID to auto-detect which system owns this prompt - dynamic_callback_params: Dynamic callback parameters for should_run_prompt_management checks - - Returns: - A CustomLogger instance if one is found, None otherwise - """ - # First check if model starts with a known custom logger compatible callback - # This takes precedence for backward compatibility - for callback_name in litellm._known_custom_logger_compatible_callbacks: - if model.startswith(callback_name): - custom_logger = _init_custom_logger_compatible_class( - logging_integration=callback_name, - internal_usage_cache=None, - llm_router=None, - ) - if custom_logger is not None: - self.model_call_details["prompt_integration"] = model.split("/")[0] - return custom_logger - - # If prompt_id is provided, try to auto-detect which system has this prompt - if prompt_id and dynamic_callback_params is not None: - auto_detected_logger = self._auto_detect_prompt_management_logger( - prompt_id=prompt_id, - prompt_spec=prompt_spec, - dynamic_callback_params=dynamic_callback_params, - ) - if auto_detected_logger is not None: - return auto_detected_logger - - # Then check for any registered CustomPromptManagement loggers (fallback) - prompt_management_loggers = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=CustomPromptManagement - ) - ) - - if prompt_management_loggers: - logger = prompt_management_loggers[0] - self.model_call_details["prompt_integration"] = logger.__class__.__name__ - return logger - - if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( - non_default_params - ): - self.model_call_details[ - "prompt_integration" - ] = anthropic_cache_control_logger.__class__.__name__ - return anthropic_cache_control_logger - - ######################################################### - # Vector Store / Knowledge Base hooks - ######################################################### - if litellm.vector_store_registry is not None: - vector_store_custom_logger = _init_custom_logger_compatible_class( - logging_integration="vector_store_pre_call_hook", - internal_usage_cache=None, - llm_router=None, - ) - self.model_call_details[ - "prompt_integration" - ] = vector_store_custom_logger.__class__.__name__ - # Add to global callbacks so post-call hooks are invoked - if ( - vector_store_custom_logger - and vector_store_custom_logger not in litellm.callbacks - ): - litellm.logging_callback_manager.add_litellm_callback( - vector_store_custom_logger - ) - return vector_store_custom_logger - - return None - - def get_custom_logger_for_anthropic_cache_control_hook( - self, non_default_params: Dict - ) -> Optional[CustomLogger]: - if non_default_params.get("cache_control_injection_points", None): - custom_logger = _init_custom_logger_compatible_class( - logging_integration="anthropic_cache_control_hook", - internal_usage_cache=None, - llm_router=None, - ) - return custom_logger - return None - - def _get_raw_request_body(self, data: Optional[Union[dict, str]]) -> dict: - if data is None: - return {"error": "Received empty dictionary for raw request body"} - if isinstance(data, str): - try: - return json.loads(data) - except Exception: - return { - "error": "Unable to parse raw request body. Got - {}".format(data) - } - return data - - def _get_masked_api_base(self, api_base: str) -> str: - if "key=" in api_base: - # Find the position of "key=" in the string - key_index = api_base.find("key=") + 4 - # Mask the last 5 characters after "key=" - masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:] - else: - masked_api_base = api_base - return str(masked_api_base) - - def _pre_call(self, input, api_key, model=None, additional_args={}): - """ - Common helper function across the sync + async pre-call function - """ - - self.model_call_details["input"] = input - self.model_call_details["api_key"] = api_key - self.model_call_details["additional_args"] = additional_args - self.model_call_details["log_event_type"] = "pre_api_call" - if ( - model - ): # if model name was changes pre-call, overwrite the initial model call name with the new one - self.model_call_details["model"] = model - self.model_call_details["litellm_params"][ - "api_base" - ] = self._get_masked_api_base(additional_args.get("api_base", "")) - - def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 - # Log the exact input to the LLM API - litellm.error_logs["PRE_CALL"] = locals() - try: - self._pre_call( - input=input, - api_key=api_key, - model=model, - additional_args=additional_args, - ) - - # User Logging -> if you pass in a custom logging function - self._print_llm_call_debugging_log( - api_base=additional_args.get("api_base", ""), - headers=additional_args.get("headers", {}), - additional_args=additional_args, - ) - # log raw request to provider (like LangFuse) -- if opted in. - if ( - self.log_raw_request_response is True - or log_raw_request_response is True - ): - _litellm_params = self.model_call_details.get("litellm_params", {}) - _metadata = _litellm_params.get("metadata", {}) or {} - try: - # [Non-blocking Extra Debug Information in metadata] - if turn_off_message_logging is True: - _metadata[ - "raw_request" - ] = "redacted by litellm. \ - 'litellm.turn_off_message_logging=True'" - else: - curl_command = self._get_request_curl_command( - api_base=additional_args.get("api_base", ""), - headers=additional_args.get("headers", {}), - additional_args=additional_args, - data=additional_args.get("complete_input_dict", {}), - ) - - _metadata["raw_request"] = str(curl_command) - # split up, so it's easier to parse in the UI - self.model_call_details[ - "raw_request_typed_dict" - ] = RawRequestTypedDict( - raw_request_api_base=str( - additional_args.get("api_base") or "" - ), - raw_request_body=self._get_raw_request_body( - additional_args.get("complete_input_dict", {}) - ), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, - ) - except Exception as e: - self.model_call_details[ - "raw_request_typed_dict" - ] = RawRequestTypedDict( - error=str(e), - ) - _metadata[ - "raw_request" - ] = "Unable to Log \ - raw request: {}".format( - str(e) - ) - if getattr(self, "logger_fn", None) and callable(self.logger_fn): - try: - self.logger_fn( - self.model_call_details - ) # Expectation: any logger function passed in by the user should accept a dict object - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) - ) - - self.model_call_details["api_call_start_time"] = datetime.datetime.now() - # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks = litellm.input_callback + (self.dynamic_input_callbacks or []) - for callback in callbacks: - try: - if callback == "supabase" and supabaseClient is not None: - verbose_logger.debug("reaches supabase for logging!") - model = self.model_call_details["model"] - messages = self.model_call_details["input"] - verbose_logger.debug(f"supabaseClient: {supabaseClient}") - supabaseClient.input_log_event( - model=model, - messages=messages, - end_user=self.model_call_details.get("user", "default"), - litellm_call_id=self.litellm_params["litellm_call_id"], - print_verbose=print_verbose, - ) - elif callback == "sentry" and add_breadcrumb: - try: - details_to_log = copy.deepcopy(self.model_call_details) - except Exception: - details_to_log = self.model_call_details - if litellm.turn_off_message_logging: - # make a copy of the _model_Call_details and log it - details_to_log.pop("messages", None) - details_to_log.pop("input", None) - details_to_log.pop("prompt", None) - - add_breadcrumb( - category="litellm.llm_call", - message=f"Model Call Details pre-call: {details_to_log}", - level="info", - ) - - elif isinstance(callback, CustomLogger): # custom logger class - callback.log_pre_api_call( - model=self.model, - messages=self.messages, - kwargs=self.model_call_details, - ) - elif ( - callable(callback) and customLogger is not None - ): # custom logger functions - customLogger.log_input_event( - model=self.model, - messages=self.messages, - kwargs=self.model_call_details, - print_verbose=print_verbose, - callback_func=callback, - ) - except Exception as e: - verbose_logger.exception( - "litellm.Logging.pre_call(): Exception occured - {}".format( - str(e) - ) - ) - verbose_logger.debug( - f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" - ) - if capture_exception: # log this error to sentry for debugging - capture_exception(e) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) - ) - verbose_logger.error( - f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" - ) - if capture_exception: # log this error to sentry for debugging - capture_exception(e) - - def _print_llm_call_debugging_log( - self, - api_base: str, - headers: dict, - additional_args: dict, - ): - """ - Internal debugging helper function - - Prints the RAW curl command sent from LiteLLM - """ - if _is_debugging_on() or self.litellm_request_debug: - if json_logs: - masked_headers = self._get_masked_headers(headers) - if self.litellm_request_debug: - verbose_logger.warning( # .warning ensures this shows up in all environments - "POST Request Sent from LiteLLM", - extra={"api_base": {api_base}, **masked_headers}, - ) - else: - verbose_logger.debug( - "POST Request Sent from LiteLLM", - extra={"api_base": {api_base}, **masked_headers}, - ) - else: - headers = additional_args.get("headers", {}) - if headers is None: - headers = {} - data = additional_args.get("complete_input_dict", {}) - api_base = str(additional_args.get("api_base", "")) - curl_command = self._get_request_curl_command( - api_base=api_base, - headers=headers, - additional_args=additional_args, - data=data, - ) - if self.litellm_request_debug: - verbose_logger.warning( - f"\033[92m{curl_command}\033[0m\n" - ) # .warning ensures this shows up in all environments - else: - verbose_logger.debug(f"\033[92m{curl_command}\033[0m\n") - - def _get_request_body(self, data: dict) -> str: - return str(data) - - def _get_request_curl_command( - self, api_base: str, headers: Optional[dict], additional_args: dict, data: dict - ) -> str: - masked_api_base = self._get_masked_api_base(api_base) - if headers is None: - headers = {} - curl_command = "\n\nPOST Request Sent from LiteLLM:\n" - curl_command += "curl -X POST \\\n" - curl_command += f"{masked_api_base} \\\n" - masked_headers = self._get_masked_headers(headers) - formatted_headers = " ".join( - [f"-H '{k}: {v}'" for k, v in masked_headers.items()] - ) - curl_command += ( - f"{formatted_headers} \\\n" if formatted_headers.strip() != "" else "" - ) - curl_command += f"-d '{self._get_request_body(data)}'\n" - if additional_args.get("request_str", None) is not None: - # print the sagemaker / bedrock client request - curl_command = "\nRequest Sent from LiteLLM:\n" - request_str = additional_args.get("request_str", "") - curl_command += request_str - elif api_base == "": - curl_command = str(self.model_call_details) - return curl_command - - def _get_masked_headers( - self, headers: dict, ignore_sensitive_headers: bool = False - ) -> dict: - """ - Internal debugging helper function - - Masks the headers of the request sent from LiteLLM - """ - return _get_masked_values( - headers, ignore_sensitive_values=ignore_sensitive_headers - ) - - def post_call( - self, original_response, input=None, api_key=None, additional_args={} - ): - # Log the exact result from the LLM API, for streaming - log the type of response received - litellm.error_logs["POST_CALL"] = locals() - if isinstance(original_response, dict): - original_response = json.dumps(original_response) - try: - self.model_call_details["input"] = input - self.model_call_details["api_key"] = api_key - self.model_call_details["original_response"] = original_response - self.model_call_details["additional_args"] = additional_args - self.model_call_details["log_event_type"] = "post_api_call" - - if self.litellm_request_debug: - attr = "warning" - else: - attr = "debug" - - if json_logs: - callattr = getattr(verbose_logger, attr) - callattr( - "RAW RESPONSE:\n{}\n\n".format( - self.model_call_details.get( - "original_response", self.model_call_details - ) - ), - ) - else: - callattr = getattr(verbose_logger, attr) - callattr( - "RAW RESPONSE:\n{}\n\n".format( - self.model_call_details.get( - "original_response", self.model_call_details - ) - ) - ) - if getattr(self, "logger_fn", None) and callable(self.logger_fn): - try: - self.logger_fn( - self.model_call_details - ) # Expectation: any logger function passed in by the user should accept a dict object - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) - ) - original_response = redact_message_input_output_from_logging( - model_call_details=( - self.model_call_details - if hasattr(self, "model_call_details") - else {} - ), - result=original_response, - ) - # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - - callbacks = litellm.input_callback + (self.dynamic_input_callbacks or []) - for callback in callbacks: - try: - if callback == "sentry" and add_breadcrumb: - verbose_logger.debug("reaches sentry breadcrumbing") - try: - details_to_log = copy.deepcopy(self.model_call_details) - except Exception: - details_to_log = self.model_call_details - if litellm.turn_off_message_logging: - # make a copy of the _model_Call_details and log it - details_to_log.pop("messages", None) - details_to_log.pop("input", None) - details_to_log.pop("prompt", None) - - add_breadcrumb( - category="litellm.llm_call", - message=f"Model Call Details post-call: {details_to_log}", - level="info", - ) - elif isinstance(callback, CustomLogger): # custom logger class - callback.log_post_api_call( - kwargs=self.model_call_details, - response_obj=None, - start_time=self.start_time, - end_time=None, - ) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while post-call logging with integrations {}".format( - str(e) - ) - ) - verbose_logger.debug( - f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" - ) - if capture_exception: # log this error to sentry for debugging - capture_exception(e) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) - ) - - async def async_post_mcp_tool_call_hook( - self, - kwargs: dict, - response_obj: Any, - start_time: datetime.datetime, - end_time: datetime.datetime, - ): - """ - Post MCP Tool Call Hook - - Use this to modify the MCP tool call response before it is returned to the user. - """ - from litellm.types.llms.base import HiddenParams - from litellm.types.mcp import MCPPostCallResponseObject - - callbacks = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_success_callbacks, - global_callbacks=litellm.success_callback, - ) - post_mcp_tool_call_response_obj: MCPPostCallResponseObject = ( - MCPPostCallResponseObject( - mcp_tool_call_response=response_obj, hidden_params=HiddenParams() - ) - ) - for callback in callbacks: - try: - if isinstance(callback, CustomLogger): - response: Optional[ - MCPPostCallResponseObject - ] = await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, - ) - ###################################################################### - # if any of the callbacks modify the response, use the modified response - # current implementation returns the first modified response - ###################################################################### - if response is not None: - response_obj = self._parse_post_mcp_call_hook_response( - response=response - ) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( - str(e) - ) - ) - return response_obj - - def _parse_post_mcp_call_hook_response( - self, response: Optional[MCPPostCallResponseObject] - ) -> Any: - """ - Parse the response from the post_mcp_tool_call_hook - - 1. Unpack the mcp_tool_call_response - 2. save the updated response_cost to the model_call_details - """ - if response is None: - return None - self.model_call_details["response_cost"] = response.hidden_params.response_cost - return response.mcp_tool_call_response - - def get_response_ms(self) -> float: - return ( - self.model_call_details.get("end_time", datetime.datetime.now()) - - self.model_call_details.get("start_time", datetime.datetime.now()) - ).total_seconds() * 1000 - - def set_cost_breakdown( - self, - input_cost: float, - output_cost: float, - total_cost: float, - cost_for_built_in_tools_cost_usd_dollar: float, - additional_costs: Optional[dict] = None, - original_cost: Optional[float] = None, - discount_percent: Optional[float] = None, - discount_amount: Optional[float] = None, - margin_percent: Optional[float] = None, - margin_fixed_amount: Optional[float] = None, - margin_total_amount: Optional[float] = None, - ) -> None: - """ - Helper method to store cost breakdown in the logging object. - - Args: - input_cost: Cost of input/prompt tokens - output_cost: Cost of output/completion tokens - cost_for_built_in_tools_cost_usd_dollar: Cost of built-in tools - total_cost: Total cost of request - additional_costs: Free-form additional costs dict (e.g., {"azure_model_router_flat_cost": 0.00014}) - original_cost: Cost before discount - discount_percent: Discount percentage (0.05 = 5%) - discount_amount: Discount amount in USD - margin_percent: Margin percentage applied (0.10 = 10%) - margin_fixed_amount: Fixed margin amount in USD - margin_total_amount: Total margin added in USD - """ - - self.cost_breakdown = CostBreakdown( - input_cost=input_cost, - output_cost=output_cost, - total_cost=total_cost, - tool_usage_cost=cost_for_built_in_tools_cost_usd_dollar, - ) - - # Store additional costs if provided (free-form dict for extensibility) - if ( - additional_costs - and isinstance(additional_costs, dict) - and len(additional_costs) > 0 - ): - self.cost_breakdown["additional_costs"] = additional_costs - - # Store discount information if provided - if original_cost is not None: - self.cost_breakdown["original_cost"] = original_cost - if discount_percent is not None: - self.cost_breakdown["discount_percent"] = discount_percent - if discount_amount is not None: - self.cost_breakdown["discount_amount"] = discount_amount - - # Store margin information if provided - if margin_percent is not None: - self.cost_breakdown["margin_percent"] = margin_percent - if margin_fixed_amount is not None: - self.cost_breakdown["margin_fixed_amount"] = margin_fixed_amount - if margin_total_amount is not None: - self.cost_breakdown["margin_total_amount"] = margin_total_amount - - def _response_cost_calculator( - self, - result: Union[ - ModelResponse, - ModelResponseStream, - EmbeddingResponse, - ImageResponse, - TranscriptionResponse, - TextCompletionResponse, - HttpxBinaryResponseContent, - RerankResponse, - Batch, - FineTuningJob, - ResponsesAPIResponse, - ResponseCompletedEvent, - OpenAIFileObject, - LiteLLMRealtimeStreamLoggingObject, - OpenAIModerationResponse, - "SearchResponse", - ], - cache_hit: Optional[bool] = None, - litellm_model_name: Optional[str] = None, - router_model_id: Optional[str] = None, - ) -> Optional[float]: - """ - Calculate response cost using result + logging object variables. - - used for consistent cost calculation across response headers + logging integrations. - """ - - if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"): - hidden_params = getattr(result, "_hidden_params", {}) - if ( - "response_cost" in hidden_params - and hidden_params["response_cost"] is not None - ): # use cost if already calculated - return hidden_params["response_cost"] - elif ( - router_model_id is None and "model_id" in hidden_params - ): # use model_id if not already set - router_model_id = hidden_params["model_id"] - - ## RESPONSE COST ## - custom_pricing = use_custom_pricing_for_model( - litellm_params=( - self.litellm_params if hasattr(self, "litellm_params") else None - ) - ) - - prompt = "" # use for tts cost calc - _input = self.model_call_details.get("input", None) - if _input is not None and isinstance(_input, str): - prompt = _input - - if cache_hit is None: - cache_hit = self.model_call_details.get("cache_hit", False) - - try: - response_cost_calculator_kwargs = { - "response_object": result, - "model": litellm_model_name or self.model, - "cache_hit": cache_hit, - "custom_llm_provider": self.model_call_details.get( - "custom_llm_provider", None - ), - "base_model": _get_base_model_from_metadata( - model_call_details=self.model_call_details - ), - "call_type": self.call_type, - "optional_params": self.optional_params, - "custom_pricing": custom_pricing, - "prompt": prompt, - "standard_built_in_tools_params": self.standard_built_in_tools_params, - "router_model_id": router_model_id, - "litellm_logging_obj": self, - "service_tier": ( - self.optional_params.get("service_tier") - if self.optional_params - else None - ), - } - except Exception as e: # error creating kwargs for cost calculation - debug_info = StandardLoggingModelCostFailureDebugInformation( - error_str=str(e), - traceback_str=_get_traceback_str_for_error(str(e)), - ) - verbose_logger.debug( - f"response_cost_failure_debug_information: {debug_info}" - ) - self.model_call_details[ - "response_cost_failure_debug_information" - ] = debug_info - return None - - try: - response_cost = litellm.response_cost_calculator( - **response_cost_calculator_kwargs - ) - - verbose_logger.debug(f"response_cost: {response_cost}") - return response_cost - except Exception as e: # error calculating cost - debug_info = StandardLoggingModelCostFailureDebugInformation( - error_str=str(e), - traceback_str=_get_traceback_str_for_error(str(e)), - model=response_cost_calculator_kwargs["model"], - cache_hit=response_cost_calculator_kwargs["cache_hit"], - custom_llm_provider=response_cost_calculator_kwargs[ - "custom_llm_provider" - ], - base_model=response_cost_calculator_kwargs["base_model"], - call_type=response_cost_calculator_kwargs["call_type"], - custom_pricing=response_cost_calculator_kwargs["custom_pricing"], - ) - verbose_logger.debug( - f"response_cost_failure_debug_information: {debug_info}" - ) - self.model_call_details[ - "response_cost_failure_debug_information" - ] = debug_info - - return None - - async def _response_cost_calculator_async( - self, - result: Union[ - ModelResponse, - ModelResponseStream, - EmbeddingResponse, - ImageResponse, - TranscriptionResponse, - TextCompletionResponse, - HttpxBinaryResponseContent, - RerankResponse, - Batch, - FineTuningJob, - ], - cache_hit: Optional[bool] = None, - ) -> Optional[float]: - return self._response_cost_calculator(result=result, cache_hit=cache_hit) - - def should_run_logging( - self, - event_type: Literal[ - "async_success", "sync_success", "async_failure", "sync_failure" - ], - stream: bool = False, - ) -> bool: - try: - if self.model_call_details.get(f"has_logged_{event_type}", False) is True: - return False - - return True - except Exception: - return True - - def has_run_logging( - self, - event_type: Literal[ - "async_success", "sync_success", "async_failure", "sync_failure" - ], - ) -> None: - if self.stream is not None and self.stream is True: - """ - Ignore check on stream, as there can be multiple chunks - """ - return - self.model_call_details[f"has_logged_{event_type}"] = True - return - - def should_run_callback( - self, callback: litellm.CALLBACK_TYPES, litellm_params: dict, event_hook: str - ) -> bool: - if litellm.global_disable_no_log_param: - return True - - if litellm_params.get("no-log", False) is True: - # proxy cost tracking cal backs should run - - if not ( - isinstance(callback, CustomLogger) - and "_PROXY_" in callback.__class__.__name__ - ): - verbose_logger.debug( - f"no-log request, skipping logging for {event_hook} event" - ) - return False - - # Check for dynamically disabled callbacks via headers - if ( - EnterpriseCallbackControls is not None - and EnterpriseCallbackControls.is_callback_disabled_dynamically( - callback=callback, - litellm_params=litellm_params, - standard_callback_dynamic_params=self.standard_callback_dynamic_params, - ) - ): - verbose_logger.debug( - f"Callback {callback} disabled via x-litellm-disable-callbacks header for {event_hook} event" - ) - return False - - return True - - def _update_completion_start_time(self, completion_start_time: datetime.datetime): - self.completion_start_time = completion_start_time - self.model_call_details["completion_start_time"] = self.completion_start_time - - def normalize_logging_result(self, result: Any) -> Any: - """ - Some endpoints return a different type of result than what is expected by the logging system. - This function is used to normalize the result to the expected type. - """ - logging_result = result - if self.call_type == CallTypes.arealtime.value and isinstance(result, list): - combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( - results=result - ) - logging_result = ( - RealtimeAPITokenUsageProcessor.create_logging_realtime_object( - usage=combined_usage_object, - results=result, - ) - ) - - elif ( - self.call_type == CallTypes.llm_passthrough_route.value - or self.call_type == CallTypes.allm_passthrough_route.value - ) and isinstance(result, Response): - from litellm.utils import ProviderConfigManager - - provider_config = ProviderConfigManager.get_provider_passthrough_config( - provider=self.model_call_details.get("custom_llm_provider", ""), - model=self.model, - ) - if provider_config is not None: - logging_result = provider_config.logging_non_streaming_response( - model=self.model, - custom_llm_provider=self.model_call_details.get( - "custom_llm_provider", "" - ), - httpx_response=result, - request_data=self.model_call_details.get("request_data", {}), - logging_obj=self, - endpoint=self.model_call_details.get("endpoint", ""), - ) - return logging_result - - def _process_hidden_params_and_response_cost( - self, - logging_result, - start_time, - end_time, - ): - hidden_params = getattr(logging_result, "_hidden_params", {}) - if hidden_params: - if self.model_call_details.get("litellm_params") is not None: - self.model_call_details["litellm_params"].setdefault("metadata", {}) - if self.model_call_details["litellm_params"]["metadata"] is None: - self.model_call_details["litellm_params"]["metadata"] = {} - self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = getattr(logging_result, "_hidden_params", {}) # type: ignore - - if "response_cost" in hidden_params: - self.model_call_details["response_cost"] = hidden_params["response_cost"] - else: - self.model_call_details["response_cost"] = self._response_cost_calculator( - result=logging_result - ) - - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=logging_result, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) - - def _transform_usage_objects(self, result): - if isinstance(result, ResponsesAPIResponse): - result = result.model_copy() - transformed_usage = ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - result.usage - ) - ) - setattr(result, "usage", transformed_usage) - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: - response_dict = ( - result.model_dump() - if hasattr(result, "model_dump") - else dict(result) - ) - # Ensure usage is properly included with transformed chat format - if transformed_usage is not None: - response_dict["usage"] = ( - transformed_usage.model_dump() - if hasattr(transformed_usage, "model_dump") - else dict(transformed_usage) - ) - standard_logging_payload["response"] = response_dict - elif isinstance(result, TranscriptionResponse): - from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( - TranscriptionUsageObjectTransformation, - ) - - result = result.model_copy() - transformed_usage = TranscriptionUsageObjectTransformation.transform_transcription_usage_object(result.usage) # type: ignore - setattr(result, "usage", transformed_usage) - return result - - def _success_handler_helper_fn( - self, - result=None, - start_time=None, - end_time=None, - cache_hit=None, - standard_logging_object: Optional[StandardLoggingPayload] = None, - ): - try: - if start_time is None: - start_time = self.start_time - if end_time is None: - end_time = datetime.datetime.now() - if self.completion_start_time is None: - self.completion_start_time = end_time - self.model_call_details[ - "completion_start_time" - ] = self.completion_start_time - - self.model_call_details["log_event_type"] = "successful_api_call" - self.model_call_details["end_time"] = end_time - self.model_call_details["cache_hit"] = cache_hit - - if self.call_type == CallTypes.anthropic_messages.value: - result = self._handle_anthropic_messages_response_logging(result=result) - elif ( - self.call_type == CallTypes.generate_content.value - or self.call_type == CallTypes.agenerate_content.value - ): - result = self._handle_non_streaming_google_genai_generate_content_response_logging( - result=result - ) - elif ( - self.call_type == CallTypes.asend_message.value - or self.call_type == CallTypes.send_message.value - ): - result = self._handle_a2a_response_logging(result=result) - - logging_result = self.normalize_logging_result(result=result) - - if ( - standard_logging_object is None - and result is not None - and self.stream is not True - ): - if self._is_recognized_call_type_for_logging( - logging_result=logging_result - ): - self._process_hidden_params_and_response_cost( - logging_result=logging_result, - start_time=start_time, - end_time=end_time, - ) - elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=result, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) - elif standard_logging_object is not None: - self.model_call_details[ - "standard_logging_object" - ] = standard_logging_object - else: - self.model_call_details["response_cost"] = None - - result = self._transform_usage_objects(result=result) - - if ( - litellm.max_budget - and self.stream is False - and result is not None - and isinstance(result, dict) - and "content" in result - ): - time_diff = (end_time - start_time).total_seconds() - float_diff = float(time_diff) - litellm._current_cost += litellm.completion_cost( - model=self.model, - prompt="", - completion=getattr(result, "content", ""), - total_time=float_diff, - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) - - return start_time, end_time, result - except Exception as e: - raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {str(e)}") - - def _is_recognized_call_type_for_logging( - self, - logging_result: Any, - ): - """ - Returns True if the call type is recognized for logging (eg. ModelResponse, ModelResponseStream, etc.) - """ - if ( - isinstance(logging_result, ModelResponse) - or isinstance(logging_result, ModelResponseStream) - or isinstance(logging_result, EmbeddingResponse) - or isinstance(logging_result, ImageResponse) - or isinstance(logging_result, TranscriptionResponse) - or isinstance(logging_result, TextCompletionResponse) - or isinstance(logging_result, HttpxBinaryResponseContent) # tts - or isinstance(logging_result, RerankResponse) - or isinstance(logging_result, FineTuningJob) - or isinstance(logging_result, LiteLLMBatch) - or isinstance(logging_result, ResponsesAPIResponse) - or isinstance(logging_result, OpenAIFileObject) - or isinstance(logging_result, LiteLLMRealtimeStreamLoggingObject) - or isinstance(logging_result, OpenAIModerationResponse) - or isinstance(logging_result, OCRResponse) # OCR - or isinstance(logging_result, SearchResponse) # Search API - or isinstance(logging_result, dict) - and logging_result.get("object") == "vector_store.search_results.page" - or isinstance(logging_result, dict) - and logging_result.get("object") == "search" # Search API (dict format) - or isinstance(logging_result, VideoObject) - or isinstance(logging_result, ContainerObject) - or isinstance(logging_result, LiteLLMSendMessageResponse) # A2A - or (self.call_type == CallTypes.call_mcp_tool.value) - ): - return True - return False - - def _flush_passthrough_collected_chunks_helper( - self, - raw_bytes: List[bytes], - provider_config: "BasePassthroughConfig", - ) -> Optional["CostResponseTypes"]: - all_chunks = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) - complete_streaming_response = provider_config.handle_logging_collected_chunks( - all_chunks=all_chunks, - litellm_logging_obj=self, - model=self.model, - custom_llm_provider=self.model_call_details.get("custom_llm_provider", ""), - endpoint=self.model_call_details.get("endpoint", ""), - ) - return complete_streaming_response - - def flush_passthrough_collected_chunks( - self, - raw_bytes: List[bytes], - provider_config: "BasePassthroughConfig", - ): - """ - Flush collected chunks from the logging object - This is used to log the collected chunks once streaming is done on passthrough endpoints - - 1. Decode the raw bytes to string lines - 2. Get the complete streaming response from the provider config - 3. Log the complete streaming response (trigger success handler) - This is used for passthrough endpoints - """ - complete_streaming_response = self._flush_passthrough_collected_chunks_helper( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - - if complete_streaming_response is not None: - self.success_handler(result=complete_streaming_response) - return - - async def async_flush_passthrough_collected_chunks( - self, - raw_bytes: List[bytes], - provider_config: "BasePassthroughConfig", - ): - complete_streaming_response = self._flush_passthrough_collected_chunks_helper( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - - if complete_streaming_response is not None: - await self.async_success_handler(result=complete_streaming_response) - return - - def success_handler( # noqa: PLR0915 - self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs - ): - verbose_logger.debug( - f"Logging Details LiteLLM-Success Call: Cache_hit={cache_hit}" - ) - if not self.should_run_logging( - event_type="sync_success" - ): # prevent double logging - return - start_time, end_time, result = self._success_handler_helper_fn( - start_time=start_time, - end_time=end_time, - result=result, - cache_hit=cache_hit, - standard_logging_object=kwargs.get("standard_logging_object", None), - ) - litellm_params = self.model_call_details.get("litellm_params", {}) - is_sync_request = ( - litellm_params.get(CallTypes.acompletion.value, False) is not True - and litellm_params.get(CallTypes.aresponses.value, False) is not True - and litellm_params.get(CallTypes.aembedding.value, False) is not True - and litellm_params.get(CallTypes.aimage_generation.value, False) is not True - and litellm_params.get(CallTypes.atranscription.value, False) is not True - ) - try: - ## BUILD COMPLETE STREAMED RESPONSE - complete_streaming_response: Optional[ - Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse] - ] = None - if "complete_streaming_response" in self.model_call_details: - return # break out of this. - complete_streaming_response = self._get_assembled_streaming_response( - result=result, - start_time=start_time, - end_time=end_time, - is_async=False, - streaming_chunks=self.sync_streaming_chunks, - ) - if complete_streaming_response is not None: - verbose_logger.debug( - "Logging Details LiteLLM-Success Call streaming complete" - ) - self.model_call_details[ - "complete_streaming_response" - ] = complete_streaming_response - self.model_call_details[ - "response_cost" - ] = self._response_cost_calculator(result=complete_streaming_response) - ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=complete_streaming_response, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: - # Only emit for sync requests (async_success_handler handles async) - if is_sync_request: - emit_standard_logging_payload(standard_logging_payload) - callbacks = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_success_callbacks, - global_callbacks=litellm.success_callback, - ) - - ## REDACT MESSAGES ## - result = redact_message_input_output_from_logging( - model_call_details=( - self.model_call_details - if hasattr(self, "model_call_details") - else {} - ), - result=result, - ) - ## LOGGING HOOK ## - for callback in callbacks: - if isinstance(callback, CustomLogger): - self.model_call_details, result = callback.logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - - self.has_run_logging(event_type="sync_success") - for callback in callbacks: - try: - should_run = self.should_run_callback( - callback=callback, - litellm_params=litellm_params, - event_hook="success_handler", - ) - if not should_run: - continue - if callback == "promptlayer" and promptLayerLogger is not None: - print_verbose("reaches promptlayer for logging!") - promptLayerLogger.log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "supabase" and supabaseClient is not None: - print_verbose("reaches supabase for logging!") - kwargs = self.model_call_details - - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - if "complete_streaming_response" not in kwargs: - continue - else: - print_verbose("reaches supabase for streaming logging!") - result = kwargs["complete_streaming_response"] - - model = kwargs["model"] - messages = kwargs["messages"] - optional_params = kwargs.get("optional_params", {}) - litellm_params = kwargs.get("litellm_params", {}) - supabaseClient.log_event( - model=model, - messages=messages, - end_user=optional_params.get("user", "default"), - response_obj=result, - start_time=start_time, - end_time=end_time, - litellm_call_id=( - current_call_id - if ( - current_call_id := litellm_params.get( - "litellm_call_id" - ) - ) - is not None - else str(uuid.uuid4()) - ), - print_verbose=print_verbose, - ) - if callback == "wandb" and weightsBiasesLogger is not None: - print_verbose("reaches wandb for logging!") - weightsBiasesLogger.log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "logfire" and logfireLogger is not None: - verbose_logger.debug("reaches logfire for success logging!") - kwargs = {} - for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - if "complete_streaming_response" not in kwargs: - continue - else: - print_verbose("reaches logfire for streaming logging!") - result = kwargs["complete_streaming_response"] - - logfireLogger.log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - level=LogfireLevel.INFO.value, # type: ignore - ) - - if callback == "lunary" and lunaryLogger is not None: - print_verbose("reaches lunary for logging!") - model = self.model - kwargs = self.model_call_details - - input = kwargs.get("messages", kwargs.get("input", None)) - - type = ( - "embed" - if self.call_type == CallTypes.embedding.value - else "llm" - ) - - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - if "complete_streaming_response" not in kwargs: - continue - else: - result = kwargs["complete_streaming_response"] - - lunaryLogger.log_event( - type=type, - kwargs=kwargs, - event="end", - model=model, - input=input, - user_id=kwargs.get("user", None), - # user_props=self.model_call_details.get("user_props", None), - extra=kwargs.get("optional_params", {}), - response_obj=result, - start_time=start_time, - end_time=end_time, - run_id=self.litellm_call_id, - print_verbose=print_verbose, - ) - if callback == "helicone" and heliconeLogger is not None: - print_verbose("reaches helicone for logging!") - model = self.model - messages = self.model_call_details["input"] - kwargs = self.model_call_details - - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - if "complete_streaming_response" not in kwargs: - continue - else: - print_verbose("reaches helicone for streaming logging!") - result = kwargs["complete_streaming_response"] - - heliconeLogger.log_success( - model=model, - messages=messages, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - kwargs=kwargs, - ) - if callback == "langfuse": - global langFuseLogger - print_verbose("reaches langfuse for success logging!") - kwargs = {} - for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - verbose_logger.debug( - f"is complete_streaming_response in kwargs: {kwargs.get('complete_streaming_response', None)}" - ) - if complete_streaming_response is None: - continue - else: - print_verbose("reaches langfuse for streaming logging!") - result = kwargs["complete_streaming_response"] - - langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( - globalLangfuseLogger=langFuseLogger, - standard_callback_dynamic_params=self.standard_callback_dynamic_params, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) - if langfuse_logger_to_use is not None: - _response = langfuse_logger_to_use.log_event_on_langfuse( - kwargs=kwargs, - response_obj=result, - start_time=start_time, - end_time=end_time, - user_id=kwargs.get("user", None), - ) - if _response is not None and isinstance(_response, dict): - _trace_id = _response.get("trace_id", None) - if _trace_id is not None: - in_memory_trace_id_cache.set_cache( - litellm_call_id=self.litellm_call_id, - service_name="langfuse", - trace_id=_trace_id, - ) - if callback == "greenscale" and greenscaleLogger is not None: - kwargs = {} - for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - verbose_logger.debug( - f"is complete_streaming_response in kwargs: {kwargs.get('complete_streaming_response', None)}" - ) - if complete_streaming_response is None: - continue - else: - print_verbose( - "reaches greenscale for streaming logging!" - ) - result = kwargs["complete_streaming_response"] - - greenscaleLogger.log_event( - kwargs=kwargs, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "athina" and athinaLogger is not None: - deep_copy = {} - for k, v in self.model_call_details.items(): - deep_copy[k] = v - athinaLogger.log_event( - kwargs=deep_copy, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "traceloop": - deep_copy = {} - for k, v in self.model_call_details.items(): - if k != "original_response": - deep_copy[k] = v - traceloopLogger.log_event( - kwargs=deep_copy, - response_obj=result, - start_time=start_time, - end_time=end_time, - user_id=kwargs.get("user", None), - print_verbose=print_verbose, - ) - if callback == "s3": - global s3Logger - if s3Logger is None: - s3Logger = S3Logger() - if self.stream: - if "complete_streaming_response" in self.model_call_details: - print_verbose( - "S3Logger Logger: Got Stream Event - Completed Stream Response" - ) - s3Logger.log_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details[ - "complete_streaming_response" - ], - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - else: - print_verbose( - "S3Logger Logger: Got Stream Event - No complete stream response as yet" - ) - else: - s3Logger.log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - - if callback == "openmeter" and is_sync_request: - global openMeterLogger - if openMeterLogger is None: - print_verbose("Instantiates openmeter client") - openMeterLogger = OpenMeterLogger() - if self.stream and complete_streaming_response is None: - openMeterLogger.log_stream_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - else: - if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} - ) - result = self.model_call_details["complete_response"] - openMeterLogger.log_success_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - if ( - isinstance(callback, CustomLogger) - and is_sync_request - and self.call_type - != CallTypes.pass_through.value # pass-through endpoints call async_log_success_event - ): # custom logger class - if self.stream and complete_streaming_response is None: - callback.log_stream_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - else: - if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} - ) - result = self.model_call_details["complete_response"] - - callback.log_success_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - if ( - callable(callback) is True - and is_sync_request - and customLogger is not None - ): # custom logger functions - print_verbose( - "success callbacks: Running Custom Callback Function - {}".format( - callback - ) - ) - - customLogger.log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - - except Exception as e: - print_verbose( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging with integrations {traceback.format_exc()}" - ) - print_verbose( - f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" - ) - if capture_exception: # log this error to sentry for debugging - capture_exception(e) - # Track callback logging failures in Prometheus - try: - self._handle_callback_failure(callback=callback) - except Exception: - pass - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {}".format( - str(e) - ), - ) - - async def async_success_handler( # noqa: PLR0915 - self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs - ): - """ - Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. - """ - print_verbose( - "Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit) - ) - if not self.should_run_logging( - event_type="async_success" - ): # prevent double logging - return - - ## CALCULATE COST FOR BATCH JOBS - if self.call_type == CallTypes.aretrieve_batch.value and isinstance( - result, LiteLLMBatch - ): - litellm_params = self.litellm_params or {} - litellm_metadata = litellm_params.get("litellm_metadata") or {} - if ( - litellm_metadata.get("batch_ignore_default_logging", False) is True - ): # polling job will query these frequently, don't spam db logs - return - - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) - - # check if file id is a unified file id - is_base64_unified_file_id = _is_base64_encoded_unified_file_id(result.id) - - batch_cost = kwargs.get("batch_cost", None) - batch_usage = kwargs.get("batch_usage", None) - batch_models = kwargs.get("batch_models", None) - has_explicit_batch_data = all( - x is not None for x in (batch_cost, batch_usage, batch_models) - ) - - should_compute_batch_data = ( - not is_base64_unified_file_id - or not has_explicit_batch_data - and result.status == "completed" - ) - if has_explicit_batch_data: - result._hidden_params["response_cost"] = batch_cost - result._hidden_params["batch_models"] = batch_models - result.usage = batch_usage - - elif should_compute_batch_data: - ( - response_cost, - batch_usage, - batch_models, - ) = await _handle_completed_batch( - batch=result, - custom_llm_provider=self.custom_llm_provider, - litellm_params=self.litellm_params, - ) - - result._hidden_params["response_cost"] = response_cost - result._hidden_params["batch_models"] = batch_models - result.usage = batch_usage - - start_time, end_time, result = self._success_handler_helper_fn( - start_time=start_time, - end_time=end_time, - result=result, - cache_hit=cache_hit, - standard_logging_object=kwargs.get("standard_logging_object", None), - ) - - ## BUILD COMPLETE STREAMED RESPONSE - if "async_complete_streaming_response" in self.model_call_details: - return # break out of this. - complete_streaming_response: Optional[ - Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse] - ] = self._get_assembled_streaming_response( - result=result, - start_time=start_time, - end_time=end_time, - is_async=True, - streaming_chunks=self.streaming_chunks, - ) - - if complete_streaming_response is not None: - print_verbose("Async success callbacks: Got a complete streaming response") - - self.model_call_details[ - "async_complete_streaming_response" - ] = complete_streaming_response - - try: - if self.model_call_details.get("cache_hit", False) is True: - self.model_call_details["response_cost"] = 0.0 - else: - # check if base_model set on azure - _get_base_model_from_metadata( - model_call_details=self.model_call_details - ) - # base_model defaults to None if not set on model_info - self.model_call_details[ - "response_cost" - ] = self._response_cost_calculator( - result=complete_streaming_response - ) - - verbose_logger.debug( - f"Model={self.model}; cost={self.model_call_details['response_cost']}" - ) - except litellm.NotFoundError: - verbose_logger.warning( - f"Model={self.model} not found in completion cost map. Setting 'response_cost' to None" - ) - self.model_call_details["response_cost"] = None - - ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=complete_streaming_response, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) - - # print standard logging payload - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: - emit_standard_logging_payload(standard_logging_payload) - elif self.call_type == "pass_through_endpoint": - print_verbose( - "Async success callbacks: Got a pass-through endpoint response" - ) - - self.model_call_details["async_complete_streaming_response"] = result - - # cost calculation not possible for pass-through - self.model_call_details["response_cost"] = None - - ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=result, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) - - # print standard logging payload - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: - emit_standard_logging_payload(standard_logging_payload) - callbacks = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_async_success_callbacks, - global_callbacks=litellm._async_success_callback, - ) - - result = redact_message_input_output_from_logging( - model_call_details=( - self.model_call_details if hasattr(self, "model_call_details") else {} - ), - result=result, - ) - - ## LOGGING HOOK ## - - for callback in callbacks: - if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks - - if ( - callback.should_run_guardrail( - data=self.model_call_details, - event_type=GuardrailEventHooks.logging_only, - ) - is not True - ): - continue - - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - elif isinstance(callback, CustomLogger): - result = redact_message_input_output_from_custom_logger( - result=result, litellm_logging_obj=self, custom_logger=callback - ) - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - - self.has_run_logging(event_type="async_success") - - for callback in callbacks: - # check if callback can run for this request - litellm_params = self.model_call_details.get("litellm_params", {}) - should_run = self.should_run_callback( - callback=callback, - litellm_params=litellm_params, - event_hook="async_success_handler", - ) - if not should_run: - continue - try: - if callback == "openmeter" and openMeterLogger is not None: - if self.stream is True: - if ( - "async_complete_streaming_response" - in self.model_call_details - ): - await openMeterLogger.async_log_success_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details[ - "async_complete_streaming_response" - ], - start_time=start_time, - end_time=end_time, - ) - else: - await openMeterLogger.async_log_stream_event( # [TODO]: move this to being an async log stream event function - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - else: - await openMeterLogger.async_log_success_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - - if isinstance(callback, CustomLogger): # custom logger class - model_call_details: Dict = self.model_call_details - ################################## - # call redaction hook for custom logger - model_call_details = callback.redact_standard_logging_payload_from_model_call_details( - model_call_details=model_call_details - ) - ################################## - if self.stream is True: - if "async_complete_streaming_response" in model_call_details: - await callback.async_log_success_event( - kwargs=model_call_details, - response_obj=model_call_details[ - "async_complete_streaming_response" - ], - start_time=start_time, - end_time=end_time, - ) - else: - await callback.async_log_stream_event( # [TODO]: move this to being an async log stream event function - kwargs=model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - else: - await callback.async_log_success_event( - kwargs=model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - if callable(callback): # custom logger functions - global customLogger - if customLogger is None: - customLogger = CustomLogger() - if self.stream: - if ( - "async_complete_streaming_response" - in self.model_call_details - ): - await customLogger.async_log_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details[ - "async_complete_streaming_response" - ], - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - else: - await customLogger.async_log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - if callback == "dynamodb": - global dynamoLogger - if dynamoLogger is None: - dynamoLogger = DyanmoDBLogger() - if self.stream: - if ( - "async_complete_streaming_response" - in self.model_call_details - ): - print_verbose( - "DynamoDB Logger: Got Stream Event - Completed Stream Response" - ) - await dynamoLogger._async_log_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details[ - "async_complete_streaming_response" - ], - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - else: - print_verbose( - "DynamoDB Logger: Got Stream Event - No complete stream response as yet" - ) - else: - await dynamoLogger._async_log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - except Exception: - verbose_logger.error( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {traceback.format_exc()}" - ) - self._handle_callback_failure(callback=callback) - pass - - def _handle_callback_failure(self, callback: Any): - """ - Handle callback logging failures by incrementing Prometheus metrics. - - Works for both sync and async contexts since Prometheus counter increment is synchronous. - - Args: - callback: The callback that failed - """ - try: - callback_name = self._get_callback_name(callback) - - all_callbacks = litellm.logging_callback_manager._get_all_callbacks() - - for callback_obj in all_callbacks: - if hasattr(callback_obj, "increment_callback_logging_failure"): - callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore - break # Only increment once - - except Exception as e: - verbose_logger.debug(f"Error in _handle_callback_failure: {str(e)}") - - def _failure_handler_helper_fn( - self, exception, traceback_exception, start_time=None, end_time=None - ): - if start_time is None: - start_time = self.start_time - if end_time is None: - end_time = datetime.datetime.now() - - # on some exceptions, model_call_details is not always initialized, this ensures that we still log those exceptions - if not hasattr(self, "model_call_details"): - self.model_call_details = {} - - self.model_call_details["log_event_type"] = "failed_api_call" - self.model_call_details["exception"] = exception - self.model_call_details["traceback_exception"] = traceback_exception - self.model_call_details["end_time"] = end_time - self.model_call_details.setdefault("original_response", None) - self.model_call_details["response_cost"] = 0 - - if hasattr(exception, "headers") and isinstance(exception.headers, dict): - self.model_call_details.setdefault("litellm_params", {}) - metadata = ( - self.model_call_details["litellm_params"].get("metadata", {}) or {} - ) - metadata.update(exception.headers) - - ## STANDARDIZED LOGGING PAYLOAD - - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj={}, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="failure", - error_str=str(exception), - original_exception=exception, - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) - return start_time, end_time - - async def special_failure_handlers(self, exception: Exception): - """ - Custom events, emitted for specific failures. - - Currently just for router model group rate limit error - """ - from litellm.types.router import RouterErrors - - litellm_params: dict = self.model_call_details.get("litellm_params") or {} - metadata = litellm_params.get("metadata") or {} - - ## BASE CASE ## check if rate limit error for model group size 1 - is_base_case = False - if metadata.get("model_group_size") is not None: - model_group_size = metadata.get("model_group_size") - if isinstance(model_group_size, int) and model_group_size == 1: - is_base_case = True - ## check if special error ## - if ( - RouterErrors.no_deployments_available.value not in str(exception) - and is_base_case is False - ): - return - - ## get original model group ## - - model_group = metadata.get("model_group") or None - for callback in litellm._async_failure_callback: - if isinstance(callback, CustomLogger): # custom logger class - await callback.log_model_group_rate_limit_error( - exception=exception, - original_model_group=model_group, - kwargs=self.model_call_details, - ) # type: ignore - - def failure_handler( # noqa: PLR0915 - self, exception, traceback_exception, start_time=None, end_time=None - ): - verbose_logger.debug( - f"Logging Details LiteLLM-Failure Call: {litellm.failure_callback}" - ) - if not self.should_run_logging( - event_type="sync_failure" - ): # prevent double logging - return - litellm_params = self.model_call_details.get("litellm_params", {}) - is_sync_request = ( - litellm_params.get(CallTypes.acompletion.value, False) is not True - and litellm_params.get(CallTypes.aresponses.value, False) is not True - and litellm_params.get(CallTypes.aembedding.value, False) is not True - and litellm_params.get(CallTypes.aimage_generation.value, False) is not True - and litellm_params.get(CallTypes.atranscription.value, False) is not True - ) - - try: - start_time, end_time = self._failure_handler_helper_fn( - exception=exception, - traceback_exception=traceback_exception, - start_time=start_time, - end_time=end_time, - ) - callbacks = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_failure_callbacks, - global_callbacks=litellm.failure_callback, - ) - - result = None # result sent to all loggers, init this to None incase it's not created - - result = redact_message_input_output_from_logging( - model_call_details=( - self.model_call_details - if hasattr(self, "model_call_details") - else {} - ), - result=result, - ) - self.has_run_logging(event_type="sync_failure") - for callback in callbacks: - try: - should_run = self.should_run_callback( - callback=callback, - litellm_params=litellm_params, - event_hook="failure_handler", - ) - if not should_run: - continue - if callback == "lunary" and lunaryLogger is not None: - print_verbose("reaches lunary for logging error!") - - model = self.model - - input = self.model_call_details["input"] - - _type = ( - "embed" - if self.call_type == CallTypes.embedding.value - else "llm" - ) - - lunaryLogger.log_event( - kwargs=self.model_call_details, - type=_type, - event="error", - user_id=self.model_call_details.get("user", "default"), - model=model, - input=input, - error=traceback_exception, - run_id=self.litellm_call_id, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "sentry": - print_verbose("sending exception to sentry") - if capture_exception: - capture_exception(exception) - else: - print_verbose( - f"capture exception not initialized: {capture_exception}" - ) - elif callback == "supabase" and supabaseClient is not None: - print_verbose("reaches supabase for logging!") - print_verbose(f"supabaseClient: {supabaseClient}") - supabaseClient.log_event( - model=self.model if hasattr(self, "model") else "", - messages=self.messages, - end_user=self.model_call_details.get("user", "default"), - response_obj=result, - start_time=start_time, - end_time=end_time, - litellm_call_id=self.model_call_details["litellm_call_id"], - print_verbose=print_verbose, - ) - if ( - callable(callback) and customLogger is not None - ): # custom logger functions - customLogger.log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - if ( - isinstance(callback, CustomLogger) and is_sync_request - ): # custom logger class - callback.log_failure_event( - start_time=start_time, - end_time=end_time, - response_obj=result, - kwargs=self.model_call_details, - ) - if callback == "langfuse": - global langFuseLogger - verbose_logger.debug("reaches langfuse for logging failure") - kwargs = {} - for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( - globalLangfuseLogger=langFuseLogger, - standard_callback_dynamic_params=self.standard_callback_dynamic_params, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) - _response = langfuse_logger_to_use.log_event_on_langfuse( - start_time=start_time, - end_time=end_time, - response_obj=None, - user_id=kwargs.get("user", None), - status_message=str(exception), - level="ERROR", - kwargs=self.model_call_details, - ) - if _response is not None and isinstance(_response, dict): - _trace_id = _response.get("trace_id", None) - if _trace_id is not None: - in_memory_trace_id_cache.set_cache( - litellm_call_id=self.litellm_call_id, - service_name="langfuse", - trace_id=_trace_id, - ) - if callback == "traceloop": - traceloopLogger.log_event( - start_time=start_time, - end_time=end_time, - response_obj=None, - user_id=self.model_call_details.get("user", None), - print_verbose=print_verbose, - status_message=str(exception), - level="ERROR", - kwargs=self.model_call_details, - ) - if callback == "logfire" and logfireLogger is not None: - verbose_logger.debug("reaches logfire for failure logging!") - kwargs = {} - for k, v in self.model_call_details.items(): - if ( - k != "original_response" - ): # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - kwargs["exception"] = exception - - logfireLogger.log_event( - kwargs=kwargs, - response_obj=result, - start_time=start_time, - end_time=end_time, - level=LogfireLevel.ERROR.value, # type: ignore - print_verbose=print_verbose, - ) - - except Exception as e: - print_verbose( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging with integrations {str(e)}" - ) - print_verbose( - f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" - ) - if capture_exception: # log this error to sentry for debugging - capture_exception(e) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging {}".format( - str(e) - ) - ) - - async def async_failure_handler( - self, exception, traceback_exception, start_time=None, end_time=None - ): - """ - Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. - """ - await self.special_failure_handlers(exception=exception) - if not self.should_run_logging( - event_type="async_failure" - ): # prevent double logging - return - start_time, end_time = self._failure_handler_helper_fn( - exception=exception, - traceback_exception=traceback_exception, - start_time=start_time, - end_time=end_time, - ) - - callbacks = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_async_failure_callbacks, - global_callbacks=litellm._async_failure_callback, - ) - - result = None # result sent to all loggers, init this to None incase it's not created - - self.has_run_logging(event_type="async_failure") - for callback in callbacks: - try: - litellm_params = self.model_call_details.get("litellm_params", {}) - should_run = self.should_run_callback( - callback=callback, - litellm_params=litellm_params, - event_hook="async_failure_handler", - ) - if not should_run: - continue - if isinstance(callback, CustomLogger): # custom logger class - await callback.async_log_failure_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) # type: ignore - if ( - callable(callback) and customLogger is not None - ): # custom logger functions - await customLogger.async_log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure \ - logging {}\nCallback={}".format( - str(e), callback - ) - ) - # Track callback logging failures in Prometheus - self._handle_callback_failure(callback=callback) - - def _get_trace_id(self, service_name: Literal["langfuse"]) -> Optional[str]: - """ - For the given service (e.g. langfuse), return the trace_id actually logged. - - Used for constructing the url in slack alerting. - - Returns: - - str: The logged trace id - - None: If trace id not yet emitted. - """ - trace_id: Optional[str] = None - if service_name == "langfuse": - trace_id = in_memory_trace_id_cache.get_cache( - litellm_call_id=self.litellm_call_id, service_name=service_name - ) - - return trace_id - - def _get_callback_object(self, service_name: Literal["langfuse"]) -> Optional[Any]: - """ - Return dynamic callback object. - - Meant to solve issue when doing key-based/team-based logging - """ - global langFuseLogger - - if service_name == "langfuse": - if langFuseLogger is None or ( - ( - self.standard_callback_dynamic_params.get("langfuse_public_key") - is not None - and self.standard_callback_dynamic_params.get("langfuse_public_key") - != langFuseLogger.public_key - ) - or ( - self.standard_callback_dynamic_params.get("langfuse_public_key") - is not None - and self.standard_callback_dynamic_params.get("langfuse_public_key") - != langFuseLogger.public_key - ) - or ( - self.standard_callback_dynamic_params.get("langfuse_host") - is not None - and self.standard_callback_dynamic_params.get("langfuse_host") - != langFuseLogger.langfuse_host - ) - ): - return LangFuseLogger( - langfuse_public_key=self.standard_callback_dynamic_params.get( - "langfuse_public_key" - ), - langfuse_secret=self.standard_callback_dynamic_params.get( - "langfuse_secret" - ), - langfuse_host=self.standard_callback_dynamic_params.get( - "langfuse_host" - ), - ) - return langFuseLogger - - return None - - def handle_sync_success_callbacks_for_async_calls( - self, - result: Any, - start_time: datetime.datetime, - end_time: datetime.datetime, - cache_hit: Optional[Any] = None, - ) -> None: - """ - Handles calling success callbacks for Async calls. - - Why: Some callbacks - `langfuse`, `s3` are sync callbacks. We need to call them in the executor. - """ - if self._should_run_sync_callbacks_for_async_calls() is False: - return - - executor.submit( - self.success_handler, - result, - start_time, - end_time, - cache_hit, - ) - - def _should_run_sync_callbacks_for_async_calls(self) -> bool: - """ - Returns: - - bool: True if sync callbacks should be run for async calls. eg. `langfuse`, `s3` - """ - _combined_sync_callbacks = self.get_combined_callback_list( - dynamic_success_callbacks=self.dynamic_success_callbacks, - global_callbacks=litellm.success_callback, - ) - _filtered_success_callbacks = self._remove_internal_custom_logger_callbacks( - _combined_sync_callbacks - ) - _filtered_success_callbacks = self._remove_internal_litellm_callbacks( - _filtered_success_callbacks - ) - return len(_filtered_success_callbacks) > 0 - - def get_combined_callback_list( - self, dynamic_success_callbacks: Optional[List], global_callbacks: List - ) -> List: - if dynamic_success_callbacks is None: - return list(global_callbacks) - return list(set(dynamic_success_callbacks + global_callbacks)) - - def _remove_internal_litellm_callbacks(self, callbacks: List) -> List: - """ - Creates a filtered list of callbacks, excluding internal LiteLLM callbacks. - - Args: - callbacks: List of callback functions/strings to filter - - Returns: - List of filtered callbacks with internal ones removed - """ - filtered = [ - cb for cb in callbacks if not self._is_internal_litellm_proxy_callback(cb) - ] - - verbose_logger.debug(f"Filtered callbacks: {filtered}") - return filtered - - def _get_callback_name(self, cb) -> str: - """ - Helper to get the name of a callback function - - Args: - cb: The callback object/function/string to get the name of - - Returns: - The name of the callback - """ - if isinstance(cb, str): - return cb - if hasattr(cb, "__name__"): - return cb.__name__ - if hasattr(cb, "__func__"): - return cb.__func__.__name__ - if hasattr(cb, "__class__"): - return cb.__class__.__name__ - return str(cb) - - def _is_internal_litellm_proxy_callback(self, cb) -> bool: - """Helper to check if a callback is internal""" - INTERNAL_PREFIXES = [ - "_PROXY", - "_service_logger.ServiceLogging", - "sync_deployment_callback_on_success", - ] - if isinstance(cb, str): - return False - - if not callable(cb): - return True - - cb_name = self._get_callback_name(cb) - return any(prefix in cb_name for prefix in INTERNAL_PREFIXES) - - def _remove_internal_custom_logger_callbacks(self, callbacks: List) -> List: - """ - Removes internal custom logger callbacks from the list. - """ - _new_callbacks = [] - for _c in callbacks: - if isinstance(_c, CustomLogger): - continue - elif ( - isinstance(_c, str) - and _c in litellm._known_custom_logger_compatible_callbacks - ): - continue - _new_callbacks.append(_c) - return _new_callbacks - - def _get_assembled_streaming_response( - self, - result: Union[ - ModelResponse, - TextCompletionResponse, - ModelResponseStream, - ResponseCompletedEvent, - Any, - ], - start_time: datetime.datetime, - end_time: datetime.datetime, - is_async: bool, - streaming_chunks: List[Any], - ) -> Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]]: - if isinstance(result, ModelResponse): - return result - elif isinstance(result, TextCompletionResponse): - return result - elif isinstance(result, ResponseCompletedEvent): - ## return unified Usage object - if isinstance(result.response.usage, ResponseAPIUsage): - transformed_usage = ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - result.response.usage - ) - ) - # Set as dict instead of Usage object so model_dump() serializes it correctly - setattr( - result.response, - "usage", - ( - transformed_usage.model_dump() - if hasattr(transformed_usage, "model_dump") - else dict(transformed_usage) - ), - ) - return result.response - else: - return None - return None - - def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse: - """ - Handles logging for Anthropic messages responses. - - Args: - result: The response object from the model call - - Returns: - The the response object from the model call - - - For Non-streaming responses, we need to transform the response to a ModelResponse object. - - For streaming responses, anthropic_messages handler calls success_handler with a assembled ModelResponse. - """ - import httpx - - if self.stream and isinstance(result, ModelResponse): - return result - elif isinstance(result, ModelResponse): - return result - - httpx_response = self.model_call_details.get("httpx_response", None) - if httpx_response and isinstance(httpx_response, httpx.Response): - result = litellm.AnthropicConfig().transform_response( - raw_response=httpx_response, - model_response=litellm.ModelResponse(), - model=self.model, - messages=[], - logging_obj=self, - optional_params={}, - api_key="", - request_data={}, - encoding=litellm.encoding, - json_mode=False, - litellm_params={}, - ) - else: - from litellm.types.llms.anthropic import AnthropicResponse - - pydantic_result = AnthropicResponse.model_validate(result) - import httpx - - result = litellm.AnthropicConfig().transform_parsed_response( - completion_response=pydantic_result.model_dump(), - raw_response=httpx.Response( - status_code=200, - headers={}, - ), - model_response=litellm.ModelResponse(), - json_mode=None, - ) - return result - - def _handle_non_streaming_google_genai_generate_content_response_logging( - self, result: Any - ) -> ModelResponse: - """ - Handles logging for Google GenAI generate content responses. - """ - import httpx - - httpx_response = self.model_call_details.get("httpx_response", None) - if httpx_response is None: - raise ValueError("Google GenAI Generate Content: httpx_response is None") - dict_result = httpx_response.json() - result = litellm.VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response( - completion_response=dict_result, - model_response=litellm.ModelResponse(), - model=self.model, - logging_obj=self, - raw_response=httpx.Response( - status_code=200, - headers={}, - ), - ) - return result - - def _handle_a2a_response_logging(self, result: Any) -> Any: - """ - Handles logging for A2A (Agent-to-Agent) responses. - - Adds usage from model_call_details to the result if available. - Uses Pydantic's model_copy to avoid modifying the original response. - - Args: - result: The LiteLLMSendMessageResponse from the A2A call - - Returns: - The response object with usage added if available - """ - # Get usage from model_call_details (set by asend_message) - usage = self.model_call_details.get("usage") - if usage is None: - return result - - # Deep copy result and add usage - result_copy = result.model_copy(deep=True) - result_copy.usage = ( - usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) - ) - return result_copy - - -def _get_masked_values( - sensitive_object: dict, - ignore_sensitive_values: bool = False, - mask_all_values: bool = False, - unmasked_length: int = 4, - number_of_asterisks: Optional[int] = 4, -) -> dict: - """ - Internal debugging helper function - - Masks the headers of the request sent from LiteLLM - - Args: - masked_length: Optional length for the masked portion (number of *). If set, will use exactly this many * - regardless of original string length. The total length will be unmasked_length + masked_length. - """ - sensitive_keywords = [ - "authorization", - "token", - "key", - "secret", - "vertex_credentials", - ] - return { - k: ( - # If ignore_sensitive_values is True, or if this key doesn't contain sensitive keywords, return original value - v - if ignore_sensitive_values - or not any( - sensitive_keyword in k.lower() - for sensitive_keyword in sensitive_keywords - ) - else ( - # Apply masking to sensitive keys - ( - v[: unmasked_length // 2] - + "*" * number_of_asterisks - + v[-unmasked_length // 2 :] - ) - if ( - isinstance(v, str) - and len(v) > unmasked_length - and number_of_asterisks is not None - ) - else ( - ( - v[: unmasked_length // 2] - + "*" * (len(v) - unmasked_length) - + v[-unmasked_length // 2 :] - ) - if (isinstance(v, str) and len(v) > unmasked_length) - else ("*****" if isinstance(v, str) else v) - ) - ) - ) - for k, v in sensitive_object.items() - } - - -def set_callbacks(callback_list, function_id=None): # noqa: PLR0915 - """ - Globally sets the callback client - """ - global sentry_sdk_instance, capture_exception, add_breadcrumb, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, supabaseClient, lunaryLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, logfireLogger, dynamoLogger, s3Logger, dataDogLogger, prometheusLogger, greenscaleLogger, openMeterLogger, deepevalLogger - - try: - for callback in callback_list: - if callback == "sentry": - try: - import sentry_sdk - except ImportError: - print_verbose("Package 'sentry_sdk' is missing. Installing it...") - subprocess.check_call( - [sys.executable, "-m", "pip", "install", "sentry_sdk"] - ) - import sentry_sdk - from sentry_sdk.scrubber import EventScrubber - - sentry_sdk_instance = sentry_sdk - sentry_trace_rate = ( - os.environ.get("SENTRY_API_TRACE_RATE") - if "SENTRY_API_TRACE_RATE" in os.environ - else "1.0" - ) - sentry_sample_rate = ( - os.environ.get("SENTRY_API_SAMPLE_RATE") - if "SENTRY_API_SAMPLE_RATE" in os.environ - else "1.0" - ) - sentry_sdk_instance.init( - dsn=os.environ.get("SENTRY_DSN"), - traces_sample_rate=float(sentry_trace_rate), # type: ignore - sample_rate=float( - sentry_sample_rate if sentry_sample_rate else 1.0 - ), - send_default_pii=False, # Prevent sending Personal Identifiable Information - event_scrubber=EventScrubber( - denylist=SENTRY_DENYLIST, pii_denylist=SENTRY_PII_DENYLIST - ), - environment=os.environ.get("SENTRY_ENVIRONMENT", "production"), - ) - capture_exception = sentry_sdk_instance.capture_exception - add_breadcrumb = sentry_sdk_instance.add_breadcrumb - elif callback == "slack": - try: - from slack_bolt import App - except ImportError: - print_verbose("Package 'slack_bolt' is missing. Installing it...") - subprocess.check_call( - [sys.executable, "-m", "pip", "install", "slack_bolt"] - ) - from slack_bolt import App - slack_app = App( - token=os.environ.get("SLACK_API_TOKEN"), - signing_secret=os.environ.get("SLACK_API_SECRET"), - ) - alerts_channel = os.environ["SLACK_API_CHANNEL"] - print_verbose(f"Initialized Slack App: {slack_app}") - elif callback == "traceloop": - traceloopLogger = TraceloopLogger() - elif callback == "athina": - athinaLogger = AthinaLogger() - print_verbose("Initialized Athina Logger") - elif callback == "helicone": - heliconeLogger = HeliconeLogger() - elif callback == "lunary": - lunaryLogger = LunaryLogger() - elif callback == "promptlayer": - promptLayerLogger = PromptLayerLogger() - elif callback == "langfuse": - langFuseLogger = LangFuseLogger( - langfuse_public_key=None, langfuse_secret=None, langfuse_host=None - ) - elif callback == "openmeter": - openMeterLogger = OpenMeterLogger() - elif callback == "datadog": - dataDogLogger = DataDogLogger() - elif callback == "dynamodb": - dynamoLogger = DyanmoDBLogger() - elif callback == "s3": - s3Logger = S3Logger() - elif callback == "wandb": - from litellm.integrations.weights_biases import WeightsBiasesLogger - - weightsBiasesLogger = WeightsBiasesLogger() - elif callback == "logfire": - logfireLogger = LogfireLogger() - elif callback == "supabase": - print_verbose("instantiating supabase") - supabaseClient = Supabase() - elif callback == "greenscale": - greenscaleLogger = GreenscaleLogger() - print_verbose("Initialized Greenscale Logger") - elif callable(callback): - customLogger = CustomLogger() - except Exception as e: - raise e - return None - - -def _init_custom_logger_compatible_class( # noqa: PLR0915 - logging_integration: _custom_logger_compatible_callbacks_literal, - internal_usage_cache: Optional[DualCache], - llm_router: Optional[ - Any - ], # expect litellm.Router, but typing errors due to circular import - custom_logger_init_args: Optional[dict] = {}, -) -> Optional[CustomLogger]: - """ - Initialize a custom logger compatible class - """ - try: - custom_logger_init_args = custom_logger_init_args or {} - if logging_integration == "agentops": # Add AgentOps initialization - for callback in _in_memory_loggers: - if isinstance(callback, AgentOps): - return callback # type: ignore - - agentops_logger = AgentOps() - _in_memory_loggers.append(agentops_logger) - return agentops_logger # type: ignore - elif logging_integration == "lago": - for callback in _in_memory_loggers: - if isinstance(callback, LagoLogger): - return callback # type: ignore - - lago_logger = LagoLogger() - _in_memory_loggers.append(lago_logger) - return lago_logger # type: ignore - elif logging_integration == "openmeter": - for callback in _in_memory_loggers: - if isinstance(callback, OpenMeterLogger): - return callback # type: ignore - - _openmeter_logger = OpenMeterLogger() - _in_memory_loggers.append(_openmeter_logger) - return _openmeter_logger # type: ignore - elif logging_integration == "posthog": - for callback in _in_memory_loggers: - if isinstance(callback, PostHogLogger): - return callback # type: ignore - - _posthog_logger = PostHogLogger() - _in_memory_loggers.append(_posthog_logger) - return _posthog_logger # type: ignore - elif logging_integration == "braintrust": - from litellm.integrations.braintrust_logging import BraintrustLogger - - for callback in _in_memory_loggers: - if isinstance(callback, BraintrustLogger): - return callback # type: ignore - - braintrust_logger = BraintrustLogger() - _in_memory_loggers.append(braintrust_logger) - return braintrust_logger # type: ignore - elif logging_integration == "langsmith": - for callback in _in_memory_loggers: - if isinstance(callback, LangsmithLogger): - return callback # type: ignore - - _langsmith_logger = LangsmithLogger() - _in_memory_loggers.append(_langsmith_logger) - return _langsmith_logger # type: ignore - elif logging_integration == "argilla": - for callback in _in_memory_loggers: - if isinstance(callback, ArgillaLogger): - return callback # type: ignore - - _argilla_logger = ArgillaLogger() - _in_memory_loggers.append(_argilla_logger) - return _argilla_logger # type: ignore - elif logging_integration == "literalai": - for callback in _in_memory_loggers: - if isinstance(callback, LiteralAILogger): - return callback # type: ignore - - _literalai_logger = LiteralAILogger() - _in_memory_loggers.append(_literalai_logger) - return _literalai_logger # type: ignore - elif logging_integration == "prometheus": - PrometheusLogger = _get_cached_prometheus_logger() - - for callback in _in_memory_loggers: - if isinstance(callback, PrometheusLogger): - return callback # type: ignore - - _prometheus_logger = PrometheusLogger() - _in_memory_loggers.append(_prometheus_logger) - return _prometheus_logger # type: ignore - elif logging_integration == "datadog": - for callback in _in_memory_loggers: - if isinstance(callback, DataDogLogger): - return callback # type: ignore - - _datadog_logger = DataDogLogger() - _in_memory_loggers.append(_datadog_logger) - return _datadog_logger # type: ignore - elif logging_integration == "datadog_llm_observability": - _datadog_llm_obs_logger = DataDogLLMObsLogger() - _in_memory_loggers.append(_datadog_llm_obs_logger) - return _datadog_llm_obs_logger # type: ignore - elif logging_integration == "azure_sentinel": - for callback in _in_memory_loggers: - if isinstance(callback, AzureSentinelLogger): - return callback # type: ignore - - _azure_sentinel_logger = AzureSentinelLogger() - _in_memory_loggers.append(_azure_sentinel_logger) - return _azure_sentinel_logger # type: ignore - elif logging_integration == "gcs_bucket": - for callback in _in_memory_loggers: - if isinstance(callback, GCSBucketLogger): - return callback # type: ignore - - _gcs_bucket_logger = GCSBucketLogger() - _in_memory_loggers.append(_gcs_bucket_logger) - return _gcs_bucket_logger # type: ignore - elif logging_integration == "s3_v2": - for callback in _in_memory_loggers: - if isinstance(callback, S3V2Logger): - return callback # type: ignore - - _s3_v2_logger = S3V2Logger() - _in_memory_loggers.append(_s3_v2_logger) - return _s3_v2_logger # type: ignore - elif logging_integration == "aws_sqs": - for callback in _in_memory_loggers: - if isinstance(callback, SQSLogger): - return callback # type: ignore - - _aws_sqs_logger = SQSLogger() - _in_memory_loggers.append(_aws_sqs_logger) - return _aws_sqs_logger # type: ignore - elif logging_integration == "azure_storage": - for callback in _in_memory_loggers: - if isinstance(callback, AzureBlobStorageLogger): - return callback # type: ignore - - _azure_storage_logger = AzureBlobStorageLogger() - _in_memory_loggers.append(_azure_storage_logger) - return _azure_storage_logger # type: ignore - elif logging_integration == "opik": - for callback in _in_memory_loggers: - if isinstance(callback, OpikLogger): - return callback # type: ignore - - _opik_logger = OpikLogger() - _in_memory_loggers.append(_opik_logger) - return _opik_logger # type: ignore - elif logging_integration == "arize": - from litellm.integrations.opentelemetry import ( - OpenTelemetry, - OpenTelemetryConfig, - ) - - arize_config = ArizeLogger.get_arize_config() - if arize_config.endpoint is None: - raise ValueError( - "No valid endpoint found for Arize, please set 'ARIZE_ENDPOINT' to your GRPC endpoint or 'ARIZE_HTTP_ENDPOINT' to your HTTP endpoint" - ) - otel_config = OpenTelemetryConfig( - exporter=arize_config.protocol, - endpoint=arize_config.endpoint, - service_name=arize_config.project_name, - ) - - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" - for callback in _in_memory_loggers: - if ( - isinstance(callback, ArizeLogger) - and callback.callback_name == "arize" - ): - return callback # type: ignore - _arize_otel_logger = ArizeLogger(config=otel_config, callback_name="arize") - _in_memory_loggers.append(_arize_otel_logger) - return _arize_otel_logger # type: ignore - elif logging_integration == "arize_phoenix": - from litellm.integrations.opentelemetry import ( - OpenTelemetry, - OpenTelemetryConfig, - ) - - arize_phoenix_config = ArizePhoenixLogger.get_arize_phoenix_config() - otel_config = OpenTelemetryConfig( - exporter=arize_phoenix_config.protocol, - endpoint=arize_phoenix_config.endpoint, - headers=arize_phoenix_config.otlp_auth_headers, - ) - if arize_phoenix_config.project_name: - existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") - # Add openinference.project.name attribute - if existing_attrs: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" - else: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"openinference.project.name={arize_phoenix_config.project_name}" - - # Set Phoenix project name from environment variable - phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) - if phoenix_project_name: - existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") - # Add openinference.project.name attribute - if existing_attrs: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}" - else: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"openinference.project.name={phoenix_project_name}" - - # auth can be disabled on local deployments of arize phoenix - if arize_phoenix_config.otlp_auth_headers is not None: - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = arize_phoenix_config.otlp_auth_headers - - for callback in _in_memory_loggers: - if ( - isinstance(callback, ArizePhoenixLogger) - and callback.callback_name == "arize_phoenix" - ): - return callback # type: ignore - _arize_phoenix_otel_logger = ArizePhoenixLogger( - config=otel_config, callback_name="arize_phoenix" - ) - _in_memory_loggers.append(_arize_phoenix_otel_logger) - return _arize_phoenix_otel_logger # type: ignore - elif logging_integration == "levo": - from litellm.integrations.levo.levo import LevoLogger - from litellm.integrations.opentelemetry import ( - OpenTelemetry, - OpenTelemetryConfig, - ) - - levo_config = LevoLogger.get_levo_config() - otel_config = OpenTelemetryConfig( - exporter=levo_config.protocol, - endpoint=levo_config.endpoint, - headers=levo_config.otlp_auth_headers, - ) - - # Check if LevoLogger instance already exists - for callback in _in_memory_loggers: - if ( - isinstance(callback, LevoLogger) - and callback.callback_name == "levo" - ): - return callback # type: ignore - - _levo_otel_logger = LevoLogger(config=otel_config, callback_name="levo") - _in_memory_loggers.append(_levo_otel_logger) - return _levo_otel_logger # type: ignore - elif logging_integration == "otel": - from litellm.integrations.opentelemetry import OpenTelemetry - - for callback in _in_memory_loggers: - if type(callback) is OpenTelemetry: - return callback # type: ignore - otel_logger = OpenTelemetry( - **_get_custom_logger_settings_from_proxy_server( - callback_name=logging_integration - ) - ) - _in_memory_loggers.append(otel_logger) - return otel_logger # type: ignore - - elif logging_integration == "galileo": - for callback in _in_memory_loggers: - if isinstance(callback, GalileoObserve): - return callback # type: ignore - - galileo_logger = GalileoObserve() - _in_memory_loggers.append(galileo_logger) - return galileo_logger # type: ignore - elif logging_integration == "cloudzero": - from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger - - for callback in _in_memory_loggers: - if isinstance(callback, CloudZeroLogger): - return callback # type: ignore - cloudzero_logger = CloudZeroLogger() - _in_memory_loggers.append(cloudzero_logger) - return cloudzero_logger # type: ignore - elif logging_integration == "focus": - from litellm.integrations.focus.focus_logger import FocusLogger - - for callback in _in_memory_loggers: - if isinstance(callback, FocusLogger): - return callback # type: ignore - focus_logger = FocusLogger() - _in_memory_loggers.append(focus_logger) - return focus_logger # type: ignore - elif logging_integration == "deepeval": - for callback in _in_memory_loggers: - if isinstance(callback, DeepEvalLogger): - return callback # type: ignore - deepeval_logger = DeepEvalLogger() - _in_memory_loggers.append(deepeval_logger) - return deepeval_logger # type: ignore - - elif logging_integration == "logfire": - if "LOGFIRE_TOKEN" not in os.environ: - raise ValueError("LOGFIRE_TOKEN not found in environment variables") - from litellm.integrations.opentelemetry import ( - OpenTelemetry, - OpenTelemetryConfig, - ) - - logfire_base_url = os.getenv( - "LOGFIRE_BASE_URL", "https://logfire-api.pydantic.dev" - ) - otel_config = OpenTelemetryConfig( - exporter="otlp_http", - endpoint=f"{logfire_base_url.rstrip('/')}/v1/traces", - headers=f"Authorization={os.getenv('LOGFIRE_TOKEN')}", - ) - for callback in _in_memory_loggers: - if isinstance(callback, OpenTelemetry): - return callback # type: ignore - _otel_logger = OpenTelemetry(config=otel_config) - _in_memory_loggers.append(_otel_logger) - return _otel_logger # type: ignore - elif logging_integration == "dynamic_rate_limiter": - from litellm.proxy.hooks.dynamic_rate_limiter import ( - _PROXY_DynamicRateLimitHandler, - ) - - for callback in _in_memory_loggers: - if isinstance(callback, _PROXY_DynamicRateLimitHandler): - return callback # type: ignore - - if internal_usage_cache is None: - raise Exception( - "Internal Error: Cache cannot be empty - internal_usage_cache={}".format( - internal_usage_cache - ) - ) - - dynamic_rate_limiter_obj = _PROXY_DynamicRateLimitHandler( - internal_usage_cache=internal_usage_cache - ) - - if llm_router is not None and isinstance(llm_router, litellm.Router): - dynamic_rate_limiter_obj.update_variables(llm_router=llm_router) - _in_memory_loggers.append(dynamic_rate_limiter_obj) - return dynamic_rate_limiter_obj # type: ignore - elif logging_integration == "dynamic_rate_limiter_v3": - from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( - _PROXY_DynamicRateLimitHandlerV3, - ) - - for callback in _in_memory_loggers: - if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3): - return callback # type: ignore - - if internal_usage_cache is None: - raise Exception( - "Internal Error: Cache cannot be empty - internal_usage_cache={}".format( - internal_usage_cache - ) - ) - - dynamic_rate_limiter_obj_v3 = _PROXY_DynamicRateLimitHandlerV3( - internal_usage_cache=internal_usage_cache - ) - - if llm_router is not None and isinstance(llm_router, litellm.Router): - dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router) - _in_memory_loggers.append(dynamic_rate_limiter_obj_v3) - return dynamic_rate_limiter_obj_v3 # type: ignore - elif logging_integration == "langtrace": - if "LANGTRACE_API_KEY" not in os.environ: - raise ValueError("LANGTRACE_API_KEY not found in environment variables") - - from litellm.integrations.opentelemetry import ( - OpenTelemetry, - OpenTelemetryConfig, - ) - - otel_config = OpenTelemetryConfig( - exporter="otlp_http", - endpoint="https://langtrace.ai/api/trace", - ) - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" - for callback in _in_memory_loggers: - if ( - isinstance(callback, OpenTelemetry) - and callback.callback_name == "langtrace" - ): - return callback # type: ignore - _otel_logger = OpenTelemetry(config=otel_config, callback_name="langtrace") - _in_memory_loggers.append(_otel_logger) - return _otel_logger # type: ignore - - elif logging_integration == "mlflow": - for callback in _in_memory_loggers: - if isinstance(callback, MlflowLogger): - return callback # type: ignore - - _mlflow_logger = MlflowLogger() - _in_memory_loggers.append(_mlflow_logger) - return _mlflow_logger # type: ignore - elif logging_integration == "langfuse": - for callback in _in_memory_loggers: - if isinstance(callback, LangfusePromptManagement): - return callback - - langfuse_logger = LangfusePromptManagement() - _in_memory_loggers.append(langfuse_logger) - return langfuse_logger # type: ignore - elif logging_integration == "langfuse_otel": - from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger - - for callback in _in_memory_loggers: - if ( - isinstance(callback, LangfuseOtelLogger) - and callback.callback_name == "langfuse_otel" - ): - return callback # type: ignore - # Allow LangfuseOtelLogger to initialize its own config safely - # This prevents startup crashes if LANGFUSE keys are not in env (e.g. for dynamic usage) - _otel_logger = LangfuseOtelLogger( - config=None, callback_name="langfuse_otel" - ) - _in_memory_loggers.append(_otel_logger) - return _otel_logger # type: ignore - elif logging_integration == "weave_otel": - from litellm.integrations.opentelemetry import OpenTelemetryConfig - from litellm.integrations.weave.weave_otel import ( - WeaveOtelLogger, - get_weave_otel_config, - ) - - weave_otel_config = get_weave_otel_config() - - otel_config = OpenTelemetryConfig( - exporter=weave_otel_config.protocol, - endpoint=weave_otel_config.endpoint, - headers=weave_otel_config.otlp_auth_headers, - ) - - for callback in _in_memory_loggers: - if ( - isinstance(callback, WeaveOtelLogger) - and callback.callback_name == "weave_otel" - ): - return callback # type: ignore - _otel_logger = WeaveOtelLogger( - config=otel_config, callback_name="weave_otel" - ) - _in_memory_loggers.append(_otel_logger) - return _otel_logger # type: ignore - elif logging_integration == "pagerduty": - for callback in _in_memory_loggers: - if isinstance(callback, PagerDutyAlerting): - return callback - pagerduty_logger = PagerDutyAlerting(**custom_logger_init_args) - _in_memory_loggers.append(pagerduty_logger) - return pagerduty_logger # type: ignore - elif logging_integration == "anthropic_cache_control_hook": - for callback in _in_memory_loggers: - if isinstance(callback, AnthropicCacheControlHook): - return callback - anthropic_cache_control_hook = AnthropicCacheControlHook() - _in_memory_loggers.append(anthropic_cache_control_hook) - return anthropic_cache_control_hook # type: ignore - elif logging_integration == "vector_store_pre_call_hook": - from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( - VectorStorePreCallHook, - ) - - for callback in _in_memory_loggers: - if isinstance(callback, VectorStorePreCallHook): - return callback - vector_store_pre_call_hook = VectorStorePreCallHook() - _in_memory_loggers.append(vector_store_pre_call_hook) - return vector_store_pre_call_hook # type: ignore - elif logging_integration == "gcs_pubsub": - for callback in _in_memory_loggers: - if isinstance(callback, GcsPubSubLogger): - return callback - _gcs_pubsub_logger = GcsPubSubLogger() - _in_memory_loggers.append(_gcs_pubsub_logger) - return _gcs_pubsub_logger # type: ignore - elif logging_integration == "generic_api": - for callback in _in_memory_loggers: - if isinstance(callback, GenericAPILogger): - return callback - generic_api_logger = GenericAPILogger() - _in_memory_loggers.append(generic_api_logger) - return generic_api_logger # type: ignore - elif logging_integration == "resend_email": - for callback in _in_memory_loggers: - if isinstance(callback, ResendEmailLogger): - return callback - resend_email_logger = ResendEmailLogger() - _in_memory_loggers.append(resend_email_logger) - return resend_email_logger # type: ignore - elif logging_integration == "sendgrid_email": - for callback in _in_memory_loggers: - if isinstance(callback, SendGridEmailLogger): - return callback - sendgrid_email_logger = SendGridEmailLogger() - _in_memory_loggers.append(sendgrid_email_logger) - return sendgrid_email_logger # type: ignore - elif logging_integration == "smtp_email": - for callback in _in_memory_loggers: - if isinstance(callback, SMTPEmailLogger): - return callback - smtp_email_logger = SMTPEmailLogger() - _in_memory_loggers.append(smtp_email_logger) - return smtp_email_logger # type: ignore - elif logging_integration == "humanloop": - for callback in _in_memory_loggers: - if isinstance(callback, HumanloopLogger): - return callback - - humanloop_logger = HumanloopLogger() - _in_memory_loggers.append(humanloop_logger) - return humanloop_logger # type: ignore - elif logging_integration == "dotprompt": - for callback in _in_memory_loggers: - if isinstance(callback, DotpromptManager): - return callback - - dotprompt_logger = DotpromptManager() - _in_memory_loggers.append(dotprompt_logger) - return dotprompt_logger # type: ignore - elif logging_integration == "bitbucket": - from litellm.integrations.bitbucket.bitbucket_prompt_manager import ( - BitBucketPromptManager, - ) - - for callback in _in_memory_loggers: - if isinstance(callback, BitBucketPromptManager): - return callback - - # Get global BitBucket config - bitbucket_config = getattr(litellm, "global_bitbucket_config", None) - if bitbucket_config is None: - raise ValueError( - "BitBucket configuration not found. Please set litellm.global_bitbucket_config first." - ) - - 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( - f"[Non-Blocking Error] Error initializing custom logger: {e}" - ) - return None - return None - - -def get_custom_logger_compatible_class( # noqa: PLR0915 - logging_integration: _custom_logger_compatible_callbacks_literal, -) -> Optional[CustomLogger]: - try: - if logging_integration == "lago": - for callback in _in_memory_loggers: - if isinstance(callback, LagoLogger): - return callback - elif logging_integration == "openmeter": - for callback in _in_memory_loggers: - if isinstance(callback, OpenMeterLogger): - return callback - elif logging_integration == "braintrust": - from litellm.integrations.braintrust_logging import BraintrustLogger - - for callback in _in_memory_loggers: - if isinstance(callback, BraintrustLogger): - return callback - elif logging_integration == "galileo": - for callback in _in_memory_loggers: - if isinstance(callback, GalileoObserve): - return callback - elif logging_integration == "cloudzero": - from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger - - for callback in _in_memory_loggers: - if isinstance(callback, CloudZeroLogger): - return callback - elif logging_integration == "focus": - from litellm.integrations.focus.focus_logger import FocusLogger - - for callback in _in_memory_loggers: - if isinstance(callback, FocusLogger): - return callback - elif logging_integration == "deepeval": - for callback in _in_memory_loggers: - if isinstance(callback, DeepEvalLogger): - return callback - elif logging_integration == "langsmith": - for callback in _in_memory_loggers: - if isinstance(callback, LangsmithLogger): - return callback - elif logging_integration == "argilla": - for callback in _in_memory_loggers: - if isinstance(callback, ArgillaLogger): - return callback - elif logging_integration == "literalai": - for callback in _in_memory_loggers: - if isinstance(callback, LiteralAILogger): - return callback - elif logging_integration == "prometheus": - PrometheusLogger = _get_cached_prometheus_logger() - for callback in _in_memory_loggers: - if isinstance(callback, PrometheusLogger): - return callback - elif logging_integration == "datadog": - for callback in _in_memory_loggers: - if isinstance(callback, DataDogLogger): - return callback - elif logging_integration == "datadog_llm_observability": - for callback in _in_memory_loggers: - if isinstance(callback, DataDogLLMObsLogger): - return callback - elif logging_integration == "azure_sentinel": - for callback in _in_memory_loggers: - if isinstance(callback, AzureSentinelLogger): - return callback - elif logging_integration == "gcs_bucket": - for callback in _in_memory_loggers: - if isinstance(callback, GCSBucketLogger): - return callback - elif logging_integration == "s3_v2": - for callback in _in_memory_loggers: - if isinstance(callback, S3V2Logger): - return callback - elif logging_integration == "aws_sqs": - for callback in _in_memory_loggers: - if isinstance(callback, SQSLogger): - return callback - _aws_sqs_logger = SQSLogger() - _in_memory_loggers.append(_aws_sqs_logger) - return _aws_sqs_logger # type: ignore - elif logging_integration == "azure_storage": - for callback in _in_memory_loggers: - if isinstance(callback, AzureBlobStorageLogger): - return callback - elif logging_integration == "opik": - for callback in _in_memory_loggers: - if isinstance(callback, OpikLogger): - return callback - elif logging_integration == "langfuse": - for callback in _in_memory_loggers: - if isinstance(callback, LangfusePromptManagement): - return callback - elif logging_integration == "otel": - from litellm.integrations.opentelemetry import OpenTelemetry - - for callback in _in_memory_loggers: - if isinstance(callback, OpenTelemetry): - return callback - elif logging_integration == "arize": - if "ARIZE_API_KEY" not in os.environ: - raise ValueError("ARIZE_API_KEY not found in environment variables") - for callback in _in_memory_loggers: - if ( - isinstance(callback, ArizeLogger) - and callback.callback_name == "arize" - ): - return callback - elif logging_integration == "logfire": - if "LOGFIRE_TOKEN" not in os.environ: - raise ValueError("LOGFIRE_TOKEN not found in environment variables") - from litellm.integrations.opentelemetry import OpenTelemetry - - for callback in _in_memory_loggers: - if isinstance(callback, OpenTelemetry): - return callback # type: ignore - - elif logging_integration == "dynamic_rate_limiter": - from litellm.proxy.hooks.dynamic_rate_limiter import ( - _PROXY_DynamicRateLimitHandler, - ) - - for callback in _in_memory_loggers: - if isinstance(callback, _PROXY_DynamicRateLimitHandler): - return callback # type: ignore - elif logging_integration == "dynamic_rate_limiter_v3": - from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( - _PROXY_DynamicRateLimitHandlerV3, - ) - - for callback in _in_memory_loggers: - if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3): - return callback # type: ignore - - elif logging_integration == "langtrace": - from litellm.integrations.opentelemetry import OpenTelemetry - - if "LANGTRACE_API_KEY" not in os.environ: - raise ValueError("LANGTRACE_API_KEY not found in environment variables") - - for callback in _in_memory_loggers: - if ( - isinstance(callback, OpenTelemetry) - and callback.callback_name == "langtrace" - ): - return callback - - elif logging_integration == "mlflow": - for callback in _in_memory_loggers: - if isinstance(callback, MlflowLogger): - return callback - elif logging_integration == "pagerduty": - for callback in _in_memory_loggers: - if isinstance(callback, PagerDutyAlerting): - return callback - elif logging_integration == "anthropic_cache_control_hook": - for callback in _in_memory_loggers: - if isinstance(callback, AnthropicCacheControlHook): - return callback - elif logging_integration == "vector_store_pre_call_hook": - from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( - VectorStorePreCallHook, - ) - - for callback in _in_memory_loggers: - if isinstance(callback, VectorStorePreCallHook): - return callback - elif logging_integration == "gcs_pubsub": - for callback in _in_memory_loggers: - if isinstance(callback, GcsPubSubLogger): - return callback - elif logging_integration == "generic_api": - for callback in _in_memory_loggers: - if isinstance(callback, GenericAPILogger): - return callback - elif logging_integration == "resend_email": - for callback in _in_memory_loggers: - if isinstance(callback, ResendEmailLogger): - return callback - elif logging_integration == "sendgrid_email": - for callback in _in_memory_loggers: - if isinstance(callback, SendGridEmailLogger): - return callback - elif logging_integration == "smtp_email": - for callback in _in_memory_loggers: - if isinstance(callback, SMTPEmailLogger): - return callback - return None - - except Exception as e: - verbose_logger.exception( - f"[Non-Blocking Error] Error getting custom logger: {e}" - ) - return None - - -def _get_custom_logger_settings_from_proxy_server(callback_name: str) -> Dict: - """ - Get the settings for a custom logger from the proxy server config.yaml - - Proxy server config.yaml defines callback_settings as: - - callback_settings: - otel: - message_logging: False - """ - if litellm.callback_settings: - return dict(litellm.callback_settings.get(callback_name, {})) - return {} - - -def use_custom_pricing_for_model(litellm_params: Optional[dict]) -> bool: - """ - Check if the model uses custom pricing - - Returns True if any of `SPECIAL_MODEL_INFO_PARAMS` are present in `litellm_params` or `model_info` - """ - if litellm_params is None: - return False - - # Check litellm_params using set intersection (only check keys that exist in both) - matching_keys = _CUSTOM_PRICING_KEYS & litellm_params.keys() - for key in matching_keys: - if litellm_params.get(key) is not None: - return True - - # Check model_info - metadata: dict = litellm_params.get("metadata", {}) or {} - model_info: dict = metadata.get("model_info", {}) or {} - - if model_info: - matching_keys = _CUSTOM_PRICING_KEYS & model_info.keys() - for key in matching_keys: - if model_info.get(key) is not None: - return True - - return False - - -def is_valid_sha256_hash(value: str) -> bool: - # Check if the value is a valid SHA-256 hash (64 hexadecimal characters) - return bool(re.fullmatch(r"[a-fA-F0-9]{64}", value)) - - -class StandardLoggingPayloadSetup: - @staticmethod - def cleanup_timestamps( - start_time: Union[dt_object, float], - end_time: Union[dt_object, float], - completion_start_time: Union[dt_object, float], - ) -> Tuple[float, float, float]: - """ - Convert datetime objects to floats - - Args: - start_time: Union[dt_object, float] - end_time: Union[dt_object, float] - completion_start_time: Union[dt_object, float] - - Returns: - Tuple[float, float, float]: A tuple containing the start time, end time, and completion start time as floats. - """ - - if isinstance(start_time, datetime.datetime): - start_time_float = start_time.timestamp() - elif isinstance(start_time, float): - start_time_float = start_time - else: - raise ValueError( - f"start_time is required, got={start_time} of type {type(start_time)}" - ) - - if isinstance(end_time, datetime.datetime): - end_time_float = end_time.timestamp() - elif isinstance(end_time, float): - end_time_float = end_time - else: - raise ValueError( - f"end_time is required, got={end_time} of type {type(end_time)}" - ) - - if isinstance(completion_start_time, datetime.datetime): - completion_start_time_float = completion_start_time.timestamp() - elif isinstance(completion_start_time, float): - completion_start_time_float = completion_start_time - else: - completion_start_time_float = end_time_float - - return start_time_float, end_time_float, completion_start_time_float - - @staticmethod - def append_system_prompt_messages( - kwargs: Optional[Dict] = None, messages: Optional[Any] = None - ): - """ - Append system prompt messages to the messages - """ - if kwargs is not None: - if kwargs.get("system") is not None and isinstance( - kwargs.get("system"), str - ): - if messages is None: - return [{"role": "system", "content": kwargs.get("system")}] - elif isinstance(messages, list): - if len(messages) == 0: - return [{"role": "system", "content": kwargs.get("system")}] - # check for duplicates - if messages[0].get("role") == "system" and messages[0].get( - "content" - ) == kwargs.get("system"): - return messages - messages = [ - {"role": "system", "content": kwargs.get("system")} - ] + messages - elif isinstance(messages, str): - messages = [ - {"role": "system", "content": kwargs.get("system")}, - {"role": "user", "content": messages}, - ] - return messages - - return messages - - @staticmethod - def merge_litellm_metadata(litellm_params: dict) -> dict: - """ - Merge both litellm_metadata and metadata from litellm_params. - - litellm_metadata contains model-related fields, metadata contains user API key fields. - We need both for complete standard logging payload. - - Args: - litellm_params: Dictionary containing metadata and litellm_metadata - - Returns: - dict: Merged metadata with user API key fields taking precedence - """ - merged_metadata: dict = {} - - # Start with metadata (user API key fields) - but skip non-serializable objects - if litellm_params.get("metadata") and isinstance( - litellm_params.get("metadata"), dict - ): - for key, value in litellm_params["metadata"].items(): - # Skip non-serializable objects like UserAPIKeyAuth - if key == "user_api_key_auth": - continue - merged_metadata[key] = value - - # Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys - if litellm_params.get("litellm_metadata") and isinstance( - litellm_params.get("litellm_metadata"), dict - ): - for key, value in litellm_params["litellm_metadata"].items(): - if ( - key not in merged_metadata - ): # Don't overwrite existing keys from metadata - merged_metadata[key] = value - - return merged_metadata - - @staticmethod - def get_standard_logging_metadata( - metadata: Optional[Dict[str, Any]], - litellm_params: Optional[dict] = None, - prompt_integration: Optional[str] = None, - applied_guardrails: Optional[List[str]] = None, - mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = None, - vector_store_request_metadata: Optional[ - List[StandardLoggingVectorStoreRequest] - ] = None, - usage_object: Optional[dict] = None, - proxy_server_request: Optional[dict] = None, - start_time: Optional[dt_object] = None, - response_id: Optional[str] = None, - ) -> StandardLoggingMetadata: - """ - Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata. - - Args: - metadata (Optional[Dict[str, Any]]): The original metadata dictionary. - - Returns: - StandardLoggingMetadata: A StandardLoggingMetadata object containing the cleaned metadata. - - Note: - - If the input metadata is None or not a dictionary, an empty StandardLoggingMetadata object is returned. - - If 'user_api_key' is present in metadata and is a valid SHA256 hash, it's stored as 'user_api_key_hash'. - """ - - prompt_management_metadata: Optional[ - StandardLoggingPromptManagementMetadata - ] = None - if litellm_params is not None: - prompt_id = cast(Optional[str], litellm_params.get("prompt_id", None)) - prompt_variables = cast( - Optional[dict], litellm_params.get("prompt_variables", None) - ) - - if prompt_id is not None and prompt_integration is not None: - prompt_management_metadata = StandardLoggingPromptManagementMetadata( - prompt_id=prompt_id, - prompt_variables=prompt_variables, - prompt_integration=prompt_integration, - ) - - # Initialize with default values - clean_metadata = StandardLoggingMetadata( - user_api_key_hash=None, - user_api_key_alias=None, - user_api_key_spend=None, - user_api_key_max_budget=None, - user_api_key_budget_reset_at=None, - user_api_key_team_id=None, - user_api_key_org_id=None, - user_api_key_user_id=None, - user_api_key_team_alias=None, - user_api_key_user_email=None, - user_api_key_end_user_id=None, - user_api_key_request_route=None, - spend_logs_metadata=None, - requester_ip_address=None, - user_agent=None, - requester_metadata=None, - prompt_management_metadata=prompt_management_metadata, - applied_guardrails=applied_guardrails, - mcp_tool_call_metadata=mcp_tool_call_metadata, - vector_store_request_metadata=vector_store_request_metadata, - usage_object=usage_object, - requester_custom_headers=None, - cold_storage_object_key=None, - user_api_key_auth_metadata=None, - team_alias=None, - team_id=None, - ) - if isinstance(metadata, dict): - for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS: - clean_metadata[key] = metadata[key] # type: ignore - - user_api_key = metadata.get("user_api_key") - if ( - user_api_key - and isinstance(user_api_key, str) - and is_valid_sha256_hash(user_api_key) - ): - clean_metadata["user_api_key_hash"] = user_api_key - _potential_requester_metadata = metadata.get( - "metadata", None - ) # check if user passed metadata in the sdk request - e.g. metadata for langsmith logging - https://docs.litellm.ai/docs/observability/langsmith_integration#set-langsmith-fields - if ( - clean_metadata["requester_metadata"] is None - and _potential_requester_metadata is not None - and isinstance(_potential_requester_metadata, dict) - ): - clean_metadata["requester_metadata"] = _potential_requester_metadata - - if ( - EnterpriseStandardLoggingPayloadSetupVAR - and proxy_server_request is not None - ): - clean_metadata = EnterpriseStandardLoggingPayloadSetupVAR.apply_enterprise_specific_metadata( - standard_logging_metadata=clean_metadata, - proxy_server_request=proxy_server_request, - ) - - # Generate cold storage object key if cold storage is configured - if start_time is not None and response_id is not None: - cold_storage_object_key = ( - StandardLoggingPayloadSetup._generate_cold_storage_object_key( - start_time=start_time, - response_id=response_id, - team_alias=clean_metadata.get("user_api_key_team_alias"), - ) - ) - if cold_storage_object_key: - clean_metadata["cold_storage_object_key"] = cold_storage_object_key - - return clean_metadata - - @staticmethod - def get_usage_from_response_obj( - response_obj: Optional[dict], combined_usage_object: Optional[Usage] = None - ) -> Usage: - ## BASE CASE ## - if combined_usage_object is not None: - return combined_usage_object - if response_obj is None: - return Usage( - prompt_tokens=0, - completion_tokens=0, - total_tokens=0, - ) - - usage = response_obj.get("usage", None) or {} - if usage is None or ( - not isinstance(usage, dict) and not isinstance(usage, Usage) - ): - return Usage( - prompt_tokens=0, - completion_tokens=0, - total_tokens=0, - ) - elif isinstance(usage, Usage): - return usage - elif isinstance(usage, ResponseAPIUsage): - return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) - elif isinstance(usage, dict): - if ResponseAPILoggingUtils._is_response_api_usage(usage): - return ( - ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) - ) - return Usage(**usage) - - raise ValueError(f"usage is required, got={usage} of type {type(usage)}") - - @staticmethod - def get_model_cost_information( - base_model: Optional[str], - custom_pricing: Optional[bool], - custom_llm_provider: Optional[str], - init_response_obj: Union[Any, BaseModel, dict], - ) -> StandardLoggingModelInformation: - model_cost_name = _select_model_name_for_cost_calc( - model=None, - completion_response=init_response_obj, # type: ignore - base_model=base_model, - custom_pricing=custom_pricing, - ) - if model_cost_name is None: - model_cost_information = StandardLoggingModelInformation( - model_map_key="", model_map_value=None - ) - else: - try: - _model_cost_information = litellm.get_model_info( - model=model_cost_name, custom_llm_provider=custom_llm_provider - ) - model_cost_information = StandardLoggingModelInformation( - model_map_key=model_cost_name, - model_map_value=_model_cost_information, - ) - except Exception: - verbose_logger.debug( # keep in debug otherwise it will trigger on every call - "Model={} is not mapped in model cost map. Defaulting to None model_cost_information for standard_logging_payload".format( - model_cost_name - ) - ) - model_cost_information = StandardLoggingModelInformation( - model_map_key=model_cost_name, model_map_value=None - ) - return model_cost_information - - @staticmethod - def get_final_response_obj( - response_obj: dict, init_response_obj: Union[Any, BaseModel, dict], kwargs: dict - ) -> Optional[Union[dict, str, list]]: - """ - Get final response object after redacting the message input/output from logging - """ - if response_obj: - final_response_obj: Optional[Union[dict, str, list]] = response_obj - elif isinstance(init_response_obj, list) or isinstance(init_response_obj, str): - final_response_obj = init_response_obj - else: - final_response_obj = {} - - modified_final_response_obj = redact_message_input_output_from_logging( - model_call_details=kwargs, - result=final_response_obj, - ) - - if modified_final_response_obj is not None and isinstance( - modified_final_response_obj, BaseModel - ): - final_response_obj = modified_final_response_obj.model_dump() - else: - final_response_obj = modified_final_response_obj - - return final_response_obj - - @staticmethod - def get_additional_headers( - additiona_headers: Optional[dict], - ) -> Optional[StandardLoggingAdditionalHeaders]: - if additiona_headers is None: - return None - - additional_logging_headers: StandardLoggingAdditionalHeaders = {} - - for key in StandardLoggingAdditionalHeaders.__annotations__.keys(): - _key = key.lower() - _key = _key.replace("_", "-") - if _key in additiona_headers: - try: - additional_logging_headers[key] = int(additiona_headers[_key]) # type: ignore - except (ValueError, TypeError): - verbose_logger.debug( - f"Could not convert {additiona_headers[_key]} to int for key {key}." - ) - return additional_logging_headers - - @staticmethod - def get_hidden_params( - hidden_params: Optional[dict], - ) -> StandardLoggingHiddenParams: - clean_hidden_params = StandardLoggingHiddenParams( - model_id=None, - cache_key=None, - api_base=None, - response_cost=None, - additional_headers=None, - litellm_overhead_time_ms=None, - batch_models=None, - litellm_model_name=None, - usage_object=None, - ) - if hidden_params is not None: - for key in StandardLoggingHiddenParams.__annotations__.keys(): - if key in hidden_params: - if key == "additional_headers": - clean_hidden_params[ - "additional_headers" - ] = StandardLoggingPayloadSetup.get_additional_headers( - hidden_params[key] - ) - else: - clean_hidden_params[key] = hidden_params[key] # type: ignore - return clean_hidden_params - - @staticmethod - def strip_trailing_slash(api_base: Optional[str]) -> Optional[str]: - if api_base: - if api_base.endswith("//"): - return api_base.rstrip("/") - if api_base[-1] == "/": - return api_base[:-1] - return api_base - - @staticmethod - def _generate_cold_storage_object_key( - start_time: dt_object, - response_id: str, - team_alias: Optional[str] = None, - ) -> Optional[str]: - """ - Generate cold storage object key in the same format as S3Logger. - - Args: - start_time: The start time of the request - response_id: The response ID - team_alias: Optional team alias for team-based prefixing - - Returns: - Optional[str]: The generated object key or None if cold storage not configured - """ - # Generate object key in same format as S3Logger - from litellm.integrations.s3 import get_s3_object_key - - # Only generate object key if cold storage is configured - cold_storage_custom_logger = litellm.cold_storage_custom_logger - if cold_storage_custom_logger is None: - return None - - try: - # Generate file name in same format as litellm.utils.get_logging_id - s3_file_name = f"time-{start_time.strftime('%H-%M-%S-%f')}_{response_id}" - - # Get the actual s3_path from the configured cold storage logger instance - s3_path = "" # default value - - # Try to get the actual logger instance from the logger name - try: - custom_logger = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name( - cold_storage_custom_logger - ) - if ( - custom_logger - and hasattr(custom_logger, "s3_path") - and getattr(custom_logger, "s3_path") - ): - s3_path = getattr(custom_logger, "s3_path") - except Exception: - # If any error occurs in getting the logger instance, use default empty s3_path - pass - - s3_object_key = get_s3_object_key( - s3_path=s3_path, # Use actual s3_path from logger configuration - prefix="", # Don't split by team alias for cold storage - start_time=start_time, - s3_file_name=s3_file_name, - ) - - return s3_object_key - except Exception: - # If any error occurs in generating the key, return None - return None - - @staticmethod - def get_error_information( - original_exception: Optional[Exception], - traceback_str: Optional[str] = None, - ) -> StandardLoggingPayloadErrorInformation: - from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG - - # Check for 'code' first (used by ProxyException), then fall back to 'status_code' (used by LiteLLM exceptions) - # Ensure error_code is always a string for Prisma Python JSON field compatibility - error_code_attr = getattr(original_exception, "code", None) - if error_code_attr is not None and str(error_code_attr) not in ("", "None"): - error_status: str = str(error_code_attr) - else: - status_code_attr = getattr(original_exception, "status_code", None) - error_status = str(status_code_attr) if status_code_attr is not None else "" - error_class: str = ( - str(original_exception.__class__.__name__) if original_exception else "" - ) - _llm_provider_in_exception = getattr(original_exception, "llm_provider", "") - - # Get traceback information (first 100 lines) - traceback_info = traceback_str or "" - if original_exception: - tb = getattr(original_exception, "__traceback__", None) - if tb: - tb_lines = traceback.format_tb(tb) - traceback_info += "".join( - tb_lines[:MAXIMUM_TRACEBACK_LINES_TO_LOG] - ) # Limit to first 100 lines - - # Get additional error details - error_message = str(original_exception) - - return StandardLoggingPayloadErrorInformation( - error_code=error_status, - error_class=error_class, - llm_provider=_llm_provider_in_exception, - traceback=traceback_info, - error_message=error_message if original_exception else "", - ) - - @staticmethod - def get_response_time( - start_time_float: float, - end_time_float: float, - completion_start_time_float: float, - stream: bool, - ) -> float: - """ - Get the response time for the LLM response - - Args: - start_time_float: float - start time of the LLM call - end_time_float: float - end time of the LLM call - completion_start_time_float: float - time to first token of the LLM response (for streaming responses) - stream: bool - True when a stream response is returned - - Returns: - float: The response time for the LLM response - """ - if stream is True: - return completion_start_time_float - start_time_float - else: - return end_time_float - start_time_float - - @staticmethod - def _get_standard_logging_payload_trace_id( - logging_obj: Logging, - litellm_params: dict, - ) -> str: - """ - Returns the `litellm_trace_id` for this request - - This helps link sessions when multiple requests are made in a single session - """ - dynamic_litellm_session_id = litellm_params.get("litellm_session_id") - dynamic_litellm_trace_id = litellm_params.get("litellm_trace_id") - - # Note: we recommend using `litellm_session_id` for session tracking - # `litellm_trace_id` is an internal litellm param - if dynamic_litellm_session_id: - return str(dynamic_litellm_session_id) - elif dynamic_litellm_trace_id: - return str(dynamic_litellm_trace_id) - else: - return logging_obj.litellm_trace_id - - @staticmethod - def _get_user_agent_tags(proxy_server_request: dict) -> Optional[List[str]]: - """ - Return the user agent tags from the proxy server request for spend tracking - """ - if litellm.disable_add_user_agent_to_request_tags is True: - return None - user_agent_tags: Optional[List[str]] = None - headers = proxy_server_request.get("headers", {}) - if headers is not None and isinstance(headers, dict): - if "user-agent" in headers: - user_agent = headers["user-agent"] - if user_agent is not None: - if user_agent_tags is None: - user_agent_tags = [] - user_agent_part: Optional[str] = None - if "/" in user_agent: - user_agent_part = user_agent.split("/")[0] - if user_agent_part is not None: - user_agent_tags.append("User-Agent: " + user_agent_part) - if user_agent is not None: - user_agent_tags.append("User-Agent: " + user_agent) - return user_agent_tags - - @staticmethod - def _get_extra_header_tags(proxy_server_request: dict) -> Optional[List[str]]: - """ - Extract additional header tags for spend tracking based on config. - """ - extra_headers: List[str] = ( - getattr(litellm, "extra_spend_tag_headers", None) or [] - ) - if not extra_headers: - return None - - headers = proxy_server_request.get("headers", {}) - if not isinstance(headers, dict): - return None - - header_tags = [] - for header_name in extra_headers: - header_value = headers.get(header_name) - if header_value: - header_tags.append(f"{header_name}: {header_value}") - - return header_tags if header_tags else None - - @staticmethod - def _get_request_tags( - litellm_params: dict, proxy_server_request: dict - ) -> List[str]: - # check for 'tags' in both 'metadata' and 'litellm_metadata' - metadata = litellm_params.get("metadata") or {} - litellm_metadata = litellm_params.get("litellm_metadata") or {} - if metadata.get("tags", []): - request_tags = metadata.get("tags", []).copy() - elif litellm_metadata.get("tags", []): - request_tags = litellm_metadata.get("tags", []).copy() - else: - request_tags = [] - user_agent_tags = StandardLoggingPayloadSetup._get_user_agent_tags( - proxy_server_request - ) - additional_header_tags = StandardLoggingPayloadSetup._get_extra_header_tags( - proxy_server_request - ) - if user_agent_tags is not None: - request_tags.extend(user_agent_tags) - if additional_header_tags is not None: - request_tags.extend(additional_header_tags) - return request_tags - - -def _get_status_fields( - status: StandardLoggingPayloadStatus, - guardrail_information: Optional[List[dict]], - error_str: Optional[str], -) -> "StandardLoggingPayloadStatusFields": - """ - Determine status fields based on request status and guardrail information. - - Args: - status: Overall request status ("success" or "failure") - guardrail_information: Guardrail information from metadata - error_str: Error string if any - - Returns: - StandardLoggingPayloadStatusFields with llm_api_status and guardrail_status - """ - # Mapping for legacy guardrail status values to new GuardrailStatus values - GUARDRAIL_STATUS_MAP: Dict[str, GuardrailStatus] = { - "success": "success", - "blocked": "guardrail_intervened", # legacy - "guardrail_intervened": "guardrail_intervened", # direct - "failure": "guardrail_failed_to_respond", # legacy - "guardrail_failed_to_respond": "guardrail_failed_to_respond", # direct - "not_run": "not_run", - } - - # Set LLM API status - llm_api_status: StandardLoggingPayloadStatus = status - - ######################################################### - # Map - guardrail_information.guardrail_status to guardrail_status - ######################################################### - guardrail_status: GuardrailStatus = "not_run" - if guardrail_information and isinstance(guardrail_information, list): - for information in guardrail_information: - if isinstance(information, dict): - raw_status = information.get("guardrail_status", "not_run") - if raw_status != "not_run": - guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") - break - - return StandardLoggingPayloadStatusFields( - llm_api_status=llm_api_status, guardrail_status=guardrail_status - ) - - -def _extract_response_obj_and_hidden_params( - init_response_obj: Union[Any, BaseModel, dict], - original_exception: Optional[Exception], -) -> Tuple[dict, Optional[dict]]: - """Extract response_obj and hidden_params from init_response_obj.""" - hidden_params: Optional[dict] = None - if init_response_obj is None: - response_obj = {} - elif isinstance(init_response_obj, BaseModel): - response_obj = init_response_obj.model_dump() - hidden_params = getattr(init_response_obj, "_hidden_params", None) - elif isinstance(init_response_obj, dict): - response_obj = init_response_obj - else: - response_obj = {} - - if original_exception is not None and hidden_params is None: - response_headers = _get_response_headers(original_exception) - if response_headers is not None: - hidden_params = dict( - StandardLoggingHiddenParams( - additional_headers=StandardLoggingPayloadSetup.get_additional_headers( - dict(response_headers) - ), - model_id=None, - cache_key=None, - api_base=None, - response_cost=None, - litellm_overhead_time_ms=None, - batch_models=None, - litellm_model_name=None, - usage_object=None, - ) - ) - - return response_obj, hidden_params - - -def get_standard_logging_object_payload( - kwargs: Optional[dict], - init_response_obj: Union[Any, BaseModel, dict], - start_time: dt_object, - end_time: dt_object, - logging_obj: Logging, - status: StandardLoggingPayloadStatus, - error_str: Optional[str] = None, - original_exception: Optional[Exception] = None, - standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None, -) -> Optional[StandardLoggingPayload]: - try: - kwargs = kwargs or {} - - response_obj, hidden_params = _extract_response_obj_and_hidden_params( - init_response_obj, original_exception - ) - - # standardize this function to be used across, s3, dynamoDB, langfuse logging - litellm_params = kwargs.get("litellm_params", {}) or {} - proxy_server_request = litellm_params.get("proxy_server_request") or {} - - # Merge both litellm_metadata and metadata to get complete metadata - metadata: dict = StandardLoggingPayloadSetup.merge_litellm_metadata( - litellm_params - ) - - completion_start_time = kwargs.get("completion_start_time", end_time) - call_type = kwargs.get("call_type") - cache_hit = kwargs.get("cache_hit", False) - usage = StandardLoggingPayloadSetup.get_usage_from_response_obj( - response_obj=response_obj, - combined_usage_object=cast( - Optional[Usage], kwargs.get("combined_usage_object") - ), - ) - - id = response_obj.get("id", kwargs.get("litellm_call_id")) - - _model_id = metadata.get("model_info", {}).get("id", "") - _model_group = metadata.get("model_group", "") - - request_tags = StandardLoggingPayloadSetup._get_request_tags( - litellm_params=litellm_params, proxy_server_request=proxy_server_request - ) - - # cleanup timestamps - ( - start_time_float, - end_time_float, - completion_start_time_float, - ) = StandardLoggingPayloadSetup.cleanup_timestamps( - start_time=start_time, - end_time=end_time, - completion_start_time=completion_start_time, - ) - response_time = StandardLoggingPayloadSetup.get_response_time( - start_time_float=start_time_float, - end_time_float=end_time_float, - completion_start_time_float=completion_start_time_float, - stream=kwargs.get("stream", False), - ) - # clean up litellm hidden params - clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params( - hidden_params - ) - - # clean up litellm metadata - clean_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( - metadata=metadata, - litellm_params=litellm_params, - prompt_integration=kwargs.get("prompt_integration", None), - applied_guardrails=kwargs.get("applied_guardrails", None), - mcp_tool_call_metadata=kwargs.get("mcp_tool_call_metadata", None), - vector_store_request_metadata=kwargs.get( - "vector_store_request_metadata", None - ), - usage_object=usage.model_dump(), - proxy_server_request=proxy_server_request, - start_time=start_time, - response_id=id, - ) - _request_body = proxy_server_request.get("body", {}) - end_user_id = clean_metadata["user_api_key_end_user_id"] or _request_body.get( - "user", None - ) # maintain backwards compatibility with old request body check - - saved_cache_cost: float = 0.0 - if cache_hit is True: - id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id - saved_cache_cost = ( - logging_obj._response_cost_calculator( - result=init_response_obj, cache_hit=False # type: ignore - ) - or 0.0 - ) - - ## Get model cost information ## - base_model = _get_base_model_from_metadata(model_call_details=kwargs) - custom_pricing = use_custom_pricing_for_model(litellm_params=litellm_params) - - model_cost_information = StandardLoggingPayloadSetup.get_model_cost_information( - base_model=base_model, - custom_pricing=custom_pricing, - custom_llm_provider=kwargs.get("custom_llm_provider"), - init_response_obj=init_response_obj, - ) - response_cost: float = kwargs.get("response_cost", 0) or 0.0 - - error_information = StandardLoggingPayloadSetup.get_error_information( - original_exception=original_exception, - ) - - ## get final response object ## - final_response_obj = StandardLoggingPayloadSetup.get_final_response_obj( - response_obj=response_obj, - init_response_obj=init_response_obj, - kwargs=kwargs, - ) - - stream: Optional[bool] = None - if ( - kwargs.get("complete_streaming_response") is not None - or kwargs.get("async_complete_streaming_response") is not None - ) and kwargs.get("stream") is True: - stream = True - - # Reconstruct full model name with provider prefix for logging - # This ensures Bedrock models like "us.anthropic.claude-3-5-sonnet-20240620-v1:0" - # are logged as "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0" - custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) - model_name = reconstruct_model_name( - kwargs.get("model", "") or "", custom_llm_provider, metadata - ) - - payload: StandardLoggingPayload = StandardLoggingPayload( - id=str(id), - trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( - logging_obj=logging_obj, - litellm_params=litellm_params, - ), - call_type=call_type or "", - cache_hit=cache_hit, - stream=stream, - status=status, - status_fields=_get_status_fields( - status=status, - guardrail_information=metadata.get( - "standard_logging_guardrail_information", None - ), - error_str=error_str, - ), - custom_llm_provider=custom_llm_provider, - saved_cache_cost=saved_cache_cost, - startTime=start_time_float, - endTime=end_time_float, - completionStartTime=completion_start_time_float, - response_time=response_time, - model=model_name, - metadata=clean_metadata, - cache_key=clean_hidden_params["cache_key"], - response_cost=response_cost, - cost_breakdown=logging_obj.cost_breakdown, - total_tokens=usage.total_tokens, - prompt_tokens=usage.prompt_tokens, - completion_tokens=usage.completion_tokens, - request_tags=request_tags, - end_user=end_user_id or "", - api_base=StandardLoggingPayloadSetup.strip_trailing_slash( - litellm_params.get("api_base", "") - ) - or "", - model_group=_model_group, - model_id=_model_id, - requester_ip_address=clean_metadata.get("requester_ip_address", None), - user_agent=clean_metadata.get("user_agent", None), - messages=StandardLoggingPayloadSetup.append_system_prompt_messages( - kwargs=kwargs, messages=kwargs.get("messages") - ), - response=final_response_obj, - model_parameters=ModelParamHelper.get_standard_logging_model_parameters( - kwargs.get("optional_params", None) or {} - ), - hidden_params=clean_hidden_params, - model_map_information=model_cost_information, - error_str=error_str, - error_information=error_information, - response_cost_failure_debug_info=kwargs.get( - "response_cost_failure_debug_information" - ), - guardrail_information=metadata.get( - "standard_logging_guardrail_information", None - ), - standard_built_in_tools_params=standard_built_in_tools_params, - ) - - # emit_standard_logging_payload(payload) - Moved to success_handler to prevent double emitting - - return payload - except Exception as e: - verbose_logger.exception( - "Error creating standard logging object - {}".format(str(e)) - ) - return None - - -def emit_standard_logging_payload(payload: StandardLoggingPayload): - if os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD"): - print(json.dumps(payload, indent=4)) # noqa - - -def get_standard_logging_metadata( - metadata: Optional[Dict[str, Any]], -) -> StandardLoggingMetadata: - """ - Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata. - - Args: - metadata (Optional[Dict[str, Any]]): The original metadata dictionary. - - Returns: - StandardLoggingMetadata: A StandardLoggingMetadata object containing the cleaned metadata. - - Note: - - If the input metadata is None or not a dictionary, an empty StandardLoggingMetadata object is returned. - - If 'user_api_key' is present in metadata and is a valid SHA256 hash, it's stored as 'user_api_key_hash'. - """ - # Initialize with default values - clean_metadata = StandardLoggingMetadata( - user_api_key_hash=None, - user_api_key_alias=None, - user_api_key_spend=None, - user_api_key_max_budget=None, - user_api_key_budget_reset_at=None, - user_api_key_team_id=None, - user_api_key_org_id=None, - user_api_key_user_id=None, - user_api_key_user_email=None, - user_api_key_team_alias=None, - spend_logs_metadata=None, - requester_ip_address=None, - user_agent=None, - requester_metadata=None, - user_api_key_end_user_id=None, - prompt_management_metadata=None, - applied_guardrails=None, - mcp_tool_call_metadata=None, - vector_store_request_metadata=None, - usage_object=None, - requester_custom_headers=None, - user_api_key_request_route=None, - cold_storage_object_key=None, - user_api_key_auth_metadata=None, - team_alias=None, - team_id=None, - ) - if isinstance(metadata, dict): - # Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields - for key in StandardLoggingMetadata.__annotations__.keys(): - if key in metadata: - clean_metadata[key] = metadata[key] # type: ignore - - if metadata.get("user_api_key") is not None: - if is_valid_sha256_hash(str(metadata.get("user_api_key"))): - clean_metadata["user_api_key_hash"] = metadata.get( - "user_api_key" - ) # this is the hash - return clean_metadata - - -def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): - if litellm_params is None: - litellm_params = {} - - metadata = litellm_params.get("metadata", {}) or {} - - ## Extract provider-specific callable values (like langfuse_masking_function) - ## Store them separately so only the intended logger can access them - ## This prevents callables from leaking to other logging integrations - if "langfuse_masking_function" in metadata: - masking_fn = metadata.pop("langfuse_masking_function", None) - if callable(masking_fn): - litellm_params["_langfuse_masking_function"] = masking_fn - litellm_params["metadata"] = metadata - - ## check user_api_key_metadata for sensitive logging keys - cleaned_user_api_key_metadata = {} - if "user_api_key_metadata" in metadata and isinstance( - metadata["user_api_key_metadata"], dict - ): - for k, v in metadata["user_api_key_metadata"].items(): - if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[ - k - ] = "scrubbed_by_litellm_for_sensitive_keys" - else: - cleaned_user_api_key_metadata[k] = v - - metadata["user_api_key_metadata"] = cleaned_user_api_key_metadata - litellm_params["metadata"] = metadata - - return litellm_params - - -# integration helper function -def modify_integration(integration_name, integration_params): - global supabaseClient - if integration_name == "supabase": - if "table_name" in integration_params: - Supabase.supabase_table_name = integration_params["table_name"] - - -@lru_cache(maxsize=16) -def _get_traceback_str_for_error(error_str: str) -> str: - """ - function wrapped with lru_cache to limit the number of times `traceback.format_exc()` is called - """ - return traceback.format_exc() - - -from decimal import Decimal - -# used for unit testing -from typing import Any, Dict, List, Optional, Union - - -def create_dummy_standard_logging_payload() -> StandardLoggingPayload: - # First create the nested objects with proper typing - model_info = StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ) - - metadata = StandardLoggingMetadata( # type: ignore - user_api_key_hash=str("test_hash"), - user_api_key_alias=str("test_alias"), - user_api_key_team_id=str("test_team"), - user_api_key_user_id=str("test_user"), - user_api_key_team_alias=str("test_team_alias"), - user_api_key_org_id=None, - spend_logs_metadata=None, - requester_ip_address=str("127.0.0.1"), - requester_metadata=None, - user_api_key_end_user_id=str("test_end_user"), - ) - - hidden_params = StandardLoggingHiddenParams( - model_id=None, - cache_key=None, - api_base=None, - response_cost=None, - additional_headers=None, - litellm_overhead_time_ms=None, - batch_models=None, - litellm_model_name=None, - usage_object=None, - ) - - # Convert numeric values to appropriate types - response_cost = Decimal("0.1") - start_time = Decimal("1234567890.0") - end_time = Decimal("1234567891.0") - completion_start_time = Decimal("1234567890.5") - saved_cache_cost = Decimal("0.0") - - # Create messages and response with proper typing - messages: List[Dict[str, str]] = [{"role": "user", "content": "Hello, world!"}] - response: Dict[str, List[Dict[str, Dict[str, str]]]] = { - "choices": [{"message": {"content": "Hi there!"}}] - } - - # Main payload initialization - return StandardLoggingPayload( # type: ignore - id=str("test_id"), - call_type=str("completion"), - stream=bool(False), - response_cost=response_cost, - response_cost_failure_debug_info=None, - status=str("success"), - total_tokens=int( - DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT - + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT - ), - prompt_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT), - completion_tokens=int(DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT), - startTime=start_time, - endTime=end_time, - completionStartTime=completion_start_time, - model_map_information=model_info, - model=str("gpt-3.5-turbo"), - model_id=str("model-123"), - model_group=str("openai-gpt"), - custom_llm_provider=str("openai"), - api_base=str("https://api.openai.com"), - metadata=metadata, - cache_hit=bool(False), - cache_key=None, - saved_cache_cost=saved_cache_cost, - request_tags=[], - end_user=None, - requester_ip_address=str("127.0.0.1"), - messages=messages, - response=response, - error_str=None, - model_parameters={"stream": True}, - hidden_params=hidden_params, - ) +# What is this? +## Common Utility file for Logging handler +# Logging function -> log the exact model details + what's being sent | Non-Blocking +import copy +import datetime +import json +import os +import re +import subprocess +import sys +import time +import traceback +from datetime import datetime as dt_object +from functools import lru_cache +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + Tuple, + Type, + Union, + cast, +) + +from httpx import Response +from pydantic import BaseModel + +import litellm +from litellm import ( + _custom_logger_compatible_callbacks_literal, + json_logs, + log_raw_request_response, + turn_off_message_logging, +) +from litellm._logging import _is_debugging_on, verbose_logger +from litellm._uuid import uuid +from litellm.batches.batch_utils import _handle_completed_batch +from litellm.caching.caching import DualCache, InMemoryCache +from litellm.caching.caching_handler import LLMCachingHandler +from litellm.constants import ( + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, + DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, + SENTRY_DENYLIST, + SENTRY_PII_DENYLIST, +) +from litellm.cost_calculator import ( + RealtimeAPITokenUsageProcessor, + _select_model_name_for_cost_calc, +) +from litellm.integrations.agentops import AgentOps +from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook +from litellm.integrations.arize.arize import ArizeLogger +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.deepeval.deepeval import DeepEvalLogger +from litellm.integrations.mlflow import MlflowLogger +from litellm.integrations.sqs import SQSLogger +from litellm.litellm_core_utils.core_helpers import reconstruct_model_name +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, +) +from litellm.litellm_core_utils.model_param_helper import ModelParamHelper +from litellm.litellm_core_utils.redact_messages import ( + redact_message_input_output_from_custom_logger, + redact_message_input_output_from_logging, +) +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.llms.base_llm.search.transformation import SearchResponse +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.agents import LiteLLMSendMessageResponse +from litellm.types.containers.main import ContainerObject +from litellm.types.llms.openai import ( + AllMessageValues, + Batch, + FineTuningJob, + HttpxBinaryResponseContent, + OpenAIFileObject, + OpenAIModerationResponse, + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIResponse, +) +from litellm.types.mcp import MCPPostCallResponseObject +from litellm.types.prompts.init_prompts import PromptSpec +from litellm.types.rerank import RerankResponse +from litellm.types.utils import ( + CachingDetails, + CallTypes, + CostBreakdown, + CostResponseTypes, + CustomPricingLiteLLMParams, + DynamicPromptManagementParamLiteral, + EmbeddingResponse, + GuardrailStatus, + ImageResponse, + LiteLLMBatch, + LiteLLMLoggingBaseClass, + LiteLLMRealtimeStreamLoggingObject, + ModelResponse, + ModelResponseStream, + RawRequestTypedDict, + StandardBuiltInToolsParams, + StandardCallbackDynamicParams, + StandardLoggingAdditionalHeaders, + StandardLoggingHiddenParams, + StandardLoggingMCPToolCall, + StandardLoggingMetadata, + StandardLoggingModelCostFailureDebugInformation, + StandardLoggingModelInformation, + StandardLoggingPayload, + StandardLoggingPayloadErrorInformation, + StandardLoggingPayloadStatus, + StandardLoggingPayloadStatusFields, + StandardLoggingPromptManagementMetadata, + StandardLoggingVectorStoreRequest, + TextCompletionResponse, + TranscriptionResponse, + Usage, +) +from litellm.types.videos.main import VideoObject +from litellm.utils import _get_base_model_from_metadata, executor, print_verbose + +from ..integrations.argilla import ArgillaLogger +from ..integrations.arize.arize_phoenix import ArizePhoenixLogger +from ..integrations.athina import AthinaLogger +from ..integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger +from ..integrations.azure_storage.azure_storage import AzureBlobStorageLogger +from ..integrations.custom_prompt_management import CustomPromptManagement +from ..integrations.datadog.datadog import DataDogLogger +from ..integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger +from ..integrations.dotprompt import DotpromptManager +from ..integrations.dynamodb import DyanmoDBLogger +from ..integrations.galileo import GalileoObserve +from ..integrations.gcs_bucket.gcs_bucket import GCSBucketLogger +from ..integrations.gcs_pubsub.pub_sub import GcsPubSubLogger +from ..integrations.greenscale import GreenscaleLogger +from ..integrations.helicone import HeliconeLogger +from ..integrations.humanloop import HumanloopLogger +from ..integrations.lago import LagoLogger +from ..integrations.langfuse.langfuse import LangFuseLogger +from ..integrations.langfuse.langfuse_handler import LangFuseHandler +from ..integrations.langfuse.langfuse_prompt_management import LangfusePromptManagement +from ..integrations.langsmith import LangsmithLogger +from ..integrations.literal_ai import LiteralAILogger +from ..integrations.logfire_logger import LogfireLevel, LogfireLogger +from ..integrations.lunary import LunaryLogger +from ..integrations.openmeter import OpenMeterLogger +from ..integrations.opik.opik import OpikLogger +from ..integrations.posthog import PostHogLogger +from ..integrations.prompt_layer import PromptLayerLogger +from ..integrations.s3 import S3Logger +from ..integrations.s3_v2 import S3Logger as S3V2Logger +from ..integrations.supabase import Supabase +from ..integrations.traceloop import TraceloopLogger +from .exception_mapping_utils import _get_response_headers +from .initialize_dynamic_callback_params import ( + initialize_standard_callback_dynamic_params as _initialize_standard_callback_dynamic_params, +) +from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache + +if TYPE_CHECKING: + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig +try: + from litellm_enterprise.enterprise_callbacks.callback_controls import ( + EnterpriseCallbackControls, + ) + from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import ( + PagerDutyAlerting, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ( + ResendEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import ( + SendGridEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import ( + SMTPEmailLogger, + ) + from litellm_enterprise.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup as EnterpriseStandardLoggingPayloadSetup, + ) + + from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger + + EnterpriseStandardLoggingPayloadSetupVAR: Optional[ + Type[EnterpriseStandardLoggingPayloadSetup] + ] = EnterpriseStandardLoggingPayloadSetup +except Exception as e: + verbose_logger.debug( + f"[Non-Blocking] Unable to import GenericAPILogger - LiteLLM Enterprise Feature - {str(e)}" + ) + GenericAPILogger = CustomLogger # type: ignore + ResendEmailLogger = CustomLogger # type: ignore + SendGridEmailLogger = CustomLogger # type: ignore + SMTPEmailLogger = CustomLogger # type: ignore + PagerDutyAlerting = CustomLogger # type: ignore + EnterpriseCallbackControls = None # type: ignore + EnterpriseStandardLoggingPayloadSetupVAR = None +_in_memory_loggers: List[Any] = [] + +_STANDARD_LOGGING_METADATA_KEYS: frozenset = frozenset( + StandardLoggingMetadata.__annotations__.keys() +) + +### GLOBAL VARIABLES ### + +# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys +_CUSTOM_PRICING_KEYS: frozenset = frozenset( + CustomPricingLiteLLMParams.model_fields.keys() +) + +sentry_sdk_instance = None +capture_exception = None +add_breadcrumb = None +slack_app = None +alerts_channel = None +heliconeLogger = None +athinaLogger = None +promptLayerLogger = None +logfireLogger = None +weightsBiasesLogger = None +customLogger = None +langFuseLogger = None +openMeterLogger = None +lagoLogger = None +dataDogLogger = None +prometheusLogger = None +dynamoLogger = None +s3Logger = None +greenscaleLogger = None +lunaryLogger = None +supabaseClient = None +deepevalLogger = None +callback_list: Optional[List[str]] = [] +user_logger_fn = None +additional_details: Optional[Dict[str, str]] = {} +local_cache: Optional[Dict[str, str]] = {} +last_fetched_at = None +last_fetched_at_keys = None + + +#### +class ServiceTraceIDCache: + def __init__(self) -> None: + self.cache = InMemoryCache() + + def get_cache(self, litellm_call_id: str, service_name: str) -> Optional[str]: + key_name = "{}:{}".format(service_name, litellm_call_id) + response = self.cache.get_cache(key=key_name) + return response + + def set_cache(self, litellm_call_id: str, service_name: str, trace_id: str) -> None: + key_name = "{}:{}".format(service_name, litellm_call_id) + self.cache.set_cache(key=key_name, value=trace_id) + return None + + +in_memory_trace_id_cache = ServiceTraceIDCache() +in_memory_dynamic_logger_cache = DynamicLoggingCache() + +# Cached lazy import for PrometheusLogger +# Module-level cache to avoid repeated imports while preserving memory benefits +_PrometheusLogger = None + + +def _get_cached_prometheus_logger(): + """ + Get cached PrometheusLogger class. + Lazy imports on first call to avoid loading prometheus.py and utils.py at import time (60MB saved). + Subsequent calls use cached class for better performance. + """ + global _PrometheusLogger + if _PrometheusLogger is None: + from litellm.integrations.prometheus import PrometheusLogger + + _PrometheusLogger = PrometheusLogger + return _PrometheusLogger + + +class Logging(LiteLLMLoggingBaseClass): + global supabaseClient, promptLayerLogger, weightsBiasesLogger, logfireLogger, capture_exception, add_breadcrumb, lunaryLogger, logfireLogger, prometheusLogger, slack_app + custom_pricing: bool = False + stream_options = None + litellm_request_debug: bool = False + + def __init__( + self, + model: str, + messages, + stream, + call_type, + start_time, + litellm_call_id: str, + function_id: str, + litellm_trace_id: Optional[str] = None, + dynamic_input_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = None, + dynamic_success_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = None, + dynamic_async_success_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = None, + dynamic_failure_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = None, + dynamic_async_failure_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = None, + applied_guardrails: Optional[List[str]] = None, + kwargs: Optional[Dict] = None, + log_raw_request_response: bool = False, + ): + _input: Optional[str] = messages # save original value of messages + if messages is not None: + if isinstance(messages, str): + messages = [ + {"role": "user", "content": messages} + ] # convert text completion input to the chat completion format + elif ( + isinstance(messages, list) + and len(messages) > 0 + and isinstance(messages[0], str) + ): + new_messages = [] + for m in messages: + new_messages.append({"role": "user", "content": m}) + messages = new_messages + + self.model = model + self.messages = copy.deepcopy(messages) if messages is not None else None + self.stream = stream + self.start_time = start_time # log the call start time + self.call_type = call_type + self.litellm_call_id = litellm_call_id + self.litellm_trace_id: str = ( + litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) + ) + self.function_id = function_id + self.streaming_chunks: List[Any] = [] # for generating complete stream response + self.sync_streaming_chunks: List[ + Any + ] = [] # for generating complete stream response + self.log_raw_request_response = log_raw_request_response + + # Initialize dynamic callbacks + self.dynamic_input_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = dynamic_input_callbacks + self.dynamic_success_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = dynamic_success_callbacks + self.dynamic_async_success_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = dynamic_async_success_callbacks + self.dynamic_failure_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = dynamic_failure_callbacks + self.dynamic_async_failure_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = dynamic_async_failure_callbacks + + # Process dynamic callbacks + self.process_dynamic_callbacks() + + ## DYNAMIC LANGFUSE / GCS / logging callback KEYS ## + self.standard_callback_dynamic_params: StandardCallbackDynamicParams = ( + self.initialize_standard_callback_dynamic_params(kwargs) + ) + self.standard_built_in_tools_params: StandardBuiltInToolsParams = ( + self.initialize_standard_built_in_tools_params(kwargs) + ) + ## TIME TO FIRST TOKEN LOGGING ## + self.completion_start_time: Optional[datetime.datetime] = None + self._llm_caching_handler: Optional[LLMCachingHandler] = None + + # INITIAL LITELLM_PARAMS + litellm_params = {} + if kwargs is not None: + litellm_params = get_litellm_params(**kwargs) + litellm_params = scrub_sensitive_keys_in_metadata(litellm_params) + + self.litellm_params = litellm_params + + # Initialize cost breakdown field + self.cost_breakdown: Optional[CostBreakdown] = None + + # Init Caching related details + self.caching_details: Optional[CachingDetails] = None + + # Passthrough endpoint guardrails config for field targeting + self.passthrough_guardrails_config: Optional[Dict[str, Any]] = None + + self.model_call_details: Dict[str, Any] = { + "litellm_trace_id": litellm_trace_id, + "litellm_call_id": litellm_call_id, + "input": _input, + "litellm_params": litellm_params, + "applied_guardrails": applied_guardrails, + "model": model, + } + + def process_dynamic_callbacks(self): + """ + Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks + + If a callback is in litellm._known_custom_logger_compatible_callbacks, it needs to be intialized and added to the respective dynamic_* callback list. + """ + # Process input callbacks + self.dynamic_input_callbacks = self._process_dynamic_callback_list( + self.dynamic_input_callbacks, dynamic_callbacks_type="input" + ) + + # Process failure callbacks + self.dynamic_failure_callbacks = self._process_dynamic_callback_list( + self.dynamic_failure_callbacks, dynamic_callbacks_type="failure" + ) + + # Process async failure callbacks + self.dynamic_async_failure_callbacks = self._process_dynamic_callback_list( + self.dynamic_async_failure_callbacks, dynamic_callbacks_type="async_failure" + ) + + # Process success callbacks + self.dynamic_success_callbacks = self._process_dynamic_callback_list( + self.dynamic_success_callbacks, dynamic_callbacks_type="success" + ) + + # Process async success callbacks + self.dynamic_async_success_callbacks = self._process_dynamic_callback_list( + self.dynamic_async_success_callbacks, dynamic_callbacks_type="async_success" + ) + + def _process_dynamic_callback_list( + self, + callback_list: Optional[List[Union[str, Callable, CustomLogger]]], + dynamic_callbacks_type: Literal[ + "input", "success", "failure", "async_success", "async_failure" + ], + ) -> Optional[List[Union[str, Callable, CustomLogger]]]: + """ + Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks + + - If a callback is in litellm._known_custom_logger_compatible_callbacks, + replace the string with the initialized callback class. + - If dynamic callback is a "success" callback that is a known_custom_logger_compatible_callbacks then add it to dynamic_async_success_callbacks + - If dynamic callback is a "failure" callback that is a known_custom_logger_compatible_callbacks then add it to dynamic_failure_callbacks + """ + if callback_list is None: + return None + + processed_list: List[Union[str, Callable, CustomLogger]] = [] + for callback in callback_list: + if ( + isinstance(callback, str) + and callback in litellm._known_custom_logger_compatible_callbacks + ): + callback_class = _init_custom_logger_compatible_class( + callback, internal_usage_cache=None, llm_router=None # type: ignore + ) + if callback_class is not None: + processed_list.append(callback_class) + + # If processing dynamic_success_callbacks, add to dynamic_async_success_callbacks + if dynamic_callbacks_type == "success": + if self.dynamic_async_success_callbacks is None: + self.dynamic_async_success_callbacks = [] + self.dynamic_async_success_callbacks.append(callback_class) + elif dynamic_callbacks_type == "failure": + if self.dynamic_async_failure_callbacks is None: + self.dynamic_async_failure_callbacks = [] + self.dynamic_async_failure_callbacks.append(callback_class) + else: + processed_list.append(callback) + return processed_list + + def initialize_standard_callback_dynamic_params( + self, kwargs: Optional[Dict] = None + ) -> StandardCallbackDynamicParams: + """ + Initialize the standard callback dynamic params from the kwargs + + checks if langfuse_secret_key, gcs_bucket_name in kwargs and sets the corresponding attributes in StandardCallbackDynamicParams + """ + + return _initialize_standard_callback_dynamic_params(kwargs) + + def initialize_standard_built_in_tools_params( + self, kwargs: Optional[Dict] = None + ) -> StandardBuiltInToolsParams: + """ + Initialize the standard built-in tools params from the kwargs + + checks if web_search_options in kwargs or tools and sets the corresponding attribute in StandardBuiltInToolsParams + """ + return StandardBuiltInToolsParams( + web_search_options=StandardBuiltInToolCostTracking._get_web_search_options( + kwargs or {} + ), + file_search=StandardBuiltInToolCostTracking._get_file_search_tool_call( + kwargs or {} + ), + ) + + def update_environment_variables( + self, + litellm_params: Dict, + optional_params: Dict, + model: Optional[str] = None, + user: Optional[str] = None, + **additional_params, + ): + self.optional_params = optional_params + if model is not None: + self.model = model + self.user = user + self.litellm_params = { + **self.litellm_params, + **scrub_sensitive_keys_in_metadata(litellm_params), + } + self.litellm_request_debug = litellm_params.get("litellm_request_debug", False) + self.logger_fn = litellm_params.get("logger_fn", None) + if _is_debugging_on() or self.litellm_request_debug: + verbose_logger.debug(f"self.optional_params: {self.optional_params}") + + self.model_call_details.update( + { + "model": self.model, + "messages": self.messages, + "optional_params": self.optional_params, + "litellm_params": self.litellm_params, + "start_time": self.start_time, + "stream": self.stream, + "user": user, + "call_type": str(self.call_type), + "litellm_call_id": self.litellm_call_id, + "completion_start_time": self.completion_start_time, + "standard_callback_dynamic_params": self.standard_callback_dynamic_params, + **self.optional_params, + **additional_params, + } + ) + + ## check if stream options is set ## - used by CustomStreamWrapper for easy instrumentation + if "stream_options" in additional_params: + self.stream_options = additional_params["stream_options"] + ## check if custom pricing set ## + if any( + litellm_params.get(key) is not None + for key in _CUSTOM_PRICING_KEYS & litellm_params.keys() + ): + self.custom_pricing = True + + if "custom_llm_provider" in self.model_call_details: + self.custom_llm_provider = self.model_call_details["custom_llm_provider"] + + def update_messages(self, messages: List[AllMessageValues]): + """ + Update the logged value of the messages in the model_call_details + + Allows pre-call hooks to update the messages before the call is made + """ + self.messages = messages + self.model_call_details["messages"] = messages + + def should_run_prompt_management_hooks( + self, + non_default_params: Dict, + prompt_id: Optional[str] = None, + tools: Optional[List[Dict]] = None, + ) -> bool: + """ + Return True if prompt management hooks should be run + """ + if prompt_id: + return True + + if self._should_run_prompt_management_hooks_without_prompt_id( + non_default_params=non_default_params, + tools=tools, + ): + return True + + return False + + def _should_run_prompt_management_hooks_without_prompt_id( + self, + non_default_params: Dict, + tools: Optional[List[Dict]] = None, + ) -> bool: + """ + Certain prompt management hooks don't need a `prompt_id` to be passed in, they are triggered by dynamic params + + eg. AnthropicCacheControlHook and BedrockKnowledgeBaseHook both don't require a `prompt_id` to be passed in, they are triggered by dynamic params + """ + for param in non_default_params: + if param in DynamicPromptManagementParamLiteral.list_all_params(): + return True + + ############################################################################# + # Check if Vector Store / Knowledge Base hooks should be applied to the prompt + ############################################################################# + if litellm.vector_store_registry is not None: + if litellm.vector_store_registry.get_vector_store_to_run( + non_default_params=non_default_params, tools=tools + ): + return True + return False + + def get_chat_completion_prompt( + self, + model: str, + messages: List[AllMessageValues], + non_default_params: Dict, + prompt_variables: Optional[dict], + prompt_id: Optional[str] = None, + prompt_spec: Optional[PromptSpec] = None, + prompt_management_logger: Optional[CustomLogger] = None, + prompt_label: Optional[str] = None, + prompt_version: Optional[int] = None, + ) -> Tuple[str, List[AllMessageValues], dict]: + custom_logger = ( + prompt_management_logger + or self.get_custom_logger_for_prompt_management( + model=model, + non_default_params=non_default_params, + prompt_id=prompt_id, + prompt_spec=prompt_spec, + dynamic_callback_params=self.standard_callback_dynamic_params, + ) + ) + + if custom_logger: + ( + model, + messages, + non_default_params, + ) = custom_logger.get_chat_completion_prompt( + model=model, + messages=messages, + non_default_params=non_default_params or {}, + prompt_id=prompt_id, + prompt_spec=prompt_spec, + prompt_variables=prompt_variables, + dynamic_callback_params=self.standard_callback_dynamic_params, + prompt_label=prompt_label, + prompt_version=prompt_version, + ) + self.messages = messages + return model, messages, non_default_params + + async def async_get_chat_completion_prompt( + self, + model: str, + messages: List[AllMessageValues], + non_default_params: Dict, + prompt_variables: Optional[dict], + prompt_id: Optional[str] = None, + prompt_spec: Optional[PromptSpec] = None, + prompt_management_logger: Optional[CustomLogger] = None, + tools: Optional[List[Dict]] = None, + prompt_label: Optional[str] = None, + prompt_version: Optional[int] = None, + ) -> Tuple[str, List[AllMessageValues], dict]: + custom_logger = ( + prompt_management_logger + or self.get_custom_logger_for_prompt_management( + model=model, + tools=tools, + non_default_params=non_default_params, + prompt_id=prompt_id, + prompt_spec=prompt_spec, + dynamic_callback_params=self.standard_callback_dynamic_params, + ) + ) + + if custom_logger: + ( + model, + messages, + non_default_params, + ) = await custom_logger.async_get_chat_completion_prompt( + model=model, + messages=messages, + non_default_params=non_default_params or {}, + prompt_id=prompt_id, + prompt_spec=prompt_spec, + prompt_variables=prompt_variables, + dynamic_callback_params=self.standard_callback_dynamic_params, + litellm_logging_obj=self, + tools=tools, + prompt_label=prompt_label, + prompt_version=prompt_version, + ) + self.messages = messages + return model, messages, non_default_params + + def _auto_detect_prompt_management_logger( + self, + prompt_id: str, + prompt_spec: Optional[PromptSpec], + dynamic_callback_params: StandardCallbackDynamicParams, + ) -> Optional[CustomLogger]: + """ + Auto-detect which prompt management system owns the given prompt_id. + + This allows a user to just pass prompt_id in the completion call and it will be auto-detected which system owns this prompt. + + Args: + prompt_id: The prompt ID to check + dynamic_callback_params: Dynamic callback parameters for should_run_prompt_management checks + + Returns: + A CustomLogger instance if a matching prompt management system is found, None otherwise + """ + prompt_management_loggers = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=CustomPromptManagement + ) + ) + + for logger in prompt_management_loggers: + if isinstance(logger, CustomPromptManagement): + try: + if logger.should_run_prompt_management( + prompt_id=prompt_id, + prompt_spec=prompt_spec, + dynamic_callback_params=dynamic_callback_params, + ): + self.model_call_details[ + "prompt_integration" + ] = logger.__class__.__name__ + return logger + except Exception: + # If check fails, continue to next logger + continue + + return None + + def get_custom_logger_for_prompt_management( + self, + model: str, + non_default_params: Dict, + tools: Optional[List[Dict]] = None, + prompt_id: Optional[str] = None, + prompt_spec: Optional[PromptSpec] = None, + dynamic_callback_params: Optional[StandardCallbackDynamicParams] = None, + ) -> Optional[CustomLogger]: + """ + Get a custom logger for prompt management based on model name or available callbacks. + + Args: + model: The model name to check for prompt management integration + non_default_params: Non-default parameters passed to the completion call + tools: Optional tools passed to the completion call + prompt_id: Optional prompt ID to auto-detect which system owns this prompt + dynamic_callback_params: Dynamic callback parameters for should_run_prompt_management checks + + Returns: + A CustomLogger instance if one is found, None otherwise + """ + # First check if model starts with a known custom logger compatible callback + # This takes precedence for backward compatibility + for callback_name in litellm._known_custom_logger_compatible_callbacks: + if model.startswith(callback_name): + custom_logger = _init_custom_logger_compatible_class( + logging_integration=callback_name, + internal_usage_cache=None, + llm_router=None, + ) + if custom_logger is not None: + self.model_call_details["prompt_integration"] = model.split("/")[0] + return custom_logger + + # If prompt_id is provided, try to auto-detect which system has this prompt + if prompt_id and dynamic_callback_params is not None: + auto_detected_logger = self._auto_detect_prompt_management_logger( + prompt_id=prompt_id, + prompt_spec=prompt_spec, + dynamic_callback_params=dynamic_callback_params, + ) + if auto_detected_logger is not None: + return auto_detected_logger + + # Then check for any registered CustomPromptManagement loggers (fallback) + prompt_management_loggers = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=CustomPromptManagement + ) + ) + + if prompt_management_loggers: + logger = prompt_management_loggers[0] + self.model_call_details["prompt_integration"] = logger.__class__.__name__ + return logger + + if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( + non_default_params + ): + self.model_call_details[ + "prompt_integration" + ] = anthropic_cache_control_logger.__class__.__name__ + return anthropic_cache_control_logger + + ######################################################### + # Vector Store / Knowledge Base hooks + ######################################################### + if litellm.vector_store_registry is not None: + vector_store_custom_logger = _init_custom_logger_compatible_class( + logging_integration="vector_store_pre_call_hook", + internal_usage_cache=None, + llm_router=None, + ) + self.model_call_details[ + "prompt_integration" + ] = vector_store_custom_logger.__class__.__name__ + # Add to global callbacks so post-call hooks are invoked + if ( + vector_store_custom_logger + and vector_store_custom_logger not in litellm.callbacks + ): + litellm.logging_callback_manager.add_litellm_callback( + vector_store_custom_logger + ) + return vector_store_custom_logger + + return None + + def get_custom_logger_for_anthropic_cache_control_hook( + self, non_default_params: Dict + ) -> Optional[CustomLogger]: + if non_default_params.get("cache_control_injection_points", None): + custom_logger = _init_custom_logger_compatible_class( + logging_integration="anthropic_cache_control_hook", + internal_usage_cache=None, + llm_router=None, + ) + return custom_logger + return None + + def _get_raw_request_body(self, data: Optional[Union[dict, str]]) -> dict: + if data is None: + return {"error": "Received empty dictionary for raw request body"} + if isinstance(data, str): + try: + return json.loads(data) + except Exception: + return { + "error": "Unable to parse raw request body. Got - {}".format(data) + } + return data + + def _get_masked_api_base(self, api_base: str) -> str: + if "key=" in api_base: + # Find the position of "key=" in the string + key_index = api_base.find("key=") + 4 + # Mask the last 5 characters after "key=" + masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:] + else: + masked_api_base = api_base + return str(masked_api_base) + + def _pre_call(self, input, api_key, model=None, additional_args={}): + """ + Common helper function across the sync + async pre-call function + """ + + self.model_call_details["input"] = input + self.model_call_details["api_key"] = api_key + self.model_call_details["additional_args"] = additional_args + self.model_call_details["log_event_type"] = "pre_api_call" + if ( + model + ): # if model name was changes pre-call, overwrite the initial model call name with the new one + self.model_call_details["model"] = model + self.model_call_details["litellm_params"][ + "api_base" + ] = self._get_masked_api_base(additional_args.get("api_base", "")) + + def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 + # Log the exact input to the LLM API + litellm.error_logs["PRE_CALL"] = locals() + try: + self._pre_call( + input=input, + api_key=api_key, + model=model, + additional_args=additional_args, + ) + + # User Logging -> if you pass in a custom logging function + self._print_llm_call_debugging_log( + api_base=additional_args.get("api_base", ""), + headers=additional_args.get("headers", {}), + additional_args=additional_args, + ) + # log raw request to provider (like LangFuse) -- if opted in. + if ( + self.log_raw_request_response is True + or log_raw_request_response is True + ): + _litellm_params = self.model_call_details.get("litellm_params", {}) + _metadata = _litellm_params.get("metadata", {}) or {} + try: + # [Non-blocking Extra Debug Information in metadata] + if turn_off_message_logging is True: + _metadata[ + "raw_request" + ] = "redacted by litellm. \ + 'litellm.turn_off_message_logging=True'" + else: + curl_command = self._get_request_curl_command( + api_base=additional_args.get("api_base", ""), + headers=additional_args.get("headers", {}), + additional_args=additional_args, + data=additional_args.get("complete_input_dict", {}), + ) + + _metadata["raw_request"] = str(curl_command) + # split up, so it's easier to parse in the UI + self.model_call_details[ + "raw_request_typed_dict" + ] = RawRequestTypedDict( + raw_request_api_base=str( + additional_args.get("api_base") or "" + ), + raw_request_body=self._get_raw_request_body( + additional_args.get("complete_input_dict", {}) + ), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, + ) + except Exception as e: + self.model_call_details[ + "raw_request_typed_dict" + ] = RawRequestTypedDict( + error=str(e), + ) + _metadata[ + "raw_request" + ] = "Unable to Log \ + raw request: {}".format( + str(e) + ) + if getattr(self, "logger_fn", None) and callable(self.logger_fn): + try: + self.logger_fn( + self.model_call_details + ) # Expectation: any logger function passed in by the user should accept a dict object + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( + str(e) + ) + ) + + self.model_call_details["api_call_start_time"] = datetime.datetime.now() + # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made + callbacks = litellm.input_callback + (self.dynamic_input_callbacks or []) + for callback in callbacks: + try: + if callback == "supabase" and supabaseClient is not None: + verbose_logger.debug("reaches supabase for logging!") + model = self.model_call_details["model"] + messages = self.model_call_details["input"] + verbose_logger.debug(f"supabaseClient: {supabaseClient}") + supabaseClient.input_log_event( + model=model, + messages=messages, + end_user=self.model_call_details.get("user", "default"), + litellm_call_id=self.litellm_params["litellm_call_id"], + print_verbose=print_verbose, + ) + elif callback == "sentry" and add_breadcrumb: + try: + details_to_log = copy.deepcopy(self.model_call_details) + except Exception: + details_to_log = self.model_call_details + if litellm.turn_off_message_logging: + # make a copy of the _model_Call_details and log it + details_to_log.pop("messages", None) + details_to_log.pop("input", None) + details_to_log.pop("prompt", None) + + add_breadcrumb( + category="litellm.llm_call", + message=f"Model Call Details pre-call: {details_to_log}", + level="info", + ) + + elif isinstance(callback, CustomLogger): # custom logger class + callback.log_pre_api_call( + model=self.model, + messages=self.messages, + kwargs=self.model_call_details, + ) + elif ( + callable(callback) and customLogger is not None + ): # custom logger functions + customLogger.log_input_event( + model=self.model, + messages=self.messages, + kwargs=self.model_call_details, + print_verbose=print_verbose, + callback_func=callback, + ) + except Exception as e: + verbose_logger.exception( + "litellm.Logging.pre_call(): Exception occured - {}".format( + str(e) + ) + ) + verbose_logger.debug( + f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" + ) + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( + str(e) + ) + ) + verbose_logger.error( + f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" + ) + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + + def _print_llm_call_debugging_log( + self, + api_base: str, + headers: dict, + additional_args: dict, + ): + """ + Internal debugging helper function + + Prints the RAW curl command sent from LiteLLM + """ + if _is_debugging_on() or self.litellm_request_debug: + if json_logs: + masked_headers = self._get_masked_headers(headers) + if self.litellm_request_debug: + verbose_logger.warning( # .warning ensures this shows up in all environments + "POST Request Sent from LiteLLM", + extra={"api_base": {api_base}, **masked_headers}, + ) + else: + verbose_logger.debug( + "POST Request Sent from LiteLLM", + extra={"api_base": {api_base}, **masked_headers}, + ) + else: + headers = additional_args.get("headers", {}) + if headers is None: + headers = {} + data = additional_args.get("complete_input_dict", {}) + api_base = str(additional_args.get("api_base", "")) + curl_command = self._get_request_curl_command( + api_base=api_base, + headers=headers, + additional_args=additional_args, + data=data, + ) + if self.litellm_request_debug: + verbose_logger.warning( + f"\033[92m{curl_command}\033[0m\n" + ) # .warning ensures this shows up in all environments + else: + verbose_logger.debug(f"\033[92m{curl_command}\033[0m\n") + + def _get_request_body(self, data: dict) -> str: + return str(data) + + def _get_request_curl_command( + self, api_base: str, headers: Optional[dict], additional_args: dict, data: dict + ) -> str: + masked_api_base = self._get_masked_api_base(api_base) + if headers is None: + headers = {} + curl_command = "\n\nPOST Request Sent from LiteLLM:\n" + curl_command += "curl -X POST \\\n" + curl_command += f"{masked_api_base} \\\n" + masked_headers = self._get_masked_headers(headers) + formatted_headers = " ".join( + [f"-H '{k}: {v}'" for k, v in masked_headers.items()] + ) + curl_command += ( + f"{formatted_headers} \\\n" if formatted_headers.strip() != "" else "" + ) + curl_command += f"-d '{self._get_request_body(data)}'\n" + if additional_args.get("request_str", None) is not None: + # print the sagemaker / bedrock client request + curl_command = "\nRequest Sent from LiteLLM:\n" + request_str = additional_args.get("request_str", "") + curl_command += request_str + elif api_base == "": + curl_command = str(self.model_call_details) + return curl_command + + def _get_masked_headers( + self, headers: dict, ignore_sensitive_headers: bool = False + ) -> dict: + """ + Internal debugging helper function + + Masks the headers of the request sent from LiteLLM + """ + return _get_masked_values( + headers, ignore_sensitive_values=ignore_sensitive_headers + ) + + def post_call( + self, original_response, input=None, api_key=None, additional_args={} + ): + # Log the exact result from the LLM API, for streaming - log the type of response received + litellm.error_logs["POST_CALL"] = locals() + if isinstance(original_response, dict): + original_response = json.dumps(original_response) + try: + self.model_call_details["input"] = input + self.model_call_details["api_key"] = api_key + self.model_call_details["original_response"] = original_response + self.model_call_details["additional_args"] = additional_args + self.model_call_details["log_event_type"] = "post_api_call" + + if self.litellm_request_debug: + attr = "warning" + else: + attr = "debug" + + if json_logs: + callattr = getattr(verbose_logger, attr) + callattr( + "RAW RESPONSE:\n{}\n\n".format( + self.model_call_details.get( + "original_response", self.model_call_details + ) + ), + ) + else: + callattr = getattr(verbose_logger, attr) + callattr( + "RAW RESPONSE:\n{}\n\n".format( + self.model_call_details.get( + "original_response", self.model_call_details + ) + ) + ) + if getattr(self, "logger_fn", None) and callable(self.logger_fn): + try: + self.logger_fn( + self.model_call_details + ) # Expectation: any logger function passed in by the user should accept a dict object + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( + str(e) + ) + ) + original_response = redact_message_input_output_from_logging( + model_call_details=( + self.model_call_details + if hasattr(self, "model_call_details") + else {} + ), + result=original_response, + ) + # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made + + callbacks = litellm.input_callback + (self.dynamic_input_callbacks or []) + for callback in callbacks: + try: + if callback == "sentry" and add_breadcrumb: + verbose_logger.debug("reaches sentry breadcrumbing") + try: + details_to_log = copy.deepcopy(self.model_call_details) + except Exception: + details_to_log = self.model_call_details + if litellm.turn_off_message_logging: + # make a copy of the _model_Call_details and log it + details_to_log.pop("messages", None) + details_to_log.pop("input", None) + details_to_log.pop("prompt", None) + + add_breadcrumb( + category="litellm.llm_call", + message=f"Model Call Details post-call: {details_to_log}", + level="info", + ) + elif isinstance(callback, CustomLogger): # custom logger class + callback.log_post_api_call( + kwargs=self.model_call_details, + response_obj=None, + start_time=self.start_time, + end_time=None, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while post-call logging with integrations {}".format( + str(e) + ) + ) + verbose_logger.debug( + f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" + ) + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( + str(e) + ) + ) + + async def async_post_mcp_tool_call_hook( + self, + kwargs: dict, + response_obj: Any, + start_time: datetime.datetime, + end_time: datetime.datetime, + ): + """ + Post MCP Tool Call Hook + + Use this to modify the MCP tool call response before it is returned to the user. + """ + from litellm.types.llms.base import HiddenParams + from litellm.types.mcp import MCPPostCallResponseObject + + callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_success_callbacks, + global_callbacks=litellm.success_callback, + ) + post_mcp_tool_call_response_obj: MCPPostCallResponseObject = ( + MCPPostCallResponseObject( + mcp_tool_call_response=response_obj, hidden_params=HiddenParams() + ) + ) + for callback in callbacks: + try: + if isinstance(callback, CustomLogger): + response: Optional[ + MCPPostCallResponseObject + ] = await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=post_mcp_tool_call_response_obj, + start_time=start_time, + end_time=end_time, + ) + ###################################################################### + # if any of the callbacks modify the response, use the modified response + # current implementation returns the first modified response + ###################################################################### + if response is not None: + response_obj = self._parse_post_mcp_call_hook_response( + response=response + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format( + str(e) + ) + ) + return response_obj + + def _parse_post_mcp_call_hook_response( + self, response: Optional[MCPPostCallResponseObject] + ) -> Any: + """ + Parse the response from the post_mcp_tool_call_hook + + 1. Unpack the mcp_tool_call_response + 2. save the updated response_cost to the model_call_details + """ + if response is None: + return None + self.model_call_details["response_cost"] = response.hidden_params.response_cost + return response.mcp_tool_call_response + + def get_response_ms(self) -> float: + return ( + self.model_call_details.get("end_time", datetime.datetime.now()) + - self.model_call_details.get("start_time", datetime.datetime.now()) + ).total_seconds() * 1000 + + def set_cost_breakdown( + self, + input_cost: float, + output_cost: float, + total_cost: float, + cost_for_built_in_tools_cost_usd_dollar: float, + additional_costs: Optional[dict] = None, + original_cost: Optional[float] = None, + discount_percent: Optional[float] = None, + discount_amount: Optional[float] = None, + margin_percent: Optional[float] = None, + margin_fixed_amount: Optional[float] = None, + margin_total_amount: Optional[float] = None, + ) -> None: + """ + Helper method to store cost breakdown in the logging object. + + Args: + input_cost: Cost of input/prompt tokens + output_cost: Cost of output/completion tokens + cost_for_built_in_tools_cost_usd_dollar: Cost of built-in tools + total_cost: Total cost of request + additional_costs: Free-form additional costs dict (e.g., {"azure_model_router_flat_cost": 0.00014}) + original_cost: Cost before discount + discount_percent: Discount percentage (0.05 = 5%) + discount_amount: Discount amount in USD + margin_percent: Margin percentage applied (0.10 = 10%) + margin_fixed_amount: Fixed margin amount in USD + margin_total_amount: Total margin added in USD + """ + + self.cost_breakdown = CostBreakdown( + input_cost=input_cost, + output_cost=output_cost, + total_cost=total_cost, + tool_usage_cost=cost_for_built_in_tools_cost_usd_dollar, + ) + + # Store additional costs if provided (free-form dict for extensibility) + if ( + additional_costs + and isinstance(additional_costs, dict) + and len(additional_costs) > 0 + ): + self.cost_breakdown["additional_costs"] = additional_costs + + # Store discount information if provided + if original_cost is not None: + self.cost_breakdown["original_cost"] = original_cost + if discount_percent is not None: + self.cost_breakdown["discount_percent"] = discount_percent + if discount_amount is not None: + self.cost_breakdown["discount_amount"] = discount_amount + + # Store margin information if provided + if margin_percent is not None: + self.cost_breakdown["margin_percent"] = margin_percent + if margin_fixed_amount is not None: + self.cost_breakdown["margin_fixed_amount"] = margin_fixed_amount + if margin_total_amount is not None: + self.cost_breakdown["margin_total_amount"] = margin_total_amount + + def _response_cost_calculator( + self, + result: Union[ + ModelResponse, + ModelResponseStream, + EmbeddingResponse, + ImageResponse, + TranscriptionResponse, + TextCompletionResponse, + HttpxBinaryResponseContent, + RerankResponse, + Batch, + FineTuningJob, + ResponsesAPIResponse, + ResponseCompletedEvent, + OpenAIFileObject, + LiteLLMRealtimeStreamLoggingObject, + OpenAIModerationResponse, + "SearchResponse", + ], + cache_hit: Optional[bool] = None, + litellm_model_name: Optional[str] = None, + router_model_id: Optional[str] = None, + ) -> Optional[float]: + """ + Calculate response cost using result + logging object variables. + + used for consistent cost calculation across response headers + logging integrations. + """ + + if isinstance(result, BaseModel) and hasattr(result, "_hidden_params"): + hidden_params = getattr(result, "_hidden_params", {}) + if ( + "response_cost" in hidden_params + and hidden_params["response_cost"] is not None + ): # use cost if already calculated + return hidden_params["response_cost"] + elif ( + router_model_id is None and "model_id" in hidden_params + ): # use model_id if not already set + router_model_id = hidden_params["model_id"] + + ## RESPONSE COST ## + custom_pricing = use_custom_pricing_for_model( + litellm_params=( + self.litellm_params if hasattr(self, "litellm_params") else None + ) + ) + + prompt = "" # use for tts cost calc + _input = self.model_call_details.get("input", None) + if _input is not None and isinstance(_input, str): + prompt = _input + + if cache_hit is None: + cache_hit = self.model_call_details.get("cache_hit", False) + + try: + response_cost_calculator_kwargs = { + "response_object": result, + "model": litellm_model_name or self.model, + "cache_hit": cache_hit, + "custom_llm_provider": self.model_call_details.get( + "custom_llm_provider", None + ), + "base_model": _get_base_model_from_metadata( + model_call_details=self.model_call_details + ), + "call_type": self.call_type, + "optional_params": self.optional_params, + "custom_pricing": custom_pricing, + "prompt": prompt, + "standard_built_in_tools_params": self.standard_built_in_tools_params, + "router_model_id": router_model_id, + "litellm_logging_obj": self, + "service_tier": ( + self.optional_params.get("service_tier") + if self.optional_params + else None + ), + } + except Exception as e: # error creating kwargs for cost calculation + debug_info = StandardLoggingModelCostFailureDebugInformation( + error_str=str(e), + traceback_str=_get_traceback_str_for_error(str(e)), + ) + verbose_logger.debug( + f"response_cost_failure_debug_information: {debug_info}" + ) + self.model_call_details[ + "response_cost_failure_debug_information" + ] = debug_info + return None + + try: + response_cost = litellm.response_cost_calculator( + **response_cost_calculator_kwargs + ) + + verbose_logger.debug(f"response_cost: {response_cost}") + return response_cost + except Exception as e: # error calculating cost + debug_info = StandardLoggingModelCostFailureDebugInformation( + error_str=str(e), + traceback_str=_get_traceback_str_for_error(str(e)), + model=response_cost_calculator_kwargs["model"], + cache_hit=response_cost_calculator_kwargs["cache_hit"], + custom_llm_provider=response_cost_calculator_kwargs[ + "custom_llm_provider" + ], + base_model=response_cost_calculator_kwargs["base_model"], + call_type=response_cost_calculator_kwargs["call_type"], + custom_pricing=response_cost_calculator_kwargs["custom_pricing"], + ) + verbose_logger.debug( + f"response_cost_failure_debug_information: {debug_info}" + ) + self.model_call_details[ + "response_cost_failure_debug_information" + ] = debug_info + + return None + + async def _response_cost_calculator_async( + self, + result: Union[ + ModelResponse, + ModelResponseStream, + EmbeddingResponse, + ImageResponse, + TranscriptionResponse, + TextCompletionResponse, + HttpxBinaryResponseContent, + RerankResponse, + Batch, + FineTuningJob, + ], + cache_hit: Optional[bool] = None, + ) -> Optional[float]: + return self._response_cost_calculator(result=result, cache_hit=cache_hit) + + def should_run_logging( + self, + event_type: Literal[ + "async_success", "sync_success", "async_failure", "sync_failure" + ], + stream: bool = False, + ) -> bool: + try: + if self.model_call_details.get(f"has_logged_{event_type}", False) is True: + return False + + return True + except Exception: + return True + + def has_run_logging( + self, + event_type: Literal[ + "async_success", "sync_success", "async_failure", "sync_failure" + ], + ) -> None: + if self.stream is not None and self.stream is True: + """ + Ignore check on stream, as there can be multiple chunks + """ + return + self.model_call_details[f"has_logged_{event_type}"] = True + return + + def should_run_callback( + self, callback: litellm.CALLBACK_TYPES, litellm_params: dict, event_hook: str + ) -> bool: + if litellm.global_disable_no_log_param: + return True + + if litellm_params.get("no-log", False) is True: + # proxy cost tracking cal backs should run + + if not ( + isinstance(callback, CustomLogger) + and "_PROXY_" in callback.__class__.__name__ + ): + verbose_logger.debug( + f"no-log request, skipping logging for {event_hook} event" + ) + return False + + # Check for dynamically disabled callbacks via headers + if ( + EnterpriseCallbackControls is not None + and EnterpriseCallbackControls.is_callback_disabled_dynamically( + callback=callback, + litellm_params=litellm_params, + standard_callback_dynamic_params=self.standard_callback_dynamic_params, + ) + ): + verbose_logger.debug( + f"Callback {callback} disabled via x-litellm-disable-callbacks header for {event_hook} event" + ) + return False + + return True + + def _update_completion_start_time(self, completion_start_time: datetime.datetime): + self.completion_start_time = completion_start_time + self.model_call_details["completion_start_time"] = self.completion_start_time + + def normalize_logging_result(self, result: Any) -> Any: + """ + Some endpoints return a different type of result than what is expected by the logging system. + This function is used to normalize the result to the expected type. + """ + logging_result = result + if self.call_type == CallTypes.arealtime.value and isinstance(result, list): + combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=result + ) + logging_result = ( + RealtimeAPITokenUsageProcessor.create_logging_realtime_object( + usage=combined_usage_object, + results=result, + ) + ) + + elif ( + self.call_type == CallTypes.llm_passthrough_route.value + or self.call_type == CallTypes.allm_passthrough_route.value + ) and isinstance(result, Response): + from litellm.utils import ProviderConfigManager + + provider_config = ProviderConfigManager.get_provider_passthrough_config( + provider=self.model_call_details.get("custom_llm_provider", ""), + model=self.model, + ) + if provider_config is not None: + logging_result = provider_config.logging_non_streaming_response( + model=self.model, + custom_llm_provider=self.model_call_details.get( + "custom_llm_provider", "" + ), + httpx_response=result, + request_data=self.model_call_details.get("request_data", {}), + logging_obj=self, + endpoint=self.model_call_details.get("endpoint", ""), + ) + return logging_result + + def _process_hidden_params_and_response_cost( + self, + logging_result, + start_time, + end_time, + ): + hidden_params = getattr(logging_result, "_hidden_params", {}) + if hidden_params: + if self.model_call_details.get("litellm_params") is not None: + self.model_call_details["litellm_params"].setdefault("metadata", {}) + if self.model_call_details["litellm_params"]["metadata"] is None: + self.model_call_details["litellm_params"]["metadata"] = {} + self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = getattr(logging_result, "_hidden_params", {}) # type: ignore + + if "response_cost" in hidden_params: + self.model_call_details["response_cost"] = hidden_params["response_cost"] + else: + self.model_call_details["response_cost"] = self._response_cost_calculator( + result=logging_result + ) + + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=logging_result, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) + + def _transform_usage_objects(self, result): + if isinstance(result, ResponsesAPIResponse): + result = result.model_copy() + transformed_usage = ( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + result.usage + ) + ) + setattr(result, "usage", transformed_usage) + if ( + standard_logging_payload := self.model_call_details.get( + "standard_logging_object" + ) + ) is not None: + response_dict = ( + result.model_dump() + if hasattr(result, "model_dump") + else dict(result) + ) + # Ensure usage is properly included with transformed chat format + if transformed_usage is not None: + response_dict["usage"] = ( + transformed_usage.model_dump() + if hasattr(transformed_usage, "model_dump") + else dict(transformed_usage) + ) + standard_logging_payload["response"] = response_dict + elif isinstance(result, TranscriptionResponse): + from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( + TranscriptionUsageObjectTransformation, + ) + + result = result.model_copy() + transformed_usage = TranscriptionUsageObjectTransformation.transform_transcription_usage_object(result.usage) # type: ignore + setattr(result, "usage", transformed_usage) + return result + + def _success_handler_helper_fn( + self, + result=None, + start_time=None, + end_time=None, + cache_hit=None, + standard_logging_object: Optional[StandardLoggingPayload] = None, + ): + try: + if start_time is None: + start_time = self.start_time + if end_time is None: + end_time = datetime.datetime.now() + if self.completion_start_time is None: + self.completion_start_time = end_time + self.model_call_details[ + "completion_start_time" + ] = self.completion_start_time + + self.model_call_details["log_event_type"] = "successful_api_call" + self.model_call_details["end_time"] = end_time + self.model_call_details["cache_hit"] = cache_hit + + if self.call_type == CallTypes.anthropic_messages.value: + result = self._handle_anthropic_messages_response_logging(result=result) + elif ( + self.call_type == CallTypes.generate_content.value + or self.call_type == CallTypes.agenerate_content.value + ): + result = self._handle_non_streaming_google_genai_generate_content_response_logging( + result=result + ) + elif ( + self.call_type == CallTypes.asend_message.value + or self.call_type == CallTypes.send_message.value + ): + result = self._handle_a2a_response_logging(result=result) + + logging_result = self.normalize_logging_result(result=result) + + if ( + standard_logging_object is None + and result is not None + and self.stream is not True + ): + if self._is_recognized_call_type_for_logging( + logging_result=logging_result + ): + self._process_hidden_params_and_response_cost( + logging_result=logging_result, + start_time=start_time, + end_time=end_time, + ) + elif isinstance(result, dict) or isinstance(result, list): + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=result, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) + elif standard_logging_object is not None: + self.model_call_details[ + "standard_logging_object" + ] = standard_logging_object + else: + self.model_call_details["response_cost"] = None + + result = self._transform_usage_objects(result=result) + + if ( + litellm.max_budget + and self.stream is False + and result is not None + and isinstance(result, dict) + and "content" in result + ): + time_diff = (end_time - start_time).total_seconds() + float_diff = float(time_diff) + litellm._current_cost += litellm.completion_cost( + model=self.model, + prompt="", + completion=getattr(result, "content", ""), + total_time=float_diff, + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) + + return start_time, end_time, result + except Exception as e: + raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {str(e)}") + + def _is_recognized_call_type_for_logging( + self, + logging_result: Any, + ): + """ + Returns True if the call type is recognized for logging (eg. ModelResponse, ModelResponseStream, etc.) + """ + if ( + isinstance(logging_result, ModelResponse) + or isinstance(logging_result, ModelResponseStream) + or isinstance(logging_result, EmbeddingResponse) + or isinstance(logging_result, ImageResponse) + or isinstance(logging_result, TranscriptionResponse) + or isinstance(logging_result, TextCompletionResponse) + or isinstance(logging_result, HttpxBinaryResponseContent) # tts + or isinstance(logging_result, RerankResponse) + or isinstance(logging_result, FineTuningJob) + or isinstance(logging_result, LiteLLMBatch) + or isinstance(logging_result, ResponsesAPIResponse) + or isinstance(logging_result, OpenAIFileObject) + or isinstance(logging_result, LiteLLMRealtimeStreamLoggingObject) + or isinstance(logging_result, OpenAIModerationResponse) + or isinstance(logging_result, OCRResponse) # OCR + or isinstance(logging_result, SearchResponse) # Search API + or isinstance(logging_result, dict) + and logging_result.get("object") == "vector_store.search_results.page" + or isinstance(logging_result, dict) + and logging_result.get("object") == "search" # Search API (dict format) + or isinstance(logging_result, VideoObject) + or isinstance(logging_result, ContainerObject) + or isinstance(logging_result, LiteLLMSendMessageResponse) # A2A + or (self.call_type == CallTypes.call_mcp_tool.value) + ): + return True + return False + + def _flush_passthrough_collected_chunks_helper( + self, + raw_bytes: List[bytes], + provider_config: "BasePassthroughConfig", + ) -> Optional["CostResponseTypes"]: + all_chunks = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) + complete_streaming_response = provider_config.handle_logging_collected_chunks( + all_chunks=all_chunks, + litellm_logging_obj=self, + model=self.model, + custom_llm_provider=self.model_call_details.get("custom_llm_provider", ""), + endpoint=self.model_call_details.get("endpoint", ""), + ) + return complete_streaming_response + + def flush_passthrough_collected_chunks( + self, + raw_bytes: List[bytes], + provider_config: "BasePassthroughConfig", + ): + """ + Flush collected chunks from the logging object + This is used to log the collected chunks once streaming is done on passthrough endpoints + + 1. Decode the raw bytes to string lines + 2. Get the complete streaming response from the provider config + 3. Log the complete streaming response (trigger success handler) + This is used for passthrough endpoints + """ + complete_streaming_response = self._flush_passthrough_collected_chunks_helper( + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + + if complete_streaming_response is not None: + self.success_handler(result=complete_streaming_response) + return + + async def async_flush_passthrough_collected_chunks( + self, + raw_bytes: List[bytes], + provider_config: "BasePassthroughConfig", + ): + complete_streaming_response = self._flush_passthrough_collected_chunks_helper( + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + + if complete_streaming_response is not None: + await self.async_success_handler(result=complete_streaming_response) + return + + def success_handler( # noqa: PLR0915 + self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs + ): + verbose_logger.debug( + f"Logging Details LiteLLM-Success Call: Cache_hit={cache_hit}" + ) + if not self.should_run_logging( + event_type="sync_success" + ): # prevent double logging + return + start_time, end_time, result = self._success_handler_helper_fn( + start_time=start_time, + end_time=end_time, + result=result, + cache_hit=cache_hit, + standard_logging_object=kwargs.get("standard_logging_object", None), + ) + litellm_params = self.model_call_details.get("litellm_params", {}) + is_sync_request = ( + litellm_params.get(CallTypes.acompletion.value, False) is not True + and litellm_params.get(CallTypes.aresponses.value, False) is not True + and litellm_params.get(CallTypes.aembedding.value, False) is not True + and litellm_params.get(CallTypes.aimage_generation.value, False) is not True + and litellm_params.get(CallTypes.atranscription.value, False) is not True + ) + try: + ## BUILD COMPLETE STREAMED RESPONSE + complete_streaming_response: Optional[ + Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse] + ] = None + if "complete_streaming_response" in self.model_call_details: + return # break out of this. + complete_streaming_response = self._get_assembled_streaming_response( + result=result, + start_time=start_time, + end_time=end_time, + is_async=False, + streaming_chunks=self.sync_streaming_chunks, + ) + if complete_streaming_response is not None: + verbose_logger.debug( + "Logging Details LiteLLM-Success Call streaming complete" + ) + self.model_call_details[ + "complete_streaming_response" + ] = complete_streaming_response + self.model_call_details[ + "response_cost" + ] = self._response_cost_calculator(result=complete_streaming_response) + ## STANDARDIZED LOGGING PAYLOAD + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=complete_streaming_response, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) + if ( + standard_logging_payload := self.model_call_details.get( + "standard_logging_object" + ) + ) is not None: + # Only emit for sync requests (async_success_handler handles async) + if is_sync_request: + emit_standard_logging_payload(standard_logging_payload) + callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_success_callbacks, + global_callbacks=litellm.success_callback, + ) + + ## REDACT MESSAGES ## + result = redact_message_input_output_from_logging( + model_call_details=( + self.model_call_details + if hasattr(self, "model_call_details") + else {} + ), + result=result, + ) + ## LOGGING HOOK ## + for callback in callbacks: + if isinstance(callback, CustomLogger): + self.model_call_details, result = callback.logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + + self.has_run_logging(event_type="sync_success") + for callback in callbacks: + try: + should_run = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="success_handler", + ) + if not should_run: + continue + if callback == "promptlayer" and promptLayerLogger is not None: + print_verbose("reaches promptlayer for logging!") + promptLayerLogger.log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "supabase" and supabaseClient is not None: + print_verbose("reaches supabase for logging!") + kwargs = self.model_call_details + + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + if "complete_streaming_response" not in kwargs: + continue + else: + print_verbose("reaches supabase for streaming logging!") + result = kwargs["complete_streaming_response"] + + model = kwargs["model"] + messages = kwargs["messages"] + optional_params = kwargs.get("optional_params", {}) + litellm_params = kwargs.get("litellm_params", {}) + supabaseClient.log_event( + model=model, + messages=messages, + end_user=optional_params.get("user", "default"), + response_obj=result, + start_time=start_time, + end_time=end_time, + litellm_call_id=( + current_call_id + if ( + current_call_id := litellm_params.get( + "litellm_call_id" + ) + ) + is not None + else str(uuid.uuid4()) + ), + print_verbose=print_verbose, + ) + if callback == "wandb" and weightsBiasesLogger is not None: + print_verbose("reaches wandb for logging!") + weightsBiasesLogger.log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "logfire" and logfireLogger is not None: + verbose_logger.debug("reaches logfire for success logging!") + kwargs = {} + for k, v in self.model_call_details.items(): + if ( + k != "original_response" + ): # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + if "complete_streaming_response" not in kwargs: + continue + else: + print_verbose("reaches logfire for streaming logging!") + result = kwargs["complete_streaming_response"] + + logfireLogger.log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + level=LogfireLevel.INFO.value, # type: ignore + ) + + if callback == "lunary" and lunaryLogger is not None: + print_verbose("reaches lunary for logging!") + model = self.model + kwargs = self.model_call_details + + input = kwargs.get("messages", kwargs.get("input", None)) + + type = ( + "embed" + if self.call_type == CallTypes.embedding.value + else "llm" + ) + + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + if "complete_streaming_response" not in kwargs: + continue + else: + result = kwargs["complete_streaming_response"] + + lunaryLogger.log_event( + type=type, + kwargs=kwargs, + event="end", + model=model, + input=input, + user_id=kwargs.get("user", None), + # user_props=self.model_call_details.get("user_props", None), + extra=kwargs.get("optional_params", {}), + response_obj=result, + start_time=start_time, + end_time=end_time, + run_id=self.litellm_call_id, + print_verbose=print_verbose, + ) + if callback == "helicone" and heliconeLogger is not None: + print_verbose("reaches helicone for logging!") + model = self.model + messages = self.model_call_details["input"] + kwargs = self.model_call_details + + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + if "complete_streaming_response" not in kwargs: + continue + else: + print_verbose("reaches helicone for streaming logging!") + result = kwargs["complete_streaming_response"] + + heliconeLogger.log_success( + model=model, + messages=messages, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + kwargs=kwargs, + ) + if callback == "langfuse": + global langFuseLogger + print_verbose("reaches langfuse for success logging!") + kwargs = {} + for k, v in self.model_call_details.items(): + if ( + k != "original_response" + ): # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + verbose_logger.debug( + f"is complete_streaming_response in kwargs: {kwargs.get('complete_streaming_response', None)}" + ) + if complete_streaming_response is None: + continue + else: + print_verbose("reaches langfuse for streaming logging!") + result = kwargs["complete_streaming_response"] + + langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( + globalLangfuseLogger=langFuseLogger, + standard_callback_dynamic_params=self.standard_callback_dynamic_params, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + if langfuse_logger_to_use is not None: + _response = langfuse_logger_to_use.log_event_on_langfuse( + kwargs=kwargs, + response_obj=result, + start_time=start_time, + end_time=end_time, + user_id=kwargs.get("user", None), + ) + if _response is not None and isinstance(_response, dict): + _trace_id = _response.get("trace_id", None) + if _trace_id is not None: + in_memory_trace_id_cache.set_cache( + litellm_call_id=self.litellm_call_id, + service_name="langfuse", + trace_id=_trace_id, + ) + if callback == "greenscale" and greenscaleLogger is not None: + kwargs = {} + for k, v in self.model_call_details.items(): + if ( + k != "original_response" + ): # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + verbose_logger.debug( + f"is complete_streaming_response in kwargs: {kwargs.get('complete_streaming_response', None)}" + ) + if complete_streaming_response is None: + continue + else: + print_verbose( + "reaches greenscale for streaming logging!" + ) + result = kwargs["complete_streaming_response"] + + greenscaleLogger.log_event( + kwargs=kwargs, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "athina" and athinaLogger is not None: + deep_copy = {} + for k, v in self.model_call_details.items(): + deep_copy[k] = v + athinaLogger.log_event( + kwargs=deep_copy, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "traceloop": + deep_copy = {} + for k, v in self.model_call_details.items(): + if k != "original_response": + deep_copy[k] = v + traceloopLogger.log_event( + kwargs=deep_copy, + response_obj=result, + start_time=start_time, + end_time=end_time, + user_id=kwargs.get("user", None), + print_verbose=print_verbose, + ) + if callback == "s3": + global s3Logger + if s3Logger is None: + s3Logger = S3Logger() + if self.stream: + if "complete_streaming_response" in self.model_call_details: + print_verbose( + "S3Logger Logger: Got Stream Event - Completed Stream Response" + ) + s3Logger.log_event( + kwargs=self.model_call_details, + response_obj=self.model_call_details[ + "complete_streaming_response" + ], + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + else: + print_verbose( + "S3Logger Logger: Got Stream Event - No complete stream response as yet" + ) + else: + s3Logger.log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + + if callback == "openmeter" and is_sync_request: + global openMeterLogger + if openMeterLogger is None: + print_verbose("Instantiates openmeter client") + openMeterLogger = OpenMeterLogger() + if self.stream and complete_streaming_response is None: + openMeterLogger.log_stream_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + else: + if self.stream and complete_streaming_response: + self.model_call_details[ + "complete_response" + ] = self.model_call_details.get( + "complete_streaming_response", {} + ) + result = self.model_call_details["complete_response"] + openMeterLogger.log_success_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + if ( + isinstance(callback, CustomLogger) + and is_sync_request + and self.call_type + != CallTypes.pass_through.value # pass-through endpoints call async_log_success_event + ): # custom logger class + if self.stream and complete_streaming_response is None: + callback.log_stream_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + else: + if self.stream and complete_streaming_response: + self.model_call_details[ + "complete_response" + ] = self.model_call_details.get( + "complete_streaming_response", {} + ) + result = self.model_call_details["complete_response"] + + callback.log_success_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + if ( + callable(callback) is True + and is_sync_request + and customLogger is not None + ): # custom logger functions + print_verbose( + "success callbacks: Running Custom Callback Function - {}".format( + callback + ) + ) + + customLogger.log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + + except Exception as e: + print_verbose( + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging with integrations {traceback.format_exc()}" + ) + print_verbose( + f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" + ) + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + # Track callback logging failures in Prometheus + try: + self._handle_callback_failure(callback=callback) + except Exception: + pass + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {}".format( + str(e) + ), + ) + + async def async_success_handler( # noqa: PLR0915 + self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs + ): + """ + Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. + """ + print_verbose( + "Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit) + ) + if not self.should_run_logging( + event_type="async_success" + ): # prevent double logging + return + + ## CALCULATE COST FOR BATCH JOBS + if self.call_type == CallTypes.aretrieve_batch.value and isinstance( + result, LiteLLMBatch + ): + litellm_params = self.litellm_params or {} + litellm_metadata = litellm_params.get("litellm_metadata") or {} + if ( + litellm_metadata.get("batch_ignore_default_logging", False) is True + ): # polling job will query these frequently, don't spam db logs + return + + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ) + + # check if file id is a unified file id + is_base64_unified_file_id = _is_base64_encoded_unified_file_id(result.id) + + batch_cost = kwargs.get("batch_cost", None) + batch_usage = kwargs.get("batch_usage", None) + batch_models = kwargs.get("batch_models", None) + has_explicit_batch_data = all( + x is not None for x in (batch_cost, batch_usage, batch_models) + ) + + should_compute_batch_data = ( + not is_base64_unified_file_id + or not has_explicit_batch_data + and result.status == "completed" + ) + if has_explicit_batch_data: + result._hidden_params["response_cost"] = batch_cost + result._hidden_params["batch_models"] = batch_models + result.usage = batch_usage + + elif should_compute_batch_data: + ( + response_cost, + batch_usage, + batch_models, + ) = await _handle_completed_batch( + batch=result, + custom_llm_provider=self.custom_llm_provider, + litellm_params=self.litellm_params, + ) + + result._hidden_params["response_cost"] = response_cost + result._hidden_params["batch_models"] = batch_models + result.usage = batch_usage + + start_time, end_time, result = self._success_handler_helper_fn( + start_time=start_time, + end_time=end_time, + result=result, + cache_hit=cache_hit, + standard_logging_object=kwargs.get("standard_logging_object", None), + ) + + ## BUILD COMPLETE STREAMED RESPONSE + if "async_complete_streaming_response" in self.model_call_details: + return # break out of this. + complete_streaming_response: Optional[ + Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse] + ] = self._get_assembled_streaming_response( + result=result, + start_time=start_time, + end_time=end_time, + is_async=True, + streaming_chunks=self.streaming_chunks, + ) + + if complete_streaming_response is not None: + print_verbose("Async success callbacks: Got a complete streaming response") + + self.model_call_details[ + "async_complete_streaming_response" + ] = complete_streaming_response + + try: + if self.model_call_details.get("cache_hit", False) is True: + self.model_call_details["response_cost"] = 0.0 + else: + # check if base_model set on azure + _get_base_model_from_metadata( + model_call_details=self.model_call_details + ) + # base_model defaults to None if not set on model_info + self.model_call_details[ + "response_cost" + ] = self._response_cost_calculator( + result=complete_streaming_response + ) + + verbose_logger.debug( + f"Model={self.model}; cost={self.model_call_details['response_cost']}" + ) + except litellm.NotFoundError: + verbose_logger.warning( + f"Model={self.model} not found in completion cost map. Setting 'response_cost' to None" + ) + self.model_call_details["response_cost"] = None + + ## STANDARDIZED LOGGING PAYLOAD + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=complete_streaming_response, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) + + # print standard logging payload + if ( + standard_logging_payload := self.model_call_details.get( + "standard_logging_object" + ) + ) is not None: + emit_standard_logging_payload(standard_logging_payload) + elif self.call_type == "pass_through_endpoint": + print_verbose( + "Async success callbacks: Got a pass-through endpoint response" + ) + + self.model_call_details["async_complete_streaming_response"] = result + + # cost calculation not possible for pass-through + self.model_call_details["response_cost"] = None + + ## STANDARDIZED LOGGING PAYLOAD + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=result, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) + + # print standard logging payload + if ( + standard_logging_payload := self.model_call_details.get( + "standard_logging_object" + ) + ) is not None: + emit_standard_logging_payload(standard_logging_payload) + callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_async_success_callbacks, + global_callbacks=litellm._async_success_callback, + ) + + result = redact_message_input_output_from_logging( + model_call_details=( + self.model_call_details if hasattr(self, "model_call_details") else {} + ), + result=result, + ) + + ## LOGGING HOOK ## + + for callback in callbacks: + if isinstance(callback, CustomGuardrail): + from litellm.types.guardrails import GuardrailEventHooks + + if ( + callback.should_run_guardrail( + data=self.model_call_details, + event_type=GuardrailEventHooks.logging_only, + ) + is not True + ): + continue + + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + elif isinstance(callback, CustomLogger): + result = redact_message_input_output_from_custom_logger( + result=result, litellm_logging_obj=self, custom_logger=callback + ) + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + + self.has_run_logging(event_type="async_success") + + for callback in callbacks: + # check if callback can run for this request + litellm_params = self.model_call_details.get("litellm_params", {}) + should_run = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="async_success_handler", + ) + if not should_run: + continue + try: + if callback == "openmeter" and openMeterLogger is not None: + if self.stream is True: + if ( + "async_complete_streaming_response" + in self.model_call_details + ): + await openMeterLogger.async_log_success_event( + kwargs=self.model_call_details, + response_obj=self.model_call_details[ + "async_complete_streaming_response" + ], + start_time=start_time, + end_time=end_time, + ) + else: + await openMeterLogger.async_log_stream_event( # [TODO]: move this to being an async log stream event function + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + else: + await openMeterLogger.async_log_success_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + + if isinstance(callback, CustomLogger): # custom logger class + model_call_details: Dict = self.model_call_details + ################################## + # call redaction hook for custom logger + model_call_details = callback.redact_standard_logging_payload_from_model_call_details( + model_call_details=model_call_details + ) + ################################## + if self.stream is True: + if "async_complete_streaming_response" in model_call_details: + await callback.async_log_success_event( + kwargs=model_call_details, + response_obj=model_call_details[ + "async_complete_streaming_response" + ], + start_time=start_time, + end_time=end_time, + ) + else: + await callback.async_log_stream_event( # [TODO]: move this to being an async log stream event function + kwargs=model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + else: + await callback.async_log_success_event( + kwargs=model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + if callable(callback): # custom logger functions + global customLogger + if customLogger is None: + customLogger = CustomLogger() + if self.stream: + if ( + "async_complete_streaming_response" + in self.model_call_details + ): + await customLogger.async_log_event( + kwargs=self.model_call_details, + response_obj=self.model_call_details[ + "async_complete_streaming_response" + ], + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + else: + await customLogger.async_log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + if callback == "dynamodb": + global dynamoLogger + if dynamoLogger is None: + dynamoLogger = DyanmoDBLogger() + if self.stream: + if ( + "async_complete_streaming_response" + in self.model_call_details + ): + print_verbose( + "DynamoDB Logger: Got Stream Event - Completed Stream Response" + ) + await dynamoLogger._async_log_event( + kwargs=self.model_call_details, + response_obj=self.model_call_details[ + "async_complete_streaming_response" + ], + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + else: + print_verbose( + "DynamoDB Logger: Got Stream Event - No complete stream response as yet" + ) + else: + await dynamoLogger._async_log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + except Exception: + verbose_logger.error( + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging {traceback.format_exc()}" + ) + self._handle_callback_failure(callback=callback) + pass + + def _handle_callback_failure(self, callback: Any): + """ + Handle callback logging failures by incrementing Prometheus metrics. + + Works for both sync and async contexts since Prometheus counter increment is synchronous. + + Args: + callback: The callback that failed + """ + try: + callback_name = self._get_callback_name(callback) + + all_callbacks = litellm.logging_callback_manager._get_all_callbacks() + + for callback_obj in all_callbacks: + if hasattr(callback_obj, "increment_callback_logging_failure"): + callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore + break # Only increment once + + except Exception as e: + verbose_logger.debug(f"Error in _handle_callback_failure: {str(e)}") + + def _failure_handler_helper_fn( + self, exception, traceback_exception, start_time=None, end_time=None + ): + if start_time is None: + start_time = self.start_time + if end_time is None: + end_time = datetime.datetime.now() + + # on some exceptions, model_call_details is not always initialized, this ensures that we still log those exceptions + if not hasattr(self, "model_call_details"): + self.model_call_details = {} + + self.model_call_details["log_event_type"] = "failed_api_call" + self.model_call_details["exception"] = exception + self.model_call_details["traceback_exception"] = traceback_exception + self.model_call_details["end_time"] = end_time + self.model_call_details.setdefault("original_response", None) + self.model_call_details["response_cost"] = 0 + + if hasattr(exception, "headers") and isinstance(exception.headers, dict): + self.model_call_details.setdefault("litellm_params", {}) + metadata = ( + self.model_call_details["litellm_params"].get("metadata", {}) or {} + ) + metadata.update(exception.headers) + + ## STANDARDIZED LOGGING PAYLOAD + + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj={}, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="failure", + error_str=str(exception), + original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) + return start_time, end_time + + async def special_failure_handlers(self, exception: Exception): + """ + Custom events, emitted for specific failures. + + Currently just for router model group rate limit error + """ + from litellm.types.router import RouterErrors + + litellm_params: dict = self.model_call_details.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + + ## BASE CASE ## check if rate limit error for model group size 1 + is_base_case = False + if metadata.get("model_group_size") is not None: + model_group_size = metadata.get("model_group_size") + if isinstance(model_group_size, int) and model_group_size == 1: + is_base_case = True + ## check if special error ## + if ( + RouterErrors.no_deployments_available.value not in str(exception) + and is_base_case is False + ): + return + + ## get original model group ## + + model_group = metadata.get("model_group") or None + for callback in litellm._async_failure_callback: + if isinstance(callback, CustomLogger): # custom logger class + await callback.log_model_group_rate_limit_error( + exception=exception, + original_model_group=model_group, + kwargs=self.model_call_details, + ) # type: ignore + + def failure_handler( # noqa: PLR0915 + self, exception, traceback_exception, start_time=None, end_time=None + ): + verbose_logger.debug( + f"Logging Details LiteLLM-Failure Call: {litellm.failure_callback}" + ) + if not self.should_run_logging( + event_type="sync_failure" + ): # prevent double logging + return + litellm_params = self.model_call_details.get("litellm_params", {}) + is_sync_request = ( + litellm_params.get(CallTypes.acompletion.value, False) is not True + and litellm_params.get(CallTypes.aresponses.value, False) is not True + and litellm_params.get(CallTypes.aembedding.value, False) is not True + and litellm_params.get(CallTypes.aimage_generation.value, False) is not True + and litellm_params.get(CallTypes.atranscription.value, False) is not True + ) + + try: + start_time, end_time = self._failure_handler_helper_fn( + exception=exception, + traceback_exception=traceback_exception, + start_time=start_time, + end_time=end_time, + ) + callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_failure_callbacks, + global_callbacks=litellm.failure_callback, + ) + + result = None # result sent to all loggers, init this to None incase it's not created + + result = redact_message_input_output_from_logging( + model_call_details=( + self.model_call_details + if hasattr(self, "model_call_details") + else {} + ), + result=result, + ) + self.has_run_logging(event_type="sync_failure") + for callback in callbacks: + try: + should_run = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="failure_handler", + ) + if not should_run: + continue + if callback == "lunary" and lunaryLogger is not None: + print_verbose("reaches lunary for logging error!") + + model = self.model + + input = self.model_call_details["input"] + + _type = ( + "embed" + if self.call_type == CallTypes.embedding.value + else "llm" + ) + + lunaryLogger.log_event( + kwargs=self.model_call_details, + type=_type, + event="error", + user_id=self.model_call_details.get("user", "default"), + model=model, + input=input, + error=traceback_exception, + run_id=self.litellm_call_id, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "sentry": + print_verbose("sending exception to sentry") + if capture_exception: + capture_exception(exception) + else: + print_verbose( + f"capture exception not initialized: {capture_exception}" + ) + elif callback == "supabase" and supabaseClient is not None: + print_verbose("reaches supabase for logging!") + print_verbose(f"supabaseClient: {supabaseClient}") + supabaseClient.log_event( + model=self.model if hasattr(self, "model") else "", + messages=self.messages, + end_user=self.model_call_details.get("user", "default"), + response_obj=result, + start_time=start_time, + end_time=end_time, + litellm_call_id=self.model_call_details["litellm_call_id"], + print_verbose=print_verbose, + ) + if ( + callable(callback) and customLogger is not None + ): # custom logger functions + customLogger.log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + if ( + isinstance(callback, CustomLogger) and is_sync_request + ): # custom logger class + callback.log_failure_event( + start_time=start_time, + end_time=end_time, + response_obj=result, + kwargs=self.model_call_details, + ) + if callback == "langfuse": + global langFuseLogger + verbose_logger.debug("reaches langfuse for logging failure") + kwargs = {} + for k, v in self.model_call_details.items(): + if ( + k != "original_response" + ): # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( + globalLangfuseLogger=langFuseLogger, + standard_callback_dynamic_params=self.standard_callback_dynamic_params, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + _response = langfuse_logger_to_use.log_event_on_langfuse( + start_time=start_time, + end_time=end_time, + response_obj=None, + user_id=kwargs.get("user", None), + status_message=str(exception), + level="ERROR", + kwargs=self.model_call_details, + ) + if _response is not None and isinstance(_response, dict): + _trace_id = _response.get("trace_id", None) + if _trace_id is not None: + in_memory_trace_id_cache.set_cache( + litellm_call_id=self.litellm_call_id, + service_name="langfuse", + trace_id=_trace_id, + ) + if callback == "traceloop": + traceloopLogger.log_event( + start_time=start_time, + end_time=end_time, + response_obj=None, + user_id=self.model_call_details.get("user", None), + print_verbose=print_verbose, + status_message=str(exception), + level="ERROR", + kwargs=self.model_call_details, + ) + if callback == "logfire" and logfireLogger is not None: + verbose_logger.debug("reaches logfire for failure logging!") + kwargs = {} + for k, v in self.model_call_details.items(): + if ( + k != "original_response" + ): # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + kwargs["exception"] = exception + + logfireLogger.log_event( + kwargs=kwargs, + response_obj=result, + start_time=start_time, + end_time=end_time, + level=LogfireLevel.ERROR.value, # type: ignore + print_verbose=print_verbose, + ) + + except Exception as e: + print_verbose( + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging with integrations {str(e)}" + ) + print_verbose( + f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}" + ) + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging {}".format( + str(e) + ) + ) + + async def async_failure_handler( + self, exception, traceback_exception, start_time=None, end_time=None + ): + """ + Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. + """ + await self.special_failure_handlers(exception=exception) + if not self.should_run_logging( + event_type="async_failure" + ): # prevent double logging + return + start_time, end_time = self._failure_handler_helper_fn( + exception=exception, + traceback_exception=traceback_exception, + start_time=start_time, + end_time=end_time, + ) + + callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_async_failure_callbacks, + global_callbacks=litellm._async_failure_callback, + ) + + result = None # result sent to all loggers, init this to None incase it's not created + + self.has_run_logging(event_type="async_failure") + for callback in callbacks: + try: + litellm_params = self.model_call_details.get("litellm_params", {}) + should_run = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="async_failure_handler", + ) + if not should_run: + continue + if isinstance(callback, CustomLogger): # custom logger class + await callback.async_log_failure_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) # type: ignore + if ( + callable(callback) and customLogger is not None + ): # custom logger functions + await customLogger.async_log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure \ + logging {}\nCallback={}".format( + str(e), callback + ) + ) + # Track callback logging failures in Prometheus + self._handle_callback_failure(callback=callback) + + def _get_trace_id(self, service_name: Literal["langfuse"]) -> Optional[str]: + """ + For the given service (e.g. langfuse), return the trace_id actually logged. + + Used for constructing the url in slack alerting. + + Returns: + - str: The logged trace id + - None: If trace id not yet emitted. + """ + trace_id: Optional[str] = None + if service_name == "langfuse": + trace_id = in_memory_trace_id_cache.get_cache( + litellm_call_id=self.litellm_call_id, service_name=service_name + ) + + return trace_id + + def _get_callback_object(self, service_name: Literal["langfuse"]) -> Optional[Any]: + """ + Return dynamic callback object. + + Meant to solve issue when doing key-based/team-based logging + """ + global langFuseLogger + + if service_name == "langfuse": + if langFuseLogger is None or ( + ( + self.standard_callback_dynamic_params.get("langfuse_public_key") + is not None + and self.standard_callback_dynamic_params.get("langfuse_public_key") + != langFuseLogger.public_key + ) + or ( + self.standard_callback_dynamic_params.get("langfuse_public_key") + is not None + and self.standard_callback_dynamic_params.get("langfuse_public_key") + != langFuseLogger.public_key + ) + or ( + self.standard_callback_dynamic_params.get("langfuse_host") + is not None + and self.standard_callback_dynamic_params.get("langfuse_host") + != langFuseLogger.langfuse_host + ) + ): + return LangFuseLogger( + langfuse_public_key=self.standard_callback_dynamic_params.get( + "langfuse_public_key" + ), + langfuse_secret=self.standard_callback_dynamic_params.get( + "langfuse_secret" + ), + langfuse_host=self.standard_callback_dynamic_params.get( + "langfuse_host" + ), + ) + return langFuseLogger + + return None + + def handle_sync_success_callbacks_for_async_calls( + self, + result: Any, + start_time: datetime.datetime, + end_time: datetime.datetime, + cache_hit: Optional[Any] = None, + ) -> None: + """ + Handles calling success callbacks for Async calls. + + Why: Some callbacks - `langfuse`, `s3` are sync callbacks. We need to call them in the executor. + """ + if self._should_run_sync_callbacks_for_async_calls() is False: + return + + executor.submit( + self.success_handler, + result, + start_time, + end_time, + cache_hit, + ) + + def _should_run_sync_callbacks_for_async_calls(self) -> bool: + """ + Returns: + - bool: True if sync callbacks should be run for async calls. eg. `langfuse`, `s3` + """ + _combined_sync_callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_success_callbacks, + global_callbacks=litellm.success_callback, + ) + _filtered_success_callbacks = self._remove_internal_custom_logger_callbacks( + _combined_sync_callbacks + ) + _filtered_success_callbacks = self._remove_internal_litellm_callbacks( + _filtered_success_callbacks + ) + return len(_filtered_success_callbacks) > 0 + + def get_combined_callback_list( + self, dynamic_success_callbacks: Optional[List], global_callbacks: List + ) -> List: + if dynamic_success_callbacks is None: + return list(global_callbacks) + return list(set(dynamic_success_callbacks + global_callbacks)) + + def _remove_internal_litellm_callbacks(self, callbacks: List) -> List: + """ + Creates a filtered list of callbacks, excluding internal LiteLLM callbacks. + + Args: + callbacks: List of callback functions/strings to filter + + Returns: + List of filtered callbacks with internal ones removed + """ + filtered = [ + cb for cb in callbacks if not self._is_internal_litellm_proxy_callback(cb) + ] + + verbose_logger.debug(f"Filtered callbacks: {filtered}") + return filtered + + def _get_callback_name(self, cb) -> str: + """ + Helper to get the name of a callback function + + Args: + cb: The callback object/function/string to get the name of + + Returns: + The name of the callback + """ + if isinstance(cb, str): + return cb + if hasattr(cb, "__name__"): + return cb.__name__ + if hasattr(cb, "__func__"): + return cb.__func__.__name__ + if hasattr(cb, "__class__"): + return cb.__class__.__name__ + return str(cb) + + def _is_internal_litellm_proxy_callback(self, cb) -> bool: + """Helper to check if a callback is internal""" + INTERNAL_PREFIXES = [ + "_PROXY", + "_service_logger.ServiceLogging", + "sync_deployment_callback_on_success", + ] + if isinstance(cb, str): + return False + + if not callable(cb): + return True + + cb_name = self._get_callback_name(cb) + return any(prefix in cb_name for prefix in INTERNAL_PREFIXES) + + def _remove_internal_custom_logger_callbacks(self, callbacks: List) -> List: + """ + Removes internal custom logger callbacks from the list. + """ + _new_callbacks = [] + for _c in callbacks: + if isinstance(_c, CustomLogger): + continue + elif ( + isinstance(_c, str) + and _c in litellm._known_custom_logger_compatible_callbacks + ): + continue + _new_callbacks.append(_c) + return _new_callbacks + + def _get_assembled_streaming_response( + self, + result: Union[ + ModelResponse, + TextCompletionResponse, + ModelResponseStream, + ResponseCompletedEvent, + Any, + ], + start_time: datetime.datetime, + end_time: datetime.datetime, + is_async: bool, + streaming_chunks: List[Any], + ) -> Optional[Union[ModelResponse, TextCompletionResponse, ResponsesAPIResponse]]: + if isinstance(result, ModelResponse): + return result + elif isinstance(result, TextCompletionResponse): + return result + elif isinstance(result, ResponseCompletedEvent): + ## return unified Usage object + if isinstance(result.response.usage, ResponseAPIUsage): + transformed_usage = ( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + result.response.usage + ) + ) + # Set as dict instead of Usage object so model_dump() serializes it correctly + setattr( + result.response, + "usage", + ( + transformed_usage.model_dump() + if hasattr(transformed_usage, "model_dump") + else dict(transformed_usage) + ), + ) + return result.response + else: + return None + return None + + def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse: + """ + Handles logging for Anthropic messages responses. + + Args: + result: The response object from the model call + + Returns: + The the response object from the model call + + - For Non-streaming responses, we need to transform the response to a ModelResponse object. + - For streaming responses, anthropic_messages handler calls success_handler with a assembled ModelResponse. + """ + import httpx + + if self.stream and isinstance(result, ModelResponse): + return result + elif isinstance(result, ModelResponse): + return result + + httpx_response = self.model_call_details.get("httpx_response", None) + if httpx_response and isinstance(httpx_response, httpx.Response): + result = litellm.AnthropicConfig().transform_response( + raw_response=httpx_response, + model_response=litellm.ModelResponse(), + model=self.model, + messages=[], + logging_obj=self, + optional_params={}, + api_key="", + request_data={}, + encoding=litellm.encoding, + json_mode=False, + litellm_params={}, + ) + else: + from litellm.types.llms.anthropic import AnthropicResponse + + pydantic_result = AnthropicResponse.model_validate(result) + import httpx + + result = litellm.AnthropicConfig().transform_parsed_response( + completion_response=pydantic_result.model_dump(), + raw_response=httpx.Response( + status_code=200, + headers={}, + ), + model_response=litellm.ModelResponse(), + json_mode=None, + ) + return result + + def _handle_non_streaming_google_genai_generate_content_response_logging( + self, result: Any + ) -> ModelResponse: + """ + Handles logging for Google GenAI generate content responses. + """ + import httpx + + httpx_response = self.model_call_details.get("httpx_response", None) + if httpx_response is None: + raise ValueError("Google GenAI Generate Content: httpx_response is None") + dict_result = httpx_response.json() + result = litellm.VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response( + completion_response=dict_result, + model_response=litellm.ModelResponse(), + model=self.model, + logging_obj=self, + raw_response=httpx.Response( + status_code=200, + headers={}, + ), + ) + return result + + def _handle_a2a_response_logging(self, result: Any) -> Any: + """ + Handles logging for A2A (Agent-to-Agent) responses. + + Adds usage from model_call_details to the result if available. + Uses Pydantic's model_copy to avoid modifying the original response. + + Args: + result: The LiteLLMSendMessageResponse from the A2A call + + Returns: + The response object with usage added if available + """ + # Get usage from model_call_details (set by asend_message) + usage = self.model_call_details.get("usage") + if usage is None: + return result + + # Deep copy result and add usage + result_copy = result.model_copy(deep=True) + result_copy.usage = ( + usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) + ) + return result_copy + + +def _get_masked_values( + sensitive_object: dict, + ignore_sensitive_values: bool = False, + mask_all_values: bool = False, + unmasked_length: int = 4, + number_of_asterisks: Optional[int] = 4, +) -> dict: + """ + Internal debugging helper function + + Masks the headers of the request sent from LiteLLM + + Args: + masked_length: Optional length for the masked portion (number of *). If set, will use exactly this many * + regardless of original string length. The total length will be unmasked_length + masked_length. + """ + sensitive_keywords = [ + "authorization", + "token", + "key", + "secret", + "vertex_credentials", + ] + return { + k: ( + # If ignore_sensitive_values is True, or if this key doesn't contain sensitive keywords, return original value + v + if ignore_sensitive_values + or not any( + sensitive_keyword in k.lower() + for sensitive_keyword in sensitive_keywords + ) + else ( + # Apply masking to sensitive keys + ( + v[: unmasked_length // 2] + + "*" * number_of_asterisks + + v[-unmasked_length // 2 :] + ) + if ( + isinstance(v, str) + and len(v) > unmasked_length + and number_of_asterisks is not None + ) + else ( + ( + v[: unmasked_length // 2] + + "*" * (len(v) - unmasked_length) + + v[-unmasked_length // 2 :] + ) + if (isinstance(v, str) and len(v) > unmasked_length) + else ("*****" if isinstance(v, str) else v) + ) + ) + ) + for k, v in sensitive_object.items() + } + + +def set_callbacks(callback_list, function_id=None): # noqa: PLR0915 + """ + Globally sets the callback client + """ + global sentry_sdk_instance, capture_exception, add_breadcrumb, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, supabaseClient, lunaryLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, logfireLogger, dynamoLogger, s3Logger, dataDogLogger, prometheusLogger, greenscaleLogger, openMeterLogger, deepevalLogger + + try: + for callback in callback_list: + if callback == "sentry": + try: + import sentry_sdk + except ImportError: + print_verbose("Package 'sentry_sdk' is missing. Installing it...") + subprocess.check_call( + [sys.executable, "-m", "pip", "install", "sentry_sdk"] + ) + import sentry_sdk + from sentry_sdk.scrubber import EventScrubber + + sentry_sdk_instance = sentry_sdk + sentry_trace_rate = ( + os.environ.get("SENTRY_API_TRACE_RATE") + if "SENTRY_API_TRACE_RATE" in os.environ + else "1.0" + ) + sentry_sample_rate = ( + os.environ.get("SENTRY_API_SAMPLE_RATE") + if "SENTRY_API_SAMPLE_RATE" in os.environ + else "1.0" + ) + sentry_sdk_instance.init( + dsn=os.environ.get("SENTRY_DSN"), + traces_sample_rate=float(sentry_trace_rate), # type: ignore + sample_rate=float( + sentry_sample_rate if sentry_sample_rate else 1.0 + ), + send_default_pii=False, # Prevent sending Personal Identifiable Information + event_scrubber=EventScrubber( + denylist=SENTRY_DENYLIST, pii_denylist=SENTRY_PII_DENYLIST + ), + environment=os.environ.get("SENTRY_ENVIRONMENT", "production"), + ) + capture_exception = sentry_sdk_instance.capture_exception + add_breadcrumb = sentry_sdk_instance.add_breadcrumb + elif callback == "slack": + try: + from slack_bolt import App + except ImportError: + print_verbose("Package 'slack_bolt' is missing. Installing it...") + subprocess.check_call( + [sys.executable, "-m", "pip", "install", "slack_bolt"] + ) + from slack_bolt import App + slack_app = App( + token=os.environ.get("SLACK_API_TOKEN"), + signing_secret=os.environ.get("SLACK_API_SECRET"), + ) + alerts_channel = os.environ["SLACK_API_CHANNEL"] + print_verbose(f"Initialized Slack App: {slack_app}") + elif callback == "traceloop": + traceloopLogger = TraceloopLogger() + elif callback == "athina": + athinaLogger = AthinaLogger() + print_verbose("Initialized Athina Logger") + elif callback == "helicone": + heliconeLogger = HeliconeLogger() + elif callback == "lunary": + lunaryLogger = LunaryLogger() + elif callback == "promptlayer": + promptLayerLogger = PromptLayerLogger() + elif callback == "langfuse": + langFuseLogger = LangFuseLogger( + langfuse_public_key=None, langfuse_secret=None, langfuse_host=None + ) + elif callback == "openmeter": + openMeterLogger = OpenMeterLogger() + elif callback == "datadog": + dataDogLogger = DataDogLogger() + elif callback == "dynamodb": + dynamoLogger = DyanmoDBLogger() + elif callback == "s3": + s3Logger = S3Logger() + elif callback == "wandb": + from litellm.integrations.weights_biases import WeightsBiasesLogger + + weightsBiasesLogger = WeightsBiasesLogger() + elif callback == "logfire": + logfireLogger = LogfireLogger() + elif callback == "supabase": + print_verbose("instantiating supabase") + supabaseClient = Supabase() + elif callback == "greenscale": + greenscaleLogger = GreenscaleLogger() + print_verbose("Initialized Greenscale Logger") + elif callable(callback): + customLogger = CustomLogger() + except Exception as e: + raise e + return None + + +def _init_custom_logger_compatible_class( # noqa: PLR0915 + logging_integration: _custom_logger_compatible_callbacks_literal, + internal_usage_cache: Optional[DualCache], + llm_router: Optional[ + Any + ], # expect litellm.Router, but typing errors due to circular import + custom_logger_init_args: Optional[dict] = {}, +) -> Optional[CustomLogger]: + """ + Initialize a custom logger compatible class + """ + try: + custom_logger_init_args = custom_logger_init_args or {} + if logging_integration == "agentops": # Add AgentOps initialization + for callback in _in_memory_loggers: + if isinstance(callback, AgentOps): + return callback # type: ignore + + agentops_logger = AgentOps() + _in_memory_loggers.append(agentops_logger) + return agentops_logger # type: ignore + elif logging_integration == "lago": + for callback in _in_memory_loggers: + if isinstance(callback, LagoLogger): + return callback # type: ignore + + lago_logger = LagoLogger() + _in_memory_loggers.append(lago_logger) + return lago_logger # type: ignore + elif logging_integration == "openmeter": + for callback in _in_memory_loggers: + if isinstance(callback, OpenMeterLogger): + return callback # type: ignore + + _openmeter_logger = OpenMeterLogger() + _in_memory_loggers.append(_openmeter_logger) + return _openmeter_logger # type: ignore + elif logging_integration == "posthog": + for callback in _in_memory_loggers: + if isinstance(callback, PostHogLogger): + return callback # type: ignore + + _posthog_logger = PostHogLogger() + _in_memory_loggers.append(_posthog_logger) + return _posthog_logger # type: ignore + elif logging_integration == "braintrust": + from litellm.integrations.braintrust_logging import BraintrustLogger + + for callback in _in_memory_loggers: + if isinstance(callback, BraintrustLogger): + return callback # type: ignore + + braintrust_logger = BraintrustLogger() + _in_memory_loggers.append(braintrust_logger) + return braintrust_logger # type: ignore + elif logging_integration == "langsmith": + for callback in _in_memory_loggers: + if isinstance(callback, LangsmithLogger): + return callback # type: ignore + + _langsmith_logger = LangsmithLogger() + _in_memory_loggers.append(_langsmith_logger) + return _langsmith_logger # type: ignore + elif logging_integration == "argilla": + for callback in _in_memory_loggers: + if isinstance(callback, ArgillaLogger): + return callback # type: ignore + + _argilla_logger = ArgillaLogger() + _in_memory_loggers.append(_argilla_logger) + return _argilla_logger # type: ignore + elif logging_integration == "literalai": + for callback in _in_memory_loggers: + if isinstance(callback, LiteralAILogger): + return callback # type: ignore + + _literalai_logger = LiteralAILogger() + _in_memory_loggers.append(_literalai_logger) + return _literalai_logger # type: ignore + elif logging_integration == "prometheus": + PrometheusLogger = _get_cached_prometheus_logger() + + for callback in _in_memory_loggers: + if isinstance(callback, PrometheusLogger): + return callback # type: ignore + + _prometheus_logger = PrometheusLogger() + _in_memory_loggers.append(_prometheus_logger) + return _prometheus_logger # type: ignore + elif logging_integration == "datadog": + for callback in _in_memory_loggers: + if isinstance(callback, DataDogLogger): + return callback # type: ignore + + _datadog_logger = DataDogLogger() + _in_memory_loggers.append(_datadog_logger) + return _datadog_logger # type: ignore + elif logging_integration == "datadog_llm_observability": + _datadog_llm_obs_logger = DataDogLLMObsLogger() + _in_memory_loggers.append(_datadog_llm_obs_logger) + return _datadog_llm_obs_logger # type: ignore + elif logging_integration == "azure_sentinel": + for callback in _in_memory_loggers: + if isinstance(callback, AzureSentinelLogger): + return callback # type: ignore + + _azure_sentinel_logger = AzureSentinelLogger() + _in_memory_loggers.append(_azure_sentinel_logger) + return _azure_sentinel_logger # type: ignore + elif logging_integration == "gcs_bucket": + for callback in _in_memory_loggers: + if isinstance(callback, GCSBucketLogger): + return callback # type: ignore + + _gcs_bucket_logger = GCSBucketLogger() + _in_memory_loggers.append(_gcs_bucket_logger) + return _gcs_bucket_logger # type: ignore + elif logging_integration == "s3_v2": + for callback in _in_memory_loggers: + if isinstance(callback, S3V2Logger): + return callback # type: ignore + + _s3_v2_logger = S3V2Logger() + _in_memory_loggers.append(_s3_v2_logger) + return _s3_v2_logger # type: ignore + elif logging_integration == "aws_sqs": + for callback in _in_memory_loggers: + if isinstance(callback, SQSLogger): + return callback # type: ignore + + _aws_sqs_logger = SQSLogger() + _in_memory_loggers.append(_aws_sqs_logger) + return _aws_sqs_logger # type: ignore + elif logging_integration == "azure_storage": + for callback in _in_memory_loggers: + if isinstance(callback, AzureBlobStorageLogger): + return callback # type: ignore + + _azure_storage_logger = AzureBlobStorageLogger() + _in_memory_loggers.append(_azure_storage_logger) + return _azure_storage_logger # type: ignore + elif logging_integration == "opik": + for callback in _in_memory_loggers: + if isinstance(callback, OpikLogger): + return callback # type: ignore + + _opik_logger = OpikLogger() + _in_memory_loggers.append(_opik_logger) + return _opik_logger # type: ignore + elif logging_integration == "arize": + from litellm.integrations.opentelemetry import ( + OpenTelemetry, + OpenTelemetryConfig, + ) + + arize_config = ArizeLogger.get_arize_config() + if arize_config.endpoint is None: + raise ValueError( + "No valid endpoint found for Arize, please set 'ARIZE_ENDPOINT' to your GRPC endpoint or 'ARIZE_HTTP_ENDPOINT' to your HTTP endpoint" + ) + otel_config = OpenTelemetryConfig( + exporter=arize_config.protocol, + endpoint=arize_config.endpoint, + service_name=arize_config.project_name, + ) + + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" + for callback in _in_memory_loggers: + if ( + isinstance(callback, ArizeLogger) + and callback.callback_name == "arize" + ): + return callback # type: ignore + _arize_otel_logger = ArizeLogger(config=otel_config, callback_name="arize") + _in_memory_loggers.append(_arize_otel_logger) + return _arize_otel_logger # type: ignore + elif logging_integration == "arize_phoenix": + from litellm.integrations.opentelemetry import ( + OpenTelemetry, + OpenTelemetryConfig, + ) + + arize_phoenix_config = ArizePhoenixLogger.get_arize_phoenix_config() + otel_config = OpenTelemetryConfig( + exporter=arize_phoenix_config.protocol, + endpoint=arize_phoenix_config.endpoint, + headers=arize_phoenix_config.otlp_auth_headers, + ) + if arize_phoenix_config.project_name: + existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") + # Add openinference.project.name attribute + if existing_attrs: + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" + else: + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"openinference.project.name={arize_phoenix_config.project_name}" + + # Set Phoenix project name from environment variable + phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) + if phoenix_project_name: + existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") + # Add openinference.project.name attribute + if existing_attrs: + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}" + else: + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"openinference.project.name={phoenix_project_name}" + + # auth can be disabled on local deployments of arize phoenix + if arize_phoenix_config.otlp_auth_headers is not None: + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = arize_phoenix_config.otlp_auth_headers + + for callback in _in_memory_loggers: + if ( + isinstance(callback, ArizePhoenixLogger) + and callback.callback_name == "arize_phoenix" + ): + return callback # type: ignore + _arize_phoenix_otel_logger = ArizePhoenixLogger( + config=otel_config, callback_name="arize_phoenix" + ) + _in_memory_loggers.append(_arize_phoenix_otel_logger) + return _arize_phoenix_otel_logger # type: ignore + elif logging_integration == "levo": + from litellm.integrations.levo.levo import LevoLogger + from litellm.integrations.opentelemetry import ( + OpenTelemetry, + OpenTelemetryConfig, + ) + + levo_config = LevoLogger.get_levo_config() + otel_config = OpenTelemetryConfig( + exporter=levo_config.protocol, + endpoint=levo_config.endpoint, + headers=levo_config.otlp_auth_headers, + ) + + # Check if LevoLogger instance already exists + for callback in _in_memory_loggers: + if ( + isinstance(callback, LevoLogger) + and callback.callback_name == "levo" + ): + return callback # type: ignore + + _levo_otel_logger = LevoLogger(config=otel_config, callback_name="levo") + _in_memory_loggers.append(_levo_otel_logger) + return _levo_otel_logger # type: ignore + elif logging_integration == "otel": + from litellm.integrations.opentelemetry import OpenTelemetry + + for callback in _in_memory_loggers: + if type(callback) is OpenTelemetry: + return callback # type: ignore + otel_logger = OpenTelemetry( + **_get_custom_logger_settings_from_proxy_server( + callback_name=logging_integration + ) + ) + _in_memory_loggers.append(otel_logger) + return otel_logger # type: ignore + + elif logging_integration == "galileo": + for callback in _in_memory_loggers: + if isinstance(callback, GalileoObserve): + return callback # type: ignore + + galileo_logger = GalileoObserve() + _in_memory_loggers.append(galileo_logger) + return galileo_logger # type: ignore + elif logging_integration == "cloudzero": + from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger + + for callback in _in_memory_loggers: + if isinstance(callback, CloudZeroLogger): + return callback # type: ignore + cloudzero_logger = CloudZeroLogger() + _in_memory_loggers.append(cloudzero_logger) + return cloudzero_logger # type: ignore + elif logging_integration == "focus": + from litellm.integrations.focus.focus_logger import FocusLogger + + for callback in _in_memory_loggers: + if isinstance(callback, FocusLogger): + return callback # type: ignore + focus_logger = FocusLogger() + _in_memory_loggers.append(focus_logger) + return focus_logger # type: ignore + elif logging_integration == "deepeval": + for callback in _in_memory_loggers: + if isinstance(callback, DeepEvalLogger): + return callback # type: ignore + deepeval_logger = DeepEvalLogger() + _in_memory_loggers.append(deepeval_logger) + return deepeval_logger # type: ignore + + elif logging_integration == "logfire": + if "LOGFIRE_TOKEN" not in os.environ: + raise ValueError("LOGFIRE_TOKEN not found in environment variables") + from litellm.integrations.opentelemetry import ( + OpenTelemetry, + OpenTelemetryConfig, + ) + + logfire_base_url = os.getenv( + "LOGFIRE_BASE_URL", "https://logfire-api.pydantic.dev" + ) + otel_config = OpenTelemetryConfig( + exporter="otlp_http", + endpoint=f"{logfire_base_url.rstrip('/')}/v1/traces", + headers=f"Authorization={os.getenv('LOGFIRE_TOKEN')}", + ) + for callback in _in_memory_loggers: + if isinstance(callback, OpenTelemetry): + return callback # type: ignore + _otel_logger = OpenTelemetry(config=otel_config) + _in_memory_loggers.append(_otel_logger) + return _otel_logger # type: ignore + elif logging_integration == "dynamic_rate_limiter": + from litellm.proxy.hooks.dynamic_rate_limiter import ( + _PROXY_DynamicRateLimitHandler, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_DynamicRateLimitHandler): + return callback # type: ignore + + if internal_usage_cache is None: + raise Exception( + "Internal Error: Cache cannot be empty - internal_usage_cache={}".format( + internal_usage_cache + ) + ) + + dynamic_rate_limiter_obj = _PROXY_DynamicRateLimitHandler( + internal_usage_cache=internal_usage_cache + ) + + if llm_router is not None and isinstance(llm_router, litellm.Router): + dynamic_rate_limiter_obj.update_variables(llm_router=llm_router) + _in_memory_loggers.append(dynamic_rate_limiter_obj) + return dynamic_rate_limiter_obj # type: ignore + elif logging_integration == "dynamic_rate_limiter_v3": + from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3): + return callback # type: ignore + + if internal_usage_cache is None: + raise Exception( + "Internal Error: Cache cannot be empty - internal_usage_cache={}".format( + internal_usage_cache + ) + ) + + dynamic_rate_limiter_obj_v3 = _PROXY_DynamicRateLimitHandlerV3( + internal_usage_cache=internal_usage_cache + ) + + if llm_router is not None and isinstance(llm_router, litellm.Router): + dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router) + _in_memory_loggers.append(dynamic_rate_limiter_obj_v3) + return dynamic_rate_limiter_obj_v3 # type: ignore + elif logging_integration == "langtrace": + if "LANGTRACE_API_KEY" not in os.environ: + raise ValueError("LANGTRACE_API_KEY not found in environment variables") + + from litellm.integrations.opentelemetry import ( + OpenTelemetry, + OpenTelemetryConfig, + ) + + otel_config = OpenTelemetryConfig( + exporter="otlp_http", + endpoint="https://langtrace.ai/api/trace", + ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" + for callback in _in_memory_loggers: + if ( + isinstance(callback, OpenTelemetry) + and callback.callback_name == "langtrace" + ): + return callback # type: ignore + _otel_logger = OpenTelemetry(config=otel_config, callback_name="langtrace") + _in_memory_loggers.append(_otel_logger) + return _otel_logger # type: ignore + + elif logging_integration == "mlflow": + for callback in _in_memory_loggers: + if isinstance(callback, MlflowLogger): + return callback # type: ignore + + _mlflow_logger = MlflowLogger() + _in_memory_loggers.append(_mlflow_logger) + return _mlflow_logger # type: ignore + elif logging_integration == "langfuse": + for callback in _in_memory_loggers: + if isinstance(callback, LangfusePromptManagement): + return callback + + langfuse_logger = LangfusePromptManagement() + _in_memory_loggers.append(langfuse_logger) + return langfuse_logger # type: ignore + elif logging_integration == "langfuse_otel": + from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger + + for callback in _in_memory_loggers: + if ( + isinstance(callback, LangfuseOtelLogger) + and callback.callback_name == "langfuse_otel" + ): + return callback # type: ignore + # Allow LangfuseOtelLogger to initialize its own config safely + # This prevents startup crashes if LANGFUSE keys are not in env (e.g. for dynamic usage) + _otel_logger = LangfuseOtelLogger( + config=None, callback_name="langfuse_otel" + ) + _in_memory_loggers.append(_otel_logger) + return _otel_logger # type: ignore + elif logging_integration == "weave_otel": + from litellm.integrations.opentelemetry import OpenTelemetryConfig + from litellm.integrations.weave.weave_otel import ( + WeaveOtelLogger, + get_weave_otel_config, + ) + + weave_otel_config = get_weave_otel_config() + + otel_config = OpenTelemetryConfig( + exporter=weave_otel_config.protocol, + endpoint=weave_otel_config.endpoint, + headers=weave_otel_config.otlp_auth_headers, + ) + + for callback in _in_memory_loggers: + if ( + isinstance(callback, WeaveOtelLogger) + and callback.callback_name == "weave_otel" + ): + return callback # type: ignore + _otel_logger = WeaveOtelLogger( + config=otel_config, callback_name="weave_otel" + ) + _in_memory_loggers.append(_otel_logger) + return _otel_logger # type: ignore + elif logging_integration == "pagerduty": + for callback in _in_memory_loggers: + if isinstance(callback, PagerDutyAlerting): + return callback + pagerduty_logger = PagerDutyAlerting(**custom_logger_init_args) + _in_memory_loggers.append(pagerduty_logger) + return pagerduty_logger # type: ignore + elif logging_integration == "anthropic_cache_control_hook": + for callback in _in_memory_loggers: + if isinstance(callback, AnthropicCacheControlHook): + return callback + anthropic_cache_control_hook = AnthropicCacheControlHook() + _in_memory_loggers.append(anthropic_cache_control_hook) + return anthropic_cache_control_hook # type: ignore + elif logging_integration == "vector_store_pre_call_hook": + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( + VectorStorePreCallHook, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, VectorStorePreCallHook): + return callback + vector_store_pre_call_hook = VectorStorePreCallHook() + _in_memory_loggers.append(vector_store_pre_call_hook) + return vector_store_pre_call_hook # type: ignore + elif logging_integration == "gcs_pubsub": + for callback in _in_memory_loggers: + if isinstance(callback, GcsPubSubLogger): + return callback + _gcs_pubsub_logger = GcsPubSubLogger() + _in_memory_loggers.append(_gcs_pubsub_logger) + return _gcs_pubsub_logger # type: ignore + elif logging_integration == "generic_api": + for callback in _in_memory_loggers: + if isinstance(callback, GenericAPILogger): + return callback + generic_api_logger = GenericAPILogger() + _in_memory_loggers.append(generic_api_logger) + return generic_api_logger # type: ignore + elif logging_integration == "resend_email": + for callback in _in_memory_loggers: + if isinstance(callback, ResendEmailLogger): + return callback + resend_email_logger = ResendEmailLogger() + _in_memory_loggers.append(resend_email_logger) + return resend_email_logger # type: ignore + elif logging_integration == "sendgrid_email": + for callback in _in_memory_loggers: + if isinstance(callback, SendGridEmailLogger): + return callback + sendgrid_email_logger = SendGridEmailLogger() + _in_memory_loggers.append(sendgrid_email_logger) + return sendgrid_email_logger # type: ignore + elif logging_integration == "smtp_email": + for callback in _in_memory_loggers: + if isinstance(callback, SMTPEmailLogger): + return callback + smtp_email_logger = SMTPEmailLogger() + _in_memory_loggers.append(smtp_email_logger) + return smtp_email_logger # type: ignore + elif logging_integration == "humanloop": + for callback in _in_memory_loggers: + if isinstance(callback, HumanloopLogger): + return callback + + humanloop_logger = HumanloopLogger() + _in_memory_loggers.append(humanloop_logger) + return humanloop_logger # type: ignore + elif logging_integration == "dotprompt": + for callback in _in_memory_loggers: + if isinstance(callback, DotpromptManager): + return callback + + dotprompt_logger = DotpromptManager() + _in_memory_loggers.append(dotprompt_logger) + return dotprompt_logger # type: ignore + elif logging_integration == "bitbucket": + from litellm.integrations.bitbucket.bitbucket_prompt_manager import ( + BitBucketPromptManager, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, BitBucketPromptManager): + return callback + + # Get global BitBucket config + bitbucket_config = getattr(litellm, "global_bitbucket_config", None) + if bitbucket_config is None: + raise ValueError( + "BitBucket configuration not found. Please set litellm.global_bitbucket_config first." + ) + + 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( + f"[Non-Blocking Error] Error initializing custom logger: {e}" + ) + return None + return None + + +def get_custom_logger_compatible_class( # noqa: PLR0915 + logging_integration: _custom_logger_compatible_callbacks_literal, +) -> Optional[CustomLogger]: + try: + if logging_integration == "lago": + for callback in _in_memory_loggers: + if isinstance(callback, LagoLogger): + return callback + elif logging_integration == "openmeter": + for callback in _in_memory_loggers: + if isinstance(callback, OpenMeterLogger): + return callback + elif logging_integration == "braintrust": + from litellm.integrations.braintrust_logging import BraintrustLogger + + for callback in _in_memory_loggers: + if isinstance(callback, BraintrustLogger): + return callback + elif logging_integration == "galileo": + for callback in _in_memory_loggers: + if isinstance(callback, GalileoObserve): + return callback + elif logging_integration == "cloudzero": + from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger + + for callback in _in_memory_loggers: + if isinstance(callback, CloudZeroLogger): + return callback + elif logging_integration == "focus": + from litellm.integrations.focus.focus_logger import FocusLogger + + for callback in _in_memory_loggers: + if isinstance(callback, FocusLogger): + return callback + elif logging_integration == "deepeval": + for callback in _in_memory_loggers: + if isinstance(callback, DeepEvalLogger): + return callback + elif logging_integration == "langsmith": + for callback in _in_memory_loggers: + if isinstance(callback, LangsmithLogger): + return callback + elif logging_integration == "argilla": + for callback in _in_memory_loggers: + if isinstance(callback, ArgillaLogger): + return callback + elif logging_integration == "literalai": + for callback in _in_memory_loggers: + if isinstance(callback, LiteralAILogger): + return callback + elif logging_integration == "prometheus": + PrometheusLogger = _get_cached_prometheus_logger() + for callback in _in_memory_loggers: + if isinstance(callback, PrometheusLogger): + return callback + elif logging_integration == "datadog": + for callback in _in_memory_loggers: + if isinstance(callback, DataDogLogger): + return callback + elif logging_integration == "datadog_llm_observability": + for callback in _in_memory_loggers: + if isinstance(callback, DataDogLLMObsLogger): + return callback + elif logging_integration == "azure_sentinel": + for callback in _in_memory_loggers: + if isinstance(callback, AzureSentinelLogger): + return callback + elif logging_integration == "gcs_bucket": + for callback in _in_memory_loggers: + if isinstance(callback, GCSBucketLogger): + return callback + elif logging_integration == "s3_v2": + for callback in _in_memory_loggers: + if isinstance(callback, S3V2Logger): + return callback + elif logging_integration == "aws_sqs": + for callback in _in_memory_loggers: + if isinstance(callback, SQSLogger): + return callback + _aws_sqs_logger = SQSLogger() + _in_memory_loggers.append(_aws_sqs_logger) + return _aws_sqs_logger # type: ignore + elif logging_integration == "azure_storage": + for callback in _in_memory_loggers: + if isinstance(callback, AzureBlobStorageLogger): + return callback + elif logging_integration == "opik": + for callback in _in_memory_loggers: + if isinstance(callback, OpikLogger): + return callback + elif logging_integration == "langfuse": + for callback in _in_memory_loggers: + if isinstance(callback, LangfusePromptManagement): + return callback + elif logging_integration == "otel": + from litellm.integrations.opentelemetry import OpenTelemetry + + for callback in _in_memory_loggers: + if isinstance(callback, OpenTelemetry): + return callback + elif logging_integration == "arize": + if "ARIZE_API_KEY" not in os.environ: + raise ValueError("ARIZE_API_KEY not found in environment variables") + for callback in _in_memory_loggers: + if ( + isinstance(callback, ArizeLogger) + and callback.callback_name == "arize" + ): + return callback + elif logging_integration == "logfire": + if "LOGFIRE_TOKEN" not in os.environ: + raise ValueError("LOGFIRE_TOKEN not found in environment variables") + from litellm.integrations.opentelemetry import OpenTelemetry + + for callback in _in_memory_loggers: + if isinstance(callback, OpenTelemetry): + return callback # type: ignore + + elif logging_integration == "dynamic_rate_limiter": + from litellm.proxy.hooks.dynamic_rate_limiter import ( + _PROXY_DynamicRateLimitHandler, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_DynamicRateLimitHandler): + return callback # type: ignore + elif logging_integration == "dynamic_rate_limiter_v3": + from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3): + return callback # type: ignore + + elif logging_integration == "langtrace": + from litellm.integrations.opentelemetry import OpenTelemetry + + if "LANGTRACE_API_KEY" not in os.environ: + raise ValueError("LANGTRACE_API_KEY not found in environment variables") + + for callback in _in_memory_loggers: + if ( + isinstance(callback, OpenTelemetry) + and callback.callback_name == "langtrace" + ): + return callback + + elif logging_integration == "mlflow": + for callback in _in_memory_loggers: + if isinstance(callback, MlflowLogger): + return callback + elif logging_integration == "pagerduty": + for callback in _in_memory_loggers: + if isinstance(callback, PagerDutyAlerting): + return callback + elif logging_integration == "anthropic_cache_control_hook": + for callback in _in_memory_loggers: + if isinstance(callback, AnthropicCacheControlHook): + return callback + elif logging_integration == "vector_store_pre_call_hook": + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( + VectorStorePreCallHook, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, VectorStorePreCallHook): + return callback + elif logging_integration == "gcs_pubsub": + for callback in _in_memory_loggers: + if isinstance(callback, GcsPubSubLogger): + return callback + elif logging_integration == "generic_api": + for callback in _in_memory_loggers: + if isinstance(callback, GenericAPILogger): + return callback + elif logging_integration == "resend_email": + for callback in _in_memory_loggers: + if isinstance(callback, ResendEmailLogger): + return callback + elif logging_integration == "sendgrid_email": + for callback in _in_memory_loggers: + if isinstance(callback, SendGridEmailLogger): + return callback + elif logging_integration == "smtp_email": + for callback in _in_memory_loggers: + if isinstance(callback, SMTPEmailLogger): + return callback + return None + + except Exception as e: + verbose_logger.exception( + f"[Non-Blocking Error] Error getting custom logger: {e}" + ) + return None + + +def _get_custom_logger_settings_from_proxy_server(callback_name: str) -> Dict: + """ + Get the settings for a custom logger from the proxy server config.yaml + + Proxy server config.yaml defines callback_settings as: + + callback_settings: + otel: + message_logging: False + """ + if litellm.callback_settings: + return dict(litellm.callback_settings.get(callback_name, {})) + return {} + + +def use_custom_pricing_for_model(litellm_params: Optional[dict]) -> bool: + """ + Check if the model uses custom pricing + + Returns True if any of `SPECIAL_MODEL_INFO_PARAMS` are present in `litellm_params` or `model_info` + """ + if litellm_params is None: + return False + + # Check litellm_params using set intersection (only check keys that exist in both) + matching_keys = _CUSTOM_PRICING_KEYS & litellm_params.keys() + for key in matching_keys: + if litellm_params.get(key) is not None: + return True + + # Check model_info + metadata: dict = litellm_params.get("metadata", {}) or {} + model_info: dict = metadata.get("model_info", {}) or {} + + if model_info: + matching_keys = _CUSTOM_PRICING_KEYS & model_info.keys() + for key in matching_keys: + if model_info.get(key) is not None: + return True + + return False + + +def is_valid_sha256_hash(value: str) -> bool: + # Check if the value is a valid SHA-256 hash (64 hexadecimal characters) + return bool(re.fullmatch(r"[a-fA-F0-9]{64}", value)) + + +class StandardLoggingPayloadSetup: + @staticmethod + def cleanup_timestamps( + start_time: Union[dt_object, float], + end_time: Union[dt_object, float], + completion_start_time: Union[dt_object, float], + ) -> Tuple[float, float, float]: + """ + Convert datetime objects to floats + + Args: + start_time: Union[dt_object, float] + end_time: Union[dt_object, float] + completion_start_time: Union[dt_object, float] + + Returns: + Tuple[float, float, float]: A tuple containing the start time, end time, and completion start time as floats. + """ + + if isinstance(start_time, datetime.datetime): + start_time_float = start_time.timestamp() + elif isinstance(start_time, float): + start_time_float = start_time + else: + raise ValueError( + f"start_time is required, got={start_time} of type {type(start_time)}" + ) + + if isinstance(end_time, datetime.datetime): + end_time_float = end_time.timestamp() + elif isinstance(end_time, float): + end_time_float = end_time + else: + raise ValueError( + f"end_time is required, got={end_time} of type {type(end_time)}" + ) + + if isinstance(completion_start_time, datetime.datetime): + completion_start_time_float = completion_start_time.timestamp() + elif isinstance(completion_start_time, float): + completion_start_time_float = completion_start_time + else: + completion_start_time_float = end_time_float + + return start_time_float, end_time_float, completion_start_time_float + + @staticmethod + def append_system_prompt_messages( + kwargs: Optional[Dict] = None, messages: Optional[Any] = None + ): + """ + Append system prompt messages to the messages + """ + if kwargs is not None: + if kwargs.get("system") is not None and isinstance( + kwargs.get("system"), str + ): + if messages is None: + return [{"role": "system", "content": kwargs.get("system")}] + elif isinstance(messages, list): + if len(messages) == 0: + return [{"role": "system", "content": kwargs.get("system")}] + # check for duplicates + if messages[0].get("role") == "system" and messages[0].get( + "content" + ) == kwargs.get("system"): + return messages + messages = [ + {"role": "system", "content": kwargs.get("system")} + ] + messages + elif isinstance(messages, str): + messages = [ + {"role": "system", "content": kwargs.get("system")}, + {"role": "user", "content": messages}, + ] + return messages + + return messages + + @staticmethod + def merge_litellm_metadata(litellm_params: dict) -> dict: + """ + Merge both litellm_metadata and metadata from litellm_params. + + litellm_metadata contains model-related fields, metadata contains user API key fields. + We need both for complete standard logging payload. + + Args: + litellm_params: Dictionary containing metadata and litellm_metadata + + Returns: + dict: Merged metadata with user API key fields taking precedence + """ + merged_metadata: dict = {} + + # Start with metadata (user API key fields) - but skip non-serializable objects + if litellm_params.get("metadata") and isinstance( + litellm_params.get("metadata"), dict + ): + for key, value in litellm_params["metadata"].items(): + # Skip non-serializable objects like UserAPIKeyAuth + if key == "user_api_key_auth": + continue + merged_metadata[key] = value + + # Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys + if litellm_params.get("litellm_metadata") and isinstance( + litellm_params.get("litellm_metadata"), dict + ): + for key, value in litellm_params["litellm_metadata"].items(): + if ( + key not in merged_metadata + ): # Don't overwrite existing keys from metadata + merged_metadata[key] = value + + return merged_metadata + + @staticmethod + def get_standard_logging_metadata( + metadata: Optional[Dict[str, Any]], + litellm_params: Optional[dict] = None, + prompt_integration: Optional[str] = None, + applied_guardrails: Optional[List[str]] = None, + mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = None, + vector_store_request_metadata: Optional[ + List[StandardLoggingVectorStoreRequest] + ] = None, + usage_object: Optional[dict] = None, + proxy_server_request: Optional[dict] = None, + start_time: Optional[dt_object] = None, + response_id: Optional[str] = None, + ) -> StandardLoggingMetadata: + """ + Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata. + + Args: + metadata (Optional[Dict[str, Any]]): The original metadata dictionary. + + Returns: + StandardLoggingMetadata: A StandardLoggingMetadata object containing the cleaned metadata. + + Note: + - If the input metadata is None or not a dictionary, an empty StandardLoggingMetadata object is returned. + - If 'user_api_key' is present in metadata and is a valid SHA256 hash, it's stored as 'user_api_key_hash'. + """ + + prompt_management_metadata: Optional[ + StandardLoggingPromptManagementMetadata + ] = None + if litellm_params is not None: + prompt_id = cast(Optional[str], litellm_params.get("prompt_id", None)) + prompt_variables = cast( + Optional[dict], litellm_params.get("prompt_variables", None) + ) + + if prompt_id is not None and prompt_integration is not None: + prompt_management_metadata = StandardLoggingPromptManagementMetadata( + prompt_id=prompt_id, + prompt_variables=prompt_variables, + prompt_integration=prompt_integration, + ) + + # Initialize with default values + clean_metadata = StandardLoggingMetadata( + user_api_key_hash=None, + user_api_key_alias=None, + user_api_key_spend=None, + user_api_key_max_budget=None, + user_api_key_budget_reset_at=None, + user_api_key_team_id=None, + user_api_key_org_id=None, + user_api_key_project_id=None, + user_api_key_user_id=None, + user_api_key_team_alias=None, + user_api_key_user_email=None, + user_api_key_end_user_id=None, + user_api_key_request_route=None, + spend_logs_metadata=None, + requester_ip_address=None, + user_agent=None, + requester_metadata=None, + prompt_management_metadata=prompt_management_metadata, + applied_guardrails=applied_guardrails, + mcp_tool_call_metadata=mcp_tool_call_metadata, + vector_store_request_metadata=vector_store_request_metadata, + usage_object=usage_object, + requester_custom_headers=None, + cold_storage_object_key=None, + user_api_key_auth_metadata=None, + team_alias=None, + team_id=None, + ) + if isinstance(metadata, dict): + for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS: + clean_metadata[key] = metadata[key] # type: ignore + + user_api_key = metadata.get("user_api_key") + if ( + user_api_key + and isinstance(user_api_key, str) + and is_valid_sha256_hash(user_api_key) + ): + clean_metadata["user_api_key_hash"] = user_api_key + _potential_requester_metadata = metadata.get( + "metadata", None + ) # check if user passed metadata in the sdk request - e.g. metadata for langsmith logging - https://docs.litellm.ai/docs/observability/langsmith_integration#set-langsmith-fields + if ( + clean_metadata["requester_metadata"] is None + and _potential_requester_metadata is not None + and isinstance(_potential_requester_metadata, dict) + ): + clean_metadata["requester_metadata"] = _potential_requester_metadata + + if ( + EnterpriseStandardLoggingPayloadSetupVAR + and proxy_server_request is not None + ): + clean_metadata = EnterpriseStandardLoggingPayloadSetupVAR.apply_enterprise_specific_metadata( + standard_logging_metadata=clean_metadata, + proxy_server_request=proxy_server_request, + ) + + # Generate cold storage object key if cold storage is configured + if start_time is not None and response_id is not None: + cold_storage_object_key = ( + StandardLoggingPayloadSetup._generate_cold_storage_object_key( + start_time=start_time, + response_id=response_id, + team_alias=clean_metadata.get("user_api_key_team_alias"), + ) + ) + if cold_storage_object_key: + clean_metadata["cold_storage_object_key"] = cold_storage_object_key + + return clean_metadata + + @staticmethod + def get_usage_from_response_obj( + response_obj: Optional[dict], combined_usage_object: Optional[Usage] = None + ) -> Usage: + ## BASE CASE ## + if combined_usage_object is not None: + return combined_usage_object + if response_obj is None: + return Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + ) + + usage = response_obj.get("usage", None) or {} + if usage is None or ( + not isinstance(usage, dict) and not isinstance(usage, Usage) + ): + return Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + ) + elif isinstance(usage, Usage): + return usage + elif isinstance(usage, ResponseAPIUsage): + return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + elif isinstance(usage, dict): + if ResponseAPILoggingUtils._is_response_api_usage(usage): + return ( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + ) + return Usage(**usage) + + raise ValueError(f"usage is required, got={usage} of type {type(usage)}") + + @staticmethod + def get_model_cost_information( + base_model: Optional[str], + custom_pricing: Optional[bool], + custom_llm_provider: Optional[str], + init_response_obj: Union[Any, BaseModel, dict], + ) -> StandardLoggingModelInformation: + model_cost_name = _select_model_name_for_cost_calc( + model=None, + completion_response=init_response_obj, # type: ignore + base_model=base_model, + custom_pricing=custom_pricing, + ) + if model_cost_name is None: + model_cost_information = StandardLoggingModelInformation( + model_map_key="", model_map_value=None + ) + else: + try: + _model_cost_information = litellm.get_model_info( + model=model_cost_name, custom_llm_provider=custom_llm_provider + ) + model_cost_information = StandardLoggingModelInformation( + model_map_key=model_cost_name, + model_map_value=_model_cost_information, + ) + except Exception: + verbose_logger.debug( # keep in debug otherwise it will trigger on every call + "Model={} is not mapped in model cost map. Defaulting to None model_cost_information for standard_logging_payload".format( + model_cost_name + ) + ) + model_cost_information = StandardLoggingModelInformation( + model_map_key=model_cost_name, model_map_value=None + ) + return model_cost_information + + @staticmethod + def get_final_response_obj( + response_obj: dict, init_response_obj: Union[Any, BaseModel, dict], kwargs: dict + ) -> Optional[Union[dict, str, list]]: + """ + Get final response object after redacting the message input/output from logging + """ + if response_obj: + final_response_obj: Optional[Union[dict, str, list]] = response_obj + elif isinstance(init_response_obj, list) or isinstance(init_response_obj, str): + final_response_obj = init_response_obj + else: + final_response_obj = {} + + modified_final_response_obj = redact_message_input_output_from_logging( + model_call_details=kwargs, + result=final_response_obj, + ) + + if modified_final_response_obj is not None and isinstance( + modified_final_response_obj, BaseModel + ): + final_response_obj = modified_final_response_obj.model_dump() + else: + final_response_obj = modified_final_response_obj + + return final_response_obj + + @staticmethod + def get_additional_headers( + additiona_headers: Optional[dict], + ) -> Optional[StandardLoggingAdditionalHeaders]: + if additiona_headers is None: + return None + + additional_logging_headers: StandardLoggingAdditionalHeaders = {} + + for key in StandardLoggingAdditionalHeaders.__annotations__.keys(): + _key = key.lower() + _key = _key.replace("_", "-") + if _key in additiona_headers: + try: + additional_logging_headers[key] = int(additiona_headers[_key]) # type: ignore + except (ValueError, TypeError): + verbose_logger.debug( + f"Could not convert {additiona_headers[_key]} to int for key {key}." + ) + return additional_logging_headers + + @staticmethod + def get_hidden_params( + hidden_params: Optional[dict], + ) -> StandardLoggingHiddenParams: + clean_hidden_params = StandardLoggingHiddenParams( + model_id=None, + cache_key=None, + api_base=None, + response_cost=None, + additional_headers=None, + litellm_overhead_time_ms=None, + batch_models=None, + litellm_model_name=None, + usage_object=None, + ) + if hidden_params is not None: + for key in StandardLoggingHiddenParams.__annotations__.keys(): + if key in hidden_params: + if key == "additional_headers": + clean_hidden_params[ + "additional_headers" + ] = StandardLoggingPayloadSetup.get_additional_headers( + hidden_params[key] + ) + else: + clean_hidden_params[key] = hidden_params[key] # type: ignore + return clean_hidden_params + + @staticmethod + def strip_trailing_slash(api_base: Optional[str]) -> Optional[str]: + if api_base: + if api_base.endswith("//"): + return api_base.rstrip("/") + if api_base[-1] == "/": + return api_base[:-1] + return api_base + + @staticmethod + def _generate_cold_storage_object_key( + start_time: dt_object, + response_id: str, + team_alias: Optional[str] = None, + ) -> Optional[str]: + """ + Generate cold storage object key in the same format as S3Logger. + + Args: + start_time: The start time of the request + response_id: The response ID + team_alias: Optional team alias for team-based prefixing + + Returns: + Optional[str]: The generated object key or None if cold storage not configured + """ + # Generate object key in same format as S3Logger + from litellm.integrations.s3 import get_s3_object_key + + # Only generate object key if cold storage is configured + cold_storage_custom_logger = litellm.cold_storage_custom_logger + if cold_storage_custom_logger is None: + return None + + try: + # Generate file name in same format as litellm.utils.get_logging_id + s3_file_name = f"time-{start_time.strftime('%H-%M-%S-%f')}_{response_id}" + + # Get the actual s3_path from the configured cold storage logger instance + s3_path = "" # default value + + # Try to get the actual logger instance from the logger name + try: + custom_logger = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name( + cold_storage_custom_logger + ) + if ( + custom_logger + and hasattr(custom_logger, "s3_path") + and getattr(custom_logger, "s3_path") + ): + s3_path = getattr(custom_logger, "s3_path") + except Exception: + # If any error occurs in getting the logger instance, use default empty s3_path + pass + + s3_object_key = get_s3_object_key( + s3_path=s3_path, # Use actual s3_path from logger configuration + prefix="", # Don't split by team alias for cold storage + start_time=start_time, + s3_file_name=s3_file_name, + ) + + return s3_object_key + except Exception: + # If any error occurs in generating the key, return None + return None + + @staticmethod + def get_error_information( + original_exception: Optional[Exception], + traceback_str: Optional[str] = None, + ) -> StandardLoggingPayloadErrorInformation: + from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG + + # Check for 'code' first (used by ProxyException), then fall back to 'status_code' (used by LiteLLM exceptions) + # Ensure error_code is always a string for Prisma Python JSON field compatibility + error_code_attr = getattr(original_exception, "code", None) + if error_code_attr is not None and str(error_code_attr) not in ("", "None"): + error_status: str = str(error_code_attr) + else: + status_code_attr = getattr(original_exception, "status_code", None) + error_status = str(status_code_attr) if status_code_attr is not None else "" + error_class: str = ( + str(original_exception.__class__.__name__) if original_exception else "" + ) + _llm_provider_in_exception = getattr(original_exception, "llm_provider", "") + + # Get traceback information (first 100 lines) + traceback_info = traceback_str or "" + if original_exception: + tb = getattr(original_exception, "__traceback__", None) + if tb: + tb_lines = traceback.format_tb(tb) + traceback_info += "".join( + tb_lines[:MAXIMUM_TRACEBACK_LINES_TO_LOG] + ) # Limit to first 100 lines + + # Get additional error details + error_message = str(original_exception) + + return StandardLoggingPayloadErrorInformation( + error_code=error_status, + error_class=error_class, + llm_provider=_llm_provider_in_exception, + traceback=traceback_info, + error_message=error_message if original_exception else "", + ) + + @staticmethod + def get_response_time( + start_time_float: float, + end_time_float: float, + completion_start_time_float: float, + stream: bool, + ) -> float: + """ + Get the response time for the LLM response + + Args: + start_time_float: float - start time of the LLM call + end_time_float: float - end time of the LLM call + completion_start_time_float: float - time to first token of the LLM response (for streaming responses) + stream: bool - True when a stream response is returned + + Returns: + float: The response time for the LLM response + """ + if stream is True: + return completion_start_time_float - start_time_float + else: + return end_time_float - start_time_float + + @staticmethod + def _get_standard_logging_payload_trace_id( + logging_obj: Logging, + litellm_params: dict, + ) -> str: + """ + Returns the `litellm_trace_id` for this request + + This helps link sessions when multiple requests are made in a single session + """ + dynamic_litellm_session_id = litellm_params.get("litellm_session_id") + dynamic_litellm_trace_id = litellm_params.get("litellm_trace_id") + + # Note: we recommend using `litellm_session_id` for session tracking + # `litellm_trace_id` is an internal litellm param + if dynamic_litellm_session_id: + return str(dynamic_litellm_session_id) + elif dynamic_litellm_trace_id: + return str(dynamic_litellm_trace_id) + else: + return logging_obj.litellm_trace_id + + @staticmethod + def _get_user_agent_tags(proxy_server_request: dict) -> Optional[List[str]]: + """ + Return the user agent tags from the proxy server request for spend tracking + """ + if litellm.disable_add_user_agent_to_request_tags is True: + return None + user_agent_tags: Optional[List[str]] = None + headers = proxy_server_request.get("headers", {}) + if headers is not None and isinstance(headers, dict): + if "user-agent" in headers: + user_agent = headers["user-agent"] + if user_agent is not None: + if user_agent_tags is None: + user_agent_tags = [] + user_agent_part: Optional[str] = None + if "/" in user_agent: + user_agent_part = user_agent.split("/")[0] + if user_agent_part is not None: + user_agent_tags.append("User-Agent: " + user_agent_part) + if user_agent is not None: + user_agent_tags.append("User-Agent: " + user_agent) + return user_agent_tags + + @staticmethod + def _get_extra_header_tags(proxy_server_request: dict) -> Optional[List[str]]: + """ + Extract additional header tags for spend tracking based on config. + """ + extra_headers: List[str] = ( + getattr(litellm, "extra_spend_tag_headers", None) or [] + ) + if not extra_headers: + return None + + headers = proxy_server_request.get("headers", {}) + if not isinstance(headers, dict): + return None + + header_tags = [] + for header_name in extra_headers: + header_value = headers.get(header_name) + if header_value: + header_tags.append(f"{header_name}: {header_value}") + + return header_tags if header_tags else None + + @staticmethod + def _get_request_tags( + litellm_params: dict, proxy_server_request: dict + ) -> List[str]: + # check for 'tags' in both 'metadata' and 'litellm_metadata' + metadata = litellm_params.get("metadata") or {} + litellm_metadata = litellm_params.get("litellm_metadata") or {} + if metadata.get("tags", []): + request_tags = metadata.get("tags", []).copy() + elif litellm_metadata.get("tags", []): + request_tags = litellm_metadata.get("tags", []).copy() + else: + request_tags = [] + user_agent_tags = StandardLoggingPayloadSetup._get_user_agent_tags( + proxy_server_request + ) + additional_header_tags = StandardLoggingPayloadSetup._get_extra_header_tags( + proxy_server_request + ) + if user_agent_tags is not None: + request_tags.extend(user_agent_tags) + if additional_header_tags is not None: + request_tags.extend(additional_header_tags) + return request_tags + + +def _get_status_fields( + status: StandardLoggingPayloadStatus, + guardrail_information: Optional[List[dict]], + error_str: Optional[str], +) -> "StandardLoggingPayloadStatusFields": + """ + Determine status fields based on request status and guardrail information. + + Args: + status: Overall request status ("success" or "failure") + guardrail_information: Guardrail information from metadata + error_str: Error string if any + + Returns: + StandardLoggingPayloadStatusFields with llm_api_status and guardrail_status + """ + # Mapping for legacy guardrail status values to new GuardrailStatus values + GUARDRAIL_STATUS_MAP: Dict[str, GuardrailStatus] = { + "success": "success", + "blocked": "guardrail_intervened", # legacy + "guardrail_intervened": "guardrail_intervened", # direct + "failure": "guardrail_failed_to_respond", # legacy + "guardrail_failed_to_respond": "guardrail_failed_to_respond", # direct + "not_run": "not_run", + } + + # Set LLM API status + llm_api_status: StandardLoggingPayloadStatus = status + + ######################################################### + # Map - guardrail_information.guardrail_status to guardrail_status + ######################################################### + guardrail_status: GuardrailStatus = "not_run" + if guardrail_information and isinstance(guardrail_information, list): + for information in guardrail_information: + if isinstance(information, dict): + raw_status = information.get("guardrail_status", "not_run") + if raw_status != "not_run": + guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") + break + + return StandardLoggingPayloadStatusFields( + llm_api_status=llm_api_status, guardrail_status=guardrail_status + ) + + +def _extract_response_obj_and_hidden_params( + init_response_obj: Union[Any, BaseModel, dict], + original_exception: Optional[Exception], +) -> Tuple[dict, Optional[dict]]: + """Extract response_obj and hidden_params from init_response_obj.""" + hidden_params: Optional[dict] = None + if init_response_obj is None: + response_obj = {} + elif isinstance(init_response_obj, BaseModel): + response_obj = init_response_obj.model_dump() + hidden_params = getattr(init_response_obj, "_hidden_params", None) + elif isinstance(init_response_obj, dict): + response_obj = init_response_obj + else: + response_obj = {} + + if original_exception is not None and hidden_params is None: + response_headers = _get_response_headers(original_exception) + if response_headers is not None: + hidden_params = dict( + StandardLoggingHiddenParams( + additional_headers=StandardLoggingPayloadSetup.get_additional_headers( + dict(response_headers) + ), + model_id=None, + cache_key=None, + api_base=None, + response_cost=None, + litellm_overhead_time_ms=None, + batch_models=None, + litellm_model_name=None, + usage_object=None, + ) + ) + + return response_obj, hidden_params + + +def get_standard_logging_object_payload( + kwargs: Optional[dict], + init_response_obj: Union[Any, BaseModel, dict], + start_time: dt_object, + end_time: dt_object, + logging_obj: Logging, + status: StandardLoggingPayloadStatus, + error_str: Optional[str] = None, + original_exception: Optional[Exception] = None, + standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None, +) -> Optional[StandardLoggingPayload]: + try: + kwargs = kwargs or {} + + response_obj, hidden_params = _extract_response_obj_and_hidden_params( + init_response_obj, original_exception + ) + + # standardize this function to be used across, s3, dynamoDB, langfuse logging + litellm_params = kwargs.get("litellm_params", {}) or {} + proxy_server_request = litellm_params.get("proxy_server_request") or {} + + # Merge both litellm_metadata and metadata to get complete metadata + metadata: dict = StandardLoggingPayloadSetup.merge_litellm_metadata( + litellm_params + ) + + completion_start_time = kwargs.get("completion_start_time", end_time) + call_type = kwargs.get("call_type") + cache_hit = kwargs.get("cache_hit", False) + usage = StandardLoggingPayloadSetup.get_usage_from_response_obj( + response_obj=response_obj, + combined_usage_object=cast( + Optional[Usage], kwargs.get("combined_usage_object") + ), + ) + + id = response_obj.get("id", kwargs.get("litellm_call_id")) + + _model_id = metadata.get("model_info", {}).get("id", "") + _model_group = metadata.get("model_group", "") + + request_tags = StandardLoggingPayloadSetup._get_request_tags( + litellm_params=litellm_params, proxy_server_request=proxy_server_request + ) + + # cleanup timestamps + ( + start_time_float, + end_time_float, + completion_start_time_float, + ) = StandardLoggingPayloadSetup.cleanup_timestamps( + start_time=start_time, + end_time=end_time, + completion_start_time=completion_start_time, + ) + response_time = StandardLoggingPayloadSetup.get_response_time( + start_time_float=start_time_float, + end_time_float=end_time_float, + completion_start_time_float=completion_start_time_float, + stream=kwargs.get("stream", False), + ) + # clean up litellm hidden params + clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params( + hidden_params + ) + + # clean up litellm metadata + clean_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( + metadata=metadata, + litellm_params=litellm_params, + prompt_integration=kwargs.get("prompt_integration", None), + applied_guardrails=kwargs.get("applied_guardrails", None), + mcp_tool_call_metadata=kwargs.get("mcp_tool_call_metadata", None), + vector_store_request_metadata=kwargs.get( + "vector_store_request_metadata", None + ), + usage_object=usage.model_dump(), + proxy_server_request=proxy_server_request, + start_time=start_time, + response_id=id, + ) + _request_body = proxy_server_request.get("body", {}) + end_user_id = clean_metadata["user_api_key_end_user_id"] or _request_body.get( + "user", None + ) # maintain backwards compatibility with old request body check + + saved_cache_cost: float = 0.0 + if cache_hit is True: + id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id + saved_cache_cost = ( + logging_obj._response_cost_calculator( + result=init_response_obj, cache_hit=False # type: ignore + ) + or 0.0 + ) + + ## Get model cost information ## + base_model = _get_base_model_from_metadata(model_call_details=kwargs) + custom_pricing = use_custom_pricing_for_model(litellm_params=litellm_params) + + model_cost_information = StandardLoggingPayloadSetup.get_model_cost_information( + base_model=base_model, + custom_pricing=custom_pricing, + custom_llm_provider=kwargs.get("custom_llm_provider"), + init_response_obj=init_response_obj, + ) + response_cost: float = kwargs.get("response_cost", 0) or 0.0 + + error_information = StandardLoggingPayloadSetup.get_error_information( + original_exception=original_exception, + ) + + ## get final response object ## + final_response_obj = StandardLoggingPayloadSetup.get_final_response_obj( + response_obj=response_obj, + init_response_obj=init_response_obj, + kwargs=kwargs, + ) + + stream: Optional[bool] = None + if ( + kwargs.get("complete_streaming_response") is not None + or kwargs.get("async_complete_streaming_response") is not None + ) and kwargs.get("stream") is True: + stream = True + + # Reconstruct full model name with provider prefix for logging + # This ensures Bedrock models like "us.anthropic.claude-3-5-sonnet-20240620-v1:0" + # are logged as "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0" + custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) + model_name = reconstruct_model_name( + kwargs.get("model", "") or "", custom_llm_provider, metadata + ) + + payload: StandardLoggingPayload = StandardLoggingPayload( + id=str(id), + trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + logging_obj=logging_obj, + litellm_params=litellm_params, + ), + call_type=call_type or "", + cache_hit=cache_hit, + stream=stream, + status=status, + status_fields=_get_status_fields( + status=status, + guardrail_information=metadata.get( + "standard_logging_guardrail_information", None + ), + error_str=error_str, + ), + custom_llm_provider=custom_llm_provider, + saved_cache_cost=saved_cache_cost, + startTime=start_time_float, + endTime=end_time_float, + completionStartTime=completion_start_time_float, + response_time=response_time, + model=model_name, + metadata=clean_metadata, + cache_key=clean_hidden_params["cache_key"], + response_cost=response_cost, + cost_breakdown=logging_obj.cost_breakdown, + total_tokens=usage.total_tokens, + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + request_tags=request_tags, + end_user=end_user_id or "", + api_base=StandardLoggingPayloadSetup.strip_trailing_slash( + litellm_params.get("api_base", "") + ) + or "", + model_group=_model_group, + model_id=_model_id, + requester_ip_address=clean_metadata.get("requester_ip_address", None), + user_agent=clean_metadata.get("user_agent", None), + messages=StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=kwargs, messages=kwargs.get("messages") + ), + response=final_response_obj, + model_parameters=ModelParamHelper.get_standard_logging_model_parameters( + kwargs.get("optional_params", None) or {} + ), + hidden_params=clean_hidden_params, + model_map_information=model_cost_information, + error_str=error_str, + error_information=error_information, + response_cost_failure_debug_info=kwargs.get( + "response_cost_failure_debug_information" + ), + guardrail_information=metadata.get( + "standard_logging_guardrail_information", None + ), + standard_built_in_tools_params=standard_built_in_tools_params, + ) + + # emit_standard_logging_payload(payload) - Moved to success_handler to prevent double emitting + + return payload + except Exception as e: + verbose_logger.exception( + "Error creating standard logging object - {}".format(str(e)) + ) + return None + + +def emit_standard_logging_payload(payload: StandardLoggingPayload): + if os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD"): + print(json.dumps(payload, indent=4)) # noqa + + +def get_standard_logging_metadata( + metadata: Optional[Dict[str, Any]], +) -> StandardLoggingMetadata: + """ + Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata. + + Args: + metadata (Optional[Dict[str, Any]]): The original metadata dictionary. + + Returns: + StandardLoggingMetadata: A StandardLoggingMetadata object containing the cleaned metadata. + + Note: + - If the input metadata is None or not a dictionary, an empty StandardLoggingMetadata object is returned. + - If 'user_api_key' is present in metadata and is a valid SHA256 hash, it's stored as 'user_api_key_hash'. + """ + # Initialize with default values + clean_metadata = StandardLoggingMetadata( + user_api_key_hash=None, + user_api_key_alias=None, + user_api_key_spend=None, + user_api_key_max_budget=None, + user_api_key_budget_reset_at=None, + user_api_key_team_id=None, + user_api_key_org_id=None, + user_api_key_project_id=None, + user_api_key_user_id=None, + user_api_key_user_email=None, + user_api_key_team_alias=None, + spend_logs_metadata=None, + requester_ip_address=None, + user_agent=None, + requester_metadata=None, + user_api_key_end_user_id=None, + prompt_management_metadata=None, + applied_guardrails=None, + mcp_tool_call_metadata=None, + vector_store_request_metadata=None, + usage_object=None, + requester_custom_headers=None, + user_api_key_request_route=None, + cold_storage_object_key=None, + user_api_key_auth_metadata=None, + team_alias=None, + team_id=None, + ) + if isinstance(metadata, dict): + # Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields + for key in StandardLoggingMetadata.__annotations__.keys(): + if key in metadata: + clean_metadata[key] = metadata[key] # type: ignore + + if metadata.get("user_api_key") is not None: + if is_valid_sha256_hash(str(metadata.get("user_api_key"))): + clean_metadata["user_api_key_hash"] = metadata.get( + "user_api_key" + ) # this is the hash + return clean_metadata + + +def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): + if litellm_params is None: + litellm_params = {} + + metadata = litellm_params.get("metadata", {}) or {} + + ## Extract provider-specific callable values (like langfuse_masking_function) + ## Store them separately so only the intended logger can access them + ## This prevents callables from leaking to other logging integrations + if "langfuse_masking_function" in metadata: + masking_fn = metadata.pop("langfuse_masking_function", None) + if callable(masking_fn): + litellm_params["_langfuse_masking_function"] = masking_fn + litellm_params["metadata"] = metadata + + ## check user_api_key_metadata for sensitive logging keys + cleaned_user_api_key_metadata = {} + if "user_api_key_metadata" in metadata and isinstance( + metadata["user_api_key_metadata"], dict + ): + for k, v in metadata["user_api_key_metadata"].items(): + if k == "logging": # prevent logging user logging keys + cleaned_user_api_key_metadata[ + k + ] = "scrubbed_by_litellm_for_sensitive_keys" + else: + cleaned_user_api_key_metadata[k] = v + + metadata["user_api_key_metadata"] = cleaned_user_api_key_metadata + litellm_params["metadata"] = metadata + + return litellm_params + + +# integration helper function +def modify_integration(integration_name, integration_params): + global supabaseClient + if integration_name == "supabase": + if "table_name" in integration_params: + Supabase.supabase_table_name = integration_params["table_name"] + + +@lru_cache(maxsize=16) +def _get_traceback_str_for_error(error_str: str) -> str: + """ + function wrapped with lru_cache to limit the number of times `traceback.format_exc()` is called + """ + return traceback.format_exc() + + +from decimal import Decimal + +# used for unit testing +from typing import Any, Dict, List, Optional, Union + + +def create_dummy_standard_logging_payload() -> StandardLoggingPayload: + # First create the nested objects with proper typing + model_info = StandardLoggingModelInformation( + model_map_key="gpt-3.5-turbo", model_map_value=None + ) + + metadata = StandardLoggingMetadata( # type: ignore + user_api_key_hash=str("test_hash"), + user_api_key_alias=str("test_alias"), + user_api_key_team_id=str("test_team"), + user_api_key_user_id=str("test_user"), + user_api_key_team_alias=str("test_team_alias"), + user_api_key_org_id=None, + spend_logs_metadata=None, + requester_ip_address=str("127.0.0.1"), + requester_metadata=None, + user_api_key_end_user_id=str("test_end_user"), + ) + + hidden_params = StandardLoggingHiddenParams( + model_id=None, + cache_key=None, + api_base=None, + response_cost=None, + additional_headers=None, + litellm_overhead_time_ms=None, + batch_models=None, + litellm_model_name=None, + usage_object=None, + ) + + # Convert numeric values to appropriate types + response_cost = Decimal("0.1") + start_time = Decimal("1234567890.0") + end_time = Decimal("1234567891.0") + completion_start_time = Decimal("1234567890.5") + saved_cache_cost = Decimal("0.0") + + # Create messages and response with proper typing + messages: List[Dict[str, str]] = [{"role": "user", "content": "Hello, world!"}] + response: Dict[str, List[Dict[str, Dict[str, str]]]] = { + "choices": [{"message": {"content": "Hi there!"}}] + } + + # Main payload initialization + return StandardLoggingPayload( # type: ignore + id=str("test_id"), + call_type=str("completion"), + stream=bool(False), + response_cost=response_cost, + response_cost_failure_debug_info=None, + status=str("success"), + total_tokens=int( + DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT + ), + prompt_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT), + completion_tokens=int(DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT), + startTime=start_time, + endTime=end_time, + completionStartTime=completion_start_time, + model_map_information=model_info, + model=str("gpt-3.5-turbo"), + model_id=str("model-123"), + model_group=str("openai-gpt"), + custom_llm_provider=str("openai"), + api_base=str("https://api.openai.com"), + metadata=metadata, + cache_hit=bool(False), + cache_key=None, + saved_cache_cost=saved_cache_cost, + request_tags=[], + end_user=None, + requester_ip_address=str("127.0.0.1"), + messages=messages, + response=response, + error_str=None, + model_parameters={"stream": True}, + hidden_params=hidden_params, + ) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index a5f8fe22a2c..364126d822e 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -191,6 +191,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "user", "web_search_options", "speed", + "context_management", ] if "claude-3-7-sonnet" in model or supports_reasoning( @@ -825,6 +826,62 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return hosted_web_search_tool + @staticmethod + def map_openai_context_management_to_anthropic( + context_management: Union[List[Dict[str, Any]], Dict[str, Any]] + ) -> Optional[Dict[str, Any]]: + """ + OpenAI format: [{"type": "compaction", "compact_threshold": 200000}] + Anthropic format: { + "edits": [ + { + "type": "compact_20260112", + "trigger": {"type": "input_tokens", "value": 150000} + } + ] + } + + Args: + context_management: OpenAI or Anthropic context_management parameter + + Returns: + Anthropic-formatted context_management dict, or None if invalid + """ + # If already in Anthropic format (dict with 'edits'), pass through + if isinstance(context_management, dict) and "edits" in context_management: + return context_management + + # If in OpenAI format (list), transform to Anthropic format + if isinstance(context_management, list): + anthropic_edits = [] + for entry in context_management: + if not isinstance(entry, dict): + continue + + entry_type = entry.get("type") + if entry_type == "compaction": + anthropic_edit: Dict[str, Any] = { + "type": "compact_20260112" + } + compact_threshold = entry.get("compact_threshold") + # Rewrite to 'trigger' with correct nesting if threshold exists + if compact_threshold is not None and isinstance(compact_threshold, (int, float)): + anthropic_edit["trigger"] = { + "type": "input_tokens", + "value": int(compact_threshold) + } + # Map any other keys by passthrough except handled ones + for k in entry: + if k not in {"type", "compact_threshold"}: # only passthrough other keys + anthropic_edit[k] = entry[k] + + anthropic_edits.append(anthropic_edit) + + if anthropic_edits: + return {"edits": anthropic_edits} + + return None + def map_openai_params( # noqa: PLR0915 self, non_default_params: dict, @@ -931,9 +988,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) elif param == "extra_headers": optional_params["extra_headers"] = value - elif param == "context_management" and isinstance(value, dict): - # Pass through Anthropic-specific context_management parameter - optional_params["context_management"] = value + elif param == "context_management": + # Supports both OpenAI list format and Anthropic dict format + if isinstance(value, (list, dict)): + anthropic_context_management = self.map_openai_context_management_to_anthropic(value) + if anthropic_context_management is not None: + optional_params["context_management"] = anthropic_context_management elif param == "speed" and isinstance(value, str): # Pass through Anthropic-specific speed parameter for fast mode optional_params["speed"] = value @@ -1094,32 +1154,39 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): headers["anthropic-beta"] = f"{existing_beta}, {beta_value}" def _ensure_context_management_beta_header( - self, headers: dict, context_management: dict + self, headers: dict, context_management: object ) -> None: """ Add appropriate beta headers based on context_management edits. - - If any edit has type "compact_20260112", add compact-2026-01-12 header - - For all other edits, add context-management-2025-06-27 header """ - edits = context_management.get("edits", []) - + edits = [] + # If anthropic format (dict with "edits" key) + if isinstance(context_management, dict) and "edits" in context_management: + edits = context_management.get("edits", []) + # If OpenAI format: list of context management entries + elif isinstance(context_management, list): + edits = context_management + # Defensive: ignore/fallback if context_management not valid + else: + return + has_compact = False has_other = False - + for edit in edits: edit_type = edit.get("type", "") - if edit_type == "compact_20260112": + if edit_type == "compact_20260112" or edit_type == "compaction": has_compact = True else: has_other = True - - # Add compact header if any compact edits exist + + # Add compact header if any compact edits/entries exist if has_compact: self._ensure_beta_header( headers, ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value ) - - # Add context management header if any other edits exist + + # Add context management header if any other edits/entries exist if has_other: self._ensure_beta_header( headers, ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 8275ba2b3e1..e8d7a0383fb 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -164,6 +164,17 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # Remove system parameter if all content was filtered out anthropic_messages_optional_request_params.pop("system", None) + # Transform context_management from OpenAI format to Anthropic format if needed + context_management_param = anthropic_messages_optional_request_params.get("context_management") + if context_management_param is not None: + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + transformed_context_management = AnthropicConfig.map_openai_context_management_to_anthropic( + context_management_param + ) + if transformed_context_management is not None: + anthropic_messages_optional_request_params["context_management"] = transformed_context_management + ####### get required params for all anthropic messages requests ###### verbose_logger.debug(f"TRANSFORMATION DEBUG - Messages: {messages}") anthropic_messages_request: AnthropicMessagesRequest = AnthropicMessagesRequest( diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 18dad503a59..69eda95be1b 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -106,6 +106,7 @@ class AzureOpenAIConfig(BaseConfig): "audio", "web_search_options", "prompt_cache_key", + "store", ] def _is_response_format_supported_model(self, model: str) -> bool: @@ -158,7 +159,6 @@ class AzureOpenAIConfig(BaseConfig): api_version: str = "", ) -> dict: supported_openai_params = self.get_supported_openai_params(model) - api_version_times = api_version.split("-") if len(api_version_times) >= 3: @@ -245,7 +245,6 @@ class AzureOpenAIConfig(BaseConfig): optional_params["tools"].extend(value) elif param in supported_openai_params: optional_params[param] = value - return optional_params def transform_request( diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 155d8c9ec27..cc5cf991826 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -4,9 +4,6 @@ Translates from OpenAI's `/v1/chat/completions` to DashScope's `/v1/chat/complet from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, overload -from litellm.litellm_core_utils.prompt_templates.common_utils import ( - handle_messages_with_content_list_to_str_conversion, -) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -32,10 +29,6 @@ class DashScopeChatConfig(OpenAIGPTConfig): def _transform_messages( self, messages: List[AllMessageValues], model: str, is_async: bool = False ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: - """ - DashScope does not support content in list format. - """ - messages = handle_messages_with_content_list_to_str_conversion(messages) if is_async: return super()._transform_messages( messages=messages, model=model, is_async=True diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index e955800b947..35dfa8a3851 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -137,10 +137,29 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): self, messages: List[AllMessageValues], model: str, is_async: bool = False ) -> Union[List[AllMessageValues], Coroutine[Any, Any, List[AllMessageValues]]]: """ - Support translating video files from file_id or file_data to video_url + Support translating: + - video files from file_id or file_data to video_url + - thinking_blocks on assistant messages to content blocks """ for message in messages: - if message["role"] == "user": + if message["role"] == "assistant": + thinking_blocks = message.pop("thinking_blocks", None) # type: ignore + if thinking_blocks: + new_content: list = [ + {"type": block["type"], "thinking": block.get("thinking", "")} + if block.get("type") == "thinking" + else {"type": block["type"], "data": block.get("data", "")} + for block in thinking_blocks + ] + existing_content = message.get("content") + if isinstance(existing_content, str): + new_content.append( + {"type": "text", "text": existing_content} + ) + elif isinstance(existing_content, list): + new_content.extend(existing_content) + message["content"] = new_content # type: ignore + elif message["role"] == "user": message_content = message.get("content") if message_content and isinstance(message_content, list): replaced_content_items: List[ diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 59f52e2b81c..aa8471a5973 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -162,6 +162,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "service_tier", "safety_identifier", "prompt_cache_key", + "store", ] # works across all models model_specific_params = [] diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index daa82a46bdc..cf3461a9960 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -767,14 +767,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif reasoning_effort == "low": return {"thinkingLevel": "low", "includeThoughts": True} elif reasoning_effort == "medium": - # For gemini-3-flash-preview, medium maps to "medium", otherwise "high" - if is_gemini3flash: return {"thinkingLevel": "medium", "includeThoughts": True} - else: - return { - "thinkingLevel": "high", - "includeThoughts": True, - } # medium is not out yet for other models elif reasoning_effort == "high": return {"thinkingLevel": "high", "includeThoughts": True} elif reasoning_effort == "disable": diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 21782fc6fbf..aa2dee354cf 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple +from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple, Union import httpx @@ -11,9 +11,18 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import Choices, ModelResponse, Usage, PromptTokensDetailsWrapper +from litellm.types.utils import ( + Choices, + ModelResponse, + ModelResponseStream, + PromptTokensDetailsWrapper, + Usage, +) -from ...openai.chat.gpt_transformation import OpenAIGPTConfig +from ...openai.chat.gpt_transformation import ( + OpenAIChatCompletionStreamingHandler, + OpenAIGPTConfig, +) class XAIChatConfig(OpenAIGPTConfig): @@ -119,6 +128,18 @@ class XAIChatConfig(OpenAIGPTConfig): optional_params[param] = value return optional_params + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + return XAIChatCompletionStreamingHandler( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + def transform_request( self, model: str, @@ -225,3 +246,25 @@ class XAIChatConfig(OpenAIGPTConfig): usage.prompt_tokens_details.web_search_requests = int(num_sources_used) setattr(usage, "num_sources_used", int(num_sources_used)) verbose_logger.debug(f"X.AI web search sources used: {num_sources_used}") + + +class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): + def chunk_parser(self, chunk: dict) -> ModelResponseStream: + """ + Handle xAI-specific streaming behavior. + + xAI Grok sends a final chunk with empty choices array but with usage data + when stream_options={"include_usage": True} is set. + + Example from xAI API: + {"id":"...","object":"chat.completion.chunk","created":...,"model":"grok-4-1-fast-non-reasoning", + "choices":[],"usage":{"prompt_tokens":171,"completion_tokens":2,"total_tokens":173,...}} + """ + # Handle chunks with empty choices but with usage data + choices = chunk.get("choices", []) + if len(choices) == 0 and "usage" in chunk: + # xAI sends usage in a chunk with empty choices array + # Add a dummy choice with empty delta to ensure proper processing + chunk["choices"] = [{"index": 0, "delta": {}, "finish_reason": None}] + + return super().chunk_parser(chunk) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2c89e53ada9..8ed45ddd90c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14696,6 +14696,108 @@ "supports_web_search": true, "supports_native_streaming": true }, + "gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_image": 0.00012, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true + }, + "gemini-3.1-pro-preview-customtools": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_image": 0.00012, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true + }, "vertex_ai/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -14789,6 +14891,108 @@ "supports_web_search": true, "supports_native_streaming": true }, + "vertex_ai/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_image": 0.00012, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true + }, + "vertex_ai/gemini-3.1-pro-preview-customtools": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_image": 0.00012, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true + }, "gemini-2.5-pro-exp-03-25": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -16751,6 +16955,108 @@ "supports_native_streaming": true, "tpm": 800000 }, + "gemini/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-3.1-pro-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true, + "tpm": 800000 + }, + "gemini/gemini-3.1-pro-preview-customtools": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-3.1-pro-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true, + "tpm": 800000 + }, "gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, "input_cost_per_audio_token": 1e-06, diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index fbbf9cd2581..fe1ecad96c2 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Union +from typing import Dict, List, Mapping, Optional, Union from urllib.parse import parse_qs import httpx @@ -9,7 +9,9 @@ from litellm.constants import PASS_THROUGH_HEADER_PREFIX class BasePassthroughUtils: @staticmethod def get_merged_query_parameters( - existing_url: httpx.URL, request_query_params: Dict[str, Union[str, list]] + existing_url: httpx.URL, + request_query_params: Mapping[str, Union[str, list]], + default_query_params: Optional[Dict[str, Union[str, list]]] = None ) -> Dict[str, Union[str, List[str]]]: # Get the existing query params from the target URL existing_query_string = existing_url.query.decode("utf-8") @@ -19,8 +21,19 @@ class BasePassthroughUtils: updated_existing_query_params = { k: v[0] if len(v) == 1 else v for k, v in existing_query_params.items() } - # Merge the query params, giving priority to the existing ones - return {**request_query_params, **updated_existing_query_params} + + # Start with default query params (lowest priority) + merged_params = {} + if default_query_params: + merged_params.update(default_query_params) + + # Override with existing URL query params (medium priority) + merged_params.update(updated_existing_query_params) + + # Override with request query params (highest priority - client can override anything) + merged_params.update(request_query_params) + + return merged_params @staticmethod def forward_headers_from_request( diff --git a/litellm/policy_templates_backup.json b/litellm/policy_templates_backup.json index 2b1ae041f85..f6b589c3b94 100644 --- a/litellm/policy_templates_backup.json +++ b/litellm/policy_templates_backup.json @@ -1581,7 +1581,13 @@ "guardrails": [ "aviation-ops-data-protection", "aviation-safety-topic-filter", - "airline-brand-protection-filter" + "airline-brand-protection-filter", + "competitor-name-input-blocker", + "competitor-name-output-blocker", + "competitor-recommendation-input-filter", + "competitor-recommendation-output-filter", + "competitor-comparison-input-filter", + "competitor-comparison-output-filter" ], "complexity": "High", "parameters": [ @@ -1590,9 +1596,14 @@ "label": "Your Airline / Brand Name", "type": "text", "required": true, - "placeholder": "e.g. Emirates" + "placeholder": "e.g. Acme Airlines" } ], + "llm_enrichment": { + "parameter": "brand_name", + "prompt": "List the top 30 direct competitors of {{brand_name}} in the airline industry. Include major international carriers, regional competitors, and low-cost carriers that operate on overlapping routes. Return ONLY airline/brand names, one per line, no numbering, no explanations.", + "result_key": "competitors" + }, "guardrailDefinitions": [ { "guardrail_name": "aviation-ops-data-protection", @@ -1734,6 +1745,72 @@ "guardrail_info": { "description": "Blocks AI-generated fake incident reports, unauthorized statements, and reputation-damaging content about your brand (runs on output)" } + }, + { + "guardrail_name": "competitor-name-input-blocker", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitors_blocked_words}}" + }, + "guardrail_info": { + "description": "Blocks user inputs that mention competitor names (pre_call)" + } + }, + { + "guardrail_name": "competitor-name-output-blocker", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "blocked_words": "{{competitors_blocked_words}}" + }, + "guardrail_info": { + "description": "Blocks AI outputs that mention competitor names (post_call)" + } + }, + { + "guardrail_name": "competitor-recommendation-input-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitor_recommendation_words}}" + }, + "guardrail_info": { + "description": "Blocks user requests asking to recommend competitors (pre_call)" + } + }, + { + "guardrail_name": "competitor-recommendation-output-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "blocked_words": "{{competitor_recommendation_words}}" + }, + "guardrail_info": { + "description": "Blocks AI from recommending or suggesting competitor services (post_call)" + } + }, + { + "guardrail_name": "competitor-comparison-input-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitor_comparison_words}}" + }, + "guardrail_info": { + "description": "Blocks user inputs requesting unfavorable brand comparisons (pre_call)" + } + }, + { + "guardrail_name": "competitor-comparison-output-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "blocked_words": "{{competitor_comparison_words}}" + }, + "guardrail_info": { + "description": "Blocks AI outputs with unfavorable brand comparisons (post_call)" + } } ], "templateData": { @@ -1742,7 +1819,13 @@ "guardrails_add": [ "aviation-ops-data-protection", "aviation-safety-topic-filter", - "airline-brand-protection-filter" + "airline-brand-protection-filter", + "competitor-name-input-blocker", + "competitor-name-output-blocker", + "competitor-recommendation-input-filter", + "competitor-recommendation-output-filter", + "competitor-comparison-input-filter", + "competitor-comparison-output-filter" ], "guardrails_remove": [] }, @@ -1883,9 +1966,12 @@ "iconColor": "text-orange-500", "iconBg": "bg-orange-50", "guardrails": [ + "competitor-input-blocker", "competitor-output-blocker", - "competitor-recommendation-filter", - "competitor-comparison-filter" + "competitor-recommendation-input-filter", + "competitor-recommendation-output-filter", + "competitor-comparison-input-filter", + "competitor-comparison-output-filter" ], "complexity": "Medium", "parameters": [ @@ -1894,15 +1980,26 @@ "label": "Your Brand Name", "type": "text", "required": true, - "placeholder": "e.g. Emirates" + "placeholder": "e.g. Acme Airlines" } ], "llm_enrichment": { "parameter": "brand_name", - "prompt": "List the top 10 direct competitors of {{brand_name}} in the same industry. Return ONLY company/brand names, one per line, no numbering, no explanations.", + "prompt": "List the top 30 direct competitors of {{brand_name}} in the same industry. Return ONLY company/brand names, one per line, no numbering, no explanations.", "result_key": "competitors" }, "guardrailDefinitions": [ + { + "guardrail_name": "competitor-input-blocker", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitors_blocked_words}}" + }, + "guardrail_info": { + "description": "Blocks user inputs that mention competitor brands (pre_call)" + } + }, { "guardrail_name": "competitor-output-blocker", "litellm_params": { @@ -1911,39 +2008,64 @@ "blocked_words": "{{competitors_blocked_words}}" }, "guardrail_info": { - "description": "Blocks AI outputs that mention or promote competitor brands (auto-discovered via LLM)" + "description": "Blocks AI outputs that mention competitor brands (post_call)" } }, { - "guardrail_name": "competitor-recommendation-filter", + "guardrail_name": "competitor-recommendation-input-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitor_recommendation_words}}" + }, + "guardrail_info": { + "description": "Blocks user requests asking to recommend competitors (pre_call)" + } + }, + { + "guardrail_name": "competitor-recommendation-output-filter", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "post_call", "blocked_words": "{{competitor_recommendation_words}}" }, "guardrail_info": { - "description": "Blocks AI from recommending, suggesting, or directing users to competitor services" + "description": "Blocks AI from recommending or suggesting competitor services (post_call)" } }, { - "guardrail_name": "competitor-comparison-filter", + "guardrail_name": "competitor-comparison-input-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitor_comparison_words}}" + }, + "guardrail_info": { + "description": "Blocks user inputs requesting unfavorable brand comparisons (pre_call)" + } + }, + { + "guardrail_name": "competitor-comparison-output-filter", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "post_call", "blocked_words": "{{competitor_comparison_words}}" }, "guardrail_info": { - "description": "Blocks unfavorable comparisons between your brand and competitors in AI outputs" + "description": "Blocks AI outputs with unfavorable brand comparisons (post_call)" } } ], "templateData": { "policy_name": "competitor-mention-detection", - "description": "Detects and blocks competitor mentions in AI outputs. Uses LLM-powered competitor discovery based on your brand name.", + "description": "Detects and blocks competitor mentions in both inputs and outputs. Uses LLM-powered competitor discovery based on your brand name.", "guardrails_add": [ + "competitor-input-blocker", "competitor-output-blocker", - "competitor-recommendation-filter", - "competitor-comparison-filter" + "competitor-recommendation-input-filter", + "competitor-recommendation-output-filter", + "competitor-comparison-input-filter", + "competitor-comparison-output-filter" ], "guardrails_remove": [] }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 612cb0e1e76..231b9fdd1f6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -198,6 +198,7 @@ class Litellm_EntityType(enum.Enum): TEAM = "team" TEAM_MEMBER = "team_member" ORGANIZATION = "organization" + PROJECT = "project" TAG = "tag" # global proxy level entity @@ -237,6 +238,9 @@ class KeyManagementRoutes(str, enum.Enum): # list routes KEY_LIST = "/key/list" + # team usage routes + TEAM_DAILY_ACTIVITY = "/team/daily/activity" + class LiteLLMRoutes(enum.Enum): openai_route_names = [ @@ -505,6 +509,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.KEY_BLOCK.value, KeyManagementRoutes.KEY_UNBLOCK.value, KeyManagementRoutes.KEY_BULK_UPDATE.value, + KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value, ] management_routes = [ @@ -925,6 +930,7 @@ class GenerateKeyRequest(KeyRequestBase): description="How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True", ) organization_id: Optional[str] = None + project_id: Optional[str] = None class GenerateKeyResponse(KeyRequestBase): @@ -934,6 +940,7 @@ class GenerateKeyResponse(KeyRequestBase): user_id: Optional[str] = None token_id: Optional[str] = None organization_id: Optional[str] = None + project_id: Optional[str] = None litellm_budget_table: Optional[Any] = None token: Optional[str] = None created_by: Optional[str] = None @@ -1909,6 +1916,10 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase): default={}, description="Key-value pairs of headers to be forwarded with the request. You can set any key value pair here and it will be forwarded to your target endpoint", ) + default_query_params: dict = Field( + default={}, + description="Key-value pairs of default query parameters to be sent with every request to this endpoint. These can be overridden by client-provided query parameters. For example: {'key': 'default_value', 'api_version': '2023-01'}", + ) include_subpath: bool = Field( default=False, description="If True, requests to subpaths of the path will be forwarded to the target endpoint. For example, if the path is /bria and include_subpath is True, requests to /bria/v1/text-to-image/base/2.3 will be forwarded to the target endpoint.", @@ -1929,6 +1940,10 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase): default=False, description="True if this endpoint is defined in the config file, False if from DB. Config-defined endpoints cannot be edited via the UI.", ) + methods: Optional[List[str]] = Field( + default=None, + description="List of HTTP methods this endpoint handles (e.g., ['GET', 'POST']). If None or empty, all methods (GET, POST, PUT, DELETE, PATCH) are supported for backward compatibility. This allows the same path to have different targets for different HTTP methods.", + ) class PassThroughEndpointResponse(LiteLLMPydanticObjectBase): @@ -2171,6 +2186,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): config: Dict = {} user_id: Optional[str] = None team_id: Optional[str] = None + project_id: Optional[str] = None max_parallel_requests: Optional[int] = None metadata: Dict = {} tpm_limit: Optional[int] = None @@ -2190,6 +2206,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): created_by: Optional[str] = None updated_at: Optional[datetime] = None updated_by: Optional[str] = None + last_active: Optional[datetime] = None object_permission_id: Optional[str] = None object_permission: Optional[LiteLLM_ObjectPermissionTable] = None access_group_ids: Optional[List[str]] = None @@ -2521,6 +2538,116 @@ class NewOrganizationResponse(LiteLLM_OrganizationTable): updated_at: datetime +### PROJECT MANAGEMENT TYPES ### + + +class ProjectBase(LiteLLMPydanticObjectBase): + """Base fields shared by project create/update requests""" + + project_id: Optional[str] = None + project_alias: Optional[str] = None + team_id: Optional[str] = None + metadata: Optional[dict] = None + models: Optional[List[str]] = None + blocked: bool = False + + +class NewProjectRequest(LiteLLM_BudgetTable): + """Request model for POST /project/new""" + + project_id: Optional[str] = None + project_alias: Optional[str] = None + description: Optional[str] = None + team_id: str + budget_id: Optional[str] = None + metadata: Optional[dict] = None + models: List[str] = [] + model_rpm_limit: Optional[dict] = None + model_tpm_limit: Optional[dict] = None + blocked: bool = False + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if values.get(field) is not None: + if values.get("metadata") is None: + values.update({"metadata": {}}) + values["metadata"][field] = values.get(field) + values.pop(field) + return values + + +class UpdateProjectRequest(LiteLLM_BudgetTable): + """Request model for POST /project/update""" + + project_id: str + project_alias: Optional[str] = None + description: Optional[str] = None + team_id: Optional[str] = None + metadata: Optional[dict] = None + models: Optional[List[str]] = None + model_rpm_limit: Optional[dict] = None + model_tpm_limit: Optional[dict] = None + blocked: Optional[bool] = None + budget_id: Optional[str] = None + object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if values.get(field) is not None: + if values.get("metadata") is None: + values.update({"metadata": {}}) + values["metadata"][field] = values.get(field) + values.pop(field) + return values + + +class DeleteProjectRequest(LiteLLMPydanticObjectBase): + """Request model for DELETE /project/delete""" + + project_ids: List[str] + + +class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): + """Database model representation for project""" + + project_id: str + project_alias: Optional[str] = None + description: Optional[str] = None + team_id: Optional[str] = None + budget_id: Optional[str] = None + metadata: Optional[dict] = None + models: List[str] = [] + spend: float = 0.0 + model_spend: Optional[dict] = None + model_rpm_limit: Optional[dict] = None + model_tpm_limit: Optional[dict] = None + blocked: bool = False + object_permission_id: Optional[str] = None + created_by: str + updated_by: str + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + + +class NewProjectResponse(LiteLLM_ProjectTable): + """Response model for POST /project/new""" + + project_id: str + created_at: datetime + updated_at: datetime + + +class LiteLLM_ProjectTableCachedObj(LiteLLM_ProjectTable): + """Cached version for auth checks. Mirrors LiteLLM_TeamTableCachedObj pattern.""" + + last_refreshed_at: Optional[float] = None + + class LiteLLM_UserTableFiltered(BaseModel): # done to avoid exposing sensitive data user_id: str user_email: Optional[str] = None @@ -2892,6 +3019,7 @@ class SpendLogsMetadata(TypedDict): user_api_key: Optional[str] user_api_key_alias: Optional[str] user_api_key_team_id: Optional[str] + user_api_key_project_id: Optional[str] user_api_key_org_id: Optional[str] user_api_key_user_id: Optional[str] user_api_key_team_alias: Optional[str] @@ -3129,6 +3257,11 @@ class ProxyErrorTypes(str, enum.Enum): Organization does not have access to the model """ + project_model_access_denied = "project_model_access_denied" + """ + Project does not have access to the model + """ + expired_key = "expired_key" """ Key has expired @@ -3191,7 +3324,7 @@ class ProxyErrorTypes(str, enum.Enum): @classmethod def get_model_access_error_type_for_object( - cls, object_type: Literal["key", "user", "team", "org"] + cls, object_type: Literal["key", "user", "team", "org", "project"] ) -> "ProxyErrorTypes": """ Get the model access error type for object_type @@ -3204,6 +3337,8 @@ class ProxyErrorTypes(str, enum.Enum): return cls.user_model_access_denied elif object_type == "org": return cls.org_model_access_denied + elif object_type == "project": + return cls.project_model_access_denied @classmethod def get_vector_store_access_error_type_for_object( @@ -3966,8 +4101,8 @@ class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): file_object: Optional[OpenAIFileObject] = None model_mappings: Dict[str, str] flat_model_file_ids: List[str] - created_by: Optional[str] - updated_by: Optional[str] + created_by: Optional[str] = None + updated_by: Optional[str] = None storage_backend: Optional[str] = None storage_url: Optional[str] = None @@ -3985,8 +4120,8 @@ class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): resource_object: Optional[Any] = None # VectorStoreCreateResponse model_mappings: Dict[str, str] flat_model_resource_ids: List[str] - created_by: Optional[str] - updated_by: Optional[str] + created_by: Optional[str] = None + updated_by: Optional[str] = None storage_backend: Optional[str] = None storage_url: Optional[str] = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5a8a9f4e4df..0e097b689e1 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -11,8 +11,7 @@ Run checks for: import asyncio import re import time -from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, - cast) +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -21,27 +20,42 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME, - DEFAULT_ACCESS_GROUP_CACHE_TTL, - DEFAULT_IN_MEMORY_TTL, - DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, - DEFAULT_MAX_RECURSE_DEPTH, - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) +from litellm.constants import ( + CLI_JWT_EXPIRATION_HOURS, + CLI_JWT_TOKEN_NAME, + DEFAULT_ACCESS_GROUP_CACHE_TTL, + DEFAULT_IN_MEMORY_TTL, + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + DEFAULT_MAX_RECURSE_DEPTH, + EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, +) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.proxy._types import (RBAC_ROLES, CallInfo, - LiteLLM_AccessGroupTable, - LiteLLM_BudgetTable, LiteLLM_EndUserTable, - Litellm_EntityType, LiteLLM_JWTAuth, - LiteLLM_ObjectPermissionTable, - LiteLLM_OrganizationMembershipTable, - LiteLLM_OrganizationTable, LiteLLM_TagTable, - LiteLLM_TeamMembership, LiteLLM_TeamTable, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, LiteLLMRoutes, - LitellmUserRoles, NewTeamRequest, - ProxyErrorTypes, ProxyException, - RoleBasedPermissions, SpecialModelNames, - UserAPIKeyAuth) +from litellm.proxy._types import ( + RBAC_ROLES, + CallInfo, + LiteLLM_AccessGroupTable, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + Litellm_EntityType, + LiteLLM_JWTAuth, + LiteLLM_ObjectPermissionTable, + LiteLLM_OrganizationMembershipTable, + LiteLLM_OrganizationTable, + LiteLLM_TagTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_ProjectTableCachedObj, + LiteLLM_UserTable, + LiteLLMRoutes, + LitellmUserRoles, + NewTeamRequest, + ProxyErrorTypes, + ProxyException, + RoleBasedPermissions, + SpecialModelNames, + UserAPIKeyAuth, +) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics @@ -64,6 +78,7 @@ db_cache_expiry = DEFAULT_IN_MEMORY_TTL # refresh every 5s all_routes = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value + def _log_budget_lookup_failure(entity: str, error: Exception) -> None: """ Log a warning when budget lookup fails; cache will not be populated. @@ -81,38 +96,41 @@ def _log_budget_lookup_failure(entity: str, error: Exception) -> None: x in err_str for x in ("column", "schema", "does not exist", "prisma", "migrate") ): - hint = " Run `prisma db push` or `prisma migrate deploy` to fix schema mismatches." + hint = ( + " Run `prisma db push` or `prisma migrate deploy` to fix schema mismatches." + ) verbose_proxy_logger.error( f"Budget lookup failed for {entity}; cache will not be populated. " f"Each request will hit the database. Error: {error}.{hint}" ) + def _is_model_cost_zero( model: Optional[Union[str, List[str]]], llm_router: Optional[Router] ) -> bool: """ Check if a model has zero cost (no configured pricing). - + Uses the router's get_model_group_info method to get pricing information. - + Args: model: The model name or list of model names llm_router: The LiteLLM router instance - + Returns: bool: True if all costs for the model are zero, False otherwise """ if model is None or llm_router is None: return False - + # Handle list of models model_list = [model] if isinstance(model, str) else model - + for model_name in model_list: try: # Use router's get_model_group_info method directly for better reliability model_group_info = llm_router.get_model_group_info(model_group=model_name) - + if model_group_info is None: # Model not found or no pricing info available # Conservative approach: assume it has cost @@ -120,42 +138,87 @@ def _is_model_cost_zero( f"No model group info found for {model_name}, assuming it has cost" ) return False - + # Check costs for this model # Only allow bypass if BOTH costs are explicitly set to 0 (not None) input_cost = model_group_info.input_cost_per_token output_cost = model_group_info.output_cost_per_token - + # If costs are not explicitly configured (None), assume it has cost if input_cost is None or output_cost is None: verbose_proxy_logger.debug( f"Model {model_name} has undefined cost (input: {input_cost}, output: {output_cost}), assuming it has cost" ) return False - + # If either cost is non-zero, return False if input_cost > 0 or output_cost > 0: verbose_proxy_logger.debug( f"Model {model_name} has non-zero cost (input: {input_cost}, output: {output_cost})" ) return False - + # This model has zero cost explicitly configured verbose_proxy_logger.debug( f"Model {model_name} has zero cost explicitly configured (input: {input_cost}, output: {output_cost})" ) - + except Exception as e: # If we can't determine the cost, assume it has cost (conservative approach) verbose_proxy_logger.debug( f"Error checking cost for model {model_name}: {str(e)}, assuming it has cost" ) return False - + # All models checked have zero cost return True +async def _run_project_checks( + project_object: Optional[LiteLLM_ProjectTableCachedObj], + _model: Optional[Union[str, List[str]]], + llm_router: Optional[Router], + skip_budget_checks: bool, + valid_token: Optional[UserAPIKeyAuth], + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Run all project-level checks: blocked, model access, budget, soft budget. + Extracted from common_checks() to keep statement count manageable. + """ + if project_object is None: + return + + # 1.1. If project is blocked + if project_object.blocked is True: + raise Exception( + f"Project={project_object.project_id} is blocked. Update via `/project/update` if you're an admin." + ) + + # 2.2 If project can call model + if _model and len(project_object.models) > 0: + can_project_access_model( + model=_model, + project_object=project_object, + llm_router=llm_router, + ) + + if not skip_budget_checks: + # 3.0.2. If project is in budget + await _project_max_budget_check( + project_object=project_object, + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + + # 3.0.3. If project is over soft budget (alert only, doesn't block) + await _project_soft_budget_check( + project_object=project_object, + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + + async def common_checks( request_body: dict, team_object: Optional[LiteLLM_TeamTable], @@ -169,13 +232,18 @@ async def common_checks( valid_token: Optional[UserAPIKeyAuth], request: Request, skip_budget_checks: bool = False, + project_object: Optional[LiteLLM_ProjectTableCachedObj] = None, ) -> bool: """ Common checks across jwt + key-based auth. 1. If team is blocked + 1.1. If project is blocked 2. If team can call model + 2.2 If project can call model 3. If team is in budget + 3.0.2. If project is in budget + 3.0.3. If project is over soft budget (alert only) 4. If user passed in (JWT or key.user_id) - is in budget 5. If end_user (either via JWT or 'user' passed to /chat/completions, /embeddings endpoint) is in budget 6. [OPTIONAL] If 'enforce_end_user' enabled - did developer pass in 'user' param for openai endpoints @@ -220,6 +288,16 @@ async def common_checks( user_object=user_object, ) + # 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget) + await _run_project_checks( + project_object=project_object, + _model=_model, + llm_router=llm_router, + skip_budget_checks=skip_budget_checks, + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + # If this is a free model, skip all budget checks if not skip_budget_checks: # 3. If team is in budget @@ -279,7 +357,10 @@ async def common_checks( ) # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget - if end_user_object is not None and end_user_object.litellm_budget_table is not None: + if ( + end_user_object is not None + and end_user_object.litellm_budget_table is not None + ): end_user_budget = end_user_object.litellm_budget_table.max_budget if end_user_budget is not None and end_user_object.spend > end_user_budget: raise litellm.BudgetExceededError( @@ -353,8 +434,7 @@ async def common_checks( _request_metadata: dict = request_body.get("metadata", {}) or {} if _request_metadata.get("guardrails"): # check if team allowed to modify guardrails - from litellm.proxy.guardrails.guardrail_helpers import \ - can_modify_guardrails + from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails can_modify: bool = can_modify_guardrails(team_object) if can_modify is False: @@ -529,11 +609,7 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: def allowed_routes_check( - user_role: Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.TEAM, - LitellmUserRoles.INTERNAL_USER, - ], + user_role: LitellmUserRoles, user_route: str, litellm_proxy_roles: LiteLLM_JWTAuth, ) -> bool: @@ -1358,7 +1434,7 @@ async def _get_team_object_from_user_api_key_cache( raise Exception _response = LiteLLM_TeamTableCachedObj(**response.dict()) - + # Load object_permission if object_permission_id exists but object_permission is not loaded if _response.object_permission_id and not _response.object_permission: try: @@ -1373,7 +1449,7 @@ async def _get_team_object_from_user_api_key_cache( verbose_proxy_logger.debug( f"Failed to load object_permission for team {team_id} with object_permission_id={_response.object_permission_id}: {e}" ) - + # save the team object to cache await _cache_team_object( team_id=team_id, @@ -1800,8 +1876,9 @@ class ExperimentalUIJWTToken: def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import \ - encrypt_value_helper + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + encrypt_value_helper, + ) if user_info.user_role is None: raise Exception("User role is required for experimental UI login") @@ -1847,8 +1924,9 @@ class ExperimentalUIJWTToken: """ from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import \ - encrypt_value_helper + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + encrypt_value_helper, + ) if user_info.user_role is None: raise Exception("User role is required for CLI JWT login") @@ -1887,8 +1965,9 @@ class ExperimentalUIJWTToken: import json from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.proxy.common_utils.encrypt_decrypt_utils import \ - decrypt_value_helper + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + ) decrypted_token = decrypt_value_helper( hashed_token, key="ui_hash_key", exception_type="debug" @@ -2136,10 +2215,8 @@ async def _get_resources_from_access_groups( # Lazy import to avoid circular imports if prisma_client is None or user_api_key_cache is None: from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import \ - proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import \ - user_api_key_cache as _user_api_key_cache + from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj + from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache prisma_client = prisma_client or _prisma_client user_api_key_cache = user_api_key_cache or _user_api_key_cache @@ -2280,7 +2357,7 @@ def _can_object_call_model( models: List[str], team_model_aliases: Optional[Dict[str, str]] = None, team_id: Optional[str] = None, - object_type: Literal["user", "team", "key", "org"] = "user", + object_type: Literal["user", "team", "key", "org", "project"] = "user", fallback_depth: int = 0, ) -> Literal[True]: """ @@ -2474,6 +2551,24 @@ async def can_team_access_model( raise +def can_project_access_model( + model: Union[str, List[str]], + project_object: LiteLLM_ProjectTableCachedObj, + llm_router: Optional[Router], +) -> Literal[True]: + """ + Returns True if the project can access a specific model. + + Raises ProxyException if access is denied. + """ + return _can_object_call_model( + model=model, + llm_router=llm_router, + models=project_object.models if project_object else [], + object_type="project", + ) + + async def can_user_call_model( model: Union[str, List[str]], llm_router: Optional[Router], @@ -2774,14 +2869,26 @@ async def _team_soft_budget_check( if valid_token: # Extract alert emails from team metadata alert_emails: Optional[List[str]] = None - if team_object.metadata is not None and isinstance(team_object.metadata, dict): - soft_budget_alert_emails = team_object.metadata.get("soft_budget_alerting_emails") + if team_object.metadata is not None and isinstance( + team_object.metadata, dict + ): + soft_budget_alert_emails = team_object.metadata.get( + "soft_budget_alerting_emails" + ) if soft_budget_alert_emails is not None: if isinstance(soft_budget_alert_emails, list): - alert_emails = [email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip()] + alert_emails = [ + email + for email in soft_budget_alert_emails + if isinstance(email, str) and email.strip() + ] elif isinstance(soft_budget_alert_emails, str): # Handle comma-separated string - alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()] + alert_emails = [ + email.strip() + for email in soft_budget_alert_emails.split(",") + if email.strip() + ] # Filter out empty strings if alert_emails: alert_emails = [email for email in alert_emails if email] @@ -2820,6 +2927,150 @@ async def _team_soft_budget_check( ) +async def _project_max_budget_check( + project_object: Optional[LiteLLM_ProjectTableCachedObj], + valid_token: Optional[UserAPIKeyAuth], + proxy_logging_obj: ProxyLogging, +): + """ + Check if the project is over its max budget. + + Raises: + BudgetExceededError if the project is over its max budget. + Triggers a budget alert if the project is over its max budget. + """ + if project_object is None: + return + + max_budget = None + if project_object.litellm_budget_table is not None: + max_budget = project_object.litellm_budget_table.max_budget + + if ( + max_budget is not None + and project_object.spend is not None + and project_object.spend > max_budget + ): + if valid_token: + call_info = CallInfo( + token=valid_token.token, + spend=project_object.spend, + max_budget=max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.PROJECT, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="project_budget", + user_info=call_info, + ) + ) + + raise litellm.BudgetExceededError( + current_cost=project_object.spend, + max_budget=max_budget, + message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}", + ) + + +async def _project_soft_budget_check( + project_object: Optional[LiteLLM_ProjectTableCachedObj], + valid_token: Optional[UserAPIKeyAuth], + proxy_logging_obj: ProxyLogging, +): + """ + Triggers a budget alert if the project is over its soft budget. + + Mirrors _team_soft_budget_check() pattern. + """ + if project_object is None: + return + + soft_budget = None + if project_object.litellm_budget_table is not None: + soft_budget = project_object.litellm_budget_table.soft_budget + + if ( + soft_budget is not None + and project_object.spend is not None + and project_object.spend >= soft_budget + ): + verbose_proxy_logger.debug( + "Crossed Soft Budget for project %s, spend %s, soft_budget %s", + project_object.project_id, + project_object.spend, + soft_budget, + ) + if valid_token: + call_info = CallInfo( + token=valid_token.token, + spend=project_object.spend, + max_budget=None, + soft_budget=soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.PROJECT, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=call_info, + ) + ) + + +async def get_project_object( + project_id: str, + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + proxy_logging_obj: Optional[ProxyLogging] = None, +) -> Optional[LiteLLM_ProjectTableCachedObj]: + """ + Fetch project object from cache or DB. + + Follows get_team_object() caching pattern with TTL and last_refreshed_at. + + Returns LiteLLM_ProjectTableCachedObj or None if not found. + """ + if prisma_client is None: + return None + + # Check cache first + cache_key = "project_id:{}".format(project_id) + cached_obj = await user_api_key_cache.async_get_cache(key=cache_key) + if cached_obj is not None: + if isinstance(cached_obj, dict): + return LiteLLM_ProjectTableCachedObj(**cached_obj) + elif isinstance(cached_obj, LiteLLM_ProjectTableCachedObj): + return cached_obj + + # Fetch from DB + project_row = await prisma_client.db.litellm_projecttable.find_unique( + where={"project_id": project_id}, + include={"litellm_budget_table": True}, + ) + if project_row is None: + return None + + project_obj = LiteLLM_ProjectTableCachedObj(**project_row.model_dump()) + + # Cache with TTL following _cache_management_object pattern + project_obj.last_refreshed_at = time.time() + await _cache_management_object( + key=cache_key, + value=project_obj, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + return project_obj + + async def _organization_max_budget_check( valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], @@ -2921,8 +3172,7 @@ async def _tag_max_budget_check( BudgetExceededError if any tag is over its max budget. Triggers a budget alert if any tag is over its max budget. """ - from litellm.proxy.common_utils.http_parsing_utils import \ - get_tags_from_request_body + from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body if prisma_client is None: return diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index 6a8df823bc8..b81109b77c8 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -28,6 +28,7 @@ class LicenseCheck: self.license_str = os.getenv("LITELLM_LICENSE", None) verbose_proxy_logger.debug("License Str value - {}".format(self.license_str)) self.http_handler = HTTPHandler(timeout=NON_LLM_CONNECTION_TIMEOUT) + self._premium_check_logged = False self.public_key = None self.read_public_key() self.airgapped_license_data: Optional["EnterpriseLicenseData"] = None @@ -99,20 +100,23 @@ class LicenseCheck: 2. _verify: checks if license is valid calling litellm API. This is the old way we were generating/validating license """ try: - verbose_proxy_logger.debug( - "litellm.proxy.auth.litellm_license.py::is_premium() - ENTERING 'IS_PREMIUM' - LiteLLM License={}".format( - self.license_str + if not self._premium_check_logged: + verbose_proxy_logger.debug( + "litellm.proxy.auth.litellm_license.py::is_premium() - ENTERING 'IS_PREMIUM' - LiteLLM License={}".format( + self.license_str + ) ) - ) if self.license_str is None: self.license_str = os.getenv("LITELLM_LICENSE", None) - verbose_proxy_logger.debug( - "litellm.proxy.auth.litellm_license.py::is_premium() - Updated 'self.license_str' - {}".format( - self.license_str + if not self._premium_check_logged: + verbose_proxy_logger.debug( + "litellm.proxy.auth.litellm_license.py::is_premium() - Updated 'self.license_str' - {}".format( + self.license_str + ) ) - ) + self._premium_check_logged = True if self.license_str is None: return False diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e2c7dbdaba5..138f9bab2c4 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -36,6 +36,7 @@ from litellm.proxy.auth.auth_checks import ( common_checks, get_end_user_object, get_key_object, + get_project_object, get_team_object, get_user_object, is_valid_fallback_model, @@ -120,12 +121,12 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str: # Handle AWS Signature V4 format from LangChain # Format: AWS4-HMAC-SHA256 Credential=Bearer sk-12345/date/region/service/aws4_request, SignedHeaders=..., Signature=... # Extract the Bearer token from the Credential field - match = re.search(r'Credential=Bearer\s+([^/\s,]+)', api_key) + match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) if match: api_key = match.group(1) else: # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r'Credential=([^/\s,]+)', api_key) + match = re.search(r"Credential=([^/\s,]+)", api_key) if match: api_key = match.group(1) @@ -145,12 +146,12 @@ def _get_bearer_token( # Handle AWS Signature V4 format from LangChain # Format: AWS4-HMAC-SHA256 Credential=Bearer sk-12345/date/region/service/aws4_request, SignedHeaders=..., Signature=... # Extract the Bearer token from the Credential field - match = re.search(r'Credential=Bearer\s+([^/\s,]+)', api_key) + match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) if match: api_key = match.group(1) else: # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r'Credential=([^/\s,]+)', api_key) + match = re.search(r"Credential=([^/\s,]+)", api_key) if match: api_key = match.group(1) else: @@ -274,7 +275,9 @@ async def get_global_proxy_spend( proxy_logging_obj: ProxyLogging, ) -> Optional[float]: global_proxy_spend = None - if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget + if ( + litellm.max_budget > 0 and prisma_client is not None + ): # user set proxy max budget # Use event-driven coordination to prevent cache stampede cache_key = "{}:spend".format(litellm_proxy_admin_name) global_proxy_spend = await _fetch_global_spend_with_event_coordination( @@ -650,7 +653,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 skip_budget_checks = False if model is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import _is_model_cost_zero - + skip_budget_checks = _is_model_cost_zero( model=model, llm_router=llm_router ) @@ -658,7 +661,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 verbose_proxy_logger.info( f"Skipping all budget checks for zero-cost model: {model}" ) - + + # Fetch project object for JWT path if project_id is set + _jwt_project_obj = None + if valid_token.project_id is not None: + _jwt_project_obj = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # run through common checks _ = await common_checks( request=request, @@ -673,6 +686,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, skip_budget_checks=skip_budget_checks, + project_object=_jwt_project_obj, ) # return UserAPIKeyAuth object @@ -1072,7 +1086,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 skip_budget_checks = False if model is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import _is_model_cost_zero - + skip_budget_checks = _is_model_cost_zero( model=model, llm_router=llm_router ) @@ -1217,6 +1231,16 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 key=valid_token.team_id, value=_team_obj ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py + # Fetch project object if key belongs to a project + _project_obj = None + if valid_token.project_id is not None: + _project_obj = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + global_proxy_spend = None if ( litellm.max_budget > 0 and prisma_client is not None @@ -1256,6 +1280,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, skip_budget_checks=skip_budget_checks, + project_object=_project_obj, ) # Token passed all checks if valid_token is None: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7dfa3bb239f..b8866fb88aa 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -619,6 +619,23 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_call_id"] = request.headers.get( "x-litellm-call-id", str(uuid.uuid4()) ) + + ### AUTO STREAM USAGE TRACKING ### + # If always_include_stream_usage is enabled and this is a streaming request + # automatically add stream_options={'include_usage': True} if not already set + if ( + general_settings.get("always_include_stream_usage", False) is True + and self.data.get("stream", False) is True + ): + # Only set if stream_options is not already provided by the client + if "stream_options" not in self.data: + self.data["stream_options"] = {"include_usage": True} + elif ( + isinstance(self.data["stream_options"], dict) + and "include_usage" not in self.data["stream_options"] + ): + self.data["stream_options"]["include_usage"] = True + ### CALL HOOKS ### - modify/reject incoming data before calling the model ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 9675b82b145..03628fda47f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,7 +12,7 @@ import os import random import time import traceback -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast, overload import litellm @@ -792,7 +792,10 @@ class DBSpendUpdateWriter: ) in key_list_transactions.items(): batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists where={"token": token}, - data={"spend": {"increment": response_cost}}, + data={ + "spend": {"increment": response_cost}, + "last_active": datetime.now(timezone.utc), + }, ) break except DB_CONNECTION_ERROR_TYPES as e: diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py index 02b38fe04eb..237364f9714 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/__init__.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, cast from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -14,7 +14,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" # Default to always-on. Only disable if the user explicitly sets default_on: false. # We check the raw guardrail dict because LitellmParams normalizes None → False, # making it impossible to distinguish "not set" from "explicitly false" via litellm_params. - _raw_default_on = guardrail.get("litellm_params", {}).get("default_on") + _raw_default_on = cast(Dict[str, Any], guardrail).get("litellm_params", {}).get("default_on") _default_on = False if _raw_default_on is False else True _callback = MCPEndUserPermissionGuardrail( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index d903ce0d9d7..815d64f22ad 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,287 +1,295 @@ -import asyncio -import traceback -from datetime import datetime -from typing import Any, List, Optional, Union, cast - -import litellm -from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - get_litellm_metadata_from_kwargs, -) -from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import log_db_metrics -from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.utils import ProxyUpdateSpend -from litellm.types.utils import ( - StandardLoggingPayload, - StandardLoggingUserAPIKeyMetadata, -) -from litellm.utils import get_end_user_id_for_cost_tracking - - -class _ProxyDBLogger(CustomLogger): - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - await self._PROXY_track_cost_callback( - kwargs, response_obj, start_time, end_time - ) - - async def async_post_call_failure_hook( - self, - request_data: dict, - original_exception: Exception, - user_api_key_dict: UserAPIKeyAuth, - traceback_str: Optional[str] = None, - ): - request_route = user_api_key_dict.request_route - if _ProxyDBLogger._should_track_errors_in_db() is False: - return - elif request_route is not None and not RouteChecks.is_llm_api_route( - route=request_route - ): - return - - from litellm.proxy.proxy_server import proxy_logging_obj - - _metadata = dict( - StandardLoggingUserAPIKeyMetadata( - user_api_key_hash=user_api_key_dict.api_key, - user_api_key_alias=user_api_key_dict.key_alias, - 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_user_email=user_api_key_dict.user_email, - user_api_key_user_id=user_api_key_dict.user_id, - user_api_key_team_id=user_api_key_dict.team_id, - user_api_key_org_id=user_api_key_dict.org_id, - user_api_key_team_alias=user_api_key_dict.team_alias, - user_api_key_end_user_id=user_api_key_dict.end_user_id, - user_api_key_request_route=user_api_key_dict.request_route, - user_api_key_auth_metadata=user_api_key_dict.metadata, - ) - ) - _metadata["user_api_key"] = user_api_key_dict.api_key - _metadata["status"] = "failure" - _metadata["error_information"] = ( - StandardLoggingPayloadSetup.get_error_information( - original_exception=original_exception, - traceback_str=traceback_str, - ) - ) - - existing_metadata: dict = request_data.get("metadata", None) or {} - existing_metadata.update(_metadata) - - if "litellm_params" not in request_data: - request_data["litellm_params"] = {} - - existing_litellm_params = request_data.get("litellm_params", {}) - existing_litellm_metadata = existing_litellm_params.get("metadata", {}) or {} - - # Preserve tags from existing metadata - if existing_litellm_metadata.get("tags"): - existing_metadata["tags"] = existing_litellm_metadata.get("tags") - - request_data["litellm_params"]["proxy_server_request"] = ( - request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} - ) - request_data["litellm_params"]["metadata"] = existing_metadata - - # Preserve model name and custom_llm_provider - if "model" not in request_data: - request_data["model"] = existing_litellm_params.get("model") or request_data.get("model", "") - if "custom_llm_provider" not in request_data: - request_data["custom_llm_provider"] = existing_litellm_params.get("custom_llm_provider") or request_data.get("custom_llm_provider", "") - - await proxy_logging_obj.db_spend_update_writer.update_database( - token=user_api_key_dict.api_key, - response_cost=0.0, - user_id=user_api_key_dict.user_id, - end_user_id=user_api_key_dict.end_user_id, - team_id=user_api_key_dict.team_id, - kwargs=request_data, - completion_response=original_exception, - start_time=datetime.now(), - end_time=datetime.now(), - org_id=user_api_key_dict.org_id, - ) - - @log_db_metrics - async def _PROXY_track_cost_callback( - self, - kwargs, # kwargs to completion - completion_response: Optional[ - Union[litellm.ModelResponse, Any] - ], # response from completion - start_time=None, - end_time=None, # start/end time for completion - ): - from litellm.proxy.proxy_server import proxy_logging_obj, update_cache - - verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") - try: - verbose_proxy_logger.debug( - f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}" - ) - parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=kwargs) - litellm_params = kwargs.get("litellm_params", {}) or {} - end_user_id = get_end_user_id_for_cost_tracking(litellm_params) - metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) - user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) - team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) - org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) - key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None)) - end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) - sl_object: Optional[StandardLoggingPayload] = kwargs.get( - "standard_logging_object", None - ) - response_cost = ( - sl_object.get("response_cost", None) - if sl_object is not None - else kwargs.get("response_cost", None) - ) - tags: Optional[List[str]] = ( - sl_object.get("request_tags", None) if sl_object is not None else None - ) - - if response_cost is not None: - user_api_key = metadata.get("user_api_key", None) - if kwargs.get("cache_hit", False) is True: - response_cost = 0.0 - verbose_proxy_logger.debug( - f"Cache Hit: response_cost {response_cost}, for user_id {user_id}" - ) - - verbose_proxy_logger.debug( - f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" - ) - if _should_track_cost_callback( - user_api_key=user_api_key, - user_id=user_id, - team_id=team_id, - end_user_id=end_user_id, - ): - ## UPDATE DATABASE - await proxy_logging_obj.db_spend_update_writer.update_database( - token=user_api_key, - response_cost=response_cost, - user_id=user_id, - end_user_id=end_user_id, - team_id=team_id, - kwargs=kwargs, - completion_response=completion_response, - start_time=start_time, - end_time=end_time, - org_id=org_id, - ) - - # update cache - asyncio.create_task( - update_cache( - token=user_api_key, - user_id=user_id, - end_user_id=end_user_id, - response_cost=response_cost, - team_id=team_id, - parent_otel_span=parent_otel_span, - tags=tags, - ) - ) - - await proxy_logging_obj.slack_alerting_instance.customer_spend_alert( - token=user_api_key, - key_alias=key_alias, - end_user_id=end_user_id, - response_cost=response_cost, - max_budget=end_user_max_budget, - ) - else: - # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. - # Use .get() for "stream" to avoid KeyError on health checks. - if sl_object is None and not kwargs.get("model"): - verbose_proxy_logger.warning( - "Cost tracking - skipping, no standard_logging_object and no model for call_type=%s", - kwargs.get("call_type", "unknown"), - ) - return - if kwargs.get("stream") is not True or ( - kwargs.get("stream") is True and "complete_streaming_response" in kwargs - ): - if sl_object is not None: - cost_tracking_failure_debug_info: Union[dict, str] = ( - sl_object["response_cost_failure_debug_info"] # type: ignore - or "response_cost_failure_debug_info is None in standard_logging_object" - ) - else: - cost_tracking_failure_debug_info = ( - "standard_logging_object not found" - ) - model = kwargs.get("model") - raise Exception( - f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" - ) - except Exception as e: - error_msg = f"Error in tracking cost callback - {str(e)}\n Traceback:{traceback.format_exc()}" - model = kwargs.get("model", "") - metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) - litellm_metadata = kwargs.get("litellm_params", {}).get( - "litellm_metadata", {} - ) - old_metadata = kwargs.get("litellm_params", {}).get("metadata", {}) - call_type = kwargs.get("call_type", "") - error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n chosen_metadata: {metadata}\n litellm_metadata: {litellm_metadata}\n old_metadata: {old_metadata}\n call_type: {call_type}\n" - asyncio.create_task( - proxy_logging_obj.failed_tracking_alert( - error_message=error_msg, - failing_model=model, - ) - ) - - verbose_proxy_logger.exception( - "Error in tracking cost callback - %s", str(e) - ) - - @staticmethod - def _should_track_errors_in_db(): - """ - Returns True if errors should be tracked in the database - - By default, errors are tracked in the database - - If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings - """ - from litellm.proxy.proxy_server import general_settings - - if general_settings.get("disable_error_logs") is True: - return False - return - - -def _should_track_cost_callback( - user_api_key: Optional[str], - user_id: Optional[str], - team_id: Optional[str], - end_user_id: Optional[str], -) -> bool: - """ - Determine if the cost callback should be tracked based on the kwargs - """ - - # don't run track cost callback if user opted into disabling spend - if ProxyUpdateSpend.disable_spend_updates() is True: - return False - - if ( - user_api_key is not None - or user_id is not None - or team_id is not None - or end_user_id is not None - ): - return True - return False +import asyncio +import traceback +from datetime import datetime +from typing import Any, List, Optional, Union, cast + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + get_litellm_metadata_from_kwargs, +) +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import log_db_metrics +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.utils import ProxyUpdateSpend +from litellm.types.utils import ( + StandardLoggingPayload, + StandardLoggingUserAPIKeyMetadata, +) +from litellm.utils import get_end_user_id_for_cost_tracking + + +class _ProxyDBLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await self._PROXY_track_cost_callback( + kwargs, response_obj, start_time, end_time + ) + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: Optional[str] = None, + ): + request_route = user_api_key_dict.request_route + if _ProxyDBLogger._should_track_errors_in_db() is False: + return + elif request_route is not None and not RouteChecks.is_llm_api_route( + route=request_route + ): + return + + from litellm.proxy.proxy_server import proxy_logging_obj + + _metadata = dict( + StandardLoggingUserAPIKeyMetadata( + user_api_key_hash=user_api_key_dict.api_key, + user_api_key_alias=user_api_key_dict.key_alias, + 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_user_email=user_api_key_dict.user_email, + user_api_key_user_id=user_api_key_dict.user_id, + user_api_key_team_id=user_api_key_dict.team_id, + user_api_key_org_id=user_api_key_dict.org_id, + user_api_key_project_id=user_api_key_dict.project_id, + user_api_key_team_alias=user_api_key_dict.team_alias, + user_api_key_end_user_id=user_api_key_dict.end_user_id, + user_api_key_request_route=user_api_key_dict.request_route, + user_api_key_auth_metadata=user_api_key_dict.metadata, + ) + ) + _metadata["user_api_key"] = user_api_key_dict.api_key + _metadata["status"] = "failure" + _metadata[ + "error_information" + ] = StandardLoggingPayloadSetup.get_error_information( + original_exception=original_exception, + traceback_str=traceback_str, + ) + + existing_metadata: dict = request_data.get("metadata", None) or {} + existing_metadata.update(_metadata) + + if "litellm_params" not in request_data: + request_data["litellm_params"] = {} + + existing_litellm_params = request_data.get("litellm_params", {}) + existing_litellm_metadata = existing_litellm_params.get("metadata", {}) or {} + + # Preserve tags from existing metadata + if existing_litellm_metadata.get("tags"): + existing_metadata["tags"] = existing_litellm_metadata.get("tags") + + request_data["litellm_params"]["proxy_server_request"] = ( + request_data.get("proxy_server_request") + or existing_litellm_params.get("proxy_server_request") + or {} + ) + request_data["litellm_params"]["metadata"] = existing_metadata + + # Preserve model name and custom_llm_provider + if "model" not in request_data: + request_data["model"] = existing_litellm_params.get( + "model" + ) or request_data.get("model", "") + if "custom_llm_provider" not in request_data: + request_data["custom_llm_provider"] = existing_litellm_params.get( + "custom_llm_provider" + ) or request_data.get("custom_llm_provider", "") + + await proxy_logging_obj.db_spend_update_writer.update_database( + token=user_api_key_dict.api_key, + response_cost=0.0, + user_id=user_api_key_dict.user_id, + end_user_id=user_api_key_dict.end_user_id, + team_id=user_api_key_dict.team_id, + kwargs=request_data, + completion_response=original_exception, + start_time=datetime.now(), + end_time=datetime.now(), + org_id=user_api_key_dict.org_id, + ) + + @log_db_metrics + async def _PROXY_track_cost_callback( + self, + kwargs, # kwargs to completion + completion_response: Optional[ + Union[litellm.ModelResponse, Any] + ], # response from completion + start_time=None, + end_time=None, # start/end time for completion + ): + from litellm.proxy.proxy_server import proxy_logging_obj, update_cache + + verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") + try: + verbose_proxy_logger.debug( + f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}" + ) + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=kwargs) + litellm_params = kwargs.get("litellm_params", {}) or {} + end_user_id = get_end_user_id_for_cost_tracking(litellm_params) + metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) + user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) + team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) + org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) + key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None)) + end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) + sl_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object", None + ) + response_cost = ( + sl_object.get("response_cost", None) + if sl_object is not None + else kwargs.get("response_cost", None) + ) + tags: Optional[List[str]] = ( + sl_object.get("request_tags", None) if sl_object is not None else None + ) + + if response_cost is not None: + user_api_key = metadata.get("user_api_key", None) + if kwargs.get("cache_hit", False) is True: + response_cost = 0.0 + verbose_proxy_logger.debug( + f"Cache Hit: response_cost {response_cost}, for user_id {user_id}" + ) + + verbose_proxy_logger.debug( + f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" + ) + if _should_track_cost_callback( + user_api_key=user_api_key, + user_id=user_id, + team_id=team_id, + end_user_id=end_user_id, + ): + ## UPDATE DATABASE + await proxy_logging_obj.db_spend_update_writer.update_database( + token=user_api_key, + response_cost=response_cost, + user_id=user_id, + end_user_id=end_user_id, + team_id=team_id, + kwargs=kwargs, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, + org_id=org_id, + ) + + # update cache + asyncio.create_task( + update_cache( + token=user_api_key, + user_id=user_id, + end_user_id=end_user_id, + response_cost=response_cost, + team_id=team_id, + parent_otel_span=parent_otel_span, + tags=tags, + ) + ) + + await proxy_logging_obj.slack_alerting_instance.customer_spend_alert( + token=user_api_key, + key_alias=key_alias, + end_user_id=end_user_id, + response_cost=response_cost, + max_budget=end_user_max_budget, + ) + else: + # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. + # Use .get() for "stream" to avoid KeyError on health checks. + if sl_object is None and not kwargs.get("model"): + verbose_proxy_logger.warning( + "Cost tracking - skipping, no standard_logging_object and no model for call_type=%s", + kwargs.get("call_type", "unknown"), + ) + return + if kwargs.get("stream") is not True or ( + kwargs.get("stream") is True + and "complete_streaming_response" in kwargs + ): + if sl_object is not None: + cost_tracking_failure_debug_info: Union[dict, str] = ( + sl_object["response_cost_failure_debug_info"] # type: ignore + or "response_cost_failure_debug_info is None in standard_logging_object" + ) + else: + cost_tracking_failure_debug_info = ( + "standard_logging_object not found" + ) + model = kwargs.get("model") + raise Exception( + f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" + ) + except Exception as e: + error_msg = f"Error in tracking cost callback - {str(e)}\n Traceback:{traceback.format_exc()}" + model = kwargs.get("model", "") + metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) + litellm_metadata = kwargs.get("litellm_params", {}).get( + "litellm_metadata", {} + ) + old_metadata = kwargs.get("litellm_params", {}).get("metadata", {}) + call_type = kwargs.get("call_type", "") + error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n chosen_metadata: {metadata}\n litellm_metadata: {litellm_metadata}\n old_metadata: {old_metadata}\n call_type: {call_type}\n" + asyncio.create_task( + proxy_logging_obj.failed_tracking_alert( + error_message=error_msg, + failing_model=model, + ) + ) + + verbose_proxy_logger.exception( + "Error in tracking cost callback - %s", str(e) + ) + + @staticmethod + def _should_track_errors_in_db(): + """ + Returns True if errors should be tracked in the database + + By default, errors are tracked in the database + + If users want to disable error tracking, they can set the disable_error_logs flag in the general_settings + """ + from litellm.proxy.proxy_server import general_settings + + if general_settings.get("disable_error_logs") is True: + return False + return + + +def _should_track_cost_callback( + user_api_key: Optional[str], + user_id: Optional[str], + team_id: Optional[str], + end_user_id: Optional[str], +) -> bool: + """ + Determine if the cost callback should be tracked based on the kwargs + """ + + # don't run track cost callback if user opted into disabling spend + if ProxyUpdateSpend.disable_spend_updates() is True: + return False + + if ( + user_api_key is not None + or user_id is not None + or team_id is not None + or end_user_id is not None + ): + return True + return False diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3f50aad355d..dc9710c2ce2 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -591,6 +591,7 @@ class LiteLLMProxyRequestSetup: user_api_key_spend=user_api_key_dict.spend, user_api_key_max_budget=user_api_key_dict.max_budget, user_api_key_team_id=user_api_key_dict.team_id, + user_api_key_project_id=user_api_key_dict.project_id, user_api_key_user_id=user_api_key_dict.user_id, user_api_key_org_id=user_api_key_dict.org_id, user_api_key_team_alias=user_api_key_dict.team_alias, diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 34a17f6c274..999c53e6b82 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -7,6 +7,7 @@ from litellm.proxy._types import ( LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_OrganizationTable, + LiteLLM_ProjectTable, LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, @@ -37,6 +38,25 @@ def _is_user_team_admin( return False +def _team_member_has_permission( + user_api_key_dict: UserAPIKeyAuth, + team_obj: LiteLLM_TeamTable, + permission: str, +) -> bool: + """Check if a non-admin team member has a specific permission on a team.""" + if not team_obj.team_member_permissions: + return False + if permission not in team_obj.team_member_permissions: + return False + for member in team_obj.members_with_roles: + if ( + member.user_id is not None + and member.user_id == user_api_key_dict.user_id + ): + return True + return False + + async def _user_has_admin_privileges( user_api_key_dict: UserAPIKeyAuth, prisma_client: Optional["PrismaClient"] = None, @@ -262,6 +282,7 @@ def _set_object_metadata_field( LiteLLM_TeamTable, KeyRequestBase, LiteLLM_OrganizationTable, + LiteLLM_ProjectTable, ], field_name: str, value: Any, @@ -270,7 +291,7 @@ def _set_object_metadata_field( Helper function to set metadata fields that require premium user checks Args: - object_data: The team data object to modify + object_data: The team/key/organization/project data object to modify field_name: Name of the metadata field to set value: Value to set for the field """ diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 57b6453ac43..36fc65de810 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -573,7 +573,7 @@ def get_user_id_from_request(request: Request) -> Optional[str]: "/user/info", tags=["Internal User management"], dependencies=[Depends(user_api_key_auth)], - # response_model=UserInfoResponse, + response_model=UserInfoResponse, ) @management_endpoint_wrapper async def user_info( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 007e558cffc..12adfc9c89b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy.auth.auth_checks import ( can_team_access_model, get_key_object, get_org_object, + get_project_object, get_team_object, ) from litellm.proxy.auth.auth_utils import abbreviate_api_key @@ -890,6 +891,61 @@ async def _check_team_key_limits( ) +async def _check_project_key_limits( + project_id: str, + data: Union[GenerateKeyRequest, UpdateKeyRequest], + prisma_client: PrismaClient, + user_api_key_cache: DualCache, +) -> None: + """ + Validate that key's models and budget respect its project's limits. + + - Key models must be a subset of project models + - Key max_budget must be <= project max_budget + """ + project_obj = await get_project_object( + project_id=project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + if project_obj is None: + raise HTTPException( + status_code=404, + detail={"error": f"Project not found, project_id={project_id}"}, + ) + + # Validate key models are a subset of project models + if data.models and len(project_obj.models) > 0: + for m in data.models: + if m not in project_obj.models: + raise HTTPException( + status_code=400, + detail={ + "error": f"Model '{m}' not in project's allowed models. Project allowed models={project_obj.models}. Project: {project_id}" + }, + ) + + # Validate key max_budget <= project max_budget + project_max_budget = None + if project_obj.litellm_budget_table is not None: + project_max_budget = getattr( + project_obj.litellm_budget_table, "max_budget", None + ) + + if ( + data.max_budget is not None + and project_max_budget is not None + and data.max_budget > project_max_budget + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Key max_budget ({data.max_budget}) exceeds project's max_budget ({project_max_budget}). Project: {project_id}" + }, + ) + + def check_org_key_model_specific_limits( keys: List[LiteLLM_VerificationToken], org_table: LiteLLM_OrganizationTable, @@ -1145,6 +1201,15 @@ async def generate_key_fn( prisma_client=prisma_client, ) + # Validate key against project limits if project_id is set + if data.project_id is not None: + await _check_project_key_limits( + project_id=data.project_id, + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return await _common_key_generation_helper( data=data, user_api_key_dict=user_api_key_dict, @@ -1820,6 +1885,20 @@ async def update_key_fn( prisma_client=prisma_client, ) + # Validate key against project limits if project_id is being set + _project_id_to_check = getattr(data, "project_id", None) or getattr( + existing_key_row, "project_id", None + ) + if _project_id_to_check is not None and ( + data.models is not None or data.max_budget is not None + ): + await _check_project_key_limits( + project_id=_project_id_to_check, + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + # if team change - check if this is possible if is_different_team(data=data, existing_key_row=existing_key_row): if llm_router is None: @@ -2475,6 +2554,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 prompts: Optional[list] = None, teams: Optional[list] = None, organization_id: Optional[str] = None, + project_id: Optional[str] = None, table_name: Optional[Literal["key", "user"]] = None, send_invite_email: Optional[bool] = None, created_by: Optional[str] = None, @@ -2588,6 +2668,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 "max_budget": key_max_budget, "user_id": user_id, "team_id": team_id, + "project_id": project_id, "max_parallel_requests": max_parallel_requests, "metadata": metadata_json, "tpm_limit": tpm_limit, diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index c400771936f..4aa33f0a473 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -9,14 +9,29 @@ All /policy management endpoints /policy/templates - Get policy templates (GitHub with local fallback) """ +import copy import json import os -from typing import TYPE_CHECKING, List, Literal, Optional, TypedDict, cast +from typing import ( + TYPE_CHECKING, + AsyncIterator, + List, + Literal, + Optional, + TypedDict, + cast, +) from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + COMPETITOR_LLM_TEMPERATURE, + DEFAULT_COMPETITOR_DISCOVERY_MODEL, + MAX_COMPETITOR_NAMES, +) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -512,6 +527,48 @@ async def get_policy_templates( class EnrichTemplateRequest(BaseModel): template_id: str parameters: dict + model: Optional[str] = None + competitors: Optional[List[str]] = Field( + default=None, + max_length=MAX_COMPETITOR_NAMES, + description="Optional list of competitor names", + ) + instruction: Optional[str] = Field( + default=None, + description="Refinement instruction for modifying the competitor list (e.g. 'add 10 more from Asia')", + ) + + +def _validate_enrichment_request(data: EnrichTemplateRequest) -> tuple[dict, dict, str]: + """ + Validate enrichment request and return (template, llm_enrichment, brand_name). + + Raises HTTPException on validation failure. + """ + templates = _load_policy_templates_from_local_backup() + template = next((t for t in templates if t.get("id") == data.template_id), None) + if template is None: + raise HTTPException(status_code=404, detail=f"Template '{data.template_id}' not found") + + llm_enrichment = template.get("llm_enrichment") + if llm_enrichment is None: + raise HTTPException(status_code=400, detail="Template does not support LLM enrichment") + + # Validate competitors list size if provided + if data.competitors and len(data.competitors) > MAX_COMPETITOR_NAMES: + raise HTTPException( + status_code=400, + detail=f"competitors list exceeds maximum of {MAX_COMPETITOR_NAMES}", + ) + + brand_name = data.parameters.get(llm_enrichment["parameter"], "") + if not brand_name: + raise HTTPException( + status_code=400, + detail=f"Parameter '{llm_enrichment['parameter']}' is required", + ) + + return template, llm_enrichment, brand_name @router.post( @@ -531,108 +588,311 @@ async def enrich_policy_template( Calls an onboarded LLM to discover competitors for the given brand name, then returns enriched guardrailDefinitions with the discovered data populated. """ - templates = _load_policy_templates_from_local_backup() - template = next((t for t in templates if t.get("id") == data.template_id), None) - if template is None: - raise HTTPException(status_code=404, detail=f"Template '{data.template_id}' not found") + template, llm_enrichment, brand_name = _validate_enrichment_request(data) + model = data.model or DEFAULT_COMPETITOR_DISCOVERY_MODEL - llm_enrichment = template.get("llm_enrichment") - if llm_enrichment is None: - raise HTTPException( - status_code=400, - detail="Template does not support LLM enrichment", + if data.competitors: + competitors = data.competitors + else: + prompt = llm_enrichment["prompt"].replace( + "{{" + llm_enrichment["parameter"] + "}}", brand_name ) + competitors = await _discover_competitors_via_llm(prompt, model=model) - brand_name = data.parameters.get(llm_enrichment["parameter"], "") - if not brand_name: - raise HTTPException( - status_code=400, - detail=f"Parameter '{llm_enrichment['parameter']}' is required", - ) - - prompt = llm_enrichment["prompt"].replace( - "{{" + llm_enrichment["parameter"] + "}}", brand_name - ) - - competitors = await _discover_competitors_via_llm(prompt) + variations_map = await _generate_competitor_variations(competitors, model=model) enriched_definitions = _build_competitor_guardrail_definitions( template.get("guardrailDefinitions", []), competitors, brand_name, + variations_map, ) - return {"guardrailDefinitions": enriched_definitions, "competitors": competitors} + return { + "guardrailDefinitions": enriched_definitions, + "competitors": competitors, + "competitor_variations": variations_map, + } -async def _discover_competitors_via_llm(prompt: str) -> list: - """Call an onboarded LLM to discover competitor names.""" - import litellm +def _build_refinement_prompt( + instruction: str, + existing_competitors: list[str], + brand_name: str, +) -> str: + """Build a prompt for refining the competitor list based on user instruction.""" + existing_list = ", ".join(existing_competitors) + return ( + f"I have a brand called '{brand_name}' and the following competitor list:\n" + f"{existing_list}\n\n" + f"User instruction: {instruction}\n\n" + "Return ONLY the NEW names to add (not the existing ones), one per line, " + "no numbering, no explanations. If the instruction asks to remove names, " + "return nothing." + ) + + +async def _stream_llm_competitor_names( + prompt: str, + model: str, + existing: list[str], +) -> AsyncIterator[tuple[Optional[str], bool]]: + """ + Stream competitor names from LLM. Yields (name, is_error) tuples. + + Deduplicates against existing names (case-insensitive). + """ + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + raise ValueError("LLM router not initialized") + + existing_lower = {n.lower() for n in existing} + response = await llm_router.acompletion( + model=model, + messages=[{"role": "user", "content": prompt}], + temperature=COMPETITOR_LLM_TEMPERATURE, + stream=True, + ) + buffer = "" + count = len(existing) + async for chunk in response: # type: ignore[union-attr] + delta = chunk.choices[0].delta.content or "" + buffer += delta + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + name = _clean_competitor_line(line) + if name and name.lower() not in existing_lower and count < MAX_COMPETITOR_NAMES: + existing_lower.add(name.lower()) + count += 1 + yield name, False + # Handle remaining buffer + name = _clean_competitor_line(buffer) + if name and name.lower() not in existing_lower and count < MAX_COMPETITOR_NAMES: + yield name, False + + +async def _stream_competitor_events( + data: EnrichTemplateRequest, + template: dict, + llm_enrichment: dict, + brand_name: str, + model: str, +) -> AsyncIterator[str]: + """Stream competitor names as SSE events, then emit a final 'done' event.""" + competitors: list[str] = list(data.competitors or []) + + if data.instruction and competitors: + # Refinement mode: keep existing, stream only new names + for comp in competitors: + yield f"data: {json.dumps({'type': 'competitor', 'name': comp})}\n\n" + + refinement_prompt = _build_refinement_prompt( + data.instruction, competitors, brand_name + ) + try: + async for name, _ in _stream_llm_competitor_names( + refinement_prompt, model, competitors + ): + if name: + competitors.append(name) + yield f"data: {json.dumps({'type': 'competitor', 'name': name})}\n\n" + except Exception as e: + verbose_proxy_logger.error("LLM competitor refinement failed: %s", e) + yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n" + return + elif data.competitors and not data.instruction: + # Free-form mode (no instruction): just emit existing + for comp in competitors: + yield f"data: {json.dumps({'type': 'competitor', 'name': comp})}\n\n" + else: + # Initial discovery mode + prompt = llm_enrichment["prompt"].replace( + "{{" + llm_enrichment["parameter"] + "}}", brand_name + ) + try: + async for name, _ in _stream_llm_competitor_names( + prompt, model, [] + ): + if name: + competitors.append(name) + yield f"data: {json.dumps({'type': 'competitor', 'name': name})}\n\n" + except Exception as e: + verbose_proxy_logger.error("LLM competitor streaming failed: %s", e) + yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n" + return + + yield f"data: {json.dumps({'type': 'status', 'message': f'Generating alternate spellings for {len(competitors)} competitors...'})}\n\n" + variations_map = await _generate_competitor_variations(competitors, model=model) + + total_variations = sum(len(v) for v in variations_map.values()) + yield f"data: {json.dumps({'type': 'status', 'message': f'Building guardrail definitions with {total_variations} variations...'})}\n\n" + enriched_definitions = _build_competitor_guardrail_definitions( + template.get("guardrailDefinitions", []), + competitors, + brand_name, + variations_map, + ) + + yield f"data: {json.dumps({'type': 'done', 'competitors': competitors, 'competitor_variations': variations_map, 'guardrailDefinitions': enriched_definitions})}\n\n" + + +@router.post( + "/policy/templates/enrich/stream", + tags=["policy management"], + dependencies=[Depends(user_api_key_auth)], +) +async def enrich_policy_template_stream( + data: EnrichTemplateRequest, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Stream competitor names as SSE events as the LLM generates them. + + Events: + - data: {"type": "competitor", "name": "..."} — each competitor as discovered + - data: {"type": "done", "competitors": [...], "competitor_variations": {...}, "guardrailDefinitions": [...]} + """ + template, llm_enrichment, brand_name = _validate_enrichment_request(data) + model = data.model or DEFAULT_COMPETITOR_DISCOVERY_MODEL + + return StreamingResponse( + _stream_competitor_events(data, template, llm_enrichment, brand_name, model), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + ) + + +def _clean_competitor_line(line: str) -> Optional[str]: + """Strip numbering, bullets, and whitespace from a competitor name line.""" + name = line.strip().strip(".-) ").strip() + return name if name and len(name) > 1 else None + + +async def _generate_competitor_variations( + competitors: list, model: str = DEFAULT_COMPETITOR_DISCOVERY_MODEL +) -> dict: + """Generate common misspellings, abbreviations, and alternate names for each competitor.""" + if not competitors: + return {} + + # Cap the list to prevent oversized prompts + capped = competitors[:MAX_COMPETITOR_NAMES] + names_list = "\n".join(capped) + prompt = ( + "For each company/brand name below, list 3-5 common misspellings, abbreviations, " + "and alternate names that people might type. Include typos, missing spaces, " + "wrong suffixes (e.g. 'Airlines' vs 'Airways' vs 'Airline'), and common shortcuts.\n\n" + f"Names:\n{names_list}\n\n" + "Return the result as one line per variation in the format:\n" + "OriginalName: variation1, variation2, variation3\n" + "Use the EXACT original name before the colon. No numbering, no extra text." + ) try: - response = await litellm.acompletion( - model="gpt-4o-mini", + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + raise ValueError("LLM router not initialized") + response = await llm_router.acompletion( + model=model, messages=[{"role": "user", "content": prompt}], - temperature=0.3, + temperature=COMPETITOR_LLM_TEMPERATURE, + ) + raw = response.choices[0].message.content or "" # type: ignore + return _parse_variations_response(raw, capped) + except Exception as e: + verbose_proxy_logger.error("LLM competitor variation generation failed: %s", e) + return {} + + +def _parse_variations_response(raw: str, competitors: list) -> dict[str, list[str]]: + """Parse the LLM response for competitor variations into a name -> variations map.""" + # Build a lowercase lookup for case-insensitive matching + lower_to_canonical = {comp.lower(): comp for comp in competitors} + variations_map: dict[str, list[str]] = {} + + for line in raw.strip().split("\n"): + if ":" not in line: + continue + name, _, variations_str = line.partition(":") + canonical = lower_to_canonical.get(name.strip().lower()) + if canonical is None: + continue + variations = [ + v.strip() + for v in variations_str.split(",") + if v.strip() and v.strip().lower() != canonical.lower() + ] + variations_map[canonical] = variations + + return variations_map + + +async def _discover_competitors_via_llm( + prompt: str, model: str = DEFAULT_COMPETITOR_DISCOVERY_MODEL +) -> list: + """Call an onboarded LLM to discover competitor names.""" + try: + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + raise ValueError("LLM router not initialized") + response = await llm_router.acompletion( + model=model, + messages=[{"role": "user", "content": prompt}], + temperature=COMPETITOR_LLM_TEMPERATURE, ) raw = response.choices[0].message.content or "" # type: ignore competitors = [ - line.strip().strip(".-) ").strip() + name for line in raw.strip().split("\n") - if line.strip() and len(line.strip()) > 1 + if (name := _clean_competitor_line(line)) is not None ] - return competitors[:15] + return competitors[:MAX_COMPETITOR_NAMES] except Exception as e: verbose_proxy_logger.error("LLM competitor discovery failed: %s", e) return [] +def _build_all_names_per_competitor( + competitors: list[str], variations_map: dict[str, list[str]] +) -> dict[str, list[str]]: + """Build canonical + variation name lists for each competitor.""" + return { + comp: [comp] + variations_map.get(comp, []) + for comp in competitors + } + + def _build_competitor_guardrail_definitions( definitions: list, competitors: list, brand_name: str, + variations_map: Optional[dict] = None, ) -> list: - """Build enriched guardrailDefinitions with competitor names populated.""" - import copy - + """Build enriched guardrailDefinitions with competitor names and variations populated.""" + variations_map = variations_map or {} enriched = copy.deepcopy(definitions) + all_names = _build_all_names_per_competitor(competitors, variations_map) - output_blocked = [ - {"keyword": comp, "action": "BLOCK", "description": f"Competitor: {comp}"} - for comp in competitors - ] - - recommendation_blocked = [] - for comp in competitors: - recommendation_blocked.append( - {"keyword": f"try {comp}", "action": "BLOCK", "description": "Recommendation to competitor"} - ) - recommendation_blocked.append( - {"keyword": f"use {comp}", "action": "BLOCK", "description": "Recommendation to competitor"} - ) - recommendation_blocked.append( - {"keyword": f"switch to {comp}", "action": "BLOCK", "description": "Recommendation to competitor"} - ) - recommendation_blocked.append( - {"keyword": f"consider {comp}", "action": "BLOCK", "description": "Recommendation to competitor"} - ) - - comparison_blocked = [] - for comp in competitors: - comparison_blocked.append( - {"keyword": f"{comp} is better", "action": "BLOCK", "description": "Unfavorable comparison"} - ) - comparison_blocked.append( - {"keyword": f"better than {brand_name}", "action": "BLOCK", "description": "Unfavorable comparison"} - ) - comparison_blocked.append( - {"keyword": f"{brand_name} is worse", "action": "BLOCK", "description": "Unfavorable comparison"} - ) + output_blocked = _build_name_blocked_words(competitors, all_names) + recommendation_blocked = _build_recommendation_blocked_words(competitors, all_names) + comparison_blocked = _build_comparison_blocked_words(competitors, all_names, brand_name) blocked_words_map = { "competitor-output-blocker": output_blocked, + "competitor-input-blocker": output_blocked, + "competitor-name-blocker": output_blocked, + "competitor-name-input-blocker": output_blocked, + "competitor-name-output-blocker": output_blocked, "competitor-recommendation-filter": recommendation_blocked, + "competitor-recommendation-input-filter": recommendation_blocked, + "competitor-recommendation-output-filter": recommendation_blocked, "competitor-comparison-filter": comparison_blocked, + "competitor-comparison-input-filter": comparison_blocked, + "competitor-comparison-output-filter": comparison_blocked, } for defn in enriched: @@ -643,6 +903,62 @@ def _build_competitor_guardrail_definitions( return enriched +def _build_name_blocked_words( + competitors: list[str], all_names: dict[str, list[str]] +) -> list[dict]: + """Build blocked word entries for direct competitor name mentions.""" + result = [] + for comp in competitors: + for name in all_names[comp]: + desc = f"Competitor: {comp}" if name == comp else f"Competitor variation ({comp}): {name}" + result.append({"keyword": name, "action": "BLOCK", "description": desc}) + return result + + +def _build_recommendation_blocked_words( + competitors: list[str], all_names: dict[str, list[str]] +) -> list[dict]: + """Build blocked word entries for competitor recommendations.""" + result = [] + for comp in competitors: + for name in all_names[comp]: + for prefix in ["try", "use", "switch to", "consider"]: + result.append({ + "keyword": f"{prefix} {name}", + "action": "BLOCK", + "description": f"Recommendation to competitor ({comp})", + }) + return result + + +def _build_comparison_blocked_words( + competitors: list[str], all_names: dict[str, list[str]], brand_name: str +) -> list[dict]: + """Build blocked word entries for unfavorable competitor comparisons.""" + result = [] + for comp in competitors: + for name in all_names[comp]: + result.append({ + "keyword": f"{name} is better", + "action": "BLOCK", + "description": f"Unfavorable comparison ({comp})", + }) + + # Brand-level comparisons (only need one entry each, not per-competitor) + result.append({ + "keyword": f"better than {brand_name}", + "action": "BLOCK", + "description": "Unfavorable comparison", + }) + result.append({ + "keyword": f"{brand_name} is worse", + "action": "BLOCK", + "description": "Unfavorable comparison", + }) + + return result + + class SuggestTemplatesRequest(BaseModel): attack_examples: List[str] = Field(default_factory=list) description: str = Field(default="") diff --git a/litellm/proxy/management_endpoints/project_endpoints.py b/litellm/proxy/management_endpoints/project_endpoints.py new file mode 100644 index 00000000000..ba3238ebfd5 --- /dev/null +++ b/litellm/proxy/management_endpoints/project_endpoints.py @@ -0,0 +1,896 @@ +""" +Endpoints for /project operations + +/project/new +/project/update +/project/delete +/project/info +/project/list +""" + +#### PROJECT MANAGEMENT #### + +import json +from typing import List, Optional, Union + +from fastapi import APIRouter, Depends, HTTPException, Request + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.proxy._types import * +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field +from litellm.proxy.management_helpers.utils import ( + management_endpoint_wrapper, +) +from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy + +router = APIRouter() + + +async def _check_user_permission_for_project( + user_api_key_dict: UserAPIKeyAuth, + team_id: Optional[str], + prisma_client: PrismaClient, + require_admin: bool = False, + team_object: Optional[LiteLLM_TeamTable] = None, +) -> bool: + """ + Check if user has permission to manage a project. + + Returns True if user is proxy admin or team admin (when team_id provided). + If require_admin=True, only proxy admins are allowed. + + If team_object is provided, it will be used instead of fetching from DB + (avoids duplicate DB queries when team was already fetched for validation). + """ + is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + + if require_admin: + return is_proxy_admin + + if is_proxy_admin: + return True + + if not team_id or not user_api_key_dict.user_id: + return False + + team = team_object + if team is None: + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + + if team and team.admins: + return user_api_key_dict.user_id in team.admins + + return False + + +async def _validate_team_exists( + team_id: str, + prisma_client: PrismaClient, +): + """Validate that a team exists. Returns the team row.""" + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + ) + + if team is None: + raise ProxyException( + message=f"Team not found, team_id={team_id}", + type="not_found", + code=404, + param="team_id", + ) + + return team + + +def _check_team_project_limits( + team_object: LiteLLM_TeamTable, + data: Union[NewProjectRequest, UpdateProjectRequest], +) -> None: + """ + Check that project limits respect its parent Team's limits. + + Mirrors _check_org_team_limits() from team_endpoints.py. + + Validates: + - Project models are a subset of Team models + - Project max_budget <= Team max_budget + - Project tpm_limit <= Team tpm_limit + - Project rpm_limit <= Team rpm_limit + - Budget values are non-negative + - soft_budget < max_budget + """ + # --- Budget non-negativity checks --- + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, + ) + if data.soft_budget is not None and data.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, + ) + + # --- soft_budget < max_budget --- + if data.soft_budget is not None and data.max_budget is not None: + if data.soft_budget >= data.max_budget: + raise HTTPException( + status_code=400, + detail={ + "error": f"soft_budget ({data.soft_budget}) must be strictly lower than max_budget ({data.max_budget})" + }, + ) + + # --- Validate project models are a subset of team models --- + project_models = getattr(data, "models", None) + team_models = team_object.models or [] + if project_models and len(team_models) > 0: + # If team has 'all-proxy-models', skip validation as it allows all models + if SpecialModelNames.all_proxy_models.value not in team_models: + for m in project_models: + if m not in team_models: + raise HTTPException( + status_code=400, + detail={ + "error": f"Model '{m}' not in team's allowed models. Team allowed models={team_models}. Team: {team_object.team_id}" + }, + ) + + # --- Validate project max_budget <= team max_budget --- + # Team stores budget fields directly (max_budget, tpm_limit, rpm_limit) + # unlike Project which uses a separate LiteLLM_BudgetTable relation + if ( + data.max_budget is not None + and team_object.max_budget is not None + and data.max_budget > team_object.max_budget + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Project max_budget ({data.max_budget}) exceeds team's max_budget ({team_object.max_budget}). Team: {team_object.team_id}" + }, + ) + + # --- Validate project tpm_limit <= team tpm_limit --- + if ( + data.tpm_limit is not None + and team_object.tpm_limit is not None + and data.tpm_limit > team_object.tpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Project tpm_limit ({data.tpm_limit}) exceeds team's tpm_limit ({team_object.tpm_limit}). Team: {team_object.team_id}" + }, + ) + + # --- Validate project rpm_limit <= team rpm_limit --- + if ( + data.rpm_limit is not None + and team_object.rpm_limit is not None + and data.rpm_limit > team_object.rpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Project rpm_limit ({data.rpm_limit}) exceeds team's rpm_limit ({team_object.rpm_limit}). Team: {team_object.team_id}" + }, + ) + + +async def _create_budget_for_project( + data: NewProjectRequest, + user_id: Optional[str], + litellm_proxy_admin_name: str, + prisma_client: PrismaClient, +) -> str: + """Create a budget for the project and return budget_id.""" + budget_params = LiteLLM_BudgetTable.model_fields.keys() + _json_data = data.json(exclude_none=True) + _budget_data = {k: v for k, v in _json_data.items() if k in budget_params} + budget_row = LiteLLM_BudgetTable(**_budget_data) + + new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + + _budget = await prisma_client.db.litellm_budgettable.create( + data={ + **new_budget, + "created_by": user_id or litellm_proxy_admin_name, + "updated_by": user_id or litellm_proxy_admin_name, + } + ) + + return _budget.budget_id + + +async def _set_project_object_permission( + data: NewProjectRequest, + prisma_client: Optional[PrismaClient], +) -> Optional[str]: + """ + Creates the LiteLLM_ObjectPermissionTable record for the project. + Returns the object_permission_id if created, otherwise None. + """ + if prisma_client is None: + return None + + if data.object_permission is not None: + created_object_permission = ( + await prisma_client.db.litellm_objectpermissiontable.create( + data=data.object_permission.model_dump(exclude_none=True), + ) + ) + del data.object_permission + return created_object_permission.object_permission_id + return None + + +def _remove_budget_fields_from_project_data(project_data: dict) -> dict: + """ + Remove budget fields from project data. + Budget fields belong to LiteLLM_BudgetTable, not LiteLLM_ProjectTable. + Keep budget_id as it's a foreign key. + + Following the pattern from organization_endpoints.py + """ + budget_fields = LiteLLM_BudgetTable.model_fields.keys() + for field in list(budget_fields): + if field != "budget_id": # Keep the foreign key + project_data.pop(field, None) + return project_data + + +@router.post( + "/project/new", + tags=["project management"], + dependencies=[Depends(user_api_key_auth)], + response_model=NewProjectResponse, +) +@management_endpoint_wrapper +async def new_project( + data: NewProjectRequest, + http_request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Create a new project. Projects sit between teams and keys in the hierarchy. + + Only admins or team admins can create projects. + + # Parameters + + - project_alias: *Optional[str]* - The name of the project. + - description: *Optional[str]* - Description of the project's purpose and use case. + - team_id: *str* - The team id that this project belongs to. Required. + - models: *List* - The models the project has access to. + - budget_id: *Optional[str]* - The id for a budget (tpm/rpm/max budget) for the project. + ### IF NO BUDGET ID - CREATE ONE WITH THESE PARAMS ### + - max_budget: *Optional[float]* - Max budget for project + - tpm_limit: *Optional[int]* - Max tpm limit for project + - rpm_limit: *Optional[int]* - Max rpm limit for project + - max_parallel_requests: *Optional[int]* - Max parallel requests for project + - soft_budget: *Optional[float]* - Get a slack alert when this soft budget is reached. Don't block requests. + - model_max_budget: *Optional[dict]* - Max budget for a specific model. Example: {"gpt-4": 100.0, "gpt-3.5-turbo": 50.0} + - model_rpm_limit: *Optional[dict]* - RPM limits per model. Example: {"gpt-4": 1000, "gpt-3.5-turbo": 5000} + - model_tpm_limit: *Optional[dict]* - TPM limits per model. Example: {"gpt-4": 50000, "gpt-3.5-turbo": 100000} + - budget_duration: *Optional[str]* - Frequency of reseting project budget + - metadata: *Optional[dict]* - Metadata for project, store information for project. Example metadata - {"use_case_id": "SNOW-12345", "responsible_ai_id": "RAI-67890"} + - blocked: *bool* - Flag indicating if the project is blocked or not - will stop all calls from keys with this project_id. + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - project-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission. + + Example 1: Create new project **without** a budget_id, with model-specific limits + + ```bash + curl --location 'http://0.0.0.0:4000/project/new' \\ + --header 'Authorization: Bearer sk-1234' \\ + --header 'Content-Type: application/json' \\ + --data '{ + "project_alias": "flight-search-assistant", + "description": "AI-powered flight search and booking assistant", + "team_id": "team-123", + "models": ["gpt-4", "gpt-3.5-turbo"], + "max_budget": 100, + "model_rpm_limit": { + "gpt-4": 1000, + "gpt-3.5-turbo": 5000 + }, + "model_tpm_limit": { + "gpt-4": 50000, + "gpt-3.5-turbo": 100000 + }, + "metadata": { + "use_case_id": "SNOW-12345", + "responsible_ai_id": "RAI-67890" + } + }' + ``` + + Example 2: Create new project **with** a budget_id + + ```bash + curl --location 'http://0.0.0.0:4000/project/new' \\ + --header 'Authorization: Bearer sk-1234' \\ + --header 'Content-Type: application/json' \\ + --data '{ + "project_alias": "hotel-recommendations", + "description": "Personalized hotel recommendation engine", + "team_id": "team-123", + "models": ["claude-3-sonnet"], + "budget_id": "428eeaa8-f3ac-4e85-a8fb-7dc8d7aa8689", + "metadata": { + "use_case_id": "SNOW-54321" + } + }' + ``` + """ + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + premium_user, + prisma_client, + ) + + try: + if not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Project management is an enterprise feature. " + + CommonProxyErrors.not_premium_user.value + }, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Validate team exists and get team object with budget + team_object = await _validate_team_exists( + team_id=data.team_id, prisma_client=prisma_client + ) + + # Validate project limits against team limits + _check_team_project_limits( + team_object=LiteLLM_TeamTable(**team_object.model_dump()), + data=data, + ) + + # Check if user has permission to create projects for this team + # only team admins can create projects for their team + has_permission = await _check_user_permission_for_project( + user_api_key_dict=user_api_key_dict, + team_id=data.team_id, + prisma_client=prisma_client, + team_object=LiteLLM_TeamTable(**team_object.model_dump()), + ) + + if not has_permission: + raise HTTPException( + status_code=403, + detail={ + "error": f"Only admins or team admins can create projects. Your role is {user_api_key_dict.user_role}" + }, + ) + + # Generate project_id if not provided + if data.project_id is None: + data.project_id = str(uuid.uuid4()) + else: + # Check if project_id already exists + existing_project = await prisma_client.db.litellm_projecttable.find_unique( + where={"project_id": data.project_id} + ) + if existing_project is not None: + raise ProxyException( + message=f"Project id = {data.project_id} already exists. Please use a different project id.", + type="bad_request", + code=400, + param="project_id", + ) + + # Create budget if not provided + if data.budget_id is None: + data.budget_id = await _create_budget_for_project( + data=data, + user_id=user_api_key_dict.user_id, + litellm_proxy_admin_name=litellm_proxy_admin_name, + prisma_client=prisma_client, + ) + + ## Handle Object Permission - MCP, Vector Stores etc. + object_permission_id = await _set_project_object_permission( + data=data, + prisma_client=prisma_client, + ) + + # Create project row (following organization_endpoints.py pattern) + project_row = LiteLLM_ProjectTable( + **data.json(exclude_none=True), + object_permission_id=object_permission_id, + created_by=user_api_key_dict.user_id or litellm_proxy_admin_name, + updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, + ) + + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=project_row, + field_name=field, + value=getattr(data, field), + ) + + new_project_row = prisma_client.jsonify_object( + project_row.json(exclude_none=True) + ) + + # Remove budget fields (following organization_endpoints.py pattern) + new_project_row = _remove_budget_fields_from_project_data(new_project_row) + + verbose_proxy_logger.info( + f"new_project_row: {json.dumps(new_project_row, indent=2)}" + ) + response = await prisma_client.db.litellm_projecttable.create( + data={ + **new_project_row, # type: ignore + }, + include={"litellm_budget_table": True}, + ) + + return response + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.management_endpoints.project_endpoints.new_project(): Exception occured - {}".format( + str(e) + ) + ) + raise handle_exception_on_proxy(e) + + +@router.post( + "/project/update", + tags=["project management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_ProjectTable, +) +@management_endpoint_wrapper +async def update_project( + data: UpdateProjectRequest, + http_request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update a project + + Parameters: + - project_id: *str* - The project id to update. Required. + - project_alias: *Optional[str]* - Updated name for the project + - description: *Optional[str]* - Updated description for the project + - team_id: *Optional[str]* - Updated team_id for the project + - metadata: *Optional[dict]* - Updated metadata for project + - models: *Optional[list]* - Updated list of models for the project + - blocked: *Optional[bool]* - Updated blocked status + - max_budget: *Optional[float]* - Updated max budget + - tpm_limit: *Optional[int]* - Updated tpm limit + - rpm_limit: *Optional[int]* - Updated rpm limit + - model_rpm_limit: *Optional[dict]* - Updated RPM limits per model + - model_tpm_limit: *Optional[dict]* - Updated TPM limits per model + - budget_duration: *Optional[str]* - Updated budget duration + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - Updated object permission + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/project/update' \\ + --header 'Authorization: Bearer sk-1234' \\ + --header 'Content-Type: application/json' \\ + --data '{ + "project_id": "project-123", + "description": "Updated flight search system with enhanced capabilities", + "max_budget": 200, + "model_rpm_limit": { + "gpt-4": 2000, + "gpt-3.5-turbo": 10000 + }, + "metadata": { + "use_case_id": "SNOW-12345", + "status": "active" + } + }' + ``` + """ + from litellm.proxy.proxy_server import ( + litellm_proxy_admin_name, + premium_user, + prisma_client, + ) + + try: + if not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Project management is an enterprise feature. " + + CommonProxyErrors.not_premium_user.value + }, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if data.project_id is None: + raise HTTPException( + status_code=400, + detail={"error": "project_id is required"}, + ) + + # Fetch existing project + existing_project = await prisma_client.db.litellm_projecttable.find_unique( + where={"project_id": data.project_id} + ) + + if existing_project is None: + raise ProxyException( + message=f"Project not found, project_id={data.project_id}", + type="not_found", + code=404, + param="project_id", + ) + + # Validate team exists and get team object for limit + permission checks + team_id_to_check = data.team_id or existing_project.team_id + team_obj_for_checks = None + if team_id_to_check is not None: + team_obj_for_checks = await _validate_team_exists( + team_id=team_id_to_check, prisma_client=prisma_client + ) + + # Check if user has permission to update this project + has_permission = await _check_user_permission_for_project( + user_api_key_dict=user_api_key_dict, + team_id=existing_project.team_id, + prisma_client=prisma_client, + team_object=LiteLLM_TeamTable(**team_obj_for_checks.model_dump()) + if team_obj_for_checks + else None, + ) + + if not has_permission: + raise HTTPException( + status_code=403, + detail={"error": "Only admins or team admins can update projects"}, + ) + + # Validate project limits against team limits + if team_obj_for_checks is not None: + _check_team_project_limits( + team_object=LiteLLM_TeamTable(**team_obj_for_checks.model_dump()), + data=data, + ) + + # Prepare update data + update_data = data.json(exclude_none=True, exclude={"project_id"}) + update_data = prisma_client.jsonify_object(update_data) + update_data["updated_by"] = ( + user_api_key_dict.user_id or litellm_proxy_admin_name + ) + + # Handle budget updates + budget_fields = LiteLLM_BudgetTable.model_fields.keys() + budget_updates = {k: v for k, v in update_data.items() if k in budget_fields} + + if budget_updates and existing_project.budget_id: + # Update existing budget + await prisma_client.db.litellm_budgettable.update( + where={"budget_id": existing_project.budget_id}, + data={ + **budget_updates, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + }, + ) + # Remove budget fields from project update + for field in budget_updates.keys(): + update_data.pop(field, None) + + # Handle object permissions + if "object_permission" in update_data: + object_permission_data = update_data.pop("object_permission") + if object_permission_data: + if existing_project.object_permission_id: + # Update existing permission + await prisma_client.db.litellm_objectpermissiontable.update( + where={ + "object_permission_id": existing_project.object_permission_id + }, + data=object_permission_data, + ) + else: + # Create new permission + created_permission = ( + await prisma_client.db.litellm_objectpermissiontable.create( + data=object_permission_data, + ) + ) + update_data[ + "object_permission_id" + ] = created_permission.object_permission_id + + # Handle metadata fields + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if field in update_data: + if update_data.get("metadata") is None: + update_data["metadata"] = {} + update_data["metadata"][field] = update_data.pop(field) + + # Remove budget fields (following organization_endpoints.py pattern) + update_data = _remove_budget_fields_from_project_data(update_data) + + # Update project + updated_project = await prisma_client.db.litellm_projecttable.update( + where={"project_id": data.project_id}, + data=update_data, + include={"litellm_budget_table": True, "object_permission": True}, + ) + + return updated_project + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.management_endpoints.project_endpoints.update_project(): Exception occured - {}".format( + str(e) + ) + ) + raise handle_exception_on_proxy(e) + + +@router.delete( + "/project/delete", + tags=["project management"], + dependencies=[Depends(user_api_key_auth)], + response_model=List[LiteLLM_ProjectTable], +) +@management_endpoint_wrapper +async def delete_project( + data: DeleteProjectRequest, + http_request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete projects + + Parameters: + - project_ids: *List[str]* - List of project ids to delete + + Example: + ```bash + curl --location --request DELETE 'http://0.0.0.0:4000/project/delete' \\ + --header 'Authorization: Bearer sk-1234' \\ + --header 'Content-Type: application/json' \\ + --data '{ + "project_ids": ["project-123", "project-456"] + }' + ``` + """ + from litellm.proxy.proxy_server import premium_user, prisma_client + + try: + if not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Project management is an enterprise feature. " + + CommonProxyErrors.not_premium_user.value + }, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Check if user is admin (only admins can delete projects) + has_permission = await _check_user_permission_for_project( + user_api_key_dict=user_api_key_dict, + team_id=None, + prisma_client=prisma_client, + require_admin=True, + ) + + if not has_permission: + raise HTTPException( + status_code=403, + detail={"error": "Only admins can delete projects"}, + ) + + deleted_projects = [] + + for project_id in data.project_ids: + # Check if project exists + existing_project = await prisma_client.db.litellm_projecttable.find_unique( + where={"project_id": project_id} + ) + + if existing_project is None: + raise ProxyException( + message=f"Project not found, project_id={project_id}", + type="not_found", + code=404, + param="project_ids", + ) + + # Check if there are any keys associated with this project + associated_keys = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"project_id": project_id} + ) + ) + + if len(associated_keys) > 0: + raise ProxyException( + message=f"Cannot delete project {project_id}. {len(associated_keys)} key(s) are associated with it. Please delete or reassign the keys first.", + type="bad_request", + code=400, + param="project_ids", + ) + + # Delete the project + deleted_project = await prisma_client.db.litellm_projecttable.delete( + where={"project_id": project_id} + ) + + deleted_projects.append(deleted_project) + + return deleted_projects + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.management_endpoints.project_endpoints.delete_project(): Exception occured - {}".format( + str(e) + ) + ) + raise handle_exception_on_proxy(e) + + +@router.get( + "/project/info", + tags=["project management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_ProjectTable, +) +async def project_info( + project_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get information about a specific project + + Parameters: + - project_id: *str* - The project id to fetch info for + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/project/info?project_id=project-123' \\ + --header 'Authorization: Bearer sk-1234' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + try: + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Fetch project + project = await prisma_client.db.litellm_projecttable.find_unique( + where={"project_id": project_id}, + include={"litellm_budget_table": True, "object_permission": True}, + ) + + if project is None: + raise ProxyException( + message=f"Project not found, project_id={project_id}", + type="not_found", + code=404, + param="project_id", + ) + + # Check if user has access to this project (admin or team member) + is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + is_team_member = False + + if project.team_id and user_api_key_dict.user_id: + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": project.team_id} + ) + if team: + is_team_member = ( + user_api_key_dict.user_id in team.admins + or user_api_key_dict.user_id in team.members + ) + + if not (is_admin or is_team_member): + raise HTTPException( + status_code=403, + detail={"error": "You don't have access to this project"}, + ) + + return project + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.management_endpoints.project_endpoints.project_info(): Exception occured - {}".format( + str(e) + ) + ) + raise handle_exception_on_proxy(e) + + +@router.get( + "/project/list", + tags=["project management"], + dependencies=[Depends(user_api_key_auth)], + response_model=List[LiteLLM_ProjectTable], +) +async def list_projects( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + List all projects that the user has access to + + Example: + ```bash + curl --location 'http://0.0.0.0:4000/project/list' \\ + --header 'Authorization: Bearer sk-1234' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + try: + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + # If proxy admin, get all projects + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + projects = await prisma_client.db.litellm_projecttable.find_many( + include={"litellm_budget_table": True, "object_permission": True} + ) + else: + # Get projects for teams the user belongs to + user_teams = await prisma_client.db.litellm_teamtable.find_many( + where={ + "OR": [ + {"members": {"has": user_api_key_dict.user_id}}, + {"admins": {"has": user_api_key_dict.user_id}}, + ] + } + ) + + team_ids = [team.team_id for team in user_teams] + + projects = await prisma_client.db.litellm_projecttable.find_many( + where={"team_id": {"in": team_ids}}, + include={"litellm_budget_table": True, "object_permission": True}, + ) + + return projects + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.management_endpoints.project_endpoints.list_projects(): Exception occured - {}".format( + str(e) + ) + ) + raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f5aabdfcedb..80d50f31a17 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -72,6 +72,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _set_object_metadata_field, + _team_member_has_permission, _update_metadata_fields, _upsert_budget_and_membership, _user_has_admin_view, @@ -3971,22 +3972,29 @@ async def get_team_daily_activity( t.team_id: {"team_alias": t.team_alias} for t in team_aliases } - # Check if user is team admin for any requested teams + # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys user_api_keys: Optional[List[str]] = None if not _user_has_admin_view(user_api_key_dict) and team_ids_list and team_aliases: - # Check if user is team admin for any of the teams - is_team_admin_for_any = False + # Check if user is team admin or has usage view permission for any team + has_full_team_view = False for team_alias in team_aliases: team_obj = LiteLLM_TeamTable(**team_alias.model_dump()) if _is_user_team_admin( user_api_key_dict=user_api_key_dict, team_obj=team_obj ): - is_team_admin_for_any = True + has_full_team_view = True + break + if _team_member_has_permission( + user_api_key_dict=user_api_key_dict, + team_obj=team_obj, + permission="/team/daily/activity", + ): + has_full_team_view = True break - # If user is not a team admin for any team, filter by their API keys - if not is_team_admin_for_any: + # If user does not have full team view, filter by their API keys + if not has_full_team_view: # Get all API keys for this user user_keys = await prisma_client.db.litellm_verificationtoken.find_many( where={"user_id": user_api_key_dict.user_id} diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 56b513554a8..d1c4b8b3403 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1,2744 +1,2820 @@ -import ast -import asyncio -import copy -import json -import traceback -from base64 import b64encode -from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple, Union, cast -from urllib.parse import urlencode, urlparse - -import httpx -from fastapi import ( - APIRouter, - Depends, - FastAPI, - HTTPException, - Request, - Response, - UploadFile, - WebSocket, - status, -) -from fastapi.responses import StreamingResponse -from starlette.datastructures import UploadFile as StarletteUploadFile -from starlette.websockets import WebSocketState -from websockets.asyncio.client import connect -from websockets.exceptions import ( - ConnectionClosedError, - ConnectionClosedOK, - InvalidStatus, -) - -import litellm -from litellm._logging import verbose_proxy_logger -from litellm._uuid import uuid -from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.passthrough import BasePassthroughUtils -from litellm.proxy._types import ( - ConfigFieldInfo, - ConfigFieldUpdate, - LiteLLMRoutes, - PassThroughEndpointResponse, - PassThroughGenericEndpoint, - ProxyException, - UserAPIKeyAuth, -) -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from litellm.proxy.common_utils.http_parsing_utils import _read_request_body -from litellm.proxy.utils import get_server_root_path -from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.custom_http import httpxSpecialProvider -from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - EndpointType, - PassthroughStandardLoggingPayload, -) -from litellm.types.utils import StandardLoggingUserAPIKeyMetadata - -from .streaming_handler import PassThroughStreamingHandler -from .success_handler import PassThroughEndpointLogging - -router = APIRouter() - -pass_through_endpoint_logging = PassThroughEndpointLogging() - -# Global registry to track registered pass-through routes and prevent memory leaks -_registered_pass_through_routes: Dict[str, Dict[str, Union[str, Dict[str, Any]]]] = {} - - -def get_response_body(response: httpx.Response) -> Optional[dict]: - try: - return response.json() - except Exception: - return None - - -async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optional[dict]: - """ - checks if any headers on config.yaml are defined as os.environ/COHERE_API_KEY etc - - only runs for headers defined on config.yaml - - example header can be - - {"Authorization": "Bearer os.environ/COHERE_API_KEY"} - """ - if custom_headers is None: - return None - headers = {} - for key, value in custom_headers.items(): - # langfuse Api requires base64 encoded headers - it's simpleer to just ask litellm users to set their langfuse public and secret keys - # we can then get the b64 encoded keys here - if key == "LANGFUSE_PUBLIC_KEY" or key == "LANGFUSE_SECRET_KEY": - # 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/"): - _langfuse_public_key = get_secret_str(_langfuse_public_key) - 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") - ).decode("ascii") - else: - # 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" - ) - # get string section that is os.environ/ - start_index = value.find("os.environ/") - _variable_name = value[start_index:] - - verbose_proxy_logger.debug( - "pass through endpoint - getting secret for variable name: %s", - _variable_name, - ) - _secret_value = get_secret_str(_variable_name) - if _secret_value is not None: - new_value = value.replace(_variable_name, _secret_value) - headers[key] = new_value - return headers - - -async def chat_completion_pass_through_endpoint( # noqa: PLR0915 - fastapi_response: Response, - request: Request, - adapter_id: str, - user_api_key_dict: UserAPIKeyAuth, -): - from litellm.proxy.proxy_server import ( - add_litellm_data_to_request, - general_settings, - llm_router, - proxy_config, - proxy_logging_obj, - user_api_base, - user_max_tokens, - user_model, - user_request_timeout, - user_temperature, - version, - ) - - data = {} - try: - body = await request.body() - body_str = body.decode() - try: - data = ast.literal_eval(body_str) - except Exception: - data = json.loads(body_str) - - data["adapter_id"] = adapter_id - - verbose_proxy_logger.debug( - "Request received by LiteLLM:\n{}".format(json.dumps(data, indent=4)), - ) - data["model"] = ( - general_settings.get("completion_model", None) # server default - or user_model # model name passed via cli args - or data.get("model", None) # default passed in http request - ) - if user_model: - data["model"] = user_model - - data = await add_litellm_data_to_request( - data=data, # type: ignore - request=request, - general_settings=general_settings, - user_api_key_dict=user_api_key_dict, - version=version, - proxy_config=proxy_config, - ) - - # override with user settings, these are params passed via cli - if user_temperature: - data["temperature"] = user_temperature - if user_request_timeout: - data["request_timeout"] = user_request_timeout - if user_max_tokens: - data["max_tokens"] = user_max_tokens - if user_api_base: - data["api_base"] = user_api_base - - ### MODEL ALIAS MAPPING ### - # check if model name in model alias map - # get the actual model name - if data["model"] in litellm.model_alias_map: - data["model"] = litellm.model_alias_map[data["model"]] - - # Check key-specific aliases - if ( - isinstance(data["model"], str) - and user_api_key_dict.aliases - and isinstance(user_api_key_dict.aliases, dict) - and data["model"] in user_api_key_dict.aliases - ): - data["model"] = user_api_key_dict.aliases[data["model"]] - - ### CALL HOOKS ### - modify incoming data before calling the model - data = await proxy_logging_obj.pre_call_hook( # type: ignore - user_api_key_dict=user_api_key_dict, data=data, call_type="text_completion" - ) - - ### ROUTE THE REQUESTs ### - router_model_names = llm_router.model_names if llm_router is not None else [] - # 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 - llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) - elif ( - llm_router is not None - and llm_router.model_group_alias is not None - and data["model"] in llm_router.model_group_alias - ): # model set in model_group_alias - llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) - elif llm_router is not None and llm_router.has_model_id( - data["model"] - ): # model in router model list - llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) - elif ( - llm_router is not None - and data["model"] not in router_model_names - and ( - llm_router.default_deployment is not None - or len(llm_router.pattern_router.patterns) > 0 - ) - ): # check for wildcard routes or default deployment before checking deployment_names - llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) - elif ( - llm_router is not None and data["model"] in llm_router.deployment_names - ): # model in router deployments, calling a specific deployment on the router (lowest priority) - llm_response = asyncio.create_task( - llm_router.aadapter_completion(**data, specific_deployment=True) - ) - elif user_model is not None: # `litellm --model ` - llm_response = asyncio.create_task(litellm.aadapter_completion(**data)) - else: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "completion: Invalid model name passed in model=" - + data.get("model", "") - }, - ) - - # Await the llm_response task - response = await llm_response - - hidden_params = getattr(response, "_hidden_params", {}) or {} - model_id = hidden_params.get("model_id", None) or "" - cache_key = hidden_params.get("cache_key", None) or "" - api_base = hidden_params.get("api_base", None) or "" - response_cost = hidden_params.get("response_cost", None) or "" - - ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) - ) - - verbose_proxy_logger.debug("final response: %s", response) - - fastapi_response.headers.update( - ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - ) - ) - - verbose_proxy_logger.debug("\nResponse from Litellm:\n{}".format(response)) - return response - except Exception as e: - 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) - ) - ) - error_msg = f"{str(e)}" - raise ProxyException( - message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), - ) - - -class HttpPassThroughEndpointHelpers(BasePassthroughUtils): - @staticmethod - def get_response_headers( - headers: httpx.Headers, - litellm_call_id: Optional[str] = None, - custom_headers: Optional[dict] = None, - ) -> dict: - excluded_headers = {"transfer-encoding", "content-encoding"} - - 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: - return_headers.update(custom_headers) - - return return_headers - - @staticmethod - def get_endpoint_type(url: str) -> EndpointType: - parsed_url = urlparse(url) - if ( - ("generateContent") in url - or ("streamGenerateContent") in url - or ("rawPredict") in url - or ("streamRawPredict") in url - ): - return EndpointType.VERTEX_AI - elif parsed_url.hostname == "api.anthropic.com": - return EndpointType.ANTHROPIC - elif ( - parsed_url.hostname == "api.openai.com" - or parsed_url.hostname == "openai.azure.com" - or (parsed_url.hostname and "openai.com" in parsed_url.hostname) - ): - return EndpointType.OPENAI - return EndpointType.GENERIC - - @staticmethod - async def _make_non_streaming_http_request( - request: Request, - async_client: httpx.AsyncClient, - url: str, - headers: dict, - requested_query_params: Optional[dict] = None, - custom_body: Optional[dict] = None, - ) -> httpx.Response: - """ - Make a non-streaming HTTP request - - If request is GET, don't include a JSON body - """ - if request.method == "GET": - response = await async_client.request( - method=request.method, - url=url, - headers=headers, - params=requested_query_params, - ) - else: - response = await async_client.request( - method=request.method, - url=url, - headers=headers, - params=requested_query_params, - json=custom_body, - ) - return response - - @staticmethod - async def non_streaming_http_request_handler( - request: Request, - async_client: httpx.AsyncClient, - url: httpx.URL, - headers: dict, - requested_query_params: Optional[dict] = None, - _parsed_body: Optional[dict] = None, - ) -> httpx.Response: - """ - Handle non-streaming HTTP requests - - Handles special cases when GET requests, multipart/form-data requests, and generic httpx requests - """ - if request.method == "GET": - response = await async_client.request( - method=request.method, - url=url, - headers=headers, - params=requested_query_params, - ) - elif HttpPassThroughEndpointHelpers.is_multipart(request) is True: - return await HttpPassThroughEndpointHelpers.make_multipart_http_request( - request=request, - async_client=async_client, - url=url, - headers=headers, - requested_query_params=requested_query_params, - ) - else: - # Generic httpx method - response = await async_client.request( - method=request.method, - url=url, - headers=headers, - params=requested_query_params, - json=_parsed_body, - ) - return response - - @staticmethod - def is_multipart(request: Request) -> bool: - """Check if the request is a multipart/form-data request""" - return "multipart/form-data" in request.headers.get("content-type", "") - - @staticmethod - async def _build_request_files_from_upload_file( - upload_file: Union[UploadFile, StarletteUploadFile], - ) -> Tuple[Optional[str], bytes, Optional[str]]: - """Build a request files dict from an UploadFile object""" - file_content = await upload_file.read() - return (upload_file.filename, file_content, upload_file.content_type) - - @staticmethod - async def make_multipart_http_request( - request: Request, - async_client: httpx.AsyncClient, - url: httpx.URL, - headers: dict, - requested_query_params: Optional[dict] = None, - ) -> httpx.Response: - """Process multipart/form-data requests, handling both files and form fields""" - form_data = await request.form() - files = {} - form_data_dict = {} - - 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 - ) - else: - form_data_dict[field_name] = field_value - - # Remove content-type header - httpx will set it correctly with the new boundary - # when it creates the multipart body from files/data parameters - headers_copy = headers.copy() - headers_copy.pop("content-type", None) - - response = await async_client.request( - method=request.method, - url=url, - headers=headers_copy, - params=requested_query_params, - files=files, - data=form_data_dict, - ) - return response - - @staticmethod - def _init_kwargs_for_pass_through_endpoint( - request: Request, - user_api_key_dict: UserAPIKeyAuth, - passthrough_logging_payload: PassthroughStandardLoggingPayload, - logging_obj: LiteLLMLoggingObj, - _parsed_body: Optional[dict] = None, - litellm_call_id: Optional[str] = None, - ) -> dict: - """ - Filter out litellm params from the request body - """ - from litellm.types.utils import all_litellm_params - - _parsed_body = _parsed_body or {} - - litellm_params_in_body = {} - for k in all_litellm_params: - if k in _parsed_body: - litellm_params_in_body[k] = _parsed_body.pop(k, None) - - _metadata = dict( - StandardLoggingUserAPIKeyMetadata( - user_api_key_hash=user_api_key_dict.api_key, - user_api_key_alias=user_api_key_dict.key_alias, - user_api_key_user_email=user_api_key_dict.user_email, - user_api_key_user_id=user_api_key_dict.user_id, - user_api_key_team_id=user_api_key_dict.team_id, - user_api_key_org_id=user_api_key_dict.org_id, - user_api_key_team_alias=user_api_key_dict.team_alias, - user_api_key_end_user_id=user_api_key_dict.end_user_id, - user_api_key_request_route=user_api_key_dict.request_route, - 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_auth_metadata=user_api_key_dict.metadata, - ) - ) - - _metadata["user_api_key"] = user_api_key_dict.api_key - - litellm_metadata = litellm_params_in_body.pop("litellm_metadata", None) - metadata = litellm_params_in_body.pop("metadata", None) - if litellm_metadata: - _metadata.update(litellm_metadata) - if metadata: - _metadata.update(metadata) - - _metadata = _update_metadata_with_tags_in_header( - request=request, - metadata=_metadata, - ) - - kwargs = { - "litellm_params": { - **litellm_params_in_body, # type: ignore - "metadata": _metadata, - "proxy_server_request": { - "url": str(request.url), - "method": request.method, - "body": copy.copy(_parsed_body), # use copy instead of deepcopy - "headers": request.headers, - }, - }, - "call_type": "pass_through_endpoint", - "litellm_call_id": litellm_call_id, - "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: - """ - Helper function to construct the full target URL with subpath handling. - - Args: - base_target: The base target URL - subpath: The captured subpath from the request - include_subpath: Whether to include the subpath in the target URL - - Returns: - The constructed full target URL - """ - if not include_subpath: - return base_target - - if not subpath: - return base_target - - # Ensure base_target ends with / and subpath doesn't start with / - if not base_target.endswith("/"): - base_target = base_target + "/" - if subpath.startswith("/"): - subpath = subpath[1:] - - return base_target + subpath - - @staticmethod - def _update_stream_param_based_on_request_body( - parsed_body: dict, - stream: Optional[bool] = None, - ) -> Optional[bool]: - """ - If stream is provided in the request body, use it. - Otherwise, use the stream parameter passed to the `pass_through_request` function - """ - if "stream" in parsed_body: - return parsed_body.get("stream", stream) - return stream - - -async def pass_through_request( # noqa: PLR0915 - request: Request, - target: str, - custom_headers: dict, - user_api_key_dict: UserAPIKeyAuth, - custom_body: Optional[dict] = None, - forward_headers: Optional[bool] = False, - merge_query_params: Optional[bool] = False, - query_params: Optional[dict] = None, - stream: Optional[bool] = None, - cost_per_request: Optional[float] = None, - custom_llm_provider: Optional[str] = None, - guardrails_config: Optional[dict] = None, -): - """ - Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called - - Args: - request: The incoming request - target: The target URL - custom_headers: The custom headers - user_api_key_dict: The user API key dictionary - custom_body: The custom body - forward_headers: Whether to forward headers - merge_query_params: Whether to merge query params - query_params: The query params - stream: Whether to stream the response - cost_per_request: Optional field - cost per request to the target endpoint - custom_llm_provider: Optional field - custom LLM provider for the endpoint - guardrails_config: Optional field - guardrails configuration for passthrough endpoint - """ - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy.pass_through_endpoints.passthrough_guardrails import ( - PassthroughGuardrailHandler, - ) - from litellm.proxy.proxy_server import proxy_logging_obj - - ######################################################### - # Initialize variables - ######################################################### - litellm_call_id = str(uuid.uuid4()) - url: Optional[httpx.URL] = None - - # parsed request body - _parsed_body: Optional[dict] = None - # kwargs for pass through endpoint, contains metadata, litellm_params, call_type, litellm_call_id, passthrough_logging_payload - kwargs: Optional[dict] = None - - ######################################################### - try: - url = httpx.URL(target) - headers = custom_headers - headers = HttpPassThroughEndpointHelpers.forward_headers_from_request( - request_headers=dict(request.headers), - headers=headers, - forward_headers=forward_headers, - ) - - if merge_query_params: - # Create a new URL with the merged query params - url = url.copy_with( - query=urlencode( - HttpPassThroughEndpointHelpers.get_merged_query_parameters( - existing_url=url, - request_query_params=dict(request.query_params), - ) - ).encode("ascii") - ) - - endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type( - str(url) - ) - - if custom_body: - _parsed_body = custom_body - else: - _parsed_body = await _read_request_body(request) - verbose_proxy_logger.debug( - "Pass through endpoint sending request to \nURL {}\nheaders: {}\nbody: {}\n".format( - url, headers, _parsed_body - ) - ) - - ### COLLECT GUARDRAILS FOR PASSTHROUGH ENDPOINT ### - # Passthrough endpoints are opt-in only for guardrails - # When enabled, collect guardrails from org/team/key levels + passthrough-specific - guardrails_to_run = PassthroughGuardrailHandler.collect_guardrails( - user_api_key_dict=user_api_key_dict, - passthrough_guardrails_config=guardrails_config, - ) - - # Add guardrails to metadata if any should run - if guardrails_to_run and len(guardrails_to_run) > 0: - if _parsed_body is None: - _parsed_body = {} - if "metadata" not in _parsed_body: - _parsed_body["metadata"] = {} - _parsed_body["metadata"]["guardrails"] = guardrails_to_run - verbose_proxy_logger.debug( - f"Added guardrails to passthrough request metadata: {guardrails_to_run}" - ) - - ## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it - start_time = datetime.now() - logging_obj = Logging( - model="unknown", - messages=[{"role": "user", "content": safe_dumps(_parsed_body)}], - stream=False, - call_type="pass_through_endpoint", - start_time=start_time, - litellm_call_id=litellm_call_id, - function_id="1245", - ) - - # Store passthrough guardrails config on logging_obj for field targeting - logging_obj.passthrough_guardrails_config = guardrails_config - - # Store logging_obj in data so guardrails can access it - if _parsed_body is None: - _parsed_body = {} - _parsed_body["litellm_logging_obj"] = logging_obj - - ### CALL HOOKS ### - modify incoming data / reject request before calling the model - _parsed_body = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, - data=_parsed_body, - call_type="pass_through_endpoint", - ) - async_client_obj = get_async_httpx_client( - llm_provider=httpxSpecialProvider.PassThroughEndpoint, - params={"timeout": 600}, - ) - async_client = async_client_obj.client - passthrough_logging_payload = PassthroughStandardLoggingPayload( - url=str(url), - request_body=_parsed_body, - request_method=getattr(request, "method", None), - cost_per_request=cost_per_request, - ) - kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( - user_api_key_dict=user_api_key_dict, - _parsed_body=_parsed_body, - passthrough_logging_payload=passthrough_logging_payload, - litellm_call_id=litellm_call_id, - request=request, - logging_obj=logging_obj, - ) - - # Store custom_llm_provider in kwargs and logging object if provided - if custom_llm_provider: - logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider - logging_obj.model_call_details["litellm_params"] = kwargs.get( - "litellm_params", {} - ) - - # done for supporting 'parallel_request_limiter.py' with pass-through endpoints - logging_obj.update_environment_variables( - model="unknown", - user="unknown", - optional_params={}, - litellm_params=kwargs["litellm_params"], - call_type="pass_through_endpoint", - ) - logging_obj.model_call_details["litellm_call_id"] = litellm_call_id - - # combine url with query params for logging - requested_query_params: Optional[dict] = query_params or dict( - request.query_params - ) - - 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() - ) - - logging_url = str(url) - if requested_query_params_str: - if "?" in str(url): - logging_url = str(url) + "&" + requested_query_params_str - else: - logging_url = str(url) + "?" + requested_query_params_str - - logging_obj.pre_call( - input=[{"role": "user", "content": safe_dumps(_parsed_body)}], - api_key="", - additional_args={ - "complete_input_dict": _parsed_body, - "api_base": str(logging_url), - "headers": headers, - }, - ) - stream = ( - HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( - parsed_body=_parsed_body, - stream=stream, - ) - ) - - if stream: - req = async_client.build_request( - "POST", - url, - json=_parsed_body, - params=requested_query_params, - headers=headers, - ) - - response = await async_client.send(req, stream=stream) - - try: - response.raise_for_status() - except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=await e.response.aread() - ) - - return StreamingResponse( - PassThroughStreamingHandler.chunk_processor( - response=response, - request_body=_parsed_body, - litellm_logging_obj=logging_obj, - endpoint_type=endpoint_type, - start_time=start_time, - passthrough_success_handler_obj=pass_through_endpoint_logging, - url_route=str(url), - ), - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, - litellm_call_id=litellm_call_id, - ), - status_code=response.status_code, - ) - - 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) - - if _is_streaming_response(response) is True: - try: - response.raise_for_status() - except httpx.HTTPStatusError as e: - raise HTTPException( - status_code=e.response.status_code, detail=await e.response.aread() - ) - - return StreamingResponse( - PassThroughStreamingHandler.chunk_processor( - response=response, - request_body=_parsed_body, - litellm_logging_obj=logging_obj, - endpoint_type=endpoint_type, - start_time=start_time, - passthrough_success_handler_obj=pass_through_endpoint_logging, - url_route=str(url), - ), - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, - litellm_call_id=litellm_call_id, - ), - status_code=response.status_code, - ) - - try: - response.raise_for_status() - except httpx.HTTPStatusError as e: - 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) - - content = await response.aread() - - ## LOG SUCCESS - response_body: Optional[dict] = get_response_body(response) - passthrough_logging_payload["response_body"] = response_body - end_time = datetime.now() - asyncio.create_task( - pass_through_endpoint_logging.pass_through_async_success_handler( - httpx_response=response, - response_body=response_body, - url_route=str(url), - result="", - start_time=start_time, - end_time=end_time, - logging_obj=logging_obj, - cache_hit=False, - request_body=_parsed_body, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - ) - - ## CUSTOM HEADERS - `x-litellm-*` - custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=litellm_call_id, - model_id=None, - cache_key=None, - api_base=str(url._uri_reference), - ) - - return Response( - content=content, - status_code=response.status_code, - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, - custom_headers=custom_headers, - ), - ) - except Exception as e: - custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=litellm_call_id, - model_id=None, - cache_key=None, - 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) - ) - ) - - ######################################################### - # Monitoring: Trigger post_call_failure_hook - # for pass through endpoint failure - ######################################################### - request_payload: dict = _parsed_body or {} - # add user_api_key_dict, litellm_call_id, passthrough_logging_payloa for logging - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - - if ( - "model" not in request_payload - and _parsed_body - and isinstance(_parsed_body, dict) - ): - request_payload["model"] = _parsed_body.get("model", "") - if "custom_llm_provider" not in request_payload and custom_llm_provider: - request_payload["custom_llm_provider"] = custom_llm_provider - - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=request_payload, - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), - ) - - ######################################################### - - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - headers=custom_headers, - ) - else: - error_msg = f"{str(e)}" - raise ProxyException( - message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), - headers=custom_headers, - ) - - -def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> dict: - """ - If tags are in the request headers, add them to the metadata - - Used for google and vertex JS SDKs, and Azure passthrough - Checks both 'tags' and 'x-litellm-tags' headers - """ - tags_to_add = [] - - # Check for 'tags' header first - _tags = request.headers.get("tags") - if _tags: - tags_to_add.extend([tag.strip() for tag in _tags.split(",")]) - - _tags = request.headers.get("x-litellm-tags") - if _tags: - tags_to_add.extend([tag.strip() for tag in _tags.split(",")]) - - # Only add tags key if there are tags to add - if tags_to_add: - if "tags" not in metadata: - metadata["tags"] = [] - metadata["tags"].extend(tags_to_add) - - return metadata - - -async def _parse_request_data_by_content_type( - request: Request, -) -> Tuple[Optional[Any], Optional[Any], Optional[Any], Optional[Any]]: - """ - Parse request data based on content type. - - Handles JSON, multipart/form-data, and URL-encoded form data. - - Returns: - Tuple of (query_params_data, custom_body_data, file_data, stream) - """ - content_type = request.headers.get("content-type", "") - - query_params_data = None - custom_body_data = None - file_data = None - stream = None - - if "application/json" in content_type: - # ✅ Handle JSON - try: - body = await request.json() - query_params_data = body.get("query_params") - custom_body_data = body.get("custom_body") - stream = body.get("stream") - except json.JSONDecodeError: - # Handle requests with no body (e.g., DELETE requests) - pass - elif "multipart/form-data" in content_type: - # ✅ Handle multipart form-data - form = await request.form() - if "query_params" in form: - form_value = form["query_params"] - if isinstance(form_value, str): - try: - query_params_data = json.loads(form_value) - except Exception: - query_params_data = form_value - else: - query_params_data = form_value - - if "custom_body" in form: - form_value = form["custom_body"] - if isinstance(form_value, str): - try: - custom_body_data = json.loads(form_value) - except Exception: - custom_body_data = form_value - else: - custom_body_data = form_value - - if "file" in form: - file_data = form["file"] # this is a Starlette UploadFile object - - elif "application/x-www-form-urlencoded" in content_type: - # ✅ Handle URL-encoded form data - form = await request.form() - query_params_data = form.get("query_params") - custom_body_data = form.get("custom_body") - - else: - # ✅ Fallback: maybe no body, just query params - query_params_data = dict(request.query_params) or None - - return query_params_data, custom_body_data, file_data, stream - - -def create_pass_through_route( - endpoint, - target: str, - custom_headers: Optional[dict] = None, - _forward_headers: Optional[bool] = False, - _merge_query_params: Optional[bool] = False, - dependencies: Optional[List] = None, - include_subpath: Optional[bool] = False, - cost_per_request: Optional[float] = None, - custom_llm_provider: Optional[str] = None, - is_streaming_request: Optional[bool] = False, - query_params: Optional[dict] = None, - guardrails: Optional[Dict[str, Any]] = None, -): - # check if target is an adapter.py or a url - from litellm._uuid import uuid - from litellm.proxy.types_utils.utils import get_instance_fn - - try: - if isinstance(target, CustomLogger): - adapter = target - else: - adapter = get_instance_fn(value=target) - adapter_id = str(uuid.uuid4()) - litellm.adapters = [{"id": adapter_id, "adapter": adapter}] - - async def endpoint_func( # type: ignore - request: Request, - fastapi_response: Response, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - subpath: str = "", # captures sub-paths when include_subpath=True - custom_body: Optional[dict] = None, # accepted for signature compatibility with URL-based path; not forwarded because chat_completion_pass_through_endpoint does not support it - ): - return await chat_completion_pass_through_endpoint( - fastapi_response=fastapi_response, - request=request, - adapter_id=adapter_id, - user_api_key_dict=user_api_key_dict, - ) - - except Exception: - verbose_proxy_logger.debug("Defaulting to target being a url.") - - async def endpoint_func( # type: ignore - request: Request, - fastapi_response: Response, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - subpath: str = "", # captures sub-paths when include_subpath=True - custom_body: Optional[dict] = None, - ): - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - InitPassThroughEndpointHelpers, - ) - - path = request.url.path - - # Parse request data based on content type - ( - query_params_data, - custom_body_data, - file_data, - stream, - ) = await _parse_request_data_by_content_type(request) - - if not InitPassThroughEndpointHelpers.is_registered_pass_through_route( - route=path - ): - raise HTTPException( - status_code=404, - detail=f"Pass-through endpoint {endpoint} not found. This could have been deleted or not yet added to the proxy.", - ) - - passthrough_params = ( - InitPassThroughEndpointHelpers.get_registered_pass_through_route( - route=path - ) - ) - target_params = { - "target": target, - "custom_headers": custom_headers, - "forward_headers": _forward_headers, - "merge_query_params": _merge_query_params, - "cost_per_request": cost_per_request, - "guardrails": None, - } - - if passthrough_params is not None: - target_params.update(passthrough_params.get("passthrough_params", {})) - - # Extract and cast parameters with proper types - param_target = target_params.get("target") or target - param_custom_headers = target_params.get("custom_headers", custom_headers) - param_forward_headers = target_params.get( - "forward_headers", _forward_headers - ) - param_merge_query_params = target_params.get( - "merge_query_params", _merge_query_params - ) - param_cost_per_request = target_params.get( - "cost_per_request", cost_per_request - ) - param_guardrails = target_params.get("guardrails", None) - - # Construct the full target URL with subpath if needed - full_target = ( - HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( - base_target=cast(str, param_target), - subpath=subpath, - include_subpath=include_subpath, - ) - ) - - # Ensure custom_headers is a dict - headers_dict = ( - param_custom_headers if isinstance(param_custom_headers, dict) else {} - ) - - # Ensure query_params and custom_body are dicts or None - final_query_params = ( - query_params_data if isinstance(query_params_data, dict) else {} - ) - if query_params: - final_query_params.update(query_params) - # When a caller (e.g. bedrock_proxy_route) supplies a pre-built - # body, use it instead of the body parsed from the raw request. - final_custom_body: Optional[dict] = None - if custom_body is not None: - final_custom_body = custom_body - elif isinstance(custom_body_data, dict): - final_custom_body = custom_body_data - - return await pass_through_request( # type: ignore - request=request, - target=full_target, - custom_headers=headers_dict, - user_api_key_dict=user_api_key_dict, - forward_headers=cast(Optional[bool], param_forward_headers), - merge_query_params=cast(Optional[bool], param_merge_query_params), - query_params=final_query_params, - stream=is_streaming_request or stream, - custom_body=final_custom_body, - cost_per_request=cast(Optional[float], param_cost_per_request), - custom_llm_provider=custom_llm_provider, - guardrails_config=cast(Optional[dict], param_guardrails), - ) - - return endpoint_func - - -def create_websocket_passthrough_route( - endpoint: str, - target: str, - custom_headers: Optional[dict] = None, - _forward_headers: Optional[bool] = False, - dependencies: Optional[List] = None, - cost_per_request: Optional[float] = None, -): - """ - Create a WebSocket passthrough route function. - - Args: - endpoint: The endpoint path (for logging purposes) - target: The target WebSocket URL (e.g., "wss://api.example.com/ws") - custom_headers: Custom headers to include in the WebSocket connection - _forward_headers: Whether to forward incoming headers - dependencies: FastAPI dependencies to inject - - Returns: - A WebSocket passthrough function that can be registered with app.websocket() - """ - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket - - async def websocket_endpoint_func( - websocket: WebSocket, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), - **kwargs, # For additional query parameters - ): - """ - WebSocket passthrough endpoint function. - - This function handles the WebSocket connection by: - 1. Accepting the incoming WebSocket connection - 2. Establishing a connection to the target WebSocket - 3. Forwarding messages bidirectionally - 4. Handling connection cleanup - """ - return await websocket_passthrough_request( - websocket=websocket, - target=target, - custom_headers=custom_headers or {}, - user_api_key_dict=user_api_key_dict, - forward_headers=_forward_headers, - endpoint=endpoint, - cost_per_request=cost_per_request, - accept_websocket=True, # Generic usage should accept the WebSocket - ) - - return websocket_endpoint_func - - -async def websocket_passthrough_request( # noqa: PLR0915 - websocket: WebSocket, - target: str, - custom_headers: dict, - user_api_key_dict: UserAPIKeyAuth, - forward_headers: Optional[bool] = False, - endpoint: Optional[str] = None, - cost_per_request: Optional[float] = None, - accept_websocket: bool = True, -): - """ - WebSocket passthrough request handler. - - Args: - websocket: The incoming WebSocket connection - target: The target WebSocket URL - custom_headers: Custom headers to include in the connection - user_api_key_dict: The user API key dictionary - forward_headers: Whether to forward incoming headers - endpoint: The endpoint path (for logging purposes) - cost_per_request: Optional field - cost per request to the target endpoint - """ - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.proxy.proxy_server import proxy_logging_obj - from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - PassthroughStandardLoggingPayload, - ) - - # Initialize tracking variables - start_time = datetime.now() - websocket_messages: list[dict[str, Any]] = [] - litellm_call_id = str(uuid.uuid4()) - - verbose_proxy_logger.info( - f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}" - ) - - # Only accept the WebSocket if requested (for generic usage) - if accept_websocket: - await websocket.accept() - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): WebSocket connection accepted" - ) - - # Prepare headers for the upstream connection - upstream_headers = custom_headers.copy() - - if forward_headers: - # Forward relevant headers from the incoming request - incoming_headers = dict(websocket.headers) - for header_name, header_value in incoming_headers.items(): - # Only forward certain headers to avoid conflicts - if header_name.lower() in [ - "authorization", - "x-api-key", - "x-goog-user-project", - ]: - upstream_headers[header_name] = header_value - - # Initialize logging object similar to HTTP passthrough - logging_obj = Logging( - model="unknown", - messages=[{"role": "user", "content": "WebSocket connection"}], - stream=True, # WebSockets are inherently streaming - call_type="pass_through_endpoint", - start_time=start_time, - litellm_call_id=litellm_call_id, - function_id="websocket_passthrough", - ) - - # Create passthrough logging payload - passthrough_logging_payload = PassthroughStandardLoggingPayload( - url=target, - request_body={}, # WebSocket doesn't have a traditional request body - request_method="WEBSOCKET", - cost_per_request=cost_per_request, - ) - - # Create a dummy request object for WebSocket connections to maintain compatibility - # with the existing _init_kwargs_for_pass_through_endpoint function - class DummyRequest: - def __init__( - self, url: str, method: str = "WEBSOCKET", headers: Optional[dict] = None - ): - self.url = url - self.method = method - self.headers = headers or {} - - def __str__(self): - return f"DummyRequest(url={self.url}, method={self.method})" - - dummy_request = DummyRequest( - url=target, - method="WEBSOCKET", - headers=dict(websocket.headers) if hasattr(websocket, "headers") else {}, - ) - - # Initialize kwargs for logging using the same pattern as HTTP passthrough - kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( - user_api_key_dict=user_api_key_dict, - _parsed_body={}, # WebSocket doesn't have a traditional request body - passthrough_logging_payload=passthrough_logging_payload, - litellm_call_id=litellm_call_id, - request=dummy_request, # type: ignore - logging_obj=logging_obj, - ) - - # Update logging environment variables - logging_obj.update_environment_variables( - model="unknown", - user="unknown", - optional_params={}, - litellm_params=dict(kwargs.get("litellm_params", {})), - call_type="pass_through_endpoint", - ) - logging_obj.model_call_details["litellm_call_id"] = litellm_call_id - - # Pre-call logging - logging_obj.pre_call( - input=[{"role": "user", "content": "WebSocket connection"}], - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": target, - "headers": upstream_headers, - }, - ) - - ### CALL HOOKS ### - modify incoming data / reject request before calling the model - websocket_data: dict[str, Any] = {} - websocket_data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, - data=websocket_data, - call_type="pass_through_endpoint", - ) - - try: - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Establishing upstream connection to {target}" - ) - async with connect( - target, - additional_headers=upstream_headers, - ) as upstream_ws: - verbose_proxy_logger.info( - f"WebSocket passthrough ({endpoint}): Upstream connection established successfully" - ) - - async def forward_client_to_upstream() -> None: - """Forward messages from client to upstream WebSocket""" - try: - while True: - message = await websocket.receive() - message_type = message.get("type") - if message_type == "websocket.disconnect": - await upstream_ws.close() - break - - text_data = message.get("text") - bytes_data = message.get("bytes") - - if text_data is not None: - # Try to extract model from client setup message for Vertex AI Live - if endpoint and "/vertex_ai/live" in endpoint: - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Processing client message for model extraction" - ) - try: - client_message = json.loads(text_data) - if ( - isinstance(client_message, dict) - and "setup" in client_message - ): - setup_data = client_message["setup"] - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Found setup data in client message: {setup_data}" - ) - if ( - isinstance(setup_data, dict) - and "model" in setup_data - ): - extracted_model = ( - _extract_model_from_vertex_ai_setup( - setup_data - ) - ) - if extracted_model: - kwargs["model"] = extracted_model - kwargs[ - "custom_llm_provider" - ] = "vertex_ai-language-models" - # Update logging object with correct model - logging_obj.model = extracted_model - logging_obj.model_call_details[ - "model" - ] = extracted_model - logging_obj.model_call_details[ - "custom_llm_provider" - ] = "vertex_ai" - verbose_proxy_logger.info( - f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from client setup message" - ) - else: - verbose_proxy_logger.warning( - f"WebSocket passthrough ({endpoint}): Failed to extract model from client setup data: {setup_data}" - ) - else: - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Setup data does not contain model field: {setup_data}" - ) - else: - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Client message does not contain setup data" - ) - except (json.JSONDecodeError, KeyError, TypeError) as e: - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Client message is not a valid setup message: {e}" - ) - pass # Not a JSON message or doesn't contain setup data - - await upstream_ws.send(text_data) - elif bytes_data is not None: - await upstream_ws.send(bytes_data) - except asyncio.CancelledError: - raise - except Exception: - verbose_proxy_logger.exception( - f"WebSocket passthrough ({endpoint}): error forwarding client message" - ) - await upstream_ws.close() - - async def forward_upstream_to_client() -> None: - """Forward messages from upstream to client WebSocket""" - try: - # Wait for the first response from upstream - raw_response = await upstream_ws.recv(decode=False) - # Ensure raw_response is bytes before decoding - if isinstance(raw_response, str): - raw_response = raw_response.encode("ascii") - setup_response = json.loads(raw_response.decode("ascii")) - verbose_proxy_logger.debug(f"Setup response: {setup_response}") - - # Extract model and provider from setup response for Vertex AI Live - if endpoint and "/vertex_ai/live" in endpoint: - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Processing server setup response for model extraction" - ) - extracted_model = _extract_model_from_vertex_ai_setup( - setup_response - ) - if extracted_model: - kwargs["model"] = extracted_model - kwargs["custom_llm_provider"] = "vertex_ai_language_models" - # Update logging object with correct model - logging_obj.model = extracted_model - logging_obj.model_call_details["model"] = extracted_model - logging_obj.model_call_details[ - "custom_llm_provider" - ] = "vertex_ai_language_models" - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from server setup response" - ) - else: - verbose_proxy_logger.warning( - f"WebSocket passthrough ({endpoint}): Failed to extract model from server setup response: {setup_response}" - ) - else: - verbose_proxy_logger.debug( - f"WebSocket passthrough ({endpoint}): Not a Vertex AI Live endpoint, skipping model extraction" - ) - - # Send the setup response to the client - await websocket.send_text(json.dumps(setup_response)) - - # Now continuously forward messages from upstream to client - async for upstream_message in upstream_ws: - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - # Parse and collect for cost tracking - try: - message_data = json.loads(upstream_message.decode()) - websocket_messages.append(message_data) - except (json.JSONDecodeError, UnicodeDecodeError): - pass - else: - await websocket.send_text(upstream_message) - # Parse and collect for cost tracking - try: - message_data = json.loads(upstream_message) - websocket_messages.append(message_data) - except json.JSONDecodeError: - pass - - except (ConnectionClosedOK, ConnectionClosedError) as e: - verbose_proxy_logger.debug( - f"Upstream WebSocket connection closed: {e}" - ) - pass - except asyncio.CancelledError: - verbose_proxy_logger.debug( - "asyncio.CancelledError in forward_upstream_to_client" - ) - raise - except Exception as e: - verbose_proxy_logger.debug( - f"Exception in forward_upstream_to_client: {e}" - ) - verbose_proxy_logger.exception( - f"WebSocket passthrough ({endpoint}): error forwarding upstream message" - ) - raise - - # Create tasks for bidirectional message forwarding - tasks = [ - asyncio.create_task(forward_client_to_upstream()), - asyncio.create_task(forward_upstream_to_client()), - ] - - done, pending = await asyncio.wait( - tasks, return_when=asyncio.FIRST_COMPLETED - ) - - # Cancel remaining tasks - for task in pending: - task.cancel() - try: - await task - except asyncio.CancelledError: - pass - - # Check for exceptions in completed tasks - for task in done: - exception = task.exception() - if exception is not None: - raise exception - - end_time = datetime.now() - - # Update passthrough logging payload with response data - passthrough_logging_payload["response_body"] = websocket_messages # type: ignore - passthrough_logging_payload["end_time"] = end_time # type: ignore - - # Remove logging_obj from kwargs to avoid duplicate keyword argument - success_kwargs = kwargs.copy() - success_kwargs.pop("logging_obj", None) - - # # Add user authentication context for database logging - # if user_api_key_dict: - # success_kwargs.setdefault('litellm_params', {}) - # success_kwargs['litellm_params'].update({ - # 'proxy_server_request': { - # 'body': { - # 'user': user_api_key_dict.user_id, - # 'team_id': user_api_key_dict.team_id, - # 'end_user_id': user_api_key_dict.end_user_id, - # } - # } - # }) - # # Also add the user_api_key for direct access - # success_kwargs['user_api_key'] = user_api_key_dict.api_key - - # Create a dummy httpx.Response for WebSocket connections - class MockWebSocketResponse: - def __init__(self, target_url: str): - self.status_code = 200 - self.text = "WebSocket connection successful" - self.headers: dict[str, str] = {} - self.request = MockWebSocketRequest(target_url) - - class MockWebSocketRequest: - def __init__(self, target_url: str): - self.method = "WEBSOCKET" - self.url = target_url - - mock_response = MockWebSocketResponse(target) - - # Use the same success handler as HTTP passthrough endpoints - asyncio.create_task( - pass_through_endpoint_logging.pass_through_async_success_handler( - httpx_response=mock_response, # type: ignore - response_body=websocket_messages, # type: ignore - url_route=endpoint or "", - result="websocket_connection_successful", - start_time=start_time, - end_time=end_time, - logging_obj=logging_obj, - cache_hit=False, - request_body={}, - **success_kwargs, - ) - ) - - # Call the proxy logging success hook - if proxy_logging_obj: - await proxy_logging_obj.post_call_success_hook( - data={}, - user_api_key_dict=user_api_key_dict, - response={"status": "websocket_connection_successful"}, # type: ignore - ) - - except InvalidStatus as exc: - verbose_proxy_logger.exception( - f"WebSocket passthrough ({endpoint}): upstream rejected WebSocket connection" - ) - - # Prepare request payload for logging - request_payload = {} - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - - # Log the connection failure using the same pattern as HTTP - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=exc, - request_data=request_payload, - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), - ) - - if websocket.client_state != WebSocketState.DISCONNECTED: - await websocket.close( - code=getattr(exc, "status_code", 1011), - reason="Upstream connection rejected", - ) - except Exception as e: - verbose_proxy_logger.exception( - f"WebSocket passthrough ({endpoint}): unexpected error while proxying WebSocket" - ) - - # Prepare request payload for logging - request_payload = {} - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - - # Log the unexpected error using the same pattern as HTTP - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=request_payload, - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), - ) - - if websocket.client_state != WebSocketState.DISCONNECTED: - await websocket.close(code=1011, reason="WebSocket passthrough error") - finally: - if websocket.client_state != WebSocketState.DISCONNECTED: - await websocket.close() - - -def _is_streaming_response(response: httpx.Response) -> bool: - _content_type = response.headers.get("content-type") - if _content_type is not None and "text/event-stream" in _content_type: - return True - return False - - -def _extract_model_from_vertex_ai_setup(setup_response: dict) -> Optional[str]: - """ - Extract the model name from Vertex AI Live setup response. - - The setup response can contain a model field in two formats: - 1. Direct: {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"} - 2. Nested: {"setup": {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}} - - We extract just the model name: "gemini-2.0-flash-live-preview-04-09" - """ - try: - # Handle both direct model field and nested setup.model field - model_path = None - if isinstance(setup_response, dict): - if "model" in setup_response: - model_path = setup_response["model"] - elif ( - "setup" in setup_response - and isinstance(setup_response["setup"], dict) - and "model" in setup_response["setup"] - ): - model_path = setup_response["setup"]["model"] - - if isinstance(model_path, str) and "/models/" in model_path: - # Extract the model name after the last "/models/" - model_name = model_path.split("/models/")[-1] - return model_name - except Exception as e: - verbose_proxy_logger.debug(f"Error extracting model from setup response: {e}") - return None - - -class SafeRouteAdder: - """ - Wrapper class for adding routes to FastAPI app. - Only adds routes if they don't already exist on the app. - """ - - @staticmethod - def _is_path_registered(app: FastAPI, path: str, methods: List[str]) -> bool: - """ - Check if a path with any of the specified methods is already registered on the app. - - Args: - app: The FastAPI application instance - path: The path to check (e.g., "/v1/chat/completions") - methods: List of HTTP methods to check (e.g., ["GET", "POST"]) - - Returns: - True if the path is already registered with any of the methods, False otherwise - """ - for route in app.routes: - # Use getattr to safely access route attributes - route_path = getattr(route, "path", None) - route_methods = getattr(route, "methods", None) - - if route_path == path and route_methods is not None: - # Check if any of the methods overlap - if any(method in route_methods for method in methods): - return True - return False - - @staticmethod - def add_api_route_if_not_exists( - app: FastAPI, - path: str, - endpoint: Any, - methods: List[str], - dependencies: Optional[List] = None, - ) -> bool: - """ - Add an API route to the app only if it doesn't already exist. - - Args: - app: The FastAPI application instance - path: The path for the route - endpoint: The endpoint function/callable - methods: List of HTTP methods - dependencies: Optional list of dependencies - - Returns: - True if route was added, False if it already existed - """ - if SafeRouteAdder._is_path_registered(app=app, path=path, methods=methods): - verbose_proxy_logger.debug( - "Skipping route registration - path %s with methods %s already registered on app", - path, - methods, - ) - return False - - app.add_api_route( - path=path, - endpoint=endpoint, - methods=methods, - dependencies=dependencies, - ) - verbose_proxy_logger.debug( - "Successfully added route: %s with methods %s", - path, - methods, - ) - return True - - -class InitPassThroughEndpointHelpers: - @staticmethod - def add_exact_path_route( - app: FastAPI, - path: str, - target: str, - custom_headers: Optional[dict], - forward_headers: Optional[bool], - merge_query_params: Optional[bool], - dependencies: Optional[List], - cost_per_request: Optional[float], - endpoint_id: str, - guardrails: Optional[dict] = None, - ): - """Add exact path route for pass-through endpoint""" - route_key = f"{endpoint_id}:exact:{path}" - - # Check if this exact route is already registered - if route_key in _registered_pass_through_routes: - verbose_proxy_logger.debug( - "Updating duplicate exact pass through endpoint: %s (already registered)", - path, - ) - - verbose_proxy_logger.debug( - "adding exact pass through endpoint: %s, dependencies: %s", - path, - dependencies, - ) - - # Use SafeRouteAdder to only add route if it doesn't exist on the app - SafeRouteAdder.add_api_route_if_not_exists( - app=app, - path=path, - endpoint=create_pass_through_route( # type: ignore - path, - target, - custom_headers, - forward_headers, - merge_query_params, - dependencies, - cost_per_request=cost_per_request, - guardrails=guardrails, - ), - methods=["GET", "POST", "PUT", "DELETE", "PATCH"], - dependencies=dependencies, - ) - - # Always register/update the route metadata (headers, target) even if FastAPI route exists - _registered_pass_through_routes[route_key] = { - "endpoint_id": endpoint_id, - "path": path, - "type": "exact", - "passthrough_params": { - "target": target, - "custom_headers": custom_headers, - "forward_headers": forward_headers, - "merge_query_params": merge_query_params, - "dependencies": dependencies, - "cost_per_request": cost_per_request, - "guardrails": guardrails, - }, - } - - @staticmethod - def add_subpath_route( - app: FastAPI, - path: str, - target: str, - custom_headers: Optional[dict], - forward_headers: Optional[bool], - merge_query_params: Optional[bool], - dependencies: Optional[List], - cost_per_request: Optional[float], - endpoint_id: str, - guardrails: Optional[dict] = None, - ): - """Add wildcard route for sub-paths""" - wildcard_path = f"{path}/{{subpath:path}}" - route_key = f"{endpoint_id}:subpath:{path}" - - # Check if this subpath route is already registered - if route_key in _registered_pass_through_routes: - verbose_proxy_logger.debug( - "Updating duplicate wildcard pass through endpoint: %s (already registered)", - wildcard_path, - ) - - verbose_proxy_logger.debug( - "adding wildcard pass through endpoint: %s, dependencies: %s", - wildcard_path, - dependencies, - ) - - # Use SafeRouteAdder to only add route if it doesn't exist on the app - SafeRouteAdder.add_api_route_if_not_exists( - app=app, - path=wildcard_path, - endpoint=create_pass_through_route( # type: ignore - path, - target, - custom_headers, - forward_headers, - merge_query_params, - dependencies, - include_subpath=True, - cost_per_request=cost_per_request, - guardrails=guardrails, - ), - methods=["GET", "POST", "PUT", "DELETE", "PATCH"], - dependencies=dependencies, - ) - - # Register the route to prevent duplicates only if it was added - _registered_pass_through_routes[route_key] = { - "endpoint_id": endpoint_id, - "path": path, - "type": "subpath", - "passthrough_params": { - "target": target, - "custom_headers": custom_headers, - "forward_headers": forward_headers, - "merge_query_params": merge_query_params, - "dependencies": dependencies, - "cost_per_request": cost_per_request, - "guardrails": guardrails, - }, - } - - @staticmethod - 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 - ] - for key in keys_to_remove: - del _registered_pass_through_routes[key] - verbose_proxy_logger.debug( - "Removed pass-through route from registry: %s", key - ) - - @staticmethod - def clear_all_pass_through_routes(): - """Clear all pass-through routes from the registry""" - _registered_pass_through_routes.clear() - - @staticmethod - def get_all_registered_pass_through_routes() -> List[str]: - """Get all registered pass-through endpoints from the registry""" - return list(_registered_pass_through_routes.keys()) - - @staticmethod - def _build_full_path_with_root(path: str) -> str: - """ - Build full path by prepending server root path if needed. - - Args: - path: The relative path to build - - Returns: - Full path with server root prepended (if root is not "/") - """ - root_path = get_server_root_path() - if root_path == "/": - return path - return f"{root_path}{path}" - - @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 - """ - ## CHECK IF MAPPED PASS THROUGH ENDPOINT - for mapped_route in LiteLLMRoutes.mapped_pass_through_routes.value: - if route.startswith(mapped_route): - return True - - # 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 = InitPassThroughEndpointHelpers._build_full_path_with_root( - 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 - - @staticmethod - def get_registered_pass_through_route(route: str) -> Optional[Dict[str, Any]]: - """Get passthrough params for a given route""" - 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 = InitPassThroughEndpointHelpers._build_full_path_with_root( - parts[2] - ) - - if route_type == "exact" and route == registered_path: - return _registered_pass_through_routes[key] - elif route_type == "subpath": - if route == registered_path or route.startswith( - registered_path + "/" - ): - return _registered_pass_through_routes[key] - - return None - - -def _get_combined_pass_through_endpoints( - pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]], - config_pass_through_endpoints: List[Dict], -): - """Get combined pass-through endpoints from db + config""" - return pass_through_endpoints + config_pass_through_endpoints - - -async def initialize_pass_through_endpoints( - pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]], -): - """ - 1. Create a global list of pass-through endpoints (db + config) - 2. Clear all existing pass-through endpoints from the FastAPI app routes - 3. Add new endpoints to the in-memory registry - - Initialize a list of pass-through endpoints by adding them to the FastAPI app routes - - Args: - pass_through_endpoints: List of pass-through endpoints to initialize - - Returns: - None - """ - from litellm._uuid import uuid - - verbose_proxy_logger.debug("initializing pass through endpoints") - from litellm.proxy._types import CommonProxyErrors, LiteLLMRoutes - from litellm.proxy.proxy_server import ( - app, - config_passthrough_endpoints, - premium_user, - ) - - ## get combined pass-through endpoints from db + config - combined_pass_through_endpoints: List[Union[Dict, PassThroughGenericEndpoint]] - - if config_passthrough_endpoints is not None: - combined_pass_through_endpoints = _get_combined_pass_through_endpoints( # type: ignore - pass_through_endpoints, config_passthrough_endpoints - ) - else: - combined_pass_through_endpoints = pass_through_endpoints # type: ignore - - ## clear all existing pass-through endpoints from the FastAPI app routes - # InitPassThroughEndpointHelpers.clear_all_pass_through_routes() - - # get a list of all registered pass-through endpoints - # mark the ones that are visited in the list - # remove the ones that are not visited from the list - registered_pass_through_endpoints = ( - InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() - ) - - visited_endpoints = set() - - for endpoint in combined_pass_through_endpoints: - if isinstance(endpoint, PassThroughGenericEndpoint): - endpoint = endpoint.model_dump() - - # Auto-generate ID for backwards compatibility if not present - if endpoint.get("id") is None: - endpoint["id"] = str(uuid.uuid4()) - - # Get the endpoint_id as a string (guaranteed to be set at this point) - endpoint_id: str = endpoint["id"] - - _target = endpoint.get("target", None) - _path: Optional[str] = endpoint.get("path", None) - 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 - ) - _forward_headers = endpoint.get("forward_headers", None) - _merge_query_params = endpoint.get("merge_query_params", None) - _auth = endpoint.get("auth", None) - _dependencies = None - if _auth is not None and str(_auth).lower() == "true": - if premium_user is not True: - raise ValueError( - "Error Setting Authentication on Pass Through Endpoint: {}".format( - CommonProxyErrors.not_premium_user.value - ) - ) - _dependencies = [Depends(user_api_key_auth)] - LiteLLMRoutes.openai_routes.value.append(_path) - - if _target is None: - continue - - # Get guardrails config if present - _guardrails = endpoint.get("guardrails", None) - - # Add exact path route - verbose_proxy_logger.debug( - "Initializing pass through endpoint: %s (ID: %s)", _path, endpoint_id - ) - InitPassThroughEndpointHelpers.add_exact_path_route( - app=app, - path=_path, - target=_target, - custom_headers=_custom_headers, - forward_headers=_forward_headers, - merge_query_params=_merge_query_params, - dependencies=_dependencies, - cost_per_request=endpoint.get("cost_per_request", None), - endpoint_id=endpoint_id, - guardrails=_guardrails, - ) - - visited_endpoints.add(f"{endpoint_id}:exact:{_path}") - - # Add wildcard route for sub-paths - if endpoint.get("include_subpath", False) is True: - InitPassThroughEndpointHelpers.add_subpath_route( - app=app, - path=_path, - target=_target, - custom_headers=_custom_headers, - forward_headers=_forward_headers, - merge_query_params=_merge_query_params, - dependencies=_dependencies, - cost_per_request=endpoint.get("cost_per_request", None), - endpoint_id=endpoint_id, - guardrails=_guardrails, - ) - - visited_endpoints.add(f"{endpoint_id}:subpath:{_path}") - - verbose_proxy_logger.debug( - "Added new pass through endpoint: %s (ID: %s)", _path, endpoint_id - ) - - # remove the ones that are not visited from the list - for endpoint_key in registered_pass_through_endpoints: - if endpoint_key not in visited_endpoints: - InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_key) - - -def _get_pass_through_endpoints_from_config() -> List[PassThroughGenericEndpoint]: - """ - Get pass-through endpoints defined in the config file. - These are read-only and cannot be edited via the UI. - Malformed endpoints are logged and skipped; they do not crash the function. - """ - from pydantic import ValidationError - - from litellm.proxy.proxy_server import config_passthrough_endpoints - - if config_passthrough_endpoints is None or len(config_passthrough_endpoints) == 0: - return [] - - returned_endpoints: List[PassThroughGenericEndpoint] = [] - for endpoint in config_passthrough_endpoints: - try: - if isinstance(endpoint, dict): - endpoint_dict = dict(endpoint) - endpoint_dict["is_from_config"] = True - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) - elif isinstance(endpoint, PassThroughGenericEndpoint): - # Create a copy with is_from_config=True - endpoint_dict = endpoint.model_dump() - endpoint_dict["is_from_config"] = True - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) - except ValidationError as e: - verbose_proxy_logger.warning( - "Skipping malformed pass-through endpoint from config: %s", - e, - exc_info=False, - ) - - return returned_endpoints - - -async def _get_pass_through_endpoints_from_db( - endpoint_id: Optional[str] = None, - user_api_key_dict: Optional[UserAPIKeyAuth] = None, -) -> List[PassThroughGenericEndpoint]: - from litellm.proxy._types import LitellmUserRoles - from litellm.proxy.proxy_server import get_config_general_settings - - try: - if user_api_key_dict is None: - user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) - response: ConfigFieldInfo = await get_config_general_settings( - field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict - ) - except Exception: - return [] - - pass_through_endpoint_data: Optional[List] = response.field_value - if pass_through_endpoint_data is None: - return [] - - returned_endpoints: List[PassThroughGenericEndpoint] = [] - if endpoint_id is None: - # Return all endpoints from DB, mark as not from config - for endpoint in pass_through_endpoint_data: - if isinstance(endpoint, dict): - endpoint_dict = dict(endpoint) - endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) - elif isinstance(endpoint, PassThroughGenericEndpoint): - endpoint_dict = endpoint.model_dump() - endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) - else: - # Find specific endpoint by ID - found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) - if found_endpoint is not None: - endpoint_dict = ( - found_endpoint.model_dump() - if isinstance(found_endpoint, PassThroughGenericEndpoint) - else dict(found_endpoint) - ) - endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) - - return returned_endpoints - - -async def _filter_endpoints_by_team_allowed_routes( - team_id: str, - pass_through_endpoints: List[PassThroughGenericEndpoint], - prisma_client, -) -> List[PassThroughGenericEndpoint]: - """ - Filter pass-through endpoints based on team's allowed_passthrough_routes metadata. - - Args: - team_id: The team ID to check permissions for - pass_through_endpoints: List of endpoints to filter - prisma_client: Database client - - Returns: - Filtered list of endpoints based on team permissions - - Raises: - HTTPException: If team is not found - """ - # retrieve team from db - team = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - ) - if team is None: - raise HTTPException( - status_code=404, - detail={"error": "Team not found"}, - ) - - # retrieve team metadata - team_metadata = team.metadata - if ( - team_metadata is not None - and team_metadata.get("allowed_passthrough_routes") is not None - ): - ## FILTER pass_through_endpoints by allowed_passthrough_routes - pass_through_endpoints = [ - endpoint - for endpoint in pass_through_endpoints - if endpoint.path in team_metadata.get("allowed_passthrough_routes") - ] - - return pass_through_endpoints - - -@router.get( - "/config/pass_through_endpoint", - dependencies=[Depends(user_api_key_auth)], - response_model=PassThroughEndpointResponse, -) -@router.get( - "/config/pass_through_endpoint/team/{team_id}", - dependencies=[Depends(user_api_key_auth)], - response_model=PassThroughEndpointResponse, -) -async def get_pass_through_endpoints( - endpoint_id: Optional[str] = None, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - team_id: Optional[str] = None, -): - """ - GET configured pass through endpoint. - - If no endpoint_id given, return all configured endpoints. - """ ## Get existing pass-through endpoint field value - from litellm.proxy._types import CommonProxyErrors - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={"error": CommonProxyErrors.db_not_connected_error.value}, - ) - - # Get endpoints from DB (editable via UI) - db_endpoints = await _get_pass_through_endpoints_from_db( - endpoint_id=endpoint_id, user_api_key_dict=user_api_key_dict - ) - - # Get endpoints from config file (read-only, not editable via UI) - config_endpoints = _get_pass_through_endpoints_from_config() - - # Merge: config endpoints not in DB + all DB endpoints (DB overrides config for same path) - db_paths = {ep.path for ep in db_endpoints} - config_only_endpoints = [ - ep for ep in config_endpoints if ep.path not in db_paths - ] - if endpoint_id is not None: - # When filtering by endpoint_id, only return if found in DB (config endpoints use generated IDs) - pass_through_endpoints = db_endpoints - else: - pass_through_endpoints = config_only_endpoints + db_endpoints - - if team_id is not None: - pass_through_endpoints = await _filter_endpoints_by_team_allowed_routes( - team_id=team_id, - pass_through_endpoints=pass_through_endpoints, - prisma_client=prisma_client, - ) - - return PassThroughEndpointResponse(endpoints=pass_through_endpoints) - - -@router.post( - "/config/pass_through_endpoint/{endpoint_id}", - dependencies=[Depends(user_api_key_auth)], -) -async def update_pass_through_endpoints( - endpoint_id: str, - data: PassThroughGenericEndpoint, - request: Request, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Update a pass-through endpoint by ID. - """ - from litellm.proxy.proxy_server import ( - get_config_general_settings, - update_config_general_settings, - ) - - ## Get existing pass-through endpoint field value - try: - response: ConfigFieldInfo = await get_config_general_settings( - field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict - ) - except Exception: - raise HTTPException( - status_code=404, - detail={"error": "No pass-through endpoints found"}, - ) - - pass_through_endpoint_data: Optional[List] = response.field_value - if pass_through_endpoint_data is None: - raise HTTPException( - status_code=404, - detail={"error": "No pass-through endpoints found"}, - ) - - # Find the endpoint to update - found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) - - if found_endpoint is None: - raise HTTPException( - status_code=404, - detail={"error": f"Endpoint with ID '{endpoint_id}' not found"}, - ) - - # 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 - ) - if _endpoint.id == endpoint_id: - endpoint_index = idx - break - - if endpoint_index is None: - raise HTTPException( - status_code=404, - 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 - # Exclude is_from_config as it's a response-only field (computed at read time) - update_data = data.model_dump(exclude_none=True, exclude={"is_from_config"}) - - # Start with existing endpoint data - endpoint_dict = found_endpoint.model_dump() - - # Update with new data (only non-None values) - endpoint_dict.update(update_data) - - # Preserve existing ID if not provided in update and endpoint has ID - if "id" not in update_data and found_endpoint.id is not None: - endpoint_dict["id"] = found_endpoint.id - - # Remove is_from_config before saving - it's a response-only field (computed at read time) - endpoint_dict.pop("is_from_config", None) - - # Create updated endpoint object - updated_endpoint = PassThroughGenericEndpoint(**endpoint_dict) - - # Update the list - pass_through_endpoint_data[endpoint_index] = endpoint_dict - - # Remove old routes from registry before they get re-registered - InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id) - - ## Update db - updated_data = ConfigFieldUpdate( - field_name="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 - ) - - # Re-register the route with updated headers - _custom_headers: Optional[dict] = updated_endpoint.headers or {} - _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) - - if updated_endpoint.include_subpath: - InitPassThroughEndpointHelpers.add_subpath_route( - app=request.app, - path=updated_endpoint.path, - target=updated_endpoint.target, - custom_headers=_custom_headers, - forward_headers=None, # Defaults not available in model? assuming None logic handles it - merge_query_params=None, - dependencies=None, - cost_per_request=updated_endpoint.cost_per_request, - endpoint_id=updated_endpoint.id or endpoint_id or "", - guardrails=getattr(updated_endpoint, "guardrails", None), - ) - else: - InitPassThroughEndpointHelpers.add_exact_path_route( - app=request.app, - path=updated_endpoint.path, - target=updated_endpoint.target, - custom_headers=_custom_headers, - forward_headers=None, - merge_query_params=None, - dependencies=None, - cost_per_request=updated_endpoint.cost_per_request, - endpoint_id=updated_endpoint.id or endpoint_id or "", - guardrails=getattr(updated_endpoint, "guardrails", None), - ) - - return PassThroughEndpointResponse( - endpoints=[updated_endpoint] if updated_endpoint else [] - ) - - -@router.post( - "/config/pass_through_endpoint", - dependencies=[Depends(user_api_key_auth)], -) -async def create_pass_through_endpoints( - data: PassThroughGenericEndpoint, - request: Request, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Create new pass-through endpoint - """ - from litellm._uuid import uuid - from litellm.proxy.proxy_server import ( - get_config_general_settings, - update_config_general_settings, - ) - - ## Get existing pass-through endpoint field value - - try: - response: ConfigFieldInfo = await get_config_general_settings( - 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 - ) - - ## Auto-generate ID if not provided - # Exclude is_from_config as it's a response-only field (computed at read time) - data_dict = data.model_dump(exclude={"is_from_config"}) - if data_dict.get("id") is None: - data_dict["id"] = str(uuid.uuid4()) - - if response.field_value is None: - response.field_value = [data_dict] - elif isinstance(response.field_value, List): - response.field_value.append(data_dict) - - ## Update db - updated_data = ConfigFieldUpdate( - field_name="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 - ) - - # Return the created endpoint with the generated ID - created_endpoint = PassThroughGenericEndpoint(**data_dict) - - # Register the new route - _custom_headers: Optional[dict] = created_endpoint.headers or {} - _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) - - if created_endpoint.include_subpath: - InitPassThroughEndpointHelpers.add_subpath_route( - app=request.app, - path=created_endpoint.path, - target=created_endpoint.target, - custom_headers=_custom_headers, - forward_headers=None, - merge_query_params=None, - dependencies=None, - cost_per_request=created_endpoint.cost_per_request, - endpoint_id=created_endpoint.id or "", - guardrails=getattr(created_endpoint, "guardrails", None), - ) - else: - InitPassThroughEndpointHelpers.add_exact_path_route( - app=request.app, - path=created_endpoint.path, - target=created_endpoint.target, - custom_headers=_custom_headers, - forward_headers=None, - merge_query_params=None, - dependencies=None, - cost_per_request=created_endpoint.cost_per_request, - endpoint_id=created_endpoint.id or "", - guardrails=getattr(created_endpoint, "guardrails", None), - ) - - return PassThroughEndpointResponse(endpoints=[created_endpoint]) - - -@router.delete( - "/config/pass_through_endpoint", - dependencies=[Depends(user_api_key_auth)], - response_model=PassThroughEndpointResponse, -) -async def delete_pass_through_endpoints( - endpoint_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Delete a pass-through endpoint by ID. - - Returns - the deleted endpoint - """ - from litellm.proxy.proxy_server import ( - get_config_general_settings, - update_config_general_settings, - ) - - ## Get existing pass-through endpoint field value - - try: - response: ConfigFieldInfo = await get_config_general_settings( - 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 - ) - - ## Update field by removing endpoint - pass_through_endpoint_data: Optional[List] = response.field_value - if response.field_value is None or pass_through_endpoint_data is None: - raise HTTPException( - status_code=400, - detail={"error": "There are no pass-through endpoints setup."}, - ) - - # Find the endpoint to delete - found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) - - 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 - ) - }, - ) - - # 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 - ) - if _endpoint.id == endpoint_id: - endpoint_index = idx - break - - if endpoint_index is None: - raise HTTPException( - status_code=400, - detail={ - "error": f"Could not find index for endpoint with ID '{endpoint_id}'" - }, - ) - - # Remove the endpoint - pass_through_endpoint_data.pop(endpoint_index) - response_obj = found_endpoint - - # Remove routes from registry - InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id) - - ## Update db - updated_data = ConfigFieldUpdate( - field_name="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 - ) - - return PassThroughEndpointResponse(endpoints=[response_obj]) - - -def _find_endpoint_by_id( - endpoints_data: List, - endpoint_id: str, -) -> Optional[PassThroughGenericEndpoint]: - """ - Find an endpoint by ID. - - Args: - endpoints_data: List of endpoint data (dicts or PassThroughGenericEndpoint objects) - endpoint_id: ID to search for - - Returns: - Found endpoint or None if not found - """ - for endpoint in endpoints_data: - _endpoint: Optional[PassThroughGenericEndpoint] = None - if isinstance(endpoint, dict): - _endpoint = PassThroughGenericEndpoint(**endpoint) - elif isinstance(endpoint, PassThroughGenericEndpoint): - _endpoint = endpoint - - # Only compare IDs to IDs - if _endpoint is not None and _endpoint.id == endpoint_id: - return _endpoint - - return None - - -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 - ) +import ast +import asyncio +import copy +import json +import traceback +from base64 import b64encode +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple, Union, cast +from urllib.parse import urlencode, urlparse + +import httpx +from fastapi import ( + APIRouter, + Depends, + FastAPI, + HTTPException, + Request, + Response, + UploadFile, + WebSocket, + status, +) +from fastapi.responses import StreamingResponse +from starlette.datastructures import UploadFile as StarletteUploadFile +from starlette.websockets import WebSocketState +from websockets.asyncio.client import connect +from websockets.exceptions import ( + ConnectionClosedError, + ConnectionClosedOK, + InvalidStatus, +) + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.passthrough import BasePassthroughUtils +from litellm.proxy._types import ( + ConfigFieldInfo, + ConfigFieldUpdate, + LiteLLMRoutes, + PassThroughEndpointResponse, + PassThroughGenericEndpoint, + ProxyException, + UserAPIKeyAuth, +) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.utils import get_server_root_path +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.custom_http import httpxSpecialProvider +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + PassthroughStandardLoggingPayload, +) +from litellm.types.utils import StandardLoggingUserAPIKeyMetadata + +from .streaming_handler import PassThroughStreamingHandler +from .success_handler import PassThroughEndpointLogging + +router = APIRouter() + +pass_through_endpoint_logging = PassThroughEndpointLogging() + +# Global registry to track registered pass-through routes and prevent memory leaks +_registered_pass_through_routes: Dict[ + str, Dict[str, Union[str, List[str], Dict[str, Any]]] +] = {} + + +def get_response_body(response: httpx.Response) -> Optional[dict]: + try: + return response.json() + except Exception: + return None + + +async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optional[dict]: + """ + checks if any headers on config.yaml are defined as os.environ/COHERE_API_KEY etc + + only runs for headers defined on config.yaml + + example header can be + + {"Authorization": "Bearer os.environ/COHERE_API_KEY"} + """ + if custom_headers is None: + return None + headers = {} + for key, value in custom_headers.items(): + # langfuse Api requires base64 encoded headers - it's simpleer to just ask litellm users to set their langfuse public and secret keys + # we can then get the b64 encoded keys here + if key == "LANGFUSE_PUBLIC_KEY" or key == "LANGFUSE_SECRET_KEY": + # 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/"): + _langfuse_public_key = get_secret_str(_langfuse_public_key) + 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") + ).decode("ascii") + else: + # 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" + ) + # get string section that is os.environ/ + start_index = value.find("os.environ/") + _variable_name = value[start_index:] + + verbose_proxy_logger.debug( + "pass through endpoint - getting secret for variable name: %s", + _variable_name, + ) + _secret_value = get_secret_str(_variable_name) + if _secret_value is not None: + new_value = value.replace(_variable_name, _secret_value) + headers[key] = new_value + return headers + + +async def chat_completion_pass_through_endpoint( # noqa: PLR0915 + fastapi_response: Response, + request: Request, + adapter_id: str, + user_api_key_dict: UserAPIKeyAuth, +): + from litellm.proxy.proxy_server import ( + add_litellm_data_to_request, + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + data = {} + try: + body = await request.body() + body_str = body.decode() + try: + data = ast.literal_eval(body_str) + except Exception: + data = json.loads(body_str) + + data["adapter_id"] = adapter_id + + verbose_proxy_logger.debug( + "Request received by LiteLLM:\n{}".format(json.dumps(data, indent=4)), + ) + data["model"] = ( + general_settings.get("completion_model", None) # server default + or user_model # model name passed via cli args + or data.get("model", None) # default passed in http request + ) + if user_model: + data["model"] = user_model + + data = await add_litellm_data_to_request( + data=data, # type: ignore + request=request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_config=proxy_config, + ) + + # override with user settings, these are params passed via cli + if user_temperature: + data["temperature"] = user_temperature + if user_request_timeout: + data["request_timeout"] = user_request_timeout + if user_max_tokens: + data["max_tokens"] = user_max_tokens + if user_api_base: + data["api_base"] = user_api_base + + ### MODEL ALIAS MAPPING ### + # check if model name in model alias map + # get the actual model name + if data["model"] in litellm.model_alias_map: + data["model"] = litellm.model_alias_map[data["model"]] + + # Check key-specific aliases + if ( + isinstance(data["model"], str) + and user_api_key_dict.aliases + and isinstance(user_api_key_dict.aliases, dict) + and data["model"] in user_api_key_dict.aliases + ): + data["model"] = user_api_key_dict.aliases[data["model"]] + + ### CALL HOOKS ### - modify incoming data before calling the model + data = await proxy_logging_obj.pre_call_hook( # type: ignore + user_api_key_dict=user_api_key_dict, data=data, call_type="text_completion" + ) + + ### ROUTE THE REQUESTs ### + router_model_names = llm_router.model_names if llm_router is not None else [] + # 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 + llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) + elif ( + llm_router is not None + and llm_router.model_group_alias is not None + and data["model"] in llm_router.model_group_alias + ): # model set in model_group_alias + llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) + elif llm_router is not None and llm_router.has_model_id( + data["model"] + ): # model in router model list + llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) + elif ( + llm_router is not None + and data["model"] not in router_model_names + and ( + llm_router.default_deployment is not None + or len(llm_router.pattern_router.patterns) > 0 + ) + ): # check for wildcard routes or default deployment before checking deployment_names + llm_response = asyncio.create_task(llm_router.aadapter_completion(**data)) + elif ( + llm_router is not None and data["model"] in llm_router.deployment_names + ): # model in router deployments, calling a specific deployment on the router (lowest priority) + llm_response = asyncio.create_task( + llm_router.aadapter_completion(**data, specific_deployment=True) + ) + elif user_model is not None: # `litellm --model ` + llm_response = asyncio.create_task(litellm.aadapter_completion(**data)) + else: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "completion: Invalid model name passed in model=" + + data.get("model", "") + }, + ) + + # Await the llm_response task + response = await llm_response + + hidden_params = getattr(response, "_hidden_params", {}) or {} + model_id = hidden_params.get("model_id", None) or "" + cache_key = hidden_params.get("cache_key", None) or "" + api_base = hidden_params.get("api_base", None) or "" + response_cost = hidden_params.get("response_cost", None) or "" + + ### ALERTING ### + asyncio.create_task( + proxy_logging_obj.update_request_status( + litellm_call_id=data.get("litellm_call_id", ""), status="success" + ) + ) + + verbose_proxy_logger.debug("final response: %s", response) + + fastapi_response.headers.update( + ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + response_cost=response_cost, + ) + ) + + verbose_proxy_logger.debug("\nResponse from Litellm:\n{}".format(response)) + return response + except Exception as e: + 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) + ) + ) + error_msg = f"{str(e)}" + raise ProxyException( + message=getattr(e, "message", error_msg), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", 500), + ) + + +class HttpPassThroughEndpointHelpers(BasePassthroughUtils): + @staticmethod + def get_response_headers( + headers: httpx.Headers, + litellm_call_id: Optional[str] = None, + custom_headers: Optional[dict] = None, + ) -> dict: + excluded_headers = {"transfer-encoding", "content-encoding"} + + 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: + return_headers.update(custom_headers) + + return return_headers + + @staticmethod + def get_endpoint_type(url: str) -> EndpointType: + parsed_url = urlparse(url) + if ( + ("generateContent") in url + or ("streamGenerateContent") in url + or ("rawPredict") in url + or ("streamRawPredict") in url + ): + return EndpointType.VERTEX_AI + elif parsed_url.hostname == "api.anthropic.com": + return EndpointType.ANTHROPIC + elif ( + parsed_url.hostname == "api.openai.com" + or parsed_url.hostname == "openai.azure.com" + or (parsed_url.hostname and "openai.com" in parsed_url.hostname) + ): + return EndpointType.OPENAI + return EndpointType.GENERIC + + @staticmethod + async def _make_non_streaming_http_request( + request: Request, + async_client: httpx.AsyncClient, + url: str, + headers: dict, + requested_query_params: Optional[dict] = None, + custom_body: Optional[dict] = None, + ) -> httpx.Response: + """ + Make a non-streaming HTTP request + + If request is GET, don't include a JSON body + """ + if request.method == "GET": + response = await async_client.request( + method=request.method, + url=url, + headers=headers, + params=requested_query_params, + ) + else: + response = await async_client.request( + method=request.method, + url=url, + headers=headers, + params=requested_query_params, + json=custom_body, + ) + return response + + @staticmethod + async def non_streaming_http_request_handler( + request: Request, + async_client: httpx.AsyncClient, + url: httpx.URL, + headers: dict, + requested_query_params: Optional[dict] = None, + _parsed_body: Optional[dict] = None, + ) -> httpx.Response: + """ + Handle non-streaming HTTP requests + + Handles special cases when GET requests, multipart/form-data requests, and generic httpx requests + """ + if request.method == "GET": + response = await async_client.request( + method=request.method, + url=url, + headers=headers, + params=requested_query_params, + ) + elif HttpPassThroughEndpointHelpers.is_multipart(request) is True: + return await HttpPassThroughEndpointHelpers.make_multipart_http_request( + request=request, + async_client=async_client, + url=url, + headers=headers, + requested_query_params=requested_query_params, + ) + else: + # Generic httpx method + response = await async_client.request( + method=request.method, + url=url, + headers=headers, + params=requested_query_params, + json=_parsed_body, + ) + return response + + @staticmethod + def is_multipart(request: Request) -> bool: + """Check if the request is a multipart/form-data request""" + return "multipart/form-data" in request.headers.get("content-type", "") + + @staticmethod + async def _build_request_files_from_upload_file( + upload_file: Union[UploadFile, StarletteUploadFile], + ) -> Tuple[Optional[str], bytes, Optional[str]]: + """Build a request files dict from an UploadFile object""" + file_content = await upload_file.read() + return (upload_file.filename, file_content, upload_file.content_type) + + @staticmethod + async def make_multipart_http_request( + request: Request, + async_client: httpx.AsyncClient, + url: httpx.URL, + headers: dict, + requested_query_params: Optional[dict] = None, + ) -> httpx.Response: + """Process multipart/form-data requests, handling both files and form fields""" + form_data = await request.form() + files = {} + form_data_dict = {} + + 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 + ) + else: + form_data_dict[field_name] = field_value + + # Remove content-type header - httpx will set it correctly with the new boundary + # when it creates the multipart body from files/data parameters + headers_copy = headers.copy() + headers_copy.pop("content-type", None) + + response = await async_client.request( + method=request.method, + url=url, + headers=headers_copy, + params=requested_query_params, + files=files, + data=form_data_dict, + ) + return response + + @staticmethod + def _init_kwargs_for_pass_through_endpoint( + request: Request, + user_api_key_dict: UserAPIKeyAuth, + passthrough_logging_payload: PassthroughStandardLoggingPayload, + logging_obj: LiteLLMLoggingObj, + _parsed_body: Optional[dict] = None, + litellm_call_id: Optional[str] = None, + ) -> dict: + """ + Filter out litellm params from the request body + """ + from litellm.types.utils import all_litellm_params + + _parsed_body = _parsed_body or {} + + litellm_params_in_body = {} + for k in all_litellm_params: + if k in _parsed_body: + litellm_params_in_body[k] = _parsed_body.pop(k, None) + + _metadata = dict( + StandardLoggingUserAPIKeyMetadata( + user_api_key_hash=user_api_key_dict.api_key, + user_api_key_alias=user_api_key_dict.key_alias, + user_api_key_user_email=user_api_key_dict.user_email, + user_api_key_user_id=user_api_key_dict.user_id, + user_api_key_team_id=user_api_key_dict.team_id, + user_api_key_org_id=user_api_key_dict.org_id, + user_api_key_project_id=user_api_key_dict.project_id, + user_api_key_team_alias=user_api_key_dict.team_alias, + user_api_key_end_user_id=user_api_key_dict.end_user_id, + user_api_key_request_route=user_api_key_dict.request_route, + 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_auth_metadata=user_api_key_dict.metadata, + ) + ) + + _metadata["user_api_key"] = user_api_key_dict.api_key + + litellm_metadata = litellm_params_in_body.pop("litellm_metadata", None) + metadata = litellm_params_in_body.pop("metadata", None) + if litellm_metadata: + _metadata.update(litellm_metadata) + if metadata: + _metadata.update(metadata) + + _metadata = _update_metadata_with_tags_in_header( + request=request, + metadata=_metadata, + ) + + kwargs = { + "litellm_params": { + **litellm_params_in_body, # type: ignore + "metadata": _metadata, + "proxy_server_request": { + "url": str(request.url), + "method": request.method, + "body": copy.copy(_parsed_body), # use copy instead of deepcopy + "headers": request.headers, + }, + }, + "call_type": "pass_through_endpoint", + "litellm_call_id": litellm_call_id, + "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: + """ + Helper function to construct the full target URL with subpath handling. + + Args: + base_target: The base target URL + subpath: The captured subpath from the request + include_subpath: Whether to include the subpath in the target URL + + Returns: + The constructed full target URL + """ + if not include_subpath: + return base_target + + if not subpath: + return base_target + + # Ensure base_target ends with / and subpath doesn't start with / + if not base_target.endswith("/"): + base_target = base_target + "/" + if subpath.startswith("/"): + subpath = subpath[1:] + + return base_target + subpath + + @staticmethod + def _update_stream_param_based_on_request_body( + parsed_body: dict, + stream: Optional[bool] = None, + ) -> Optional[bool]: + """ + If stream is provided in the request body, use it. + Otherwise, use the stream parameter passed to the `pass_through_request` function + """ + if "stream" in parsed_body: + return parsed_body.get("stream", stream) + return stream + + +async def pass_through_request( # noqa: PLR0915 + request: Request, + target: str, + custom_headers: dict, + user_api_key_dict: UserAPIKeyAuth, + custom_body: Optional[dict] = None, + forward_headers: Optional[bool] = False, + merge_query_params: Optional[bool] = False, + query_params: Optional[dict] = None, + default_query_params: Optional[dict] = None, + stream: Optional[bool] = None, + cost_per_request: Optional[float] = None, + custom_llm_provider: Optional[str] = None, + guardrails_config: Optional[dict] = None, +): + """ + Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called + + Args: + request: The incoming request + target: The target URL + custom_headers: The custom headers + user_api_key_dict: The user API key dictionary + custom_body: The custom body + forward_headers: Whether to forward headers + merge_query_params: Whether to merge query params + query_params: The query params + default_query_params: The default query params to be applied if not overridden by client + stream: Whether to stream the response + cost_per_request: Optional field - cost per request to the target endpoint + custom_llm_provider: Optional field - custom LLM provider for the endpoint + guardrails_config: Optional field - guardrails configuration for passthrough endpoint + """ + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.pass_through_endpoints.passthrough_guardrails import ( + PassthroughGuardrailHandler, + ) + from litellm.proxy.proxy_server import proxy_logging_obj + + ######################################################### + # Initialize variables + ######################################################### + litellm_call_id = str(uuid.uuid4()) + url: Optional[httpx.URL] = None + + # parsed request body + _parsed_body: Optional[dict] = None + # kwargs for pass through endpoint, contains metadata, litellm_params, call_type, litellm_call_id, passthrough_logging_payload + kwargs: Optional[dict] = None + + ######################################################### + try: + url = httpx.URL(target) + headers = custom_headers + headers = HttpPassThroughEndpointHelpers.forward_headers_from_request( + request_headers=dict(request.headers), + headers=headers, + forward_headers=forward_headers, + ) + + # Apply default query parameters if provided, regardless of merge_query_params setting + if default_query_params or merge_query_params: + # Determine what to merge based on settings + request_params = dict(request.query_params) if merge_query_params else {} + + # Create a new URL with the merged query params + url = url.copy_with( + query=urlencode( + HttpPassThroughEndpointHelpers.get_merged_query_parameters( + existing_url=url, + request_query_params=request_params, + default_query_params=default_query_params, + ) + ).encode("ascii") + ) + + endpoint_type: EndpointType = HttpPassThroughEndpointHelpers.get_endpoint_type( + str(url) + ) + + if custom_body: + _parsed_body = custom_body + else: + _parsed_body = await _read_request_body(request) + verbose_proxy_logger.debug( + "Pass through endpoint sending request to \nURL {}\nheaders: {}\nbody: {}\n".format( + url, headers, _parsed_body + ) + ) + + ### COLLECT GUARDRAILS FOR PASSTHROUGH ENDPOINT ### + # Passthrough endpoints are opt-in only for guardrails + # When enabled, collect guardrails from org/team/key levels + passthrough-specific + guardrails_to_run = PassthroughGuardrailHandler.collect_guardrails( + user_api_key_dict=user_api_key_dict, + passthrough_guardrails_config=guardrails_config, + ) + + # Add guardrails to metadata if any should run + if guardrails_to_run and len(guardrails_to_run) > 0: + if _parsed_body is None: + _parsed_body = {} + if "metadata" not in _parsed_body: + _parsed_body["metadata"] = {} + _parsed_body["metadata"]["guardrails"] = guardrails_to_run + verbose_proxy_logger.debug( + f"Added guardrails to passthrough request metadata: {guardrails_to_run}" + ) + + ## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it + start_time = datetime.now() + logging_obj = Logging( + model="unknown", + messages=[{"role": "user", "content": safe_dumps(_parsed_body)}], + stream=False, + call_type="pass_through_endpoint", + start_time=start_time, + litellm_call_id=litellm_call_id, + function_id="1245", + ) + + # Store passthrough guardrails config on logging_obj for field targeting + logging_obj.passthrough_guardrails_config = guardrails_config + + # Store logging_obj in data so guardrails can access it + if _parsed_body is None: + _parsed_body = {} + _parsed_body["litellm_logging_obj"] = logging_obj + + ### CALL HOOKS ### - modify incoming data / reject request before calling the model + _parsed_body = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=_parsed_body, + call_type="pass_through_endpoint", + ) + async_client_obj = get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": 600}, + ) + async_client = async_client_obj.client + passthrough_logging_payload = PassthroughStandardLoggingPayload( + url=str(url), + request_body=_parsed_body, + request_method=getattr(request, "method", None), + cost_per_request=cost_per_request, + ) + kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + user_api_key_dict=user_api_key_dict, + _parsed_body=_parsed_body, + passthrough_logging_payload=passthrough_logging_payload, + litellm_call_id=litellm_call_id, + request=request, + logging_obj=logging_obj, + ) + + # Store custom_llm_provider in kwargs and logging object if provided + if custom_llm_provider: + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + logging_obj.model_call_details["litellm_params"] = kwargs.get( + "litellm_params", {} + ) + + # done for supporting 'parallel_request_limiter.py' with pass-through endpoints + logging_obj.update_environment_variables( + model="unknown", + user="unknown", + optional_params={}, + litellm_params=kwargs["litellm_params"], + call_type="pass_through_endpoint", + ) + logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + + # combine url with query params for logging + requested_query_params: Optional[dict] = query_params or dict( + request.query_params + ) + + 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() + ) + + logging_url = str(url) + if requested_query_params_str: + if "?" in str(url): + logging_url = str(url) + "&" + requested_query_params_str + else: + logging_url = str(url) + "?" + requested_query_params_str + + logging_obj.pre_call( + input=[{"role": "user", "content": safe_dumps(_parsed_body)}], + api_key="", + additional_args={ + "complete_input_dict": _parsed_body, + "api_base": str(logging_url), + "headers": headers, + }, + ) + stream = ( + HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( + parsed_body=_parsed_body, + stream=stream, + ) + ) + + if stream: + req = async_client.build_request( + "POST", + url, + json=_parsed_body, + params=requested_query_params, + headers=headers, + ) + + response = await async_client.send(req, stream=stream) + + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + raise HTTPException( + status_code=e.response.status_code, detail=await e.response.aread() + ) + + return StreamingResponse( + PassThroughStreamingHandler.chunk_processor( + response=response, + request_body=_parsed_body, + litellm_logging_obj=logging_obj, + endpoint_type=endpoint_type, + start_time=start_time, + passthrough_success_handler_obj=pass_through_endpoint_logging, + url_route=str(url), + ), + headers=HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + ), + status_code=response.status_code, + ) + + 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) + + if _is_streaming_response(response) is True: + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + raise HTTPException( + status_code=e.response.status_code, detail=await e.response.aread() + ) + + return StreamingResponse( + PassThroughStreamingHandler.chunk_processor( + response=response, + request_body=_parsed_body, + litellm_logging_obj=logging_obj, + endpoint_type=endpoint_type, + start_time=start_time, + passthrough_success_handler_obj=pass_through_endpoint_logging, + url_route=str(url), + ), + headers=HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + ), + status_code=response.status_code, + ) + + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + 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) + + content = await response.aread() + + ## LOG SUCCESS + response_body: Optional[dict] = get_response_body(response) + passthrough_logging_payload["response_body"] = response_body + end_time = datetime.now() + asyncio.create_task( + pass_through_endpoint_logging.pass_through_async_success_handler( + httpx_response=response, + response_body=response_body, + url_route=str(url), + result="", + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + cache_hit=False, + request_body=_parsed_body, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + ) + + ## CUSTOM HEADERS - `x-litellm-*` + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=litellm_call_id, + model_id=None, + cache_key=None, + api_base=str(url._uri_reference), + ) + + return Response( + content=content, + status_code=response.status_code, + headers=HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + custom_headers=custom_headers, + ), + ) + except Exception as e: + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=litellm_call_id, + model_id=None, + cache_key=None, + 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) + ) + ) + + ######################################################### + # Monitoring: Trigger post_call_failure_hook + # for pass through endpoint failure + ######################################################### + request_payload: dict = _parsed_body or {} + # add user_api_key_dict, litellm_call_id, passthrough_logging_payloa for logging + if kwargs: + for key, value in kwargs.items(): + request_payload[key] = value + + if ( + "model" not in request_payload + and _parsed_body + and isinstance(_parsed_body, dict) + ): + request_payload["model"] = _parsed_body.get("model", "") + if "custom_llm_provider" not in request_payload and custom_llm_provider: + request_payload["custom_llm_provider"] = custom_llm_provider + + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_payload, + traceback_str=traceback.format_exc( + limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, + ), + ) + + ######################################################### + + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "message", str(getattr(e, "detail", str(e)))), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + headers=custom_headers, + ) + else: + error_msg = f"{str(e)}" + raise ProxyException( + message=getattr(e, "message", error_msg), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", 500), + headers=custom_headers, + ) + + +def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> dict: + """ + If tags are in the request headers, add them to the metadata + + Used for google and vertex JS SDKs, and Azure passthrough + Checks both 'tags' and 'x-litellm-tags' headers + """ + tags_to_add = [] + + # Check for 'tags' header first + _tags = request.headers.get("tags") + if _tags: + tags_to_add.extend([tag.strip() for tag in _tags.split(",")]) + + _tags = request.headers.get("x-litellm-tags") + if _tags: + tags_to_add.extend([tag.strip() for tag in _tags.split(",")]) + + # Only add tags key if there are tags to add + if tags_to_add: + if "tags" not in metadata: + metadata["tags"] = [] + metadata["tags"].extend(tags_to_add) + + return metadata + + +async def _parse_request_data_by_content_type( + request: Request, +) -> Tuple[Optional[Any], Optional[Any], Optional[Any], Optional[Any]]: + """ + Parse request data based on content type. + + Handles JSON, multipart/form-data, and URL-encoded form data. + + Returns: + Tuple of (query_params_data, custom_body_data, file_data, stream) + """ + content_type = request.headers.get("content-type", "") + + query_params_data = None + custom_body_data = None + file_data = None + stream = None + + if "application/json" in content_type: + # ✅ Handle JSON + try: + body = await request.json() + query_params_data = body.get("query_params") + custom_body_data = body.get("custom_body") + stream = body.get("stream") + except json.JSONDecodeError: + # Handle requests with no body (e.g., DELETE requests) + pass + elif "multipart/form-data" in content_type: + # ✅ Handle multipart form-data + form = await request.form() + if "query_params" in form: + form_value = form["query_params"] + if isinstance(form_value, str): + try: + query_params_data = json.loads(form_value) + except Exception: + query_params_data = form_value + else: + query_params_data = form_value + + if "custom_body" in form: + form_value = form["custom_body"] + if isinstance(form_value, str): + try: + custom_body_data = json.loads(form_value) + except Exception: + custom_body_data = form_value + else: + custom_body_data = form_value + + if "file" in form: + file_data = form["file"] # this is a Starlette UploadFile object + + elif "application/x-www-form-urlencoded" in content_type: + # ✅ Handle URL-encoded form data + form = await request.form() + query_params_data = form.get("query_params") + custom_body_data = form.get("custom_body") + + else: + # ✅ Fallback: maybe no body, just query params + query_params_data = dict(request.query_params) or None + + return query_params_data, custom_body_data, file_data, stream + + +def create_pass_through_route( + endpoint, + target: str, + custom_headers: Optional[dict] = None, + _forward_headers: Optional[bool] = False, + _merge_query_params: Optional[bool] = False, + dependencies: Optional[List] = None, + include_subpath: Optional[bool] = False, + cost_per_request: Optional[float] = None, + custom_llm_provider: Optional[str] = None, + is_streaming_request: Optional[bool] = False, + query_params: Optional[dict] = None, + default_query_params: Optional[dict] = None, + guardrails: Optional[Dict[str, Any]] = None, +): + # check if target is an adapter.py or a url + from litellm._uuid import uuid + from litellm.proxy.types_utils.utils import get_instance_fn + + try: + if isinstance(target, CustomLogger): + adapter = target + else: + adapter = get_instance_fn(value=target) + adapter_id = str(uuid.uuid4()) + litellm.adapters = [{"id": adapter_id, "adapter": adapter}] + + async def endpoint_func( # type: ignore + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + subpath: str = "", # captures sub-paths when include_subpath=True + custom_body: Optional[ + dict + ] = None, # accepted for signature compatibility with URL-based path; not forwarded because chat_completion_pass_through_endpoint does not support it + ): + return await chat_completion_pass_through_endpoint( + fastapi_response=fastapi_response, + request=request, + adapter_id=adapter_id, + user_api_key_dict=user_api_key_dict, + ) + + except Exception: + verbose_proxy_logger.debug("Defaulting to target being a url.") + + async def endpoint_func( # type: ignore + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + subpath: str = "", # captures sub-paths when include_subpath=True + custom_body: Optional[dict] = None, + ): + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + path = request.url.path + + # Parse request data based on content type + ( + query_params_data, + custom_body_data, + file_data, + stream, + ) = await _parse_request_data_by_content_type(request) + + if not InitPassThroughEndpointHelpers.is_registered_pass_through_route( + route=path + ): + raise HTTPException( + status_code=404, + detail=f"Pass-through endpoint {endpoint} not found. This could have been deleted or not yet added to the proxy.", + ) + + passthrough_params = ( + InitPassThroughEndpointHelpers.get_registered_pass_through_route( + route=path, method=request.method + ) + ) + target_params = { + "target": target, + "custom_headers": custom_headers, + "forward_headers": _forward_headers, + "merge_query_params": _merge_query_params, + "cost_per_request": cost_per_request, + "guardrails": None, + } + + if passthrough_params is not None: + target_params.update(passthrough_params.get("passthrough_params", {})) + + # Extract and cast parameters with proper types + param_target = target_params.get("target") or target + param_custom_headers = target_params.get("custom_headers", custom_headers) + param_forward_headers = target_params.get( + "forward_headers", _forward_headers + ) + param_merge_query_params = target_params.get( + "merge_query_params", _merge_query_params + ) + param_cost_per_request = target_params.get( + "cost_per_request", cost_per_request + ) + param_guardrails = target_params.get("guardrails", None) + param_default_query_params = target_params.get("default_query_params", None) + + # Construct the full target URL with subpath if needed + full_target = ( + HttpPassThroughEndpointHelpers.construct_target_url_with_subpath( + base_target=cast(str, param_target), + subpath=subpath, + include_subpath=include_subpath, + ) + ) + + # Ensure custom_headers is a dict + headers_dict = ( + param_custom_headers if isinstance(param_custom_headers, dict) else {} + ) + + # Ensure query_params and custom_body are dicts or None + final_query_params = ( + query_params_data if isinstance(query_params_data, dict) else {} + ) + if query_params: + final_query_params.update(query_params) + # When a caller (e.g. bedrock_proxy_route) supplies a pre-built + # body, use it instead of the body parsed from the raw request. + final_custom_body: Optional[dict] = None + if custom_body is not None: + final_custom_body = custom_body + elif isinstance(custom_body_data, dict): + final_custom_body = custom_body_data + + return await pass_through_request( # type: ignore + request=request, + target=full_target, + custom_headers=headers_dict, + user_api_key_dict=user_api_key_dict, + forward_headers=cast(Optional[bool], param_forward_headers), + merge_query_params=cast(Optional[bool], param_merge_query_params), + query_params=final_query_params, + default_query_params=cast(Optional[dict], param_default_query_params), + stream=is_streaming_request or stream, + custom_body=final_custom_body, + cost_per_request=cast(Optional[float], param_cost_per_request), + custom_llm_provider=custom_llm_provider, + guardrails_config=cast(Optional[dict], param_guardrails), + ) + + return endpoint_func + + +def create_websocket_passthrough_route( + endpoint: str, + target: str, + custom_headers: Optional[dict] = None, + _forward_headers: Optional[bool] = False, + dependencies: Optional[List] = None, + cost_per_request: Optional[float] = None, +): + """ + Create a WebSocket passthrough route function. + + Args: + endpoint: The endpoint path (for logging purposes) + target: The target WebSocket URL (e.g., "wss://api.example.com/ws") + custom_headers: Custom headers to include in the WebSocket connection + _forward_headers: Whether to forward incoming headers + dependencies: FastAPI dependencies to inject + + Returns: + A WebSocket passthrough function that can be registered with app.websocket() + """ + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket + + async def websocket_endpoint_func( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), + **kwargs, # For additional query parameters + ): + """ + WebSocket passthrough endpoint function. + + This function handles the WebSocket connection by: + 1. Accepting the incoming WebSocket connection + 2. Establishing a connection to the target WebSocket + 3. Forwarding messages bidirectionally + 4. Handling connection cleanup + """ + return await websocket_passthrough_request( + websocket=websocket, + target=target, + custom_headers=custom_headers or {}, + user_api_key_dict=user_api_key_dict, + forward_headers=_forward_headers, + endpoint=endpoint, + cost_per_request=cost_per_request, + accept_websocket=True, # Generic usage should accept the WebSocket + ) + + return websocket_endpoint_func + + +async def websocket_passthrough_request( # noqa: PLR0915 + websocket: WebSocket, + target: str, + custom_headers: dict, + user_api_key_dict: UserAPIKeyAuth, + forward_headers: Optional[bool] = False, + endpoint: Optional[str] = None, + cost_per_request: Optional[float] = None, + accept_websocket: bool = True, +): + """ + WebSocket passthrough request handler. + + Args: + websocket: The incoming WebSocket connection + target: The target WebSocket URL + custom_headers: Custom headers to include in the connection + user_api_key_dict: The user API key dictionary + forward_headers: Whether to forward incoming headers + endpoint: The endpoint path (for logging purposes) + cost_per_request: Optional field - cost per request to the target endpoint + """ + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.proxy_server import proxy_logging_obj + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + PassthroughStandardLoggingPayload, + ) + + # Initialize tracking variables + start_time = datetime.now() + websocket_messages: list[dict[str, Any]] = [] + litellm_call_id = str(uuid.uuid4()) + + verbose_proxy_logger.info( + f"WebSocket passthrough ({endpoint}): Starting WebSocket connection to {target}" + ) + + # Only accept the WebSocket if requested (for generic usage) + if accept_websocket: + await websocket.accept() + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): WebSocket connection accepted" + ) + + # Prepare headers for the upstream connection + upstream_headers = custom_headers.copy() + + if forward_headers: + # Forward relevant headers from the incoming request + incoming_headers = dict(websocket.headers) + for header_name, header_value in incoming_headers.items(): + # Only forward certain headers to avoid conflicts + if header_name.lower() in [ + "authorization", + "x-api-key", + "x-goog-user-project", + ]: + upstream_headers[header_name] = header_value + + # Initialize logging object similar to HTTP passthrough + logging_obj = Logging( + model="unknown", + messages=[{"role": "user", "content": "WebSocket connection"}], + stream=True, # WebSockets are inherently streaming + call_type="pass_through_endpoint", + start_time=start_time, + litellm_call_id=litellm_call_id, + function_id="websocket_passthrough", + ) + + # Create passthrough logging payload + passthrough_logging_payload = PassthroughStandardLoggingPayload( + url=target, + request_body={}, # WebSocket doesn't have a traditional request body + request_method="WEBSOCKET", + cost_per_request=cost_per_request, + ) + + # Create a dummy request object for WebSocket connections to maintain compatibility + # with the existing _init_kwargs_for_pass_through_endpoint function + class DummyRequest: + def __init__( + self, url: str, method: str = "WEBSOCKET", headers: Optional[dict] = None + ): + self.url = url + self.method = method + self.headers = headers or {} + + def __str__(self): + return f"DummyRequest(url={self.url}, method={self.method})" + + dummy_request = DummyRequest( + url=target, + method="WEBSOCKET", + headers=dict(websocket.headers) if hasattr(websocket, "headers") else {}, + ) + + # Initialize kwargs for logging using the same pattern as HTTP passthrough + kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + user_api_key_dict=user_api_key_dict, + _parsed_body={}, # WebSocket doesn't have a traditional request body + passthrough_logging_payload=passthrough_logging_payload, + litellm_call_id=litellm_call_id, + request=dummy_request, # type: ignore + logging_obj=logging_obj, + ) + + # Update logging environment variables + logging_obj.update_environment_variables( + model="unknown", + user="unknown", + optional_params={}, + litellm_params=dict(kwargs.get("litellm_params", {})), + call_type="pass_through_endpoint", + ) + logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + + # Pre-call logging + logging_obj.pre_call( + input=[{"role": "user", "content": "WebSocket connection"}], + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": target, + "headers": upstream_headers, + }, + ) + + ### CALL HOOKS ### - modify incoming data / reject request before calling the model + websocket_data: dict[str, Any] = {} + websocket_data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=websocket_data, + call_type="pass_through_endpoint", + ) + + try: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Establishing upstream connection to {target}" + ) + async with connect( + target, + additional_headers=upstream_headers, + ) as upstream_ws: + verbose_proxy_logger.info( + f"WebSocket passthrough ({endpoint}): Upstream connection established successfully" + ) + + async def forward_client_to_upstream() -> None: + """Forward messages from client to upstream WebSocket""" + try: + while True: + message = await websocket.receive() + message_type = message.get("type") + if message_type == "websocket.disconnect": + await upstream_ws.close() + break + + text_data = message.get("text") + bytes_data = message.get("bytes") + + if text_data is not None: + # Try to extract model from client setup message for Vertex AI Live + if endpoint and "/vertex_ai/live" in endpoint: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Processing client message for model extraction" + ) + try: + client_message = json.loads(text_data) + if ( + isinstance(client_message, dict) + and "setup" in client_message + ): + setup_data = client_message["setup"] + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Found setup data in client message: {setup_data}" + ) + if ( + isinstance(setup_data, dict) + and "model" in setup_data + ): + extracted_model = ( + _extract_model_from_vertex_ai_setup( + setup_data + ) + ) + if extracted_model: + kwargs["model"] = extracted_model + kwargs[ + "custom_llm_provider" + ] = "vertex_ai-language-models" + # Update logging object with correct model + logging_obj.model = extracted_model + logging_obj.model_call_details[ + "model" + ] = extracted_model + logging_obj.model_call_details[ + "custom_llm_provider" + ] = "vertex_ai" + verbose_proxy_logger.info( + f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from client setup message" + ) + else: + verbose_proxy_logger.warning( + f"WebSocket passthrough ({endpoint}): Failed to extract model from client setup data: {setup_data}" + ) + else: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Setup data does not contain model field: {setup_data}" + ) + else: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Client message does not contain setup data" + ) + except (json.JSONDecodeError, KeyError, TypeError) as e: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Client message is not a valid setup message: {e}" + ) + pass # Not a JSON message or doesn't contain setup data + + await upstream_ws.send(text_data) + elif bytes_data is not None: + await upstream_ws.send(bytes_data) + except asyncio.CancelledError: + raise + except Exception: + verbose_proxy_logger.exception( + f"WebSocket passthrough ({endpoint}): error forwarding client message" + ) + await upstream_ws.close() + + async def forward_upstream_to_client() -> None: + """Forward messages from upstream to client WebSocket""" + try: + # Wait for the first response from upstream + raw_response = await upstream_ws.recv(decode=False) + # Ensure raw_response is bytes before decoding + if isinstance(raw_response, str): + raw_response = raw_response.encode("ascii") + setup_response = json.loads(raw_response.decode("ascii")) + verbose_proxy_logger.debug(f"Setup response: {setup_response}") + + # Extract model and provider from setup response for Vertex AI Live + if endpoint and "/vertex_ai/live" in endpoint: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Processing server setup response for model extraction" + ) + extracted_model = _extract_model_from_vertex_ai_setup( + setup_response + ) + if extracted_model: + kwargs["model"] = extracted_model + kwargs["custom_llm_provider"] = "vertex_ai_language_models" + # Update logging object with correct model + logging_obj.model = extracted_model + logging_obj.model_call_details["model"] = extracted_model + logging_obj.model_call_details[ + "custom_llm_provider" + ] = "vertex_ai_language_models" + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Successfully extracted model '{extracted_model}' and set provider to 'vertex_ai' from server setup response" + ) + else: + verbose_proxy_logger.warning( + f"WebSocket passthrough ({endpoint}): Failed to extract model from server setup response: {setup_response}" + ) + else: + verbose_proxy_logger.debug( + f"WebSocket passthrough ({endpoint}): Not a Vertex AI Live endpoint, skipping model extraction" + ) + + # Send the setup response to the client + await websocket.send_text(json.dumps(setup_response)) + + # Now continuously forward messages from upstream to client + async for upstream_message in upstream_ws: + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + # Parse and collect for cost tracking + try: + message_data = json.loads(upstream_message.decode()) + websocket_messages.append(message_data) + except (json.JSONDecodeError, UnicodeDecodeError): + pass + else: + await websocket.send_text(upstream_message) + # Parse and collect for cost tracking + try: + message_data = json.loads(upstream_message) + websocket_messages.append(message_data) + except json.JSONDecodeError: + pass + + except (ConnectionClosedOK, ConnectionClosedError) as e: + verbose_proxy_logger.debug( + f"Upstream WebSocket connection closed: {e}" + ) + pass + except asyncio.CancelledError: + verbose_proxy_logger.debug( + "asyncio.CancelledError in forward_upstream_to_client" + ) + raise + except Exception as e: + verbose_proxy_logger.debug( + f"Exception in forward_upstream_to_client: {e}" + ) + verbose_proxy_logger.exception( + f"WebSocket passthrough ({endpoint}): error forwarding upstream message" + ) + raise + + # Create tasks for bidirectional message forwarding + tasks = [ + asyncio.create_task(forward_client_to_upstream()), + asyncio.create_task(forward_upstream_to_client()), + ] + + done, pending = await asyncio.wait( + tasks, return_when=asyncio.FIRST_COMPLETED + ) + + # Cancel remaining tasks + for task in pending: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Check for exceptions in completed tasks + for task in done: + exception = task.exception() + if exception is not None: + raise exception + + end_time = datetime.now() + + # Update passthrough logging payload with response data + passthrough_logging_payload["response_body"] = websocket_messages # type: ignore + passthrough_logging_payload["end_time"] = end_time # type: ignore + + # Remove logging_obj from kwargs to avoid duplicate keyword argument + success_kwargs = kwargs.copy() + success_kwargs.pop("logging_obj", None) + + # # Add user authentication context for database logging + # if user_api_key_dict: + # success_kwargs.setdefault('litellm_params', {}) + # success_kwargs['litellm_params'].update({ + # 'proxy_server_request': { + # 'body': { + # 'user': user_api_key_dict.user_id, + # 'team_id': user_api_key_dict.team_id, + # 'end_user_id': user_api_key_dict.end_user_id, + # } + # } + # }) + # # Also add the user_api_key for direct access + # success_kwargs['user_api_key'] = user_api_key_dict.api_key + + # Create a dummy httpx.Response for WebSocket connections + class MockWebSocketResponse: + def __init__(self, target_url: str): + self.status_code = 200 + self.text = "WebSocket connection successful" + self.headers: dict[str, str] = {} + self.request = MockWebSocketRequest(target_url) + + class MockWebSocketRequest: + def __init__(self, target_url: str): + self.method = "WEBSOCKET" + self.url = target_url + + mock_response = MockWebSocketResponse(target) + + # Use the same success handler as HTTP passthrough endpoints + asyncio.create_task( + pass_through_endpoint_logging.pass_through_async_success_handler( + httpx_response=mock_response, # type: ignore + response_body=websocket_messages, # type: ignore + url_route=endpoint or "", + result="websocket_connection_successful", + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + cache_hit=False, + request_body={}, + **success_kwargs, + ) + ) + + # Call the proxy logging success hook + if proxy_logging_obj: + await proxy_logging_obj.post_call_success_hook( + data={}, + user_api_key_dict=user_api_key_dict, + response={"status": "websocket_connection_successful"}, # type: ignore + ) + + except InvalidStatus as exc: + verbose_proxy_logger.exception( + f"WebSocket passthrough ({endpoint}): upstream rejected WebSocket connection" + ) + + # Prepare request payload for logging + request_payload = {} + if kwargs: + for key, value in kwargs.items(): + request_payload[key] = value + + # Log the connection failure using the same pattern as HTTP + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=exc, + request_data=request_payload, + traceback_str=traceback.format_exc( + limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, + ), + ) + + if websocket.client_state != WebSocketState.DISCONNECTED: + await websocket.close( + code=getattr(exc, "status_code", 1011), + reason="Upstream connection rejected", + ) + except Exception as e: + verbose_proxy_logger.exception( + f"WebSocket passthrough ({endpoint}): unexpected error while proxying WebSocket" + ) + + # Prepare request payload for logging + request_payload = {} + if kwargs: + for key, value in kwargs.items(): + request_payload[key] = value + + # Log the unexpected error using the same pattern as HTTP + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_payload, + traceback_str=traceback.format_exc( + limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, + ), + ) + + if websocket.client_state != WebSocketState.DISCONNECTED: + await websocket.close(code=1011, reason="WebSocket passthrough error") + finally: + if websocket.client_state != WebSocketState.DISCONNECTED: + await websocket.close() + + +def _is_streaming_response(response: httpx.Response) -> bool: + _content_type = response.headers.get("content-type") + if _content_type is not None and "text/event-stream" in _content_type: + return True + return False + + +def _extract_model_from_vertex_ai_setup(setup_response: dict) -> Optional[str]: + """ + Extract the model name from Vertex AI Live setup response. + + The setup response can contain a model field in two formats: + 1. Direct: {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"} + 2. Nested: {"setup": {"model": "projects/.../models/gemini-2.0-flash-live-preview-04-09"}} + + We extract just the model name: "gemini-2.0-flash-live-preview-04-09" + """ + try: + # Handle both direct model field and nested setup.model field + model_path = None + if isinstance(setup_response, dict): + if "model" in setup_response: + model_path = setup_response["model"] + elif ( + "setup" in setup_response + and isinstance(setup_response["setup"], dict) + and "model" in setup_response["setup"] + ): + model_path = setup_response["setup"]["model"] + + if isinstance(model_path, str) and "/models/" in model_path: + # Extract the model name after the last "/models/" + model_name = model_path.split("/models/")[-1] + return model_name + except Exception as e: + verbose_proxy_logger.debug(f"Error extracting model from setup response: {e}") + return None + + +class SafeRouteAdder: + """ + Wrapper class for adding routes to FastAPI app. + Only adds routes if they don't already exist on the app. + """ + + @staticmethod + def _is_path_registered(app: FastAPI, path: str, methods: List[str]) -> bool: + """ + Check if a path with any of the specified methods is already registered on the app. + + Args: + app: The FastAPI application instance + path: The path to check (e.g., "/v1/chat/completions") + methods: List of HTTP methods to check (e.g., ["GET", "POST"]) + + Returns: + True if the path is already registered with any of the methods, False otherwise + """ + for route in app.routes: + # Use getattr to safely access route attributes + route_path = getattr(route, "path", None) + route_methods = getattr(route, "methods", None) + + if route_path == path and route_methods is not None: + # Check if any of the methods overlap + if any(method in route_methods for method in methods): + return True + return False + + @staticmethod + def add_api_route_if_not_exists( + app: FastAPI, + path: str, + endpoint: Any, + methods: List[str], + dependencies: Optional[List] = None, + ) -> bool: + """ + Add an API route to the app only if it doesn't already exist. + + Args: + app: The FastAPI application instance + path: The path for the route + endpoint: The endpoint function/callable + methods: List of HTTP methods + dependencies: Optional list of dependencies + + Returns: + True if route was added, False if it already existed + """ + if SafeRouteAdder._is_path_registered(app=app, path=path, methods=methods): + verbose_proxy_logger.debug( + "Skipping route registration - path %s with methods %s already registered on app", + path, + methods, + ) + return False + + app.add_api_route( + path=path, + endpoint=endpoint, + methods=methods, + dependencies=dependencies, + ) + verbose_proxy_logger.debug( + "Successfully added route: %s with methods %s", + path, + methods, + ) + return True + + +class InitPassThroughEndpointHelpers: + @staticmethod + def add_exact_path_route( + app: FastAPI, + path: str, + target: str, + custom_headers: Optional[dict], + forward_headers: Optional[bool], + merge_query_params: Optional[bool], + dependencies: Optional[List], + cost_per_request: Optional[float], + endpoint_id: str, + guardrails: Optional[dict] = None, + methods: Optional[List[str]] = None, + default_query_params: Optional[dict] = None, + ): + """Add exact path route for pass-through endpoint""" + # Default to all methods if none specified (backward compatibility) + if methods is None or len(methods) == 0: + methods = ["GET", "POST", "PUT", "DELETE", "PATCH"] + + # Create route key that includes methods for uniqueness + methods_str = ",".join(sorted(methods)) + route_key = f"{endpoint_id}:exact:{path}:{methods_str}" + + # Check if this exact route is already registered + if route_key in _registered_pass_through_routes: + verbose_proxy_logger.debug( + "Updating duplicate exact pass through endpoint: %s with methods %s (already registered)", + path, + methods, + ) + + verbose_proxy_logger.debug( + "adding exact pass through endpoint: %s, methods: %s, dependencies: %s", + path, + methods, + dependencies, + ) + + # Use SafeRouteAdder to only add route if it doesn't exist on the app + SafeRouteAdder.add_api_route_if_not_exists( + app=app, + path=path, + endpoint=create_pass_through_route( # type: ignore + path, + target, + custom_headers, + forward_headers, + merge_query_params, + dependencies, + cost_per_request=cost_per_request, + default_query_params=default_query_params, + guardrails=guardrails, + ), + methods=methods, + dependencies=dependencies, + ) + + # Always register/update the route metadata (headers, target) even if FastAPI route exists + _registered_pass_through_routes[route_key] = { + "endpoint_id": endpoint_id, + "path": path, + "type": "exact", + "methods": methods, + "passthrough_params": { + "target": target, + "custom_headers": custom_headers, + "forward_headers": forward_headers, + "merge_query_params": merge_query_params, + "default_query_params": default_query_params, + "dependencies": dependencies, + "cost_per_request": cost_per_request, + "guardrails": guardrails, + }, + } + + @staticmethod + def add_subpath_route( + app: FastAPI, + path: str, + target: str, + custom_headers: Optional[dict], + forward_headers: Optional[bool], + merge_query_params: Optional[bool], + dependencies: Optional[List], + cost_per_request: Optional[float], + endpoint_id: str, + guardrails: Optional[dict] = None, + methods: Optional[List[str]] = None, + default_query_params: Optional[dict] = None, + ): + """Add wildcard route for sub-paths""" + # Default to all methods if none specified (backward compatibility) + if methods is None or len(methods) == 0: + methods = ["GET", "POST", "PUT", "DELETE", "PATCH"] + + wildcard_path = f"{path}/{{subpath:path}}" + methods_str = ",".join(sorted(methods)) + route_key = f"{endpoint_id}:subpath:{path}:{methods_str}" + + # Check if this subpath route is already registered + if route_key in _registered_pass_through_routes: + verbose_proxy_logger.debug( + "Updating duplicate wildcard pass through endpoint: %s with methods %s (already registered)", + wildcard_path, + methods, + ) + + verbose_proxy_logger.debug( + "adding wildcard pass through endpoint: %s, methods: %s, dependencies: %s", + wildcard_path, + methods, + dependencies, + ) + + # Use SafeRouteAdder to only add route if it doesn't exist on the app + SafeRouteAdder.add_api_route_if_not_exists( + app=app, + path=wildcard_path, + endpoint=create_pass_through_route( # type: ignore + path, + target, + custom_headers, + forward_headers, + merge_query_params, + dependencies, + include_subpath=True, + cost_per_request=cost_per_request, + default_query_params=default_query_params, + guardrails=guardrails, + ), + methods=methods, + dependencies=dependencies, + ) + + # Register the route to prevent duplicates only if it was added + _registered_pass_through_routes[route_key] = { + "endpoint_id": endpoint_id, + "path": path, + "type": "subpath", + "methods": methods, + "passthrough_params": { + "target": target, + "custom_headers": custom_headers, + "forward_headers": forward_headers, + "merge_query_params": merge_query_params, + "default_query_params": default_query_params, + "dependencies": dependencies, + "cost_per_request": cost_per_request, + "guardrails": guardrails, + }, + } + + @staticmethod + 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 + ] + for key in keys_to_remove: + del _registered_pass_through_routes[key] + verbose_proxy_logger.debug( + "Removed pass-through route from registry: %s", key + ) + + @staticmethod + def clear_all_pass_through_routes(): + """Clear all pass-through routes from the registry""" + _registered_pass_through_routes.clear() + + @staticmethod + def get_all_registered_pass_through_routes() -> List[str]: + """Get all registered pass-through endpoints from the registry""" + return list(_registered_pass_through_routes.keys()) + + @staticmethod + def _build_full_path_with_root(path: str) -> str: + """ + Build full path by prepending server root path if needed. + + Args: + path: The relative path to build + + Returns: + Full path with server root prepended (if root is not "/") + """ + root_path = get_server_root_path() + if root_path == "/": + return path + return f"{root_path}{path}" + + @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 + """ + ## CHECK IF MAPPED PASS THROUGH ENDPOINT + for mapped_route in LiteLLMRoutes.mapped_pass_through_routes.value: + if route.startswith(mapped_route): + return True + + # Fast path: check if any registered route key contains this path + # Keys are in format: "{endpoint_id}:exact:{path}:{methods}" or "{endpoint_id}:subpath:{path}:{methods}" + # For backward compatibility, also support old 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(":", 3) # Split into [endpoint_id, type, path, methods?] + if len(parts) >= 3: + route_type = parts[1] + registered_path = ( + InitPassThroughEndpointHelpers._build_full_path_with_root(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 + + @staticmethod + def get_registered_pass_through_route( + route: str, method: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """Get passthrough params for a given route and optionally filter by HTTP method""" + for key in _registered_pass_through_routes.keys(): + parts = key.split(":", 3) # Split into [endpoint_id, type, path, methods?] + if len(parts) >= 3: + route_type = parts[1] + registered_path = ( + InitPassThroughEndpointHelpers._build_full_path_with_root(parts[2]) + ) + + # Get the methods for this route + route_methods = _registered_pass_through_routes[key].get("methods", []) + + # Check if path matches + path_matches = False + if route_type == "exact" and route == registered_path: + path_matches = True + elif route_type == "subpath": + if route == registered_path or route.startswith( + registered_path + "/" + ): + path_matches = True + + # If path matches and method filter is provided, check if method is allowed + if path_matches: + if ( + method is None + or not route_methods + or method in route_methods + ): + return _registered_pass_through_routes[key] + + return None + + +def _get_combined_pass_through_endpoints( + pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]], + config_pass_through_endpoints: List[Dict], +): + """Get combined pass-through endpoints from db + config""" + return pass_through_endpoints + config_pass_through_endpoints + + +async def initialize_pass_through_endpoints( + pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]], +): + """ + 1. Create a global list of pass-through endpoints (db + config) + 2. Clear all existing pass-through endpoints from the FastAPI app routes + 3. Add new endpoints to the in-memory registry + + Initialize a list of pass-through endpoints by adding them to the FastAPI app routes + + Args: + pass_through_endpoints: List of pass-through endpoints to initialize + + Returns: + None + """ + from litellm._uuid import uuid + + verbose_proxy_logger.debug("initializing pass through endpoints") + from litellm.proxy._types import CommonProxyErrors, LiteLLMRoutes + from litellm.proxy.proxy_server import ( + app, + config_passthrough_endpoints, + premium_user, + ) + + ## get combined pass-through endpoints from db + config + combined_pass_through_endpoints: List[Union[Dict, PassThroughGenericEndpoint]] + + if config_passthrough_endpoints is not None: + combined_pass_through_endpoints = _get_combined_pass_through_endpoints( # type: ignore + pass_through_endpoints, config_passthrough_endpoints + ) + else: + combined_pass_through_endpoints = pass_through_endpoints # type: ignore + + ## clear all existing pass-through endpoints from the FastAPI app routes + # InitPassThroughEndpointHelpers.clear_all_pass_through_routes() + + # get a list of all registered pass-through endpoints + # mark the ones that are visited in the list + # remove the ones that are not visited from the list + registered_pass_through_endpoints = ( + InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes() + ) + + visited_endpoints = set() + + for endpoint in combined_pass_through_endpoints: + if isinstance(endpoint, PassThroughGenericEndpoint): + endpoint = endpoint.model_dump() + + # Auto-generate ID for backwards compatibility if not present + if endpoint.get("id") is None: + endpoint["id"] = str(uuid.uuid4()) + + # Get the endpoint_id as a string (guaranteed to be set at this point) + endpoint_id: str = endpoint["id"] + + _target = endpoint.get("target", None) + _path: Optional[str] = endpoint.get("path", None) + 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 + ) + _forward_headers = endpoint.get("forward_headers", None) + _merge_query_params = endpoint.get("merge_query_params", None) + _default_query_params = endpoint.get("default_query_params", None) + _auth = endpoint.get("auth", None) + _dependencies = None + if _auth is not None and str(_auth).lower() == "true": + if premium_user is not True: + raise ValueError( + "Error Setting Authentication on Pass Through Endpoint: {}".format( + CommonProxyErrors.not_premium_user.value + ) + ) + _dependencies = [Depends(user_api_key_auth)] + LiteLLMRoutes.openai_routes.value.append(_path) + + if _target is None: + continue + + # Get guardrails config if present + _guardrails = endpoint.get("guardrails", None) + + # Get methods list if present (None means all methods for backward compatibility) + _methods = endpoint.get("methods", None) + + # Add exact path route + verbose_proxy_logger.debug( + "Initializing pass through endpoint: %s (ID: %s)", _path, endpoint_id + ) + InitPassThroughEndpointHelpers.add_exact_path_route( + app=app, + path=_path, + target=_target, + custom_headers=_custom_headers, + forward_headers=_forward_headers, + merge_query_params=_merge_query_params, + dependencies=_dependencies, + cost_per_request=endpoint.get("cost_per_request", None), + endpoint_id=endpoint_id, + guardrails=_guardrails, + methods=_methods, + default_query_params=_default_query_params, + ) + + # Generate route key with methods for tracking + methods_for_key = ( + _methods if _methods else ["GET", "POST", "PUT", "DELETE", "PATCH"] + ) + methods_str = ",".join(sorted(methods_for_key)) + visited_endpoints.add(f"{endpoint_id}:exact:{_path}:{methods_str}") + + # Add wildcard route for sub-paths + if endpoint.get("include_subpath", False) is True: + InitPassThroughEndpointHelpers.add_subpath_route( + app=app, + path=_path, + target=_target, + custom_headers=_custom_headers, + forward_headers=_forward_headers, + merge_query_params=_merge_query_params, + dependencies=_dependencies, + cost_per_request=endpoint.get("cost_per_request", None), + endpoint_id=endpoint_id, + guardrails=_guardrails, + methods=_methods, + default_query_params=_default_query_params, + ) + + visited_endpoints.add(f"{endpoint_id}:subpath:{_path}:{methods_str}") + + verbose_proxy_logger.debug( + "Added new pass through endpoint: %s (ID: %s)", _path, endpoint_id + ) + + # remove the ones that are not visited from the list + for endpoint_key in registered_pass_through_endpoints: + if endpoint_key not in visited_endpoints: + InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_key) + + +def _get_pass_through_endpoints_from_config() -> List[PassThroughGenericEndpoint]: + """ + Get pass-through endpoints defined in the config file. + These are read-only and cannot be edited via the UI. + Malformed endpoints are logged and skipped; they do not crash the function. + """ + from pydantic import ValidationError + + from litellm.proxy.proxy_server import config_passthrough_endpoints + + if config_passthrough_endpoints is None or len(config_passthrough_endpoints) == 0: + return [] + + returned_endpoints: List[PassThroughGenericEndpoint] = [] + for endpoint in config_passthrough_endpoints: + try: + if isinstance(endpoint, dict): + endpoint_dict = dict(endpoint) + endpoint_dict["is_from_config"] = True + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + elif isinstance(endpoint, PassThroughGenericEndpoint): + # Create a copy with is_from_config=True + endpoint_dict = endpoint.model_dump() + endpoint_dict["is_from_config"] = True + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + except ValidationError as e: + verbose_proxy_logger.warning( + "Skipping malformed pass-through endpoint from config: %s", + e, + exc_info=False, + ) + + return returned_endpoints + + +async def _get_pass_through_endpoints_from_db( + endpoint_id: Optional[str] = None, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, +) -> List[PassThroughGenericEndpoint]: + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.proxy_server import get_config_general_settings + + try: + if user_api_key_dict is None: + user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + response: ConfigFieldInfo = await get_config_general_settings( + field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict + ) + except Exception: + return [] + + pass_through_endpoint_data: Optional[List] = response.field_value + if pass_through_endpoint_data is None: + return [] + + returned_endpoints: List[PassThroughGenericEndpoint] = [] + if endpoint_id is None: + # Return all endpoints from DB, mark as not from config + for endpoint in pass_through_endpoint_data: + if isinstance(endpoint, dict): + endpoint_dict = dict(endpoint) + endpoint_dict["is_from_config"] = False + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + elif isinstance(endpoint, PassThroughGenericEndpoint): + endpoint_dict = endpoint.model_dump() + endpoint_dict["is_from_config"] = False + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + else: + # Find specific endpoint by ID + found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) + if found_endpoint is not None: + endpoint_dict = ( + found_endpoint.model_dump() + if isinstance(found_endpoint, PassThroughGenericEndpoint) + else dict(found_endpoint) + ) + endpoint_dict["is_from_config"] = False + returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + + return returned_endpoints + + +async def _filter_endpoints_by_team_allowed_routes( + team_id: str, + pass_through_endpoints: List[PassThroughGenericEndpoint], + prisma_client, +) -> List[PassThroughGenericEndpoint]: + """ + Filter pass-through endpoints based on team's allowed_passthrough_routes metadata. + + Args: + team_id: The team ID to check permissions for + pass_through_endpoints: List of endpoints to filter + prisma_client: Database client + + Returns: + Filtered list of endpoints based on team permissions + + Raises: + HTTPException: If team is not found + """ + # retrieve team from db + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + ) + if team is None: + raise HTTPException( + status_code=404, + detail={"error": "Team not found"}, + ) + + # retrieve team metadata + team_metadata = team.metadata + if ( + team_metadata is not None + and team_metadata.get("allowed_passthrough_routes") is not None + ): + ## FILTER pass_through_endpoints by allowed_passthrough_routes + pass_through_endpoints = [ + endpoint + for endpoint in pass_through_endpoints + if endpoint.path in team_metadata.get("allowed_passthrough_routes") + ] + + return pass_through_endpoints + + +@router.get( + "/config/pass_through_endpoint", + dependencies=[Depends(user_api_key_auth)], + response_model=PassThroughEndpointResponse, +) +@router.get( + "/config/pass_through_endpoint/team/{team_id}", + dependencies=[Depends(user_api_key_auth)], + response_model=PassThroughEndpointResponse, +) +async def get_pass_through_endpoints( + endpoint_id: Optional[str] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + team_id: Optional[str] = None, +): + """ + GET configured pass through endpoint. + + If no endpoint_id given, return all configured endpoints. + """ ## Get existing pass-through endpoint field value + from litellm.proxy._types import CommonProxyErrors + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + # Get endpoints from DB (editable via UI) + db_endpoints = await _get_pass_through_endpoints_from_db( + endpoint_id=endpoint_id, user_api_key_dict=user_api_key_dict + ) + + # Get endpoints from config file (read-only, not editable via UI) + config_endpoints = _get_pass_through_endpoints_from_config() + + # Merge: config endpoints not in DB + all DB endpoints (DB overrides config for same path) + db_paths = {ep.path for ep in db_endpoints} + config_only_endpoints = [ep for ep in config_endpoints if ep.path not in db_paths] + if endpoint_id is not None: + # When filtering by endpoint_id, only return if found in DB (config endpoints use generated IDs) + pass_through_endpoints = db_endpoints + else: + pass_through_endpoints = config_only_endpoints + db_endpoints + + if team_id is not None: + pass_through_endpoints = await _filter_endpoints_by_team_allowed_routes( + team_id=team_id, + pass_through_endpoints=pass_through_endpoints, + prisma_client=prisma_client, + ) + + return PassThroughEndpointResponse(endpoints=pass_through_endpoints) + + +@router.post( + "/config/pass_through_endpoint/{endpoint_id}", + dependencies=[Depends(user_api_key_auth)], +) +async def update_pass_through_endpoints( + endpoint_id: str, + data: PassThroughGenericEndpoint, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update a pass-through endpoint by ID. + """ + from litellm.proxy.proxy_server import ( + get_config_general_settings, + update_config_general_settings, + ) + + ## Get existing pass-through endpoint field value + try: + response: ConfigFieldInfo = await get_config_general_settings( + field_name="pass_through_endpoints", user_api_key_dict=user_api_key_dict + ) + except Exception: + raise HTTPException( + status_code=404, + detail={"error": "No pass-through endpoints found"}, + ) + + pass_through_endpoint_data: Optional[List] = response.field_value + if pass_through_endpoint_data is None: + raise HTTPException( + status_code=404, + detail={"error": "No pass-through endpoints found"}, + ) + + # Find the endpoint to update + found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) + + if found_endpoint is None: + raise HTTPException( + status_code=404, + detail={"error": f"Endpoint with ID '{endpoint_id}' not found"}, + ) + + # 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 + ) + if _endpoint.id == endpoint_id: + endpoint_index = idx + break + + if endpoint_index is None: + raise HTTPException( + status_code=404, + 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 + # Exclude is_from_config as it's a response-only field (computed at read time) + update_data = data.model_dump(exclude_none=True, exclude={"is_from_config"}) + + # Start with existing endpoint data + endpoint_dict = found_endpoint.model_dump() + + # Update with new data (only non-None values) + endpoint_dict.update(update_data) + + # Preserve existing ID if not provided in update and endpoint has ID + if "id" not in update_data and found_endpoint.id is not None: + endpoint_dict["id"] = found_endpoint.id + + # Remove is_from_config before saving - it's a response-only field (computed at read time) + endpoint_dict.pop("is_from_config", None) + + # Create updated endpoint object + updated_endpoint = PassThroughGenericEndpoint(**endpoint_dict) + + # Update the list + pass_through_endpoint_data[endpoint_index] = endpoint_dict + + # Remove old routes from registry before they get re-registered + InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id) + + ## Update db + updated_data = ConfigFieldUpdate( + field_name="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 + ) + + # Re-register the route with updated headers + _custom_headers: Optional[dict] = updated_endpoint.headers or {} + _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) + + if updated_endpoint.include_subpath: + InitPassThroughEndpointHelpers.add_subpath_route( + app=request.app, + path=updated_endpoint.path, + target=updated_endpoint.target, + custom_headers=_custom_headers, + forward_headers=None, # Defaults not available in model? assuming None logic handles it + merge_query_params=None, + dependencies=None, + cost_per_request=updated_endpoint.cost_per_request, + endpoint_id=updated_endpoint.id or endpoint_id or "", + guardrails=getattr(updated_endpoint, "guardrails", None), + methods=updated_endpoint.methods, + default_query_params=updated_endpoint.default_query_params, + ) + else: + InitPassThroughEndpointHelpers.add_exact_path_route( + app=request.app, + path=updated_endpoint.path, + target=updated_endpoint.target, + custom_headers=_custom_headers, + forward_headers=None, + merge_query_params=None, + dependencies=None, + cost_per_request=updated_endpoint.cost_per_request, + endpoint_id=updated_endpoint.id or endpoint_id or "", + guardrails=getattr(updated_endpoint, "guardrails", None), + methods=updated_endpoint.methods, + default_query_params=updated_endpoint.default_query_params, + ) + + return PassThroughEndpointResponse( + endpoints=[updated_endpoint] if updated_endpoint else [] + ) + + +@router.post( + "/config/pass_through_endpoint", + dependencies=[Depends(user_api_key_auth)], +) +async def create_pass_through_endpoints( + data: PassThroughGenericEndpoint, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Create new pass-through endpoint + """ + from litellm._uuid import uuid + from litellm.proxy.proxy_server import ( + get_config_general_settings, + update_config_general_settings, + ) + + ## Get existing pass-through endpoint field value + + try: + response: ConfigFieldInfo = await get_config_general_settings( + 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 + ) + + ## Auto-generate ID if not provided + # Exclude is_from_config as it's a response-only field (computed at read time) + data_dict = data.model_dump(exclude={"is_from_config"}) + if data_dict.get("id") is None: + data_dict["id"] = str(uuid.uuid4()) + + if response.field_value is None: + response.field_value = [data_dict] + elif isinstance(response.field_value, List): + response.field_value.append(data_dict) + + ## Update db + updated_data = ConfigFieldUpdate( + field_name="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 + ) + + # Return the created endpoint with the generated ID + created_endpoint = PassThroughGenericEndpoint(**data_dict) + + # Register the new route + _custom_headers: Optional[dict] = created_endpoint.headers or {} + _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) + + if created_endpoint.include_subpath: + InitPassThroughEndpointHelpers.add_subpath_route( + app=request.app, + path=created_endpoint.path, + target=created_endpoint.target, + custom_headers=_custom_headers, + forward_headers=None, + merge_query_params=None, + dependencies=None, + cost_per_request=created_endpoint.cost_per_request, + endpoint_id=created_endpoint.id or "", + guardrails=getattr(created_endpoint, "guardrails", None), + methods=created_endpoint.methods, + default_query_params=created_endpoint.default_query_params, + ) + else: + InitPassThroughEndpointHelpers.add_exact_path_route( + app=request.app, + path=created_endpoint.path, + target=created_endpoint.target, + custom_headers=_custom_headers, + forward_headers=None, + merge_query_params=None, + dependencies=None, + cost_per_request=created_endpoint.cost_per_request, + endpoint_id=created_endpoint.id or "", + guardrails=getattr(created_endpoint, "guardrails", None), + methods=created_endpoint.methods, + default_query_params=created_endpoint.default_query_params, + ) + + return PassThroughEndpointResponse(endpoints=[created_endpoint]) + + +@router.delete( + "/config/pass_through_endpoint", + dependencies=[Depends(user_api_key_auth)], + response_model=PassThroughEndpointResponse, +) +async def delete_pass_through_endpoints( + endpoint_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete a pass-through endpoint by ID. + + Returns - the deleted endpoint + """ + from litellm.proxy.proxy_server import ( + get_config_general_settings, + update_config_general_settings, + ) + + ## Get existing pass-through endpoint field value + + try: + response: ConfigFieldInfo = await get_config_general_settings( + 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 + ) + + ## Update field by removing endpoint + pass_through_endpoint_data: Optional[List] = response.field_value + if response.field_value is None or pass_through_endpoint_data is None: + raise HTTPException( + status_code=400, + detail={"error": "There are no pass-through endpoints setup."}, + ) + + # Find the endpoint to delete + found_endpoint = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) + + 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 + ) + }, + ) + + # 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 + ) + if _endpoint.id == endpoint_id: + endpoint_index = idx + break + + if endpoint_index is None: + raise HTTPException( + status_code=400, + detail={ + "error": f"Could not find index for endpoint with ID '{endpoint_id}'" + }, + ) + + # Remove the endpoint + pass_through_endpoint_data.pop(endpoint_index) + response_obj = found_endpoint + + # Remove routes from registry + InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id) + + ## Update db + updated_data = ConfigFieldUpdate( + field_name="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 + ) + + return PassThroughEndpointResponse(endpoints=[response_obj]) + + +def _find_endpoint_by_id( + endpoints_data: List, + endpoint_id: str, +) -> Optional[PassThroughGenericEndpoint]: + """ + Find an endpoint by ID. + + Args: + endpoints_data: List of endpoint data (dicts or PassThroughGenericEndpoint objects) + endpoint_id: ID to search for + + Returns: + Found endpoint or None if not found + """ + for endpoint in endpoints_data: + _endpoint: Optional[PassThroughGenericEndpoint] = None + if isinstance(endpoint, dict): + _endpoint = PassThroughGenericEndpoint(**endpoint) + elif isinstance(endpoint, PassThroughGenericEndpoint): + _endpoint = endpoint + + # Only compare IDs to IDs + if _endpoint is not None and _endpoint.id == endpoint_id: + return _endpoint + + return None + + +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 + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6657adf9657..0e702abfc6c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -388,6 +388,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( from litellm.proxy.management_endpoints.organization_endpoints import ( router as organization_router, ) +from litellm.proxy.management_endpoints.project_endpoints import ( + router as project_router, +) from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router from litellm.proxy.management_endpoints.router_settings_endpoints import ( router as router_settings_router, @@ -11361,6 +11364,7 @@ async def get_config_list( "maximum_spend_logs_retention_period": {"type": "String"}, "mcp_internal_ip_ranges": {"type": "List"}, "mcp_trusted_proxy_ranges": {"type": "List"}, + "always_include_stream_usage": {"type": "Boolean"}, } return_val = [] @@ -12478,6 +12482,7 @@ app.include_router(team_router) app.include_router(ui_sso_router) app.include_router(scim_router) app.include_router(organization_router) +app.include_router(project_router) app.include_router(customer_router) app.include_router(spend_management_router) app.include_router(cloudzero_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8bd46672ae6..d483e92e528 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -24,6 +24,7 @@ model LiteLLM_BudgetTable { updated_at DateTime @default(now()) @updatedAt @map("updated_at") updated_by String organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget + projects LiteLLM_ProjectTable[] // multiple projects can have the same budget keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget @@ -135,6 +136,34 @@ model LiteLLM_TeamTable { litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + projects LiteLLM_ProjectTable[] +} + +// Projects sit between teams and keys for use-case management +model LiteLLM_ProjectTable { + project_id String @id @default(uuid()) + project_alias String? + description String? + team_id String? + budget_id String? + metadata Json @default("{}") + models String[] + spend Float @default(0.0) + model_spend Json @default("{}") + model_rpm_limit Json @default("{}") + model_tpm_limit Json @default("{}") + blocked Boolean @default(false) + object_permission_id String? + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + + // Relations + litellm_team_table LiteLLM_TeamTable? @relation(fields: [team_id], references: [team_id]) + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + keys LiteLLM_VerificationToken[] + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) } // Audit table for deleted teams - preserves spend and team information for historical tracking @@ -230,6 +259,7 @@ model LiteLLM_ObjectPermissionTable { agents String[] @default([]) agent_access_groups String[] @default([]) teams LiteLLM_TeamTable[] + projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] @@ -284,6 +314,7 @@ model LiteLLM_VerificationToken { router_settings Json? @default("{}") user_id String? team_id String? + project_id String? permissions Json @default("{}") max_parallel_requests Int? metadata Json @default("{}") @@ -306,6 +337,7 @@ model LiteLLM_VerificationToken { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + last_active DateTime? // When this key was last used rotation_count Int? @default(0) // Number of times key has been rotated auto_rotate Boolean? @default(false) // Whether this key should be auto-rotated rotation_interval String? // How often to rotate (e.g., "30d", "90d") @@ -313,6 +345,7 @@ model LiteLLM_VerificationToken { key_rotation_at DateTime? // When this key should next be rotated litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) + litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" @@ -353,6 +386,7 @@ model LiteLLM_DeletedVerificationToken { config Json @default("{}") user_id String? team_id String? + project_id String? permissions Json @default("{}") max_parallel_requests Int? metadata Json @default("{}") @@ -376,6 +410,7 @@ model LiteLLM_DeletedVerificationToken { created_by String? // Original creator updated_at DateTime? // Last update timestamp before deletion updated_by String? // Last user who updated before deletion + last_active DateTime? // When this key was last used before deletion rotation_count Int? @default(0) auto_rotate Boolean? @default(false) rotation_interval String? @@ -448,7 +483,7 @@ model LiteLLM_SpendLogs { custom_llm_provider String? @default("") // litellm used custom_llm_provider api_base String? @default("") user String? @default("") - metadata Json? @default("{}") + metadata Json? @default("{}") // project_id stored here cache_hit String? @default("") cache_key String? @default("") request_tags Json? @default("[]") diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8224425a135..c36e50eb97a 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -67,6 +67,7 @@ def _get_spend_logs_metadata( user_api_key=None, user_api_key_alias=None, user_api_key_team_id=None, + user_api_key_project_id=None, user_api_key_org_id=None, user_api_key_user_id=None, user_api_key_team_alias=None, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 94471f0e32b..1a1764324a3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1238,7 +1238,8 @@ class ProxyLogging: if result.terminal_action == "modify_response": raise ModifyResponseException( - message=result.modify_response_message or "Response modified by pipeline", + message=result.modify_response_message + or "Response modified by pipeline", model=data.get("model", "unknown"), request_data=data, guardrail_name=f"pipeline:{policy_name}", @@ -1321,7 +1322,6 @@ class ProxyLogging: metadata = data.get("metadata", data.get("litellm_metadata", {})) or {} pipeline_managed: set = metadata.get("_pipeline_managed_guardrails", set()) - for callback in litellm.callbacks: start_time = time.time() _callback = None @@ -1337,7 +1337,10 @@ class ProxyLogging: and data is not None ): # Skip guardrails managed by a pipeline - if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed: + if ( + _callback.guardrail_name + and _callback.guardrail_name in pipeline_managed + ): continue result = await self._process_guardrail_callback( @@ -1491,6 +1494,7 @@ class ProxyLogging: "organization_budget", "proxy_budget", "projected_limit_exceeded", + "project_budget", ], user_info: CallInfo, ): @@ -1885,7 +1889,6 @@ class ProxyLogging: from litellm.types.guardrails import GuardrailEventHooks - guardrail_callbacks: List[CustomGuardrail] = [] other_callbacks: List[CustomLogger] = [] try: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 609e19b2443..19845d7c493 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -211,6 +211,7 @@ class LiteLLMCompletionResponsesConfig: "web_search_options": web_search_options, "response_format": response_format, "reasoning_effort": reasoning_effort, + "context_management": responses_api_request.get("context_management"), # litellm specific params "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, @@ -1349,7 +1350,7 @@ class LiteLLMCompletionResponsesConfig: result.append(tool) # type: ignore continue if tool.get("type") == "function": - fn = tool.get("function") or {} + fn = cast(Dict[str, Any], tool.get("function") or {}) parameters = dict(fn.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 1fc0e55dab3..9228b25b03e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2408,6 +2408,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_budget_reset_at: Optional[str] user_api_key_org_id: Optional[str] user_api_key_team_id: Optional[str] + user_api_key_project_id: Optional[str] user_api_key_user_id: Optional[str] user_api_key_user_email: Optional[str] user_api_key_team_alias: Optional[str] diff --git a/litellm/utils.py b/litellm/utils.py index 6a18fcc9e35..241b9d217b7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5294,6 +5294,9 @@ def _check_provider_match(model_info: dict, custom_llm_provider: Optional[str]) # as a last attempt if the model is not on Azure AI, Azure then fallback to OpenAI cost # tracking the cost is better than attributing 0 cost to it. return True + elif custom_llm_provider == "github": + # Allow github/ aliases to reuse existing provider metadata. + return True else: return False diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2c89e53ada9..8ed45ddd90c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14696,6 +14696,108 @@ "supports_web_search": true, "supports_native_streaming": true }, + "gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_image": 0.00012, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true + }, + "gemini-3.1-pro-preview-customtools": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_image": 0.00012, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true + }, "vertex_ai/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -14789,6 +14891,108 @@ "supports_web_search": true, "supports_native_streaming": true }, + "vertex_ai/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_image": 0.00012, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true + }, + "vertex_ai/gemini-3.1-pro-preview-customtools": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_image": 0.00012, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true + }, "gemini-2.5-pro-exp-03-25": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -16751,6 +16955,108 @@ "supports_native_streaming": true, "tpm": 800000 }, + "gemini/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-3.1-pro-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true, + "tpm": 800000 + }, + "gemini/gemini-3.1-pro-preview-customtools": { + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/models#gemini-3.1-pro-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_url_context": true, + "supports_native_streaming": true, + "tpm": 800000 + }, "gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, "input_cost_per_audio_token": 1e-06, diff --git a/policy_templates.json b/policy_templates.json index f27f20187f9..0dfa070aaa7 100644 --- a/policy_templates.json +++ b/policy_templates.json @@ -1522,7 +1522,13 @@ "guardrails": [ "aviation-ops-data-protection", "aviation-safety-topic-filter", - "airline-brand-protection-filter" + "airline-brand-protection-filter", + "competitor-name-input-blocker", + "competitor-name-output-blocker", + "competitor-recommendation-input-filter", + "competitor-recommendation-output-filter", + "competitor-comparison-input-filter", + "competitor-comparison-output-filter" ], "complexity": "High", "parameters": [ @@ -1531,9 +1537,14 @@ "label": "Your Airline / Brand Name", "type": "text", "required": true, - "placeholder": "e.g. Emirates" + "placeholder": "e.g. Acme Airlines" } ], + "llm_enrichment": { + "parameter": "brand_name", + "prompt": "List the top 30 direct competitors of {{brand_name}} in the airline industry. Include major international carriers, regional competitors, and low-cost carriers that operate on overlapping routes. Return ONLY airline/brand names, one per line, no numbering, no explanations.", + "result_key": "competitors" + }, "guardrailDefinitions": [ { "guardrail_name": "aviation-ops-data-protection", @@ -1675,6 +1686,72 @@ "guardrail_info": { "description": "Blocks AI-generated fake incident reports, unauthorized statements, and reputation-damaging content about your brand (runs on output)" } + }, + { + "guardrail_name": "competitor-name-input-blocker", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitors_blocked_words}}" + }, + "guardrail_info": { + "description": "Blocks user inputs that mention competitor names (pre_call)" + } + }, + { + "guardrail_name": "competitor-name-output-blocker", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "blocked_words": "{{competitors_blocked_words}}" + }, + "guardrail_info": { + "description": "Blocks AI outputs that mention competitor names (post_call)" + } + }, + { + "guardrail_name": "competitor-recommendation-input-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitor_recommendation_words}}" + }, + "guardrail_info": { + "description": "Blocks user requests asking to recommend competitors (pre_call)" + } + }, + { + "guardrail_name": "competitor-recommendation-output-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "blocked_words": "{{competitor_recommendation_words}}" + }, + "guardrail_info": { + "description": "Blocks AI from recommending or suggesting competitor services (post_call)" + } + }, + { + "guardrail_name": "competitor-comparison-input-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitor_comparison_words}}" + }, + "guardrail_info": { + "description": "Blocks user inputs requesting unfavorable brand comparisons (pre_call)" + } + }, + { + "guardrail_name": "competitor-comparison-output-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "blocked_words": "{{competitor_comparison_words}}" + }, + "guardrail_info": { + "description": "Blocks AI outputs with unfavorable brand comparisons (post_call)" + } } ], "templateData": { @@ -1683,7 +1760,13 @@ "guardrails_add": [ "aviation-ops-data-protection", "aviation-safety-topic-filter", - "airline-brand-protection-filter" + "airline-brand-protection-filter", + "competitor-name-input-blocker", + "competitor-name-output-blocker", + "competitor-recommendation-input-filter", + "competitor-recommendation-output-filter", + "competitor-comparison-input-filter", + "competitor-comparison-output-filter" ], "guardrails_remove": [] }, @@ -1812,9 +1895,12 @@ "iconColor": "text-orange-500", "iconBg": "bg-orange-50", "guardrails": [ + "competitor-input-blocker", "competitor-output-blocker", - "competitor-recommendation-filter", - "competitor-comparison-filter" + "competitor-recommendation-input-filter", + "competitor-recommendation-output-filter", + "competitor-comparison-input-filter", + "competitor-comparison-output-filter" ], "complexity": "Medium", "parameters": [ @@ -1823,15 +1909,26 @@ "label": "Your Brand Name", "type": "text", "required": true, - "placeholder": "e.g. Emirates" + "placeholder": "e.g. Acme Airlines" } ], "llm_enrichment": { "parameter": "brand_name", - "prompt": "List the top 10 direct competitors of {{brand_name}} in the same industry. Return ONLY company/brand names, one per line, no numbering, no explanations.", + "prompt": "List the top 30 direct competitors of {{brand_name}} in the same industry. Return ONLY company/brand names, one per line, no numbering, no explanations.", "result_key": "competitors" }, "guardrailDefinitions": [ + { + "guardrail_name": "competitor-input-blocker", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitors_blocked_words}}" + }, + "guardrail_info": { + "description": "Blocks user inputs that mention competitor brands (pre_call)" + } + }, { "guardrail_name": "competitor-output-blocker", "litellm_params": { @@ -1840,39 +1937,64 @@ "blocked_words": "{{competitors_blocked_words}}" }, "guardrail_info": { - "description": "Blocks AI outputs that mention or promote competitor brands (auto-discovered via LLM)" + "description": "Blocks AI outputs that mention competitor brands (post_call)" } }, { - "guardrail_name": "competitor-recommendation-filter", + "guardrail_name": "competitor-recommendation-input-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitor_recommendation_words}}" + }, + "guardrail_info": { + "description": "Blocks user requests asking to recommend competitors (pre_call)" + } + }, + { + "guardrail_name": "competitor-recommendation-output-filter", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "post_call", "blocked_words": "{{competitor_recommendation_words}}" }, "guardrail_info": { - "description": "Blocks AI from recommending, suggesting, or directing users to competitor services" + "description": "Blocks AI from recommending or suggesting competitor services (post_call)" } }, { - "guardrail_name": "competitor-comparison-filter", + "guardrail_name": "competitor-comparison-input-filter", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "blocked_words": "{{competitor_comparison_words}}" + }, + "guardrail_info": { + "description": "Blocks user inputs requesting unfavorable brand comparisons (pre_call)" + } + }, + { + "guardrail_name": "competitor-comparison-output-filter", "litellm_params": { "guardrail": "litellm_content_filter", "mode": "post_call", "blocked_words": "{{competitor_comparison_words}}" }, "guardrail_info": { - "description": "Blocks unfavorable comparisons between your brand and competitors in AI outputs" + "description": "Blocks AI outputs with unfavorable brand comparisons (post_call)" } } ], "templateData": { "policy_name": "competitor-mention-detection", - "description": "Detects and blocks competitor mentions in AI outputs. Uses LLM-powered competitor discovery based on your brand name.", + "description": "Detects and blocks competitor mentions in both inputs and outputs. Uses LLM-powered competitor discovery based on your brand name.", "guardrails_add": [ + "competitor-input-blocker", "competitor-output-blocker", - "competitor-recommendation-filter", - "competitor-comparison-filter" + "competitor-recommendation-input-filter", + "competitor-recommendation-output-filter", + "competitor-comparison-input-filter", + "competitor-comparison-output-filter" ], "guardrails_remove": [] }, diff --git a/pyproject.toml b/pyproject.toml index 31131ac4aac..e301d57912a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ boto3 = { version = "1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "0.4.40", optional = true} +litellm-proxy-extras = {version = "0.4.44", optional = true} rich = {version = "13.7.1", optional = true} litellm-enterprise = {version = "0.1.32", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index f7108ddd201..87b2d733051 100644 --- a/requirements.txt +++ b/requirements.txt @@ -55,7 +55,7 @@ grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.40 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.44 # for proxy extras - e.g. prisma migrations llm-sandbox==0.3.31 # for skill execution in sandbox ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env diff --git a/schema.prisma b/schema.prisma index 8bd46672ae6..d483e92e528 100644 --- a/schema.prisma +++ b/schema.prisma @@ -24,6 +24,7 @@ model LiteLLM_BudgetTable { updated_at DateTime @default(now()) @updatedAt @map("updated_at") updated_by String organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget + projects LiteLLM_ProjectTable[] // multiple projects can have the same budget keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget @@ -135,6 +136,34 @@ model LiteLLM_TeamTable { litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_model_table LiteLLM_ModelTable? @relation(fields: [model_id], references: [id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + projects LiteLLM_ProjectTable[] +} + +// Projects sit between teams and keys for use-case management +model LiteLLM_ProjectTable { + project_id String @id @default(uuid()) + project_alias String? + description String? + team_id String? + budget_id String? + metadata Json @default("{}") + models String[] + spend Float @default(0.0) + model_spend Json @default("{}") + model_rpm_limit Json @default("{}") + model_tpm_limit Json @default("{}") + blocked Boolean @default(false) + object_permission_id String? + created_at DateTime @default(now()) @map("created_at") + created_by String + updated_at DateTime @default(now()) @updatedAt @map("updated_at") + updated_by String + + // Relations + litellm_team_table LiteLLM_TeamTable? @relation(fields: [team_id], references: [team_id]) + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + keys LiteLLM_VerificationToken[] + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) } // Audit table for deleted teams - preserves spend and team information for historical tracking @@ -230,6 +259,7 @@ model LiteLLM_ObjectPermissionTable { agents String[] @default([]) agent_access_groups String[] @default([]) teams LiteLLM_TeamTable[] + projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] @@ -284,6 +314,7 @@ model LiteLLM_VerificationToken { router_settings Json? @default("{}") user_id String? team_id String? + project_id String? permissions Json @default("{}") max_parallel_requests Int? metadata Json @default("{}") @@ -306,6 +337,7 @@ model LiteLLM_VerificationToken { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + last_active DateTime? // When this key was last used rotation_count Int? @default(0) // Number of times key has been rotated auto_rotate Boolean? @default(false) // Whether this key should be auto-rotated rotation_interval String? // How often to rotate (e.g., "30d", "90d") @@ -313,6 +345,7 @@ model LiteLLM_VerificationToken { key_rotation_at DateTime? // When this key should next be rotated litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) + litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" @@ -353,6 +386,7 @@ model LiteLLM_DeletedVerificationToken { config Json @default("{}") user_id String? team_id String? + project_id String? permissions Json @default("{}") max_parallel_requests Int? metadata Json @default("{}") @@ -376,6 +410,7 @@ model LiteLLM_DeletedVerificationToken { created_by String? // Original creator updated_at DateTime? // Last update timestamp before deletion updated_by String? // Last user who updated before deletion + last_active DateTime? // When this key was last used before deletion rotation_count Int? @default(0) auto_rotate Boolean? @default(false) rotation_interval String? @@ -448,7 +483,7 @@ model LiteLLM_SpendLogs { custom_llm_provider String? @default("") // litellm used custom_llm_provider api_base String? @default("") user String? @default("") - metadata Json? @default("{}") + metadata Json? @default("{}") // project_id stored here cache_hit String? @default("") cache_key String? @default("") request_tags Json? @default("[]") diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index 6ecac7b36a2..a521e8e43f8 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -2045,3 +2045,42 @@ def test_store_in_openai_chat_completion_params(): result = get_standard_openai_params({"store": True, "temperature": 0.7}) assert "store" in result assert result["store"] is True + + +def test_store_param_passed_through_openai_azure(): + """ + Test that the `store` parameter is correctly passed through to OpenAI + and Azure OpenAI providers when using get_optional_params(). + + This verifies the fix for the regression where `store` was being filtered + out by get_non_default_completion_params() due to architectural issues + in parameter processing pipeline. + + Ref: https://github.com/BerriAI/litellm/issues/19700 + """ + # Test OpenAI provider + optional_params_openai = get_optional_params( + model="gpt-4o", + custom_llm_provider="openai", + store=True, + ) + assert "store" in optional_params_openai + assert optional_params_openai["store"] is True + + # Test Azure OpenAI provider + optional_params_azure = get_optional_params( + model="gpt-4.1-2025-04-14", + custom_llm_provider="azure", + store=True, + ) + assert "store" in optional_params_azure + assert optional_params_azure["store"] is True + + # Test with store=False + optional_params_false = get_optional_params( + model="gpt-4o", + custom_llm_provider="openai", + store=False, + ) + assert "store" in optional_params_false + assert optional_params_false["store"] is False diff --git a/tests/llm_translation/test_xai.py b/tests/llm_translation/test_xai.py index c71b9e3fe51..1abbaa214ac 100644 --- a/tests/llm_translation/test_xai.py +++ b/tests/llm_translation/test_xai.py @@ -201,3 +201,55 @@ class TestXAIChat(BaseLLMChatTest): ) assert response is not None + + +def test_xai_streaming_with_include_usage(): + """ + Test that xAI streaming correctly handles usage in the last chunk + when stream_options={"include_usage": True} is set. + + xAI sends usage in a chunk with empty choices array, which should be + handled by XAIChatCompletionStreamingHandler. + """ + try: + response = completion( + model="xai/grok-4-1-fast-non-reasoning", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Say hello in one word"} + ], + stream=True, + stream_options={"include_usage": True}, + max_tokens=10, + ) + + chunks = [] + usage_chunk = None + + for chunk in response: + chunks.append(chunk) + if hasattr(chunk, "usage") and chunk.usage is not None: + usage_chunk = chunk + + # Verify we got chunks + assert len(chunks) > 0, "Should receive streaming chunks" + + # Verify usage was included in one of the chunks + assert usage_chunk is not None, "Should receive usage in streaming chunks" + + # Verify usage has expected fields + assert hasattr(usage_chunk.usage, "prompt_tokens"), "Usage should have prompt_tokens" + assert hasattr(usage_chunk.usage, "completion_tokens"), "Usage should have completion_tokens" + assert hasattr(usage_chunk.usage, "total_tokens"), "Usage should have total_tokens" + + # Verify usage values are positive + assert usage_chunk.usage.prompt_tokens > 0, "prompt_tokens should be positive" + assert usage_chunk.usage.completion_tokens > 0, "completion_tokens should be positive" + assert usage_chunk.usage.total_tokens > 0, "total_tokens should be positive" + + print(f"✓ Successfully received usage in streaming chunk: {usage_chunk.usage}") + + except Exception as e: + if "API key" in str(e) or "authentication" in str(e).lower(): + pytest.skip(f"Skipping test due to API key issue: {str(e)}") + raise diff --git a/tests/proxy_unit_tests/test_project_endpoints_prisma.py b/tests/proxy_unit_tests/test_project_endpoints_prisma.py new file mode 100644 index 00000000000..c2366139b01 --- /dev/null +++ b/tests/proxy_unit_tests/test_project_endpoints_prisma.py @@ -0,0 +1,789 @@ +import os +import sys +import traceback +from litellm._uuid import uuid +from unittest import mock + +from dotenv import load_dotenv +from fastapi import Request + +load_dotenv() +import time + +sys.path.insert(0, os.path.abspath("../..")) +import logging + +import pytest + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy.management_endpoints.team_endpoints import ( + new_team, +) +from litellm.proxy.management_endpoints.project_endpoints import ( + new_project, + update_project, + delete_project, + project_info, +) +from litellm.proxy.proxy_server import ( + LitellmUserRoles, +) +from litellm.proxy.utils import PrismaClient, ProxyLogging + +verbose_proxy_logger.setLevel(level=logging.DEBUG) + + +from litellm.caching.caching import DualCache +from litellm.proxy._types import ( + NewProjectRequest, + UpdateProjectRequest, + DeleteProjectRequest, + NewTeamRequest, + UserAPIKeyAuth, +) + +proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + + +@pytest.fixture +def prisma_client(): + from litellm.proxy.proxy_cli import append_query_params + + ### add connection pool + pool timeout args + params = {"connection_limit": 100, "pool_timeout": 60} + database_url = os.getenv("DATABASE_URL") + modified_url = append_query_params(database_url, params) + os.environ["DATABASE_URL"] = modified_url + + # Assuming PrismaClient is a class that needs to be instantiated + prisma_client = PrismaClient( + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj + ) + + # Reset litellm.proxy.proxy_server.prisma_client to None + litellm.proxy.proxy_server.litellm_proxy_budget_name = ( + f"litellm-proxy-budget-{time.time()}" + ) + litellm.proxy.proxy_server.user_custom_key_generate = None + + # Enable premium_user for project management tests + setattr(litellm.proxy.proxy_server, "premium_user", True) + + return prisma_client + + +@pytest.mark.asyncio +async def test_new_project(prisma_client): + """ + Test creating a new project with budget, models, and metadata. + """ + try: + print("prisma client=", prisma_client) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + + await litellm.proxy.proxy_server.prisma_client.connect() + + # Create a team first + _team_id = f"project-test-team_{uuid.uuid4()}" + await new_team( + NewTeamRequest( + team_id=_team_id, + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + # Create a project + project_data = NewProjectRequest( + project_alias="test-project", + description="Test project for unit testing", + team_id=_team_id, + metadata={"use_case_id": "TEST-001", "responsible_ai_id": "RAI-001"}, + models=["gpt-4", "gpt-3.5-turbo"], + max_budget=100.0, + model_rpm_limit={"gpt-4": 100}, + model_tpm_limit={"gpt-4": 1000}, + ) + + response = await new_project( + data=project_data, + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + print("New project response:", response) + + # Assertions + assert response.project_id is not None + assert response.project_alias == "test-project" + assert response.description == "Test project for unit testing" + assert response.team_id == _team_id + assert response.models == ["gpt-4", "gpt-3.5-turbo"] + # model_rpm_limit and model_tpm_limit are stored in metadata + assert response.metadata["use_case_id"] == "TEST-001" + assert response.metadata["responsible_ai_id"] == "RAI-001" + assert response.metadata["model_rpm_limit"] == {"gpt-4": 100} + assert response.metadata["model_tpm_limit"] == {"gpt-4": 1000} + assert response.litellm_budget_table is not None + assert response.litellm_budget_table.max_budget == 100.0 + + except Exception as e: + print("Got Exception", e) + traceback.print_exc() + pytest.fail(f"Got exception {e}") + + +@pytest.mark.asyncio +async def test_update_project(prisma_client): + """ + Test updating an existing project's budget, models, and metadata. + """ + try: + print("prisma client=", prisma_client) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + + await litellm.proxy.proxy_server.prisma_client.connect() + + # Create a team first + _team_id = f"project-test-team_{uuid.uuid4()}" + await new_team( + NewTeamRequest( + team_id=_team_id, + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + # Create a project + project_data = NewProjectRequest( + project_alias="test-project-update", + description="Original description", + team_id=_team_id, + metadata={ + "use_case_id": "TEST-002", + }, + models=["gpt-4"], + max_budget=50.0, + ) + + create_response = await new_project( + data=project_data, + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + print("Created project:", create_response) + project_id = create_response.project_id + + # Update the project + update_data = UpdateProjectRequest( + project_id=project_id, + project_alias="test-project-updated", + description="Updated description", + metadata={ + "use_case_id": "TEST-002-UPDATED", + "additional_field": "new_value", + }, + models=["gpt-4", "gpt-3.5-turbo", "claude-3"], + max_budget=200.0, + model_rpm_limit={"gpt-4": 200, "claude-3": 50}, + model_tpm_limit={"gpt-4": 2000, "claude-3": 500}, + ) + + update_response = await update_project( + data=update_data, + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + print("Updated project response:", update_response) + + # Assertions + assert update_response.project_id == project_id + assert update_response.project_alias == "test-project-updated" + assert update_response.description == "Updated description" + assert update_response.models == ["gpt-4", "gpt-3.5-turbo", "claude-3"] + # model_rpm_limit and model_tpm_limit are stored in metadata + assert update_response.metadata["use_case_id"] == "TEST-002-UPDATED" + assert update_response.metadata["additional_field"] == "new_value" + assert update_response.metadata["model_rpm_limit"] == { + "gpt-4": 200, + "claude-3": 50, + } + assert update_response.metadata["model_tpm_limit"] == { + "gpt-4": 2000, + "claude-3": 500, + } + assert update_response.litellm_budget_table is not None + assert update_response.litellm_budget_table.max_budget == 200.0 + + except Exception as e: + print("Got Exception", e) + traceback.print_exc() + pytest.fail(f"Got exception {e}") + + +@pytest.mark.asyncio +async def test_delete_project(prisma_client): + """ + Test deleting a project. + """ + try: + print("prisma client=", prisma_client) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + + await litellm.proxy.proxy_server.prisma_client.connect() + + # Create a team first + _team_id = f"project-test-team_{uuid.uuid4()}" + await new_team( + NewTeamRequest( + team_id=_team_id, + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + # Create a project + project_data = NewProjectRequest( + project_alias="test-project-delete", + team_id=_team_id, + models=["gpt-4"], + max_budget=50.0, + ) + + create_response = await new_project( + data=project_data, + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + print("Created project:", create_response) + project_id = create_response.project_id + + # Delete the project + delete_data = DeleteProjectRequest(project_ids=[project_id]) + + delete_response = await delete_project( + data=delete_data, + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + print("Delete project response:", delete_response) + + # Assertions - delete_project returns a list of deleted project objects + assert isinstance(delete_response, list) + assert len(delete_response) == 1 + assert delete_response[0].project_id == project_id + + # Try to get info on the deleted project - should fail or return None + try: + await project_info( + project_id=project_id, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + pytest.fail("Expected to fail when fetching deleted project") + except Exception as e: + print("Expected error when fetching deleted project:", e) + # This is expected behavior + + except Exception as e: + print("Got Exception", e) + traceback.print_exc() + pytest.fail(f"Got exception {e}") + + +@pytest.mark.asyncio +async def test_project_info(prisma_client): + """ + Test getting project info. + """ + try: + print("prisma client=", prisma_client) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + + await litellm.proxy.proxy_server.prisma_client.connect() + + # Create a team first + _team_id = f"project-test-team_{uuid.uuid4()}" + await new_team( + NewTeamRequest( + team_id=_team_id, + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + # Create a project + project_data = NewProjectRequest( + project_alias="test-project-info", + description="Test project info endpoint", + team_id=_team_id, + metadata={"use_case_id": "TEST-003", "cost_center": "engineering"}, + models=["gpt-4", "claude-3"], + max_budget=150.0, + model_rpm_limit={"gpt-4": 150}, + model_tpm_limit={"gpt-4": 1500}, + ) + + create_response = await new_project( + data=project_data, + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + print("Created project:", create_response) + project_id = create_response.project_id + + # Get project info + info_response = await project_info( + project_id=project_id, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + print("Project info response:", info_response) + + # Assertions - project_info returns the project object directly + assert info_response.project_id == project_id + assert info_response.project_alias == "test-project-info" + assert info_response.description == "Test project info endpoint" + assert info_response.team_id == _team_id + assert info_response.models == ["gpt-4", "claude-3"] + # model_rpm_limit and model_tpm_limit are stored in metadata + assert info_response.metadata["use_case_id"] == "TEST-003" + assert info_response.metadata["cost_center"] == "engineering" + assert info_response.metadata["model_rpm_limit"] == {"gpt-4": 150} + assert info_response.metadata["model_tpm_limit"] == {"gpt-4": 1500} + assert info_response.litellm_budget_table is not None + assert info_response.litellm_budget_table.max_budget == 150.0 + + except Exception as e: + print("Got Exception", e) + traceback.print_exc() + pytest.fail(f"Got exception {e}") + + +### VALIDATION TESTS ### + + +def test_check_team_project_limits_models_not_in_team(): + """ + Test that creating a project with models not in the team raises an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4", "gpt-3.5-turbo"], + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4", "claude-3"], # claude-3 not in team + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "claude-3" in str(exc_info.value.detail) + assert "not in team's allowed models" in str(exc_info.value.detail) + + +def test_check_team_project_limits_budget_exceeds_team(): + """ + Test that creating a project with budget > team budget raises an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4"], + max_budget=100.0, + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4"], + max_budget=150.0, # exceeds team's 100.0 + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "exceeds team's max_budget" in str(exc_info.value.detail) + + +def test_check_team_project_limits_valid_subset(): + """ + Test that a valid project (models subset, budget within limit) passes. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4", "gpt-3.5-turbo", "claude-3"], + max_budget=1000.0, + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4", "gpt-3.5-turbo"], + max_budget=500.0, + ) + + # Should not raise + _check_team_project_limits(team_object=team, data=data) + + +def test_check_team_project_limits_all_proxy_models(): + """ + Test that team with 'all-proxy-models' allows any project models. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["all-proxy-models"], + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4", "claude-3", "anything-goes"], + ) + + # Should not raise - team allows all models + _check_team_project_limits(team_object=team, data=data) + + +def test_check_team_project_limits_tpm_exceeds_team(): + """ + Test that project tpm_limit exceeding team tpm_limit raises an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4"], + tpm_limit=10000, + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4"], + tpm_limit=20000, # exceeds team's 10000 + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "exceeds team's tpm_limit" in str(exc_info.value.detail) + + +def test_check_team_project_limits_negative_budget(): + """ + Test that negative budget values raise an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4"], + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4"], + max_budget=-10.0, + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "cannot be negative" in str(exc_info.value.detail) + + +def test_check_team_project_limits_soft_budget_gte_max(): + """ + Test that soft_budget >= max_budget raises an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4"], + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4"], + max_budget=100.0, + soft_budget=100.0, # equal to max, should fail + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "must be strictly lower" in str(exc_info.value.detail) + + +def test_premium_user_gate(): + """ + Test that project endpoints require premium_user=True. + """ + + # This test just validates the premium_user check exists + # The actual endpoint test would need prisma, but we can verify + # the import path works + setattr(litellm.proxy.proxy_server, "premium_user", False) + + # Verify that CommonProxyErrors.not_premium_user exists + from litellm.proxy._types import CommonProxyErrors + + assert hasattr(CommonProxyErrors, "not_premium_user") + + # Reset + setattr(litellm.proxy.proxy_server, "premium_user", True) + + +def test_project_model_access_denied_error_type(): + """ + Test that ProxyErrorTypes.project_model_access_denied exists. + """ + from litellm.proxy._types import ProxyErrorTypes + + assert hasattr(ProxyErrorTypes, "project_model_access_denied") + assert ( + ProxyErrorTypes.project_model_access_denied.value + == "project_model_access_denied" + ) + + # Test the classmethod resolves correctly + result = ProxyErrorTypes.get_model_access_error_type_for_object("project") + assert result == ProxyErrorTypes.project_model_access_denied + + +def test_project_cached_obj_has_last_refreshed_at(): + """ + Test that LiteLLM_ProjectTableCachedObj has last_refreshed_at field + matching LiteLLM_TeamTableCachedObj pattern. + """ + from litellm.proxy._types import ( + LiteLLM_ProjectTableCachedObj, + LiteLLM_ProjectTable, + ) + + # Verify inheritance + assert issubclass(LiteLLM_ProjectTableCachedObj, LiteLLM_ProjectTable) + + # Verify last_refreshed_at field exists and defaults to None + obj = LiteLLM_ProjectTableCachedObj( + project_id="test", + created_by="admin", + updated_by="admin", + ) + assert obj.last_refreshed_at is None + + # Verify it can be set + obj.last_refreshed_at = 1234567890.0 + assert obj.last_refreshed_at == 1234567890.0 + + +@pytest.mark.asyncio +async def test_project_max_budget_check_fires_alert(): + """ + Test that _project_max_budget_check fires a budget alert + when project exceeds its max budget (matches _team_max_budget_check pattern). + """ + from litellm.proxy.auth.auth_checks import _project_max_budget_check + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ProjectTableCachedObj, + ) + + project = LiteLLM_ProjectTableCachedObj( + project_id="test-project", + spend=150.0, + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="user-1", + team_id="team-1", + ) + + mock_proxy_logging = mock.AsyncMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = mock.AsyncMock() + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _project_max_budget_check( + project_object=project, + valid_token=valid_token, + proxy_logging_obj=mock_proxy_logging, + ) + + assert "Project=test-project" in str(exc_info.value) + assert "150.0" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_project_soft_budget_check(): + """ + Test that _project_soft_budget_check triggers alert when soft budget is exceeded. + """ + from litellm.proxy.auth.auth_checks import _project_soft_budget_check + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ProjectTableCachedObj, + ) + + project = LiteLLM_ProjectTableCachedObj( + project_id="test-project", + spend=80.0, + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(soft_budget=75.0), + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="user-1", + team_id="team-1", + ) + + mock_proxy_logging = mock.AsyncMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = mock.AsyncMock() + + # Should not raise (soft budget only alerts, doesn't block) + await _project_soft_budget_check( + project_object=project, + valid_token=valid_token, + proxy_logging_obj=mock_proxy_logging, + ) + + +@pytest.mark.asyncio +async def test_project_soft_budget_check_no_alert_under_budget(): + """ + Test that _project_soft_budget_check does NOT trigger alert when under soft budget. + """ + from litellm.proxy.auth.auth_checks import _project_soft_budget_check + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ProjectTableCachedObj, + ) + + project = LiteLLM_ProjectTableCachedObj( + project_id="test-project", + spend=50.0, + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(soft_budget=75.0), + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="user-1", + team_id="team-1", + ) + + mock_proxy_logging = mock.AsyncMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = mock.AsyncMock() + + # Should not raise and should not alert + await _project_soft_budget_check( + project_object=project, + valid_token=valid_token, + proxy_logging_obj=mock_proxy_logging, + ) + + +def test_litellm_entity_type_has_project(): + """ + Test that Litellm_EntityType has PROJECT member for budget alerts. + """ + from litellm.proxy._types import Litellm_EntityType + + assert hasattr(Litellm_EntityType, "PROJECT") + assert Litellm_EntityType.PROJECT.value == "project" diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 77ff5dbb97e..10d3323a255 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -478,6 +478,7 @@ def test_max_langfuse_clients_limit(): mock_langfuse = MagicMock() mock_langfuse.version.__version__ = "3.0.0" # Set max clients to 2 for testing + original_initialized_langfuse_clients = litellm.initialized_langfuse_clients with patch.dict("sys.modules", {"langfuse": mock_langfuse}), patch.object( langfuse_module, "MAX_LANGFUSE_INITIALIZED_CLIENTS", 2 ): @@ -513,3 +514,5 @@ def test_max_langfuse_clients_limit(): # Counter should still be 2 (third client failed to initialize) assert litellm.initialized_langfuse_clients == 2 + + litellm.initialized_langfuse_clients = original_initialized_langfuse_clients diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index 5abecb46c99..85ed5dddee1 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -100,3 +100,176 @@ async def test_internal_flags_filtered_from_followup_kwargs(): # Verify regular kwargs are preserved assert kwargs_for_followup["temperature"] == 0.7 assert kwargs_for_followup["max_tokens"] == 1024 + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_provider_from_top_level_kwargs(): + """Test that async_pre_call_deployment_hook finds custom_llm_provider at top-level kwargs. + + Regression test for bug where the hook only checked kwargs["litellm_params"]["custom_llm_provider"] + but the router places custom_llm_provider at the top level of kwargs. + """ + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + + # Simulate kwargs as they arrive from the router path: + # custom_llm_provider is at the TOP LEVEL (not nested under litellm_params) + kwargs = { + "model": "anthropic.claude-3-5-sonnet-20241022-v2:0", + "messages": [{"role": "user", "content": "Search the web for LiteLLM"}], + "tools": [ + {"type": "web_search_20250305", "name": "web_search", "max_uses": 3}, + {"type": "function", "function": {"name": "other_tool", "parameters": {}}}, + ], + "custom_llm_provider": "bedrock", + "api_key": "fake-key", + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + + # Should NOT be None — the hook should have triggered + assert result is not None + # The web_search tool should be converted to litellm_web_search (OpenAI format) + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" + for t in result["tools"] + ) + # The non-web-search tool should be preserved + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "other_tool" + for t in result["tools"] + ) + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_returns_full_kwargs(): + """Test that async_pre_call_deployment_hook returns the full kwargs dict, not a partial one. + + Regression test for bug where the hook returned {"tools": converted_tools} instead of + the full kwargs dict, causing model/messages/api_key/etc. to be lost. + """ + logger = WebSearchInterceptionLogger(enabled_providers=["openai"]) + + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Search for something"}], + "tools": [ + {"type": "web_search_20250305", "name": "web_search"}, + ], + "custom_llm_provider": "openai", + "api_key": "sk-fake", + "temperature": 0.7, + "metadata": {"user": "test"}, + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + + assert result is not None + # All original keys must be preserved + assert result["model"] == "gpt-4o" + assert result["messages"] == [{"role": "user", "content": "Search for something"}] + assert result["api_key"] == "sk-fake" + assert result["temperature"] == 0.7 + assert result["metadata"] == {"user": "test"} + assert result["custom_llm_provider"] == "openai" + # Tools should be converted + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" + for t in result["tools"] + ) + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_skips_disabled_provider(): + """Test that the hook returns None for providers not in enabled_providers.""" + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "test"}], + "tools": [{"type": "web_search_20250305", "name": "web_search"}], + "custom_llm_provider": "openai", # Not in enabled_providers + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + assert result is None + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_skips_no_websearch_tools(): + """Test that the hook returns None when no web search tools are present.""" + logger = WebSearchInterceptionLogger(enabled_providers=["openai"]) + + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "test"}], + "tools": [ + {"type": "function", "function": {"name": "calculator", "parameters": {}}}, + ], + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + assert result is None + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_nested_litellm_params_fallback(): + """Test that the hook still works when custom_llm_provider is in nested litellm_params. + + This is the Anthropic experimental pass-through path where litellm_params is + explicitly constructed with custom_llm_provider inside it. + """ + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + + kwargs = { + "model": "anthropic.claude-3-5-sonnet-20241022-v2:0", + "messages": [{"role": "user", "content": "test"}], + "tools": [{"type": "web_search_20250305", "name": "web_search"}], + "litellm_params": { + "custom_llm_provider": "bedrock", + }, + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + + assert result is not None + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" + for t in result["tools"] + ) + # Full kwargs preserved + assert result["model"] == "anthropic.claude-3-5-sonnet-20241022-v2:0" + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_provider_derived_from_model_name(): + """Test that async_pre_call_deployment_hook derives custom_llm_provider from the model name. + + Regression test for the router _acompletion path where custom_llm_provider is NOT + in kwargs at all — neither at top-level nor in litellm_params. The hook must derive + the provider from the model name (e.g., "openai/gpt-4o-mini" → "openai"). + """ + logger = WebSearchInterceptionLogger(enabled_providers=["openai"]) + + # Simulate kwargs as they arrive from router._acompletion: + # NO custom_llm_provider key anywhere — only model name contains the provider + kwargs = { + "model": "openai/gpt-4o-mini", + "messages": [{"role": "user", "content": "Search the web for LiteLLM"}], + "tools": [ + {"type": "web_search_20250305", "name": "web_search", "max_uses": 3}, + ], + "api_key": "fake-key", + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + + # Should NOT be None — the hook should derive "openai" from "openai/gpt-4o-mini" + assert result is not None + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" + for t in result["tools"] + ) + # Full kwargs preserved + assert result["model"] == "openai/gpt-4o-mini" + assert result["api_key"] == "fake-key" diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 50e948c1a27..4b15d30ccec 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -2582,6 +2582,94 @@ def test_compaction_block_with_other_content_types(): assert tool_calls[0]["function"]["name"] == "get_weather" +def test_map_openai_context_management_to_anthropic(): + """ + Test mapping OpenAI Responses API context_management format to Anthropic format. + """ + config = AnthropicConfig() + + # Test OpenAI list format with compaction + openai_format = [{"type": "compaction", "compact_threshold": 200000}] + result = config.map_openai_context_management_to_anthropic(openai_format) + + assert result is not None + assert "edits" in result + assert len(result["edits"]) == 1 + assert result["edits"][0]["type"] == "compact_20260112" + assert result["edits"][0]["trigger"]["type"] == "input_tokens" + assert result["edits"][0]["trigger"]["value"] == 200000 + + # Test OpenAI format with instructions + openai_format_with_instructions = [{ + "type": "compaction", + "compact_threshold": 150000, + "instructions": "Focus on preserving code snippets" + }] + result = config.map_openai_context_management_to_anthropic(openai_format_with_instructions) + + assert result is not None + assert result["edits"][0]["trigger"]["value"] == 150000 + assert result["edits"][0]["instructions"] == "Focus on preserving code snippets" + + # Test Anthropic format (should pass through) + anthropic_format = { + "edits": [{ + "type": "compact_20260112", + "trigger": {"type": "input_tokens", "value": 150000} + }] + } + result = config.map_openai_context_management_to_anthropic(anthropic_format) + + assert result == anthropic_format + + +def test_map_openai_params_with_context_management(): + """ + Test that map_openai_params correctly transforms context_management from OpenAI to Anthropic format. + """ + config = AnthropicConfig() + + # Test with OpenAI list format + non_default_params = { + "context_management": [{"type": "compaction", "compact_threshold": 200000}] + } + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="claude-opus-4-6", + drop_params=False + ) + + assert "context_management" in result + assert "edits" in result["context_management"] + assert result["context_management"]["edits"][0]["type"] == "compact_20260112" + assert result["context_management"]["edits"][0]["trigger"]["value"] == 200000 + + # Test with Anthropic dict format (should pass through) + non_default_params_anthropic = { + "context_management": { + "edits": [{ + "type": "compact_20260112", + "trigger": {"type": "input_tokens", "value": 150000}, + "instructions": "Focus on preserving code" + }] + } + } + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params_anthropic, + optional_params=optional_params, + model="claude-opus-4-6", + drop_params=False + ) + + assert "context_management" in result + assert result["context_management"] == non_default_params_anthropic["context_management"] + + def test_compaction_block_empty_list_not_added(): """ Test that empty compaction_blocks list is not added to provider_specific_fields. diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py b/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py index ff2302749f1..b5f656c71f8 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py @@ -12,6 +12,7 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +from litellm.types.llms.openai import AllMessageValues import pytest import litellm @@ -111,3 +112,35 @@ class TestDashScopeConfig: # Check for specific content in the response assert "```python" in response.choices[0].message.content assert "Hey from LiteLLM" in response.choices[0].message.content + + def test_dashscope_no_longer_transforms_content_list(self): + """ + Test that DashScopeChatConfig does not transform content lists to strings. + This ensures that the transformation logic specific to content lists is not applied, + as DashScope should handle content in list format natively. + """ + config = DashScopeChatConfig() + + # Create a message with content in list format + messages: list[AllMessageValues] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello"}, + {"type": "text", "text": "World"}, + ], + } + ] + + # Call the _transform_messages method directly + transformed_messages = config._transform_messages( + messages=messages, model="qwen-turbo", is_async=False + ) + + # Verify that the content is still in list format and has not been transformed to a string + assert isinstance(transformed_messages[0]["content"], list) + assert len(transformed_messages[0]["content"]) == 2 + assert transformed_messages[0]["content"][0]["type"] == "text" + assert transformed_messages[0]["content"][0]["text"] == "Hello" + assert transformed_messages[0]["content"][1]["type"] == "text" + assert transformed_messages[0]["content"][1]["text"] == "World" diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index ddc01db0ec7..4dc7c7a4054 100644 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -165,3 +165,95 @@ def test_hosted_vllm_supports_thinking(): drop_params=False, ) assert optional_params["reasoning_effort"] == "low" + + +def test_hosted_vllm_thinking_blocks_prepended_to_assistant_content(): + """ + Test that thinking_blocks on assistant messages are converted to content + blocks prepended before the existing content. + """ + config = HostedVLLMChatConfig() + messages = [ + { + "role": "user", + "content": "Hello", + }, + { + "role": "assistant", + "content": "Here is my answer.", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "Let me reason about this...", + "signature": "abc123", + } + ], + }, + { + "role": "user", + "content": "Follow up question", + }, + ] + transformed = config.transform_request( + model="hosted_vllm/llama-3.1-70b-instruct", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assistant_msg = transformed["messages"][1] + assert assistant_msg["role"] == "assistant" + assert isinstance(assistant_msg["content"], list) + assert assistant_msg["content"][0] == { + "type": "thinking", + "thinking": "Let me reason about this...", + } + assert assistant_msg["content"][1] == { + "type": "text", + "text": "Here is my answer.", + } + assert "thinking_blocks" not in assistant_msg + + +def test_hosted_vllm_thinking_blocks_with_list_content(): + """ + Test thinking_blocks prepended when assistant content is already a list. + """ + config = HostedVLLMChatConfig() + messages = [ + { + "role": "assistant", + "content": [{"type": "text", "text": "Response text"}], + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "Step 1 reasoning", + "signature": "sig1", + }, + { + "type": "thinking", + "thinking": "Step 2 reasoning", + "signature": "sig2", + }, + ], + }, + ] + transformed = config.transform_request( + model="hosted_vllm/llama-3.1-70b-instruct", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assistant_msg = transformed["messages"][0] + assert len(assistant_msg["content"]) == 3 + assert assistant_msg["content"][0] == { + "type": "thinking", + "thinking": "Step 1 reasoning", + } + assert assistant_msg["content"][1] == { + "type": "thinking", + "thinking": "Step 2 reasoning", + } + assert assistant_msg["content"][2] == {"type": "text", "text": "Response text"} + assert "thinking_blocks" not in assistant_msg diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 6047da66b6d..bef3a69bb90 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1972,7 +1972,7 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): model=model, drop_params=False, ) - assert result["thinkingConfig"]["thinkingLevel"] == "high" + assert result["thinkingConfig"]["thinkingLevel"] == "medium" assert result["thinkingConfig"]["includeThoughts"] is True # Test high -> high + includeThoughts=True @@ -2061,7 +2061,7 @@ def test_reasoning_effort_dict_format_gemini_3(): model=model, drop_params=False, ) - assert result["thinkingConfig"]["thinkingLevel"] == "high" + assert result["thinkingConfig"]["thinkingLevel"] == "medium" assert result["thinkingConfig"]["includeThoughts"] is True # Test dict format without effort key - should fall back to Gemini 3 default (low) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py index b4a5a8ca19b..e93638df441 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py @@ -184,6 +184,7 @@ async def test_e2e_jwt_team_mcp_permissions_enforced(monkeypatch): proxy_server_module.prisma_client = MagicMock() # Mock prisma client proxy_server_module.user_api_key_cache = DualCache() proxy_server_module.proxy_logging_obj = MagicMock() + proxy_server_module.general_settings = {} monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) # Team "ABC" has MCP servers assigned via object_permission @@ -389,6 +390,7 @@ async def test_e2e_jwt_team_mcp_key_intersection(monkeypatch): proxy_server_module.prisma_client = MagicMock() proxy_server_module.user_api_key_cache = DualCache() proxy_server_module.proxy_logging_obj = MagicMock() + proxy_server_module.general_settings = {} monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) # Team MCP servers diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index d630ba6aaa5..b5f2197407a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -332,7 +332,7 @@ async def test_mcp_get_prompt_success(): mock_manager.get_prompt_from_server = AsyncMock(return_value=prompt_result) result = await mcp_get_prompt( - name="hello", + name="server_a-hello", # prefixed name since server prefixes are always added arguments={"foo": "bar"}, user_api_key_auth=user_api_key_auth, ) @@ -1006,7 +1006,7 @@ async def test_oauth2_headers_passed_to_mcp_client(): @pytest.mark.asyncio async def test_list_tools_single_server_unprefixed_names(): - """When only one MCP server is allowed, list tools should return unprefixed names.""" + """When only one MCP server is allowed, list tools should return prefixed names (server prefix is always added).""" try: from litellm.proxy._experimental.mcp_server.server import ( _get_tools_from_mcp_servers, @@ -1063,9 +1063,9 @@ async def test_list_tools_single_server_unprefixed_names(): mcp_server_auth_headers=None, ) - # Should be unprefixed since only one server is allowed + # Server prefix is always added regardless of number of allowed servers assert len(tools) == 1 - assert tools[0].name == "toolA" + assert tools[0].name == "zapier-toolA" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 1dd5cba2c4b..0fa0d4cf10b 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -7,7 +7,7 @@ sys.path.insert( ) # Adds the parent directory to the system path -from datetime import datetime +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch, call import pytest @@ -1000,3 +1000,79 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): table_name="litellm_dailyuserspend", unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) + + +@pytest.mark.asyncio +async def test_commit_key_spend_updates_includes_last_active(): + """ + Test that _commit_spend_updates_to_db sets last_active alongside spend + when updating the key table. + """ + db_writer = DBSpendUpdateWriter() + + # Create mock prisma client with transaction support + mock_batcher = MagicMock() + mock_batcher.litellm_verificationtoken = MagicMock() + mock_batcher.litellm_verificationtoken.update_many = MagicMock() + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock(return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + )) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + + # Also mock the other table batchers to avoid errors + mock_batcher.litellm_usertable = MagicMock() + mock_batcher.litellm_usertable.update_many = MagicMock() + mock_batcher.litellm_teamtable = MagicMock() + mock_batcher.litellm_teamtable.update_many = MagicMock() + mock_batcher.litellm_organizationtable = MagicMock() + mock_batcher.litellm_organizationtable.update_many = MagicMock() + + mock_proxy_logging = MagicMock() + + db_spend_update_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {"hashed_token_abc": 0.05}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + } + + before_call = datetime.now(timezone.utc) + + with patch( + "litellm.proxy.utils._raise_failed_update_spend_exception" + ): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=mock_proxy_logging, + db_spend_update_transactions=db_spend_update_transactions, + ) + + after_call = datetime.now(timezone.utc) + + # Verify update_many was called on the key table + mock_batcher.litellm_verificationtoken.update_many.assert_called_once() + call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] + + # Verify the where clause targets the correct token + assert call_kwargs["where"] == {"token": "hashed_token_abc"} + + # Verify data includes both spend increment and last_active + assert call_kwargs["data"]["spend"] == {"increment": 0.05} + assert "last_active" in call_kwargs["data"] + + # Verify last_active is a datetime within the expected range + last_active = call_kwargs["data"]["last_active"] + assert isinstance(last_active, datetime) + assert before_call <= last_active <= after_call diff --git a/tests/test_litellm/proxy/management_endpoints/test_policy_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_policy_endpoints.py index 1e4f161c6fc..14d7b8a9367 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_policy_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_policy_endpoints.py @@ -587,3 +587,256 @@ class TestApplyPoliciesDirectGuardrailNames: # Sorted order: direct_guardrail then from_policy; final output is from_policy assert result["inputs"] == first_output assert result["guardrail_errors"] == [] + + +# --------------------------------------------------------------------------- +# Tests for competitor enrichment helper functions +# --------------------------------------------------------------------------- +from litellm.proxy.management_endpoints.policy_endpoints import ( + _build_all_names_per_competitor, + _build_comparison_blocked_words, + _build_competitor_guardrail_definitions, + _build_name_blocked_words, + _build_recommendation_blocked_words, + _build_refinement_prompt, + _clean_competitor_line, + _parse_variations_response, +) + + +class TestCleanCompetitorLine: + """Tests for _clean_competitor_line.""" + + def test_strips_bullets_and_dashes(self): + assert _clean_competitor_line("- United Airlines") == "United Airlines" + assert _clean_competitor_line(" - JetBlue ") == "JetBlue" + + def test_strips_trailing_punctuation(self): + assert _clean_competitor_line("Delta Airlines.") == "Delta Airlines" + assert _clean_competitor_line("Southwest)") == "Southwest" + + def test_returns_none_for_empty(self): + assert _clean_competitor_line("") is None + assert _clean_competitor_line(" ") is None + + def test_returns_none_for_single_char(self): + assert _clean_competitor_line("A") is None + assert _clean_competitor_line(" - ") is None + + def test_plain_name(self): + assert _clean_competitor_line("Qatar Airways") == "Qatar Airways" + + +class TestParseVariationsResponse: + """Tests for _parse_variations_response.""" + + def test_parses_standard_format(self): + raw = "Delta Airlines: Delta Air Lines, DeltaAirlines, Delta\nUnited Airlines: United, UAL" + competitors = ["Delta Airlines", "United Airlines"] + result = _parse_variations_response(raw, competitors) + assert "Delta Airlines" in result + assert "Delta Air Lines" in result["Delta Airlines"] + assert "United" in result["United Airlines"] + + def test_case_insensitive_matching(self): + raw = "delta airlines: Delta Air Lines, DeltaAirlines" + competitors = ["Delta Airlines"] + result = _parse_variations_response(raw, competitors) + assert "Delta Airlines" in result + assert len(result["Delta Airlines"]) == 2 + + def test_skips_lines_without_colon(self): + raw = "This is a header\nDelta Airlines: Delta Air Lines" + competitors = ["Delta Airlines"] + result = _parse_variations_response(raw, competitors) + assert len(result) == 1 + + def test_skips_unknown_competitors(self): + raw = "Unknown Corp: Foo, Bar\nDelta Airlines: Delta" + competitors = ["Delta Airlines"] + result = _parse_variations_response(raw, competitors) + assert "Unknown Corp" not in result + assert "Delta Airlines" in result + + def test_filters_out_self_reference(self): + raw = "Delta Airlines: Delta Airlines, Delta Air Lines" + competitors = ["Delta Airlines"] + result = _parse_variations_response(raw, competitors) + # "Delta Airlines" should be filtered out (same as canonical) + assert "Delta Airlines" not in result["Delta Airlines"] + assert "Delta Air Lines" in result["Delta Airlines"] + + def test_empty_input(self): + assert _parse_variations_response("", []) == {} + + +class TestBuildRefinementPrompt: + """Tests for _build_refinement_prompt.""" + + def test_includes_brand_name(self): + prompt = _build_refinement_prompt("add 10 more", ["Delta"], "Emirates") + assert "Emirates" in prompt + + def test_includes_existing_competitors(self): + prompt = _build_refinement_prompt("add more", ["Delta", "United"], "Emirates") + assert "Delta" in prompt + assert "United" in prompt + + def test_includes_instruction(self): + prompt = _build_refinement_prompt("add 10 from Asia", ["Delta"], "Emirates") + assert "add 10 from Asia" in prompt + + def test_asks_for_new_names_only(self): + prompt = _build_refinement_prompt("add more", ["Delta"], "Emirates") + assert "NEW" in prompt + + +class TestBuildAllNamesPerCompetitor: + """Tests for _build_all_names_per_competitor.""" + + def test_includes_canonical_and_variations(self): + result = _build_all_names_per_competitor( + ["Delta Airlines"], {"Delta Airlines": ["Delta", "DeltaAir"]} + ) + assert result["Delta Airlines"] == ["Delta Airlines", "Delta", "DeltaAir"] + + def test_no_variations(self): + result = _build_all_names_per_competitor(["Delta Airlines"], {}) + assert result["Delta Airlines"] == ["Delta Airlines"] + + def test_multiple_competitors(self): + result = _build_all_names_per_competitor( + ["Delta", "United"], + {"Delta": ["DL"], "United": ["UA"]}, + ) + assert len(result) == 2 + assert result["Delta"] == ["Delta", "DL"] + assert result["United"] == ["United", "UA"] + + +class TestBuildNameBlockedWords: + """Tests for _build_name_blocked_words.""" + + def test_basic_output(self): + all_names = {"Delta": ["Delta", "DL"]} + result = _build_name_blocked_words(["Delta"], all_names) + keywords = [r["keyword"] for r in result] + assert "Delta" in keywords + assert "DL" in keywords + assert all(r["action"] == "BLOCK" for r in result) + + def test_descriptions_differ_for_variations(self): + all_names = {"Delta": ["Delta", "DL"]} + result = _build_name_blocked_words(["Delta"], all_names) + descs = {r["keyword"]: r["description"] for r in result} + assert "Competitor: Delta" == descs["Delta"] + assert "variation" in descs["DL"].lower() + + +class TestBuildRecommendationBlockedWords: + """Tests for _build_recommendation_blocked_words.""" + + def test_generates_prefix_combinations(self): + all_names = {"Delta": ["Delta"]} + result = _build_recommendation_blocked_words(["Delta"], all_names) + keywords = [r["keyword"] for r in result] + assert "try Delta" in keywords + assert "use Delta" in keywords + assert "switch to Delta" in keywords + assert "consider Delta" in keywords + + def test_includes_variations(self): + all_names = {"Delta": ["Delta", "DL"]} + result = _build_recommendation_blocked_words(["Delta"], all_names) + keywords = [r["keyword"] for r in result] + assert "try DL" in keywords + + +class TestBuildComparisonBlockedWords: + """Tests for _build_comparison_blocked_words.""" + + def test_generates_competitor_comparisons(self): + all_names = {"Delta": ["Delta"]} + result = _build_comparison_blocked_words(["Delta"], all_names, "Emirates") + keywords = [r["keyword"] for r in result] + assert "Delta is better" in keywords + + def test_generates_brand_comparisons_once(self): + all_names = {"Delta": ["Delta"], "United": ["United"]} + result = _build_comparison_blocked_words(["Delta", "United"], all_names, "Emirates") + keywords = [r["keyword"] for r in result] + # Brand-level entries should appear exactly once + assert keywords.count("better than Emirates") == 1 + assert keywords.count("Emirates is worse") == 1 + + def test_includes_variation_comparisons(self): + all_names = {"Delta": ["Delta", "DL"]} + result = _build_comparison_blocked_words(["Delta"], all_names, "Emirates") + keywords = [r["keyword"] for r in result] + assert "DL is better" in keywords + + +class TestBuildCompetitorGuardrailDefinitions: + """Tests for _build_competitor_guardrail_definitions.""" + + def test_populates_blocked_words_for_known_guardrail_names(self): + definitions = [ + { + "guardrail_name": "competitor-name-blocker", + "litellm_params": {"blocked_words": []}, + }, + { + "guardrail_name": "competitor-recommendation-filter", + "litellm_params": {"blocked_words": []}, + }, + ] + result = _build_competitor_guardrail_definitions( + definitions, ["Delta"], "Emirates", {"Delta": ["DL"]} + ) + # Name blocker should have entries + name_blocker = next(d for d in result if d["guardrail_name"] == "competitor-name-blocker") + assert len(name_blocker["litellm_params"]["blocked_words"]) > 0 + + # Recommendation filter should have entries + rec_filter = next(d for d in result if d["guardrail_name"] == "competitor-recommendation-filter") + assert len(rec_filter["litellm_params"]["blocked_words"]) > 0 + + def test_does_not_modify_unknown_guardrail_names(self): + definitions = [ + { + "guardrail_name": "some-other-guardrail", + "litellm_params": {"blocked_words": ["original"]}, + }, + ] + result = _build_competitor_guardrail_definitions( + definitions, ["Delta"], "Emirates" + ) + assert result[0]["litellm_params"]["blocked_words"] == ["original"] + + def test_does_not_mutate_original_definitions(self): + definitions = [ + { + "guardrail_name": "competitor-name-blocker", + "litellm_params": {"blocked_words": []}, + }, + ] + _build_competitor_guardrail_definitions(definitions, ["Delta"], "Emirates") + # Original should be unchanged + assert definitions[0]["litellm_params"]["blocked_words"] == [] + + def test_handles_input_and_output_blocker_variants(self): + definitions = [ + { + "guardrail_name": "competitor-name-input-blocker", + "litellm_params": {"blocked_words": []}, + }, + { + "guardrail_name": "competitor-name-output-blocker", + "litellm_params": {"blocked_words": []}, + }, + ] + result = _build_competitor_guardrail_definitions( + definitions, ["Delta"], "Emirates" + ) + for defn in result: + assert len(defn["litellm_params"]["blocked_words"]) > 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index c4c953b75fb..b6ac974e2cf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -5496,6 +5496,190 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) assert False, "API keys should not be fetched for team admin users" +@pytest.mark.asyncio +async def test_get_team_daily_activity_member_with_permission_sees_all_spend( + mock_db_client, +): + """ + Test that non-admin team members with /team/daily/activity permission + can see all team spend (no API key filtering), same as team admins. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + get_team_daily_activity, + ) + + # Create a non-admin user + user_id = "test_user_with_perm_123" + team_id = "test_team_789" + user_api_key_dict = UserAPIKeyAuth( + user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER + ) + + # Mock user info + mock_user_info = LiteLLM_UserTable( + user_id=user_id, + teams=[team_id], + max_budget=1000.0, + spend=0.0, + user_email="member@example.com", + user_role="internal_user", + ) + + # Mock team with user as non-admin member AND /team/daily/activity permission + mock_team_member = Member(user_id=user_id, role="user") + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.team_id = team_id + mock_team.team_alias = "Test Team" + mock_team.members_with_roles = [mock_team_member] + mock_team.team_member_permissions = ["/team/daily/activity"] + mock_team.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Test Team", + "members_with_roles": [{"user_id": user_id, "role": "user"}], + "team_member_permissions": ["/team/daily/activity"], + } + + # Setup mocks + mock_db_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[mock_team] + ) + + # Mock get_user_object + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + ) as mock_get_user_object: + mock_get_user_object.return_value = mock_user_info + + # Mock get_daily_activity to capture the api_key parameter + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + new_callable=AsyncMock, + ) as mock_get_daily_activity: + mock_get_daily_activity.return_value = MagicMock() + + # Call the endpoint + await get_team_daily_activity( + team_ids=team_id, + start_date="2024-01-01", + end_date="2024-01-02", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=user_api_key_dict, + ) + + # Verify get_daily_activity was called WITHOUT API key filtering + mock_get_daily_activity.assert_called_once() + call_kwargs = mock_get_daily_activity.call_args[1] + assert call_kwargs["api_key"] is None + assert call_kwargs["entity_id"] == [team_id] + + # Verify user's API keys were NOT fetched + if hasattr( + mock_db_client.db.litellm_verificationtoken, "find_many" + ) and mock_db_client.db.litellm_verificationtoken.find_many.called: + assert ( + False + ), "API keys should not be fetched for members with /team/daily/activity permission" + + +@pytest.mark.asyncio +async def test_get_team_daily_activity_member_without_permission_filters_by_keys( + mock_db_client, +): + """ + Test that non-admin team members WITHOUT /team/daily/activity permission + still have their results filtered by their own API keys. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + get_team_daily_activity, + ) + + # Create a non-admin user + user_id = "test_user_no_perm_123" + team_id = "test_team_789" + user_api_key_dict = UserAPIKeyAuth( + user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER + ) + + # Mock user info + mock_user_info = LiteLLM_UserTable( + user_id=user_id, + teams=[team_id], + max_budget=1000.0, + spend=0.0, + user_email="member@example.com", + user_role="internal_user", + ) + + # Mock team with user as non-admin member and NO usage permission + mock_team_member = Member(user_id=user_id, role="user") + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.team_id = team_id + mock_team.team_alias = "Test Team" + mock_team.members_with_roles = [mock_team_member] + mock_team.team_member_permissions = ["/key/info"] + mock_team.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Test Team", + "members_with_roles": [{"user_id": user_id, "role": "user"}], + "team_member_permissions": ["/key/info"], + } + + # Mock user's API keys + user_api_key_1 = MagicMock() + user_api_key_1.token = "user_key_abc" + user_api_key_2 = MagicMock() + user_api_key_2.token = "user_key_def" + + # Setup mocks + mock_db_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[mock_team] + ) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[user_api_key_1, user_api_key_2] + ) + + # Mock get_user_object + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + ) as mock_get_user_object: + mock_get_user_object.return_value = mock_user_info + + # Mock get_daily_activity to capture the api_key parameter + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + new_callable=AsyncMock, + ) as mock_get_daily_activity: + mock_get_daily_activity.return_value = MagicMock() + + # Call the endpoint + await get_team_daily_activity( + team_ids=team_id, + start_date="2024-01-01", + end_date="2024-01-02", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=user_api_key_dict, + ) + + # Verify get_daily_activity was called WITH API key filtering + mock_get_daily_activity.assert_called_once() + call_kwargs = mock_get_daily_activity.call_args[1] + assert call_kwargs["api_key"] == ["user_key_abc", "user_key_def"] + assert call_kwargs["entity_id"] == [team_id] + + # Verify user's API keys were fetched + mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once() + + @pytest.mark.asyncio async def test_update_team_with_router_settings(mock_db_client, mock_admin_auth): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 511979b781a..3a8d17ccb45 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3750,11 +3750,10 @@ async def test_sso_role_preserved_without_role_mappings(): ), "user_defined_values should retain the SSO role after insert_sso_user" finally: - if original_default_params is not None: - litellm.default_internal_user_params = original_default_params - else: - if hasattr(litellm, "default_internal_user_params"): - delattr(litellm, "default_internal_user_params") + # Restore original default_internal_user_params (always assign, never delattr — + # deleting the attribute causes AttributeError in subsequent tests because + # litellm.__getattr__ has no handler for this name) + litellm.default_internal_user_params = original_default_params class TestSSOReadinessEndpoint: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_method_specific_routing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_method_specific_routing.py new file mode 100644 index 00000000000..7e5b9ff6403 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_method_specific_routing.py @@ -0,0 +1,155 @@ +""" +Test method-specific routing for pass-through endpoints. + +This test demonstrates the ability to configure different targets +for the same path but different HTTP methods. +""" + +import pytest + +from litellm.proxy._types import PassThroughGenericEndpoint + + +def test_pass_through_endpoint_with_methods(): + """Test creating pass-through endpoints with specific methods""" + + # Create endpoint for GET /azure/kb + get_endpoint = PassThroughGenericEndpoint( + id="get-azure-kb", + path="/azure/kb", + target="https://api1.example.com/knowledge-base", + methods=["GET"], + headers={"Authorization": "Bearer token1"}, + ) + + assert get_endpoint.path == "/azure/kb" + assert get_endpoint.methods == ["GET"] + assert get_endpoint.target == "https://api1.example.com/knowledge-base" + + # Create endpoint for POST /azure/kb + post_endpoint = PassThroughGenericEndpoint( + id="post-azure-kb", + path="/azure/kb", + target="https://api2.example.com/knowledge-base", + methods=["POST"], + headers={"Authorization": "Bearer token2"}, + ) + + assert post_endpoint.path == "/azure/kb" + assert post_endpoint.methods == ["POST"] + assert post_endpoint.target == "https://api2.example.com/knowledge-base" + + # These should be different endpoints despite same path + assert get_endpoint.id != post_endpoint.id + assert get_endpoint.target != post_endpoint.target + + +def test_pass_through_endpoint_multiple_methods(): + """Test creating endpoint with multiple methods""" + + endpoint = PassThroughGenericEndpoint( + id="multi-method", + path="/azure/kb", + target="https://api.example.com/kb", + methods=["GET", "POST", "PUT"], + headers={}, + ) + + assert len(endpoint.methods) == 3 + assert "GET" in endpoint.methods + assert "POST" in endpoint.methods + assert "PUT" in endpoint.methods + + +def test_pass_through_endpoint_no_methods_backward_compatibility(): + """Test that endpoints without methods field work (backward compatibility)""" + + # When methods is None, all methods should be supported + endpoint = PassThroughGenericEndpoint( + id="all-methods", + path="/azure/kb", + target="https://api.example.com/kb", + headers={}, + ) + + assert endpoint.methods is None # Default is None for backward compatibility + + +def test_pass_through_endpoint_serialization(): + """Test that endpoints with methods can be serialized/deserialized""" + + endpoint = PassThroughGenericEndpoint( + id="test-endpoint", + path="/test", + target="https://api.example.com", + methods=["GET", "POST"], + headers={"key": "value"}, + cost_per_request=0.5, + ) + + # Serialize to dict + endpoint_dict = endpoint.model_dump() + assert endpoint_dict["methods"] == ["GET", "POST"] + + # Deserialize from dict + restored_endpoint = PassThroughGenericEndpoint(**endpoint_dict) + assert restored_endpoint.methods == ["GET", "POST"] + assert restored_endpoint.path == "/test" + assert restored_endpoint.target == "https://api.example.com" + + +def test_route_key_generation_with_methods(): + """Test that route keys include methods for uniqueness""" + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + # Simulate how route keys are generated + endpoint_id_1 = "endpoint-1" + path = "/azure/kb" + methods_1 = ["GET"] + methods_str_1 = ",".join(sorted(methods_1)) + route_key_1 = f"{endpoint_id_1}:exact:{path}:{methods_str_1}" + + endpoint_id_2 = "endpoint-2" + methods_2 = ["POST"] + methods_str_2 = ",".join(sorted(methods_2)) + route_key_2 = f"{endpoint_id_2}:exact:{path}:{methods_str_2}" + + # Keys should be different even though path is the same + assert route_key_1 != route_key_2 + assert route_key_1 == "endpoint-1:exact:/azure/kb:GET" + assert route_key_2 == "endpoint-2:exact:/azure/kb:POST" + + +def test_config_yaml_example(): + """ + Example configuration for config.yaml showing method-specific routing: + + general_settings: + pass_through_endpoints: + # GET endpoint for retrieving knowledge base + - id: "get-azure-kb" + path: "/azure/kb" + target: "https://read-api.example.com/kb" + methods: ["GET"] + headers: + Authorization: "bearer os.environ/READ_API_KEY" + + # POST endpoint for creating knowledge base entries + - id: "post-azure-kb" + path: "/azure/kb" + target: "https://write-api.example.com/kb" + methods: ["POST"] + headers: + Authorization: "bearer os.environ/WRITE_API_KEY" + + # PUT endpoint for updating knowledge base + - id: "put-azure-kb" + path: "/azure/kb" + target: "https://update-api.example.com/kb" + methods: ["PUT"] + headers: + Authorization: "bearer os.environ/UPDATE_API_KEY" + """ + pass diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 731ecd7379d..79b5e34022f 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1999,22 +1999,27 @@ class TestPriceDataReloadIntegration: "gpt-4": {"input_cost_per_token": 0.03, "output_cost_per_token": 0.06}, } - with patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map: - mock_get_map.return_value = mock_cost_map + original_model_cost = litellm.model_cost.copy() + try: + with patch( + "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" + ) as mock_get_map: + mock_get_map.return_value = mock_cost_map - # Mock the database connection - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + # Mock the database connection + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - # Test reload endpoint - response = client_with_auth.post("/reload/model_cost_map") - assert response.status_code == 200 + # Test reload endpoint + response = client_with_auth.post("/reload/model_cost_map") + assert response.status_code == 200 - # Test get endpoint - response = client_with_auth.get("/public/litellm_model_cost_map") - assert response.status_code == 200 + # Test get endpoint + response = client_with_auth.get("/public/litellm_model_cost_map") + assert response.status_code == 200 + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() def test_distributed_reload_check_function(self): """Test the _check_and_reload_model_cost_map function""" @@ -2054,23 +2059,28 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - with patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map: - mock_get_map.return_value = { - "gpt-3.5-turbo": {"input_cost_per_token": 0.001} - } + original_model_cost = litellm.model_cost.copy() + try: + with patch( + "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" + ) as mock_get_map: + mock_get_map.return_value = { + "gpt-3.5-turbo": {"input_cost_per_token": 0.001} + } - # Should reload due to force flag - asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + # Should reload due to force flag + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - # Verify force_reload was reset to False - mock_prisma.db.litellm_config.upsert.assert_called() - call_args = mock_prisma.db.litellm_config.upsert.call_args - # The param_value is now a JSON string, so we need to parse it - param_value_json = call_args[1]["data"]["update"]["param_value"] - param_value_dict = json.loads(param_value_json) - assert param_value_dict["force_reload"] == False + # Verify force_reload was reset to False + mock_prisma.db.litellm_config.upsert.assert_called() + call_args = mock_prisma.db.litellm_config.upsert.call_args + # The param_value is now a JSON string, so we need to parse it + param_value_json = call_args[1]["data"]["update"]["param_value"] + param_value_dict = json.loads(param_value_json) + assert param_value_dict["force_reload"] == False + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() def test_config_file_parsing(self): """Test parsing of config file with reload settings""" diff --git a/tests/test_litellm/test_cost_calculation_log_level.py b/tests/test_litellm/test_cost_calculation_log_level.py index 8ee9ad95cd0..61f1e716e4d 100644 --- a/tests/test_litellm/test_cost_calculation_log_level.py +++ b/tests/test_litellm/test_cost_calculation_log_level.py @@ -67,7 +67,7 @@ def test_cost_calculation_uses_debug_level(): # Find the cost calculation log records cost_calc_records = [ record for record in handler.records - if "selected model name for cost calculation" in record.message + if "selected model name for cost calculation" in record.getMessage() ] # Verify that cost calculation logs are at DEBUG level @@ -126,7 +126,7 @@ def test_batch_cost_calculation_uses_debug_level(): # Find batch cost calculation log records batch_cost_records = [ record for record in handler.records - if "Calculating batch cost per token" in record.message + if "Calculating batch cost per token" in record.getMessage() ] # Verify logs exist and are at DEBUG level diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 7b29a4d90aa..6170f6b6f55 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -56,6 +56,55 @@ def test_check_provider_match_azure_ai_allows_openai_and_azure(): ) is False +def test_check_provider_match_github_allows_upstream_provider_metadata(): + """ + Test that github provider can match upstream provider metadata. + GitHub Models can provide models from multiple providers. + """ + assert _check_provider_match( + model_info={"litellm_provider": "openai"}, + custom_llm_provider="github", + ) is True + + assert _check_provider_match( + model_info={"litellm_provider": "github"}, + custom_llm_provider="github", + ) is True + + assert _check_provider_match( + model_info={"litellm_provider": "anthropic"}, + custom_llm_provider="github", + ) is True + + +def test_supports_function_calling_github_openai_alias(): + assert litellm.utils.supports_function_calling(model="github/gpt-4o-mini") is True + assert ( + litellm.utils.supports_function_calling( + model="gpt-4o-mini", custom_llm_provider="github" + ) + is True + ) + + +def test_supports_function_calling_github_anthropic_alias(): + assert ( + litellm.utils.supports_function_calling( + model="github/claude-3-5-sonnet-latest" + ) + is True + ) + + +def test_supports_function_calling_unknown_github_alias_returns_false(): + assert ( + litellm.utils.supports_function_calling( + model="github/non-existent-model-for-capability-check" + ) + is False + ) + + def test_get_optional_params_image_gen(): from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts new file mode 100644 index 00000000000..6ff784ebd90 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts @@ -0,0 +1,147 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useStoreModelInDB } from "./useStoreModelInDB"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => ""), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), +})); + +describe("useStoreModelInDB", () => { + let queryClient: QueryClient; + let fetchSpy: ReturnType; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + mutations: { retry: false }, + }, + }); + + fetchSpy = vi.fn(); + global.fetch = fetchSpy; + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should send correct request body to /config/field/update", async () => { + fetchSpy.mockResolvedValue({ + ok: true, + json: async () => ({ message: "Success" }), + }); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: true }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(fetchSpy).toHaveBeenCalledWith( + "/config/field/update", + expect.objectContaining({ + method: "POST", + body: JSON.stringify({ + field_name: "store_model_in_db", + field_value: true, + config_type: "general_settings", + }), + }) + ); + }); + + it("should handle setting store_model_in_db to false", async () => { + fetchSpy.mockResolvedValue({ + ok: true, + json: async () => ({ message: "Success" }), + }); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: false }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(fetchSpy).toHaveBeenCalledWith( + "/config/field/update", + expect.objectContaining({ + body: JSON.stringify({ + field_name: "store_model_in_db", + field_value: false, + config_type: "general_settings", + }), + }) + ); + }); + + it("should throw error when access token is missing", async () => { + vi.spyOn( + await import("../useAuthorized"), + "default" + ).mockReturnValue({ + accessToken: null, + userRole: null, + userId: null, + token: null, + userEmail: null, + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + } as any); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: true }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error?.message).toBe("Access token is required"); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("should handle API error response", async () => { + fetchSpy.mockResolvedValue({ + ok: false, + json: async () => ({ detail: "Unauthorized" }), + }); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: true }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error?.message).toBe("Unauthorized"); + }); + + it("should use fallback error message when API returns empty error", async () => { + fetchSpy.mockResolvedValue({ + ok: false, + json: async () => ({}), + }); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: true }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error?.message).toBe("Failed to update model storage settings"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts new file mode 100644 index 00000000000..e6efbd724cd --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts @@ -0,0 +1,59 @@ +import { useMutation, UseMutationResult } from "@tanstack/react-query"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; +import useAuthorized from "../useAuthorized"; + +export interface StoreModelInDBParams { + store_model_in_db: boolean; +} + +export interface StoreModelInDBResponse { + message: string; +} + +const performStoreModelInDB = async ( + accessToken: string, + params: StoreModelInDBParams +): Promise => { + const proxyBaseUrl = getProxyBaseUrl(); + const url = proxyBaseUrl ? `${proxyBaseUrl}/config/field/update` : `/config/field/update`; + + const response = await fetch(url, { + method: "POST", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + field_name: "store_model_in_db", + field_value: params.store_model_in_db, + config_type: "general_settings", + }), + }); + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})); + const errorMessage = + errorData?.error?.message || errorData?.message || errorData?.detail || "Failed to update model storage settings"; + throw new Error(errorMessage); + } + + const data = await response.json(); + return data; +}; + +export const useStoreModelInDB = (): UseMutationResult< + StoreModelInDBResponse, + Error, + StoreModelInDBParams +> => { + const { accessToken } = useAuthorized(); + + return useMutation({ + mutationFn: async (params: StoreModelInDBParams) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return await performStoreModelInDB(accessToken, params); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index e252f273316..36948630d8a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -5,10 +5,11 @@ import { Team } from "@/components/key_team_helpers/key_list"; import { AllModelsDataTable } from "@/components/model_dashboard/all_models_table"; import { columns } from "@/components/molecules/models/columns"; import { getDisplayModelName } from "@/components/view_model/model_name_display"; -import { InfoCircleOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined, SettingOutlined } from "@ant-design/icons"; import { PaginationState, SortingState } from "@tanstack/react-table"; import { Grid, TabPanel } from "@tremor/react"; -import { Badge, Select, Skeleton, Space, Typography } from "antd"; +import { Badge, Button, Select, Skeleton, Space, Typography } from "antd"; +import ModelSettingsModal from "@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal"; import debounce from "lodash/debounce"; import { useEffect, useMemo, useState } from "react"; import { useModelsInfo } from "../../hooks/models/useModels"; @@ -51,6 +52,7 @@ const AllModelsTab = ({ pageSize: 50, }); const [sorting, setSorting] = useState([]); + const [isModelSettingsModalVisible, setIsModelSettingsModalVisible] = useState(false); // Debounce search input const debouncedUpdateSearch = useMemo( @@ -326,62 +328,71 @@ const AllModelsTab = ({
{/* Search and Filter Controls */} -
- {/* Model Name Search */} -
- setModelNameSearch(e.target.value)} - /> - - +
+ {/* Model Name Search */} +
+ setModelNameSearch(e.target.value)} /> - + + + +
+ + {/* Filter Button */} + + + {/* Reset Filters Button */} +
- {/* Filter Button */} - - - {/* Reset Filters Button */} - + {/* Model Settings Button */} +
{/* Additional Filters */} @@ -505,6 +516,11 @@ const AllModelsTab = ({
+ setIsModelSettingsModalVisible(false)} + onSuccess={() => setIsModelSettingsModalVisible(false)} + /> ); }; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 749396c82f3..9dc468c6a9f 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -99,6 +99,7 @@ const mockKey: KeyResponse = { created_at: "2024-11-01T10:00:00Z", created_by: "user-1", updated_at: "2024-11-15T10:00:00Z", + last_active: "2024-11-20T14:30:00Z", team_spend: 5.5, team_alias: "Test Team", team_tpm_limit: 5000, @@ -625,3 +626,78 @@ it("should render table without crashing when models is undefined", async () => expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); }); + +it("should render Last Active column header with info icon", () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + expect(screen.getByText("Last Active")).toBeInTheDocument(); +}); + +it("should display formatted date for last_active when value exists", async () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + await waitFor(() => { + const expectedDate = new Date("2024-11-20T14:30:00Z").toLocaleDateString(); + expect(screen.getByText(expectedDate)).toBeInTheDocument(); + }); +}); + +it("should display 'Unknown' for last_active when value is null", async () => { + const keyWithNullLastActive = { + ...mockKey, + last_active: null, + }; + + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [keyWithNullLastActive], + allKeyAliases: ["test-key-alias"], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Unknown")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index f7c47943e7e..ff996a0434a 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -24,8 +24,9 @@ import { TableRow, Text, } from "@tremor/react"; -import { Skeleton, Tooltip } from "antd"; -import React, { useEffect, useState } from "react"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { Popover, Skeleton, Tooltip } from "antd"; +import React, { useEffect, useMemo, useState } from "react"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; import { useFilterLogic } from "../key_team_helpers/filter_logic"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; @@ -112,7 +113,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo } }, [refetch]); - const columns: ColumnDef[] = [ + const columns: ColumnDef[] = useMemo(() => [ { id: "expander", header: () => null, @@ -292,6 +293,33 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo return value ? new Date(value as string).toLocaleDateString() : "Never"; }, }, + { + id: "last_active", + accessorKey: "last_active", + header: () => ( + + Last Active + + + + + ), + size: 130, + enableSorting: false, + cell: (info) => { + const value = info.getValue(); + if (!value) return "Unknown"; + const date = new Date(value as string); + return ( + + {date.toLocaleDateString()} + + ); + }, + }, { id: "expires", accessorKey: "expires", @@ -437,7 +465,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo ); }, }, - ]; + ], []); const filterOptions: FilterOption[] = [ { diff --git a/ui/litellm-dashboard/src/components/add_pass_through.tsx b/ui/litellm-dashboard/src/components/add_pass_through.tsx index ff724e42eb8..af55a7919a9 100644 --- a/ui/litellm-dashboard/src/components/add_pass_through.tsx +++ b/ui/litellm-dashboard/src/components/add_pass_through.tsx @@ -23,6 +23,7 @@ import { ApiOutlined, } from "@ant-design/icons"; import KeyValueInput from "./key_value_input"; +import QueryParamInput from "./query_param_input"; import { passThroughItem } from "./pass_through_settings"; import RoutePreview from "./route_preview"; import NotificationsManager from "./molecules/notifications_manager"; @@ -30,6 +31,8 @@ import PassThroughSecuritySection from "./common_components/PassThroughSecurityS import PassThroughGuardrailsSection from "./common_components/PassThroughGuardrailsSection"; const { Option } = Select2; +const HTTP_METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH"]; + interface AddFallbacksProps { // models: string[] | undefined; accessToken: string; @@ -52,12 +55,14 @@ const AddPassThroughEndpoint: React.FC = ({ const [targetValue, setTargetValue] = useState(""); const [includeSubpath, setIncludeSubpath] = useState(true); const [authEnabled, setAuthEnabled] = useState(false); + const [selectedMethods, setSelectedMethods] = useState([]); const [guardrails, setGuardrails] = useState>({}); const handleCancel = () => { form.resetFields(); setPathValue(""); setTargetValue(""); setIncludeSubpath(true); + setSelectedMethods([]); setGuardrails({}); setIsModalVisible(false); }; @@ -86,6 +91,11 @@ const AddPassThroughEndpoint: React.FC = ({ formValues.guardrails = guardrails; } + // Add methods to formValues (only if specific methods are selected) + if (selectedMethods && selectedMethods.length > 0) { + formValues.methods = selectedMethods; + } + console.log(`formValues: ${JSON.stringify(formValues)}`); const response = await createPassThroughEndpoint(accessToken, formValues); @@ -101,6 +111,7 @@ const AddPassThroughEndpoint: React.FC = ({ setPathValue(""); setTargetValue(""); setIncludeSubpath(true); + setSelectedMethods([]); setGuardrails({}); setIsModalVisible(false); } catch (error) { @@ -204,6 +215,41 @@ const AddPassThroughEndpoint: React.FC = ({ /> + + HTTP Methods (Optional) + + + + + } + name="methods" + extra={ +
+ {selectedMethods.length === 0 + ? "All HTTP methods supported (default)" + : `Only ${selectedMethods.join(", ")} requests will be routed to this endpoint`} +
+ } + className="mb-4" + > + + {HTTP_METHODS.map((method) => ( + + ))} + + +
Include Subpaths
@@ -250,6 +296,34 @@ const AddPassThroughEndpoint: React.FC = ({ + {/* Default Query Parameters Section */} + + Default Query Parameters + + Add query parameters that will be automatically sent with every request to the target API + + + + Default Query Parameters (Optional) + + + + + } + name="default_query_params" + extra={ +
+
Parameters are sent with all GET, POST, PUT, PATCH requests
+
Client parameters override defaults. Examples: version=v1, format=json, key=default
+
+ } + > + + + + {/* Security Section */} = ({ onChange, value, @@ -21,7 +26,7 @@ const PassThroughRoutesSelector: React.FC = ({ disabled = false, teamId, }) => { - const [passThroughRoutes, setPassThroughRoutes] = useState([]); + const [passThroughRoutes, setPassThroughRoutes] = useState>([]); const [loading, setLoading] = useState(false); useEffect(() => { @@ -32,7 +37,24 @@ const PassThroughRoutesSelector: React.FC = ({ try { const response = await getPassThroughEndpointsCall(accessToken, teamId); if (response.endpoints) { - const routes = response.endpoints.map((route: { path: string }) => route.path); + const routes = response.endpoints.flatMap((endpoint: PassThroughEndpoint) => { + const path = endpoint.path; + const methods = endpoint.methods; + + // If methods are specified, create one entry per method + if (methods && methods.length > 0) { + return methods.map((method) => ({ + label: `${method} ${path}`, + value: path, // Keep value as path for backward compatibility + })); + } + + // If no methods specified, show just the path (all methods supported) + return [{ + label: path, + value: path, + }]; + }); setPassThroughRoutes(routes); } } catch (error) { @@ -54,10 +76,7 @@ const PassThroughRoutesSelector: React.FC = ({ loading={loading} className={className} allowClear - options={passThroughRoutes.map((route) => ({ - label: route, - value: route, - }))} + options={passThroughRoutes} optionFilterProp="label" showSearch style={{ width: "100%" }} diff --git a/ui/litellm-dashboard/src/components/general_settings.tsx b/ui/litellm-dashboard/src/components/general_settings.tsx index 18f891705e1..22c8d38a379 100644 --- a/ui/litellm-dashboard/src/components/general_settings.tsx +++ b/ui/litellm-dashboard/src/components/general_settings.tsx @@ -11,6 +11,7 @@ import { Text, Button, Icon, + Switch, } from "@tremor/react"; import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react"; import { @@ -163,7 +164,12 @@ const GeneralSettings: React.FC = ({ accessToken, user handleInputChange(value.field_name, newValue)} // Handle value change + onChange={(newValue) => handleInputChange(value.field_name, newValue)} + /> + ) : value.field_type == "Boolean" ? ( + handleInputChange(value.field_name, checked)} /> ) : null} diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index b54eb21a0ae..5512809ba3f 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -48,6 +48,7 @@ export interface KeyResponse { organization_id: string | null; created_at: string; updated_at: string; + last_active: string | null; team_spend: number; team_alias: string; team_tpm_limit: number; diff --git a/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.test.tsx new file mode 100644 index 00000000000..d4103a1b38e --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.test.tsx @@ -0,0 +1,306 @@ +import { useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"; +import { useStoreModelInDB } from "@/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; +import ModelSettingsModal from "./ModelSettingsModal"; + +vi.mock("@/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB"); +vi.mock("@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"); +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn(), +})); + +const mockUseStoreModelInDB = vi.mocked(useStoreModelInDB); +const mockUseProxyConfig = vi.mocked(useProxyConfig); +const mockNotificationsManager = vi.mocked(NotificationsManager); +const mockParseErrorMessage = vi.mocked(parseErrorMessage); + +describe("ModelSettingsModal", () => { + const mockOnCancel = vi.fn(); + const mockOnSuccess = vi.fn(); + const mockMutateAsync = vi.fn(); + const mockRefetch = vi.fn(); + + const defaultProps = { + isVisible: true, + onCancel: mockOnCancel, + onSuccess: mockOnSuccess, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockUseStoreModelInDB.mockReturnValue({ + mutateAsync: mockMutateAsync, + isPending: false, + } as any); + mockUseProxyConfig.mockReturnValue({ + data: [], + isLoading: false, + refetch: mockRefetch, + } as any); + mockParseErrorMessage.mockImplementation((error: any) => error?.message || String(error)); + }); + + it("should render the modal", () => { + renderWithProviders(); + expect(screen.getByRole("dialog")).toBeInTheDocument(); + expect(screen.getByText("Model Settings")).toBeInTheDocument(); + }); + + it("should render form field with initial values", () => { + renderWithProviders(); + expect(screen.getByText("Store Model in DB")).toBeInTheDocument(); + }); + + it("should render cancel and save buttons", () => { + renderWithProviders(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Save Settings" })).toBeInTheDocument(); + }); + + it("should call onCancel when cancel button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + await user.click(cancelButton); + + expect(mockOnCancel).toHaveBeenCalledTimes(1); + }); + + it("should call onCancel when modal close button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const closeButton = screen.getByRole("button", { name: /close/i }); + await user.click(closeButton); + + expect(mockOnCancel).toHaveBeenCalledTimes(1); + }); + + it("should toggle store model switch", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + expect(switchElement).not.toBeChecked(); + + await user.click(switchElement); + + await waitFor(() => { + expect(switchElement).toBeChecked(); + }); + }); + + it("should submit form with store_model_in_db enabled", async () => { + const user = userEvent.setup(); + mockMutateAsync.mockImplementation(async (params, options) => { + await Promise.resolve(); + options?.onSuccess?.(); + return { message: "Success" }; + }); + + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + await user.click(switchElement); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockMutateAsync).toHaveBeenCalledWith( + { store_model_in_db: true }, + expect.any(Object) + ); + }); + }); + + it("should submit form with store_model_in_db disabled", async () => { + const user = userEvent.setup(); + mockMutateAsync.mockImplementation(async (params, options) => { + await Promise.resolve(); + options?.onSuccess?.(); + return { message: "Success" }; + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockMutateAsync).toHaveBeenCalledWith( + { store_model_in_db: false }, + expect.any(Object) + ); + }); + }); + + it("should show success notification and call onSuccess on successful submission", async () => { + const user = userEvent.setup(); + mockMutateAsync.mockImplementation(async (params, options) => { + await Promise.resolve(); + options?.onSuccess?.(); + return { message: "Success" }; + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("Model storage settings updated successfully"); + expect(mockRefetch).toHaveBeenCalled(); + expect(mockOnSuccess).toHaveBeenCalledTimes(1); + }); + }); + + it("should show error notification when submission fails", async () => { + const user = userEvent.setup(); + const error = new Error("Network error"); + mockMutateAsync.mockRejectedValue(error); + mockParseErrorMessage.mockReturnValue("Network error"); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to save model storage settings: Network error"); + }); + }); + + it("should show error notification from onError callback", async () => { + const user = userEvent.setup(); + const error = new Error("Backend error"); + mockMutateAsync.mockImplementation((params, options) => { + options?.onError?.(error); + return Promise.reject(error); + }); + mockParseErrorMessage.mockReturnValue("Backend error"); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to save model storage settings: Backend error"); + }); + }); + + it("should disable cancel button when pending", () => { + mockUseStoreModelInDB.mockReturnValue({ + mutateAsync: mockMutateAsync, + isPending: true, + } as any); + + renderWithProviders(); + + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + expect(cancelButton).toBeDisabled(); + }); + + it("should disable cancel button when loading config", () => { + mockUseProxyConfig.mockReturnValue({ + data: undefined, + isLoading: true, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + expect(cancelButton).toBeDisabled(); + }); + + it("should show loading state on save button when pending", () => { + mockUseStoreModelInDB.mockReturnValue({ + mutateAsync: mockMutateAsync, + isPending: true, + } as any); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: /Saving/i }); + expect(saveButton).toBeInTheDocument(); + expect(saveButton.className).toContain("ant-btn-loading"); + }); + + it("should not render modal when isVisible is false", () => { + renderWithProviders(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + + it("should call refetch when modal opens", () => { + renderWithProviders(); + expect(mockRefetch).toHaveBeenCalledTimes(1); + }); + + it("should render form with initial values from config data", () => { + mockUseProxyConfig.mockReturnValue({ + data: [ + { + field_name: "store_model_in_db", + field_type: "bool", + field_description: "Store model in DB", + field_value: true, + stored_in_db: true, + field_default_value: false, + }, + ], + isLoading: false, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + expect(switchElement).toBeChecked(); + }); + + it("should show skeleton loader when config is loading", () => { + mockUseProxyConfig.mockReturnValue({ + data: undefined, + isLoading: true, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + expect(screen.queryByRole("switch")).not.toBeInTheDocument(); + const skeletons = document.querySelectorAll(".ant-skeleton"); + expect(skeletons.length).toBeGreaterThan(0); + }); + + it("should not call onSuccess when it is not provided", async () => { + const user = userEvent.setup(); + mockMutateAsync.mockImplementation(async (params, options) => { + await Promise.resolve(); + options?.onSuccess?.(); + return { message: "Success" }; + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalled(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.tsx b/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.tsx new file mode 100644 index 00000000000..106f82585dd --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.tsx @@ -0,0 +1,104 @@ +"use client"; + +import { ConfigType, useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"; +import { StoreModelInDBParams, useStoreModelInDB } from "@/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { Button, Form, Modal, Skeleton, Space, Switch, Typography } from "antd"; +import React, { useEffect, useMemo } from "react"; + +interface ModelSettingsModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess?: () => void; +} + +const ModelSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { + const [form] = Form.useForm(); + const { mutateAsync, isPending } = useStoreModelInDB(); + const { data: proxyConfigData, isLoading: isLoadingConfig, refetch } = useProxyConfig(ConfigType.GENERAL_SETTINGS); + + // Refetch config when modal opens to ensure we have the latest values + useEffect(() => { + if (isVisible) { + refetch(); + } + }, [isVisible, refetch]); + + // Compute initial values from fetched config data + const initialValues = useMemo(() => { + if (!proxyConfigData) { + return { + store_model_in_db: false, + }; + } + + const storeModelField = proxyConfigData.find(field => field.field_name === 'store_model_in_db'); + + return { + store_model_in_db: storeModelField?.field_value ?? false, + }; + }, [proxyConfigData]); + + const handleFormSubmit = async (formValues: StoreModelInDBParams) => { + try { + await mutateAsync(formValues, { + onSuccess: () => { + NotificationsManager.success("Model storage settings updated successfully"); + refetch(); + onSuccess?.(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save model storage settings: " + parseErrorMessage(error)); + }, + }); + } catch (error) { + NotificationsManager.fromBackend("Failed to save model storage settings: " + parseErrorMessage(error)); + } + }; + + const handleCancel = () => { + form.resetFields(); + onCancel(); + }; + + return ( + Model Settings} + open={isVisible} + footer={ + + + + + } + onCancel={handleCancel} + > +
+ f.field_name === 'store_model_in_db')?.field_description || + "If enabled, models and config are stored in and loaded from the database." + } + valuePropName="checked" + > + {isLoadingConfig ? : } + +
+ + ); +}; + +export default ModelSettingsModal; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 132ac020a3d..b0596c025a5 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2428,7 +2428,7 @@ export const tagsSpendLogsCall = async ( // if tags, convert the list to a comma separated string if (tags) { - url += `${url}&tags=${tags.join(",")}`; + url += `&tags=${tags.join(",")}`; } console.log("in tagsSpendLogsCall:", url); @@ -5555,19 +5555,24 @@ export const getPolicyTemplates = async (accessToken: string) => { export const enrichPolicyTemplate = async ( accessToken: string, templateId: string, - parameters: Record + parameters: Record, + model?: string, + competitors?: string[] ) => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/policy/templates/enrich` : `/policy/templates/enrich`; + const body: any = { template_id: templateId, parameters }; + if (model) body.model = model; + if (competitors) body.competitors = competitors; const response = await fetch(url, { method: "POST", headers: { [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, - body: JSON.stringify({ template_id: templateId, parameters }), + body: JSON.stringify(body), }); if (!response.ok) { @@ -5620,6 +5625,78 @@ export const suggestPolicyTemplates = async ( } }; +export const enrichPolicyTemplateStream = async ( + accessToken: string, + templateId: string, + parameters: Record, + model: string, + onCompetitor: (name: string) => void, + onDone: (result: { + competitors: string[]; + competitor_variations: Record; + guardrailDefinitions: any[]; + }) => void, + onError?: (error: string) => void, + options?: { instruction?: string; existingCompetitors?: string[] }, + onStatus?: (message: string) => void +) => { + const url = proxyBaseUrl + ? `${proxyBaseUrl}/policy/templates/enrich/stream` + : `/policy/templates/enrich/stream`; + const body: any = { template_id: templateId, parameters, model }; + if (options?.instruction) body.instruction = options.instruction; + if (options?.existingCompetitors) body.competitors = options.existingCompetitors; + + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(body), + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + const reader = response.body?.getReader(); + if (!reader) throw new Error("No response body"); + + const decoder = new TextDecoder(); + let buffer = ""; + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + buffer += decoder.decode(value, { stream: true }); + const lines = buffer.split("\n"); + buffer = lines.pop() || ""; + + for (const line of lines) { + if (!line.startsWith("data: ")) continue; + try { + const event = JSON.parse(line.slice(6)); + if (event.type === "competitor") { + onCompetitor(event.name); + } else if (event.type === "status") { + onStatus?.(event.message); + } else if (event.type === "done") { + onDone(event); + } else if (event.type === "error") { + onError?.(event.message); + } + } catch { + // skip malformed events + } + } + } +}; + export const createPolicyCall = async (accessToken: string, policyData: any) => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/policies` : `/policies`; diff --git a/ui/litellm-dashboard/src/components/pass_through_info.tsx b/ui/litellm-dashboard/src/components/pass_through_info.tsx index fe3b204b507..9a8d2a3a07f 100644 --- a/ui/litellm-dashboard/src/components/pass_through_info.tsx +++ b/ui/litellm-dashboard/src/components/pass_through_info.tsx @@ -13,7 +13,7 @@ import { TabPanels, TextInput, } from "@tremor/react"; -import { Button, Form, Input, Switch, InputNumber } from "antd"; +import { Button, Form, Input, Switch, InputNumber, Select } from "antd"; import { updatePassThroughEndpoint, deletePassThroughEndpointsCall } from "./networking"; import { Eye, EyeOff } from "lucide-react"; import RoutePreview from "./route_preview"; @@ -21,6 +21,9 @@ import NotificationsManager from "./molecules/notifications_manager"; import PassThroughSecuritySection from "./common_components/PassThroughSecuritySection"; import PassThroughGuardrailsSection from "./common_components/PassThroughGuardrailsSection"; +const HTTP_METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH"]; +const { Option } = Select; + export interface PassThroughInfoProps { endpointData: PassThroughEndpoint; onClose: () => void; @@ -38,6 +41,7 @@ interface PassThroughEndpoint { include_subpath?: boolean; cost_per_request?: number; auth?: boolean; + methods?: string[]; guardrails?: Record; } @@ -70,6 +74,7 @@ const PassThroughInfoView: React.FC = ({ const [loading, setLoading] = useState(false); const [isEditing, setIsEditing] = useState(false); const [authEnabled, setAuthEnabled] = useState(initialEndpointData?.auth || false); + const [selectedMethods, setSelectedMethods] = useState(initialEndpointData?.methods || []); const [guardrails, setGuardrails] = useState>( initialEndpointData?.guardrails || {} ); @@ -97,6 +102,7 @@ const PassThroughInfoView: React.FC = ({ include_subpath: values.include_subpath, cost_per_request: values.cost_per_request, auth: premiumUser ? values.auth : undefined, + methods: selectedMethods && selectedMethods.length > 0 ? selectedMethods : undefined, guardrails: guardrails && Object.keys(guardrails).length > 0 ? guardrails : undefined, }; @@ -191,6 +197,23 @@ const PassThroughInfoView: React.FC = ({ {endpointData.auth ? "Auth Required" : "No Auth"}
+ {endpointData.methods && endpointData.methods.length > 0 && ( +
+ HTTP Methods: +
+ {endpointData.methods.map((method) => ( + + {method} + + ))} +
+
+ )} + {(!endpointData.methods || endpointData.methods.length === 0) && ( +
+ All HTTP methods supported +
+ )} {endpointData.cost_per_request !== undefined && (
Cost per request: ${endpointData.cost_per_request} @@ -277,6 +300,7 @@ const PassThroughInfoView: React.FC = ({ include_subpath: endpointData.include_subpath || false, cost_per_request: endpointData.cost_per_request, auth: endpointData.auth || false, + methods: endpointData.methods || [], }} layout="vertical" > @@ -295,6 +319,31 @@ const PassThroughInfoView: React.FC = ({ /> + + + + diff --git a/ui/litellm-dashboard/src/components/pass_through_settings.tsx b/ui/litellm-dashboard/src/components/pass_through_settings.tsx index b37d3df7e41..d60caa54667 100644 --- a/ui/litellm-dashboard/src/components/pass_through_settings.tsx +++ b/ui/litellm-dashboard/src/components/pass_through_settings.tsx @@ -39,7 +39,9 @@ export interface passThroughItem { include_subpath?: boolean; cost_per_request?: number; auth?: boolean; + methods?: string[]; guardrails?: Record; + default_query_params?: Record; } // Password field component for headers @@ -147,6 +149,32 @@ const PassThroughSettings: React.FC = ({ accessToken, accessorKey: "target", cell: (info: any) => {info.getValue()}, }, + { + header: () => ( +
+ Methods + + + +
+ ), + accessorKey: "methods", + cell: (info: any) => { + const methods = info.getValue(); + if (!methods || methods.length === 0) { + return ALL; + } + return ( +
+ {methods.map((method: string) => ( + + {method} + + ))} +
+ ); + }, + }, { header: () => (
diff --git a/ui/litellm-dashboard/src/components/policies/guardrail_selection_modal.tsx b/ui/litellm-dashboard/src/components/policies/guardrail_selection_modal.tsx index 2984478ba5a..adf1b14a4a6 100644 --- a/ui/litellm-dashboard/src/components/policies/guardrail_selection_modal.tsx +++ b/ui/litellm-dashboard/src/components/policies/guardrail_selection_modal.tsx @@ -222,6 +222,31 @@ const GuardrailSelectionModal: React.FC = ({
)} + {/* Discovered Competitors */} + {template?.discoveredCompetitors?.length > 0 && ( + <> + +
+
+ ✨ + + AI-Discovered Competitors ({template.discoveredCompetitors.length}) + +
+
+ {template.discoveredCompetitors.map((name: string) => ( + + {name} + + ))} +
+

+ These competitor names will be automatically blocked by the competitor-name-blocker guardrail. +

+
+ + )} + {/* Selected Summary */} diff --git a/ui/litellm-dashboard/src/components/policies/index.tsx b/ui/litellm-dashboard/src/components/policies/index.tsx index 185c17db157..db2da493bd7 100644 --- a/ui/litellm-dashboard/src/components/policies/index.tsx +++ b/ui/litellm-dashboard/src/components/policies/index.tsx @@ -231,7 +231,10 @@ const PoliciesPanel: React.FC = ({ return JSON.parse(templateStr); }; - const handleParameterConfirm = async (parameters: Record) => { + const handleParameterConfirm = async ( + parameters: Record, + enrichmentOptions?: { model?: string; competitors?: string[] } + ) => { if (!accessToken || !pendingTemplate) return; setIsEnrichingTemplate(true); @@ -240,14 +243,20 @@ const PoliciesPanel: React.FC = ({ let enrichedTemplate = pendingTemplate; if (pendingTemplate.llm_enrichment) { - // Call backend to enrich template with LLM-discovered data + // Call backend to enrich template with LLM-discovered data (or user-provided competitors) const enrichResult = await enrichPolicyTemplate( accessToken, pendingTemplate.id, - parameters + parameters, + enrichmentOptions?.model, + enrichmentOptions?.competitors ); - // The backend returns the enriched guardrailDefinitions - enrichedTemplate = { ...pendingTemplate, guardrailDefinitions: enrichResult.guardrailDefinitions }; + // The backend returns the enriched guardrailDefinitions + discovered competitors + enrichedTemplate = { + ...pendingTemplate, + guardrailDefinitions: enrichResult.guardrailDefinitions, + discoveredCompetitors: enrichResult.competitors || [], + }; } // Substitute parameters in template @@ -499,6 +508,7 @@ const PoliciesPanel: React.FC = ({ onConfirm={handleParameterConfirm} onCancel={handleParameterCancel} isLoading={isEnrichingTemplate} + accessToken={accessToken || ""} /> diff --git a/ui/litellm-dashboard/src/components/policies/template_parameter_modal.tsx b/ui/litellm-dashboard/src/components/policies/template_parameter_modal.tsx index ce1e5bc22e9..1bbb73db9ef 100644 --- a/ui/litellm-dashboard/src/components/policies/template_parameter_modal.tsx +++ b/ui/litellm-dashboard/src/components/policies/template_parameter_modal.tsx @@ -1,6 +1,7 @@ import React, { useState, useEffect } from "react"; -import { Modal, Spin } from "antd"; +import { Modal, Spin, Radio, Select } from "antd"; import { Button, TextInput } from "@tremor/react"; +import { modelHubCall, enrichPolicyTemplateStream } from "../networking"; interface TemplateParameter { name: string; @@ -13,9 +14,13 @@ interface TemplateParameter { interface TemplateParameterModalProps { visible: boolean; template: any; - onConfirm: (parameters: Record) => void; + onConfirm: ( + parameters: Record, + enrichmentOptions?: { model?: string; competitors?: string[] } + ) => void; onCancel: () => void; isLoading?: boolean; + accessToken: string; } const TemplateParameterModal: React.FC = ({ @@ -24,10 +29,28 @@ const TemplateParameterModal: React.FC = ({ onConfirm, onCancel, isLoading = false, + accessToken, }) => { const [parameterValues, setParameterValues] = useState>({}); + const [competitorMode, setCompetitorMode] = useState<"ai" | "manual">("ai"); + const [selectedModel, setSelectedModel] = useState(undefined); + const [availableModels, setAvailableModels] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + const [competitorTags, setCompetitorTags] = useState([]); + const [variationsMap, setVariationsMap] = useState>({}); + const [isGenerating, setIsGenerating] = useState(false); + const [refinementInput, setRefinementInput] = useState(""); + const [isRefining, setIsRefining] = useState(false); + const [hasGenerated, setHasGenerated] = useState(false); + const [statusMessage, setStatusMessage] = useState(""); const parameters: TemplateParameter[] = template?.parameters || []; + const hasEnrichment = !!template?.llm_enrichment; + const enrichmentParam = hasEnrichment ? template.llm_enrichment.parameter : null; + + const nonEnrichmentParams = hasEnrichment + ? parameters.filter((p) => p.name !== enrichmentParam) + : parameters; useEffect(() => { if (visible && template) { @@ -36,15 +59,136 @@ const TemplateParameterModal: React.FC = ({ initial[p.name] = ""; }); setParameterValues(initial); + setCompetitorMode("ai"); + setSelectedModel(undefined); + setCompetitorTags([]); + setVariationsMap({}); + setIsGenerating(false); + setRefinementInput(""); + setIsRefining(false); + setHasGenerated(false); + setStatusMessage(""); } }, [visible, template]); - const allRequiredFilled = parameters + useEffect(() => { + if (visible && hasEnrichment && competitorMode === "ai" && availableModels.length === 0) { + loadModels(); + } + }, [visible, hasEnrichment, competitorMode]); + + const loadModels = async () => { + if (!accessToken) return; + setIsLoadingModels(true); + try { + const fetchedModels = await modelHubCall(accessToken); + if (fetchedModels?.data?.length > 0) { + const models = fetchedModels.data + .map((item: any) => item.model_group as string) + .sort(); + setAvailableModels(models); + } + } catch (error) { + console.error("Error fetching models:", error); + } finally { + setIsLoadingModels(false); + } + }; + + const handleGenerateNames = async () => { + if (!accessToken || !selectedModel || !template) return; + const brandName = (parameterValues[enrichmentParam || "brand_name"] || "").trim(); + if (!brandName) return; + + setIsGenerating(true); + setCompetitorTags([]); + setVariationsMap({}); + setStatusMessage(""); + try { + await enrichPolicyTemplateStream( + accessToken, + template.id, + parameterValues, + selectedModel, + (name) => { + setCompetitorTags((prev) => [...prev, name]); + }, + (result) => { + setCompetitorTags(result.competitors); + setVariationsMap(result.competitor_variations || {}); + setIsGenerating(false); + setHasGenerated(true); + setStatusMessage(""); + }, + (error) => { + console.error("Streaming error:", error); + setIsGenerating(false); + setStatusMessage(""); + }, + undefined, + (status) => setStatusMessage(status), + ); + } catch (error) { + console.error("Error generating competitor names:", error); + setIsGenerating(false); + } + }; + + const handleRefine = async () => { + if (!accessToken || !selectedModel || !template || !refinementInput.trim()) return; + + setIsRefining(true); + setStatusMessage(""); + try { + await enrichPolicyTemplateStream( + accessToken, + template.id, + parameterValues, + selectedModel, + (name) => { + setCompetitorTags((prev) => { + if (prev.some((t) => t.toLowerCase() === name.toLowerCase())) return prev; + return [...prev, name]; + }); + }, + (result) => { + setCompetitorTags(result.competitors); + setVariationsMap(result.competitor_variations || {}); + setIsRefining(false); + setRefinementInput(""); + setStatusMessage(""); + }, + (error) => { + console.error("Refinement error:", error); + setIsRefining(false); + setStatusMessage(""); + }, + { + instruction: refinementInput.trim(), + existingCompetitors: competitorTags, + }, + (status) => setStatusMessage(status), + ); + } catch (error) { + console.error("Error refining competitor names:", error); + setIsRefining(false); + } + }; + + const allNonEnrichmentFilled = nonEnrichmentParams .filter((p) => p.required) .every((p) => (parameterValues[p.name] || "").trim().length > 0); + const brandNameFilled = enrichmentParam + ? (parameterValues[enrichmentParam] || "").trim().length > 0 + : true; + + const canContinue = hasEnrichment + ? allNonEnrichmentFilled && brandNameFilled && competitorTags.length > 0 + : allNonEnrichmentFilled && brandNameFilled; + const handleConfirm = () => { - onConfirm(parameterValues); + onConfirm(parameterValues, { competitors: competitorTags }); }; return ( @@ -53,15 +197,13 @@ const TemplateParameterModal: React.FC = ({

{template?.title}

- {template?.llm_enrichment - ? "Enter your brand name to auto-discover competitors and configure guardrails" - : "Configure template parameters"} + Configure competitor blocking for your brand

} open={visible} onCancel={onCancel} - width={500} + width={700} footer={[ , ]} >
- {parameters.map((param) => ( + {nonEnrichmentParams.map((param) => (
))} - {template?.llm_enrichment && ( -
-

- This template uses AI to automatically discover your competitors and configure - guardrails. An onboarded LLM will be called to identify competitor names. -

-
+ {hasEnrichment && ( + <> +
+ + setCompetitorMode(e.target.value)} + className="w-full" + > +
+ + ✨ Use AI + + + Enter Manually + +
+ +
+ + {/* Brand Name */} +
+ + + setParameterValues((prev) => ({ + ...prev, + [enrichmentParam || "brand_name"]: e.target.value, + })) + } + /> +
+ + {competitorMode === "ai" && ( + <> +
+ + setCompetitorTags(values)} + tokenSeparators={[","]} + open={false} + suffixIcon={null} + /> +

+ Type a name and press Enter to add. Click ✕ to remove. +

+ {statusMessage && ( +
+ + {statusMessage} +
+ )} + {Object.keys(variationsMap).length > 0 && !statusMessage && ( +

+ ✓ {Object.values(variationsMap).flat().length} alternate spellings & variations auto-generated for guardrail matching +

+ )} +
+ + {/* Refinement input — shown after initial generation in AI mode */} + {competitorMode === "ai" && hasGenerated && competitorTags.length > 0 && ( +
+ +
+ setRefinementInput(e.target.value)} + onKeyDown={(e) => { + if (e.key === "Enter" && refinementInput.trim() && !isRefining) { + handleRefine(); + } + }} + disabled={isRefining} + /> + +
+

+ Give instructions to add, remove, or change competitors. Press Enter to send. +

+
+ )} + )} - {isLoading && ( -
- - - {template?.llm_enrichment - ? "Using AI to discover competitors..." - : "Processing template..."} - -
- )} + {!hasEnrichment && + parameters.map((param) => ( +
+ + + setParameterValues((prev) => ({ + ...prev, + [param.name]: e.target.value, + })) + } + /> +
+ ))}
); diff --git a/ui/litellm-dashboard/src/components/query_param_input.tsx b/ui/litellm-dashboard/src/components/query_param_input.tsx new file mode 100644 index 00000000000..4a98ab48bca --- /dev/null +++ b/ui/litellm-dashboard/src/components/query_param_input.tsx @@ -0,0 +1,57 @@ +import React, { useState } from "react"; +import { Button, Space } from "antd"; +import { MinusCircleOutlined, PlusOutlined } from "@ant-design/icons"; +import { TextInput } from "@tremor/react"; + +interface QueryParamInputProps { + value?: Record; + onChange?: (value: Record) => void; +} + +const QueryParamInput: React.FC = ({ value = {}, onChange }) => { + const [pairs, setPairs] = useState<[string, string][]>(Object.entries(value)); + + const handleAdd = () => { + setPairs([...pairs, ["", ""]]); + }; + + const handleRemove = (index: number) => { + const newPairs = pairs.filter((_, i) => i !== index); + setPairs(newPairs); + onChange?.(Object.fromEntries(newPairs)); + }; + + const handleChange = (index: number, key: string, val: string) => { + const newPairs = [...pairs]; + newPairs[index] = [key, val]; + setPairs(newPairs); + onChange?.(Object.fromEntries(newPairs)); + }; + + return ( +
+ {pairs.map(([key, val], index) => ( + + handleChange(index, e.target.value, val)} + /> + handleChange(index, key, e.target.value)} + /> +
+ handleRemove(index)} style={{ cursor: "pointer" }} /> +
+ + ))} + +
+ ); +}; + +export default QueryParamInput; \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/team/member_permissions.test.tsx b/ui/litellm-dashboard/src/components/team/member_permissions.test.tsx index cee2b8d587a..10c78331c68 100644 --- a/ui/litellm-dashboard/src/components/team/member_permissions.test.tsx +++ b/ui/litellm-dashboard/src/components/team/member_permissions.test.tsx @@ -97,6 +97,20 @@ describe("MemberPermissions", () => { } }); + it("should render team daily activity permission with correct method and description", async () => { + vi.mocked(networking.getTeamPermissionsCall).mockResolvedValue({ + all_available_permissions: ["/key/generate", "/team/daily/activity"], + team_member_permissions: [], + }); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("/team/daily/activity")).toBeInTheDocument(); + expect(screen.getByText("Member can view all team usage data (not just their own)")).toBeInTheDocument(); + }); + }); + it("should not show save button when canEditTeam is false", async () => { vi.mocked(networking.getTeamPermissionsCall).mockResolvedValue({ all_available_permissions: ["/key/generate", "/key/list"], diff --git a/ui/litellm-dashboard/src/components/team/permission_definitions.test.tsx b/ui/litellm-dashboard/src/components/team/permission_definitions.test.tsx index 2d8aeabf383..a85ed8353d2 100644 --- a/ui/litellm-dashboard/src/components/team/permission_definitions.test.tsx +++ b/ui/litellm-dashboard/src/components/team/permission_definitions.test.tsx @@ -11,6 +11,10 @@ describe("permission_definitions", () => { expect(getMethodForEndpoint("/key/list")).toBe("GET"); }); + it("should return GET for activity endpoints", () => { + expect(getMethodForEndpoint("/team/daily/activity")).toBe("GET"); + }); + it("should return POST for other endpoints", () => { expect(getMethodForEndpoint("/key/generate")).toBe("POST"); expect(getMethodForEndpoint("/key/update")).toBe("POST"); @@ -48,6 +52,14 @@ describe("permission_definitions", () => { expect(result.description).toBe(PERMISSION_DESCRIPTIONS["/key/service-account/generate"]); }); + it("should return correct info for team daily activity permission", () => { + const result = getPermissionInfo("/team/daily/activity"); + expect(result.method).toBe("GET"); + expect(result.endpoint).toBe("/team/daily/activity"); + expect(result.description).toBe(PERMISSION_DESCRIPTIONS["/team/daily/activity"]); + expect(result.route).toBe("/team/daily/activity"); + }); + it("should return fallback description for unknown permission", () => { const result = getPermissionInfo("/unknown/endpoint"); expect(result.method).toBe("POST"); @@ -56,4 +68,11 @@ describe("permission_definitions", () => { expect(result.route).toBe("/unknown/endpoint"); }); }); + + describe("PERMISSION_DESCRIPTIONS", () => { + it("should include team daily activity permission", () => { + expect(PERMISSION_DESCRIPTIONS["/team/daily/activity"]).toBeDefined(); + expect(PERMISSION_DESCRIPTIONS["/team/daily/activity"]).toContain("team usage"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/permission_definitions.tsx b/ui/litellm-dashboard/src/components/team/permission_definitions.tsx index 6e7bdaa739c..1c1e795d219 100644 --- a/ui/litellm-dashboard/src/components/team/permission_definitions.tsx +++ b/ui/litellm-dashboard/src/components/team/permission_definitions.tsx @@ -20,13 +20,15 @@ export const PERMISSION_DESCRIPTIONS: Record = { "/key/list": "Member can list virtual keys belonging to this team", "/key/block": "Member can block a virtual key belonging to this team", "/key/unblock": "Member can unblock a virtual key belonging to this team", + "/team/daily/activity": + "Member can view all team usage data (not just their own)", }; /** * Determines the HTTP method for a given permission endpoint */ export const getMethodForEndpoint = (endpoint: string): string => { - if (endpoint.includes("/info") || endpoint.includes("/list")) { + if (endpoint.includes("/info") || endpoint.includes("/list") || endpoint.includes("/activity")) { return "GET"; } return "POST";