scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

Merge branch 'litellm_internal_staging' into litellm_fix_bedrock_deepseek_thinking_leak

Browse source
This commit is contained in:
mateo-berri 2026-09-09 19:01:17 -07:00
commit 263c34fa26
240 changed files with 20006 additions and 3188 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -105,7 +105,7 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38271
"limit": 38269
},
"reportUnknownParameterType": {
"limit": 19584

13
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -1801,7 +1801,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Remove conflicting keys from data to avoid duplicate keyword arguments
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
for model_id, model_file_id in specific_model_file_id_mapping.items():
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
@ -1812,7 +1821,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prom_logger.record_managed_file_deleted(result="success")
if stored_file_object:
return stored_file_object
return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id})
elif delete_response:
delete_response.id = file_id
return delete_response

15
litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql Normal file
View file

@ -0,0 +1,15 @@
-- DropForeignKey
DO $$
BEGIN
IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey";
END IF;
END $$;
-- 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 CASCADE ON UPDATE CASCADE;
END IF;
END $$;

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])

4
litellm/__init__.py
View file

@ -546,7 +546,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None
#### PII MASKING ####
output_parse_pii: bool = False
#############################################
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete
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
@ -2405,3 +2405,5 @@ def __getattr__(name: str) -> Any:
# ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time
mark_litellm_import_complete()

24
litellm/_internal_context.py
View file

@ -6,9 +6,33 @@ be settable from user input. Context variables are scoped to the current
asyncio task and cannot be injected via HTTP request bodies.
"""
from collections.abc import Generator
from contextlib import contextmanager
from contextvars import ContextVar
from datetime import datetime, timezone
from typing import Final
# When True, suppresses async logging and billing for internal sub-calls
# (e.g., emulated file-search steps that make nested LLM calls).
is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False)
# One request prices its totals, its per-token-type lines and the rates it reports on
# separate code paths. Each reads the clock for off-peak pricing, so without a pinned
# moment they can land on either side of a window boundary and disagree with each other.
_billing_time: Final[ContextVar[datetime | None]] = ContextVar("billing_time", default=None)
@contextmanager
def pinned_billing_time(moment: datetime) -> Generator[None]:
"""Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read."""
token: Final = _billing_time.set(moment)
try:
yield
finally:
_billing_time.reset(token)
def current_billing_time() -> datetime:
"""The pinned billing moment, or now in UTC outside a pinned block."""
pinned: Final = _billing_time.get()
return pinned if pinned is not None else datetime.now(timezone.utc)

8
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -227,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
provider_specific_fields: Mapping[str, object]
def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
"""Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
string payload in ``input`` rather than ``arguments``; both map to
@ -755,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Tool calls accumulate into the single trailing tool_calls choice
# like the typed branches above; a choice per call would hide every
# call after choices[0] from chat clients
accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index))
accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
tool_call_index += 1
elif handle_raw_dict_callback is not None:
choice, index = handle_raw_dict_callback(item=raw_item, index=index)
@ -1409,7 +1409,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
# New output item added
output_item = parsed_chunk.get("item", {})
if output_item.get("type") in ("function_call", "custom_tool_call"):
converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
provider_specific_fields: Final = converted.get("provider_specific_fields")
function_chunk: Final = ChatCompletionToolCallFunctionChunk(
@ -1484,7 +1484,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
index=0,
delta=Delta(
tool_calls=(
_tool_call_dict_from_output_item(
tool_call_dict_from_output_item(
output_item, parsed_chunk.get("output_index", 0)
),
)

3
litellm/constants.py
View file

@ -143,6 +143,7 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float(
os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3)
)
MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150))
MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048
DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000
@ -197,6 +198,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [
"x-litellm-adaptive-router-model",
"x-litellm-applied-guardrails",
"x-litellm-guardrail-scan-id",
"x-litellm-guardrail-scan-metadata",
"x-litellm-cache-key",
]
@ -333,6 +335,7 @@ DEFAULT_SSL_CIPHERS: Final = os.getenv(
########### v2 Architecture constants for managing writing updates to the database ###########
REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer"
REDIS_GATEWAY_REQUESTS_BUFFER_KEY: Final = "litellm_gateway_requests_buffer"
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer"
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer"
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer"

13
litellm/cost_calculator.py
View file

@ -25,6 +25,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import
TranscriptionUsageObjectTransformation,
)
from litellm.litellm_core_utils.llm_cost_calc.utils import (
BilledTokenRates,
CostCalculatorUtils,
_generic_cost_per_character,
_get_regional_uplift_multiplier,
@ -45,6 +46,9 @@ from litellm.llms.azure.cost_calculation import (
from litellm.llms.azure_ai.cost_calculator import (
cost_per_token as azure_ai_cost_per_token,
)
from litellm.llms.azure_ai.cost_calculator import (
is_azure_model_router as azure_ai_is_model_router_name,
)
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.llms.bedrock.cost_calculation import (
cost_per_token as bedrock_cost_per_token,
@ -1122,6 +1126,7 @@ def _store_cost_breakdown_in_logging_obj(
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
billed_token_rates: BilledTokenRates | None = None,
) -> None:
"""
Helper function to store cost breakdown in the logging object.
@ -1166,6 +1171,7 @@ def _store_cost_breakdown_in_logging_obj(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
billed_token_rates=billed_token_rates,
)
except Exception as breakdown_error:
@ -1659,11 +1665,10 @@ def completion_cost(
data_residency=data_residency,
vertex_location=vertex_location,
response=completion_response,
request_model=request_model_for_cost,
)
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
if custom_llm_provider == "azure_ai":
if custom_llm_provider == "azure_ai" and not azure_ai_is_model_router_name(model):
model_for_additional_costs = request_model_for_cost
if completion_response is not None:
hidden_params = getattr(completion_response, "_hidden_params", None) or {}
@ -1735,6 +1740,7 @@ def completion_cost(
_reasoning_cost: float | None = None
_cache_read_cost: float | None = None
_cache_creation_cost: float | None = None
_billed_token_rates: BilledTokenRates | None = None
if cost_per_token_usage_object is not None and model:
_breakdown_provider: str | None = (
custom_llm_provider if isinstance(custom_llm_provider, str) else None
@ -1746,10 +1752,12 @@ def completion_cost(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
custom_cost_per_token=custom_cost_per_token,
)
_reasoning_cost = _token_type_breakdown.reasoning_cost
_cache_read_cost = _token_type_breakdown.cache_read_cost
_cache_creation_cost = _token_type_breakdown.cache_creation_cost
_billed_token_rates = _token_type_breakdown.rates
_store_cost_breakdown_in_logging_obj(
litellm_logging_obj=litellm_logging_obj,
prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar,
@ -1769,6 +1777,7 @@ def completion_cost(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
billed_token_rates=_billed_token_rates,
)
return _final_cost

30
litellm/experimental_mcp_client/client.py
View file

@ -18,6 +18,7 @@ from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServ
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from typing_extensions import Unpack
_TransportStreams: TypeAlias = tuple[
@ -56,10 +57,13 @@ def missing_streamable_http_client_error() -> ImportError:
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import (
ClientResult,
GetPromptRequestParams,
GetPromptResult,
Prompt,
ResourceTemplate,
ServerNotification,
ServerRequest,
TextContent,
)
from mcp.types import Tool as MCPTool
@ -146,8 +150,8 @@ _SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT)
otherwise carries JSON-RPC error codes."""
def _as_read_timeout(exc: BaseException) -> TimeoutError | None:
"""The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``.
def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None:
"""Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``.
The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a
field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error
@ -442,6 +446,18 @@ class MCPClient:
in_flight_error: BaseException | None = None
try:
read_stream, write_stream = transport[0], transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
async def receive_message(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None:
if not isinstance(message, (ValueError, httpx.RequestError, OSError)):
return
if not stream_error.done():
stream_error.set_result(message)
# The SDK closes pending requests when its message handler raises.
raise RuntimeError("MCP response stream failed")
# Build session kwargs with optional callbacks
session_kwargs: Final[dict[str, Any]] = {}
if self._sampling_callback is not None:
@ -456,6 +472,7 @@ class MCPClient:
read_stream,
write_stream,
read_timeout_seconds=timedelta(seconds=self.timeout),
message_handler=receive_message,
**session_kwargs,
)
session: Final = await session_ctx.__aenter__()
@ -467,6 +484,10 @@ class MCPClient:
if isinstance(ins, str) and ins.strip():
self._last_initialize_instructions = ins.strip()
return await operation(session)
except McpError:
if stream_error.done():
raise stream_error.result()
raise
finally:
try:
await session_ctx.__aexit__(None, None, None)
@ -501,11 +522,10 @@ class MCPClient:
transport_ctx, http_client = self._create_transport_context()
return await self._execute_session_operation(transport_ctx, operation)
except Exception as e:
read_timeout: Final = _as_read_timeout(e)
read_timeout: Final = as_mcp_read_timeout(e)
if read_timeout is not None:
verbose_logger.warning(
"MCP client timed out after %ss waiting for %s to answer; the server accepted the "
"request and ended its response stream without a JSON-RPC reply",
"MCP client timed out after %ss waiting for a valid MCP response from %s",
self.timeout,
self.server_url or "stdio",
)

2
litellm/files/main.py
View file

@ -31,7 +31,7 @@ FileCreateProvider = Literal[
FileRetrieveProvider = Literal[
"openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic"
]
FileDeleteProvider = Literal["openai", "azure", "gemini", "litellm_proxy", "manus", "anthropic"]
FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"]
FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"]
import litellm
from litellm import get_secret_str

4
litellm/integrations/azure_sentinel/azure_sentinel.py
View file

@ -21,6 +21,8 @@ from types import MappingProxyType
from typing import Final, TypeVar
from urllib.parse import urlparse
import httpx
from litellm._logging import verbose_logger
from litellm.integrations.batch_utils import (
BatchSendCancelled,
@ -418,7 +420,7 @@ class AzureSentinelLogger(CustomBatchLogger):
"Content-Type": "application/json",
}
async def _send_batch(batch: Sequence[_QueuedPayload]):
async def _send_batch(batch: Sequence[_QueuedPayload]) -> httpx.Response:
body: Final = safe_dumps(batch)
return await self.async_httpx_client.post(
url=api_endpoint,

30
litellm/integrations/custom_guardrail.py
View file

@ -850,20 +850,24 @@ class CustomGuardrail(CustomLogger):
if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True:
return None
# CHECK IF GUARDRAIL REJECTS THE REQUEST
target: Final = self._deployment_hook_target()
hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data
result: Final = await target.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=request_data.get("user_api_key_user_id"),
team_id=request_data.get("user_api_key_team_id"),
end_user_id=request_data.get("user_api_key_end_user_id"),
api_key=request_data.get("user_api_key_hash"),
request_route=request_data.get("user_api_key_request_route"),
),
data=hook_request_data,
response=response,
)
try:
if target is not self:
request_data["guardrail_to_apply"] = self # rebind-ok: dispatch consumes this key
result: Final = await target.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=request_data.get("user_api_key_user_id"),
team_id=request_data.get("user_api_key_team_id"),
end_user_id=request_data.get("user_api_key_end_user_id"),
api_key=request_data.get("user_api_key_hash"),
request_route=request_data.get("user_api_key_request_route"),
),
data=request_data,
response=response,
)
finally:
if target is not self:
request_data.pop("guardrail_to_apply", None)
if not self._is_valid_response_type(result):
return None

59
litellm/integrations/langfuse/langfuse.py
View file

@ -2,6 +2,7 @@
# On success, logs events to Langfuse
import inspect
import os
import re
import traceback
from collections.abc import Callable, Iterable, Mapping
from datetime import datetime
@ -63,6 +64,44 @@ def _object_mapping(value: object) -> Mapping[str, object] | None:
return value if isinstance(value, dict) else None
def _widened_items(mapping: Mapping[str, object]) -> Iterable[tuple[object, object]]:
"""Header pairs with the key type widened back to what a caller-supplied dict can actually hold."""
return mapping.items()
def _is_session_header_trace(trace_id: object, session_id: object, proxy_server_request: object) -> bool:
if not isinstance(trace_id, str) or not isinstance(session_id, str):
return False
request: Final = _object_mapping(proxy_server_request)
raw_headers: Final = _object_mapping(request.get("headers")) if request is not None else None
if raw_headers is None:
return False
headers: Final = MappingProxyType(
{key.lower(): value for key, value in _widened_items(raw_headers) if isinstance(key, str)}
)
if headers.get("x-litellm-trace-id"):
return False
if headers.get("langfuse_trace_id") is not None:
return False
if trace_id != session_id and headers.get("langfuse_session_id") != session_id:
return False
if headers.get("x-litellm-session-id") == trace_id:
return True
if re.fullmatch(r"[a-zA-Z0-9_\-]{8,}", trace_id) is None:
return False
user_agent: Final = headers.get("user-agent")
codex: Final = isinstance(user_agent, str) and re.match(r"^codex[-_ /]", user_agent, re.IGNORECASE) is not None
return any(
value == trace_id
and (
key == "x-session-id"
or re.fullmatch(r"x-.+-session-id", key) is not None
or (codex and key in ("session-id", "session_id", "thread-id", "conversation_id"))
)
for key, value in headers.items()
)
class _UsageObject(Protocol):
"""Token-count surface the Langfuse logger reads off a response usage payload."""
@ -609,6 +648,18 @@ class LangFuseLogger:
# This allows continuing an existing trace while still returning the correct trace_id
if existing_trace_id is not None:
trace_id = existing_trace_id
resolved_trace_id: Final = (
litellm_call_id or trace_id
if existing_trace_id is None
and _is_session_header_trace(trace_id, session_id, litellm_params.get("proxy_server_request"))
else trace_id
)
if resolved_trace_id != trace_id:
verbose_logger.debug(
"Langfuse: trace_id %s came from a session header; using call id %s so each call gets its own trace",
trace_id,
resolved_trace_id,
)
requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ()))
update_trace_keys: Final = (
requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else ()
@ -663,7 +714,7 @@ class LangFuseLogger:
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
else: # don't overwrite an existing trace
trace_params = {
"id": trace_id,
"id": resolved_trace_id,
"name": trace_name,
"session_id": session_id,
"input": masked_input if not mask_input else "redacted-by-litellm",
@ -845,13 +896,13 @@ class LangFuseLogger:
# Verify langfuse accepted our trace_id; if it differs, log a warning but still return our intended value
# to match expected test behavior
if hasattr(generation_client, "trace_id") and generation_client.trace_id:
if generation_client.trace_id != trace_id:
if generation_client.trace_id != resolved_trace_id:
verbose_logger.warning(
"Langfuse trace_id mismatch: set %s, but langfuse returned %s. Using our intended trace_id for consistency.",
trace_id,
resolved_trace_id,
generation_client.trace_id,
)
return trace_id, generation_id
return resolved_trace_id, generation_id
except Exception:
verbose_logger.error("Langfuse Layer Error - %s", traceback.format_exc())
return None, None

128
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -2,6 +2,7 @@
Pulls the cost + context window + provider route for known models from https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json
This can be disabled by setting the LITELLM_LOCAL_MODEL_COST_MAP environment variable to True.
The ``lite`` and ``litellm-proxy`` CLI entry points also use the bundled map without fetching.
```
export LITELLM_LOCAL_MODEL_COST_MAP=True
@ -13,11 +14,14 @@ import hashlib
import json
import os
import random
import sys
import threading
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, replace
from datetime import datetime, timezone
from importlib.resources import files
from pathlib import Path
from typing import Final, Protocol
import httpx
@ -33,6 +37,12 @@ from litellm.litellm_core_utils.fallback_generalizations import (
)
FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations"
_CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"})
def _is_cli_process() -> bool:
return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES
# Reserved top-level keys that are not model entries. They must be excluded
# from the model-count integrity check so a real upstream shrink can't be masked.
@ -176,6 +186,11 @@ class GetModelCostMap:
RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504})
MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3
MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0
_litellm_import_complete = threading.Event()
def mark_litellm_import_complete() -> None:
_litellm_import_complete.set()
@dataclass(frozen=True, slots=True)
@ -314,12 +329,13 @@ async def _fetch_remote_model_cost_map_with_retry(
def _fetch_remote_model_cost_map_with_retry_sync(
url: str,
timeout: int,
max_attempts: int,
attempts: range,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
) -> ModelCostMapReloadResult:
for attempt in range(1, max_attempts + 1):
max_attempts: Final = attempts.stop - 1
for attempt in attempts:
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
if not isinstance(outcome, _FetchAttemptRetryable):
return outcome
@ -520,6 +536,68 @@ def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMa
return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map))
def adopt_model_cost_map(
new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract
) -> int:
import litellm
from litellm import utils
litellm.model_cost = new_model_cost_map
utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation
litellm.add_known_models(model_cost_map=new_model_cost_map)
fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
utils.reapply_runtime_model_cost_registrations()
return fetched_model_count
def _retry_remote_fetch_in_background(
url: str,
timeout: int,
max_attempts: int,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
first_outcome: _FetchAttemptRetryable,
) -> None:
try:
first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng)
if isinstance(first_wait, ModelCostMapReloadUnavailable):
return
sleep(first_wait)
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
attempts=range(2, max_attempts + 1),
sleep=sleep,
rng=rng,
client=client,
)
if isinstance(result, ModelCostMapReloadUnavailable):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup",
url,
max_attempts,
)
return
_litellm_import_complete.wait()
if not GetModelCostMap.validate_model_cost_map(
fetched_map=result.model_cost_map,
backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache
):
verbose_logger.warning(
"LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s",
url,
)
return
finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map
_cost_map_source_info.source = "remote"
_cost_map_source_info.fallback_reason = None
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
adopt_model_cost_map(finalized)
except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays
verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e)
def get_model_cost_map(
url: str,
timeout: int = 5,
@ -531,10 +609,12 @@ def get_model_cost_map(
"""
Public entry point — returns the model cost map dict.
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only.
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set or this is a ``lite`` /
``litellm-proxy`` CLI process, uses the local backup only.
2. Otherwise fetches from ``url``, retrying transient HTTP errors
(429/5xx/transport) with Retry-After-aware backoff, validates
integrity, and falls back to the local backup on any failure.
(429/5xx/transport) with Retry-After-aware backoff in a background
thread, validates integrity, and falls back to the local backup on any
failure.
Only the backup model count is cached (a single int) for validation.
The full backup dict is only parsed when it must be *returned* as a
@ -543,7 +623,7 @@ def get_model_cost_map(
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
# Note: can't use get_secret_bool here — this runs during litellm.__init__
# before litellm._key_management_settings is set.
if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true":
if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process():
_cost_map_source_info.source = "local"
_cost_map_source_info.url = None
_cost_map_source_info.is_env_forced = True
@ -553,24 +633,34 @@ def get_model_cost_map(
_cost_map_source_info.url = url
_cost_map_source_info.is_env_forced = False
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
max_attempts=max_attempts,
sleep=sleep,
rng=rng if rng is not None else random.Random(),
client=client if client is not None else httpx,
)
if isinstance(result, ModelCostMapReloadUnavailable):
fetch_client: Final = client if client is not None else httpx
fetch_rng: Final = rng if rng is not None else random.Random()
outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout)
if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1:
threading.Thread(
target=_retry_remote_fetch_in_background,
kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping
"url": url,
"timeout": timeout,
"max_attempts": max_attempts,
"sleep": sleep,
"rng": fetch_rng,
"client": fetch_client,
"first_outcome": outcome,
},
name="litellm-model-cost-map-retry",
daemon=True,
).start()
if not isinstance(outcome, ModelCostMapReloaded):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",
url,
result.reason,
outcome.reason,
)
_cost_map_source_info.source = "local"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}"
return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map
content: Final = result.model_cost_map
content: Final = outcome.model_cost_map
# Validate using cached count (cheap int comparison, no file I/O)
if not GetModelCostMap.validate_model_cost_map(
@ -587,4 +677,4 @@ def get_model_cost_map(
_cost_map_source_info.source = "remote"
_cost_map_source_info.fallback_reason = None
return _finalize_loaded_model_cost_map(result).model_cost_map
return _finalize_loaded_model_cost_map(outcome).model_cost_map

5
litellm/litellm_core_utils/litellm_logging.py
View file

@ -203,6 +203,7 @@ if TYPE_CHECKING:
from litellm.integrations.otel.logger import OpenTelemetryV2
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
try:
from litellm_enterprise.enterprise_callbacks.callback_controls import (
@ -590,6 +591,7 @@ class Logging(LiteLLMLoggingBaseClass):
# Initialize cost breakdown field
self.cost_breakdown: CostBreakdown | None = None
self.billed_token_rates: BilledTokenRates | None = None
# Init Caching related details
self.caching_details: CachingDetails | None = None
@ -1587,6 +1589,7 @@ class Logging(LiteLLMLoggingBaseClass):
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
billed_token_rates: "BilledTokenRates | None" = None,
) -> None:
"""
Helper method to store cost breakdown in the logging object.
@ -1606,8 +1609,10 @@ class Logging(LiteLLMLoggingBaseClass):
service_tier: Tier the costs above were priced on, already resolved
data_residency: Region uplift the costs above were priced on, already resolved
vertex_location: Vertex AI location the costs above were priced on, already resolved
billed_token_rates: Per-token rates the costs above were billed at, already resolved
"""
self.billed_token_rates = billed_token_rates
self.cost_breakdown = CostBreakdown(
input_cost=input_cost,
output_cost=output_cost,

247
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -10,6 +10,7 @@ from typing import Any, Final, Literal, TypedDict, cast
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
import litellm
from litellm._internal_context import current_billing_time
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import (
select_tier_for_input,
@ -19,6 +20,7 @@ from litellm.types.utils import (
CacheCreationTokenDetails,
CallTypes,
CompletionTokensDetailsWrapper,
CostPerToken,
DataResidency,
ImageResponse,
ModelInfo,
@ -305,7 +307,7 @@ def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_
than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(),
or every window shifts by the host's offset.
"""
reference: Final = current_time if current_time is not None else datetime.now(timezone.utc)
reference: Final = current_time if current_time is not None else current_billing_time()
now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time()
windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc
for window in windows:
@ -392,7 +394,7 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None =
rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose
hours apply only on its weekdays.
"""
reference: Final = current_time if current_time is not None else datetime.now(timezone.utc)
reference: Final = current_time if current_time is not None else current_billing_time()
reference_utc: Final = (
reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc)
)
@ -1195,7 +1197,7 @@ def generic_cost_per_token(
usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0
)
billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc)
billing_time: Final = current_time if current_time is not None else current_billing_time()
(
prompt_base_cost,
completion_base_cost,
@ -1309,42 +1311,90 @@ def _coerce_token_count(value: object) -> int:
return value if isinstance(value, int) and value > 0 else 0
@dataclass(frozen=True, slots=True)
class BilledTokenRates:
"""Per-token rates one request's usage bills at, after token tiers, off-peak windows and the
regional multipliers the totals apply, so each cost line equals its token count times its rate."""
input_cost_per_token: float
output_cost_per_token: float
cache_read_input_token_cost: float
cache_creation_input_token_cost: float
cache_creation_input_token_cost_above_1hr: float
output_cost_per_reasoning_token: float
def scaled(self, multiplier: float) -> "BilledTokenRates":
if multiplier == 1.0:
return self
return BilledTokenRates(
input_cost_per_token=self.input_cost_per_token * multiplier,
output_cost_per_token=self.output_cost_per_token * multiplier,
cache_read_input_token_cost=self.cache_read_input_token_cost * multiplier,
cache_creation_input_token_cost=self.cache_creation_input_token_cost * multiplier,
cache_creation_input_token_cost_above_1hr=self.cache_creation_input_token_cost_above_1hr * multiplier,
output_cost_per_reasoning_token=self.output_cost_per_reasoning_token * multiplier,
)
@dataclass(frozen=True, slots=True)
class TokenTypeCostBreakdown:
reasoning_cost: float
cache_read_cost: float
cache_creation_cost: float
rates: BilledTokenRates | None = None
"""Rates these lines were billed at, so a caller reporting both cannot resolve them a second,
differently-argued way. None when the model's pricing could not be resolved."""
def get_token_type_cost_breakdown(
model: str,
custom_llm_provider: str | None,
def _reasoning_token_count(usage: Usage) -> int:
parsed: Final = (
parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0
)
return parsed or _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
def _cache_token_counts(usage: Usage) -> tuple[int, int, CacheCreationTokenDetails | None]:
"""(cache read tokens, cache creation tokens, cache creation details): read from prompt_tokens_details
first, then the private top-level counters the Usage constructor mirrors cache tokens onto for
providers/callers that bypass the details."""
parsed: Final = parse_prompt_tokens_details(usage) if usage.prompt_tokens_details is not None else None
parsed_read: Final = parsed["cache_hit_tokens"] if parsed is not None else 0
parsed_creation: Final = parsed["cache_creation_tokens"] if parsed is not None else 0
return (
parsed_read or _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0)),
parsed_creation or _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0)),
parsed["cache_creation_token_details"] if parsed is not None else None,
)
def _custom_pricing_rates(custom_cost_per_token: CostPerToken) -> BilledTokenRates:
"""Flat custom pricing has no tiers, uplifts or reasoning rate: cache tokens bill at the configured
cache rates (else the input rate) and reasoning at the output rate, as _cost_per_token_custom_pricing_helper does."""
input_rate: Final = custom_cost_per_token["input_cost_per_token"]
output_rate: Final = custom_cost_per_token["output_cost_per_token"]
cache_creation_rate: Final = custom_cost_per_token.get("cache_creation_input_token_cost", input_rate)
return BilledTokenRates(
input_cost_per_token=input_rate,
output_cost_per_token=output_rate,
cache_read_input_token_cost=custom_cost_per_token.get("cache_read_input_token_cost", input_rate),
cache_creation_input_token_cost=cache_creation_rate,
cache_creation_input_token_cost_above_1hr=cache_creation_rate,
output_cost_per_reasoning_token=output_rate,
)
def _cost_map_billed_rates(
model_info: ModelInfo,
usage: Usage,
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
current_time: datetime | None = None,
) -> TokenTypeCostBreakdown:
"""
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
object and model pricing alone.
This works for every provider, including Perplexity/Cerebras/Dashscope whose
cost calculators bypass ``generic_cost_per_token``, because cache tokens always
land on ``prompt_tokens_details`` (via the Usage constructor and provider
transformations) and reasoning tokens on ``completion_tokens_details``. It reuses
the same rate-resolution primitives as the total-cost path so the breakdown can
never drift from the totals. Returns zeros (never raises) when the model or its
pricing cannot be resolved.
"""
try:
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception:
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)
billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc)
custom_llm_provider: str | None,
service_tier: str | None,
data_residency: str | None,
vertex_location: str | None,
current_time: datetime | None,
) -> BilledTokenRates:
billing_time: Final = current_time if current_time is not None else current_billing_time()
(
_prompt_base_cost,
prompt_base_cost,
completion_base_cost,
cache_creation_cost_rate,
cache_creation_cost_above_1hr_rate,
@ -1356,13 +1406,6 @@ def get_token_type_cost_breakdown(
current_time=billing_time,
threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
)
reasoning_tokens = (
parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0
)
if not reasoning_tokens:
reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
reasoning_rate: Final = _resolve_billed_reasoning_rate(
model_info=model_info,
usage=usage,
@ -1370,57 +1413,103 @@ def get_token_type_cost_breakdown(
completion_base_cost=completion_base_cost,
current_time=billing_time,
)
reasoning_cost = float(reasoning_tokens) * reasoning_rate
multiplier: Final = (
_get_regional_uplift_multiplier(model_info, data_residency)
* get_vertex_regional_endpoint_uplift(model_info, vertex_location)
* get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
)
return BilledTokenRates(
input_cost_per_token=prompt_base_cost,
output_cost_per_token=completion_base_cost,
cache_read_input_token_cost=cache_read_cost_rate,
cache_creation_input_token_cost=cache_creation_cost_rate,
cache_creation_input_token_cost_above_1hr=cache_creation_cost_above_1hr_rate,
output_cost_per_reasoning_token=reasoning_rate,
).scaled(multiplier)
cache_read_tokens = 0
cache_creation_tokens = 0
cache_creation_token_details: CacheCreationTokenDetails | None = None
if usage.prompt_tokens_details is not None:
prompt_tokens_details: Final = parse_prompt_tokens_details(usage)
cache_read_tokens = prompt_tokens_details["cache_hit_tokens"]
cache_creation_tokens = prompt_tokens_details["cache_creation_tokens"]
cache_creation_token_details = prompt_tokens_details["cache_creation_token_details"]
# Fall back to the private top-level counters the Usage constructor mirrors cache
# tokens onto, so providers/callers that bypass prompt_tokens_details are covered.
if not cache_read_tokens:
cache_read_tokens = _coerce_token_count(getattr(usage, "_cache_read_input_tokens", 0))
if not cache_creation_tokens:
cache_creation_tokens = _coerce_token_count(getattr(usage, "_cache_creation_input_tokens", 0))
cache_read_cost = float(cache_read_tokens) * cache_read_cost_rate
cache_creation_cost = calculate_cache_writing_cost(
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_token_details,
cache_creation_cost_above_1hr=cache_creation_cost_above_1hr_rate,
cache_creation_cost=cache_creation_cost_rate,
def get_billed_token_rates(
model: str,
custom_llm_provider: str | None,
usage: Usage,
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
) -> BilledTokenRates | None:
"""Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type
breakdown resolve them. None when the model's pricing cannot be resolved."""
if custom_cost_per_token is not None:
return _custom_pricing_rates(custom_cost_per_token)
try:
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates
return None
return _cost_map_billed_rates(
model_info=model_info,
usage=usage,
custom_llm_provider=custom_llm_provider,
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
current_time=current_time,
)
# Apply the same flat regional-processing uplift the totals get, so per-type
# costs stay reconciled with input_cost/output_cost for regionalized OpenAI hosts.
uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency)
if uplift != 1.0:
reasoning_cost *= uplift
cache_read_cost *= uplift
cache_creation_cost *= uplift
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
if vertex_uplift != 1.0:
reasoning_cost *= vertex_uplift
cache_read_cost *= vertex_uplift
cache_creation_cost *= vertex_uplift
def get_token_type_cost_breakdown(
model: str,
custom_llm_provider: str | None,
usage: Usage,
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
) -> TokenTypeCostBreakdown:
"""
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
object and model pricing alone.
# Mirror the provider-specific geo uplift (e.g. Anthropic us: 1.1) the totals
# apply, so cache and reasoning line items stay reconciled with them.
geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
if geo_multiplier != 1.0:
reasoning_cost *= geo_multiplier
cache_read_cost *= geo_multiplier
cache_creation_cost *= geo_multiplier
This works for every provider, including Perplexity/Cerebras/Dashscope whose
cost calculators bypass ``generic_cost_per_token``, because cache tokens always
land on ``prompt_tokens_details`` (via the Usage constructor and provider
transformations) and reasoning tokens on ``completion_tokens_details``. It reuses
the same rate resolution as the total-cost path (``get_billed_token_rates``) so the
breakdown can never drift from the totals. A deployment billed by
``custom_cost_per_token`` is priced from those flat rates instead of the cost map and,
like its totals, bills cache writes flat rather than by their 5m/1h split.
Returns zeros (never raises) when the model or its pricing cannot be resolved.
"""
rates: Final = get_billed_token_rates(
model=model,
custom_llm_provider=custom_llm_provider,
usage=usage,
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
current_time=current_time,
custom_cost_per_token=custom_cost_per_token,
)
if rates is None:
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)
cache_read_tokens, cache_creation_tokens, cache_creation_token_details = _cache_token_counts(usage)
cache_creation_cost: Final = (
float(cache_creation_tokens) * rates.cache_creation_input_token_cost
if custom_cost_per_token is not None
else calculate_cache_writing_cost(
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_token_details,
cache_creation_cost_above_1hr=rates.cache_creation_input_token_cost_above_1hr,
cache_creation_cost=rates.cache_creation_input_token_cost,
)
)
return TokenTypeCostBreakdown(
reasoning_cost=reasoning_cost,
cache_read_cost=cache_read_cost,
reasoning_cost=float(_reasoning_token_count(usage)) * rates.output_cost_per_reasoning_token,
cache_read_cost=float(cache_read_tokens) * rates.cache_read_input_token_cost,
cache_creation_cost=cache_creation_cost,
rates=rates,
)

31
litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
View file

@ -3,7 +3,7 @@ import json
import re
import time
import traceback
from collections.abc import Iterable, Sequence
from collections.abc import Mapping, Sequence
from typing import Final, Literal, cast
import litellm
@ -151,6 +151,16 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
del choice.enhancements
def _invalid_choices_message(response_object: Mapping[str, object]) -> str:
raw_keys: Final = list(response_object.keys())
if "choices" not in response_object:
return f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {raw_keys}"
return (
f"LiteLLM: provider returned 'choices' that is not a list ({type(response_object['choices']).__name__}). "
f"Raw keys: {raw_keys}"
)
async def convert_to_streaming_response_async(
response_object: dict | None = None,
):
@ -179,14 +189,12 @@ async def convert_to_streaming_response_async(
choice_list: Final[list[StreamingChoices]] = []
if not response_object.get("choices"):
if not isinstance(response_object.get("choices"), list):
from litellm.exceptions import APIError
raise APIError(
status_code=500,
message=(
f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}"
),
message=_invalid_choices_message(response_object),
llm_provider="",
model="",
)
@ -287,14 +295,12 @@ def convert_to_streaming_response(
model_response_object: Final = ModelResponseStream()
choice_list: Final[list[StreamingChoices]] = []
if not response_object.get("choices"):
if not isinstance(response_object.get("choices"), list):
from litellm.exceptions import APIError
raise APIError(
status_code=500,
message=(
f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(response_object.keys())}"
),
message=_invalid_choices_message(response_object),
llm_provider="",
model="",
)
@ -623,15 +629,12 @@ def convert_to_model_response_object(
return convert_to_streaming_response(response_object=response_object)
choice_list: Final[list[Choices]] = []
if not response_object.get("choices") or not isinstance(response_object["choices"], Iterable):
if not isinstance(response_object.get("choices"), list):
from litellm.exceptions import APIError
raise APIError(
status_code=500,
message=(
"LiteLLM: provider returned a response with no 'choices'. "
f"Raw keys: {list(response_object.keys())}"
),
message=_invalid_choices_message(response_object),
llm_provider="",
model="",
)

11
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1473,17 +1473,14 @@ class CustomStreamWrapper:
self.received_finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider == "cached_response":
cached_chunk: Final = cast(ModelResponseStream, chunk)
chunk_finish_reason: Final = cached_chunk.choices[0].finish_reason
cached_choice: Final = cached_chunk.choices[0] if cached_chunk.choices else None
chunk_finish_reason: Final = cached_choice.finish_reason if cached_choice is not None else None
response_obj = {
"text": cached_chunk.choices[0].delta.content,
"text": cached_choice.delta.content if cached_choice is not None else None,
"is_finished": chunk_finish_reason is not None,
"finish_reason": chunk_finish_reason,
"original_chunk": cached_chunk,
"tool_calls": (
cached_chunk.choices[0].delta.tool_calls
if hasattr(cached_chunk.choices[0].delta, "tool_calls")
else None
),
"tool_calls": (getattr(cached_choice.delta, "tool_calls", None) if cached_choice is not None else None),
}
completion_obj["content"] = response_obj["text"]

200
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -13,10 +13,11 @@ Pattern Overview:
"""
import json
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Mapping, MutableSequence, Sequence
from copy import deepcopy
from dataclasses import dataclass
from itertools import chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from typing_extensions import ReadOnly, TypedDict, assert_never
@ -41,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
merge_guardrailed_scoped_messages,
merge_returned_tools_into_request_tools,
scoped_structured_message_indices,
stream_item_field,
stream_item_fingerprint,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
@ -153,6 +155,46 @@ class ExtractedInput:
EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
@dataclass(frozen=True, slots=True)
class _ToolCallShape:
name: str | None
arguments: str
@dataclass(frozen=True, slots=True)
class _SSEFieldRewrite:
"""One field of one nested section of a buffered SSE event, rewritten."""
section: str
field: str
value: object
class _SSEEventRewriter(Protocol):
def __call__(self, event: Mapping[str, object]) -> _SSEFieldRewrite | None: ...
def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewriter) -> Mapping[str, object]:
rewrite: Final = rewrite_event(event)
section: Final = None if rewrite is None else event.get(rewrite.section)
if rewrite is None or not isinstance(section, Mapping):
return event
return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} # mutable-ok: json.dumps needs a dict
def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]:
"""The guardrail-visible shape of each tool call, whether the guardrail handed
back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts."""
functions: Final = tuple(stream_item_field(tool_call, "function") for tool_call in tool_calls)
return tuple(
_ToolCallShape(
name=name if isinstance(name := stream_item_field(function, "name"), str) else None,
arguments=arguments if isinstance(arguments := stream_item_field(function, "arguments"), str) else "",
)
for function in functions
)
class _AnthropicSSEDelta(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[str]
@ -170,12 +212,18 @@ class AnthropicMessagesHandler(BaseTranslation):
them through guardrail rewrites; downstream provider handling is out of scope.
"""
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def __init__(self):
super().__init__()
self.adapter = LiteLLMAnthropicMessagesAdapter()
def post_call_hook_response(self, response: object) -> object:
if not isinstance(response, ModelResponse):
return response
return self.adapter.translate_openai_response_to_anthropic(response)
@staticmethod
def _build_streaming_usage_response(
responses_so_far: Sequence[object],
@ -1050,6 +1098,7 @@ class AnthropicMessagesHandler(BaseTranslation):
first_choice.message.tool_calls,
)
string_so_far = first_choice.message.content
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ())
guardrail_inputs: Final = GenericGuardrailAPIInputs()
if string_so_far:
guardrail_inputs["texts"] = [string_so_far]
@ -1084,6 +1133,19 @@ class AnthropicMessagesHandler(BaseTranslation):
and guardrailed_texts[0] != string_so_far
):
self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
if deliver_ended_stream_rewrites:
returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
self._write_ended_stream_tool_call_rewrites(
responses_so_far,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
post_guardrail_tool_calls=_tool_call_shapes(
returned_tool_calls
if isinstance(returned_tool_calls, list)
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
else tool_calls_list or ()
),
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
else:
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
return responses_so_far
@ -1206,44 +1268,124 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _write_ended_stream_text_rewrite(
responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
rewritten_text: str,
) -> None:
"""Deliver an ended-stream guardrail text rewrite by rewriting the
buffered chunks in place: the first ``text_delta`` carries the full
rewritten text and every later one is blanked, leaving the surrounding
message and content-block framing untouched. Handles both chunk formats
this stream carries (parsed event dicts and raw SSE bytes)."""
message and content-block framing untouched."""
replacements: Final = chain((rewritten_text,), repeat(""))
for idx, item in enumerate(responses_so_far):
if isinstance(item, dict):
delta = item.get("delta")
if item.get("type") == "content_block_delta" and isinstance(delta, dict):
if delta.get("type") == "text_delta":
delta["text"] = next(replacements)
elif isinstance(item, (bytes, bytearray)):
responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer
AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements)
)
def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
delta: Final = event.get("delta")
if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping):
return None
if delta.get("type") != "text_delta":
return None
return _SSEFieldRewrite("delta", "text", next(replacements))
AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
@classmethod
def _write_ended_stream_tool_call_rewrites(
cls,
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
*,
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
guardrail_name: str,
) -> None:
"""Deliver ended-stream guardrail tool-call rewrites by rewriting the
buffered chunks in place: the rebuilt response lists tool calls in the
order of the stream's ``tool_use`` blocks, so the nth rewritten call lands
on the nth block, its first ``input_json_delta`` carrying the full rewritten
arguments, every later one blanked, and ``content_block_start`` carrying the
rewritten name. Blocks that do not line up with the rebuilt tool calls make
the rewrite undeliverable, so the pipeline executor discards it and releases
the original chunks."""
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
block_indices: Final = tuple(
index
for item in responses_so_far
for event in cls._iter_sse_events(item)
if event.get("type") == "content_block_start"
and isinstance(block := event.get("content_block"), Mapping)
and block.get("type") == "tool_use"
and isinstance(index := event.get("index"), int)
)
if len(block_indices) != len(post_guardrail_tool_calls):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
rewrites_by_block: Final = MappingProxyType(
{
index: after
for index, before, after in zip(block_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if after != before
}
)
argument_replacements: Final = MappingProxyType(
{index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()}
)
def rewrite_tool_use(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
index: Final = event.get("index")
if not isinstance(index, int) or index not in rewrites_by_block:
return None
match event.get("type"):
case "content_block_start":
name: Final = rewrites_by_block[index].name
if name is None:
return None
return _SSEFieldRewrite("content_block", "name", name)
case "content_block_delta":
delta: Final = event.get("delta")
if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta":
return None
return _SSEFieldRewrite("delta", "partial_json", next(argument_replacements[index]))
case _:
return None
cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use)
@staticmethod
def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes:
"""Rewrite every ``text_delta`` data line in one SSE chunk with the next
replacement text, leaving all other events and framing byte-identical."""
def _rewrite_ended_stream_events(
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
rewrite_event: _SSEEventRewriter,
) -> None:
"""Replace every buffered event ``rewrite_event`` returns a rewrite for, in
both chunk formats this stream carries (parsed event dicts and raw SSE
bytes), leaving every other event and the framing untouched."""
rewritten_items: Final = tuple(
AnthropicMessagesHandler._rewrite_buffered_item(item, rewrite_event) for item in responses_so_far
)
responses_so_far[:] = rewritten_items # rebind-ok: delivers the rewrites into the caller's buffer
@staticmethod
def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object:
if isinstance(item, dict):
return _rewritten_event(_as_str_mapping(item), rewrite_event)
if isinstance(item, (bytes, bytearray)):
return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event)
return item
@staticmethod
def _rewrite_sse_events(sse_bytes: bytes, rewrite_event: _SSEEventRewriter) -> bytes:
"""Rewrite the data lines of one SSE chunk that ``rewrite_event`` rewrites,
leaving all other events and framing byte-identical."""
try:
decoded: Final = sse_bytes.decode("utf-8")
except UnicodeDecodeError:
return sse_bytes
return "\n\n".join(
AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n")
"\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, rewrite_event) for line in block.split("\n"))
for block in decoded.split("\n\n")
).encode("utf-8")
@staticmethod
def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str:
return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n"))
@staticmethod
def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str:
def _rewrite_sse_line(line: str, rewrite_event: _SSEEventRewriter) -> str:
if not line.startswith("data:"):
return line
try:
@ -1252,14 +1394,10 @@ class AnthropicMessagesHandler(BaseTranslation):
)
except json.JSONDecodeError:
return line
if not isinstance(data, dict) or data.get("type") != "content_block_delta":
if not isinstance(data, dict):
return line
delta: Final = data.get("delta")
if not isinstance(delta, dict) or delta.get("type") != "text_delta":
return line
return "data: " + json.dumps(
{**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts
)
rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event)
return line if rewritten is data else "data: " + json.dumps(rewritten)
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)

3
litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
View file

@ -1487,8 +1487,9 @@ class LiteLLMAnthropicMessagesAdapter:
anthropic_content.insert(0, polyfill_result.compaction_block)
## extract finish reason
openai_finish_reason: Final = response.choices[0].finish_reason if response.choices else "stop"
translated_finish_reason: Final = self._translate_openai_finish_reason_to_anthropic(
openai_finish_reason=response.choices[0].finish_reason
openai_finish_reason=openai_finish_reason
)
anthropic_finish_reason: Final = (
"refusal"

7
litellm/llms/azure/audio_transcriptions.py
View file

@ -37,6 +37,7 @@ class AzureAudioTranscription(AzureChatCompletion):
azure_ad_token: str | None = None,
atranscription: bool = False,
litellm_params: dict | None = None,
custom_llm_provider: str = "azure",
) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]:
data: Final = {"model": model, "file": audio_file, **optional_params}
@ -53,6 +54,7 @@ class AzureAudioTranscription(AzureChatCompletion):
logging_obj=logging_obj,
model=model,
litellm_params=litellm_params,
custom_llm_provider=custom_llm_provider,
)
azure_client: Final = self.get_azure_openai_client(
@ -99,7 +101,7 @@ class AzureAudioTranscription(AzureChatCompletion):
additional_args={"complete_input_dict": data},
original_response=stringified_response,
)
hidden_params: Final = {"model": model, "custom_llm_provider": "azure"}
hidden_params: Final = {"model": model, "custom_llm_provider": custom_llm_provider}
final_response: Final[TranscriptionResponse] = convert_to_model_response_object(
response_object=stringified_response,
model_response_object=model_response,
@ -122,6 +124,7 @@ class AzureAudioTranscription(AzureChatCompletion):
client=None,
max_retries=None,
litellm_params: dict | None = None,
custom_llm_provider: str = "azure",
) -> TranscriptionResponse:
response = None
try:
@ -178,7 +181,7 @@ class AzureAudioTranscription(AzureChatCompletion):
},
original_response=stringified_response,
)
hidden_params: Final = {"model": model, "custom_llm_provider": "azure"}
hidden_params: Final = {"model": model, "custom_llm_provider": custom_llm_provider}
response = convert_to_model_response_object(
_response_headers=headers,
response_object=stringified_response,

7
litellm/llms/azure/common_utils.py
View file

@ -14,6 +14,7 @@ from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.constants import DEFAULT_MAX_RETRIES
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.common_utils import BaseOpenAILLM
from litellm.secret_managers.get_azure_ad_token_provider import (
@ -582,7 +583,8 @@ class BaseAzureLLM(BaseOpenAILLM):
if scope is None:
scope = "https://cognitiveservices.azure.com/.default"
max_retries: Final = litellm_params.get("max_retries")
configured_max_retries: Final = litellm_params.get("max_retries")
max_retries: Final = DEFAULT_MAX_RETRIES if configured_max_retries is None else configured_max_retries
timeout: Final = litellm_params.get("timeout")
if not api_key and azure_ad_token_provider is None and tenant_id and client_id and client_secret:
verbose_logger.debug("Using Azure AD Token Provider from Entra ID for Azure Auth")
@ -642,8 +644,7 @@ class BaseAzureLLM(BaseOpenAILLM):
else:
azure_client_params["http_client"] = self._get_sync_http_client()
if max_retries is not None:
azure_client_params["max_retries"] = max_retries
azure_client_params["max_retries"] = max_retries
if timeout is not None:
azure_client_params["timeout"] = timeout

116
litellm/llms/azure_ai/cost_calculator.py
View file

@ -11,7 +11,7 @@ from litellm.types.utils import Usage
from litellm.utils import get_model_info
def _is_azure_model_router(model: str) -> bool:
def is_azure_model_router(model: str) -> bool:
"""
Check if the model is Azure AI Foundry Model Router.
@ -31,6 +31,18 @@ def _is_azure_model_router(model: str) -> bool:
return "model-router" in model_lower or "model_router" in model_lower or model_lower == "azure-model-router"
ROUTER_FEE_ENTRY_NAMES: Final = frozenset({"model-router", "model_router"})
def is_router_fee_entry(model: str) -> bool:
return model.lower().removeprefix("azure_ai/") in ROUTER_FEE_ENTRY_NAMES
def _router_fee_entry_name(model: str) -> str:
entry_name: Final = model.lower().removeprefix("azure_ai/")
return entry_name if entry_name in ROUTER_FEE_ENTRY_NAMES else "model_router"
def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> float:
"""
Calculate the flat cost for Azure AI Foundry Model Router.
@ -42,20 +54,39 @@ def calculate_azure_model_router_flat_cost(model: str, prompt_tokens: int) -> fl
Returns:
float: The flat cost in USD, or 0.0 if not applicable
"""
if not _is_azure_model_router(model):
if not is_azure_model_router(model):
return 0.0
# Get the model router pricing from model_prices_and_context_window.json
# Use "model_router" as the key (without actual model name suffix)
model_info: Final = get_model_info(model="model_router", custom_llm_provider="azure_ai")
model_info: Final = get_model_info(model=_router_fee_entry_name(model), custom_llm_provider="azure_ai")
router_flat_cost_per_token: Final = model_info.get("input_cost_per_token", 0)
if router_flat_cost_per_token and router_flat_cost_per_token > 0:
return prompt_tokens * router_flat_cost_per_token
return 0.0
def _response_model_cost(model: str, usage: Usage, service_tier: str | None) -> tuple[float, float]:
try:
return generic_cost_per_token(
model=model, usage=usage, custom_llm_provider="azure_ai", service_tier=service_tier
)
except Exception as e:
if not is_azure_model_router(model):
raise
verbose_logger.debug(
"Azure AI Model Router: model '%s' not in cost map, only the routing fee applies. Error: %s", model, e
)
return 0.0, 0.0
def _router_fee_name(model: str, request_model: str | None) -> str | None:
if is_router_fee_entry(model):
return None
if is_azure_model_router(model):
return model
if request_model is not None and is_azure_model_router(request_model):
return request_model
return None
def cost_per_token(
model: str,
usage: Usage,
@ -64,68 +95,31 @@ def cost_per_token(
service_tier: str | None = None,
) -> tuple[float, float]:
"""
Calculate the cost per token for Azure AI models.
Price the response model's own tokens for Azure AI, plus the Model Router fee exactly once when either the
priced name or request_model is a Model Router name.
For Azure AI Foundry Model Router:
- Adds a flat cost of $0.14 per million input tokens (from model_prices_and_context_window.json)
- Plus the cost of the actual model used (handled by generic_cost_per_token)
A response priced as the router entry itself already carries the fee, so nothing is added on top of it. A
router deployment name that is missing from the cost map prices at the fee alone.
completion_cost passes only the priced name: when that name is a routed model it adds the fee itself through
AzureModelRouterConfig.calculate_additional_costs as the "Azure Model Router Flat Cost" line of the cost
breakdown, and when the name is router-shaped the fee is already in the prompt cost returned here.
Args:
model: str, the model name without provider prefix (from response)
usage: LiteLLM Usage block
response_time_ms: Optional response time in milliseconds
request_model: Optional[str], the original request model name (to detect router usage)
request_model: Optional[str], the original request model name; a Model Router name adds the routing fee
service_tier: Optional service tier the request was priced on
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
Raises:
ValueError: If the model is not found in the cost map and cost cannot be calculated
(except for Model Router models where we return just the routing flat cost)
ValueError: If a model that is not a Model Router name is missing from the cost map
"""
prompt_cost = 0.0
completion_cost = 0.0
# Determine if this was a model router request
# Check both the response model and the request model
is_router_request: Final = _is_azure_model_router(model) or (
request_model is not None and _is_azure_model_router(request_model)
)
# Calculate base cost using generic cost calculator
# This may raise an exception if the model is not in the cost map
try:
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider="azure_ai",
service_tier=service_tier,
)
except Exception as e:
# For Model Router, the model name (e.g., "azure-model-router") may not be in the cost map
# because it's a routing service, not an actual model. In this case, we continue
# to calculate just the routing flat cost.
if not _is_azure_model_router(model):
# Re-raise for non-router models - they should have pricing defined
raise
verbose_logger.debug(
"Azure AI Model Router: model '%s' not in cost map, calculating routing flat cost only. Error: %s", model, e
)
# Add flat cost for Azure Model Router
# The flat cost is defined in model_prices_and_context_window.json for azure_ai/model_router
if is_router_request:
# Use the request model for flat cost calculation if available, otherwise use response model
router_model_for_calc: Final = request_model if request_model else model
router_flat_cost: Final = calculate_azure_model_router_flat_cost(router_model_for_calc, usage.prompt_tokens)
if router_flat_cost > 0:
verbose_logger.debug(
f"Azure AI Model Router flat cost: ${router_flat_cost:.6f} "
f"({usage.prompt_tokens} tokens × ${router_flat_cost / usage.prompt_tokens:.9f}/token)"
)
# Add flat cost to prompt cost
prompt_cost += router_flat_cost
return prompt_cost, completion_cost
prompt_cost, completion_cost = _response_model_cost(model=model, usage=usage, service_tier=service_tier)
fee_name: Final = _router_fee_name(model=model, request_model=request_model)
if fee_name is None:
return prompt_cost, completion_cost
return prompt_cost + calculate_azure_model_router_flat_cost(fee_name, usage.prompt_tokens), completion_cost

31
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -52,13 +52,28 @@ class StreamingScanKey:
class BaseTranslation(ABC):
delivers_ended_stream_text_rewrites: ClassVar[bool] = False
delivers_ended_stream_rewrites: ClassVar[bool] = False
"""Whether ``process_output_streaming_response`` accepts
``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered)
stream, writes guardrail text rewrites back across ``responses_so_far`` so
a buffered pipeline can release rewritten chunks. Tool-call rewrites, and
text rewrites on every other translation, are undeliverable: the pipeline
executor discards them and releases the original chunks."""
stream, writes guardrail text and tool-call rewrites back across
``responses_so_far`` so a buffered pipeline can release rewritten chunks,
raising ``UndeliverableStreamRewrite`` for a shape it cannot place. Rewrites
on every other translation are undeliverable: the pipeline executor
discards them and releases the original chunks."""
assembles_streamed_response: ClassVar[bool] = False
"""Whether ``process_output_streaming_response`` stores the assembled response of an
ended stream under ``request_data["response"]`` before scanning it, the way the chat,
Responses, and Messages translations do. A streaming pipeline runs a guardrail that only
has the legacy post-call hook against that response, so on a translation without it such
a guardrail keeps running on its own."""
def post_call_hook_response(self, response: object) -> object:
"""The ``response`` this endpoint's non-streaming post-call hooks receive, derived from
the object the translation stores under ``request_data["response"]`` while scanning an
ended stream. Chat and Responses scan that shape already; a translation that scans a
different one (Messages scans an OpenAI-shaped ModelResponse) overrides this."""
return response
@staticmethod
def transform_user_api_key_dict_to_metadata(
@ -175,9 +190,9 @@ class BaseTranslation(ABC):
transformations (see ``StreamTransformSink``); base handlers ignore it.
``deliver_ended_stream_rewrites`` is passed True only when the caller
holds the whole buffered stream and the subclass declares
``delivers_ended_stream_text_rewrites``: the handler then writes
guardrail text rewrites back across ``responses_so_far`` instead of
discarding them.
``delivers_ended_stream_rewrites``: the handler then writes
guardrail text and tool-call rewrites back across ``responses_so_far``
instead of discarding them.
"""
return responses_so_far

27
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -4,6 +4,7 @@ Translating between OpenAI's `/chat/completion` format and Amazon's `/converse`
import copy
import json
import re
import time
import types
from collections.abc import Mapping
@ -293,6 +294,10 @@ class AmazonConverseConfig(BaseConfig):
llm_provider="bedrock",
)
@staticmethod
def _is_openai_gpt_reasoning_model(model: str) -> bool:
return re.search(r"openai\.gpt-\d", model) is not None
def _is_nova_2_model(self, model: str) -> bool:
"""
Check if the model is a Nova 2 model that supports reasoningConfig.
@ -423,14 +428,14 @@ class AmazonConverseConfig(BaseConfig):
Handle the reasoning_effort parameter based on the model type.
- GPT-OSS and DeepSeek V3 models: passed through unchanged via additionalModelRequestFields.
- OpenAI GPT-5.x models: mapped to ``reasoning.effort`` via additionalModelRequestFields.
- OpenAI GPT-5.x and GPT-6 models: mapped to ``reasoning.effort`` via additionalModelRequestFields.
- Nova 2 models: transformed to reasoningConfig.
- Anthropic models: mapped to ``thinking`` (and ``output_config.effort`` on
adaptive Claude 4.6 / 4.7).
"""
if "gpt-oss" in model or "deepseek" in model:
optional_params["reasoning_effort"] = reasoning_effort
elif "openai.gpt-5" in model:
elif self._is_openai_gpt_reasoning_model(model):
reasoning: Final[BedrockConverseGptReasoningEffortBlock] = {"effort": reasoning_effort}
optional_params["reasoning"] = reasoning
elif self._is_nova_2_model(model):
@ -594,7 +599,11 @@ class AmazonConverseConfig(BaseConfig):
# only anthropic and mistral support tool choice config. otherwise (E.g. cohere) will fail the call - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
supported_params.append("tool_choice")
if "gpt-oss" in model or "openai.gpt-5" in model or "openai.gpt-5" in base_model:
if (
"gpt-oss" in model
or self._is_openai_gpt_reasoning_model(model)
or self._is_openai_gpt_reasoning_model(base_model)
):
supported_params.append("reasoning_effort")
elif self._is_deepseek_model(model=model, base_model=base_model):
if not self._is_deepseek_r1_model(model=model, base_model=base_model):
@ -953,7 +962,7 @@ class AmazonConverseConfig(BaseConfig):
"Dropping unsupported `thinking` param for Bedrock model=%s; it reasons natively.",
model,
)
elif param == "thinking" and "openai.gpt-5" not in model:
elif param == "thinking" and not self._is_openai_gpt_reasoning_model(model):
if (
isinstance(value, dict)
and value.get("type") == "adaptive"
@ -1843,6 +1852,7 @@ class AmazonConverseConfig(BaseConfig):
data=request_data,
messages=messages,
encoding=encoding,
json_mode=json_mode,
)
def _transform_reasoning_content(self, reasoning_content_blocks: list[BedrockConverseReasoningContentBlock]) -> str:
@ -2275,6 +2285,7 @@ class AmazonConverseConfig(BaseConfig):
data: dict | str,
messages: list,
encoding,
json_mode: bool | None = None,
) -> ModelResponse:
## LOGGING
if logging_obj is not None:
@ -2285,7 +2296,9 @@ class AmazonConverseConfig(BaseConfig):
additional_args={"complete_input_dict": data},
)
json_mode: Final[bool | None] = optional_params.get("json_mode", None)
resolved_json_mode: Final[bool | None] = (
json_mode if json_mode is not None else optional_params.get("json_mode", None)
)
## RESPONSE OBJECT
try:
completion_response: Final = ConverseResponseBlock(**response.json())
@ -2377,7 +2390,7 @@ class AmazonConverseConfig(BaseConfig):
chat_completion_message["thinking_blocks"] = self._transform_thinking_blocks(reasoningContentBlocks)
chat_completion_message["content"] = content_str
filtered_tools: Final = self._filter_json_mode_tools(
json_mode=json_mode,
json_mode=resolved_json_mode,
tools=tools,
chat_completion_message=chat_completion_message,
)
@ -2401,7 +2414,7 @@ class AmazonConverseConfig(BaseConfig):
# When json_mode filtered out all synthetic tool calls the response
# is plain content, not a pending tool invocation. Fix finish_reason
# so callers (e.g. OpenAI SDK) don't misinterpret it.
if json_mode and not filtered_tools and tools:
if resolved_json_mode and not filtered_tools and tools:
initial_finish_reason = "stop"
(

1
litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
View file

@ -340,6 +340,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
json_mode=json_mode,
)
elif provider == "twelvelabs":
return litellm.AmazonTwelveLabsPegasusConfig().transform_response(

72
litellm/llms/bedrock/files/transformation.py
View file

@ -7,13 +7,13 @@ from contextlib import suppress
from functools import cache
from itertools import chain
from types import MappingProxyType
from typing import Any, Final, TypeAlias, TypedDict
from typing import Any, Final, Literal, TypeAlias, TypedDict
from urllib.parse import unquote
import httpx
from httpx import Headers, Response
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, ConfigDict, TypeAdapter
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from typing_extensions import ReadOnly
from litellm._logging import verbose_logger
@ -60,11 +60,12 @@ from litellm.utils import get_llm_provider
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id
# litellm_params key used to hand the SigV4-signed GET headers from
# `transform_file_content_request` to `validate_environment` (the only hook
# the shared file-content HTTP handler exposes for setting request headers).
# Same pattern as the `upload_url` handoff in `transform_create_file_request`.
S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers"
S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers"
class _S3DeleteContext(BaseModel):
file_id: str = Field(min_length=1)
# litellm_params key carrying the size of the body uploaded to S3, handed from
# `transform_create_file_request` to `transform_create_file_response`.
@ -291,7 +292,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
) -> dict:
result: Final[dict[str, object]] = {}
result.update(headers)
signed_headers: Final = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None)
signed_headers: Final = litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM, None)
if isinstance(signed_headers, Mapping):
result.update(signed_headers) # any-ok: untyped handoff headers
# otherwise no extra headers - AWS credentials are handled by BaseAWSLLM
@ -1187,18 +1188,27 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
def transform_delete_file_request(
self,
file_id: str,
optional_params: dict,
litellm_params: dict,
) -> tuple[str, dict]:
raise NotImplementedError("BedrockFilesConfig does not support file deletion")
optional_params: Mapping[str, object],
litellm_params: MutableMapping[str, object],
) -> tuple[str, dict[str, str]]:
return self._transform_s3_file_request(
file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params
)
def transform_delete_file_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
litellm_params: Mapping[str, object],
) -> FileDeleted:
raise NotImplementedError("BedrockFilesConfig does not support file deletion")
if raw_response.status_code != 204:
raise BedrockError(
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
message=raw_response.text or f"S3 file deletion returned HTTP {raw_response.status_code}",
headers=raw_response.headers,
)
context: Final = _S3DeleteContext.model_validate(logging_obj.model_call_details.get("additional_args"))
return FileDeleted(id=context.file_id, deleted=True, object="file")
def transform_list_files_request(
self,
@ -1233,6 +1243,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
if not file_id:
raise ValueError("file_id is required for Bedrock file content retrieval")
return self._transform_s3_file_request(
file_id=file_id, method="GET", optional_params=optional_params, litellm_params=litellm_params
)
def _transform_s3_file_request(
self,
*,
file_id: str,
method: Literal["GET", "DELETE"],
optional_params: Mapping[str, object],
litellm_params: MutableMapping[str, object],
) -> tuple[str, dict[str, str]]:
s3_uri: Final = extract_s3_uri_from_file_id(file_id)
bucket_name, object_key = _validate_file_id_against_configured_buckets(
s3_uri=s3_uri,
@ -1240,40 +1262,32 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
)
# The shared file-content handler passes optional_params={}, so AWS
# credentials/region arrive via litellm_params here (unlike the upload
# path). s3_region_name wins over aws_region_name, same priority as
# get_complete_file_url above.
merged_params: Final[dict[str, object]] = {}
merged_params.update(litellm_params)
merged_params.update(optional_params)
request_params: Final = _BedrockS3RequestParams.model_validate(merged_params)
request_params: Final = _BedrockS3RequestParams.model_validate({**litellm_params, **optional_params})
region_preference: Final = request_params.s3_region_name or request_params.aws_region_name
region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference}
aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="")
s3_endpoint_url = (
s3_endpoint_url: Final = (
request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
).rstrip("/")
url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request(
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = self._sign_s3_request_without_body(
api_base=url,
aws_region_name=aws_region_name,
request_params=request_params,
method=method,
)
return url, {}
def _sign_s3_get_request(
def _sign_s3_request_without_body(
self,
api_base: str,
aws_region_name: str,
request_params: _BedrockS3RequestParams,
method: Literal["GET", "DELETE"] = "GET",
) -> dict[str, str]:
"""
SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT).
"""
try:
import hashlib
@ -1297,7 +1311,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
empty_body_hash: Final = hashlib.sha256(b"").hexdigest()
aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped
method="GET",
method=method,
url=api_base,
headers={"x-amz-content-sha256": empty_body_hash},
)

4
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -579,7 +579,7 @@ class BaseLLMHTTPHandler:
data: dict[str, object], # mutable-ok: async_completion takes dict
signed_headers: dict[str, object], # mutable-ok: async_completion takes dict
signed_json_body: bytes | None,
):
) -> Coroutine[object, object, ModelResponse | CustomStreamWrapper]:
async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None
if stream is True:
return self.acompletion_stream_function(
@ -626,7 +626,7 @@ class BaseLLMHTTPHandler:
if acompletion is True and provider_config.uses_async_transform_request:
async def transform_then_dispatch():
async def transform_then_dispatch() -> ModelResponse | CustomStreamWrapper:
transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict
"dict[str, object]",
await provider_config.async_transform_request(

39
litellm/llms/databricks/chat/transformation.py
View file

@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
_should_convert_tool_call_to_json_mode,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
_extract_reasoning_content, # pyright: ignore[reportPrivateUsage] # same import as the OpenAI transformation
strip_litellm_internal_message_fields,
strip_name_from_message,
)
@ -23,7 +24,9 @@ from litellm.types.llms.anthropic import AllAnthropicToolsValues
from litellm.types.llms.databricks import (
AllDatabricksContentValues,
DatabricksChoice,
DatabricksDelta,
DatabricksFunction,
DatabricksMessage,
DatabricksResponse,
DatabricksTool,
)
@ -247,8 +250,10 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
litellm_params: dict,
stream: bool | None = None,
) -> str:
api_base = self._get_api_base(api_base)
complete_url: Final = f"{api_base}/chat/completions"
use_ai_gateway: Final = model.removeprefix("databricks/").count(".") >= 2
api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway)
url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base
complete_url: Final = f"{url_base}/chat/completions"
return complete_url
def get_supported_openai_params(self, model: str | None = None) -> list:
@ -534,6 +539,19 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
thinking_blocks.append(thinking_block)
return reasoning_content, thinking_blocks
@staticmethod
def extract_top_level_reasoning_content(delta: DatabricksDelta) -> str | None:
return delta.get("reasoning_content")
@staticmethod
def resolve_reasoning_and_content(
message: DatabricksMessage, block_reasoning_content: str | None
) -> tuple[str | None, str | None]:
content_str: Final = DatabricksConfig.extract_content_str(message["content"])
if block_reasoning_content is not None:
return block_reasoning_content, content_str
return _extract_reasoning_content({**message, "content": content_str})
@staticmethod
def extract_citations(
content: AllDatabricksContentValues | None,
@ -577,14 +595,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
finish_reason = "stop"
if translated_message is None:
## get the content str
content_str = DatabricksConfig.extract_content_str(choice["message"]["content"])
## get the reasoning content
(
reasoning_content,
block_reasoning_content,
thinking_blocks,
) = DatabricksConfig.extract_reasoning_content(choice["message"].get("content"))
reasoning_content, content_str = DatabricksConfig.resolve_reasoning_and_content(
choice["message"], block_reasoning_content
)
citations = DatabricksConfig.extract_citations(choice["message"].get("content"))
@ -738,12 +755,16 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator):
# extract the reasoning content
(
reasoning_content,
block_reasoning_content,
thinking_blocks,
) = DatabricksConfig.extract_reasoning_content(choice["delta"].get("content"))
choice["delta"]["content"] = content_str
choice["delta"]["reasoning_content"] = reasoning_content
choice["delta"]["reasoning_content"] = (
block_reasoning_content
if block_reasoning_content is not None
else DatabricksConfig.extract_top_level_reasoning_content(choice["delta"])
)
choice["delta"]["thinking_blocks"] = thinking_blocks
translated_choices.append(choice)
return ModelResponseStream(

20
litellm/llms/databricks/common_utils.py
View file

@ -177,19 +177,13 @@ class DatabricksBase:
# Default: just litellm
return f"litellm/{version}"
def _get_api_base(self, api_base: str | None) -> str:
"""
Get the Databricks API base URL.
If not provided, attempts to get it from the Databricks SDK.
"""
def _get_api_base(self, api_base: str | None, use_ai_gateway: bool = False) -> str:
if api_base is None:
try:
from databricks.sdk import WorkspaceClient
databricks_client: Final = WorkspaceClient()
api_base = f"{databricks_client.config.host}/serving-endpoints"
return api_base
except ImportError:
raise DatabricksException(
status_code=400,
@ -198,6 +192,18 @@ class DatabricksBase:
"or install the databricks-sdk Python library."
),
)
if not use_ai_gateway:
return api_base
normalized_api_base: Final = api_base.rstrip("/")
if normalized_api_base.endswith("/ai-gateway/mlflow/v1"):
return normalized_api_base
if normalized_api_base.endswith("/serving-endpoints"):
return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1"
api_base_parts: Final = urlsplit(normalized_api_base)
if api_base_parts.path in ("", "/"):
return f"{normalized_api_base}/ai-gateway/mlflow/v1"
return api_base
def _get_oauth_m2m_token(

9
litellm/llms/hosted_vllm/image_edit/__init__.py Normal file
View file

@ -0,0 +1,9 @@
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from .transformation import HostedVLLMImageEditConfig
__all__ = ("HostedVLLMImageEditConfig",)
def get_hosted_vllm_image_edit_config(model: str) -> BaseImageEditConfig:
return HostedVLLMImageEditConfig()

43
litellm/llms/hosted_vllm/image_edit/transformation.py Normal file
View file

@ -0,0 +1,43 @@
from typing import Final
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
from litellm.secret_managers.main import get_secret_str
PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT: Final = frozenset({"mask", "quality", "input_fidelity"})
class HostedVLLMImageEditConfig(OpenAIImageEditConfig):
def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseImageEditConfig contract
return [ # mutable-ok: BaseImageEditConfig returns list
param
for param in super().get_supported_openai_params(model)
if param not in PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT
]
def validate_environment(
self,
headers: dict, # mutable-ok: BaseImageEditConfig contract
model: str,
api_key: str | None = None,
litellm_params: dict | None = None, # mutable-ok: BaseImageEditConfig contract
api_base: str | None = None,
) -> dict: # mutable-ok: BaseImageEditConfig contract
resolved_key: Final = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key"
return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: dict, # mutable-ok: BaseImageEditConfig contract
) -> str:
resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE")
if resolved_api_base is None:
raise ValueError(
"api_base not set for Hosted VLLM images edits API. "
"Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable"
)
trimmed: Final = resolved_api_base.rstrip("/")
if trimmed.endswith("/v1"):
return f"{trimmed}/images/edits"
return f"{trimmed}/v1/images/edits"

97
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -49,6 +49,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import
coerce_stream_holdback_value,
)
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
ChatCompletionMessageToolCall,
Choices,
GenericGuardrailAPIInputs,
ModelResponse,
@ -78,7 +80,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
@ -610,13 +613,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
deliver_ended_stream_rewrites: bool,
) -> None:
"""Ended-stream path: rebuild the full response, run the non-streaming
output guardrail against it, and (when opted in) write any text rewrite
back across the buffered chunks."""
output guardrail against it, and (when opted in) write any text or
tool-call rewrite back across the buffered chunks."""
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response)
await self.process_output_response(
response=model_response,
guardrail_to_apply=guardrail_to_apply,
@ -624,13 +628,21 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
if deliver_ended_stream_rewrites:
await self._write_ended_stream_text_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_texts=pre_guardrail_texts,
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
if not deliver_ended_stream_rewrites:
return
guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown"
await self._write_ended_stream_text_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_texts=pre_guardrail_texts,
guardrail_name=guardrail_name,
)
self._write_ended_stream_tool_call_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
guardrail_name=guardrail_name,
)
def build_stream_error_items(
self,
@ -1043,6 +1055,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists
)
@staticmethod
def _function_tool_call_shapes(response: "ModelResponse") -> tuple[tuple[str | None, str], ...]:
return tuple(
(tool_call.function.name, tool_call.function.arguments)
for choice in response.choices
for tool_call in choice.message.tool_calls or ()
if isinstance(tool_call, ChatCompletionMessageToolCall)
)
@staticmethod
def _function_tool_call_fragments(
responses_so_far: Sequence["ModelResponseStream"],
) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]:
"""Group the stream's function tool-call fragments by their tool-call index, in
the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping
only the indices the builder keeps (an id and a name somewhere in the stream)."""
fragments: Final = tuple(
tool_call
for response in responses_so_far
for choice in response.choices
for tool_call in choice.delta.tool_calls or ()
if isinstance(tool_call, ChatCompletionDeltaToolCall)
)
identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id)
named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name)
return tuple(
tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named)
)
def _write_ended_stream_tool_call_rewrites(
self,
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrailed_response: "ModelResponse",
pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail tool-call rewrites back across the buffered
chunks: the rewritten name and full arguments land in the tool call's first
fragment and the arguments of its later fragments are blanked, mirroring the
text write-back. A rewrite on a stream carrying more than one distinct choice
index, or whose fragments do not line up with the rebuilt tool calls, is
reported as undeliverable, so the pipeline executor discards it and releases
the original chunks."""
post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response)
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
stream_choice_indices: Final = frozenset(
choice.index for response in responses_so_far for choice in response.choices
)
fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far)
if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
for before, (name, arguments), fragments in zip(
pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call
):
if (name, arguments) == before:
continue
head, *tail = fragments
head.function.name = name
head.function.arguments = arguments
for fragment in tail:
fragment.function.arguments = ""
async def _apply_guardrail_responses_to_output_streaming(
self,
responses: list["ModelResponseStream"],

308
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -37,14 +37,14 @@ from itertools import accumulate, chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from pydantic import BaseModel, TypeAdapter
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
OpenAiResponsesToChatCompletionStreamIterator,
tool_call_dict_from_output_item,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
@ -84,7 +84,6 @@ from litellm.types.llms.openai import (
)
from litellm.types.responses.main import (
GenericResponseOutputItem,
OutputFunctionToolCall,
OutputText,
)
from litellm.types.utils import GenericGuardrailAPIInputs
@ -101,6 +100,72 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import ResponseInputParam
class _ToolCallShape(NamedTuple):
name: str | None
arguments: str
class _ToolCallFunctionFields(BaseModel):
model_config = ConfigDict(frozen=True)
name: str | None = None
arguments: str = ""
class _ToolCallFields(BaseModel):
model_config = ConfigDict(frozen=True)
function: _ToolCallFunctionFields
def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]:
return tuple(
_ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", ""))
for tool_call in tool_calls
)
def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None:
payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
try:
fields: Final = _ToolCallFields.model_validate(payload)
except ValidationError:
return None
return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments)
def _post_guardrail_tool_call_shapes(
returned_tool_calls: Sequence[object] | None,
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
guardrail_name: str | None,
) -> tuple[_ToolCallShape, ...]:
if not pre_guardrail_tool_calls:
return pre_guardrail_tool_calls
if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls):
verbose_proxy_logger.warning(
"OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, "
"leaving the tool call output items unchanged",
guardrail_name,
"no" if returned_tool_calls is None else len(returned_tool_calls),
len(pre_guardrail_tool_calls),
)
return pre_guardrail_tool_calls
returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls)
validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None)
if len(validated_shapes) != len(returned_shapes):
verbose_proxy_logger.warning(
"OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, "
"leaving the tool call output items unchanged",
guardrail_name,
)
return pre_guardrail_tool_calls
return validated_shapes
def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape:
return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments)
class ResponseOutputEnvelope(TypedDict, total=False):
"""Dict form of a Responses API response, as far as guardrail write-back reads it."""
@ -128,6 +193,20 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
)
_TOOL_CALL_ITEM_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
_TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call": "arguments", "custom_tool_call": "input"}
)
_TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES: Final = frozenset(
{"response.function_call_arguments.delta", "response.custom_tool_call_input.delta"}
)
_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"response.function_call_arguments.done": "arguments", "response.custom_tool_call_input.done": "input"}
)
_TOOL_CALL_PAYLOAD_EVENT_TYPES: Final = _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES | frozenset(
_TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS
)
_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call_output": "output", "message": "content"}
)
@ -164,8 +243,20 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp
return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts
def _is_function_call_item(item: object) -> bool:
return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call")
def _is_tool_call_item(item: object) -> bool:
return isinstance(item, Mapping) and item.get("type") in _TOOL_CALL_ITEM_TYPES
def _tool_call_output_item_mapping(item: object) -> Mapping[str, object] | None:
if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES:
return None
if isinstance(item, Mapping):
return cast("Mapping[str, object]", item) # cast-ok: output items are str-keyed JSON objects
return item.model_dump() if isinstance(item, BaseModel) else None
def _is_tool_call_output_item(item: object) -> bool:
return _tool_call_output_item_mapping(item) is not None
def _last_message_role(messages: Sequence[object]) -> str | None:
@ -189,7 +280,7 @@ def _provenance_unit_bounds(
start_indexes: Final = tuple(
index
for index in range(len(raw_input))
if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
if index == 0 or not (_is_tool_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant")
)
return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input))))
@ -340,7 +431,8 @@ class OpenAIResponsesHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
@ -587,7 +679,7 @@ class OpenAIResponsesHandler(BaseTranslation):
- response.output is a list of output items
- Each output item can be:
* GenericResponseOutputItem with a content list of OutputText objects
* ResponseFunctionToolCall with tool call data
* ResponseFunctionToolCall or CustomToolCallOutputItem with tool call data
- Each OutputText object has a text field
"""
@ -652,6 +744,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if response_model:
inputs["model"] = response_model
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
@ -660,6 +753,11 @@ class OpenAIResponsesHandler(BaseTranslation):
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
guardrail_name=guardrail_to_apply.guardrail_name,
)
# Step 3: Map guardrail responses back to original response structure
await self._apply_guardrail_responses_to_output(
@ -667,6 +765,11 @@ class OpenAIResponsesHandler(BaseTranslation):
responses=guardrailed_texts,
task_mappings=task_mappings,
)
self._write_tool_call_rewrites_to_output(
tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)),
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
post_guardrail_tool_calls=post_guardrail_tool_calls,
)
verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response)
@ -754,6 +857,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if response_model:
inputs["model"] = response_model
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
@ -762,6 +866,11 @@ class OpenAIResponsesHandler(BaseTranslation):
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes(
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
guardrail_name=guardrail_to_apply.guardrail_name,
)
# Write guardrailed texts back into the output items in-place.
# final_chunk is a reference into responses_so_far so this
@ -784,6 +893,13 @@ class OpenAIResponsesHandler(BaseTranslation):
stream_events=responses_so_far[:-1],
rewrites_by_position=rewrites_by_position,
)
self._deliver_ended_stream_tool_call_rewrites(
responses_so_far=responses_so_far,
outputs=outputs,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
post_guardrail_tool_calls=post_guardrail_tool_calls,
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
return responses_so_far
# ------------------------------------------------------------------ #
@ -894,6 +1010,148 @@ class OpenAIResponsesHandler(BaseTranslation):
continue
OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten)
def _deliver_ended_stream_tool_call_rewrites(
self,
responses_so_far: Sequence[object],
outputs: Sequence[object],
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail tool-call rewrites into the completed
envelope's ``function_call`` and ``custom_tool_call`` items and sync the
earlier stream events, keyed by ``call_id``. The guardrail sees the
envelope's tool calls in output order, which is how a rewritten call
finds its ``call_id``; the stream events find their call through the
``call_id`` on ``output_item`` events and the ``item_id`` on argument
and custom-input events, since an
event's ``output_index`` need not match the envelope's (the chat bridge
numbers tool calls from 1 while the envelope lists them after the
message). A rewrite whose calls do not line up with the envelope, or
whose events cannot be found, is reported as undeliverable, so the
pipeline executor discards it and releases the original events."""
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
tool_call_items: Final = tuple(output_item for output_item in outputs if _is_tool_call_output_item(output_item))
call_ids: Final = tuple(
call_id
for output_item in tool_call_items
if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id
)
stream_events: Final = responses_so_far[:-1]
call_id_by_item_id: Final = self._tool_call_ids_by_item_id(stream_events)
event_call_ids: Final = tuple(
self._tool_call_event_call_id(event, call_id_by_item_id) for event in stream_events
)
rewrites_by_call_id: Final = MappingProxyType(
{
call_id: _tool_call_rewrite(before, after)
for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if after != before
}
)
unresolved_argument_event: Final = any(
call_id is None and stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES
for event, call_id in zip(stream_events, event_call_ids)
)
if (
len(call_ids) != len(tool_call_items)
or len(frozenset(call_ids)) != len(call_ids)
or len(call_ids) != len(post_guardrail_tool_calls)
or unresolved_argument_event
or not rewrites_by_call_id.keys() <= frozenset(event_call_ids)
):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
for output_item, rewrite in (
(output_item, rewrites_by_call_id[call_id])
for output_item, call_id in zip(tool_call_items, call_ids)
if call_id in rewrites_by_call_id
):
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
delta_replacements: Final = MappingProxyType(
{call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()}
)
for event, call_id in zip(stream_events, event_call_ids):
if call_id not in rewrites_by_call_id:
continue
match stream_item_field(event, "type"):
case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DELTA_EVENT_TYPES:
self._write_event_field(event, "delta", next(delta_replacements[call_id]))
case str() as event_type if event_type in _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS:
self._write_event_field(
event, _TOOL_CALL_PAYLOAD_DONE_EVENT_FIELDS[event_type], rewrites_by_call_id[call_id].arguments
)
case "response.output_item.added":
self._write_tool_call_item(
stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None
)
case "response.output_item.done":
self._write_tool_call_item(
stream_item_field(event, "item"),
rewrites_by_call_id[call_id].name,
rewrites_by_call_id[call_id].arguments,
)
case _:
pass
def _write_tool_call_rewrites_to_output(
self,
tool_call_items: Sequence[object],
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
) -> None:
if len(tool_call_items) != len(post_guardrail_tool_calls):
return
for output_item, rewrite in (
(output_item, _tool_call_rewrite(before, after))
for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if after != before
):
self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments)
@staticmethod
def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
items: Final = tuple(
stream_item_field(event, "item")
for event in stream_events
if stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
)
return MappingProxyType(
{
item_id: call_id
for item in items
if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
and isinstance(item_id := stream_item_field(item, "id"), str)
and isinstance(call_id := stream_item_field(item, "call_id"), str)
}
)
@staticmethod
def _tool_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
event_type: Final = stream_item_field(event, "type")
if event_type in _TOOL_CALL_PAYLOAD_EVENT_TYPES:
item_id: Final = stream_item_field(event, "item_id")
return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None
if event_type not in _OUTPUT_ITEM_EVENT_TYPES:
return None
item: Final = stream_item_field(event, "item")
call_id: Final = stream_item_field(item, "call_id")
return (
call_id if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES and isinstance(call_id, str) else None
)
@staticmethod
def _write_tool_call_item(item: object, name: str | None, payload: str | None) -> None:
if item is None:
return
if name is not None:
OpenAIResponsesHandler._write_event_field(item, "name", name)
item_type: Final = stream_item_field(item, "type")
if payload is not None and isinstance(item_type, str) and item_type in _TOOL_CALL_PAYLOAD_FIELDS:
OpenAIResponsesHandler._write_event_field(item, _TOOL_CALL_PAYLOAD_FIELDS[item_type], payload)
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
Check if the streaming has ended.
@ -920,7 +1178,7 @@ class OpenAIResponsesHandler(BaseTranslation):
def _completed_response_scan_key(response: object) -> StreamingScanKey:
output_items: Final = stream_item_items(response, "output")
message_items: Final = tuple(
item for item in output_items if stream_item_field(item, "type") != "function_call"
item for item in output_items if stream_item_field(item, "type") not in _TOOL_CALL_ITEM_TYPES
)
return StreamingScanKey(
texts=tuple(
@ -932,7 +1190,7 @@ class OpenAIResponsesHandler(BaseTranslation):
tool_calls=tuple(
stream_item_fingerprint(item)
for item in output_items
if stream_item_field(item, "type") == "function_call"
if stream_item_field(item, "type") in _TOOL_CALL_ITEM_TYPES
),
stream_ended=True,
)
@ -1043,34 +1301,10 @@ class OpenAIResponsesHandler(BaseTranslation):
Override this method to customize text/image/tool extraction logic.
"""
# Check if this is a tool call (OutputFunctionToolCall)
if isinstance(output_item, OutputFunctionToolCall) or (
isinstance(output_item, BaseModel)
and hasattr(output_item, "type")
and getattr(output_item, "type") == "function_call"
):
tool_call_item: Final = _tool_call_output_item_mapping(output_item)
if tool_call_item is not None:
if tool_calls_to_check is not None:
tool_call_dict = (
LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
tool_call_item=output_item,
index=output_idx,
)
)
tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
return
elif isinstance(output_item, dict) and output_item.get("type") == "function_call":
# Handle dict representation of tool call
if tool_calls_to_check is not None:
# Convert dict to ResponseFunctionToolCall for processing
try:
tool_call_obj: Final = ResponseFunctionToolCall(**output_item)
tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
tool_call_item=tool_call_obj,
index=output_idx,
)
tool_calls_to_check.append(cast(ChatCompletionToolCallChunk, tool_call_dict))
except Exception:
pass
tool_calls_to_check.append(tool_call_dict_from_output_item(tool_call_item, output_idx))
return
# Handle both GenericResponseOutputItem and dict

11
litellm/main.py
View file

@ -5398,6 +5398,14 @@ def completion(
if dynamic_api_key is not None:
api_key = dynamic_api_key
# check if user passed in any of the OpenAI optional params
bridges_to_responses_api: Final = (
responses_api_model_info.get("mode") == "responses" and not skip_responses_api_bridge
)
allowed_openai_params: Final[list[str] | None] = (
[*(kwargs.get("allowed_openai_params") or []), "reasoning_effort"]
if bridges_to_responses_api
else kwargs.get("allowed_openai_params")
)
optional_param_args: Final = {
"functions": functions,
"function_call": function_call,
@ -5442,7 +5450,7 @@ def completion(
"service_tier": service_tier,
"store": store,
"prompt_cache_key": prompt_cache_key,
"allowed_openai_params": kwargs.get("allowed_openai_params"),
"allowed_openai_params": allowed_openai_params,
"base_model": base_model,
}
optional_params = get_optional_params(**optional_param_args, **non_default_params)
@ -7805,6 +7813,7 @@ def transcription(
azure_ad_token=azure_ad_token,
max_retries=max_retries,
litellm_params=litellm_params_dict,
custom_llm_provider=custom_llm_provider,
)
elif custom_llm_provider == "openai" or (custom_llm_provider in litellm.openai_compatible_providers):
api_base = (

845
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

82
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -3,12 +3,15 @@ import importlib
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from traceback import walk_tb
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal
from uuid import uuid4
import anyio
import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from pydantic import ValidationError
from starlette.datastructures import Headers
from litellm._logging import verbose_logger
@ -30,6 +33,8 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
list_fault_http_status,
outcome_wire_value,
)
from litellm.proxy._experimental.mcp_server.faults.traversal import iter_exception_tree
from litellm.proxy._experimental.mcp_server.oauth_utils import _redact_mcp_resource_url
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
acting_user_auth,
build_effective_auth_contexts,
@ -78,11 +83,39 @@ _MCP_GUARDRAIL_REJECTIONS: Final = (
def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str:
reference: Final = uuid4().hex
verbose_logger.error(
"MCP connection test failed (reference=%s): %s",
reference,
tuple(
(
type(cause).__name__,
tuple(
(frame.f_code.co_filename, lineno, frame.f_code.co_name)
for frame, lineno in walk_tb(cause.__traceback__)
),
)
for cause in iter_exception_tree(exc)
),
)
return next(
(
message
for cause in iter_exception_tree(exc)
if (message := _known_connection_error_message(cause, url, timeout_seconds)) is not None
),
"An unexpected error occurred while testing the MCP connection. "
f"Retry; if it persists, share reference {reference} with your gateway administrator.",
)
def _known_connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str | None:
if isinstance(exc, MCPServerURLCredentialsError):
return str(exc.detail)
if isinstance(exc, TimeoutError):
return (
f"Failed to connect to MCP server: no response from {url or 'the server'} "
"Failed to connect to MCP server: no valid MCP response received from "
f"{_redact_mcp_resource_url(url) or 'the server'} "
f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL "
"from its network (DNS, egress rules, firewalls) and that the server answers MCP requests."
)
@ -99,13 +132,45 @@ def _connection_error_message(exc: BaseException, url: str | None, timeout_secon
return "Failed to connect to MCP server: the connection timed out."
if isinstance(exc, httpx.HTTPStatusError):
return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}."
return "Failed to connect to MCP server. Check proxy logs for details."
if isinstance(exc, (httpx.NetworkError, httpx.RemoteProtocolError, ConnectionError)):
return (
"Failed to connect to MCP server: the connection was interrupted. "
"Check the server and network connection, then retry."
)
if isinstance(exc, ValueError) and str(exc).startswith("Unexpected content type:"):
return (
"Failed to connect to MCP server: the endpoint returned an unsupported content type. "
"Check that the URL is an MCP endpoint, not a web page, and matches the selected transport."
)
if isinstance(exc, ValidationError) and exc.title in ("JSONRPCMessage", "InitializeResult", "ListToolsResult"):
return (
"Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. "
"Check the MCP endpoint URL and the server's protocol implementation."
)
if MCP_AVAILABLE and isinstance(exc, McpError):
if exc.error.code == -32000 and exc.error.message == "Connection closed":
return (
"Failed to connect to MCP server: the connection was closed before the request completed. "
"Check that the server stays running and returns a complete MCP response, then retry."
)
if exc.error.code == 32600 and exc.error.message == "Session terminated":
return (
"Failed to connect to MCP server: the MCP session was terminated. "
"Check that the URL points to an MCP endpoint and matches the selected transport, "
"then retry to start a new session."
)
return (
f"Failed to connect to MCP server: the MCP request failed (JSON-RPC code {exc.error.code}). "
"Check that the endpoint supports MCP initialization and tool listing, and check the upstream server logs."
)
return None
if MCP_AVAILABLE:
from mcp.shared.exceptions import McpError
from mcp.types import Tool as MCPTool
from litellm.experimental_mcp_client.client import MCPClient
from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout
from litellm.llms.litellm_proxy.skills.skill_search import (
DEFAULT_SKILL_SEARCH_TOP_K,
)
@ -1342,11 +1407,18 @@ if MCP_AVAILABLE:
except (KeyboardInterrupt, SystemExit, asyncio.CancelledError):
raise
except BaseException as e:
verbose_logger.error("Error in MCP operation: %s", e, exc_info=True)
effective_timeout: Final = (
min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds)
if any(
isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None
for cause in iter_exception_tree(e)
)
else timeout_seconds
)
return {
"status": "error",
"error": True,
"message": _connection_error_message(e, request.url, timeout_seconds),
"message": _connection_error_message(e, request.url, effective_timeout),
}
async def _preview_openapi_tools(spec_path: str) -> dict:

2
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -1053,7 +1053,7 @@ if MCP_AVAILABLE:
route="/mcp/call_tool",
traceback_str=failure_traceback,
)
except Exception:
except Exception: # noqa: BLE001 # a failing failure hook must not mask the tool call's own error
verbose_logger.exception("Error logging failed MCP proxy tool call")
raise
if proxy_logging_obj is not None:

16
litellm/proxy/_experimental/mcp_server/toolset_db.py
View file

@ -132,10 +132,18 @@ async def update_mcp_toolset(
data: UpdateMCPToolsetRequest,
touched_by: str,
) -> MCPToolset | None:
data_dict: Final = data.model_dump(exclude_none=True, exclude={"toolset_id"})
if "tools" in data_dict:
data_dict["tools"] = json.dumps(data_dict["tools"])
data_dict["updated_by"] = touched_by
"""A partial update: absent keeps, null clears. A toolset always has a name and a
tool list, so a null ``toolset_name`` or ``tools`` is a no-op rather than a clear;
emptying the tool selection is an explicit ``[]``, which cannot be mistaken for a
caller that left the field out."""
data_dict: Final = dict( # mutable-ok: Prisma requires a plain dict for JSON query serialization
(
(field, json.dumps(value) if field == "tools" else value)
for field, value in data.model_dump(exclude_unset=True).items()
if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None)
),
updated_by=touched_by,
)
try:
row: Final = await _toolset_table(prisma_client).update(
where={"toolset_id": data.toolset_id},

62
litellm/proxy/_types.py
View file

@ -3,6 +3,7 @@ import json
import os
from collections.abc import Callable, Mapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
import httpx
@ -1294,6 +1295,13 @@ class UpdateKeyRequest(KeyRequestBase):
rotation_interval: str | None = None
organization_id: str | None = None
@model_validator(mode="before")
@classmethod
def drop_blank_team_id(cls, values: object) -> object:
if isinstance(values, Mapping) and values.get("team_id") == "":
return MappingProxyType({k: v for k, v in values.items() if k != "team_id"})
return values
@field_validator("organization_id", mode="before")
@classmethod
def treat_cleared_organization_id_as_unset(cls, v: object) -> object:
@ -2828,6 +2836,18 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"UI username/password login. Default is False."
),
)
disable_env_credential_login: bool | None = Field(
None,
description=(
"If True, disables signing in to the Admin UI with the environment credentials: "
"UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback "
"means env-credential login is always live by default). Database users with passwords "
"are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password "
"before enabling, or nobody can sign in to the UI. A locked-out admin can still "
"administer the proxy over the API with the master key, and can unset this setting "
"and restart the proxy to restore env-credential login. Default is False."
),
)
disable_budget_reservation: bool | None = Field(
None,
description=(
@ -5148,9 +5168,26 @@ class CostEstimateRequest(LiteLLMPydanticObjectBase):
model: str = Field(description="Model name (from /model_group/info)")
input_tokens: int = Field(description="Expected input tokens per request", ge=0)
output_tokens: int = Field(description="Expected output tokens per request", ge=0)
cache_read_input_tokens: int = Field(
default=0, description="Input tokens read from the prompt cache; counted within input_tokens", ge=0
)
cache_creation_input_tokens: int = Field(
default=0, description="Input tokens written to the prompt cache; counted within input_tokens", ge=0
)
reasoning_tokens: int = Field(
default=0, description="Reasoning tokens the model emits; counted within output_tokens", ge=0
)
num_requests_per_day: int | None = Field(default=None, description="Number of requests per day", ge=0)
num_requests_per_month: int | None = Field(default=None, description="Number of requests per month", ge=0)
@model_validator(mode="after")
def validate_token_subsets(self) -> "CostEstimateRequest":
if self.cache_read_input_tokens + self.cache_creation_input_tokens > self.input_tokens:
raise ValueError("cache_read_input_tokens plus cache_creation_input_tokens cannot exceed input_tokens")
if self.reasoning_tokens > self.output_tokens:
raise ValueError("reasoning_tokens cannot exceed output_tokens")
return self
class CostEstimateResponse(LiteLLMPydanticObjectBase):
"""Response body for /cost/estimate endpoint."""
@ -5158,6 +5195,9 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase):
model: str
input_tokens: int
output_tokens: int
cache_read_input_tokens: int = 0
cache_creation_input_tokens: int = 0
reasoning_tokens: int = 0
num_requests_per_day: int | None = None
num_requests_per_month: int | None = None
# Per-request costs
@ -5165,17 +5205,33 @@ class CostEstimateResponse(LiteLLMPydanticObjectBase):
input_cost_per_request: float = Field(description="Input token cost per request (before margin)")
output_cost_per_request: float = Field(description="Output token cost per request (before margin)")
margin_cost_per_request: float = Field(default=0.0, description="Margin/fee added per request")
cache_read_cost_per_request: float = Field(default=0.0, description="Cache-read share of input_cost_per_request")
cache_creation_cost_per_request: float = Field(
default=0.0, description="Cache-write share of input_cost_per_request"
)
reasoning_cost_per_request: float = Field(default=0.0, description="Reasoning share of output_cost_per_request")
# Daily costs (if num_requests_per_day provided)
daily_cost: float | None = Field(default=None, description="Total daily cost (includes margin)")
daily_input_cost: float | None = Field(default=None, description="Daily input token cost")
daily_output_cost: float | None = Field(default=None, description="Daily output token cost")
daily_margin_cost: float | None = Field(default=None, description="Daily margin/fee")
daily_cache_read_cost: float | None = Field(default=None, description="Cache-read share of daily_input_cost")
daily_cache_creation_cost: float | None = Field(default=None, description="Cache-write share of daily_input_cost")
daily_reasoning_cost: float | None = Field(default=None, description="Reasoning share of daily_output_cost")
# Monthly costs (if num_requests_per_month provided)
monthly_cost: float | None = Field(default=None, description="Total monthly cost (includes margin)")
monthly_input_cost: float | None = Field(default=None, description="Monthly input token cost")
monthly_output_cost: float | None = Field(default=None, description="Monthly output token cost")
monthly_margin_cost: float | None = Field(default=None, description="Monthly margin/fee")
# Pricing info
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
monthly_cache_read_cost: float | None = Field(default=None, description="Cache-read share of monthly_input_cost")
monthly_cache_creation_cost: float | None = Field(
default=None, description="Cache-write share of monthly_input_cost"
)
monthly_reasoning_cost: float | None = Field(default=None, description="Reasoning share of monthly_output_cost")
# Pricing info: the rates this request's usage bills at, after token tiers and regional multipliers
input_cost_per_token: float | None = Field(default=None, description="Rate billed per input token")
output_cost_per_token: float | None = Field(default=None, description="Rate billed per output token")
cache_read_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-read token")
cache_creation_input_token_cost: float | None = Field(default=None, description="Rate billed per cache-write token")
output_cost_per_reasoning_token: float | None = Field(default=None, description="Rate billed per reasoning token")
provider: str | None = None

9
litellm/proxy/auth/auth_checks.py
View file

@ -475,6 +475,7 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
_TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True})
def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool:
@ -2858,7 +2859,9 @@ class TeamNotFoundError(HTTPException):
async def _get_team_db_check(
team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None
) -> "_PrismaTeamRow | None":
response = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id})
response = await _team_table(TeamRepository(prisma_client)).find_unique(
where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS
)
if response is None and team_id_upsert:
from litellm.proxy.management_endpoints.team_endpoints import new_team
@ -3158,7 +3161,9 @@ async def get_team_object_by_alias(
# Query database by team_alias
try:
teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(where={"team_alias": team_alias})
teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(
where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS
)
if not teams:
raise HTTPException(

5
litellm/proxy/auth/handle_jwt.py
View file

@ -53,6 +53,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_checks import can_team_access_model
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.team_grants import team_model_aliases
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
@ -1595,7 +1596,7 @@ class JWTAuthManager:
model=requested_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=None,
team_model_aliases=team_model_aliases(team_object),
)
):
is_allowed = allowed_routes_check(
@ -2132,7 +2133,7 @@ class JWTAuthManager:
model=requested_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=None,
team_model_aliases=team_model_aliases(team_object),
)
except ProxyException:
continue

39
litellm/proxy/auth/login_utils.py
View file

@ -85,6 +85,29 @@ def get_ui_credentials(master_key: str | None) -> tuple[str, str]:
return ui_username, ui_password
def _matches_env_credentials(username: str, password: str, master_key: str | None) -> bool:
ui_username, ui_password = get_ui_credentials(master_key)
return secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest(
password.encode("utf-8"), ui_password.encode("utf-8")
)
def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool:
"""Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed.
Two settings can turn it off: `disable_env_credential_login` unconditionally, and
`disable_password_login_when_sso_enabled` as a side effect, since its gate rejects
every username/password login before the env comparison runs. Feeds both the
`authenticate_user` gate and the Admin UI warning banner, so the banner never nags
about a login path that is already unreachable.
"""
if general_settings.get("disable_env_credential_login") is True:
return False
if general_settings.get("disable_password_login_when_sso_enabled") is True and is_sso_provider_fully_configured():
return False
return True
class LoginResult:
"""Result object containing authentication data from login."""
@ -129,7 +152,8 @@ async def authenticate_user(
master_key: Master key for the proxy (required)
prisma_client: Prisma database client (optional)
general_settings: Proxy general_settings, checked for
`disable_password_login_when_sso_enabled`
`disable_password_login_when_sso_enabled` and
`disable_env_credential_login`
Returns:
LoginResult: Object containing authentication data
@ -170,8 +194,6 @@ async def authenticate_user(
code=500,
)
ui_username, ui_password = get_ui_credentials(master_key)
# Check if we can find the `username` in the db. On the UI, users can enter username=their email
_user_row: LiteLLM_UserTable | None = None
user_role: (
@ -197,8 +219,8 @@ async def authenticate_user(
- Login with UI_USERNAME and UI_PASSWORD
- Login with Invite Link `user_email` and `password` combination
"""
if secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest(
password.encode("utf-8"), ui_password.encode("utf-8")
if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials(
username, password, master_key
):
# Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin
user_role = LitellmUserRoles.PROXY_ADMIN
@ -340,8 +362,13 @@ async def authenticate_user(
code=401,
)
else:
env_credentials_hint: Final = (
"\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file"
if is_env_credential_login_enabled(general_settings)
else ""
)
raise ProxyException(
message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file",
message=f"Invalid credentials used to access UI.{env_credentials_hint}",
type=ProxyErrorTypes.auth_error,
param="invalid_credentials",
code=401,

122
litellm/proxy/auth/team_grants.py Normal file
View file

@ -0,0 +1,122 @@
"""Project a team row (plus the caller's membership in it) onto the ``team_*`` fields of ``UserAPIKeyAuth``.
The virtual-key path gets these fields for free from the combined-view SQL join. Every other auth path
starts from a ``LiteLLM_TeamTable`` object instead and has to copy them over by hand, which is how JWT
callers kept losing grants (aliases, permissions, limits) one field at a time. Build the badge through
``team_grants`` and the two paths cannot drift.
"""
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Annotated, Final
from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError
from pydantic.main import IncEx
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
Member,
)
_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str])
_JSON_COLUMNS: Final[Mapping[str, IncEx | bool]] = MappingProxyType(
{"metadata": True, "litellm_model_table": MappingProxyType({"model_aliases": True})}
)
def _decode_model_aliases(value: object) -> object:
"""``LiteLLM_ModelTable.model_aliases`` is typed ``str | dict``; writers hand Prisma ``json.dumps(...)``, so take both."""
if not isinstance(value, str):
return value
try:
return _MODEL_ALIASES_ADAPTER.validate_json(value)
except ValidationError:
return None
class TeamModelAliasTable(BaseModel):
model_config = ConfigDict(protected_namespaces=())
model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None
class _TeamJsonColumns(BaseModel):
"""The two loosely typed columns on ``LiteLLM_TeamTable``, re-read with the shape the badge needs."""
metadata: Mapping[str, object] | None = None
litellm_model_table: TeamModelAliasTable | None = None
class TeamGrants(TypedDict, total=False):
"""Keyword arguments for ``UserAPIKeyAuth``. Empty when the caller has no team, so the model's own defaults apply."""
team_alias: ReadOnly[str | None]
team_tpm_limit: ReadOnly[int | None]
team_rpm_limit: ReadOnly[int | None]
team_max_budget: ReadOnly[float | None]
team_soft_budget: ReadOnly[float | None]
team_spend: ReadOnly[float | None]
team_models: ReadOnly[Sequence[str]]
team_blocked: ReadOnly[bool]
team_metadata: ReadOnly[Mapping[str, object] | None]
team_model_aliases: ReadOnly[Mapping[str, str] | None]
team_object_permission_id: ReadOnly[str | None]
team_object_permission: ReadOnly[LiteLLM_ObjectPermissionTable | None]
team_member: ReadOnly[Member | None]
team_member_spend: ReadOnly[float | None]
team_member_tpm_limit: ReadOnly[int | None]
team_member_rpm_limit: ReadOnly[int | None]
def _json_columns(team_object: LiteLLM_TeamTable) -> _TeamJsonColumns:
try:
return _TeamJsonColumns.model_validate(team_object.model_dump(include=_JSON_COLUMNS))
except ValidationError:
return _TeamJsonColumns()
def team_model_aliases(team_object: LiteLLM_TeamTable | None) -> Mapping[str, str] | None:
if team_object is None:
return None
alias_table: Final = _json_columns(team_object).litellm_model_table
return alias_table.model_aliases if alias_table is not None else None
def team_grants(
team_object: LiteLLM_TeamTable | None,
team_membership: LiteLLM_TeamMembership | None,
user_id: str | None,
) -> TeamGrants:
if team_object is None:
return TeamGrants()
json_columns: Final = _json_columns(team_object)
return TeamGrants(
team_alias=team_object.team_alias,
team_tpm_limit=team_object.tpm_limit,
team_rpm_limit=team_object.rpm_limit,
team_max_budget=team_object.max_budget,
team_soft_budget=team_object.soft_budget,
team_spend=team_object.spend,
team_models=tuple(team_object.models),
team_blocked=team_object.blocked,
team_metadata=json_columns.metadata,
team_model_aliases=(
json_columns.litellm_model_table.model_aliases if json_columns.litellm_model_table is not None else None
),
team_object_permission_id=team_object.object_permission_id,
team_object_permission=team_object.object_permission,
team_member=next(
(m for m in team_object.members_with_roles if user_id is not None and m.user_id == user_id),
None,
),
team_member_spend=team_membership.spend if team_membership is not None else None,
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None
),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None
),
)

22
litellm/proxy/auth/user_api_key_auth.py
View file

@ -82,6 +82,7 @@ from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
from litellm.proxy.auth.resolvers import CredentialRef, Principal
from litellm.proxy.auth.resolvers.store import IdentityStore
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.team_grants import team_grants
from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs
from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
from litellm.proxy.common_utils.http_parsing_utils import (
@ -1476,24 +1477,16 @@ async def _user_api_key_auth_builder(
user_id=user_id,
user_email=user_email,
team_id=team_id,
team_alias=(team_object.team_alias if team_object is not None else None),
team_tpm_limit=(team_object.tpm_limit if team_object is not None else None),
team_rpm_limit=(team_object.rpm_limit if team_object is not None else None),
team_models=(team_object.models if team_object is not None else []),
team_metadata=(team_object.metadata if team_object is not None else None),
org_id=org_id,
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
jwt_claims=jwt_claims,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
valid_token = UserAPIKeyAuth(
api_key=None,
team_id=team_id,
team_alias=(team_object.team_alias if team_object is not None else None),
team_tpm_limit=(team_object.tpm_limit if team_object is not None else None),
team_rpm_limit=(team_object.rpm_limit if team_object is not None else None),
team_models=(team_object.models if team_object is not None else []),
user_role=(
LitellmUserRoles(user_object.user_role)
if user_object is not None and user_object.user_role is not None
@ -1507,17 +1500,8 @@ async def _user_api_key_auth_builder(
user_tpm_limit=(user_object.tpm_limit if user_object is not None else None),
user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None
),
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None
),
team_metadata=(team_object.metadata if team_object is not None else None),
jwt_claims=jwt_claims,
)
valid_token.team_object_permission = (
team_object.object_permission if team_object is not None else None
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
# AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key.

22
litellm/proxy/client/cli/README.md
View file

@ -508,7 +508,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_
### Route Every Claude Code Session Through the Proxy
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you.
@ -532,12 +532,28 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi
lite --base-url https://your-proxy.example.com login --config-claude
```
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: no foreground process to keep alive, and `lite unconfigure claude` restores what it changed (see below). Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`.
Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops.
#### Configuring Claude Code Once, With a Virtual Key or Your Login
`lite configure claude` wires Claude Code up persistently and `lite unconfigure claude` puts things back. It is what `lite login --config-claude` does, plus a pinned model and an undo, and it also takes a long-lived virtual key when that is what you have:
```bash
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
lite --base-url https://your-proxy.example.com configure claude --api-key sk-... --model claude-auto
claude
```
With `--api-key` (or `lite --api-key` / `LITELLM_PROXY_API_KEY`) the key is written into `env.ANTHROPIC_AUTH_TOKEN`. Without one, your `lite login` credential is used the way `--config-claude` uses it, through `apiKeyHelper`, so a later `lite login` (or a `--pkce` renewal) picks up on its own and nothing secret lands in the file; a missing or stale login is refreshed first. Either way the command checks the key against `GET /v1/models`, then patches `~/.claude/settings.json`: `env.ANTHROPIC_BASE_URL`, the credential, and `env.ENABLE_TOOL_SEARCH` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` when those are missing, so Claude Code's `/model` picker lists the proxy's models (the ones whose id contains `claude` or `anthropic`) and you pick between them as usual. Claude Code keeps its own default model until you switch, so that id has to exist on the proxy for the first message to go through; `--model` (or the interactive prompt below) sets the model Claude Code starts on instead, as the top-level `model` key, which has to be on `/v1/models` for the key. Nothing forces Claude Code's sub-agent or background tiers onto a proxy model, so those built-in ids need to exist on the proxy too; `lite autoroute up` is the mode that pins every tier to one group. Claude Code treats a name it does not know as an unknown model: it prints a one-line `unrecognized_model` note, assumes a 200k context window and sends no thinking parameters for it, so either name the group like a Claude model id or append `[1m]` to opt into the 1M window. The other credential slots (`env.ANTHROPIC_API_KEY`, a stale `env.ANTHROPIC_AUTH_TOKEN` or `apiKeyHelper`) are removed so they cannot fight the one written. Every other setting is preserved and the file is written atomically with owner-only permissions; if `settings.json` is a symlink into a dotfiles repository, the key is written through to that target and the command says so, so keep it out of version control
Plain `lite configure`, with no agent named, asks the same things interactively: which agents to wire (Claude Code today) and which of the proxy's models to start on, picked from `/v1/models` with a type-to-filter prompt
What the command changed is recorded in `~/.litellm/claude_configure_state.json` (previous values plus fingerprints of what was written, never a second copy of the key). `lite unconfigure claude` restores each of those keys only if it still holds what `configure` wrote, so anything you changed since is left alone and named in the output; a `settings.json` or `env` object that only existed because of `configure` is removed again. Ownership moves only by a write: running `configure` again (a re-login is one) refreshes the record only for the keys its merge changed, keeps the original snapshot of a key that still holds what it wrote, and snapshots afresh a key you changed in between, so `unconfigure` brings back whatever the repeat displaced and never adopts your edit as its own. A credential (`env.ANTHROPIC_API_KEY`, `env.ANTHROPIC_AUTH_TOKEN`, `apiKeyHelper`) is put back only when the restored file points at the `ANTHROPIC_BASE_URL` it was captured next to; otherwise it stays removed, the output says which server it belonged to, and the receipt is kept so pointing the URL back and running `unconfigure` again finishes the job. It also undoes `lite login --config-claude`, which writes through the same path. Like `--config-claude`, both refuse to run while a `lite up` or `lite autoroute up` session holds a backup, and that check comes before any login prompt or request
### QA Complexity-Based Auto-Routing Against Your Real Proxy
`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session.
@ -584,7 +600,7 @@ An interactive wizard. It runs the same model-group discovery as above, splits t
The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/<model-name>` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key.
You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets the top-level `model` and `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, `ANTHROPIC_DEFAULT_OPUS_MODEL` and `ANTHROPIC_DEFAULT_FABLE_MODEL` to `autorouter` in `~/.claude/settings.json` (and `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when missing, like every other wiring), so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first.

5
litellm/proxy/client/cli/commands/agents.py
View file

@ -16,6 +16,7 @@ from .cmd_quoting import quote_for_cmd
from .pi import (
LITELLM_PROXY_API_KEY_ENV,
PI_PROVIDER_NAME,
ListingFailure,
PiSyncError,
fetch_model_ids,
fetch_model_limits,
@ -165,7 +166,9 @@ def prepare_pi(
"""
ids: Final = fetch_model_ids(base_url, api_key, get=get)
if isinstance(ids, PiSyncError):
raise AgentRunError(ids.message)
raise AgentRunError(
f"{ids.message} pi would have nothing to run." if ids.kind is ListingFailure.EMPTY else ids.message
)
limits: Final = fetch_model_limits(base_url, api_key, get=get)
path: Final = models_json_path(base_env)
error: Final = sync_models_json(path, base_url, ids, limits)

29
litellm/proxy/client/cli/commands/auth.py
View file

@ -41,9 +41,15 @@ from litellm.litellm_core_utils.cli_token_utils import (
from .claude_settings import (
CLAUDE_SETTINGS_PATH,
CONFIGURE_STATE_PATH,
SETTINGS_FILE_OWNERS,
STARTING_MODEL_ROLE,
ApiKeyHelper,
ClaudeSettingsError,
write_claude_settings,
KeepModel,
configure_claude_settings,
refuse_while_owned,
resolve_api_key_helper,
)
from .pkce_login import (
Http,
@ -778,13 +784,23 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None:
def _configure_claude_code(base_url: str) -> None:
"""Point Claude Code at base_url by patching ~/.claude/settings.json."""
"""Point Claude Code at base_url by patching ~/.claude/settings.json, undoable with `lite unconfigure claude`."""
try:
write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
configure_claude_settings(
base_url,
ApiKeyHelper(resolve_api_key_helper(base_url)),
KeepModel(),
CLAUDE_SETTINGS_PATH,
CONFIGURE_STATE_PATH,
SETTINGS_FILE_OWNERS,
)
except ClaudeSettingsError as e:
raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}")
click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.")
click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.")
click.echo(
"Your other Claude Code settings were left untouched. Restart Claude Code to pick this up. "
f"Undo with `lite unconfigure claude`; `lite configure claude --model` sets {STARTING_MODEL_ROLE}."
)
def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: SecretSave) -> None:
@ -853,6 +869,11 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None:
ctx_obj: Final[CliContextObj] = ctx.obj
base_url: Final = ctx_obj["base_url"]
if config_claude:
try:
refuse_while_owned(CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
except ClaudeSettingsError as e:
raise click.ClickException(f"Cannot configure Claude Code, so not logging in: {e}")
try:
if pkce:

9
litellm/proxy/client/cli/commands/autoroute/commands.py
View file

@ -14,11 +14,13 @@ from ..claude_settings import (
AUTOROUTE_BACKUP_PATH,
CLAUDE_SETTINGS_PATH,
ClaudeSettingsError,
StaticToken,
load_json_or_empty,
merge_claude_settings,
)
from ..up import BackupRecord as ClaudeBackupRecord
from ..up import restore_claude_settings, write_backup
from .config import master_key_from_config
from .config import AUTOROUTER_MODEL_NAME, master_key_from_config
from .process import (
CONFIG_PATH,
DEFAULT_AUTOROUTE_PORT,
@ -37,7 +39,6 @@ from .process import (
terminate,
write_pid_record,
)
from .settings import merge_claude_settings_static_token
from .wizard import run_configure_wizard
_GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
@ -156,7 +157,9 @@ def up(port: int) -> None:
ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None),
AUTOROUTE_BACKUP_PATH,
)
merged: Final = merge_claude_settings_static_token(original_settings, base_url, master_key)
merged: Final = merge_claude_settings(
original_settings, base_url, StaticToken(master_key), AUTOROUTER_MODEL_NAME, AUTOROUTER_MODEL_NAME
)
CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
with secure_create(CLAUDE_SETTINGS_PATH) as f:
json.dump(merged, f, indent=2)

51
litellm/proxy/client/cli/commands/autoroute/settings.py
View file

@ -1,51 +0,0 @@
from typing import Final
from pydantic import JsonValue
from .config import AUTOROUTER_MODEL_NAME
ENV_KEY: Final = "env"
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
# Force every one of Claude Code's own model tiers to request the auto-router by name.
# Router's auto-router registry is keyed by the literal requested model string
# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*"
# model_name can never work as a catch-all -- these overrides are what actually makes
# Claude Code send "autorouter" regardless of /model or its own version-specific defaults.
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = (
"ANTHROPIC_DEFAULT_SONNET_MODEL",
"ANTHROPIC_DEFAULT_HAIKU_MODEL",
"ANTHROPIC_DEFAULT_OPUS_MODEL",
)
def merge_claude_settings_static_token(
settings: dict[str, JsonValue], base_url: str, auth_token: str
) -> dict[str, JsonValue]:
"""Return a new settings dict wired to a local ephemeral proxy with a static token.
Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real
remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key is the
locally persisted autoroute master key, so a plain env var is simpler and correct. Any
existing apiKeyHelper is cleared so it can't fight with the static token.
"""
raw_env: Final = settings.get(ENV_KEY, {})
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final[dict[str, JsonValue]] = {
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
**base_env,
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
ANTHROPIC_AUTH_TOKEN_KEY: auth_token,
**{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS},
}
env.pop(ANTHROPIC_API_KEY_KEY, None)
merged: Final[dict[str, JsonValue]] = {**settings, ENV_KEY: env}
merged.pop(API_KEY_HELPER_KEY, None)
return merged
__all__ = ["merge_claude_settings_static_token"]

517
litellm/proxy/client/cli/commands/claude_settings.py
View file

@ -1,37 +1,70 @@
"""Shared handling of Claude Code's ~/.claude/settings.json.
`lite up` patches this file temporarily and restores it on exit; `lite login
--config-claude` patches it persistently. Both need the same merge and the same
apiKeyHelper command, and `up` already imports from `auth`, so the shared parts
live here rather than in either command module.
`lite up` and `lite autoroute up` patch this file temporarily and restore it on
exit; `lite login --config-claude` and `lite configure claude` patch it
persistently and record how to undo it. All of them need the same merge and the
same apiKeyHelper command, and `up` already imports from `auth`, so the shared
parts live here rather than in any one command module.
"""
import hashlib
import json
import shlex
import shutil
import sys
from collections.abc import Mapping, Sequence
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from functools import reduce
from itertools import chain
from pathlib import Path
from typing import Final
from types import MappingProxyType
from typing import Final, TypeAlias
from pydantic import JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from litellm.litellm_core_utils.private_json import write_private_json
from litellm.litellm_core_utils.private_json import (
commit_staged_json,
discard_staged_json,
ensure_private_dir,
stage_private_json,
)
from .cmd_quoting import quote_for_cmd
ENV_KEY: Final = "env"
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
MODEL_KEY: Final = "model"
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"
ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1"
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS: Final = (
"ANTHROPIC_DEFAULT_SONNET_MODEL",
"ANTHROPIC_DEFAULT_HAIKU_MODEL",
"ANTHROPIC_DEFAULT_OPUS_MODEL",
"ANTHROPIC_DEFAULT_FABLE_MODEL",
)
OWNED_ENV_KEYS: Final = (
ENABLE_TOOL_SEARCH_KEY,
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY,
ANTHROPIC_BASE_URL_KEY,
ANTHROPIC_AUTH_TOKEN_KEY,
ANTHROPIC_API_KEY_KEY,
)
OWNED_TOP_LEVEL_KEYS: Final = (API_KEY_HELPER_KEY, MODEL_KEY)
OWNED_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in OWNED_ENV_KEYS), *OWNED_TOP_LEVEL_KEYS)
_CREDENTIAL_ENV_KEYS: Final = frozenset((ANTHROPIC_API_KEY_KEY, ANTHROPIC_AUTH_TOKEN_KEY))
_CREDENTIAL_PATHS: Final = (*(f"{ENV_KEY}.{key}" for key in sorted(_CREDENTIAL_ENV_KEYS)), API_KEY_HELPER_KEY)
_BASE_URL_PATH: Final = f"{ENV_KEY}.{ANTHROPIC_BASE_URL_KEY}"
STARTING_MODEL_ROLE: Final = "the /model picker's default row, the model Claude Code starts on"
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json"
CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json"
@dataclass(frozen=True, slots=True)
@ -55,6 +88,102 @@ class ClaudeSettingsError(Exception):
"""Raised for any user-actionable failure while reading or writing Claude Code settings."""
@dataclass(frozen=True, slots=True)
class StaticToken:
"""A long-lived virtual key, written into env.ANTHROPIC_AUTH_TOKEN."""
token: str
@dataclass(frozen=True, slots=True)
class ApiKeyHelper:
"""A `lite auth print-token` command Claude Code runs per request, so a login renews in place."""
command: str
ClaudeCredential: TypeAlias = StaticToken | ApiKeyHelper
@dataclass(frozen=True, slots=True)
class KeepModel:
"""Leave the top-level `model` as it is, the user's or an earlier configure's (a re-login)."""
@dataclass(frozen=True, slots=True)
class UnpinModel:
"""Let go of a `model` an earlier configure pinned; one the user set themselves stays."""
@dataclass(frozen=True, slots=True)
class StartOn:
"""Pin the top-level `model`, the row Claude Code starts on."""
model: str
ModelChoice: TypeAlias = KeepModel | UnpinModel | StartOn
class OwnedValue(BaseModel):
"""What one key held at a moment in time; `present=False` is an absent key, not a null one."""
model_config = ConfigDict(frozen=True)
present: bool
value: JsonValue = None
class ConfigureReceipt(BaseModel):
"""What `lite configure claude` found and what it owns, keyed by dotted path (`env.X` or a top-level key).
Ownership moves only by a write: `written` fingerprints the keys some configure changed, at the
value it wrote; a repeat configure refreshes a fingerprint only for a key its merge changed and
carries the earlier one otherwise, so a key the user edited in between stops matching and is left
alone. `previous` is what each key held before configure took it over; a repeat keeps the earlier
snapshot while the key still holds our value and snapshots afresh otherwise, so whatever the
repeat displaces is what comes back. `endpoints` is the ANTHROPIC_BASE_URL each credential slot
was captured beside, so a credential is only ever put back next to the server it was issued for.
No fingerprint is a second copy of a token.
"""
model_config = ConfigDict(frozen=True)
file_existed: bool
env_present: bool
env_was_object: bool
previous: Mapping[str, OwnedValue]
written: Mapping[str, str]
endpoints: Mapping[str, OwnedValue]
@dataclass(frozen=True, slots=True)
class WithheldCredential:
"""A credential left removed: captured beside `endpoint`, while the restored file points elsewhere."""
key: str
endpoint: str
@dataclass(frozen=True, slots=True)
class UnconfigureOutcome:
"""Keys whose value unconfigure changed back, keys the user changed since and so were left as they
are, credentials withheld (the receipt is kept for them, so a later unconfigure can finish once the
URL points back), and whether no settings file remains."""
restored: tuple[str, ...]
kept: tuple[str, ...]
withheld: tuple[WithheldCredential, ...] = ()
file_removed: bool = False
@dataclass(frozen=True, slots=True)
class _Claim:
previous: OwnedValue
written: str | None
endpoint: OwnedValue | None
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
try:
content: Final = path.read_bytes() if path.exists() else b""
@ -70,29 +199,104 @@ def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
)
def merge_claude_settings(
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
) -> dict[str, JsonValue]:
"""Return a new settings dict wired to route Claude Code through the proxy.
def _env_object(settings: Mapping[str, JsonValue], path: Path) -> Mapping[str, JsonValue]:
raw_env: Final = settings.get(ENV_KEY)
if raw_env is None:
return MappingProxyType({})
if not isinstance(raw_env, dict):
raise ClaudeSettingsError(
f'{path} has a non-object "{ENV_KEY}" value, which this would discard. Fix or remove it, then retry.'
)
return raw_env
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH
defaults to true because Claude Code turns tool search off when
ANTHROPIC_BASE_URL is not a first-party Anthropic host, and
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker
is filled from the proxy's /v1/models; existing values of both are left
alone. Every other key is preserved untouched.
def refuse_while_owned(settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
"""Refuse while `lite up` or `lite autoroute up` holds a backup it will restore over any write; a
purely local check, so commands run it before any login prompt or request."""
for owner in owners:
if owner.backup_path.exists():
raise ClaudeSettingsError(
f"`{owner.start_command}` is currently managing {settings_path} (backup at "
f"{owner.backup_path}) and will restore it when it stops. "
f"Run `{owner.stop_command}` first, then retry."
)
def _write_target(settings_path: Path) -> Path:
"""Write through a symlinked settings.json rather than replacing the link, which would silently
detach a file symlinked into a dotfiles repo."""
try:
return settings_path.resolve() if settings_path.is_symlink() else settings_path
except OSError as e:
raise ClaudeSettingsError(f"Could not resolve {settings_path}: {e}") from e
def _stage(path: Path, document: Mapping[str, object]) -> str:
try:
return stage_private_json(str(path), document)
except OSError as e:
raise ClaudeSettingsError(f"Could not write {path}: {e}") from e
def _land(
path: Path,
staged: str | None,
also_discard: Sequence[str | None] = (),
commit: Callable[[str, str], None] = commit_staged_json,
) -> None:
"""Commit a staged file to `path`, or remove `path` when nothing is staged for it. The one place a
filesystem error becomes a ClaudeSettingsError; on failure the operation's other staged files are
discarded, so no temp file holding a token is left behind."""
try:
if staged is None:
path.unlink(missing_ok=True)
else:
commit(staged, str(path))
except OSError as e:
for other in also_discard:
if other is not None:
discard_staged_json(other)
raise ClaudeSettingsError(f"Could not {'remove' if staged is None else 'write'} {path}: {e}") from e
def merge_claude_settings(
settings: Mapping[str, JsonValue],
base_url: str,
credential: ClaudeCredential,
default_model: str | None = None,
tier_model: str | None = None,
) -> Mapping[str, JsonValue]:
"""Return a new settings mapping wired to route Claude Code through the proxy.
A StaticToken lands in env.ANTHROPIC_AUTH_TOKEN, an ApiKeyHelper in the top-level apiKeyHelper;
the other credential slots are removed either way, since Claude Code given two credentials may
send the wrong one. ENABLE_TOOL_SEARCH and CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY get their
defaults only when missing. `default_model` is the top-level `model`, the row Claude Code starts
on; `tier_model` is `lite autoroute up`'s knob that points every ANTHROPIC_DEFAULT_*_MODEL at one
group. Apart from those tier keys, exactly OWNED_PATHS are touched.
"""
raw_env: Final = settings.get(ENV_KEY, {})
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final = {
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE,
**{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY},
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
}
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
current_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final = dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
chain(
(
(ENABLE_TOOL_SEARCH_KEY, ENABLE_TOOL_SEARCH_VALUE),
(ENABLE_GATEWAY_MODEL_DISCOVERY_KEY, ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE),
),
((key, value) for key, value in current_env.items() if key not in _CREDENTIAL_ENV_KEYS),
((ANTHROPIC_BASE_URL_KEY, base_url.rstrip("/")),),
((ANTHROPIC_AUTH_TOKEN_KEY, credential.token),) if isinstance(credential, StaticToken) else (),
((key, tier_model) for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS if tier_model is not None),
)
)
return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
chain(
((key, value) for key, value in settings.items() if key not in (API_KEY_HELPER_KEY, ENV_KEY)),
((ENV_KEY, env),),
((API_KEY_HELPER_KEY, credential.command),) if isinstance(credential, ApiKeyHelper) else (),
((MODEL_KEY, default_model),) if default_model is not None else (),
)
)
def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
@ -121,56 +325,255 @@ def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str:
return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token"))
def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
"""Persistently point Claude Code at base_url, preserving every unrelated setting.
def _owned(container: Mapping[str, JsonValue], key: str) -> OwnedValue:
return OwnedValue(present=key in container, value=container.get(key))
Refuses while any owner holds a backup: each restores its backup when it
stops, which would silently undo this write.
"""
for owner in owners:
if owner.backup_path.exists():
raise ClaudeSettingsError(
f"`{owner.start_command}` is currently managing {settings_path} (backup at "
f"{owner.backup_path}) and will restore it when it stops. "
f"Run `{owner.stop_command}` first, then retry."
)
normalized_base_url: Final = base_url.rstrip("/")
api_key_helper: Final = resolve_api_key_helper(normalized_base_url)
existing: Final = load_json_or_empty(settings_path)
raw_env: Final = existing.get(ENV_KEY)
if raw_env is not None and not isinstance(raw_env, dict):
raise ClaudeSettingsError(
f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. '
"Fix or remove it, then retry."
)
merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper)
# os.replace() swaps the symlink itself for a regular file, silently detaching a
# settings.json that is symlinked into a dotfiles repo. There is no backup to undo
# that here, unlike `lite up`, so write through to the link's target instead.
target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path
def _fingerprint(owned: OwnedValue) -> str:
return hashlib.sha256(json.dumps(owned.model_dump(mode="json"), sort_keys=True).encode()).hexdigest()
def _env(settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
raw_env: Final = settings.get(ENV_KEY)
return raw_env if isinstance(raw_env, dict) else MappingProxyType({})
def _lookup(settings: Mapping[str, JsonValue], path: str) -> OwnedValue:
section, _, key = path.rpartition(".")
return _owned(_env(settings) if section else settings, key)
def _with_key(container: Mapping[str, JsonValue], key: str, owned: OwnedValue) -> Mapping[str, JsonValue]:
return dict( # mutable-ok: JSON document handed to json.dump, which rejects a read-only mapping
chain(((k, v) for k, v in container.items() if k != key), ((key, owned.value),) if owned.present else ())
)
def _with(settings: Mapping[str, JsonValue], path: str, owned: OwnedValue) -> Mapping[str, JsonValue]:
"""`settings` with the key at `path` set (or removed when `owned` is absent); nothing else changes."""
section, _, key = path.rpartition(".")
if not section:
return _with_key(settings, key, owned)
return _with_key(settings, section, OwnedValue(present=True, value=_with_key(_env(settings), key, owned)))
def _with_all(settings: Mapping[str, JsonValue], updates: Mapping[str, OwnedValue]) -> Mapping[str, JsonValue]:
return reduce(lambda acc, item: _with(acc, *item), updates.items(), settings)
def _ours(settings: Mapping[str, JsonValue], path: str, receipt: ConfigureReceipt) -> bool:
"""Whether the key still holds what a configure wrote (a key no configure ever changed is never ours)."""
return receipt.written.get(path) == _fingerprint(_lookup(settings, path))
def _claim(
path: str,
current: Mapping[str, JsonValue],
merged: Mapping[str, JsonValue],
earlier: ConfigureReceipt | None,
url_now: OwnedValue,
) -> _Claim:
"""What this configure records for one key; see ConfigureReceipt for the rules."""
before, after = _lookup(current, path), _lookup(merged, path)
carried: Final = earlier if earlier is not None and _ours(current, path, earlier) else None
return _Claim(
previous=before if carried is None else carried.previous.get(path, before),
written=_fingerprint(after) if before != after else (None if earlier is None else earlier.written.get(path)),
endpoint=None
if path not in _CREDENTIAL_PATHS
else (url_now if carried is None else carried.endpoints.get(path, url_now)),
)
def _receipt(
current: Mapping[str, JsonValue],
merged: Mapping[str, JsonValue],
earlier: ConfigureReceipt | None,
file_exists: bool,
) -> ConfigureReceipt:
url_now: Final = _lookup(current, _BASE_URL_PATH)
claims: Final = MappingProxyType({path: _claim(path, current, merged, earlier, url_now) for path in OWNED_PATHS})
return ConfigureReceipt(
file_existed=file_exists if earlier is None else earlier.file_existed,
env_present=ENV_KEY in current if earlier is None else earlier.env_present,
env_was_object=isinstance(current.get(ENV_KEY), dict) if earlier is None else earlier.env_was_object,
previous=MappingProxyType({path: claim.previous for path, claim in claims.items()}),
written=MappingProxyType({path: claim.written for path, claim in claims.items() if claim.written is not None}),
endpoints=MappingProxyType(
{path: claim.endpoint for path, claim in claims.items() if claim.endpoint is not None}
),
)
def read_configure_receipt(state_path: Path) -> ConfigureReceipt | None:
if not state_path.exists():
return None
try:
write_private_json(str(target), merged)
return ConfigureReceipt.model_validate_json(state_path.read_bytes())
except (OSError, ValidationError) as e:
raise ClaudeSettingsError(
f"{state_path} is not a readable `lite configure claude` receipt ({e}). "
"Remove it and edit Claude Code's settings by hand if they still point at the proxy."
) from e
def configure_claude_settings(
base_url: str,
credential: ClaudeCredential,
model: ModelChoice,
settings_path: Path,
state_path: Path,
owners: Sequence[SettingsFileOwner],
commit: Callable[[str, str], None] = commit_staged_json,
) -> None:
"""Persistently route Claude Code through base_url, recording how to undo it.
Both files are staged before either is committed, so a full disk or a read-only directory fails
before anything changes. The two commits are still two renames: a receipt rename that fails
discards the staged settings, and a settings rename that fails after the receipt landed puts the
earlier receipt back (or removes the new one), so the receipt on disk never describes settings
that were not written. `model`: StartOn pins the starting model, UnpinModel lets go of a pin an
earlier configure made (never of the user's own), KeepModel leaves it alone (a re-login).
"""
refuse_while_owned(settings_path, owners)
current: Final = load_json_or_empty(settings_path)
_env_object(current, settings_path)
earlier: Final = read_configure_receipt(state_path)
existing: Final = (
_with(current, MODEL_KEY, earlier.previous[MODEL_KEY])
if isinstance(model, UnpinModel) and earlier is not None and _ours(current, MODEL_KEY, earlier)
else current
)
merged: Final = merge_claude_settings(
existing, base_url, credential, model.model if isinstance(model, StartOn) else None
)
receipt: Final = _receipt(current, merged, earlier, settings_path.exists())
target: Final = _write_target(settings_path)
try:
ensure_private_dir(state_path.parent)
except OSError as e:
raise ClaudeSettingsError(f"Could not write {target}: {e}") from e
raise ClaudeSettingsError(f"Could not write {state_path}: {e}") from e
staged_receipt: Final = _stage(state_path, receipt.model_dump(mode="json"))
try:
staged_settings: Final = _stage(target, merged)
except ClaudeSettingsError:
discard_staged_json(staged_receipt)
raise
_land(state_path, staged_receipt, (staged_settings,), commit)
try:
_land(target, staged_settings, commit=commit)
except ClaudeSettingsError as settings_error:
try:
_land(state_path, None if earlier is None else _stage(state_path, earlier.model_dump(mode="json")))
except ClaudeSettingsError as receipt_error:
raise ClaudeSettingsError(
f"{settings_error} The receipt at {state_path} now describes settings that were not written and "
f"could not be put back either ({receipt_error}); remove it before retrying."
) from settings_error
raise
def _endpoint_text(endpoint: OwnedValue) -> str:
if not endpoint.present:
return f"no {ANTHROPIC_BASE_URL_KEY} (Anthropic's default endpoint)"
return endpoint.value if isinstance(endpoint.value, str) else json.dumps(endpoint.value)
def unconfigure_claude_settings(
settings_path: Path, state_path: Path, owners: Sequence[SettingsFileOwner]
) -> UnconfigureOutcome:
"""Undo `lite configure claude`: put back every key still holding what configure wrote, leave the
rest alone, and withhold a credential the restored file would send to a different server than it
was issued for (the receipt stays, owning only those slots, so a later unconfigure can finish)."""
refuse_while_owned(settings_path, owners)
receipt: Final = read_configure_receipt(state_path)
if receipt is None:
raise ClaudeSettingsError(
f"Claude Code is not configured by `lite configure claude` (no receipt at {state_path}); nothing to undo."
)
current: Final = load_json_or_empty(settings_path)
_env_object(current, settings_path)
ours: Final = tuple(path for path in receipt.written if _ours(current, path, receipt))
kept: Final = tuple(path for path in receipt.written if path not in ours and _lookup(current, path).present)
put_back: Final = _with_all(current, MappingProxyType({path: receipt.previous[path] for path in ours}))
url_after: Final = _lookup(put_back, _BASE_URL_PATH)
withheld: Final = tuple(
WithheldCredential(path, _endpoint_text(receipt.endpoints[path]))
for path in _CREDENTIAL_PATHS
if path in ours and receipt.previous[path].present and receipt.endpoints[path] != url_after
)
absent: Final = OwnedValue(present=False)
trimmed: Final = _with_all(put_back, MappingProxyType({item.key: absent for item in withheld}))
settings: Final = (
trimmed
if _env(trimmed) or receipt.env_was_object
else _with_key(trimmed, ENV_KEY, OwnedValue(present=receipt.env_present, value=None))
)
target: Final = _write_target(settings_path)
file_removed: Final = not settings and not (receipt.file_existed and target.exists())
kept_receipt: Final = ( # mutable-ok: pydantic serializes the update as given and rejects a mappingproxy
receipt.model_copy(update={"written": {item.key: _fingerprint(absent) for item in withheld}})
if withheld
else None
)
staged_settings: Final = None if file_removed else _stage(target, settings)
try:
staged_receipt: Final = (
None if kept_receipt is None else _stage(state_path, kept_receipt.model_dump(mode="json"))
)
except ClaudeSettingsError:
if staged_settings is not None:
discard_staged_json(staged_settings)
raise
_land(target, staged_settings, (staged_receipt,))
_land(state_path, staged_receipt)
return UnconfigureOutcome(
restored=tuple(path for path in ours if _lookup(current, path) != _lookup(settings, path)),
kept=kept,
withheld=withheld,
file_removed=file_removed,
)
__all__ = (
"ANTHROPIC_API_KEY_KEY",
"ANTHROPIC_AUTH_TOKEN_KEY",
"ANTHROPIC_BASE_URL_KEY",
"ANTHROPIC_DEFAULT_MODEL_ENV_KEYS",
"API_KEY_HELPER_KEY",
"AUTOROUTE_BACKUP_PATH",
"BACKUP_PATH",
"CLAUDE_SETTINGS_PATH",
"CONFIGURE_STATE_PATH",
"ENABLE_GATEWAY_MODEL_DISCOVERY_KEY",
"ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE",
"ENABLE_TOOL_SEARCH_KEY",
"ENABLE_TOOL_SEARCH_VALUE",
"ENV_KEY",
"MODEL_KEY",
"OWNED_ENV_KEYS",
"OWNED_PATHS",
"OWNED_TOP_LEVEL_KEYS",
"SETTINGS_FILE_OWNERS",
"STARTING_MODEL_ROLE",
"ApiKeyHelper",
"ClaudeCredential",
"ClaudeSettingsError",
"ConfigureReceipt",
"KeepModel",
"ModelChoice",
"OwnedValue",
"SettingsFileOwner",
"StartOn",
"StaticToken",
"UnconfigureOutcome",
"UnpinModel",
"WithheldCredential",
"configure_claude_settings",
"load_json_or_empty",
"merge_claude_settings",
"read_configure_receipt",
"refuse_while_owned",
"resolve_api_key_helper",
"write_claude_settings",
"unconfigure_claude_settings",
)

252
litellm/proxy/client/cli/commands/configure.py Normal file
View file

@ -0,0 +1,252 @@
"""`lite configure claude` and `lite unconfigure claude`: persistent Claude Code wiring, undoable."""
import re
import sys
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Final
import click
from InquirerPy import inquirer
from InquirerPy.base.control import Choice
from .auth import CliContextObj, context_secret_vault, get_stored_api_key
from .claude_settings import (
CLAUDE_SETTINGS_PATH,
CONFIGURE_STATE_PATH,
SETTINGS_FILE_OWNERS,
STARTING_MODEL_ROLE,
ApiKeyHelper,
ClaudeCredential,
ClaudeSettingsError,
ModelChoice,
StartOn,
StaticToken,
UnconfigureOutcome,
UnpinModel,
configure_claude_settings,
refuse_while_owned,
resolve_api_key_helper,
unconfigure_claude_settings,
)
from .pi import ListingFailure, PiSyncError, fetch_model_ids
from .up import ensure_fresh_login
_LISTED_MODELS_SHOWN: Final = 20
_CLAUDE_TARGET: Final = "claude"
_TARGETS: Final = ((_CLAUDE_TARGET, "Claude Code (CLI)"),)
_KEEP_DEFAULT_MODEL: Final = "Keep Claude Code's own default"
_CLAUDE_CODE_PICKER_FILTER: Final = re.compile(r"claude|anthropic", re.IGNORECASE)
_MODEL_OPTION_HELP: Final = (
f"Proxy model to set as {STARTING_MODEL_ROLE}. Must be listed on /v1/models for the key; without it, "
"Claude Code keeps its own default and a pin an earlier configure made is let go of. Nothing pins Claude "
"Code's sub-agent or background tiers; `lite autoroute up` is the mode that does."
)
def resolve_credential(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, str]:
"""The credential to write and the key to check the proxy with.
An explicit key (--api-key, `lite --api-key`, LITELLM_PROXY_API_KEY) is long-lived and goes
into settings.json as a static token. Without one, the stored `lite login` credential is used
the way `lite login --config-claude` uses it, through apiKeyHelper, since it expires within a
day and renews in place there; a missing or stale login is refreshed first, as `lite up` does.
"""
ctx_obj: Final[CliContextObj] = ctx.obj
explicit: Final = api_key or (None if ctx_obj.get("api_key_from_token_file") else ctx_obj.get("api_key"))
if explicit:
return StaticToken(explicit), explicit
base_url: Final = ctx_obj["base_url"]
ensure_fresh_login(ctx)
stored: Final = get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx))
if not stored:
raise ClaudeSettingsError("Login did not produce a usable token.")
return ApiKeyHelper(resolve_api_key_helper(base_url)), stored
def _start(ctx: click.Context, api_key: str | None) -> tuple[ClaudeCredential, tuple[str, ...]]:
"""Every configure path begins the same way: the local ownership check first, so a `lite up`
session is refused before any login prompt or request, then the credential, then the listing."""
try:
refuse_while_owned(CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
credential, key = resolve_credential(ctx, api_key)
except ClaudeSettingsError as e:
raise click.ClickException(str(e))
return credential, _listed_models(ctx.obj["base_url"], key)
def _listing_error(base_url: str, error: PiSyncError) -> str:
"""The hint that fits how the listing failed: only an unreachable proxy gets the "is it running" question."""
if error.kind is ListingFailure.REJECTED:
return f"LiteLLM rejected your key (HTTP {error.status}). Run `lite login` to refresh it, or pass a valid --api-key."
if error.kind is ListingFailure.UNREACHABLE:
return f"{error.message} Is the proxy at {base_url} running, and is --base-url (or LITELLM_PROXY_URL) correct?"
if error.kind is ListingFailure.EMPTY:
return f"{error.message} Claude Code would have nothing to run; give the key access to at least one model."
return f"{error.message} The proxy at {base_url} answered, so check that it is a LiteLLM proxy and is healthy."
def _listed_models(base_url: str, key: str) -> tuple[str, ...]:
listed: Final = fetch_model_ids(base_url, key)
if isinstance(listed, PiSyncError):
raise click.ClickException(_listing_error(base_url, listed))
return listed
def _model_choice(model: str | None) -> ModelChoice:
return StartOn(model) if model is not None else UnpinModel()
def _apply_claude(ctx: click.Context, credential: ClaudeCredential, listed: Sequence[str], model: str | None) -> None:
ctx_obj: Final[CliContextObj] = ctx.obj
base_url: Final = ctx_obj["base_url"]
if model is not None and model not in listed:
shown: Final = ", ".join(listed[:_LISTED_MODELS_SHOWN])
more: Final = f", and {len(listed) - _LISTED_MODELS_SHOWN} more" if len(listed) > _LISTED_MODELS_SHOWN else ""
raise click.ClickException(
f"{model!r} is not served by {base_url} for this key. /v1/models lists: {shown}{more}."
)
try:
configure_claude_settings(
base_url, credential, _model_choice(model), CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, SETTINGS_FILE_OWNERS
)
except ClaudeSettingsError as e:
raise click.ClickException(str(e))
in_picker: Final = sum(1 for listed_model in listed if _CLAUDE_CODE_PICKER_FILTER.search(listed_model))
click.echo(f"Configured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url}.")
click.echo(
"Credential: your virtual key, stored in the file as ANTHROPIC_AUTH_TOKEN."
if isinstance(credential, StaticToken)
else "Credential: your `lite login`, read through apiKeyHelper on every request, so a later login renews it."
)
click.echo(
f"Starting model: {model} ({STARTING_MODEL_ROLE}); switch any time with /model."
if model is not None
else "Starting model: not pinned (Claude Code's default, or a model you set yourself); switch with /model, or "
"pass --model to start on a proxy model."
)
click.echo(
f"/model will list {in_picker} of the proxy's {len(listed)} models (Claude Code shows only ids containing "
"'claude' or 'anthropic')."
)
click.echo("Start `claude` from any terminal. Undo with `lite unconfigure claude`.")
if isinstance(credential, StaticToken) and CLAUDE_SETTINGS_PATH.is_symlink():
click.echo(
f"Note: {CLAUDE_SETTINGS_PATH} is a symlink to {CLAUDE_SETTINGS_PATH.resolve()}, so your key now lives in "
"that file; keep it out of version control.",
err=True,
)
def _pick_targets() -> tuple[str, ...]:
picked: Final = inquirer.checkbox(
message="Which agents should route through LiteLLM?",
choices=[Choice(value, name=label, enabled=True) for value, label in _TARGETS],
validate=lambda chosen: len(chosen) > 0,
invalid_message="Pick at least one.",
).execute()
return tuple(str(value) for value in picked)
def _pick_model(listed: Sequence[str]) -> str | None:
picked: Final = inquirer.fuzzy(
message="Model Claude Code starts on (type to filter; /model switches any time):",
choices=[_KEEP_DEFAULT_MODEL, *listed],
).execute()
return None if picked == _KEEP_DEFAULT_MODEL else str(picked)
def interactive_configure(
ctx: click.Context,
pick_targets: Callable[[], tuple[str, ...]] = _pick_targets,
pick_model: Callable[[Sequence[str]], str | None] = _pick_model,
) -> None:
"""`lite configure` with no agent named: ask which agents to wire and which model to pin."""
targets: Final = pick_targets()
if _CLAUDE_TARGET not in targets:
return
credential, listed = _start(ctx, None)
_apply_claude(ctx, credential, listed, pick_model(listed))
@click.group(name="configure", invoke_without_command=True)
@click.pass_context
def configure_group(ctx: click.Context) -> None:
"""Persistently route a coding agent through your LiteLLM proxy.
With no agent named, asks which agents to wire and which proxy model to pin.
"""
if ctx.invoked_subcommand is not None:
return
if not sys.stdin.isatty():
raise click.ClickException(
"`lite configure` asks questions, so it needs a terminal. Non-interactively, run "
"`lite configure claude --api-key <key> --model <model>`."
)
interactive_configure(ctx)
@click.group(name="unconfigure")
def unconfigure_group() -> None:
"""Undo `lite configure` for a coding agent."""
@configure_group.command(name="claude")
@click.option(
"--api-key",
"api_key",
default=None,
help="Long-lived LiteLLM virtual key written into Claude Code's settings. Defaults to the `lite --api-key` / "
"LITELLM_PROXY_API_KEY value; with neither, your `lite login` credential is used through apiKeyHelper.",
)
@click.option("--model", default=None, help=_MODEL_OPTION_HELP)
@click.pass_context
def configure_claude(ctx: click.Context, api_key: str | None, model: str | None) -> None:
"""Route every Claude Code session through your LiteLLM proxy until `lite unconfigure claude`.
Patches ~/.claude/settings.json in place: the proxy URL, your credential (a virtual key as a
static token, or your `lite login` through apiKeyHelper), and gateway model discovery so
/model lists the proxy's models; --model picks the one Claude Code starts on. Every other
setting is kept, and what changed is recorded so `lite unconfigure claude` can put it back.
Assumes the proxy is already running.
"""
credential, listed = _start(ctx, api_key)
_apply_claude(ctx, credential, listed, model)
@unconfigure_group.command(name="claude")
def unconfigure_claude() -> None:
"""Return Claude Code's settings to what they were before `lite configure claude`.
Also undoes `lite login --config-claude`. Only keys still holding what configure wrote are
put back; anything you changed since is left as it is and named in the output.
"""
try:
outcome: Final = unconfigure_claude_settings(CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, SETTINGS_FILE_OWNERS)
except ClaudeSettingsError as e:
raise click.ClickException(str(e))
_report_unconfigure(CLAUDE_SETTINGS_PATH, CONFIGURE_STATE_PATH, outcome)
def _report_unconfigure(settings_path: Path, state_path: Path, outcome: UnconfigureOutcome) -> None:
"""Say what unconfigure did, naming only keys whose value it changed."""
if outcome.file_removed:
click.echo(
f"No settings file remains at {settings_path}; it held nothing but `lite configure claude`'s own keys."
)
elif outcome.restored:
click.echo(f"Restored in {settings_path}: {', '.join(outcome.restored)}.")
else:
click.echo(f"Nothing in {settings_path} was still ours to restore.")
if outcome.kept:
click.echo(f"Left as you changed them since: {', '.join(outcome.kept)}.")
if outcome.withheld:
click.echo(
"Left removed, since the file now points at a different server than they were issued for: "
+ "; ".join(f"{item.key} (captured with {item.endpoint})" for item in outcome.withheld)
+ f". They stay in {state_path}: point env.ANTHROPIC_BASE_URL back and run `lite unconfigure claude` "
"again to put them back, or delete that file to drop them."
)
__all__ = ("configure_group", "interactive_configure", "resolve_credential", "unconfigure_group")

31
litellm/proxy/client/cli/commands/pi.py
View file

@ -10,6 +10,7 @@ import os
import tempfile
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from enum import StrEnum
from pathlib import Path
from types import MappingProxyType
from typing import Final
@ -20,11 +21,28 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR"
PI_PROVIDER_NAME: Final = "litellm"
LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
_REJECTED_STATUSES: Final = frozenset((401, 403))
class ListingFailure(StrEnum):
"""Why a proxy could not be listed, decided once where the HTTP outcome is classified.
`unreachable` means no response at all; the other kinds prove the proxy answered, so callers
must not suggest checking whether it is running.
"""
UNREACHABLE = "unreachable"
REJECTED = "rejected"
BAD_BODY = "bad_body"
EMPTY = "empty"
OTHER = "other"
@dataclass(frozen=True, slots=True)
class PiSyncError:
message: str
status: int | None = None
kind: ListingFailure | None = None
@dataclass(frozen=True, slots=True)
@ -65,16 +83,20 @@ def fetch_model_ids(
timeout=10,
)
except requests.RequestException as e:
return PiSyncError(f"Could not list models from the proxy: {e}")
return PiSyncError(f"Could not list models from the proxy: {e}", kind=ListingFailure.UNREACHABLE)
if resp.status_code != 200:
return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.")
return PiSyncError(
f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot list models.",
resp.status_code,
ListingFailure.REJECTED if resp.status_code in _REJECTED_STATUSES else ListingFailure.OTHER,
)
try:
listing: Final = _ModelList.model_validate(resp.json())
except (ValueError, ValidationError) as e:
return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}")
return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY)
ids: Final = tuple(dict.fromkeys(model.id for model in listing.data))
if not ids:
return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.")
return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY)
return ids
@ -200,6 +222,7 @@ __all__ = (
"LITELLM_PROXY_API_KEY_ENV",
"PI_CONFIG_DIR_ENV",
"PI_PROVIDER_NAME",
"ListingFailure",
"ModelLimits",
"PiSyncError",
"fetch_model_ids",

9
litellm/proxy/client/cli/commands/up.py
View file

@ -23,6 +23,7 @@ from .auth import CliContextObj, context_secret_vault, get_stored_api_key, load_
from .claude_settings import (
BACKUP_PATH,
CLAUDE_SETTINGS_PATH,
ApiKeyHelper,
ClaudeSettingsError,
load_json_or_empty,
merge_claude_settings,
@ -123,7 +124,7 @@ def _stored_login_is_pkce(vault: SecretVault) -> bool:
return token_data is not None and token_data.get("refresh_token") is not None
def _ensure_fresh_login(ctx: click.Context) -> None:
def ensure_fresh_login(ctx: click.Context) -> None:
ctx_obj: Final[CliContextObj] = ctx.obj
base_url: Final = ctx_obj["base_url"].rstrip("/")
vault: Final = context_secret_vault(ctx)
@ -141,7 +142,7 @@ def _ensure_fresh_login(ctx: click.Context) -> None:
click.echo("No fresh LiteLLM login found for this proxy; starting login...")
ctx.invoke(login, pkce=pkce)
if not _usable_login(get_stored_api_key(expected_base_url=base_url, vault=vault), vault):
raise UpError("Login did not produce a usable token; cannot start `lite up`.")
raise UpError("Login did not produce a usable token.")
def _restore_and_report() -> None:
@ -169,7 +170,7 @@ def up(ctx: click.Context) -> None:
base_url: Final = ctx.obj["base_url"]
try:
_ensure_fresh_login(ctx)
ensure_fresh_login(ctx)
api_key: Final = resolve_api_key(ctx)
verify_proxy_key(base_url, api_key)
@ -190,7 +191,7 @@ def up(ctx: click.Context) -> None:
)
CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True)
merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper)
merged: Final = merge_claude_settings(original_settings, base_url, ApiKeyHelper(api_key_helper))
with open(CLAUDE_SETTINGS_PATH, "w") as f:
json.dump(merged, f, indent=2)
except (AgentRunError, ClaudeSettingsError) as e:

4
litellm/proxy/client/cli/main.py
View file

@ -13,6 +13,7 @@ from .commands.auth import auth_group, context_secret_vault, get_stored_api_key,
from .commands.autoroute.commands import autoroute_group
from .commands.chat import chat
from .commands.config import config_commands, get_config_value, hidden_command_names
from .commands.configure import configure_group, unconfigure_group
from .commands.credentials import credentials
from .commands.debug import debug
from .commands.encryption import encryption
@ -162,6 +163,9 @@ cli.add_command(model_groups)
# Add the autoroute command group (QA auto-routing against your real proxy)
cli.add_command(autoroute_group, name="autoroute")
cli.add_command(config_commands)
# Add configure/unconfigure (persistently wire a coding agent to the proxy with a virtual key)
cli.add_command(configure_group)
cli.add_command(unconfigure_group)
if __name__ == "__main__":

23
litellm/proxy/common_request_processing.py
View file

@ -190,6 +190,7 @@ from litellm.proxy.litellm_pre_call_utils import (
refresh_proxy_server_request_body_snapshot,
reject_url_valued_destination,
)
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -1849,7 +1850,6 @@ class ProxyBaseLLMRequestProcessing:
data=self.data,
user_api_key_dict=user_api_key_dict,
)
# Calculate request queue time after add_litellm_data_to_request
# which sets arrival_time in proxy_server_request. Ends at start_time
# (not a freshly captured time.time() here) so this window is exactly
@ -1997,6 +1997,12 @@ class ProxyBaseLLMRequestProcessing:
data=self.data,
call_type=route_type,
)
if route_type == "aget_responses":
attach_post_call_pipelines_to_retrieval(
data=self.data,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
# Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may
# have mutated `self.data` in place, and the audit-trail snapshot taken in
@ -3294,9 +3300,10 @@ class ProxyBaseLLMRequestProcessing:
has completed.
Guardrails routed through unified_guardrail are skipped, since they already ran
via its streaming iterator. Guardrails that override
async_post_call_success_hook directly run here, including those that implement
apply_guardrail but keep their native lifecycle hooks.
via its streaming iterator, and so are guardrails a post_call policy pipeline
manages, since the pipeline ran them against the buffered stream. Guardrails
that override async_post_call_success_hook directly run here, including those
that implement apply_guardrail but keep their native lifecycle hooks.
This is audit-only — content has already been delivered to the client.
@ -3306,12 +3313,18 @@ class ProxyBaseLLMRequestProcessing:
_response = assembled_response
try:
from litellm.proxy.proxy_server import llm_router as _global_llm_router
from litellm.proxy.utils import _check_and_merge_model_level_guardrails
from litellm.proxy.utils import (
_check_and_merge_model_level_guardrails,
stream_gated_guardrail_names,
)
guardrail_data = _check_and_merge_model_level_guardrails(data=captured_data, llm_router=_global_llm_router)
stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict)
for cb in litellm.callbacks:
if not isinstance(cb, CustomGuardrail):
continue
if cb.guardrail_name in stream_gated:
continue
if not cb.should_run_guardrail(
data=guardrail_data,
event_type=GuardrailEventHooks.post_call,

62
litellm/proxy/common_utils/callback_utils.py
View file

@ -1,10 +1,12 @@
import copy
import json
import os
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from itertools import accumulate
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias
from typing_extensions import assert_never
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
from litellm import get_secret
@ -12,6 +14,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
@ -28,6 +31,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.types_utils.utils import get_instance_fn
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
StandardLoggingGuardrailInformation,
StandardLoggingPayload,
@ -52,6 +56,15 @@ reset_color_code: Final = "\033[0m"
TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY: Final = "_pillar_response_headers_trusted"
GUARDRAIL_SCAN_IDS_METADATA_KEY: Final = "guardrail_scan_ids"
GUARDRAIL_SCAN_METADATA_METADATA_KEY: Final = "guardrail_scan_metadata"
class GuardrailScanMetadata(TypedDict):
guardrail: ReadOnly[str | None]
stage: ReadOnly[str]
provider: ReadOnly[str]
scan_id: ReadOnly[str]
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
@ -450,6 +463,16 @@ def get_remaining_tokens_and_requests_from_request_data(data: dict) -> dict[str,
return headers
def _serialize_scan_metadata_header(entries: Iterable[object], *, max_length: int) -> str | None:
"""Compact JSON list of scan metadata entries, dropping trailing entries so the header fits in max_length."""
encoded: Final = tuple(json.dumps(entry, separators=(",", ":")) for entry in entries)
lengths: Final = tuple(accumulate(len(item) + 1 for item in encoded))
kept: Final = sum(1 for length in lengths if length + 1 <= max_length)
if kept == 0:
return None
return f"[{','.join(encoded[:kept])}]"
def get_logging_caching_headers(request_data: dict) -> dict | None:
_metadata: Final[dict] = {}
metadata_bucket: Final = request_data.get("metadata")
@ -468,6 +491,15 @@ def get_logging_caching_headers(request_data: dict) -> dict | None:
if scan_ids:
headers["x-litellm-guardrail-scan-id"] = ",".join(scan_ids)
scan_metadata: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY)
scan_metadata_header: Final = (
_serialize_scan_metadata_header(scan_metadata, max_length=MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH)
if isinstance(scan_metadata, (list, tuple))
else None
)
if scan_metadata_header:
headers["x-litellm-guardrail-scan-metadata"] = scan_metadata_header
if "applied_policies" in _metadata:
headers["x-litellm-applied-policies"] = ",".join(_metadata["applied_policies"])
@ -501,6 +533,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
"applied_policies",
"applied_guardrails",
GUARDRAIL_SCAN_IDS_METADATA_KEY,
GUARDRAIL_SCAN_METADATA_METADATA_KEY,
"policy_sources",
"guardrails",
"guardrail_config",
@ -565,21 +598,40 @@ def add_guardrail_to_applied_guardrails_header(request_data: dict, guardrail_nam
_metadata["applied_guardrails"] = [guardrail_name]
def add_guardrail_scan_id(request_data: dict, scan_id: str | None) -> None:
def add_guardrail_scan_id(
request_data: dict[str, object],
scan_id: str | None,
*,
guardrail_name: str | None,
provider: str,
stage: GuardrailEventHooks,
) -> None:
"""
Record a provider scan id so it can be surfaced to the caller.
Record a provider scan id, keyed to the guardrail execution that produced it, so it can be surfaced to the caller.
Guardrails only return scan details to the client when they block, so allowed requests carry no
audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header.
audit trail. Ids recorded here become the x-litellm-guardrail-scan-id response header, and the
(guardrail, stage, provider, scan_id) entries become the x-litellm-guardrail-scan-metadata header.
"""
if not scan_id:
return
_, _metadata = get_or_create_metadata_bucket(request_data)
existing: Final = _metadata.get(GUARDRAIL_SCAN_IDS_METADATA_KEY)
scan_ids: Final = tuple(existing) if isinstance(existing, (list, tuple)) else ()
scan_ids: Final[tuple[object, ...]] = tuple(existing) if isinstance(existing, (list, tuple)) else ()
if scan_id not in scan_ids:
_metadata[GUARDRAIL_SCAN_IDS_METADATA_KEY] = (*scan_ids, scan_id)
entry: Final[GuardrailScanMetadata] = {
"guardrail": guardrail_name,
"stage": stage.value,
"provider": provider,
"scan_id": scan_id,
}
existing_entries: Final = _metadata.get(GUARDRAIL_SCAN_METADATA_METADATA_KEY)
entries: Final[tuple[object, ...]] = tuple(existing_entries) if isinstance(existing_entries, (list, tuple)) else ()
if entry not in entries:
_metadata[GUARDRAIL_SCAN_METADATA_METADATA_KEY] = (*entries, entry)
def add_policy_to_applied_policies_header(request_data: dict, policy_name: str | None):
"""

107
litellm/proxy/db/db_url_settings.py
View file

@ -34,12 +34,20 @@ writer's connection params (pool size, timeouts, pgbouncer mode) for the
ones the reader URL does not pin itself.
"""
import _ssl
import hashlib
import os
import socket
import ssl
import struct
import sys
import tempfile
import urllib.parse
from collections.abc import Mapping
from collections.abc import Callable, Mapping, Sequence
from functools import partial
from pathlib import Path
from types import MappingProxyType
from typing import Annotated, Final, cast
from typing import Annotated, Final, Protocol, TypeAlias, cast
from pydantic import AliasChoices, BeforeValidator, Field
from pydantic_settings import BaseSettings, SettingsConfigDict
@ -126,21 +134,100 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float])
LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"})
PEM_CERT_HEADER: Final = b"-----BEGIN CERTIFICATE-----"
PG_SSL_REQUEST: Final = struct.pack("!ii", 8, 80877103)
TLS_PROBE_TIMEOUT_SECONDS: Final = 10.0
RootCertResolver: TypeAlias = Callable[[str, str, int], str] # mutable-ok: Callable parameter syntax
def translate_libpq_ssl_params(url: str) -> str:
class _VerifiedChainSource(Protocol):
def get_verified_chain(self) -> Sequence[_ssl.Certificate] | None: ...
def _verified_chain_der(tls: ssl.SSLSocket) -> tuple[bytes, ...]:
if sys.version_info >= (3, 13):
return tuple(tls.get_verified_chain())
legacy: Final = cast( # cast-ok: the stub omits _sslobj, the C object has get_verified_chain since 3.10
"_VerifiedChainSource | None",
tls._sslobj, # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # public API only from 3.13
)
chain: Final = () if legacy is None else legacy.get_verified_chain() or ()
return tuple(cert.public_bytes(_ssl.ENCODING_DER) for cert in chain)
def _server_trust_anchor(cafile: str, host: str, port: int) -> bytes | None:
try:
context: Final = ssl.create_default_context(cafile=cafile)
with socket.create_connection((host, port), timeout=TLS_PROBE_TIMEOUT_SECONDS) as raw:
raw.sendall(PG_SSL_REQUEST)
if raw.recv(1) != b"S":
return None
with context.wrap_socket(raw, server_hostname=host) as tls:
chain: Final = _verified_chain_der(tls)
except (OSError, ValueError):
return None
return chain[-1] if chain else None
def pin_bundle_root(cert_path: str, host: str, port: int) -> str:
"""Reduce a multi-root CA bundle to the one root that verifies ``host``.
Prisma's ``sslcert`` loads a single PEM certificate (native-tls
``Certificate::from_pem``), so pointing it at a bundle such as the AWS RDS
global bundle trusts only the first of its 108 regional roots and the
handshake fails with "unable to get local issuer certificate" for every
other region. A single-certificate file is returned as is. For a bundle,
one verifying handshake (chain and hostname, whole bundle as trust store)
identifies the trust anchor the server actually chains to, which is
written to a single-certificate file for Prisma. If the probe fails the
bundle path is returned unchanged, so Prisma fails closed exactly as
before rather than trusting anything the bundle would not.
"""
try:
if Path(cert_path).read_bytes().count(PEM_CERT_HEADER) < 2:
return cert_path
except OSError:
return cert_path
root: Final = _server_trust_anchor(cert_path, host, port)
if root is None:
return cert_path
pinned: Final = Path(tempfile.gettempdir()) / f"litellm-sslcert-{hashlib.sha256(root).hexdigest()[:16]}.pem"
return str(pinned) if _replace_file(pinned, ssl.DER_cert_to_PEM_cert(root)) else cert_path
def _replace_file(target: Path, content: str) -> bool:
"""Write ``content`` to a private temp file and rename it over ``target``, so
readers never see a partial file and a symlink planted at ``target`` is
replaced rather than followed."""
try:
fd, staged = tempfile.mkstemp(dir=target.parent, prefix=f"{target.name}.")
except OSError:
return False
try:
with os.fdopen(fd, "w") as handle:
handle.write(content)
os.replace(staged, target)
except OSError:
Path(staged).unlink(missing_ok=True)
return False
return True
def translate_libpq_ssl_params(url: str, resolve_root_cert: RootCertResolver = pin_bundle_root) -> str:
"""Rewrite libpq's certificate-verification params into Prisma's dialect.
Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert``
(the CA bundle) and ``sslaccept=strict``. It silently discards
(a single CA certificate) and ``sslaccept=strict``. It silently discards
``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to
``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no
certificate check at all. ``verify-ca`` and ``verify-full`` both become
``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes
``sslcert``, and either one turns on ``sslaccept=strict`` (chain and
hostname), matching libpq where a root cert makes ``require`` verify.
Prisma params the operator pinned themselves win; anything else is left
untouched.
``sslcert`` (run through ``resolve_root_cert``, which pins a multi-root
bundle down to the server's root), and either one turns on
``sslaccept=strict`` (chain and hostname), matching libpq where a root
cert makes ``require`` verify. Prisma params the operator pinned
themselves win; anything else is left untouched.
"""
parsed: Final = urllib.parse.urlsplit(url)
pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
@ -154,7 +241,9 @@ def translate_libpq_ssl_params(url: str) -> str:
if key != "sslrootcert"
)
root_cert: Final = tuple(
("sslcert", value) for key, value in pairs if key == "sslrootcert" and "sslcert" not in keys
("sslcert", resolve_root_cert(value, parsed.hostname or "", parsed.port or int(DEFAULT_POSTGRES_PORT)))
for key, value in pairs
if key == "sslrootcert" and "sslcert" not in keys
)
strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),)
query: Final = urllib.parse.urlencode(translated + root_cert + strict)

220
litellm/proxy/db/gateway_request_tracking.py
View file

@ -10,13 +10,28 @@ strings rather than passing the raw path through. Nothing a caller sends can
add a key, so the fold and the table it commits to are bounded by (days x
routes) however much traffic arrives, and the response path carries no
unbounded queue that would block once full.
A flush commits its whole snapshot as one multi-row ``INSERT ... ON CONFLICT DO
UPDATE`` rather than one upsert per key, so a worker costs the primary one
statement per interval however many routes it served. With
``use_redis_transaction_buffer`` on, workers instead push their snapshot to a
Redis list and one lock-holding pod folds every entry and writes the table, so
the deployment as a whole costs the primary one statement per interval.
"""
from dataclasses import asdict
import json
from collections.abc import AsyncIterator, Iterable
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory
from litellm.types.proxy.gateway_requests import (
GatewayRequestCounts,
@ -28,6 +43,15 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_EMPTY: Final = GatewayRequestCounts(successful_requests=0, failed_requests=0)
_TABLE: Final = '"LiteLLM_DailyGatewayRequests"'
_COLUMNS_PER_ROW: Final = 5
_UTC_NOW: Final = "(NOW() AT TIME ZONE 'UTC')"
GATEWAY_REQUESTS_JOB_NAME: Final = "update_gateway_requests_job"
_BufferedRows: TypeAlias = tuple[tuple[str, str, str, int, int], ...]
_BUFFERED_ROWS: Final = TypeAdapter(_BufferedRows)
_BUFFERED_ENTRIES: Final = TypeAdapter(tuple[str | bytes, ...])
_NO_COUNTS: Final[GatewayRequestSnapshot] = MappingProxyType({})
def _utc_date() -> str:
@ -59,20 +83,54 @@ class GatewayRequestAccumulator:
route) however long the database is unreachable.
This buys at-least-once, not exactly-once, and the cost is worth stating.
The batch commits inside its context manager's ``__aexit__``, so a failure
raised after the transaction committed (a connection dropped while reading
the acknowledgement) restores counts that are already persisted, and the
next flush increments them a second time. Exactly-once would need a dedup
key the upserts could ignore on replay. For a traffic-volume metric a rare
The statement commits on the server before its acknowledgement is read, so
a failure raised after the commit (a connection dropped while reading the
acknowledgement) restores counts that are already persisted, and the next
flush increments them a second time. Exactly-once would need a dedup key
the upsert could ignore on replay. For a traffic-volume metric a rare
overcount on a dropped acknowledgement beats losing a whole interval to
every database blip, so the trade is deliberate.
"""
for key, counts in snapshot.items():
existing = self._counts.get(key, _EMPTY)
self._counts[key] = GatewayRequestCounts(
successful_requests=existing.successful_requests + counts.successful_requests,
failed_requests=existing.failed_requests + counts.failed_requests,
)
self._counts = dict(fold_counts(chain(self._counts.items(), snapshot.items()))) # mutable-ok: fold replaced
def fold_counts(items: Iterable[tuple[GatewayRequestKey, GatewayRequestCounts]]) -> GatewayRequestSnapshot:
"""Sum counts key-wise; the result stays bounded by (date x category x route)."""
folded: Final[dict[GatewayRequestKey, GatewayRequestCounts]] = {} # mutable-ok: local fold returned once
for key, counts in items:
existing = folded.get(key, _EMPTY)
folded[key] = GatewayRequestCounts(
successful_requests=existing.successful_requests + counts.successful_requests,
failed_requests=existing.failed_requests + counts.failed_requests,
)
return folded
def build_gateway_requests_upsert(snapshot: GatewayRequestSnapshot) -> tuple[str, tuple[str | int, ...]]:
"""
One ``INSERT ... ON CONFLICT DO UPDATE`` that increments every (date, category,
route) in the snapshot. Rows are ordered by the conflict key so concurrent
writers lock rows in the same order and cannot deadlock.
"""
ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route))
rows: Final = ", ".join(
f"(${base + 1}::text, ${base + 2}::text, ${base + 3}::text, ${base + 4}::bigint, ${base + 5}::bigint, {_UTC_NOW})"
for base in range(0, len(ordered) * _COLUMNS_PER_ROW, _COLUMNS_PER_ROW)
)
sql: Final = (
f'INSERT INTO {_TABLE} ("date", "category", "route", "successful_requests", "failed_requests", "updated_at")\n'
f"VALUES {rows}\n"
'ON CONFLICT ("date", "category", "route") DO UPDATE SET\n'
f' "successful_requests" = {_TABLE}."successful_requests" + EXCLUDED."successful_requests",\n'
f' "failed_requests" = {_TABLE}."failed_requests" + EXCLUDED."failed_requests",\n'
f' "updated_at" = {_UTC_NOW}'
)
params: Final[tuple[str | int, ...]] = tuple(
value
for key, counts in ordered
for value in (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests)
)
return sql, params
async def commit_gateway_requests_to_db(
@ -80,50 +138,130 @@ async def commit_gateway_requests_to_db(
prisma_client: "PrismaClient",
snapshot: GatewayRequestSnapshot,
) -> None:
"""Upsert one incrementing row per (date, category, route)."""
"""Increment every (date, category, route) in the snapshot with a single statement."""
if not snapshot:
return
ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route))
sql, params = build_gateway_requests_upsert(snapshot)
await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client
# pyright: ignore[reportAny] on both lines -- prisma's generated client is untyped,
# so .db and every table action off it resolve to Any at this boundary. The dict
# literals below are the shape prisma's generated inputs require.
async with prisma_client.db.batch_() as batcher: # pyright: ignore[reportAny] # untyped prisma client
for key, counts in ordered:
columns = asdict(key)
batcher.litellm_dailygatewayrequests.upsert( # pyright: ignore[reportAny] # untyped prisma client
where={"date_category_route": columns}, # mutable-ok: prisma input is dict-shaped
data={ # mutable-ok: prisma input is dict-shaped
"create": { # mutable-ok: prisma input is dict-shaped
**columns,
"successful_requests": counts.successful_requests,
"failed_requests": counts.failed_requests,
},
"update": { # mutable-ok: prisma input is dict-shaped
"successful_requests": {"increment": counts.successful_requests}, # mutable-ok: as above
"failed_requests": {"increment": counts.failed_requests}, # mutable-ok: as above
},
},
verbose_proxy_logger.debug(
"Gateway request tracking - committed %d aggregated rows in one statement", len(snapshot)
)
class GatewayRequestRedisBuffer:
"""
Folds every worker's snapshot through one Redis list so a single pod per
interval writes the table, mirroring the spend writer's transaction buffer.
Each entry is one worker's snapshot as JSON rows; the lock holder pops them,
sums them, and commits one statement. A commit failure pushes the summed
rows back so the next holder retries, keeping the at-least-once guarantee.
If that push fails too, the rows go back to the holder's own accumulator so
they ride along with its next flush instead of vanishing with the pop.
"""
def __init__(self, *, redis_cache: RedisCache, pod_lock_manager: PodLockManager) -> None:
self._redis_cache: Final = redis_cache
self._pod_lock_manager: Final = pod_lock_manager
async def push(self, snapshot: GatewayRequestSnapshot) -> None:
if not snapshot:
return
rows: Final[_BufferedRows] = tuple(
(key.date, key.category, key.route, counts.successful_requests, counts.failed_requests)
for key, counts in snapshot.items()
)
await self._redis_cache.async_rpush(key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, values=(json.dumps(rows),))
async def _pop_batch(self) -> tuple[str | bytes, ...]:
popped: Final[object] = await self._redis_cache.async_lpop( # pyright: ignore[reportAny] # redis returns Any
key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT
)
if not popped:
return ()
return _BUFFERED_ENTRIES.validate_python(popped if isinstance(popped, list) else (popped,))
async def _pop_all(self) -> AsyncIterator[str | bytes]:
while True:
batch = await self._pop_batch()
for entry in batch:
yield entry
if len(batch) < MAX_REDIS_BUFFER_DEQUEUE_COUNT:
return
async def pop(self) -> GatewayRequestSnapshot:
entries: Final = tuple([entry async for entry in self._pop_all()])
return fold_counts(
(
GatewayRequestKey(date=date, category=category, route=route),
GatewayRequestCounts(successful_requests=succeeded, failed_requests=failed),
)
for entry in entries
for date, category, route, succeeded, failed in _BUFFERED_ROWS.validate_json(entry)
)
verbose_proxy_logger.debug("Gateway request tracking - committed %d aggregated rows", len(ordered))
async def commit_if_leader(self, prisma_client: "PrismaClient") -> GatewayRequestSnapshot:
"""
Drain the list and write it as one statement, but only on the pod holding the job lock.
The lock is a lease, never released: the holder re-enters it on every flush and
keeps committing alone until the TTL lapses, so the primary sees one statement
per flush interval deployment-wide instead of one per worker.
Returns the popped rows that could be neither committed nor re-queued, for the
caller to keep in memory. Empty on success.
"""
if not await self._pod_lock_manager.acquire_lock(cronjob_id=GATEWAY_REQUESTS_JOB_NAME):
return _NO_COUNTS
buffered: Final = await self.pop()
try:
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=buffered)
except Exception: # noqa: BLE001 -- a failed commit must not stop the scheduler
verbose_proxy_logger.warning(
"Gateway request tracking - failed to commit %d buffered rows, re-queuing to Redis for the next flush",
len(buffered),
exc_info=True,
)
return await self._requeue(buffered)
return _NO_COUNTS
async def _requeue(self, snapshot: GatewayRequestSnapshot) -> GatewayRequestSnapshot:
try:
await self.push(snapshot)
except Exception: # noqa: BLE001 -- the rows go back to the caller's accumulator instead
verbose_proxy_logger.warning(
"Gateway request tracking - Redis re-queue failed, keeping %d rows in memory for the next flush",
len(snapshot),
exc_info=True,
)
return snapshot
return _NO_COUNTS
async def flush_gateway_requests(
prisma_client: "PrismaClient",
accumulator: GatewayRequestAccumulator,
redis_buffer: GatewayRequestRedisBuffer | None = None,
) -> None:
"""
Scheduler entrypoint. Never raises: a metering failure must not kill the job.
With ``redis_buffer`` the snapshot goes to Redis and only the lease holder
writes to Postgres. Shutdown passes no buffer so a departing worker writes its
own counts directly instead of parking them behind a lease it may not hold.
``CancelledError`` is deliberately not caught, so a flush cancelled during
shutdown drops its snapshot rather than restoring counts onto an accumulator
the process is about to discard.
"""
snapshot: Final = accumulator.drain()
try:
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot)
if redis_buffer is None:
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot)
else:
await redis_buffer.push(snapshot)
except Exception: # noqa: BLE001 -- a failed flush must not stop the scheduler
accumulator.restore(snapshot)
verbose_proxy_logger.warning(
@ -131,3 +269,13 @@ async def flush_gateway_requests(
len(snapshot),
exc_info=True,
)
return
if redis_buffer is None:
return
try:
accumulator.restore(await redis_buffer.commit_if_leader(prisma_client))
except Exception: # noqa: BLE001 -- entries still in Redis are drained by the next flush
verbose_proxy_logger.warning(
"Gateway request tracking - leader drain failed, buffered rows stay in Redis for the next flush",
exc_info=True,
)

25
litellm/proxy/guardrails/auto_router_compression.py
View file

@ -80,17 +80,19 @@ def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRout
def policy_for_model(
llm_router: "Router | None",
model_alias: str,
team_id: str | None,
request_kwargs: Mapping[str, object],
request_tags: Sequence[str],
) -> AutoRouterCompressionPolicy | None:
"""The compression policy of the auto router marker `model_alias` resolves to.
"""The compression policy of the auto router marker `model_alias` resolves to for this caller.
Pre-call arming and the routing hook both resolve through here, so an alias with
several tag-scoped markers cannot suppress under one and then route under another.
Pre-call arming and the routing hook both resolve through here, and here resolves through the
router's own request-scoped deployment lookup, so an alias with several tag-scoped markers
cannot suppress under one and then route under another, and a team router reached by its
public name carries its policy for every principal that can reach it.
"""
if llm_router is None:
return None
deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or ()
deployments: Final = llm_router.deployments_for_request(model_alias, request_kwargs)
markers: Final = tuple(
litellm_params
for deployment in deployments
@ -108,17 +110,6 @@ def policy_for_model(
return next((policy for policy in candidates if policy is not None), None)
def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None:
"""The caller's team id, from whichever metadata bucket this surface writes to."""
for meta_key in ("metadata", "litellm_metadata"):
meta = request_kwargs.get(meta_key)
if isinstance(meta, Mapping):
team_id = meta.get("user_api_key_team_id")
if isinstance(team_id, str):
return team_id
return None
def _compression_guardrail_classes() -> tuple[type, ...]:
"""The registered guardrail classes whose provider compresses prompts."""
from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry
@ -172,7 +163,7 @@ async def arm_pre_call(
policy: Final = policy_for_model(
llm_router=llm_router,
model_alias=model_alias,
team_id=team_id_from_request(data),
request_kwargs=data,
request_tags=_get_tags_from_request_kwargs(data),
)
if policy is None:

10
litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py
View file

@ -17,7 +17,8 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.proxy.common_utils.callback_utils import add_guardrail_scan_id
from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations
from litellm.types.utils import (
GenericGuardrailAPIInputs,
GuardrailStatus,
@ -218,6 +219,13 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
metadata: Final = request_data.get("metadata") or {}
request_data["metadata"] = metadata
metadata["_openai_moderation_response"] = moderation_response.model_dump()
add_guardrail_scan_id(
request_data=request_data,
scan_id=moderation_response.id,
guardrail_name=self.guardrail_name,
provider=SupportedGuardrailIntegrations.OPENAI_MODERATION.value,
stage=GuardrailEventHooks.post_call if input_type == "response" else GuardrailEventHooks.pre_call,
)
# Check if content is flagged and raise exception if needed
self._check_moderation_result(moderation_response)

34
litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
View file

@ -721,10 +721,18 @@ class PanwPrismaAirsHandler(CustomGuardrail):
}
}
def _record_scan_id(self, request_data: dict[str, object], scan_result: Mapping[str, object]) -> None:
def _record_scan_id(
self, request_data: dict[str, object], scan_result: Mapping[str, object], stage: GuardrailEventHooks
) -> None:
"""Surface the AIRS scan id on the response, so allowed calls are auditable too."""
scan_id: Final = scan_result.get("scan_id")
add_guardrail_scan_id(request_data=request_data, scan_id=str(scan_id) if scan_id else None)
add_guardrail_scan_id(
request_data=request_data,
scan_id=str(scan_id) if scan_id else None,
guardrail_name=self.guardrail_name,
provider=self._PROVIDER_NAME,
stage=stage,
)
def _handle_api_error_with_logging(
self,
@ -948,7 +956,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
event_type=GuardrailEventHooks.post_call,
)
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
self._record_scan_id(request_data, scan_result)
self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call)
def _check_and_mark_scanned(self, data: dict, scan_type: str) -> bool:
"""
@ -1078,7 +1086,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
duration=(end_time - start_time).total_seconds(),
event_type=GuardrailEventHooks.pre_call,
)
self._record_scan_id(data, scan_result)
self._record_scan_id(data, scan_result, GuardrailEventHooks.pre_call)
action: Final = scan_result.get("action", "block")
category: Final = scan_result.get("category", "unknown")
@ -1199,7 +1207,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
duration=(end_time - start_time).total_seconds(),
event_type=GuardrailEventHooks.post_call,
)
self._record_scan_id(data, scan_result)
self._record_scan_id(data, scan_result, GuardrailEventHooks.post_call)
action: Final = scan_result.get("action", "block")
category: Final = scan_result.get("category", "unknown")
@ -1401,7 +1409,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
duration=(end_time - start_time).total_seconds(),
event_type=GuardrailEventHooks.post_call,
)
self._record_scan_id(request_data, scan_result)
self._record_scan_id(request_data, scan_result, GuardrailEventHooks.post_call)
# Add guardrail to applied guardrails header for observability
add_guardrail_to_applied_guardrails_header(
@ -1475,7 +1483,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
)
continue
self._record_scan_id(request_data, scan_result)
self._record_scan_id(
request_data,
scan_result,
GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call,
)
action = scan_result.get("action", "block")
masked_args = self._masked_tool_call_arguments(
@ -1829,7 +1841,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
new_texts.append(text)
continue
self._record_scan_id(request_data, scan_result)
self._record_scan_id(
request_data,
scan_result,
GuardrailEventHooks.post_call if is_response else GuardrailEventHooks.pre_call,
)
action = scan_result.get("action", "block")
masked_text = self._get_masked_text(scan_result, is_response=is_response)
@ -1901,7 +1917,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
)
# If we reach here, fallback_on_error="allow"
else:
self._record_scan_id(request_data, mcp_scan_result)
self._record_scan_id(request_data, mcp_scan_result, GuardrailEventHooks.pre_call)
action = mcp_scan_result.get("action", "block")
masked_text = self._get_masked_text(mcp_scan_result, is_response=False)
if action == "allow":

10
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -1582,6 +1582,13 @@ async def _show_no_redis_warning() -> bool:
return await count_live_proxy_workers(prisma_client) != 1
def _show_env_credential_login_warning() -> bool:
from litellm.proxy.auth.login_utils import is_env_credential_login_enabled
from litellm.proxy.proxy_server import general_settings
return is_env_credential_login_enabled(general_settings)
async def _get_health_readiness_details(
response: Response | None = None,
) -> dict[str, Any]:
@ -1623,6 +1630,7 @@ async def _get_health_readiness_details(
log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel())
is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG)
show_no_redis_warning: Final = await _show_no_redis_warning()
show_env_credential_login_warning: Final = _show_env_credential_login_warning()
# check DB
if prisma_client is not None: # if db passed in, check if it's connected
@ -1650,6 +1658,7 @@ async def _get_health_readiness_details(
"log_level": log_level_name,
"is_detailed_debug": is_detailed_debug,
"show_no_redis_warning": show_no_redis_warning,
"show_env_credential_login_warning": show_env_credential_login_warning,
}
else:
return {
@ -1662,6 +1671,7 @@ async def _get_health_readiness_details(
"log_level": log_level_name,
"is_detailed_debug": is_detailed_debug,
"show_no_redis_warning": show_no_redis_warning,
"show_env_credential_login_warning": show_env_credential_login_warning,
}
except Exception as e:
raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e})")

18
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -32,6 +32,10 @@ from litellm.proxy.hooks.rate_limiter_utils import (
resolve_llm_provider_for_rate_limit,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.router_utils.add_retry_fallback_headers import (
ensure_response_additional_headers,
response_has_hidden_params,
)
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
@ -659,22 +663,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
data=data, user_api_key_dict=user_api_key_dict, response=response
)
# Add additional priority-specific headers
if isinstance(response, ModelResponse):
if response_has_hidden_params(response):
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)
# Get existing additional headers
additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {}
# Add priority information
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers["x-litellm-priority"] = priority or "default"
additional_headers["x-litellm-rate-limiter-version"] = "v3"
# Update response
if not hasattr(response, "_hidden_params"):
response._hidden_params = {}
response._hidden_params["additional_headers"] = additional_headers
return response
except Exception as e:

33
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -52,6 +52,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
canonical_provider_batch_id,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.router_utils.add_retry_fallback_headers import (
ensure_response_additional_headers,
response_has_hidden_params,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
from litellm.types.utils import (
@ -4677,34 +4681,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
Post-call hook to update rate limit headers in the response.
"""
try:
from pydantic import BaseModel
stash: Final = get_request_stash()
litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None
if litellm_proxy_rate_limit_response is not None:
# Update response headers
if hasattr(response, "_hidden_params"):
_hidden_params = getattr(response, "_hidden_params")
else:
_hidden_params = None
if _hidden_params is not None and (
isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict)
):
if isinstance(_hidden_params, BaseModel):
_hidden_params = _hidden_params.model_dump()
_additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers(
additional_headers=_hidden_params.get("additional_headers", {}) or {},
if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response):
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers.update(
self._merge_ratelimit_statuses_into_additional_headers(
additional_headers={},
statuses=litellm_proxy_rate_limit_response["statuses"],
)
setattr(
response,
"_hidden_params",
{**_hidden_params, "additional_headers": _additional_headers},
)
)
except Exception as e:
verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e)

16
litellm/proxy/litellm_pre_call_utils.py
View file

@ -235,6 +235,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
"applied_policies",
"policy_sources",
"guardrail_scan_ids",
"guardrail_scan_metadata",
"routing_decision",
GATEWAY_INJECTED_CACHE_METADATA_KEY,
"pillar_response_headers",
@ -291,6 +292,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
"applied_policies",
"policy_sources",
"guardrail_scan_ids",
"guardrail_scan_metadata",
"routing_decision",
GATEWAY_INJECTED_CACHE_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
@ -771,6 +773,16 @@ def apply_missing_session_id_policy(
return
if policy == "omit":
metadata[SESSION_ID_OMITTED_METADATA_KEY] = True
requester_metadata: Final = data.get("metadata")
requester_session_id: Final = (
requester_metadata.get("session_id") if isinstance(requester_metadata, dict) else None
)
if (
(body_session_id := data.get("litellm_session_id"))
and not metadata.get("session_id")
and not requester_session_id
):
metadata["session_id"] = body_session_id
return
if data.get("litellm_session_id") or metadata.get("session_id"):
return
@ -1748,7 +1760,9 @@ class LiteLLMProxyRequestSetup:
callback_vars_dict.pop("success_callback", None)
callback_vars_dict.pop("failure_callback", None)
callback_vars_dict = {
key: (litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else value)
key: (
litellm.utils.get_secret(value, default_value=value) or value if isinstance(value, str) else str(value)
)
for key, value in callback_vars_dict.items()
}

258
litellm/proxy/management_endpoints/cost_tracking_settings.py
View file

@ -18,6 +18,7 @@ from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
import litellm
from litellm._internal_context import current_billing_time, pinned_billing_time
from litellm._logging import verbose_proxy_logger
from litellm.cost_calculator import completion_cost
from litellm.proxy._types import (
@ -27,7 +28,15 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.utils import CostPerToken, LlmProvidersSet, ModelInfo
from litellm.types.utils import (
CostBreakdown,
CostPerToken,
LlmProvidersSet,
ModelInfo,
ModelResponse,
PromptTokensDetailsWrapper,
Usage,
)
router: Final = APIRouter()
@ -46,13 +55,15 @@ def _configured_price(key: str, sources: tuple[Mapping[str, object], ...]) -> fl
def _extract_custom_pricing(
litellm_params: Mapping[str, object], model_info: Mapping[str, object]
litellm_params: Mapping[str, object], model_info: Mapping[str, object], builtin: ModelInfo | None
) -> CostPerToken | None:
"""
Pull per-token pricing configured on a deployment so on-prem / self-hosted
models (absent from the public cost map) still estimate a real cost.
Pricing may live on ``litellm_params`` or ``model_info``; ``litellm_params``
wins, matching the router's cost-map registration precedence.
wins, matching the router's cost-map registration precedence. Cache rates the
deployment leaves unset come from the backend model's built-in entry, then its
own input rate, again matching what the router registers for live billing.
"""
sources: Final = (litellm_params, model_info)
input_price: Final = _configured_price("input_cost_per_token", sources)
@ -61,15 +72,21 @@ def _extract_custom_pricing(
if input_price is None and output_price is None:
return None
input_rate: Final = input_price or 0.0
cache_sources: Final = sources if builtin is None else (*sources, builtin)
cache_read_price: Final = _configured_price("cache_read_input_token_cost", cache_sources)
cache_creation_price: Final = _configured_price("cache_creation_input_token_cost", cache_sources)
return CostPerToken(
input_cost_per_token=input_price or 0.0,
input_cost_per_token=input_rate,
output_cost_per_token=output_price or 0.0,
cache_read_input_token_cost=input_rate if cache_read_price is None else cache_read_price,
cache_creation_input_token_cost=input_rate if cache_creation_price is None else cache_creation_price,
)
def _lookup_model_info(model: str) -> ModelInfo | None:
def _lookup_model_info(model: str, custom_llm_provider: str | None = None) -> ModelInfo | None:
try:
return litellm.get_model_info(model=model)
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception:
return None
@ -98,17 +115,14 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel:
model_info: Final = first_deployment.get("model_info", {})
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
provider: Final = str(custom_llm_provider) if custom_llm_provider is not None else None
custom_cost_per_token: Final = _extract_custom_pricing(litellm_params, model_info)
# Check base_model first (needed for Azure custom deployment names)
# base_model wins (needed for Azure custom deployment names)
base_model: Final = model_info.get("base_model") or litellm_params.get("base_model")
if base_model:
verbose_proxy_logger.debug("Resolved model '%s' to base_model '%s' from router", model, base_model)
return ResolvedCostModel(str(base_model), provider, custom_cost_per_token)
resolved_model: Final = litellm_params.get("model")
resolved_model: Final = base_model or litellm_params.get("model")
if resolved_model:
verbose_proxy_logger.debug("Resolved model '%s' to '%s' from router", model, resolved_model)
custom_cost_per_token: Final = _extract_custom_pricing(
litellm_params, model_info, _lookup_model_info(str(resolved_model), provider)
)
return ResolvedCostModel(str(resolved_model), provider, custom_cost_per_token)
except Exception as e:
verbose_proxy_logger.debug("Could not resolve model '%s' from router: %s", model, e)
@ -117,19 +131,59 @@ def _resolve_model_for_cost_lookup(model: str) -> ResolvedCostModel:
return ResolvedCostModel(model, None, None)
def _calculate_period_costs(num_requests, cost_per_request, input_cost, output_cost, margin_cost):
"""
Calculate costs for a given number of requests.
@dataclass(frozen=True, slots=True)
class CostLines:
"""Cost of one request split the way the spend logs split it: the cache lines are
shares of input_cost and the reasoning line is a share of output_cost."""
Returns tuple of (total_cost, input_cost, output_cost, margin_cost) or all None if num_requests is None/0.
"""
if not num_requests:
return None, None, None, None
return (
cost_per_request * num_requests,
input_cost * num_requests,
output_cost * num_requests,
margin_cost * num_requests,
total_cost: float
input_cost: float
output_cost: float
margin_cost: float
cache_read_cost: float
cache_creation_cost: float
reasoning_cost: float
def times(self, num_requests: int | None) -> "CostLines | None":
if not num_requests:
return None
return CostLines(
total_cost=self.total_cost * num_requests,
input_cost=self.input_cost * num_requests,
output_cost=self.output_cost * num_requests,
margin_cost=self.margin_cost * num_requests,
cache_read_cost=self.cache_read_cost * num_requests,
cache_creation_cost=self.cache_creation_cost * num_requests,
reasoning_cost=self.reasoning_cost * num_requests,
)
def _cost_lines(cost_per_request: float, cost_breakdown: CostBreakdown | None) -> CostLines:
breakdown: Final = cost_breakdown if cost_breakdown is not None else CostBreakdown()
return CostLines(
total_cost=cost_per_request,
input_cost=breakdown.get("input_cost", 0.0),
output_cost=breakdown.get("output_cost", 0.0),
margin_cost=breakdown.get("margin_total_amount", 0.0),
cache_read_cost=breakdown.get("cache_read_cost", 0.0),
cache_creation_cost=breakdown.get("cache_creation_cost", 0.0),
reasoning_cost=breakdown.get("reasoning_cost", 0.0),
)
def _usage_for_estimate(request: CostEstimateRequest) -> Usage:
cache_tokens: Final = request.cache_read_input_tokens + request.cache_creation_input_tokens
return Usage(
prompt_tokens=request.input_tokens,
completion_tokens=request.output_tokens,
total_tokens=request.input_tokens + request.output_tokens,
reasoning_tokens=request.reasoning_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=request.cache_read_input_tokens,
cache_creation_tokens=request.cache_creation_input_tokens,
)
if cache_tokens
else None,
)
@ -530,11 +584,14 @@ async def estimate_cost(
- model: Model name (e.g., "gpt-4", "claude-3-opus")
- input_tokens: Expected input tokens per request
- output_tokens: Expected output tokens per request
- cache_read_input_tokens: Cache-read tokens per request, counted within input_tokens (optional)
- cache_creation_input_tokens: Cache-write tokens per request, counted within input_tokens (optional)
- reasoning_tokens: Reasoning tokens per request, counted within output_tokens (optional)
- num_requests_per_day: Number of requests per day (optional)
- num_requests_per_month: Number of requests per month (optional)
Returns cost breakdown including:
- Per-request costs (input, output, margin)
- Per-request costs (input, output, margin, plus the cache-read, cache-write and reasoning shares)
- Daily costs (if num_requests_per_day provided)
- Monthly costs (if num_requests_per_month provided)
@ -543,14 +600,15 @@ async def estimate_cost(
{
"model": "gpt-4",
"input_tokens": 1000,
"cache_read_input_tokens": 800,
"output_tokens": 500,
"reasoning_tokens": 200,
"num_requests_per_day": 100,
"num_requests_per_month": 3000
}
```
"""
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import ModelResponse, Usage
# Resolve model name (handles router aliases like 'e-model-router' -> 'azure_ai/gpt-4')
resolved: Final = _resolve_model_for_cost_lookup(request.model)
@ -559,15 +617,8 @@ async def estimate_cost(
verbose_proxy_logger.debug("Cost estimate: request.model='%s' resolved to '%s'", request.model, resolved_model)
# Create a mock response with usage for completion_cost
mock_response: Final = ModelResponse(
model=resolved_model,
usage=Usage(
prompt_tokens=request.input_tokens,
completion_tokens=request.output_tokens,
total_tokens=request.input_tokens + request.output_tokens,
),
)
usage: Final = _usage_for_estimate(request)
mock_response: Final = ModelResponse(model=resolved_model, usage=usage)
# Create a logging object to capture cost breakdown
litellm_logging_obj: Final = LiteLLMLoggingObj(
@ -580,92 +631,73 @@ async def estimate_cost(
function_id="cost-estimate",
)
# Use completion_cost which handles all the logic including margins/discounts
try:
cost_per_request: Final = completion_cost(
completion_response=mock_response,
model=resolved_model,
custom_llm_provider=resolved_provider,
custom_cost_per_token=resolved.custom_cost_per_token,
litellm_logging_obj=litellm_logging_obj,
)
except Exception as e:
raise HTTPException(
status_code=404,
detail={
"error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}"
},
)
# Pinning one moment keeps an off-peak window that opens mid-quote from pricing the totals on
# one side of it and the reported rates on the other.
with pinned_billing_time(current_billing_time()):
# Use completion_cost which handles all the logic including margins/discounts
try:
cost_per_request: Final = completion_cost(
completion_response=mock_response,
model=resolved_model,
custom_llm_provider=resolved_provider,
custom_cost_per_token=resolved.custom_cost_per_token,
litellm_logging_obj=litellm_logging_obj,
)
except Exception as e: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriceable model
raise HTTPException(
status_code=404,
detail={
"error": f"Could not calculate cost for model '{request.model}' (resolved to '{resolved_model}'): {e}"
},
)
# Get cost breakdown from the logging object
cost_breakdown: Final = litellm_logging_obj.cost_breakdown
# The rates come back from the pricing call itself rather than a second lookup, so they are the
# ones the cost lines above billed at even when completion_cost infers a provider this endpoint
# never resolved (an unrouted "xai/grok-4" prices on xai's inclusive tier thresholds; a lookup
# here without that provider would report the sub-200k rate for a line billed above it).
rates: Final = litellm_logging_obj.billed_token_rates
per_request: Final = _cost_lines(cost_per_request, litellm_logging_obj.cost_breakdown)
daily: Final = per_request.times(request.num_requests_per_day)
monthly: Final = per_request.times(request.num_requests_per_month)
input_cost: Final = cost_breakdown.get("input_cost", 0.0) if cost_breakdown else 0.0
output_cost: Final = cost_breakdown.get("output_cost", 0.0) if cost_breakdown else 0.0
margin_cost: Final = cost_breakdown.get("margin_total_amount", 0.0) if cost_breakdown else 0.0
model_info: Final = _lookup_model_info(resolved_model)
mapped_input_price: Final = model_info.get("input_cost_per_token") if model_info is not None else None
mapped_output_price: Final = model_info.get("output_cost_per_token") if model_info is not None else None
model_info: Final = _lookup_model_info(resolved_model, resolved_provider)
mapped_provider: Final = model_info.get("litellm_provider") if model_info is not None else None
input_cost_per_token: Final = (
resolved.custom_cost_per_token["input_cost_per_token"]
if resolved.custom_cost_per_token is not None
else mapped_input_price
)
output_cost_per_token: Final = (
resolved.custom_cost_per_token["output_cost_per_token"]
if resolved.custom_cost_per_token is not None
else mapped_output_price
)
custom_llm_provider: Final = mapped_provider if mapped_provider is not None else resolved_provider
# Calculate daily and monthly costs
(
daily_cost,
daily_input_cost,
daily_output_cost,
daily_margin_cost,
) = _calculate_period_costs(
num_requests=request.num_requests_per_day,
cost_per_request=cost_per_request,
input_cost=input_cost,
output_cost=output_cost,
margin_cost=margin_cost,
)
(
monthly_cost,
monthly_input_cost,
monthly_output_cost,
monthly_margin_cost,
) = _calculate_period_costs(
num_requests=request.num_requests_per_month,
cost_per_request=cost_per_request,
input_cost=input_cost,
output_cost=output_cost,
margin_cost=margin_cost,
)
return CostEstimateResponse(
model=request.model,
input_tokens=request.input_tokens,
output_tokens=request.output_tokens,
cache_read_input_tokens=request.cache_read_input_tokens,
cache_creation_input_tokens=request.cache_creation_input_tokens,
reasoning_tokens=request.reasoning_tokens,
num_requests_per_day=request.num_requests_per_day,
num_requests_per_month=request.num_requests_per_month,
cost_per_request=cost_per_request,
input_cost_per_request=input_cost,
output_cost_per_request=output_cost,
margin_cost_per_request=margin_cost,
daily_cost=daily_cost,
daily_input_cost=daily_input_cost,
daily_output_cost=daily_output_cost,
daily_margin_cost=daily_margin_cost,
monthly_cost=monthly_cost,
monthly_input_cost=monthly_input_cost,
monthly_output_cost=monthly_output_cost,
monthly_margin_cost=monthly_margin_cost,
input_cost_per_token=input_cost_per_token,
output_cost_per_token=output_cost_per_token,
cost_per_request=per_request.total_cost,
input_cost_per_request=per_request.input_cost,
output_cost_per_request=per_request.output_cost,
margin_cost_per_request=per_request.margin_cost,
cache_read_cost_per_request=per_request.cache_read_cost,
cache_creation_cost_per_request=per_request.cache_creation_cost,
reasoning_cost_per_request=per_request.reasoning_cost,
daily_cost=daily.total_cost if daily is not None else None,
daily_input_cost=daily.input_cost if daily is not None else None,
daily_output_cost=daily.output_cost if daily is not None else None,
daily_margin_cost=daily.margin_cost if daily is not None else None,
daily_cache_read_cost=daily.cache_read_cost if daily is not None else None,
daily_cache_creation_cost=daily.cache_creation_cost if daily is not None else None,
daily_reasoning_cost=daily.reasoning_cost if daily is not None else None,
monthly_cost=monthly.total_cost if monthly is not None else None,
monthly_input_cost=monthly.input_cost if monthly is not None else None,
monthly_output_cost=monthly.output_cost if monthly is not None else None,
monthly_margin_cost=monthly.margin_cost if monthly is not None else None,
monthly_cache_read_cost=monthly.cache_read_cost if monthly is not None else None,
monthly_cache_creation_cost=monthly.cache_creation_cost if monthly is not None else None,
monthly_reasoning_cost=monthly.reasoning_cost if monthly is not None else None,
input_cost_per_token=rates.input_cost_per_token if rates is not None else None,
output_cost_per_token=rates.output_cost_per_token if rates is not None else None,
cache_read_input_token_cost=rates.cache_read_input_token_cost if rates is not None else None,
cache_creation_input_token_cost=rates.cache_creation_input_token_cost if rates is not None else None,
output_cost_per_reasoning_token=rates.output_cost_per_reasoning_token if rates is not None else None,
provider=custom_llm_provider,
)

19
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -4559,6 +4559,23 @@ async def delete_verification_tokens(
litellm_changed_by=litellm_changed_by,
)
# Snapshot before the delete: the FK cascade drops the mapping rows, but their
# cached jwt_key_mapping entries still resolve to the now-dead token (LIT-5380).
jwt_mapping_cache_keys: Final[tuple[str, ...]] = tuple(
cache_key
for keys_for_token in await asyncio.gather(
*(
get_jwt_key_mapping_cache_keys_for_token(
hashed_token=key.token,
prisma_client=prisma_client,
)
for key in authorized_keys
if key.token is not None
)
)
for cache_key in keys_for_token
)
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
deleted_tokens = await prisma_client.delete_data(tokens=tokens)
if deleted_tokens is not None and len(deleted_tokens) != len(tokens):
@ -4571,6 +4588,8 @@ async def delete_verification_tokens(
if len(deleted_tokens) != len(tokens):
failed_tokens = [token for token in tokens if token not in deleted_tokens]
await evict_and_broadcast(cache_keys=jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache)
else:
raise Exception("DB not connected. prisma_client is None")
except Exception as e:

4
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -2673,6 +2673,8 @@ if MCP_AVAILABLE:
"""
Updates the MCP Server in the db.
Partial update: a field left out of the payload keeps its stored value, and a field sent as null is cleared.
Parameters:
- payload: UpdateMCPServerRequest - Required. The updated mcp server data.
```
@ -3098,6 +3100,8 @@ if MCP_AVAILABLE:
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
litellm_changed_by: str | None = Header(None),
):
"""Partial update: a field left out keeps its stored value, and a field sent as null is cleared, except
``toolset_name`` and ``tools``, which a toolset always has; empty the tool selection with an explicit []."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
raise HTTPException(

4
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -5601,7 +5601,7 @@ async def team_model_add(
updated_team: Final = await _team_db(prisma_client).update(
where={"team_id": data.team_id},
data={"updated_at": datetime.now(timezone.utc)},
include={"object_permission": True},
include={"litellm_model_table": True, "object_permission": True},
)
if updated_team is None:
raise HTTPException(
@ -5688,7 +5688,7 @@ async def team_model_delete(
updated_team: Final = await _team_db(prisma_client).update(
where={"team_id": data.team_id},
data={"models": updated_models},
include={"object_permission": True},
include={"litellm_model_table": True, "object_permission": True},
)
if updated_team is None:
raise HTTPException(

23
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -22,7 +22,6 @@ from html import escape
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
Annotated,
Any,
Final,
Literal,
@ -42,7 +41,7 @@ if TYPE_CHECKING:
import jwt
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
from fastapi.responses import RedirectResponse
from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -95,6 +94,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.team_grants import TeamModelAliasTable
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.admin_ui_utils import (
admin_ui_disabled,
@ -209,31 +209,14 @@ def _team_detail_db(repo: TeamRepository) -> "TableActions[_TeamDetailRow]":
return repo.table
_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str])
_SSO_TOKEN_CLAIMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _decode_model_aliases(value: object) -> object:
"""``/team/new`` stores team model aliases as a JSON-encoded string in the Json column."""
if not isinstance(value, str):
return value
try:
return _MODEL_ALIASES_ADAPTER.validate_json(value)
except ValidationError:
return None
class _TeamModelAliasTable(BaseModel):
model_config = ConfigDict(protected_namespaces=())
model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None
class _TeamRowGrants(BaseModel):
team_id: str
team_alias: str | None = None
models: tuple[str, ...] = ()
litellm_model_table: _TeamModelAliasTable | None = None
litellm_model_table: TeamModelAliasTable | None = None
class CliSsoTeamDetail(BaseModel):

168
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -70,6 +70,10 @@ def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
return None if texts is None else tuple(texts)
def _scanned_texts(texts: Sequence[str] | None) -> tuple[str, ...]:
return tuple(texts or ())
def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None:
return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls)
@ -78,6 +82,10 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non
return sent is not None and returned is not None and returned != sent
def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool:
return sent is not None and returned is not None and len(returned) != len(sent)
_GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object])
@ -89,10 +97,11 @@ def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT:
class _StreamRewriteObserver(CustomGuardrail):
"""Stand-in handed to the endpoint translation in place of a streaming pipeline step's
guardrail. It records whether the guardrail returned different output than it was given,
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text
rewrites are deliverable on translations that write them back across the buffered chunks
(``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any
other translation are discarded by the executor, which releases the original chunks.
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and
tool-call rewrites are deliverable on translations that write them back across the
buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation,
and a rewrite that drops or adds a tool call on any translation, are discarded by the
executor, which releases the original chunks.
The inner guardrail's ``apply_guardrail`` already records the guardrail information
and span, so the observer's stays out of ``log_guardrail_information``."""
@ -101,6 +110,7 @@ class _StreamRewriteObserver(CustomGuardrail):
self.inner: Final = inner
self.rewrote_texts = False
self.rewrote_tool_calls = False
self.changed_tool_call_count = False
def structured_messages_cover_full_request(self) -> bool:
return self.inner.structured_messages_cover_full_request()
@ -118,13 +128,103 @@ class _StreamRewriteObserver(CustomGuardrail):
outputs: Final = await self.inner.apply_guardrail(
inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
)
returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls"))
self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts")))
self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(
sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls"))
self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes)
self.changed_tool_call_count = self.changed_tool_call_count or _changed_count(
sent_tool_shapes, returned_tool_shapes
)
return outputs
class _ScannedTextRecorder(CustomGuardrail):
def __init__(self, guardrail_name: str) -> None:
super().__init__(guardrail_name=guardrail_name)
self.inputs: GenericGuardrailAPIInputs | None = None
@_logged_by_inner_guardrail
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
self.inputs = inputs
return inputs
class _LegacyHookStreamAdapter(CustomGuardrail):
"""Runs a guardrail that only implements the legacy post-call hook (no unified
``apply_guardrail``, or ``use_native_lifecycle_hooks``) as a streaming pipeline step. The
endpoint translation hands it the texts it scanned plus the assembled response under
``request_data["response"]``; the hook gets that response in the shape its route gives
non-streaming hooks, an exception it raises ends the stream through the executor's
fail/error classification, and the response it hands back, or the one it changed in place
and returned ``None`` for, is re-scanned by the same translation so its texts reach the
client through the translation's ended-stream write-back. A
replacement whose scanned texts do not line up with the originals, or whose tool calls
differ from them, is undeliverable, so the executor releases the original chunks. A stream
that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as
long as the hook left the tool calls alone."""
def __init__(
self,
inner: CustomGuardrail,
endpoint_translation: "BaseTranslation",
user_api_key_dict: "UserAPIKeyAuth",
) -> None:
super().__init__(guardrail_name=inner.guardrail_name)
self.inner: Final = inner
self.endpoint_translation: Final = endpoint_translation
self.user_api_key_dict: Final = user_api_key_dict
def structured_messages_cover_full_request(self) -> bool:
return self.inner.structured_messages_cover_full_request()
@_logged_by_inner_guardrail
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
hooked: Final = self.endpoint_translation.post_call_hook_response(request_data.get("response"))
replacement: Final = await self.inner.async_post_call_success_hook(
data=request_data,
user_api_key_dict=self.user_api_key_dict,
response=hooked,
)
rewrite: Final = hooked if replacement is None else replacement
if rewrite is None:
return inputs
rescanned: Final = await self._rescan(rewrite, logging_obj)
if rescanned is None:
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
rewritten: Final = rescanned.get("texts")
if len(_scanned_texts(rewritten)) != len(_scanned_texts(inputs.get("texts"))):
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
if _tool_call_shapes(rescanned.get("tool_calls")) != _tool_call_shapes(inputs.get("tool_calls")):
raise UndeliverableStreamRewrite(self.guardrail_name or "unknown")
if not rewritten:
return inputs
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
return rewritten_inputs
async def _rescan(
self, response: object, logging_obj: "LiteLLMLoggingObj | None"
) -> GenericGuardrailAPIInputs | None:
recorder: Final = _ScannedTextRecorder(self.guardrail_name or "unknown")
await self.endpoint_translation.process_output_response(
response=response,
guardrail_to_apply=recorder,
litellm_logging_obj=logging_obj,
user_api_key_dict=self.user_api_key_dict,
)
return recorder.inputs
def _prepare_hook_input(
step: PipelineStep,
callback: CustomGuardrail,
@ -292,18 +392,29 @@ class PipelineExecutor:
endpoint_translation: "BaseTranslation",
streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place
hook_input: dict[str, object], # mutable-ok: same request-payload shape as data
user_api_key_dict: "UserAPIKeyAuth | None",
user_api_key_dict: "UserAPIKeyAuth",
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> None:
"""Run one streaming post_call step through the endpoint translation, delivering
text rewrites on translations that support ended-stream write-back. A rewrite that
cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation
without write-back, or one the translation refused with
``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the
originals and the step passes, so the client gets the stream the merge base sent."""
observer: Final = _StreamRewriteObserver(callback)
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites
text and tool-call rewrites on translations that support ended-stream write-back. A
guardrail without the unified interface runs its legacy post-call hook against the
assembled response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the
client yet (one on a translation without write-back, one that drops or adds a tool call,
or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is
discarded: the buffered chunks go back to the originals and the step passes, so the
client gets the stream the merge base sent, and the guardrail stays out of the
applied-guardrails header since its output never reached the client. The response an
earlier step's translation stored under ``request_data["response"]`` is dropped first,
so this step's hook sees the stream as the steps before it left it."""
scanner: Final = (
callback
if PipelineExecutor.supports_unified_execution(callback)
else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict)
)
observer: Final = _StreamRewriteObserver(scanner)
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites
originals: Final = copy.deepcopy(streaming_chunks)
hook_input.pop("response", None) # rebind-ok: an earlier step's stored response goes so this step's is stored
try:
if deliver_rewrites:
await endpoint_translation.process_output_streaming_response(
@ -324,9 +435,12 @@ class PipelineExecutor:
)
except UndeliverableStreamRewrite:
_release_original_chunks(step.guardrail, streaming_chunks, originals)
else:
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
return
if observer.changed_tool_call_count or (
not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
return
if not callback.records_own_guardrail_information:
add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)
@ -386,11 +500,11 @@ class PipelineExecutor:
if isinstance(response, dict):
callback.mark_pre_call_hook_ran(response)
elif mode == "post_call" and streaming_chunks is not None:
if not use_unified or endpoint_translation is None:
if endpoint_translation is None:
return (
"error",
None,
f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
f"Guardrail '{step.guardrail}' cannot run on a stream without an endpoint translation",
None,
)
await PipelineExecutor._run_streaming_step(
@ -446,10 +560,22 @@ class PipelineExecutor:
@staticmethod
def supports_unified_execution(callback: CustomGuardrail) -> bool:
"""Whether this guardrail runs through the unified apply_guardrail path,
the interface streaming pipeline execution requires."""
"""Whether this guardrail runs through the unified apply_guardrail path."""
return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
@staticmethod
def supports_streaming_execution(callback: CustomGuardrail) -> bool:
"""Whether a streaming pipeline step can run this guardrail against the buffered
stream: through the unified path, or through its post-call hook on the assembled
response when that hook is its only streaming path. A guardrail with its own
streaming iterator hook, or with neither hook, keeps running on its own."""
callback_type: Final = type(callback)
return PipelineExecutor.supports_unified_execution(callback) or (
callback_type.async_post_call_success_hook is not CustomLogger.async_post_call_success_hook
and callback_type.async_post_call_streaming_iterator_hook
is CustomLogger.async_post_call_streaming_iterator_hook
)
@staticmethod
def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
"""Look up an initialized guardrail callback by name from litellm.callbacks."""

152
litellm/proxy/policy_engine/response_retrieval.py Normal file
View file

@ -0,0 +1,152 @@
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
add_policy_sources_to_metadata,
add_policy_to_applied_policies_header,
)
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.types.proxy.policy_engine import PolicyMatchContext
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...]
_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines)
@dataclass(frozen=True, slots=True)
class UngovernedRetrieval:
reason: Literal["no router", "response id names no deployment", "deployment no longer in the router"]
def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | UngovernedRetrieval:
if llm_router is None:
return UngovernedRetrieval("no router")
model_id: Final = (
ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) if isinstance(response_id, str) else None
)
if model_id is None:
return UngovernedRetrieval("response id names no deployment")
deployment: Final = llm_router.get_deployment(model_id)
if deployment is None:
return UngovernedRetrieval("deployment no longer in the router")
hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias)
if hidden_by is not None:
verbose_proxy_logger.warning(
"Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), "
"so a policy attached to the model name it was submitted as does not run on it",
response_id,
deployment.model_name,
hidden_by,
)
return deployment.model_name
def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None:
if "*" in model_group:
return "a wildcard deployment"
aliases: Final = tuple(
alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group
)
if not aliases:
return None
return f"the target of model_group_alias {', '.join(aliases)}"
def _retrieval_context(
data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str
) -> PolicyMatchContext:
team_alias: Final = user_api_key_dict.team_alias
key_alias: Final = user_api_key_dict.key_alias
return PolicyMatchContext(
team_alias=team_alias if isinstance(team_alias, str) else None,
key_alias=key_alias if isinstance(key_alias, str) else None,
model=model_group,
tags=get_tags_from_request_body(data) or None,
)
def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]:
matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context)
if not matches:
return (), MappingProxyType({})
applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions(
policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list
context=context,
)
post_call_pipelines: Final = tuple(
(policy_name, pipeline)
for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context(
context=context, policy_names=applied_policy_names
)
if pipeline.mode == "post_call"
)
return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches})
def attach_post_call_pipelines_to_retrieval(
data: dict[str, object], # mutable-ok: request-state dict the policy engine hooks all write in place
user_api_key_dict: "UserAPIKeyAuth",
llm_router: "Router | None",
) -> None:
if not get_policy_registry().is_initialized():
return
model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router)
if isinstance(model_group, UngovernedRetrieval):
verbose_proxy_logger.warning(
"Policy engine: background response %s is retrieved without its post_call policy pipelines (%s)",
data.get("response_id"),
model_group.reason,
)
return
context: Final = _retrieval_context(data, user_api_key_dict, model_group)
post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context)
_, bucket = get_or_create_metadata_bucket(data)
already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ())
attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached)
added: Final = tuple(
(policy_name, pipeline)
for policy_name, pipeline in post_call_pipelines
if policy_name not in attached_policy_names
)
if not added:
return
pipelines: Final = (*already_attached, *added)
bucket["_guardrail_pipelines"] = pipelines
bucket["_pipeline_managed_guardrails"] = frozenset(
step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps
)
for policy_name, _pipeline in added:
add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name)
for _policy_name, pipeline in added:
for step in pipeline.steps:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail)
add_policy_sources_to_metadata(
request_data=data,
policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict
policy_name: policy_sources[policy_name] for policy_name, _pipeline in added
},
)
verbose_proxy_logger.debug(
"Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s",
data.get("response_id"),
model_group,
", ".join(policy_name for policy_name, _pipeline in added),
)

368
litellm/proxy/proxy_server.py
View file

@ -17,6 +17,7 @@ import time
import traceback
import warnings
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from types import MappingProxyType, UnionType
from typing import (
@ -131,6 +132,7 @@ from litellm.router_utils.auto_router_tuning_baseline import (
snapshot_tuning_baselines,
tuning_limit_violation,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -138,11 +140,7 @@ from litellm.types.utils import (
TextCompletionResponse,
TokenCountResponse,
)
from litellm.utils import (
_invalidate_model_cost_lowercase_map,
load_credentials_from_list,
reapply_runtime_model_cost_registrations,
)
from litellm.utils import load_credentials_from_list
if TYPE_CHECKING:
from aiohttp import ClientSession
@ -426,6 +424,7 @@ from litellm.proxy.db.exception_handler import (
)
from litellm.proxy.db.gateway_request_tracking import (
GatewayRequestAccumulator,
GatewayRequestRedisBuffer,
flush_gateway_requests,
)
from litellm.proxy.db.proxy_worker_heartbeat import (
@ -2357,6 +2356,17 @@ open_telemetry_logger: OpenTelemetry | None = None
gateway_request_accumulator: Final = GatewayRequestAccumulator()
### INITIALIZE GLOBAL LOGGING OBJECT ###
proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user)
def _gateway_request_redis_buffer() -> GatewayRequestRedisBuffer | None:
"""Shares the spend writer's transaction-buffer Redis and pod lock when use_redis_transaction_buffer is on."""
writer: Final = proxy_logging_obj.db_spend_update_writer
redis_cache: Final = writer.redis_update_buffer.redis_cache
if redis_cache is None or not writer.redis_update_buffer._should_commit_spend_updates_to_redis():
return None
return GatewayRequestRedisBuffer(redis_cache=redis_cache, pod_lock_manager=writer.pod_lock_manager)
### REDIS QUEUE ###
async_result: Final = None
celery_app_conn: Final = None
@ -2707,6 +2717,12 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float)
return fallback_spend, False
@dataclass(frozen=True, slots=True)
class _PendingSpendIncrement:
counter_key: str
increment: float
async def increment_spend_counters(
token: str | None,
team_id: str | None,
@ -2741,7 +2757,7 @@ async def increment_spend_counters(
cost: Final[float] = response_cost
async def _key_scope(key_token: str) -> None:
async def _key_scope(key_token: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
# key_token arrives pre-hashed from metadata["user_api_key"] (auth flow
# hashes raw "sk-..." keys before they reach the callback). The
# startswith("sk-") check is a safety net matching update_cache —
@ -2752,30 +2768,29 @@ async def increment_spend_counters(
hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token
)
key_counter_key: Final = f"spend:key:{hashed_token}"
if key_counter_key not in reserved_counter_keys:
await _init_and_increment_spend_counter(
counter_key=key_counter_key,
source_cache_key=hashed_token,
increment=cost,
key_pending: Final[tuple[_PendingSpendIncrement, ...]] = (
()
if key_counter_key in reserved_counter_keys
else (
await _prepare_spend_counter_increment(
counter_key=key_counter_key,
source_cache_key=hashed_token,
increment=cost,
),
)
key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token)
if key_obj is None:
return
key_budget_limits = getattr(key_obj, "budget_limits", None) or (
key_obj.get("budget_limits") if isinstance(key_obj, dict) else None
)
if isinstance(key_budget_limits, str):
key_budget_limits = json.loads(key_budget_limits)
if not isinstance(key_budget_limits, list):
return
for window in key_budget_limits:
duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration
key_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at
key_window_counter = f"spend:key:{hashed_token}:window:{duration}"
async def _key_window_increment(window: object) -> _PendingSpendIncrement | None:
duration = (
window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None)
)
key_window_reset_at = (
window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
)
key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}"
key_window_start = get_budget_window_start(window)
if key_window_counter not in reserved_counter_keys:
await _init_and_increment_window_spend_counter(
pending_window: Final = (
await _prepare_window_spend_counter_increment(
counter_key=key_window_counter,
entity_type="Key",
entity_id=hashed_token,
@ -2783,6 +2798,9 @@ async def increment_spend_counters(
window_start=key_window_start,
increment=cost,
)
if key_window_counter not in reserved_counter_keys
else None
)
await _enqueue_window_spend_row_update(
entity_type=Litellm_EntityType.KEY,
entity_id=hashed_token,
@ -2792,33 +2810,48 @@ async def increment_spend_counters(
increment=cost,
request_started_at=request_started_at,
)
return pending_window
async def _team_scope(scope_team_id: str) -> None:
team_counter_key: Final = f"spend:team:{scope_team_id}"
if team_counter_key not in reserved_counter_keys:
await _init_and_increment_spend_counter(
counter_key=team_counter_key,
source_cache_key=f"team_id:{scope_team_id}",
increment=cost,
)
team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}")
if team_obj is None:
return
team_budget_limits = getattr(team_obj, "budget_limits", None) or (
team_obj.get("budget_limits") if isinstance(team_obj, dict) else None
key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token)
if key_obj is None:
return key_pending
key_budget_limits = getattr(key_obj, "budget_limits", None) or (
key_obj.get("budget_limits") if isinstance(key_obj, dict) else None
)
if isinstance(team_budget_limits, str):
team_budget_limits = json.loads(team_budget_limits)
if not isinstance(team_budget_limits, list):
return
for window in team_budget_limits:
duration = window["budget_duration"] if isinstance(window, dict) else window.budget_duration
team_window_reset_at = window.get("reset_at") if isinstance(window, dict) else window.reset_at
team_window_counter = f"spend:team:{scope_team_id}:window:{duration}"
if isinstance(key_budget_limits, str):
key_budget_limits = json.loads(key_budget_limits)
if not isinstance(key_budget_limits, list):
return key_pending
window_pending: Final = await asyncio.gather(
*(_key_window_increment(window) for window in key_budget_limits), return_exceptions=True
)
return key_pending + tuple(item for item in window_pending if item is not None)
async def _team_scope(scope_team_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
team_counter_key: Final = f"spend:team:{scope_team_id}"
team_pending: Final[tuple[_PendingSpendIncrement, ...]] = (
()
if team_counter_key in reserved_counter_keys
else (
await _prepare_spend_counter_increment(
counter_key=team_counter_key,
source_cache_key=f"team_id:{scope_team_id}",
increment=cost,
),
)
)
async def _team_window_increment(window: object) -> _PendingSpendIncrement | None:
duration = (
window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None)
)
team_window_reset_at = (
window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
)
team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}"
team_window_start = get_budget_window_start(window)
if team_window_counter not in reserved_counter_keys:
await _init_and_increment_window_spend_counter(
pending_window: Final = (
await _prepare_window_spend_counter_increment(
counter_key=team_window_counter,
entity_type="Team",
entity_id=scope_team_id,
@ -2826,6 +2859,9 @@ async def increment_spend_counters(
window_start=team_window_start,
increment=cost,
)
if team_window_counter not in reserved_counter_keys
else None
)
await _enqueue_window_spend_row_update(
entity_type=Litellm_EntityType.TEAM,
entity_id=scope_team_id,
@ -2835,25 +2871,47 @@ async def increment_spend_counters(
increment=cost,
request_started_at=request_started_at,
)
return pending_window
async def _team_member_scope(scope_user_id: str, scope_team_id: str) -> None:
team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}")
if team_obj is None:
return team_pending
team_budget_limits = getattr(team_obj, "budget_limits", None) or (
team_obj.get("budget_limits") if isinstance(team_obj, dict) else None
)
if isinstance(team_budget_limits, str):
team_budget_limits = json.loads(team_budget_limits)
if not isinstance(team_budget_limits, list):
return team_pending
window_pending: Final = await asyncio.gather(
*(_team_window_increment(window) for window in team_budget_limits), return_exceptions=True
)
return team_pending + tuple(item for item in window_pending if item is not None)
async def _team_member_scope(
scope_user_id: str, scope_team_id: str
) -> tuple[_PendingSpendIncrement | BaseException, ...]:
team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}"
if team_member_counter_key in reserved_counter_keys:
return
await _init_and_increment_spend_counter(
counter_key=team_member_counter_key,
source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}",
increment=cost,
return ()
return (
await _prepare_spend_counter_increment(
counter_key=team_member_counter_key,
source_cache_key=f"team_membership:{scope_user_id}:{scope_team_id}",
increment=cost,
),
)
async def _user_scope(scope_user_id: str) -> None:
async def _user_scope(scope_user_id: str) -> tuple[_PendingSpendIncrement | BaseException, ...]:
user_counter_key: Final = f"spend:user:{scope_user_id}"
if user_counter_key in reserved_counter_keys:
return
await _init_and_increment_spend_counter(
counter_key=user_counter_key,
source_cache_key=scope_user_id,
increment=cost,
return ()
return (
await _prepare_spend_counter_increment(
counter_key=user_counter_key,
source_cache_key=scope_user_id,
increment=cost,
),
)
scope_coros: Final = tuple(
@ -2863,7 +2921,7 @@ async def increment_spend_counters(
_team_scope(team_id) if team_id is not None else None,
_team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None,
_user_scope(user_id) if user_id is not None else None,
_increment_end_user_and_tag_spend_counters(
_prepare_end_user_and_tag_spend_increments(
end_user_id=end_user_id,
tags=tags,
response_cost=cost,
@ -2871,14 +2929,14 @@ async def increment_spend_counters(
)
if end_user_id is not None or tags is not None
else None,
_increment_model_access_group_spend_counters(
_prepare_model_access_group_spend_increments(
model_access_groups=model_access_groups,
response_cost=cost,
reserved_counter_keys=reserved_counter_keys,
)
if model_access_groups
else None,
_increment_org_spend_counter(
_prepare_org_spend_increment(
org_id=org_id,
response_cost=cost,
reserved_counter_keys=reserved_counter_keys,
@ -2893,7 +2951,20 @@ async def increment_spend_counters(
# as orphaned tasks that race the caller's reservation-counter invalidation;
# all scopes settle, then the first error propagates as before.
scope_results: Final = await asyncio.gather(*scope_coros, return_exceptions=True)
scope_errors: Final = [r for r in scope_results if isinstance(r, BaseException)]
scope_errors: Final = tuple(
item
for scope in scope_results
for item in (scope if isinstance(scope, tuple) else (scope,))
if isinstance(item, BaseException)
)
pending: Final = tuple(
item
for scope in scope_results
if not isinstance(scope, BaseException)
for item in scope
if not isinstance(item, BaseException)
)
await _apply_spend_counter_increments(pending=pending)
if scope_errors:
raise scope_errors[0]
@ -2936,41 +3007,49 @@ async def _reconcile_budget_reservation_for_counter_update(
return reserved_counter_keys
async def _increment_end_user_and_tag_spend_counters(
async def _prepare_end_user_and_tag_spend_increments(
end_user_id: str | None,
tags: list[str] | None,
response_cost: float,
reserved_counter_keys: set[str],
) -> None:
if end_user_id is not None:
await _init_and_increment_unreserved_spend_counter(
counter_key=f"spend:end_user:{end_user_id}",
source_cache_key=end_user_cache_key(end_user_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
if tags is None:
return
seen_tags: Final[set[str]] = set()
for tag_name in tags:
if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags:
continue
seen_tags.add(tag_name)
await _init_and_increment_unreserved_spend_counter(
counter_key=f"spend:tag:{tag_name}",
source_cache_key=tag_cache_key(tag_name),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
) -> tuple[_PendingSpendIncrement | BaseException, ...]:
unique_tags: Final = (
tuple(dict.fromkeys(tag for tag in tags if tag and isinstance(tag, str))) if tags is not None else ()
)
results: Final = await asyncio.gather(
*(
coro
for coro in (
_prepare_unreserved_spend_counter_increment(
counter_key=f"spend:end_user:{end_user_id}",
source_cache_key=end_user_cache_key(end_user_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
if end_user_id is not None
else None,
*(
_prepare_unreserved_spend_counter_increment(
counter_key=f"spend:tag:{tag_name}",
source_cache_key=tag_cache_key(tag_name),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
for tag_name in unique_tags
),
)
if coro is not None
),
return_exceptions=True,
)
return tuple(item for item in results if item is not None)
async def _increment_model_access_group_spend_counters(
async def _prepare_model_access_group_spend_increments(
model_access_groups: Sequence[object],
response_cost: float,
reserved_counter_keys: set[str],
) -> None:
) -> tuple[_PendingSpendIncrement | BaseException, ...]:
"""Charge the model access groups that authorized this request.
Without this the counter auth reads is written only by the reservation path, so
@ -2984,55 +3063,63 @@ async def _increment_model_access_group_spend_counters(
unique_groups: Final = tuple(
dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str))
)
for group in unique_groups:
await _init_and_increment_unreserved_spend_counter(
counter_key=model_access_group_spend_counter_key(group),
source_cache_key=model_access_group_cache_key(group),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
results: Final = await asyncio.gather(
*(
_prepare_unreserved_spend_counter_increment(
counter_key=model_access_group_spend_counter_key(group),
source_cache_key=model_access_group_cache_key(group),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
for group in unique_groups
),
return_exceptions=True,
)
return tuple(item for item in results if item is not None)
async def _increment_org_spend_counter(
async def _prepare_org_spend_increment(
org_id: str | None,
response_cost: float,
reserved_counter_keys: set[str],
) -> None:
) -> tuple[_PendingSpendIncrement, ...]:
if org_id is None:
return
return ()
await _init_and_increment_unreserved_spend_counter(
pending: Final = await _prepare_unreserved_spend_counter_increment(
counter_key=f"spend:org:{org_id}",
source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"],
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
return (pending,) if pending is not None else ()
async def _init_and_increment_unreserved_spend_counter(
async def _prepare_unreserved_spend_counter_increment(
counter_key: str,
source_cache_key: str | list[str],
increment: float,
reserved_counter_keys: set[str],
) -> None:
) -> _PendingSpendIncrement | None:
if counter_key in reserved_counter_keys:
return
return None
await _init_and_increment_spend_counter(
return await _prepare_spend_counter_increment(
counter_key=counter_key,
source_cache_key=source_cache_key,
increment=increment,
)
async def _init_and_increment_spend_counter(
async def _prepare_spend_counter_increment(
counter_key: str,
source_cache_key: str | list[str],
increment: float,
):
) -> _PendingSpendIncrement:
"""
Initialize counter from the authoritative DB spend value if not yet
set, then atomically increment in both in-memory and Redis.
set, then return the pending increment for the caller to apply in one
pipelined Redis call.
On first access per pod:
1. Check spend_counter_cache (in-memory -> Redis via DualCache)
@ -3044,13 +3131,13 @@ async def _init_and_increment_spend_counter(
the counter as absent and seed it. Using increment means the worst case
is over-counting (conservative, blocks slightly early) rather than
under-counting (would allow overspend).
4. Increment atomically (both in-memory + Redis)
4. Increment is returned for the caller to apply via pipeline
"""
await _ensure_spend_counter_initialized(
counter_key=counter_key,
source_cache_key=source_cache_key,
)
await _increment_spend_counter_cache(counter_key=counter_key, increment=increment)
return _PendingSpendIncrement(counter_key=counter_key, increment=increment)
async def _enqueue_window_spend_row_update(
@ -3102,20 +3189,20 @@ async def _enqueue_window_spend_row_update(
)
async def _init_and_increment_window_spend_counter(
async def _prepare_window_spend_counter_increment(
counter_key: str,
entity_type: str,
entity_id: str,
window_duration: str | None,
window_start: datetime | None,
increment: float,
):
) -> _PendingSpendIncrement | None:
if window_start is None:
verbose_proxy_logger.warning(
"Skipping spend counter increment for invalid budget window %s",
counter_key,
)
return
return None
initialized: Final = await _ensure_window_spend_counter_initialized(
counter_key=counter_key,
@ -3125,8 +3212,8 @@ async def _init_and_increment_window_spend_counter(
window_start=window_start,
)
if initialized is False:
return
await _increment_spend_counter_cache(counter_key=counter_key, increment=increment)
return None
return _PendingSpendIncrement(counter_key=counter_key, increment=increment)
async def _ensure_spend_counter_initialized(
@ -3259,6 +3346,32 @@ async def _invalidate_spend_counter(counter_key: str):
)
async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncrement]) -> None:
if not pending:
return
redis_cache: Final = spend_counter_cache.redis_cache
if redis_cache is None:
for item in pending:
await spend_counter_cache.async_increment_cache(
key=item.counter_key,
value=item.increment,
refresh_ttl=True,
)
return
ttl: Final = redis_cache.get_ttl()
increment_list: Final = [ # mutable-ok: async_increment_pipeline signature requires list[RedisPipelineIncrementOperation]
RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl)
for item in pending
]
try:
results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list)
except Exception:
await asyncio.gather(*(_invalidate_spend_counter(counter_key=item.counter_key) for item in pending))
raise
for item, current_value in zip(pending, results or ()):
spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value)
async def update_cache(
token: str | None,
user_id: str | None,
@ -4436,20 +4549,9 @@ def resolve_classifier_plugin(
def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
"""Adopt a freshly fetched cost map into this process's litellm state, return the model count"""
litellm.model_cost = new_model_cost_map
# Invalidate case-insensitive lookup map since model_cost was replaced
_invalidate_model_cost_lowercase_map()
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
# wildcard patterns like "anthropic/*" include any newly added models.
litellm.add_known_models(model_cost_map=new_model_cost_map)
# Counted before the re-apply below, which writes into this same dict, so the
# number reported describes the fetched price data alone.
fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
# The swap discards everything registered at runtime (deployment model_info,
# register_model overrides), so put it back on top of the fresh catalog.
reapply_runtime_model_cost_registrations()
return fetched_model_count
from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map
return adopt_model_cost_map(new_model_cost_map)
def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:
@ -9543,7 +9645,7 @@ class ProxyStartupEvent:
flush_gateway_requests,
"interval",
seconds=batch_writing_interval,
args=(prisma_client, gateway_request_accumulator),
args=(prisma_client, gateway_request_accumulator, _gateway_request_redis_buffer()),
id="update_gateway_requests_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,

2
litellm/proxy/schema.prisma
View file

@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])

264
litellm/proxy/utils.py
View file

@ -94,6 +94,7 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert
from litellm.litellm_core_utils.core_helpers import (
coerce_token_limit,
get_or_create_metadata_bucket,
independent_snapshot,
is_expected_client_error,
)
@ -157,6 +158,8 @@ from litellm.proxy.hooks.sensitive_data_routing import (
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.table_repositories import (
@ -172,6 +175,7 @@ from litellm.repositories.verification_token_repository import (
)
from litellm.secret_managers.main import str_to_bool
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.mcp import (
MCPDuringCallResponseObject,
MCPPreCallRequestObject,
@ -193,6 +197,7 @@ if TYPE_CHECKING:
from prisma.types import HttpConfig
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.models.team import LiteLLM_TeamTableCachedObj
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
@ -455,7 +460,7 @@ def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipe
return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps)
def _pipeline_managed_guardrail_names(
def pipeline_managed_guardrail_names(
data: Mapping[str, object], mode: Literal["pre_call", "post_call"]
) -> frozenset[str]:
return _pipeline_step_guardrail_names(
@ -518,9 +523,17 @@ def _merge_pipeline_metadata_writes(
_merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool:
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
if callback is None:
return False
if PipelineExecutor.supports_unified_execution(callback):
return True
return (
translation is not None
and type(translation).assembles_streamed_response
and PipelineExecutor.supports_streaming_execution(callback)
)
def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
@ -529,50 +542,174 @@ def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardr
)
def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None:
if data.get("background") is not True:
return
policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data))
if not policy_names:
return
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines do not run on background responses yet; "
"the response is released ungoverned by them: %s",
", ".join(policy_names),
_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress"))
def _is_pending_background_response(response: LLMResponseTypes) -> bool:
return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES
def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]:
resolved: Final = PolicyResolver.resolve_policy_guardrails(
policy_name=policy_name, policies=get_policy_registry().get_all_policies()
)
return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps)
def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]:
return frozenset(
callback.guardrail_name
for callback in litellm.callbacks
if isinstance(callback, CustomGuardrail)
and callback.guardrail_name is not None
and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
)
def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool:
unsupported: Final = tuple(
def _without_names(
bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write
slot: str,
names: frozenset[str],
) -> None:
claimed: Final = bucket.get(slot)
if not isinstance(claimed, list):
return
remaining: Final = [ # mutable-ok: the slot stays a list, the shape every applied_* header writer appends to
name for name in claimed if name not in names
]
if remaining:
bucket[slot] = remaining # rebind-ok: the slot lives in the shared request-state dict, rewritten in place
else:
bucket.pop(slot)
def _withdraw_deferred_claims(
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
deferred: Sequence[tuple[str, "GuardrailPipeline"]],
) -> None:
outside_by_policy: Final = MappingProxyType(
{policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred}
)
running_elsewhere: Final = pipeline_managed_guardrail_names(data, "pre_call").union(
_guardrails_run_standalone_pre_call(data), *outside_by_policy.values()
)
withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside)
withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere
_, bucket = get_or_create_metadata_bucket(data)
_without_names(bucket, "applied_policies", withdrawn_policies)
_without_names(bucket, "applied_guardrails", withdrawn_guardrails)
sources: Final = bucket.get("policy_sources")
if not isinstance(sources, dict):
return
remaining_sources: Final = { # mutable-ok: policy_sources stays a dict, the shape its writer updates in place
name: reason for name, reason in sources.items() if name not in withdrawn_policies
}
if remaining_sources:
bucket["policy_sources"] = remaining_sources
else:
bucket.pop("policy_sources")
def _defer_post_call_pipelines(
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
response: ResponsesAPIResponse,
) -> None:
deferred: Final = _post_call_pipelines(data)
if not deferred:
return
verbose_proxy_logger.debug(
"Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s",
response.id,
response.status,
", ".join(policy_name for policy_name, _pipeline in deferred),
)
tag_matched: Final = _tag_matched_deferrals(data, deferred)
if tag_matched:
verbose_proxy_logger.warning(
"Policy engine: background response %s matched post_call policies through a request tag at submit; "
"retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body "
"does not govern the completed response: %s",
response.id,
", ".join(tag_matched),
)
body_selected: Final = _body_selected_deferrals(data, deferred)
if body_selected:
verbose_proxy_logger.warning(
"Policy engine: background response %s matched post_call policies through the request body's policies "
"list at submit; retrieval carries no request body, so those policies do not govern the completed "
"response: %s",
response.id,
", ".join(body_selected),
)
_withdraw_deferred_claims(data, deferred)
def _tag_matched_deferrals(
data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]]
) -> tuple[str, ...]:
sources: Final = _policy_state_metadata(data).get("policy_sources")
if not isinstance(sources, dict):
return ()
return tuple(
policy_name
for policy_name, _pipeline in deferred
if policy_name in sources and "tag:" in str(sources[policy_name])
)
def _body_selected_deferrals(
data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]]
) -> tuple[str, ...]:
sources: Final = _policy_state_metadata(data).get("policy_sources")
attributed: Final = frozenset(sources) if isinstance(sources, dict) else frozenset()
return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed)
def _pipeline_unsupported_streaming_guardrails(
pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
) -> tuple[str, ...]:
return tuple(
dict.fromkeys(
step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail)
step.guardrail
for step in pipeline.steps
if not _pipeline_step_supports_streaming(step.guardrail, translation)
)
)
def _pipeline_is_streamable(
policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
) -> bool:
unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation)
if not unsupported:
return True
verbose_proxy_logger.warning(
"Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, "
"which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s",
"Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they "
"need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a "
"route whose translation assembles the streamed response. The stream skips the pipeline and its "
"guardrails run on their own: %s",
policy_name,
", ".join(unsupported),
)
return False
def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool:
return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None
def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None":
resolved: Final = resolve_endpoint_translation(user_api_key_dict, None)
return None if resolved is None else resolved[1]
def _stream_gated_guardrail_names(
def stream_gated_guardrail_names(
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> frozenset[str]:
if not _route_supports_streaming_pipelines(user_api_key_dict):
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
if translation is None:
return frozenset()
return _pipeline_step_guardrail_names(
tuple(
(policy_name, pipeline)
for policy_name, pipeline in _post_call_pipelines(request_data)
if all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps)
if not _pipeline_unsupported_streaming_guardrails(pipeline, translation)
)
)
@ -584,16 +721,19 @@ def _streamable_post_call_pipelines(
The post_call pipelines a streaming response can be gated through.
Streaming pipelines scan the buffered stream through the endpoint guardrail
translation of the request route, so every step's guardrail needs the
unified apply_guardrail interface and the route needs a translation. A
pipeline that cannot be run that way yet is left out and its guardrails
run on the stream on their own, the way they did before pipelines ran on
streams at all, with a warning naming the pipeline.
translation of the request route, so every step's guardrail needs either the
unified apply_guardrail interface or, on a route whose translation assembles
the streamed response, a post-call hook that is its only streaming path, and
the route needs a translation. A pipeline that
cannot be run that way yet is left out and its guardrails run on the stream
on their own, the way they did before pipelines ran on streams at all, with
a warning naming the pipeline.
"""
post_call_pipelines: Final = _post_call_pipelines(request_data)
if not post_call_pipelines:
return ()
if not _route_supports_streaming_pipelines(user_api_key_dict):
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
if translation is None:
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
"(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run "
@ -605,7 +745,7 @@ def _streamable_post_call_pipelines(
return tuple(
(policy_name, pipeline)
for policy_name, pipeline in post_call_pipelines
if _pipeline_is_streamable(policy_name, pipeline)
if _pipeline_is_streamable(policy_name, pipeline, translation)
)
@ -1985,8 +2125,6 @@ class ProxyLogging:
)
try:
_warn_background_skips_post_call_pipelines(data)
# Execute guardrail pipelines before the normal callback loop
data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below
data=data,
@ -1997,7 +2135,7 @@ class ProxyLogging:
)
# Get pipeline-managed guardrails to skip in normal loop
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call")
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call")
caps: Final = ProxyLogging._callback_capabilities()
# Skip the per-request callback walk entirely when nothing in
@ -2956,6 +3094,24 @@ class ProxyLogging:
daemon=True,
).start()
async def _run_post_call_pipelines(
self,
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
user_api_key_dict: UserAPIKeyAuth,
response: LLMResponseTypes,
) -> LLMResponseTypes | None:
if _is_pending_background_response(response):
_defer_post_call_pipelines(data, response)
return None
_, pipeline_response = await self._maybe_execute_pipelines(
data=data,
user_api_key_dict=user_api_key_dict,
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
event_hook="post_call",
response=response,
)
return pipeline_response
async def post_call_success_hook(
self,
data: dict,
@ -2975,17 +3131,15 @@ class ProxyLogging:
from litellm.proxy.proxy_server import llm_router
from litellm.types.guardrails import GuardrailEventHooks
_, pipeline_response = await self._maybe_execute_pipelines(
pipeline_response: Final = await self._run_post_call_pipelines(
data=data,
user_api_key_dict=user_api_key_dict,
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
event_hook="post_call",
response=response,
)
if pipeline_response is not None:
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call")
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call")
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
try:
# Merge model-level guardrails before checking which guardrails to run
@ -3301,7 +3455,7 @@ class ProxyLogging:
_cached_guardrail_data: dict | None = None
_guardrail_data_computed = False
pipeline_gated: Final = (
_stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
)
for callback in litellm.callbacks:
@ -3436,12 +3590,16 @@ class ProxyLogging:
),
)
if post_call_pipelines:
pipeline_translation: Final = (
resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None
)
if pipeline_translation is not None:
current_response = self._pipeline_gated_stream(
response=current_response,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
pipelines=post_call_pipelines,
translation=pipeline_translation,
)
try:
@ -3465,6 +3623,7 @@ class ProxyLogging:
user_api_key_dict: UserAPIKeyAuth,
request_data: dict, # mutable-ok: same request-payload shape the hooks mutate
pipelines: "tuple[tuple[str, GuardrailPipeline], ...]",
translation: "tuple[str, BaseTranslation]",
) -> "AsyncGenerator[Any, None]":
"""
Execute post_call policy pipelines against a streamed response.
@ -3474,14 +3633,13 @@ class ProxyLogging:
assembled output through the endpoint guardrail translation, the same
machinery flat post_call guardrails use at end of stream. An allow
releases the buffered chunks: verbatim when no guardrail rewrote the
output, rewritten in place when one rewrote text and the translation
delivers ended-stream rewrites (later steps then re-scan the rewritten
chunks, so rewrites chain). A rewrite the translation cannot deliver
yet (a tool-call rewrite, or a text rewrite on a route without
write-back) is discarded by the executor and the original chunks are
released, as is a buffered shape no translation resolves; a block or
modify_response terminates with the translation's block chunks or the
raised error.
output, rewritten in place when one rewrote text or a tool call and the
translation delivers ended-stream rewrites (later steps then re-scan the
rewritten chunks, so rewrites chain). A rewrite the translation cannot
deliver yet (one on a route without write-back, or a shape the route
refuses) is discarded by the executor and the original chunks are
released; a block or modify_response terminates with the translation's
block chunks or the raised error.
"""
buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict
async for item in response:
@ -3489,17 +3647,7 @@ class ProxyLogging:
if not buffered:
return
resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0])
if resolved is None:
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; "
"the stream is released ungoverned by them: %s",
", ".join(policy_name for policy_name, _pipeline in pipelines),
)
for buffered_item in buffered:
yield buffered_item
return
call_type, endpoint_translation = resolved
call_type, endpoint_translation = translation
for policy_name, pipeline in pipelines:
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(

11
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -437,14 +437,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
response_created_event_data["temperature"] = self.responses_api_request["temperature"]
if "text" in self.responses_api_request:
response_created_event_data["text"] = self.responses_api_request["text"]
if "tool_choice" in self.responses_api_request:
# Transform tool_choice from dict format (e.g., {"type": "auto"}) to string format
response_created_event_data["tool_choice"] = (
LiteLLMCompletionResponsesConfig._transform_tool_choice(self.responses_api_request["tool_choice"])
or "auto"
response_created_event_data["tool_choice"] = (
LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response(
self.responses_api_request.get("tool_choice")
)
else:
response_created_event_data["tool_choice"] = "auto"
)
if "tools" in self.responses_api_request:
response_created_event_data["tools"] = self.responses_api_request["tools"]
else:

31
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -27,8 +27,10 @@ from openai.types.chat.chat_completion_named_tool_choice_param import (
)
from openai.types.responses import ResponseFunctionToolCall
from openai.types.responses.response_create_params import ResponseInputParam
from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam
from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam
from openai.types.responses.tool_param import FunctionToolParam
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
@ -68,6 +70,7 @@ from litellm.types.llms.openai import (
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
ResponsesAPIStatus,
ToolChoice,
ValidChatCompletionMessageContentTypes,
ValidChatCompletionMessageContentTypesLiteral,
)
@ -126,6 +129,7 @@ _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]])
_TEXT_ADAPTER: Final = TypeAdapter(str)
_RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice)
@runtime_checkable
@ -267,6 +271,27 @@ class LiteLLMCompletionResponsesConfig:
# Return as-is for unknown formats
return tool_choice
@staticmethod
def _transform_tool_choice_for_responses_api_response(tool_choice: object) -> ToolChoice:
if tool_choice is None:
return "auto"
try:
return _RESPONSES_API_TOOL_CHOICE_ADAPTER.validate_python(tool_choice)
except ValidationError:
return LiteLLMCompletionResponsesConfig._chat_tool_choice_as_responses_api_tool_choice(tool_choice)
@staticmethod
def _chat_tool_choice_as_responses_api_tool_choice(tool_choice: object) -> ToolChoice:
match tool_choice, LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice):
case {"type": "custom"}, {"function": {"name": str(custom_name)}}:
return ToolChoiceCustomParam(type="custom", name=custom_name)
case _, {"type": "function", "function": {"name": str(function_name)}}:
return ToolChoiceFunctionParam(type="function", name=function_name)
case _, "none" | "auto" | "required" as normalized:
return normalized
case _, _:
return "auto"
@staticmethod
def _should_drop_derived_web_search_options(model: str, custom_llm_provider: str | None) -> bool:
"""
@ -2263,7 +2288,9 @@ class LiteLLMCompletionResponsesConfig:
),
parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False),
temperature=getattr(chat_completion_response, "temperature", 0),
tool_choice=getattr(chat_completion_response, "tool_choice", "auto"),
tool_choice=LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response(
responses_api_request.get("tool_choice")
),
tools=getattr(chat_completion_response, "tools", []),
top_p=getattr(chat_completion_response, "top_p", None),
max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None),

56
litellm/responses/streaming_iterator.py
View file

@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runti
import httpx
from openai._streaming import SSEDecoder
from pydantic import BaseModel, ValidationError
from typing_extensions import TypeIs
import litellm
@ -438,18 +439,7 @@ class BaseResponsesAPIStreamingIterator:
if self._persist_completed_response_before_logging:
self._persist_completed_response_to_cache(is_async=is_async)
# Create a copy for logging to avoid modifying the response object that will be returned to the user
# The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
logging_response = self.completed_response
if self.completed_response is not None and hasattr(self.completed_response, "model_dump"):
try:
logging_response = type(self.completed_response).model_validate(self.completed_response.model_dump())
except Exception:
# Fallback to original if serialization fails
pass
logging_response: Final[object] = _logging_copy(self.completed_response)
self._restore_provider_response_headers(logging_response)
end_time: Final = datetime.now()
@ -488,10 +478,10 @@ class BaseResponsesAPIStreamingIterator:
def _restore_provider_response_headers(self, logging_response: object) -> None:
"""Re-apply the provider's response headers to the copy handed to logging callbacks.
``model_validate(model_dump())`` above drops pydantic private attributes, so the
``model_validate(model_dump())`` in ``_logging_copy`` drops pydantic private attributes, so the
``_hidden_params`` the provider transform set on the nested response are lost. Returns early
when that copy fell back to the original event, so logging-only state never lands on the
object the caller is iterating.
when the event was not a pydantic model and logging got the original, so logging-only state
never lands on the object the caller is iterating.
"""
if logging_response is self.completed_response:
return
@ -544,7 +534,7 @@ class BaseResponsesAPIStreamingIterator:
def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None:
if response_obj is None or self.logging_obj is None:
return
usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None))
if usage_obj is None:
return
try:
@ -1293,14 +1283,46 @@ def _add_text_like_part_events(
)
def _logging_copy(event: object) -> object:
"""Hand logging callbacks a copy, so their usage rewrite (Responses shape to chat shape) never
reaches the event the caller is iterating. The round trip through ``model_dump`` sidesteps the
deepcopy pickle errors of #17192; when a provider payload fails validation (LIT-7391), shallow
copies of the event and its nested response still keep the caller's ``usage`` attribute separate."""
if not isinstance(event, BaseModel):
return event
try:
return type(event).model_validate(event.model_dump())
except Exception:
return _detached_shallow_copy(event)
def _detached_shallow_copy(event: BaseModel) -> BaseModel:
nested: Final[object] = getattr(event, "response", None)
if isinstance(nested, BaseModel):
return event.model_copy(update={"response": nested.model_copy()})
return event.model_copy()
def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
if isinstance(usage, ResponseAPIUsage):
return usage
if not isinstance(usage, dict):
return None
try:
return ResponseAPIUsage.model_validate(usage)
except ValidationError:
return None
def _stamp_responses_usage_cost(
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
) -> None:
if response_obj is None or logging_obj is None:
return
usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None)
usage_obj: Final[ResponseAPIUsage | None] = _usage_as_model(getattr(response_obj, "usage", None))
if usage_obj is None:
return
response_obj.usage = usage_obj # rebind-ok: the stamped cost has to ride on the response the client receives
if isinstance(getattr(usage_obj, "cost", None), (int, float)):
return
try:

43
litellm/responses/utils.py
View file

@ -543,6 +543,49 @@ class ResponsesAPIRequestUtils:
return request_input
@staticmethod
def strip_encrypted_reasoning_from_input(request_input: object) -> None:
"""Drop reasoning items the routed deployment cannot decrypt, keeping their readable summary.
Mutates ``request_input`` in place: the router's fallback snapshot shares this
list object, so a rebound list would replay the stripped items on the fallback hop.
"""
if not isinstance(request_input, list):
return
items: Final = cast(list[object], request_input) # cast-ok: untyped client json
stripped: Final = tuple(ResponsesAPIRequestUtils._without_encrypted_reasoning(item) for item in items)
items[:] = (item for item in stripped if item is not None) # rebind-ok: list shared with fallback snapshot
@staticmethod
def _without_encrypted_reasoning(item: object) -> object | None:
if not isinstance(item, dict):
return item
reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json
if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"):
return reasoning
readable: Final = any(
ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content")
)
if not readable:
return None
kept: Final[dict[str, object]] = { # mutable-ok: request item rebuilt without the undecryptable keys
key: value for key, value in reasoning.items() if key not in ("encrypted_content", "id")
}
return kept
@staticmethod
def _has_readable_text(value: object) -> bool:
"""A reasoning item's ``summary``/``content`` carries readable text: a non-empty string, or a
list holding at least one block with a non-empty ``text`` field (summary_text / output_text)."""
if isinstance(value, str):
return bool(value.strip())
if isinstance(value, list):
return any(
isinstance(block, dict) and bool(cast(Mapping[str, object], block).get("text")) # cast-ok: untyped json
for block in value
)
return False
@staticmethod
def _build_responses_api_response_id(
custom_llm_provider: str | None,

228
litellm/router.py
View file

@ -148,6 +148,7 @@ from litellm.router_utils.common_utils import (
_is_proxy_admin_request,
filter_team_based_models,
filter_web_search_deployments,
get_request_team_id,
resolve_model_group_alias,
truncate_fallback_error_detail,
warn_on_provider_credential_mismatch,
@ -1317,6 +1318,43 @@ class Router:
if isinstance(litellm.input_callback, list):
litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids]
def _apply_updated_routing_strategy_args(self) -> None:
"""
Re-link the default group's selector to the current `routing_strategy_args`.
Selectors freeze their `RoutingArgs` at construction, so a runtime args
update would otherwise keep serving the boot-time values until restart.
Latency/usage state survives the rebuild: it lives in the shared router
cache, not on the selector.
"""
strategy: Final = self._normalize_strategy(self.routing_strategy)
if strategy == "lar1":
from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy
apply_lar1_routing_strategy(self, self.routing_strategy_args)
return
attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "")
current: Final = getattr(self, attr, None) if attr is not None else None
if attr is None or current is None:
return
try:
rebuilt: Final = self._build_strategy_selector(
strategy=strategy or "",
routing_strategy_args=self.routing_strategy_args,
)
except (TypeError, ValidationError):
verbose_router_logger.exception(
"Invalid routing_strategy_args %s for '%s'; keeping the previous ones",
self.routing_strategy_args,
strategy,
)
return
self._unregister_router_selectors((current,))
setattr(self, attr, rebuilt)
def routing_strategy_init(self, routing_strategy: RoutingStrategy | str, routing_strategy_args: dict):
verbose_router_logger.info("Routing strategy: %s", routing_strategy)
self._validate_routing_strategy(routing_strategy)
@ -11130,6 +11168,43 @@ class Router:
return ids
def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]:
"""
Deployment ids that could serve ``model`` for ``team_id``, following the same
precedence ``_common_checks_available_deployment`` uses to build a candidate pool:
``model_group_alias``, then a routing group, then the first matching early-resolve
path for a name that is not a ``model_name`` (team route, wildcard pattern via
``get_deployments_by_pattern``, team pattern router, default deployment), then the
``model_name`` and team indexes. Delegating to the router's own resolvers keeps this
aligned with how a route actually resolves rather than re-deriving it, and unlike
``_common_checks_available_deployment`` it is read-only: it does not apply request
fallbacks and (with ``include_team_models`` left off) does not raise. Lets a pre-call
check tell a genuine cross-group route from same-group unavailability without leaking
deployment ids into request kwargs bound for the provider.
"""
resolved: Final = self._get_model_from_alias(model=model) or model
routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id)
if routing_group_members is not None:
return self._deployment_ids(routing_group_members)
early: Final = self._try_early_resolve_deployments_for_model_not_in_names(
model=resolved, request_team_id=team_id
)
if early is not None:
early_deployments: Final = early[1]
return self._deployment_ids(
(early_deployments,) if isinstance(early_deployments, Mapping) else early_deployments
)
return self._deployment_ids(self._get_all_deployments(model_name=resolved, team_id=team_id))
@staticmethod
def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]:
return frozenset(
str(model_info["id"])
for deployment in deployments
for model_info in (deployment.get("model_info"),)
if isinstance(model_info, Mapping) and model_info.get("id") is not None
)
def has_model_id(self, candidate_id: str) -> bool:
"""
O(1) membership check for a deployment ID without allocating large lists.
@ -11847,7 +11922,7 @@ class Router:
_existing_router_settings: Final = self.get_settings()
rebuild_routing_groups = False
relink_lar1_from_args = False
routing_args_updated = False
for var in kwargs:
if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS:
if var in _int_settings:
@ -11886,15 +11961,13 @@ class Router:
)
rebuild_routing_groups = True
elif var == "routing_strategy_args":
relink_lar1_from_args = True
routing_args_updated = True
setattr(self, var, value)
else:
verbose_router_logger.debug("Setting %s is not allowed", var)
if relink_lar1_from_args and self._normalize_strategy(self.routing_strategy) == "lar1":
from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy
apply_lar1_routing_strategy(self, self.routing_strategy_args)
if routing_args_updated:
self._apply_updated_routing_strategy_args()
if rebuild_routing_groups:
self._init_routing_groups(self._routing_groups_input)
@ -12268,27 +12341,7 @@ class Router:
if team_deployments:
return model, team_deployments
elif include_team_models:
team_deployments = [
self.model_list[index]
for (_, public_model_name), indices in self.team_model_to_deployment_indices.items()
if public_model_name == model
for index in indices
]
team_ids: Final = {
team_id
for deployment in team_deployments
for team_id in [(deployment.get("model_info") or {}).get("team_id")]
if team_id is not None
}
if len(team_ids) > 1:
raise litellm.BadRequestError(
message=(
f"Model name '{model}' matches deployments from multiple teams. "
"Specify the deployment ID directly to disambiguate."
),
model=model,
llm_provider="",
)
team_deployments = self._team_deployments_across_teams(model)
if team_deployments:
return model, team_deployments
@ -12315,6 +12368,45 @@ class Router:
return None
def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]:
"""Every team's deployments under public name `model`, for a proxy admin calling without a team."""
team_deployments: Final = [
self.model_list[index]
for (_, public_model_name), indices in self.team_model_to_deployment_indices.items()
if public_model_name == model
for index in indices
]
team_ids: Final = {
team_id
for deployment in team_deployments
for team_id in [(deployment.get("model_info") or {}).get("team_id")]
if team_id is not None
}
if len(team_ids) > 1:
raise litellm.BadRequestError(
message=(
f"Model name '{model}' matches deployments from multiple teams. "
"Specify the deployment ID directly to disambiguate."
),
model=model,
llm_provider="",
)
return team_deployments
def deployments_for_request(
self, model: str, request_kwargs: Mapping[str, object]
) -> Sequence[DeploymentTypedDict]:
"""The deployments `model` names for this caller, through the same alias, then team-first, then
global, then admin-across-teams resolution `_common_checks_available_deployment` applies, so
strategy selection and compression policy can never disagree with deployment selection about
which marker a name means."""
registered_name: Final = self._get_model_from_alias(model=model) or model
team_id: Final = get_request_team_id(request_kwargs)
deployments: Final = self._get_all_deployments(model_name=registered_name, team_id=team_id)
if deployments or team_id is not None or not _is_proxy_admin_request(request_kwargs):
return deployments
return self._team_deployments_across_teams(registered_name)
@staticmethod
def _is_strategy_marker_deployment(deployment: Mapping[str, object]) -> bool:
litellm_params: Final = deployment.get("litellm_params")
@ -12342,11 +12434,7 @@ class Router:
- Dict, if specific model chosen
"""
request_team_id: str | None = None
if request_kwargs is not None:
metadata: Final = request_kwargs.get("metadata") or {}
litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {}
request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id")
request_team_id: Final = get_request_team_id(request_kwargs)
# check if aliases set on litellm model alias map
if specific_deployment is True:
return model, self._get_deployment_by_litellm_model(model=model)
@ -12371,7 +12459,9 @@ class Router:
include_team_models=_is_proxy_admin_request(request_kwargs),
)
if early is not None:
return early
if not isinstance(early[1], list):
return early
return early[0], self._drop_strategy_markers(early[0], early[1])
## get healthy deployments
### get all deployments
@ -12448,19 +12538,22 @@ class Router:
model
] # update the model to the actual value if an alias has been passed in
marker_flags: Final = tuple(self._is_strategy_marker_deployment(d) for d in healthy_deployments)
if not any(marker_flags):
return model, healthy_deployments
selectable: Final = [ # mutable-ok: matches this function's list contract expected by downstream filters
d for d, is_marker in zip(healthy_deployments, marker_flags, strict=True) if not is_marker
return model, self._drop_strategy_markers(model, healthy_deployments)
def _drop_strategy_markers(
self, model: str, deployments: Sequence[DeploymentTypedDict]
) -> list[DeploymentTypedDict]:
"""A strategy marker is never a callable deployment, whichever resolution arm produced it."""
selectable: Final = [ # mutable-ok: matches _common_checks_available_deployment's list contract
d for d in deployments if not self._is_strategy_marker_deployment(d)
]
if not selectable:
if deployments and not selectable:
raise litellm.BadRequestError(
message=f"You passed in model={model}. {RouterErrors.only_strategy_marker_deployments.value}",
model=model,
llm_provider="",
)
return model, selectable
return selectable
def _filter_deployments_by_model_access_groups(
self,
@ -13150,12 +13243,8 @@ class Router:
return filtered
def _model_name_has_plain_deployments(self, model: str) -> bool:
indices: Final = self.model_name_to_deployment_indices.get(model) or ()
return any(not self._is_strategy_marker_deployment(self.model_list[idx]) for idx in indices)
def _select_pre_routing_strategy(
self, model: str, request_kwargs: dict
self, model: str, request_kwargs: Mapping[str, object]
) -> "TaggedPreRoutingStrategy[PreRoutingStrategy] | None":
"""
Resolve the pre-routing strategy for `model`, disambiguating deployments
@ -13166,6 +13255,12 @@ class Router:
deployment the strategy was registered from via its (model_name, tags)
pair.
The registries are keyed by the marker deployment's own `model_name`, which
for a team-scoped router is the internal `model_name_{team}_{uuid}` while
the caller sends the team's public name. So the names looked up are the
`model_name`s of whatever deployments this caller's request resolves `model`
to, and `model` itself when it resolves to none.
With tag filtering enabled, router-wide or by the request's
enable_tag_filtering (which the proxy sets from key/team
router_settings), strategies that all carry real tags matching none of
@ -13173,12 +13268,14 @@ class Router:
deployments: returning None hands the request to ordinary tag-aware
deployment selection.
"""
candidates: Final[list[TaggedPreRoutingStrategy[PreRoutingStrategy]]] = [
*self.auto_routers.get(model, []),
*self.complexity_routers.get(model, []),
*self.adaptive_routers.get(model, []),
*self.quality_routers.get(model, []),
]
registries: Final = (self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers)
if not any(registries):
return None
deployments: Final = self.deployments_for_request(model, request_kwargs)
registered_names: Final = tuple(dict.fromkeys(str(d["model_name"]) for d in deployments)) or (model,)
candidates: Final = tuple(
tagged for registry in registries for name in registered_names for tagged in registry.get(name, [])
)
if not candidates:
return None
@ -13196,7 +13293,7 @@ class Router:
if (
(self.enable_tag_filtering or request_scoped_filtering)
and all(tagged.tags for tagged in candidates)
and self._model_name_has_plain_deployments(model)
and any(not self._is_strategy_marker_deployment(d) for d in deployments)
):
return None
return candidates[0]
@ -13308,11 +13405,12 @@ class Router:
Used for the litellm auto-router to modify the request before the routing decision is made.
`model` is whatever the caller asked for, which may be a `model_group_alias` key, while the
strategy registries and the marker deployment are keyed by the marker's own `model_name`, so
every lookup below resolves the alias first. Only the lookups: the caller-facing name stays
the alias, since spend metadata is stamped before routing and the response carries the tier
group the strategy picked.
`model` is whatever the caller asked for, which may be a `model_group_alias` key or a team's
public model name, while the strategy registries and the marker deployment are keyed by the
marker's own `model_name`, so every lookup below resolves the alias first and the team name
through the deployment path. Only the lookups: the caller-facing name stays the alias, since
spend metadata is stamped before routing and the response carries the tier group the
strategy picked.
"""
requested_registered_model_name: Final = self._get_model_from_alias(model=model) or model
registered_model_name: Final = await self._resolve_claude_code_session_router(
@ -13349,7 +13447,6 @@ class Router:
messages_for_routing,
model_hop_compression_armed,
policy_for_model,
team_id_from_request,
)
# Same tag-aware lookup the proxy's pre-call arming used, so an alias with
@ -13357,7 +13454,7 @@ class Router:
compression_policy: Final = policy_for_model(
llm_router=self,
model_alias=registered_model_name,
team_id=team_id_from_request(request_kwargs),
request_kwargs=request_kwargs,
request_tags=_get_tags_from_request_kwargs(request_kwargs),
)
# Shared compression already ran in the pre-call hook, so reuse it rather than
@ -13426,7 +13523,9 @@ class Router:
# Per-tier `litellm_params` on the hook response are deliberate overrides
# the caller applies on top, so those keys are never forwarded here.
marker_params: Final = (
self._forwardable_alias_marker_params(model=registered_model_name, strategy_tags=selected_strategy.tags)
self._forwardable_alias_marker_params(
model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs
)
if pre_routing_hook_response is not None
else ()
)
@ -13444,13 +13543,14 @@ class Router:
return pre_routing_hook_response
def _forwardable_alias_marker_params(
self, model: str, strategy_tags: tuple[str, ...]
self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object]
) -> tuple[tuple[str, object], ...]:
marker_params: Final = tuple(
litellm_params
for idx in self.model_name_to_deployment_indices.get(model, ())
if isinstance(litellm_params := self.model_list[idx].get("litellm_params", {}), dict)
and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX)
for deployment in self.deployments_for_request(model, request_kwargs)
if str((litellm_params := deployment["litellm_params"]).get("model", "")).startswith(
AUTO_ROUTER_MODEL_PREFIX
)
)
tag_matched: Final = tuple(
params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags

21
litellm/router_strategy/lowest_latency.py
View file

@ -3,8 +3,11 @@
import random
from collections.abc import Sequence
from datetime import datetime, timedelta
from math import ceil
from typing import TYPE_CHECKING, Any, Final
from pydantic import Field
import litellm
from litellm import ModelResponse, token_counter, verbose_logger
from litellm.caching.caching import DualCache
@ -24,6 +27,7 @@ class RoutingArgs(LiteLLMPydanticObjectBase):
ttl: float = 1 * 60 * 60 # 1 hour
lowest_latency_buffer: float = 0
max_latency_list_size: int = 10
ttft_percentile: float | None = Field(default=None, gt=0, le=1)
def _average_latency(samples: Sequence[float]) -> float:
@ -32,6 +36,12 @@ def _average_latency(samples: Sequence[float]) -> float:
return sum(samples) / len(samples)
def _percentile_latency(samples: Sequence[float], percentile: float) -> float:
values: Final = sorted(samples)
index: Final = ceil(len(values) * percentile) - 1
return values[index]
def _ttft_seconds(elapsed: timedelta | float) -> float:
if isinstance(elapsed, timedelta):
return elapsed.total_seconds()
@ -427,14 +437,17 @@ class LowestLatencyLoggingHandler(CustomLogger):
item_rpm = item_map.get(precise_minute, {}).get("rpm", 0)
item_tpm = item_map.get(precise_minute, {}).get("tpm", 0)
# get average latency or average ttft (depending on streaming/non-streaming)
use_ttft = (
request_kwargs is not None
and request_kwargs.get("stream", None) is not None
and request_kwargs["stream"] is True
and len(item_ttft_latency) > 0
)
average_latency = _average_latency(item_ttft_latency if use_ttft else item_latency)
selected_latency = (
_percentile_latency(item_ttft_latency, self.routing_args.ttft_percentile)
if use_ttft and self.routing_args.ttft_percentile is not None
else _average_latency(item_ttft_latency if use_ttft else item_latency)
)
# -------------- #
# Debugging Logic
@ -443,7 +456,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
# this helps a user to debug why the router picked a specfic deployment #
_deployment_api_base = _deployment.get("litellm_params", {}).get("api_base", "")
if _deployment_api_base is not None:
_latency_per_deployment[_deployment_api_base] = average_latency
_latency_per_deployment[_deployment_api_base] = selected_latency
# -------------- #
# End of Debugging Logic
# -------------- #
@ -453,7 +466,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
): # if user passed in tpm / rpm in the model_list
continue
else:
potential_deployments.append((_deployment, average_latency))
potential_deployments.append((_deployment, selected_latency))
if len(potential_deployments) == 0:
return None

6
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -50,6 +50,12 @@ def prepare_response_for_header_attachment(response: object) -> object | None:
return response
def response_has_hidden_params(response: object) -> bool:
if isinstance(response, dict):
return "_hidden_params" in response
return hasattr(response, "_hidden_params")
def ensure_response_additional_headers(response: object) -> dict[str, object]:
hidden_params: Final = get_hidden_params_dict(response, create=isinstance(response, dict))
_write_hidden_params(response, hidden_params)

14
litellm/router_utils/common_utils.py
View file

@ -26,6 +26,18 @@ def _is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool
return getattr(user_api_key_auth, "user_role", None) == "proxy_admin"
def get_request_team_id(request_kwargs: Mapping[str, object] | None) -> str | None:
"""The caller's team id, from whichever metadata bucket this surface writes to."""
if request_kwargs is None:
return None
for bucket_name in ("metadata", "litellm_metadata"):
bucket = request_kwargs.get(bucket_name)
team_id = bucket.get("user_api_key_team_id") if isinstance(bucket, Mapping) else None
if isinstance(team_id, str) and team_id:
return team_id
return None
def resolve_model_group_alias(model_group_alias: object, model: str) -> str | None:
"""
Resolve ``model`` through a ``model_group_alias`` map.
@ -110,7 +122,7 @@ def filter_team_based_models(
metadata: Final = request_kwargs.get("metadata") or {}
litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {}
request_team_id: Final = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id")
request_team_id: Final = get_request_team_id(request_kwargs)
if request_team_id is None and _is_proxy_admin_request(request_kwargs) and isinstance(healthy_deployments, list):
requested_model: Final = (
request_kwargs.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group")

73
litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py
View file

@ -37,13 +37,13 @@ Safe to enable globally:
"""
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
import httpx
from litellm._logging import verbose_router_logger
from litellm.exceptions import (
BadRequestError,
RateLimitError,
ServiceUnavailableError,
)
@ -158,6 +158,23 @@ class EncryptedContentAffinityCheck(CustomLogger):
return deployment
return None
@staticmethod
def _request_team_id(request_kwargs: Mapping[str, object]) -> str | None:
containers: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata"))
team_ids: Final = (c.get("user_api_key_team_id") for c in containers if isinstance(c, Mapping))
return next((tid for tid in team_ids if isinstance(tid, str)), None)
def _routed_group_candidate_model_ids(self, request_kwargs: Mapping[str, object], model: str) -> frozenset[str]:
"""
Deployment ids that could serve this turn's routed ``model``, as the router
resolves a route (model_group_alias / routing group / model_name / team /
pattern). Delegates to the router so the full precedence is not re-derived here
and no deployment ids are written into request kwargs bound for the provider.
"""
if self.router is None:
return frozenset()
return self.router.get_candidate_model_ids_for_route(model=model, team_id=self._request_team_id(request_kwargs))
@staticmethod
def _encryption_boundary_key(
litellm_params: object,
@ -225,10 +242,14 @@ class EncryptedContentAffinityCheck(CustomLogger):
"""
If the request ``input`` contains litellm-encoded item IDs, decode the
embedded ``model_id`` and pin the request to that deployment. Raises
``RateLimitError`` / ``ServiceUnavailableError`` / ``BadRequestError``
when the originating deployment is unavailable and no encryption-boundary
peer exists, rather than dispatching a doomed request to a non-peer
deployment. The 429/503 split mirrors the originating cooldown's status:
``RateLimitError`` / ``ServiceUnavailableError`` when the originating
deployment is a member of the routed model group but currently unavailable
and no encryption-boundary peer exists, rather than dispatching a doomed
request to a non-peer deployment. When the origin is not a member of the
routed group (an auto-router tier change, a model switch with no peer, a
removed deployment, or an unknown/forged marker), the encrypted reasoning is
stripped and the request dispatches with its readable history instead. The
429/503 split mirrors the originating cooldown's status:
a 429-induced cooldown surfaces as 429 (with ``Retry-After`` set to the
remaining cooldown window) so OpenAI-compatible clients back off and
retry after the deployment is eligible again.
@ -285,12 +306,34 @@ class EncryptedContentAffinityCheck(CustomLogger):
request_kwargs["_encrypted_content_affinity_pinned"] = True
return boundary_matches
# Dispatching to a non-peer would guarantee an upstream
# `invalid_encrypted_content` 400, so fail fast with a clearer error.
# The origin cannot serve this turn's routed group and no peer shares the boundary, so its
# encrypted reasoning can never decrypt here. Strip it, keep the readable history, and dispatch
# to the routed group instead of failing. Membership is tested by deployment id against the set
# the router actually resolved for this route, not by model-group name, so an alias, a
# provider-qualified spelling, a team-public name, or a pattern route of the same group is not
# mistaken for a tier change. An unknown origin (a removed deployment, or a forged marker) is
# treated the same as a cross-group one, which also denies an authenticated caller a
# deployment-id existence oracle: a real cross-group id and a nonexistent id both strip and
# dispatch rather than returning distinguishable responses. Only a genuine same-group member
# that is currently unavailable falls through to the fail-fast, preserving the cooldown contract.
routed_group_model_ids: Final = (
self._routed_group_candidate_model_ids(request_kwargs, model) if originating is not None else frozenset()
)
if str(model_id) not in routed_group_model_ids:
verbose_router_logger.debug(
"EncryptedContentAffinityCheck: model_id=%s is not a candidate for the routed group %s; "
"forwarding without its encrypted reasoning",
model_id,
model,
)
ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input)
return typed_healthy_deployments
# The origin is a member of the routed group but currently unavailable (cooled down); fail fast
# rather than dispatching to a non-peer, which would guarantee an upstream 400.
raise await self._unavailable_origin_error(
model=model,
model_id=model_id,
originating=originating,
parent_otel_span=parent_otel_span,
)
@ -298,25 +341,11 @@ class EncryptedContentAffinityCheck(CustomLogger):
self,
model: str,
model_id: str,
originating: Deployment | None,
parent_otel_span: Span | None,
) -> Exception:
# Public error messages intentionally omit the originating ``model_id`` so
# an authenticated caller forging encrypted-content markers cannot use the
# error surface to enumerate which deployment IDs exist on this router.
if originating is None:
return BadRequestError(
message=(
"The deployment that produced this encrypted_content is no "
"longer configured on this router, and no deployment on the "
"same encryption boundary is available. Re-issue the request "
"without the stale encrypted_content items, or restore the "
"originating deployment."
),
model=model,
llm_provider="",
)
cooldown: Final = await self._get_origin_cooldown(model_id=model_id, parent_otel_span=parent_otel_span)
if cooldown is not None and str(cooldown.get("status_code")) == "429":

9
litellm/types/llms/databricks.py
View file

@ -2,6 +2,7 @@ from typing import Any, Literal
from pydantic import BaseModel
from typing_extensions import (
ReadOnly,
Required,
TypedDict,
)
@ -57,6 +58,14 @@ class DatabricksMessage(TypedDict, total=False):
role: Required[str]
content: Required[AllDatabricksContentValues]
tool_calls: list[DatabricksTool] | None
reasoning_content: ReadOnly[str | None]
reasoning: ReadOnly[str | None]
class DatabricksDelta(TypedDict, total=False):
role: ReadOnly[str]
content: ReadOnly[AllDatabricksContentValues | None]
reasoning_content: ReadOnly[str | None]
class DatabricksChoice(TypedDict, total=False):

1
litellm/types/router.py
View file

@ -525,6 +525,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
input_cost_per_second: float | None
output_cost_per_second: float | None
output_cost_per_second_480p: ReadOnly[float | None]
output_cost_per_second_720p: ReadOnly[float | None]
output_cost_per_second_1080p: float | None
output_cost_per_second_4k: ReadOnly[float | None]
num_retries: int | None

2
litellm/types/utils.py
View file

@ -318,6 +318,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
float | None
) # video_generation tier: key output_cost_per_second_<resolution> (e.g. 1080p, 720p)
output_cost_per_second_480p: ReadOnly[float | None]
output_cost_per_second_720p: ReadOnly[float | None]
output_cost_per_second_4k: ReadOnly[float | None]
ocr_cost_per_page: float | None # for OCR models
ocr_cost_per_credit: float | None # for OCR models priced by credit
@ -3522,6 +3523,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
output_cost_per_second: float | None = None
output_cost_per_second_1080p: float | None = None
output_cost_per_second_480p: float | None = None
output_cost_per_second_720p: float | None = None
output_cost_per_second_4k: float | None = None
input_cost_per_pixel: float | None = None
output_cost_per_pixel: float | None = None

7
litellm/utils.py
View file

@ -3079,7 +3079,7 @@ def register_model(
# Convert stringified numbers to appropriate numeric types
loaded_model_cost = model_cost
elif isinstance(model_cost, str):
loaded_model_cost = litellm.get_model_cost_map(url=model_cost)
loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1)
if persist_across_reloads:
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost
@ -5913,6 +5913,7 @@ def _get_model_info_helper(
output_cost_per_second=_model_info.get("output_cost_per_second", None),
output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None),
output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None),
output_cost_per_second_720p=_model_info.get("output_cost_per_second_720p", None),
output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None),
output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None),
output_cost_per_image=_model_info.get("output_cost_per_image", None),
@ -9241,6 +9242,10 @@ class ProviderConfigManager:
from litellm.llms.openai.image_edit import get_openai_image_edit_config
return get_openai_image_edit_config(model=model)
elif LlmProviders.HOSTED_VLLM == provider:
from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config
return get_hosted_vllm_image_edit_config(model=model)
elif LlmProviders.AZURE == provider:
from litellm.llms.azure.image_edit.transformation import (
AzureImageEditConfig,

845
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

4
model_prices_and_context_window.schema.json
View file

@ -478,6 +478,10 @@
"type": "number",
"minimum": 0
},
"output_cost_per_second_720p": {
"type": "number",
"minimum": 0
},
"output_cost_per_token": {
"type": "number",
"minimum": 0,

2
schema.prisma
View file

@ -492,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])

1
tests/code_coverage_tests/router_code_coverage.py
View file

@ -87,6 +87,7 @@ ignored_function_names = [
"_delete_claude_code_session_router_binding", # Tested through Redis cleanup failure in test_router.py
"_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py
"_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py
"_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name)
]

30
tests/e2e/batches/COVERAGE.md
View file

@ -120,6 +120,36 @@ create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin:
nested managed ids round-trip retrieve. This self-chaining only needs the proxy to
reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage.
## Cleanup
Batch teardown cancels active batches before deleting their input files and keys.
Raw file IDs from both `model_param` and `provider_fallback` uploads use the upload
provider when deleted. Model-encoded and managed file IDs route themselves
File deletion and batch cancellation check their responses and retry transient
failures up to three times. Teardown attempts every registered cleanup before
reporting failures as test errors. Already deleted files and batches that are
terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes
before input deletion: the ten-minute provider window plus a propagation margin.
Accepted cancellation may still report validating or in_progress while the provider
updates its state. Raw and model-encoded batches are polled until cancelling or
terminal before input deletion. OpenAI and Azure lifecycle cleanup also deletes
output and error files returned by terminal batches. Bedrock deletion uses a signed S3 DELETE
restricted to the configured storage buckets and managed file prefixes. The low-RPM
test submits with its restricted key and cleans up with the test administrator key
Managed deletion forwards the deployment's trusted bucket configuration and returns
the requested managed file ID even when stored output metadata carries a provider ID
Azure input uploads request `expires_after` anchored to `created_at` with
`seconds=1209600`, and the lifecycle tests check the returned expiry. This is a
fallback for interrupted runs: immediate deletion remains the normal cleanup.
Azure's minimum supported native expiry is 14 days, so a three-day expiry cannot
be requested through its Files API
The Azure entry in `files_settings` must use `api_version: 2025-04-01-preview`
for raw uploads to honor expiry, matching the batch deployment's API version
## Terminal state + cost write-back (cross-run marker baton)
The 24h completion window rules out submit-and-wait inside one run, so

140
tests/e2e/batches/batch_cleanup.py Normal file
View file

@ -0,0 +1,140 @@
from builtins import ExceptionGroup
from collections.abc import Callable
from itertools import count
from time import monotonic, sleep
from typing import Final, Protocol
from batch_client import BatchObject, FileDeleteResponse
from capabilities import is_managed_id
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
from pydantic import BaseModel
CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0)
BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"})
BATCH_PENDING_STATUSES: Final = frozenset({"validating", "in_progress", "finalizing", "cancelling"})
BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0
BATCH_CANCEL_POLL_SECONDS: Final = 10.0
class BatchCleanupClient(Protocol):
def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]: ...
def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...
def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]: ...
def cleanup_result[R: BaseModel](
action: Callable[[], Result[R]], *, wait: Callable[[float], None] = sleep
) -> Result[R]:
for delay, result in ((delay, action()) for delay in CLEANUP_DELAYS):
match result:
case NetworkError() | RateLimitedError():
wait(delay)
case UnknownApiError(status_code=code) if code in {408, 429, 500, 502, 503, 504}:
wait(delay)
case _:
return result
return action()
def _require_cleanup_success[R: BaseModel](result: Result[R], operation: str) -> R:
match result:
case Success(data=data):
return data
case UnknownApiError(status_code=code):
raise AssertionError(f"{operation} failed: HTTP {code}")
case _:
raise AssertionError(f"{operation} failed: {result.kind}")
def cleanup_file(client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None = None) -> None:
result: Final = cleanup_result(lambda: client.delete_file(file_id, key=key, provider=provider))
if isinstance(result, UnknownApiError) and result.status_code == 404:
return
deleted: Final = _require_cleanup_success(result, f"Delete file {file_id}")
assert deleted.deleted is True or (
deleted.deleted is None and is_managed_id(file_id) and deleted.id == file_id and deleted.object == "file"
), f"Delete file {file_id} did not confirm deletion"
def cleanup_batch(
client: BatchCleanupClient,
batch_id: str,
*,
key: str,
provider: str | None = None,
delete_output_files: bool = False,
wait: Callable[[float], None] = sleep,
clock: Callable[[], float] = monotonic,
) -> None:
needs_terminal_state: Final = is_managed_id(batch_id)
fetched: Final = _require_cleanup_success(
cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)),
f"Retrieve batch {batch_id} for cleanup",
)
if fetched.status in BATCH_TERMINAL_STATUSES:
if delete_output_files:
_cleanup_batch_outputs(client, fetched, key=key, provider=provider)
return
if fetched.status == "cancelling" and not needs_terminal_state:
return
result: Final = (
Success(status_code=200, data=fetched)
if fetched.status == "cancelling"
else cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider))
)
conflicted: Final = isinstance(result, UnknownApiError) and result.status_code in {400, 409}
if not conflicted:
cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}")
assert cancelled.status in BATCH_TERMINAL_STATUSES | BATCH_PENDING_STATUSES, (
f"Cancel batch {batch_id} left status {cancelled.status}"
)
if cancelled.status in BATCH_TERMINAL_STATUSES:
if delete_output_files:
_cleanup_batch_outputs(client, cancelled, key=key, provider=provider)
return
if cancelled.status == "cancelling" and not needs_terminal_state:
return
deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS
for current in (
_require_cleanup_success(
cleanup_result(lambda: client.retrieve_batch(batch_id, key=key, provider=provider)),
f"Retrieve batch {batch_id} after cancellation",
)
for _ in count()
):
if current.status in BATCH_TERMINAL_STATUSES:
if delete_output_files:
_cleanup_batch_outputs(client, current, key=key, provider=provider)
return
assert current.status in ({"cancelling"} if conflicted else BATCH_PENDING_STATUSES), (
f"Cancel batch {batch_id} left status {current.status}"
)
if current.status == "cancelling" and not needs_terminal_state:
return
assert clock() < deadline, (
f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s"
)
wait(BATCH_CANCEL_POLL_SECONDS)
def _cleanup_batch_outputs(client: BatchCleanupClient, batch: BatchObject, *, key: str, provider: str | None) -> None:
errors: Final = tuple(
error
for file_id in dict.fromkeys((batch.output_file_id, batch.error_file_id))
if file_id is not None and file_id != batch.input_file_id
if (error := _output_cleanup_error(client, file_id, key=key, provider=provider)) is not None
)
if errors:
raise ExceptionGroup(f"Batch {batch.id} output cleanup failed", errors)
def _output_cleanup_error(
client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None
) -> Exception | None:
try:
cleanup_file(client, file_id, key=key, provider=provider)
except Exception as error:
return error
return None

18
tests/e2e/batches/batch_client.py
View file

@ -13,8 +13,9 @@ co-located here because only this suite uses them.
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, Literal
from pydantic import BaseModel
from pydantic import BaseModel, Field
from proxy_client import ProxyClient
from e2e_http import (
@ -27,6 +28,18 @@ from e2e_http import (
from models import LiteLLMParamsBody
UPLOAD_FILENAME = "batch_input.jsonl"
AZURE_FILE_EXPIRY_SECONDS: Final = 14 * 24 * 60 * 60
class ExpiringFileUploadForm(FileUploadForm):
expires_after_anchor: Literal["created_at"] = Field(default="created_at", alias="expires_after[anchor]")
expires_after_seconds: int = Field(default=AZURE_FILE_EXPIRY_SECONDS, alias="expires_after[seconds]")
def batch_upload_form(provider: str, *, target_model_names: str | None = None) -> FileUploadForm:
if provider == "azure":
return ExpiringFileUploadForm(target_model_names=target_model_names)
return FileUploadForm(target_model_names=target_model_names)
class FileObject(BaseModel):
@ -37,6 +50,7 @@ class FileObject(BaseModel):
bytes: int | None = None
status: str | None = None
created_at: int | None = None
expires_at: int | None = None
class FileList(BaseModel):
@ -85,7 +99,7 @@ class BatchList(BaseModel):
class FileDeleteResponse(BaseModel):
id: str
object: str | None = None
deleted: bool
deleted: bool | None = None
class BatchCreateBody(BaseModel):

4
tests/e2e/batches/capabilities.py
View file

@ -108,6 +108,10 @@ class Capability:
def id(self) -> str:
return f"{self.provider}-{self.scenario}"
@property
def file_provider(self) -> str | None:
return self.provider if self.scenario in {"model_param", "provider_fallback"} else None
@property
def jsonl_model(self) -> str:
# Always the provider deployment name. Unified routes via

10
tests/e2e/batches/conftest.py
View file

@ -13,7 +13,7 @@ the proxy config.
from __future__ import annotations
import os
from typing import Iterator
from typing import Final, Iterator
import pytest
@ -21,6 +21,7 @@ from batch_client import BatchClient, build_client
from capabilities import PROVIDERS
from e2e_config import MANAGED_FILES_OPT_IN_ENV
from e2e_http import NoBody
from lifecycle import ResourceManager
from proxy_client import ProxyClient
@ -52,6 +53,13 @@ def client(proxy: ProxyClient) -> BatchClient:
return build_client(proxy)
@pytest.fixture
def resources(client: BatchClient) -> Iterator[ResourceManager]:
manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True)
yield manager
manager.teardown()
@pytest.fixture(scope="session")
def batch_deployments(client: BatchClient) -> Iterator[None]:
probe = client.proxy.probe("/health/liveliness", params=NoBody())

313
tests/e2e/batches/test_batch_cleanup.py Normal file
View file

@ -0,0 +1,313 @@
from builtins import ExceptionGroup
from collections.abc import Callable
from typing import Final
from unittest.mock import Mock, call
import pytest
from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result
from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form
from capabilities import CAPABILITIES, Capability
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
from lifecycle import ResourceManager
from models import KeyGenerateBody
MANAGED_FILE_ID: Final = "bGl0ZWxsbV9wcm94eTtmaWxlLTE="
MANAGED_BATCH_ID: Final = "bGl0ZWxsbV9wcm94eTtiYXRjaC0x"
class ExpectedCalls[T]:
def __init__(self, values: tuple[T, ...]) -> None:
self.values: Final = values
self.recorder: Final = Mock()
def __call__(self, value: T) -> None:
self.recorder(value)
def assert_done(self) -> None:
assert tuple(self.recorder.call_args_list) == tuple(call(value) for value in self.values)
class CleanupClient:
def __init__(
self,
*,
calls: ExpectedCalls[str],
files: tuple[Result[FileDeleteResponse], ...] = (),
batches: tuple[Result[BatchObject], ...] = (),
cancellations: tuple[Result[BatchObject], ...] = (),
) -> None:
self.calls: Final = calls
self.file_response: Final[Callable[[], Result[FileDeleteResponse]]] = Mock(side_effect=files)
self.batch_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=batches)
self.cancel_response: Final[Callable[[], Result[BatchObject]]] = Mock(side_effect=cancellations)
def delete_file(self, file_id: str, *, key: str, provider: str | None = None) -> Result[FileDeleteResponse]:
self.calls(f"delete {provider} {file_id}")
return self.file_response()
def retrieve_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
self.calls(f"retrieve {provider} {batch_id}")
return self.batch_response()
def cancel_batch(self, batch_id: str, *, key: str, provider: str | None = None) -> Result[BatchObject]:
self.calls(f"cancel {provider} {batch_id}")
return self.cancel_response()
def generate_key(self, body: KeyGenerateBody) -> str:
return "test-key"
def delete_key(self, key: str) -> None:
self.calls(f"delete key {key}")
def delete_customers(self, user_ids: list[str]) -> None:
self.calls(f"delete customers {user_ids}")
def batch(status: str) -> Success[BatchObject]:
return Success(status_code=200, data=BatchObject(id="batch-1", status=status))
def deleted_file(*, deleted: bool = True) -> Success[FileDeleteResponse]:
return Success(status_code=200, data=FileDeleteResponse(id="file-1", deleted=deleted))
class TestFileCleanup:
def test_managed_delete_accepts_the_deleted_file_object(self) -> None:
response: Final = Success(
status_code=200, data=FileDeleteResponse.model_validate({"id": MANAGED_FILE_ID, "object": "file"})
)
client: Final = CleanupClient(calls=ExpectedCalls((f"delete None {MANAGED_FILE_ID}",)), files=(response,))
cleanup_file(client, MANAGED_FILE_ID, key="test-key")
client.calls.assert_done()
@pytest.mark.parametrize("file_id", ["file-1", MANAGED_FILE_ID])
def test_a_success_status_without_a_deletion_confirmation_is_rejected(self, file_id: str) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls((f"delete None {file_id}",)),
files=(Success(status_code=200, data=FileDeleteResponse(id=file_id)),),
)
with pytest.raises(AssertionError, match="did not confirm deletion"):
cleanup_file(client, file_id, key="test-key")
client.calls.assert_done()
@pytest.mark.parametrize("cap", CAPABILITIES, ids=[cap.id for cap in CAPABILITIES])
def test_deletes_raw_files_through_the_upload_provider(self, cap: Capability) -> None:
expected_provider: Final = cap.provider if cap.scenario in {"model_param", "provider_fallback"} else None
client: Final = CleanupClient(
calls=ExpectedCalls((f"delete {expected_provider} file-1",)), files=(deleted_file(),)
)
cleanup_file(client, "file-1", key="test-key", provider=cap.file_provider)
client.calls.assert_done()
def test_failed_delete_is_reported_after_remaining_resources_are_cleaned(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete azure file-1", "delete key test-key")),
files=(UnknownApiError(status_code=403, body="secret response"),),
)
manager: Final = ResourceManager(client=client, strict_cleanup=True)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="azure"))
with pytest.raises(ExceptionGroup) as caught:
manager.teardown()
client.calls.assert_done()
assert len(caught.value.exceptions) == 1
assert str(caught.value.exceptions[0]) == "Delete file file-1 failed: HTTP 403"
def test_success_response_must_confirm_deletion(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete None file-1",)), files=(deleted_file(deleted=False),)
)
with pytest.raises(AssertionError, match="did not confirm deletion"):
cleanup_file(client, "file-1", key="test-key")
client.calls.assert_done()
def test_cleanup_is_idempotent_when_file_is_already_deleted(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete azure file-1",)),
files=(UnknownApiError(status_code=404, body="missing"),),
)
cleanup_file(client, "file-1", key="test-key", provider="azure")
client.calls.assert_done()
def test_default_resource_cleanup_keeps_existing_best_effort_behavior(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("delete None file-1", "delete key test-key")),
files=(UnknownApiError(status_code=403, body="forbidden"),),
)
manager: Final = ResourceManager(client=client)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key))
manager.teardown()
client.calls.assert_done()
class TestCleanupRetries:
@pytest.mark.parametrize(
"failure",
[NetworkError(message="offline"), RateLimitedError(), UnknownApiError(status_code=503, body="unavailable")],
)
def test_transient_error_retries_and_returns_success(self, failure: Result[FileDeleteResponse]) -> None:
responses: Final = (failure, deleted_file())
outcomes: Final = Mock(side_effect=responses)
delays: Final = ExpectedCalls((1.0,))
result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays)
assert isinstance(result, Success) and result.data.deleted
delays.assert_done()
def test_persistent_error_has_bounded_retries(self) -> None:
failure: Final = UnknownApiError(status_code=503, body="unavailable")
outcomes: Final = Mock(return_value=failure)
delays: Final = ExpectedCalls(CLEANUP_DELAYS)
result: Final[Result[FileDeleteResponse]] = cleanup_result(outcomes, wait=delays)
assert result is failure
delays.assert_done()
assert outcomes.call_count == len(CLEANUP_DELAYS) + 1
def test_permanent_error_is_not_retried(self) -> None:
failure: Final = UnknownApiError(status_code=403, body="forbidden")
responses: Final = (failure, deleted_file())
outcomes: Final = Mock(side_effect=responses)
delays: Final = ExpectedCalls[float](())
assert cleanup_result(outcomes, wait=delays) is failure
delays.assert_done()
assert outcomes.call_count == 1
class TestBatchCancellation:
def test_cancelling_batch_is_polled_until_terminal_without_cancelling_again(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls((f"retrieve None {MANAGED_BATCH_ID}",) * 3),
batches=(batch("cancelling"), batch("cancelling"), batch("cancelled")),
)
delays: Final = ExpectedCalls((10.0,))
cleanup_batch(client, MANAGED_BATCH_ID, key="test-key", wait=delays)
client.calls.assert_done()
delays.assert_done()
def test_cancellation_timeout_is_reported_but_file_and_key_cleanup_still_run(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(
(
f"retrieve None {MANAGED_BATCH_ID}",
f"retrieve None {MANAGED_BATCH_ID}",
"delete None file-1",
"delete key test-key",
)
),
batches=(batch("cancelling"), batch("cancelling")),
files=(deleted_file(),),
)
times: Final = (0.0, BATCH_CANCEL_TIMEOUT_SECONDS)
ticks: Final[Callable[[], float]] = Mock(side_effect=times)
manager: Final = ResourceManager(client=client, strict_cleanup=True)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key))
manager.defer(lambda: cleanup_batch(client, MANAGED_BATCH_ID, key=key, clock=ticks))
with pytest.raises(ExceptionGroup) as caught:
manager.teardown()
assert "cancellation did not finish" in str(caught.value.exceptions[0])
client.calls.assert_done()
@pytest.mark.parametrize("status", ["completed", "failed", "expired", "cancelled"])
def test_inactive_batch_needs_no_cancellation(self, status: str) -> None:
client: Final = CleanupClient(calls=ExpectedCalls(("retrieve None batch-1",)), batches=(batch(status),))
cleanup_batch(client, "batch-1", key="test-key")
client.calls.assert_done()
def test_active_batch_is_cancelled_through_its_provider(self) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("retrieve azure batch-1", "cancel azure batch-1")),
batches=(batch("in_progress"), batch("cancelled")),
cancellations=(batch("cancelling"),),
)
cleanup_batch(client, "batch-1", key="test-key", provider="azure")
client.calls.assert_done()
@pytest.mark.parametrize("batch_id", ["batch-1", MANAGED_BATCH_ID])
@pytest.mark.parametrize("pending_status", ["validating", "in_progress"])
def test_accepted_cancellation_waits_through_stale_provider_status(
self, batch_id: str, pending_status: str
) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(
(
f"retrieve vertex_ai {batch_id}",
f"cancel vertex_ai {batch_id}",
f"retrieve vertex_ai {batch_id}",
f"retrieve vertex_ai {batch_id}",
f"retrieve vertex_ai {batch_id}",
"delete vertex_ai file-1",
"delete key test-key",
)
),
batches=(batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled")),
cancellations=(batch(pending_status),),
files=(deleted_file(),),
)
delays: Final = ExpectedCalls((10.0, 10.0))
manager: Final = ResourceManager(client=client, strict_cleanup=True)
key: Final = manager.key()
manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai"))
manager.defer(lambda: cleanup_batch(client, batch_id, key=key, provider="vertex_ai", wait=delays))
manager.teardown()
client.calls.assert_done()
delays.assert_done()
@pytest.mark.parametrize("output_delete_fails", [False, True])
def test_batch_that_completed_before_cleanup_deletes_output_and_error_files(
self, output_delete_fails: bool
) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error")),
batches=(
Success(
status_code=200,
data=BatchObject(
id="batch-1",
status="completed",
input_file_id="file-input",
output_file_id="file-output",
error_file_id="file-error",
),
),
),
files=(
UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(),
deleted_file(),
),
)
if output_delete_fails:
with pytest.raises(ExceptionGroup, match="output cleanup failed"):
cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True)
else:
cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True)
client.calls.assert_done()
@pytest.mark.parametrize("status", ["completed", "in_progress"])
def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None:
client: Final = CleanupClient(
calls=ExpectedCalls(("retrieve None batch-1", "cancel None batch-1", "retrieve None batch-1")),
batches=(batch("in_progress"), batch(status)),
cancellations=(UnknownApiError(status_code=409, body="conflict"),),
)
if status == "completed":
cleanup_batch(client, "batch-1", key="test-key")
else:
with pytest.raises(AssertionError, match="Cancel batch batch-1 left status in_progress"):
cleanup_batch(client, "batch-1", key="test-key")
client.calls.assert_done()
class TestAzureFileExpiry:
def test_azure_form_serializes_native_expiry_for_the_proxy(self) -> None:
form: Final = batch_upload_form("azure", target_model_names="azure-test")
assert form.model_dump(by_alias=True, exclude_none=True) == {
"purpose": "batch",
"target_model_names": "azure-test",
"expires_after[anchor]": "created_at",
"expires_after[seconds]": AZURE_FILE_EXPIRY_SECONDS,
}
@pytest.mark.parametrize("provider", ["openai", "vertex_ai", "bedrock"])
def test_other_providers_keep_their_existing_upload_fields(self, provider: str) -> None:
assert batch_upload_form(provider).model_dump(by_alias=True, exclude_none=True) == {"purpose": "batch"}

Some files were not shown because too many files have changed in this diff Show more