Central service (LLM Gateway) to access multiple LLMs
+
Use LiteLLM directly in your Python code
+
+
+
Who Uses It?
+
Gen AI Enablement / ML Platform Teams
+
Developers building LLM projects
+
+
+
Key Features
+
• Centralized API gateway with authentication & authorization • Multi-tenant cost tracking and spend management per project/user • Per-project customization (logging, guardrails, caching) • Virtual keys for secure access control • Admin dashboard UI for monitoring and management
+
• Direct Python library integration in your codebase • Router with retry/fallback logic across multiple deployments (e.g. Azure/OpenAI) - Router • Application-level load balancing and cost tracking • Exception handling with OpenAI-compatible errors • Observability callbacks (Lunary, MLflow, Langfuse, etc.)
+
+
+
## **LiteLLM Python SDK**
@@ -67,7 +66,7 @@ import os
os.environ["OPENAI_API_KEY"] = "your-api-key"
response = completion(
- model="gpt-3.5-turbo",
+ model="openai/gpt-5",
messages=[{ "content": "Hello, how are you?","role": "user"}]
)
```
@@ -83,13 +82,27 @@ import os
os.environ["ANTHROPIC_API_KEY"] = "your-api-key"
response = completion(
- model="claude-2",
+ model="anthropic/claude-sonnet-4-5-20250929",
messages=[{ "content": "Hello, how are you?","role": "user"}]
)
```
+
+```python
+from litellm import completion
+import os
+
+## set ENV variables
+os.environ["XAI_API_KEY"] = "your-api-key"
+
+response = completion(
+ model="xai/grok-2-latest",
+ messages=[{ "content": "Hello, how are you?","role": "user"}]
+)
+```
+
```python
@@ -97,11 +110,11 @@ from litellm import completion
import os
# auth: run 'gcloud auth application-default'
-os.environ["VERTEX_PROJECT"] = "hardy-device-386718"
-os.environ["VERTEX_LOCATION"] = "us-central1"
+os.environ["VERTEXAI_PROJECT"] = "hardy-device-386718"
+os.environ["VERTEXAI_LOCATION"] = "us-central1"
response = completion(
- model="chat-bison",
+ model="vertex_ai/gemini-1.5-pro",
messages=[{ "content": "Hello, how are you?","role": "user"}]
)
```
@@ -212,8 +225,61 @@ response = completion(
+
+
+```python
+from litellm import completion
+import os
+
+## set ENV variables. Visit https://vercel.com/docs/ai-gateway#using-the-ai-gateway-with-an-api-key for instructions on obtaining a key
+os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-vercel-api-key"
+
+response = completion(
+ model="vercel_ai_gateway/openai/gpt-5",
+ messages=[{ "content": "Hello, how are you?","role": "user"}]
+)
+```
+
+
+
+### Response Format (OpenAI Chat Completions Format)
+
+```json
+{
+ "id": "chatcmpl-565d891b-a42e-4c39-8d14-82a1f5208885",
+ "created": 1734366691,
+ "model": "gpt-5",
+ "object": "chat.completion",
+ "system_fingerprint": null,
+ "choices": [
+ {
+ "finish_reason": "stop",
+ "index": 0,
+ "message": {
+ "content": "Hello! As an AI language model, I don't have feelings, but I'm operating properly and ready to assist you with any questions or tasks you may have. How can I help you today?",
+ "role": "assistant",
+ "tool_calls": null,
+ "function_call": null
+ }
+ }
+ ],
+ "usage": {
+ "completion_tokens": 43,
+ "prompt_tokens": 13,
+ "total_tokens": 56,
+ "completion_tokens_details": null,
+ "prompt_tokens_details": {
+ "audio_tokens": null,
+ "cached_tokens": 0
+ },
+ "cache_creation_input_tokens": 0,
+ "cache_read_input_tokens": 0
+ }
+}
+```
+
### Responses API
Use `litellm.responses()` for advanced models that support reasoning content like GPT-5, o3, etc.
@@ -265,11 +331,11 @@ from litellm import responses
import os
# auth: run 'gcloud auth application-default'
-os.environ["VERTEX_PROJECT"] = "jr-smith-386718"
-os.environ["VERTEX_LOCATION"] = "us-central1"
+os.environ["VERTEXAI_PROJECT"] = "jr-smith-386718"
+os.environ["VERTEXAI_LOCATION"] = "us-central1"
response = responses(
- model="chat-bison",
+ model="vertex_ai/gemini-1.5-pro",
messages=[{ "content": "What is the capital of France?","role": "user"}]
)
```
@@ -314,7 +380,7 @@ import os
os.environ["OPENAI_API_KEY"] = "your-api-key"
response = completion(
- model="gpt-3.5-turbo",
+ model="openai/gpt-5",
messages=[{ "content": "Hello, how are you?","role": "user"}],
stream=True,
)
@@ -331,14 +397,29 @@ import os
os.environ["ANTHROPIC_API_KEY"] = "your-api-key"
response = completion(
- model="claude-2",
+ model="anthropic/claude-sonnet-4-5-20250929",
messages=[{ "content": "Hello, how are you?","role": "user"}],
stream=True,
)
```
+
+```python
+from litellm import completion
+import os
+
+## set ENV variables
+os.environ["XAI_API_KEY"] = "your-api-key"
+
+response = completion(
+ model="xai/grok-2-latest",
+ messages=[{ "content": "Hello, how are you?","role": "user"}],
+ stream=True,
+)
+```
+
```python
@@ -346,11 +427,11 @@ from litellm import completion
import os
# auth: run 'gcloud auth application-default'
-os.environ["VERTEX_PROJECT"] = "hardy-device-386718"
-os.environ["VERTEX_LOCATION"] = "us-central1"
+os.environ["VERTEXAI_PROJECT"] = "hardy-device-386718"
+os.environ["VERTEXAI_LOCATION"] = "us-central1"
response = completion(
- model="chat-bison",
+ model="vertex_ai/gemini-1.5-pro",
messages=[{ "content": "Hello, how are you?","role": "user"}],
stream=True,
)
@@ -370,7 +451,7 @@ os.environ["NVIDIA_NIM_API_BASE"] = "nvidia_nim_endpoint_url"
response = completion(
model="nvidia_nim/",
- messages=[{ "content": "Hello, how are you?","role": "user"}]
+ messages=[{ "content": "Hello, how are you?","role": "user"}],
stream=True,
)
```
@@ -466,22 +547,74 @@ response = completion(
```
+
+
+
+```python
+from litellm import completion
+import os
+
+## set ENV variables. Visit https://vercel.com/docs/ai-gateway#using-the-ai-gateway-with-an-api-key for instructions on obtaining a key
+os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-vercel-api-key"
+
+response = completion(
+ model="vercel_ai_gateway/openai/gpt-5",
+ messages = [{ "content": "Hello, how are you?","role": "user"}],
+ stream=True,
+)
+```
+
+
+
+### Streaming Response Format (OpenAI Format)
+
+```json
+{
+ "id": "chatcmpl-2be06597-eb60-4c70-9ec5-8cd2ab1b4697",
+ "created": 1734366925,
+ "model": "claude-sonnet-4-5-20250929",
+ "object": "chat.completion.chunk",
+ "system_fingerprint": null,
+ "choices": [
+ {
+ "finish_reason": null,
+ "index": 0,
+ "delta": {
+ "content": "Hello",
+ "role": "assistant",
+ "function_call": null,
+ "tool_calls": null,
+ "audio": null
+ },
+ "logprobs": null
+ }
+ ]
+}
+```
+
### Exception handling
LiteLLM maps exceptions across all supported providers to the OpenAI exceptions. All our exceptions inherit from OpenAI's exception types, so any error-handling you have for that, should work out of the box with LiteLLM.
```python
-from openai.error import OpenAIError
+import litellm
from litellm import completion
+import os
os.environ["ANTHROPIC_API_KEY"] = "bad-key"
try:
- # some code
- completion(model="claude-instant-1", messages=[{"role": "user", "content": "Hey, how's it going?"}])
-except OpenAIError as e:
- print(e)
+ completion(model="anthropic/claude-instant-1", messages=[{"role": "user", "content": "Hey, how's it going?"}])
+except litellm.AuthenticationError as e:
+ # Thrown when the API key is invalid
+ print(f"Authentication failed: {e}")
+except litellm.RateLimitError as e:
+ # Thrown when you've exceeded your rate limit
+ print(f"Rate limited: {e}")
+except litellm.APIError as e:
+ # Thrown for general API errors
+ print(f"API error: {e}")
```
### Logging Observability - Log LLM Input/Output ([Docs](https://docs.litellm.ai/docs/observability/callbacks))
@@ -502,7 +635,7 @@ os.environ["OPENAI_API_KEY"]
litellm.success_callback = ["lunary", "mlflow", "langfuse", "helicone"] # log input/output to lunary, mlflow, langfuse, helicone
#openai call
-response = completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}])
+response = completion(model="openai/gpt-5", messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}])
```
### Track Costs, Usage, Latency for streaming
@@ -527,7 +660,7 @@ litellm.success_callback = [track_cost_callback] # set custom callback function
# litellm.completion() call
response = completion(
- model="gpt-3.5-turbo",
+ model="openai/gpt-5",
messages=[
{
"role": "user",
@@ -584,7 +717,7 @@ Example `litellm_config.yaml`
```yaml
model_list:
- - model_name: gpt-3.5-turbo
+ - model_name: gpt-5
litellm_params:
model: azure/
api_base: os.environ/AZURE_API_BASE # runs os.getenv("AZURE_API_BASE")
@@ -621,7 +754,7 @@ docker run \
import openai # openai v1.0.0+
client = openai.OpenAI(api_key="anything",base_url="http://0.0.0.0:4000") # set proxy to base_url
# request sent to model set on litellm proxy, `litellm --model`
-response = client.chat.completions.create(model="gpt-3.5-turbo", messages = [
+response = client.chat.completions.create(model="gpt-5", messages = [
{
"role": "user",
"content": "this is a test request, write a short poem"
diff --git a/docs/my-website/src/pages/troubleshoot.md b/docs/my-website/src/pages/troubleshoot.md
deleted file mode 100644
index 05dbf56caae..00000000000
--- a/docs/my-website/src/pages/troubleshoot.md
+++ /dev/null
@@ -1,11 +0,0 @@
-# Troubleshooting
-
-## Stable Version
-
-If you're running into problems with installation / Usage
-Use the stable version of litellm
-
-```
-pip install litellm==0.1.345
-```
-
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/LICENSE.md b/enterprise/LICENSE.md
index 5cd298ce658..c14a2a0c487 100644
--- a/enterprise/LICENSE.md
+++ b/enterprise/LICENSE.md
@@ -7,7 +7,7 @@ With regard to the BerriAI Software:
This software and associated documentation files (the "Software") may only be
used in production, if you (and any entity that you represent) have agreed to,
and are in compliance with, the BerriAI Subscription Terms of Service, available
-via [call](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat) or email (info@berri.ai) (the "Enterprise Terms"), or other
+via [call](https://calendly.com/d/cx9p-5yf-2nm/litellm-introductions) or email (info@berri.ai) (the "Enterprise Terms"), or other
agreement governing the use of the Software, as agreed by you and BerriAI,
and otherwise have a valid BerriAI Enterprise license for the
correct number of user seats. Subject to the foregoing sentence, you are free to
diff --git a/enterprise/README.md b/enterprise/README.md
index d5c27bab679..3b2ada6dd82 100644
--- a/enterprise/README.md
+++ b/enterprise/README.md
@@ -4,6 +4,6 @@ Code in this folder is licensed under a commercial license. Please review the [L
**These features are covered under the LiteLLM Enterprise contract**
-👉 **Using in an Enterprise / Need specific features ?** Meet with us [here](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat?month=2024-02)
+👉 **Using in an Enterprise / Need specific features ?** Meet with us [here](https://calendly.com/d/cx9p-5yf-2nm/litellm-introductions?month=2024-02)
See all Enterprise Features here 👉 [Docs](https://docs.litellm.ai/docs/proxy/enterprise)
diff --git a/enterprise/dist/litellm_enterprise-0.1.30-py3-none-any.whl b/enterprise/dist/litellm_enterprise-0.1.30-py3-none-any.whl
new file mode 100644
index 00000000000..0165bb096c0
Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.30-py3-none-any.whl differ
diff --git a/enterprise/dist/litellm_enterprise-0.1.30.tar.gz b/enterprise/dist/litellm_enterprise-0.1.30.tar.gz
new file mode 100644
index 00000000000..2bb7510e5d3
Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.30.tar.gz differ
diff --git a/enterprise/dist/litellm_enterprise-0.1.31-py3-none-any.whl b/enterprise/dist/litellm_enterprise-0.1.31-py3-none-any.whl
new file mode 100644
index 00000000000..03cadbd9023
Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.31-py3-none-any.whl differ
diff --git a/enterprise/dist/litellm_enterprise-0.1.31.tar.gz b/enterprise/dist/litellm_enterprise-0.1.31.tar.gz
new file mode 100644
index 00000000000..1ba1a717f62
Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.31.tar.gz differ
diff --git a/enterprise/dist/litellm_enterprise-0.1.32-py3-none-any.whl b/enterprise/dist/litellm_enterprise-0.1.32-py3-none-any.whl
new file mode 100644
index 00000000000..0c87c72c989
Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.32-py3-none-any.whl differ
diff --git a/enterprise/dist/litellm_enterprise-0.1.32.tar.gz b/enterprise/dist/litellm_enterprise-0.1.32.tar.gz
new file mode 100644
index 00000000000..4f0ac1a9b20
Binary files /dev/null and b/enterprise/dist/litellm_enterprise-0.1.32.tar.gz differ
diff --git a/enterprise/enterprise_hooks/__init__.py b/enterprise/enterprise_hooks/__init__.py
index 9eb1c8960a6..e93c8c9150a 100644
--- a/enterprise/enterprise_hooks/__init__.py
+++ b/enterprise/enterprise_hooks/__init__.py
@@ -1,11 +1,15 @@
from typing import Dict, Literal, Type, Union
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
+from litellm_enterprise.proxy.hooks.managed_vector_stores import (
+ _PROXY_LiteLLMManagedVectorStores,
+)
from litellm.integrations.custom_logger import CustomLogger
ENTERPRISE_PROXY_HOOKS: Dict[str, Type[CustomLogger]] = {
"managed_files": _PROXY_LiteLLMManagedFiles,
+ "managed_vector_stores": _PROXY_LiteLLMManagedVectorStores,
}
@@ -13,6 +17,7 @@ def get_enterprise_proxy_hook(
hook_name: Union[
Literal[
"managed_files",
+ "managed_vector_stores",
"max_parallel_requests",
],
str,
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/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
index 61e0745bab1..2f2e444850a 100644
--- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
+++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
@@ -16,6 +16,10 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import (
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
+from litellm.constants import (
+ EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
+ EMAIL_BUDGET_ALERT_TTL,
+)
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER
from litellm.integrations.email_templates.key_created_email import (
@@ -24,20 +28,23 @@ from litellm.integrations.email_templates.key_created_email import (
from litellm.integrations.email_templates.key_rotated_email import (
KEY_ROTATED_EMAIL_TEMPLATE,
)
-from litellm.integrations.email_templates.user_invitation_email import (
- USER_INVITATION_EMAIL_TEMPLATE,
-)
from litellm.integrations.email_templates.templates import (
MAX_BUDGET_ALERT_EMAIL_TEMPLATE,
SOFT_BUDGET_ALERT_EMAIL_TEMPLATE,
+ TEAM_SOFT_BUDGET_ALERT_EMAIL_TEMPLATE,
+)
+from litellm.integrations.email_templates.user_invitation_email import (
+ USER_INVITATION_EMAIL_TEMPLATE,
+)
+from litellm.proxy._types import (
+ CallInfo,
+ InvitationNew,
+ Litellm_EntityType,
+ UserAPIKeyAuth,
+ WebhookEvent,
)
-from litellm.proxy._types import CallInfo, InvitationNew, UserAPIKeyAuth, WebhookEvent
from litellm.secret_managers.main import get_secret_bool
from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL
-from litellm.constants import (
- EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
- EMAIL_BUDGET_ALERT_TTL,
-)
class BaseEmailLogger(CustomLogger):
@@ -114,10 +121,16 @@ class BaseEmailLogger(CustomLogger):
)
# Check if API key should be included in email
- include_api_key = get_secret_bool(secret_name="EMAIL_INCLUDE_API_KEY", default_value=True)
+ include_api_key = get_secret_bool(
+ secret_name="EMAIL_INCLUDE_API_KEY", default_value=True
+ )
if include_api_key is None:
include_api_key = True # Default to True if not set
- key_token_display = send_key_created_email_event.virtual_key if include_api_key else "[Key hidden for security - retrieve from dashboard]"
+ key_token_display = (
+ send_key_created_email_event.virtual_key
+ if include_api_key
+ else "[Key hidden for security - retrieve from dashboard]"
+ )
email_html_content = KEY_CREATED_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
@@ -155,10 +168,16 @@ class BaseEmailLogger(CustomLogger):
)
# Check if API key should be included in email
- include_api_key = get_secret_bool(secret_name="EMAIL_INCLUDE_API_KEY", default_value=True)
+ include_api_key = get_secret_bool(
+ secret_name="EMAIL_INCLUDE_API_KEY", default_value=True
+ )
if include_api_key is None:
include_api_key = True # Default to True if not set
- key_token_display = send_key_rotated_email_event.virtual_key if include_api_key else "[Key hidden for security - retrieve from dashboard]"
+ key_token_display = (
+ send_key_rotated_email_event.virtual_key
+ if include_api_key
+ else "[Key hidden for security - retrieve from dashboard]"
+ )
email_html_content = KEY_ROTATED_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
@@ -194,7 +213,9 @@ class BaseEmailLogger(CustomLogger):
)
# Format budget values
- soft_budget_str = f"${event.soft_budget}" if event.soft_budget is not None else "N/A"
+ soft_budget_str = (
+ f"${event.soft_budget}" if event.soft_budget is not None else "N/A"
+ )
spend_str = f"${event.spend}" if event.spend is not None else "$0.00"
max_budget_info = ""
if event.max_budget is not None:
@@ -217,6 +238,80 @@ class BaseEmailLogger(CustomLogger):
)
pass
+ async def send_team_soft_budget_alert_email(self, event: WebhookEvent):
+ """
+ Send email to team members when team soft budget is crossed
+ Supports multiple recipients via alert_emails field from team metadata
+ """
+ # Collect all recipient emails
+ recipient_emails: List[str] = []
+
+ # Add additional alert emails from team metadata.soft_budget_alert_emails
+ if hasattr(event, "alert_emails") and event.alert_emails:
+ for email in event.alert_emails:
+ if email and email not in recipient_emails: # Avoid duplicates
+ recipient_emails.append(email)
+
+ # If no recipients found, skip sending
+ if not recipient_emails:
+ verbose_proxy_logger.warning(
+ f"No recipient emails found for team soft budget alert. event={event.model_dump(exclude_none=True)}"
+ )
+ return
+
+ # Validate that we have at least one valid email address
+ first_recipient_email = recipient_emails[0]
+ if not first_recipient_email or not first_recipient_email.strip():
+ verbose_proxy_logger.warning(
+ f"Invalid recipient email found for team soft budget alert. event={event.model_dump(exclude_none=True)}"
+ )
+ return
+
+ verbose_proxy_logger.debug(
+ f"send_team_soft_budget_alert_email_event: {json.dumps(event.model_dump(exclude_none=True), indent=4, default=str)}"
+ )
+
+ # Get email params using the first recipient email (for template formatting)
+ # For team alerts with alert_emails, we don't need user_id lookup since we already have email addresses
+ # Pass user_id=None to prevent _get_email_params from trying to look up email from a potentially None user_id
+ email_params = await self._get_email_params(
+ email_event=EmailEvent.soft_budget_crossed,
+ user_id=None, # Team alerts don't require user_id when alert_emails are provided
+ user_email=first_recipient_email,
+ event_message=event.event_message,
+ )
+
+ # Format budget values
+ soft_budget_str = (
+ f"${event.soft_budget}" if event.soft_budget is not None else "N/A"
+ )
+ spend_str = f"${event.spend}" if event.spend is not None else "$0.00"
+ max_budget_info = ""
+ if event.max_budget is not None:
+ max_budget_info = f"Maximum Budget: ${event.max_budget} "
+
+ # Use team alias or generic greeting
+ team_alias = event.team_alias or "Team"
+
+ email_html_content = TEAM_SOFT_BUDGET_ALERT_EMAIL_TEMPLATE.format(
+ email_logo_url=email_params.logo_url,
+ team_alias=team_alias,
+ soft_budget=soft_budget_str,
+ spend=spend_str,
+ max_budget_info=max_budget_info,
+ base_url=email_params.base_url,
+ email_support_contact=email_params.support_contact,
+ )
+
+ # Send email to all recipients
+ await self.send_email(
+ from_email=self.DEFAULT_LITELLM_EMAIL,
+ to_email=recipient_emails,
+ subject=email_params.subject,
+ html_body=email_html_content,
+ )
+ pass
+
async def send_max_budget_alert_email(self, event: WebhookEvent):
"""
Send email to user when max budget alert threshold is reached
@@ -234,11 +329,17 @@ class BaseEmailLogger(CustomLogger):
# Format budget values
spend_str = f"${event.spend}" if event.spend is not None else "$0.00"
- max_budget_str = f"${event.max_budget}" if event.max_budget is not None else "N/A"
-
+ max_budget_str = (
+ f"${event.max_budget}" if event.max_budget is not None else "N/A"
+ )
+
# Calculate percentage and alert threshold
percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100)
- alert_threshold_str = f"${event.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE:.2f}" if event.max_budget is not None else "N/A"
+ alert_threshold_str = (
+ f"${event.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE:.2f}"
+ if event.max_budget is not None
+ else "N/A"
+ )
email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
@@ -285,17 +386,41 @@ class BaseEmailLogger(CustomLogger):
# - Don't re-alert, if alert already sent
_cache: DualCache = self.internal_usage_cache
- # percent of max_budget left to spend
- if user_info.max_budget is None and user_info.soft_budget is None:
- return
-
# For soft_budget alerts, check if we've already sent an alert
if type == "soft_budget":
- if user_info.soft_budget is not None and user_info.spend >= user_info.soft_budget:
+ # For team soft budget alerts, we only need team soft_budget to be set
+ # For other entity types, we need either max_budget or soft_budget
+ if user_info.event_group == Litellm_EntityType.TEAM:
+ if user_info.soft_budget is None:
+ return
+ # For team soft budget alerts, require alert_emails to be configured
+ # Team soft budget alerts are sent via metadata.soft_budget_alerting_emails
+ if user_info.alert_emails is None or len(user_info.alert_emails) == 0:
+ verbose_proxy_logger.debug(
+ "Skipping team soft budget email alert: no alert_emails configured",
+ )
+ return
+ else:
+ # For non-team alerts, require either max_budget or soft_budget
+ if user_info.max_budget is None and user_info.soft_budget is None:
+ return
+ if (
+ user_info.soft_budget is not None
+ and user_info.spend >= user_info.soft_budget
+ ):
# Generate cache key based on event type and identifier
- _id = user_info.token or user_info.user_id or "default_id"
+ # Use appropriate ID based on event_group to ensure unique cache keys per entity type
+ if user_info.event_group == Litellm_EntityType.TEAM:
+ _id = user_info.team_id or "default_id"
+ elif user_info.event_group == Litellm_EntityType.ORGANIZATION:
+ _id = user_info.organization_id or "default_id"
+ elif user_info.event_group == Litellm_EntityType.USER:
+ _id = user_info.user_id or "default_id"
+ else:
+ # For KEY and other types, use token or user_id
+ _id = user_info.token or user_info.user_id or "default_id"
_cache_key = f"email_budget_alerts:soft_budget_crossed:{_id}"
-
+
# Check if we've already sent this alert
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
@@ -318,11 +443,16 @@ class BaseEmailLogger(CustomLogger):
projected_exceeded_date=user_info.projected_exceeded_date,
projected_spend=user_info.projected_spend,
event_group=user_info.event_group,
+ alert_emails=user_info.alert_emails,
)
-
+
try:
- await self.send_soft_budget_alert_email(webhook_event)
-
+ # Use team-specific function for team alerts, otherwise use standard function
+ if user_info.event_group == Litellm_EntityType.TEAM:
+ await self.send_team_soft_budget_alert_email(webhook_event)
+ else:
+ await self.send_soft_budget_alert_email(webhook_event)
+
# Cache the alert to prevent duplicate sends
await _cache.async_set_cache(
key=_cache_key,
@@ -339,20 +469,27 @@ class BaseEmailLogger(CustomLogger):
# For max_budget_alert, check if we've already sent an alert
if type == "max_budget_alert":
if user_info.max_budget is not None and user_info.spend is not None:
- alert_threshold = user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
-
+ alert_threshold = (
+ user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
+ )
+
# Only alert if we've crossed the threshold but haven't exceeded max_budget yet
- if user_info.spend >= alert_threshold and user_info.spend < user_info.max_budget:
+ if (
+ user_info.spend >= alert_threshold
+ and user_info.spend < user_info.max_budget
+ ):
# Generate cache key based on event type and identifier
_id = user_info.token or user_info.user_id or "default_id"
_cache_key = f"email_budget_alerts:max_budget_alert:{_id}"
-
+
# Check if we've already sent this alert
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
# Calculate percentage
- percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100)
-
+ percentage = int(
+ EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100
+ )
+
# Create WebhookEvent for max budget alert
event_message = f"Max Budget Alert - {percentage}% of Maximum Budget Reached"
webhook_event = WebhookEvent(
@@ -373,10 +510,10 @@ class BaseEmailLogger(CustomLogger):
projected_spend=user_info.projected_spend,
event_group=user_info.event_group,
)
-
+
try:
await self.send_max_budget_alert_email(webhook_event)
-
+
# Cache the alert to prevent duplicate sends
await _cache.async_set_cache(
key=_cache_key,
@@ -420,9 +557,14 @@ class BaseEmailLogger(CustomLogger):
unused_custom_fields = []
# Function to safely get custom value or default
- def get_custom_or_default(custom_value: Optional[str], default_value: str, field_name: str) -> str:
- if custom_value is not None: # Only check premium if trying to use custom value
+ def get_custom_or_default(
+ custom_value: Optional[str], default_value: str, field_name: str
+ ) -> str:
+ if (
+ custom_value is not None
+ ): # Only check premium if trying to use custom value
from litellm.proxy.proxy_server import premium_user
+
if premium_user is not True:
unused_custom_fields.append(field_name)
return default_value
@@ -431,38 +573,48 @@ class BaseEmailLogger(CustomLogger):
# Get parameters, falling back to defaults if custom values aren't allowed
logo_url = get_custom_or_default(custom_logo, LITELLM_LOGO_URL, "logo URL")
- support_contact = get_custom_or_default(custom_support, self.DEFAULT_SUPPORT_EMAIL, "support contact")
- base_url = os.getenv("PROXY_BASE_URL", "http://0.0.0.0:4000") # Not a premium feature
- signature = get_custom_or_default(custom_signature, EMAIL_FOOTER, "email signature")
+ support_contact = get_custom_or_default(
+ custom_support, self.DEFAULT_SUPPORT_EMAIL, "support contact"
+ )
+ base_url = os.getenv(
+ "PROXY_BASE_URL", "http://0.0.0.0:4000"
+ ) # Not a premium feature
+ signature = get_custom_or_default(
+ custom_signature, EMAIL_FOOTER, "email signature"
+ )
# Get custom subject template based on email event type
if email_event == EmailEvent.new_user_invitation:
subject_template = get_custom_or_default(
custom_subject_invitation,
self.DEFAULT_SUBJECT_TEMPLATES[EmailEvent.new_user_invitation],
- "invitation subject template"
+ "invitation subject template",
)
elif email_event == EmailEvent.virtual_key_created:
subject_template = get_custom_or_default(
custom_subject_key_created,
self.DEFAULT_SUBJECT_TEMPLATES[EmailEvent.virtual_key_created],
- "key created subject template"
+ "key created subject template",
)
elif email_event == EmailEvent.virtual_key_rotated:
custom_subject_key_rotated = os.getenv("EMAIL_SUBJECT_KEY_ROTATED", None)
subject_template = get_custom_or_default(
custom_subject_key_rotated,
self.DEFAULT_SUBJECT_TEMPLATES[EmailEvent.virtual_key_rotated],
- "key rotated subject template"
+ "key rotated subject template",
)
else:
subject_template = "LiteLLM: {event_message}"
- subject = subject_template.format(event_message=event_message) if event_message else "LiteLLM Notification"
+ subject = (
+ subject_template.format(event_message=event_message)
+ if event_message
+ else "LiteLLM Notification"
+ )
- recipient_email: Optional[
- str
- ] = user_email or await self._lookup_user_email_from_db(user_id=user_id)
+ recipient_email: Optional[str] = (
+ user_email or await self._lookup_user_email_from_db(user_id=user_id)
+ )
if recipient_email is None:
raise ValueError(
f"User email not found for user_id: {user_id}. User email is required to send email."
@@ -480,11 +632,9 @@ class BaseEmailLogger(CustomLogger):
warning_msg = (
f"Email sent with default values instead of custom values for: {fields_str}. "
"This is an Enterprise feature. To use custom email fields, please upgrade to LiteLLM Enterprise. "
- "Schedule a meeting here: https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat"
- )
- verbose_proxy_logger.warning(
- f"{warning_msg}"
+ "Schedule a meeting here: https://calendly.com/d/cx9p-5yf-2nm/litellm-introductions"
)
+ verbose_proxy_logger.warning(f"{warning_msg}")
return EmailParams(
logo_url=logo_url,
@@ -531,44 +681,49 @@ class BaseEmailLogger(CustomLogger):
if not user_id:
verbose_proxy_logger.debug("No user_id provided for invitation link")
return base_url
-
+
if not await self._is_prisma_client_available():
return base_url
-
+
# Wait for any concurrent invitation creation to complete
await self._wait_for_invitation_creation()
-
+
# Get or create invitation
invitation = await self._get_or_create_invitation(user_id)
if not invitation:
- verbose_proxy_logger.warning(f"Failed to get/create invitation for user_id: {user_id}")
+ verbose_proxy_logger.warning(
+ f"Failed to get/create invitation for user_id: {user_id}"
+ )
return base_url
-
+
return self._construct_invitation_link(invitation.id, base_url)
async def _is_prisma_client_available(self) -> bool:
"""Check if Prisma client is available"""
from litellm.proxy.proxy_server import prisma_client
-
+
if prisma_client is None:
- verbose_proxy_logger.debug("Prisma client not found. Unable to lookup invitation")
+ verbose_proxy_logger.debug(
+ "Prisma client not found. Unable to lookup invitation"
+ )
return False
return True
async def _wait_for_invitation_creation(self) -> None:
"""
Wait for any concurrent invitation creation to complete.
-
+
The UI calls /invitation/new to generate the invitation link.
We wait to ensure any pending invitation creation is completed.
"""
import asyncio
+
await asyncio.sleep(10)
async def _get_or_create_invitation(self, user_id: str):
"""
Get existing invitation or create a new one for the user
-
+
Returns:
Invitation object with id attribute, or None if failed
"""
@@ -576,31 +731,41 @@ class BaseEmailLogger(CustomLogger):
create_invitation_for_user,
)
from litellm.proxy.proxy_server import prisma_client
-
+
if prisma_client is None:
- verbose_proxy_logger.error("Prisma client is None in _get_or_create_invitation")
+ verbose_proxy_logger.error(
+ "Prisma client is None in _get_or_create_invitation"
+ )
return None
-
+
try:
# Try to get existing invitation
- existing_invitations = await prisma_client.db.litellm_invitationlink.find_many(
- where={"user_id": user_id},
- order={"created_at": "desc"},
+ existing_invitations = (
+ await prisma_client.db.litellm_invitationlink.find_many(
+ where={"user_id": user_id},
+ order={"created_at": "desc"},
+ )
)
-
+
if existing_invitations and len(existing_invitations) > 0:
- verbose_proxy_logger.debug(f"Found existing invitation for user_id: {user_id}")
+ verbose_proxy_logger.debug(
+ f"Found existing invitation for user_id: {user_id}"
+ )
return existing_invitations[0]
-
+
# Create new invitation if none exists
- verbose_proxy_logger.debug(f"Creating new invitation for user_id: {user_id}")
+ verbose_proxy_logger.debug(
+ f"Creating new invitation for user_id: {user_id}"
+ )
return await create_invitation_for_user(
data=InvitationNew(user_id=user_id),
user_api_key_dict=UserAPIKeyAuth(user_id=user_id),
)
-
+
except Exception as e:
- verbose_proxy_logger.error(f"Error getting/creating invitation for user_id {user_id}: {e}")
+ verbose_proxy_logger.error(
+ f"Error getting/creating invitation for user_id {user_id}: {e}"
+ )
return None
def _construct_invitation_link(self, invitation_id: str, base_url: str) -> str:
diff --git a/enterprise/litellm_enterprise/integrations/custom_guardrail.py b/enterprise/litellm_enterprise/integrations/custom_guardrail.py
index b165d788f35..f07752d5c18 100644
--- a/enterprise/litellm_enterprise/integrations/custom_guardrail.py
+++ b/enterprise/litellm_enterprise/integrations/custom_guardrail.py
@@ -10,10 +10,15 @@ class EnterpriseCustomGuardrailHelper:
event_hook: Optional[
Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]
],
+ event_type: Optional[GuardrailEventHooks] = None,
) -> Optional[bool]:
"""
- Assumes check for event match is done in `should_run_guardrail`
- Returns True if the guardrail should be run by tag
+ Returns True if the guardrail should be run for this request and event_type.
+
+ Logic:
+ - If a request tag matches a Mode tag key, only run if event_type matches
+ the tag's value (the mode for that tag).
+ - If no request tag matches, fall back to default mode(s).
"""
from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
@@ -36,11 +41,31 @@ class EnterpriseCustomGuardrailHelper:
proxy_server_request=proxy_server_request,
)
- if request_tags and any(tag in event_hook.tags for tag in request_tags):
- return True
- elif event_hook.default and any(
- tag in event_hook.default for tag in request_tags
- ):
+ # Check if any request tag matches a Mode tag key
+ matched_mode = None
+ if request_tags:
+ for tag in request_tags:
+ if tag in event_hook.tags:
+ matched_mode = event_hook.tags[tag]
+ break
+
+ if matched_mode is not None:
+ # Tag matched: only run if event_type matches the tag's mode value(s)
+ if event_type is not None:
+ if isinstance(matched_mode, list):
+ return event_type.value in matched_mode
+ return event_type.value == matched_mode
return True
+ # No tag matched: fall back to default mode(s)
+ if event_hook.default is not None:
+ if event_type is not None:
+ default_list = (
+ event_hook.default
+ if isinstance(event_hook.default, list)
+ else [event_hook.default]
+ )
+ return event_type.value in default_list
+ return False
+
return False
diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
index d1b00420d31..18ac29b9781 100644
--- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
+++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
@@ -1,13 +1,13 @@
"""
AUDIT LOGGING
-All /audit logging endpoints. Attempting to write these as CRUD endpoints.
+All /audit logging endpoints. Attempting to write these as CRUD endpoints.
GET - /audit/{id} - Get audit log by id
GET - /audit - Get all audit logs
"""
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
#### AUDIT LOGGING ####
from fastapi import APIRouter, Depends, HTTPException, Query
@@ -22,6 +22,27 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
router = APIRouter()
+def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]:
+ """
+ Build an OR condition that matches a value inside a JSON column at the
+ given key, checking both before_value and updated_values.
+
+ Uses Prisma's JSON path filtering (PostgreSQL only).
+
+ Example result (team_id="t1"):
+ {"OR": [
+ {"before_value": {"path": ["team_id"], "string_contains": "t1"}},
+ {"updated_values": {"path": ["team_id"], "string_contains": "t1"}},
+ ]}
+ """
+ return {
+ "OR": [
+ {"before_value": {"path": [json_key], "string_contains": value}},
+ {"updated_values": {"path": [json_key], "string_contains": value}},
+ ]
+ }
+
+
@router.get(
"/audit",
tags=["Audit Logging"],
@@ -49,6 +70,14 @@ async def get_audit_logs(
),
start_date: Optional[str] = Query(None, description="Filter logs after this date"),
end_date: Optional[str] = Query(None, description="Filter logs before this date"),
+ object_team_id: Optional[str] = Query(
+ None,
+ description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)",
+ ),
+ object_key_hash: Optional[str] = Query(
+ None,
+ description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
+ ),
# Sorting parameters
sort_by: Optional[str] = Query(
None,
@@ -60,6 +89,9 @@ async def get_audit_logs(
Get all audit logs with filtering and pagination.
Returns a paginated response of audit logs matching the specified filters.
+
+ Note: object_team_id and object_key_hash use Prisma JSON path filtering,
+ which requires PostgreSQL.
"""
from litellm.proxy.proxy_server import prisma_client
@@ -82,18 +114,29 @@ async def get_audit_logs(
if object_id:
where_conditions["object_id"] = object_id
if start_date or end_date:
- date_filter = {}
+ date_filter: Dict[str, Any] = {}
if start_date:
date_filter["gte"] = start_date
if end_date:
date_filter["lte"] = end_date
where_conditions["updated_at"] = date_filter
+ # JSON field filters (PostgreSQL only) — each filter is AND'd with the
+ # others, but checks both before_value and updated_values internally (OR).
+ if object_team_id:
+ where_conditions["AND"] = where_conditions.get("AND", []) + [
+ _build_json_field_or_condition("team_id", object_team_id)
+ ]
+ if object_key_hash:
+ where_conditions["AND"] = where_conditions.get("AND", []) + [
+ _build_json_field_or_condition("token", object_key_hash)
+ ]
+
# Build sort conditions
- order_by = {}
+ order_by: Dict[str, Any] = {}
if sort_by and isinstance(sort_by, str):
order_by[sort_by] = sort_order
- elif sort_order and isinstance(sort_order, str):
+ else:
order_by["updated_at"] = sort_order # Default sort by updated_at
# Get paginated results
diff --git a/enterprise/litellm_enterprise/proxy/auth/route_checks.py b/enterprise/litellm_enterprise/proxy/auth/route_checks.py
index 6f7cf9143f4..fc57292a8d2 100644
--- a/enterprise/litellm_enterprise/proxy/auth/route_checks.py
+++ b/enterprise/litellm_enterprise/proxy/auth/route_checks.py
@@ -41,6 +41,10 @@ class EnterpriseRouteChecks:
return get_secret_bool("DISABLE_ADMIN_ENDPOINTS") is True
+ # Routes that should remain accessible even when LLM API endpoints are disabled.
+ # These are read-only model listing routes needed by the Admin UI.
+ LLM_API_EXEMPT_ROUTES = ["/models", "/v1/models"]
+
@staticmethod
def should_call_route(route: str):
"""
@@ -58,6 +62,7 @@ class EnterpriseRouteChecks:
)
elif (
RouteChecks.is_llm_api_route(route=route)
+ and route not in EnterpriseRouteChecks.LLM_API_EXEMPT_ROUTES
and EnterpriseRouteChecks.is_llm_api_route_disabled()
):
raise HTTPException(
diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index bb25e4f0626..cbe8d449b42 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -2,17 +2,24 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""
-from litellm._uuid import uuid
-from datetime import datetime
-from typing import TYPE_CHECKING, Optional, cast
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Optional
from litellm._logging import verbose_proxy_logger
+from litellm._uuid import uuid
+from litellm.constants import (
+ MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
+ MAX_OBJECTS_PER_POLL_CYCLE,
+)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
+CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
+
+
class CheckBatchCost:
def __init__(
self,
@@ -26,6 +33,69 @@ class CheckBatchCost:
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
+ # Cached after the first poll cycle. Once we know the column is absent we skip
+ # the guaranteed-failing primary query on every subsequent cycle.
+ self._has_batch_processed_column: bool = True
+
+ async def _get_user_info(self, batch_id, user_id) -> dict:
+ """
+ Look up user email and key alias by user_id for enriching the S3 callback metadata.
+ Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
+ """
+ try:
+ user_row = await self.prisma_client.db.litellm_usertable.find_unique(
+ where={"user_id": user_id}
+ )
+ if user_row is None:
+ return {}
+ return {
+ "user_api_key_user_email": getattr(user_row, "user_email", None),
+ "user_api_key_alias": getattr(user_row, "user_alias", None),
+ }
+ except Exception as e:
+ verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
+ return {}
+
+ async def _cleanup_stale_managed_objects(self) -> None:
+ """
+ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
+ in non-terminal states as 'stale_expired'. These will never complete and
+ should not be polled.
+ """
+ cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
+ result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
+ where={
+ "file_purpose": "batch",
+ "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
+ "created_at": {"lt": cutoff},
+ },
+ data={"status": "stale_expired"},
+ )
+ if result > 0:
+ verbose_proxy_logger.warning(
+ f"CheckBatchCost: marked {result} stale managed objects "
+ f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
+ )
+
+ async def _fallback_find_jobs(self) -> list:
+ """Query batch jobs without the batch_processed filter (for older schemas)."""
+ return await self.prisma_client.db.litellm_managedobjecttable.find_many(
+ where={
+ "file_purpose": "batch",
+ "status": {
+ "not_in": [
+ "failed",
+ "expired",
+ "cancelled",
+ "complete",
+ "completed",
+ "stale_expired",
+ ]
+ },
+ },
+ take=MAX_OBJECTS_PER_POLL_CYCLE,
+ order={"created_at": "asc"},
+ )
async def check_batch_cost(self):
"""
@@ -35,14 +105,11 @@ class CheckBatchCost:
- if not, return False
- if so, return True
"""
- from litellm_enterprise.proxy.hooks.managed_files import (
- _PROXY_LiteLLMManagedFiles,
- )
-
from litellm.batches.batch_utils import (
_get_file_content_as_dictionary,
calculate_batch_cost_and_usage,
)
+ from litellm.files.main import afile_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.proxy.openai_files_endpoints.common_utils import (
@@ -51,14 +118,50 @@ class CheckBatchCost:
get_model_id_from_unified_batch_id,
)
- jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
- where={
- "status": {"in": ["validating", "in_progress", "finalizing"]},
- "file_purpose": "batch",
- }
- )
- completed_jobs = []
+ try:
+ await self._cleanup_stale_managed_objects()
+ except Exception as cleanup_err:
+ verbose_proxy_logger.warning(
+ f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}"
+ )
+ # Look for all batches that have not yet been processed by CheckBatchCost.
+ # self._has_batch_processed_column is cached after the first probe so that
+ # older schemas don't pay a guaranteed-failing primary query + warning on
+ # every subsequent poll cycle.
+ if self._has_batch_processed_column:
+ try:
+ # Include "complete"/"completed" batches: the retrieve_batch
+ # endpoint may transition a batch to "complete" before
+ # CheckBatchCost runs. The batch_processed=False filter
+ # already prevents reprocessing finished batches.
+ jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
+ where={
+ "file_purpose": "batch",
+ "batch_processed": False,
+ "status": {
+ "not_in": [
+ "failed",
+ "expired",
+ "cancelled",
+ "stale_expired",
+ ]
+ },
+ },
+ take=MAX_OBJECTS_PER_POLL_CYCLE,
+ order={"created_at": "asc"},
+ )
+ except Exception as query_err:
+ if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower():
+ raise
+ # Permanent schema gap — cache the result so future cycles skip straight to fallback
+ self._has_batch_processed_column = False
+ verbose_proxy_logger.warning(
+ "CheckBatchCost: batch_processed column not found, querying without it"
+ )
+ jobs = await self._fallback_find_jobs()
+ else:
+ jobs = await self._fallback_find_jobs()
for job in jobs:
# get the model from the job
unified_object_id = job.unified_object_id
@@ -102,31 +205,56 @@ class CheckBatchCost:
continue
## RETRIEVE THE BATCH JOB OUTPUT FILE
- managed_files_obj = cast(
- Optional[_PROXY_LiteLLMManagedFiles],
- self.proxy_logging_obj.get_proxy_hook("managed_files"),
- )
if (
response.status == "completed"
and response.output_file_id is not None
- and managed_files_obj is not None
):
verbose_proxy_logger.info(
f"Batch ID: {batch_id} is complete, tracking cost and usage"
)
- # track cost
- model_file_id_mapping = {
- response.output_file_id: {model_id: response.output_file_id}
- }
- _file_content = await managed_files_obj.afile_content(
- file_id=response.output_file_id,
- litellm_parent_otel_span=None,
- llm_router=self.llm_router,
- model_file_id_mapping=model_file_id_mapping,
+
+ # aretrieve_batch is called with the raw provider batch ID, so response.id
+ # is the raw provider value (e.g. "batch_20260223-0518.234"). We need the
+ # unified base64 ID in the S3 log so downstream consumers can correlate it
+ # back to the batch they submitted via the proxy.
+ #
+ # CheckBatchCost builds its own LiteLLMLogging object (logging_obj below) and
+ # calls async_success_handler(result=response) directly. That handler calls
+ # _build_standard_logging_payload(response, ...) which reads response.id at
+ # that point — so setting response.id here is sufficient.
+ #
+ # The HTTP endpoint does this substitution via the managed files hook
+ # (async_post_call_success_hook). CheckBatchCost bypasses that hook entirely,
+ # so we do it explicitly here.
+ response.id = job.unified_object_id
+
+ # This background job runs as default_user_id, so going through the HTTP endpoint
+ # would trigger check_managed_file_id_access and get 403. Instead, extract the raw
+ # provider file ID and call afile_content directly with deployment credentials.
+ raw_output_file_id = response.output_file_id
+ decoded = _is_base64_encoded_unified_file_id(raw_output_file_id)
+ if decoded:
+ try:
+ raw_output_file_id = decoded.split("llm_output_file_id,")[1].split(";")[0]
+ except (IndexError, AttributeError):
+ pass
+
+ credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
+ _file_content = await afile_content(
+ file_id=raw_output_file_id,
+ **credentials,
)
+ # Access content - handle both direct attribute and method call
+ if hasattr(_file_content, 'content'):
+ content_bytes = _file_content.content # type: ignore[union-attr]
+ elif hasattr(_file_content, 'read'):
+ content_bytes = await _file_content.read() # type: ignore[misc]
+ else:
+ content_bytes = _file_content # type: ignore[assignment]
+
file_content_as_dict = _get_file_content_as_dictionary(
- _file_content.content
+ content_bytes # type: ignore[arg-type]
)
deployment_info = self.llm_router.get_deployment(model_id=model_id)
@@ -143,11 +271,15 @@ class CheckBatchCost:
custom_llm_provider=custom_llm_provider,
)
+ # Pass deployment model_info so custom batch pricing
+ # (input_cost_per_token_batches etc.) is used for cost calc
+ deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {}
batch_cost, batch_usage, batch_models = (
await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=llm_provider, # type: ignore
model_name=model_name,
+ model_info=deployment_model_info, # type: ignore[arg-type]
)
)
logging_obj = LiteLLMLogging(
@@ -160,11 +292,21 @@ class CheckBatchCost:
function_id=str(uuid.uuid4()),
)
+ creator_user_id = job.created_by
+ user_info = await self._get_user_info(batch_id, job.created_by)
+
logging_obj.update_environment_variables(
litellm_params={
+ # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
+ "proxy_server_request": {
+ "headers": {
+ "user-agent": CHECK_BATCH_COST_USER_AGENT,
+ }
+ },
"metadata": {
- "user_api_key_user_id": job.created_by or "default-user-id",
- }
+ "user_api_key_user_id": creator_user_id,
+ **user_info,
+ },
},
optional_params={},
)
@@ -177,11 +319,18 @@ class CheckBatchCost:
)
# mark the job as complete
- completed_jobs.append(job)
-
- if len(completed_jobs) > 0:
- # mark the jobs as complete
- await self.prisma_client.db.litellm_managedobjecttable.update_many(
- where={"id": {"in": [job.id for job in completed_jobs]}},
- data={"status": "complete"},
- )
+ try:
+ update_data: dict = {
+ "status": "complete",
+ "file_object": response.model_dump_json(),
+ }
+ if self._has_batch_processed_column:
+ update_data["batch_processed"] = True
+ await self.prisma_client.db.litellm_managedobjecttable.update(
+ where={"id": job.id},
+ data=update_data,
+ )
+ except Exception as db_err:
+ verbose_proxy_logger.error(
+ f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
+ )
diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
index 4ee6a89cc98..54fbc7abcc5 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
@@ -3,10 +3,15 @@ Polls LiteLLM_ManagedObjectTable to check if the response is complete.
Cost tracking is handled automatically by litellm.aget_responses().
"""
+from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING
import litellm
from litellm._logging import verbose_proxy_logger
+from litellm.constants import (
+ MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
+ MAX_OBJECTS_PER_POLL_CYCLE,
+)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
@@ -27,6 +32,27 @@ class CheckResponsesCost:
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
+ async def _cleanup_stale_managed_objects(self) -> None:
+ """
+ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
+ in non-terminal states as 'stale_expired'. These will never complete and
+ should not be polled.
+ """
+ cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
+ result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
+ where={
+ "file_purpose": "response",
+ "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
+ "created_at": {"lt": cutoff},
+ },
+ data={"status": "stale_expired"},
+ )
+ if result > 0:
+ verbose_proxy_logger.warning(
+ f"CheckResponsesCost: marked {result} stale managed objects "
+ f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
+ )
+
async def check_responses_cost(self):
"""
Check if background responses are complete and track their cost.
@@ -35,11 +61,20 @@ class CheckResponsesCost:
- Cost is automatically tracked by litellm.aget_responses()
- Mark completed/failed/cancelled responses as complete in the database
"""
+ try:
+ await self._cleanup_stale_managed_objects()
+ except Exception as cleanup_err:
+ verbose_proxy_logger.warning(
+ f"CheckResponsesCost: stale cleanup failed (poll will continue): {cleanup_err}"
+ )
+
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
where={
"status": {"in": ["queued", "in_progress"]},
"file_purpose": "response",
- }
+ },
+ take=MAX_OBJECTS_PER_POLL_CYCLE,
+ order={"created_at": "asc"},
)
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py
index 569ea17f6d8..5530054170c 100644
--- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py
+++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py
@@ -26,6 +26,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
get_batch_id_from_unified_batch_id,
get_content_type_from_file_object,
get_model_id_from_unified_batch_id,
+ get_models_from_unified_file_id,
normalize_mime_type_for_provider,
)
from litellm.types.llms.openai import (
@@ -230,12 +231,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if managed_file:
return managed_file.created_by == user_id
- return False
+ raise HTTPException(
+ status_code=404,
+ detail=f"File not found: {unified_file_id}",
+ )
async def can_user_call_unified_object_id(
self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth
) -> bool:
- ## check if the user has access to the unified object id
## check if the user has access to the unified object id
user_id = user_api_key_dict.user_id
managed_object = (
@@ -246,7 +249,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if managed_object:
return managed_object.created_by == user_id
- return True # don't raise error if managed object is not found
+ raise HTTPException(
+ status_code=404,
+ detail=f"Object not found: {unified_object_id}",
+ )
async def list_user_batches(
self,
@@ -584,7 +590,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
model_file_id_mapping = cast(
Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")
)
+ # model_info may be at top-level or nested under litellm_metadata
+ # (batch/file operations use litellm_metadata)
model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
+ if model_id is None:
+ model_id = cast(
+ Optional[str],
+ kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
+ )
mapped_file_id: Optional[str] = None
if input_file_id and model_file_id_mapping and model_id:
mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(
@@ -892,6 +905,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
) # managed batch id
model_id = cast(Optional[str], response._hidden_params.get("model_id"))
model_name = cast(Optional[str], response._hidden_params.get("model_name"))
+ resolved_model_name = model_name
+
+ # Some providers (e.g. Vertex batch retrieve) do not set model_name on
+ # the response. In that case, recover target_model_names from the input
+ # managed file metadata so unified output IDs preserve routing metadata.
+ if not resolved_model_name and isinstance(unified_file_id, str):
+ decoded_unified_file_id = (
+ _is_base64_encoded_unified_file_id(unified_file_id)
+ or unified_file_id
+ )
+ target_model_names = get_models_from_unified_file_id(
+ decoded_unified_file_id
+ )
+ if target_model_names:
+ resolved_model_name = ",".join(target_model_names)
original_response_id = response.id
if (unified_batch_id or unified_file_id) and model_id:
@@ -899,49 +927,58 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
batch_id=response.id, model_id=model_id
)
- if (
- response.output_file_id and model_id
- ): # return a file id with the model_id and output_file_id
- original_output_file_id = response.output_file_id
- response.output_file_id = self.get_unified_output_file_id(
- output_file_id=response.output_file_id,
- model_id=model_id,
- model_name=model_name,
- )
-
- # Fetch the actual file object for the output file
- file_object = None
- try:
- # Use litellm to retrieve the file object from the provider
- from litellm import afile_retrieve
- file_object = await afile_retrieve(
- custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai",
- file_id=original_output_file_id
+ # Handle both output_file_id and error_file_id
+ for file_attr in ["output_file_id", "error_file_id"]:
+ file_id_value = getattr(response, file_attr, None)
+ if file_id_value and model_id:
+ original_file_id = file_id_value
+ unified_file_id = self.get_unified_output_file_id(
+ output_file_id=original_file_id,
+ model_id=model_id,
+ model_name=resolved_model_name,
)
- verbose_logger.debug(
- f"Successfully retrieved file object for output_file_id={original_output_file_id}"
+ setattr(response, file_attr, unified_file_id)
+
+ # Use llm_router credentials when available. Without credentials,
+ # Azure and other auth-required providers return 500/401.
+ file_object = None
+ try:
+ # Import module and use getattr for better testability with mocks
+ import litellm.proxy.proxy_server as proxy_server_module
+ _llm_router = getattr(proxy_server_module, 'llm_router', None)
+ if _llm_router is not None and model_id:
+ _creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {}
+ file_object = await litellm.afile_retrieve(
+ file_id=original_file_id,
+ **_creds,
+ )
+ else:
+ file_object = await litellm.afile_retrieve(
+ custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai",
+ file_id=original_file_id,
+ )
+ verbose_logger.debug(
+ f"Successfully retrieved file object for {file_attr}={original_file_id}"
+ )
+ except Exception as e:
+ verbose_logger.warning(
+ f"Failed to retrieve file object for {file_attr}={original_file_id}: {str(e)}. Storing with None and will fetch on-demand."
+ )
+
+ await self.store_unified_file_id(
+ file_id=unified_file_id,
+ file_object=file_object,
+ litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
+ model_mappings={model_id: original_file_id},
+ user_api_key_dict=user_api_key_dict,
)
- except Exception as e:
- verbose_logger.warning(
- f"Failed to retrieve file object for output_file_id={original_output_file_id}: {str(e)}. Storing with None and will fetch on-demand."
- )
-
- await self.store_unified_file_id(
- file_id=response.output_file_id,
- file_object=file_object,
- litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
- model_mappings={model_id: original_output_file_id},
- user_api_key_dict=user_api_key_dict,
- )
- asyncio.create_task(
- self.store_unified_object_id(
- unified_object_id=response.id,
- file_object=response,
- litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
- model_object_id=original_response_id,
- file_purpose="batch",
- user_api_key_dict=user_api_key_dict,
- )
+ await self.store_unified_object_id(
+ unified_object_id=response.id,
+ file_object=response,
+ litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
+ model_object_id=original_response_id,
+ file_purpose="batch",
+ user_api_key_dict=user_api_key_dict,
)
elif isinstance(response, LiteLLMFineTuningJob):
## Check if unified_file_id is in the response
@@ -958,15 +995,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
response.id = self.get_unified_generic_response_id(
model_id=model_id, generic_response_id=response.id
)
- asyncio.create_task(
- self.store_unified_object_id(
- unified_object_id=response.id,
- file_object=response,
- litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
- model_object_id=original_response_id,
- file_purpose="fine-tune",
- user_api_key_dict=user_api_key_dict,
- )
+ await self.store_unified_object_id(
+ unified_object_id=response.id,
+ file_object=response,
+ litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
+ model_object_id=original_response_id,
+ file_purpose="fine-tune",
+ user_api_key_dict=user_api_key_dict,
)
elif isinstance(response, AsyncCursorPage):
"""
@@ -1006,8 +1041,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
# Case 2: Managed file and the file object exists in the database
+ # The stored file_object has the raw provider ID. Replace with the unified ID
+ # so callers see a consistent ID (matching Case 3 which does response.id = file_id).
if stored_file_object and stored_file_object.file_object:
- return stored_file_object.file_object
+ # Use model_copy to ensure the ID update persists (Pydantic v2 compatibility)
+ response = stored_file_object.file_object.model_copy(update={"id": file_id})
+ return response
# Case 3: Managed file exists in the database but not the file object (for. e.g the batch task might not have run)
# So we fetch the file object from the provider. We deliberately do not store the result to avoid interfering with batch cost tracking code.
@@ -1035,6 +1074,166 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"""Handled in files_endpoints.py"""
return []
+ def _is_batch_polling_enabled(self) -> bool:
+ """
+ Check if batch cost tracking is actually enabled and running.
+ Returns:
+ bool: True if batch cost tracking is active, False otherwise
+ """
+ try:
+ # Import here to avoid circular dependencies
+ import litellm.proxy.proxy_server as proxy_server_module
+
+ # Check if the scheduler has the batch cost checking job registered
+ scheduler = getattr(proxy_server_module, 'scheduler', None)
+ if scheduler is None:
+ return False
+
+ # Check if the check_batch_cost_job exists in the scheduler
+ try:
+ job = scheduler.get_job('check_batch_cost_job')
+ if job is not None:
+ return True
+ except Exception:
+ # Job not found or scheduler doesn't support get_job
+ pass
+
+ return False
+ except Exception as e:
+ verbose_logger.warning(
+ f"Error checking batch polling configuration: {e}. Assuming disabled."
+ )
+ return False
+
+ async def _get_batches_referencing_file(
+ self, file_id: str
+ ) -> List[Dict[str, Any]]:
+ """
+ Find batches that reference this file and still need cost tracking.
+ Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost.
+ Args:
+ file_id: The unified file ID to check
+
+ Returns:
+ List of batch objects referencing this file in non-terminal state
+ (max 10 for error message display)
+ """
+ # Prepare list of file IDs to check (both unified and provider IDs)
+ file_ids_to_check = [file_id]
+
+ # Get model-specific file IDs for this unified file ID if it's a managed file
+ try:
+ model_file_id_mapping = await self.get_model_file_id_mapping(
+ [file_id], litellm_parent_otel_span=None
+ )
+
+ if model_file_id_mapping and file_id in model_file_id_mapping:
+ # Add all provider file IDs for this unified file
+ provider_file_ids = list(model_file_id_mapping[file_id].values())
+ file_ids_to_check.extend(provider_file_ids)
+ except Exception as e:
+ verbose_logger.debug(
+ f"Could not get model file ID mapping for {file_id}: {e}. "
+ f"Will only check unified file ID."
+ )
+ MAX_MATCHES_TO_RETURN = 10
+
+ batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
+ where={
+ "file_purpose": "batch",
+ "batch_processed": False,
+ "status": {"not_in": ["failed", "expired", "cancelled"]}
+ },
+ take=MAX_MATCHES_TO_RETURN,
+ order={"created_at": "desc"},
+ )
+
+ referencing_batches = []
+ for batch in batches:
+ try:
+ # Parse the batch file_object to check for file references
+ batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object
+
+ # Extract file IDs from batch
+ # Batches typically reference the unified file ID in input_file_id
+ # Output and error files are generated by the provider
+ input_file_id = batch_data.get("input_file_id")
+ output_file_id = batch_data.get("output_file_id")
+ error_file_id = batch_data.get("error_file_id")
+
+ referenced_file_ids = [fid for fid in [input_file_id, output_file_id, error_file_id] if fid]
+
+ # Check if any referenced file ID matches the file we're trying to delete
+ if any(ref_id in file_ids_to_check for ref_id in referenced_file_ids):
+ referencing_batches.append({
+ "batch_id": batch.unified_object_id,
+ "status": batch.status,
+ "created_at": batch.created_at,
+ })
+ except Exception as e:
+ verbose_logger.warning(
+ f"Error parsing batch object {batch.unified_object_id}: {e}"
+ )
+ continue
+
+ return referencing_batches
+
+ async def _check_file_deletion_allowed(self, file_id: str) -> None:
+ """
+ Check if file deletion should be blocked due to batch references.
+
+ Blocks deletion if:
+ 1. File is referenced by any batch in non-terminal state, AND
+ 2. Batch polling is configured (user wants cost tracking)
+
+ Args:
+ file_id: The unified file ID to check
+
+ Raises:
+ HTTPException: If file deletion should be blocked
+ """
+ # Check if batch polling is enabled
+ if not self._is_batch_polling_enabled():
+ # Batch polling not configured, allow deletion
+ return
+
+ # Check if file is referenced by any non-terminal batches
+ referencing_batches = await self._get_batches_referencing_file(file_id)
+
+ if referencing_batches:
+ # File is referenced by non-terminal batches and polling is enabled
+ MAX_BATCHES_IN_ERROR = 5 # Limit batches shown in error message for readability
+
+ # Show up to MAX_BATCHES_IN_ERROR in the error message
+ batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR]
+ batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show]
+
+ # Determine the count message
+ count_message = f"{len(referencing_batches)}"
+ if len(referencing_batches) >= 10: # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file
+ count_message = "10+"
+
+ error_message = (
+ f"Cannot delete file {file_id}. "
+ f"The file is referenced by {count_message} batch(es) in non-terminal state"
+ )
+
+ # Add specific batch details if not too many
+ if len(referencing_batches) <= MAX_BATCHES_IN_ERROR:
+ error_message += f": {', '.join(batch_statuses)}. "
+ else:
+ error_message += f" (showing {MAX_BATCHES_IN_ERROR} most recent): {', '.join(batch_statuses)}. "
+
+ error_message += (
+ f"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. "
+ f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)."
+ )
+
+ raise HTTPException(
+ status_code=400,
+ detail=error_message,
+ )
+
async def afile_delete(
self,
file_id: str,
@@ -1043,6 +1242,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
**data: Dict,
) -> OpenAIFileObject:
+ # Check if file deletion should be blocked due to batch references
+ await self._check_file_deletion_allowed(file_id)
+
# file_id = convert_b64_uid_to_unified_uid(file_id)
model_file_id_mapping = await self.get_model_file_id_mapping(
[file_id], litellm_parent_otel_span
diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py
new file mode 100644
index 00000000000..254d816039c
--- /dev/null
+++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py
@@ -0,0 +1,464 @@
+# What is this?
+## This hook is used to manage vector stores with target_model_names support
+## It allows creating vector stores across multiple models and managing them with unified IDs
+
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
+
+from fastapi import HTTPException
+
+import litellm
+from litellm import Router, verbose_logger
+from litellm._uuid import uuid
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.llms.base_llm.managed_resources import BaseManagedResource
+from litellm.llms.base_llm.managed_resources.utils import (
+ generate_unified_id_string,
+ is_base64_encoded_unified_id,
+)
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.vector_stores import (
+ VectorStoreCreateOptionalRequestParams,
+ VectorStoreCreateResponse,
+)
+
+if TYPE_CHECKING:
+ from opentelemetry.trace import Span as _Span
+
+ from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
+ from litellm.proxy.utils import PrismaClient as _PrismaClient
+
+ Span = Union[_Span, Any]
+ InternalUsageCache = _InternalUsageCache
+ PrismaClient = _PrismaClient
+else:
+ Span = Any
+ InternalUsageCache = Any
+ PrismaClient = Any
+
+
+class _PROXY_LiteLLMManagedVectorStores(
+ CustomLogger, BaseManagedResource[VectorStoreCreateResponse]
+):
+ """
+ Managed vector stores with target_model_names support.
+
+ This class provides functionality to:
+ - Create vector stores across multiple models
+ - Retrieve vector stores by unified ID
+ - Delete vector stores from all models
+ - List vector stores created by a user
+ """
+
+ def __init__(
+ self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient
+ ):
+ CustomLogger.__init__(self)
+ BaseManagedResource.__init__(self, internal_usage_cache, prisma_client)
+
+ # ============================================================================
+ # ABSTRACT METHOD IMPLEMENTATIONS
+ # ============================================================================
+
+ @property
+ def resource_type(self) -> str:
+ """Return the resource type identifier."""
+ return "vector_store"
+
+ @property
+ def table_name(self) -> str:
+ """Return the database table name for vector stores."""
+ # Prisma converts model name LiteLLM_ManagedVectorStoreTable to litellm_managedvectorstoretable
+ return "litellm_managedvectorstoretable"
+
+ def get_unified_resource_id_format(
+ self,
+ resource_object: VectorStoreCreateResponse,
+ target_model_names_list: List[str],
+ ) -> str:
+ """
+ Generate the format string for the unified vector store ID.
+
+ Format:
+ litellm_proxy:vector_store;unified_id,;target_model_names,;resource_id,;model_id,
+ """
+ # VectorStoreCreateResponse is a TypedDict, so resource_object is a dictionary
+ # Extract provider resource ID from the response
+ provider_resource_id = resource_object.get("id", "")
+
+ # Model ID is stored in hidden params if the response object supports it
+ # For TypedDict responses, we need to check if _hidden_params was added
+ hidden_params: Dict[str, Any] = {}
+ if hasattr(resource_object, "_hidden_params"):
+ hidden_params = getattr(resource_object, "_hidden_params", {}) or {}
+ model_id = hidden_params.get("model_id", "")
+
+ return generate_unified_id_string(
+ resource_type=self.resource_type,
+ unified_uuid=str(uuid.uuid4()),
+ target_model_names=target_model_names_list,
+ provider_resource_id=provider_resource_id,
+ model_id=model_id,
+ )
+
+ async def create_resource_for_model(
+ self,
+ llm_router: Router,
+ model: str,
+ request_data: Dict[str, Any],
+ litellm_parent_otel_span: Span,
+ ) -> VectorStoreCreateResponse:
+ """
+ Create a vector store for a specific model.
+
+ Args:
+ llm_router: LiteLLM router instance
+ model: Model name to create vector store for
+ request_data: Request data for vector store creation
+ litellm_parent_otel_span: OpenTelemetry span for tracing
+
+ Returns:
+ VectorStoreCreateResponse from the provider
+ """
+ # Use the router to create the vector store
+ response = await llm_router.avector_store_create(
+ model=model, **request_data
+ )
+ return response
+
+ # ============================================================================
+ # VECTOR STORE CRUD OPERATIONS
+ # ============================================================================
+
+ async def acreate_vector_store(
+ self,
+ create_request: VectorStoreCreateOptionalRequestParams,
+ llm_router: Router,
+ target_model_names_list: List[str],
+ litellm_parent_otel_span: Span,
+ user_api_key_dict: UserAPIKeyAuth,
+ ) -> VectorStoreCreateResponse:
+ """
+ Create a vector store across multiple models.
+
+ Args:
+ create_request: Vector store creation request parameters
+ llm_router: LiteLLM router instance
+ target_model_names_list: List of target model names
+ litellm_parent_otel_span: OpenTelemetry span for tracing
+ user_api_key_dict: User API key authentication details
+
+ Returns:
+ VectorStoreCreateResponse with unified ID
+ """
+ verbose_logger.info(
+ f"Creating managed vector store for models: {target_model_names_list}"
+ )
+
+ # Create vector store for each model
+ # Convert TypedDict to Dict[str, Any] for base class compatibility
+ request_data_dict: Dict[str, Any] = dict(create_request)
+ responses = await self.create_resource_for_each_model(
+ llm_router=llm_router,
+ request_data=request_data_dict,
+ target_model_names_list=target_model_names_list,
+ litellm_parent_otel_span=litellm_parent_otel_span,
+ )
+
+ # Generate unified ID
+ unified_id = self.generate_unified_resource_id(
+ resource_objects=responses,
+ target_model_names_list=target_model_names_list,
+ )
+
+ # Extract model mappings from responses
+ model_mappings: Dict[str, str] = {}
+ for response in responses:
+ hidden_params = getattr(response, "_hidden_params", {}) or {}
+ model_id = hidden_params.get("model_id")
+ if model_id:
+ # VectorStoreCreateResponse is a TypedDict, use dict access
+ model_mappings[model_id] = response["id"]
+
+ verbose_logger.debug(
+ f"Created vector stores with model mappings: {model_mappings}"
+ )
+
+ # Store in database
+ await self.store_unified_resource_id(
+ unified_resource_id=unified_id,
+ resource_object=responses[0], # Store first response as template
+ litellm_parent_otel_span=litellm_parent_otel_span,
+ model_mappings=model_mappings,
+ user_api_key_dict=user_api_key_dict,
+ )
+
+ # Return response with unified ID
+ # VectorStoreCreateResponse is a TypedDict, so we need to create a new dict with the unified ID
+ response = responses[0].copy()
+ response["id"] = unified_id
+
+ verbose_logger.info(
+ f"Successfully created managed vector store with unified ID: {unified_id}"
+ )
+
+ return response
+
+ async def alist_vector_stores(
+ self,
+ user_api_key_dict: UserAPIKeyAuth,
+ limit: Optional[int] = None,
+ after: Optional[str] = None,
+ order: Optional[str] = None,
+ ) -> Dict[str, Any]:
+ """
+ List vector stores created by a user.
+
+ Args:
+ user_api_key_dict: User API key authentication details
+ limit: Maximum number of vector stores to return
+ after: Cursor for pagination
+ order: Sort order ('asc' or 'desc')
+
+ Returns:
+ Dictionary with list of vector stores and pagination info
+ """
+ # Use the base class method
+ return await self.list_user_resources(
+ user_api_key_dict=user_api_key_dict,
+ limit=limit,
+ after=after,
+ )
+
+ # ============================================================================
+ # ACCESS CONTROL
+ # ============================================================================
+
+ async def check_vector_store_access(
+ self, vector_store_id: str, user_api_key_dict: UserAPIKeyAuth
+ ) -> bool:
+ """
+ Check if user has access to a vector store.
+
+ Args:
+ vector_store_id: The unified vector store ID
+ user_api_key_dict: User API key authentication details
+
+ Returns:
+ True if user has access, False otherwise
+ """
+ is_unified_id = is_base64_encoded_unified_id(vector_store_id)
+
+ if is_unified_id:
+ # Check access for managed vector store
+ return await self.can_user_access_unified_resource_id(
+ vector_store_id,
+ user_api_key_dict,
+ )
+
+ # Not a managed vector store, allow access
+ return True
+
+ async def check_managed_vector_store_access(
+ self, data: Dict, user_api_key_dict: UserAPIKeyAuth
+ ) -> bool:
+ """
+ Check if user has access to a managed vector store in request data.
+
+ Args:
+ data: Request data containing vector_store_id
+ user_api_key_dict: User API key authentication details
+
+ Returns:
+ True if this is a managed vector store and user has access
+
+ Raises:
+ HTTPException: If user doesn't have access
+ """
+ vector_store_id = cast(Optional[str], data.get("vector_store_id"))
+ is_unified_id = (
+ is_base64_encoded_unified_id(vector_store_id)
+ if vector_store_id
+ else False
+ )
+
+ if is_unified_id and vector_store_id:
+ if await self.can_user_access_unified_resource_id(
+ vector_store_id, user_api_key_dict
+ ):
+ return True
+ else:
+ raise HTTPException(
+ status_code=403,
+ detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}",
+ )
+
+ return False
+
+ # ============================================================================
+ # PRE-CALL HOOK (For Router Integration)
+ # ============================================================================
+
+ async def async_pre_call_hook(
+ self,
+ user_api_key_dict: UserAPIKeyAuth,
+ cache: Any,
+ data: Dict,
+ call_type: str,
+ ) -> Union[Exception, str, Dict, None]:
+ """
+ Pre-call hook to handle vector store operations.
+
+ This hook intercepts vector store requests and:
+ - Validates access for managed vector stores
+ - Transforms unified IDs to provider-specific IDs
+ - Adds model routing information
+
+ Args:
+ user_api_key_dict: User API key authentication details
+ cache: Cache instance
+ data: Request data
+ call_type: Type of call being made
+
+ Returns:
+ Modified request data or None
+ """
+ from litellm.llms.base_llm.managed_resources.utils import (
+ is_base64_encoded_unified_id,
+ parse_unified_id,
+ )
+
+ # Handle vector store search operations
+ if call_type == "avector_store_search":
+ vector_store_id = data.get("vector_store_id")
+
+ if vector_store_id:
+ # Check if it's a managed vector store ID
+ decoded_id = is_base64_encoded_unified_id(vector_store_id)
+
+ if decoded_id:
+ verbose_logger.debug(
+ f"Processing managed vector store search: {vector_store_id}"
+ )
+
+ # Check access
+ has_access = await self.can_user_access_unified_resource_id(
+ vector_store_id, user_api_key_dict
+ )
+
+ if not has_access:
+ raise HTTPException(
+ status_code=403,
+ detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}",
+ )
+
+ # Parse the unified ID to extract components
+ parsed_id = parse_unified_id(vector_store_id)
+
+ if parsed_id:
+ # Extract the model ID and provider resource ID
+ model_id = parsed_id.get("model_id")
+ provider_resource_id = parsed_id.get("provider_resource_id")
+ target_model_names = parsed_id.get("target_model_names", [])
+
+ verbose_logger.debug(
+ f"Decoded vector store - model_id: {model_id}, provider_resource_id: {provider_resource_id}, target_model_names: {target_model_names}"
+ )
+
+ # Determine which model to use for routing
+ # Priority: model_id (deployment ID) > first target_model_name
+ routing_model = None
+ if model_id:
+ routing_model = model_id
+ elif target_model_names and len(target_model_names) > 0:
+ routing_model = target_model_names[0]
+
+ # Set the model for routing
+ if routing_model:
+ data["model"] = routing_model
+ verbose_logger.info(
+ f"Routing vector store search to model: {routing_model}"
+ )
+
+ # Replace the unified ID with the provider-specific ID
+ if provider_resource_id:
+ data["vector_store_id"] = provider_resource_id
+ verbose_logger.debug(
+ f"Replaced unified ID with provider resource ID: {provider_resource_id}"
+ )
+
+ # Handle vector store retrieve/delete operations
+ elif call_type in ("avector_store_retrieve", "avector_store_delete"):
+ await self.check_managed_vector_store_access(data, user_api_key_dict)
+
+ # If it's a managed vector store, we'll handle it in the endpoint
+ # No need to transform here as the endpoint will route to the hook
+
+ return data
+
+ # ============================================================================
+ # POST-CALL HOOK (For Response Transformation)
+ # ============================================================================
+
+ async def async_post_call_success_hook(
+ self,
+ data: Dict,
+ user_api_key_dict: UserAPIKeyAuth,
+ response: Any,
+ ) -> Any:
+ """
+ Post-call hook to transform responses.
+
+ This hook can be used to transform responses if needed.
+ For now, it just passes through the response.
+
+ Args:
+ data: Request data
+ user_api_key_dict: User API key authentication details
+ response: Response from the provider
+
+ Returns:
+ Potentially modified response
+ """
+ # Currently no transformation needed
+ return response
+
+ # ============================================================================
+ # DEPLOYMENT FILTERING
+ # ============================================================================
+
+ async def async_filter_deployments( # type: ignore[override]
+ self,
+ model: str,
+ healthy_deployments: List,
+ messages: Optional[List] = None,
+ request_kwargs: Optional[Dict] = None,
+ parent_otel_span: Optional[Span] = None,
+ ) -> List[Dict]:
+ """
+ Filter deployments based on vector store availability.
+
+ This is used by the router to select only deployments that have
+ the vector store available.
+
+ Note: This method signature is a compromise between CustomLogger and BaseManagedResource
+ parent classes which have incompatible signatures. The type: ignore[override] is necessary
+ due to this multiple inheritance conflict.
+
+ Args:
+ model: Model name
+ healthy_deployments: List of healthy deployments
+ messages: Messages (unused for vector stores, required by CustomLogger interface)
+ request_kwargs: Request kwargs containing vector_store_id and mappings
+ parent_otel_span: OpenTelemetry span for tracing
+
+ Returns:
+ Filtered list of deployments
+ """
+ return await BaseManagedResource.async_filter_deployments(
+ self,
+ model=model,
+ healthy_deployments=healthy_deployments,
+ request_kwargs=request_kwargs,
+ parent_otel_span=parent_otel_span,
+ resource_id_key="vector_store_id",
+ )
diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml
index c5aaa0a3407..515885944f0 100644
--- a/enterprise/pyproject.toml
+++ b/enterprise/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-enterprise"
-version = "0.1.29"
+version = "0.1.34"
description = "Package for LiteLLM Enterprise features"
authors = ["BerriAI"]
readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
-version = "0.1.29"
+version = "0.1.33"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-enterprise==",
diff --git a/license_cache.json b/license_cache.json
new file mode 100644
index 00000000000..575554c49b4
--- /dev/null
+++ b/license_cache.json
@@ -0,0 +1,9 @@
+{
+ "tornado:6.5.3": "Apache-2.0",
+ "redisvl:0.4.1": "MIT",
+ "google-cloud-iam:2.19.1": "Apache 2.0",
+ "google-genai:1.37.0": "Apache-2.0",
+ "azure-keyvault:4.2.0": "MIT License",
+ "soundfile:0.12.1": "BSD 3-Clause License",
+ "openapi-core:0.21.0": "BSD-3-Clause"
+}
\ No newline at end of file
diff --git a/litellm-js/spend-logs/package-lock.json b/litellm-js/spend-logs/package-lock.json
index 1a13a76820e..b24ff0a4940 100644
--- a/litellm-js/spend-logs/package-lock.json
+++ b/litellm-js/spend-logs/package-lock.json
@@ -6,7 +6,7 @@
"": {
"dependencies": {
"@hono/node-server": "^1.10.1",
- "hono": "^4.10.3"
+ "hono": "^4.12.7"
},
"devDependencies": {
"@types/node": "^20.11.17",
@@ -548,9 +548,9 @@
}
},
"node_modules/hono": {
- "version": "4.10.6",
- "resolved": "https://registry.npmjs.org/hono/-/hono-4.10.6.tgz",
- "integrity": "sha512-BIdolzGpDO9MQ4nu3AUuDwHZZ+KViNm+EZ75Ae55eMXMqLVhDFqEMXxtUe9Qh8hjL+pIna/frs2j6Y2yD5Ua/g==",
+ "version": "4.12.7",
+ "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.7.tgz",
+ "integrity": "sha512-jq9l1DM0zVIvsm3lv9Nw9nlJnMNPOcAtsbsgiUhWcFzPE99Gvo6yRTlszSLLYacMeQ6quHD6hMfId8crVHvexw==",
"license": "MIT",
"engines": {
"node": ">=16.9.0"
diff --git a/litellm-js/spend-logs/package.json b/litellm-js/spend-logs/package.json
index 9c1c2d4f6dc..a40b0fc2a83 100644
--- a/litellm-js/spend-logs/package.json
+++ b/litellm-js/spend-logs/package.json
@@ -4,13 +4,27 @@
},
"dependencies": {
"@hono/node-server": "^1.10.1",
- "hono": "^4.10.3"
+ "hono": "^4.12.7"
},
"devDependencies": {
"@types/node": "^20.11.17",
"tsx": "^4.7.1"
},
"overrides": {
- "glob": ">=11.1.0"
+ "glob": ">=11.1.0",
+ "tar": ">=7.5.10",
+ "minimatch": ">=10.2.4",
+ "diff": ">=8.0.3",
+ "@isaacs/brace-expansion": ">=5.0.1",
+ "@babel/traverse": ">=7.23.2",
+ "ws": ">=7.5.10",
+ "http-proxy-middleware": ">=2.0.9",
+ "tar-fs": ">=2.1.4",
+ "webpack-dev-middleware": ">=5.3.4",
+ "braces": ">=3.0.3",
+ "axios": ">=0.30.2",
+ "webpack": ">=5.94.0",
+ "serve-static": ">=1.16.0",
+ "path-to-regexp": ">=0.1.12"
}
-}
+}
\ No newline at end of file
diff --git a/litellm-proxy-extras/build_and_publish.md b/litellm-proxy-extras/build_and_publish.md
new file mode 100644
index 00000000000..6bf16b99466
--- /dev/null
+++ b/litellm-proxy-extras/build_and_publish.md
@@ -0,0 +1,127 @@
+# Build & Publish `litellm-proxy-extras`
+
+This runbook covers building and publishing a new version of the `litellm-proxy-extras` PyPI package. For use by litellm engineers only.
+
+## Prerequisites
+
+- All `schema.prisma` files are in sync (see [migration_runbook.md](./migration_runbook.md) Step 0)
+- Migration has been generated and committed
+- You are in the `litellm-proxy-extras/` directory
+
+## Step 1: Bump the Version
+
+### Option A: Automatic Version Bump (Recommended)
+
+Use commitizen to automatically bump the version across all files:
+
+```bash
+cd litellm-proxy-extras
+cz bump --increment patch
+```
+
+This will automatically:
+- Bump the version in `pyproject.toml` (both `[tool.poetry].version` and `[tool.commitizen].version`)
+- Update the version in `../requirements.txt`
+- Update the version in `../pyproject.toml` (root)
+- Create a git commit with the version bump
+
+Then skip to Step 3 (Install Build Dependencies).
+
+### Option B: Manual Version Bump
+
+Update the version in `pyproject.toml`:
+
+```bash
+cd litellm-proxy-extras
+
+# Check current version
+grep 'version' pyproject.toml
+```
+
+Edit `pyproject.toml` and bump the version (both `[tool.poetry].version` and `[tool.commitizen].version`).
+
+#### Step 2: Update Version in Root Package Files (Manual Only)
+
+After bumping the version in `litellm-proxy-extras/pyproject.toml`, you **must** also update the version reference in the root-level files:
+
+| File | Line to update |
+|------|---------------|
+| `requirements.txt` | `litellm-proxy-extras==X.Y.Z` |
+| `pyproject.toml` (root) | `litellm-proxy-extras = {version = "X.Y.Z", optional = true}` |
+
+```bash
+# From the repo root — replace OLD with NEW version
+sed -i '' 's/litellm-proxy-extras==OLD/litellm-proxy-extras==NEW/' requirements.txt
+sed -i '' 's/litellm-proxy-extras = {version = "OLD"/litellm-proxy-extras = {version = "NEW"/' pyproject.toml
+```
+
+> **Do NOT skip this step.** The main `litellm` package pins the extras version — if you don't update these, users will install the old version.
+
+## Step 3: Install Build Dependencies
+
+```bash
+pip install build twine
+```
+
+## Step 4: Clean Old Artifacts
+
+```bash
+rm -rf dist/ build/ *.egg-info
+```
+
+## Step 5: Build the Package
+
+```bash
+python3 -m build
+```
+
+This creates `.tar.gz` and `.whl` files in the `dist/` directory.
+
+Verify the build output:
+
+```bash
+ls -la dist/
+```
+
+## Step 6: Upload to PyPI
+
+```bash
+twine upload dist/*
+```
+
+You will be prompted for your PyPI API token:
+
+```
+Enter your API token: pypi-...
+```
+
+> Use `__token__` as the username and your PyPI API token as the password.
+
+## Quick Reference (Copy-Paste)
+
+```bash
+cd litellm-proxy-extras
+rm -rf dist/ build/ *.egg-info
+python3 -m build
+twine upload dist/*
+```
+
+---
+
+## Do you want to build and publish a new `litellm-proxy-extras` package? (y/n)
+
+If **yes**, run the following commands in order:
+
+```bash
+cd litellm-proxy-extras
+pip install build twine
+rm -rf dist/ build/ *.egg-info
+python3 -m build
+twine upload dist/*
+```
+
+When `twine upload` runs, enter your PyPI credentials:
+- **Username:** `__token__`
+- **Password:** *(paste your PyPI API key)*
+
+If **no**, you're done — no package publish needed.
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31-py3-none-any.whl
new file mode 100644
index 00000000000..90b36bd78ac
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31.tar.gz
new file mode 100644
index 00000000000..64607235479
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.31.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.32-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.32-py3-none-any.whl
new file mode 100644
index 00000000000..deb9653aa78
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.32-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.32.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.32.tar.gz
new file mode 100644
index 00000000000..212194e31e2
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.32.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.33-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.33-py3-none-any.whl
new file mode 100644
index 00000000000..a4872243ae6
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.33-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.33.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.33.tar.gz
new file mode 100644
index 00000000000..643be22aa42
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.33.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34-py3-none-any.whl
new file mode 100644
index 00000000000..175d84543ec
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34.tar.gz
new file mode 100644
index 00000000000..e1fcc0c603f
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.34.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.35-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.35-py3-none-any.whl
new file mode 100644
index 00000000000..8a443f38ef5
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.35-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.35.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.35.tar.gz
new file mode 100644
index 00000000000..4dde13b32e2
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.35.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36-py3-none-any.whl
new file mode 100644
index 00000000000..c98d9cfcfac
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36.tar.gz
new file mode 100644
index 00000000000..c8c33404620
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.37-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.37-py3-none-any.whl
new file mode 100644
index 00000000000..695dc102c72
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.37-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.37.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.37.tar.gz
new file mode 100644
index 00000000000..d3ecef1752e
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.37.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.40-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.40-py3-none-any.whl
new file mode 100644
index 00000000000..9f2ad8fd317
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.40-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.40.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.40.tar.gz
new file mode 100644
index 00000000000..fdab43c01a3
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.40.tar.gz differ
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/dist/litellm_proxy_extras-0.4.45-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.45-py3-none-any.whl
new file mode 100644
index 00000000000..f658eef665d
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.45-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.45.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.45.tar.gz
new file mode 100644
index 00000000000..5680b26dbff
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.45.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.47-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.47-py3-none-any.whl
new file mode 100644
index 00000000000..9db37609bd1
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.47-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.47.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.47.tar.gz
new file mode 100644
index 00000000000..37c1775e66c
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.47.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.48-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.48-py3-none-any.whl
new file mode 100644
index 00000000000..8dc2d8e136d
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.48-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.48.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.48.tar.gz
new file mode 100644
index 00000000000..65bf8c3718e
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.48.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49-py3-none-any.whl
new file mode 100644
index 00000000000..e44b58f8e63
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49.tar.gz
new file mode 100644
index 00000000000..2c8549ad069
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.49.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.51-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.51-py3-none-any.whl
new file mode 100644
index 00000000000..f3b69199c87
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.51-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.51.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.51.tar.gz
new file mode 100644
index 00000000000..a1ea473b8b1
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.51.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.52-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.52-py3-none-any.whl
new file mode 100644
index 00000000000..d13dbf15536
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.52-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.52.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.52.tar.gz
new file mode 100644
index 00000000000..1c9ade9aa1c
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.52.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.53-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.53-py3-none-any.whl
new file mode 100644
index 00000000000..019b21ccdf2
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.53-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.53.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.53.tar.gz
new file mode 100644
index 00000000000..773a40d38d3
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.53.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl
new file mode 100644
index 00000000000..9a5c185de28
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz
new file mode 100644
index 00000000000..3e4be95b519
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.54.tar.gz differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.56-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.56-py3-none-any.whl
new file mode 100644
index 00000000000..fceb3b04cee
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.56-py3-none-any.whl differ
diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.56.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.56.tar.gz
new file mode 100644
index 00000000000..5e7841ab0da
Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.56.tar.gz differ
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql
index fb8a44814f0..7b33d58899e 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326162113_baseline/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_BudgetTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetTable" (
"budget_id" TEXT NOT NULL,
"max_budget" DOUBLE PRECISION,
"soft_budget" DOUBLE PRECISION,
@@ -18,7 +18,7 @@ CREATE TABLE "LiteLLM_BudgetTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_CredentialsTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_CredentialsTable" (
"credential_id" TEXT NOT NULL,
"credential_name" TEXT NOT NULL,
"credential_values" JSONB NOT NULL,
@@ -32,7 +32,7 @@ CREATE TABLE "LiteLLM_CredentialsTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_ProxyModelTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_ProxyModelTable" (
"model_id" TEXT NOT NULL,
"model_name" TEXT NOT NULL,
"litellm_params" JSONB NOT NULL,
@@ -46,7 +46,7 @@ CREATE TABLE "LiteLLM_ProxyModelTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_OrganizationTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_OrganizationTable" (
"organization_id" TEXT NOT NULL,
"organization_alias" TEXT NOT NULL,
"budget_id" TEXT NOT NULL,
@@ -63,7 +63,7 @@ CREATE TABLE "LiteLLM_OrganizationTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_ModelTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_ModelTable" (
"id" SERIAL NOT NULL,
"aliases" JSONB,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
@@ -75,7 +75,7 @@ CREATE TABLE "LiteLLM_ModelTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_TeamTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_TeamTable" (
"team_id" TEXT NOT NULL,
"team_alias" TEXT,
"organization_id" TEXT,
@@ -102,7 +102,7 @@ CREATE TABLE "LiteLLM_TeamTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_UserTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_UserTable" (
"user_id" TEXT NOT NULL,
"user_alias" TEXT,
"team_id" TEXT,
@@ -131,7 +131,7 @@ CREATE TABLE "LiteLLM_UserTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_VerificationToken" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_VerificationToken" (
"token" TEXT NOT NULL,
"key_name" TEXT,
"key_alias" TEXT,
@@ -166,7 +166,7 @@ CREATE TABLE "LiteLLM_VerificationToken" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_EndUserTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_EndUserTable" (
"user_id" TEXT NOT NULL,
"alias" TEXT,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
@@ -179,7 +179,7 @@ CREATE TABLE "LiteLLM_EndUserTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_Config" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_Config" (
"param_name" TEXT NOT NULL,
"param_value" JSONB,
@@ -187,7 +187,7 @@ CREATE TABLE "LiteLLM_Config" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_SpendLogs" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs" (
"request_id" TEXT NOT NULL,
"call_type" TEXT NOT NULL,
"api_key" TEXT NOT NULL DEFAULT '',
@@ -218,7 +218,7 @@ CREATE TABLE "LiteLLM_SpendLogs" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_ErrorLogs" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_ErrorLogs" (
"request_id" TEXT NOT NULL,
"startTime" TIMESTAMP(3) NOT NULL,
"endTime" TIMESTAMP(3) NOT NULL,
@@ -235,7 +235,7 @@ CREATE TABLE "LiteLLM_ErrorLogs" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_UserNotifications" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_UserNotifications" (
"request_id" TEXT NOT NULL,
"user_id" TEXT NOT NULL,
"models" TEXT[],
@@ -246,7 +246,7 @@ CREATE TABLE "LiteLLM_UserNotifications" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_TeamMembership" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_TeamMembership" (
"user_id" TEXT NOT NULL,
"team_id" TEXT NOT NULL,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
@@ -256,7 +256,7 @@ CREATE TABLE "LiteLLM_TeamMembership" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_OrganizationMembership" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_OrganizationMembership" (
"user_id" TEXT NOT NULL,
"organization_id" TEXT NOT NULL,
"user_role" TEXT,
@@ -269,7 +269,7 @@ CREATE TABLE "LiteLLM_OrganizationMembership" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_InvitationLink" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_InvitationLink" (
"id" TEXT NOT NULL,
"user_id" TEXT NOT NULL,
"is_accepted" BOOLEAN NOT NULL DEFAULT false,
@@ -284,7 +284,7 @@ CREATE TABLE "LiteLLM_InvitationLink" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_AuditLog" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_AuditLog" (
"id" TEXT NOT NULL,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"changed_by" TEXT NOT NULL DEFAULT '',
@@ -299,62 +299,132 @@ CREATE TABLE "LiteLLM_AuditLog" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_CredentialsTable_credential_name_key" ON "LiteLLM_CredentialsTable"("credential_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_CredentialsTable_credential_name_key" ON "LiteLLM_CredentialsTable"("credential_name");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_TeamTable_model_id_key" ON "LiteLLM_TeamTable"("model_id");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_TeamTable_model_id_key" ON "LiteLLM_TeamTable"("model_id");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_UserTable_sso_user_id_key" ON "LiteLLM_UserTable"("sso_user_id");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_UserTable_sso_user_id_key" ON "LiteLLM_UserTable"("sso_user_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_SpendLogs_startTime_idx" ON "LiteLLM_SpendLogs"("startTime");
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx" ON "LiteLLM_SpendLogs"("startTime");
-- CreateIndex
-CREATE INDEX "LiteLLM_SpendLogs_end_user_idx" ON "LiteLLM_SpendLogs"("end_user");
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx" ON "LiteLLM_SpendLogs"("end_user");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_OrganizationMembership_user_id_organization_id_key" ON "LiteLLM_OrganizationMembership"("user_id", "organization_id");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_OrganizationMembership_user_id_organization_id_key" ON "LiteLLM_OrganizationMembership"("user_id", "organization_id");
-- AddForeignKey
-ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationTable_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamTable_organization_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_model_id_fkey" FOREIGN KEY ("model_id") REFERENCES "LiteLLM_ModelTable"("id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamTable_model_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_model_id_fkey" FOREIGN KEY ("model_id") REFERENCES "LiteLLM_ModelTable"("id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_UserTable_organization_id_fkey') THEN
+ ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerificationToken_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerificationToken_organization_id_fkey') THEN
+ ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_EndUserTable_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationMembership_user_id_fkey') THEN
+ ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationMembership_organization_id_fkey') THEN
+ ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationMembership_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_InvitationLink_user_id_fkey') THEN
+ ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_created_by_fkey" FOREIGN KEY ("created_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_InvitationLink_created_by_fkey') THEN
+ ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_created_by_fkey" FOREIGN KEY ("created_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_updated_by_fkey" FOREIGN KEY ("updated_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_InvitationLink_updated_by_fkey') THEN
+ ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_updated_by_fkey" FOREIGN KEY ("updated_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;
+ END IF;
+END $$;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql
index 3379d8e9fda..52f20ee0f28 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250326171002_add_daily_user_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_DailyUserSpend" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_DailyUserSpend" (
"id" TEXT NOT NULL,
"user_id" TEXT NOT NULL,
"date" TEXT NOT NULL,
@@ -17,17 +17,17 @@ CREATE TABLE "LiteLLM_DailyUserSpend" (
);
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyUserSpend_date_idx" ON "LiteLLM_DailyUserSpend"("date");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_date_idx" ON "LiteLLM_DailyUserSpend"("date");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyUserSpend_user_id_idx" ON "LiteLLM_DailyUserSpend"("user_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_idx" ON "LiteLLM_DailyUserSpend"("user_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyUserSpend_api_key_idx" ON "LiteLLM_DailyUserSpend"("api_key");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_api_key_idx" ON "LiteLLM_DailyUserSpend"("api_key");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyUserSpend_model_idx" ON "LiteLLM_DailyUserSpend"("model");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_model_idx" ON "LiteLLM_DailyUserSpend"("model");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql
index e7c5ab566a9..3865194ce88 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250327180120_add_api_requests_to_daily_user_table/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "api_requests" INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "api_requests" INTEGER NOT NULL DEFAULT 0;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql
index e7ea2e9015a..ba3000f7512 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250329084805_new_cron_job_table/migration.sql
@@ -2,7 +2,7 @@
CREATE TYPE "JobStatus" AS ENUM ('ACTIVE', 'INACTIVE');
-- CreateTable
-CREATE TABLE "LiteLLM_CronJob" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_CronJob" (
"cronjob_id" TEXT NOT NULL,
"pod_id" TEXT NOT NULL,
"status" "JobStatus" NOT NULL DEFAULT 'INACTIVE',
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql
index 9f1693500d0..47962aec772 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250331215456_track_success_and_failed_requests_daily_agg_table/migration.sql
@@ -1,4 +1,4 @@
-- AlterTable
-ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "failed_requests" INTEGER NOT NULL DEFAULT 0,
-ADD COLUMN "successful_requests" INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "failed_requests" INTEGER NOT NULL DEFAULT 0,
+ADD COLUMN IF NOT EXISTS "successful_requests" INTEGER NOT NULL DEFAULT 0;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql
index d14a6294581..602cb5ada37 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250411215431_add_managed_file_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_ManagedFileTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedFileTable" (
"id" TEXT NOT NULL,
"unified_file_id" TEXT NOT NULL,
"file_object" JSONB NOT NULL,
@@ -11,8 +11,8 @@ CREATE TABLE "LiteLLM_ManagedFileTable" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_ManagedFileTable_unified_file_id_key" ON "LiteLLM_ManagedFileTable"("unified_file_id");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedFileTable_unified_file_id_key" ON "LiteLLM_ManagedFileTable"("unified_file_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_ManagedFileTable_unified_file_id_idx" ON "LiteLLM_ManagedFileTable"("unified_file_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedFileTable_unified_file_id_idx" ON "LiteLLM_ManagedFileTable"("unified_file_id");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql
index c07df813796..334b421ca31 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250412081753_team_member_permissions/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "team_member_permissions" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "team_member_permissions" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql
index f47e1c2e91b..1a39ee3579d 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415151647_add_cache_read_write_tokens_daily_spend_transactions/migration.sql
@@ -1,4 +1,4 @@
-- AlterTable
-ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0,
-ADD COLUMN "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0,
+ADD COLUMN IF NOT EXISTS "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql
index a6eb461bc2f..c8ef4eff7d3 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250415191926_add_daily_team_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_DailyTeamSpend" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_DailyTeamSpend" (
"id" TEXT NOT NULL,
"team_id" TEXT NOT NULL,
"date" TEXT NOT NULL,
@@ -20,17 +20,17 @@ CREATE TABLE "LiteLLM_DailyTeamSpend" (
);
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTeamSpend_date_idx" ON "LiteLLM_DailyTeamSpend"("date");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_date_idx" ON "LiteLLM_DailyTeamSpend"("date");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTeamSpend_team_id_idx" ON "LiteLLM_DailyTeamSpend"("team_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_idx" ON "LiteLLM_DailyTeamSpend"("team_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTeamSpend_api_key_idx" ON "LiteLLM_DailyTeamSpend"("api_key");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_api_key_idx" ON "LiteLLM_DailyTeamSpend"("api_key");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTeamSpend_model_idx" ON "LiteLLM_DailyTeamSpend"("model");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_model_idx" ON "LiteLLM_DailyTeamSpend"("model");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql
index 8c3cea70937..0a528348144 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416115320_add_tag_table_to_db/migration.sql
@@ -1,9 +1,9 @@
-- AlterTable
-ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0,
-ADD COLUMN "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0;
+ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "cache_creation_input_tokens" INTEGER NOT NULL DEFAULT 0,
+ADD COLUMN IF NOT EXISTS "cache_read_input_tokens" INTEGER NOT NULL DEFAULT 0;
-- CreateTable
-CREATE TABLE "LiteLLM_DailyTagSpend" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_DailyTagSpend" (
"id" TEXT NOT NULL,
"tag" TEXT NOT NULL,
"date" TEXT NOT NULL,
@@ -26,20 +26,20 @@ CREATE TABLE "LiteLLM_DailyTagSpend" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_key" ON "LiteLLM_DailyTagSpend"("tag");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_key" ON "LiteLLM_DailyTagSpend"("tag");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTagSpend_date_idx" ON "LiteLLM_DailyTagSpend"("date");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_date_idx" ON "LiteLLM_DailyTagSpend"("date");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTagSpend_tag_idx" ON "LiteLLM_DailyTagSpend"("tag");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_idx" ON "LiteLLM_DailyTagSpend"("tag");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTagSpend_api_key_idx" ON "LiteLLM_DailyTagSpend"("api_key");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_api_key_idx" ON "LiteLLM_DailyTagSpend"("api_key");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTagSpend_model_idx" ON "LiteLLM_DailyTagSpend"("model");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_model_idx" ON "LiteLLM_DailyTagSpend"("model");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql
index 5c27b84efbf..a5793331eac 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416151339_drop_tag_uniqueness_requirement/migration.sql
@@ -1,3 +1,3 @@
-- DropIndex
-DROP INDEX "LiteLLM_DailyTagSpend_tag_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_key";
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql
index 2ee7838dcfa..fd8997813f4 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250416185146_add_allowed_routes_litellm_verification_token/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "allowed_routes" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "allowed_routes" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql
index 751c75e5f24..f0ad3886eef 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250425182129_add_session_id/migration.sql
@@ -1,4 +1,4 @@
-- AlterTable
-ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "proxy_server_request" JSONB DEFAULT '{}',
-ADD COLUMN "session_id" TEXT;
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "proxy_server_request" JSONB DEFAULT '{}',
+ADD COLUMN IF NOT EXISTS "session_id" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql
index 39e7f2f3b20..ae73e40cd59 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250430193429_add_managed_vector_stores/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_ManagedVectorStoresTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedVectorStoresTable" (
"vector_store_id" TEXT NOT NULL,
"custom_llm_provider" TEXT NOT NULL,
"vector_store_name" TEXT,
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql
index 6b8adc6e7e8..95f26d5cb55 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161526_add_mcp_table_to_db/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_MCPServerTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_MCPServerTable" (
"server_id" TEXT NOT NULL,
"alias" TEXT,
"description" TEXT,
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161527_add_health_check_fields_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161527_add_health_check_fields_to_mcp_servers/migration.sql
index d5c206d1929..71f595173b7 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161527_add_health_check_fields_to_mcp_servers/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507161527_add_health_check_fields_to_mcp_servers/migration.sql
@@ -1,4 +1,4 @@
-- Add health check fields to MCP server table
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "status" TEXT DEFAULT 'unknown';
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "last_health_check" TIMESTAMP(3);
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "health_check_error" TEXT;
\ No newline at end of file
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "status" TEXT DEFAULT 'unknown';
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "last_health_check" TIMESTAMP(3);
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "health_check_error" TEXT;
\ No newline at end of file
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql
index dcfce07a487..9937857875e 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250507184818_add_mcp_key_team_permission_mgmt/migration.sql
@@ -1,17 +1,17 @@
-- AlterTable
-ALTER TABLE "LiteLLM_OrganizationTable" ADD COLUMN "object_permission_id" TEXT;
+ALTER TABLE "LiteLLM_OrganizationTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT;
-- AlterTable
-ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "object_permission_id" TEXT;
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT;
-- AlterTable
-ALTER TABLE "LiteLLM_UserTable" ADD COLUMN "object_permission_id" TEXT;
+ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT;
-- AlterTable
-ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "object_permission_id" TEXT;
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT;
-- CreateTable
-CREATE TABLE "LiteLLM_ObjectPermissionTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_ObjectPermissionTable" (
"object_permission_id" TEXT NOT NULL,
"mcp_servers" TEXT[] DEFAULT ARRAY[]::TEXT[],
@@ -19,14 +19,34 @@ CREATE TABLE "LiteLLM_ObjectPermissionTable" (
);
-- AddForeignKey
-ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_OrganizationTable_object_permission_id_fkey') THEN
+ ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamTable_object_permission_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_UserTable_object_permission_id_fkey') THEN
+ ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerificationToken_object_permission_id_fkey') THEN
+ ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql
index 8f6c68aa67e..a001977be55 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250508072103_add_status_to_spendlogs/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "status" TEXT;
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "status" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql
index eda055d6e56..a647f386d9e 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250510142544_add_session_id_index_spend_logs/migration.sql
@@ -1,3 +1,3 @@
-- CreateIndex
-CREATE INDEX "LiteLLM_SpendLogs_session_id_idx" ON "LiteLLM_SpendLogs"("session_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx" ON "LiteLLM_SpendLogs"("session_id");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250514142245_add_guardrails_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250514142245_add_guardrails_table/migration.sql
index fa99e3be637..41d8c188947 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250514142245_add_guardrails_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250514142245_add_guardrails_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_GuardrailsTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_GuardrailsTable" (
"guardrail_id" TEXT NOT NULL,
"guardrail_name" TEXT NOT NULL,
"litellm_params" JSONB NOT NULL,
@@ -11,5 +11,5 @@ CREATE TABLE "LiteLLM_GuardrailsTable" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_GuardrailsTable_guardrail_name_key" ON "LiteLLM_GuardrailsTable"("guardrail_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_GuardrailsTable_guardrail_name_key" ON "LiteLLM_GuardrailsTable"("guardrail_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250522223020_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250522223020_managed_object_table/migration.sql
index 95fb8372458..564811655f3 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250522223020_managed_object_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250522223020_managed_object_table/migration.sql
@@ -1,10 +1,10 @@
-- AlterTable
-ALTER TABLE "LiteLLM_ManagedFileTable" ADD COLUMN "created_by" TEXT,
-ADD COLUMN "flat_model_file_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
-ADD COLUMN "updated_by" TEXT;
+ALTER TABLE "LiteLLM_ManagedFileTable" ADD COLUMN IF NOT EXISTS "created_by" TEXT,
+ADD COLUMN IF NOT EXISTS "flat_model_file_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ADD COLUMN IF NOT EXISTS "updated_by" TEXT;
-- CreateTable
-CREATE TABLE "LiteLLM_ManagedObjectTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedObjectTable" (
"id" TEXT NOT NULL,
"unified_object_id" TEXT NOT NULL,
"model_object_id" TEXT NOT NULL,
@@ -19,14 +19,14 @@ CREATE TABLE "LiteLLM_ManagedObjectTable" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_ManagedObjectTable_unified_object_id_key" ON "LiteLLM_ManagedObjectTable"("unified_object_id");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedObjectTable_unified_object_id_key" ON "LiteLLM_ManagedObjectTable"("unified_object_id");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_ManagedObjectTable_model_object_id_key" ON "LiteLLM_ManagedObjectTable"("model_object_id");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedObjectTable_model_object_id_key" ON "LiteLLM_ManagedObjectTable"("model_object_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_ManagedObjectTable_unified_object_id_idx" ON "LiteLLM_ManagedObjectTable"("unified_object_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedObjectTable_unified_object_id_idx" ON "LiteLLM_ManagedObjectTable"("unified_object_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_ManagedObjectTable_model_object_id_idx" ON "LiteLLM_ManagedObjectTable"("model_object_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedObjectTable_model_object_id_idx" ON "LiteLLM_ManagedObjectTable"("model_object_id");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250528185438_add_vector_stores_to_object_permissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250528185438_add_vector_stores_to_object_permissions/migration.sql
index 39db701056e..f1f76153496 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250528185438_add_vector_stores_to_object_permissions/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250528185438_add_vector_stores_to_object_permissions/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "vector_stores" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "vector_stores" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql
index 3d36e42577c..ef53fbf3f4a 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250603210143_cascade_budget_changes/migration.sql
@@ -1,6 +1,16 @@
-- DropForeignKey
-ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey";
+DO $$
+BEGIN
+ IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey";
+ END IF;
+END $$;
-- AddForeignKey
-ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE CASCADE ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE CASCADE ON UPDATE CASCADE;
+ END IF;
+END $$;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250618225828_add_health_check_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250618225828_add_health_check_table/migration.sql
index da6f4c23c81..2b4dd1d17d3 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250618225828_add_health_check_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250618225828_add_health_check_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_HealthCheckTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_HealthCheckTable" (
"health_check_id" TEXT NOT NULL,
"model_name" TEXT NOT NULL,
"model_id" TEXT,
@@ -18,11 +18,11 @@ CREATE TABLE "LiteLLM_HealthCheckTable" (
);
-- CreateIndex
-CREATE INDEX "LiteLLM_HealthCheckTable_model_name_idx" ON "LiteLLM_HealthCheckTable"("model_name");
+CREATE INDEX IF NOT EXISTS "LiteLLM_HealthCheckTable_model_name_idx" ON "LiteLLM_HealthCheckTable"("model_name");
-- CreateIndex
-CREATE INDEX "LiteLLM_HealthCheckTable_checked_at_idx" ON "LiteLLM_HealthCheckTable"("checked_at");
+CREATE INDEX IF NOT EXISTS "LiteLLM_HealthCheckTable_checked_at_idx" ON "LiteLLM_HealthCheckTable"("checked_at");
-- CreateIndex
-CREATE INDEX "LiteLLM_HealthCheckTable_status_idx" ON "LiteLLM_HealthCheckTable"("status");
+CREATE INDEX IF NOT EXISTS "LiteLLM_HealthCheckTable_status_idx" ON "LiteLLM_HealthCheckTable"("status");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql
index 51461b82058..75d7e0e74b1 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625145206_cascade_budget_and_loosen_managed_file_json/migration.sql
@@ -1,9 +1,19 @@
-- DropForeignKey
-ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey";
+DO $$
+BEGIN
+ IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TeamMembership" DROP CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey";
+ END IF;
+END $$;
-- AlterTable
ALTER TABLE "LiteLLM_ManagedFileTable" ALTER COLUMN "file_object" DROP NOT NULL;
-- AddForeignKey
-ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TeamMembership_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625213625_add_status_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625213625_add_status_to_managed_object_table/migration.sql
index 7ca7b2c3705..a4f0ca07944 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625213625_add_status_to_managed_object_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250625213625_add_status_to_managed_object_table/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN "status" TEXT;
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "status" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707212517_add_mcp_info_column_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707212517_add_mcp_info_column_mcp_servers/migration.sql
index efe68ff4792..e942383ded7 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707212517_add_mcp_info_column_mcp_servers/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707212517_add_mcp_info_column_mcp_servers/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "mcp_info" JSONB DEFAULT '{}';
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "mcp_info" JSONB DEFAULT '{}';
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707230009_add_mcp_namespaced_tool_name/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707230009_add_mcp_namespaced_tool_name/migration.sql
index 3130619a773..963b29e1875 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707230009_add_mcp_namespaced_tool_name/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250707230009_add_mcp_namespaced_tool_name/migration.sql
@@ -1,42 +1,42 @@
-- DropIndex
-DROP INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key";
-- DropIndex
-DROP INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key";
-- DropIndex
-DROP INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key";
-- AlterTable
-ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN "mcp_namespaced_tool_name" TEXT,
+ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "mcp_namespaced_tool_name" TEXT,
ALTER COLUMN "model" DROP NOT NULL;
-- AlterTable
-ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN "mcp_namespaced_tool_name" TEXT,
+ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "mcp_namespaced_tool_name" TEXT,
ALTER COLUMN "model" DROP NOT NULL;
-- AlterTable
-ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "mcp_namespaced_tool_name" TEXT,
+ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "mcp_namespaced_tool_name" TEXT,
ALTER COLUMN "model" DROP NOT NULL;
-- AlterTable
-ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "mcp_namespaced_tool_name" TEXT;
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "mcp_namespaced_tool_name" TEXT;
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTagSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyTagSpend"("mcp_namespaced_tool_name");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyTagSpend"("mcp_namespaced_tool_name");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTeamSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyTeamSpend"("mcp_namespaced_tool_name");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyTeamSpend"("mcp_namespaced_tool_name");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyUserSpend"("mcp_namespaced_tool_name");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyUserSpend"("mcp_namespaced_tool_name");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250711220620_add_stdio_mcp/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250711220620_add_stdio_mcp/migration.sql
index ebe7a6adb58..685dc6c1d39 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250711220620_add_stdio_mcp/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250711220620_add_stdio_mcp/migration.sql
@@ -1,10 +1,10 @@
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "args" TEXT[] DEFAULT ARRAY[]::TEXT[],
-ADD COLUMN "command" TEXT,
-ADD COLUMN "env" JSONB DEFAULT '{}',
-ADD COLUMN "mcp_access_groups" TEXT[],
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "args" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ADD COLUMN IF NOT EXISTS "command" TEXT,
+ADD COLUMN IF NOT EXISTS "env" JSONB DEFAULT '{}',
+ADD COLUMN IF NOT EXISTS "mcp_access_groups" TEXT[],
ALTER COLUMN "url" DROP NOT NULL;
-- AlterTable
-ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "mcp_access_groups" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "mcp_access_groups" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250718125714_add_litellm_params_to_vector_stores/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250718125714_add_litellm_params_to_vector_stores/migration.sql
index ef9956ddd5f..5f502a374e2 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250718125714_add_litellm_params_to_vector_stores/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250718125714_add_litellm_params_to_vector_stores/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_ManagedVectorStoresTable" ADD COLUMN "litellm_params" JSONB;
+ALTER TABLE "LiteLLM_ManagedVectorStoresTable" ADD COLUMN IF NOT EXISTS "litellm_params" JSONB;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250802162330_prompt_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250802162330_prompt_table/migration.sql
index e5c00ef4adb..81b3574499b 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250802162330_prompt_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250802162330_prompt_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_PromptTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_PromptTable" (
"id" TEXT NOT NULL,
"prompt_id" TEXT NOT NULL,
"litellm_params" JSONB NOT NULL,
@@ -11,5 +11,5 @@ CREATE TABLE "LiteLLM_PromptTable" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_PromptTable_prompt_id_key" ON "LiteLLM_PromptTable"("prompt_id");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_PromptTable_prompt_id_key" ON "LiteLLM_PromptTable"("prompt_id");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql
index 472e2ea1e0c..5b4db0e500a 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250918083359_drop_spec_version_column_from_mcp_table/migration.sql
@@ -5,4 +5,4 @@
*/
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN "spec_version";
+ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN IF EXISTS "spec_version";
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250926194702_unnamed_migration/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250926194702_unnamed_migration/migration.sql
index ea28db19662..a40c88cc85d 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20250926194702_unnamed_migration/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20250926194702_unnamed_migration/migration.sql
@@ -1,7 +1,7 @@
-- AlterTable
-ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "auto_rotate" BOOLEAN DEFAULT false,
-ADD COLUMN "key_rotation_at" TIMESTAMP(3),
-ADD COLUMN "last_rotation_at" TIMESTAMP(3),
-ADD COLUMN "rotation_count" INTEGER DEFAULT 0,
-ADD COLUMN "rotation_interval" TEXT;
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "auto_rotate" BOOLEAN DEFAULT false,
+ADD COLUMN IF NOT EXISTS "key_rotation_at" TIMESTAMP(3),
+ADD COLUMN IF NOT EXISTS "last_rotation_at" TIMESTAMP(3),
+ADD COLUMN IF NOT EXISTS "rotation_count" INTEGER DEFAULT 0,
+ADD COLUMN IF NOT EXISTS "rotation_interval" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003165142_add_allowed_tools_to_mcp/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003165142_add_allowed_tools_to_mcp/migration.sql
index bdac1e42bc2..d5206e8d3a5 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003165142_add_allowed_tools_to_mcp/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003165142_add_allowed_tools_to_mcp/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "allowed_tools" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "allowed_tools" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003190954_extra_headers_to_mcp_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003190954_extra_headers_to_mcp_table/migration.sql
index 1cfcf062eb1..b6275c6421c 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003190954_extra_headers_to_mcp_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251003190954_extra_headers_to_mcp_table/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251006143948_add_mcp_tool_permissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251006143948_add_mcp_tool_permissions/migration.sql
index 51f3be87582..70584a96286 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251006143948_add_mcp_tool_permissions/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251006143948_add_mcp_tool_permissions/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "mcp_tool_permissions" JSONB;
+ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251011084309_add_tag_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251011084309_add_tag_table/migration.sql
index 541c70c7e48..c1005a75a4a 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251011084309_add_tag_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251011084309_add_tag_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_TagTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_TagTable" (
"tag_name" TEXT NOT NULL,
"description" TEXT,
"models" TEXT[],
@@ -14,5 +14,10 @@ CREATE TABLE "LiteLLM_TagTable" (
);
-- AddForeignKey
-ALTER TABLE "LiteLLM_TagTable" ADD CONSTRAINT "LiteLLM_TagTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_TagTable_budget_id_fkey') THEN
+ ALTER TABLE "LiteLLM_TagTable" ADD CONSTRAINT "LiteLLM_TagTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql
index 4cbe4a7184f..3cd2a4c4899 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251023141814_add_search_tool_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_SearchToolsTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_SearchToolsTable" (
"search_tool_id" TEXT NOT NULL,
"search_tool_name" TEXT NOT NULL,
"litellm_params" JSONB NOT NULL,
@@ -11,5 +11,5 @@ CREATE TABLE "LiteLLM_SearchToolsTable" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_SearchToolsTable_search_tool_name_key" ON "LiteLLM_SearchToolsTable"("search_tool_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_SearchToolsTable_search_tool_name_key" ON "LiteLLM_SearchToolsTable"("search_tool_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251031181430_add_cache_config_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251031181430_add_cache_config_table/migration.sql
index 705a6fd4d9b..bb475b03d25 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251031181430_add_cache_config_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251031181430_add_cache_config_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_SSOConfig" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_SSOConfig" (
"id" TEXT NOT NULL DEFAULT 'sso_config',
"sso_settings" JSONB NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
@@ -9,7 +9,7 @@ CREATE TABLE "LiteLLM_SSOConfig" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_CacheConfig" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_CacheConfig" (
"id" TEXT NOT NULL DEFAULT 'cache_config',
"cache_settings" JSONB NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251101131415_add_managed_vector_store_index_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251101131415_add_managed_vector_store_index_table/migration.sql
index af13500d1c7..0997e569171 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251101131415_add_managed_vector_store_index_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251101131415_add_managed_vector_store_index_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_ManagedVectorStoreIndexTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedVectorStoreIndexTable" (
"id" TEXT NOT NULL,
"index_name" TEXT NOT NULL,
"litellm_params" JSONB NOT NULL,
@@ -13,5 +13,5 @@ CREATE TABLE "LiteLLM_ManagedVectorStoreIndexTable" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_ManagedVectorStoreIndexTable_index_name_key" ON "LiteLLM_ManagedVectorStoreIndexTable"("index_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoreIndexTable_index_name_key" ON "LiteLLM_ManagedVectorStoreIndexTable"("index_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251103072422_add_static_headers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251103072422_add_static_headers/migration.sql
index 0bedac76313..452a0b73a51 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251103072422_add_static_headers/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251103072422_add_static_headers/migration.sql
@@ -1,2 +1,2 @@
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "static_headers" JSONB DEFAULT '{}';
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "static_headers" JSONB DEFAULT '{}';
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251104220043_add_credentials_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251104220043_add_credentials_to_mcp_servers/migration.sql
index 800c96f18b7..c86c3f1f4b7 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251104220043_add_credentials_to_mcp_servers/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251104220043_add_credentials_to_mcp_servers/migration.sql
@@ -1,2 +1,2 @@
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "credentials" JSONB DEFAULT '{}';
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "credentials" JSONB DEFAULT '{}';
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..9769699b33f
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251113000000_add_project_table/migration.sql
@@ -0,0 +1,55 @@
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "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
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ProjectTable_team_id_fkey') THEN
+ 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;
+ END IF;
+END $$;
+
+-- AddForeignKey
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ProjectTable_budget_id_fkey') THEN
+ 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;
+ END IF;
+END $$;
+
+-- AddForeignKey
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ProjectTable_object_permission_id_fkey') THEN
+ 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;
+ END IF;
+END $$;
+
+-- AlterTable: Add project_id to LiteLLM_VerificationToken
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "project_id" TEXT;
+
+-- AddForeignKey
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerificationToken_project_id_fkey') THEN
+ 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;
+ END IF;
+END $$;
+
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..69dbb6790c7
--- /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 IF NOT EXISTS "description" TEXT;
+ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN IF NOT EXISTS "model_rpm_limit" JSONB NOT NULL DEFAULT '{}';
+ALTER TABLE "LiteLLM_ProjectTable" ADD COLUMN IF NOT EXISTS "model_tpm_limit" JSONB NOT NULL DEFAULT '{}';
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114173537_add_request_id_to_daily_tag_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114173537_add_request_id_to_daily_tag_spend/migration.sql
index 6871e27a28a..6d66ef36ae2 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114173537_add_request_id_to_daily_tag_spend/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114173537_add_request_id_to_daily_tag_spend/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN "request_id" TEXT;
+ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "request_id" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql
index 74e0eea3134..387b40461a8 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_DailyOrganizationSpend" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_DailyOrganizationSpend" (
"id" TEXT NOT NULL,
"organization_id" TEXT,
"date" TEXT NOT NULL,
@@ -23,20 +23,20 @@ CREATE TABLE "LiteLLM_DailyOrganizationSpend" (
);
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyOrganizationSpend_date_idx" ON "LiteLLM_DailyOrganizationSpend"("date");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_date_idx" ON "LiteLLM_DailyOrganizationSpend"("date");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_idx" ON "LiteLLM_DailyOrganizationSpend"("organization_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_idx" ON "LiteLLM_DailyOrganizationSpend"("organization_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyOrganizationSpend_api_key_idx" ON "LiteLLM_DailyOrganizationSpend"("api_key");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_api_key_idx" ON "LiteLLM_DailyOrganizationSpend"("api_key");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyOrganizationSpend_model_idx" ON "LiteLLM_DailyOrganizationSpend"("model");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_model_idx" ON "LiteLLM_DailyOrganizationSpend"("model");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyOrganizationSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyOrganizationSpend"("mcp_namespaced_tool_name");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyOrganizationSpend"("mcp_namespaced_tool_name");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114182247_agents_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114182247_agents_table/migration.sql
index 28760dcfe48..35595c7eedf 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114182247_agents_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114182247_agents_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_AgentsTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_AgentsTable" (
"agent_id" TEXT NOT NULL,
"agent_name" TEXT NOT NULL,
"litellm_params" JSONB,
@@ -13,5 +13,5 @@ CREATE TABLE "LiteLLM_AgentsTable" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_AgentsTable_agent_name_key" ON "LiteLLM_AgentsTable"("agent_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentsTable_agent_name_key" ON "LiteLLM_AgentsTable"("agent_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql
index 43eb2401422..f604dcadbd9 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251119131227_add_prompt_versioning/migration.sql
@@ -3,10 +3,10 @@ DROP INDEX IF EXISTS "LiteLLM_PromptTable_prompt_id_key";
-- AlterTable
ALTER TABLE "LiteLLM_PromptTable"
-ADD COLUMN "version" INTEGER NOT NULL DEFAULT 1;
+ADD COLUMN IF NOT EXISTS "version" INTEGER NOT NULL DEFAULT 1;
-- CreateIndex
-CREATE INDEX "LiteLLM_PromptTable_prompt_id_idx" ON "LiteLLM_PromptTable" ("prompt_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_PromptTable_prompt_id_idx" ON "LiteLLM_PromptTable" ("prompt_id");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_PromptTable_prompt_id_version_key" ON "LiteLLM_PromptTable" ("prompt_id", "version");
\ No newline at end of file
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_PromptTable_prompt_id_version_key" ON "LiteLLM_PromptTable" ("prompt_id", "version");
\ No newline at end of file
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql
index 4ea082f2750..bc483aace76 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "organization_id" TEXT;
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "organization_id" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql
index c4234785c54..3544768cc16 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251204124859_add_end_user_spend_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_DailyEndUserSpend" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_DailyEndUserSpend" (
"id" TEXT NOT NULL,
"end_user_id" TEXT,
"date" TEXT NOT NULL,
@@ -23,20 +23,20 @@ CREATE TABLE "LiteLLM_DailyEndUserSpend" (
);
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyEndUserSpend_date_idx" ON "LiteLLM_DailyEndUserSpend"("date");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_date_idx" ON "LiteLLM_DailyEndUserSpend"("date");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_idx" ON "LiteLLM_DailyEndUserSpend"("end_user_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_idx" ON "LiteLLM_DailyEndUserSpend"("end_user_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyEndUserSpend_api_key_idx" ON "LiteLLM_DailyEndUserSpend"("api_key");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_api_key_idx" ON "LiteLLM_DailyEndUserSpend"("api_key");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyEndUserSpend_model_idx" ON "LiteLLM_DailyEndUserSpend"("model");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_model_idx" ON "LiteLLM_DailyEndUserSpend"("model");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyEndUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyEndUserSpend"("mcp_namespaced_tool_name");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyEndUserSpend"("mcp_namespaced_tool_name");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251209112246_add_ui_settings_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251209112246_add_ui_settings_table/migration.sql
index 1719ce646d4..01e2c9fa761 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251209112246_add_ui_settings_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251209112246_add_ui_settings_table/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_UISettings" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_UISettings" (
"id" TEXT NOT NULL DEFAULT 'ui_settings',
"ui_settings" JSONB NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql
index 964904c14c1..4cc894f9b22 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql
@@ -1,8 +1,8 @@
-- AlterTable
-ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "agent_id" TEXT;
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "agent_id" TEXT;
-- CreateTable
-CREATE TABLE "LiteLLM_DailyAgentSpend" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_DailyAgentSpend" (
"id" TEXT NOT NULL,
"agent_id" TEXT,
"date" TEXT NOT NULL,
@@ -26,20 +26,20 @@ CREATE TABLE "LiteLLM_DailyAgentSpend" (
);
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyAgentSpend_date_idx" ON "LiteLLM_DailyAgentSpend"("date");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_date_idx" ON "LiteLLM_DailyAgentSpend"("date");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyAgentSpend_agent_id_idx" ON "LiteLLM_DailyAgentSpend"("agent_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_idx" ON "LiteLLM_DailyAgentSpend"("agent_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyAgentSpend_api_key_idx" ON "LiteLLM_DailyAgentSpend"("api_key");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_api_key_idx" ON "LiteLLM_DailyAgentSpend"("api_key");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyAgentSpend_model_idx" ON "LiteLLM_DailyAgentSpend"("model");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_model_idx" ON "LiteLLM_DailyAgentSpend"("model");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyAgentSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyAgentSpend"("mcp_namespaced_tool_name");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyAgentSpend"("mcp_namespaced_tool_name");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql
index 6ca66ddaad2..8acff490bb7 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_DeletedTeamTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_DeletedTeamTable" (
"id" TEXT NOT NULL,
"team_id" TEXT NOT NULL,
"team_alias" TEXT,
@@ -33,7 +33,7 @@ CREATE TABLE "LiteLLM_DeletedTeamTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_DeletedVerificationToken" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_DeletedVerificationToken" (
"id" TEXT NOT NULL,
"token" TEXT NOT NULL,
"key_name" TEXT,
@@ -80,38 +80,38 @@ CREATE TABLE "LiteLLM_DeletedVerificationToken" (
);
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedTeamTable_team_id_idx" ON "LiteLLM_DeletedTeamTable"("team_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_team_id_idx" ON "LiteLLM_DeletedTeamTable"("team_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedTeamTable_deleted_at_idx" ON "LiteLLM_DeletedTeamTable"("deleted_at");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_deleted_at_idx" ON "LiteLLM_DeletedTeamTable"("deleted_at");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedTeamTable_organization_id_idx" ON "LiteLLM_DeletedTeamTable"("organization_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_organization_id_idx" ON "LiteLLM_DeletedTeamTable"("organization_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedTeamTable_team_alias_idx" ON "LiteLLM_DeletedTeamTable"("team_alias");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_team_alias_idx" ON "LiteLLM_DeletedTeamTable"("team_alias");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedTeamTable_created_at_idx" ON "LiteLLM_DeletedTeamTable"("created_at");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedTeamTable_created_at_idx" ON "LiteLLM_DeletedTeamTable"("created_at");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedVerificationToken_token_idx" ON "LiteLLM_DeletedVerificationToken"("token");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_token_idx" ON "LiteLLM_DeletedVerificationToken"("token");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedVerificationToken_deleted_at_idx" ON "LiteLLM_DeletedVerificationToken"("deleted_at");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_deleted_at_idx" ON "LiteLLM_DeletedVerificationToken"("deleted_at");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedVerificationToken_user_id_idx" ON "LiteLLM_DeletedVerificationToken"("user_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_user_id_idx" ON "LiteLLM_DeletedVerificationToken"("user_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedVerificationToken_team_id_idx" ON "LiteLLM_DeletedVerificationToken"("team_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_team_id_idx" ON "LiteLLM_DeletedVerificationToken"("team_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedVerificationToken_organization_id_idx" ON "LiteLLM_DeletedVerificationToken"("organization_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_organization_id_idx" ON "LiteLLM_DeletedVerificationToken"("organization_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedVerificationToken_key_alias_idx" ON "LiteLLM_DeletedVerificationToken"("key_alias");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_key_alias_idx" ON "LiteLLM_DeletedVerificationToken"("key_alias");
-- CreateIndex
-CREATE INDEX "LiteLLM_DeletedVerificationToken_created_at_idx" ON "LiteLLM_DeletedVerificationToken"("created_at");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeletedVerificationToken_created_at_idx" ON "LiteLLM_DeletedVerificationToken"("created_at");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251220144550_schema_update/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251220144550_schema_update/migration.sql
index b40defec309..c79ce17b9ee 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251220144550_schema_update/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251220144550_schema_update/migration.sql
@@ -1,5 +1,5 @@
-- CreateTable
-CREATE TABLE "LiteLLM_SkillsTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_SkillsTable" (
"skill_id" TEXT NOT NULL,
"display_title" TEXT,
"description" TEXT,
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql
index 8eebb797e2c..a854693d57a 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql
@@ -1,5 +1,5 @@
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "authorization_url" TEXT,
-ADD COLUMN "registration_url" TEXT,
-ADD COLUMN "token_url" TEXT;
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "authorization_url" TEXT,
+ADD COLUMN IF NOT EXISTS "registration_url" TEXT,
+ADD COLUMN IF NOT EXISTS "token_url" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql
index 8d3e02bd051..26319c5ed84 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260105151539_add_allow_all_keys_to_mcp_servers/migration.sql
@@ -1,3 +1,3 @@
-- AlterTable
-ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "allow_all_keys" BOOLEAN NOT NULL DEFAULT false;
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "allow_all_keys" BOOLEAN NOT NULL DEFAULT false;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260106155622_add_endpoint_to_daily_activity_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260106155622_add_endpoint_to_daily_activity_tables/migration.sql
index 4ed7feb9ca0..3e94f736ea3 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260106155622_add_endpoint_to_daily_activity_tables/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260106155622_add_endpoint_to_daily_activity_tables/migration.sql
@@ -1,72 +1,72 @@
-- DropIndex
-DROP INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key";
+DROP INDEX IF EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key";
-- DropIndex
-DROP INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key";
-- DropIndex
-DROP INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key";
-- DropIndex
-DROP INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key";
-- DropIndex
-DROP INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key";
-- DropIndex
-DROP INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key";
+DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key";
-- AlterTable
-ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN "endpoint" TEXT;
+ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT;
-- AlterTable
-ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN "endpoint" TEXT;
+ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT;
-- AlterTable
-ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN "endpoint" TEXT;
+ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT;
-- AlterTable
-ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN "endpoint" TEXT;
+ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT;
-- AlterTable
-ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN "endpoint" TEXT;
+ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT;
-- AlterTable
-ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN "endpoint" TEXT;
+ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "endpoint" TEXT;
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyAgentSpend_endpoint_idx" ON "LiteLLM_DailyAgentSpend"("endpoint");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_endpoint_idx" ON "LiteLLM_DailyAgentSpend"("endpoint");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyEndUserSpend_endpoint_idx" ON "LiteLLM_DailyEndUserSpend"("endpoint");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_endpoint_idx" ON "LiteLLM_DailyEndUserSpend"("endpoint");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_api_key_model_cu_key" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyOrganizationSpend_endpoint_idx" ON "LiteLLM_DailyOrganizationSpend"("endpoint");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_endpoint_idx" ON "LiteLLM_DailyOrganizationSpend"("endpoint");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTagSpend_endpoint_idx" ON "LiteLLM_DailyTagSpend"("endpoint");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_endpoint_idx" ON "LiteLLM_DailyTagSpend"("endpoint");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_api_key_model_custom_llm_pro_key" ON "LiteLLM_DailyTagSpend"("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyTeamSpend_endpoint_idx" ON "LiteLLM_DailyTeamSpend"("endpoint");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_endpoint_idx" ON "LiteLLM_DailyTeamSpend"("endpoint");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyTeamSpend"("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
-- CreateIndex
-CREATE INDEX "LiteLLM_DailyUserSpend_endpoint_idx" ON "LiteLLM_DailyUserSpend"("endpoint");
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_endpoint_idx" ON "LiteLLM_DailyUserSpend"("endpoint");
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260107111013_add_router_settings_to_keys_teams/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260107111013_add_router_settings_to_keys_teams/migration.sql
index 95566950118..07cf2f95f90 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260107111013_add_router_settings_to_keys_teams/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260107111013_add_router_settings_to_keys_teams/migration.sql
@@ -1,6 +1,6 @@
-- AlterTable
-ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "router_settings" JSONB DEFAULT '{}';
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "router_settings" JSONB DEFAULT '{}';
-- AlterTable
-ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "router_settings" JSONB DEFAULT '{}';
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "router_settings" JSONB DEFAULT '{}';
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql
index add80b39e7f..561c36530d9 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260108_add_user_email_lower_idx/migration.sql
@@ -6,4 +6,4 @@
-- With this index, PostgreSQL can use an Index Scan for O(log n) performance.
--
-- Related: GitHub Issue #18411
-CREATE INDEX "LiteLLM_UserTable_user_email_lower_idx" ON "LiteLLM_UserTable"(LOWER("user_email"));
+CREATE INDEX IF NOT EXISTS "LiteLLM_UserTable_user_email_lower_idx" ON "LiteLLM_UserTable"(LOWER("user_email"));
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260116142756_update_deleted_keys_teams_table_routing_settings/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260116142756_update_deleted_keys_teams_table_routing_settings/migration.sql
index 9426bed0da2..b948ffcad66 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260116142756_update_deleted_keys_teams_table_routing_settings/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260116142756_update_deleted_keys_teams_table_routing_settings/migration.sql
@@ -1,6 +1,6 @@
-- AlterTable
-ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "router_settings" JSONB DEFAULT '{}';
+ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "router_settings" JSONB DEFAULT '{}';
-- AlterTable
-ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "router_settings" JSONB DEFAULT '{}';
+ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "router_settings" JSONB DEFAULT '{}';
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260123131407_add_policy_tables_and_policies_field/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260123131407_add_policy_tables_and_policies_field/migration.sql
index 595d8f4a0c5..8b6ac7a7f3e 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260123131407_add_policy_tables_and_policies_field/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260123131407_add_policy_tables_and_policies_field/migration.sql
@@ -1,20 +1,20 @@
-- AlterTable
-ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- AlterTable
-ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- AlterTable
-ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- AlterTable
-ALTER TABLE "LiteLLM_UserTable" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- AlterTable
-ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "policies" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- CreateTable
-CREATE TABLE "LiteLLM_PolicyTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_PolicyTable" (
"policy_id" TEXT NOT NULL,
"policy_name" TEXT NOT NULL,
"inherit" TEXT,
@@ -31,7 +31,7 @@ CREATE TABLE "LiteLLM_PolicyTable" (
);
-- CreateTable
-CREATE TABLE "LiteLLM_PolicyAttachmentTable" (
+CREATE TABLE IF NOT EXISTS "LiteLLM_PolicyAttachmentTable" (
"attachment_id" TEXT NOT NULL,
"policy_name" TEXT NOT NULL,
"scope" TEXT,
@@ -47,5 +47,5 @@ CREATE TABLE "LiteLLM_PolicyAttachmentTable" (
);
-- CreateIndex
-CREATE UNIQUE INDEX "LiteLLM_PolicyTable_policy_name_key" ON "LiteLLM_PolicyTable"("policy_name");
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_PolicyTable_policy_name_key" ON "LiteLLM_PolicyTable"("policy_name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260131150814_add_team_user_to_vector_stores/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260131150814_add_team_user_to_vector_stores/migration.sql
index 2032f76a5de..1f5dc311bd6 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260131150814_add_team_user_to_vector_stores/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260131150814_add_team_user_to_vector_stores/migration.sql
@@ -1,10 +1,13 @@
-- AlterTable
-ALTER TABLE "LiteLLM_ManagedVectorStoresTable" ADD COLUMN "team_id" TEXT,
-ADD COLUMN "user_id" TEXT;
+ALTER TABLE "LiteLLM_ManagedVectorStoresTable"
+ ADD COLUMN IF NOT EXISTS "team_id" TEXT,
+ ADD COLUMN IF NOT EXISTS "user_id" TEXT;
-- CreateIndex
-CREATE INDEX "LiteLLM_ManagedVectorStoresTable_team_id_idx" ON "LiteLLM_ManagedVectorStoresTable"("team_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoresTable_team_id_idx"
+ ON "LiteLLM_ManagedVectorStoresTable"("team_id");
-- CreateIndex
-CREATE INDEX "LiteLLM_ManagedVectorStoresTable_user_id_idx" ON "LiteLLM_ManagedVectorStoresTable"("user_id");
+CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoresTable_user_id_idx"
+ ON "LiteLLM_ManagedVectorStoresTable"("user_id");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql
new file mode 100644
index 00000000000..58b2d4048c2
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql
@@ -0,0 +1,19 @@
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_DeprecatedVerificationToken" (
+ "id" TEXT NOT NULL,
+ "token" TEXT NOT NULL,
+ "active_token_id" TEXT NOT NULL,
+ "revoke_at" TIMESTAMP(3) NOT NULL,
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+ CONSTRAINT "LiteLLM_DeprecatedVerificationToken_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DeprecatedVerificationToken_token_key" ON "LiteLLM_DeprecatedVerificationToken"("token");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeprecatedVerificationToken_token_revoke_at_idx" ON "LiteLLM_DeprecatedVerificationToken"("token", "revoke_at");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DeprecatedVerificationToken_revoke_at_idx" ON "LiteLLM_DeprecatedVerificationToken"("revoke_at");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205091235_allow_team_guardrail_config/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205091235_allow_team_guardrail_config/migration.sql
index 000b96b3b87..53ef243f7d0 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205091235_allow_team_guardrail_config/migration.sql
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205091235_allow_team_guardrail_config/migration.sql
@@ -1,6 +1,6 @@
-- AlterTable
-ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false;
+ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false;
-- AlterTable
-ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false;
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "allow_team_guardrail_config" BOOLEAN NOT NULL DEFAULT false;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql
new file mode 100644
index 00000000000..6a3c3aa66db
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260205144610_add_soft_budget_to_team_table/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "soft_budget" DOUBLE PRECISION;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207093506_add_available_on_public_internet_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207093506_add_available_on_public_internet_to_mcp_servers/migration.sql
new file mode 100644
index 00000000000..9b969c1e3b7
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207093506_add_available_on_public_internet_to_mcp_servers/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "available_on_public_internet" BOOLEAN NOT NULL DEFAULT false;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207110613_add_soft_budget_to_deleted_teams_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207110613_add_soft_budget_to_deleted_teams_table/migration.sql
new file mode 100644
index 00000000000..229032e2b13
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260207110613_add_soft_budget_to_deleted_teams_table/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "soft_budget" DOUBLE PRECISION;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql
new file mode 100644
index 00000000000..92baf0d25e4
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260209085821_add_verificationtoken_indexes/migration.sql
@@ -0,0 +1,8 @@
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_VerificationToken_user_id_team_id_idx" ON "LiteLLM_VerificationToken"("user_id", "team_id");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_VerificationToken_team_id_idx" ON "LiteLLM_VerificationToken"("team_id");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_VerificationToken_budget_reset_at_expires_idx" ON "LiteLLM_VerificationToken"("budget_reset_at", "expires");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212103349_adjust_tags_policy_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212103349_adjust_tags_policy_table/migration.sql
new file mode 100644
index 00000000000..de03c010bf4
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212103349_adjust_tags_policy_table/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "tags" TEXT[] DEFAULT ARRAY[]::TEXT[];
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql
new file mode 100644
index 00000000000..b3ae05e202b
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql
@@ -0,0 +1,33 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
+
+-- AlterTable
+ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
+
+-- AlterTable
+ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
+
+-- AlterTable
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_AccessGroupTable" (
+ "access_group_id" TEXT NOT NULL,
+ "access_group_name" TEXT NOT NULL,
+ "description" TEXT,
+ "access_model_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ "access_mcp_server_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ "access_agent_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ "assigned_team_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ "assigned_key_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "created_by" TEXT,
+ "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "updated_by" TEXT,
+
+ CONSTRAINT "LiteLLM_AccessGroupTable_pkey" PRIMARY KEY ("access_group_id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AccessGroupTable_access_group_name_key" ON "LiteLLM_AccessGroupTable"("access_group_name");
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213105436_add_managed_vector_store_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213105436_add_managed_vector_store_table/migration.sql
new file mode 100644
index 00000000000..eb3dd90602d
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213105436_add_managed_vector_store_table/migration.sql
@@ -0,0 +1,22 @@
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedVectorStoreTable" (
+ "id" TEXT NOT NULL,
+ "unified_resource_id" TEXT NOT NULL,
+ "resource_object" JSONB,
+ "model_mappings" JSONB NOT NULL,
+ "flat_model_resource_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ "storage_backend" TEXT,
+ "storage_url" TEXT,
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "created_by" TEXT,
+ "updated_at" TIMESTAMP(3) NOT NULL,
+ "updated_by" TEXT,
+
+ CONSTRAINT "LiteLLM_ManagedVectorStoreTable_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoreTable_unified_resource_id_key" ON "LiteLLM_ManagedVectorStoreTable"("unified_resource_id");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_ManagedVectorStoreTable_unified_resource_id_idx" ON "LiteLLM_ManagedVectorStoreTable"("unified_resource_id");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213170952_access_group_change_to_model_name/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213170952_access_group_change_to_model_name/migration.sql
new file mode 100644
index 00000000000..921bee97752
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260213170952_access_group_change_to_model_name/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_AccessGroupTable" DROP COLUMN IF EXISTS "access_model_ids",
+ADD COLUMN IF NOT EXISTS "access_model_names" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214094754_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214094754_schema_sync/migration.sql
new file mode 100644
index 00000000000..31e4320bb17
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214094754_schema_sync/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN IF NOT EXISTS "team_id" TEXT;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214163027_add_pipeline_to_policy_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214163027_add_pipeline_to_policy_table/migration.sql
new file mode 100644
index 00000000000..bb121466c9a
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214163027_add_pipeline_to_policy_table/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_PolicyTable" ADD COLUMN IF NOT EXISTS "pipeline" JSONB;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql
new file mode 100644
index 00000000000..4b909471e02
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260214185341_object_permissions_for_end_users/migration.sql
@@ -0,0 +1,11 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_EndUserTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT;
+
+-- AddForeignKey
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_EndUserTable_object_permission_id_fkey') THEN
+ ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
+
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..bfd50fe33be
--- /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 IF NOT EXISTS "last_active" TIMESTAMP(3);
+
+-- AlterTable
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "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..ef1ae258562
--- /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 IF NOT EXISTS "project_id" TEXT;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219181415_baseline_diff/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219181415_baseline_diff/migration.sql
new file mode 100644
index 00000000000..13be156db8e
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260219181415_baseline_diff/migration.sql
@@ -0,0 +1,60 @@
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGuardrailMetrics" (
+ "guardrail_id" TEXT NOT NULL,
+ "date" TEXT NOT NULL,
+ "requests_evaluated" BIGINT NOT NULL DEFAULT 0,
+ "passed_count" BIGINT NOT NULL DEFAULT 0,
+ "blocked_count" BIGINT NOT NULL DEFAULT 0,
+ "flagged_count" BIGINT NOT NULL DEFAULT 0,
+ "avg_score" DOUBLE PRECISION,
+ "avg_latency_ms" DOUBLE PRECISION,
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "updated_at" TIMESTAMP(3) NOT NULL,
+
+ CONSTRAINT "LiteLLM_DailyGuardrailMetrics_pkey" PRIMARY KEY ("guardrail_id","date")
+);
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_DailyPolicyMetrics" (
+ "policy_id" TEXT NOT NULL,
+ "date" TEXT NOT NULL,
+ "requests_evaluated" BIGINT NOT NULL DEFAULT 0,
+ "passed_count" BIGINT NOT NULL DEFAULT 0,
+ "blocked_count" BIGINT NOT NULL DEFAULT 0,
+ "flagged_count" BIGINT NOT NULL DEFAULT 0,
+ "avg_score" DOUBLE PRECISION,
+ "avg_latency_ms" DOUBLE PRECISION,
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "updated_at" TIMESTAMP(3) NOT NULL,
+
+ CONSTRAINT "LiteLLM_DailyPolicyMetrics_pkey" PRIMARY KEY ("policy_id","date")
+);
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogGuardrailIndex" (
+ "request_id" TEXT NOT NULL,
+ "guardrail_id" TEXT NOT NULL,
+ "policy_id" TEXT,
+ "start_time" TIMESTAMP(3) NOT NULL,
+
+ CONSTRAINT "LiteLLM_SpendLogGuardrailIndex_pkey" PRIMARY KEY ("request_id","guardrail_id")
+);
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGuardrailMetrics_date_idx" ON "LiteLLM_DailyGuardrailMetrics"("date");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGuardrailMetrics_guardrail_id_idx" ON "LiteLLM_DailyGuardrailMetrics"("guardrail_id");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyPolicyMetrics_date_idx" ON "LiteLLM_DailyPolicyMetrics"("date");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyPolicyMetrics_policy_id_idx" ON "LiteLLM_DailyPolicyMetrics"("policy_id");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogGuardrailIndex_guardrail_id_start_time_idx" ON "LiteLLM_SpendLogGuardrailIndex"("guardrail_id", "start_time");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogGuardrailIndex_policy_id_start_time_idx" ON "LiteLLM_SpendLogGuardrailIndex"("policy_id", "start_time");
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220124742_add_spec_path_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220124742_add_spec_path_to_mcp_servers/migration.sql
new file mode 100644
index 00000000000..55c12d9476f
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220124742_add_spec_path_to_mcp_servers/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "spec_path" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220153844_add_composite_index_aggregate_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220153844_add_composite_index_aggregate_tables/migration.sql
new file mode 100644
index 00000000000..a95a3c5c9c5
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260220153844_add_composite_index_aggregate_tables/migration.sql
@@ -0,0 +1,36 @@
+-- DropIndex
+DROP INDEX IF EXISTS "LiteLLM_DailyAgentSpend_agent_id_idx";
+
+-- DropIndex
+DROP INDEX IF EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_idx";
+
+-- DropIndex
+DROP INDEX IF EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_idx";
+
+-- DropIndex
+DROP INDEX IF EXISTS "LiteLLM_DailyTagSpend_tag_idx";
+
+-- DropIndex
+DROP INDEX IF EXISTS "LiteLLM_DailyTeamSpend_team_id_idx";
+
+-- DropIndex
+DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_idx";
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_agent_id_date_idx" ON "LiteLLM_DailyAgentSpend"("agent_id", "date");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_end_user_id_date_idx" ON "LiteLLM_DailyEndUserSpend"("end_user_id", "date");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_organization_id_date_idx" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_tag_date_idx" ON "LiteLLM_DailyTagSpend"("tag", "date");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_team_id_date_idx" ON "LiteLLM_DailyTeamSpend"("team_id", "date");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_user_id_date_idx" ON "LiteLLM_DailyUserSpend"("user_id", "date");
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221000000_ensure_project_id_verification_token/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221000000_ensure_project_id_verification_token/migration.sql
new file mode 100644
index 00000000000..697928c85d2
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221000000_ensure_project_id_verification_token/migration.sql
@@ -0,0 +1,5 @@
+-- Ensure project_id column exists in LiteLLM_VerificationToken.
+-- The original migration (20251113000000_add_project_table) adds this column,
+-- but if it failed partway through (e.g. LiteLLM_ProjectTable already existed)
+-- and was resolved as idempotent, the ALTER TABLE step may have been skipped.
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "project_id" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221183800_add_policy_versioning/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221183800_add_policy_versioning/migration.sql
new file mode 100644
index 00000000000..67647da66b0
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260221183800_add_policy_versioning/migration.sql
@@ -0,0 +1,17 @@
+-- DropIndex
+DROP INDEX IF EXISTS "LiteLLM_PolicyTable_policy_name_key";
+
+-- AlterTable
+ALTER TABLE "LiteLLM_PolicyTable" ADD COLUMN IF NOT EXISTS "is_latest" BOOLEAN NOT NULL DEFAULT true,
+ADD COLUMN IF NOT EXISTS "parent_version_id" TEXT,
+ADD COLUMN IF NOT EXISTS "production_at" TIMESTAMP(3),
+ADD COLUMN IF NOT EXISTS "published_at" TIMESTAMP(3),
+ADD COLUMN IF NOT EXISTS "version_number" INTEGER NOT NULL DEFAULT 1,
+ADD COLUMN IF NOT EXISTS "version_status" TEXT NOT NULL DEFAULT 'production';
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_PolicyTable_policy_name_version_status_idx" ON "LiteLLM_PolicyTable"("policy_name", "version_status");
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_PolicyTable_policy_name_version_number_key" ON "LiteLLM_PolicyTable"("policy_name", "version_number");
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql
new file mode 100644
index 00000000000..361a4705e62
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql
@@ -0,0 +1,3 @@
+-- Add batch_processed column to LiteLLM_ManagedObjectTable
+-- Set to true by CheckBatchCost after cost has been computed for a completed batch
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "batch_processed" BOOLEAN NOT NULL DEFAULT false;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224201417_spend_logs_request_duration/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224201417_spend_logs_request_duration/migration.sql
new file mode 100644
index 00000000000..aa0237eae1f
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224201417_spend_logs_request_duration/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "request_duration_ms" INTEGER;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql
new file mode 100644
index 00000000000..ecfe47a44d3
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260224203854_add_agent_object_permissions_table/migration.sql
@@ -0,0 +1,45 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "object_permission_id" TEXT;
+
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN IF EXISTS "spec_path";
+
+-- AlterTable
+ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "agent_id" TEXT;
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_ToolTable" (
+ "tool_id" TEXT NOT NULL,
+ "tool_name" TEXT NOT NULL,
+ "origin" TEXT,
+ "call_policy" TEXT NOT NULL DEFAULT 'untrusted',
+ "call_count" INTEGER NOT NULL DEFAULT 0,
+ "assignments" JSONB DEFAULT '{}',
+ "key_hash" TEXT,
+ "team_id" TEXT,
+ "key_alias" TEXT,
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "created_by" TEXT,
+ "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "updated_by" TEXT,
+
+ CONSTRAINT "LiteLLM_ToolTable_pkey" PRIMARY KEY ("tool_id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ToolTable_tool_name_key" ON "LiteLLM_ToolTable"("tool_name");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_ToolTable_call_policy_idx" ON "LiteLLM_ToolTable"("call_policy");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_ToolTable_team_id_idx" ON "LiteLLM_ToolTable"("team_id");
+
+-- AddForeignKey
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_AgentsTable_object_permission_id_fkey') THEN
+ ALTER TABLE "LiteLLM_AgentsTable" ADD CONSTRAINT "LiteLLM_AgentsTable_object_permission_id_fkey" FOREIGN KEY ("object_permission_id") REFERENCES "LiteLLM_ObjectPermissionTable"("object_permission_id") ON DELETE SET NULL ON UPDATE CASCADE;
+ END IF;
+END $$;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql
new file mode 100644
index 00000000000..24724cb18e0
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226000000_add_blocked_tools_to_object_permission/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "blocked_tools" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql
new file mode 100644
index 00000000000..d9c234696c8
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226120000_add_spend_log_tool_index/migration.sql
@@ -0,0 +1,11 @@
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogToolIndex" (
+ "request_id" TEXT NOT NULL,
+ "tool_name" TEXT NOT NULL,
+ "start_time" TIMESTAMP(3) NOT NULL,
+
+ CONSTRAINT "LiteLLM_SpendLogToolIndex_pkey" PRIMARY KEY ("request_id","tool_name")
+);
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_tool_name_start_time_idx" ON "LiteLLM_SpendLogToolIndex"("tool_name", "start_time");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql
new file mode 100644
index 00000000000..a41160dac6a
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260226202727_add_agent_id_to_delete_keys/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "agent_id" TEXT;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql
new file mode 100644
index 00000000000..7fa354d9d80
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228000000_add_claude_code_plugin_table/migration.sql
@@ -0,0 +1,18 @@
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginTable" (
+ "id" TEXT NOT NULL,
+ "name" TEXT NOT NULL,
+ "version" TEXT,
+ "description" TEXT,
+ "manifest_json" TEXT,
+ "files_json" TEXT DEFAULT '{}',
+ "enabled" BOOLEAN NOT NULL DEFAULT true,
+ "created_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP,
+ "updated_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP,
+ "created_by" TEXT,
+
+ CONSTRAINT "LiteLLM_ClaudeCodePluginTable_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginTable_name_key" ON "LiteLLM_ClaudeCodePluginTable"("name");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql
new file mode 100644
index 00000000000..66d792fe386
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228100000_add_spend_logs_composite_index/migration.sql
@@ -0,0 +1,2 @@
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" ON "LiteLLM_SpendLogs"("startTime", "request_id");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228110000_mcp_default_public_internet_true/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228110000_mcp_default_public_internet_true/migration.sql
new file mode 100644
index 00000000000..dd286464141
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228110000_mcp_default_public_internet_true/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" ALTER COLUMN "available_on_public_internet" SET DEFAULT true;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228170127_support_team_based_guardrails/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228170127_support_team_based_guardrails/migration.sql
new file mode 100644
index 00000000000..44d079ad194
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260228170127_support_team_based_guardrails/migration.sql
@@ -0,0 +1,8 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN IF NOT EXISTS "reviewed_at" TIMESTAMP(3),
+ADD COLUMN IF NOT EXISTS "status" TEXT NOT NULL DEFAULT 'active',
+ADD COLUMN IF NOT EXISTS "submitted_at" TIMESTAMP(3);
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_GuardrailsTable_status_idx" ON "LiteLLM_GuardrailsTable"("status");
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260303000000_update_tool_table_policies/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260303000000_update_tool_table_policies/migration.sql
new file mode 100644
index 00000000000..7aa329c6bb1
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260303000000_update_tool_table_policies/migration.sql
@@ -0,0 +1,25 @@
+-- Rename call_policy to input_policy (only if the old name still exists)
+DO $$
+BEGIN
+ IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'LiteLLM_ToolTable' AND column_name = 'call_policy') THEN
+ ALTER TABLE "LiteLLM_ToolTable" RENAME COLUMN "call_policy" TO "input_policy";
+ END IF;
+END $$;
+
+-- Add output_policy column
+ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN IF NOT EXISTS "output_policy" TEXT NOT NULL DEFAULT 'untrusted';
+
+-- Add user_agent column
+ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN IF NOT EXISTS "user_agent" TEXT;
+
+-- Add last_used_at column
+ALTER TABLE "LiteLLM_ToolTable" ADD COLUMN IF NOT EXISTS "last_used_at" TIMESTAMP(3);
+
+-- Drop old index on call_policy
+DROP INDEX IF EXISTS "LiteLLM_ToolTable_call_policy_idx";
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_ToolTable_input_policy_idx" ON "LiteLLM_ToolTable"("input_policy");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_ToolTable_output_policy_idx" ON "LiteLLM_ToolTable"("output_policy");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql
new file mode 100644
index 00000000000..a045b7d1d66
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql
new file mode 100644
index 00000000000..acb35baba96
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql
@@ -0,0 +1,5 @@
+-- Add static_headers and extra_headers to LiteLLM_AgentsTable
+
+ALTER TABLE "LiteLLM_AgentsTable"
+ ADD COLUMN IF NOT EXISTS "static_headers" JSONB DEFAULT '{}',
+ ADD COLUMN IF NOT EXISTS "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[];
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql
new file mode 100644
index 00000000000..c1556822c1a
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql
@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "tpm_limit" INTEGER;
+ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "rpm_limit" INTEGER;
+ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "session_tpm_limit" INTEGER;
+ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "session_rpm_limit" INTEGER;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306175056_add_configs_override_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306175056_add_configs_override_table/migration.sql
new file mode 100644
index 00000000000..616463a9e2e
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306175056_add_configs_override_table/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "spec_path" TEXT;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306233848_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306233848_schema_sync/migration.sql
new file mode 100644
index 00000000000..4f3c0b7485b
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306233848_schema_sync/migration.sql
@@ -0,0 +1,62 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "byok_api_key_help_url" TEXT,
+ADD COLUMN IF NOT EXISTS "byok_description" TEXT[] DEFAULT ARRAY[]::TEXT[],
+ADD COLUMN IF NOT EXISTS "is_byok" BOOLEAN NOT NULL DEFAULT false,
+ADD COLUMN IF NOT EXISTS "tool_name_to_description" JSONB DEFAULT '{}',
+ADD COLUMN IF NOT EXISTS "tool_name_to_display_name" JSONB DEFAULT '{}';
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_MCPUserCredentials" (
+ "id" TEXT NOT NULL,
+ "user_id" TEXT NOT NULL,
+ "server_id" TEXT NOT NULL,
+ "credential_b64" TEXT NOT NULL,
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+ CONSTRAINT "LiteLLM_MCPUserCredentials_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_JWTKeyMapping" (
+ "id" TEXT NOT NULL,
+ "jwt_claim_name" TEXT NOT NULL,
+ "jwt_claim_value" TEXT NOT NULL,
+ "token" TEXT NOT NULL,
+ "description" TEXT,
+ "is_active" BOOLEAN NOT NULL DEFAULT true,
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "created_by" TEXT,
+ "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "updated_by" TEXT,
+
+ CONSTRAINT "LiteLLM_JWTKeyMapping_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_ConfigOverrides" (
+ "config_type" TEXT NOT NULL,
+ "config_value" JSONB NOT NULL,
+ "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+ "updated_at" TIMESTAMP(3) NOT NULL,
+
+ CONSTRAINT "LiteLLM_ConfigOverrides_pkey" PRIMARY KEY ("config_type")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MCPUserCredentials_user_id_server_id_key" ON "LiteLLM_MCPUserCredentials"("user_id", "server_id");
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value", "is_active");
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value");
+
+-- AddForeignKey
+DO $$
+BEGIN
+ IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
+ ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE RESTRICT ON UPDATE CASCADE;
+ END IF;
+END $$;
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000000_add_mcp_approval_status/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000000_add_mcp_approval_status/migration.sql
new file mode 100644
index 00000000000..184caef0809
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000000_add_mcp_approval_status/migration.sql
@@ -0,0 +1,11 @@
+-- AlterTable: Add BYOM approval workflow fields to LiteLLM_MCPServerTable
+ALTER TABLE "LiteLLM_MCPServerTable"
+ ADD COLUMN IF NOT EXISTS "approval_status" TEXT DEFAULT 'active',
+ ADD COLUMN IF NOT EXISTS "submitted_by" TEXT,
+ ADD COLUMN IF NOT EXISTS "submitted_at" TIMESTAMP(3),
+ ADD COLUMN IF NOT EXISTS "reviewed_at" TIMESTAMP(3),
+ ADD COLUMN IF NOT EXISTS "review_notes" TEXT;
+
+-- CreateIndex
+CREATE INDEX IF NOT EXISTS "LiteLLM_MCPServerTable_approval_status_idx"
+ ON "LiteLLM_MCPServerTable"("approval_status");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000001_add_mcp_source_url/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000001_add_mcp_source_url/migration.sql
new file mode 100644
index 00000000000..dc468b82061
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260309000001_add_mcp_source_url/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable: Add source_url field to LiteLLM_MCPServerTable for GitHub/docs link
+ALTER TABLE "LiteLLM_MCPServerTable"
+ ADD COLUMN IF NOT EXISTS "source_url" TEXT;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql
new file mode 100644
index 00000000000..84eb70ce097
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260311180521_schema_sync/migration.sql
@@ -0,0 +1,11 @@
+-- DropIndex
+DROP INDEX IF EXISTS "LiteLLM_MCPServerTable_approval_status_idx";
+
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" DROP COLUMN IF EXISTS "approval_status",
+DROP COLUMN IF EXISTS "review_notes",
+DROP COLUMN IF EXISTS "reviewed_at",
+DROP COLUMN IF EXISTS "source_url",
+DROP COLUMN IF EXISTS "submitted_at",
+DROP COLUMN IF EXISTS "submitted_by";
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql
new file mode 100644
index 00000000000..cc48a742f20
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260312124619_schema_sync/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "models" TEXT[] DEFAULT ARRAY[]::TEXT[];
+
diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
index dc49036cb15..ce79c2b3d52 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
@@ -62,7 +63,16 @@ model LiteLLM_AgentsTable {
agent_name String @unique
litellm_params Json?
agent_card_params Json
+ static_headers Json? @default("{}")
+ extra_headers String[] @default([])
agent_access_groups String[] @default([])
+ object_permission_id String?
+ object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
+ spend Float @default(0.0)
+ tpm_limit Int?
+ rpm_limit Int?
+ session_tpm_limit Int?
+ session_rpm_limit Int?
created_at DateTime @default(now()) @map("created_at")
created_by String
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
@@ -113,6 +123,7 @@ model LiteLLM_TeamTable {
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?
@@ -127,12 +138,41 @@ model LiteLLM_TeamTable {
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? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
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
@@ -147,6 +187,7 @@ model LiteLLM_DeletedTeamTable {
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?
@@ -159,6 +200,7 @@ model LiteLLM_DeletedTeamTable {
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)
@@ -225,10 +267,15 @@ model LiteLLM_ObjectPermissionTable {
vector_stores String[] @default([])
agents String[] @default([])
agent_access_groups String[] @default([])
+ models String[] @default([])
+ blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission
teams LiteLLM_TeamTable[]
+ projects LiteLLM_ProjectTable[]
verification_tokens LiteLLM_VerificationToken[]
organizations LiteLLM_OrganizationTable[]
users LiteLLM_UserTable[]
+ end_users LiteLLM_EndUserTable[]
+ agents_table LiteLLM_AgentsTable[]
}
// Holds the MCP server configuration
@@ -238,6 +285,7 @@ model LiteLLM_MCPServerTable {
alias String?
description String?
url String?
+ spec_path String?
transport String @default("sse")
auth_type String?
credentials Json? @default("{}")
@@ -248,6 +296,8 @@ model LiteLLM_MCPServerTable {
mcp_info Json? @default("{}")
mcp_access_groups String[]
allowed_tools String[] @default([])
+ tool_name_to_display_name Json? @default("{}")
+ tool_name_to_description Json? @default("{}")
extra_headers String[] @default([])
static_headers Json? @default("{}")
// Health check status
@@ -262,6 +312,27 @@ model LiteLLM_MCPServerTable {
token_url String?
registration_url String?
allow_all_keys Boolean @default(false)
+ available_on_public_internet Boolean @default(true)
+ is_byok Boolean @default(false)
+ byok_description String[] @default([])
+ byok_api_key_help_url String?
+ approval_status String @default("approved")
+ submitted_by String?
+ submitted_at DateTime?
+ reviewed_at DateTime?
+ review_notes String?
+}
+
+// Per-user BYOK credentials for MCP servers
+model LiteLLM_MCPUserCredentials {
+ id String @id @default(uuid())
+ user_id String
+ server_id String
+ credential_b64 String
+ created_at DateTime @default(now()) @map("created_at")
+ updated_at DateTime @default(now()) @updatedAt @map("updated_at")
+
+ @@unique([user_id, server_id])
}
// Generate Tokens for Proxy
@@ -278,6 +349,8 @@ model LiteLLM_VerificationToken {
router_settings Json? @default("{}")
user_id String?
team_id String?
+ agent_id String?
+ project_id String?
permissions Json @default("{}")
max_parallel_requests Int?
metadata Json @default("{}")
@@ -290,6 +363,7 @@ model LiteLLM_VerificationToken {
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
policies String[] @default([])
+ access_group_ids String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
budget_id String?
@@ -299,6 +373,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")
@@ -306,7 +381,50 @@ 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])
+ jwt_key_mappings LiteLLM_JWTKeyMapping[]
+
+ // 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"
+ // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2
+ @@index([user_id, team_id])
+
+ // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2
+ @@index([team_id])
+
+ // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3
+ @@index([budget_reset_at, expires])
+}
+
+model LiteLLM_JWTKeyMapping {
+ id String @id @default(uuid())
+ jwt_claim_name String // e.g. "sub", "email"
+ jwt_claim_value String // The claim value to match
+ token String // Hashed virtual key (FK)
+ description String?
+ is_active Boolean @default(true)
+ created_at DateTime @default(now())
+ created_by String?
+ updated_at DateTime @default(now()) @updatedAt
+ updated_by String?
+
+ litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
+
+ @@unique([jwt_claim_name, jwt_claim_value])
+ @@index([jwt_claim_name, jwt_claim_value, is_active])
+}
+
+// Deprecated keys during grace period - allows old key to work until revoke_at
+model LiteLLM_DeprecatedVerificationToken {
+ id String @id @default(uuid())
+ token String // Hashed old key
+ active_token_id String // Current token hash in LiteLLM_VerificationToken
+ revoke_at DateTime // When the old key stops working
+ created_at DateTime @default(now()) @map("created_at")
+
+ @@unique([token])
+ @@index([token, revoke_at])
+ @@index([revoke_at])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
@@ -323,6 +441,8 @@ model LiteLLM_DeletedVerificationToken {
config Json @default("{}")
user_id String?
team_id String?
+ agent_id String?
+ project_id String?
permissions Json @default("{}")
max_parallel_requests Int?
metadata Json @default("{}")
@@ -335,6 +455,7 @@ model LiteLLM_DeletedVerificationToken {
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
policies String[] @default([])
+ access_group_ids String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
@@ -345,6 +466,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?
@@ -373,7 +495,9 @@ model LiteLLM_EndUserTable {
allowed_model_region String? // require all user requests to use models in this specific region
default_model String? // use along with 'allowed_model_region'. if no available model in region, default to this model.
budget_id String?
+ object_permission_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
+ object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
blocked Boolean @default(false)
}
@@ -408,6 +532,7 @@ model LiteLLM_SpendLogs {
completion_tokens Int @default(0)
startTime DateTime // Assuming start_time is a DateTime field
endTime DateTime // Assuming end_time is a DateTime field
+ request_duration_ms Int?
completionStartTime DateTime? // Assuming completionStartTime is a DateTime field
model String @default("")
model_id String? @default("") // the model id stored in proxy model db
@@ -415,7 +540,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("[]")
@@ -431,6 +556,7 @@ model LiteLLM_SpendLogs {
agent_id String?
proxy_server_request Json? @default("{}")
@@index([startTime])
+ @@index([startTime, request_id])
@@index([end_user])
@@index([session_id])
}
@@ -545,7 +671,7 @@ model LiteLLM_DailyUserSpend {
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
- @@index([user_id])
+ @@index([user_id, date])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@ -576,7 +702,7 @@ model LiteLLM_DailyOrganizationSpend {
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
- @@index([organization_id])
+ @@index([organization_id, date])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@ -606,7 +732,7 @@ model LiteLLM_DailyEndUserSpend {
updated_at DateTime @updatedAt
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
- @@index([end_user_id])
+ @@index([end_user_id, date])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@ -636,7 +762,7 @@ model LiteLLM_DailyAgentSpend {
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
- @@index([agent_id])
+ @@index([agent_id, date])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@ -667,7 +793,7 @@ model LiteLLM_DailyTeamSpend {
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
- @@index([team_id])
+ @@index([team_id, date])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@ -699,7 +825,7 @@ model LiteLLM_DailyTagSpend {
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
- @@index([tag])
+ @@index([tag, date])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
@@ -744,6 +870,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
file_object Json // Stores the OpenAIFileObject
file_purpose String // either 'batch' or 'fine-tune'
status String? // check if batch cost has been tracked
+ batch_processed Boolean @default(false) // set to true by CheckBatchCost after cost is computed
created_at DateTime @default(now())
created_by String?
updated_at DateTime @updatedAt
@@ -753,6 +880,22 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
@@index([model_object_id])
}
+model LiteLLM_ManagedVectorStoreTable {
+ id String @id @default(uuid())
+ unified_resource_id String @unique // The base64 encoded unified vector store ID
+ resource_object Json? // Stores the VectorStoreCreateResponse
+ model_mappings Json // Maps model_id -> provider_vector_store_id
+ flat_model_resource_ids String[] @default([]) // Flat list of provider vector store IDs for faster querying
+ storage_backend String? // Storage backend name (if applicable)
+ storage_url String? // Storage URL (if applicable)
+ created_at DateTime @default(now())
+ created_by String?
+ updated_at DateTime @updatedAt
+ updated_by String?
+
+ @@index([unified_resource_id])
+}
+
model LiteLLM_ManagedVectorStoresTable {
vector_store_id String @id
custom_llm_provider String
@@ -779,6 +922,71 @@ model LiteLLM_GuardrailsTable {
team_id String?
created_at DateTime @default(now())
updated_at DateTime @updatedAt
+ // Submission lifecycle. Possible values: pending_review (team-registered, awaiting approval), active (approved), rejected
+ status String @default("active")
+ submitted_at DateTime?
+ reviewed_at DateTime?
+ // submitted_by_user_id and submitted_by_email live in guardrail_info JSON
+
+ @@index([status])
+}
+
+// Daily guardrail metrics for usage dashboard (one row per guardrail per day)
+model LiteLLM_DailyGuardrailMetrics {
+ guardrail_id String // logical id; may not FK if guardrail from config
+ date String // YYYY-MM-DD
+ requests_evaluated BigInt @default(0)
+ passed_count BigInt @default(0)
+ blocked_count BigInt @default(0)
+ flagged_count BigInt @default(0)
+ avg_score Float?
+ avg_latency_ms Float?
+ created_at DateTime @default(now())
+ updated_at DateTime @updatedAt
+
+ @@id([guardrail_id, date])
+ @@index([date])
+ @@index([guardrail_id])
+}
+
+// Daily policy metrics for usage dashboard (one row per policy per day)
+model LiteLLM_DailyPolicyMetrics {
+ policy_id String
+ date String // YYYY-MM-DD
+ requests_evaluated BigInt @default(0)
+ passed_count BigInt @default(0)
+ blocked_count BigInt @default(0)
+ flagged_count BigInt @default(0)
+ avg_score Float?
+ avg_latency_ms Float?
+ created_at DateTime @default(now())
+ updated_at DateTime @updatedAt
+
+ @@id([policy_id, date])
+ @@index([date])
+ @@index([policy_id])
+}
+
+// Index for fast "last N logs for guardrail/policy" from SpendLogs
+model LiteLLM_SpendLogGuardrailIndex {
+ request_id String
+ guardrail_id String
+ policy_id String? // set when run as part of a policy pipeline
+ start_time DateTime
+
+ @@id([request_id, guardrail_id])
+ @@index([guardrail_id, start_time])
+ @@index([policy_id, start_time])
+}
+
+// Index for fast "last N logs for tool" from SpendLogs – see how a tool is called in production
+model LiteLLM_SpendLogToolIndex {
+ request_id String
+ tool_name String // matches LiteLLM_ToolTable.tool_name; join for input_policy/output_policy etc.
+ start_time DateTime
+
+ @@id([request_id, tool_name])
+ @@index([tool_name, start_time])
}
// Prompt table for storing prompt configurations
@@ -860,6 +1068,14 @@ model LiteLLM_UISettings {
updated_at DateTime @updatedAt
}
+// Generic config overrides table - one row per config_type
+model LiteLLM_ConfigOverrides {
+ config_type String @id
+ config_value Json
+ created_at DateTime @default(now())
+ updated_at DateTime @updatedAt
+}
+
// Skills table for storing LiteLLM-managed skills
model LiteLLM_SkillsTable {
skill_id String @id @default(uuid())
@@ -878,19 +1094,29 @@ model LiteLLM_SkillsTable {
updated_by String?
}
-// Policy table for storing guardrail policies
+// Policy table for storing guardrail policies (versioned)
model LiteLLM_PolicyTable {
- policy_id String @id @default(uuid())
- policy_name String @unique
- inherit String? // Name of parent policy to inherit from
- description String?
- guardrails_add String[] @default([])
- guardrails_remove String[] @default([])
- condition Json? @default("{}") // Policy conditions (e.g., model matching)
- created_at DateTime @default(now())
- created_by String?
- updated_at DateTime @default(now()) @updatedAt
- updated_by String?
+ policy_id String @id @default(uuid())
+ policy_name String // No longer @unique; use @@unique([policy_name, version_number])
+ version_number Int @default(1)
+ version_status String @default("production") // "draft" | "published" | "production"
+ parent_version_id String?
+ is_latest Boolean @default(true)
+ published_at DateTime?
+ production_at DateTime?
+ inherit String? // Name of parent policy to inherit from
+ description String?
+ guardrails_add String[] @default([])
+ guardrails_remove String[] @default([])
+ condition Json? @default("{}") // Policy conditions (e.g., model matching)
+ pipeline Json? // Optional guardrail pipeline (mode + steps[])
+ created_at DateTime @default(now())
+ created_by String?
+ updated_at DateTime @default(now()) @updatedAt
+ updated_by String?
+
+ @@unique([policy_name, version_number])
+ @@index([policy_name, version_status])
}
// Policy attachment table for defining where policies apply
@@ -901,8 +1127,69 @@ model LiteLLM_PolicyAttachmentTable {
teams String[] @default([]) // Team aliases or patterns
keys String[] @default([]) // Key aliases or patterns
models String[] @default([]) // Model names or patterns
+ tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
+
+// Global tool registry - auto-discovered from LLM responses; admins set input/output policies here
+model LiteLLM_ToolTable {
+ tool_id String @id @default(uuid())
+ tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space"
+ origin String? // MCP server name or "user_defined"
+ input_policy String @default("untrusted") // "trusted" | "untrusted" | "blocked"
+ output_policy String @default("untrusted") // "trusted" | "untrusted"
+ call_count Int @default(0) // cumulative number of times this tool was seen
+ assignments Json? @default("{}")
+ key_hash String? // hash of the virtual key that first called this tool
+ team_id String? // team that first called this tool
+ key_alias String? // human-readable alias of the virtual key
+ user_agent String? // user-agent of the first request that discovered this tool
+ last_used_at DateTime? // timestamp of the most recent call
+ created_at DateTime @default(now())
+ created_by String?
+ updated_at DateTime @default(now()) @updatedAt
+ updated_by String?
+
+ @@index([input_policy])
+ @@index([output_policy])
+ @@index([team_id])
+}
+
+// Per-(tool, team/key) policy overrides. When present, override replaces global tool policy for that scope.
+//Unified Access Groups table for storing unified access groups
+model LiteLLM_AccessGroupTable {
+ access_group_id String @id @default(uuid())
+ access_group_name String @unique
+ description String?
+
+ // Resource memberships - explicit arrays per type
+ access_model_names String[] @default([])
+ access_mcp_server_ids String[] @default([])
+ access_agent_ids String[] @default([])
+
+ assigned_team_ids String[] @default([])
+ assigned_key_ids String[] @default([])
+
+ created_at DateTime @default(now())
+ created_by String?
+ updated_at DateTime @default(now()) @updatedAt
+ updated_by String?
+}
+// Claude Code Plugin Marketplace table
+model LiteLLM_ClaudeCodePluginTable {
+ id String @id @default(uuid())
+ name String @unique
+ version String?
+ description String?
+ manifest_json String?
+ files_json String? @default("{}")
+ enabled Boolean @default(true)
+ created_at DateTime? @default(now())
+ updated_at DateTime? @default(now()) @updatedAt
+ created_by String?
+
+ @@map("LiteLLM_ClaudeCodePluginTable")
+}
diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py
index f3155722187..7eff0c00f75 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/utils.py
+++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py
@@ -415,16 +415,26 @@ class ProxyExtrasDBManager:
logger.info(
f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied"
)
- ProxyExtrasDBManager._roll_back_migration(
- failed_migration
- )
- ProxyExtrasDBManager._resolve_specific_migration(
- failed_migration
- )
- logger.info(
- f"✅ Migration {failed_migration} resolved."
- )
- return True
+ try:
+ ProxyExtrasDBManager._roll_back_migration(
+ failed_migration
+ )
+ except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as rollback_err:
+ logger.warning(
+ f"Failed to roll back migration {failed_migration}: {rollback_err}. "
+ f"It may already be in a rolled-back state."
+ )
+ try:
+ ProxyExtrasDBManager._resolve_specific_migration(
+ failed_migration
+ )
+ logger.info(
+ f"✅ Migration {failed_migration} resolved, retrying to apply remaining migrations"
+ )
+ except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as resolve_err:
+ logger.warning(
+ f"Failed to resolve migration {failed_migration}: {resolve_err}"
+ )
else:
logger.info(
f"Found failed migration: {failed_migration}, marking as rolled back"
@@ -514,20 +524,34 @@ class ProxyExtrasDBManager:
)
if migration_match:
migration_name = migration_match.group(1)
- logger.info(
- f"Rolling back migration {migration_name}"
- )
- ProxyExtrasDBManager._roll_back_migration(
- migration_name
- )
- logger.info(
- f"Resolving migration {migration_name} that failed "
- f"due to existing schema objects"
- )
- ProxyExtrasDBManager._resolve_specific_migration(
- migration_name
- )
- logger.info("✅ Migration resolved.")
+ try:
+ logger.info(
+ f"Rolling back migration {migration_name}"
+ )
+ ProxyExtrasDBManager._roll_back_migration(
+ migration_name
+ )
+ except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as rollback_err:
+ logger.warning(
+ f"Failed to roll back migration {migration_name}: {rollback_err}. "
+ f"It may already be in a rolled-back state."
+ )
+ try:
+ logger.info(
+ f"Resolving migration {migration_name} that failed "
+ f"due to existing schema objects"
+ )
+ ProxyExtrasDBManager._resolve_specific_migration(
+ migration_name
+ )
+ logger.info(
+ f"✅ Migration {migration_name} resolved, "
+ f"retrying to apply remaining migrations"
+ )
+ except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as resolve_err:
+ logger.warning(
+ f"Failed to resolve migration {migration_name}: {resolve_err}"
+ )
else:
# Unknown P3018 error - log and re-raise for safety
logger.warning(
diff --git a/litellm-proxy-extras/migration_runbook.md b/litellm-proxy-extras/migration_runbook.md
index 93948f24b13..3310b1626a8 100644
--- a/litellm-proxy-extras/migration_runbook.md
+++ b/litellm-proxy-extras/migration_runbook.md
@@ -2,7 +2,35 @@
This is a runbook for creating and running database migrations for the LiteLLM proxy. For use for litellm engineers only.
-## Quick Start
+## Step 0: Sync All `schema.prisma` Files
+
+Before doing anything else, make sure all `schema.prisma` files in the repo are in sync. There are multiple copies that must match:
+
+| File | Purpose |
+|------|---------|
+| `schema.prisma` (repo root) | Source of truth |
+| `litellm/proxy/schema.prisma` | Used by the proxy server |
+| `litellm-proxy-extras/litellm_proxy_extras/schema.prisma` | Used for migration generation |
+
+**Sync process:**
+
+```bash
+# 1. Diff all schema files against the root source of truth
+diff schema.prisma litellm/proxy/schema.prisma
+diff schema.prisma litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+
+# 2. If there are differences, copy the root schema to all locations
+cp schema.prisma litellm/proxy/schema.prisma
+cp schema.prisma litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+
+# 3. Verify all files are now identical
+diff schema.prisma litellm/proxy/schema.prisma && echo "proxy schema in sync" || echo "MISMATCH"
+diff schema.prisma litellm-proxy-extras/litellm_proxy_extras/schema.prisma && echo "extras schema in sync" || echo "MISMATCH"
+```
+
+> **Do NOT proceed to migration generation until all schema files are identical.**
+
+## Step 1: Quick Start — Generate Migration
```bash
# Install deps (one time)
@@ -43,8 +71,13 @@ rm -rf litellm-proxy-extras/litellm_proxy_extras/migrations/[empty_dir]
## Rules
-- Update `schema.prisma` first
+- Sync all `schema.prisma` files first (Step 0)
+- Update `schema.prisma` at the repo root first, then sync copies
- Review generated SQL before committing
- Use descriptive migration names
- Never edit existing migration files
- Commit schema + migration together
+
+---
+
+**Done with migration?** See [build_and_publish.md](./build_and_publish.md) to publish a new `litellm-proxy-extras` package.
diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml
index d43b591686c..b65dbe45233 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.30"
+version = "0.4.56"
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.30"
+version = "0.4.56"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",
diff --git a/litellm/__init__.py b/litellm/__init__.py
index 8174b9d2655..51c66838613 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -12,6 +12,13 @@ warnings.filterwarnings(
### INIT VARIABLES #########################
import threading
import os
+
+# Load .env before any other litellm imports so env vars (e.g. LITELLM_UI_SESSION_DURATION) are available
+import dotenv as _dotenv
+
+if os.getenv("LITELLM_MODE", "DEV") == "DEV":
+ _dotenv.load_dotenv()
+
from typing import (
Callable,
List,
@@ -74,12 +81,10 @@ from litellm.constants import (
DEFAULT_ALLOWED_FAILS,
)
import httpx
-import dotenv
+
# register_async_client_cleanup is lazy-loaded and called on first access
litellm_mode = os.getenv("LITELLM_MODE", "DEV") # "PRODUCTION", "DEV"
-if litellm_mode == "DEV":
- dotenv.load_dotenv()
####################################################
@@ -98,12 +103,14 @@ _custom_logger_compatible_callbacks_literal = Literal[
"openmeter",
"logfire",
"literalai",
+ "litellm_agent",
"dynamic_rate_limiter",
"dynamic_rate_limiter_v3",
"langsmith",
"prometheus",
"otel",
"datadog",
+ "datadog_metrics",
"datadog_llm_observability",
"galileo",
"braintrust",
@@ -137,6 +144,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"gitlab",
"cloudzero",
"focus",
+ "vantage",
"posthog",
"levo",
]
@@ -146,7 +154,9 @@ _known_custom_logger_compatible_callbacks: List = list(
get_args(_custom_logger_compatible_callbacks_literal)
)
callbacks: List[
- Union[Callable, _custom_logger_compatible_callbacks_literal, "CustomLogger"] # CustomLogger is lazy-loaded
+ Union[
+ Callable, _custom_logger_compatible_callbacks_literal, "CustomLogger"
+ ] # CustomLogger is lazy-loaded
] = []
callback_settings: Dict[str, Dict[str, Any]] = {}
initialized_langfuse_clients: int = 0
@@ -156,45 +166,57 @@ prometheus_initialize_budget_metrics: Optional[bool] = False
require_auth_for_metrics_endpoint: Optional[bool] = False
argilla_batch_size: Optional[int] = None
datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload.
-gcs_pub_sub_use_v1: Optional[bool] = (
- False # if you want to use v1 gcs pubsub logged payload
-)
-generic_api_use_v1: Optional[bool] = (
- False # if you want to use v1 generic api logged payload
-)
+gcs_pub_sub_use_v1: Optional[
+ bool
+] = False # if you want to use v1 gcs pubsub logged payload
+generic_api_use_v1: Optional[
+ bool
+] = False # if you want to use v1 generic api logged payload
argilla_transformation_object: Optional[Dict[str, Any]] = None
-_async_input_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded
+_async_input_callback: List[
+ Union[str, Callable, "CustomLogger"]
+] = ( # CustomLogger is lazy-loaded
[]
) # internal variable - async custom callbacks are routed here.
-_async_success_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded
+_async_success_callback: List[
+ Union[str, Callable, "CustomLogger"]
+] = ( # CustomLogger is lazy-loaded
[]
) # internal variable - async custom callbacks are routed here.
-_async_failure_callback: List[Union[str, Callable, "CustomLogger"]] = ( # CustomLogger is lazy-loaded
+_async_failure_callback: List[
+ Union[str, Callable, "CustomLogger"]
+] = ( # CustomLogger is lazy-loaded
[]
) # internal variable - async custom callbacks are routed here.
pre_call_rules: List[Callable] = []
post_call_rules: List[Callable] = []
turn_off_message_logging: Optional[bool] = False
+standard_logging_payload_excluded_fields: Optional[
+ List[str]
+] = None # Fields to exclude from StandardLoggingPayload before callbacks receive it
log_raw_request_response: bool = False
redact_messages_in_exceptions: Optional[bool] = False
redact_user_api_key_info: Optional[bool] = False
filter_invalid_headers: Optional[bool] = False
-add_user_information_to_llm_headers: Optional[bool] = (
- None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers
-)
+add_user_information_to_llm_headers: Optional[
+ bool
+] = None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers
store_audit_logs = False # Enterprise feature, allow users to see audit logs
### end of callbacks #############
-email: Optional[str] = (
- None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
-)
-token: Optional[str] = (
- None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
-)
+email: Optional[
+ str
+] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
+token: Optional[
+ str
+] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
telemetry = True
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False))
modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
+use_chat_completions_url_for_anthropic_messages: bool = bool(
+ os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)
+) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API
retry = True
### AUTH ###
api_key: Optional[str] = None
@@ -249,9 +271,9 @@ use_client: bool = False
ssl_verify: Union[str, bool] = True
ssl_security_level: Optional[str] = None
ssl_certificate: Optional[str] = None
-ssl_ecdh_curve: Optional[str] = (
- None # Set to 'X25519' to disable PQC and improve performance
-)
+ssl_ecdh_curve: Optional[
+ str
+] = None # Set to 'X25519' to disable PQC and improve performance
disable_streaming_logging: bool = False
disable_token_counter: bool = False
disable_add_transform_inline_image_block: bool = False
@@ -295,30 +317,29 @@ return_response_headers: bool = (
False # get response headers from LLM Api providers - example x-remaining-requests,
)
enable_json_schema_validation: bool = False
+enable_key_alias_format_validation: bool = (
+ False # opt-in validation of key_alias format on /key/generate and /key/update
+)
####################
logging: bool = True
enable_loadbalancing_on_batch_endpoints: Optional[bool] = None
enable_caching_on_provider_specific_optional_params: bool = (
False # feature-flag for caching on optional params - e.g. 'top_k'
)
-caching: bool = (
- False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
-)
-caching_with_models: bool = (
- False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
-)
-cache: Optional["Cache"] = (
- None # cache object <- use this - https://docs.litellm.ai/docs/caching
-)
+caching: bool = False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
+caching_with_models: bool = False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
+cache: Optional[
+ "Cache"
+] = None # cache object <- use this - https://docs.litellm.ai/docs/caching
default_in_memory_ttl: Optional[float] = None
default_redis_ttl: Optional[float] = None
default_redis_batch_cache_expiry: Optional[float] = None
model_alias_map: Dict[str, str] = {}
model_group_settings: Optional["ModelGroupSettings"] = None
max_budget: float = 0.0 # set the max budget across all providers
-budget_duration: Optional[str] = (
- None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
-)
+budget_duration: Optional[
+ str
+] = None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
default_soft_budget: float = (
DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0
)
@@ -327,9 +348,7 @@ forward_traceparent_to_llm_provider: bool = False
_current_cost = 0.0 # private variable, used if max budget is set
error_logs: Dict = {}
-add_function_to_prompt: bool = (
- False # if function calling not supported by api, append function call details to system prompt
-)
+add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt
client_session: Optional[httpx.Client] = None
aclient_session: Optional[httpx.AsyncClient] = None
model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks'
@@ -337,6 +356,14 @@ model_cost_map_url: str = os.getenv(
"LITELLM_MODEL_COST_MAP_URL",
"https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json",
)
+blog_posts_url: str = os.getenv(
+ "LITELLM_BLOG_POSTS_URL",
+ "https://docs.litellm.ai/blog/rss.xml",
+)
+anthropic_beta_headers_url: str = os.getenv(
+ "LITELLM_ANTHROPIC_BETA_HEADERS_URL",
+ "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/anthropic_beta_headers_config.json",
+)
suppress_debug_info = False
dynamodb_table_name: Optional[str] = None
s3_callback_params: Optional[Dict] = None
@@ -364,12 +391,11 @@ enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
custom_prometheus_metadata_labels: List[str] = []
custom_prometheus_tags: List[str] = []
prometheus_metrics_config: Optional[List] = None
+prometheus_emit_stream_label: bool = False
disable_add_prefix_to_prompt: bool = (
False # used by anthropic, to disable adding prefix to prompt
)
-disable_copilot_system_to_assistant: bool = (
- False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior.
-)
+disable_copilot_system_to_assistant: bool = False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior.
public_mcp_servers: Optional[List[str]] = None
public_model_groups: Optional[List[str]] = None
public_agent_groups: Optional[List[str]] = None
@@ -388,17 +414,14 @@ if TYPE_CHECKING:
######## Networking Settings ########
-use_aiohttp_transport: bool = (
- True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead.
-)
+use_aiohttp_transport: bool = True # Older variable, aiohttp is now the default. use disable_aiohttp_transport instead.
aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings
disable_aiohttp_transport: bool = False # Set this to true to use httpx instead
disable_aiohttp_trust_env: bool = (
False # When False, aiohttp will respect HTTP(S)_PROXY env vars
)
-force_ipv4: bool = (
- False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6.
-)
+force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6.
+network_mock: bool = False # When True, use mock transport — no real network calls
####### STOP SEQUENCE LIMIT #######
disable_stop_sequence_limit: bool = False # when True, stop sequence limit is disabled
@@ -412,13 +435,13 @@ context_window_fallbacks: Optional[List] = None
content_policy_fallbacks: Optional[List] = None
allowed_fails: int = 3
allow_dynamic_callback_disabling: bool = True
-num_retries_per_request: Optional[int] = (
- None # for the request overall (incl. fallbacks + model retries)
-)
+num_retries_per_request: Optional[
+ int
+] = None # for the request overall (incl. fallbacks + model retries)
####### SECRET MANAGERS #####################
-secret_manager_client: Optional[Any] = (
- None # list of instantiated key management clients - e.g. azure kv, infisical, etc.
-)
+secret_manager_client: Optional[
+ Any
+] = None # list of instantiated key management clients - e.g. azure kv, infisical, etc.
_google_kms_resource_name: Optional[str] = None
_key_management_system: Optional["KeyManagementSystem"] = None
# Note: KeyManagementSettings must be eagerly imported because _key_management_settings
@@ -431,12 +454,12 @@ output_parse_pii: bool = False
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
model_cost = get_model_cost_map(url=model_cost_map_url)
-cost_discount_config: Dict[str, float] = (
- {}
-) # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
-cost_margin_config: Dict[str, Union[float, Dict[str, float]]] = (
- {}
-) # Provider-specific or global cost margins. Examples:
+cost_discount_config: Dict[
+ str, float
+] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
+cost_margin_config: Dict[
+ str, Union[float, Dict[str, float]]
+] = {} # Provider-specific or global cost margins. Examples:
# Percentage: {"openai": 0.10} = 10% margin
# Fixed: {"openai": {"fixed_amount": 0.001}} = $0.001 per request
# Global: {"global": 0.05} = 5% global margin on all providers
@@ -555,6 +578,7 @@ v0_models: Set = set()
morph_models: Set = set()
lambda_ai_models: Set = set()
hyperbolic_models: Set = set()
+black_forest_labs_models: Set = set()
recraft_models: Set = set()
cometapi_models: Set = set()
oci_models: Set = set()
@@ -573,6 +597,7 @@ minimax_models: Set = set()
aws_polly_models: Set = set()
gigachat_models: Set = set()
llamagate_models: Set = set()
+bedrock_mantle_models: Set = set()
def is_bedrock_pricing_only_model(key: str) -> bool:
@@ -608,8 +633,9 @@ def is_openai_finetune_model(key: str) -> bool:
return key.startswith("ft:") and not key.count(":") > 1
-def add_known_models():
- for key, value in model_cost.items():
+def add_known_models(model_cost_map: Optional[Dict] = None):
+ _map = model_cost_map if model_cost_map is not None else model_cost
+ for key, value in _map.items():
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(
key
):
@@ -800,6 +826,8 @@ def add_known_models():
lambda_ai_models.add(key)
elif value.get("litellm_provider") == "hyperbolic":
hyperbolic_models.add(key)
+ elif value.get("litellm_provider") == "black_forest_labs":
+ black_forest_labs_models.add(key)
elif value.get("litellm_provider") == "recraft":
recraft_models.add(key)
elif value.get("litellm_provider") == "cometapi":
@@ -834,6 +862,8 @@ def add_known_models():
gigachat_models.add(key)
elif value.get("litellm_provider") == "llamagate":
llamagate_models.add(key)
+ elif value.get("litellm_provider") == "bedrock_mantle":
+ bedrock_mantle_models.add(key)
add_known_models()
@@ -931,6 +961,7 @@ model_list = list(
| v0_models
| morph_models
| lambda_ai_models
+ | black_forest_labs_models
| recraft_models
| cometapi_models
| oci_models
@@ -941,6 +972,7 @@ model_list = list(
| ovhcloud_models
| lemonade_models
| docker_model_runner_models
+ | bedrock_mantle_models
| set(clarifai_models)
)
@@ -1028,6 +1060,7 @@ models_by_provider: dict = {
"morph": morph_models,
"lambda_ai": lambda_ai_models,
"hyperbolic": hyperbolic_models,
+ "black_forest_labs": black_forest_labs_models,
"recraft": recraft_models,
"cometapi": cometapi_models,
"oci": oci_models,
@@ -1044,6 +1077,7 @@ models_by_provider: dict = {
"aws_polly": aws_polly_models,
"gigachat": gigachat_models,
"llamagate": llamagate_models,
+ "bedrock_mantle": bedrock_mantle_models,
}
# mapping for those models which have larger equivalents
@@ -1094,10 +1128,12 @@ openai_video_generation_models = ["sora-2"]
# Import KeyManagementSettings here (before utils import) because _key_management_settings
# is accessed during import time in secret_managers/main.py (via dd_tracing -> datadog -> _service_logger -> utils)
from litellm.types.secret_managers.main import KeyManagementSettings
+
_key_management_settings: KeyManagementSettings = KeyManagementSettings()
# client must be imported immediately as it's used as a decorator at function definition time
from .utils import client
+
# Note: Most other utils imports are lazy-loaded via __getattr__ to avoid loading utils.py
# (which imports tiktoken) at import time
@@ -1126,6 +1162,7 @@ from .llms.topaz.common_utils import TopazModelInfo
# OpenAIOSeriesConfig is lazy loaded - openaiOSeriesConfig will be created on first access
# OpenAIGPTConfig, OpenAIGPT5Config, etc. are lazy loaded - instances will be created on first access
from .llms.xai.common_utils import XAIModelInfo
+
# PublicAI now uses JSON-based configuration (see litellm/llms/openai_like/providers.json)
# All remaining configs are now lazy loaded - see _lazy_imports_registry.py
@@ -1147,6 +1184,28 @@ from .skills.main import (
delete_skill,
adelete_skill,
)
+from .evals.main import (
+ create_eval,
+ acreate_eval,
+ list_evals,
+ alist_evals,
+ get_eval,
+ aget_eval,
+ delete_eval,
+ adelete_eval,
+ cancel_eval,
+ acancel_eval,
+ create_run,
+ acreate_run,
+ list_runs,
+ alist_runs,
+ get_run,
+ aget_run,
+ delete_run,
+ adelete_run,
+ cancel_run,
+ acancel_run,
+)
from .integrations import *
from .llms.custom_httpx.async_client_cleanup import close_litellm_async_clients
from .exceptions import (
@@ -1155,6 +1214,7 @@ from .exceptions import (
BadRequestError,
ImageFetchError,
NotFoundError,
+ PermissionDeniedError,
RateLimitError,
ServiceUnavailableError,
BadGatewayError,
@@ -1184,6 +1244,7 @@ from .batch_completion.main import * # type: ignore
from .rerank_api.main import *
from .llms.anthropic.experimental_pass_through.messages.handler import *
from .responses.main import *
+
# Interactions API is available as litellm.interactions module
# Usage: litellm.interactions.create(), litellm.interactions.get(), etc.
from . import interactions
@@ -1201,7 +1262,12 @@ from .containers.main import *
from .ocr.main import *
from .rag.main import *
from .search.main import *
-from .realtime_api.main import _arealtime
+from .realtime_api.main import (
+ _arealtime,
+ acreate_realtime_client_secret,
+ arealtime_calls,
+)
+from .responses.main import _aresponses_websocket
from .fine_tuning.main import *
from .files.main import *
from .vector_store_files.main import (
@@ -1242,12 +1308,12 @@ from . import rag
from .types.llms.custom_llm import CustomLLMItem
custom_provider_map: List[CustomLLMItem] = []
-_custom_providers: List[str] = (
- []
-) # internal helper util, used to track names of custom providers
-disable_hf_tokenizer_download: Optional[bool] = (
- None # disable huggingface tokenizer download. Defaults to openai clk100
-)
+_custom_providers: List[
+ str
+] = [] # internal helper util, used to track names of custom providers
+disable_hf_tokenizer_download: Optional[
+ bool
+] = None # disable huggingface tokenizer download. Defaults to openai clk100
global_disable_no_log_param: bool = False
### CLI UTILITIES ###
@@ -1286,125 +1352,324 @@ if TYPE_CHECKING:
from litellm.caching.caching import Cache
# Type stubs for lazy-loaded configs to help mypy
- from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig as AmazonConverseConfig
- from .llms.openai_like.chat.handler import OpenAILikeChatConfig as OpenAILikeChatConfig
- from .llms.galadriel.chat.transformation import GaladrielChatConfig as GaladrielChatConfig
+ from .llms.bedrock.chat.converse_transformation import (
+ AmazonConverseConfig as AmazonConverseConfig,
+ )
+ from .llms.openai_like.chat.handler import (
+ OpenAILikeChatConfig as OpenAILikeChatConfig,
+ )
+ from .llms.galadriel.chat.transformation import (
+ GaladrielChatConfig as GaladrielChatConfig,
+ )
from .llms.github.chat.transformation import GithubChatConfig as GithubChatConfig
- from .llms.azure_ai.anthropic.transformation import AzureAnthropicConfig as AzureAnthropicConfig
+ from .llms.azure_ai.anthropic.transformation import (
+ AzureAnthropicConfig as AzureAnthropicConfig,
+ )
from .llms.bytez.chat.transformation import BytezChatConfig as BytezChatConfig
- from .llms.compactifai.chat.transformation import CompactifAIChatConfig as CompactifAIChatConfig
+ from .llms.compactifai.chat.transformation import (
+ CompactifAIChatConfig as CompactifAIChatConfig,
+ )
from .llms.empower.chat.transformation import EmpowerChatConfig as EmpowerChatConfig
from .llms.minimax.chat.transformation import MinimaxChatConfig as MinimaxChatConfig
- from .llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig as AiohttpOpenAIChatConfig
- from .llms.huggingface.chat.transformation import HuggingFaceChatConfig as HuggingFaceChatConfig
- from .llms.huggingface.embedding.transformation import HuggingFaceEmbeddingConfig as HuggingFaceEmbeddingConfig
+ from .llms.aiohttp_openai.chat.transformation import (
+ AiohttpOpenAIChatConfig as AiohttpOpenAIChatConfig,
+ )
+ from .llms.huggingface.chat.transformation import (
+ HuggingFaceChatConfig as HuggingFaceChatConfig,
+ )
+ from .llms.huggingface.embedding.transformation import (
+ HuggingFaceEmbeddingConfig as HuggingFaceEmbeddingConfig,
+ )
from .llms.oobabooga.chat.transformation import OobaboogaConfig as OobaboogaConfig
from .llms.maritalk import MaritalkConfig as MaritalkConfig
- from .llms.openrouter.chat.transformation import OpenrouterConfig as OpenrouterConfig
+ from .llms.openrouter.chat.transformation import (
+ OpenrouterConfig as OpenrouterConfig,
+ )
from .llms.datarobot.chat.transformation import DataRobotConfig as DataRobotConfig
from .llms.anthropic.chat.transformation import AnthropicConfig as AnthropicConfig
- from .llms.anthropic.completion.transformation import AnthropicTextConfig as AnthropicTextConfig
+ from .llms.anthropic.completion.transformation import (
+ AnthropicTextConfig as AnthropicTextConfig,
+ )
from .llms.groq.stt.transformation import GroqSTTConfig as GroqSTTConfig
from .llms.triton.completion.transformation import TritonConfig as TritonConfig
- from .llms.triton.completion.transformation import TritonGenerateConfig as TritonGenerateConfig
- from .llms.triton.completion.transformation import TritonInferConfig as TritonInferConfig
- from .llms.triton.embedding.transformation import TritonEmbeddingConfig as TritonEmbeddingConfig
- from .llms.huggingface.rerank.transformation import HuggingFaceRerankConfig as HuggingFaceRerankConfig
- from .llms.databricks.chat.transformation import DatabricksConfig as DatabricksConfig
- from .llms.databricks.embed.transformation import DatabricksEmbeddingConfig as DatabricksEmbeddingConfig
+ from .llms.triton.completion.transformation import (
+ TritonGenerateConfig as TritonGenerateConfig,
+ )
+ from .llms.triton.completion.transformation import (
+ TritonInferConfig as TritonInferConfig,
+ )
+ from .llms.triton.embedding.transformation import (
+ TritonEmbeddingConfig as TritonEmbeddingConfig,
+ )
+ from .llms.huggingface.rerank.transformation import (
+ HuggingFaceRerankConfig as HuggingFaceRerankConfig,
+ )
+ from .llms.databricks.chat.transformation import (
+ DatabricksConfig as DatabricksConfig,
+ )
+ from .llms.databricks.embed.transformation import (
+ DatabricksEmbeddingConfig as DatabricksEmbeddingConfig,
+ )
from .llms.predibase.chat.transformation import PredibaseConfig as PredibaseConfig
from .llms.replicate.chat.transformation import ReplicateConfig as ReplicateConfig
from .llms.snowflake.chat.transformation import SnowflakeConfig as SnowflakeConfig
- from .llms.cohere.rerank.transformation import CohereRerankConfig as CohereRerankConfig
- from .llms.cohere.rerank_v2.transformation import CohereRerankV2Config as CohereRerankV2Config
- from .llms.azure_ai.rerank.transformation import AzureAIRerankConfig as AzureAIRerankConfig
- from .llms.infinity.rerank.transformation import InfinityRerankConfig as InfinityRerankConfig
- from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig as JinaAIRerankConfig
- from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig as DeepinfraRerankConfig
- from .llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig as HostedVLLMRerankConfig
- from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig as NvidiaNimRerankConfig
- from .llms.nvidia_nim.rerank.ranking_transformation import NvidiaNimRankingConfig as NvidiaNimRankingConfig
- from .llms.vertex_ai.rerank.transformation import VertexAIRerankConfig as VertexAIRerankConfig
- from .llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig as FireworksAIRerankConfig
- from .llms.voyage.rerank.transformation import VoyageRerankConfig as VoyageRerankConfig
+ from .llms.cohere.rerank.transformation import (
+ CohereRerankConfig as CohereRerankConfig,
+ )
+ from .llms.cohere.rerank_v2.transformation import (
+ CohereRerankV2Config as CohereRerankV2Config,
+ )
+ from .llms.azure_ai.rerank.transformation import (
+ AzureAIRerankConfig as AzureAIRerankConfig,
+ )
+ from .llms.infinity.rerank.transformation import (
+ InfinityRerankConfig as InfinityRerankConfig,
+ )
+ from .llms.jina_ai.rerank.transformation import (
+ JinaAIRerankConfig as JinaAIRerankConfig,
+ )
+ from .llms.deepinfra.rerank.transformation import (
+ DeepinfraRerankConfig as DeepinfraRerankConfig,
+ )
+ from .llms.hosted_vllm.rerank.transformation import (
+ HostedVLLMRerankConfig as HostedVLLMRerankConfig,
+ )
+ from .llms.nvidia_nim.rerank.transformation import (
+ NvidiaNimRerankConfig as NvidiaNimRerankConfig,
+ )
+ from .llms.nvidia_nim.rerank.ranking_transformation import (
+ NvidiaNimRankingConfig as NvidiaNimRankingConfig,
+ )
+ from .llms.vertex_ai.rerank.transformation import (
+ VertexAIRerankConfig as VertexAIRerankConfig,
+ )
+ from .llms.fireworks_ai.rerank.transformation import (
+ FireworksAIRerankConfig as FireworksAIRerankConfig,
+ )
+ from .llms.voyage.rerank.transformation import (
+ VoyageRerankConfig as VoyageRerankConfig,
+ )
+ from .llms.watsonx.rerank.transformation import (
+ IBMWatsonXRerankConfig as IBMWatsonXRerankConfig,
+ )
from .llms.clarifai.chat.transformation import ClarifaiConfig as ClarifaiConfig
from .llms.ai21.chat.transformation import AI21ChatConfig as AI21ChatConfig
from .llms.meta_llama.chat.transformation import LlamaAPIConfig as LlamaAPIConfig
- from .llms.together_ai.completion.transformation import TogetherAITextCompletionConfig as TogetherAITextCompletionConfig
- from .llms.cloudflare.chat.transformation import CloudflareChatConfig as CloudflareChatConfig
+ from .llms.together_ai.completion.transformation import (
+ TogetherAITextCompletionConfig as TogetherAITextCompletionConfig,
+ )
+ from .llms.cloudflare.chat.transformation import (
+ CloudflareChatConfig as CloudflareChatConfig,
+ )
from .llms.novita.chat.transformation import NovitaConfig as NovitaConfig
from .llms.petals.completion.transformation import PetalsConfig as PetalsConfig
from .llms.ollama.chat.transformation import OllamaChatConfig as OllamaChatConfig
from .llms.ollama.completion.transformation import OllamaConfig as OllamaConfig
from .llms.sagemaker.completion.transformation import SagemakerConfig as SagemakerConfig
from .llms.sagemaker.chat.transformation import SagemakerChatConfig as SagemakerChatConfig
+ from .llms.sagemaker.nova.transformation import SagemakerNovaConfig as SagemakerNovaConfig
from .llms.cohere.chat.transformation import CohereChatConfig as CohereChatConfig
- from .llms.anthropic.experimental_pass_through.messages.transformation import AnthropicMessagesConfig as AnthropicMessagesConfig
- from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig
+ from .llms.anthropic.experimental_pass_through.messages.transformation import (
+ AnthropicMessagesConfig as AnthropicMessagesConfig,
+ )
+ from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
+ AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig,
+ )
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig
- from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig as VertexGeminiConfig
- from .llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig as GoogleAIStudioGeminiConfig
- from .llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig as VertexAIAnthropicConfig
- from .llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import VertexAILlama3Config as VertexAILlama3Config
- from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import VertexAIAi21Config as VertexAIAi21Config
- from .llms.bedrock.chat.invoke_handler import AmazonCohereChatConfig as AmazonCohereChatConfig
- from .llms.bedrock.common_utils import AmazonBedrockGlobalConfig as AmazonBedrockGlobalConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_ai21_transformation import AmazonAI21Config as AmazonAI21Config
- from .llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import AmazonInvokeNovaConfig as AmazonInvokeNovaConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformation import AmazonQwen2Config as AmazonQwen2Config
- from .llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import AmazonQwen3Config as AmazonQwen3Config
- from .llms.bedrock.chat.invoke_transformations.anthropic_claude2_transformation import AmazonAnthropicConfig as AmazonAnthropicConfig
- from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import AmazonAnthropicClaudeConfig as AmazonAnthropicClaudeConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_cohere_transformation import AmazonCohereConfig as AmazonCohereConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_llama_transformation import AmazonLlamaConfig as AmazonLlamaConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import AmazonDeepSeekR1Config as AmazonDeepSeekR1Config
- from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import AmazonMistralConfig as AmazonMistralConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import AmazonMoonshotConfig as AmazonMoonshotConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import AmazonTitanConfig as AmazonTitanConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import AmazonTwelveLabsPegasusConfig as AmazonTwelveLabsPegasusConfig
- from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig as AmazonInvokeConfig
- from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig as AmazonBedrockOpenAIConfig
- from .llms.bedrock.image_generation.amazon_stability1_transformation import AmazonStabilityConfig as AmazonStabilityConfig
- from .llms.bedrock.image_generation.amazon_stability3_transformation import AmazonStability3Config as AmazonStability3Config
- from .llms.bedrock.image_generation.amazon_nova_canvas_transformation import AmazonNovaCanvasConfig as AmazonNovaCanvasConfig
- from .llms.bedrock.embed.amazon_titan_g1_transformation import AmazonTitanG1Config as AmazonTitanG1Config
- from .llms.bedrock.embed.amazon_titan_multimodal_transformation import AmazonTitanMultimodalEmbeddingG1Config as AmazonTitanMultimodalEmbeddingG1Config
- from .llms.cohere.chat.v2_transformation import CohereV2ChatConfig as CohereV2ChatConfig
- from .llms.bedrock.embed.cohere_transformation import BedrockCohereEmbeddingConfig as BedrockCohereEmbeddingConfig
- from .llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig as TwelveLabsMarengoEmbeddingConfig
- from .llms.bedrock.embed.amazon_nova_transformation import AmazonNovaEmbeddingConfig as AmazonNovaEmbeddingConfig
- from .llms.openai.openai import OpenAIConfig as OpenAIConfig, MistralEmbeddingConfig as MistralEmbeddingConfig
- from .llms.openai.image_variations.transformation import OpenAIImageVariationConfig as OpenAIImageVariationConfig
- from .llms.deepgram.audio_transcription.transformation import DeepgramAudioTranscriptionConfig as DeepgramAudioTranscriptionConfig
- from .llms.topaz.image_variations.transformation import TopazImageVariationConfig as TopazImageVariationConfig
- from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig as OpenAITextCompletionConfig
+ from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+ VertexGeminiConfig as VertexGeminiConfig,
+ )
+ from .llms.gemini.chat.transformation import (
+ GoogleAIStudioGeminiConfig as GoogleAIStudioGeminiConfig,
+ )
+ from .llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import (
+ VertexAIAnthropicConfig as VertexAIAnthropicConfig,
+ )
+ from .llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import (
+ VertexAILlama3Config as VertexAILlama3Config,
+ )
+ from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import (
+ VertexAIAi21Config as VertexAIAi21Config,
+ )
+ from .llms.bedrock.chat.invoke_handler import (
+ AmazonCohereChatConfig as AmazonCohereChatConfig,
+ )
+ from .llms.bedrock.common_utils import (
+ AmazonBedrockGlobalConfig as AmazonBedrockGlobalConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_ai21_transformation import (
+ AmazonAI21Config as AmazonAI21Config,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import (
+ AmazonInvokeNovaConfig as AmazonInvokeNovaConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformation import (
+ AmazonQwen2Config as AmazonQwen2Config,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import (
+ AmazonQwen3Config as AmazonQwen3Config,
+ )
+ from .llms.bedrock.chat.invoke_transformations.anthropic_claude2_transformation import (
+ AmazonAnthropicConfig as AmazonAnthropicConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
+ AmazonAnthropicClaudeConfig as AmazonAnthropicClaudeConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_cohere_transformation import (
+ AmazonCohereConfig as AmazonCohereConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_llama_transformation import (
+ AmazonLlamaConfig as AmazonLlamaConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_deepseek_transformation import (
+ AmazonDeepSeekR1Config as AmazonDeepSeekR1Config,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_mistral_transformation import (
+ AmazonMistralConfig as AmazonMistralConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
+ AmazonMoonshotConfig as AmazonMoonshotConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation import (
+ AmazonTitanConfig as AmazonTitanConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import (
+ AmazonTwelveLabsPegasusConfig as AmazonTwelveLabsPegasusConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
+ AmazonInvokeConfig as AmazonInvokeConfig,
+ )
+ from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import (
+ AmazonBedrockOpenAIConfig as AmazonBedrockOpenAIConfig,
+ )
+ from .llms.bedrock.image_generation.amazon_stability1_transformation import (
+ AmazonStabilityConfig as AmazonStabilityConfig,
+ )
+ from .llms.bedrock.image_generation.amazon_stability3_transformation import (
+ AmazonStability3Config as AmazonStability3Config,
+ )
+ from .llms.bedrock.image_generation.amazon_nova_canvas_transformation import (
+ AmazonNovaCanvasConfig as AmazonNovaCanvasConfig,
+ )
+ from .llms.bedrock.embed.amazon_titan_g1_transformation import (
+ AmazonTitanG1Config as AmazonTitanG1Config,
+ )
+ from .llms.bedrock.embed.amazon_titan_multimodal_transformation import (
+ AmazonTitanMultimodalEmbeddingG1Config as AmazonTitanMultimodalEmbeddingG1Config,
+ )
+ from .llms.cohere.chat.v2_transformation import (
+ CohereV2ChatConfig as CohereV2ChatConfig,
+ )
+ from .llms.bedrock.embed.cohere_transformation import (
+ BedrockCohereEmbeddingConfig as BedrockCohereEmbeddingConfig,
+ )
+ from .llms.bedrock.embed.twelvelabs_marengo_transformation import (
+ TwelveLabsMarengoEmbeddingConfig as TwelveLabsMarengoEmbeddingConfig,
+ )
+ from .llms.bedrock.embed.amazon_nova_transformation import (
+ AmazonNovaEmbeddingConfig as AmazonNovaEmbeddingConfig,
+ )
+ from .llms.openai.openai import (
+ OpenAIConfig as OpenAIConfig,
+ MistralEmbeddingConfig as MistralEmbeddingConfig,
+ )
+ from .llms.openai.image_variations.transformation import (
+ OpenAIImageVariationConfig as OpenAIImageVariationConfig,
+ )
+ from .llms.deepgram.audio_transcription.transformation import (
+ DeepgramAudioTranscriptionConfig as DeepgramAudioTranscriptionConfig,
+ )
+ from .llms.topaz.image_variations.transformation import (
+ TopazImageVariationConfig as TopazImageVariationConfig,
+ )
+ from litellm.llms.openai.completion.transformation import (
+ OpenAITextCompletionConfig as OpenAITextCompletionConfig,
+ )
from .llms.groq.chat.transformation import GroqChatConfig as GroqChatConfig
+ from .llms.bedrock_mantle.chat.transformation import (
+ BedrockMantleChatConfig as BedrockMantleChatConfig,
+ )
from .llms.a2a.chat.transformation import A2AConfig as A2AConfig
- from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig as VoyageEmbeddingConfig
- from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig
- from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig as InfinityEmbeddingConfig
- from .llms.azure_ai.chat.transformation import AzureAIStudioConfig as AzureAIStudioConfig
+ from .llms.voyage.embedding.transformation import (
+ VoyageEmbeddingConfig as VoyageEmbeddingConfig,
+ )
+ from .llms.voyage.embedding.transformation_contextual import (
+ VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig,
+ )
+ from .llms.infinity.embedding.transformation import (
+ InfinityEmbeddingConfig as InfinityEmbeddingConfig,
+ )
+ from .llms.perplexity.embedding.transformation import (
+ PerplexityEmbeddingConfig as PerplexityEmbeddingConfig,
+ )
+ from .llms.azure_ai.chat.transformation import (
+ AzureAIStudioConfig as AzureAIStudioConfig,
+ )
from .llms.mistral.chat.transformation import MistralConfig as MistralConfig
- from .llms.openai.responses.transformation import OpenAIResponsesAPIConfig as OpenAIResponsesAPIConfig
- from .llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig as AzureOpenAIResponsesAPIConfig
- from .llms.azure.responses.o_series_transformation import AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig
- from .llms.xai.responses.transformation import XAIResponsesAPIConfig as XAIResponsesAPIConfig
- from .llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig
- from .llms.volcengine.responses.transformation import VolcEngineResponsesAPIConfig as VolcEngineResponsesAPIConfig
- from .llms.manus.responses.transformation import ManusResponsesAPIConfig as ManusResponsesAPIConfig
- from .llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig
- from .llms.openai.chat.o_series_transformation import OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config
- from .llms.anthropic.skills.transformation import AnthropicSkillsConfig as AnthropicSkillsConfig
- from .llms.base_llm.skills.transformation import BaseSkillsAPIConfig as BaseSkillsAPIConfig
- from .llms.gradient_ai.chat.transformation import GradientAIConfig as GradientAIConfig
+ from .llms.openai.responses.transformation import (
+ OpenAIResponsesAPIConfig as OpenAIResponsesAPIConfig,
+ )
+ from .llms.azure.responses.transformation import (
+ AzureOpenAIResponsesAPIConfig as AzureOpenAIResponsesAPIConfig,
+ )
+ from .llms.azure.responses.o_series_transformation import (
+ AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig,
+ )
+ from .llms.xai.responses.transformation import (
+ XAIResponsesAPIConfig as XAIResponsesAPIConfig,
+ )
+ from .llms.litellm_proxy.responses.transformation import (
+ LiteLLMProxyResponsesAPIConfig as LiteLLMProxyResponsesAPIConfig,
+ )
+ from .llms.volcengine.responses.transformation import (
+ VolcEngineResponsesAPIConfig as VolcEngineResponsesAPIConfig,
+ )
+ from .llms.manus.responses.transformation import (
+ ManusResponsesAPIConfig as ManusResponsesAPIConfig,
+ )
+ from .llms.perplexity.responses.transformation import (
+ PerplexityResponsesConfig as PerplexityResponsesConfig,
+ )
+ from .llms.databricks.responses.transformation import (
+ DatabricksResponsesAPIConfig as DatabricksResponsesAPIConfig,
+ )
+ from .llms.openrouter.responses.transformation import (
+ OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig,
+ )
+ from .llms.gemini.interactions.transformation import (
+ GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
+ )
+ from .llms.openai.chat.o_series_transformation import (
+ OpenAIOSeriesConfig as OpenAIOSeriesConfig,
+ OpenAIOSeriesConfig as OpenAIO1Config,
+ )
+ from .llms.anthropic.skills.transformation import (
+ AnthropicSkillsConfig as AnthropicSkillsConfig,
+ )
+ from .llms.base_llm.skills.transformation import (
+ BaseSkillsAPIConfig as BaseSkillsAPIConfig,
+ )
+ from .llms.gradient_ai.chat.transformation import (
+ GradientAIConfig as GradientAIConfig,
+ )
from .llms.openai.chat.gpt_transformation import OpenAIGPTConfig as OpenAIGPTConfig
- from .llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config as OpenAIGPT5Config
- from .llms.openai.transcriptions.whisper_transformation import OpenAIWhisperAudioTranscriptionConfig as OpenAIWhisperAudioTranscriptionConfig
- from .llms.openai.transcriptions.gpt_transformation import OpenAIGPTAudioTranscriptionConfig as OpenAIGPTAudioTranscriptionConfig
- from .llms.openai.chat.gpt_audio_transformation import OpenAIGPTAudioConfig as OpenAIGPTAudioConfig
+ from .llms.openai.chat.gpt_5_transformation import (
+ OpenAIGPT5Config as OpenAIGPT5Config,
+ )
+ from .llms.openai.transcriptions.whisper_transformation import (
+ OpenAIWhisperAudioTranscriptionConfig as OpenAIWhisperAudioTranscriptionConfig,
+ )
+ from .llms.openai.transcriptions.gpt_transformation import (
+ OpenAIGPTAudioTranscriptionConfig as OpenAIGPTAudioTranscriptionConfig,
+ )
+ from .llms.openai.chat.gpt_audio_transformation import (
+ OpenAIGPTAudioConfig as OpenAIGPTAudioConfig,
+ )
from .llms.nvidia_nim.chat.transformation import NvidiaNimConfig as NvidiaNimConfig
- from .llms.nvidia_nim.embed import NvidiaNimEmbeddingConfig as NvidiaNimEmbeddingConfig
+ from .llms.nvidia_nim.embed import (
+ NvidiaNimEmbeddingConfig as NvidiaNimEmbeddingConfig,
+ )
# Type stubs for lazy-loaded config instances
openaiOSeriesConfig: OpenAIOSeriesConfig
@@ -1416,21 +1681,47 @@ if TYPE_CHECKING:
# Import config classes that need type stubs (for mypy) - import with _ prefix to avoid circular reference
from .llms.vllm.completion.transformation import VLLMConfig as _VLLMConfig
- from .llms.deepseek.chat.transformation import DeepSeekChatConfig as _DeepSeekChatConfig
- from .llms.sap.chat.transformation import GenAIHubOrchestrationConfig as _GenAIHubOrchestrationConfig
- from .llms.sap.embed.transformation import GenAIHubEmbeddingConfig as _GenAIHubEmbeddingConfig
- from .llms.azure.chat.o_series_transformation import AzureOpenAIO1Config as _AzureOpenAIO1Config
- from .llms.perplexity.chat.transformation import PerplexityChatConfig as _PerplexityChatConfig
+ from .llms.deepseek.chat.transformation import (
+ DeepSeekChatConfig as _DeepSeekChatConfig,
+ )
+ from .llms.sap.chat.transformation import (
+ GenAIHubOrchestrationConfig as _GenAIHubOrchestrationConfig,
+ )
+ from .llms.sap.embed.transformation import (
+ GenAIHubEmbeddingConfig as _GenAIHubEmbeddingConfig,
+ )
+ from .llms.azure.chat.o_series_transformation import (
+ AzureOpenAIO1Config as _AzureOpenAIO1Config,
+ )
+ from .llms.perplexity.chat.transformation import (
+ PerplexityChatConfig as _PerplexityChatConfig,
+ )
from .llms.nscale.chat.transformation import NscaleConfig as _NscaleConfig
- from .llms.watsonx.chat.transformation import IBMWatsonXChatConfig as _IBMWatsonXChatConfig
- from .llms.watsonx.completion.transformation import IBMWatsonXAIConfig as _IBMWatsonXAIConfig
- from .llms.litellm_proxy.chat.transformation import LiteLLMProxyChatConfig as _LiteLLMProxyChatConfig
+ from .llms.watsonx.chat.transformation import (
+ IBMWatsonXChatConfig as _IBMWatsonXChatConfig,
+ )
+ from .llms.watsonx.completion.transformation import (
+ IBMWatsonXAIConfig as _IBMWatsonXAIConfig,
+ )
+ from .llms.litellm_proxy.chat.transformation import (
+ LiteLLMProxyChatConfig as _LiteLLMProxyChatConfig,
+ )
from .llms.deepinfra.chat.transformation import DeepInfraConfig as _DeepInfraConfig
- from .llms.llamafile.chat.transformation import LlamafileChatConfig as _LlamafileChatConfig
- from .llms.lm_studio.chat.transformation import LMStudioChatConfig as _LMStudioChatConfig
- from .llms.lm_studio.embed.transformation import LmStudioEmbeddingConfig as _LmStudioEmbeddingConfig
- from .llms.watsonx.embed.transformation import IBMWatsonXEmbeddingConfig as _IBMWatsonXEmbeddingConfig
- from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig as _VertexGeminiConfig
+ from .llms.llamafile.chat.transformation import (
+ LlamafileChatConfig as _LlamafileChatConfig,
+ )
+ from .llms.lm_studio.chat.transformation import (
+ LMStudioChatConfig as _LMStudioChatConfig,
+ )
+ from .llms.lm_studio.embed.transformation import (
+ LmStudioEmbeddingConfig as _LmStudioEmbeddingConfig,
+ )
+ from .llms.watsonx.embed.transformation import (
+ IBMWatsonXEmbeddingConfig as _IBMWatsonXEmbeddingConfig,
+ )
+ from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+ VertexGeminiConfig as _VertexGeminiConfig,
+ )
# Type stubs for lazy-loaded config classes (to help mypy understand types)
VLLMConfig: Type[_VLLMConfig]
@@ -1450,55 +1741,125 @@ if TYPE_CHECKING:
IBMWatsonXEmbeddingConfig: Type[_IBMWatsonXEmbeddingConfig]
VertexAIConfig: Type[_VertexGeminiConfig] # Alias for VertexGeminiConfig
- from .llms.featherless_ai.chat.transformation import FeatherlessAIConfig as FeatherlessAIConfig
+ from .llms.featherless_ai.chat.transformation import (
+ FeatherlessAIConfig as FeatherlessAIConfig,
+ )
from .llms.cerebras.chat import CerebrasConfig as CerebrasConfig
from .llms.baseten.chat import BasetenConfig as BasetenConfig
from .llms.sambanova.chat import SambanovaConfig as SambanovaConfig
- from .llms.sambanova.embedding.transformation import SambaNovaEmbeddingConfig as SambaNovaEmbeddingConfig
- from .llms.fireworks_ai.chat.transformation import FireworksAIConfig as FireworksAIConfig
- from .llms.fireworks_ai.completion.transformation import FireworksAITextCompletionConfig as FireworksAITextCompletionConfig
- from .llms.fireworks_ai.audio_transcription.transformation import FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig
- from .llms.fireworks_ai.embed.fireworks_ai_transformation import FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig
- from .llms.friendliai.chat.transformation import FriendliaiChatConfig as FriendliaiChatConfig
- from .llms.jina_ai.embedding.transformation import JinaAIEmbeddingConfig as JinaAIEmbeddingConfig
+ from .llms.sambanova.embedding.transformation import (
+ SambaNovaEmbeddingConfig as SambaNovaEmbeddingConfig,
+ )
+ from .llms.fireworks_ai.chat.transformation import (
+ FireworksAIConfig as FireworksAIConfig,
+ )
+ from .llms.fireworks_ai.completion.transformation import (
+ FireworksAITextCompletionConfig as FireworksAITextCompletionConfig,
+ )
+ from .llms.fireworks_ai.audio_transcription.transformation import (
+ FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig,
+ )
+ from .llms.fireworks_ai.embed.fireworks_ai_transformation import (
+ FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig,
+ )
+ from .llms.friendliai.chat.transformation import (
+ FriendliaiChatConfig as FriendliaiChatConfig,
+ )
+ from .llms.jina_ai.embedding.transformation import (
+ JinaAIEmbeddingConfig as JinaAIEmbeddingConfig,
+ )
from .llms.xai.chat.transformation import XAIChatConfig as XAIChatConfig
from .llms.zai.chat.transformation import ZAIChatConfig as ZAIChatConfig
from .llms.aiml.chat.transformation import AIMLChatConfig as AIMLChatConfig
- from .llms.volcengine.chat.transformation import VolcEngineChatConfig as VolcEngineChatConfig, VolcEngineChatConfig as VolcEngineConfig
- from .llms.codestral.completion.transformation import CodestralTextCompletionConfig as CodestralTextCompletionConfig
- from .llms.azure.azure import AzureOpenAIAssistantsAPIConfig as AzureOpenAIAssistantsAPIConfig
+ from .llms.volcengine.chat.transformation import (
+ VolcEngineChatConfig as VolcEngineChatConfig,
+ VolcEngineChatConfig as VolcEngineConfig,
+ )
+ from .llms.codestral.completion.transformation import (
+ CodestralTextCompletionConfig as CodestralTextCompletionConfig,
+ )
+ from .llms.azure.azure import (
+ AzureOpenAIAssistantsAPIConfig as AzureOpenAIAssistantsAPIConfig,
+ )
from .llms.heroku.chat.transformation import HerokuChatConfig as HerokuChatConfig
from .llms.cometapi.chat.transformation import CometAPIConfig as CometAPIConfig
- from .llms.azure.chat.gpt_transformation import AzureOpenAIConfig as AzureOpenAIConfig
- from .llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config as AzureOpenAIGPT5Config
- from .llms.azure.completion.transformation import AzureOpenAITextConfig as AzureOpenAITextConfig
- from .llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig as HostedVLLMChatConfig
- from .llms.hosted_vllm.embedding.transformation import HostedVLLMEmbeddingConfig as HostedVLLMEmbeddingConfig
- from .llms.github_copilot.chat.transformation import GithubCopilotConfig as GithubCopilotConfig
- from .llms.github_copilot.responses.transformation import GithubCopilotResponsesAPIConfig as GithubCopilotResponsesAPIConfig
- from .llms.github_copilot.embedding.transformation import GithubCopilotEmbeddingConfig as GithubCopilotEmbeddingConfig
+ from .llms.azure.chat.gpt_transformation import (
+ AzureOpenAIConfig as AzureOpenAIConfig,
+ )
+ from .llms.azure.chat.gpt_5_transformation import (
+ AzureOpenAIGPT5Config as AzureOpenAIGPT5Config,
+ )
+ from .llms.azure.completion.transformation import (
+ AzureOpenAITextConfig as AzureOpenAITextConfig,
+ )
+ from .llms.hosted_vllm.chat.transformation import (
+ HostedVLLMChatConfig as HostedVLLMChatConfig,
+ )
+ from .llms.hosted_vllm.embedding.transformation import (
+ HostedVLLMEmbeddingConfig as HostedVLLMEmbeddingConfig,
+ )
+ from .llms.hosted_vllm.responses.transformation import (
+ HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig,
+ )
+ from .llms.github_copilot.chat.transformation import (
+ GithubCopilotConfig as GithubCopilotConfig,
+ )
+ from .llms.github_copilot.responses.transformation import (
+ GithubCopilotResponsesAPIConfig as GithubCopilotResponsesAPIConfig,
+ )
+ from .llms.github_copilot.embedding.transformation import (
+ GithubCopilotEmbeddingConfig as GithubCopilotEmbeddingConfig,
+ )
from .llms.chatgpt.chat.transformation import ChatGPTConfig as ChatGPTConfig
- from .llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig as ChatGPTResponsesAPIConfig
+ from .llms.chatgpt.responses.transformation import (
+ ChatGPTResponsesAPIConfig as ChatGPTResponsesAPIConfig,
+ )
from .llms.gigachat.chat.transformation import GigaChatConfig as GigaChatConfig
- from .llms.gigachat.embedding.transformation import GigaChatEmbeddingConfig as GigaChatEmbeddingConfig
+ from .llms.gigachat.embedding.transformation import (
+ GigaChatEmbeddingConfig as GigaChatEmbeddingConfig,
+ )
from .llms.nebius.chat.transformation import NebiusConfig as NebiusConfig
from .llms.wandb.chat.transformation import WandbConfig as WandbConfig
- from .llms.dashscope.chat.transformation import DashScopeChatConfig as DashScopeChatConfig
- from .llms.moonshot.chat.transformation import MoonshotChatConfig as MoonshotChatConfig
- from .llms.docker_model_runner.chat.transformation import DockerModelRunnerChatConfig as DockerModelRunnerChatConfig
+ from .llms.dashscope.chat.transformation import (
+ DashScopeChatConfig as DashScopeChatConfig,
+ )
+ from .llms.moonshot.chat.transformation import (
+ MoonshotChatConfig as MoonshotChatConfig,
+ )
+ from .llms.docker_model_runner.chat.transformation import (
+ DockerModelRunnerChatConfig as DockerModelRunnerChatConfig,
+ )
from .llms.v0.chat.transformation import V0ChatConfig as V0ChatConfig
from .llms.oci.chat.transformation import OCIChatConfig as OCIChatConfig
from .llms.morph.chat.transformation import MorphChatConfig as MorphChatConfig
from .llms.ragflow.chat.transformation import RAGFlowConfig as RAGFlowConfig
- from .llms.lambda_ai.chat.transformation import LambdaAIChatConfig as LambdaAIChatConfig
- from .llms.hyperbolic.chat.transformation import HyperbolicChatConfig as HyperbolicChatConfig
- from .llms.vercel_ai_gateway.chat.transformation import VercelAIGatewayConfig as VercelAIGatewayConfig
- from .llms.ovhcloud.chat.transformation import OVHCloudChatConfig as OVHCloudChatConfig
- from .llms.ovhcloud.embedding.transformation import OVHCloudEmbeddingConfig as OVHCloudEmbeddingConfig
- from .llms.cometapi.embed.transformation import CometAPIEmbeddingConfig as CometAPIEmbeddingConfig
- from .llms.lemonade.chat.transformation import LemonadeChatConfig as LemonadeChatConfig
- from .llms.snowflake.embedding.transformation import SnowflakeEmbeddingConfig as SnowflakeEmbeddingConfig
- from .llms.amazon_nova.chat.transformation import AmazonNovaChatConfig as AmazonNovaChatConfig
+ from .llms.lambda_ai.chat.transformation import (
+ LambdaAIChatConfig as LambdaAIChatConfig,
+ )
+ from .llms.hyperbolic.chat.transformation import (
+ HyperbolicChatConfig as HyperbolicChatConfig,
+ )
+ from .llms.vercel_ai_gateway.chat.transformation import (
+ VercelAIGatewayConfig as VercelAIGatewayConfig,
+ )
+ from .llms.ovhcloud.chat.transformation import (
+ OVHCloudChatConfig as OVHCloudChatConfig,
+ )
+ from .llms.ovhcloud.embedding.transformation import (
+ OVHCloudEmbeddingConfig as OVHCloudEmbeddingConfig,
+ )
+ from .llms.cometapi.embed.transformation import (
+ CometAPIEmbeddingConfig as CometAPIEmbeddingConfig,
+ )
+ from .llms.lemonade.chat.transformation import (
+ LemonadeChatConfig as LemonadeChatConfig,
+ )
+ from .llms.snowflake.embedding.transformation import (
+ SnowflakeEmbeddingConfig as SnowflakeEmbeddingConfig,
+ )
+ from .llms.amazon_nova.chat.transformation import (
+ AmazonNovaChatConfig as AmazonNovaChatConfig,
+ )
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.types.llms.bedrock import COHERE_EMBEDDING_INPUT_TYPES
from litellm.types.utils import (
@@ -1534,7 +1895,7 @@ if TYPE_CHECKING:
supports_reasoning: Callable[..., bool]
acreate: Callable[..., Any]
get_max_tokens: Callable[..., int]
- get_model_info: Callable[..., _ModelInfoType]
+ get_model_info: Callable[..., _ModelInfoType] # type: ignore[no-redef]
register_prompt_template: Callable[..., None]
validate_environment: Callable[..., dict]
check_valid_key: Callable[..., bool]
@@ -1559,6 +1920,7 @@ if TYPE_CHECKING:
# Bedrock tool name mappings instance (lazy-loaded)
from litellm.caching.caching import InMemoryCache
+
bedrock_tool_name_mappings: InMemoryCache
# Azure exception class (lazy-loaded)
@@ -1577,11 +1939,15 @@ if TYPE_CHECKING:
from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams
# Logging callback manager class and instance (lazy-loaded)
- from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
+ from litellm.litellm_core_utils.logging_callback_manager import (
+ LoggingCallbackManager,
+ )
+
logging_callback_manager: LoggingCallbackManager
# provider_list is lazy-loaded
from litellm.types.utils import LlmProviders
+
provider_list: List[Union[LlmProviders, str]]
# Note: AmazonConverseConfig and OpenAILikeChatConfig are imported above in TYPE_CHECKING block
@@ -1606,7 +1972,10 @@ def __getattr__(name: str) -> Any:
global _async_client_cleanup_registered
# Register async client cleanup on first access (only once)
if not _async_client_cleanup_registered:
- from litellm.llms.custom_httpx.async_client_cleanup import register_async_client_cleanup
+ from litellm.llms.custom_httpx.async_client_cleanup import (
+ register_async_client_cleanup,
+ )
+
register_async_client_cleanup()
_async_client_cleanup_registered = True
@@ -1623,36 +1992,45 @@ def __getattr__(name: str) -> Any:
# Lazy load encoding from main.py to avoid heavy tiktoken import
if name == "encoding":
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
# Check if already cached
if "encoding" not in _globals:
from .main import encoding as _encoding
+
_globals["encoding"] = _encoding
return _globals["encoding"]
# Lazy load bedrock_tool_name_mappings instance
if name == "bedrock_tool_name_mappings":
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
# Check if already cached
if "bedrock_tool_name_mappings" not in _globals:
- from .llms.bedrock.chat.invoke_handler import bedrock_tool_name_mappings as _bedrock_tool_name_mappings
+ from .llms.bedrock.chat.invoke_handler import (
+ bedrock_tool_name_mappings as _bedrock_tool_name_mappings,
+ )
+
_globals["bedrock_tool_name_mappings"] = _bedrock_tool_name_mappings
return _globals["bedrock_tool_name_mappings"]
# Lazy load AzureOpenAIError exception class
if name == "AzureOpenAIError":
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
# Check if already cached
if "AzureOpenAIError" not in _globals:
from .llms.azure.common_utils import AzureOpenAIError as _AzureOpenAIError
+
_globals["AzureOpenAIError"] = _AzureOpenAIError
return _globals["AzureOpenAIError"]
# Lazy load openaiOSeriesConfig instance
if name == "openaiOSeriesConfig":
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
if "openaiOSeriesConfig" not in _globals:
# Import the config class and instantiate it
@@ -1670,6 +2048,7 @@ def __getattr__(name: str) -> Any:
}
if name in _config_instances:
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
if name not in _globals:
# Import the config class and instantiate it
@@ -1684,17 +2063,20 @@ def __getattr__(name: str) -> Any:
# Lazy load provider_list
if name == "provider_list":
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
# Check if already cached
if "provider_list" not in _globals:
# LlmProviders is eagerly imported above, so we can import it directly
from litellm.types.utils import LlmProviders
+
_globals["provider_list"] = list(LlmProviders)
return _globals["provider_list"]
# Lazy load priority_reservation_settings instance
if name == "priority_reservation_settings":
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
# Check if already cached
if "priority_reservation_settings" not in _globals:
@@ -1706,6 +2088,7 @@ def __getattr__(name: str) -> Any:
# Lazy load logging_callback_manager instance
if name == "logging_callback_manager":
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
# Check if already cached
if "logging_callback_manager" not in _globals:
@@ -1717,14 +2100,68 @@ def __getattr__(name: str) -> Any:
# Lazy load _service_logger module
if name == "_service_logger":
from ._lazy_imports import _get_litellm_globals
+
_globals = _get_litellm_globals()
# Check if already cached
if "_service_logger" not in _globals:
# Import the module lazily
import litellm._service_logger
+
_globals["_service_logger"] = litellm._service_logger
return _globals["_service_logger"]
+ # Lazy load evals module functions
+ if name in [
+ "acreate_eval",
+ "alist_evals",
+ "aget_eval",
+ "aupdate_eval",
+ "adelete_eval",
+ "acancel_eval",
+ "create_eval",
+ "list_evals",
+ "get_eval",
+ "update_eval",
+ "delete_eval",
+ "cancel_eval",
+ "acreate_run",
+ "alist_runs",
+ "aget_run",
+ "acancel_run",
+ "adelete_run",
+ "create_run",
+ "list_runs",
+ "get_run",
+ "cancel_run",
+ "delete_run",
+ ]:
+ from litellm.evals.main import (
+ acreate_eval,
+ alist_evals,
+ aget_eval,
+ aupdate_eval,
+ adelete_eval,
+ acancel_eval,
+ create_eval,
+ list_evals,
+ get_eval,
+ update_eval,
+ delete_eval,
+ cancel_eval,
+ acreate_run,
+ alist_runs,
+ aget_run,
+ acancel_run,
+ adelete_run,
+ create_run,
+ list_runs,
+ get_run,
+ cancel_run,
+ delete_run,
+ )
+
+ return locals()[name]
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py
index 3bfeba2e394..3604506d406 100644
--- a/litellm/_lazy_imports.py
+++ b/litellm/_lazy_imports.py
@@ -55,7 +55,7 @@ from ._lazy_imports_registry import (
def _get_litellm_globals() -> dict:
"""
Get the globals dictionary of the litellm module.
-
+
This is where we cache imported attributes so we don't import them twice.
When you do `litellm.some_function`, it gets stored in this dictionary.
"""
@@ -65,12 +65,13 @@ def _get_litellm_globals() -> dict:
def _get_utils_globals() -> dict:
"""
Get the globals dictionary of the utils module.
-
+
This is where we cache imported attributes so we don't import them twice.
When you do `litellm.utils.some_function`, it gets stored in this dictionary.
"""
return sys.modules["litellm.utils"].__dict__
+
# These are special lazy loaders for things that are used internally
# They're separate from the main lazy import system because they have specific use cases
@@ -81,10 +82,10 @@ _default_encoding: Optional[Any] = None
def _get_default_encoding() -> Any:
"""
Lazily load and cache the default OpenAI encoding.
-
+
This avoids importing `litellm.litellm_core_utils.default_encoding` (and thus tiktoken)
at `litellm` import time. The encoding is cached after the first import.
-
+
This is used internally by utils.py functions that need the encoding but shouldn't
trigger its import during module load.
"""
@@ -103,10 +104,10 @@ _get_modified_max_tokens_func: Optional[Any] = None
def _get_modified_max_tokens() -> Any:
"""
Lazily load and cache the get_modified_max_tokens function.
-
+
This avoids importing `litellm.litellm_core_utils.token_counter` at `litellm` import time.
The function is cached after the first import.
-
+
This is used internally by utils.py functions that need the token counter but shouldn't
trigger its import during module load.
"""
@@ -127,10 +128,10 @@ _token_counter_new_func: Optional[Any] = None
def _get_token_counter_new() -> Any:
"""
Lazily load and cache the token_counter function (aliased as token_counter_new).
-
+
This avoids importing `litellm.litellm_core_utils.token_counter` at `litellm` import time.
The function is cached after the first import.
-
+
This is used internally by utils.py functions that need the token counter but shouldn't
trigger its import during module load.
"""
@@ -157,10 +158,10 @@ _LAZY_IMPORT_REGISTRY: Optional[dict[str, Callable[[str], Any]]] = None
def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
"""
Build the registry that maps attribute names to their handler functions.
-
+
This is called once, the first time someone accesses a lazy-loaded attribute.
After that, we just look up the handler function in this dictionary.
-
+
Returns:
Dictionary like {"ModelResponse": _lazy_import_utils, ...}
"""
@@ -199,17 +200,19 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
_LAZY_IMPORT_REGISTRY[name] = _lazy_import_llm_provider_logic
for name in UTILS_MODULE_NAMES:
_LAZY_IMPORT_REGISTRY[name] = _lazy_import_utils_module
-
+
return _LAZY_IMPORT_REGISTRY
-def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any:
+def _generic_lazy_import(
+ name: str, import_map: dict[str, tuple[str, str]], category: str
+) -> Any:
"""
Generic function that handles lazy importing for most attributes.
-
+
This is the workhorse function - it does the actual importing and caching.
Most handler functions just call this with their specific import map.
-
+
Steps:
1. Check if the name exists in the import map (if not, raise error)
2. Check if we've already imported it (if yes, return cached value)
@@ -218,7 +221,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
5. Get the attribute from the module
6. Cache it in _globals so we don't import again
7. Return it
-
+
Args:
name: The attribute name someone is trying to access (e.g., "ModelResponse")
import_map: Dictionary telling us where to find each attribute
@@ -228,19 +231,19 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# Step 1: Make sure this attribute exists in our map
if name not in import_map:
raise AttributeError(f"{category} lazy import: unknown attribute {name!r}")
-
+
# Step 2: Get the cache (where we store imported things)
_globals = _get_litellm_globals()
-
+
# Step 3: If we've already imported it, just return the cached version
if name in _globals:
return _globals[name]
-
+
# Step 4: Look up where to find this attribute
# The map tells us: (module_path, attribute_name)
# Example: (".utils", "ModelResponse") means "look in .utils module, get ModelResponse"
module_path, attr_name = import_map[name]
-
+
# Step 5: Import the module
# Python automatically caches modules in sys.modules, so calling this twice is fast
# If module_path starts with ".", it's a relative import (needs package="litellm")
@@ -249,14 +252,14 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
module = importlib.import_module(module_path, package="litellm")
else:
module = importlib.import_module(module_path)
-
+
# Step 6: Get the actual attribute from the module
# Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class
value = getattr(module, attr_name)
-
+
# Step 7: Cache it so we don't have to import again next time
_globals[name] = value
-
+
# Step 8: Return it
return value
@@ -268,6 +271,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# Most of them just call _generic_lazy_import with their specific import map.
# The registry (above) maps attribute names to these handler functions.
+
def _lazy_import_utils(name: str) -> Any:
"""Handler for utils module attributes (ModelResponse, token_counter, etc.)"""
return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils")
@@ -297,6 +301,7 @@ def _lazy_import_caching(name: str) -> Any:
"""Handler for caching classes (Cache, DualCache, RedisCache, etc.)"""
return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching")
+
def _lazy_import_dotprompt(name: str) -> Any:
"""Handler for dotprompt integration globals"""
return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt")
@@ -311,6 +316,7 @@ def _lazy_import_llm_configs(name: str) -> Any:
"""Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)"""
return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config")
+
def _lazy_import_litellm_logging(name: str) -> Any:
"""Handler for litellm_logging module (Logging, modify_integration)"""
return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging")
@@ -318,87 +324,91 @@ def _lazy_import_litellm_logging(name: str) -> Any:
def _lazy_import_llm_provider_logic(name: str) -> Any:
"""Handler for LLM provider logic functions (get_llm_provider, etc.)"""
- return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic")
+ return _generic_lazy_import(
+ name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic"
+ )
def _lazy_import_utils_module(name: str) -> Any:
"""
Handler for utils module lazy imports.
-
+
This uses a custom implementation because utils module needs to use
_get_utils_globals() instead of _get_litellm_globals() for caching.
"""
# Check if this attribute exists in our map
if name not in _UTILS_MODULE_IMPORT_MAP:
raise AttributeError(f"Utils module lazy import: unknown attribute {name!r}")
-
+
# Get the cache (where we store imported things) - use utils globals
_globals = _get_utils_globals()
-
+
# If we've already imported it, just return the cached version
if name in _globals:
return _globals[name]
-
+
# Look up where to find this attribute
module_path, attr_name = _UTILS_MODULE_IMPORT_MAP[name]
-
+
# Import the module
if module_path.startswith("."):
module = importlib.import_module(module_path, package="litellm")
else:
module = importlib.import_module(module_path)
-
+
# Get the actual attribute from the module
value = getattr(module, attr_name)
-
+
# Cache it so we don't have to import again next time
_globals[name] = value
-
+
# Return it
return value
+
# ============================================================================
# SPECIAL HANDLERS
# ============================================================================
# These handlers have custom logic that doesn't fit the generic pattern
+
def _lazy_import_llm_client_cache(name: str) -> Any:
"""
Handler for LLM client cache - has special logic for singleton instance.
-
+
This one is different because:
- "LLMClientCache" is the class itself
- "in_memory_llm_clients_cache" is a singleton instance of that class
So we need custom logic to handle both cases.
"""
_globals = _get_litellm_globals()
-
+
# If already cached, return it
if name in _globals:
return _globals[name]
-
+
# Import the class
module = importlib.import_module("litellm.caching.llm_caching_handler")
LLMClientCache = getattr(module, "LLMClientCache")
-
+
# If they want the class itself, return it
if name == "LLMClientCache":
_globals["LLMClientCache"] = LLMClientCache
return LLMClientCache
-
+
# If they want the singleton instance, create it (only once)
if name == "in_memory_llm_clients_cache":
instance = LLMClientCache()
_globals["in_memory_llm_clients_cache"] = instance
return instance
-
+
raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}")
def _lazy_import_http_handlers(name: str) -> Any:
"""
Handler for HTTP clients - has special logic for creating client instances.
-
+
This one is different because:
- These aren't just imports, they're actual client instances that need to be created
- They need configuration (timeout, etc.) from the module globals
@@ -413,14 +423,14 @@ def _lazy_import_http_handlers(name: str) -> Any:
# Get timeout from module config (if set)
timeout = _globals.get("request_timeout")
params = {"timeout": timeout, "client_alias": "module level aclient"}
-
+
# Create the client instance
provider_id = cast(Any, "litellm_module_level_client")
async_client = get_async_httpx_client(
llm_provider=provider_id,
params=params,
)
-
+
# Cache it so we don't create it again
_globals["module_level_aclient"] = async_client
return async_client
@@ -431,7 +441,7 @@ def _lazy_import_http_handlers(name: str) -> Any:
timeout = _globals.get("request_timeout")
sync_client = HTTPHandler(timeout=timeout)
-
+
# Cache it
_globals["module_level_client"] = sync_client
return sync_client
diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py
index a01fe9c11db..9164a3c8ae4 100644
--- a/litellm/_lazy_imports_registry.py
+++ b/litellm/_lazy_imports_registry.py
@@ -155,6 +155,7 @@ LLM_CONFIG_NAMES = (
"VertexAIRerankConfig",
"FireworksAIRerankConfig",
"VoyageRerankConfig",
+ "IBMWatsonXRerankConfig",
"ClarifaiConfig",
"AI21ChatConfig",
"LlamaAPIConfig",
@@ -166,6 +167,7 @@ LLM_CONFIG_NAMES = (
"OllamaConfig",
"SagemakerConfig",
"SagemakerChatConfig",
+ "SagemakerNovaConfig",
"CohereChatConfig",
"AnthropicMessagesConfig",
"AmazonAnthropicClaudeMessagesConfig",
@@ -213,11 +215,13 @@ LLM_CONFIG_NAMES = (
"TopazImageVariationConfig",
"OpenAITextCompletionConfig",
"GroqChatConfig",
+ "BedrockMantleChatConfig",
"A2AConfig",
"GenAIHubOrchestrationConfig",
"VoyageEmbeddingConfig",
"VoyageContextualEmbeddingConfig",
"InfinityEmbeddingConfig",
+ "PerplexityEmbeddingConfig",
"AzureAIStudioConfig",
"MistralConfig",
"OpenAIResponsesAPIConfig",
@@ -225,7 +229,11 @@ LLM_CONFIG_NAMES = (
"AzureOpenAIOSeriesResponsesAPIConfig",
"XAIResponsesAPIConfig",
"LiteLLMProxyResponsesAPIConfig",
+ "HostedVLLMResponsesAPIConfig",
"VolcEngineResponsesAPIConfig",
+ "PerplexityResponsesConfig",
+ "DatabricksResponsesAPIConfig",
+ "OpenRouterResponsesAPIConfig",
"GoogleAIStudioInteractionsConfig",
"OpenAIOSeriesConfig",
"AnthropicSkillsConfig",
@@ -670,6 +678,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
"FireworksAIRerankConfig",
),
"VoyageRerankConfig": (".llms.voyage.rerank.transformation", "VoyageRerankConfig"),
+ "IBMWatsonXRerankConfig": (
+ ".llms.watsonx.rerank.transformation",
+ "IBMWatsonXRerankConfig",
+ ),
"ClarifaiConfig": (".llms.clarifai.chat.transformation", "ClarifaiConfig"),
"AI21ChatConfig": (".llms.ai21.chat.transformation", "AI21ChatConfig"),
"LlamaAPIConfig": (".llms.meta_llama.chat.transformation", "LlamaAPIConfig"),
@@ -690,6 +702,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.sagemaker.chat.transformation",
"SagemakerChatConfig",
),
+ "SagemakerNovaConfig": (
+ ".llms.sagemaker.nova.transformation",
+ "SagemakerNovaConfig",
+ ),
"CohereChatConfig": (".llms.cohere.chat.transformation", "CohereChatConfig"),
"AnthropicMessagesConfig": (
".llms.anthropic.experimental_pass_through.messages.transformation",
@@ -851,6 +867,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
"OpenAITextCompletionConfig",
),
"GroqChatConfig": (".llms.groq.chat.transformation", "GroqChatConfig"),
+ "BedrockMantleChatConfig": (
+ ".llms.bedrock_mantle.chat.transformation",
+ "BedrockMantleChatConfig",
+ ),
"A2AConfig": (".llms.a2a.chat.transformation", "A2AConfig"),
"GenAIHubOrchestrationConfig": (
".llms.sap.chat.transformation",
@@ -868,6 +888,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.infinity.embedding.transformation",
"InfinityEmbeddingConfig",
),
+ "PerplexityEmbeddingConfig": (
+ ".llms.perplexity.embedding.transformation",
+ "PerplexityEmbeddingConfig",
+ ),
"AzureAIStudioConfig": (
".llms.azure_ai.chat.transformation",
"AzureAIStudioConfig",
@@ -893,6 +917,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.litellm_proxy.responses.transformation",
"LiteLLMProxyResponsesAPIConfig",
),
+ "HostedVLLMResponsesAPIConfig": (
+ ".llms.hosted_vllm.responses.transformation",
+ "HostedVLLMResponsesAPIConfig",
+ ),
"VolcEngineResponsesAPIConfig": (
".llms.volcengine.responses.transformation",
"VolcEngineResponsesAPIConfig",
@@ -901,6 +929,18 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.manus.responses.transformation",
"ManusResponsesAPIConfig",
),
+ "PerplexityResponsesConfig": (
+ ".llms.perplexity.responses.transformation",
+ "PerplexityResponsesConfig",
+ ),
+ "DatabricksResponsesAPIConfig": (
+ ".llms.databricks.responses.transformation",
+ "DatabricksResponsesAPIConfig",
+ ),
+ "OpenRouterResponsesAPIConfig": (
+ ".llms.openrouter.responses.transformation",
+ "OpenRouterResponsesAPIConfig",
+ ),
"GoogleAIStudioInteractionsConfig": (
".llms.gemini.interactions.transformation",
"GoogleAIStudioInteractionsConfig",
diff --git a/litellm/_logging.py b/litellm/_logging.py
index e222627e76c..5de9fbb3558 100644
--- a/litellm/_logging.py
+++ b/litellm/_logging.py
@@ -1,9 +1,14 @@
-import json
+import ast
import logging
import os
+import re
import sys
from datetime import datetime
from logging import Formatter
+from typing import Any, Dict, List, Optional
+
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
set_verbose = False
@@ -11,12 +16,155 @@ if set_verbose is True:
logging.warning(
"`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs."
)
+
+_ENABLE_SECRET_REDACTION = os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true"
+
+_REDACTED = "REDACTED"
+
+
+def _build_secret_patterns() -> re.Pattern:
+ patterns: List[str] = [
+ # AWS access key IDs
+ r"(?:AKIA|ASIA)[0-9A-Z]{16}",
+ # AWS secrets / session tokens / access key IDs (key=value)
+ r"(?:aws_secret_access_key|aws_session_token|aws_access_key_id)"
+ r"\s*[:=]\s*[A-Za-z0-9/+=]{20,}",
+ # Bearer tokens (OAuth, JWT, etc.)
+ r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*",
+ # Basic auth headers
+ r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}",
+ # OpenAI / Anthropic sk- prefixed keys
+ r"sk-[A-Za-z0-9\-_]{20,}",
+ # Generic api_key / api-key / apikey (handles 'key': 'value' dict repr)
+ r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}",
+ # x-api-key / api-key header values (handles 'key': 'value' dict repr)
+ r"(?:x-api-key|api-key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+",
+ # Anthropic internal header keys
+ r"x-ak-[A-Za-z0-9\-_]{20,}",
+ # Google API keys
+ r"AIza[0-9A-Za-z\-_]{35}",
+ # Password / secret params (handles key=value and 'key': 'value')
+ r"\w*(?:password|passwd|client_secret|secret_key|_secret)"
+ r"['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+",
+ # Database connection string credentials (scheme://user:pass@host)
+ r"(?<=://)[^\s'\"]*:[^\s'\"@]+(?=@)",
+ # Databricks personal access tokens
+ r"dapi[0-9a-f]{32}",
+ ]
+ return re.compile("|".join(patterns), re.IGNORECASE)
+
+
+_SECRET_RE = _build_secret_patterns()
+
+
+def _redact_string(value: str) -> str:
+ return _SECRET_RE.sub(_REDACTED, value)
+
+
+class SecretRedactionFilter(logging.Filter):
+ """Scrubs known secret/credential patterns from log records."""
+
+ _formatter = logging.Formatter()
+
+ def filter(self, record: logging.LogRecord) -> bool:
+ if not _ENABLE_SECRET_REDACTION:
+ return True
+
+ try:
+ record.msg = _redact_string(record.getMessage())
+ record.args = None
+ except Exception:
+ if isinstance(record.msg, str):
+ record.msg = _redact_string(record.msg)
+
+ # Redact exception tracebacks
+ if record.exc_info and record.exc_info[1] is not None:
+ try:
+ record.exc_text = _redact_string(
+ self._formatter.formatException(record.exc_info)
+ )
+ except Exception:
+ pass
+
+ # Redact extra fields passed via logger.debug("msg", extra={...})
+ for key, value in list(record.__dict__.items()):
+ if key not in _STANDARD_RECORD_ATTRS and isinstance(value, str):
+ setattr(record, key, _redact_string(value))
+
+ return True
+
+
+_secret_filter = SecretRedactionFilter()
+
+
json_logs = bool(os.getenv("JSON_LOGS", False))
# Create a handler for the logger (you may need to adapt this based on your needs)
log_level = os.getenv("LITELLM_LOG", "DEBUG")
numeric_level: str = getattr(logging, log_level.upper())
handler = logging.StreamHandler()
handler.setLevel(numeric_level)
+handler.addFilter(_secret_filter)
+
+
+def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
+ """
+ Try to parse a log message as JSON. Returns parsed dict if valid, else None.
+ Handles messages that are entirely valid JSON (e.g. json.dumps output).
+ Uses shared safe_json_loads for consistent error handling.
+ """
+ if not message or not isinstance(message, str):
+ return None
+ msg_stripped = message.strip()
+ if not (msg_stripped.startswith("{") or msg_stripped.startswith("[")):
+ return None
+ parsed = safe_json_loads(message, default=None)
+ if parsed is None or not isinstance(parsed, dict):
+ return None
+ return parsed
+
+
+def _try_parse_embedded_python_dict(message: str) -> Optional[Dict[str, Any]]:
+ """
+ Try to find and parse a Python dict repr (e.g. str(d) or repr(d)) embedded in
+ the message. Handles patterns like:
+ "get_available_deployment for model: X, Selected deployment: {'model_name': '...', ...} for model: X"
+ Uses ast.literal_eval for safe parsing. Returns the parsed dict or None.
+ """
+ if not message or not isinstance(message, str) or "{" not in message:
+ return None
+ i = 0
+ while i < len(message):
+ start = message.find("{", i)
+ if start == -1:
+ break
+ depth = 0
+ for j in range(start, len(message)):
+ c = message[j]
+ if c == "{":
+ depth += 1
+ elif c == "}":
+ depth -= 1
+ if depth == 0:
+ substr = message[start : j + 1]
+ try:
+ result = ast.literal_eval(substr)
+ if isinstance(result, dict) and len(result) > 0:
+ return result
+ except (ValueError, SyntaxError, TypeError):
+ pass
+ break
+ i = start + 1
+ return None
+
+
+# Standard LogRecord attribute names - used to identify 'extra' fields.
+# Derived at runtime so we automatically include version-specific attrs (e.g. taskName).
+def _get_standard_record_attrs() -> frozenset:
+ """Standard LogRecord attribute names - excludes extra keys from logger.debug(..., extra={...})."""
+ return frozenset(logging.LogRecord("", 0, "", 0, "", (), None).__dict__.keys())
+
+
+_STANDARD_RECORD_ATTRS = _get_standard_record_attrs()
class JsonFormatter(Formatter):
@@ -29,16 +177,31 @@ class JsonFormatter(Formatter):
return dt.isoformat()
def format(self, record):
- json_record = {
- "message": record.getMessage(),
+ message_str = record.getMessage()
+ json_record: Dict[str, Any] = {
+ "message": message_str,
"level": record.levelname,
"timestamp": self.formatTime(record),
}
- if record.exc_info:
- json_record["stacktrace"] = self.formatException(record.exc_info)
+ # Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties
+ parsed = _try_parse_json_message(message_str)
+ if parsed is None:
+ parsed = _try_parse_embedded_python_dict(message_str)
+ if parsed is not None:
+ for key, value in parsed.items():
+ if key not in json_record:
+ json_record[key] = value
- return json.dumps(json_record)
+ # Include extra attributes passed via logger.debug("msg", extra={...})
+ for key, value in record.__dict__.items():
+ if key not in _STANDARD_RECORD_ATTRS and key not in json_record:
+ json_record[key] = value
+
+ if record.exc_info:
+ json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info)
+
+ return safe_dumps(json_record)
# Function to set up exception handlers for JSON logging
@@ -46,6 +209,7 @@ def _setup_json_exception_handlers(formatter):
# Create a handler with JSON formatting for exceptions
error_handler = logging.StreamHandler()
error_handler.setFormatter(formatter)
+ error_handler.addFilter(_secret_filter)
# Setup excepthook for uncaught exceptions
def json_excepthook(exc_type, exc_value, exc_traceback):
@@ -69,6 +233,7 @@ def _setup_json_exception_handlers(formatter):
def async_json_exception_handler(loop, context):
exception = context.get("exception")
if exception:
+ exc_type = type(exception)
record = logging.LogRecord(
name="LiteLLM",
level=logging.ERROR,
@@ -76,7 +241,7 @@ def _setup_json_exception_handlers(formatter):
lineno=0,
msg=str(exception),
args=(),
- exc_info=None,
+ exc_info=(exc_type, exception, exception.__traceback__),
)
error_handler.handle(record)
else:
@@ -160,6 +325,7 @@ def _initialize_loggers_with_handler(handler: logging.Handler):
- Adds a handler to each logger
- Prevents bubbling to parent/root (critical to prevent duplicate JSON logs)
"""
+ handler.addFilter(_secret_filter)
for lg in _get_loggers_to_initialize():
lg.handlers.clear() # remove any existing handlers
lg.addHandler(handler) # add JSON formatter handler
@@ -169,15 +335,15 @@ def _initialize_loggers_with_handler(handler: logging.Handler):
def _get_uvicorn_json_log_config():
"""
Generate a uvicorn log_config dictionary that applies JSON formatting to all loggers.
-
+
This ensures that uvicorn's access logs, error logs, and all application logs
are formatted as JSON when json_logs is enabled.
"""
json_formatter_class = "litellm._logging.JsonFormatter"
-
+
# Use the module-level log_level variable for consistency
uvicorn_log_level = log_level.upper()
-
+
log_config = {
"version": 1,
"disable_existing_loggers": False,
@@ -222,7 +388,7 @@ def _get_uvicorn_json_log_config():
},
},
}
-
+
return log_config
diff --git a/litellm/_redis.py b/litellm/_redis.py
index a86ebd9ea9e..b754c1f4330 100644
--- a/litellm/_redis.py
+++ b/litellm/_redis.py
@@ -34,7 +34,12 @@ def _get_redis_kwargs():
"retry",
}
- include_args = ["url", "redis_connect_func", "gcp_service_account", "gcp_ssl_ca_certs"]
+ include_args = [
+ "url",
+ "redis_connect_func",
+ "gcp_service_account",
+ "gcp_ssl_ca_certs",
+ ]
available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
@@ -75,7 +80,9 @@ def _get_redis_cluster_kwargs(client=None):
available_args.append("ssl_cert_reqs")
available_args.append("ssl_check_hostname")
available_args.append("ssl_ca_certs")
- available_args.append("redis_connect_func") # Needed for sync clusters and IAM detection
+ available_args.append(
+ "redis_connect_func"
+ ) # Needed for sync clusters and IAM detection
available_args.append("gcp_service_account")
available_args.append("gcp_ssl_ca_certs")
available_args.append("max_connections")
@@ -103,10 +110,10 @@ def _redis_kwargs_from_environment():
def _generate_gcp_iam_access_token(service_account: str) -> str:
"""
Generate GCP IAM access token for Redis authentication.
-
+
Args:
service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com'
-
+
Returns:
Access token string for GCP IAM authentication
"""
@@ -117,11 +124,11 @@ def _generate_gcp_iam_access_token(service_account: str) -> str:
"google-cloud-iam is required for GCP IAM Redis authentication. "
"Install it with: pip install google-cloud-iam"
)
-
+
client = iam_credentials_v1.IAMCredentialsClient()
request = iam_credentials_v1.GenerateAccessTokenRequest(
name=service_account,
- scope=['https://www.googleapis.com/auth/cloud-platform'],
+ scope=["https://www.googleapis.com/auth/cloud-platform"],
)
response = client.generate_access_token(request=request)
return str(response.access_token)
@@ -133,14 +140,15 @@ def create_gcp_iam_redis_connect_func(
) -> Callable:
"""
Creates a custom Redis connection function for GCP IAM authentication.
-
+
Args:
service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com'
ssl_ca_certs: Path to SSL CA certificate file for secure connections
-
+
Returns:
A connection function that can be used with Redis clients
"""
+
def iam_connect(self):
"""Initialize the connection and authenticate using GCP IAM"""
from redis.exceptions import (
@@ -148,25 +156,25 @@ def create_gcp_iam_redis_connect_func(
AuthenticationWrongNumberOfArgsError,
)
from redis.utils import str_if_bytes
-
+
self._parser.on_connect(self)
-
+
auth_args = (_generate_gcp_iam_access_token(service_account),)
self.send_command("AUTH", *auth_args, check_health=False)
-
+
try:
auth_response = self.read_response()
except AuthenticationWrongNumberOfArgsError:
# Fallback to password auth if IAM fails
- if hasattr(self, 'password') and self.password:
+ if hasattr(self, "password") and self.password:
self.send_command("AUTH", self.password, check_health=False)
auth_response = self.read_response()
else:
raise
-
+
if str_if_bytes(auth_response) != "OK":
raise AuthenticationError("GCP IAM authentication failed")
-
+
return iam_connect
@@ -178,22 +186,20 @@ def get_redis_url_from_environment():
raise ValueError(
"Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified for Redis."
)
-
+
if "REDIS_SSL" in os.environ and os.environ["REDIS_SSL"].lower() == "true":
redis_protocol = "rediss"
else:
redis_protocol = "redis"
-
+
# Build authentication part of URL
auth_part = ""
if "REDIS_USERNAME" in os.environ and "REDIS_PASSWORD" in os.environ:
auth_part = f"{os.environ['REDIS_USERNAME']}:{os.environ['REDIS_PASSWORD']}@"
elif "REDIS_PASSWORD" in os.environ:
auth_part = f"{os.environ['REDIS_PASSWORD']}@"
-
- return (
- f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
- )
+
+ return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
def _get_redis_client_logic(**env_overrides):
@@ -241,22 +247,27 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs["service_name"] = _service_name
# Handle GCP IAM authentication
- _gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
- _gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
-
+ _gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str(
+ "REDIS_GCP_SERVICE_ACCOUNT"
+ )
+ _gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str(
+ "REDIS_GCP_SSL_CA_CERTS"
+ )
+
if _gcp_service_account is not None:
- verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
+ verbose_logger.debug(
+ "Setting up GCP IAM authentication for Redis with service account."
+ )
redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func(
- service_account=_gcp_service_account,
- ssl_ca_certs=_gcp_ssl_ca_certs
+ service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
)
# Store GCP service account in redis_connect_func for async cluster access
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
-
+
# Remove GCP-specific kwargs that shouldn't be passed to Redis client
redis_kwargs.pop("gcp_service_account", None)
redis_kwargs.pop("gcp_ssl_ca_certs", None)
-
+
# Only enable SSL if explicitly requested AND SSL CA certs are provided
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
@@ -377,10 +388,13 @@ def get_redis_client(**env_overrides):
def get_redis_async_client(
- connection_pool: Optional[async_redis.BlockingConnectionPool] = None, **env_overrides,
+ connection_pool: Optional[async_redis.BlockingConnectionPool] = None,
+ **env_overrides,
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
redis_kwargs = _get_redis_client_logic(**env_overrides)
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
+ if connection_pool is not None:
+ return async_redis.Redis(connection_pool=connection_pool)
args = _get_redis_url_kwargs(client=async_redis.Redis.from_url)
url_kwargs = {}
for arg in redis_kwargs:
@@ -409,39 +423,50 @@ def get_redis_async_client(
# Get GCP service account - first try from redis_connect_func, then from environment
gcp_service_account = None
- if redis_connect_func and hasattr(redis_connect_func, '_gcp_service_account'):
+ if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
gcp_service_account = redis_connect_func._gcp_service_account
else:
- gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
-
- verbose_logger.debug(f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}")
-
+ gcp_service_account = redis_kwargs.get(
+ "gcp_service_account"
+ ) or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
+
+ verbose_logger.debug(
+ f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}"
+ )
+
# If GCP IAM is configured (indicated by redis_connect_func), generate access token and use as password
if redis_connect_func and gcp_service_account:
- verbose_logger.debug("DEBUG: Generating IAM token for service account (value not logged for security reasons)")
+ verbose_logger.debug(
+ "DEBUG: Generating IAM token for service account (value not logged for security reasons)"
+ )
try:
# Generate IAM access token using the helper function
access_token = _generate_gcp_iam_access_token(gcp_service_account)
cluster_kwargs["password"] = access_token
- verbose_logger.debug("DEBUG: Successfully generated GCP IAM access token for async Redis cluster")
+ verbose_logger.debug(
+ "DEBUG: Successfully generated GCP IAM access token for async Redis cluster"
+ )
except Exception as e:
verbose_logger.error(f"Failed to generate GCP IAM access token: {e}")
from redis.exceptions import AuthenticationError
+
raise AuthenticationError("Failed to generate GCP IAM access token")
else:
- verbose_logger.debug(f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}")
-
+ verbose_logger.debug(
+ f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}"
+ )
+
new_startup_nodes: List[ClusterNode] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
cluster_kwargs.pop("startup_nodes", None)
-
+
# Create async RedisCluster with IAM token as password if available
cluster_client = async_redis.RedisCluster(
startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore
)
-
+
return cluster_client
# Check for Redis Sentinel
@@ -461,9 +486,19 @@ def get_redis_connection_pool(**env_overrides):
redis_kwargs = _get_redis_client_logic(**env_overrides)
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
- return async_redis.BlockingConnectionPool.from_url(
- timeout=REDIS_CONNECTION_POOL_TIMEOUT, url=redis_kwargs["url"]
- )
+ pool_kwargs = {
+ "timeout": REDIS_CONNECTION_POOL_TIMEOUT,
+ "url": redis_kwargs["url"],
+ }
+ if "max_connections" in redis_kwargs:
+ try:
+ pool_kwargs["max_connections"] = int(redis_kwargs["max_connections"])
+ except (TypeError, ValueError):
+ verbose_logger.warning(
+ "REDIS: invalid max_connections value %r, ignoring",
+ redis_kwargs["max_connections"],
+ )
+ return async_redis.BlockingConnectionPool.from_url(**pool_kwargs)
connection_class = async_redis.Connection
if "ssl" in redis_kwargs:
connection_class = async_redis.SSLConnection
@@ -474,6 +509,7 @@ def get_redis_connection_pool(**env_overrides):
timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs
)
+
def _pretty_print_redis_config(redis_kwargs: dict) -> None:
"""Pretty print the Redis configuration using rich with sensitive data masking"""
try:
@@ -483,6 +519,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
+
if not verbose_logger.isEnabledFor(logging.DEBUG):
return
@@ -490,7 +527,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
# Initialize the sensitive data masker
masker = SensitiveDataMasker()
-
+
# Mask sensitive data in redis_kwargs
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
@@ -522,7 +559,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
value_str = str(value)
else:
value_str = str(value)
-
+
config_table.add_row(key, value_str)
# Determine connection type
@@ -559,4 +596,3 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
verbose_logger.info(f"Redis configuration: {masked_redis_kwargs}")
except Exception as e:
verbose_logger.error(f"Error pretty printing Redis configuration: {e}")
-
diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py
index b67d0d86063..1a3be203fec 100644
--- a/litellm/_service_logger.py
+++ b/litellm/_service_logger.py
@@ -312,10 +312,12 @@ class ServiceLogging(CustomLogger):
_duration, type(_duration)
)
) # invalid _duration value
+ # Batch polling callbacks (check_batch_cost) don't include call_type in kwargs.
+ # Use .get() to avoid KeyError.
await self.async_service_success_hook(
service=ServiceTypes.LITELLM,
duration=_duration,
- call_type=kwargs["call_type"],
+ call_type=kwargs.get("call_type", "unknown"),
)
except Exception as e:
raise e
diff --git a/litellm/a2a_protocol/__init__.py b/litellm/a2a_protocol/__init__.py
index d8d349bb98a..85c03687e25 100644
--- a/litellm/a2a_protocol/__init__.py
+++ b/litellm/a2a_protocol/__init__.py
@@ -39,6 +39,12 @@ Example usage (class-based):
"""
from litellm.a2a_protocol.client import A2AClient
+from litellm.a2a_protocol.exceptions import (
+ A2AAgentCardError,
+ A2AConnectionError,
+ A2AError,
+ A2ALocalhostURLError,
+)
from litellm.a2a_protocol.main import (
aget_agent_card,
asend_message,
@@ -49,11 +55,19 @@ from litellm.a2a_protocol.main import (
from litellm.types.agents import LiteLLMSendMessageResponse
__all__ = [
+ # Client
"A2AClient",
+ # Functions
"asend_message",
"send_message",
"asend_message_streaming",
"aget_agent_card",
"create_a2a_client",
+ # Response types
"LiteLLMSendMessageResponse",
+ # Exceptions
+ "A2AError",
+ "A2AConnectionError",
+ "A2AAgentCardError",
+ "A2ALocalhostURLError",
]
diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py
index 7c4c5af149d..4c5dd3e3ba6 100644
--- a/litellm/a2a_protocol/card_resolver.py
+++ b/litellm/a2a_protocol/card_resolver.py
@@ -7,6 +7,7 @@ Extends the A2A SDK's card resolver to support multiple well-known paths.
from typing import TYPE_CHECKING, Any, Dict, Optional
from litellm._logging import verbose_logger
+from litellm.constants import LOCALHOST_URL_PATTERNS
if TYPE_CHECKING:
from a2a.types import AgentCard
@@ -26,15 +27,61 @@ except ImportError:
pass
+def is_localhost_or_internal_url(url: Optional[str]) -> bool:
+ """
+ Check if a URL is a localhost or internal URL.
+
+ This detects common development URLs that are accidentally left in
+ agent cards when deploying to production.
+
+ Args:
+ url: The URL to check
+
+ Returns:
+ True if the URL is localhost/internal
+ """
+ if not url:
+ return False
+
+ url_lower = url.lower()
+
+ return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS)
+
+
+def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard":
+ """
+ Fix the agent card URL if it contains a localhost/internal address.
+
+ Many A2A agents are deployed with agent cards that contain internal URLs
+ like "http://0.0.0.0:8001/" or "http://localhost:8000/". This function
+ replaces such URLs with the provided base_url.
+
+ Args:
+ agent_card: The agent card to fix
+ base_url: The base URL to use as replacement
+
+ Returns:
+ The agent card with the URL fixed if necessary
+ """
+ card_url = getattr(agent_card, "url", None)
+
+ if card_url and is_localhost_or_internal_url(card_url):
+ # Normalize base_url to ensure it ends with /
+ fixed_url = base_url.rstrip("/") + "/"
+ agent_card.url = fixed_url
+
+ return agent_card
+
+
class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
"""
Custom A2A card resolver that supports multiple well-known paths.
-
+
Extends the base A2ACardResolver to try both:
- /.well-known/agent-card.json (standard)
- /.well-known/agent.json (previous/alternative)
"""
-
+
async def get_agent_card(
self,
relative_card_path: Optional[str] = None,
@@ -42,17 +89,17 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
) -> "AgentCard":
"""
Fetch the agent card, trying multiple well-known paths.
-
+
First tries the standard path, then falls back to the previous path.
-
+
Args:
relative_card_path: Optional path to the agent card endpoint.
If None, tries both well-known paths.
http_kwargs: Optional dictionary of keyword arguments to pass to httpx.get
-
+
Returns:
AgentCard from the A2A agent
-
+
Raises:
A2AClientHTTPError or A2AClientJSONError if both paths fail
"""
@@ -62,13 +109,13 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
relative_card_path=relative_card_path,
http_kwargs=http_kwargs,
)
-
+
# Try both well-known paths
paths = [
AGENT_CARD_WELL_KNOWN_PATH,
PREV_AGENT_CARD_WELL_KNOWN_PATH,
]
-
+
last_error = None
for path in paths:
try:
@@ -85,11 +132,11 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
)
last_error = e
continue
-
+
# If we get here, all paths failed - re-raise the last error
if last_error is not None:
raise last_error
-
+
# This shouldn't happen, but just in case
raise Exception(
f"Failed to fetch agent card from {self.base_url}. "
diff --git a/litellm/a2a_protocol/client.py b/litellm/a2a_protocol/client.py
index 31f7c3b6a90..05e21284af1 100644
--- a/litellm/a2a_protocol/client.py
+++ b/litellm/a2a_protocol/client.py
@@ -103,5 +103,7 @@ class A2AClient:
from litellm.a2a_protocol.main import asend_message_streaming
a2a_client = await self._get_client()
- async for chunk in asend_message_streaming(a2a_client=a2a_client, request=request):
+ async for chunk in asend_message_streaming(
+ a2a_client=a2a_client, request=request
+ ):
yield chunk
diff --git a/litellm/a2a_protocol/cost_calculator.py b/litellm/a2a_protocol/cost_calculator.py
index f3e84c5b84d..f64174f8be5 100644
--- a/litellm/a2a_protocol/cost_calculator.py
+++ b/litellm/a2a_protocol/cost_calculator.py
@@ -97,7 +97,11 @@ class A2ACostCalculator:
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
# Calculate costs
- input_cost = prompt_tokens * (float(input_cost_per_token) if input_cost_per_token else 0.0)
- output_cost = completion_tokens * (float(output_cost_per_token) if output_cost_per_token else 0.0)
+ input_cost = prompt_tokens * (
+ float(input_cost_per_token) if input_cost_per_token else 0.0
+ )
+ output_cost = completion_tokens * (
+ float(output_cost_per_token) if output_cost_per_token else 0.0
+ )
return input_cost + output_cost
diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py
new file mode 100644
index 00000000000..49dbb22b158
--- /dev/null
+++ b/litellm/a2a_protocol/exception_mapping_utils.py
@@ -0,0 +1,203 @@
+"""
+A2A Protocol Exception Mapping Utils.
+
+Maps A2A SDK exceptions to LiteLLM A2A exception types.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+from litellm._logging import verbose_logger
+from litellm.a2a_protocol.card_resolver import (
+ fix_agent_card_url,
+ is_localhost_or_internal_url,
+)
+from litellm.a2a_protocol.exceptions import (
+ A2AAgentCardError,
+ A2AConnectionError,
+ A2AError,
+ A2ALocalhostURLError,
+)
+from litellm.constants import CONNECTION_ERROR_PATTERNS
+
+if TYPE_CHECKING:
+ from a2a.client import A2AClient as A2AClientType
+
+
+# Runtime import
+A2A_SDK_AVAILABLE = False
+try:
+ from a2a.client import A2AClient as _A2AClient # type: ignore[no-redef]
+
+ A2A_SDK_AVAILABLE = True
+except ImportError:
+ _A2AClient = None # type: ignore[assignment, misc]
+
+
+class A2AExceptionCheckers:
+ """
+ Helper class for checking various A2A error conditions.
+ """
+
+ @staticmethod
+ def is_connection_error(error_str: str) -> bool:
+ """
+ Check if an error string indicates a connection error.
+
+ Args:
+ error_str: The error string to check
+
+ Returns:
+ True if the error indicates a connection issue
+ """
+ if not isinstance(error_str, str):
+ return False
+
+ error_str_lower = error_str.lower()
+ return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS)
+
+ @staticmethod
+ def is_localhost_url(url: Optional[str]) -> bool:
+ """
+ Check if a URL is a localhost/internal URL.
+
+ Args:
+ url: The URL to check
+
+ Returns:
+ True if the URL is localhost/internal
+ """
+ return is_localhost_or_internal_url(url)
+
+ @staticmethod
+ def is_agent_card_error(error_str: str) -> bool:
+ """
+ Check if an error string indicates an agent card error.
+
+ Args:
+ error_str: The error string to check
+
+ Returns:
+ True if the error is related to agent card fetching/parsing
+ """
+ if not isinstance(error_str, str):
+ return False
+
+ error_str_lower = error_str.lower()
+ agent_card_patterns = [
+ "agent card",
+ "agent-card",
+ ".well-known",
+ "card not found",
+ "invalid agent",
+ ]
+ return any(pattern in error_str_lower for pattern in agent_card_patterns)
+
+
+def map_a2a_exception(
+ original_exception: Exception,
+ card_url: Optional[str] = None,
+ api_base: Optional[str] = None,
+ model: Optional[str] = None,
+) -> Exception:
+ """
+ Map an A2A SDK exception to a LiteLLM A2A exception type.
+
+ Args:
+ original_exception: The original exception from the A2A SDK
+ card_url: The URL from the agent card (if available)
+ api_base: The original API base URL
+ model: The model/agent name
+
+ Returns:
+ A mapped LiteLLM A2A exception
+
+ Raises:
+ A2ALocalhostURLError: If the error is a connection error to a localhost URL
+ A2AConnectionError: If the error is a general connection error
+ A2AAgentCardError: If the error is related to agent card issues
+ A2AError: For other A2A-related errors
+ """
+ error_str = str(original_exception)
+
+ # Check for localhost URL connection error (special case - retryable)
+ if (
+ card_url
+ and api_base
+ and A2AExceptionCheckers.is_localhost_url(card_url)
+ and A2AExceptionCheckers.is_connection_error(error_str)
+ ):
+ raise A2ALocalhostURLError(
+ localhost_url=card_url,
+ base_url=api_base,
+ original_error=original_exception,
+ model=model,
+ )
+
+ # Check for agent card errors
+ if A2AExceptionCheckers.is_agent_card_error(error_str):
+ raise A2AAgentCardError(
+ message=error_str,
+ url=api_base,
+ model=model,
+ )
+
+ # Check for general connection errors
+ if A2AExceptionCheckers.is_connection_error(error_str):
+ raise A2AConnectionError(
+ message=error_str,
+ url=card_url or api_base,
+ model=model,
+ )
+
+ # Default: wrap in generic A2AError
+ raise A2AError(
+ message=error_str,
+ model=model,
+ )
+
+
+def handle_a2a_localhost_retry(
+ error: A2ALocalhostURLError,
+ agent_card: Any,
+ a2a_client: "A2AClientType",
+ is_streaming: bool = False,
+) -> "A2AClientType":
+ """
+ Handle A2ALocalhostURLError by fixing the URL and creating a new client.
+
+ This is called when we catch an A2ALocalhostURLError and want to retry
+ with the corrected URL.
+
+ Args:
+ error: The localhost URL error
+ agent_card: The agent card object to fix
+ a2a_client: The current A2A client
+ is_streaming: Whether this is a streaming request (for logging)
+
+ Returns:
+ A new A2A client with the fixed URL
+
+ Raises:
+ ImportError: If the A2A SDK is not installed
+ """
+ if not A2A_SDK_AVAILABLE or _A2AClient is None:
+ raise ImportError(
+ "A2A SDK is required for localhost retry handling. "
+ "Install it with: pip install a2a"
+ )
+
+ request_type = "streaming " if is_streaming else ""
+ verbose_logger.warning(
+ f"A2A {request_type}request to '{error.localhost_url}' failed: {error.original_error}. "
+ f"Agent card contains localhost/internal URL. "
+ f"Retrying with base_url '{error.base_url}'."
+ )
+
+ # Fix the agent card URL
+ fix_agent_card_url(agent_card, error.base_url)
+
+ # Create a new client with the fixed agent card (transport caches URL)
+ return _A2AClient(
+ httpx_client=a2a_client._transport.httpx_client, # type: ignore[union-attr]
+ agent_card=agent_card,
+ )
diff --git a/litellm/a2a_protocol/exceptions.py b/litellm/a2a_protocol/exceptions.py
new file mode 100644
index 00000000000..546b23105be
--- /dev/null
+++ b/litellm/a2a_protocol/exceptions.py
@@ -0,0 +1,150 @@
+"""
+A2A Protocol Exceptions.
+
+Custom exception types for A2A protocol operations, following LiteLLM's exception pattern.
+"""
+
+from typing import Optional
+
+import httpx
+
+
+class A2AError(Exception):
+ """
+ Base exception for A2A protocol errors.
+
+ Follows the same pattern as LiteLLM's main exceptions.
+ """
+
+ def __init__(
+ self,
+ message: str,
+ status_code: int = 500,
+ llm_provider: str = "a2a_agent",
+ model: Optional[str] = None,
+ response: Optional[httpx.Response] = None,
+ litellm_debug_info: Optional[str] = None,
+ max_retries: Optional[int] = None,
+ num_retries: Optional[int] = None,
+ ):
+ self.status_code = status_code
+ self.message = f"litellm.A2AError: {message}"
+ self.llm_provider = llm_provider
+ self.model = model
+ self.litellm_debug_info = litellm_debug_info
+ self.max_retries = max_retries
+ self.num_retries = num_retries
+ self.response = response or httpx.Response(
+ status_code=self.status_code,
+ request=httpx.Request(method="POST", url="https://litellm.ai"),
+ )
+ super().__init__(self.message)
+
+ def __str__(self) -> str:
+ _message = self.message
+ if self.num_retries:
+ _message += f" LiteLLM Retried: {self.num_retries} times"
+ if self.max_retries:
+ _message += f", LiteLLM Max Retries: {self.max_retries}"
+ return _message
+
+ def __repr__(self) -> str:
+ return self.__str__()
+
+
+class A2AConnectionError(A2AError):
+ """
+ Raised when connection to an A2A agent fails.
+
+ This typically occurs when:
+ - The agent is unreachable
+ - The agent card contains a localhost/internal URL
+ - Network issues prevent connection
+ """
+
+ def __init__(
+ self,
+ message: str,
+ url: Optional[str] = None,
+ model: Optional[str] = None,
+ response: Optional[httpx.Response] = None,
+ litellm_debug_info: Optional[str] = None,
+ max_retries: Optional[int] = None,
+ num_retries: Optional[int] = None,
+ ):
+ self.url = url
+ super().__init__(
+ message=message,
+ status_code=503,
+ llm_provider="a2a_agent",
+ model=model,
+ response=response,
+ litellm_debug_info=litellm_debug_info,
+ max_retries=max_retries,
+ num_retries=num_retries,
+ )
+
+
+class A2AAgentCardError(A2AError):
+ """
+ Raised when there's an issue with the agent card.
+
+ This includes:
+ - Failed to fetch agent card
+ - Invalid agent card format
+ - Missing required fields
+ """
+
+ def __init__(
+ self,
+ message: str,
+ url: Optional[str] = None,
+ model: Optional[str] = None,
+ response: Optional[httpx.Response] = None,
+ litellm_debug_info: Optional[str] = None,
+ ):
+ self.url = url
+ super().__init__(
+ message=message,
+ status_code=404,
+ llm_provider="a2a_agent",
+ model=model,
+ response=response,
+ litellm_debug_info=litellm_debug_info,
+ )
+
+
+class A2ALocalhostURLError(A2AConnectionError):
+ """
+ Raised when an agent card contains a localhost/internal URL.
+
+ Many A2A agents are deployed with agent cards that contain internal URLs
+ like "http://0.0.0.0:8001/" or "http://localhost:8000/". This error
+ indicates that the URL needs to be corrected and the request should be retried.
+
+ Attributes:
+ localhost_url: The localhost/internal URL found in the agent card
+ base_url: The public base URL that should be used instead
+ original_error: The original connection error that was raised
+ """
+
+ def __init__(
+ self,
+ localhost_url: str,
+ base_url: str,
+ original_error: Optional[Exception] = None,
+ model: Optional[str] = None,
+ ):
+ self.localhost_url = localhost_url
+ self.base_url = base_url
+ self.original_error = original_error
+
+ message = (
+ f"Agent card contains localhost/internal URL '{localhost_url}'. "
+ f"Retrying with base URL '{base_url}'."
+ )
+ super().__init__(
+ message=message,
+ url=localhost_url,
+ model=model,
+ )
diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py
index 1916b04454a..c3d2e415237 100644
--- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py
+++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py
@@ -50,30 +50,28 @@ class A2ACompletionBridgeHandler:
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider
)
-
+
# If provider config exists, use it
if a2a_provider_config is not None:
if api_base is None:
raise ValueError(f"api_base is required for {custom_llm_provider}")
-
- verbose_logger.info(
- f"A2A: Using provider config for {custom_llm_provider}"
- )
-
+
+ verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider}")
+
response_data = await a2a_provider_config.handle_non_streaming(
request_id=request_id,
params=params,
api_base=api_base,
)
-
+
return response_data
-
+
# Extract message from params
message = params.get("message", {})
# Transform A2A message to OpenAI format
- openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(
- message
+ openai_messages = (
+ A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
@@ -100,7 +98,8 @@ class A2ACompletionBridgeHandler:
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
- k: v for k, v in litellm_params.items()
+ k: v
+ for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
@@ -109,9 +108,11 @@ class A2ACompletionBridgeHandler:
response = await litellm.acompletion(**completion_params)
# Transform response to A2A format
- a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
- response=response,
- request_id=request_id,
+ a2a_response = (
+ A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
+ response=response,
+ request_id=request_id,
+ )
)
verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}")
@@ -148,25 +149,25 @@ class A2ACompletionBridgeHandler:
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider
)
-
+
# If provider config exists, use it
if a2a_provider_config is not None:
if api_base is None:
raise ValueError(f"api_base is required for {custom_llm_provider}")
-
+
verbose_logger.info(
f"A2A: Using provider config for {custom_llm_provider} (streaming)"
)
-
+
async for chunk in a2a_provider_config.handle_streaming(
request_id=request_id,
params=params,
api_base=api_base,
):
yield chunk
-
+
return
-
+
# Extract message from params
message = params.get("message", {})
@@ -177,8 +178,8 @@ class A2ACompletionBridgeHandler:
)
# Transform A2A message to OpenAI format
- openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(
- message
+ openai_messages = (
+ A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
@@ -205,7 +206,8 @@ class A2ACompletionBridgeHandler:
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
- k: v for k, v in litellm_params.items()
+ k: v
+ for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
@@ -244,9 +246,11 @@ class A2ACompletionBridgeHandler:
# Emit artifact update with accumulated content
if accumulated_text:
- artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event(
- ctx=ctx,
- text=accumulated_text,
+ artifact_event = (
+ A2ACompletionBridgeTransformation.create_artifact_update_event(
+ ctx=ctx,
+ text=accumulated_text,
+ )
)
yield artifact_event
diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py
index bbe7daa9fc4..8a03569f689 100644
--- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py
+++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py
@@ -124,9 +124,7 @@ class A2ACompletionBridgeTransformation:
},
}
- verbose_logger.debug(
- f"OpenAI -> A2A transform: content_length={len(content)}"
- )
+ verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}")
return a2a_response
diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py
index b326f9e7ed5..c86549da77a 100644
--- a/litellm/a2a_protocol/main.py
+++ b/litellm/a2a_protocol/main.py
@@ -24,11 +24,7 @@ from litellm.utils import client
if TYPE_CHECKING:
from a2a.client import A2AClient as A2AClientType
- from a2a.types import (
- AgentCard,
- SendMessageRequest,
- SendStreamingMessageRequest,
- )
+ from a2a.types import AgentCard, SendMessageRequest, SendStreamingMessageRequest
# Runtime imports with availability check
A2A_SDK_AVAILABLE = False
@@ -44,6 +40,11 @@ except ImportError:
# Import our custom card resolver that supports multiple well-known paths
from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
+from litellm.a2a_protocol.exception_mapping_utils import (
+ handle_a2a_localhost_retry,
+ map_a2a_exception,
+)
+from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
# Use our custom resolver instead of the default A2A SDK resolver
A2ACardResolver = LiteLLMA2ACardResolver
@@ -126,6 +127,84 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str:
return agent_name
+async def _send_message_via_completion_bridge(
+ request: "SendMessageRequest",
+ custom_llm_provider: str,
+ api_base: Optional[str],
+ litellm_params: Dict[str, Any],
+) -> LiteLLMSendMessageResponse:
+ """
+ Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
+
+ Requires request; api_base is optional for providers that derive endpoint from model.
+ """
+ verbose_logger.info(
+ f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}"
+ )
+
+ from litellm.a2a_protocol.litellm_completion_bridge.handler import (
+ A2ACompletionBridgeHandler,
+ )
+
+ params = (
+ request.params.model_dump(mode="json")
+ if hasattr(request.params, "model_dump")
+ else dict(request.params)
+ )
+
+ response_dict = await A2ACompletionBridgeHandler.handle_non_streaming(
+ request_id=str(request.id),
+ params=params,
+ litellm_params=litellm_params,
+ api_base=api_base,
+ )
+
+ return LiteLLMSendMessageResponse.from_dict(response_dict)
+
+
+async def _execute_a2a_send_with_retry(
+ a2a_client: Any,
+ request: Any,
+ agent_card: Any,
+ card_url: Optional[str],
+ api_base: Optional[str],
+ agent_name: Optional[str],
+) -> Any:
+ """Send an A2A message with retry logic for localhost URL errors."""
+ a2a_response = None
+ for _ in range(2): # max 2 attempts: original + 1 retry
+ try:
+ a2a_response = await a2a_client.send_message(request)
+ break # success, exit retry loop
+ except A2ALocalhostURLError as e:
+ a2a_client = handle_a2a_localhost_retry(
+ error=e,
+ agent_card=agent_card,
+ a2a_client=a2a_client,
+ is_streaming=False,
+ )
+ card_url = agent_card.url if agent_card else None
+ except Exception as e:
+ try:
+ map_a2a_exception(e, card_url, api_base, model=agent_name)
+ except A2ALocalhostURLError as localhost_err:
+ a2a_client = handle_a2a_localhost_retry(
+ error=localhost_err,
+ agent_card=agent_card,
+ a2a_client=a2a_client,
+ is_streaming=False,
+ )
+ card_url = agent_card.url if agent_card else None
+ continue
+ except Exception:
+ raise
+ if a2a_response is None:
+ raise RuntimeError(
+ "A2A send_message failed: no response received after retry attempts."
+ )
+ return a2a_response
+
+
@client
async def asend_message(
a2a_client: Optional["A2AClientType"] = None,
@@ -133,6 +212,7 @@ async def asend_message(
api_base: Optional[str] = None,
litellm_params: Optional[Dict[str, Any]] = None,
agent_id: Optional[str] = None,
+ agent_extra_headers: Optional[Dict[str, str]] = None,
**kwargs: Any,
) -> LiteLLMSendMessageResponse:
"""
@@ -188,39 +268,21 @@ async def asend_message(
```
"""
litellm_params = litellm_params or {}
+ logging_obj = kwargs.get("litellm_logging_obj")
+ trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None
custom_llm_provider = litellm_params.get("custom_llm_provider")
# Route through completion bridge if custom_llm_provider is set
if custom_llm_provider:
if request is None:
raise ValueError("request is required for completion bridge")
- # api_base is optional for providers that derive endpoint from model (e.g., bedrock/agentcore)
-
- verbose_logger.info(
- f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}"
- )
-
- from litellm.a2a_protocol.litellm_completion_bridge.handler import (
- A2ACompletionBridgeHandler,
- )
-
- # Extract params from request
- params = (
- request.params.model_dump(mode="json")
- if hasattr(request.params, "model_dump")
- else dict(request.params)
- )
-
- response_dict = await A2ACompletionBridgeHandler.handle_non_streaming(
- request_id=str(request.id),
- params=params,
- litellm_params=litellm_params,
+ return await _send_message_via_completion_bridge(
+ request=request,
+ custom_llm_provider=custom_llm_provider,
api_base=api_base,
+ litellm_params=litellm_params,
)
- # Convert to LiteLLMSendMessageResponse
- return LiteLLMSendMessageResponse.from_dict(response_dict)
-
# Standard A2A client flow
if request is None:
raise ValueError("request is required")
@@ -231,11 +293,16 @@ async def asend_message(
raise ValueError(
"Either a2a_client or api_base is required for standard A2A flow"
)
- trace_id = str(uuid.uuid4())
- extra_headers = {"X-LiteLLM-Trace-Id": trace_id}
+ trace_id = trace_id or str(uuid.uuid4())
+ extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id}
if agent_id:
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
- a2a_client = await create_a2a_client(base_url=api_base, extra_headers=extra_headers)
+ # Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones)
+ if agent_extra_headers:
+ extra_headers.update(agent_extra_headers)
+ a2a_client = await create_a2a_client(
+ base_url=api_base, extra_headers=extra_headers
+ )
# Type assertion: a2a_client is guaranteed to be non-None here
assert a2a_client is not None
@@ -244,7 +311,29 @@ async def asend_message(
verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}")
- a2a_response = await a2a_client.send_message(request)
+ # Get agent card URL for localhost retry logic
+ agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(
+ a2a_client, "agent_card", None
+ )
+ card_url = getattr(agent_card, "url", None) if agent_card else None
+
+ context_id = trace_id or str(uuid.uuid4())
+ message = request.params.message
+ if isinstance(message, dict):
+ if message.get("context_id") is None:
+ message["context_id"] = context_id
+ else:
+ if getattr(message, "context_id", None) is None:
+ message.context_id = context_id
+
+ a2a_response = await _execute_a2a_send_with_retry(
+ a2a_client=a2a_client,
+ request=request,
+ agent_card=agent_card,
+ card_url=card_url,
+ api_base=api_base,
+ agent_name=agent_name,
+ )
verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
@@ -307,7 +396,49 @@ def send_message(
)
-async def asend_message_streaming(
+def _build_streaming_logging_obj(
+ request: "SendStreamingMessageRequest",
+ agent_name: str,
+ agent_id: Optional[str],
+ litellm_params: Optional[Dict[str, Any]],
+ metadata: Optional[Dict[str, Any]],
+ proxy_server_request: Optional[Dict[str, Any]],
+) -> Logging:
+ """Build logging object for streaming A2A requests."""
+ start_time = datetime.datetime.now()
+ model = f"a2a_agent/{agent_name}"
+
+ logging_obj = Logging(
+ model=model,
+ messages=[{"role": "user", "content": "streaming-request"}],
+ stream=False,
+ call_type="asend_message_streaming",
+ start_time=start_time,
+ litellm_call_id=str(request.id),
+ function_id=str(request.id),
+ )
+ logging_obj.model = model
+ logging_obj.custom_llm_provider = "a2a_agent"
+ logging_obj.model_call_details["model"] = model
+ logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
+ if agent_id:
+ logging_obj.model_call_details["agent_id"] = agent_id
+
+ _litellm_params = litellm_params.copy() if litellm_params else {}
+ if metadata:
+ _litellm_params["metadata"] = metadata
+ if proxy_server_request:
+ _litellm_params["proxy_server_request"] = proxy_server_request
+
+ logging_obj.litellm_params = _litellm_params
+ logging_obj.optional_params = _litellm_params
+ logging_obj.model_call_details["litellm_params"] = _litellm_params
+ logging_obj.model_call_details["metadata"] = metadata or {}
+
+ return logging_obj
+
+
+async def asend_message_streaming( # noqa: PLR0915
a2a_client: Optional["A2AClientType"] = None,
request: Optional["SendStreamingMessageRequest"] = None,
api_base: Optional[str] = None,
@@ -315,6 +446,7 @@ async def asend_message_streaming(
agent_id: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
proxy_server_request: Optional[Dict[str, Any]] = None,
+ agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Any]:
"""
Async: Send a streaming message to an A2A agent.
@@ -396,62 +528,89 @@ async def asend_message_streaming(
raise ValueError(
"Either a2a_client or api_base is required for standard A2A flow"
)
- a2a_client = await create_a2a_client(base_url=api_base)
+ # Mirror the non-streaming path: always include trace and agent-id headers
+ streaming_extra_headers: Dict[str, str] = {
+ "X-LiteLLM-Trace-Id": str(request.id),
+ }
+ if agent_id:
+ streaming_extra_headers["X-LiteLLM-Agent-Id"] = agent_id
+ if agent_extra_headers:
+ streaming_extra_headers.update(agent_extra_headers)
+ a2a_client = await create_a2a_client(
+ base_url=api_base, extra_headers=streaming_extra_headers
+ )
# Type assertion: a2a_client is guaranteed to be non-None here
assert a2a_client is not None
verbose_logger.info(f"A2A send_message_streaming request_id={request.id}")
- # Track for logging
- start_time = datetime.datetime.now()
- stream = a2a_client.send_message_streaming(request)
-
# Build logging object for streaming completion callbacks
agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(
a2a_client, "agent_card", None
)
+ card_url = getattr(agent_card, "url", None) if agent_card else None
agent_name = getattr(agent_card, "name", "unknown") if agent_card else "unknown"
- model = f"a2a_agent/{agent_name}"
- logging_obj = Logging(
- model=model,
- messages=[{"role": "user", "content": "streaming-request"}],
- stream=False, # complete response logging after stream ends
- call_type="asend_message_streaming",
- start_time=start_time,
- litellm_call_id=str(request.id),
- function_id=str(request.id),
- )
- logging_obj.model = model
- logging_obj.custom_llm_provider = "a2a_agent"
- logging_obj.model_call_details["model"] = model
- logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
- if agent_id:
- logging_obj.model_call_details["agent_id"] = agent_id
-
- # Propagate litellm_params for spend logging (includes cost_per_query, etc.)
- _litellm_params = litellm_params.copy() if litellm_params else {}
- # Merge metadata into litellm_params.metadata (required for proxy cost tracking)
- if metadata:
- _litellm_params["metadata"] = metadata
- if proxy_server_request:
- _litellm_params["proxy_server_request"] = proxy_server_request
-
- logging_obj.litellm_params = _litellm_params
- logging_obj.optional_params = _litellm_params # used by cost calc
- logging_obj.model_call_details["litellm_params"] = _litellm_params
- logging_obj.model_call_details["metadata"] = metadata or {}
-
- iterator = A2AStreamingIterator(
- stream=stream,
+ logging_obj = _build_streaming_logging_obj(
request=request,
- logging_obj=logging_obj,
agent_name=agent_name,
+ agent_id=agent_id,
+ litellm_params=litellm_params,
+ metadata=metadata,
+ proxy_server_request=proxy_server_request,
)
- async for chunk in iterator:
- yield chunk
+ # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
+ # Connection errors in streaming typically occur on first chunk iteration
+ first_chunk = True
+ for attempt in range(2): # max 2 attempts: original + 1 retry
+ stream = a2a_client.send_message_streaming(request)
+ iterator = A2AStreamingIterator(
+ stream=stream,
+ request=request,
+ logging_obj=logging_obj,
+ agent_name=agent_name,
+ )
+
+ try:
+ first_chunk = True
+ async for chunk in iterator:
+ if first_chunk:
+ first_chunk = False # connection succeeded
+ yield chunk
+ return # stream completed successfully
+ except A2ALocalhostURLError as e:
+ # Only retry on first chunk, not mid-stream
+ if first_chunk and attempt == 0:
+ a2a_client = handle_a2a_localhost_retry(
+ error=e,
+ agent_card=agent_card,
+ a2a_client=a2a_client,
+ is_streaming=True,
+ )
+ card_url = agent_card.url if agent_card else None
+ else:
+ raise
+ except Exception as e:
+ # Only map exception on first chunk
+ if first_chunk and attempt == 0:
+ try:
+ map_a2a_exception(e, card_url, api_base, model=agent_name)
+ except A2ALocalhostURLError as localhost_err:
+ # Localhost URL error - fix and retry
+ a2a_client = handle_a2a_localhost_retry(
+ error=localhost_err,
+ agent_card=agent_card,
+ a2a_client=a2a_client,
+ is_streaming=True,
+ )
+ card_url = agent_card.url if agent_card else None
+ continue
+ except Exception:
+ # Re-raise the mapped exception
+ raise
+ raise
async def create_a2a_client(
@@ -493,16 +652,29 @@ async def create_a2a_client(
verbose_logger.info(f"Creating A2A client for {base_url}")
- # Use LiteLLM's cached httpx client
- http_handler = get_async_httpx_client(
- llm_provider=httpxSpecialProvider.A2A,
- params={"timeout": timeout},
+ # Use get_async_httpx_client with per-agent params so that different agents
+ # (with different extra_headers) get separate cached clients. The params
+ # dict is hashed into the cache key, keeping agent auth isolated while
+ # still reusing connections within the same agent.
+ #
+ # Only pass params that AsyncHTTPHandler.__init__ accepts (e.g. timeout).
+ # Use "disable_aiohttp_transport" key for cache-key-only data (it's
+ # filtered out before reaching the constructor).
+ _client_params: dict = {"timeout": timeout}
+ if extra_headers:
+ # Encode headers into a cache-key-only param so each unique header
+ # set produces a distinct cache key.
+ _client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items()))
+ _async_handler = get_async_httpx_client(
+ llm_provider=httpxSpecialProvider.A2AProvider,
+ params=_client_params,
)
- httpx_client = http_handler.client
-
+ httpx_client = _async_handler.client
if extra_headers:
httpx_client.headers.update(extra_headers)
- verbose_proxy_logger.debug(f"A2A client created with extra_headers={extra_headers}")
+ verbose_proxy_logger.debug(
+ f"A2A client created with extra_headers={list(extra_headers.keys())}"
+ )
# Resolve agent card
resolver = A2ACardResolver(
diff --git a/litellm/a2a_protocol/providers/__init__.py b/litellm/a2a_protocol/providers/__init__.py
index 873a5a83749..a21fa5f8f5e 100644
--- a/litellm/a2a_protocol/providers/__init__.py
+++ b/litellm/a2a_protocol/providers/__init__.py
@@ -8,4 +8,3 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
__all__ = ["BaseA2AProviderConfig", "A2AProviderConfigManager"]
-
diff --git a/litellm/a2a_protocol/providers/base.py b/litellm/a2a_protocol/providers/base.py
index 9931076a948..a2354b3495e 100644
--- a/litellm/a2a_protocol/providers/base.py
+++ b/litellm/a2a_protocol/providers/base.py
@@ -9,7 +9,7 @@ from typing import Any, AsyncIterator, Dict
class BaseA2AProviderConfig(ABC):
"""
Base configuration class for A2A protocol providers.
-
+
Each provider should implement this interface to define how to handle
A2A requests for their specific agent type.
"""
@@ -60,4 +60,3 @@ class BaseA2AProviderConfig(ABC):
# The yield is here to make this a generator function
if False: # pragma: no cover
yield {}
-
diff --git a/litellm/a2a_protocol/providers/config_manager.py b/litellm/a2a_protocol/providers/config_manager.py
index e0703ec466b..a8b9566c171 100644
--- a/litellm/a2a_protocol/providers/config_manager.py
+++ b/litellm/a2a_protocol/providers/config_manager.py
@@ -12,7 +12,7 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
class A2AProviderConfigManager:
"""
Manager for A2A provider configurations.
-
+
Similar to ProviderConfigManager in litellm.utils but specifically for A2A providers.
"""
@@ -31,7 +31,7 @@ class A2AProviderConfigManager:
"""
if custom_llm_provider is None:
return None
-
+
if custom_llm_provider == "pydantic_ai_agents":
from litellm.a2a_protocol.providers.pydantic_ai_agents.config import (
PydanticAIProviderConfig,
@@ -45,4 +45,3 @@ class A2AProviderConfigManager:
# return AnotherProviderConfig()
return None
-
diff --git a/litellm/a2a_protocol/providers/litellm_completion/__init__.py b/litellm/a2a_protocol/providers/litellm_completion/__init__.py
index 3f2b88bfaa3..fc2fc17f54f 100644
--- a/litellm/a2a_protocol/providers/litellm_completion/__init__.py
+++ b/litellm/a2a_protocol/providers/litellm_completion/__init__.py
@@ -3,4 +3,3 @@ LiteLLM Completion bridge provider for A2A protocol.
Routes A2A requests through litellm.acompletion based on custom_llm_provider.
"""
-
diff --git a/litellm/a2a_protocol/providers/litellm_completion/handler.py b/litellm/a2a_protocol/providers/litellm_completion/handler.py
index 57388a5d0ed..730f8f6b36f 100644
--- a/litellm/a2a_protocol/providers/litellm_completion/handler.py
+++ b/litellm/a2a_protocol/providers/litellm_completion/handler.py
@@ -52,26 +52,26 @@ class A2ACompletionBridgeHandler:
if custom_llm_provider == "pydantic_ai_agents":
if api_base is None:
raise ValueError("api_base is required for Pydantic AI agents")
-
+
verbose_logger.info(
f"Pydantic AI: Routing to Pydantic AI agent at {api_base}"
)
-
+
# Send request directly to Pydantic AI agent
response_data = await PydanticAITransformation.send_non_streaming_request(
api_base=api_base,
request_id=request_id,
params=params,
)
-
+
return response_data
-
+
# Extract message from params
message = params.get("message", {})
# Transform A2A message to OpenAI format
- openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(
- message
+ openai_messages = (
+ A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
@@ -98,7 +98,8 @@ class A2ACompletionBridgeHandler:
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
- k: v for k, v in litellm_params.items()
+ k: v
+ for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
@@ -107,9 +108,11 @@ class A2ACompletionBridgeHandler:
response = await litellm.acompletion(**completion_params)
# Transform response to A2A format
- a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
- response=response,
- request_id=request_id,
+ a2a_response = (
+ A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
+ response=response,
+ request_id=request_id,
+ )
)
verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}")
@@ -146,27 +149,27 @@ class A2ACompletionBridgeHandler:
if custom_llm_provider == "pydantic_ai_agents":
if api_base is None:
raise ValueError("api_base is required for Pydantic AI agents")
-
+
verbose_logger.info(
f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}"
)
-
+
# Get non-streaming response first
response_data = await PydanticAITransformation.send_non_streaming_request(
api_base=api_base,
request_id=request_id,
params=params,
)
-
+
# Convert to fake streaming
async for chunk in PydanticAITransformation.fake_streaming_from_response(
response_data=response_data,
request_id=request_id,
):
yield chunk
-
+
return
-
+
# Extract message from params
message = params.get("message", {})
@@ -177,8 +180,8 @@ class A2ACompletionBridgeHandler:
)
# Transform A2A message to OpenAI format
- openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(
- message
+ openai_messages = (
+ A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
@@ -205,7 +208,8 @@ class A2ACompletionBridgeHandler:
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
- k: v for k, v in litellm_params.items()
+ k: v
+ for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
@@ -244,9 +248,11 @@ class A2ACompletionBridgeHandler:
# Emit artifact update with accumulated content
if accumulated_text:
- artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event(
- ctx=ctx,
- text=accumulated_text,
+ artifact_event = (
+ A2ACompletionBridgeTransformation.create_artifact_update_event(
+ ctx=ctx,
+ text=accumulated_text,
+ )
)
yield artifact_event
diff --git a/litellm/a2a_protocol/providers/litellm_completion/transformation.py b/litellm/a2a_protocol/providers/litellm_completion/transformation.py
index bbe7daa9fc4..8a03569f689 100644
--- a/litellm/a2a_protocol/providers/litellm_completion/transformation.py
+++ b/litellm/a2a_protocol/providers/litellm_completion/transformation.py
@@ -124,9 +124,7 @@ class A2ACompletionBridgeTransformation:
},
}
- verbose_logger.debug(
- f"OpenAI -> A2A transform: content_length={len(content)}"
- )
+ verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}")
return a2a_response
diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py
index 2187400b2d1..8e9cd6fc87e 100644
--- a/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py
+++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py
@@ -14,4 +14,3 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
)
__all__ = ["PydanticAIHandler", "PydanticAITransformation", "PydanticAIProviderConfig"]
-
diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py
index acf09554e5e..d4c5f6a2985 100644
--- a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py
+++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py
@@ -11,7 +11,7 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAI
class PydanticAIProviderConfig(BaseA2AProviderConfig):
"""
Provider configuration for Pydantic AI agents.
-
+
Pydantic AI agents follow A2A protocol but don't support streaming natively.
This config provides fake streaming by converting non-streaming responses into streaming chunks.
"""
@@ -48,4 +48,3 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
delay_ms=kwargs.get("delay_ms", 10),
):
yield chunk
-
diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py
index 6680a9fe487..7d4167752f8 100644
--- a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py
+++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py
@@ -16,7 +16,7 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
class PydanticAIHandler:
"""
Handler for Pydantic AI agent requests.
-
+
Provides:
- Direct non-streaming requests to Pydantic AI agents
- Fake streaming by converting non-streaming responses into streaming chunks
@@ -41,9 +41,7 @@ class PydanticAIHandler:
Returns:
A2A SendMessageResponse dict
"""
- verbose_logger.info(
- f"Pydantic AI: Routing to Pydantic AI agent at {api_base}"
- )
+ verbose_logger.info(f"Pydantic AI: Routing to Pydantic AI agent at {api_base}")
# Send request directly to Pydantic AI agent
response_data = await PydanticAITransformation.send_non_streaming_request(
@@ -102,5 +100,3 @@ class PydanticAIHandler:
delay_ms=delay_ms,
):
yield chunk
-
-
diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
index 9352eab6c8e..e73b17ac3c0 100644
--- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
+++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
@@ -10,13 +10,16 @@ from typing import Any, AsyncIterator, Dict, cast
from uuid import uuid4
from litellm._logging import verbose_logger
-from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
+from litellm.llms.custom_httpx.http_handler import (
+ AsyncHTTPHandler,
+ get_async_httpx_client,
+)
class PydanticAITransformation:
"""
Transformation layer for Pydantic AI agents.
-
+
Handles:
- Direct A2A requests to Pydantic AI endpoints
- Polling for task completion (since Pydantic AI doesn't support streaming)
@@ -27,13 +30,13 @@ class PydanticAITransformation:
def _remove_none_values(obj: Any) -> Any:
"""
Recursively remove None values from a dict/list structure.
-
+
FastA2A/Pydantic AI servers don't accept None values for optional fields -
they expect those fields to be omitted entirely.
-
+
Args:
obj: Dict, list, or other value to clean
-
+
Returns:
Cleaned object with None values removed
"""
@@ -56,10 +59,10 @@ class PydanticAITransformation:
def _params_to_dict(params: Any) -> Dict[str, Any]:
"""
Convert params to a dict, handling Pydantic models.
-
+
Args:
params: Dict or Pydantic model
-
+
Returns:
Dict representation of params
"""
@@ -86,7 +89,7 @@ class PydanticAITransformation:
) -> Dict[str, Any]:
"""
Poll for task completion using tasks/get method.
-
+
Args:
client: HTTPX async client
endpoint: API endpoint URL
@@ -94,7 +97,7 @@ class PydanticAITransformation:
request_id: JSON-RPC request ID
max_attempts: Maximum polling attempts
poll_interval: Seconds between poll attempts
-
+
Returns:
Completed task response
"""
@@ -105,7 +108,7 @@ class PydanticAITransformation:
"method": "tasks/get",
"params": {"id": task_id},
}
-
+
response = await client.post(
endpoint,
json=poll_request,
@@ -113,23 +116,25 @@ class PydanticAITransformation:
)
response.raise_for_status()
poll_data = response.json()
-
+
result = poll_data.get("result", {})
status = result.get("status", {})
state = status.get("state", "")
-
+
verbose_logger.debug(
f"Pydantic AI: Poll attempt {attempt + 1}/{max_attempts}, state={state}"
)
-
+
if state == "completed":
return poll_data
elif state in ("failed", "canceled"):
raise Exception(f"Task {task_id} ended with state: {state}")
-
+
await asyncio.sleep(poll_interval)
-
- raise TimeoutError(f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds")
+
+ raise TimeoutError(
+ f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds"
+ )
@staticmethod
async def _send_and_poll_raw(
@@ -140,7 +145,7 @@ class PydanticAITransformation:
) -> Dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
-
+
This is an internal method used by both non-streaming and streaming handlers.
Returns the raw Pydantic AI task format with history/artifacts.
@@ -155,10 +160,10 @@ class PydanticAITransformation:
"""
# Convert params to dict if it's a Pydantic model
params_dict = PydanticAITransformation._params_to_dict(params)
-
+
# Remove None values - FastA2A doesn't accept null for optional fields
params_dict = PydanticAITransformation._remove_none_values(params_dict)
-
+
# Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI
if "message" in params_dict:
params_dict["message"]["kind"] = "message"
@@ -174,9 +179,7 @@ class PydanticAITransformation:
# FastA2A uses root endpoint (/) not /messages
endpoint = api_base.rstrip("/")
- verbose_logger.info(
- f"Pydantic AI: Sending non-streaming request to {endpoint}"
- )
+ verbose_logger.info(f"Pydantic AI: Sending non-streaming request to {endpoint}")
# Send request to Pydantic AI agent using shared async HTTP client
client = get_async_httpx_client(
@@ -190,12 +193,12 @@ class PydanticAITransformation:
)
response.raise_for_status()
response_data = response.json()
-
+
# Check if task is already completed
result = response_data.get("result", {})
status = result.get("status", {})
state = status.get("state", "")
-
+
if state != "completed":
# Need to poll for completion
task_id = result.get("id")
@@ -210,7 +213,9 @@ class PydanticAITransformation:
request_id=request_id,
)
- verbose_logger.info(f"Pydantic AI: Received completed response for request_id={request_id}")
+ verbose_logger.info(
+ f"Pydantic AI: Received completed response for request_id={request_id}"
+ )
return response_data
@@ -256,7 +261,7 @@ class PydanticAITransformation:
) -> Dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
-
+
Used by streaming handler to get raw response for fake streaming.
Args:
@@ -282,7 +287,7 @@ class PydanticAITransformation:
) -> Dict[str, Any]:
"""
Transform Pydantic AI task response to standard A2A non-streaming format.
-
+
Pydantic AI returns a task with history/artifacts, but the standard A2A
non-streaming format expects:
{
@@ -296,11 +301,11 @@ class PydanticAITransformation:
}
}
}
-
+
Args:
response_data: Pydantic AI task response
request_id: Original request ID
-
+
Returns:
Standard A2A non-streaming response format
"""
@@ -308,14 +313,14 @@ class PydanticAITransformation:
full_text, message_id, parts = PydanticAITransformation._extract_response_text(
response_data
)
-
+
# Build standard A2A message
a2a_message = {
"role": "agent",
"parts": parts if parts else [{"kind": "text", "text": full_text}],
"messageId": message_id,
}
-
+
# Return standard A2A non-streaming format
return {
"jsonrpc": "2.0",
@@ -329,19 +334,19 @@ class PydanticAITransformation:
def _extract_response_text(response_data: Dict[str, Any]) -> tuple[str, str, list]:
"""
Extract response text from completed task response.
-
+
Pydantic AI returns completed tasks with:
- history: list of messages (user and agent)
- artifacts: list of result artifacts
-
+
Args:
response_data: Completed task response
-
+
Returns:
Tuple of (full_text, message_id, parts)
"""
result = response_data.get("result", {})
-
+
# Try to extract from artifacts first (preferred for results)
artifacts = result.get("artifacts", [])
if artifacts:
@@ -352,7 +357,7 @@ class PydanticAITransformation:
text = part.get("text", "")
if text:
return text, str(uuid4()), parts
-
+
# Fall back to history - get the last agent message
history = result.get("history", [])
for msg in reversed(history):
@@ -365,7 +370,7 @@ class PydanticAITransformation:
full_text += part.get("text", "")
if full_text:
return full_text, message_id, parts
-
+
# Fall back to message field (original format)
message = result.get("message", {})
if message:
@@ -376,7 +381,7 @@ class PydanticAITransformation:
if part.get("kind") == "text":
full_text += part.get("text", "")
return full_text, message_id, parts
-
+
return "", str(uuid4()), []
@staticmethod
@@ -408,7 +413,7 @@ class PydanticAITransformation:
full_text, message_id, parts = PydanticAITransformation._extract_response_text(
response_data
)
-
+
# Extract input message from raw response for history
result = response_data.get("result", {})
history = result.get("history", [])
@@ -436,7 +441,9 @@ class PydanticAITransformation:
"contextId": context_id,
"kind": "message",
"messageId": input_message_id,
- "parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
+ "parts": input_message.get(
+ "parts", [{"kind": "text", "text": ""}]
+ ),
"role": "user",
"taskId": task_id,
}
@@ -475,7 +482,7 @@ class PydanticAITransformation:
if full_text:
# Split text into chunks
for i in range(0, len(full_text), chunk_size):
- chunk_text = full_text[i:i + chunk_size]
+ chunk_text = full_text[i : i + chunk_size]
is_last_chunk = (i + chunk_size) >= len(full_text)
artifact_event = {
@@ -521,5 +528,3 @@ class PydanticAITransformation:
verbose_logger.info(
f"Pydantic AI: Fake streaming completed for request_id={request_id}"
)
-
-
diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py
index 921dc0e52e0..98d45cf2ac1 100644
--- a/litellm/a2a_protocol/streaming_iterator.py
+++ b/litellm/a2a_protocol/streaming_iterator.py
@@ -71,7 +71,11 @@ class A2AStreamingIterator:
def _collect_text_from_chunk(self, chunk: Any) -> None:
"""Extract text from a streaming chunk and add to collected parts."""
try:
- chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
+ chunk_dict = (
+ chunk.model_dump(mode="json", exclude_none=True)
+ if hasattr(chunk, "model_dump")
+ else {}
+ )
text = A2ARequestUtils.extract_text_from_response(chunk_dict)
if text:
self.collected_text_parts.append(text)
@@ -81,7 +85,11 @@ class A2AStreamingIterator:
def _is_completed_chunk(self, chunk: Any) -> bool:
"""Check if chunk indicates stream completion."""
try:
- chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
+ chunk_dict = (
+ chunk.model_dump(mode="json", exclude_none=True)
+ if hasattr(chunk, "model_dump")
+ else {}
+ )
result = chunk_dict.get("result", {})
if isinstance(result, dict):
status = result.get("status", {})
@@ -102,7 +110,9 @@ class A2AStreamingIterator:
prompt_tokens = A2ARequestUtils.count_tokens(input_text)
# Use the last (most complete) text from chunks
- output_text = self.collected_text_parts[-1] if self.collected_text_parts else ""
+ output_text = (
+ self.collected_text_parts[-1] if self.collected_text_parts else ""
+ )
completion_tokens = A2ARequestUtils.count_tokens(output_text)
total_tokens = prompt_tokens + completion_tokens
@@ -158,7 +168,9 @@ class A2AStreamingIterator:
result: Dict[str, Any] = {
"id": getattr(self.request, "id", "unknown"),
"jsonrpc": "2.0",
- "usage": usage.model_dump() if hasattr(usage, "model_dump") else dict(usage),
+ "usage": usage.model_dump()
+ if hasattr(usage, "model_dump")
+ else dict(usage),
}
# Add final chunk result if available
@@ -170,4 +182,3 @@ class A2AStreamingIterator:
pass
return result
-
diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json
new file mode 100644
index 00000000000..df8d49ac8f2
--- /dev/null
+++ b/litellm/anthropic_beta_headers_config.json
@@ -0,0 +1,182 @@
+{
+ "description": "Mapping of Anthropic beta headers for each provider. Keys are input header names, values are provider-specific header names (or null if unsupported). Only headers present in mapping keys with non-null values can be forwarded.",
+ "anthropic": {
+ "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20",
+ "bash_20241022": null,
+ "bash_20250124": null,
+ "code-execution-2025-08-25": "code-execution-2025-08-25",
+ "compact-2026-01-12": "compact-2026-01-12",
+ "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": "context-1m-2025-08-07",
+ "context-management-2025-06-27": "context-management-2025-06-27",
+ "effort-2025-11-24": "effort-2025-11-24",
+ "fast-mode-2026-02-01": "fast-mode-2026-02-01",
+ "files-api-2025-04-14": "files-api-2025-04-14",
+ "structured-output-2024-03-01": null,
+ "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
+ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
+ "mcp-client-2025-11-20": "mcp-client-2025-11-20",
+ "mcp-client-2025-04-04": "mcp-client-2025-04-04",
+ "mcp-servers-2025-12-04": null,
+ "oauth-2025-04-20": "oauth-2025-04-20",
+ "output-128k-2025-02-19": "output-128k-2025-02-19",
+ "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
+ "skills-2025-10-02": "skills-2025-10-02",
+ "structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
+ "text_editor_20241022": null,
+ "text_editor_20250124": null,
+ "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
+ "web-fetch-2025-09-10": "web-fetch-2025-09-10",
+ "web-search-2025-03-05": "web-search-2025-03-05"
+ },
+ "azure_ai": {
+ "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20",
+ "bash_20241022": null,
+ "bash_20250124": null,
+ "code-execution-2025-08-25": "code-execution-2025-08-25",
+ "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": "context-1m-2025-08-07",
+ "context-management-2025-06-27": "context-management-2025-06-27",
+ "effort-2025-11-24": "effort-2025-11-24",
+ "fast-mode-2026-02-01": null,
+ "files-api-2025-04-14": "files-api-2025-04-14",
+ "fine-grained-tool-streaming-2025-05-14": null,
+ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
+ "mcp-client-2025-11-20": "mcp-client-2025-11-20",
+ "mcp-client-2025-04-04": "mcp-client-2025-04-04",
+ "mcp-servers-2025-12-04": null,
+ "output-128k-2025-02-19": null,
+ "structured-output-2024-03-01": null,
+ "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
+ "skills-2025-10-02": "skills-2025-10-02",
+ "structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
+ "text_editor_20241022": null,
+ "text_editor_20250124": null,
+ "token-efficient-tools-2025-02-19": null,
+ "web-fetch-2025-09-10": "web-fetch-2025-09-10",
+ "web-search-2025-03-05": "web-search-2025-03-05"
+ },
+ "bedrock_converse": {
+ "advanced-tool-use-2025-11-20": null,
+ "bash_20241022": null,
+ "bash_20250124": null,
+ "code-execution-2025-08-25": null,
+ "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": "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,
+ "files-api-2025-04-14": null,
+ "fine-grained-tool-streaming-2025-05-14": null,
+ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
+ "mcp-client-2025-11-20": null,
+ "mcp-client-2025-04-04": null,
+ "mcp-servers-2025-12-04": null,
+ "output-128k-2025-02-19": null,
+ "structured-output-2024-03-01": null,
+ "prompt-caching-scope-2026-01-05": null,
+ "skills-2025-10-02": null,
+ "structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
+ "text_editor_20241022": null,
+ "text_editor_20250124": null,
+ "token-efficient-tools-2025-02-19": null,
+ "tool-search-tool-2025-10-19": null,
+ "web-fetch-2025-09-10": null,
+ "web-search-2025-03-05": null
+ },
+ "bedrock": {
+ "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19",
+ "bash_20241022": null,
+ "bash_20250124": null,
+ "code-execution-2025-08-25": null,
+ "compact-2026-01-12": "compact-2026-01-12",
+ "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": "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,
+ "files-api-2025-04-14": null,
+ "fine-grained-tool-streaming-2025-05-14": null,
+ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
+ "mcp-client-2025-11-20": null,
+ "mcp-client-2025-04-04": null,
+ "mcp-servers-2025-12-04": null,
+ "output-128k-2025-02-19": null,
+ "structured-output-2024-03-01": null,
+ "prompt-caching-scope-2026-01-05": null,
+ "skills-2025-10-02": null,
+ "structured-outputs-2025-11-13": null,
+ "text_editor_20241022": null,
+ "text_editor_20250124": null,
+ "token-efficient-tools-2025-02-19": null,
+ "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
+ "web-fetch-2025-09-10": null,
+ "web-search-2025-03-05": null
+ },
+ "vertex_ai": {
+ "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19",
+ "bash_20241022": null,
+ "bash_20250124": null,
+ "code-execution-2025-08-25": null,
+ "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": "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,
+ "files-api-2025-04-14": null,
+ "fine-grained-tool-streaming-2025-05-14": null,
+ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
+ "mcp-client-2025-11-20": null,
+ "mcp-client-2025-04-04": null,
+ "mcp-servers-2025-12-04": null,
+ "output-128k-2025-02-19": null,
+ "structured-output-2024-03-01": null,
+ "prompt-caching-scope-2026-01-05": null,
+ "skills-2025-10-02": null,
+ "structured-outputs-2025-11-13": null,
+ "text_editor_20241022": null,
+ "text_editor_20250124": null,
+ "token-efficient-tools-2025-02-19": null,
+ "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
+ "web-fetch-2025-09-10": null,
+ "web-search-2025-03-05": "web-search-2025-03-05"
+ },
+ "databricks": {
+ "advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20",
+ "bash_20241022": null,
+ "bash_20250124": null,
+ "code-execution-2025-08-25": "code-execution-2025-08-25",
+ "compact-2026-01-12": "compact-2026-01-12",
+ "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": "context-1m-2025-08-07",
+ "context-management-2025-06-27": "context-management-2025-06-27",
+ "effort-2025-11-24": "effort-2025-11-24",
+ "fast-mode-2026-02-01": "fast-mode-2026-02-01",
+ "files-api-2025-04-14": "files-api-2025-04-14",
+ "structured-output-2024-03-01": null,
+ "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
+ "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
+ "mcp-client-2025-11-20": "mcp-client-2025-11-20",
+ "mcp-client-2025-04-04": "mcp-client-2025-04-04",
+ "mcp-servers-2025-12-04": null,
+ "oauth-2025-04-20": "oauth-2025-04-20",
+ "output-128k-2025-02-19": "output-128k-2025-02-19",
+ "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
+ "skills-2025-10-02": "skills-2025-10-02",
+ "structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
+ "text_editor_20241022": null,
+ "text_editor_20250124": null,
+ "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
+ "web-fetch-2025-09-10": "web-fetch-2025-09-10",
+ "web-search-2025-03-05": "web-search-2025-03-05"
+ }
+}
diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py
new file mode 100644
index 00000000000..97d223088fa
--- /dev/null
+++ b/litellm/anthropic_beta_headers_manager.py
@@ -0,0 +1,421 @@
+"""
+Centralized manager for Anthropic beta headers across different providers.
+
+This module provides utilities to:
+1. Load beta header configuration from JSON (mapping of supported headers per provider)
+2. Filter and map beta headers based on provider support
+3. Handle provider-specific header name mappings (e.g., advanced-tool-use -> tool-search-tool)
+4. Support remote fetching and caching similar to model cost map
+
+Design:
+- JSON config contains mapping of beta headers for each provider
+- Keys are input header names, values are provider-specific header names (or null if unsupported)
+- Only headers present in mapping keys with non-null values can be forwarded
+- This enforces stricter validation than the previous unsupported list approach
+
+Configuration can be loaded from:
+- Remote URL (default): Fetches from GitHub repository
+- Local file: Set LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS=True to use bundled config only
+
+Environment Variables:
+- LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS: Set to "True" to disable remote fetching
+- LITELLM_ANTHROPIC_BETA_HEADERS_URL: Custom URL for remote config (optional)
+"""
+
+import json
+import os
+from importlib.resources import files
+from typing import Dict, List, Optional, Set
+
+import httpx
+
+from litellm.litellm_core_utils.litellm_logging import verbose_logger
+
+# Cache for the loaded configuration
+_BETA_HEADERS_CONFIG: Optional[Dict] = None
+
+
+class GetAnthropicBetaHeadersConfig:
+ """
+ Handles fetching, validating, and loading the Anthropic beta headers configuration.
+
+ Similar to GetModelCostMap, this class manages the lifecycle of the beta headers
+ configuration with support for remote fetching and local fallback.
+ """
+
+ @staticmethod
+ def load_local_beta_headers_config() -> Dict:
+ """Load the local backup beta headers config bundled with the package."""
+ try:
+ content = json.loads(
+ files("litellm")
+ .joinpath("anthropic_beta_headers_config.json")
+ .read_text(encoding="utf-8")
+ )
+ return content
+ except Exception as e:
+ verbose_logger.error(f"Failed to load local beta headers config: {e}")
+ # Return empty config as fallback
+ return {
+ "anthropic": {},
+ "azure_ai": {},
+ "bedrock": {},
+ "bedrock_converse": {},
+ "vertex_ai": {},
+ "provider_aliases": {},
+ }
+
+ @staticmethod
+ def _check_is_valid_dict(fetched_config: dict) -> bool:
+ """Check if fetched config is a non-empty dict with expected structure."""
+ if not isinstance(fetched_config, dict):
+ verbose_logger.warning(
+ "LiteLLM: Fetched beta headers config is not a dict (type=%s). "
+ "Falling back to local backup.",
+ type(fetched_config).__name__,
+ )
+ return False
+
+ if len(fetched_config) == 0:
+ verbose_logger.warning(
+ "LiteLLM: Fetched beta headers config is empty. "
+ "Falling back to local backup.",
+ )
+ return False
+
+ # Check for at least one provider key
+ provider_keys = [
+ "anthropic",
+ "azure_ai",
+ "bedrock",
+ "bedrock_converse",
+ "vertex_ai",
+ ]
+ has_provider = any(key in fetched_config for key in provider_keys)
+
+ if not has_provider:
+ verbose_logger.warning(
+ "LiteLLM: Fetched beta headers config missing provider keys. "
+ "Falling back to local backup.",
+ )
+ return False
+
+ return True
+
+ @classmethod
+ def validate_beta_headers_config(cls, fetched_config: dict) -> bool:
+ """
+ Validate the integrity of a fetched beta headers config.
+
+ Returns True if all checks pass, False otherwise.
+ """
+ return cls._check_is_valid_dict(fetched_config)
+
+ @staticmethod
+ def fetch_remote_beta_headers_config(url: str, timeout: int = 5) -> dict:
+ """
+ Fetch the beta headers config from a remote URL.
+
+ Returns the parsed JSON dict. Raises on network/parse errors
+ (caller is expected to handle).
+ """
+ response = httpx.get(url, timeout=timeout)
+ response.raise_for_status()
+ return response.json()
+
+
+def get_beta_headers_config(url: str) -> dict:
+ """
+ Public entry point — returns the beta headers config dict.
+
+ 1. If ``LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS`` is set, uses the local backup only.
+ 2. Otherwise fetches from ``url``, validates integrity, and falls back
+ to the local backup on any failure.
+
+ Args:
+ url: URL to fetch the remote beta headers configuration from
+
+ Returns:
+ Dict containing the beta headers configuration
+ """
+ # Check if local-only mode is enabled
+ if os.getenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "").lower() == "true":
+ # verbose_logger.debug("Using local Anthropic beta headers config (LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS=True)")
+ return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()
+
+ try:
+ content = GetAnthropicBetaHeadersConfig.fetch_remote_beta_headers_config(url)
+ except Exception as e:
+ verbose_logger.warning(
+ "LiteLLM: Failed to fetch remote beta headers config from %s: %s. "
+ "Falling back to local backup.",
+ url,
+ str(e),
+ )
+ return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()
+
+ # Validate the fetched config
+ if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(
+ fetched_config=content
+ ):
+ verbose_logger.warning(
+ "LiteLLM: Fetched beta headers config failed integrity check. "
+ "Using local backup instead. url=%s",
+ url,
+ )
+ return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()
+
+ return content
+
+
+def _load_beta_headers_config() -> Dict:
+ """
+ Load the beta headers configuration.
+ Uses caching to avoid repeated fetches/file reads.
+
+ This function is called by all public API functions and manages the global cache.
+
+ Returns:
+ Dict containing the beta headers configuration
+ """
+ global _BETA_HEADERS_CONFIG
+
+ if _BETA_HEADERS_CONFIG is not None:
+ return _BETA_HEADERS_CONFIG
+
+ # Get the URL from environment or use default
+ from litellm import anthropic_beta_headers_url
+
+ _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
+ verbose_logger.debug("Loaded and cached beta headers config")
+
+ return _BETA_HEADERS_CONFIG
+
+
+def reload_beta_headers_config() -> Dict:
+ """
+ Force reload the beta headers configuration from source (remote or local).
+ Clears the cache and fetches fresh configuration.
+
+ Returns:
+ Dict containing the newly loaded beta headers configuration
+ """
+ global _BETA_HEADERS_CONFIG
+ _BETA_HEADERS_CONFIG = None
+ verbose_logger.info("Reloading beta headers config (cache cleared)")
+ return _load_beta_headers_config()
+
+
+def get_provider_name(provider: str) -> str:
+ """
+ Resolve provider aliases to canonical provider names.
+
+ Args:
+ provider: Provider name (may be an alias)
+
+ Returns:
+ Canonical provider name
+ """
+ config = _load_beta_headers_config()
+ aliases = config.get("provider_aliases", {})
+ return aliases.get(provider, provider)
+
+
+def filter_and_transform_beta_headers(
+ beta_headers: List[str],
+ provider: str,
+) -> List[str]:
+ """
+ Filter and transform beta headers based on provider's mapping configuration.
+
+ This function:
+ 1. Only allows headers that are present in the provider's mapping keys
+ 2. Filters out headers with null values (unsupported)
+ 3. Maps headers to provider-specific names (e.g., advanced-tool-use -> tool-search-tool)
+
+ Args:
+ beta_headers: List of Anthropic beta header values
+ provider: Provider name (e.g., "anthropic", "bedrock", "vertex_ai")
+
+ Returns:
+ List of filtered and transformed beta headers for the provider
+ """
+ if not beta_headers:
+ return []
+
+ config = _load_beta_headers_config()
+ provider = get_provider_name(provider)
+
+ # Get the header mapping for this provider
+ provider_mapping = config.get(provider, {})
+
+ filtered_headers: Set[str] = set()
+
+ for header in beta_headers:
+ header = header.strip()
+
+ # Check if header is in the mapping
+ if header not in provider_mapping:
+ verbose_logger.debug(
+ f"Dropping unknown beta header '{header}' for provider '{provider}' (not in mapping)"
+ )
+ continue
+
+ # Get the mapped header value
+ mapped_header = provider_mapping[header]
+
+ # Skip if header is unsupported (null value)
+ if mapped_header is None:
+ verbose_logger.debug(
+ f"Dropping unsupported beta header '{header}' for provider '{provider}'"
+ )
+ continue
+
+ # Add the mapped header
+ filtered_headers.add(mapped_header)
+
+ return sorted(list(filtered_headers))
+
+
+def is_beta_header_supported(
+ beta_header: str,
+ provider: str,
+) -> bool:
+ """
+ Check if a specific beta header is supported by a provider.
+
+ Args:
+ beta_header: The Anthropic beta header value
+ provider: Provider name
+
+ Returns:
+ True if the header is in the mapping with a non-null value, False otherwise
+ """
+ config = _load_beta_headers_config()
+ provider = get_provider_name(provider)
+ provider_mapping = config.get(provider, {})
+
+ # Header is supported if it's in the mapping and has a non-null value
+ return beta_header in provider_mapping and provider_mapping[beta_header] is not None
+
+
+def get_provider_beta_header(
+ anthropic_beta_header: str,
+ provider: str,
+) -> Optional[str]:
+ """
+ Get the provider-specific beta header name for a given Anthropic beta header.
+
+ This function handles header transformations/mappings (e.g., advanced-tool-use -> tool-search-tool).
+
+ Args:
+ anthropic_beta_header: The Anthropic beta header value
+ provider: Provider name
+
+ Returns:
+ The provider-specific header name if supported, or None if unsupported/unknown
+ """
+ config = _load_beta_headers_config()
+ provider = get_provider_name(provider)
+
+ # Get the header mapping for this provider
+ provider_mapping = config.get(provider, {})
+
+ # Check if header is in the mapping
+ if anthropic_beta_header not in provider_mapping:
+ return None
+
+ # Return the mapped value (could be None if unsupported)
+ return provider_mapping[anthropic_beta_header]
+
+
+def update_headers_with_filtered_beta(
+ headers: dict,
+ provider: str,
+) -> dict:
+ """
+ Update headers dict by filtering and transforming anthropic-beta header values.
+ Modifies the headers dict in place and returns it.
+
+ Args:
+ headers: Request headers dict (will be modified in place)
+ provider: Provider name
+
+ Returns:
+ Updated headers dict
+ """
+ existing_beta = headers.get("anthropic-beta")
+ if not existing_beta:
+ return headers
+
+ # Parse existing beta headers
+ beta_values = [b.strip() for b in existing_beta.split(",") if b.strip()]
+
+ # Filter and transform based on provider
+ filtered_beta_values = filter_and_transform_beta_headers(
+ beta_headers=beta_values,
+ provider=provider,
+ )
+
+ # Update or remove the header
+ if filtered_beta_values:
+ headers["anthropic-beta"] = ",".join(filtered_beta_values)
+ else:
+ # Remove the header if no values remain
+ headers.pop("anthropic-beta", None)
+
+ return headers
+
+
+def update_request_with_filtered_beta(
+ headers: dict,
+ request_data: dict,
+ provider: str,
+) -> tuple[dict, dict]:
+ """
+ Update both headers and request body beta fields based on provider support.
+ Modifies both dicts in place and returns them.
+
+ Args:
+ headers: Request headers dict (will be modified in place)
+ request_data: Request body dict (will be modified in place)
+ provider: Provider name
+
+ Returns:
+ Tuple of (updated headers, updated request_data)
+ """
+ headers = update_headers_with_filtered_beta(headers=headers, provider=provider)
+
+ existing_body_betas = request_data.get("anthropic_beta")
+ if not existing_body_betas:
+ return headers, request_data
+
+ filtered_body_betas = filter_and_transform_beta_headers(
+ beta_headers=existing_body_betas,
+ provider=provider,
+ )
+
+ if filtered_body_betas:
+ request_data["anthropic_beta"] = filtered_body_betas
+ else:
+ request_data.pop("anthropic_beta", None)
+
+ return headers, request_data
+
+
+def get_unsupported_headers(provider: str) -> List[str]:
+ """
+ Get all beta headers that are unsupported by a provider (have null values in mapping).
+
+ Args:
+ provider: Provider name
+
+ Returns:
+ List of unsupported Anthropic beta header names
+ """
+ config = _load_beta_headers_config()
+ provider = get_provider_name(provider)
+ provider_mapping = config.get(provider, {})
+
+ # Return headers with null values
+ return [header for header, value in provider_mapping.items() if value is None]
diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py
index b8a5079a4eb..28020e763f4 100644
--- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py
+++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py
@@ -149,7 +149,9 @@ class AnthropicExceptionMapping:
parsed = None
# If parsed and already in Anthropic format - passthrough
- if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict(parsed):
+ if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict(
+ parsed
+ ):
# Optionally add request_id if provided and not present
if request_id and "request_id" not in parsed:
parsed["request_id"] = request_id
@@ -157,7 +159,9 @@ class AnthropicExceptionMapping:
# Extract message - use parsed dict if available, otherwise raw string
if parsed is not None:
- message = AnthropicExceptionMapping._extract_message_from_dict(parsed, raw_message)
+ message = AnthropicExceptionMapping._extract_message_from_dict(
+ parsed, raw_message
+ )
else:
message = raw_message
diff --git a/litellm/batch_completion/main.py b/litellm/batch_completion/main.py
index 7100fb004f8..446e3f2f990 100644
--- a/litellm/batch_completion/main.py
+++ b/litellm/batch_completion/main.py
@@ -237,17 +237,37 @@ def batch_completion_models_all_responses(*args, **kwargs):
if "model" in kwargs:
kwargs.pop("model")
if "models" in kwargs:
- models = kwargs["models"]
- kwargs.pop("models")
+ models = kwargs.pop("models")
else:
raise Exception("'models' param not in kwargs")
+ if isinstance(models, str):
+ models = [models]
+ elif isinstance(models, (list, tuple)):
+ models = list(models)
+ else:
+ raise TypeError("'models' must be a string or list of strings")
+
+ if len(models) == 0:
+ return []
+
responses = []
with concurrent.futures.ThreadPoolExecutor(max_workers=len(models)) as executor:
- for idx, model in enumerate(models):
- future = executor.submit(litellm.completion, *args, model=model, **kwargs)
- if future.result() is not None:
- responses.append(future.result())
+ futures = [
+ executor.submit(litellm.completion, *args, model=model, **kwargs)
+ for model in models
+ ]
+
+ for future in futures:
+ try:
+ result = future.result()
+ if result is not None:
+ responses.append(result)
+ except Exception as e:
+ print_verbose(
+ f"batch_completion_models_all_responses: model request failed: {str(e)}"
+ )
+ continue
return responses
diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py
index f80eae20f3b..4b965d4e635 100644
--- a/litellm/batches/batch_utils.py
+++ b/litellm/batches/batch_utils.py
@@ -1,49 +1,67 @@
import json
-import time
from typing import Any, List, Literal, Optional, Tuple
-import httpx
-
import litellm
from litellm._logging import verbose_logger
-from litellm._uuid import uuid
from litellm.types.llms.openai import Batch
-from litellm.types.utils import CallTypes, ModelResponse, Usage
+from litellm.types.utils import CallTypes, ModelInfo, Usage
from litellm.utils import token_counter
async def calculate_batch_cost_and_usage(
file_content_dictionary: List[dict],
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
+ ],
model_name: Optional[str] = None,
+ model_info: Optional[ModelInfo] = None,
) -> Tuple[float, Usage, List[str]]:
"""
- Calculate the cost and usage of a batch
+ Calculate the cost and usage of a batch.
+
+ Args:
+ model_info: Optional deployment-level model info with custom batch
+ pricing. Threaded through to batch_cost_calculator so that
+ deployment-specific pricing (e.g. input_cost_per_token_batches)
+ is used instead of the global cost map.
"""
batch_cost = _batch_cost_calculator(
custom_llm_provider=custom_llm_provider,
file_content_dictionary=file_content_dictionary,
model_name=model_name,
+ model_info=model_info,
)
batch_usage = _get_batch_job_total_usage_from_file_content(
file_content_dictionary=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
)
- batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name)
+ batch_models = _get_batch_models_from_file_content(
+ file_content_dictionary, model_name
+ )
return batch_cost, batch_usage, batch_models
async def _handle_completed_batch(
batch: Batch,
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
+ ],
model_name: Optional[str] = None,
+ litellm_params: Optional[dict] = None,
) -> Tuple[float, Usage, List[str]]:
- """Helper function to process a completed batch and handle logging"""
+ """Helper function to process a completed batch and handle logging
+
+ Args:
+ batch: The batch object
+ custom_llm_provider: The LLM provider
+ model_name: Optional model name
+ litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
+ """
# Get batch results
file_content_dictionary = await _get_batch_output_file_content_as_dictionary(
- batch, custom_llm_provider
+ batch, custom_llm_provider, litellm_params=litellm_params
)
# Calculate costs and usage
@@ -58,7 +76,9 @@ async def _handle_completed_batch(
model_name=model_name,
)
- batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name)
+ batch_models = _get_batch_models_from_file_content(
+ file_content_dictionary, model_name
+ )
return batch_cost, batch_usage, batch_models
@@ -84,22 +104,28 @@ def _get_batch_models_from_file_content(
def _batch_cost_calculator(
file_content_dictionary: List[dict],
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
+ ] = "openai",
model_name: Optional[str] = None,
+ model_info: Optional[ModelInfo] = None,
) -> float:
"""
Calculate the cost of a batch based on the output file id
"""
# Handle Vertex AI with specialized method
if custom_llm_provider == "vertex_ai" and model_name:
- batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
+ batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(
+ file_content_dictionary, model_name
+ )
verbose_logger.debug("vertex_ai_total_cost=%s", batch_cost)
return batch_cost
-
+
# For other providers, use the existing logic
total_cost = _get_batch_job_cost_from_file_content(
file_content_dictionary=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
+ model_info=model_info,
)
verbose_logger.debug("total_cost=%s", total_cost)
return total_cost
@@ -110,73 +136,61 @@ def calculate_vertex_ai_batch_cost_and_usage(
model_name: Optional[str] = None,
) -> Tuple[float, Usage]:
"""
- Calculate both cost and usage from Vertex AI batch responses
+ Calculate both cost and usage from Vertex AI batch responses.
+
+ Vertex AI batch output lines have format:
+ {"request": ..., "status": "", "response": {"candidates": [...], "usageMetadata": {...}}}
+
+ usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount.
"""
- from litellm.litellm_core_utils.litellm_logging import Logging
- from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
- VertexGeminiConfig,
- )
+ from litellm.cost_calculator import batch_cost_calculator
+
total_cost = 0.0
total_tokens = 0
prompt_tokens = 0
completion_tokens = 0
-
- for response in vertex_ai_batch_responses:
- if response.get("status") == "JOB_STATE_SUCCEEDED": # Check if response was successful
- # Transform Vertex AI response to OpenAI format if needed
+ actual_model_name = model_name or "gemini-2.0-flash-001"
- # Create required arguments for the transformation method
- model_response = ModelResponse()
-
- # Ensure model_name is not None
- actual_model_name = model_name or "gemini-2.5-flash"
-
- # Create a real LiteLLM logging object
- logging_obj = Logging(
+ for response in vertex_ai_batch_responses:
+ response_body = response.get("response")
+ if response_body is None:
+ continue
+
+ usage_metadata = response_body.get("usageMetadata", {})
+ _prompt = usage_metadata.get("promptTokenCount", 0) or 0
+ _completion = usage_metadata.get("candidatesTokenCount", 0) or 0
+ _total = usage_metadata.get("totalTokenCount", 0) or (_prompt + _completion)
+
+ line_usage = Usage(
+ prompt_tokens=_prompt,
+ completion_tokens=_completion,
+ total_tokens=_total,
+ )
+
+ try:
+ p_cost, c_cost = batch_cost_calculator(
+ usage=line_usage,
model=actual_model_name,
- messages=[{"role": "user", "content": "batch_request"}],
- stream=False,
- call_type=CallTypes.aretrieve_batch,
- start_time=time.time(),
- litellm_call_id="batch_" + str(uuid.uuid4()),
- function_id="batch_processing",
- litellm_trace_id=str(uuid.uuid4()),
- kwargs={"optional_params": {}}
- )
-
- # Add the optional_params attribute that the Vertex AI transformation expects
- logging_obj.optional_params = {}
- raw_response = httpx.Response(200) # Mock response object
-
- openai_format_response = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
- completion_response=response["response"],
- model_response=model_response,
- model=actual_model_name,
- logging_obj=logging_obj,
- raw_response=raw_response,
- )
-
- # Calculate cost using existing function
- cost = litellm.completion_cost(
- completion_response=openai_format_response,
custom_llm_provider="vertex_ai",
- call_type=CallTypes.aretrieve_batch.value,
)
- total_cost += cost
-
- # Extract usage from the transformed response
- usage_obj = getattr(openai_format_response, 'usage', None)
- if usage_obj:
- usage = usage_obj
- else:
- # Fallback: create usage from response dict
- response_dict = openai_format_response.dict() if hasattr(openai_format_response, 'dict') else {}
- usage = _get_batch_job_usage_from_response_body(response_dict)
-
- total_tokens += usage.total_tokens
- prompt_tokens += usage.prompt_tokens
- completion_tokens += usage.completion_tokens
-
+ total_cost += p_cost + c_cost
+ except Exception as e:
+ verbose_logger.debug(
+ "vertex_ai batch cost calculation error for line: %s", str(e)
+ )
+
+ prompt_tokens += _prompt
+ completion_tokens += _completion
+ total_tokens += _total
+
+ verbose_logger.info(
+ "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d",
+ total_cost,
+ prompt_tokens,
+ completion_tokens,
+ total_tokens,
+ )
+
return total_cost, Usage(
total_tokens=total_tokens,
prompt_tokens=prompt_tokens,
@@ -186,10 +200,19 @@ def calculate_vertex_ai_batch_cost_and_usage(
async def _get_batch_output_file_content_as_dictionary(
batch: Batch,
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
+ ] = "openai",
+ litellm_params: Optional[dict] = None,
) -> List[dict]:
"""
Get the batch output file content as a list of dictionaries
+
+ Args:
+ batch: The batch object
+ custom_llm_provider: The LLM provider
+ litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
+ Required for Azure and other providers that need authentication
"""
from litellm.files.main import afile_content
from litellm.proxy.openai_files_endpoints.common_utils import (
@@ -206,18 +229,68 @@ async def _get_batch_output_file_content_as_dictionary(
is_base64_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
if is_base64_unified_file_id:
try:
- file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0]
- verbose_logger.debug(f"Extracted LLM output file ID from unified file ID: {file_id}")
+ file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(
+ ";"
+ )[0]
+ verbose_logger.debug(
+ f"Extracted LLM output file ID from unified file ID: {file_id}"
+ )
except (IndexError, AttributeError) as e:
- verbose_logger.error(f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}")
+ verbose_logger.error(
+ f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}"
+ )
- _file_content = await afile_content(
- file_id=file_id,
- custom_llm_provider=custom_llm_provider,
- )
+ # Build kwargs for afile_content with credentials from litellm_params
+ file_content_kwargs = {
+ "file_id": file_id,
+ "custom_llm_provider": custom_llm_provider,
+ }
+
+ # Extract and add credentials for file access
+ credentials = _extract_file_access_credentials(litellm_params)
+ file_content_kwargs.update(credentials)
+
+ _file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
return _get_file_content_as_dictionary(_file_content.content)
+def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
+ """
+ Extract credentials from litellm_params for file access operations.
+
+ This method extracts relevant authentication and configuration parameters
+ needed for accessing files across different providers (Azure, Vertex AI, etc.).
+
+ Args:
+ litellm_params: Dictionary containing litellm parameters with credentials
+
+ Returns:
+ Dictionary containing only the credentials needed for file access
+ """
+ credentials = {}
+
+ if litellm_params:
+ # List of credential keys that should be passed to file operations
+ credential_keys = [
+ "api_key",
+ "api_base",
+ "api_version",
+ "organization",
+ "azure_ad_token",
+ "azure_ad_token_provider",
+ "vertex_project",
+ "vertex_location",
+ "vertex_credentials",
+ "timeout",
+ "max_retries",
+ ]
+ for key in credential_keys:
+ if key in litellm_params:
+ credentials[key] = litellm_params[key]
+
+ return credentials
+
+
def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
"""
Get the file content as a list of dictionaries from JSON Lines format
@@ -237,11 +310,16 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
def _get_batch_job_cost_from_file_content(
file_content_dictionary: List[dict],
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
+ ] = "openai",
+ model_info: Optional[ModelInfo] = None,
) -> float:
"""
Get the cost of a batch job from the file content
"""
+ from litellm.cost_calculator import batch_cost_calculator
+
try:
total_cost: float = 0.0
# parse the file content as json
@@ -251,11 +329,22 @@ def _get_batch_job_cost_from_file_content(
for _item in file_content_dictionary:
if _batch_response_was_successful(_item):
_response_body = _get_response_from_batch_job_output_file(_item)
- total_cost += litellm.completion_cost(
- completion_response=_response_body,
- custom_llm_provider=custom_llm_provider,
- call_type=CallTypes.aretrieve_batch.value,
- )
+ if model_info is not None:
+ usage = _get_batch_job_usage_from_response_body(_response_body)
+ model = _response_body.get("model", "")
+ prompt_cost, completion_cost = batch_cost_calculator(
+ usage=usage,
+ model=model,
+ custom_llm_provider=custom_llm_provider,
+ model_info=model_info,
+ )
+ total_cost += prompt_cost + completion_cost
+ else:
+ total_cost += litellm.completion_cost(
+ completion_response=_response_body,
+ custom_llm_provider=custom_llm_provider,
+ call_type=CallTypes.aretrieve_batch.value,
+ )
verbose_logger.debug("total_cost=%s", total_cost)
return total_cost
except Exception as e:
@@ -265,7 +354,9 @@ def _get_batch_job_cost_from_file_content(
def _get_batch_job_total_usage_from_file_content(
file_content_dictionary: List[dict],
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
+ ] = "openai",
model_name: Optional[str] = None,
) -> Usage:
"""
@@ -273,9 +364,11 @@ def _get_batch_job_total_usage_from_file_content(
"""
# Handle Vertex AI with specialized method
if custom_llm_provider == "vertex_ai" and model_name:
- _, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
+ _, batch_usage = calculate_vertex_ai_batch_cost_and_usage(
+ file_content_dictionary, model_name
+ )
return batch_usage
-
+
# For other providers, use the existing logic
total_tokens: int = 0
prompt_tokens: int = 0
@@ -293,6 +386,7 @@ def _get_batch_job_total_usage_from_file_content(
completion_tokens=completion_tokens,
)
+
def _get_batch_job_input_file_usage(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
@@ -302,25 +396,26 @@ def _get_batch_job_input_file_usage(
Count the number of tokens in the input file
Used for batch rate limiting to count the number of tokens in the input file
- """
+ """
prompt_tokens: int = 0
completion_tokens: int = 0
-
+
for _item in file_content_dictionary:
body = _item.get("body", {})
model = body.get("model", model_name or "")
messages = body.get("messages", [])
-
+
if messages:
item_tokens = token_counter(model=model, messages=messages)
prompt_tokens += item_tokens
-
+
return Usage(
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
)
+
def _get_batch_job_usage_from_response_body(response_body: dict) -> Usage:
"""
Get the tokens of a batch job from the response body
@@ -344,4 +439,4 @@ def _batch_response_was_successful(batch_job_output_file: dict) -> bool:
Check if the batch job response status == 200
"""
_response: dict = batch_job_output_file.get("response", None) or {}
- return _response.get("status_code", None) == 200
\ No newline at end of file
+ return _response.get("status_code", None) == 200
diff --git a/litellm/batches/main.py b/litellm/batches/main.py
index 25f6e284bcd..e176dc42921 100644
--- a/litellm/batches/main.py
+++ b/litellm/batches/main.py
@@ -33,11 +33,14 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
CancelBatchRequest,
CreateBatchRequest,
+ FileExpiresAfter,
RetrieveBatchRequest,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
+ LIST_BATCHES_SUPPORTED_PROVIDERS,
OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS,
+ ListBatchesSupportedProvider,
LiteLLMBatch,
LlmProviders,
)
@@ -106,10 +109,13 @@ async def acreate_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"
+ ] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
+ output_expires_after: Optional[Dict[str, Any]] = None,
**kwargs,
) -> LiteLLMBatch:
"""
@@ -131,6 +137,7 @@ async def acreate_batch(
metadata,
extra_headers,
extra_body,
+ output_expires_after,
**kwargs,
)
@@ -150,14 +157,17 @@ async def acreate_batch(
@client
-def create_batch(
+def create_batch( # noqa: PLR0915
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"
+ ] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
+ output_expires_after: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
"""
@@ -189,7 +199,8 @@ def create_batch(
)
### TIMEOUT LOGIC ###
timeout = _resolve_timeout(optional_params, kwargs, custom_llm_provider)
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model=model,
user=None,
optional_params=optional_params.model_dump(),
@@ -197,7 +208,6 @@ def create_batch(
"litellm_call_id": litellm_call_id,
"proxy_server_request": proxy_server_request,
"model_info": model_info,
- "metadata": metadata,
"preset_cache_key": None,
"stream_response": {},
**optional_params.model_dump(exclude_unset=True),
@@ -213,6 +223,10 @@ def create_batch(
extra_headers=extra_headers,
extra_body=extra_body,
)
+ if output_expires_after is not None:
+ _create_batch_request["output_expires_after"] = cast(
+ FileExpiresAfter, output_expires_after
+ )
if model is not None:
provider_config = ProviderConfigManager.get_provider_batches_config(
model=model,
@@ -356,7 +370,9 @@ def create_batch(
@client
async def aretrieve_batch(
batch_id: str,
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"
+ ] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -402,7 +418,9 @@ def _handle_retrieve_batch_providers_without_provider_config(
litellm_params: dict,
_retrieve_batch_request: RetrieveBatchRequest,
_is_async: bool,
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"
+ ] = "openai",
logging_obj: Optional[Any] = None,
):
api_base: Optional[str] = None
@@ -541,7 +559,9 @@ def _handle_retrieve_batch_providers_without_provider_config(
@client
def retrieve_batch(
batch_id: str,
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
+ custom_llm_provider: Literal[
+ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"
+ ] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -564,7 +584,8 @@ def retrieve_batch(
**kwargs,
)
if litellm_logging_obj is not None:
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model=None,
user=None,
optional_params=optional_params.model_dump(),
@@ -674,7 +695,7 @@ def retrieve_batch(
async def alist_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
- custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai",
+ custom_llm_provider: ListBatchesSupportedProvider = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -717,7 +738,7 @@ async def alist_batches(
def list_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
- custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai",
+ custom_llm_provider: ListBatchesSupportedProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -843,8 +864,9 @@ def list_batches(
)
else:
raise litellm.exceptions.BadRequestError(
- message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: openai, azure, vertex_ai.".format(
- custom_llm_provider
+ message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: {}.".format(
+ custom_llm_provider,
+ ", ".join(sorted(LIST_BATCHES_SUPPORTED_PROVIDERS)),
),
model="n/a",
llm_provider=custom_llm_provider,
@@ -920,7 +942,6 @@ def cancel_batch(
LiteLLM Equivalent of POST https://api.openai.com/v1/batches/{batch_id}/cancel
"""
try:
-
try:
if model is not None:
_, custom_llm_provider, _, _ = get_llm_provider(
@@ -1088,25 +1109,40 @@ def _handle_async_invoke_status(
"inprogress": "in_progress",
"in_progress": "in_progress",
}
- normalized_status: BatchJobStatus = status_mapping.get(aws_status_lower, "failed") # Default to "failed" if unknown status
+ normalized_status: BatchJobStatus = status_mapping.get(
+ aws_status_lower, "failed"
+ ) # Default to "failed" if unknown status
# Get output S3 URI safely
output_s3_uri = ""
try:
- output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"]
+ output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"][
+ "s3Uri"
+ ]
except (KeyError, TypeError):
pass
-
+
# Use BedrockBatchesConfig's timestamp parsing method (expects raw AWS status string)
import time
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
- created_at, in_progress_at, completed_at, failed_at, _, _ = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw)
+
+ (
+ created_at,
+ in_progress_at,
+ completed_at,
+ failed_at,
+ _,
+ _,
+ ) = BedrockBatchesConfig()._parse_timestamps_and_status(
+ status_response, aws_status_raw
+ )
result = LiteLLMBatch(
id=status_response["invocationArn"],
object="batch",
status=normalized_status,
- created_at=created_at or int(time.time()), # Provide default timestamp if None
+ created_at=created_at
+ or int(time.time()), # Provide default timestamp if None
in_progress_at=in_progress_at,
completed_at=completed_at,
failed_at=failed_at,
diff --git a/litellm/blog_posts.json b/litellm/blog_posts.json
new file mode 100644
index 00000000000..fa768b3ec57
--- /dev/null
+++ b/litellm/blog_posts.json
@@ -0,0 +1,10 @@
+{
+ "posts": [
+ {
+ "title": "Realtime WebRTC HTTP Endpoints",
+ "description": "Use the LiteLLM proxy to route OpenAI-style WebRTC realtime via HTTP: client_secrets and SDP exchange.",
+ "date": "2026-03-12",
+ "url": "https://docs.litellm.ai/blog/realtime_webrtc_http_endpoints"
+ }
+ ]
+}
diff --git a/litellm/caching/azure_blob_cache.py b/litellm/caching/azure_blob_cache.py
index 45e551bdae9..a2246640c30 100644
--- a/litellm/caching/azure_blob_cache.py
+++ b/litellm/caching/azure_blob_cache.py
@@ -22,7 +22,9 @@ class AzureBlobCache(BaseCache):
from azure.storage.blob import BlobServiceClient
from azure.core.exceptions import ResourceExistsError
from azure.identity import DefaultAzureCredential
- from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
+ from azure.identity.aio import (
+ DefaultAzureCredential as AsyncDefaultAzureCredential,
+ )
from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
self.container_client = BlobServiceClient(
@@ -50,14 +52,16 @@ class AzureBlobCache(BaseCache):
print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}")
serialized_value = json.dumps(value)
try:
- await self.async_container_client.upload_blob(key, serialized_value, overwrite=True)
+ await self.async_container_client.upload_blob(
+ key, serialized_value, overwrite=True
+ )
except Exception as e:
# NON blocking - notify users Azure Blob is throwing an exception
print_verbose(f"LiteLLM set_cache() - Got exception from Azure Blob: {e}")
def get_cache(self, key, **kwargs):
from azure.core.exceptions import ResourceNotFoundError
-
+
try:
print_verbose(f"Get Azure Blob Cache: key: {key}")
as_bytes = self.container_client.download_blob(key).readall()
@@ -74,7 +78,7 @@ class AzureBlobCache(BaseCache):
async def async_get_cache(self, key, **kwargs):
from azure.core.exceptions import ResourceNotFoundError
-
+
try:
print_verbose(f"Get Azure Blob Cache: key: {key}")
blob = await self.async_container_client.download_blob(key)
diff --git a/litellm/caching/base_cache.py b/litellm/caching/base_cache.py
index 8660e64efde..81f1d61bd0d 100644
--- a/litellm/caching/base_cache.py
+++ b/litellm/caching/base_cache.py
@@ -53,12 +53,12 @@ class BaseCache(ABC):
async def disconnect(self):
raise NotImplementedError
-
+
async def test_connection(self) -> dict:
"""
Test the cache connection.
-
+
Returns:
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
"""
- raise NotImplementedError
\ No newline at end of file
+ raise NotImplementedError
diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py
index a03bff60686..406a4f8c98a 100644
--- a/litellm/caching/caching.py
+++ b/litellm/caching/caching.py
@@ -108,6 +108,7 @@ class Cache:
qdrant_collection_name: Optional[str] = None,
qdrant_quantization_config: Optional[str] = None,
qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002",
+ qdrant_semantic_cache_vector_size: Optional[int] = None,
# GCP IAM authentication parameters
gcp_service_account: Optional[str] = None,
gcp_ssl_ca_certs: Optional[str] = None,
@@ -165,6 +166,14 @@ class Cache:
None. Cache is set as a litellm param
"""
if type == LiteLLMCacheType.REDIS:
+ # Check REDIS_CLUSTER_NODES env var if no explicit startup nodes
+ if not redis_startup_nodes:
+ _env_cluster_nodes = litellm.get_secret("REDIS_CLUSTER_NODES")
+ if _env_cluster_nodes is not None and isinstance(
+ _env_cluster_nodes, str
+ ):
+ redis_startup_nodes = json.loads(_env_cluster_nodes)
+
if redis_startup_nodes:
# Only pass GCP parameters if they are provided
cluster_kwargs = {
@@ -207,6 +216,7 @@ class Cache:
similarity_threshold=similarity_threshold,
quantization_config=qdrant_quantization_config,
embedding_model=qdrant_semantic_cache_embedding_model,
+ vector_size=qdrant_semantic_cache_vector_size,
)
elif type == LiteLLMCacheType.LOCAL:
self.cache = InMemoryCache()
diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py
index 4e97197a9de..7cdbd3fc03d 100644
--- a/litellm/caching/caching_handler.py
+++ b/litellm/caching/caching_handler.py
@@ -78,9 +78,7 @@ class CachingHandlerResponse(BaseModel):
cached_result: Optional[Any] = None
final_embedding_cached_response: Optional[EmbeddingResponse] = None
- embedding_all_elements_cache_hit: bool = (
- False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
- )
+ embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
in_memory_cache_obj = InMemoryCache()
@@ -159,7 +157,7 @@ class LLMCachingHandler:
#########################################################
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
kwargs["parent_otel_span"] = parent_otel_span
-
+
if litellm.cache is not None and self._is_call_type_supported_by_cache(
original_function=original_function
):
@@ -181,7 +179,9 @@ class LLMCachingHandler:
api_base=kwargs.get("api_base", None),
api_key=kwargs.get("api_key", None),
)
- cache_duration_ms = (cache_check_end_time - cache_check_start_time) * 1000
+ cache_duration_ms = (
+ cache_check_end_time - cache_check_start_time
+ ) * 1000
self._update_litellm_logging_obj_environment(
logging_obj=logging_obj,
model=model,
@@ -194,7 +194,6 @@ class LLMCachingHandler:
call_type = original_function.__name__
-
cached_result = self._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=call_type,
@@ -244,7 +243,7 @@ class LLMCachingHandler:
final_embedding_cached_response=final_embedding_cached_response,
embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
)
-
+
verbose_logger.debug(f"CACHE RESULT: {cached_result}")
return CachingHandlerResponse(
cached_result=cached_result,
@@ -265,9 +264,8 @@ class LLMCachingHandler:
) -> CachingHandlerResponse:
from litellm.utils import CustomStreamWrapper
-
cached_result: Optional[Any] = None
-
+
# Check if caching should be performed BEFORE doing expensive kwargs copy
if litellm.cache is not None and self._is_call_type_supported_by_cache(
original_function=original_function
@@ -325,7 +323,7 @@ class LLMCachingHandler:
result=cached_result,
start_time=start_time,
end_time=end_time,
- cache_hit=cache_hit
+ cache_hit=cache_hit,
)
cache_key = litellm.cache.get_cache_key(**kwargs)
if (
@@ -554,12 +552,18 @@ class LLMCachingHandler:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=logging_obj.async_success_handler(
- result=cached_result, start_time=start_time, end_time=end_time, cache_hit=cache_hit
+ result=cached_result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
)
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
- result=cached_result, start_time=start_time, end_time=end_time, cache_hit=cache_hit
+ result=cached_result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
)
async def _retrieve_from_cache(
@@ -728,10 +732,9 @@ class LLMCachingHandler:
response_type="audio_transcription",
hidden_params=hidden_params,
)
- elif (
- call_type == "aresponses"
- or call_type == "responses"
- ) and isinstance(cached_result, dict):
+ elif (call_type == "aresponses" or call_type == "responses") and isinstance(
+ cached_result, dict
+ ):
# Convert cached dict back to ResponsesAPIResponse object
cached_result = ResponsesAPIResponse(**cached_result)
@@ -741,7 +744,7 @@ class LLMCachingHandler:
and isinstance(cached_result._hidden_params, dict)
):
cached_result._hidden_params["cache_hit"] = True
-
+
#########################################################
# Add final timing metrics to the cached result
#########################################################
@@ -1011,9 +1014,9 @@ class LLMCachingHandler:
}
if litellm.cache is not None:
- litellm_params["preset_cache_key"] = (
- litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
- )
+ litellm_params[
+ "preset_cache_key"
+ ] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
else:
litellm_params["preset_cache_key"] = None
diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py
index 3edc3f42820..4020b8cc22e 100644
--- a/litellm/caching/dual_cache.py
+++ b/litellm/caching/dual_cache.py
@@ -12,7 +12,8 @@ import asyncio
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
-from typing import TYPE_CHECKING, Any, List, Optional, Union
+from threading import Lock
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
if TYPE_CHECKING:
from litellm.types.caching import RedisPipelineIncrementOperation
@@ -71,6 +72,7 @@ class DualCache(BaseCache):
self.last_redis_batch_access_time = LimitedSizeOrderedDict(
max_size=default_max_redis_batch_cache_size
)
+ self._last_redis_batch_access_time_lock = Lock()
self.redis_batch_cache_expiry = (
default_redis_batch_cache_expiry
or litellm.default_redis_batch_cache_expiry
@@ -236,22 +238,46 @@ class DualCache(BaseCache):
except Exception:
verbose_logger.error(traceback.format_exc())
- def get_redis_batch_keys(
+ def _reserve_redis_batch_keys(
self,
current_time: float,
keys: List[str],
result: List[Any],
- ) -> List[str]:
- sublist_keys = []
- for key, value in zip(keys, result):
- if value is None:
+ ) -> Tuple[List[str], Dict[str, Optional[float]]]:
+ """
+ Atomically choose keys to fetch from Redis and reserve their access time.
+ This prevents check-then-act races under concurrent async callers.
+ """
+ sublist_keys: List[str] = []
+ previous_access_times: Dict[str, Optional[float]] = {}
+
+ with self._last_redis_batch_access_time_lock:
+ for key, value in zip(keys, result):
+ if value is not None:
+ continue
+
if (
key not in self.last_redis_batch_access_time
or current_time - self.last_redis_batch_access_time[key]
>= self.redis_batch_cache_expiry
):
sublist_keys.append(key)
- return sublist_keys
+ previous_access_times[key] = self.last_redis_batch_access_time.get(
+ key
+ )
+ self.last_redis_batch_access_time[key] = current_time
+
+ return sublist_keys, previous_access_times
+
+ def _rollback_redis_batch_key_reservations(
+ self, previous_access_times: Dict[str, Optional[float]]
+ ) -> None:
+ with self._last_redis_batch_access_time_lock:
+ for key, previous_time in previous_access_times.items():
+ if previous_time is None:
+ self.last_redis_batch_access_time.pop(key, None)
+ else:
+ self.last_redis_batch_access_time[key] = previous_time
async def async_batch_get_cache(
self,
@@ -276,31 +302,37 @@ class DualCache(BaseCache):
- check the redis cache
"""
current_time = time.time()
- sublist_keys = self.get_redis_batch_keys(current_time, keys, result)
+ sublist_keys, previous_access_times = self._reserve_redis_batch_keys(
+ current_time, keys, result
+ )
- # Only hit Redis if the last access time was more than 5 seconds ago
+ # Only hit Redis if enough time has passed since last access.
if len(sublist_keys) > 0:
- # If not found in in-memory cache, try fetching from Redis
- redis_result = await self.redis_cache.async_batch_get_cache(
- sublist_keys, parent_otel_span=parent_otel_span
- )
-
- # Update the last access time for ALL queried keys
- # This includes keys with None values to throttle repeated Redis queries
- for key in sublist_keys:
- self.last_redis_batch_access_time[key] = current_time
-
+ try:
+ # If not found in in-memory cache, try fetching from Redis
+ redis_result = await self.redis_cache.async_batch_get_cache(
+ sublist_keys, parent_otel_span=parent_otel_span
+ )
+ except Exception:
+ # Do not throttle subsequent callers if the Redis read fails.
+ self._rollback_redis_batch_key_reservations(
+ previous_access_times
+ )
+ raise
+
# Short-circuit if redis_result is None or contains only None values
- if redis_result is None or all(v is None for v in redis_result.values()):
+ if redis_result is None or all(
+ v is None for v in redis_result.values()
+ ):
return result
# Pre-compute key-to-index mapping for O(1) lookup
key_to_index = {key: i for i, key in enumerate(keys)}
-
+
# Update both result and in-memory cache in a single loop
for key, value in redis_result.items():
result[key_to_index[key]] = value
-
+
if value is not None and self.in_memory_cache is not None:
await self.in_memory_cache.async_set_cache(
key, value, **kwargs
@@ -316,6 +348,8 @@ class DualCache(BaseCache):
)
try:
if self.in_memory_cache is not None:
+ if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
+ kwargs["ttl"] = self.default_in_memory_ttl
await self.in_memory_cache.async_set_cache(key, value, **kwargs)
if self.redis_cache is not None and local_only is False:
@@ -337,6 +371,8 @@ class DualCache(BaseCache):
)
try:
if self.in_memory_cache is not None:
+ if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
+ kwargs["ttl"] = self.default_in_memory_ttl
await self.in_memory_cache.async_set_cache_pipeline(
cache_list=cache_list, **kwargs
)
diff --git a/litellm/caching/gcs_cache.py b/litellm/caching/gcs_cache.py
index 88857ba0e70..a5bd092f154 100644
--- a/litellm/caching/gcs_cache.py
+++ b/litellm/caching/gcs_cache.py
@@ -16,13 +16,23 @@ from .base_cache import BaseCache
class GCSCache(BaseCache):
- def __init__(self, bucket_name: Optional[str] = None, path_service_account: Optional[str] = None, gcs_path: Optional[str] = None) -> None:
+ def __init__(
+ self,
+ bucket_name: Optional[str] = None,
+ path_service_account: Optional[str] = None,
+ gcs_path: Optional[str] = None,
+ ) -> None:
super().__init__()
self.bucket_name = bucket_name or GCSBucketBase(bucket_name=None).BUCKET_NAME
- self.path_service_account = path_service_account or GCSBucketBase(bucket_name=None).path_service_account_json
+ self.path_service_account = (
+ path_service_account
+ or GCSBucketBase(bucket_name=None).path_service_account_json
+ )
self.key_prefix = gcs_path.rstrip("/") + "/" if gcs_path else ""
# create httpx clients
- self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
+ self.async_client = get_async_httpx_client(
+ llm_provider=httpxSpecialProvider.LoggingCallback
+ )
self.sync_client = _get_httpx_client()
def _construct_headers(self) -> dict:
@@ -52,7 +62,9 @@ class GCSCache(BaseCache):
data = json.dumps(value)
await self.async_client.post(url=url, data=data, headers=headers)
except Exception as e:
- print_verbose(f"GCS Caching: async_set_cache() - Got exception from GCS: {e}")
+ print_verbose(
+ f"GCS Caching: async_set_cache() - Got exception from GCS: {e}"
+ )
def get_cache(self, key, **kwargs):
try:
@@ -69,7 +81,9 @@ class GCSCache(BaseCache):
return cached_response
return None
except Exception as e:
- verbose_logger.error(f"GCS Caching: get_cache() - Got exception from GCS: {e}")
+ verbose_logger.error(
+ f"GCS Caching: get_cache() - Got exception from GCS: {e}"
+ )
async def async_get_cache(self, key, **kwargs):
try:
@@ -82,7 +96,9 @@ class GCSCache(BaseCache):
return json.loads(response.text)
return None
except Exception as e:
- verbose_logger.error(f"GCS Caching: async_get_cache() - Got exception from GCS: {e}")
+ verbose_logger.error(
+ f"GCS Caching: async_get_cache() - Got exception from GCS: {e}"
+ )
def flush_cache(self):
pass
diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py
index 16eb824f4c9..c2274713bb9 100644
--- a/litellm/caching/llm_caching_handler.py
+++ b/litellm/caching/llm_caching_handler.py
@@ -8,6 +8,17 @@ from .in_memory_cache import InMemoryCache
class LLMClientCache(InMemoryCache):
+ """Cache for LLM HTTP clients (OpenAI, Azure, httpx, etc.).
+
+ IMPORTANT: This cache intentionally does NOT close clients on eviction.
+ Evicted clients may still be in use by in-flight requests. Closing them
+ eagerly causes ``RuntimeError: Cannot send a request, as the client has
+ been closed.`` errors in production after the TTL (1 hour) expires.
+
+ Clients that are no longer referenced will be garbage-collected normally.
+ For explicit shutdown cleanup, use ``close_litellm_async_clients()``.
+ """
+
def update_cache_key_with_event_loop(self, key):
"""
Add the event loop to the cache key, to prevent event loop closed errors.
diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py
index 0e77b5a6c21..5e3713e5a15 100644
--- a/litellm/caching/qdrant_semantic_cache.py
+++ b/litellm/caching/qdrant_semantic_cache.py
@@ -31,6 +31,7 @@ class QdrantSemanticCache(BaseCache):
quantization_config=None,
embedding_model="text-embedding-ada-002",
host_type=None,
+ vector_size=None,
):
import os
@@ -53,6 +54,9 @@ class QdrantSemanticCache(BaseCache):
raise Exception("similarity_threshold must be provided, passed None")
self.similarity_threshold = similarity_threshold
self.embedding_model = embedding_model
+ self.vector_size = (
+ vector_size if vector_size is not None else QDRANT_VECTOR_SIZE
+ )
headers = {}
# check if defined as os.environ/ variable
@@ -138,7 +142,7 @@ class QdrantSemanticCache(BaseCache):
new_collection_status = self.sync_client.put(
url=f"{self.qdrant_api_base}/collections/{self.collection_name}",
json={
- "vectors": {"size": QDRANT_VECTOR_SIZE, "distance": "Cosine"},
+ "vectors": {"size": self.vector_size, "distance": "Cosine"},
"quantization_config": quantization_params,
},
headers=self.headers,
diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py
index ea7e3f5a979..82794c116f2 100644
--- a/litellm/caching/redis_cache.py
+++ b/litellm/caching/redis_cache.py
@@ -22,7 +22,11 @@ from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_REDIS_MAJOR_VERSION
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
-from litellm.types.caching import RedisPipelineIncrementOperation
+from litellm.types.caching import (
+ RedisPipelineIncrementOperation,
+ RedisPipelineLpopOperation,
+ RedisPipelineRpushOperation,
+)
from litellm.types.services import ServiceTypes
from .base_cache import BaseCache
@@ -264,19 +268,19 @@ class RedisCache(BaseCache):
def _parse_redis_major_version(self) -> int:
"""
Parse Redis version to extract the major version number.
-
+
Handles multiple version formats:
- Strings: "7.0.0", "6", "7.0.0-rc1", " 7.0.0 "
- Floats: 7.0 (e.g., from AWS ElastiCache Valkey)
- Integers: 7
- Malformed: "latest", "", "Unknown" (defaults to DEFAULT_REDIS_MAJOR_VERSION)
-
+
Returns:
int: The major version number (defaults to DEFAULT_REDIS_MAJOR_VERSION if unparseable)
"""
if self.redis_version == "Unknown":
return DEFAULT_REDIS_MAJOR_VERSION
-
+
try:
version_str = str(self.redis_version).strip()
# Handle cases where there's no dot (e.g., "7" or 7)
@@ -1105,14 +1109,18 @@ class RedisCache(BaseCache):
async def disconnect(self):
await self.async_redis_conn_pool.disconnect(inuse_connections=True)
-
+ try:
+ self.redis_client.close()
+ except Exception as e:
+ verbose_logger.debug("Error closing sync Redis client: %s", e)
+
async def test_connection(self) -> dict:
"""
Test the Redis connection by creating a new client and pinging it.
-
+
This creates a fresh connection without using cached clients or connection pools
to ensure the credentials are actually valid.
-
+
Returns:
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
"""
@@ -1121,29 +1129,26 @@ class RedisCache(BaseCache):
# Create a fresh Redis client with current settings
redis_client = redis_async.Redis(**self.redis_kwargs)
-
+
# Test the connection
- ping_result = await redis_client.ping()
+ ping_result = await redis_client.ping() # type: ignore[misc]
# Close the connection
await redis_client.aclose() # type: ignore[attr-defined]
-
+
if ping_result:
return {
"status": "success",
- "message": "Redis connection test successful"
+ "message": "Redis connection test successful",
}
else:
- return {
- "status": "failed",
- "message": "Redis ping returned False"
- }
+ return {"status": "failed", "message": "Redis ping returned False"}
except Exception as e:
verbose_logger.error(f"Redis connection test failed: {str(e)}")
return {
"status": "failed",
"message": f"Redis connection failed: {str(e)}",
- "error": str(e)
+ "error": str(e),
}
async def async_delete_cache(self, key: str):
@@ -1316,6 +1321,75 @@ class RedisCache(BaseCache):
)
raise e
+ async def _pipeline_rpush_helper(
+ self,
+ pipe: pipeline,
+ rpush_list: List[RedisPipelineRpushOperation],
+ ) -> List[int]:
+ """Helper function for pipeline rpush operations"""
+ for rpush_op in rpush_list:
+ pipe.rpush(rpush_op["key"], *rpush_op["values"])
+ results = await pipe.execute()
+ # Preserve positional correspondence — raise on per-command errors
+ for r in results:
+ if isinstance(r, Exception):
+ raise r
+ return results
+
+ async def async_rpush_pipeline(
+ self,
+ rpush_list: List[RedisPipelineRpushOperation],
+ ) -> List[int]:
+ """
+ Use Redis Pipelines for bulk RPUSH operations
+
+ Args:
+ rpush_list: List of RedisPipelineRpushOperation dicts containing:
+ - key: str
+ - values: List[Any]
+
+ Returns:
+ List[int]: List lengths after each push
+ """
+ if len(rpush_list) == 0:
+ return []
+
+ _redis_client: Any = self.init_async_client()
+ start_time = time.time()
+
+ try:
+ async with _redis_client.pipeline(transaction=False) as pipe:
+ results = await self._pipeline_rpush_helper(pipe, rpush_list)
+
+ ## LOGGING ##
+ end_time = time.time()
+ _duration = end_time - start_time
+ asyncio.create_task(
+ self.service_logger_obj.async_service_success_hook(
+ service=ServiceTypes.REDIS,
+ duration=_duration,
+ call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
+ )
+ )
+ return results
+ except Exception as e:
+ ## LOGGING ##
+ end_time = time.time()
+ _duration = end_time - start_time
+ asyncio.create_task(
+ self.service_logger_obj.async_service_failure_hook(
+ service=ServiceTypes.REDIS,
+ duration=_duration,
+ error=e,
+ call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
+ )
+ )
+ verbose_logger.error(
+ "LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s",
+ str(e),
+ )
+ raise e
+
async def handle_lpop_count_for_older_redis_versions(
self, pipe: pipeline, key: str, count: int
) -> List[bytes]:
@@ -1396,3 +1470,120 @@ class RedisCache(BaseCache):
f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}"
)
raise e
+
+ async def _pipeline_lpop_helper(
+ self,
+ pipe: pipeline,
+ lpop_list: List[RedisPipelineLpopOperation],
+ ) -> List[Optional[List[str]]]:
+ """Helper function for pipeline lpop operations.
+
+ For Redis >= 7, queues one LPOP(key, count) per operation.
+ For Redis < 7, queues `count` individual LPOP(key) commands per operation.
+ """
+ major_version = self._parse_redis_major_version()
+
+ if major_version >= 7:
+ for lpop_op in lpop_list:
+ pipe.lpop(lpop_op["key"], lpop_op["count"])
+ raw_results = await pipe.execute()
+ else:
+ # For Redis < 7, LPOP doesn't support count param.
+ # Issue `count` individual LPOP commands per key, all in one pipeline.
+ counts: List[int] = []
+ for lpop_op in lpop_list:
+ count = lpop_op["count"] or 1
+ counts.append(count)
+ for _ in range(count):
+ pipe.lpop(lpop_op["key"])
+ flat_results = await pipe.execute()
+
+ # Re-group the flat results back into per-key lists
+ raw_results = []
+ offset = 0
+ for count in counts:
+ key_results = [
+ r for r in flat_results[offset : offset + count] if r is not None
+ ]
+ raw_results.append(key_results if key_results else None)
+ offset += count
+
+ # Raise on per-command errors (matches _pipeline_rpush_helper behavior)
+ for r in raw_results:
+ if isinstance(r, Exception):
+ raise r
+
+ # Decode bytes -> str for each result set
+ decoded_results: List[Optional[List[str]]] = []
+ for r in raw_results:
+ if r is None:
+ decoded_results.append(None)
+ elif isinstance(r, list):
+ try:
+ decoded_results.append(
+ [
+ item.decode("utf-8") if isinstance(item, bytes) else item
+ for item in r
+ if item is not None
+ ]
+ or None
+ )
+ except Exception:
+ decoded_results.append(r) # type: ignore
+ else:
+ decoded_results.append(None)
+ return decoded_results
+
+ async def async_lpop_pipeline(
+ self,
+ lpop_list: List[RedisPipelineLpopOperation],
+ ) -> List[Optional[List[str]]]:
+ """
+ Use Redis Pipelines for bulk LPOP operations
+
+ Args:
+ lpop_list: List of RedisPipelineLpopOperation dicts containing:
+ - key: str
+ - count: Optional[int]
+
+ Returns:
+ List[Optional[List[str]]]: Decoded results per key, None if key was empty
+ """
+ if len(lpop_list) == 0:
+ return []
+
+ _redis_client: Any = self.init_async_client()
+ start_time = time.time()
+
+ try:
+ async with _redis_client.pipeline(transaction=False) as pipe:
+ results = await self._pipeline_lpop_helper(pipe, lpop_list)
+
+ ## LOGGING ##
+ end_time = time.time()
+ _duration = end_time - start_time
+ asyncio.create_task(
+ self.service_logger_obj.async_service_success_hook(
+ service=ServiceTypes.REDIS,
+ duration=_duration,
+ call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
+ )
+ )
+ return results
+ except Exception as e:
+ ## LOGGING ##
+ end_time = time.time()
+ _duration = end_time - start_time
+ asyncio.create_task(
+ self.service_logger_obj.async_service_failure_hook(
+ service=ServiceTypes.REDIS,
+ duration=_duration,
+ error=e,
+ call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
+ )
+ )
+ verbose_logger.error(
+ "LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s",
+ str(e),
+ )
+ raise e
diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py
index 91fcf1d7288..b0f5754f58e 100644
--- a/litellm/caching/redis_cluster_cache.py
+++ b/litellm/caching/redis_cluster_cache.py
@@ -57,11 +57,11 @@ class RedisClusterCache(RedisCache):
"""
async_redis_cluster_client = self.init_async_client()
return await async_redis_cluster_client.mget_nonatomic(keys=keys) # type: ignore
-
+
async def test_connection(self) -> dict:
"""
Test the Redis Cluster connection.
-
+
Returns:
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
"""
@@ -72,37 +72,38 @@ class RedisClusterCache(RedisCache):
# Create ClusterNode objects from startup_nodes
cluster_kwargs = self.redis_kwargs.copy()
startup_nodes = cluster_kwargs.pop("startup_nodes", [])
-
+
new_startup_nodes: List[ClusterNode] = []
for item in startup_nodes:
new_startup_nodes.append(ClusterNode(**item))
-
+
# Create a fresh Redis Cluster client with current settings
redis_client = redis_async.RedisCluster(
startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore
)
-
+
# Test the connection
- ping_result = await redis_client.ping() # type: ignore[attr-defined]
+ ping_result = await redis_client.ping() # type: ignore[attr-defined, misc]
# Close the connection
await redis_client.aclose() # type: ignore[attr-defined]
-
+
if ping_result:
return {
"status": "success",
- "message": "Redis Cluster connection test successful"
+ "message": "Redis Cluster connection test successful",
}
else:
return {
"status": "failed",
- "message": "Redis Cluster ping returned False"
+ "message": "Redis Cluster ping returned False",
}
except Exception as e:
from litellm._logging import verbose_logger
+
verbose_logger.error(f"Redis Cluster connection test failed: {str(e)}")
return {
"status": "failed",
"message": f"Redis Cluster connection failed: {str(e)}",
- "error": str(e)
- }
\ No newline at end of file
+ "error": str(e),
+ }
diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py
index 180964605f6..e26fbe8981c 100644
--- a/litellm/caching/s3_cache.py
+++ b/litellm/caching/s3_cache.py
@@ -110,7 +110,9 @@ class S3Cache(BaseCache):
func = partial(self.set_cache, key, value, **kwargs)
await loop.run_in_executor(None, func)
except Exception as e:
- verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}")
+ verbose_logger.error(
+ f"S3 Caching: async_set_cache() - Got exception from S3: {e}"
+ )
def get_cache(self, key, **kwargs):
import botocore
@@ -126,7 +128,7 @@ class S3Cache(BaseCache):
if cached_response is not None:
if "Expires" in cached_response:
- expires_time = cached_response['Expires']
+ expires_time = cached_response["Expires"]
current_time = datetime.now(expires_time.tzinfo)
if current_time > expires_time:
diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py
index 5c051797e8b..2164a2c0f01 100644
--- a/litellm/completion_extras/litellm_responses_transformation/handler.py
+++ b/litellm/completion_extras/litellm_responses_transformation/handler.py
@@ -61,9 +61,7 @@ class ResponsesToCompletionBridgeHandler:
existing.setdefault(key, value)
return response
- def _collect_response_from_stream(
- self, stream_iter: Any
- ) -> "ResponsesAPIResponse":
+ def _collect_response_from_stream(self, stream_iter: Any) -> "ResponsesAPIResponse":
for _ in stream_iter:
pass
@@ -144,7 +142,9 @@ class ResponsesToCompletionBridgeHandler:
custom_llm_provider=custom_llm_provider,
)
- def completion(self, *args, **kwargs) -> Union[
+ def completion(
+ self, *args, **kwargs
+ ) -> Union[
Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]],
"ModelResponse",
"CustomStreamWrapper",
@@ -221,7 +221,9 @@ class ResponsesToCompletionBridgeHandler:
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
- return streamwrapper
+ return self._apply_post_stream_processing(
+ streamwrapper, model, custom_llm_provider
+ )
async def acompletion(
self, *args, **kwargs
@@ -300,7 +302,30 @@ class ResponsesToCompletionBridgeHandler:
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
- return streamwrapper
+ return self._apply_post_stream_processing(
+ streamwrapper, model, custom_llm_provider
+ )
+
+ @staticmethod
+ def _apply_post_stream_processing(
+ stream: "CustomStreamWrapper",
+ model: str,
+ custom_llm_provider: str,
+ ) -> Any:
+ """Apply provider-specific post-stream processing if available."""
+ from litellm.types.utils import LlmProviders
+ from litellm.utils import ProviderConfigManager
+
+ try:
+ provider_config = ProviderConfigManager.get_provider_chat_config(
+ model=model, provider=LlmProviders(custom_llm_provider)
+ )
+ except (ValueError, KeyError):
+ return stream
+
+ if provider_config is not None:
+ return provider_config.post_stream_processing(stream)
+ return stream
responses_api_bridge = ResponsesToCompletionBridgeHandler()
diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py
index 753a94295b3..53ffd3647bd 100644
--- a/litellm/completion_extras/litellm_responses_transformation/transformation.py
+++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py
@@ -49,6 +49,7 @@ if TYPE_CHECKING:
ALL_RESPONSES_API_TOOL_PARAMS,
AllMessageValues,
ChatCompletionImageObject,
+ ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
OpenAIMessageContentListBlock,
)
@@ -169,7 +170,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
"type": "message",
"role": role,
"content": self._convert_content_to_responses_format(
- content, role # type: ignore
+ content, # type: ignore[arg-type]
+ role, # type: ignore
),
}
)
@@ -186,7 +188,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif isinstance(content, list):
# Transform list content to Responses API format
tool_output = self._convert_content_to_responses_format(
- content, "user" # Use "user" role to get input_* types
+ content,
+ "user", # Use "user" role to get input_* types
)
else:
# Fallback: convert unexpected types to input_text
@@ -219,14 +222,88 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
{
"type": "message",
"role": role,
- "content": self._convert_content_to_responses_format(
- content, cast(str, role)
- ),
+ "content": self._convert_content_to_responses_format(content, cast(str, role)), # type: ignore[arg-type]
}
)
return input_items, instructions
+ def _map_optional_params_to_responses_api_request(
+ self,
+ optional_params: dict,
+ responses_api_request: "ResponsesAPIOptionalRequestParams",
+ ) -> None:
+ """Map optional_params into responses_api_request (mutates in place)."""
+ for key, value in optional_params.items():
+ if value is None:
+ continue
+ if key in ("max_tokens", "max_completion_tokens"):
+ responses_api_request["max_output_tokens"] = value
+ elif key == "tools" and value is not None:
+ responses_api_request[
+ "tools"
+ ] = self._convert_tools_to_responses_format(
+ cast(List[Dict[str, Any]], value)
+ )
+ elif key == "response_format":
+ text_format = self._transform_response_format_to_text_format(value)
+ if text_format:
+ responses_api_request["text"] = text_format # type: ignore
+ elif key in ResponsesAPIOptionalRequestParams.__annotations__.keys():
+ responses_api_request[key] = value # type: ignore
+ elif key == "previous_response_id":
+ responses_api_request["previous_response_id"] = value
+ elif key == "reasoning_effort":
+ responses_api_request["reasoning"] = self._map_reasoning_effort(value)
+ elif key == "web_search_options":
+ self._add_web_search_tool(responses_api_request, value)
+
+ def _build_sanitized_litellm_params(self, litellm_params: dict) -> Dict[str, Any]:
+ """Build sanitized litellm_params with merged metadata."""
+ responses_optional_param_keys = set(
+ ResponsesAPIOptionalRequestParams.__annotations__.keys()
+ )
+ sanitized: Dict[str, Any] = {
+ key: value
+ for key, value in litellm_params.items()
+ if key not in responses_optional_param_keys
+ }
+ legacy_metadata = litellm_params.get("metadata")
+ existing_litellm_metadata = litellm_params.get("litellm_metadata")
+ merged_litellm_metadata: Dict[str, Any] = {}
+ if isinstance(legacy_metadata, dict):
+ merged_litellm_metadata.update(legacy_metadata)
+ if isinstance(existing_litellm_metadata, dict):
+ merged_litellm_metadata.update(existing_litellm_metadata)
+ if merged_litellm_metadata:
+ sanitized["litellm_metadata"] = merged_litellm_metadata
+ else:
+ sanitized.pop("litellm_metadata", None)
+ return sanitized
+
+ def _merge_responses_api_request_into_request_data(
+ self,
+ request_data: Dict[str, Any],
+ responses_api_request: "ResponsesAPIOptionalRequestParams",
+ instructions: Optional[str],
+ ) -> None:
+ """Add non-None values from responses_api_request into request_data."""
+ for key, value in responses_api_request.items():
+ if value is None:
+ continue
+ if key == "instructions" and instructions:
+ request_data["instructions"] = instructions
+ elif key == "stream_options" and isinstance(value, dict):
+ request_data["stream_options"] = value.get("include_obfuscation")
+ elif key == "user" and isinstance(value, str):
+ # OpenAI API requires user param to be max 64 chars - truncate if longer
+ if len(value) <= 64:
+ request_data["user"] = value
+ else:
+ request_data["user"] = value[:64]
+ else:
+ request_data[key] = value
+
def transform_request(
self,
model: str,
@@ -251,36 +328,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if instructions:
responses_api_request["instructions"] = instructions
- # Map optional parameters
- for key, value in optional_params.items():
- if value is None:
- continue
- if key in ("max_tokens", "max_completion_tokens"):
- responses_api_request["max_output_tokens"] = value
- elif key == "tools" and value is not None:
- # Convert chat completion tools to responses API tools format
- responses_api_request["tools"] = (
- self._convert_tools_to_responses_format(
- cast(List[Dict[str, Any]], value)
- )
- )
- elif key == "response_format":
- # Convert response_format to text.format
- text_format = self._transform_response_format_to_text_format(value)
- if text_format:
- responses_api_request["text"] = text_format # type: ignore
- elif key in ResponsesAPIOptionalRequestParams.__annotations__.keys():
- responses_api_request[key] = value # type: ignore
- elif key == "metadata":
- responses_api_request["metadata"] = value
- elif key == "previous_response_id":
- responses_api_request["previous_response_id"] = value
- elif key == "reasoning_effort":
- responses_api_request["reasoning"] = self._map_reasoning_effort(value)
- elif key == "web_search_options":
- self._add_web_search_tool(responses_api_request, value)
+ self._map_optional_params_to_responses_api_request(
+ optional_params, responses_api_request
+ )
- # Get stream parameter from litellm_params if not in optional_params
stream = optional_params.get("stream") or litellm_params.get("stream", False)
verbose_logger.debug(f"Chat provider: Stream parameter: {stream}")
@@ -304,11 +355,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
setattr(litellm_logging_obj, "call_type", CallTypes.responses.value)
+ sanitized_litellm_params = self._build_sanitized_litellm_params(litellm_params)
+
request_data = {
"model": api_model,
"input": input_items,
"litellm_logging_obj": litellm_logging_obj,
- **litellm_params,
+ **sanitized_litellm_params,
"client": client,
}
@@ -316,18 +369,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
f"Chat provider: Final request model={api_model}, input_items={len(input_items)}"
)
- # Add non-None values from responses_api_request
- for key, value in responses_api_request.items():
- if value is not None:
- if key == "instructions" and instructions:
- request_data["instructions"] = instructions
- elif key == "stream_options" and isinstance(value, dict):
- request_data["stream_options"] = value.get("include_obfuscation")
- elif key == "user": # string can't be longer than 64 characters
- if isinstance(value, str) and len(value) <= 64:
- request_data["user"] = value
- else:
- request_data[key] = value
+ self._merge_responses_api_request_into_request_data(
+ request_data, responses_api_request, instructions
+ )
if headers:
request_data["extra_headers"] = headers
@@ -355,6 +399,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
ResponseReasoningItem,
)
+ try:
+ from openai.types.responses.response_output_item import (
+ ResponseApplyPatchToolCall,
+ )
+ except ImportError:
+ ResponseApplyPatchToolCall = None # type: ignore[assignment,misc]
+
from litellm.types.utils import Choices, Message
choices: List[Choices] = []
@@ -410,6 +461,20 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
accumulated_tool_calls.append(tool_call_dict)
tool_call_index += 1
+ elif ResponseApplyPatchToolCall is not None and isinstance(
+ item, ResponseApplyPatchToolCall
+ ):
+ from litellm.responses.litellm_completion_transformation.transformation import (
+ LiteLLMCompletionResponsesConfig,
+ )
+
+ tool_call_dict = LiteLLMCompletionResponsesConfig.convert_apply_patch_tool_call_to_chat_completion_tool_call(
+ tool_call_item=item,
+ index=tool_call_index,
+ )
+ accumulated_tool_calls.append(tool_call_dict)
+ tool_call_index += 1
+
elif isinstance(item, dict) and handle_raw_dict_callback is not None:
# Handle raw dict responses (e.g., from GPT-5 Codex)
choice, index = handle_raw_dict_callback(item=item, index=index)
@@ -486,24 +551,29 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
raw_response.usage
),
)
-
+
# Preserve hidden params from the ResponsesAPIResponse, especially the headers
# which contain important provider information like x-request-id
raw_response_hidden_params = getattr(raw_response, "_hidden_params", {})
if raw_response_hidden_params:
- if not hasattr(model_response, "_hidden_params") or model_response._hidden_params is None:
+ if (
+ not hasattr(model_response, "_hidden_params")
+ or model_response._hidden_params is None
+ ):
model_response._hidden_params = {}
# Merge the raw_response hidden params with model_response hidden params
# Preserve existing keys in model_response but add/override with raw_response params
for key, value in raw_response_hidden_params.items():
if key == "additional_headers" and key in model_response._hidden_params:
# Merge additional_headers to preserve both sets
- existing_additional_headers = model_response._hidden_params.get("additional_headers", {})
+ existing_additional_headers = model_response._hidden_params.get(
+ "additional_headers", {}
+ )
merged_headers = {**value, **existing_additional_headers}
model_response._hidden_params[key] = merged_headers
else:
model_response._hidden_params[key] = value
-
+
return model_response
def get_model_response_iterator(
@@ -558,11 +628,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def _convert_content_to_responses_format(
self,
- content: Union[
- str,
- Iterable[
- Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock"]
- ],
+ content: Optional[
+ Union[
+ str,
+ List[Any],
+ Iterable[
+ Union[
+ "OpenAIMessageContentListBlock",
+ "ChatCompletionThinkingBlock",
+ "ChatCompletionRedactedThinkingBlock",
+ ]
+ ],
+ ]
],
role: str,
) -> List[Dict[str, Any]]:
@@ -573,7 +650,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
f"Chat provider: Converting content to responses format - input type: {type(content)}"
)
- if isinstance(content, str):
+ if content is None:
+ return [self._convert_content_str_to_input_text("", role)]
+ elif isinstance(content, str):
result = [self._convert_content_str_to_input_text(content, role)]
verbose_logger.debug(f"Chat provider: String content -> {result}")
return result
@@ -721,15 +800,31 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if reasoning_effort == "none":
return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") # type: ignore
elif reasoning_effort == "high":
- return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high")
+ return (
+ Reasoning(effort="high", summary="detailed")
+ if auto_summary_enabled
+ else Reasoning(effort="high")
+ )
elif reasoning_effort == "xhigh":
return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") # type: ignore[typeddict-item]
elif reasoning_effort == "medium":
- return Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium")
+ return (
+ Reasoning(effort="medium", summary="detailed")
+ if auto_summary_enabled
+ else Reasoning(effort="medium")
+ )
elif reasoning_effort == "low":
- return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low")
+ return (
+ Reasoning(effort="low", summary="detailed")
+ if auto_summary_enabled
+ else Reasoning(effort="low")
+ )
elif reasoning_effort == "minimal":
- return Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal")
+ return (
+ Reasoning(effort="minimal", summary="detailed")
+ if auto_summary_enabled
+ else Reasoning(effort="minimal")
+ )
return None
def _add_web_search_tool(
@@ -744,7 +839,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
responses_api_request: The responses API request dict to modify
web_search_options: Web search configuration (dict or other value)
"""
- if "tools" not in responses_api_request or responses_api_request["tools"] is None:
+ if (
+ "tools" not in responses_api_request
+ or responses_api_request["tools"] is None
+ ):
responses_api_request["tools"] = []
# Get the tools list with proper type narrowing
@@ -808,7 +906,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return {"format": {"type": "text"}}
return None
-
+
@staticmethod
def _convert_annotations_to_chat_format(
annotations: Optional[List[Any]],
@@ -834,13 +932,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
annotation_dict = annotation
else:
# Skip unsupported annotation types
- verbose_logger.debug(f"Skipping unsupported annotation type: {type(annotation)}")
+ verbose_logger.debug(
+ f"Skipping unsupported annotation type: {type(annotation)}"
+ )
continue
result.append(annotation_dict) # type: ignore
except Exception as e:
# Skip malformed annotations
- verbose_logger.debug(f"Skipping malformed annotation: {annotation}, error: {e}")
+ verbose_logger.debug(
+ f"Skipping malformed annotation: {annotation}, error: {e}"
+ )
continue
return result if result else None
@@ -956,13 +1058,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
)
if provider_specific_fields:
- function_chunk["provider_specific_fields"] = (
- provider_specific_fields
- )
+ function_chunk[
+ "provider_specific_fields"
+ ] = provider_specific_fields
+ tool_call_index = parsed_chunk.get("output_index", 0)
tool_call_chunk = ChatCompletionToolCallChunk(
id=output_item.get("call_id"),
- index=0,
+ index=tool_call_index,
type="function",
function=function_chunk,
)
@@ -983,6 +1086,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
elif event_type == "response.function_call_arguments.delta":
content_part: Optional[str] = parsed_chunk.get("delta", None)
if content_part:
+ tool_call_index = parsed_chunk.get("output_index", 0)
return ModelResponseStream(
choices=[
StreamingChoices(
@@ -991,7 +1095,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
tool_calls=[
ChatCompletionToolCallChunk(
id=None,
- index=0,
+ index=tool_call_index,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=None, arguments=content_part
@@ -1029,13 +1133,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
# Add provider_specific_fields to function if present
if provider_specific_fields:
- function_chunk["provider_specific_fields"] = (
- provider_specific_fields
- )
+ function_chunk[
+ "provider_specific_fields"
+ ] = provider_specific_fields
+ tool_call_index = parsed_chunk.get("output_index", 0)
tool_call_chunk = ChatCompletionToolCallChunk(
id=output_item.get("call_id"),
- index=0,
+ index=tool_call_index,
type="function",
function=function_chunk,
)
@@ -1044,12 +1149,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
if provider_specific_fields:
tool_call_chunk.provider_specific_fields = provider_specific_fields # type: ignore
+ # Do NOT emit finish_reason here — response.completed handles the terminal
+ # finish_reason. Emitting "tool_calls" here would prematurely terminate
+ # the stream before subsequent tool calls arrive (same fix as #17246 for
+ # the message-type branch).
return ModelResponseStream(
choices=[
StreamingChoices(
index=0,
- delta=Delta(tool_calls=[tool_call_chunk]),
- finish_reason="tool_calls",
+ delta=Delta(),
+ finish_reason=None,
)
]
)
@@ -1095,14 +1204,38 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
elif event_type == "response.completed":
# Response is fully complete - now we can signal is_finished=True
# This ensures we don't prematurely end the stream before tool_calls arrive
+
+ # Check if response contains function_call items in output
+ # to determine correct finish_reason
+ response_data = parsed_chunk.get("response", {})
+ output_items = response_data.get("output", []) if response_data else []
+
+ has_function_calls = any(
+ item.get("type") == "function_call"
+ for item in output_items
+ if isinstance(item, dict)
+ )
+
+ finish_reason = "tool_calls" if has_function_calls else "stop"
+
+ usage = None
+ if response_data.get("usage"):
+ from litellm.responses.utils import ResponseAPILoggingUtils
+
+ usage = (
+ ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
+ response_data.get("usage")
+ )
+ )
return ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta=Delta(content=""),
- finish_reason="stop",
+ finish_reason=finish_reason,
)
- ]
+ ],
+ usage=usage,
)
else:
pass
diff --git a/litellm/constants.py b/litellm/constants.py
index 974a08f8483..89c59ee9326 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -2,6 +2,8 @@ import os
import sys
from typing import List, Literal
+from litellm.litellm_core_utils.env_utils import get_env_int
+
DEFAULT_HEALTH_CHECK_PROMPT = str(
os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")
)
@@ -46,6 +48,33 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int(
os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)
)
DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
+
+# Maximum wall-clock seconds a streaming response is allowed to run.
+# Streams exceeding this duration are terminated with a Timeout error.
+# None (default) = no limit. Set env var to a number of seconds to enable globally.
+_max_stream_duration_env = os.getenv("LITELLM_MAX_STREAMING_DURATION_SECONDS", None)
+LITELLM_MAX_STREAMING_DURATION_SECONDS = (
+ float(_max_stream_duration_env) if _max_stream_duration_env is not None else None
+)
+
+# Maximum number of base64 characters to keep in logging payloads.
+# Data URIs exceeding this are replaced with a size placeholder.
+# Set to 0 to disable truncation.
+MAX_BASE64_LENGTH_FOR_LOGGING = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
+
+# When true, adds detailed per-phase timing breakdown headers to responses.
+# Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms
+LITELLM_DETAILED_TIMING = (
+ os.getenv("LITELLM_DETAILED_TIMING", "false").lower() == "true"
+)
+
+# Model cost map validation constants
+MODEL_COST_MAP_MIN_MODEL_COUNT = int(
+ os.getenv("MODEL_COST_MAP_MIN_MODEL_COUNT", 50)
+) # Minimum number of models a fetched cost map must contain to be considered valid
+MODEL_COST_MAP_MAX_SHRINK_RATIO = float(
+ os.getenv("MODEL_COST_MAP_MAX_SHRINK_RATIO", 0.5)
+) # Maximum allowed shrinkage ratio vs local backup (0.5 = reject if fetched map is <50% of backup)
DEFAULT_IMAGE_WIDTH = int(os.getenv("DEFAULT_IMAGE_WIDTH", 300))
DEFAULT_IMAGE_HEIGHT = int(os.getenv("DEFAULT_IMAGE_HEIGHT", 300))
# Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit)
@@ -81,6 +110,37 @@ MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH = int(
os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)
)
+# Semantic Guard Defaults
+DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL = str(
+ os.getenv("DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL", "text-embedding-3-small")
+)
+DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(
+ os.getenv("DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD", 0.75)
+)
+
+# MCP OAuth2 Client Credentials Defaults
+MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS = int(
+ os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60")
+)
+MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE = int(
+ os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200")
+)
+MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int(
+ os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600")
+)
+
+# Default npm cache directory for STDIO MCP servers.
+# npm/npx needs a writable cache dir; in containers the default (~/.npm)
+# may not exist or be read-only. /tmp is always writable.
+MCP_NPM_CACHE_DIR = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache")
+MCP_OAUTH2_TOKEN_CACHE_MIN_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10"))
+
+# MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
+MCP_CLIENT_TIMEOUT = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"))
+MCP_TOOL_LISTING_TIMEOUT = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
+MCP_METADATA_TIMEOUT = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
+MCP_HEALTH_CHECK_TIMEOUT = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
+
LITELLM_UI_ALLOW_HEADERS = [
"x-litellm-semantic-filter",
"x-litellm-semantic-filter-tools",
@@ -99,6 +159,11 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int(
)
)
+# Maximum number of callbacks that can be registered
+# This prevents callbacks from exponentially growing and consuming CPU resources
+# Override with LITELLM_MAX_CALLBACKS env var for large deployments (e.g., many teams with guardrails)
+MAX_CALLBACKS = get_env_int("LITELLM_MAX_CALLBACKS", 100)
+
# Generic fallback for unknown models
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int(
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)
@@ -132,16 +197,20 @@ _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client fo
# Aiohttp connection pooling - prevents memory leaks from unbounded connection growth
# Set to 0 for unlimited (not recommended for production)
-AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 300))
-AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 50))
+AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000))
+AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(
+ os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500)
+)
AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120))
AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
# enable_cleanup_closed is only needed for Python versions with the SSL leak bug
# Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960)
# Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78
-AIOHTTP_NEEDS_CLEANUP_CLOSED = (
- (3, 13, 0) <= sys.version_info < (3, 13, 1) or sys.version_info < (3, 12, 7)
-)
+AIOHTTP_NEEDS_CLEANUP_CLOSED = (3, 13, 0) <= sys.version_info < (
+ 3,
+ 13,
+ 1,
+) or sys.version_info < (3, 12, 7)
# WebSocket constants
# Default to None (unlimited) to match OpenAI's official agents SDK behavior
@@ -179,11 +248,20 @@ REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer"
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer"
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_team_spend_update_buffer"
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_org_spend_update_buffer"
-REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_end_user_spend_update_buffer"
+REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = (
+ "litellm_daily_end_user_spend_update_buffer"
+)
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_buffer"
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer"
MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
-MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", 2000))
+# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth
+LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000))
+TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60))
+# Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
+# Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
+MAX_SIZE_IN_MEMORY_QUEUE = int(
+ os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))
+)
MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int(
os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000)
)
@@ -247,7 +325,9 @@ MIN_NON_ZERO_TEMPERATURE = float(os.getenv("MIN_NON_ZERO_TEMPERATURE", 0.0001))
REPEATED_STREAMING_CHUNK_LIMIT = int(
os.getenv("REPEATED_STREAMING_CHUNK_LIMIT", 100)
) # catch if model starts looping the same chunk while streaming. Uses high default to prevent false positives.
-DEFAULT_MAX_LRU_CACHE_SIZE = int(os.getenv("DEFAULT_MAX_LRU_CACHE_SIZE", 16))
+# Shared maxsize for functools.lru_cache usage across hot paths.
+# Defaulted to 64 to avoid cache thrash in multi-model production workloads.
+DEFAULT_MAX_LRU_CACHE_SIZE = int(os.getenv("DEFAULT_MAX_LRU_CACHE_SIZE", 64))
_REALTIME_BODY_CACHE_SIZE = 1000 # Keep realtime helper caches bounded; workloads rarely exceed 1k models/intents
INITIAL_RETRY_DELAY = float(os.getenv("INITIAL_RETRY_DELAY", 0.5))
MAX_RETRY_DELAY = float(os.getenv("MAX_RETRY_DELAY", 8.0))
@@ -279,6 +359,9 @@ NON_LLM_CONNECTION_TIMEOUT = int(
MAX_EXCEPTION_MESSAGE_LENGTH = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000))
MAX_STRING_LENGTH_PROMPT_IN_DB = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048))
BEDROCK_MAX_POLICY_SIZE = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75))
+BEDROCK_MIN_THINKING_BUDGET_TOKENS = int(
+ os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024)
+)
REPLICATE_POLLING_DELAY_SECONDS = float(
os.getenv("REPLICATE_POLLING_DELAY_SECONDS", 0.5)
)
@@ -305,7 +388,25 @@ MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int(
DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000))
#### Networking settings ####
request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000)) # time in seconds
-DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes
+DEFAULT_A2A_AGENT_TIMEOUT: float = float(
+ os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)
+) # 10 minutes
+# Patterns that indicate a localhost/internal URL in A2A agent cards that should be
+# replaced with the original base_url. This is a common misconfiguration where
+# developers deploy agents with development URLs in their agent cards.
+LOCALHOST_URL_PATTERNS: List[str] = [
+ "localhost",
+ "127.0.0.1",
+ "0.0.0.0",
+ "[::1]", # IPv6 localhost
+]
+# Patterns in error messages that indicate a connection failure
+CONNECTION_ERROR_PATTERNS: List[str] = [
+ "connect",
+ "connection",
+ "network",
+ "refused",
+]
STREAM_SSE_DONE_STRING: str = "[DONE]"
STREAM_SSE_DATA_PREFIX: str = "data: "
### SPEND TRACKING ###
@@ -341,8 +442,12 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv(
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
)
-EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds
-EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)) # 80% of max budget
+EMAIL_BUDGET_ALERT_TTL = int(
+ os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)
+) # 24 hours in seconds
+EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(
+ os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)
+) # 80% of max budget
############### LLM Provider Constants ###############
### ANTHROPIC CONSTANTS ###
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv(
@@ -400,6 +505,7 @@ LITELLM_CHAT_PROVIDERS = [
"azure_ai",
"sagemaker",
"sagemaker_chat",
+ "sagemaker_nova",
"bedrock",
"vllm",
"nlp_cloud",
@@ -511,6 +617,11 @@ OPENAI_CHAT_COMPLETION_PARAMS = [
"thinking",
"web_search_options",
"service_tier",
+ "prompt_cache_key",
+ "prompt_cache_retention",
+ "safety_identifier",
+ "verbosity",
+ "store",
]
OPENAI_TRANSCRIPTION_PARAMS = [
@@ -572,6 +683,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES = {
"prompt_cache_retention": None,
"store": None,
"metadata": None,
+ "context_management": None,
}
openai_compatible_endpoints: List = [
@@ -962,15 +1074,19 @@ BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[
BEDROCK_CONVERSE_MODELS = [
"qwen.qwen3-coder-480b-a35b-v1:0",
+ "qwen.qwen3-coder-next",
"qwen.qwen3-235b-a22b-2507-v1:0",
"qwen.qwen3-coder-30b-a3b-v1:0",
"qwen.qwen3-32b-v1:0",
"deepseek.v3-v1:0",
+ "deepseek.v3.2",
"openai.gpt-oss-20b-1:0",
"openai.gpt-oss-120b-1:0",
"anthropic.claude-haiku-4-5-20251001-v1:0",
"anthropic.claude-sonnet-4-5-20250929-v1:0",
"anthropic.claude-opus-4-6-v1:0",
+ "anthropic.claude-opus-4-6-v1",
+ "anthropic.claude-sonnet-4-6",
"anthropic.claude-opus-4-1-20250805-v1:0",
"anthropic.claude-opus-4-20250514-v1:0",
"anthropic.claude-sonnet-4-20250514-v1:0",
@@ -1007,6 +1123,8 @@ BEDROCK_CONVERSE_MODELS = [
"amazon.nova-pro-v1:0",
"writer.palmyra-x4-v1:0",
"writer.palmyra-x5-v1:0",
+ "minimax.minimax-m2.1",
+ "moonshotai.kimi-k2.5",
]
@@ -1091,7 +1209,13 @@ known_tokenizer_config = {
}
-OPENAI_FINISH_REASONS = ["stop", "length", "function_call", "content_filter", "null", "finish_reason_unspecified", "malformed_function_call", "guardrail_intervened", "eos"]
+OPENAI_FINISH_REASONS = [
+ "stop",
+ "length",
+ "function_call",
+ "tool_calls",
+ "content_filter",
+]
HUMANLOOP_PROMPT_CACHE_TTL_SECONDS = int(
os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60)
) # 1 minute
@@ -1113,6 +1237,11 @@ X_LITELLM_DISABLE_CALLBACKS = "x-litellm-disable-callbacks"
LITELLM_METADATA_FIELD = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD = "metadata"
LITELLM_TRUNCATED_PAYLOAD_FIELD = "litellm_truncated"
+LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE = (
+ "Truncation is a DB storage safeguard. "
+ "Full, untruncated data is logged to logging callbacks (OTEL, Datadog, etc.). "
+ "To increase the truncation limit, set `MAX_STRING_LENGTH_PROMPT_IN_DB` in your env."
+)
########################### LiteLLM Proxy Specific Constants ###########################
########################################################################################
@@ -1181,6 +1310,9 @@ LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false"
LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int(
os.getenv("LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS", 86400)
) # 24 hours default
+LITELLM_KEY_ROTATION_GRACE_PERIOD: str = os.getenv(
+ "LITELLM_KEY_ROTATION_GRACE_PERIOD", ""
+) # Duration to keep old key valid after rotation (e.g. "24h", "2d"); empty = immediate revoke (default)
UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard"
LITELLM_PROXY_ADMIN_NAME = "default_user_id"
@@ -1191,11 +1323,16 @@ CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session"
CLI_JWT_TOKEN_NAME = "cli-jwt-token"
# Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility
CLI_JWT_EXPIRATION_HOURS = int(
- os.getenv("CLI_JWT_EXPIRATION_HOURS")
- or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS")
+ os.getenv("CLI_JWT_EXPIRATION_HOURS")
+ or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS")
or 24
)
+########################### UI SESSION DURATION ###########################
+# Duration for UI login session (username/password, SSO, invitation links). Format: "30s", "30m", "24h", "7d"
+# Does NOT apply to EXPERIMENTAL_UI_LOGIN flow, which intentionally uses a fixed 10-minute expiry for security.
+LITELLM_UI_SESSION_DURATION = os.getenv("LITELLM_UI_SESSION_DURATION", "24h")
+
########################### DB CRON JOB NAMES ###########################
DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job"
PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics"
@@ -1215,6 +1352,15 @@ PROXY_BUDGET_RESCHEDULER_MIN_TIME = int(
os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
)
PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
+MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
+MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max(
+ 1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
+)
+# Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and
+# CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on
+# installations with large numbers of stale managed objects).
+_batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower()
+PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true"
PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(
os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
)
@@ -1260,6 +1406,9 @@ DEFAULT_SLACK_ALERTING_THRESHOLD = int(
os.getenv("DEFAULT_SLACK_ALERTING_THRESHOLD", 300)
)
MAX_TEAM_LIST_LIMIT = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20))
+MAX_POLICY_ESTIMATE_IMPACT_ROWS = int(
+ os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000)
+)
DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(
os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7)
)
@@ -1269,6 +1418,7 @@ SECRET_MANAGER_REFRESH_INTERVAL = int(
)
LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [
"default_internal_user_params",
+ "default_team_params",
"public_mcp_servers",
"public_agent_groups",
"public_model_groups",
@@ -1280,6 +1430,7 @@ SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(
os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)
)
+DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600))
# Sentry Scrubbing Configuration
SENTRY_DENYLIST = [
@@ -1370,12 +1521,21 @@ MICROSOFT_USER_EMAIL_ATTRIBUTE = str(
MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE = str(
os.getenv("MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "displayName")
)
-MICROSOFT_USER_ID_ATTRIBUTE = str(
- os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id")
-)
+MICROSOFT_USER_ID_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id"))
MICROSOFT_USER_FIRST_NAME_ATTRIBUTE = str(
os.getenv("MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "givenName")
)
MICROSOFT_USER_LAST_NAME_ATTRIBUTE = str(
os.getenv("MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "surname")
)
+
+# Maximum payload size (in bytes) to fully serialize for DEBUG logging.
+# Payloads larger than this are truncated to avoid multi-second json.dumps blocking the response.
+MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG = int(
+ os.getenv("MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG", 102400)
+) # 100 KB
+
+# 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/containers/__init__.py b/litellm/containers/__init__.py
index e279cb429e5..48ab5de4181 100644
--- a/litellm/containers/__init__.py
+++ b/litellm/containers/__init__.py
@@ -42,4 +42,3 @@ __all__ = [
"retrieve_container_file",
"retrieve_container_file_content",
]
-
diff --git a/litellm/containers/endpoint_factory.py b/litellm/containers/endpoint_factory.py
index 0b73a19b922..1d8e50856fe 100644
--- a/litellm/containers/endpoint_factory.py
+++ b/litellm/containers/endpoint_factory.py
@@ -43,13 +43,13 @@ def _load_endpoints_config() -> Dict:
def create_sync_endpoint_function(endpoint_config: Dict) -> Callable:
"""
Create a sync SDK function from endpoint config.
-
+
Uses the generic container handler instead of individual handler methods.
"""
endpoint_name = endpoint_config["name"]
response_type = RESPONSE_TYPES.get(endpoint_config["response_type"])
path_params = endpoint_config.get("path_params", [])
-
+
@client
def endpoint_func(
timeout: int = 600,
@@ -76,20 +76,23 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable:
# Get provider config
litellm_params = GenericLiteLLMParams(**kwargs)
- container_provider_config: Optional[BaseContainerConfig] = (
- ProviderConfigManager.get_provider_container_config(
- provider=litellm.LlmProviders(custom_llm_provider),
- )
+ container_provider_config: Optional[
+ BaseContainerConfig
+ ] = ProviderConfigManager.get_provider_container_config(
+ provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
- raise ValueError(f"Container provider config not found for: {custom_llm_provider}")
+ raise ValueError(
+ f"Container provider config not found for: {custom_llm_provider}"
+ )
# Build optional params for logging
optional_params = {k: kwargs.get(k) for k in path_params if k in kwargs}
# Pre-call logging
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model="",
optional_params=optional_params,
litellm_params={"litellm_call_id": litellm_call_id},
@@ -126,7 +129,7 @@ def create_async_endpoint_function(
endpoint_config: Dict,
) -> Callable:
"""Create an async SDK function that wraps the sync function."""
-
+
@client
async def async_endpoint_func(
timeout: int = 600,
@@ -176,21 +179,21 @@ def create_async_endpoint_function(
def generate_container_endpoints() -> Dict[str, Callable]:
"""
Generate all container endpoint functions from the JSON config.
-
+
Returns a dict mapping function names to their implementations.
"""
config = _load_endpoints_config()
endpoints = {}
-
+
for endpoint_config in config["endpoints"]:
# Create sync function
sync_func = create_sync_endpoint_function(endpoint_config)
endpoints[endpoint_config["name"]] = sync_func
-
+
# Create async function
async_func = create_async_endpoint_function(sync_func, endpoint_config)
endpoints[endpoint_config["async_name"]] = async_func
-
+
return endpoints
@@ -222,5 +225,9 @@ retrieve_container_file = _generated_endpoints.get("retrieve_container_file")
aretrieve_container_file = _generated_endpoints.get("aretrieve_container_file")
delete_container_file = _generated_endpoints.get("delete_container_file")
adelete_container_file = _generated_endpoints.get("adelete_container_file")
-retrieve_container_file_content = _generated_endpoints.get("retrieve_container_file_content")
-aretrieve_container_file_content = _generated_endpoints.get("aretrieve_container_file_content")
+retrieve_container_file_content = _generated_endpoints.get(
+ "retrieve_container_file_content"
+)
+aretrieve_container_file_content = _generated_endpoints.get(
+ "aretrieve_container_file_content"
+)
diff --git a/litellm/containers/main.py b/litellm/containers/main.py
index 105e999ffe8..916fc26351b 100644
--- a/litellm/containers/main.py
+++ b/litellm/containers/main.py
@@ -39,6 +39,7 @@ __all__ = [
"upload_container_file",
]
+
##### Container Create #######################
@client
async def acreate_container(
@@ -164,10 +165,7 @@ def create_container(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
-) -> Union[
- ContainerObject,
- Coroutine[Any, Any, ContainerObject],
-]:
+) -> Union[ContainerObject, Coroutine[Any, Any, ContainerObject],]:
"""Create a container using the OpenAI Container API.
Currently supports OpenAI
@@ -175,7 +173,7 @@ def create_container(
Example:
```python
import litellm
-
+
response = litellm.create_container(
name="My Container",
custom_llm_provider="openai",
@@ -207,19 +205,23 @@ def create_container(
**kwargs,
)
# get provider config
- container_provider_config: Optional[BaseContainerConfig] = (
- ProviderConfigManager.get_provider_container_config(
- provider=litellm.LlmProviders(custom_llm_provider),
- )
+ container_provider_config: Optional[
+ BaseContainerConfig
+ ] = ProviderConfigManager.get_provider_container_config(
+ provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
- raise ValueError(f"container operations are not supported for {custom_llm_provider}")
+ raise ValueError(
+ f"container operations are not supported for {custom_llm_provider}"
+ )
local_vars.update(kwargs)
# Get ContainerCreateOptionalRequestParams with only valid parameters
container_create_optional_params: ContainerCreateOptionalRequestParams = (
- ContainerRequestUtils.get_requested_container_create_optional_param(local_vars)
+ ContainerRequestUtils.get_requested_container_create_optional_param(
+ local_vars
+ )
)
# Get optional parameters for the container API
@@ -231,7 +233,8 @@ def create_container(
)
# Pre Call logging
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model="",
optional_params=dict(container_create_request_params),
litellm_params={
@@ -388,10 +391,7 @@ def list_containers(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
-) -> Union[
- ContainerListResponse,
- Coroutine[Any, Any, ContainerListResponse],
-]:
+) -> Union[ContainerListResponse, Coroutine[Any, Any, ContainerListResponse],]:
"""List containers using the OpenAI Container API.
Currently supports OpenAI
@@ -420,22 +420,27 @@ def list_containers(
**kwargs,
)
# get provider config
- container_provider_config: Optional[BaseContainerConfig] = (
- ProviderConfigManager.get_provider_container_config(
- provider=litellm.LlmProviders(custom_llm_provider),
- )
+ container_provider_config: Optional[
+ BaseContainerConfig
+ ] = ProviderConfigManager.get_provider_container_config(
+ provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
- raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
+ raise ValueError(
+ f"Container provider config not found for provider: {custom_llm_provider}"
+ )
# Get container list request parameters
container_list_optional_params: ContainerListOptionalRequestParams = (
- ContainerRequestUtils.get_requested_container_list_optional_param(local_vars)
+ ContainerRequestUtils.get_requested_container_list_optional_param(
+ local_vars
+ )
)
# Pre Call logging
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model="",
optional_params=dict(container_list_optional_params),
litellm_params={
@@ -582,10 +587,7 @@ def retrieve_container(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
-) -> Union[
- ContainerObject,
- Coroutine[Any, Any, ContainerObject],
-]:
+) -> Union[ContainerObject, Coroutine[Any, Any, ContainerObject],]:
"""Retrieve a container using the OpenAI Container API.
Currently supports OpenAI
@@ -614,17 +616,20 @@ def retrieve_container(
**kwargs,
)
# get provider config
- container_provider_config: Optional[BaseContainerConfig] = (
- ProviderConfigManager.get_provider_container_config(
- provider=litellm.LlmProviders(custom_llm_provider),
- )
+ container_provider_config: Optional[
+ BaseContainerConfig
+ ] = ProviderConfigManager.get_provider_container_config(
+ provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
- raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
+ raise ValueError(
+ f"Container provider config not found for provider: {custom_llm_provider}"
+ )
# Pre Call logging
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model="",
optional_params={},
litellm_params={
@@ -768,10 +773,7 @@ def delete_container(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
-) -> Union[
- DeleteContainerResult,
- Coroutine[Any, Any, DeleteContainerResult],
-]:
+) -> Union[DeleteContainerResult, Coroutine[Any, Any, DeleteContainerResult],]:
"""Delete a container using the OpenAI Container API.
Currently supports OpenAI
@@ -800,17 +802,20 @@ def delete_container(
**kwargs,
)
# get provider config
- container_provider_config: Optional[BaseContainerConfig] = (
- ProviderConfigManager.get_provider_container_config(
- provider=litellm.LlmProviders(custom_llm_provider),
- )
+ container_provider_config: Optional[
+ BaseContainerConfig
+ ] = ProviderConfigManager.get_provider_container_config(
+ provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
- raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
+ raise ValueError(
+ f"Container provider config not found for provider: {custom_llm_provider}"
+ )
# Pre Call logging
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model="",
optional_params={},
litellm_params={
@@ -968,10 +973,7 @@ def list_container_files(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
-) -> Union[
- ContainerFileListResponse,
- Coroutine[Any, Any, ContainerFileListResponse],
-]:
+) -> Union[ContainerFileListResponse, Coroutine[Any, Any, ContainerFileListResponse],]:
"""List files in a container using the OpenAI Container API.
Currently supports OpenAI
@@ -1000,19 +1002,27 @@ def list_container_files(
**kwargs,
)
# get provider config
- container_provider_config: Optional[BaseContainerConfig] = (
- ProviderConfigManager.get_provider_container_config(
- provider=litellm.LlmProviders(custom_llm_provider),
- )
+ container_provider_config: Optional[
+ BaseContainerConfig
+ ] = ProviderConfigManager.get_provider_container_config(
+ provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
- raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
+ raise ValueError(
+ f"Container provider config not found for provider: {custom_llm_provider}"
+ )
# Pre Call logging
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model="",
- optional_params={"container_id": container_id, "after": after, "limit": limit, "order": order},
+ optional_params={
+ "container_id": container_id,
+ "after": after,
+ "limit": limit,
+ "order": order,
+ },
litellm_params={
"litellm_call_id": litellm_call_id,
},
@@ -1180,10 +1190,7 @@ def upload_container_file(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
-) -> Union[
- ContainerFileObject,
- Coroutine[Any, Any, ContainerFileObject],
-]:
+) -> Union[ContainerFileObject, Coroutine[Any, Any, ContainerFileObject],]:
"""Upload a file to a container using the OpenAI Container API.
This endpoint allows uploading files directly to a container session,
@@ -1241,17 +1248,20 @@ def upload_container_file(
**kwargs,
)
# get provider config
- container_provider_config: Optional[BaseContainerConfig] = (
- ProviderConfigManager.get_provider_container_config(
- provider=litellm.LlmProviders(custom_llm_provider),
- )
+ container_provider_config: Optional[
+ BaseContainerConfig
+ ] = ProviderConfigManager.get_provider_container_config(
+ provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
- raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
+ raise ValueError(
+ f"Container provider config not found for provider: {custom_llm_provider}"
+ )
# Pre Call logging
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model="",
optional_params={"container_id": container_id},
litellm_params={
diff --git a/litellm/containers/utils.py b/litellm/containers/utils.py
index f30f1e154be..048f587fda7 100644
--- a/litellm/containers/utils.py
+++ b/litellm/containers/utils.py
@@ -1,7 +1,10 @@
from typing import Dict
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
-from litellm.types.containers.main import ContainerCreateOptionalRequestParams, ContainerListOptionalRequestParams
+from litellm.types.containers.main import (
+ ContainerCreateOptionalRequestParams,
+ ContainerListOptionalRequestParams,
+)
class ContainerRequestUtils:
diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index bef4d52ce49..e0e1e35b94e 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -1,5 +1,6 @@
# What is this?
## File for 'response_cost' calculation in Logging
+import logging
import time
from functools import lru_cache
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union, cast
@@ -73,6 +74,7 @@ from litellm.llms.vertex_ai.cost_calculator import (
from litellm.llms.vertex_ai.cost_calculator import cost_router as google_cost_router
from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_token
from litellm.responses.utils import ResponseAPILoggingUtils
+from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.llms.openai import (
HttpxBinaryResponseContent,
ImageGenerationRequestQuality,
@@ -117,6 +119,54 @@ if TYPE_CHECKING:
else:
LitellmLoggingObject = Any
+# Pre-resolved CallTypes enum values for fast membership checks
+_A2A_CALL_TYPES = frozenset(
+ {
+ CallTypes.asend_message.value,
+ CallTypes.send_message.value,
+ }
+)
+
+_VIDEO_CALL_TYPES = frozenset(
+ {
+ CallTypes.create_video.value,
+ CallTypes.acreate_video.value,
+ CallTypes.video_remix.value,
+ CallTypes.avideo_remix.value,
+ }
+)
+
+_SPEECH_CALL_TYPES = frozenset(
+ {
+ CallTypes.speech.value,
+ CallTypes.aspeech.value,
+ }
+)
+
+_TRANSCRIPTION_CALL_TYPES = frozenset(
+ {
+ CallTypes.atranscription.value,
+ CallTypes.transcription.value,
+ }
+)
+
+_RERANK_CALL_TYPES = frozenset(
+ {
+ CallTypes.rerank.value,
+ CallTypes.arerank.value,
+ }
+)
+
+_SEARCH_CALL_TYPES = frozenset(
+ {
+ CallTypes.search.value,
+ CallTypes.asearch.value,
+ }
+)
+
+_AREALTIME_CALL_TYPE = CallTypes.arealtime.value
+_MCP_CALL_TYPE = CallTypes.call_mcp_tool.value
+
def _cost_per_token_custom_pricing_helper(
prompt_tokens: float = 0,
@@ -149,32 +199,33 @@ def _get_additional_costs(
) -> Optional[dict]:
"""
Calculate additional costs beyond standard token costs.
-
+
This function delegates to provider-specific config classes to calculate
any additional costs like routing fees, infrastructure costs, etc.
-
+
Args:
model: The model name
custom_llm_provider: The provider name (optional)
prompt_tokens: Number of prompt tokens
completion_tokens: Number of completion tokens
-
+
Returns:
Optional dictionary with cost names and amounts, or None if no additional costs
"""
if not custom_llm_provider:
return None
-
+
try:
config_class = None
if custom_llm_provider == "azure_ai":
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
config_class = AzureFoundryModelInfo.get_azure_ai_config_for_model(model)
# Add more providers here as needed
# elif custom_llm_provider == "other_provider":
# config_class = get_other_provider_config(model)
-
- if config_class and hasattr(config_class, 'calculate_additional_costs'):
+
+ if config_class and hasattr(config_class, "calculate_additional_costs"):
return config_class.calculate_additional_costs(
model=model,
prompt_tokens=prompt_tokens,
@@ -182,7 +233,7 @@ def _get_additional_costs(
)
except Exception as e:
verbose_logger.debug(f"Error calculating additional costs: {e}")
-
+
return None
@@ -233,6 +284,8 @@ def cost_per_token( # noqa: PLR0915
### SERVICE TIER ###
service_tier: Optional[str] = None, # for OpenAI service tier pricing
response: Optional[Any] = None,
+ ### REQUEST MODEL ###
+ request_model: Optional[str] = None, # original request model for router detection
) -> Tuple[float, float]: # type: ignore
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -441,11 +494,14 @@ def cost_per_token( # noqa: PLR0915
model=model_without_prefix,
custom_llm_provider=custom_llm_provider,
usage=usage_block,
+ service_tier=service_tier,
)
elif custom_llm_provider == "anthropic":
return anthropic_cost_per_token(model=model, usage=usage_block)
elif custom_llm_provider == "bedrock":
- return bedrock_cost_per_token(model=model, usage=usage_block)
+ return bedrock_cost_per_token(
+ model=model, usage=usage_block, service_tier=service_tier
+ )
elif custom_llm_provider == "openai":
return openai_cost_per_token(
model=model, usage=usage_block, service_tier=service_tier
@@ -459,7 +515,9 @@ def cost_per_token( # noqa: PLR0915
model=model, usage=usage_block, response_time_ms=response_time_ms
)
elif custom_llm_provider == "gemini":
- return gemini_cost_per_token(model=model, usage=usage_block)
+ return gemini_cost_per_token(
+ model=model, usage=usage_block, service_tier=service_tier
+ )
elif custom_llm_provider == "deepseek":
return deepseek_cost_per_token(model=model, usage=usage_block)
elif custom_llm_provider == "perplexity":
@@ -476,7 +534,10 @@ def cost_per_token( # noqa: PLR0915
return dashscope_cost_per_token(model=model, usage=usage_block)
elif custom_llm_provider == "azure_ai":
return azure_ai_cost_per_token(
- model=model, usage=usage_block, response_time_ms=response_time_ms
+ model=model,
+ usage=usage_block,
+ response_time_ms=response_time_ms,
+ request_model=request_model,
)
else:
model_info = _cached_get_model_info_helper(
@@ -599,7 +660,14 @@ def _select_model_name_for_cost_calc(
if custom_pricing is True:
if router_model_id is not None and router_model_id in litellm.model_cost:
- return_model = router_model_id
+ entry = litellm.model_cost[router_model_id]
+ if (
+ entry.get("input_cost_per_token") is not None
+ or entry.get("input_cost_per_second") is not None
+ ):
+ return_model = router_model_id
+ else:
+ return_model = model
else:
return_model = model
@@ -663,6 +731,36 @@ def _get_response_model(completion_response: Any) -> Optional[str]:
return None
+_GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: dict = {
+ # ON_DEMAND_PRIORITY maps to "priority" — selects input_cost_per_token_priority, etc.
+ "ON_DEMAND_PRIORITY": "priority",
+ # FLEX / BATCH maps to "flex" — selects input_cost_per_token_flex, etc.
+ "FLEX": "flex",
+ "BATCH": "flex",
+ # ON_DEMAND is standard pricing — no service_tier suffix applied
+ "ON_DEMAND": None,
+}
+
+
+def _map_traffic_type_to_service_tier(traffic_type: Optional[str]) -> Optional[str]:
+ """
+ Map a Gemini usageMetadata.trafficType value to a LiteLLM service_tier string.
+
+ This allows the same `_priority` / `_flex` cost-key suffix logic used for
+ OpenAI/Azure to work for Gemini and Vertex AI models.
+
+ trafficType values seen in practice
+ ------------------------------------
+ ON_DEMAND -> standard pricing (service_tier = None)
+ ON_DEMAND_PRIORITY -> priority pricing (service_tier = "priority")
+ FLEX / BATCH -> batch/flex pricing (service_tier = "flex")
+ """
+ if traffic_type is None:
+ return None
+ service_tier = _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER.get(traffic_type.upper())
+ return service_tier
+
+
def _get_usage_object(
completion_response: Any,
) -> Optional[Usage]:
@@ -747,6 +845,8 @@ def _infer_call_type(
return "image_generation"
elif isinstance(completion_response, TextCompletionResponse):
return "text_completion"
+ elif isinstance(completion_response, LiteLLMSendMessageResponse):
+ return "send_message"
return call_type
@@ -774,10 +874,11 @@ def _apply_cost_discount(
discount_amount = original_cost * discount_percent
final_cost = original_cost - discount_amount
- verbose_logger.debug(
- f"Applied {discount_percent*100}% discount to {custom_llm_provider}: "
- f"${original_cost:.6f} -> ${final_cost:.6f} (saved ${discount_amount:.6f})"
- )
+ if verbose_logger.isEnabledFor(logging.DEBUG):
+ verbose_logger.debug(
+ f"Applied {discount_percent*100}% discount to {custom_llm_provider}: "
+ f"${original_cost:.6f} -> ${final_cost:.6f} (saved ${discount_amount:.6f})"
+ )
return final_cost, discount_percent, discount_amount
@@ -807,17 +908,20 @@ def _apply_cost_margin(
margin_config = None
if custom_llm_provider and custom_llm_provider in litellm.cost_margin_config:
margin_config = litellm.cost_margin_config[custom_llm_provider]
- verbose_logger.debug(
- f"Found provider-specific margin config for {custom_llm_provider}: {margin_config}"
- )
+ if verbose_logger.isEnabledFor(logging.DEBUG):
+ verbose_logger.debug(
+ f"Found provider-specific margin config for {custom_llm_provider}: {margin_config}"
+ )
elif "global" in litellm.cost_margin_config:
margin_config = litellm.cost_margin_config["global"]
- verbose_logger.debug(f"Using global margin config: {margin_config}")
+ if verbose_logger.isEnabledFor(logging.DEBUG):
+ verbose_logger.debug(f"Using global margin config: {margin_config}")
else:
- verbose_logger.debug(
- f"No margin config found. Provider: {custom_llm_provider}, "
- f"Available configs: {list(litellm.cost_margin_config.keys())}"
- )
+ if verbose_logger.isEnabledFor(logging.DEBUG):
+ verbose_logger.debug(
+ f"No margin config found. Provider: {custom_llm_provider}, "
+ f"Available configs: {list(litellm.cost_margin_config.keys())}"
+ )
if margin_config is not None:
# Handle different margin config formats
@@ -836,11 +940,12 @@ def _apply_cost_margin(
final_cost = original_cost + margin_total_amount
- verbose_logger.debug(
- f"Applied margin to {custom_llm_provider or 'global'}: "
- f"${original_cost:.6f} -> ${final_cost:.6f} "
- f"(margin: {margin_percent*100 if margin_percent > 0 else 0}% + ${margin_fixed_amount:.6f} = ${margin_total_amount:.6f})"
- )
+ if verbose_logger.isEnabledFor(logging.DEBUG):
+ verbose_logger.debug(
+ f"Applied margin to {custom_llm_provider or 'global'}: "
+ f"${original_cost:.6f} -> ${final_cost:.6f} "
+ f"(margin: {margin_percent*100 if margin_percent > 0 else 0}% + ${margin_fixed_amount:.6f} = ${margin_total_amount:.6f})"
+ )
return final_cost, margin_percent, margin_fixed_amount, margin_total_amount
@@ -1021,9 +1126,10 @@ def completion_cost( # noqa: PLR0915
for idx, model in enumerate(potential_model_names):
try:
- verbose_logger.debug(
- f"selected model name for cost calculation: {model}"
- )
+ if verbose_logger.isEnabledFor(logging.DEBUG):
+ verbose_logger.debug(
+ f"selected model name for cost calculation: {model}"
+ )
if completion_response is not None and (
isinstance(completion_response, BaseModel)
@@ -1083,7 +1189,7 @@ def completion_cost( # noqa: PLR0915
and _usage["prompt_tokens_details"] != {}
and _usage["prompt_tokens_details"]
):
- prompt_tokens_details = _usage.get("prompt_tokens_details", {})
+ prompt_tokens_details = _usage.get("prompt_tokens_details") or {}
cache_read_input_tokens = prompt_tokens_details.get(
"cached_tokens", 0
)
@@ -1096,6 +1202,20 @@ def completion_cost( # noqa: PLR0915
"custom_llm_provider", custom_llm_provider or None
)
region_name = hidden_params.get("region_name", region_name)
+
+ # For Gemini/Vertex AI responses, trafficType is stored in
+ # provider_specific_fields. Map it to the service_tier used
+ # by the cost key lookup (_priority / _flex suffixes) so that
+ # ON_DEMAND_PRIORITY requests are billed at priority prices.
+ if service_tier is None:
+ provider_specific = (
+ hidden_params.get("provider_specific_fields") or {}
+ )
+ raw_traffic_type = provider_specific.get("traffic_type")
+ if raw_traffic_type:
+ service_tier = _map_traffic_type_to_service_tier(
+ raw_traffic_type
+ )
else:
if model is None:
raise ValueError(
@@ -1108,10 +1228,7 @@ def completion_cost( # noqa: PLR0915
completion_tokens = token_counter(model=model, text=completion)
# Handle A2A calls before model check - A2A doesn't require a model
- if call_type in (
- CallTypes.asend_message.value,
- CallTypes.send_message.value,
- ):
+ if call_type in _A2A_CALL_TYPES:
from litellm.a2a_protocol.cost_calculator import A2ACostCalculator
return A2ACostCalculator.calculate_a2a_cost(
@@ -1147,13 +1264,18 @@ def completion_cost( # noqa: PLR0915
optional_params=optional_params,
call_type=call_type,
)
- elif (
- call_type == CallTypes.create_video.value
- or call_type == CallTypes.acreate_video.value
- or call_type == CallTypes.video_remix.value
- or call_type == CallTypes.avideo_remix.value
- ):
+ elif call_type in _VIDEO_CALL_TYPES:
### VIDEO GENERATION COST CALCULATION ###
+ # Extract custom model_info for deployment-specific pricing
+ _video_model_info: Optional[ModelInfo] = None
+ if custom_pricing and litellm_logging_obj is not None:
+ _litellm_params = getattr(
+ litellm_logging_obj, "litellm_params", None
+ )
+ if _litellm_params is not None:
+ _metadata = _litellm_params.get("metadata", {}) or {}
+ _video_model_info = _metadata.get("model_info", None)
+
usage_obj = getattr(completion_response, "usage", None)
if completion_response is not None and usage_obj:
# Handle both dict and Pydantic Usage object
@@ -1174,29 +1296,28 @@ def completion_cost( # noqa: PLR0915
model=model,
duration_seconds=duration_seconds,
custom_llm_provider=custom_llm_provider,
+ model_info=_video_model_info,
)
# Fallback to default video cost calculation if no duration available
return default_video_cost_calculator(
model=model,
duration_seconds=0.0, # Default to 0 if no duration available
custom_llm_provider=custom_llm_provider,
+ model_info=_video_model_info,
)
- elif (
- call_type == CallTypes.speech.value
- or call_type == CallTypes.aspeech.value
- ):
+ elif call_type in _SPEECH_CALL_TYPES:
prompt_characters = litellm.utils._count_characters(text=prompt)
- elif (
- call_type == CallTypes.atranscription.value
- or call_type == CallTypes.transcription.value
- ):
- audio_transcription_file_duration = getattr(
- completion_response, "duration", 0.0
+ elif call_type in _TRANSCRIPTION_CALL_TYPES:
+ # Check _hidden_params first (duration stored there to
+ # avoid polluting the response body), then fall back to
+ # the response attribute (for verbose_json responses that
+ # naturally include duration from the provider).
+ _hidden = getattr(completion_response, "_hidden_params", {}) or {}
+ audio_transcription_file_duration = _hidden.get(
+ "audio_transcription_duration",
+ getattr(completion_response, "duration", 0.0),
)
- elif (
- call_type == CallTypes.rerank.value
- or call_type == CallTypes.arerank.value
- ):
+ elif call_type in _RERANK_CALL_TYPES:
if completion_response is not None and isinstance(
completion_response, RerankResponse
):
@@ -1215,10 +1336,7 @@ def completion_cost( # noqa: PLR0915
billed_units.get("search_units") or 1
) # cohere charges per request by default.
completion_tokens = search_units
- elif (
- call_type == CallTypes.search.value
- or call_type == CallTypes.asearch.value
- ):
+ elif call_type in _SEARCH_CALL_TYPES:
from litellm.search import search_provider_cost_per_query
# Extract number_of_queries from optional_params or default to 1
@@ -1287,7 +1405,7 @@ def completion_cost( # noqa: PLR0915
)
return _final_cost
- elif call_type == CallTypes.arealtime.value and isinstance(
+ elif call_type == _AREALTIME_CALL_TYPE and isinstance(
completion_response, LiteLLMRealtimeStreamLoggingObject
):
if (
@@ -1306,7 +1424,7 @@ def completion_cost( # noqa: PLR0915
custom_llm_provider=custom_llm_provider,
litellm_model_name=model,
)
- elif call_type == CallTypes.call_mcp_tool.value:
+ elif call_type == _MCP_CALL_TYPE:
from litellm.proxy._experimental.mcp_server.cost_calculator import (
MCPCostCalculator,
)
@@ -1363,13 +1481,18 @@ def completion_cost( # noqa: PLR0915
text=completion_string
)
+ # Get the original request model for router detection
+ request_model_for_cost = None
+ if litellm_logging_obj is not None:
+ request_model_for_cost = litellm_logging_obj.model
+
(
prompt_tokens_cost_usd_dollar,
completion_tokens_cost_usd_dollar,
) = cost_per_token(
model=model,
- prompt_tokens=prompt_tokens,
- completion_tokens=completion_tokens,
+ prompt_tokens=prompt_tokens or 0,
+ completion_tokens=completion_tokens or 0,
custom_llm_provider=custom_llm_provider,
response_time_ms=total_time,
region_name=region_name,
@@ -1380,21 +1503,40 @@ def completion_cost( # noqa: PLR0915
cache_creation_input_tokens=cache_creation_input_tokens,
cache_read_input_tokens=cache_read_input_tokens,
usage_object=cost_per_token_usage_object,
- call_type=cast(CallTypesLiteral, call_type),
+ call_type=call_type,
audio_transcription_file_duration=audio_transcription_file_duration,
rerank_billed_units=rerank_billed_units,
service_tier=service_tier,
response=completion_response,
+ request_model=request_model_for_cost,
)
-
+
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
- additional_costs = _get_additional_costs(
- model=model,
- custom_llm_provider=custom_llm_provider,
- prompt_tokens=prompt_tokens,
- completion_tokens=completion_tokens,
- )
-
+ if custom_llm_provider == "azure_ai":
+ model_for_additional_costs = request_model_for_cost
+ if completion_response is not None:
+ hidden_params = getattr(completion_response, "_hidden_params", None) or {}
+ hidden_model = hidden_params.get("model") or hidden_params.get(
+ "litellm_model_name"
+ )
+ if hidden_model and (
+ "model_router" in (hidden_model or "").lower()
+ or "model-router" in (hidden_model or "").lower()
+ ):
+ model_for_additional_costs = hidden_model
+ elif model_for_additional_costs is None:
+ model_for_additional_costs = hidden_model
+ if model_for_additional_costs is None:
+ model_for_additional_costs = model
+ additional_costs = _get_additional_costs(
+ model=model_for_additional_costs,
+ custom_llm_provider=custom_llm_provider,
+ prompt_tokens=prompt_tokens or 0,
+ completion_tokens=completion_tokens or 0,
+ )
+ else:
+ additional_costs = None
+
_final_cost = (
prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar
)
@@ -1408,40 +1550,55 @@ def completion_cost( # noqa: PLR0915
)
)
_final_cost += cost_for_built_in_tools
+ if additional_costs:
+ _final_cost += sum(additional_costs.values())
- # Apply discount from module-level config if configured
original_cost = _final_cost
- _final_cost, discount_percent, discount_amount = _apply_cost_discount(
- base_cost=_final_cost,
- custom_llm_provider=custom_llm_provider,
- )
+ if litellm.cost_discount_config:
+ (
+ _final_cost,
+ discount_percent,
+ discount_amount,
+ ) = _apply_cost_discount(
+ base_cost=_final_cost,
+ custom_llm_provider=custom_llm_provider,
+ )
+ else:
+ discount_percent = 0.0
+ discount_amount = 0.0
# Apply margin from module-level config if configured
- (
- _final_cost,
- margin_percent,
- margin_fixed_amount,
- margin_total_amount,
- ) = _apply_cost_margin(
- base_cost=_final_cost,
- custom_llm_provider=custom_llm_provider,
- )
+ if litellm.cost_margin_config:
+ (
+ _final_cost,
+ margin_percent,
+ margin_fixed_amount,
+ margin_total_amount,
+ ) = _apply_cost_margin(
+ base_cost=_final_cost,
+ custom_llm_provider=custom_llm_provider,
+ )
+ else:
+ margin_percent = 0.0
+ margin_fixed_amount = 0.0
+ margin_total_amount = 0.0
# Store cost breakdown in logging object if available
- _store_cost_breakdown_in_logging_obj(
- litellm_logging_obj=litellm_logging_obj,
- prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar,
- completion_tokens_cost_usd_dollar=completion_tokens_cost_usd_dollar,
- cost_for_built_in_tools_cost_usd_dollar=cost_for_built_in_tools,
- total_cost_usd_dollar=_final_cost,
- additional_costs=additional_costs,
- original_cost=original_cost,
- discount_percent=discount_percent,
- discount_amount=discount_amount,
- margin_percent=margin_percent,
- margin_fixed_amount=margin_fixed_amount,
- margin_total_amount=margin_total_amount,
- )
+ if litellm_logging_obj is not None:
+ _store_cost_breakdown_in_logging_obj(
+ litellm_logging_obj=litellm_logging_obj,
+ prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar,
+ completion_tokens_cost_usd_dollar=completion_tokens_cost_usd_dollar,
+ cost_for_built_in_tools_cost_usd_dollar=cost_for_built_in_tools,
+ total_cost_usd_dollar=_final_cost,
+ original_cost=original_cost,
+ additional_costs=additional_costs,
+ discount_percent=discount_percent,
+ discount_amount=discount_amount,
+ margin_percent=margin_percent,
+ margin_fixed_amount=margin_fixed_amount,
+ margin_total_amount=margin_total_amount,
+ )
return _final_cost
except Exception as e:
@@ -1801,6 +1958,7 @@ def default_video_cost_calculator(
model: str,
duration_seconds: float,
custom_llm_provider: Optional[str] = None,
+ model_info: Optional[ModelInfo] = None,
) -> float:
"""
Default video cost calculator for video generation
@@ -1809,6 +1967,9 @@ def default_video_cost_calculator(
model (str): Model name
duration_seconds (float): Duration of the generated video in seconds
custom_llm_provider (Optional[str]): Custom LLM provider
+ model_info (Optional[ModelInfo]): Deployment-level model info containing
+ custom video pricing. When provided, used before falling back to
+ the global litellm.model_cost lookup.
Returns:
float: Cost in USD for the video generation
@@ -1816,43 +1977,46 @@ def default_video_cost_calculator(
Raises:
Exception: If model pricing not found in cost map
"""
- # Build model names for cost lookup
- base_model_name = model
- model_name_without_custom_llm_provider: Optional[str] = None
- if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"):
- model_name_without_custom_llm_provider = model.replace(
- f"{custom_llm_provider}/", ""
- )
- base_model_name = (
- f"{custom_llm_provider}/{model_name_without_custom_llm_provider}"
- )
-
- verbose_logger.debug(f"Looking up cost for video model: {base_model_name}")
-
- model_without_provider = model.split("/")[-1]
-
- # Try model with provider first, fall back to base model name
+ # Use custom model_info pricing if provided (deployment-specific pricing)
cost_info: Optional[dict] = None
- models_to_check: List[Optional[str]] = [
- base_model_name,
- model,
- model_without_provider,
- model_name_without_custom_llm_provider,
- ]
- for _model in models_to_check:
- if _model is not None and _model in litellm.model_cost:
- cost_info = litellm.model_cost[_model]
- break
+ if model_info is not None:
+ cost_info = dict(model_info)
+ else:
+ # Build model names for cost lookup
+ base_model_name = model
+ model_name_without_custom_llm_provider: Optional[str] = None
+ if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"):
+ model_name_without_custom_llm_provider = model.replace(
+ f"{custom_llm_provider}/", ""
+ )
+ base_model_name = (
+ f"{custom_llm_provider}/{model_name_without_custom_llm_provider}"
+ )
+
+ verbose_logger.debug(f"Looking up cost for video model: {base_model_name}")
+
+ model_without_provider = model.split("/")[-1]
+
+ # Try model with provider first, fall back to base model name
+ models_to_check: List[Optional[str]] = [
+ base_model_name,
+ model,
+ model_without_provider,
+ model_name_without_custom_llm_provider,
+ ]
+ for _model in models_to_check:
+ if _model is not None and _model in litellm.model_cost:
+ cost_info = litellm.model_cost[_model]
+ break
+
+ # If still not found, try with custom_llm_provider prefix
+ if cost_info is None and custom_llm_provider:
+ prefixed_model = f"{custom_llm_provider}/{model}"
+ if prefixed_model in litellm.model_cost:
+ cost_info = litellm.model_cost[prefixed_model]
- # If still not found, try with custom_llm_provider prefix
- if cost_info is None and custom_llm_provider:
- prefixed_model = f"{custom_llm_provider}/{model}"
- if prefixed_model in litellm.model_cost:
- cost_info = litellm.model_cost[prefixed_model]
if cost_info is None:
- raise Exception(
- f"Model not found in cost map. Tried checking {models_to_check}"
- )
+ raise Exception(f"Model not found in cost map for model={model}")
# Check for video-specific cost per second first
video_cost_per_second = cost_info.get("output_cost_per_video_per_second")
@@ -1875,9 +2039,16 @@ def batch_cost_calculator(
usage: Usage,
model: str,
custom_llm_provider: Optional[str] = None,
+ model_info: Optional[ModelInfo] = None,
) -> Tuple[float, float]:
"""
- Calculate the cost of a batch job
+ Calculate the cost of a batch job.
+
+ Args:
+ model_info: Optional deployment-level model info containing custom
+ batch pricing (e.g. input_cost_per_token_batches). When provided,
+ skips the global litellm.get_model_info() lookup so that
+ deployment-specific pricing is used.
"""
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
@@ -1890,12 +2061,13 @@ def batch_cost_calculator(
custom_llm_provider,
)
- try:
- model_info: Optional[ModelInfo] = litellm.get_model_info(
- model=model, custom_llm_provider=custom_llm_provider
- )
- except Exception:
- model_info = None
+ if model_info is None:
+ try:
+ model_info = litellm.get_model_info(
+ model=model, custom_llm_provider=custom_llm_provider
+ )
+ except Exception:
+ model_info = None
if not model_info:
return 0.0, 0.0
diff --git a/litellm/evals/__init__.py b/litellm/evals/__init__.py
new file mode 100644
index 00000000000..89dfb62b2b7
--- /dev/null
+++ b/litellm/evals/__init__.py
@@ -0,0 +1,33 @@
+"""
+Evals API operations
+"""
+
+from .main import (
+ acancel_eval,
+ acreate_eval,
+ adelete_eval,
+ aget_eval,
+ alist_evals,
+ aupdate_eval,
+ cancel_eval,
+ create_eval,
+ delete_eval,
+ get_eval,
+ list_evals,
+ update_eval,
+)
+
+__all__ = [
+ "acreate_eval",
+ "alist_evals",
+ "aget_eval",
+ "aupdate_eval",
+ "adelete_eval",
+ "acancel_eval",
+ "create_eval",
+ "list_evals",
+ "get_eval",
+ "update_eval",
+ "delete_eval",
+ "cancel_eval",
+]
diff --git a/litellm/evals/main.py b/litellm/evals/main.py
new file mode 100644
index 00000000000..eab909a6b11
--- /dev/null
+++ b/litellm/evals/main.py
@@ -0,0 +1,1986 @@
+"""
+Main entry point for Evals API operations
+Provides create, list, get, update, delete, and cancel operations for evals
+"""
+
+import asyncio
+import contextvars
+from functools import partial
+from typing import Any, Coroutine, Dict, List, Optional, Union
+
+import httpx
+
+import litellm
+from litellm.constants import request_timeout
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig
+from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
+from litellm.types.llms.openai_evals import (
+ CancelEvalResponse,
+ CancelRunResponse,
+ CreateEvalRequest,
+ CreateRunRequest,
+ DeleteEvalResponse,
+ Eval,
+ ListEvalsParams,
+ ListEvalsResponse,
+ ListRunsParams,
+ ListRunsResponse,
+ Run,
+ RunDeleteResponse,
+ UpdateEvalRequest,
+)
+from litellm.types.router import GenericLiteLLMParams
+from litellm.utils import ProviderConfigManager, client
+
+# Initialize HTTP handler
+base_llm_http_handler = BaseLLMHTTPHandler()
+DEFAULT_OPENAI_API_BASE = "https://api.openai.com"
+
+
+@client
+async def acreate_eval(
+ data_source_config: Dict[str, Any],
+ testing_criteria: List[Dict[str, Any]],
+ name: Optional[str] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ extra_body: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Eval:
+ """
+ Async: Create a new evaluation
+
+ Args:
+ data_source_config: Configuration for the data source
+ testing_criteria: List of graders for all eval runs
+ name: Optional name for the evaluation
+ metadata: Optional additional metadata (max 16 key-value pairs)
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ extra_body: Additional body parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Eval object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["acreate_eval"] = True
+
+ func = partial(
+ create_eval,
+ data_source_config=data_source_config,
+ testing_criteria=testing_criteria,
+ name=name,
+ metadata=metadata,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ extra_body=extra_body,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def create_eval(
+ data_source_config: Dict[str, Any],
+ testing_criteria: List[Dict[str, Any]],
+ name: Optional[str] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ extra_body: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[Eval, Coroutine[Any, Any, Eval]]:
+ """
+ Create a new evaluation
+
+ Args:
+ data_source_config: Configuration for the data source
+ testing_criteria: List of graders for all eval runs
+ name: Optional name for the evaluation
+ metadata: Optional additional metadata (max 16 key-value pairs)
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ extra_body: Additional body parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Eval object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("acreate_eval", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"CREATE eval is not supported for {custom_llm_provider}")
+
+ # Build create request
+ create_request: CreateEvalRequest = {
+ "data_source_config": data_source_config, # type: ignore
+ "testing_criteria": testing_criteria, # type: ignore
+ }
+ if name is not None:
+ create_request["name"] = name
+
+ # Merge extra_body if provided
+ if extra_body:
+ create_request.update(extra_body) # type: ignore
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ request_body = evals_api_provider_config.transform_create_eval_request(
+ create_request=create_request,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Get API base and URL
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ url = evals_api_provider_config.get_complete_url(
+ api_base=api_base, endpoint="evals"
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params=request_body,
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.create_eval_handler( # type: ignore
+ url=url,
+ request_body=request_body,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+async def alist_evals(
+ limit: Optional[int] = None,
+ after: Optional[str] = None,
+ before: Optional[str] = None,
+ order: Optional[str] = None,
+ order_by: Optional[str] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> ListEvalsResponse:
+ """
+ Async: List all evaluations
+
+ Args:
+ limit: Number of results to return per page (max 100, default 20)
+ after: Cursor for pagination - returns evals after this ID
+ before: Cursor for pagination - returns evals before this ID
+ order: Sort order ('asc' or 'desc', default 'desc')
+ order_by: Field to sort by ('created_at' or 'updated_at', default 'created_at')
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ ListEvalsResponse object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["alist_evals"] = True
+
+ func = partial(
+ list_evals,
+ limit=limit,
+ after=after,
+ before=before,
+ order=order,
+ order_by=order_by,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def list_evals(
+ limit: Optional[int] = None,
+ after: Optional[str] = None,
+ before: Optional[str] = None,
+ order: Optional[str] = None,
+ order_by: Optional[str] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[ListEvalsResponse, Coroutine[Any, Any, ListEvalsResponse]]:
+ """
+ List all evaluations
+
+ Args:
+ limit: Number of results to return per page (max 100, default 20)
+ after: Cursor for pagination - returns evals after this ID
+ before: Cursor for pagination - returns evals before this ID
+ order: Sort order ('asc' or 'desc', default 'desc')
+ order_by: Field to sort by ('created_at' or 'updated_at', default 'created_at')
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ ListEvalsResponse object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("alist_evals", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"LIST evals is not supported for {custom_llm_provider}")
+
+ # Build list parameters
+ list_params: ListEvalsParams = {}
+ if limit is not None:
+ list_params["limit"] = limit
+ if after is not None:
+ list_params["after"] = after
+ if before is not None:
+ list_params["before"] = before
+ if order is not None:
+ list_params["order"] = order # type: ignore
+ if order_by is not None:
+ list_params["order_by"] = order_by # type: ignore
+
+ # Merge extra_query if provided
+ if extra_query:
+ list_params.update(extra_query) # type: ignore
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ url, query_params = evals_api_provider_config.transform_list_evals_request(
+ list_params=list_params,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params=query_params,
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.list_evals_handler( # type: ignore
+ url=url,
+ query_params=query_params,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+async def aget_eval(
+ eval_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Eval:
+ """
+ Async: Get an evaluation by ID
+
+ Args:
+ eval_id: The ID of the evaluation to fetch
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Eval object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["aget_eval"] = True
+
+ func = partial(
+ get_eval,
+ eval_id=eval_id,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def get_eval(
+ eval_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[Eval, Coroutine[Any, Any, Eval]]:
+ """
+ Get an evaluation by ID
+
+ Args:
+ eval_id: The ID of the evaluation to fetch
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Eval object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("aget_eval", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"GET eval is not supported for {custom_llm_provider}")
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ url, headers = evals_api_provider_config.transform_get_eval_request(
+ eval_id=eval_id,
+ api_base=api_base,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params={"eval_id": eval_id},
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.get_eval_handler( # type: ignore
+ url=url,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+async def aupdate_eval(
+ eval_id: str,
+ name: Optional[str] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ extra_body: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Eval:
+ """
+ Async: Update an evaluation
+
+ Args:
+ eval_id: The ID of the evaluation to update
+ name: Updated name
+ metadata: Updated metadata
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ extra_body: Additional body parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Eval object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["aupdate_eval"] = True
+
+ func = partial(
+ update_eval,
+ eval_id=eval_id,
+ name=name,
+ metadata=metadata,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ extra_body=extra_body,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def update_eval(
+ eval_id: str,
+ name: Optional[str] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ extra_body: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[Eval, Coroutine[Any, Any, Eval]]:
+ """
+ Update an evaluation
+
+ Args:
+ eval_id: The ID of the evaluation to update
+ name: Updated name
+ metadata: Updated metadata
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ extra_body: Additional body parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Eval object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("aupdate_eval", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"UPDATE eval is not supported for {custom_llm_provider}")
+
+ # Build update request
+ update_request: UpdateEvalRequest = {}
+ if name is not None:
+ update_request["name"] = name
+
+ # Filter metadata to exclude internal LiteLLM fields
+ if metadata is not None:
+ # List of internal LiteLLM metadata keys that should NOT be sent to OpenAI
+ internal_keys = {
+ "headers",
+ "requester_metadata",
+ "user_api_key_hash",
+ "user_api_key_alias",
+ "user_api_key_spend",
+ "user_api_key_max_budget",
+ "user_api_key_team_id",
+ "user_api_key_user_id",
+ "user_api_key_org_id",
+ "user_api_key_team_alias",
+ "user_api_key_end_user_id",
+ "user_api_key_user_email",
+ "user_api_key_request_route",
+ "user_api_key_budget_reset_at",
+ "user_api_key_auth_metadata",
+ "user_api_key",
+ "user_api_end_user_max_budget",
+ "user_api_key_auth",
+ "litellm_api_version",
+ "global_max_parallel_requests",
+ "user_api_key_team_max_budget",
+ "user_api_key_team_spend",
+ "user_api_key_model_max_budget",
+ "user_api_key_user_spend",
+ "user_api_key_user_max_budget",
+ "user_api_key_metadata",
+ "endpoint",
+ "litellm_parent_otel_span",
+ "requester_ip_address",
+ "user_agent",
+ }
+ # Only include user-provided metadata keys
+ filtered_metadata = {
+ k: v for k, v in metadata.items() if k not in internal_keys
+ }
+ if filtered_metadata: # Only add if there's user metadata
+ update_request["metadata"] = filtered_metadata
+
+ # Merge extra_body if provided
+ if extra_body:
+ update_request.update(extra_body) # type: ignore
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ (
+ url,
+ headers,
+ request_body,
+ ) = evals_api_provider_config.transform_update_eval_request(
+ eval_id=eval_id,
+ update_request=update_request,
+ api_base=api_base,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params=request_body,
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.update_eval_handler( # type: ignore
+ url=url,
+ request_body=request_body,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+async def adelete_eval(
+ eval_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> DeleteEvalResponse:
+ """
+ Async: Delete an evaluation
+
+ Args:
+ eval_id: The ID of the evaluation to delete
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ DeleteEvalResponse object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["adelete_eval"] = True
+
+ func = partial(
+ delete_eval,
+ eval_id=eval_id,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def delete_eval(
+ eval_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[DeleteEvalResponse, Coroutine[Any, Any, DeleteEvalResponse]]:
+ """
+ Delete an evaluation
+
+ Args:
+ eval_id: The ID of the evaluation to delete
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ DeleteEvalResponse object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("adelete_eval", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"DELETE eval is not supported for {custom_llm_provider}")
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ url, headers = evals_api_provider_config.transform_delete_eval_request(
+ eval_id=eval_id,
+ api_base=api_base,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params={"eval_id": eval_id},
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.delete_eval_handler( # type: ignore
+ url=url,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+async def acancel_eval(
+ eval_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> CancelEvalResponse:
+ """
+ Async: Cancel a running evaluation
+
+ Args:
+ eval_id: The ID of the evaluation to cancel
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ CancelEvalResponse object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["acancel_eval"] = True
+
+ func = partial(
+ cancel_eval,
+ eval_id=eval_id,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def cancel_eval(
+ eval_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[CancelEvalResponse, Coroutine[Any, Any, CancelEvalResponse]]:
+ """
+ Cancel a running evaluation
+
+ Args:
+ eval_id: The ID of the evaluation to cancel
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ CancelEvalResponse object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("acancel_eval", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"CANCEL eval is not supported for {custom_llm_provider}")
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ (
+ url,
+ headers,
+ request_body,
+ ) = evals_api_provider_config.transform_cancel_eval_request(
+ eval_id=eval_id,
+ api_base=api_base,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params={"eval_id": eval_id},
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.cancel_eval_handler( # type: ignore
+ url=url,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+# ===================================
+# Run API Functions
+# ===================================
+
+
+@client
+async def acreate_run(
+ eval_id: str,
+ data_source: Dict[str, Any],
+ name: Optional[str] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ extra_body: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Run:
+ """
+ Async: Create a new run for an evaluation
+
+ Args:
+ eval_id: The ID of the evaluation to run
+ data_source: Data source configuration for the run (can be jsonl, completions, or responses type)
+ name: Optional name for the run
+ metadata: Optional additional metadata
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ extra_body: Additional body parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Run object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["acreate_run"] = True
+
+ func = partial(
+ create_run,
+ eval_id=eval_id,
+ data_source=data_source,
+ name=name,
+ metadata=metadata,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ extra_body=extra_body,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def create_run(
+ eval_id: str,
+ data_source: Dict[str, Any],
+ name: Optional[str] = None,
+ metadata: Optional[Dict[str, Any]] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ extra_body: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[Run, Coroutine[Any, Any, Run]]:
+ """
+ Create a new run for an evaluation
+
+ Args:
+ eval_id: The ID of the evaluation to run
+ data_source: Data source configuration for the run (can be jsonl, completions, or responses type)
+ name: Optional name for the run
+ metadata: Optional additional metadata
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ extra_body: Additional body parameters
+ timeout: Request timeout (default 600s for long-running operations)
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Run object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("acreate_run", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"CREATE run is not supported for {custom_llm_provider}")
+
+ # Build create request
+ create_request: CreateRunRequest = {
+ "data_source": data_source, # type: ignore
+ }
+ if name is not None:
+ create_request["name"] = name
+ # if metadata is not None:
+ # create_request["metadata"] = metadata
+
+ # Merge extra_body if provided
+ if extra_body:
+ create_request.update(extra_body) # type: ignore
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ url, request_body = evals_api_provider_config.transform_create_run_request(
+ eval_id=eval_id,
+ create_request=create_request,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params=request_body,
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request (default 600s timeout for long-running operations)
+ response = base_llm_http_handler.create_run_handler( # type: ignore
+ url=url,
+ request_body=request_body,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or httpx.Timeout(timeout=600.0, connect=5.0),
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+async def alist_runs(
+ eval_id: str,
+ limit: Optional[int] = None,
+ after: Optional[str] = None,
+ before: Optional[str] = None,
+ order: Optional[str] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> ListRunsResponse:
+ """
+ Async: List all runs for an evaluation
+
+ Args:
+ eval_id: The ID of the evaluation
+ limit: Number of results to return per page (max 100, default 20)
+ after: Cursor for pagination - returns runs after this ID
+ before: Cursor for pagination - returns runs before this ID
+ order: Sort order ('asc' or 'desc', default 'desc')
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ ListRunsResponse object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["alist_runs"] = True
+
+ func = partial(
+ list_runs,
+ eval_id=eval_id,
+ limit=limit,
+ after=after,
+ before=before,
+ order=order,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def list_runs(
+ eval_id: str,
+ limit: Optional[int] = None,
+ after: Optional[str] = None,
+ before: Optional[str] = None,
+ order: Optional[str] = None,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[ListRunsResponse, Coroutine[Any, Any, ListRunsResponse]]:
+ """
+ List all runs for an evaluation
+
+ Args:
+ eval_id: The ID of the evaluation
+ limit: Number of results to return per page (max 100, default 20)
+ after: Cursor for pagination - returns runs after this ID
+ before: Cursor for pagination - returns runs before this ID
+ order: Sort order ('asc' or 'desc', default 'desc')
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ ListRunsResponse object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("alist_runs", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"LIST runs is not supported for {custom_llm_provider}")
+
+ # Build list parameters
+ list_params: ListRunsParams = {}
+ if limit is not None:
+ list_params["limit"] = limit
+ if after is not None:
+ list_params["after"] = after
+ if before is not None:
+ list_params["before"] = before
+ if order is not None:
+ list_params["order"] = order # type: ignore
+
+ # Merge extra_query if provided
+ if extra_query:
+ list_params.update(extra_query) # type: ignore
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ url, query_params = evals_api_provider_config.transform_list_runs_request(
+ eval_id=eval_id,
+ list_params=list_params,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params={"eval_id": eval_id, **query_params},
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.list_runs_handler( # type: ignore
+ url=url,
+ query_params=query_params,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+async def aget_run(
+ eval_id: str,
+ run_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Run:
+ """
+ Async: Get a specific run
+
+ Args:
+ eval_id: The ID of the evaluation
+ run_id: The ID of the run to retrieve
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Run object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["aget_run"] = True
+
+ func = partial(
+ get_run,
+ eval_id=eval_id,
+ run_id=run_id,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def get_run(
+ eval_id: str,
+ run_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[Run, Coroutine[Any, Any, Run]]:
+ """
+ Get a specific run
+
+ Args:
+ eval_id: The ID of the evaluation
+ run_id: The ID of the run to retrieve
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ Run object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("aget_run", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"GET run is not supported for {custom_llm_provider}")
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ url, headers = evals_api_provider_config.transform_get_run_request(
+ eval_id=eval_id,
+ run_id=run_id,
+ api_base=api_base,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params={"eval_id": eval_id, "run_id": run_id},
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.get_run_handler( # type: ignore
+ url=url,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+async def acancel_run(
+ eval_id: str,
+ run_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> CancelRunResponse:
+ """
+ Async: Cancel a running run
+
+ Args:
+ eval_id: The ID of the evaluation
+ run_id: The ID of the run to cancel
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ CancelRunResponse object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["acancel_run"] = True
+
+ func = partial(
+ cancel_run,
+ eval_id=eval_id,
+ run_id=run_id,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def cancel_run(
+ eval_id: str,
+ run_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[CancelRunResponse, Coroutine[Any, Any, CancelRunResponse]]:
+ """
+ Cancel a running run
+
+ Args:
+ eval_id: The ID of the evaluation
+ run_id: The ID of the run to cancel
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ CancelRunResponse object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("acancel_run", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"CANCEL run is not supported for {custom_llm_provider}")
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ (
+ url,
+ headers,
+ request_body,
+ ) = evals_api_provider_config.transform_cancel_run_request(
+ eval_id=eval_id,
+ run_id=run_id,
+ api_base=api_base,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params={"eval_id": eval_id, "run_id": run_id},
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.cancel_run_handler( # type: ignore
+ url=url,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+# ===================================
+# Delete Run API Functions
+# ===================================
+
+
+@client
+async def adelete_run(
+ eval_id: str,
+ run_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> RunDeleteResponse:
+ """
+ Async: Delete a run
+
+ Args:
+ eval_id: The ID of the evaluation
+ run_id: The ID of the run to delete
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ RunDeleteResponse object
+ """
+ local_vars = locals()
+ try:
+ loop = asyncio.get_event_loop()
+ kwargs["adelete_run"] = True
+
+ func = partial(
+ delete_run,
+ eval_id=eval_id,
+ run_id=run_id,
+ extra_headers=extra_headers,
+ extra_query=extra_query,
+ timeout=timeout,
+ custom_llm_provider=custom_llm_provider,
+ **kwargs,
+ )
+
+ ctx = contextvars.copy_context()
+ func_with_context = partial(ctx.run, func)
+ init_response = await loop.run_in_executor(None, func_with_context)
+
+ if asyncio.iscoroutine(init_response):
+ response = await init_response
+ else:
+ response = init_response
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
+
+
+@client
+def delete_run(
+ eval_id: str,
+ run_id: str,
+ extra_headers: Optional[Dict[str, Any]] = None,
+ extra_query: Optional[Dict[str, Any]] = None,
+ timeout: Optional[Union[float, httpx.Timeout]] = None,
+ custom_llm_provider: Optional[str] = None,
+ **kwargs,
+) -> Union[RunDeleteResponse, Coroutine[Any, Any, RunDeleteResponse]]:
+ """
+ Delete a run
+
+ Args:
+ eval_id: The ID of the evaluation
+ run_id: The ID of the run to delete
+ extra_headers: Additional headers for the request
+ extra_query: Additional query parameters
+ timeout: Request timeout
+ custom_llm_provider: Provider name (e.g., 'openai')
+ **kwargs: Additional parameters
+
+ Returns:
+ RunDeleteResponse object
+ """
+ local_vars = locals()
+ try:
+ litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
+ litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ _is_async = kwargs.pop("adelete_run", False) is True
+
+ # Get LiteLLM parameters
+ litellm_params = GenericLiteLLMParams(**kwargs)
+
+ # Determine provider
+ if custom_llm_provider is None:
+ custom_llm_provider = "openai"
+
+ # Get provider config
+ evals_api_provider_config: Optional[
+ BaseEvalsAPIConfig
+ ] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
+ provider=litellm.LlmProviders(custom_llm_provider),
+ )
+
+ if evals_api_provider_config is None:
+ raise ValueError(f"DELETE run is not supported for {custom_llm_provider}")
+
+ # Validate environment and get headers
+ headers = extra_headers or {}
+ headers = evals_api_provider_config.validate_environment(
+ headers=headers, litellm_params=litellm_params
+ )
+
+ # Transform request
+ api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
+ (
+ url,
+ headers,
+ request_body,
+ ) = evals_api_provider_config.transform_delete_run_request(
+ eval_id=eval_id,
+ run_id=run_id,
+ api_base=api_base,
+ litellm_params=litellm_params,
+ headers=headers,
+ )
+
+ # Pre-call logging
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
+ model=None,
+ optional_params={"eval_id": eval_id, "run_id": run_id},
+ litellm_params={
+ "litellm_call_id": litellm_call_id,
+ },
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ # Make HTTP request
+ response = base_llm_http_handler.delete_run_handler( # type: ignore
+ url=url,
+ evals_api_provider_config=evals_api_provider_config,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=headers,
+ timeout=timeout or request_timeout,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ shared_session=kwargs.get("shared_session"),
+ )
+
+ return response
+ except Exception as e:
+ raise litellm.exception_type(
+ model=None,
+ custom_llm_provider=custom_llm_provider,
+ original_exception=e,
+ completion_kwargs=local_vars,
+ extra_kwargs=kwargs,
+ )
diff --git a/litellm/exceptions.py b/litellm/exceptions.py
index eb027334606..abdba09dd8d 100644
--- a/litellm/exceptions.py
+++ b/litellm/exceptions.py
@@ -25,9 +25,7 @@ def _get_minimal_error_response() -> httpx.Response:
if _MINIMAL_ERROR_RESPONSE is None:
_MINIMAL_ERROR_RESPONSE = httpx.Response(
status_code=400,
- request=httpx.Request(
- method="GET", url="https://litellm.ai"
- ),
+ request=httpx.Request(method="GET", url="https://litellm.ai"),
)
return _MINIMAL_ERROR_RESPONSE
@@ -955,7 +953,8 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
generated_content: str = "",
is_pre_first_chunk: bool = False,
):
- self.status_code = 503 # Service Unavailable
+ original_status = getattr(original_exception, "status_code", None)
+ self.status_code = int(original_status) if original_status is not None else 503
self.message = f"litellm.MidStreamFallbackError: {message}"
self.model = model
self.llm_provider = llm_provider
@@ -978,7 +977,14 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
else:
self.response = response
- # Call the parent constructor
+ # Save the original attributes before they are overridden by ServiceUnavailableError
+ _saved_response = self.response
+ _saved_request = getattr(self.response, "request", None) or httpx.Request(
+ method="POST", url=f"https://{llm_provider}.com/v1/"
+ )
+ _saved_message = self.message
+
+ # Call the parent constructor (which hardcodes status_code=503 and modifies the response object)
super().__init__(
message=self.message,
llm_provider=llm_provider,
@@ -989,6 +995,13 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
num_retries=self.num_retries,
)
+ # Restore the propagated status and original response/request objects
+ self.status_code = int(original_status) if original_status is not None else 503
+ self.response = _saved_response
+ self.request = _saved_request
+ self.message = _saved_message
+ self.args = (_saved_message,)
+
def __str__(self):
_message = self.message
if self.num_retries:
diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py
index 3e8f9bc337b..a638a28aba3 100644
--- a/litellm/experimental_mcp_client/client.py
+++ b/litellm/experimental_mcp_client/client.py
@@ -4,7 +4,18 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
import asyncio
import base64
-from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, TypeVar, Union
+from typing import (
+ Any,
+ Awaitable,
+ Callable,
+ Dict,
+ Generator,
+ List,
+ Optional,
+ Tuple,
+ TypeVar,
+ Union,
+)
import httpx
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
@@ -14,7 +25,10 @@ from mcp.client.stdio import stdio_client
streamable_http_client: Optional[Any] = None
try:
import mcp.client.streamable_http as streamable_http_module # type: ignore
- streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None)
+
+ streamable_http_client = getattr(
+ streamable_http_module, "streamable_http_client", None
+ )
except ImportError:
pass
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
@@ -30,6 +44,7 @@ from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
from litellm._logging import verbose_logger
+from litellm.constants import MCP_CLIENT_TIMEOUT
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
from litellm.types.llms.custom_http import VerifyTypes
from litellm.types.mcp import (
@@ -49,6 +64,86 @@ def to_basic_auth(auth_value: str) -> str:
TSessionResult = TypeVar("TSessionResult")
+class MCPSigV4Auth(httpx.Auth):
+ """
+ httpx Auth class that signs each request with AWS SigV4.
+
+ This is used for MCP servers that require AWS SigV4 authentication,
+ such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow()
+ for every outgoing request, enabling per-request signature computation.
+ """
+
+ requires_request_body = True
+
+ def __init__(
+ self,
+ aws_access_key_id: Optional[str] = None,
+ aws_secret_access_key: Optional[str] = None,
+ aws_session_token: Optional[str] = None,
+ aws_region_name: Optional[str] = None,
+ aws_service_name: Optional[str] = None,
+ ):
+ try:
+ from botocore.credentials import Credentials
+ except ImportError:
+ raise ImportError(
+ "Missing botocore to use AWS SigV4 authentication. "
+ "Run 'pip install boto3'."
+ )
+
+ self.service_name = aws_service_name or "bedrock-agentcore"
+ self.region_name = aws_region_name or "us-east-1"
+
+ # Note: os.environ/ prefixed values are already resolved by
+ # ProxyConfig._check_for_os_environ_vars() at config load time.
+ # Values arrive here as plain strings.
+ if aws_access_key_id and aws_secret_access_key:
+ self.credentials = Credentials(
+ access_key=aws_access_key_id,
+ secret_key=aws_secret_access_key,
+ token=aws_session_token,
+ )
+ else:
+ # Fall back to default boto3 credential chain
+ import botocore.session
+
+ session = botocore.session.get_session()
+ self.credentials = session.get_credentials()
+ if self.credentials is None:
+ raise ValueError(
+ "No AWS credentials found. Provide aws_access_key_id and "
+ "aws_secret_access_key, or configure default credentials "
+ "(env vars, ~/.aws/credentials, instance profile)."
+ )
+
+ def auth_flow(
+ self, request: httpx.Request
+ ) -> Generator[httpx.Request, httpx.Response, None]:
+ from botocore.auth import SigV4Auth
+ from botocore.awsrequest import AWSRequest
+
+ # Build AWSRequest from the httpx Request.
+ # Pass all request headers so the canonical SigV4 signature covers them.
+ aws_request = AWSRequest(
+ method=request.method,
+ url=str(request.url),
+ data=request.content,
+ headers=dict(request.headers),
+ )
+
+ # Sign the request — SigV4Auth.add_auth() adds Authorization,
+ # X-Amz-Date, and X-Amz-Security-Token (if session token present).
+ # Host header is derived automatically from the URL.
+ sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name)
+ sigv4.add_auth(aws_request)
+
+ # Copy SigV4 headers back to the httpx request
+ for header_name, header_value in aws_request.headers.items():
+ request.headers[header_name] = header_value
+
+ yield request
+
+
class MCPClient:
"""
MCP Client supporting:
@@ -63,19 +158,21 @@ class MCPClient:
transport_type: MCPTransportType = MCPTransport.http,
auth_type: MCPAuthType = None,
auth_value: Optional[Union[str, Dict[str, str]]] = None,
- timeout: float = 60.0,
+ timeout: Optional[float] = None,
stdio_config: Optional[MCPStdioConfig] = None,
extra_headers: Optional[Dict[str, str]] = None,
ssl_verify: Optional[VerifyTypes] = None,
+ aws_auth: Optional[httpx.Auth] = None,
):
self.server_url: str = server_url
self.transport_type: MCPTransport = transport_type
self.auth_type: MCPAuthType = auth_type
- self.timeout: float = timeout
+ self.timeout: float = timeout if timeout is not None else MCP_CLIENT_TIMEOUT
self._mcp_auth_value: Optional[Union[str, Dict[str, str]]] = None
self.stdio_config: Optional[MCPStdioConfig] = stdio_config
self.extra_headers: Optional[Dict[str, str]] = extra_headers
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
+ self._aws_auth: Optional[httpx.Auth] = aws_auth
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
@@ -105,12 +202,15 @@ class MCPClient:
if self.transport_type == MCPTransport.sse:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
- return sse_client(
- url=self.server_url,
- timeout=self.timeout,
- headers=headers,
- httpx_client_factory=httpx_client_factory,
- ), None
+ return (
+ sse_client(
+ url=self.server_url,
+ timeout=self.timeout,
+ headers=headers,
+ httpx_client_factory=httpx_client_factory,
+ ),
+ None,
+ )
# HTTP transport (default)
if streamable_http_client is None:
@@ -118,12 +218,10 @@ class MCPClient:
"streamable_http_client is not available. "
"Please install mcp with HTTP support."
)
-
+
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
- verbose_logger.debug(
- "litellm headers for streamable_http_client: %s", headers
- )
+ verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
http_client = httpx_client_factory(
headers=headers,
timeout=httpx.Timeout(self.timeout),
@@ -209,8 +307,15 @@ class MCPClient:
headers["X-API-Key"] = self._mcp_auth_value
elif self.auth_type == MCPAuth.authorization:
headers["Authorization"] = self._mcp_auth_value
+ elif self.auth_type == MCPAuth.oauth2:
+ headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
+ elif self.auth_type == MCPAuth.token:
+ headers["Authorization"] = f"token {self._mcp_auth_value}"
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
+ # Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
+ # signing (including the body hash), so it uses httpx.Auth flow instead
+ # of static headers. See MCPSigV4Auth and _create_httpx_client_factory().
# update the headers with the extra headers
if self.extra_headers:
@@ -243,10 +348,16 @@ class MCPClient:
f"MCP client using SSL configuration: {type(ssl_config).__name__}"
)
+ # Use SigV4 auth if configured and no explicit auth provided.
+ # The MCP SDK's sse_client and streamable_http_client call this
+ # factory without passing auth=, so self._aws_auth is used.
+ # For non-SigV4 clients, self._aws_auth is None — no behavior change.
+ effective_auth = auth if auth is not None else self._aws_auth
+
return httpx.AsyncClient(
headers=headers,
timeout=timeout,
- auth=auth,
+ auth=effective_auth,
verify=ssl_config,
follow_redirects=True,
)
@@ -296,7 +407,7 @@ class MCPClient:
async def call_tool(
self,
call_tool_request_params: MCPCallToolRequestParams,
- host_progress_callback: Optional[Callable] = None
+ host_progress_callback: Optional[Callable] = None,
) -> MCPCallToolResult:
"""
Call an MCP Tool.
@@ -305,13 +416,15 @@ class MCPClient:
f"MCP client calling tool '{call_tool_request_params.name}' with arguments: {call_tool_request_params.arguments}"
)
- async def on_progress(progress: float, total: float | None, message: str | None):
+ async def on_progress(
+ progress: float, total: float | None, message: str | None
+ ):
percentage = (progress / total * 100) if total else 0
verbose_logger.info(
f"MCP Tool '{call_tool_request_params.name}' progress: "
f"{progress}/{total} ({percentage:.0f}%) - {message or ''}"
)
-
+
# Forward to Host if callback provided
if host_progress_callback:
try:
@@ -325,8 +438,8 @@ class MCPClient:
name=call_tool_request_params.name,
arguments=call_tool_request_params.arguments,
progress_callback=on_progress,
-
)
+
try:
tool_result = await self.run_with_session(_call_tool_operation)
verbose_logger.info(
diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py
index b716e3171e7..bd42f7e7111 100644
--- a/litellm/experimental_mcp_client/tools.py
+++ b/litellm/experimental_mcp_client/tools.py
@@ -18,7 +18,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall
def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolParam:
"""Convert an MCP tool to an OpenAI tool."""
normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema)
-
+
return ChatCompletionToolParam(
type="function",
function=FunctionDefinition(
@@ -33,41 +33,39 @@ def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolPa
def _normalize_mcp_input_schema(input_schema: dict) -> dict:
"""
Normalize MCP input schema to ensure it's valid for OpenAI function calling.
-
+
OpenAI requires that function parameters have:
- type: 'object'
- properties: dict (can be empty)
- additionalProperties: false (recommended)
"""
if not input_schema:
- return {
- "type": "object",
- "properties": {},
- "additionalProperties": False
- }
-
+ return {"type": "object", "properties": {}, "additionalProperties": False}
+
# Make a copy to avoid modifying the original
normalized_schema = dict(input_schema)
-
+
# Ensure type is 'object'
if "type" not in normalized_schema:
normalized_schema["type"] = "object"
-
+
# Ensure properties exists (can be empty)
if "properties" not in normalized_schema:
normalized_schema["properties"] = {}
-
+
# Add additionalProperties if not present (recommended by OpenAI)
if "additionalProperties" not in normalized_schema:
normalized_schema["additionalProperties"] = False
-
+
return normalized_schema
-def transform_mcp_tool_to_openai_responses_api_tool(mcp_tool: MCPTool) -> FunctionToolParam:
+def transform_mcp_tool_to_openai_responses_api_tool(
+ mcp_tool: MCPTool,
+) -> FunctionToolParam:
"""Convert an MCP tool to an OpenAI Responses API tool."""
normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema)
-
+
return FunctionToolParam(
name=mcp_tool.name,
parameters=normalized_parameters,
@@ -76,6 +74,7 @@ def transform_mcp_tool_to_openai_responses_api_tool(mcp_tool: MCPTool) -> Functi
description=mcp_tool.description or "",
)
+
async def load_mcp_tools(
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
) -> Union[List[MCPTool], List[ChatCompletionToolParam]]:
diff --git a/litellm/files/main.py b/litellm/files/main.py
index 78e41bb5a68..f7c89e0ba3b 100644
--- a/litellm/files/main.py
+++ b/litellm/files/main.py
@@ -7,7 +7,6 @@ https://platform.openai.com/docs/api-reference/files
import asyncio
import contextvars
-import os
import time
import uuid as uuid_module
from functools import partial
@@ -15,15 +14,36 @@ from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast
import httpx
+# Type aliases for provider parameters
+FileCreateProvider = Literal[
+ "openai",
+ "azure",
+ "gemini",
+ "vertex_ai",
+ "bedrock",
+ "hosted_vllm",
+ "manus",
+ "anthropic",
+]
+FileRetrieveProvider = Literal[
+ "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus", "anthropic"
+]
+FileDeleteProvider = Literal["openai", "azure", "gemini", "manus", "anthropic"]
+FileListProvider = Literal["openai", "azure", "manus", "anthropic"]
+FileContentProvider = Literal[
+ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"
+]
+
import litellm
from litellm import get_secret_str
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.llms.anthropic.files.handler import AnthropicFilesHandler
+from litellm.llms.azure.common_utils import get_azure_credentials
from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI
from litellm.llms.bedrock.files.handler import BedrockFilesHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
+from litellm.llms.openai.common_utils import get_openai_credentials
from litellm.llms.openai.openai import FileDeleted, FileObject, OpenAIFilesAPI
from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler
from litellm.types.llms.openai import (
@@ -53,16 +73,15 @@ openai_files_instance = OpenAIFilesAPI()
azure_files_instance = AzureOpenAIFilesAPI()
vertex_ai_files_instance = VertexAIFilesHandler()
bedrock_files_instance = BedrockFilesHandler()
-anthropic_files_instance = AnthropicFilesHandler()
#################################################
@client
async def acreate_file(
file: FileTypes,
- purpose: Literal["assistants", "batch", "fine-tune"],
+ purpose: Literal["assistants", "batch", "fine-tune", "messages"],
expires_after: Optional[FileExpiresAfter] = None,
- custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "bedrock", "hosted_vllm", "manus"] = "openai",
+ custom_llm_provider: FileCreateProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -105,9 +124,9 @@ async def acreate_file(
@client
def create_file(
file: FileTypes,
- purpose: Literal["assistants", "batch", "fine-tune"],
+ purpose: Literal["assistants", "batch", "fine-tune", "messages"],
expires_after: Optional[FileExpiresAfter] = None,
- custom_llm_provider: Optional[Literal["openai", "azure", "gemini", "vertex_ai", "bedrock", "hosted_vllm", "manus"]] = None,
+ custom_llm_provider: Optional[FileCreateProvider] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -185,98 +204,39 @@ def create_file(
timeout=timeout,
)
elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
- # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
- api_base = (
- optional_params.api_base
- or litellm.api_base
- or os.getenv("OPENAI_BASE_URL")
- or os.getenv("OPENAI_API_BASE")
- or "https://api.openai.com/v1"
+ openai_creds = get_openai_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ organization=optional_params.organization,
)
- organization = (
- optional_params.organization
- or litellm.organization
- or os.getenv("OPENAI_ORGANIZATION", None)
- or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
- )
- # set API KEY
- api_key = (
- optional_params.api_key
- or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
- or litellm.openai_key
- or os.getenv("OPENAI_API_KEY")
- )
-
response = openai_files_instance.create_file(
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
+ api_base=openai_creds.api_base,
+ api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
- organization=organization,
+ organization=openai_creds.organization,
create_file_data=_create_file_request,
)
elif custom_llm_provider == "azure":
- api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
- api_version = (
- optional_params.api_version
- or litellm.api_version
- or get_secret_str("AZURE_API_VERSION")
- ) # type: ignore
-
- api_key = (
- optional_params.api_key
- or litellm.api_key
- or litellm.azure_key
- or get_secret_str("AZURE_OPENAI_API_KEY")
- or get_secret_str("AZURE_API_KEY")
- ) # type: ignore
-
- extra_body = optional_params.get("extra_body", {})
- if extra_body is not None:
- extra_body.pop("azure_ad_token", None)
- else:
- get_secret_str("AZURE_AD_TOKEN") # type: ignore
-
+ azure_creds = get_azure_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ api_version=optional_params.api_version,
+ )
response = azure_files_instance.create_file(
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
- api_version=api_version,
+ api_base=azure_creds.api_base,
+ api_key=azure_creds.api_key,
+ api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
create_file_data=_create_file_request,
litellm_params=litellm_params_dict,
)
- elif custom_llm_provider == "vertex_ai":
- api_base = optional_params.api_base or ""
- vertex_ai_project = (
- optional_params.vertex_project
- or litellm.vertex_project
- or get_secret_str("VERTEXAI_PROJECT")
- )
- vertex_ai_location = (
- optional_params.vertex_location
- or litellm.vertex_location
- or get_secret_str("VERTEXAI_LOCATION")
- )
- vertex_credentials = optional_params.vertex_credentials or get_secret_str(
- "VERTEXAI_CREDENTIALS"
- )
-
- response = vertex_ai_files_instance.create_file(
- _is_async=_is_async,
- api_base=api_base,
- vertex_project=vertex_ai_project,
- vertex_location=vertex_ai_location,
- vertex_credentials=vertex_credentials,
- timeout=timeout,
- max_retries=optional_params.max_retries,
- create_file_data=_create_file_request,
- )
else:
raise litellm.exceptions.BadRequestError(
- message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus'] are supported.".format(
+ message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus', 'anthropic'] are supported.".format(
custom_llm_provider
),
model="n/a",
@@ -295,7 +255,7 @@ def create_file(
@client
async def afile_retrieve(
file_id: str,
- custom_llm_provider: Literal["openai", "azure", "gemini", "hosted_vllm", "manus"] = "openai",
+ custom_llm_provider: FileRetrieveProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -336,7 +296,7 @@ async def afile_retrieve(
@client
def file_retrieve(
file_id: str,
- custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "manus"] = "openai",
+ custom_llm_provider: FileRetrieveProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -367,64 +327,31 @@ def file_retrieve(
_is_async = kwargs.pop("is_async", False) is True
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
- # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
- api_base = (
- optional_params.api_base
- or litellm.api_base
- or os.getenv("OPENAI_BASE_URL")
- or os.getenv("OPENAI_API_BASE")
- or "https://api.openai.com/v1"
+ openai_creds = get_openai_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ organization=optional_params.organization,
)
- organization = (
- optional_params.organization
- or litellm.organization
- or os.getenv("OPENAI_ORGANIZATION", None)
- or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
- )
- # set API KEY
- api_key = (
- optional_params.api_key
- or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
- or litellm.openai_key
- or os.getenv("OPENAI_API_KEY")
- )
-
response = openai_files_instance.retrieve_file(
file_id=file_id,
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
+ api_base=openai_creds.api_base,
+ api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
- organization=organization,
+ organization=openai_creds.organization,
)
elif custom_llm_provider == "azure":
- api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
- api_version = (
- optional_params.api_version
- or litellm.api_version
- or get_secret_str("AZURE_API_VERSION")
- ) # type: ignore
-
- api_key = (
- optional_params.api_key
- or litellm.api_key
- or litellm.azure_key
- or get_secret_str("AZURE_OPENAI_API_KEY")
- or get_secret_str("AZURE_API_KEY")
- ) # type: ignore
-
- extra_body = optional_params.get("extra_body", {})
- if extra_body is not None:
- extra_body.pop("azure_ad_token", None)
- else:
- get_secret_str("AZURE_AD_TOKEN") # type: ignore
-
+ azure_creds = get_azure_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ api_version=optional_params.api_version,
+ )
response = azure_files_instance.retrieve_file(
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
- api_version=api_version,
+ api_base=azure_creds.api_base,
+ api_key=azure_creds.api_key,
+ api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
file_id=file_id,
@@ -439,22 +366,25 @@ def file_retrieve(
litellm_params_dict = get_litellm_params(**kwargs)
litellm_params_dict["api_key"] = optional_params.api_key
litellm_params_dict["api_base"] = optional_params.api_base
-
+
logging_obj = kwargs.get("litellm_logging_obj")
if logging_obj is None:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
+
logging_obj = LiteLLMLoggingObj(
model="",
messages=[],
stream=False,
call_type="afile_retrieve" if _is_async else "file_retrieve",
start_time=time.time(),
- litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())),
+ litellm_call_id=kwargs.get(
+ "litellm_call_id", str(uuid_module.uuid4())
+ ),
function_id=str(kwargs.get("id") or ""),
)
-
+
client = kwargs.get("client")
response = base_llm_http_handler.retrieve_file(
file_id=file_id,
@@ -473,7 +403,7 @@ def file_retrieve(
)
else:
raise litellm.exceptions.BadRequestError(
- message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai', 'azure', and 'manus' are supported.".format(
+ message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
model="n/a",
@@ -494,7 +424,7 @@ def file_retrieve(
@client
async def afile_delete(
file_id: str,
- custom_llm_provider: Literal["openai", "azure", "gemini", "manus"] = "openai",
+ custom_llm_provider: FileDeleteProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -538,7 +468,7 @@ async def afile_delete(
def file_delete(
file_id: str,
model: Optional[str] = None,
- custom_llm_provider: Union[Literal["openai", "azure", "gemini", "manus"], str] = "openai",
+ custom_llm_provider: Union[FileDeleteProvider, str] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -576,63 +506,31 @@ def file_delete(
timeout = 600.0
_is_async = kwargs.pop("is_async", False) is True
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
- # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
- api_base = (
- optional_params.api_base
- or litellm.api_base
- or os.getenv("OPENAI_BASE_URL")
- or os.getenv("OPENAI_API_BASE")
- or "https://api.openai.com/v1"
- )
- organization = (
- optional_params.organization
- or litellm.organization
- or os.getenv("OPENAI_ORGANIZATION", None)
- or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
- )
- # set API KEY
- api_key = (
- optional_params.api_key
- or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
- or litellm.openai_key
- or os.getenv("OPENAI_API_KEY")
+ openai_creds = get_openai_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ organization=optional_params.organization,
)
response = openai_files_instance.delete_file(
file_id=file_id,
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
+ api_base=openai_creds.api_base,
+ api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
- organization=organization,
+ organization=openai_creds.organization,
)
elif custom_llm_provider == "azure":
- api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
- api_version = (
- optional_params.api_version
- or litellm.api_version
- or get_secret_str("AZURE_API_VERSION")
- ) # type: ignore
-
- api_key = (
- optional_params.api_key
- or litellm.api_key
- or litellm.azure_key
- or get_secret_str("AZURE_OPENAI_API_KEY")
- or get_secret_str("AZURE_API_KEY")
- ) # type: ignore
-
- extra_body = optional_params.get("extra_body", {})
- if extra_body is not None:
- extra_body.pop("azure_ad_token", None)
- else:
- get_secret_str("AZURE_AD_TOKEN") # type: ignore
-
+ azure_creds = get_azure_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ api_version=optional_params.api_version,
+ )
response = azure_files_instance.delete_file(
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
- api_version=api_version,
+ api_base=azure_creds.api_base,
+ api_key=azure_creds.api_key,
+ api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
file_id=file_id,
@@ -648,22 +546,25 @@ def file_delete(
if provider_config is not None:
litellm_params_dict["api_key"] = optional_params.api_key
litellm_params_dict["api_base"] = optional_params.api_base
-
+
logging_obj = kwargs.get("litellm_logging_obj")
if logging_obj is None:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
+
logging_obj = LiteLLMLoggingObj(
model="",
messages=[],
stream=False,
call_type="afile_delete" if _is_async else "file_delete",
start_time=time.time(),
- litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())),
+ litellm_call_id=kwargs.get(
+ "litellm_call_id", str(uuid_module.uuid4())
+ ),
function_id=str(kwargs.get("id") or ""),
)
-
+
response = base_llm_http_handler.delete_file(
file_id=file_id,
provider_config=provider_config,
@@ -681,7 +582,7 @@ def file_delete(
)
else:
raise litellm.exceptions.BadRequestError(
- message="LiteLLM doesn't support {} for 'file_delete'. Only 'openai', 'azure', 'gemini', and 'manus' are supported.".format(
+ message="LiteLLM doesn't support {} for 'file_delete'. Only 'openai', 'azure', 'gemini', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
model="n/a",
@@ -700,7 +601,7 @@ def file_delete(
# List files
@client
async def afile_list(
- custom_llm_provider: Literal["openai", "azure", "manus"] = "openai",
+ custom_llm_provider: FileListProvider = "openai",
purpose: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -741,7 +642,7 @@ async def afile_list(
@client
def file_list(
- custom_llm_provider: Literal["openai", "azure", "manus"] = "openai",
+ custom_llm_provider: FileListProvider = "openai",
purpose: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -771,7 +672,7 @@ def file_list(
timeout = 600.0
_is_async = kwargs.pop("is_async", False) is True
-
+
# Check if provider has a custom files config (e.g., Manus, Bedrock, Vertex AI)
provider_config = ProviderConfigManager.get_provider_files_config(
model="",
@@ -781,22 +682,25 @@ def file_list(
litellm_params_dict = get_litellm_params(**kwargs)
litellm_params_dict["api_key"] = optional_params.api_key
litellm_params_dict["api_base"] = optional_params.api_base
-
+
logging_obj = kwargs.get("litellm_logging_obj")
if logging_obj is None:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
+
logging_obj = LiteLLMLoggingObj(
model="",
messages=[],
stream=False,
call_type="afile_list" if _is_async else "file_list",
start_time=time.time(),
- litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())),
+ litellm_call_id=kwargs.get(
+ "litellm_call_id", str(uuid_module.uuid4())
+ ),
function_id=str(kwargs.get("id", "")),
)
-
+
client = kwargs.get("client")
response = base_llm_http_handler.list_files(
purpose=purpose,
@@ -815,71 +719,38 @@ def file_list(
)
return response
elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
- # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
- api_base = (
- optional_params.api_base
- or litellm.api_base
- or os.getenv("OPENAI_BASE_URL")
- or os.getenv("OPENAI_API_BASE")
- or "https://api.openai.com/v1"
+ openai_creds = get_openai_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ organization=optional_params.organization,
)
- organization = (
- optional_params.organization
- or litellm.organization
- or os.getenv("OPENAI_ORGANIZATION", None)
- or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
- )
- # set API KEY
- api_key = (
- optional_params.api_key
- or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
- or litellm.openai_key
- or os.getenv("OPENAI_API_KEY")
- )
-
response = openai_files_instance.list_files(
purpose=purpose,
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
+ api_base=openai_creds.api_base,
+ api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
- organization=organization,
+ organization=openai_creds.organization,
)
elif custom_llm_provider == "azure":
- api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
- api_version = (
- optional_params.api_version
- or litellm.api_version
- or get_secret_str("AZURE_API_VERSION")
- ) # type: ignore
-
- api_key = (
- optional_params.api_key
- or litellm.api_key
- or litellm.azure_key
- or get_secret_str("AZURE_OPENAI_API_KEY")
- or get_secret_str("AZURE_API_KEY")
- ) # type: ignore
-
- extra_body = optional_params.get("extra_body", {})
- if extra_body is not None:
- extra_body.pop("azure_ad_token", None)
- else:
- get_secret_str("AZURE_AD_TOKEN") # type: ignore
-
+ azure_creds = get_azure_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ api_version=optional_params.api_version,
+ )
response = azure_files_instance.list_files(
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
- api_version=api_version,
+ api_base=azure_creds.api_base,
+ api_key=azure_creds.api_key,
+ api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
purpose=purpose,
)
else:
raise litellm.exceptions.BadRequestError(
- message="LiteLLM doesn't support {} for 'file_list'. Only 'openai', 'azure', and 'manus' are supported.".format(
+ message="LiteLLM doesn't support {} for 'file_list'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
model="n/a",
@@ -898,7 +769,7 @@ def file_list(
@client
async def afile_content(
file_id: str,
- custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"] = "openai",
+ custom_llm_provider: FileContentProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -942,9 +813,7 @@ async def afile_content(
def file_content(
file_id: str,
model: Optional[str] = None,
- custom_llm_provider: Optional[
- Union[Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"], str]
- ] = None,
+ custom_llm_provider: Optional[Union[FileContentProvider, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -990,77 +859,72 @@ def file_content(
_is_async = kwargs.pop("afile_content", False) is True
- # Check if this is an Anthropic batch results request
- if custom_llm_provider == "anthropic":
- response = anthropic_files_instance.file_content(
- _is_async=_is_async,
+ # Check if provider has a custom files config (e.g., Anthropic, Manus)
+ provider_config = ProviderConfigManager.get_provider_files_config(
+ model="",
+ provider=LlmProviders(custom_llm_provider),
+ )
+ if provider_config is not None:
+ litellm_params_dict["api_key"] = optional_params.api_key
+ litellm_params_dict["api_base"] = optional_params.api_base
+
+ logging_obj = kwargs.get("litellm_logging_obj")
+ if logging_obj is None:
+ logging_obj = LiteLLMLoggingObj(
+ model="",
+ messages=[],
+ stream=False,
+ call_type="afile_content" if _is_async else "file_content",
+ start_time=time.time(),
+ litellm_call_id=kwargs.get(
+ "litellm_call_id", str(uuid_module.uuid4())
+ ),
+ function_id=str(kwargs.get("id") or ""),
+ )
+
+ response = base_llm_http_handler.retrieve_file_content(
file_content_request=_file_content_request,
- api_base=optional_params.api_base,
- api_key=optional_params.api_key,
+ provider_config=provider_config,
+ litellm_params=litellm_params_dict,
+ headers=extra_headers or {},
+ logging_obj=logging_obj,
+ _is_async=_is_async,
+ client=(
+ client
+ if client is not None
+ and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
+ else None
+ ),
timeout=timeout,
- max_retries=optional_params.max_retries,
)
return response
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
- # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
- api_base = (
- optional_params.api_base
- or litellm.api_base
- or os.getenv("OPENAI_BASE_URL")
- or os.getenv("OPENAI_API_BASE")
- or "https://api.openai.com/v1"
+ openai_creds = get_openai_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ organization=optional_params.organization,
)
- organization = (
- optional_params.organization
- or litellm.organization
- or os.getenv("OPENAI_ORGANIZATION", None)
- or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
- )
- # set API KEY
- api_key = (
- optional_params.api_key
- or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
- or litellm.openai_key
- or os.getenv("OPENAI_API_KEY")
- )
-
response = openai_files_instance.file_content(
_is_async=_is_async,
file_content_request=_file_content_request,
- api_base=api_base,
- api_key=api_key,
+ api_base=openai_creds.api_base,
+ api_key=openai_creds.api_key,
timeout=timeout,
max_retries=optional_params.max_retries,
- organization=organization,
+ organization=openai_creds.organization,
)
elif custom_llm_provider == "azure":
- api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
- api_version = (
- optional_params.api_version
- or litellm.api_version
- or get_secret_str("AZURE_API_VERSION")
- ) # type: ignore
-
- api_key = (
- optional_params.api_key
- or litellm.api_key
- or litellm.azure_key
- or get_secret_str("AZURE_OPENAI_API_KEY")
- or get_secret_str("AZURE_API_KEY")
- ) # type: ignore
-
- extra_body = optional_params.get("extra_body", {})
- if extra_body is not None:
- extra_body.pop("azure_ad_token", None)
- else:
- get_secret_str("AZURE_AD_TOKEN") # type: ignore
-
+ azure_creds = get_azure_credentials(
+ api_base=optional_params.api_base,
+ api_key=optional_params.api_key,
+ api_version=optional_params.api_version,
+ )
response = azure_files_instance.file_content(
_is_async=_is_async,
- api_base=api_base,
- api_key=api_key,
- api_version=api_version,
+ api_base=azure_creds.api_base,
+ api_key=azure_creds.api_key,
+ api_version=azure_creds.api_version,
timeout=timeout,
max_retries=optional_params.max_retries,
file_content_request=_file_content_request,
@@ -1104,7 +968,7 @@ def file_content(
)
else:
raise litellm.exceptions.BadRequestError(
- message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus'.".format(
+ message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.".format(
custom_llm_provider
),
model="n/a",
diff --git a/litellm/files/utils.py b/litellm/files/utils.py
index a56a29467d9..a2b9a42c154 100644
--- a/litellm/files/utils.py
+++ b/litellm/files/utils.py
@@ -8,17 +8,22 @@ class FilesAPIUtils:
"""
Utils for files API interface on litellm
"""
+
@staticmethod
- def is_batch_jsonl_file(create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData) -> bool:
+ def is_batch_jsonl_file(
+ create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData
+ ) -> bool:
"""
Check if the file is a batch jsonl file
"""
return (
create_file_data.get("purpose") == "batch"
- and FilesAPIUtils.valid_content_type(extracted_file_data.get("content_type"))
+ and FilesAPIUtils.valid_content_type(
+ extracted_file_data.get("content_type")
+ )
and extracted_file_data.get("content") is not None
)
-
+
@staticmethod
def valid_content_type(content_type: Optional[str]) -> bool:
"""
diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py
index f5b8b097026..08373cda782 100644
--- a/litellm/fine_tuning/main.py
+++ b/litellm/fine_tuning/main.py
@@ -34,6 +34,44 @@ vertex_fine_tuning_apis_instance = VertexFineTuningAPI()
#################################################
+def _prepare_azure_extra_body(
+ extra_body: Optional[Dict[str, Any]],
+ kwargs: Dict[str, Any],
+ azure_specific_hyperparams: Dict[str, Any],
+) -> Dict[str, Any]:
+ """
+ Prepare extra_body for Azure fine-tuning API by combining Azure-specific parameters.
+
+ Azure fine-tuning API accepts additional parameters beyond the standard OpenAI spec:
+ - trainingType: Type of training (e.g., 1 for supervised fine-tuning)
+ - prompt_loss_weight: Weight for prompt loss in training
+
+ These parameters must be passed in the extra_body field when calling the Azure OpenAI SDK.
+
+ Args:
+ extra_body: Optional existing extra_body dict
+ kwargs: Request kwargs that may contain Azure-specific parameters
+ azure_specific_hyperparams: Dict of Azure-specific hyperparameters already extracted
+
+ Returns:
+ Dict containing all Azure-specific parameters to be passed in extra_body
+ """
+ if extra_body is None:
+ extra_body = {}
+
+ # Azure-specific root-level parameters
+ azure_specific_params = ["trainingType"]
+ for param in azure_specific_params:
+ if param in kwargs:
+ extra_body[param] = kwargs[param]
+
+ # Add Azure-specific hyperparameters
+ if azure_specific_hyperparams:
+ extra_body.update(azure_specific_hyperparams)
+
+ return extra_body
+
+
@client
async def acreate_fine_tuning_job(
model: str,
@@ -88,6 +126,33 @@ async def acreate_fine_tuning_job(
raise e
+def _build_fine_tuning_job_data(
+ model, training_file, hyperparameters, suffix, validation_file, integrations, seed
+):
+ return FineTuningJobCreate(
+ model=model,
+ training_file=training_file,
+ hyperparameters=hyperparameters,
+ suffix=suffix,
+ validation_file=validation_file,
+ integrations=integrations,
+ seed=seed,
+ )
+
+
+def _resolve_fine_tuning_timeout(
+ timeout: Any,
+ custom_llm_provider: str,
+) -> Union[float, httpx.Timeout]:
+ """Normalise a raw timeout value to a float (seconds) or httpx.Timeout for fine-tuning calls."""
+ timeout = timeout or 600.0
+ if isinstance(timeout, httpx.Timeout):
+ if not supports_httpx_timeout(custom_llm_provider):
+ return float(timeout.read or 600)
+ return timeout
+ return float(timeout)
+
+
@client
def create_fine_tuning_job(
model: str,
@@ -114,24 +179,22 @@ def create_fine_tuning_job(
# handle hyperparameters
hyperparameters = hyperparameters or {} # original hyperparameters
+
+ # For Azure, extract Azure-specific hyperparameters before creating OpenAI-spec hyperparameters
+ azure_specific_hyperparams = {}
+ if custom_llm_provider == "azure":
+ azure_hyperparameter_keys = ["prompt_loss_weight"]
+ for key in azure_hyperparameter_keys:
+ if key in hyperparameters:
+ azure_specific_hyperparams[key] = hyperparameters.pop(key)
+
_oai_hyperparameters: Hyperparameters = Hyperparameters(
**hyperparameters
) # Typed Hyperparameters for OpenAI Spec
- ### TIMEOUT LOGIC ###
- timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
- # set timeout for 10 minutes by default
-
- if (
- timeout is not None
- and isinstance(timeout, httpx.Timeout)
- and supports_httpx_timeout(custom_llm_provider) is False
- ):
- read_timeout = timeout.read or 600
- timeout = read_timeout # default 10 min timeout
- elif timeout is not None and not isinstance(timeout, httpx.Timeout):
- timeout = float(timeout) # type: ignore
- elif timeout is None:
- timeout = 600.0
+ timeout = _resolve_fine_tuning_timeout(
+ optional_params.timeout or kwargs.get("request_timeout", 600),
+ custom_llm_provider,
+ )
# OpenAI
if custom_llm_provider == "openai":
@@ -157,19 +220,15 @@ def create_fine_tuning_job(
or os.getenv("OPENAI_API_KEY")
)
- create_fine_tuning_job_data = FineTuningJobCreate(
- model=model,
- training_file=training_file,
- hyperparameters=_oai_hyperparameters,
- suffix=suffix,
- validation_file=validation_file,
- integrations=integrations,
- seed=seed,
- )
-
- create_fine_tuning_job_data_dict = create_fine_tuning_job_data.model_dump(
- exclude_none=True
- )
+ create_fine_tuning_job_data_dict = _build_fine_tuning_job_data(
+ model,
+ training_file,
+ _oai_hyperparameters,
+ suffix,
+ validation_file,
+ integrations,
+ seed,
+ ).model_dump(exclude_none=True)
response = openai_fine_tuning_apis_instance.create_fine_tuning_job(
api_base=api_base,
@@ -207,19 +266,25 @@ def create_fine_tuning_job(
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
- create_fine_tuning_job_data = FineTuningJobCreate(
- model=model,
- training_file=training_file,
- hyperparameters=_oai_hyperparameters,
- suffix=suffix,
- validation_file=validation_file,
- integrations=integrations,
- seed=seed,
+
+ # Prepare Azure-specific parameters for extra_body
+ extra_body = _prepare_azure_extra_body(
+ extra_body, kwargs, azure_specific_hyperparams
)
- create_fine_tuning_job_data_dict = create_fine_tuning_job_data.model_dump(
- exclude_none=True
- )
+ create_fine_tuning_job_data_dict = _build_fine_tuning_job_data(
+ model,
+ training_file,
+ _oai_hyperparameters,
+ suffix,
+ validation_file,
+ integrations,
+ seed,
+ ).model_dump(exclude_none=True)
+
+ # Add extra_body if it has Azure-specific parameters
+ if extra_body:
+ create_fine_tuning_job_data_dict["extra_body"] = extra_body
response = azure_fine_tuning_apis_instance.create_fine_tuning_job(
api_base=api_base,
@@ -246,18 +311,17 @@ def create_fine_tuning_job(
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
"VERTEXAI_CREDENTIALS"
)
- create_fine_tuning_job_data = FineTuningJobCreate(
- model=model,
- training_file=training_file,
- hyperparameters=_oai_hyperparameters,
- suffix=suffix,
- validation_file=validation_file,
- integrations=integrations,
- seed=seed,
- )
response = vertex_fine_tuning_apis_instance.create_fine_tuning_job(
_is_async=_is_async,
- create_fine_tuning_job_data=create_fine_tuning_job_data,
+ create_fine_tuning_job_data=_build_fine_tuning_job_data(
+ model,
+ training_file,
+ _oai_hyperparameters,
+ suffix,
+ validation_file,
+ integrations,
+ seed,
+ ),
vertex_credentials=vertex_credentials,
vertex_project=vertex_ai_project,
vertex_location=vertex_ai_location,
diff --git a/litellm/google_genai/__init__.py b/litellm/google_genai/__init__.py
index faeb1f227d1..ca7b547c440 100644
--- a/litellm/google_genai/__init__.py
+++ b/litellm/google_genai/__init__.py
@@ -13,7 +13,7 @@ from .main import (
__all__ = [
"generate_content",
- "agenerate_content",
+ "agenerate_content",
"generate_content_stream",
"agenerate_content_stream",
-]
\ No newline at end of file
+]
diff --git a/litellm/google_genai/adapters/__init__.py b/litellm/google_genai/adapters/__init__.py
index 96ff777ebe8..bfa9e712678 100644
--- a/litellm/google_genai/adapters/__init__.py
+++ b/litellm/google_genai/adapters/__init__.py
@@ -13,7 +13,7 @@ from .handler import GenerateContentToCompletionHandler
from .transformation import GoogleGenAIAdapter, GoogleGenAIStreamWrapper
__all__ = [
- "GoogleGenAIAdapter",
+ "GoogleGenAIAdapter",
"GoogleGenAIStreamWrapper",
- "GenerateContentToCompletionHandler"
-]
\ No newline at end of file
+ "GenerateContentToCompletionHandler",
+]
diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py
index 0a296012210..c5d9fd124fa 100644
--- a/litellm/google_genai/adapters/transformation.py
+++ b/litellm/google_genai/adapters/transformation.py
@@ -770,8 +770,6 @@ class GoogleGenAIAdapter:
"content_filter": "SAFETY",
"tool_calls": "STOP",
"function_call": "STOP",
- "finish_reason_unspecified": "FINISH_REASON_UNSPECIFIED",
- "malformed_function_call": "MALFORMED_FUNCTION_CALL",
}
return mapping.get(finish_reason, "STOP")
diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py
index 9ec56c37170..bdbb483dcf6 100644
--- a/litellm/google_genai/main.py
+++ b/litellm/google_genai/main.py
@@ -168,7 +168,9 @@ class GenerateContentHelper:
)
)
# Extract systemInstruction from kwargs to pass to transform
- system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
+ system_instruction = kwargs.get("systemInstruction") or kwargs.get(
+ "system_instruction"
+ )
request_body = (
generate_content_provider_config.transform_generate_content_request(
model=model,
@@ -183,7 +185,8 @@ class GenerateContentHelper:
if litellm_logging_obj is None:
raise ValueError("litellm_logging_obj is required, but got None")
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model=model,
optional_params=dict(generate_content_config_dict),
litellm_params={
@@ -318,7 +321,9 @@ def generate_content(
)
# Extract systemInstruction from kwargs to pass to handler
- system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
+ system_instruction = kwargs.get("systemInstruction") or kwargs.get(
+ "system_instruction"
+ )
# Check if we should use the adapter (when provider config is None)
if setup_result.generate_content_provider_config is None:
@@ -407,7 +412,9 @@ async def agenerate_content_stream(
)
# Extract systemInstruction from kwargs to pass to handler
- system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
+ system_instruction = kwargs.get("systemInstruction") or kwargs.get(
+ "system_instruction"
+ )
# Check if we should use the adapter (when provider config is None)
if setup_result.generate_content_provider_config is None:
diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py
index d0fa5a0be6c..8cb2ee09370 100644
--- a/litellm/google_genai/streaming_iterator.py
+++ b/litellm/google_genai/streaming_iterator.py
@@ -17,6 +17,7 @@ else:
GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging()
+
class BaseGoogleGenAIGenerateContentStreamingIterator:
"""
Base class for Google GenAI Generate Content streaming iterators that provides common logic
@@ -42,6 +43,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
from litellm.proxy.pass_through_endpoints.streaming_handler import (
PassThroughStreamingHandler,
)
+
end_time = datetime.now()
asyncio.create_task(
PassThroughStreamingHandler._route_streaming_logging_to_handler(
@@ -58,7 +60,9 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
)
-class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator):
+class GoogleGenAIGenerateContentStreamingIterator(
+ BaseGoogleGenAIGenerateContentStreamingIterator
+):
"""
Streaming iterator specifically for Google GenAI generate content API.
"""
@@ -105,10 +109,14 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
async def __anext__(self):
# This should not be used for sync responses
# If you need async iteration, use AsyncGoogleGenAIGenerateContentStreamingIterator
- raise NotImplementedError("Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration")
+ raise NotImplementedError(
+ "Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration"
+ )
-class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator):
+class AsyncGoogleGenAIGenerateContentStreamingIterator(
+ BaseGoogleGenAIGenerateContentStreamingIterator
+):
"""
Async streaming iterator specifically for Google GenAI generate content API.
"""
@@ -148,4 +156,4 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
return chunk
except StopAsyncIteration:
await self._handle_async_streaming_logging()
- raise StopAsyncIteration
\ No newline at end of file
+ raise StopAsyncIteration
diff --git a/litellm/images/main.py b/litellm/images/main.py
index 6c4c502a7b0..a5ae154190a 100644
--- a/litellm/images/main.py
+++ b/litellm/images/main.py
@@ -40,6 +40,9 @@ from litellm.utils import exception_type, get_litellm_params
llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler()
from openai.types.audio.transcription_create_params import FileTypes # type: ignore
+# BFL handlers
+from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
+from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
from litellm.main import (
azure_chat_completions,
base_llm_aiohttp_handler,
@@ -82,7 +85,6 @@ def _get_ImageEditRequestUtils() -> "ImageEditRequestUtils":
return _ImageEditRequestUtils_cache
-
##### Image Generation #######################
@client
async def aimage_generation(*args, **kwargs) -> ImageResponse:
@@ -208,10 +210,7 @@ def image_generation( # noqa: PLR0915
api_version: Optional[str] = None,
custom_llm_provider=None,
**kwargs,
-) -> Union[
- ImageResponse,
- Coroutine[Any, Any, ImageResponse],
-]:
+) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse],]:
"""
Maps the https://api.openai.com/v1/images/generations endpoint.
@@ -297,7 +296,8 @@ def image_generation( # noqa: PLR0915
litellm_params_dict = get_litellm_params(**kwargs)
logging: Logging = litellm_logging_obj
- logging.update_environment_variables(
+ logging.update_from_kwargs(
+ kwargs=kwargs,
model=model,
user=user,
optional_params=optional_params,
@@ -308,7 +308,6 @@ def image_generation( # noqa: PLR0915
"logger_fn": logger_fn,
"proxy_server_request": proxy_server_request,
"model_info": model_info,
- "metadata": metadata,
"preset_cache_key": None,
"stream_response": {},
},
@@ -342,7 +341,7 @@ def image_generation( # noqa: PLR0915
azure_ad_token = optional_params.pop(
"azure_ad_token", None
) or get_secret_str("AZURE_AD_TOKEN")
-
+
# Create azure_ad_token_provider from tenant_id, client_id, client_secret if not already provided
if azure_ad_token_provider is None:
from litellm.llms.azure.common_utils import (
@@ -353,8 +352,11 @@ def image_generation( # noqa: PLR0915
tenant_id = litellm_params_dict.get("tenant_id")
client_id = litellm_params_dict.get("client_id")
client_secret = litellm_params_dict.get("client_secret")
- azure_scope = litellm_params_dict.get("azure_scope") or "https://cognitiveservices.azure.com/.default"
-
+ azure_scope = (
+ litellm_params_dict.get("azure_scope")
+ or "https://cognitiveservices.azure.com/.default"
+ )
+
# Create token provider if credentials are available
if tenant_id and client_id and client_secret:
azure_ad_token_provider = get_azure_ad_token_from_entra_id(
@@ -371,7 +373,7 @@ def image_generation( # noqa: PLR0915
# Azure AD authentication will use Authorization header instead
if api_key is not None:
default_headers["api-key"] = api_key
-
+
for k, v in default_headers.items():
if k not in headers:
headers[k] = v
@@ -404,7 +406,7 @@ def image_generation( # noqa: PLR0915
litellm.LlmProviders.STABILITY,
litellm.LlmProviders.RUNWAYML,
litellm.LlmProviders.VERTEX_AI,
- litellm.LlmProviders.OPENROUTER
+ litellm.LlmProviders.OPENROUTER,
):
if image_generation_config is None:
raise ValueError(
@@ -427,6 +429,22 @@ def image_generation( # noqa: PLR0915
timeout=timeout,
client=client,
)
+ elif custom_llm_provider == "black_forest_labs":
+ # Route to BFL-specific handler (polling required)
+ if model is None:
+ raise Exception("Model needs to be set for black_forest_labs")
+ return bfl_image_generation.image_generation(
+ model=model,
+ prompt=prompt,
+ model_response=model_response,
+ optional_params=optional_params,
+ litellm_params=litellm_params_dict,
+ logging_obj=litellm_logging_obj,
+ timeout=timeout,
+ extra_headers=extra_headers,
+ client=client,
+ aimg_generation=aimg_generation,
+ )
elif custom_llm_provider == "azure_ai":
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
@@ -442,7 +460,7 @@ def image_generation( # noqa: PLR0915
# Azure AD authentication will use Authorization header instead
if api_key is not None:
default_headers["api-key"] = api_key
-
+
for k, v in default_headers.items():
if k not in headers:
headers[k] = v
@@ -469,6 +487,8 @@ def image_generation( # noqa: PLR0915
or custom_llm_provider == LlmProviders.LITELLM_PROXY.value
or custom_llm_provider in litellm.openai_compatible_providers
):
+ if extra_headers is not None:
+ optional_params["extra_headers"] = extra_headers
# Forward OpenAI organization if present (set by proxy pre-call utils)
organization: Optional[str] = kwargs.get("organization", None)
model_response = openai_chat_completions.image_generation(
@@ -483,6 +503,7 @@ def image_generation( # noqa: PLR0915
organization=organization,
aimg_generation=aimg_generation,
client=client,
+ headers=headers,
)
elif custom_llm_provider == "bedrock":
if model is None:
@@ -715,7 +736,7 @@ def image_variation(
@client
def image_edit( # noqa: PLR0915
image: Optional[Union[FileTypes, List[FileTypes]]] = None,
- prompt: Optional[str]= None,
+ prompt: Optional[str] = None,
model: Optional[str] = None,
mask: Optional[str] = None,
n: Optional[int] = None,
@@ -739,23 +760,23 @@ def image_edit( # noqa: PLR0915
local_vars = locals()
try:
openai_params = [
- "user",
- "request_timeout",
- "api_base",
- "api_version",
- "api_key",
- "deployment_id",
- "organization",
- "base_url",
- "default_headers",
- "timeout",
- "max_retries",
- "n",
- "quality",
- "size",
- "style",
- "async_call",
- ]
+ "user",
+ "request_timeout",
+ "api_base",
+ "api_version",
+ "api_key",
+ "deployment_id",
+ "organization",
+ "base_url",
+ "default_headers",
+ "timeout",
+ "max_retries",
+ "n",
+ "quality",
+ "size",
+ "style",
+ "async_call",
+ ]
litellm_params_list = all_litellm_params
default_params = openai_params + litellm_params_list
non_default_params = {
@@ -763,10 +784,14 @@ def image_edit( # noqa: PLR0915
} # model-specific params - pass them straight to the model/provider
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
+ model_info = kwargs.get("model_info", None)
+ metadata = kwargs.get("metadata", {})
_is_async = kwargs.pop("async_call", False) is True
# add images / or return a single image
- images = image if isinstance(image, list) else ([image] if image is not None else [])
+ images = (
+ image if isinstance(image, list) else ([image] if image is not None else [])
+ )
headers_from_kwargs = kwargs.get("headers")
merged_extra_headers: Dict[str, Any] = {}
@@ -839,11 +864,11 @@ def image_edit( # noqa: PLR0915
)
# get provider config
- image_edit_provider_config: Optional[BaseImageEditConfig] = (
- ProviderConfigManager.get_provider_image_edit_config(
- model=model,
- provider=litellm.LlmProviders(custom_llm_provider),
- )
+ image_edit_provider_config: Optional[
+ BaseImageEditConfig
+ ] = ProviderConfigManager.get_provider_image_edit_config(
+ model=model,
+ provider=litellm.LlmProviders(custom_llm_provider),
)
if image_edit_provider_config is None:
@@ -852,7 +877,9 @@ def image_edit( # noqa: PLR0915
local_vars.update(kwargs)
# Get ImageEditOptionalRequestParams with only valid parameters
image_edit_optional_params: ImageEditOptionalRequestParams = (
- _get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars)
+ _get_ImageEditRequestUtils().get_requested_image_edit_optional_param(
+ local_vars
+ )
)
# Get optional parameters for the responses API
image_edit_request_params: Dict = (
@@ -866,13 +893,15 @@ def image_edit( # noqa: PLR0915
)
# Pre Call logging
- litellm_logging_obj.update_environment_variables(
+ litellm_logging_obj.update_from_kwargs(
+ kwargs=kwargs,
model=model,
user=user,
optional_params=dict(image_edit_request_params),
litellm_params={
- "litellm_call_id": litellm_call_id,
**image_edit_request_params,
+ "litellm_call_id": litellm_call_id,
+ "model_info": model_info,
},
custom_llm_provider=custom_llm_provider,
)
@@ -899,20 +928,37 @@ def image_edit( # noqa: PLR0915
elif custom_llm_provider == "stability":
image_edit_request_params.update(non_default_params)
return base_llm_http_handler.image_edit_handler(
- model=model,
- image=images,
- prompt=prompt,
- image_edit_provider_config=image_edit_provider_config,
- image_edit_optional_request_params=image_edit_request_params,
- custom_llm_provider=custom_llm_provider,
- litellm_params=litellm_params,
- logging_obj=litellm_logging_obj,
- extra_headers=extra_headers,
- extra_body=extra_body,
- timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
- _is_async=_is_async,
- client=kwargs.get("client"),
- )
+ model=model,
+ image=images,
+ prompt=prompt,
+ image_edit_provider_config=image_edit_provider_config,
+ image_edit_optional_request_params=image_edit_request_params,
+ custom_llm_provider=custom_llm_provider,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ extra_headers=extra_headers,
+ extra_body=extra_body,
+ timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
+ _is_async=_is_async,
+ client=kwargs.get("client"),
+ )
+ elif custom_llm_provider == "black_forest_labs":
+ # Route to BFL-specific handler (polling required)
+ if model is None:
+ raise Exception("Model needs to be set for black_forest_labs")
+ image_edit_request_params.update(non_default_params)
+ return bfl_image_edit.image_edit(
+ model=model,
+ image=images,
+ prompt=prompt,
+ image_edit_optional_request_params=image_edit_request_params,
+ litellm_params=litellm_params,
+ logging_obj=litellm_logging_obj,
+ timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
+ extra_headers=extra_headers,
+ client=kwargs.get("client"),
+ aimage_edit=_is_async,
+ )
# Call the handler with _is_async flag instead of directly calling the async handler
return base_llm_http_handler.image_edit_handler(
model=model,
diff --git a/litellm/images/utils.py b/litellm/images/utils.py
index fa271b61b6a..8d3e96f1433 100644
--- a/litellm/images/utils.py
+++ b/litellm/images/utils.py
@@ -40,9 +40,7 @@ class ImageEditRequestUtils:
filtered_optional_params.pop(param, None)
unsupported_params = [
- param
- for param in filtered_optional_params
- if param not in supported_params
+ param for param in filtered_optional_params if param not in supported_params
]
if unsupported_params:
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/hanging_request_check.py b/litellm/integrations/SlackAlerting/hanging_request_check.py
index 713e790ba90..b9c485dce82 100644
--- a/litellm/integrations/SlackAlerting/hanging_request_check.py
+++ b/litellm/integrations/SlackAlerting/hanging_request_check.py
@@ -102,10 +102,10 @@ class AlertingHangingRequestCheck:
)
for request_id in hanging_requests:
- hanging_request_data: Optional[HangingRequestData] = (
- await self.hanging_request_cache.async_get_cache(
- key=request_id,
- )
+ hanging_request_data: Optional[
+ HangingRequestData
+ ] = await self.hanging_request_cache.async_get_cache(
+ key=request_id,
)
if hanging_request_data is None:
@@ -172,4 +172,6 @@ Team Alias: `{hanging_request_data.team_alias}`"""
level="Medium",
alert_type=AlertType.llm_requests_hanging,
alerting_metadata=hanging_request_data.alerting_metadata or {},
+ request_model=hanging_request_data.model,
+ api_base=hanging_request_data.api_base,
)
diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py
index 8fb3e132ded..013cef74805 100644
--- a/litellm/integrations/SlackAlerting/slack_alerting.py
+++ b/litellm/integrations/SlackAlerting/slack_alerting.py
@@ -70,6 +70,7 @@ class SlackAlerting(CustomBatchLogger):
] = None, # if user wants to separate alerts to diff channels
alerting_args={},
default_webhook_url: Optional[str] = None,
+ alert_type_config: Optional[Dict[str, dict]] = None,
**kwargs,
):
if alerting_threshold is None:
@@ -92,6 +93,14 @@ class SlackAlerting(CustomBatchLogger):
self.hanging_request_check = AlertingHangingRequestCheck(
slack_alerting_object=self,
)
+ self.alert_type_config: Dict[str, AlertTypeConfig] = {}
+ if alert_type_config:
+ for key, val in alert_type_config.items():
+ self.alert_type_config[key] = (
+ AlertTypeConfig(**val) if isinstance(val, dict) else val
+ )
+ self.digest_buckets: Dict[str, DigestEntry] = {}
+ self.digest_lock = asyncio.Lock()
super().__init__(**kwargs, flush_lock=self.flush_lock)
def update_values(
@@ -102,6 +111,7 @@ class SlackAlerting(CustomBatchLogger):
alert_to_webhook_url: Optional[Dict[AlertType, Union[List[str], str]]] = None,
alerting_args: Optional[Dict] = None,
llm_router: Optional[Router] = None,
+ alert_type_config: Optional[Dict[str, dict]] = None,
):
if alerting is not None:
self.alerting = alerting
@@ -116,6 +126,11 @@ class SlackAlerting(CustomBatchLogger):
if not self.periodic_started:
asyncio.create_task(self.periodic_flush())
self.periodic_started = True
+ if alert_type_config is not None:
+ for key, val in alert_type_config.items():
+ self.alert_type_config[key] = (
+ AlertTypeConfig(**val) if isinstance(val, dict) else val
+ )
if alert_to_webhook_url is not None:
# update the dict
@@ -284,6 +299,8 @@ class SlackAlerting(CustomBatchLogger):
level="Low",
alert_type=AlertType.llm_too_slow,
alerting_metadata=alerting_metadata,
+ request_model=model,
+ api_base=api_base,
)
async def async_update_daily_reports(
@@ -538,6 +555,7 @@ class SlackAlerting(CustomBatchLogger):
"organization_budget",
"proxy_budget",
"projected_limit_exceeded",
+ "project_budget",
],
user_info: CallInfo,
):
@@ -1353,13 +1371,15 @@ Model Info:
return False
- async def send_alert(
+ async def send_alert( # noqa: PLR0915
self,
message: str,
level: Literal["Low", "Medium", "High"],
alert_type: AlertType,
alerting_metadata: dict,
user_info: Optional[WebhookEvent] = None,
+ request_model: Optional[str] = None,
+ api_base: Optional[str] = None,
**kwargs,
):
"""
@@ -1375,12 +1395,18 @@ Model Info:
Parameters:
level: str - Low|Medium|High - if calls might fail (Medium) or are failing (High); Currently, no alerts would be 'Low'.
message: str - what is the alert about
+ request_model: Optional[str] - model name for digest grouping
+ api_base: Optional[str] - api base for digest grouping
"""
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
@@ -1408,6 +1434,46 @@ Model Info:
from datetime import datetime
+ # Check if digest mode is enabled for this alert type
+ alert_type_name_str = getattr(alert_type, "value", str(alert_type))
+ _atc = self.alert_type_config.get(alert_type_name_str)
+ if _atc is not None and _atc.digest:
+ # Resolve webhook URL for this alert type (needed for digest entry)
+ if (
+ self.alert_to_webhook_url is not None
+ and alert_type in self.alert_to_webhook_url
+ ):
+ _digest_webhook: Optional[
+ Union[str, List[str]]
+ ] = self.alert_to_webhook_url[alert_type]
+ elif self.default_webhook_url is not None:
+ _digest_webhook = self.default_webhook_url
+ else:
+ _digest_webhook = os.getenv("SLACK_WEBHOOK_URL", None)
+ if _digest_webhook is None:
+ raise ValueError("Missing SLACK_WEBHOOK_URL from environment")
+
+ digest_key = f"{alert_type_name_str}:{request_model or ''}:{api_base or ''}"
+
+ async with self.digest_lock:
+ now = datetime.now()
+ if digest_key in self.digest_buckets:
+ self.digest_buckets[digest_key]["count"] += 1
+ self.digest_buckets[digest_key]["last_time"] = now
+ else:
+ self.digest_buckets[digest_key] = DigestEntry(
+ alert_type=alert_type_name_str,
+ request_model=request_model or "",
+ api_base=api_base or "",
+ first_message=message,
+ level=level,
+ count=1,
+ start_time=now,
+ last_time=now,
+ webhook_url=_digest_webhook,
+ )
+ return # Suppress immediate alert; will be emitted by _flush_digest_buckets
+
# Get the current timestamp
current_time = datetime.now().strftime("%H:%M:%S")
_proxy_base_url = os.getenv("PROXY_BASE_URL", None)
@@ -1483,6 +1549,82 @@ Model Info:
await asyncio.gather(*tasks)
self.log_queue.clear()
+ async def _flush_digest_buckets(self):
+ """Flush any digest buckets whose interval has expired.
+
+ For each expired bucket, formats a digest summary message and
+ appends it to the log_queue for delivery via the normal batching path.
+ """
+ from datetime import datetime
+
+ now = datetime.now()
+ flushed_keys: List[str] = []
+
+ async with self.digest_lock:
+ for key, entry in self.digest_buckets.items():
+ alert_type_name = entry["alert_type"]
+ _atc = self.alert_type_config.get(alert_type_name)
+ if _atc is None:
+ continue
+ elapsed = (now - entry["start_time"]).total_seconds()
+ if elapsed < _atc.digest_interval:
+ continue
+
+ # Build digest summary message
+ start_ts = entry["start_time"].strftime("%H:%M:%S")
+ end_ts = entry["last_time"].strftime("%H:%M:%S")
+ start_date = entry["start_time"].strftime("%Y-%m-%d")
+ end_date = entry["last_time"].strftime("%Y-%m-%d")
+ formatted_message = (
+ f"Alert type: `{alert_type_name}` (Digest)\n"
+ f"Level: `{entry['level']}`\n"
+ f"Start: `{start_date} {start_ts}`\n"
+ f"End: `{end_date} {end_ts}`\n"
+ f"Count: `{entry['count']}`\n\n"
+ f"Message: {entry['first_message']}"
+ )
+ _proxy_base_url = os.getenv("PROXY_BASE_URL", None)
+ if _proxy_base_url is not None:
+ formatted_message += f"\n\nProxy URL: `{_proxy_base_url}`"
+
+ payload = {"text": formatted_message}
+ headers = {"Content-type": "application/json"}
+ webhook_url = entry["webhook_url"]
+
+ if isinstance(webhook_url, list):
+ for url in webhook_url:
+ self.log_queue.append(
+ {
+ "url": url,
+ "headers": headers,
+ "payload": payload,
+ "alert_type": alert_type_name,
+ }
+ )
+ else:
+ self.log_queue.append(
+ {
+ "url": webhook_url,
+ "headers": headers,
+ "payload": payload,
+ "alert_type": alert_type_name,
+ }
+ )
+ flushed_keys.append(key)
+
+ for key in flushed_keys:
+ del self.digest_buckets[key]
+
+ async def periodic_flush(self):
+ """Override base periodic_flush to also flush digest buckets."""
+ while True:
+ await asyncio.sleep(self.flush_interval)
+ try:
+ await self._flush_digest_buckets()
+ except Exception as e:
+ verbose_proxy_logger.debug(f"Error flushing digest buckets: {str(e)}")
+ await self.flush_queue()
+
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""Log deployment latency"""
try:
diff --git a/litellm/integrations/_types/open_inference.py b/litellm/integrations/_types/open_inference.py
index 0fde1ff7525..3404df7495f 100644
--- a/litellm/integrations/_types/open_inference.py
+++ b/litellm/integrations/_types/open_inference.py
@@ -73,11 +73,15 @@ class SpanAttributes:
"""
Number of tokens in the prompt.
"""
- LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = "llm.token_count.prompt_details.cache_write"
+ LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = (
+ "llm.token_count.prompt_details.cache_write"
+ )
"""
Number of tokens in the prompt that were written to cache.
"""
- LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = "llm.token_count.prompt_details.cache_read"
+ LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = (
+ "llm.token_count.prompt_details.cache_read"
+ )
"""
Number of tokens in the prompt that were read from cache.
"""
@@ -89,11 +93,15 @@ class SpanAttributes:
"""
Number of tokens in the completion.
"""
- LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = "llm.token_count.completion_details.reasoning"
+ LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = (
+ "llm.token_count.completion_details.reasoning"
+ )
"""
Number of tokens used for reasoning steps in the completion.
"""
- LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = "llm.token_count.completion_details.audio"
+ LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = (
+ "llm.token_count.completion_details.audio"
+ )
"""
The number of audio input tokens generated by the model
"""
@@ -396,7 +404,7 @@ class OpenInferenceLLMProviderValues(Enum):
class ErrorAttributes:
"""
Attributes for error information in spans.
-
+
These attributes follow OpenTelemetry semantic conventions for exceptions
and are used to record error information from StandardLoggingPayloadErrorInformation.
"""
diff --git a/litellm/integrations/agentops/__init__.py b/litellm/integrations/agentops/__init__.py
index 6ad02ce0ba1..003a12a6112 100644
--- a/litellm/integrations/agentops/__init__.py
+++ b/litellm/integrations/agentops/__init__.py
@@ -1,3 +1,3 @@
from .agentops import AgentOps
-__all__ = ["AgentOps"]
\ No newline at end of file
+__all__ = ["AgentOps"]
diff --git a/litellm/integrations/agentops/agentops.py b/litellm/integrations/agentops/agentops.py
index 11e76841e99..38b91c06587 100644
--- a/litellm/integrations/agentops/agentops.py
+++ b/litellm/integrations/agentops/agentops.py
@@ -7,6 +7,7 @@ from typing import Optional, Dict, Any
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
+
@dataclass
class AgentOpsConfig:
endpoint: str = "https://otlp.agentops.cloud/v1/traces"
@@ -22,9 +23,10 @@ class AgentOpsConfig:
api_key=os.getenv("AGENTOPS_API_KEY"),
service_name=os.getenv("AGENTOPS_SERVICE_NAME", "agentops"),
deployment_environment=os.getenv("AGENTOPS_ENVIRONMENT", "production"),
- auth_endpoint="https://api.agentops.ai/v3/auth/token"
+ auth_endpoint="https://api.agentops.ai/v3/auth/token",
)
+
class AgentOps(OpenTelemetry):
"""
AgentOps integration - built on top of OpenTelemetry
@@ -32,7 +34,7 @@ class AgentOps(OpenTelemetry):
Example usage:
```python
import litellm
-
+
litellm.success_callback = ["agentops"]
response = litellm.completion(
@@ -41,6 +43,7 @@ class AgentOps(OpenTelemetry):
)
```
"""
+
def __init__(
self,
config: Optional[AgentOpsConfig] = None,
@@ -60,18 +63,13 @@ class AgentOps(OpenTelemetry):
pass
headers = f"Authorization=Bearer {jwt_token}" if jwt_token else None
-
+
otel_config = OpenTelemetryConfig(
- exporter="otlp_http",
- endpoint=config.endpoint,
- headers=headers
+ exporter="otlp_http", endpoint=config.endpoint, headers=headers
)
# Initialize OpenTelemetry with our config
- super().__init__(
- config=otel_config,
- callback_name="agentops"
- )
+ super().__init__(config=otel_config, callback_name="agentops")
# Set AgentOps-specific resource attributes
resource_attrs = {
@@ -79,20 +77,20 @@ class AgentOps(OpenTelemetry):
"deployment.environment": config.deployment_environment or "production",
"telemetry.sdk.name": "agentops",
}
-
+
if project_id:
resource_attrs["project.id"] = project_id
-
+
self.resource_attributes = resource_attrs
def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> Dict[str, Any]:
"""
Fetch JWT authentication token from AgentOps API
-
+
Args:
api_key: AgentOps API key
auth_endpoint: Authentication endpoint
-
+
Returns:
Dict containing JWT token and project ID
"""
@@ -100,19 +98,19 @@ class AgentOps(OpenTelemetry):
"Content-Type": "application/json",
"Connection": "keep-alive",
}
-
+
client = _get_httpx_client()
try:
response = client.post(
url=auth_endpoint,
headers=headers,
json={"api_key": api_key},
- timeout=10
+ timeout=10,
)
-
+
if response.status_code != 200:
raise Exception(f"Failed to fetch auth token: {response.text}")
-
+
return response.json()
finally:
- client.close()
\ No newline at end of file
+ client.close()
diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py
index 5df79580d3e..8e4d40c460e 100644
--- a/litellm/integrations/anthropic_cache_control_hook.py
+++ b/litellm/integrations/anthropic_cache_control_hook.py
@@ -82,8 +82,10 @@ class AnthropicCacheControlHook(CustomPromptManagement):
_targetted_index: Optional[Union[int, str]] = point.get("index", None)
targetted_index: Optional[int] = None
if isinstance(_targetted_index, str):
- if _targetted_index.isdigit():
+ try:
targetted_index = int(_targetted_index)
+ except ValueError:
+ pass
else:
targetted_index = _targetted_index
@@ -97,10 +99,10 @@ class AnthropicCacheControlHook(CustomPromptManagement):
targetted_index += len(messages)
if 0 <= targetted_index < len(messages):
- messages[targetted_index] = (
- AnthropicCacheControlHook._safe_insert_cache_control_in_message(
- messages[targetted_index], control
- )
+ messages[
+ targetted_index
+ ] = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
+ messages[targetted_index], control
)
else:
verbose_logger.warning(
diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py
index b75e296be47..8dfaa8b1425 100644
--- a/litellm/integrations/arize/_utils.py
+++ b/litellm/integrations/arize/_utils.py
@@ -14,12 +14,12 @@ from litellm.types.utils import StandardLoggingPayload
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm.integrations._types.open_inference import (
- MessageAttributes,
- ImageAttributes,
- SpanAttributes,
- AudioAttributes,
- EmbeddingAttributes,
- OpenInferenceSpanKindValues
+ MessageAttributes,
+ ImageAttributes,
+ SpanAttributes,
+ AudioAttributes,
+ EmbeddingAttributes,
+ OpenInferenceSpanKindValues,
)
@@ -158,7 +158,9 @@ def _set_audio_outputs(span: "Span", response_obj, audio_attrs, span_attrs):
audio_transcript = audio_item.get("transcript")
if audio_transcript:
- safe_set_attribute(span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript)
+ safe_set_attribute(
+ span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript
+ )
def _set_embedding_outputs(span: "Span", response_obj, embedding_attrs, span_attrs):
@@ -212,7 +214,9 @@ def _set_structured_outputs(span: "Span", response_obj, msg_attrs, span_attrs):
message_content = getattr(first_content, "text", "")
message_role = getattr(item, "role", "assistant")
safe_set_attribute(span, span_attrs.OUTPUT_VALUE, message_content)
- safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content)
+ safe_set_attribute(
+ span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content
+ )
safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_ROLE}", message_role)
@@ -221,16 +225,24 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs):
if not usage:
return
- safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens"))
+ safe_set_attribute(
+ span, span_attrs.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens")
+ )
completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens")
if completion_tokens:
- safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens)
+ safe_set_attribute(
+ span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens
+ )
prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens")
if prompt_tokens:
safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_PROMPT, prompt_tokens)
reasoning_tokens = usage.get("output_tokens_details", {}).get("reasoning_tokens")
if reasoning_tokens:
- safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, reasoning_tokens)
+ safe_set_attribute(
+ span,
+ span_attrs.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING,
+ reasoning_tokens,
+ )
def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
@@ -281,11 +293,15 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
):
return OpenInferenceSpanKindValues.LLM.value
- if any(keyword in lowered for keyword in ("file", "batch", "container", "fine_tuning_job")):
+ if any(
+ keyword in lowered
+ for keyword in ("file", "batch", "container", "fine_tuning_job")
+ ):
return OpenInferenceSpanKindValues.CHAIN.value
return OpenInferenceSpanKindValues.UNKNOWN.value
+
def _set_tool_attributes(
span: "Span", optional_tools: Optional[list], metadata_tools: Optional[list]
):
@@ -294,18 +310,30 @@ def _set_tool_attributes(
for idx, tool in enumerate(optional_tools):
if not isinstance(tool, dict):
continue
- function = tool.get("function") if isinstance(tool.get("function"), dict) else None
+ function = (
+ tool.get("function") if isinstance(tool.get("function"), dict) else None
+ )
if not function:
continue
tool_name = function.get("name")
if tool_name:
- safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.name", tool_name)
+ safe_set_attribute(
+ span, f"{SpanAttributes.LLM_TOOLS}.{idx}.name", tool_name
+ )
tool_description = function.get("description")
if tool_description:
- safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.description", tool_description)
+ safe_set_attribute(
+ span,
+ f"{SpanAttributes.LLM_TOOLS}.{idx}.description",
+ tool_description,
+ )
params = function.get("parameters")
if params is not None:
- safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.parameters", json.dumps(params))
+ safe_set_attribute(
+ span,
+ f"{SpanAttributes.LLM_TOOLS}.{idx}.parameters",
+ json.dumps(params),
+ )
if metadata_tools and isinstance(metadata_tools, list):
for idx, tool in enumerate(metadata_tools):
@@ -343,7 +371,11 @@ def set_attributes(
if standard_logging_payload is None:
raise ValueError("standard_logging_object not found in kwargs")
- metadata = standard_logging_payload.get("metadata") if standard_logging_payload else None
+ metadata = (
+ standard_logging_payload.get("metadata")
+ if standard_logging_payload
+ else None
+ )
_set_metadata_attributes(span, metadata, SpanAttributes)
metadata_tools = _extract_metadata_tools(metadata)
@@ -362,13 +394,19 @@ def set_attributes(
span_kind = _infer_open_inference_span_kind(call_type=call_type)
_set_tool_attributes(span, optional_tools, metadata_tools)
- if (optional_tools or metadata_tools) and span_kind != OpenInferenceSpanKindValues.TOOL.value:
+ if (
+ optional_tools or metadata_tools
+ ) and span_kind != OpenInferenceSpanKindValues.TOOL.value:
span_kind = OpenInferenceSpanKindValues.TOOL.value
safe_set_attribute(span, SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind)
attributes.set_messages(span, kwargs)
- model_params = standard_logging_payload.get("model_parameters") if standard_logging_payload else None
+ model_params = (
+ standard_logging_payload.get("model_parameters")
+ if standard_logging_payload
+ else None
+ )
_set_model_params(span, model_params, SpanAttributes)
_set_response_attributes(span=span, response_obj=response_obj)
@@ -418,17 +456,29 @@ def _set_request_attributes(
if kwargs.get("model"):
safe_set_attribute(span, span_attrs.LLM_MODEL_NAME, kwargs.get("model"))
- safe_set_attribute(span, "llm.request.type", standard_logging_payload.get("call_type"))
- safe_set_attribute(span, span_attrs.LLM_PROVIDER, litellm_params.get("custom_llm_provider", "Unknown"))
+ safe_set_attribute(
+ span, "llm.request.type", standard_logging_payload.get("call_type")
+ )
+ safe_set_attribute(
+ span,
+ span_attrs.LLM_PROVIDER,
+ litellm_params.get("custom_llm_provider", "Unknown"),
+ )
if optional_params.get("max_tokens"):
- safe_set_attribute(span, "llm.request.max_tokens", optional_params.get("max_tokens"))
+ safe_set_attribute(
+ span, "llm.request.max_tokens", optional_params.get("max_tokens")
+ )
if optional_params.get("temperature"):
- safe_set_attribute(span, "llm.request.temperature", optional_params.get("temperature"))
+ safe_set_attribute(
+ span, "llm.request.temperature", optional_params.get("temperature")
+ )
if optional_params.get("top_p"):
safe_set_attribute(span, "llm.request.top_p", optional_params.get("top_p"))
- safe_set_attribute(span, "llm.is_streaming", str(optional_params.get("stream", False)))
+ safe_set_attribute(
+ span, "llm.is_streaming", str(optional_params.get("stream", False))
+ )
if optional_params.get("user"):
safe_set_attribute(span, "llm.user", optional_params.get("user"))
@@ -443,7 +493,9 @@ def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) ->
if not model_params:
return
- safe_set_attribute(span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params))
+ safe_set_attribute(
+ span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params)
+ )
if model_params.get("user"):
user_id = model_params.get("user")
if user_id is not None:
diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py
index 9c2f0d95d4d..fe2f9f41f1b 100644
--- a/litellm/integrations/arize/arize.py
+++ b/litellm/integrations/arize/arize.py
@@ -28,6 +28,41 @@ else:
class ArizeLogger(OpenTelemetry):
+ """
+ Arize logger that sends traces to an Arize endpoint.
+
+ Creates its own dedicated TracerProvider so it can coexist with the
+ generic ``otel`` callback (or any other OTEL-based integration) without
+ fighting over the global ``opentelemetry.trace`` TracerProvider singleton.
+ """
+
+ def _init_tracing(self, tracer_provider):
+ """
+ Override to always create a *private* TracerProvider for Arize.
+
+ See ArizePhoenixLogger._init_tracing for full rationale.
+ """
+ from opentelemetry.sdk.trace import TracerProvider
+ from opentelemetry.trace import SpanKind
+
+ if tracer_provider is not None:
+ self.tracer = tracer_provider.get_tracer("litellm")
+ self.span_kind = SpanKind
+ return
+
+ provider = TracerProvider(resource=self._get_litellm_resource(self.config))
+ provider.add_span_processor(self._get_span_processor())
+ self.tracer = provider.get_tracer("litellm")
+ self.span_kind = SpanKind
+
+ def _init_otel_logger_on_litellm_proxy(self):
+ """
+ Override: Arize should NOT overwrite the proxy's
+ ``open_telemetry_logger``. That attribute is reserved for the
+ primary ``otel`` callback which handles proxy-level parent spans.
+ """
+ pass
+
def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]):
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
return
diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py
index cd345a7f76d..00bc24d4188 100644
--- a/litellm/integrations/arize/arize_phoenix.py
+++ b/litellm/integrations/arize/arize_phoenix.py
@@ -5,43 +5,276 @@ from litellm._logging import verbose_logger
from litellm.integrations.arize import _utils
from litellm.integrations.arize._utils import ArizeOTELAttributes
from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig
-from litellm.integrations.opentelemetry import OpenTelemetry
if TYPE_CHECKING:
+ from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import Span as _Span
+ from opentelemetry.trace import SpanKind
- from litellm.integrations.opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig
+ from litellm.integrations.opentelemetry import OpenTelemetry as _OpenTelemetry
+ from litellm.integrations.opentelemetry import (
+ OpenTelemetryConfig as _OpenTelemetryConfig,
+ )
from litellm.types.integrations.arize import Protocol as _Protocol
Protocol = _Protocol
OpenTelemetryConfig = _OpenTelemetryConfig
Span = Union[_Span, Any]
+ OpenTelemetry = _OpenTelemetry
else:
Protocol = Any
OpenTelemetryConfig = Any
Span = Any
+ TracerProvider = Any
+ SpanKind = Any
+ # Import OpenTelemetry at runtime
+ try:
+ from litellm.integrations.opentelemetry import OpenTelemetry
+ except ImportError:
+ OpenTelemetry = None # type: ignore
ARIZE_HOSTED_PHOENIX_ENDPOINT = "https://otlp.arize.com/v1/traces"
-class ArizePhoenixLogger(OpenTelemetry):
+class ArizePhoenixLogger(OpenTelemetry): # type: ignore
+ """
+ Arize Phoenix logger that sends traces to a Phoenix endpoint.
+
+ Creates its own dedicated TracerProvider so it can coexist with the
+ generic ``otel`` callback (or any other OTEL-based integration) without
+ fighting over the global ``opentelemetry.trace`` TracerProvider singleton.
+ """
+
+ def _init_tracing(self, tracer_provider):
+ """
+ Override to always create a *private* TracerProvider for Arize Phoenix.
+
+ The base ``OpenTelemetry._init_tracing`` falls back to the global
+ TracerProvider when one already exists. That causes whichever
+ integration initialises second to silently reuse the first one's
+ exporter, so spans only reach one destination.
+
+ By creating our own provider we guarantee Arize Phoenix always gets
+ its own exporter pipeline, regardless of initialisation order.
+ """
+ from opentelemetry.sdk.trace import TracerProvider
+ from opentelemetry.trace import SpanKind
+
+ if tracer_provider is not None:
+ # Explicitly supplied (e.g. in tests) — honour it.
+ self.tracer = tracer_provider.get_tracer("litellm")
+ self.span_kind = SpanKind
+ return
+
+ # Always create a dedicated provider — never touch the global one.
+ provider = TracerProvider(resource=self._get_litellm_resource(self.config))
+ provider.add_span_processor(self._get_span_processor())
+ self.tracer = provider.get_tracer("litellm")
+ self.span_kind = SpanKind
+ verbose_logger.debug(
+ "ArizePhoenixLogger: Created dedicated TracerProvider "
+ "(endpoint=%s, exporter=%s)",
+ self.config.endpoint,
+ self.config.exporter,
+ )
+
+ def _init_otel_logger_on_litellm_proxy(self):
+ """
+ Override: Arize Phoenix should NOT overwrite the proxy's
+ ``open_telemetry_logger``. That attribute is reserved for the
+ primary ``otel`` callback which handles proxy-level parent spans.
+ """
+ pass
+
def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]):
ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj)
return
@staticmethod
def set_arize_phoenix_attributes(span: Span, kwargs, response_obj):
+ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import (
+ safe_set_attribute,
+ )
+
_utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes)
-
- # Set project name on the span for all traces to go to custom Phoenix projects
- config = ArizePhoenixLogger.get_arize_phoenix_config()
- if config.project_name:
- from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute
- safe_set_attribute(span, "openinference.project.name", config.project_name)
-
+
+ # Dynamic project name: check metadata first, then fall back to env var config
+ dynamic_project_name = ArizePhoenixLogger._get_dynamic_project_name(kwargs)
+ if dynamic_project_name:
+ safe_set_attribute(span, "openinference.project.name", dynamic_project_name)
+ else:
+ # Fall back to static config from env var
+ config = ArizePhoenixLogger.get_arize_phoenix_config()
+ if config.project_name:
+ safe_set_attribute(
+ span, "openinference.project.name", config.project_name
+ )
+
return
+ @staticmethod
+ def _get_dynamic_project_name(kwargs) -> Optional[str]:
+ """
+ Retrieve dynamic Phoenix project name from request metadata.
+
+ Users can set `metadata.phoenix_project_name` in their request to route
+ traces to different Phoenix projects dynamically.
+ """
+ standard_logging_payload = kwargs.get("standard_logging_object")
+ if isinstance(standard_logging_payload, dict):
+ metadata = standard_logging_payload.get("metadata")
+ if isinstance(metadata, dict):
+ project_name = metadata.get("phoenix_project_name")
+ if project_name:
+ return str(project_name)
+
+ # Also check litellm_params.metadata for SDK usage
+ litellm_params = kwargs.get("litellm_params")
+ if isinstance(litellm_params, dict):
+ metadata = litellm_params.get("metadata") or {}
+ else:
+ metadata = {}
+ if isinstance(metadata, dict):
+ project_name = metadata.get("phoenix_project_name")
+ if project_name:
+ return str(project_name)
+
+ return None
+
+ def _get_phoenix_context(self, kwargs):
+ """
+ Build a trace context for Phoenix's dedicated TracerProvider.
+
+ The base ``_get_span_context`` returns parent spans from the global
+ TracerProvider (the ``otel`` callback). Those spans live on a
+ *different* TracerProvider, so they won't appear in Phoenix — using
+ them as parents just creates broken links.
+
+ Instead we:
+ 1. Honour an incoming ``traceparent`` HTTP header (distributed tracing).
+ 2. In proxy mode, create our *own* parent span on Phoenix's tracer
+ so the hierarchy is visible end-to-end inside Phoenix.
+ 3. In SDK (non-proxy) mode, just return (None, None) for a root span.
+ """
+ from opentelemetry import trace
+
+ litellm_params = kwargs.get("litellm_params", {}) or {}
+ proxy_server_request = litellm_params.get("proxy_server_request", {}) or {}
+ headers = proxy_server_request.get("headers", {}) or {}
+
+ # Propagate distributed trace context if the caller sent a traceparent
+ traceparent_ctx = (
+ self.get_traceparent_from_header(headers=headers)
+ if headers.get("traceparent")
+ else None
+ )
+
+ is_proxy_mode = bool(proxy_server_request)
+
+ if is_proxy_mode:
+ # Create a parent span on Phoenix's own tracer so both parent
+ # and child are exported to Phoenix.
+ start_time_val = kwargs.get("start_time", kwargs.get("api_call_start_time"))
+ parent_span = self.tracer.start_span(
+ name="litellm_proxy_request",
+ start_time=self._to_ns(start_time_val)
+ if start_time_val is not None
+ else None,
+ context=traceparent_ctx,
+ kind=self.span_kind.SERVER,
+ )
+ ctx = trace.set_span_in_context(parent_span)
+ return ctx, parent_span
+
+ # SDK mode — no parent span needed
+ return traceparent_ctx, None
+
+ def _handle_success(self, kwargs, response_obj, start_time, end_time):
+ """
+ Override to always create spans on ArizePhoenixLogger's dedicated TracerProvider.
+
+ The base class's ``_get_span_context`` would find the parent span created by
+ the ``otel`` callback on the *global* TracerProvider. That span is invisible
+ in Phoenix (different exporter pipeline), so we ignore it and build our own
+ hierarchy via ``_get_phoenix_context``.
+ """
+ from opentelemetry.trace import Status, StatusCode
+
+ verbose_logger.debug(
+ "ArizePhoenixLogger: Logging kwargs: %s, OTEL config settings=%s",
+ kwargs,
+ self.config,
+ )
+
+ ctx, parent_span = self._get_phoenix_context(kwargs)
+
+ # Create litellm_request span (child of our parent when in proxy mode)
+ span = self.tracer.start_span(
+ name=self._get_span_name(kwargs),
+ start_time=self._to_ns(start_time),
+ context=ctx,
+ )
+ span.set_status(Status(StatusCode.OK))
+ self.set_attributes(span, kwargs, response_obj)
+
+ # Raw-request sub-span (if enabled) — must be created before
+ # ending the parent span so the hierarchy is valid.
+ self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span)
+ span.end(end_time=self._to_ns(end_time))
+
+ # Guardrail span
+ self._create_guardrail_span(kwargs=kwargs, context=ctx)
+
+ # Annotate and close our proxy parent span
+ if parent_span is not None:
+ parent_span.set_status(Status(StatusCode.OK))
+ self.set_attributes(parent_span, kwargs, response_obj)
+ parent_span.end(end_time=self._to_ns(end_time))
+
+ # Metrics & cost recording
+ self._record_metrics(kwargs, response_obj, start_time, end_time)
+
+ # Semantic logs
+ if self.config.enable_events:
+ self._emit_semantic_logs(kwargs, response_obj, span)
+
+ def _handle_failure(self, kwargs, response_obj, start_time, end_time):
+ """
+ Override to always create failure spans on ArizePhoenixLogger's dedicated
+ TracerProvider. Mirrors ``_handle_success`` but sets ERROR status.
+ """
+ from opentelemetry.trace import Status, StatusCode
+
+ verbose_logger.debug(
+ "ArizePhoenixLogger: Failure - Logging kwargs: %s, OTEL config settings=%s",
+ kwargs,
+ self.config,
+ )
+
+ ctx, parent_span = self._get_phoenix_context(kwargs)
+
+ # Create litellm_request span (child of our parent when in proxy mode)
+ span = self.tracer.start_span(
+ name=self._get_span_name(kwargs),
+ start_time=self._to_ns(start_time),
+ context=ctx,
+ )
+ span.set_status(Status(StatusCode.ERROR))
+ self.set_attributes(span, kwargs, response_obj)
+ self._record_exception_on_span(span=span, kwargs=kwargs)
+ span.end(end_time=self._to_ns(end_time))
+
+ # Guardrail span
+ self._create_guardrail_span(kwargs=kwargs, context=ctx)
+
+ # Annotate and close our proxy parent span
+ if parent_span is not None:
+ parent_span.set_status(Status(StatusCode.ERROR))
+ self.set_attributes(parent_span, kwargs, response_obj)
+ self._record_exception_on_span(span=parent_span, kwargs=kwargs)
+ parent_span.end(end_time=self._to_ns(end_time))
+
@staticmethod
def get_arize_phoenix_config() -> ArizePhoenixConfig:
"""
@@ -63,7 +296,9 @@ class ArizePhoenixLogger(OpenTelemetry):
if collector_endpoint:
# Parse the endpoint to determine protocol
- if collector_endpoint.startswith("grpc://") or (":4317" in collector_endpoint and "/v1/traces" not in collector_endpoint):
+ if collector_endpoint.startswith("grpc://") or (
+ ":4317" in collector_endpoint and "/v1/traces" not in collector_endpoint
+ ):
endpoint = collector_endpoint
protocol = "otlp_grpc"
else:
@@ -107,11 +342,10 @@ class ArizePhoenixLogger(OpenTelemetry):
endpoint=endpoint,
project_name=project_name,
)
-
+
## cannot suppress additional proxy server spans, removed previous methods.
async def async_health_check(self):
-
config = self.get_arize_phoenix_config()
if not config.otlp_auth_headers:
@@ -123,4 +357,4 @@ class ArizePhoenixLogger(OpenTelemetry):
return {
"status": "healthy",
"message": "Arize-Phoenix credentials are configured properly",
- }
\ No newline at end of file
+ }
diff --git a/litellm/integrations/azure_sentinel/__init__.py b/litellm/integrations/azure_sentinel/__init__.py
index 46f2fed0a97..036711a80dd 100644
--- a/litellm/integrations/azure_sentinel/__init__.py
+++ b/litellm/integrations/azure_sentinel/__init__.py
@@ -1,4 +1,3 @@
from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
__all__ = ["AzureSentinelLogger"]
-
diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py
index 875432de876..dd508e6c6c2 100644
--- a/litellm/integrations/azure_sentinel/azure_sentinel.py
+++ b/litellm/integrations/azure_sentinel/azure_sentinel.py
@@ -62,18 +62,22 @@ class AzureSentinelLogger(CustomBatchLogger):
llm_provider=httpxSpecialProvider.LoggingCallback
)
- self.dcr_immutable_id = (
- dcr_immutable_id or os.getenv("AZURE_SENTINEL_DCR_IMMUTABLE_ID")
+ self.dcr_immutable_id = dcr_immutable_id or os.getenv(
+ "AZURE_SENTINEL_DCR_IMMUTABLE_ID"
)
self.stream_name = stream_name or os.getenv(
"AZURE_SENTINEL_STREAM_NAME", "Custom-LiteLLM"
)
self.endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT")
- self.tenant_id = tenant_id or os.getenv("AZURE_SENTINEL_TENANT_ID") or os.getenv(
- "AZURE_TENANT_ID"
+ self.tenant_id = (
+ tenant_id
+ or os.getenv("AZURE_SENTINEL_TENANT_ID")
+ or os.getenv("AZURE_TENANT_ID")
)
- self.client_id = client_id or os.getenv("AZURE_SENTINEL_CLIENT_ID") or os.getenv(
- "AZURE_CLIENT_ID"
+ self.client_id = (
+ client_id
+ or os.getenv("AZURE_SENTINEL_CLIENT_ID")
+ or os.getenv("AZURE_CLIENT_ID")
)
self.client_secret = (
client_secret
@@ -103,9 +107,7 @@ class AzureSentinelLogger(CustomBatchLogger):
)
# Build API endpoint: {Endpoint}/dataCollectionRules/{DCR Immutable ID}/streams/{Stream Name}?api-version=2023-01-01
- self.api_endpoint = (
- f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01"
- )
+ self.api_endpoint = f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01"
# OAuth2 scope for Azure Monitor
self.oauth_scope = "https://monitor.azure.com/.default"
@@ -139,7 +141,9 @@ class AzureSentinelLogger(CustomBatchLogger):
assert self.client_id is not None, "client_id is required"
assert self.client_secret is not None, "client_secret is required"
- token_url = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
+ token_url = (
+ f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
+ )
token_data = {
"client_id": self.client_id,
@@ -173,9 +177,7 @@ class AzureSentinelLogger(CustomBatchLogger):
return self.oauth_token
- async def async_log_success_event(
- self, kwargs, response_obj, start_time, end_time
- ):
+ async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
Async Log success events to Azure Sentinel
@@ -209,9 +211,7 @@ class AzureSentinelLogger(CustomBatchLogger):
)
pass
- async def async_log_failure_event(
- self, kwargs, response_obj, start_time, end_time
- ):
+ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
"""
Async Log failure events to Azure Sentinel
diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py
index 85f91199c1c..6fc7b9c1048 100644
--- a/litellm/integrations/azure_storage/azure_storage.py
+++ b/litellm/integrations/azure_storage/azure_storage.py
@@ -54,12 +54,12 @@ class AzureBlobStorageLogger(CustomBatchLogger):
self._service_client_timeout: Optional[float] = None
# Internal variables used for Token based authentication
- self.azure_auth_token: Optional[str] = (
- None # the Azure AD token to use for Azure Storage API requests
- )
- self.token_expiry: Optional[datetime] = (
- None # the expiry time of the currentAzure AD token
- )
+ self.azure_auth_token: Optional[
+ str
+ ] = None # the Azure AD token to use for Azure Storage API requests
+ self.token_expiry: Optional[
+ datetime
+ ] = None # the expiry time of the currentAzure AD token
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py
index 42e9680a7fc..cb1b2bc5531 100644
--- a/litellm/integrations/braintrust_logging.py
+++ b/litellm/integrations/braintrust_logging.py
@@ -41,7 +41,9 @@ class BraintrustLogger(CustomLogger):
self.is_mock_mode = should_use_braintrust_mock()
if self.is_mock_mode:
create_mock_braintrust_client()
- verbose_logger.info("[BRAINTRUST MOCK] Braintrust logger initialized in mock mode")
+ verbose_logger.info(
+ "[BRAINTRUST MOCK] Braintrust logger initialized in mock mode"
+ )
self.validate_environment(api_key=api_key)
self.api_base = api_base or os.getenv("BRAINTRUST_API_BASE") or API_BASE
self.default_project_id = None
@@ -50,9 +52,9 @@ class BraintrustLogger(CustomLogger):
"Authorization": "Bearer " + self.api_key,
"Content-Type": "application/json",
}
- self._project_id_cache: Dict[str, str] = (
- {}
- ) # Cache mapping project names to IDs
+ self._project_id_cache: Dict[
+ str, str
+ ] = {} # Cache mapping project names to IDs
self.global_braintrust_http_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
@@ -214,7 +216,7 @@ class BraintrustLogger(CustomLogger):
# Allow metadata override for span name
span_name = dynamic_metadata.get("span_name", "Chat Completion")
-
+
# Span parents is a special case
span_parents = dynamic_metadata.get("span_parents")
@@ -236,7 +238,7 @@ class BraintrustLogger(CustomLogger):
"span_attributes": {"name": span_name, "type": "llm"},
}
- # Braintrust cannot specify 'tags' for non-root spans
+ # Braintrust cannot specify 'tags' for non-root spans
if dynamic_metadata.get("root_span_id") is None:
request_data["tags"] = tags
@@ -386,7 +388,7 @@ class BraintrustLogger(CustomLogger):
"span_attributes": {"name": span_name, "type": "llm"},
}
- # Braintrust cannot specify 'tags' for non-root spans
+ # Braintrust cannot specify 'tags' for non-root spans
if dynamic_metadata.get("root_span_id") is None:
request_data["tags"] = tags
diff --git a/litellm/integrations/braintrust_mock_client.py b/litellm/integrations/braintrust_mock_client.py
index 030aa62cd0f..59e0988a10a 100644
--- a/litellm/integrations/braintrust_mock_client.py
+++ b/litellm/integrations/braintrust_mock_client.py
@@ -13,7 +13,11 @@ import time
from urllib.parse import urlparse
from litellm._logging import verbose_logger
-from litellm.integrations.mock_client_factory import MockClientConfig, MockResponse, create_mock_client_factory
+from litellm.integrations.mock_client_factory import (
+ MockClientConfig,
+ MockResponse,
+ create_mock_client_factory,
+)
# Use factory for should_use_mock and MockResponse
# Braintrust uses both HTTPHandler (sync) and AsyncHTTPHandler (async)
@@ -37,7 +41,10 @@ _config = MockClientConfig(
# Get should_use_mock and create_mock_client from factory
# We need to call the factory's create_mock_client to patch AsyncHTTPHandler.post
-create_mock_braintrust_factory_client, should_use_braintrust_mock = create_mock_client_factory(_config)
+(
+ create_mock_braintrust_factory_client,
+ should_use_braintrust_mock,
+) = create_mock_client_factory(_config)
# Store original HTTPHandler.post method (Braintrust-specific for sync calls with custom logic)
_original_http_handler_post = None
@@ -66,7 +73,19 @@ def _is_braintrust_url(url: str) -> bool:
)
-def _mock_http_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, files=None, content=None, logging_obj=None):
+def _mock_http_handler_post(
+ self,
+ url,
+ data=None,
+ json=None,
+ params=None,
+ headers=None,
+ timeout=None,
+ stream=False,
+ files=None,
+ content=None,
+ logging_obj=None,
+):
"""Monkey-patched HTTPHandler.post that intercepts Braintrust calls with endpoint-specific responses."""
# Only mock Braintrust API calls
if isinstance(url, str) and _is_braintrust_url(url):
@@ -86,46 +105,62 @@ def _mock_http_handler_post(self, url, data=None, json=None, params=None, header
status_code=_config.default_status_code,
json_data=mock_data,
url=url,
- elapsed_seconds=_MOCK_LATENCY_SECONDS
+ elapsed_seconds=_MOCK_LATENCY_SECONDS,
)
if _original_http_handler_post is not None:
- return _original_http_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, files=files, content=content, logging_obj=logging_obj)
+ return _original_http_handler_post(
+ self,
+ url=url,
+ data=data,
+ json=json,
+ params=params,
+ headers=headers,
+ timeout=timeout,
+ stream=stream,
+ files=files,
+ content=content,
+ logging_obj=logging_obj,
+ )
raise RuntimeError("Original HTTPHandler.post not available")
def create_mock_braintrust_client():
"""
Monkey-patch HTTPHandler.post to intercept Braintrust sync calls.
-
+
Braintrust uses HTTPHandler for sync calls and AsyncHTTPHandler for async calls.
HTTPHandler.post uses self.client.send(), not self.client.post(), so we need
custom patching for sync (similar to Helicone).
AsyncHTTPHandler.post is patched by the factory.
-
+
We use custom patching instead of factory's patch_http_handler because we need
endpoint-specific responses (different for /project vs /project_logs).
-
+
This function is idempotent - it only initializes mocks once, even if called multiple times.
"""
global _original_http_handler_post, _mocks_initialized
-
+
if _mocks_initialized:
return
-
+
verbose_logger.debug("[BRAINTRUST MOCK] Initializing Braintrust mock client...")
-
+
from litellm.llms.custom_httpx.http_handler import HTTPHandler
-
+
if _original_http_handler_post is None:
_original_http_handler_post = HTTPHandler.post
HTTPHandler.post = _mock_http_handler_post # type: ignore
verbose_logger.debug("[BRAINTRUST MOCK] Patched HTTPHandler.post")
-
+
# CRITICAL: Call the factory's initialization function to patch AsyncHTTPHandler.post
# This is required for async calls to be mocked
create_mock_braintrust_factory_client()
-
- verbose_logger.debug(f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms")
- verbose_logger.debug("[BRAINTRUST MOCK] Braintrust mock client initialization complete")
-
+
+ verbose_logger.debug(
+ f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms"
+ )
+ verbose_logger.debug(
+ "[BRAINTRUST MOCK] Braintrust mock client initialization complete"
+ )
+
_mocks_initialized = True
diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json
index 6a003b8c499..c2b0c4ddce9 100644
--- a/litellm/integrations/callback_configs.json
+++ b/litellm/integrations/callback_configs.json
@@ -83,6 +83,27 @@
},
"description": "Datadog Logging Integration"
},
+ {
+ "id": "datadog_metrics",
+ "displayName": "Datadog Metrics",
+ "logo": "datadog.png",
+ "supports_key_team_logging": false,
+ "dynamic_params": {
+ "dd_api_key": {
+ "type": "password",
+ "ui_name": "API Key",
+ "description": "Datadog API key for authentication",
+ "required": true
+ },
+ "dd_site": {
+ "type": "text",
+ "ui_name": "Site",
+ "description": "Datadog site URL (e.g., us5.datadoghq.com)",
+ "required": true
+ }
+ },
+ "description": "Datadog Custom Metrics Integration"
+ },
{
"id": "datadog_cost_management",
"displayName": "Datadog Cost Management",
@@ -434,4 +455,4 @@
},
"description": "SQS Queue (AWS) Logging Integration"
}
-]
\ No newline at end of file
+]
diff --git a/litellm/integrations/cloudzero/cz_resource_names.py b/litellm/integrations/cloudzero/cz_resource_names.py
index f1098d20381..20862c1c7ec 100644
--- a/litellm/integrations/cloudzero/cz_resource_names.py
+++ b/litellm/integrations/cloudzero/cz_resource_names.py
@@ -30,7 +30,9 @@ class CZEntityType(str, Enum):
class CZRNGenerator:
"""Generate CloudZero Resource Names (CZRNs) for LiteLLM resources."""
- CZRN_REGEX = re.compile(r'^czrn:([a-z0-9-]+):([a-zA-Z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):(.+)$')
+ CZRN_REGEX = re.compile(
+ r"^czrn:([a-z0-9-]+):([a-zA-Z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):(.+)$"
+ )
def __init__(self):
"""Initialize CZRN generator."""
@@ -38,9 +40,9 @@ class CZRNGenerator:
def create_from_litellm_data(self, row: dict[str, Any]) -> str:
"""Create a CZRN from LiteLLM daily spend data.
-
+
CZRN format: czrn::::::
-
+
For LiteLLM resources, we map:
- service-type: 'litellm' (the service managing the LLM calls)
- provider: The custom_llm_provider (e.g., 'openai', 'anthropic', 'azure')
@@ -49,18 +51,18 @@ class CZRNGenerator:
- resource-type: 'llm-usage' (represents LLM usage/inference)
- cloud-local-id: model
"""
- service_type = 'litellm'
- provider = self._normalize_provider(row.get('custom_llm_provider', 'unknown'))
- region = 'cross-region'
+ service_type = "litellm"
+ provider = self._normalize_provider(row.get("custom_llm_provider", "unknown"))
+ region = "cross-region"
# Use the actual entity_id (team_id or user_id) as the owner account
- team_id = row.get('team_id', 'unknown')
+ team_id = row.get("team_id", "unknown")
owner_account_id = self._normalize_component(team_id)
- resource_type = 'llm-usage'
+ resource_type = "llm-usage"
# Create a unique identifier with just the model (entity info already in owner_account_id)
- model = row.get('model', 'unknown')
+ model = row.get("model", "unknown")
cloud_local_id = model
@@ -70,7 +72,7 @@ class CZRNGenerator:
region=region,
owner_account_id=owner_account_id,
resource_type=resource_type,
- cloud_local_id=cloud_local_id
+ cloud_local_id=cloud_local_id,
)
def create_from_components(
@@ -80,7 +82,7 @@ class CZRNGenerator:
region: str,
owner_account_id: str,
resource_type: str,
- cloud_local_id: str
+ cloud_local_id: str,
) -> str:
"""Create a CZRN from individual components."""
# Normalize components to ensure they meet CZRN requirements
@@ -104,7 +106,7 @@ class CZRNGenerator:
def extract_components(self, czrn: str) -> tuple[str, str, str, str, str, str]:
"""Extract all components from a CZRN.
-
+
Returns: (service_type, provider, region, owner_account_id, resource_type, cloud_local_id)
"""
match = self.CZRN_REGEX.match(czrn)
@@ -117,42 +119,43 @@ class CZRNGenerator:
"""Normalize provider names to standard CZRN format."""
# Map common provider names to CZRN standards
provider_map = {
- litellm.LlmProviders.AZURE.value: 'azure',
- litellm.LlmProviders.AZURE_AI.value: 'azure',
- litellm.LlmProviders.ANTHROPIC.value: 'anthropic',
- litellm.LlmProviders.BEDROCK.value: 'aws',
- litellm.LlmProviders.VERTEX_AI.value: 'gcp',
- litellm.LlmProviders.GEMINI.value: 'google',
- litellm.LlmProviders.COHERE.value: 'cohere',
- litellm.LlmProviders.HUGGINGFACE.value: 'huggingface',
- litellm.LlmProviders.REPLICATE.value: 'replicate',
- litellm.LlmProviders.TOGETHER_AI.value: 'together-ai',
+ litellm.LlmProviders.AZURE.value: "azure",
+ litellm.LlmProviders.AZURE_AI.value: "azure",
+ litellm.LlmProviders.ANTHROPIC.value: "anthropic",
+ litellm.LlmProviders.BEDROCK.value: "aws",
+ litellm.LlmProviders.VERTEX_AI.value: "gcp",
+ litellm.LlmProviders.GEMINI.value: "google",
+ litellm.LlmProviders.COHERE.value: "cohere",
+ litellm.LlmProviders.HUGGINGFACE.value: "huggingface",
+ litellm.LlmProviders.REPLICATE.value: "replicate",
+ litellm.LlmProviders.TOGETHER_AI.value: "together-ai",
}
- normalized = provider.lower().replace('_', '-')
+ normalized = provider.lower().replace("_", "-")
# use litellm custom llm provider if not in provider_map
if normalized not in provider_map:
return normalized
return provider_map.get(normalized, normalized)
- def _normalize_component(self, component: str, allow_uppercase: bool = False) -> str:
+ def _normalize_component(
+ self, component: str, allow_uppercase: bool = False
+ ) -> str:
"""Normalize a CZRN component to meet format requirements."""
if not component:
- return 'unknown'
+ return "unknown"
# Convert to lowercase unless uppercase is allowed
if not allow_uppercase:
component = component.lower()
# Replace invalid characters with hyphens
- component = re.sub(r'[^a-zA-Z0-9-]', '-', component)
+ component = re.sub(r"[^a-zA-Z0-9-]", "-", component)
# Remove consecutive hyphens
- component = re.sub(r'-+', '-', component)
+ component = re.sub(r"-+", "-", component)
# Remove leading/trailing hyphens
- component = component.strip('-')
-
- return component or 'unknown'
+ component = component.strip("-")
+ return component or "unknown"
diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py
index 83b6e318ba7..d673536e72d 100644
--- a/litellm/integrations/cloudzero/cz_stream_api.py
+++ b/litellm/integrations/cloudzero/cz_stream_api.py
@@ -30,7 +30,9 @@ from rich.console import Console
class CloudZeroStreamer:
"""Stream CBF data to CloudZero AnyCost API with proper batching and timezone handling."""
- def __init__(self, api_key: str, connection_id: str, user_timezone: Optional[str] = None):
+ def __init__(
+ self, api_key: str, connection_id: str, user_timezone: Optional[str] = None
+ ):
"""Initialize CloudZero streamer with credentials."""
self.api_key = api_key
self.connection_id = connection_id
@@ -43,12 +45,16 @@ class CloudZeroStreamer:
try:
self.user_timezone = zoneinfo.ZoneInfo(user_timezone)
except zoneinfo.ZoneInfoNotFoundError:
- self.console.print(f"[yellow]Warning: Unknown timezone '{user_timezone}', using UTC[/yellow]")
+ self.console.print(
+ f"[yellow]Warning: Unknown timezone '{user_timezone}', using UTC[/yellow]"
+ )
self.user_timezone = timezone.utc
else:
self.user_timezone = timezone.utc
- def send_batched(self, data: pl.DataFrame, operation: str = "replace_hourly") -> None:
+ def send_batched(
+ self, data: pl.DataFrame, operation: str = "replace_hourly"
+ ) -> None:
"""Send CBF data in daily batches to CloudZero AnyCost API."""
if data.is_empty():
self.console.print("[yellow]No data to send to CloudZero[/yellow]")
@@ -61,7 +67,9 @@ class CloudZeroStreamer:
self.console.print("[yellow]No valid daily batches to send[/yellow]")
return
- self.console.print(f"[blue]Sending {len(daily_batches)} daily batch(es) with operation '{operation}'[/blue]")
+ self.console.print(
+ f"[blue]Sending {len(daily_batches)} daily batch(es) with operation '{operation}'[/blue]"
+ )
for batch_date, batch_data in daily_batches.items():
self._send_daily_batch(batch_date, batch_data, operation)
@@ -71,21 +79,23 @@ class CloudZeroStreamer:
daily_batches: dict[str, list[dict[str, Any]]] = {}
# Ensure we have the required columns
- if 'time/usage_start' not in data.columns:
- self.console.print("[red]Error: Missing 'time/usage_start' column for date grouping[/red]")
+ if "time/usage_start" not in data.columns:
+ self.console.print(
+ "[red]Error: Missing 'time/usage_start' column for date grouping[/red]"
+ )
return {}
-
+
timestamp_str: Optional[str] = None
for row in data.iter_rows(named=True):
try:
# Parse the timestamp and convert to UTC
- timestamp_str = row.get('time/usage_start')
+ timestamp_str = row.get("time/usage_start")
if not timestamp_str:
continue
# Parse timestamp and handle timezone conversion
dt = self._parse_and_convert_timestamp(timestamp_str)
- batch_date = dt.strftime('%Y-%m-%d')
+ batch_date = dt.strftime("%Y-%m-%d")
if batch_date not in daily_batches:
daily_batches[batch_date] = []
@@ -93,25 +103,54 @@ class CloudZeroStreamer:
daily_batches[batch_date].append(row)
except Exception as e:
- self.console.print(f"[yellow]Warning: Could not process timestamp '{timestamp_str}': {e}[/yellow]")
+ self.console.print(
+ f"[yellow]Warning: Could not process timestamp '{timestamp_str}': {e}[/yellow]"
+ )
continue
# Convert lists back to DataFrames
- return {date_key: pl.DataFrame(records) for date_key, records in daily_batches.items() if records}
+ return {
+ date_key: pl.DataFrame(records)
+ for date_key, records in daily_batches.items()
+ if records
+ }
def _parse_and_convert_timestamp(self, timestamp_str: str) -> datetime:
"""Parse timestamp string and convert to UTC."""
# Try to parse the timestamp string
try:
# Handle various ISO 8601 formats
- if timestamp_str.endswith('Z'):
- dt = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00'))
- elif '+' in timestamp_str or timestamp_str.endswith(('-00:00', '-01:00', '-02:00', '-03:00',
- '-04:00', '-05:00', '-06:00', '-07:00',
- '-08:00', '-09:00', '-10:00', '-11:00',
- '-12:00', '+01:00', '+02:00', '+03:00',
- '+04:00', '+05:00', '+06:00', '+07:00',
- '+08:00', '+09:00', '+10:00', '+11:00', '+12:00')):
+ if timestamp_str.endswith("Z"):
+ dt = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
+ elif "+" in timestamp_str or timestamp_str.endswith(
+ (
+ "-00:00",
+ "-01:00",
+ "-02:00",
+ "-03:00",
+ "-04:00",
+ "-05:00",
+ "-06:00",
+ "-07:00",
+ "-08:00",
+ "-09:00",
+ "-10:00",
+ "-11:00",
+ "-12:00",
+ "+01:00",
+ "+02:00",
+ "+03:00",
+ "+04:00",
+ "+05:00",
+ "+06:00",
+ "+07:00",
+ "+08:00",
+ "+09:00",
+ "+10:00",
+ "+11:00",
+ "+12:00",
+ )
+ ):
dt = datetime.fromisoformat(timestamp_str)
else:
# Assume user timezone if no timezone info
@@ -125,14 +164,16 @@ class CloudZeroStreamer:
except ValueError as e:
raise ValueError(f"Could not parse timestamp '{timestamp_str}': {e}")
- def _send_daily_batch(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> None:
+ def _send_daily_batch(
+ self, batch_date: str, batch_data: pl.DataFrame, operation: str
+ ) -> None:
"""Send a single daily batch to CloudZero API."""
if batch_data.is_empty():
return
headers = {
- 'Authorization': f'Bearer {self.api_key}',
- 'Content-Type': 'application/json'
+ "Authorization": f"Bearer {self.api_key}",
+ "Content-Type": "application/json",
}
# Use the correct API endpoint format from documentation
@@ -143,29 +184,39 @@ class CloudZeroStreamer:
try:
with httpx.Client(timeout=30.0) as client:
- self.console.print(f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]")
+ self.console.print(
+ f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]"
+ )
response = client.post(url, headers=headers, json=payload)
response.raise_for_status()
- self.console.print(f"[green]✓ Successfully sent batch for {batch_date} ({len(batch_data)} records)[/green]")
+ self.console.print(
+ f"[green]✓ Successfully sent batch for {batch_date} ({len(batch_data)} records)[/green]"
+ )
except httpx.RequestError as e:
- self.console.print(f"[red]✗ Network error sending batch for {batch_date}: {e}[/red]")
+ self.console.print(
+ f"[red]✗ Network error sending batch for {batch_date}: {e}[/red]"
+ )
raise
except httpx.HTTPStatusError as e:
- self.console.print(f"[red]✗ HTTP error sending batch for {batch_date}: {e.response.status_code} {e.response.text}[/red]")
+ self.console.print(
+ f"[red]✗ HTTP error sending batch for {batch_date}: {e.response.status_code} {e.response.text}[/red]"
+ )
raise
- def _prepare_batch_payload(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> dict[str, Any]:
+ def _prepare_batch_payload(
+ self, batch_date: str, batch_data: pl.DataFrame, operation: str
+ ) -> dict[str, Any]:
"""Prepare batch payload according to CloudZero AnyCost API format."""
# Convert batch_date to month for the API (YYYY-MM format)
try:
- date_obj = datetime.strptime(batch_date, '%Y-%m-%d')
- month_str = date_obj.strftime('%Y-%m')
+ date_obj = datetime.strptime(batch_date, "%Y-%m-%d")
+ month_str = date_obj.strftime("%Y-%m")
except ValueError:
# Fallback to current month
- month_str = datetime.now().strftime('%Y-%m')
+ month_str = datetime.now().strftime("%Y-%m")
# Convert DataFrame rows to API format
data_records = []
@@ -174,15 +225,13 @@ class CloudZeroStreamer:
if record:
data_records.append(record)
- payload = {
- 'month': month_str,
- 'operation': operation,
- 'data': data_records
- }
+ payload = {"month": month_str, "operation": operation, "data": data_records}
return payload
- def _convert_cbf_to_api_format(self, row: dict[str, Any]) -> Optional[dict[str, Any]]:
+ def _convert_cbf_to_api_format(
+ self, row: dict[str, Any]
+ ) -> Optional[dict[str, Any]]:
"""Convert CBF row to CloudZero API format - keeping CBF field names as CloudZero expects them."""
try:
# CloudZero expects CBF format field names directly, not converted names
@@ -196,20 +245,24 @@ class CloudZeroStreamer:
# Format floats to avoid scientific notation
if isinstance(value, float):
# Use a reasonable precision that avoids scientific notation
- api_record[key] = f"{value:.10f}".rstrip('0').rstrip('.')
+ api_record[key] = f"{value:.10f}".rstrip("0").rstrip(".")
else:
api_record[key] = str(value)
else:
api_record[key] = value
# Ensure timestamp is in UTC format
- if 'time/usage_start' in api_record:
- api_record['time/usage_start'] = self._ensure_utc_timestamp(api_record['time/usage_start'])
+ if "time/usage_start" in api_record:
+ api_record["time/usage_start"] = self._ensure_utc_timestamp(
+ api_record["time/usage_start"]
+ )
return api_record
except Exception as e:
- self.console.print(f"[yellow]Warning: Could not convert record to API format: {e}[/yellow]")
+ self.console.print(
+ f"[yellow]Warning: Could not convert record to API format: {e}[/yellow]"
+ )
return None
def _ensure_utc_timestamp(self, timestamp_str: str) -> str:
@@ -219,9 +272,7 @@ class CloudZeroStreamer:
try:
dt = self._parse_and_convert_timestamp(timestamp_str)
- return dt.isoformat().replace('+00:00', 'Z')
+ return dt.isoformat().replace("+00:00", "Z")
except Exception:
# Fallback to current time in UTC
- return datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z')
-
-
+ return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py
index e06b944a419..c1b0d5cf411 100644
--- a/litellm/integrations/cloudzero/transform.py
+++ b/litellm/integrations/cloudzero/transform.py
@@ -41,8 +41,8 @@ class CBFTransformer:
# Filter out records with zero successful_requests first
original_count = len(data)
- if 'successful_requests' in data.columns:
- filtered_data = data.filter(pl.col('successful_requests') > 0)
+ if "successful_requests" in data.columns:
+ filtered_data = data.filter(pl.col("successful_requests") > 0)
zero_requests_dropped = original_count - len(filtered_data)
else:
filtered_data = data
@@ -64,16 +64,23 @@ class CBFTransformer:
# Print summary of dropped records if any
from rich.console import Console
+
console = Console()
if zero_requests_dropped > 0:
- console.print(f"[yellow]⚠️ Dropped {zero_requests_dropped:,} of {original_count:,} records with zero successful_requests[/yellow]")
+ console.print(
+ f"[yellow]⚠️ Dropped {zero_requests_dropped:,} of {original_count:,} records with zero successful_requests[/yellow]"
+ )
if czrn_dropped_count > 0:
- console.print(f"[yellow]⚠️ Dropped {czrn_dropped_count:,} of {filtered_count:,} filtered records due to invalid CZRNs[/yellow]")
+ console.print(
+ f"[yellow]⚠️ Dropped {czrn_dropped_count:,} of {filtered_count:,} filtered records due to invalid CZRNs[/yellow]"
+ )
if len(cbf_data) > 0:
- console.print(f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]")
+ console.print(
+ f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]"
+ )
return pl.DataFrame(cbf_data)
@@ -81,87 +88,116 @@ class CBFTransformer:
"""Create a single CBF record from LiteLLM daily spend row."""
# Parse date (daily spend tables use date strings like '2025-04-19')
- usage_date = self._parse_date(row.get('date'))
+ usage_date = self._parse_date(row.get("date"))
# Calculate total tokens
- prompt_tokens = int(row.get('prompt_tokens', 0))
- completion_tokens = int(row.get('completion_tokens', 0))
+ prompt_tokens = int(row.get("prompt_tokens", 0))
+ completion_tokens = int(row.get("completion_tokens", 0))
total_tokens = prompt_tokens + completion_tokens
# Create CloudZero Resource Name (CZRN) as resource_id
resource_id = self.czrn_generator.create_from_litellm_data(row)
# Build dimensions for CloudZero
- model = str(row.get('model', ''))
- api_key_hash = str(row.get('api_key', ''))[:8] # First 8 chars for identification
-
+ model = str(row.get("model", ""))
+ api_key_hash = str(row.get("api_key", ""))[
+ :8
+ ] # First 8 chars for identification
+
# Handle team information with fallbacks
- team_id = row.get('team_id')
- team_alias = row.get('team_alias')
- user_email = row.get('user_email')
-
+ team_id = row.get("team_id")
+ team_alias = row.get("team_alias")
+ user_email = row.get("user_email")
+
# Use team_alias if available, otherwise team_id, otherwise fallback to 'unknown'
- entity_id = str(team_alias) if team_alias else (str(team_id) if team_id else 'unknown')
-
+ entity_id = (
+ str(team_alias) if team_alias else (str(team_id) if team_id else "unknown")
+ )
+
+ # Get alias fields if they exist
+ api_key_alias = row.get("api_key_alias")
+ organization_alias = row.get("organization_alias")
+ project_alias = row.get("project_alias")
+ user_alias = row.get("user_alias")
+
dimensions = {
- 'entity_type': CZEntityType.TEAM.value,
- 'entity_id': entity_id,
- 'team_id': str(team_id) if team_id else 'unknown',
- 'team_alias': str(team_alias) if team_alias else 'unknown',
- 'model': model,
- 'model_group': str(row.get('model_group', '')),
- 'provider': str(row.get('custom_llm_provider', '')),
- 'api_key_prefix': api_key_hash,
- 'api_key_alias': str(row.get('api_key_alias', '')),
- 'user_email': str(user_email) if user_email else '',
- 'api_requests': str(row.get('api_requests', 0)),
- 'successful_requests': str(row.get('successful_requests', 0)),
- 'failed_requests': str(row.get('failed_requests', 0)),
- 'cache_creation_tokens': str(row.get('cache_creation_input_tokens', 0)),
- 'cache_read_tokens': str(row.get('cache_read_input_tokens', 0)),
+ "entity_type": CZEntityType.TEAM.value,
+ "entity_id": entity_id,
+ "team_alias": str(team_alias) if team_alias else "unknown",
+ "model": model,
+ "model_group": str(row.get("model_group", "")),
+ "provider": str(row.get("custom_llm_provider", "")),
+ "api_key_prefix": api_key_hash,
+ "api_key_alias": str(row.get("api_key_alias", "")),
+ "user_email": str(user_email) if user_email else "",
+ "api_requests": str(row.get("api_requests", 0)),
+ "successful_requests": str(row.get("successful_requests", 0)),
+ "failed_requests": str(row.get("failed_requests", 0)),
+ "cache_creation_tokens": str(row.get("cache_creation_input_tokens", 0)),
+ "cache_read_tokens": str(row.get("cache_read_input_tokens", 0)),
+ "organization_alias": str(organization_alias) if organization_alias else "",
+ "project_alias": str(project_alias) if project_alias else "",
+ "user_alias": str(user_alias) if user_alias else "",
}
# Extract CZRN components to populate corresponding CBF columns
czrn_components = self.czrn_generator.extract_components(resource_id)
- service_type, provider, region, owner_account_id, resource_type, cloud_local_id = czrn_components
+ (
+ service_type,
+ provider,
+ region,
+ owner_account_id,
+ resource_type,
+ cloud_local_id,
+ ) = czrn_components
+
+ # Build resource/account as concat of api_key_alias and api_key_prefix
+ resource_account = (
+ f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash
+ )
# CloudZero CBF format with proper column names
cbf_record = {
# Required CBF fields
- 'time/usage_start': usage_date.isoformat() if usage_date else None, # Required: ISO-formatted UTC datetime
- 'cost/cost': float(row.get('spend', 0.0)), # Required: billed cost
- 'resource/id': resource_id, # Required when resource tags are present
-
+ "time/usage_start": usage_date.isoformat()
+ if usage_date
+ else None, # Required: ISO-formatted UTC datetime
+ "cost/cost": float(row.get("spend", 0.0)), # Required: billed cost
+ "resource/id": resource_id, # CZRN (CloudZero Resource Name)
# Usage metrics for token consumption
- 'usage/amount': total_tokens, # Numeric value of tokens consumed
- 'usage/units': 'tokens', # Description of token units
-
- # CBF fields that correspond to CZRN components
- 'resource/service': service_type, # Maps to CZRN service-type (litellm)
- 'resource/account': owner_account_id, # Maps to CZRN owner-account-id (entity_id)
- 'resource/region': region, # Maps to CZRN region (cross-region)
- 'resource/usage_family': resource_type, # Maps to CZRN resource-type (llm-usage)
-
+ "usage/amount": total_tokens, # Numeric value of tokens consumed
+ "usage/units": "tokens", # Description of token units
+ # CBF fields - updated per LIT-1907
+ "resource/service": str(row.get("model_group", "")), # Send model_group
+ "resource/account": resource_account, # Send api_key_alias|api_key_prefix
+ "resource/region": region, # Maps to CZRN region (cross-region)
+ "resource/usage_family": str(
+ row.get("custom_llm_provider", "")
+ ), # Send provider
+ # Action field
+ "action/operation": str(team_id) if team_id else "", # Send team_id
# Line item details
- 'lineitem/type': 'Usage', # Standard usage line item
+ "lineitem/type": "Usage", # Standard usage line item
}
# Add CZRN components that don't have direct CBF column mappings as resource tags
- cbf_record['resource/tag:provider'] = provider # CZRN provider component
- cbf_record['resource/tag:model'] = cloud_local_id # CZRN cloud-local-id component (model)
-
+ cbf_record["resource/tag:provider"] = provider # CZRN provider component
+ cbf_record[
+ "resource/tag:model"
+ ] = cloud_local_id # CZRN cloud-local-id component (model)
+
# Add resource tags for all dimensions (using resource/tag: format)
for key, value in dimensions.items():
- if value and value != 'N/A' and value != 'unknown': # Only add meaningful tags
- cbf_record[f'resource/tag:{key}'] = str(value)
+ if (
+ value and value != "N/A" and value != "unknown"
+ ): # Only add meaningful tags
+ cbf_record[f"resource/tag:{key}"] = str(value)
- # Add token breakdown as resource tags for analysis
+ # Add token breakdown as resource tags for analysis (excluding total_tokens per LIT-1907)
if prompt_tokens > 0:
- cbf_record['resource/tag:prompt_tokens'] = str(prompt_tokens)
+ cbf_record["resource/tag:prompt_tokens"] = str(prompt_tokens)
if completion_tokens > 0:
- cbf_record['resource/tag:completion_tokens'] = str(completion_tokens)
- if total_tokens > 0:
- cbf_record['resource/tag:total_tokens'] = str(total_tokens)
+ cbf_record["resource/tag:completion_tokens"] = str(completion_tokens)
return CBFRecord(cbf_record)
@@ -185,4 +221,3 @@ class CBFTransformer:
return None
return None
-
diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py
index 1652ec2aa0c..aa2a8121ee8 100644
--- a/litellm/integrations/custom_guardrail.py
+++ b/litellm/integrations/custom_guardrail.py
@@ -26,10 +26,16 @@ from litellm.types.utils import (
CallTypes,
GenericGuardrailAPIInputs,
GuardrailStatus,
+ GuardrailTracingDetail,
LLMResponseTypes,
StandardLoggingGuardrailInformation,
)
+try:
+ from fastapi.exceptions import HTTPException
+except ImportError:
+ HTTPException = None # type: ignore
+
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
dc = DualCache()
@@ -86,6 +92,9 @@ class CustomGuardrail(CustomLogger):
mask_request_content: bool = False,
mask_response_content: bool = False,
violation_message_template: Optional[str] = None,
+ end_session_after_n_fails: Optional[int] = None,
+ on_violation: Optional[str] = None,
+ realtime_violation_message: Optional[str] = None,
**kwargs,
):
"""
@@ -98,6 +107,9 @@ class CustomGuardrail(CustomLogger):
default_on: If True, the guardrail will be run by default on all requests
mask_request_content: If True, the guardrail will mask the request content
mask_response_content: If True, the guardrail will mask the response content
+ end_session_after_n_fails: For /v1/realtime sessions, end the session after this many violations
+ on_violation: For /v1/realtime sessions, 'warn' or 'end_session'
+ realtime_violation_message: Message the bot speaks aloud when a /v1/realtime guardrail fires
"""
self.guardrail_name = guardrail_name
self.supported_event_hooks = supported_event_hooks
@@ -108,6 +120,9 @@ class CustomGuardrail(CustomLogger):
self.mask_request_content: bool = mask_request_content
self.mask_response_content: bool = mask_response_content
self.violation_message_template: Optional[str] = violation_message_template
+ self.end_session_after_n_fails: Optional[int] = end_session_after_n_fails
+ self.on_violation: Optional[str] = on_violation
+ self.realtime_violation_message: Optional[str] = realtime_violation_message
if supported_event_hooks:
## validate event_hook is in supported_event_hooks
@@ -216,12 +231,23 @@ class CustomGuardrail(CustomLogger):
event_hook, supported_event_hooks
)
elif isinstance(event_hook, Mode):
+ tag_values_flat: list = []
+ for v in event_hook.tags.values():
+ if isinstance(v, list):
+ tag_values_flat.extend(v)
+ else:
+ tag_values_flat.append(v)
_validate_event_hook_list_is_in_supported_event_hooks(
- list(event_hook.tags.values()), supported_event_hooks
+ tag_values_flat, supported_event_hooks
)
if event_hook.default:
+ default_list = (
+ event_hook.default
+ if isinstance(event_hook.default, list)
+ else [event_hook.default]
+ )
_validate_event_hook_list_is_in_supported_event_hooks(
- [event_hook.default], supported_event_hooks
+ default_list, supported_event_hooks
)
elif isinstance(event_hook, GuardrailEventHooks):
if event_hook not in supported_event_hooks:
@@ -268,6 +294,7 @@ class CustomGuardrail(CustomLogger):
"""
Returns the guardrail(s) to be run from the metadata or root
"""
+
if "guardrails" in data:
return data["guardrails"]
metadata = data.get("litellm_metadata") or data.get("metadata", {})
@@ -399,7 +426,7 @@ class CustomGuardrail(CustomLogger):
"Setting tag-based guardrails is only available in litellm-enterprise. You must be a premium user to use this feature."
)
result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag(
- data, self.event_hook
+ data, self.event_hook, event_type
)
if result is not None:
return result
@@ -426,7 +453,7 @@ class CustomGuardrail(CustomLogger):
"Setting tag-based guardrails is only available in litellm-enterprise. You must be a premium user to use this feature."
)
result = EnterpriseCustomGuardrailHelper._should_run_if_mode_by_tag(
- data, self.event_hook
+ data, self.event_hook, event_type
)
if result is not None:
return result
@@ -445,7 +472,20 @@ class CustomGuardrail(CustomLogger):
if isinstance(self.event_hook, list):
return event_type.value in self.event_hook
if isinstance(self.event_hook, Mode):
- return event_type.value in self.event_hook.tags.values()
+ for tag_value in self.event_hook.tags.values():
+ if isinstance(tag_value, list):
+ if event_type.value in tag_value:
+ return True
+ elif event_type.value == tag_value:
+ return True
+ if self.event_hook.default:
+ default_list = (
+ self.event_hook.default
+ if isinstance(self.event_hook.default, list)
+ else [self.event_hook.default]
+ )
+ return event_type.value in default_list
+ return False
return self.event_hook == event_type.value
def get_guardrail_dynamic_request_body_params(self, request_data: dict) -> dict:
@@ -514,9 +554,15 @@ class CustomGuardrail(CustomLogger):
masked_entity_count: Optional[Dict[str, int]] = None,
guardrail_provider: Optional[str] = None,
event_type: Optional[GuardrailEventHooks] = None,
+ tracing_detail: Optional[GuardrailTracingDetail] = None,
) -> None:
"""
Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc.
+
+ Args:
+ tracing_detail: Optional typed dict with provider-specific tracing fields
+ (guardrail_id, policy_template, detection_method, confidence_score,
+ classification, match_details, patterns_checked, alert_recipients).
"""
if isinstance(guardrail_json_response, Exception):
guardrail_json_response = str(guardrail_json_response)
@@ -543,6 +589,16 @@ class CustomGuardrail(CustomLogger):
guardrail_json_response
)
+ # Strip secret_fields to prevent plaintext Authorization headers from
+ # being persisted to spend logs, OTEL traces, or other logging backends.
+ # This matches the pattern used by Langfuse and Arize integrations.
+ if isinstance(clean_guardrail_response, dict):
+ clean_guardrail_response.pop("secret_fields", None)
+ elif isinstance(clean_guardrail_response, list):
+ for item in clean_guardrail_response:
+ if isinstance(item, dict):
+ item.pop("secret_fields", None)
+
slg = StandardLoggingGuardrailInformation(
guardrail_name=self.guardrail_name,
guardrail_provider=guardrail_provider,
@@ -553,6 +609,7 @@ class CustomGuardrail(CustomLogger):
end_time=end_time,
duration=duration,
masked_entity_count=masked_entity_count,
+ **(tracing_detail or {}),
)
def _append_guardrail_info(container: dict) -> None:
@@ -573,9 +630,10 @@ class CustomGuardrail(CustomLogger):
elif "litellm_metadata" in request_data:
_append_guardrail_info(request_data["litellm_metadata"])
else:
- verbose_logger.warning(
- "unable to log guardrail information. No metadata found in request_data"
- )
+ # Ensure guardrail info is always logged (e.g. proxy may not have set
+ # metadata yet). Attach to "metadata" so spend log / standard logging see it.
+ request_data["metadata"] = {}
+ _append_guardrail_info(request_data["metadata"])
async def apply_guardrail(
self,
@@ -615,6 +673,7 @@ class CustomGuardrail(CustomLogger):
end_time: Optional[float] = None,
duration: Optional[float] = None,
event_type: Optional[GuardrailEventHooks] = None,
+ original_inputs: Optional[Dict] = None,
):
"""
Add StandardLoggingGuardrailInformation to the request data
@@ -622,7 +681,20 @@ class CustomGuardrail(CustomLogger):
This gets logged on downsteam Langfuse, DataDog, etc.
"""
# Convert None to empty dict to satisfy type requirements
- guardrail_response = {} if response is None else response
+ guardrail_response: Union[Dict[str, Any], str] = (
+ {} if response is None else response
+ )
+
+ # For apply_guardrail functions in custom_code_guardrail scenario,
+ # simplify the logged response to "allow", "deny", or "mask"
+ if original_inputs is not None and isinstance(response, dict):
+ # Check if inputs were modified by comparing them
+ if self._inputs_were_modified(original_inputs, response):
+ guardrail_response = "mask"
+ else:
+ guardrail_response = "allow"
+
+ verbose_logger.debug(f"Guardrail response: {response}")
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=guardrail_response,
@@ -635,6 +707,27 @@ class CustomGuardrail(CustomLogger):
)
return response
+ @staticmethod
+ def _is_guardrail_intervention(e: Exception) -> bool:
+ """
+ Returns True if the exception represents an intentional guardrail block
+ (this was logged previously as an API failure - guardrail_failed_to_respond).
+
+ Guardrails signal intentional blocks by raising:
+ - HTTPException with status 400 (content policy violation)
+ - ModifyResponseException (passthrough mode violation)
+ """
+
+ if isinstance(e, ModifyResponseException):
+ return True
+ if (
+ HTTPException is not None
+ and isinstance(e, HTTPException)
+ and e.status_code == 400
+ ):
+ return True
+ return False
+
def _process_error(
self,
e: Exception,
@@ -649,10 +742,21 @@ class CustomGuardrail(CustomLogger):
This gets logged on downsteam Langfuse, DataDog, etc.
"""
+ guardrail_status: GuardrailStatus = (
+ "guardrail_intervened"
+ if self._is_guardrail_intervention(e)
+ else "guardrail_failed_to_respond"
+ )
+ # For custom_code_guardrail scenario, log as "deny" instead of full exception
+ # Check if this is from custom_code_guardrail by checking the class name
+ guardrail_response: Union[Exception, str] = e
+ if "CustomCodeGuardrail" in self.__class__.__name__:
+ guardrail_response = "deny"
+
self.add_standard_logging_guardrail_information_to_request_data(
- guardrail_json_response=e,
+ guardrail_json_response=guardrail_response,
request_data=request_data,
- guardrail_status="guardrail_failed_to_respond",
+ guardrail_status=guardrail_status,
duration=duration,
start_time=start_time,
end_time=end_time,
@@ -660,6 +764,25 @@ class CustomGuardrail(CustomLogger):
)
raise e
+ def _inputs_were_modified(self, original_inputs: Dict, response: Dict) -> bool:
+ """
+ Compare original inputs with response to determine if content was modified.
+
+ Returns True if the inputs were modified (mask scenario), False otherwise (allow scenario).
+ """
+ # Get all keys from both dictionaries
+ all_keys = set(original_inputs.keys()) | set(response.keys())
+
+ # Compare each key's value
+ for key in all_keys:
+ original_value = original_inputs.get(key)
+ response_value = response.get(key)
+ if original_value != response_value:
+ return True
+
+ # No modifications detected
+ return False
+
def mask_content_in_string(
self,
content_string: str,
@@ -743,8 +866,8 @@ def log_guardrail_information(func):
- during_call
- post_call
"""
- import asyncio
import functools
+ import inspect
def _infer_event_type_from_function_name(
func_name: str,
@@ -767,6 +890,12 @@ def log_guardrail_information(func):
self: CustomGuardrail = args[0]
request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {}
event_type = _infer_event_type_from_function_name(func.__name__)
+
+ # Store original inputs for comparison (for apply_guardrail functions)
+ original_inputs = None
+ if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
+ original_inputs = kwargs.get("inputs")
+
try:
response = await func(*args, **kwargs)
return self._process_response(
@@ -776,6 +905,7 @@ def log_guardrail_information(func):
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
+ original_inputs=original_inputs,
)
except Exception as e:
return self._process_error(
@@ -793,6 +923,12 @@ def log_guardrail_information(func):
self: CustomGuardrail = args[0]
request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {}
event_type = _infer_event_type_from_function_name(func.__name__)
+
+ # Store original inputs for comparison (for apply_guardrail functions)
+ original_inputs = None
+ if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
+ original_inputs = kwargs.get("inputs")
+
try:
response = func(*args, **kwargs)
return self._process_response(
@@ -800,6 +936,7 @@ def log_guardrail_information(func):
request_data=request_data,
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
+ original_inputs=original_inputs,
)
except Exception as e:
return self._process_error(
@@ -811,7 +948,7 @@ def log_guardrail_information(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
- if asyncio.iscoroutinefunction(func):
+ if inspect.iscoroutinefunction(func):
return async_wrapper(*args, **kwargs)
return sync_wrapper(*args, **kwargs)
diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py
index 07d237c4758..06ba9675ca2 100644
--- a/litellm/integrations/custom_logger.py
+++ b/litellm/integrations/custom_logger.py
@@ -377,6 +377,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_headers: Optional[Dict[str, str]] = None,
+ litellm_call_info: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, str]]:
"""
Called after an LLM API call (success or failure) to allow injecting custom HTTP response headers.
@@ -386,6 +387,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
- user_api_key_dict: UserAPIKeyAuth - The user API key dictionary.
- response: Any - The response object (None for failure cases).
- request_headers: Optional[Dict[str, str]] - The original request headers.
+ - litellm_call_info: Optional[Dict[str, Any]] - Normalized routing metadata:
+ - custom_llm_provider: str - The LLM provider (e.g. "openai", "azure")
+ - model_info: dict - The model_info from router config
+ - api_base: str - The API base URL used
+ - model_id: str - The deployment model ID
Returns:
- Optional[Dict[str, str]]: A dictionary of headers to inject into the HTTP response.
@@ -665,6 +671,37 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
pass
+ async def async_should_run_chat_completion_agentic_loop(
+ self,
+ response: Any,
+ model: str,
+ messages: List[Dict],
+ tools: Optional[List[Dict]],
+ stream: bool,
+ custom_llm_provider: str,
+ kwargs: Dict,
+ ) -> Tuple[bool, Dict]:
+ """
+ Hook to determine if chat completion agentic loop should be executed.
+ """
+ return False, {}
+
+ async def async_run_chat_completion_agentic_loop(
+ self,
+ tools: Dict,
+ model: str,
+ messages: List[Dict],
+ response: Any,
+ optional_params: Dict,
+ logging_obj: "LiteLLMLoggingObj",
+ stream: bool,
+ kwargs: Dict,
+ ) -> Any:
+ """
+ Hook to execute chat completion agentic loop based on context from should_run hook.
+ """
+ pass
+
# Useful helpers for custom logger classes
def truncate_standard_logging_payload_content(
@@ -743,15 +780,17 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
self, model_call_details: Dict
) -> Dict:
"""
- Only redacts messages and responses when self.turn_off_message_logging is True
+ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it.
+ This method handles two features:
+ 1. turn_off_message_logging: When True, redacts messages and responses
+ 2. standard_logging_payload_excluded_fields: Removes specified fields entirely
- By default, self.turn_off_message_logging is False and this does nothing.
-
- Return a redacted deepcopy of the provided logging payload.
+ Return a modified copy of the provided logging payload.
This is useful for logging payloads that contain sensitive information.
"""
+ import litellm
from copy import copy
from litellm import Choices, Message, ModelResponse
@@ -759,14 +798,17 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
turn_off_message_logging: bool = getattr(
self, "turn_off_message_logging", False
)
+ excluded_fields: Optional[List[str]] = getattr(
+ litellm, "standard_logging_payload_excluded_fields", None
+ )
- if turn_off_message_logging is False:
+ # Early return if no processing needed
+ if turn_off_message_logging is False and not excluded_fields:
return model_call_details
# Only make a shallow copy of the top-level dict to avoid deepcopy issues
# with complex objects like AuthenticationError that may be present
model_call_details_copy = copy(model_call_details)
- redacted_str = "redacted-by-litellm"
standard_logging_object = model_call_details.get("standard_logging_object")
if standard_logging_object is None:
return model_call_details_copy
@@ -774,43 +816,62 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
# Make a copy of just the standard_logging_object to avoid modifying the original
standard_logging_object_copy = copy(standard_logging_object)
- if standard_logging_object_copy.get("messages") is not None:
- standard_logging_object_copy["messages"] = [
- Message(content=redacted_str).model_dump()
- ]
+ # Handle excluded fields - remove them entirely from the payload
+ if excluded_fields:
+ for field in excluded_fields:
+ if field in standard_logging_object_copy:
+ del standard_logging_object_copy[field]
- if standard_logging_object_copy.get("response") is not None:
- response = standard_logging_object_copy["response"]
- # Check if this is a ResponsesAPIResponse (has "output" field)
- if isinstance(response, dict) and "output" in response:
- # Make a copy to avoid modifying the original
- from copy import deepcopy
+ # Handle turn_off_message_logging - redact messages and responses (if not already excluded)
+ if turn_off_message_logging:
+ redacted_str = "redacted-by-litellm"
- response_copy = deepcopy(response)
- # Redact content in output array
- if isinstance(response_copy.get("output"), list):
- for output_item in response_copy["output"]:
- if isinstance(output_item, dict) and "content" in output_item:
- if isinstance(output_item["content"], list):
- # Redact text in content items
- for content_item in output_item["content"]:
- if (
- isinstance(content_item, dict)
- and "text" in content_item
- ):
- content_item["text"] = redacted_str
- standard_logging_object_copy["response"] = response_copy
- else:
- # Standard ModelResponse format
- model_response = ModelResponse(
- choices=[Choices(message=Message(content=redacted_str))]
- )
- model_response_dict = model_response.model_dump()
- standard_logging_object_copy["response"] = model_response_dict
+ if (
+ "messages" not in (excluded_fields or [])
+ and standard_logging_object_copy.get("messages") is not None
+ ):
+ standard_logging_object_copy["messages"] = [
+ Message(content=redacted_str).model_dump()
+ ]
- model_call_details_copy["standard_logging_object"] = (
- standard_logging_object_copy
- )
+ if (
+ "response" not in (excluded_fields or [])
+ and standard_logging_object_copy.get("response") is not None
+ ):
+ response = standard_logging_object_copy["response"]
+ # Check if this is a ResponsesAPIResponse (has "output" field)
+ if isinstance(response, dict) and "output" in response:
+ # Make a copy to avoid modifying the original
+ from copy import deepcopy
+
+ response_copy = deepcopy(response)
+ # Redact content in output array
+ if isinstance(response_copy.get("output"), list):
+ for output_item in response_copy["output"]:
+ if (
+ isinstance(output_item, dict)
+ and "content" in output_item
+ ):
+ if isinstance(output_item["content"], list):
+ # Redact text in content items
+ for content_item in output_item["content"]:
+ if (
+ isinstance(content_item, dict)
+ and "text" in content_item
+ ):
+ content_item["text"] = redacted_str
+ standard_logging_object_copy["response"] = response_copy
+ else:
+ # Standard ModelResponse format
+ model_response = ModelResponse(
+ choices=[Choices(message=Message(content=redacted_str))]
+ )
+ model_response_dict = model_response.model_dump()
+ standard_logging_object_copy["response"] = model_response_dict
+
+ model_call_details_copy[
+ "standard_logging_object"
+ ] = standard_logging_object_copy
return model_call_details_copy
async def get_proxy_server_request_from_cold_storage_with_object_key(
diff --git a/litellm/integrations/custom_secret_manager.py b/litellm/integrations/custom_secret_manager.py
index 2125aef2200..45ffa2e08cf 100644
--- a/litellm/integrations/custom_secret_manager.py
+++ b/litellm/integrations/custom_secret_manager.py
@@ -100,9 +100,7 @@ class CustomSecretManager(BaseSecretManager):
"""
super().__init__()
self.secret_manager_name = secret_manager_name or "custom_secret_manager"
- verbose_logger.info(
- "Initialized custom secret manager"
- )
+ verbose_logger.info("Initialized custom secret manager")
@abstractmethod
async def async_read_secret(
diff --git a/litellm/integrations/custom_sso_handler.py b/litellm/integrations/custom_sso_handler.py
index bc80966f8ca..7f60decabc3 100644
--- a/litellm/integrations/custom_sso_handler.py
+++ b/litellm/integrations/custom_sso_handler.py
@@ -13,6 +13,7 @@ class CustomSSOLoginHandler(CustomLogger):
Useful when you have an OAuth proxy in front of LiteLLM
and you want to use the headers from the proxy to sign in the user
"""
+
async def handle_custom_ui_sso_sign_in(
self,
request: Request,
@@ -26,4 +27,4 @@ class CustomSSOLoginHandler(CustomLogger):
display_name="Test",
picture="https://test.com/test.png",
provider="test",
- )
\ No newline at end of file
+ )
diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py
index 127b0e53fa8..64e0b26a8e7 100644
--- a/litellm/integrations/datadog/datadog.py
+++ b/litellm/integrations/datadog/datadog.py
@@ -45,7 +45,14 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
-from litellm.types.integrations.datadog import *
+from litellm.types.integrations.datadog import (
+ DD_ERRORS,
+ DD_MAX_BATCH_SIZE,
+ DataDogStatus,
+ DatadogInitParams,
+ DatadogPayload,
+ DatadogProxyFailureHookJsonMessage,
+)
from litellm.types.services import ServiceLoggerPayload, ServiceTypes
from litellm.types.utils import StandardLoggingPayload
@@ -85,12 +92,14 @@ class DataDogLogger(
"""
try:
verbose_logger.debug("Datadog: in init datadog logger")
-
+
self.is_mock_mode = should_use_datadog_mock()
-
+
if self.is_mock_mode:
create_mock_datadog_client()
- verbose_logger.debug("[DATADOG MOCK] Datadog logger initialized in mock mode")
+ verbose_logger.debug(
+ "[DATADOG MOCK] Datadog logger initialized in mock mode"
+ )
#########################################################
# Handle datadog_params set as litellm.datadog_params
@@ -209,6 +218,96 @@ class DataDogLogger(
)
pass
+ async def async_post_call_failure_hook(
+ self,
+ request_data: dict,
+ original_exception: Exception,
+ user_api_key_dict: Any,
+ traceback_str: Optional[str] = None,
+ ) -> Optional[Any]:
+ """
+ Log proxy-level failures (e.g. 401 auth, DB connection errors) to Datadog.
+
+ Ensures failures that occur before or outside the LLM completion flow
+ (e.g. ConnectError during auth when DB is down) are visible in Datadog
+ alongside Prometheus.
+ """
+ try:
+ from litellm.litellm_core_utils.litellm_logging import (
+ StandardLoggingPayloadSetup,
+ )
+ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+
+ error_information = StandardLoggingPayloadSetup.get_error_information(
+ original_exception=original_exception,
+ traceback_str=traceback_str,
+ )
+ _code = error_information.get("error_code") or ""
+ status_code: Optional[int] = None
+ if _code and str(_code).strip().isdigit():
+ status_code = int(_code)
+
+ # Use project-standard sanitized user context when running in proxy
+ user_context: Dict[str, Any] = {}
+ try:
+ from litellm.proxy.litellm_pre_call_utils import (
+ LiteLLMProxyRequestSetup,
+ )
+
+ _meta = (
+ LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
+ user_api_key_dict=user_api_key_dict
+ )
+ )
+ user_context = dict(_meta) if isinstance(_meta, dict) else _meta
+ except Exception:
+ # Fallback if proxy not available (e.g. SDK-only): minimal safe fields
+ if hasattr(user_api_key_dict, "request_route"):
+ user_context["request_route"] = getattr(
+ user_api_key_dict, "request_route", None
+ )
+ if hasattr(user_api_key_dict, "team_id"):
+ user_context["team_id"] = getattr(
+ user_api_key_dict, "team_id", None
+ )
+ if hasattr(user_api_key_dict, "user_id"):
+ user_context["user_id"] = getattr(
+ user_api_key_dict, "user_id", None
+ )
+ if hasattr(user_api_key_dict, "end_user_id"):
+ user_context["end_user_id"] = getattr(
+ user_api_key_dict, "end_user_id", None
+ )
+
+ message_payload: DatadogProxyFailureHookJsonMessage = {
+ "exception": error_information.get("error_message")
+ or str(original_exception),
+ "error_class": error_information.get("error_class")
+ or original_exception.__class__.__name__,
+ "status_code": status_code,
+ "traceback": error_information.get("traceback") or "",
+ "user_api_key_dict": user_context,
+ }
+
+ dd_payload = DatadogPayload(
+ ddsource=get_datadog_source(),
+ ddtags=get_datadog_tags(),
+ hostname=get_datadog_hostname(),
+ message=safe_dumps(message_payload),
+ service=get_datadog_service(),
+ status=DataDogStatus.ERROR,
+ )
+ self._add_trace_context_to_payload(dd_payload=dd_payload)
+ self.log_queue.append(dd_payload)
+
+ if len(self.log_queue) >= self.batch_size:
+ await self.async_send_batch()
+ except Exception as e:
+ verbose_logger.exception(
+ f"Datadog: async_post_call_failure_hook - {str(e)}\n{traceback.format_exc()}"
+ )
+ return None
+
async def async_send_batch(self):
"""
Sends the in memory logs queue to datadog api
@@ -230,9 +329,11 @@ class DataDogLogger(
len(self.log_queue),
self.intake_url,
)
-
+
if self.is_mock_mode:
- verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted")
+ verbose_logger.debug(
+ "[DATADOG MOCK] Mock mode enabled - API calls will be intercepted"
+ )
response = await self.async_send_compressed_data(self.log_queue)
if response.status_code == 413:
diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py
index 2eb94b59dd8..a961d4f9244 100644
--- a/litellm/integrations/datadog/datadog_cost_management.py
+++ b/litellm/integrations/datadog/datadog_cost_management.py
@@ -93,7 +93,9 @@ class DatadogCostManagementLogger(CustomBatchLogger):
Aggregates costs by Provider, Model, and Date.
Returns a list of DatadogFOCUSCostEntry.
"""
- aggregator: Dict[Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry] = {}
+ aggregator: Dict[
+ Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry
+ ] = {}
for log in logs:
try:
@@ -167,10 +169,20 @@ class DatadogCostManagementLogger(CustomBatchLogger):
metadata = log.get("metadata", {})
if metadata:
# Add user info
- if "user_api_key_alias" in metadata:
+ # Add user info
+ if metadata.get("user_api_key_alias"):
tags["user"] = str(metadata["user_api_key_alias"])
- if "user_api_key_team_alias" in metadata:
- tags["team"] = str(metadata["user_api_key_team_alias"])
+
+ # Add Team Tag
+ team_tag = (
+ metadata.get("user_api_key_team_alias")
+ or metadata.get("team_alias") # type: ignore
+ or metadata.get("user_api_key_team_id")
+ or metadata.get("team_id") # type: ignore
+ )
+
+ if team_tag:
+ tags["team"] = str(team_tag)
# model_group is not in StandardLoggingMetadata TypedDict, so we need to access it via dict.get()
model_group = metadata.get("model_group") # type: ignore[misc]
if model_group:
diff --git a/litellm/integrations/datadog/datadog_handler.py b/litellm/integrations/datadog/datadog_handler.py
index e2f30f2f614..0406f1e5d20 100644
--- a/litellm/integrations/datadog/datadog_handler.py
+++ b/litellm/integrations/datadog/datadog_handler.py
@@ -55,4 +55,15 @@ def get_datadog_tags(
request_tags = standard_logging_object.get("request_tags", []) or []
tags.extend(f"request_tag:{tag}" for tag in request_tags)
+ # Add Team Tag
+ metadata = standard_logging_object.get("metadata", {}) or {}
+ team_tag = (
+ metadata.get("user_api_key_team_alias")
+ or metadata.get("team_alias")
+ or metadata.get("user_api_key_team_id")
+ or metadata.get("team_id")
+ )
+ if team_tag:
+ tags.append(f"team:{team_tag}")
+
return ",".join(tags)
diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py
index e5ce9997491..de6cc02fa3d 100644
--- a/litellm/integrations/datadog/datadog_llm_obs.py
+++ b/litellm/integrations/datadog/datadog_llm_obs.py
@@ -48,13 +48,15 @@ class DataDogLLMObsLogger(CustomBatchLogger):
def __init__(self, **kwargs):
try:
verbose_logger.debug("DataDogLLMObs: Initializing logger")
-
+
self.is_mock_mode = should_use_datadog_mock()
-
+
if self.is_mock_mode:
create_mock_datadog_client()
- verbose_logger.debug("[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode")
-
+ verbose_logger.debug(
+ "[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode"
+ )
+
# Configure DataDog endpoint (Agent or Direct API)
# Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
# Check for agent mode FIRST - agent mode doesn't require DD_API_KEY or DD_SITE
@@ -189,9 +191,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
verbose_logger.debug(
f"DataDogLLMObs: Flushing {len(self.log_queue)} events"
)
-
+
if self.is_mock_mode:
- verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted")
+ verbose_logger.debug(
+ "[DATADOG MOCK] Mock mode enabled - API calls will be intercepted"
+ )
# Prepare the payload
payload = {
diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py
new file mode 100644
index 00000000000..fcf40701e28
--- /dev/null
+++ b/litellm/integrations/datadog/datadog_metrics.py
@@ -0,0 +1,286 @@
+import asyncio
+import gzip
+import os
+import time
+from datetime import datetime
+from typing import List, Optional, Union
+
+from litellm._logging import verbose_logger
+from litellm.integrations.custom_batch_logger import CustomBatchLogger
+from litellm.integrations.datadog.datadog_handler import (
+ get_datadog_env,
+ get_datadog_hostname,
+ get_datadog_pod_name,
+ get_datadog_service,
+)
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+from litellm.llms.custom_httpx.http_handler import (
+ get_async_httpx_client,
+ httpxSpecialProvider,
+)
+from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
+from litellm.types.integrations.datadog_metrics import (
+ DatadogMetricPoint,
+ DatadogMetricSeries,
+ DatadogMetricsPayload,
+)
+from litellm.types.utils import StandardLoggingPayload
+
+
+class DatadogMetricsLogger(CustomBatchLogger):
+ def __init__(self, start_periodic_flush: bool = True, **kwargs):
+ self.dd_api_key = os.getenv("DD_API_KEY")
+ self.dd_app_key = os.getenv("DD_APP_KEY")
+ self.dd_site = os.getenv("DD_SITE", "datadoghq.com")
+
+ if not self.dd_api_key:
+ verbose_logger.warning(
+ "Datadog Metrics: DD_API_KEY is required. Integration will not work."
+ )
+
+ self.upload_url = f"https://api.{self.dd_site}/api/v2/series"
+
+ self.async_client = get_async_httpx_client(
+ llm_provider=httpxSpecialProvider.LoggingCallback
+ )
+
+ # Initialize lock
+ self.flush_lock = asyncio.Lock()
+
+ # Only set flush_lock if not already provided by caller
+ if "flush_lock" not in kwargs:
+ kwargs["flush_lock"] = self.flush_lock
+
+ # Send metrics more quickly to datadog (every 5 seconds)
+ if "flush_interval" not in kwargs:
+ kwargs["flush_interval"] = 5
+
+ super().__init__(**kwargs)
+
+ # Start periodic flush task only if instructed
+ if start_periodic_flush:
+ asyncio.create_task(self.periodic_flush())
+
+ def _extract_tags(
+ self,
+ log: StandardLoggingPayload,
+ status_code: Optional[Union[str, int]] = None,
+ ) -> List[str]:
+ """
+ Builds the list of tags for a Datadog metric point
+ """
+ # Base tags
+ tags = [
+ f"env:{get_datadog_env()}",
+ f"service:{get_datadog_service()}",
+ f"version:{os.getenv('DD_VERSION', 'unknown')}",
+ f"HOSTNAME:{get_datadog_hostname()}",
+ f"POD_NAME:{get_datadog_pod_name()}",
+ ]
+
+ # Add metric-specific tags
+ if provider := log.get("custom_llm_provider"):
+ tags.append(f"provider:{provider}")
+
+ if model := log.get("model"):
+ tags.append(f"model_name:{model}")
+
+ if model_group := log.get("model_group"):
+ tags.append(f"model_group:{model_group}")
+
+ if status_code is not None:
+ tags.append(f"status_code:{status_code}")
+
+ # Extract team tag
+ metadata = log.get("metadata", {}) or {}
+ team_tag = (
+ metadata.get("user_api_key_team_alias")
+ or metadata.get("team_alias") # type: ignore
+ or metadata.get("user_api_key_team_id")
+ or metadata.get("team_id") # type: ignore
+ )
+
+ if team_tag:
+ tags.append(f"team:{team_tag}")
+
+ return tags
+
+ def _add_metrics_from_log(
+ self,
+ log: StandardLoggingPayload,
+ kwargs: dict,
+ status_code: Union[str, int] = "200",
+ ):
+ """
+ Extracts latencies and appends Datadog metric series to the queue
+ """
+ tags = self._extract_tags(log, status_code=status_code)
+
+ # We record metrics with the end_time as the timestamp for the point
+ end_time_dt = kwargs.get("end_time") or datetime.now()
+ timestamp = int(end_time_dt.timestamp())
+
+ # 1. Total Request Latency Metric (End to End)
+ start_time_dt = kwargs.get("start_time")
+ if start_time_dt and end_time_dt:
+ total_duration = (end_time_dt - start_time_dt).total_seconds()
+ series_total_latency: DatadogMetricSeries = {
+ "metric": "litellm.request.total_latency",
+ "type": 3, # gauge
+ "points": [{"timestamp": timestamp, "value": total_duration}],
+ "tags": tags,
+ }
+ self.log_queue.append(series_total_latency)
+
+ # 2. LLM API Latency Metric (Provider alone)
+ api_call_start_time = kwargs.get("api_call_start_time")
+ if api_call_start_time and end_time_dt:
+ llm_api_duration = (end_time_dt - api_call_start_time).total_seconds()
+ series_llm_latency: DatadogMetricSeries = {
+ "metric": "litellm.llm_api.latency",
+ "type": 3, # gauge
+ "points": [{"timestamp": timestamp, "value": llm_api_duration}],
+ "tags": tags,
+ }
+ self.log_queue.append(series_llm_latency)
+
+ # 3. Request Count / Status Code
+ series_count: DatadogMetricSeries = {
+ "metric": "litellm.llm_api.request_count",
+ "type": 1, # count
+ "points": [{"timestamp": timestamp, "value": 1.0}],
+ "tags": tags,
+ "interval": self.flush_interval,
+ }
+ self.log_queue.append(series_count)
+
+ async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+ try:
+ standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
+ "standard_logging_object", None
+ )
+
+ if standard_logging_object is None:
+ return
+
+ self._add_metrics_from_log(
+ log=standard_logging_object, kwargs=kwargs, status_code="200"
+ )
+
+ if len(self.log_queue) >= self.batch_size:
+ await self.flush_queue()
+
+ except Exception as e:
+ verbose_logger.exception(
+ f"Datadog Metrics: Error in async_log_success_event: {str(e)}"
+ )
+
+ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+ try:
+ standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
+ "standard_logging_object", None
+ )
+
+ if standard_logging_object is None:
+ return
+
+ # Extract status code from error information
+ status_code = "500" # default
+ error_information = (
+ standard_logging_object.get("error_information", {}) or {}
+ )
+ error_code = error_information.get("error_code") # type: ignore
+ if error_code is not None:
+ status_code = str(error_code)
+
+ self._add_metrics_from_log(
+ log=standard_logging_object, kwargs=kwargs, status_code=status_code
+ )
+
+ if len(self.log_queue) >= self.batch_size:
+ await self.flush_queue()
+
+ except Exception as e:
+ verbose_logger.exception(
+ f"Datadog Metrics: Error in async_log_failure_event: {str(e)}"
+ )
+
+ async def async_send_batch(self):
+ if not self.log_queue:
+ return
+
+ batch = self.log_queue.copy()
+ payload_data: DatadogMetricsPayload = {"series": batch}
+
+ try:
+ await self._upload_to_datadog(payload_data)
+ except Exception as e:
+ verbose_logger.exception(
+ f"Datadog Metrics: Error in async_send_batch: {str(e)}"
+ )
+ raise
+
+ async def _upload_to_datadog(self, payload: DatadogMetricsPayload):
+ if not self.dd_api_key:
+ return
+
+ headers = {
+ "Content-Type": "application/json",
+ "DD-API-KEY": self.dd_api_key,
+ }
+
+ if self.dd_app_key:
+ headers["DD-APPLICATION-KEY"] = self.dd_app_key
+
+ json_data = safe_dumps(payload)
+ compressed_data = gzip.compress(json_data.encode("utf-8"))
+ headers["Content-Encoding"] = "gzip"
+
+ response = await self.async_client.post(
+ self.upload_url, content=compressed_data, headers=headers # type: ignore
+ )
+
+ response.raise_for_status()
+
+ verbose_logger.debug(
+ f"Datadog Metrics: Uploaded {len(payload['series'])} metric points. Status: {response.status_code}"
+ )
+
+ async def async_health_check(self) -> IntegrationHealthCheckStatus:
+ """
+ Check if the service is healthy
+ """
+ try:
+ # Send a test metric point to Datadog
+ test_metric_point: DatadogMetricPoint = {
+ "timestamp": int(time.time()),
+ "value": 1.0,
+ }
+ test_metric_series: DatadogMetricSeries = {
+ "metric": "litellm.health_check",
+ "type": 3, # Gauge
+ "points": [test_metric_point],
+ "tags": ["env:health_check"],
+ }
+
+ payload_data: DatadogMetricsPayload = {"series": [test_metric_series]}
+
+ await self._upload_to_datadog(payload_data)
+
+ return IntegrationHealthCheckStatus(
+ status="healthy",
+ error_message=None,
+ )
+ except Exception as e:
+ return IntegrationHealthCheckStatus(
+ status="unhealthy",
+ error_message=str(e),
+ )
+
+ async def get_request_response_payload(
+ self,
+ request_id: str,
+ start_time_utc: Optional[datetime],
+ end_time_utc: Optional[datetime],
+ ) -> Optional[dict]:
+ pass
diff --git a/litellm/integrations/datadog/datadog_mock_client.py b/litellm/integrations/datadog/datadog_mock_client.py
index a0a760deb0b..7f9beab72cc 100644
--- a/litellm/integrations/datadog/datadog_mock_client.py
+++ b/litellm/integrations/datadog/datadog_mock_client.py
@@ -8,7 +8,10 @@ Usage:
Set DATADOG_MOCK=true in environment variables or config to enable mock mode.
"""
-from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory
+from litellm.integrations.mock_client_factory import (
+ MockClientConfig,
+ create_mock_client_factory,
+)
# Create mock client using factory
_config = MockClientConfig(
@@ -25,4 +28,6 @@ _config = MockClientConfig(
patch_sync_client=True,
)
-create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory(_config)
+create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory(
+ _config
+)
diff --git a/litellm/integrations/dotprompt/__init__.py b/litellm/integrations/dotprompt/__init__.py
index 3847c8fa192..394929f4a25 100644
--- a/litellm/integrations/dotprompt/__init__.py
+++ b/litellm/integrations/dotprompt/__init__.py
@@ -25,6 +25,7 @@ def set_global_prompt_directory(directory: str) -> None:
litellm.global_prompt_directory = directory # type: ignore
+
def _get_prompt_data_from_dotprompt_content(dotprompt_content: str) -> dict:
"""
Get the prompt data from the dotprompt content.
@@ -36,12 +37,10 @@ def _get_prompt_data_from_dotprompt_content(dotprompt_content: str) -> dict:
# Parse the dotprompt content to extract frontmatter and content
temp_manager = PromptManager()
metadata, content = temp_manager._parse_frontmatter(dotprompt_content)
-
+
# Convert to prompt_data format
- return {
- "content": content.strip(),
- "metadata": metadata
- }
+ return {"content": content.strip(), "metadata": metadata}
+
def prompt_initializer(
litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec"
@@ -58,7 +57,7 @@ def prompt_initializer(
)
prompt_file = getattr(litellm_params, "prompt_file", None)
-
+
# Handle dotprompt_content from database
dotprompt_content = getattr(litellm_params, "dotprompt_content", None)
if dotprompt_content and not prompt_data and not prompt_file:
@@ -74,7 +73,6 @@ def prompt_initializer(
return dot_prompt_manager
except Exception as e:
-
raise e
diff --git a/litellm/integrations/dotprompt/dotprompt_manager.py b/litellm/integrations/dotprompt/dotprompt_manager.py
index 9412ac3c842..37fdf7da693 100644
--- a/litellm/integrations/dotprompt/dotprompt_manager.py
+++ b/litellm/integrations/dotprompt/dotprompt_manager.py
@@ -128,7 +128,6 @@ class DotpromptManager(CustomPromptManagement):
raise ValueError("prompt_id is required for dotprompt manager")
try:
-
# Get the prompt template (versioned or base)
template = self.prompt_manager.get_prompt(
prompt_id=prompt_id, version=prompt_version
@@ -205,7 +204,6 @@ class DotpromptManager(CustomPromptManagement):
ignore_prompt_manager_model: Optional[bool] = False,
ignore_prompt_manager_optional_params: Optional[bool] = False,
) -> Tuple[str, List[AllMessageValues], dict]:
-
from litellm.integrations.prompt_management_base import PromptManagementBase
return PromptManagementBase.get_chat_completion_prompt(
diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py
index fc5a325ffe1..997a40d545e 100644
--- a/litellm/integrations/dotprompt/prompt_manager.py
+++ b/litellm/integrations/dotprompt/prompt_manager.py
@@ -205,7 +205,7 @@ class PromptManager:
"""
# Get the template (versioned or base)
template = self.get_prompt(prompt_id=prompt_id, version=version)
-
+
if template is None:
available_prompts = list(self.prompts.keys())
version_str = f" (version {version})" if version else ""
@@ -266,11 +266,11 @@ class PromptManager:
) -> Optional[PromptTemplate]:
"""
Get a prompt template by ID and optional version.
-
+
Args:
prompt_id: The base prompt ID
version: Optional version number. If provided, looks for {prompt_id}.v{version}
-
+
Returns:
The prompt template if found, None otherwise
"""
@@ -279,7 +279,7 @@ class PromptManager:
versioned_id = f"{prompt_id}.v{version}"
if versioned_id in self.prompts:
return self.prompts[versioned_id]
-
+
# Fall back to base prompt_id
return self.prompts.get(prompt_id)
diff --git a/litellm/integrations/email_templates/key_rotated_email.py b/litellm/integrations/email_templates/key_rotated_email.py
index dab7172dc6a..9e6dd41378f 100644
--- a/litellm/integrations/email_templates/key_rotated_email.py
+++ b/litellm/integrations/email_templates/key_rotated_email.py
@@ -222,4 +222,3 @@ response = client.chat.completions.create(