Merge stable/1.104.x (v1.104.6 prep) into patch_d351e54e

Keeps stable/1.104.x's 1.104.6 version in pyproject.toml and uv.lock, so this branch no longer bumps the version. In litellm_logging.py both async success helpers stay, the base's enqueued flag and the claim that #44508 adds, matching main. The two test files keep both sides' new tests
This commit is contained in:
Yuneng Jiang 2026-10-20 23:54:09 -07:00
commit 2b2e960848
No known key found for this signature in database
144 changed files with 35645 additions and 24214 deletions

View file

@ -91,6 +91,17 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
},
"additionalProperties": False,
},
"search_free_requests": {
"type": "object",
"description": "Free Google Search grounding requests per period, shared across models in the same shared_group.",
"required": ["count", "period", "shared_group"],
"additionalProperties": False,
"properties": {
"count": {"type": "integer", "minimum": 0},
"period": {"type": "string", "enum": ["day", "month"]},
"shared_group": {"type": "string", "pattern": "\\S"},
},
},
"guardrail_cost_per_unit": {
"type": "object",
"description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).",
@ -111,8 +122,9 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
"properties": {
"model": NON_BLANK_STRING,
"region": NON_BLANK_STRING,
"litellm_provider": {"type": "string", "enum": ["bedrock", "bedrock_mantle"]},
},
"description": "Model that counts tokens for this one, and the AWS region its count API runs in.",
"description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in.",
},
}

View file

@ -2,7 +2,9 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""
from collections.abc import Sequence
import json
import re
from collections.abc import Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
@ -27,7 +29,12 @@ if TYPE_CHECKING:
from litellm.types.utils import LiteLLMBatch
def _error_names_resource(error: Exception, resource_id: str) -> bool:
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
@ -44,6 +51,12 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
class _ManagedObjectRow(Protocol):
@property
def created_at(self) -> datetime: ...
@property
def updated_at(self) -> datetime: ...
@property
def id(self) -> str: ...
@ -98,6 +111,11 @@ class CheckBatchCost:
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
self.batch_processed_support_confirmed: bool = False
self._poll_cursor: tuple[datetime, str] | None = None
self._poll_before: datetime = datetime.now(timezone.utc)
self._arrival_cursor: tuple[datetime, str] | None = None
self._arrival_before: datetime = self._poll_before
self._poll_arrivals: bool = False
@staticmethod
def _is_missing_batch_processed_column_error(err: Exception) -> bool:
@ -269,29 +287,9 @@ class CheckBatchCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
if not self._has_batch_processed_column:
return
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"in": ["complete", "completed"]},
"created_at": {"lt": cutoff},
},
data={"batch_processed": True},
)
if retired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await _managed_object_table(self.prisma_client).find_many(
return await self._find_jobs(
where={
"file_purpose": "batch",
"status": {
@ -305,10 +303,59 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[_ManagedObjectRow]:
arrivals: Final = self._poll_arrivals
page: Final = await self._find_job_page(where, arrivals=arrivals)
self._poll_arrivals = not arrivals if page else arrivals
if page:
return page
return await self._find_job_page(where, arrivals=not arrivals)
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[_ManagedObjectRow]:
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
if cursor is None:
if arrivals:
self._arrival_before = datetime.now(timezone.utc)
else:
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
before: Final = self._arrival_before if arrivals else self._poll_before
page: Final = await _managed_object_table(self.prisma_client).find_many(
where={
**where,
"created_at": {
"lte": before,
**({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}),
},
**(
{
"OR": [
{"created_at": {"gt": cursor[0]}},
{"created_at": cursor[0], "id": {"gt": cursor[1]}},
]
}
if cursor is not None
else {}
),
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order=[{"created_at": "asc"}, {"id": "asc"}],
)
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
if arrivals:
self._arrival_cursor = next_cursor
else:
self._poll_cursor = next_cursor
return page
def _unprocessed_snapshot(self, job: _ManagedObjectRow) -> Mapping[str, object]:
return {
"id": job.id,
"updated_at": job.updated_at,
**({"batch_processed": False} if self._has_batch_processed_column else {}),
}
async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
@ -316,27 +363,23 @@ class CheckBatchCost:
once enough such rows accumulate no newer batch is ever reached. Older schemas
without batch_processed can only be excluded through the status filter.
"""
data: Final = (
{"batch_processed": True}
if self._has_batch_processed_column
else {"status": "stale_expired"}
)
data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"}
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
)
verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}")
return
if retired == 0:
return
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
"so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "_ManagedObjectRow") -> bool:
async def _claim_job_for_costing(self, job: "_ManagedObjectRow", response: "LiteLLMBatch") -> bool:
"""
Atomically flip batch_processed from false to true, returning whether this pod won
the row. Every pod and uvicorn worker schedules its own poller against the shared
@ -350,10 +393,16 @@ class CheckBatchCost:
"""
if not self._has_batch_processed_column:
return True
from litellm_enterprise.proxy.hooks.managed_files import postgres_storable_json
try:
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "batch_processed": False},
data={"batch_processed": True},
data={
"batch_processed": True,
"status": response.status if response.status != "completed" else "complete",
"file_object": postgres_storable_json(response),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
@ -362,7 +411,9 @@ class CheckBatchCost:
return False
return claimed > 0
async def _release_job_claim(self, job: "_ManagedObjectRow") -> None:
async def _release_job_claim(
self, job: "_ManagedObjectRow", original_status: str | None, original_file_object: object
) -> None:
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
Safe to match on batch_processed=True: while this poller is active the retrieve
@ -374,7 +425,11 @@ class CheckBatchCost:
try:
await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "batch_processed": True},
data={"batch_processed": False},
data={
"batch_processed": False,
"status": original_status,
"file_object": json.dumps(original_file_object),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
@ -408,13 +463,13 @@ class CheckBatchCost:
from litellm.exceptions import NotFoundError
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id)
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
404 must not retire a potentially recoverable row."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
@ -429,7 +484,7 @@ class CheckBatchCost:
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
async def _finalize_unbilled_terminal_job(
self, job: "_ManagedObjectRow", response: "LiteLLMBatch"
@ -451,15 +506,17 @@ class CheckBatchCost:
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data: Final[dict] = {
update_data: Final[dict[str, object]] = {
"status": response.status,
"file_object": response.model_dump_json(),
**({"batch_processed": True} if self._has_batch_processed_column else {}),
}
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
finalized: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=update_data,
)
if finalized == 0:
return
verbose_proxy_logger.info(
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
)
@ -671,8 +728,6 @@ class CheckBatchCost:
@staticmethod
def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
import json
from litellm.types.utils import LiteLLMBatch
file_object = job.file_object
@ -712,7 +767,11 @@ class CheckBatchCost:
from litellm.files.main import afile_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
from litellm.litellm_core_utils.litellm_logging import (
deployment_pricing_model_info,
mask_api_base_credentials,
resolve_vertex_location_for_cost,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -849,6 +908,12 @@ class CheckBatchCost:
custom_llm_provider=batch_file_provider,
model_name=model_name,
model_info=deployment_model_info,
vertex_location=resolve_vertex_location_for_cost(
custom_llm_provider=llm_provider,
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
optional_params=None,
model=model_name,
),
)
error_file_failed_requests: Final = await count_error_file_failed_requests(
response,
@ -866,8 +931,9 @@ class CheckBatchCost:
failed_requests=output_file_result.failed_requests + error_file_failed_requests,
)
)
logging_model: Final = batch_result.models[0] if batch_result.models else litellm_model_name or model_id
logging_obj = LiteLLMLogging(
model=batch_result.models[0],
model=logging_model,
messages=[{"role": "user", "content": "<retrieve_batch>"}],
stream=False,
call_type="aretrieve_batch",
@ -898,27 +964,35 @@ class CheckBatchCost:
custom_llm_provider=str(llm_provider) if llm_provider else None,
)
if not await self._claim_job_for_costing(job):
original_status: Final = job.status
original_file_object: Final = job.file_object
if not await self._claim_job_for_costing(job, response.model_copy(update={"usage": batch_result.usage})):
verbose_proxy_logger.info(
f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
"in this window, so its cost is already being tracked there"
)
return None
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job)
raise
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
with track_batch_accounting(logging_obj.litellm_call_id) as accounting:
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
finally:
if accounting.retryable_failure:
await self._release_job_claim(job, original_status, original_file_object)
if accounting.retryable_failure:
return None
# Record batch duration (completed_at - created_at)
if prom_logger and response.completed_at and response.created_at:
@ -966,7 +1040,7 @@ class CheckBatchCost:
# endpoint may transition a batch to "complete" before
# CheckBatchCost runs. The batch_processed=False filter
# already prevents reprocessing finished batches.
jobs = await _managed_object_table(self.prisma_client).find_many(
jobs = await self._find_jobs(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -979,8 +1053,6 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
self.batch_processed_support_confirmed = True
except Exception as query_err:
@ -1050,7 +1122,7 @@ class CheckBatchCost:
continue
verbose_proxy_logger.error(
f"CheckBatchCost: failed to track cost for batch {batch_id} "
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
f"(job {job.id}): {tracking_err}"
)
self._record_error(prom_logger, "cost_tracking_error")
continue
@ -1060,22 +1132,19 @@ class CheckBatchCost:
# Track this job for the final metrics summary
processed_models.append(tracked)
# mark the job as complete
try:
update_data: dict = {
"status": response.status if response.status != "completed" else "complete",
"file_object": response.model_dump_json(),
}
if self._has_batch_processed_column:
update_data["batch_processed"] = True
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data=update_data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
)
if not self._has_batch_processed_column:
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data={
"status": response.status if response.status != "completed" else "complete",
"file_object": response.model_dump_json(),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
)
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
from litellm.proxy.openai_files_endpoints.common_utils import (

View file

@ -22,7 +22,7 @@ from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import ValidationError
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
import litellm
from litellm import Router, verbose_logger
@ -205,6 +205,31 @@ def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableAct
return prisma_client.db.litellm_managedobjecttable
_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
def _storable_text(text: str) -> str:
return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "")
def _storable_json_value(value: JsonValue) -> JsonValue:
if isinstance(value, str):
return _storable_text(value)
if isinstance(value, list):
return [_storable_json_value(item) for item in value]
if isinstance(value, dict):
return {_storable_text(key): _storable_json_value(item) for key, item in value.items()}
return value
def postgres_storable_json(file_object: BaseModel) -> str:
dumped: Final = file_object.model_dump_json()
if "\\u0000" not in dumped:
return dumped
return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode()
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.71"
version = "0.1.71.post1"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.71"
version = "0.1.71.post1"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -140,11 +140,14 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig {
}
fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers {
update_headers_with_anthropic_beta(headers, request)
merge_beta_headers(
update_headers_with_anthropic_beta(headers, request),
safeguards_betas(&request.params),
)
}
}
fn update_headers_with_anthropic_beta(
pub(crate) fn update_headers_with_anthropic_beta(
headers: Headers,
request: &AnthropicMessagesRequest,
) -> Headers {
@ -199,6 +202,20 @@ fn uses_structured_output(params: &AnthropicMessagesOptionalParams) -> bool {
.is_some_and(|config| config.format.is_some())
}
fn safeguards_betas(params: &AnthropicMessagesOptionalParams) -> BetaSet {
carries_safeguards(params)
.then_some(AnthropicBeta::DangerousToolUse20260903)
.into_iter()
.collect()
}
fn carries_safeguards(params: &AnthropicMessagesOptionalParams) -> bool {
params
.extra
.get("safeguards")
.is_some_and(|value| !value.is_null())
}
fn messages_carry_output_config(messages: &[AnthropicMessage]) -> bool {
messages
.iter()
@ -1047,6 +1064,12 @@ mod tests {
json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}),
&[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")],
)]
#[case::safeguards_add_the_dangerous_tool_use_beta(
&[("anthropic-beta", "claude-code-20250219")],
json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
&[("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03")],
)]
#[case::null_safeguards_add_no_beta(&[("x-api-key", "k")], json!({"safeguards": null}), &[("x-api-key", "k")])]
fn request_headers_merge_the_feature_betas(
#[case] input: &[(&str, &str)],
#[case] fields: Value,

View file

@ -12,7 +12,9 @@ use crate::{
Error,
anthropic::{
common_utils::{API_KEY_PLACEMENT, MESSAGES_PATH_SUFFIX, non_empty},
messages::transformation::{ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig},
messages::transformation::{
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, update_headers_with_anthropic_beta,
},
},
base_llm::{
anthropic_messages::transformation::{
@ -103,7 +105,7 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig {
}
fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers {
self.anthropic.request_headers(headers, request)
update_headers_with_anthropic_beta(headers, request)
}
}

View file

@ -29,10 +29,16 @@ pub struct ModelInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_100k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_100k_tokens_batches: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_1hr: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_1hr_above_100k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_creation_input_token_cost_above_1hr_above_200k_tokens: Option<f64>,
@ -76,6 +82,10 @@ pub struct ModelInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_100k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_100k_tokens_batches: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cache_read_input_token_cost_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]
@ -183,6 +193,10 @@ pub struct ModelInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_100k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_100k_tokens_batches: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub input_cost_per_token_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]
@ -329,6 +343,10 @@ pub struct ModelInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_32k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_100k_tokens: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_100k_tokens_batches: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_cost_per_token_above_128k_tokens: Option<f64>,
/// Rate applied once the prompt exceeds the token threshold in the field name.
#[serde(skip_serializing_if = "Option::is_none")]

View file

@ -74,3 +74,20 @@ fn checked_in_catalog_and_backup_match() {
"invalid registry aliases"
);
}
#[rstest]
#[case::input("input_cost_per_token_above_100k_tokens")]
#[case::input_batches("input_cost_per_token_above_100k_tokens_batches")]
#[case::output("output_cost_per_token_above_100k_tokens")]
#[case::output_batches("output_cost_per_token_above_100k_tokens_batches")]
#[case::cache_creation("cache_creation_input_token_cost_above_100k_tokens")]
#[case::cache_creation_batches("cache_creation_input_token_cost_above_100k_tokens_batches")]
#[case::cache_creation_1hr("cache_creation_input_token_cost_above_1hr_above_100k_tokens")]
#[case::cache_read("cache_read_input_token_cost_above_100k_tokens")]
#[case::cache_read_batches("cache_read_input_token_cost_above_100k_tokens_batches")]
fn registry_validation_keeps_above_100k_tier_rates(#[case] field: &str) {
let mut entry = Map::new();
entry.insert("litellm_provider".into(), "anthropic".into());
entry.insert(field.into(), 5e-7.into());
validate_model_entry("test", &Value::Object(entry)).unwrap();
}

View file

@ -556,6 +556,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, ma
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
apply_search_free_tier: bool = False
cost_margin_config: Dict[
str, Union[float, Dict[str, float]]
] = {} # Provider-specific or global cost margins. Examples:

View file

@ -217,7 +217,7 @@
"computer-use-2025-11-24": "computer-use-2025-11-24",
"context-1m-2025-08-07": "context-1m-2025-08-07",
"context-management-2025-06-27": "context-management-2025-06-27",
"dangerous-tool-use-2026-09-03": null,
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": "fast-mode-2026-02-01",
"files-api-2025-04-14": "files-api-2025-04-14",

View file

@ -22,17 +22,28 @@ Environment Variables:
- LITELLM_ANTHROPIC_BETA_HEADERS_URL: Custom URL for remote config (optional)
"""
import asyncio
import json
import os
import threading
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from functools import cache
from importlib.resources import files
from typing import Final
from typing import Final, TypeAlias
import httpx
from pydantic import TypeAdapter
from litellm.litellm_core_utils.litellm_logging import verbose_logger
# Cache for the loaded configuration
_BETA_HEADERS_CONFIG: dict | None = None
_BETA_HEADERS_CONFIG_LOCK: Final = threading.Lock()
_BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-beta-headers-config")
BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]]
_BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig)
class GetAnthropicBetaHeadersConfig:
@ -174,13 +185,27 @@ def _load_beta_headers_config() -> dict:
if _BETA_HEADERS_CONFIG is not None:
return _BETA_HEADERS_CONFIG
# Get the URL from environment or use default
from litellm import anthropic_beta_headers_url
with _BETA_HEADERS_CONFIG_LOCK:
if _BETA_HEADERS_CONFIG is not None:
return _BETA_HEADERS_CONFIG
_BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
verbose_logger.debug("Loaded and cached beta headers config")
# Get the URL from environment or use default
from litellm import anthropic_beta_headers_url
return _BETA_HEADERS_CONFIG
_BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
verbose_logger.debug("Loaded and cached beta headers config")
return _BETA_HEADERS_CONFIG
async def aload_beta_headers_config() -> None:
if _BETA_HEADERS_CONFIG is None:
await asyncio.get_running_loop().run_in_executor(_BETA_HEADERS_CONFIG_LOADER, _load_beta_headers_config)
@cache
def _bundled_beta_headers_config() -> BetaHeadersConfig:
return _BETA_HEADERS_CONFIG_ADAPTER.validate_python(GetAnthropicBetaHeadersConfig.load_local_beta_headers_config())
def reload_beta_headers_config() -> dict:
@ -288,6 +313,14 @@ def is_beta_header_supported(
return beta_header in provider_mapping and provider_mapping[beta_header] is not None
def is_beta_header_supported_by_bundled_config(beta_header: str, provider: str) -> bool:
config: Final = _bundled_beta_headers_config()
aliases: Final = config.get("provider_aliases", {})
canonical_provider: Final = (aliases.get(provider) if isinstance(aliases, Mapping) else None) or provider
provider_mapping: Final = config.get(canonical_provider, {})
return isinstance(provider_mapping, Mapping) and provider_mapping.get(beta_header) is not None
def get_provider_beta_header(
anthropic_beta_header: str,
provider: str,

View file

@ -3,7 +3,7 @@ from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from dataclasses import replace as dataclasses_replace
from enum import Enum
from typing import Any, Final, Literal
from typing import Any, Final, Literal, cast
import litellm
from litellm._logging import verbose_logger
@ -73,6 +73,7 @@ async def calculate_batch_cost_and_usage(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""
Calculate the cost and usage of a batch.
@ -82,16 +83,21 @@ async def calculate_batch_cost_and_usage(
pricing. Threaded through to batch_cost_calculator so that
deployment-specific pricing (e.g. input_cost_per_token_batches)
is used instead of the global cost map.
vertex_location: The Vertex AI location the batch ran on; non-global
locations apply the model's regional-endpoint uplift multiplier.
"""
first_row: Final = file_content_dictionary[0] if file_content_dictionary else None
if _uses_native_vertex_output(custom_llm_provider, model_name, first_row):
return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name, model_info=model_info)
return calculate_vertex_ai_batch_cost_and_usage(
file_content_dictionary, model_name, model_info=model_info, vertex_location=vertex_location
)
return _aggregate_batch_cost_usage_models(
entries=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
vertex_location=vertex_location,
)
@ -101,6 +107,7 @@ async def _handle_completed_batch(
model_name: str | None = None,
litellm_params: dict | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""Fetch a completed batch's output file and aggregate its cost, usage, and
models in a single pass over the JSONL lines, so the parsed file content is
@ -114,6 +121,8 @@ async def _handle_completed_batch(
model_info: Optional deployment-level model info with custom pricing,
threaded through so a deployment's configured rates win over the
global cost map.
vertex_location: The Vertex AI location the batch ran on; non-global
locations apply the model's regional-endpoint uplift multiplier.
"""
# A completed batch whose request lines all failed has no output file - the
# results are written to a separate error_file_id and output_file_id is None.
@ -141,7 +150,10 @@ async def _handle_completed_batch(
output_file_result: Final = (
calculate_vertex_ai_batch_cost_and_usage(
_iter_batch_output_entries(file_content), model_name, model_info=model_info
_iter_batch_output_entries(file_content),
model_name,
model_info=model_info,
vertex_location=vertex_location,
)
if _uses_native_vertex_output(
custom_llm_provider, model_name, next(_iter_batch_output_entries(file_content), None)
@ -151,6 +163,7 @@ async def _handle_completed_batch(
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
vertex_location=vertex_location,
)
)
@ -186,6 +199,7 @@ def _classify_output_line_stats(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None,
model_info: ModelInfo | None,
vertex_location: str | None = None,
) -> Iterator[_BatchOutputLineStats | _LineOutcome]:
"""Classify every output line in a single pass, so counting failures never needs
a second read of a potentially huge output file. A line the provider reported as
@ -196,7 +210,7 @@ def _classify_output_line_stats(
if not _batch_response_was_successful(entry, custom_llm_provider):
yield _LineOutcome.PROVIDER_FAILED
continue
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info, vertex_location)
yield stats if stats is not None else _LineOutcome.UNCOSTABLE
@ -205,13 +219,14 @@ def _safe_output_line_stats(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None,
model_info: ModelInfo | None,
vertex_location: str | None = None,
) -> _BatchOutputLineStats | None:
"""Return the stats for one provider-successful batch output line, or None when
it cannot be costed, so a single bad line never aborts the whole batch's cost
accounting."""
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
try:
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info, vertex_location)
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
verbose_logger.warning(
"batch output line could not be costed, so it is billed at $0 and the rest of the batch "
@ -227,6 +242,7 @@ def _compute_output_line_stats(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None,
model_info: ModelInfo | None,
vertex_location: str | None = None,
) -> _BatchOutputLineStats:
response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
@ -241,6 +257,7 @@ def _compute_output_line_stats(
model_name=model_name,
response_model=response_model,
model_info=model_info,
vertex_location=vertex_location,
)
return _BatchOutputLineStats(
prompt_cost=line_prompt_cost,
@ -270,6 +287,7 @@ def _output_line_cost(
model_name: str | None,
response_model: str | None,
model_info: ModelInfo | None,
vertex_location: str | None = None,
) -> tuple[float, float]:
"""(prompt_cost, completion_cost) for one output line, priced at batch rates."""
from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost
@ -290,6 +308,7 @@ def _output_line_cost(
model=cost_model,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
vertex_location=vertex_location,
)
@ -298,11 +317,14 @@ def _aggregate_batch_cost_usage_models(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""Aggregate cost, usage, models, and pass/fail counts from batch output
entries in a single pass, holding one small stats record per line instead
of the parsed file."""
all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info))
all_results: Final = tuple(
_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info, vertex_location)
)
line_stats: Final = tuple(result for result in all_results if isinstance(result, _BatchOutputLineStats))
failed_requests: Final = sum(1 for result in all_results if result is _LineOutcome.PROVIDER_FAILED)
successful_requests: Final = len(all_results) - failed_requests
@ -349,6 +371,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses: Iterable[dict],
model_name: str | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""
Cost and usage of a native Vertex predictions.jsonl, one
@ -367,6 +390,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
model_info=model_info,
calculate_usage=VertexGeminiConfig._calculate_usage,
cost_calculator=batch_cost_calculator,
vertex_location=vertex_location,
)
for row in vertex_ai_batch_responses
)
@ -744,9 +768,17 @@ def _batch_response_was_successful(
message batch results lines report ``result.type == "succeeded"``; Bedrock
batch output lines report ``modelOutput`` (and no ``error``).
"""
if custom_llm_provider == "anthropic":
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded"
if custom_llm_provider == "bedrock":
return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None
_response: Final[dict] = batch_job_output_file.get("response", None) or {}
return _response.get("status_code", None) == 200
envelope_key: Final = "result" if custom_llm_provider == "anthropic" else "response"
envelope: Final = batch_job_output_file.get(envelope_key)
if not isinstance(envelope, Mapping):
if envelope is not None:
verbose_logger.warning("batch output %s must be an object, treating the line as unsuccessful", envelope_key)
return False
response: Final = cast(Mapping[object, object], envelope) # cast-ok: mapping shape validated above
return (
response.get("type") == "succeeded"
if custom_llm_provider == "anthropic"
else response.get("status_code") == 200
)

View file

@ -1871,6 +1871,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"public_model_groups",
"public_model_groups_links",
"cost_discount_config",
"apply_search_free_tier",
"cost_margin_config",
"block_requests_for_models_without_pricing",
"budget_exceeded_throttle_percentage",

View file

@ -8,7 +8,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from httpx import Response
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
@ -36,6 +36,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
get_batch_cost_rates,
get_billable_input_tokens,
get_token_type_cost_breakdown,
get_vertex_regional_endpoint_uplift,
parse_prompt_tokens_details,
select_cost_metric_for_model,
)
@ -125,6 +126,7 @@ from litellm.types.utils import (
TranscriptionUsageTokensObject,
Usage,
VectorStoreSearchResponse,
is_server_derived_pricing_key,
)
from litellm.utils import (
CallTypes,
@ -137,6 +139,8 @@ from litellm.utils import (
TextCompletionResponse,
TranscriptionResponse,
_cached_get_model_info_helper,
model_info_with_pricing_overrides,
runtime_model_cost_registrations,
token_counter,
)
@ -402,6 +406,7 @@ def cost_per_token(
### REQUEST MODEL ###
request_model: str | None = None, # original request model for router detection
custom_model_info: OCRPricing | None = None,
bedrock_model_info: ModelInfo | None = None,
) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@ -627,6 +632,7 @@ def cost_per_token(
model=model,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
vertex_location=vertex_location,
)
elif call_type == "atranscription" or call_type == "transcription":
if _transcription_usage_has_token_details(usage_block):
@ -688,6 +694,15 @@ def cost_per_token(
elif custom_llm_provider == "anthropic":
return anthropic_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
elif custom_llm_provider == "bedrock":
resolved_bedrock_info: Final = bedrock_model_info or _bedrock_pricing_info(model, None, "bedrock", False)
if resolved_bedrock_info is not None:
return generic_cost_per_token(
model=model,
usage=usage_block,
custom_llm_provider="bedrock",
service_tier=service_tier,
model_info=resolved_bedrock_info,
)
return bedrock_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
elif custom_llm_provider == "openai":
return openai_cost_per_token(
@ -801,6 +816,55 @@ def _cost_map_entry_prices_anything(entry: Mapping[str, object]) -> bool:
)
_LOGGED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
def _bedrock_registration_has_prices(entry: Mapping[str, object]) -> bool:
return _cost_map_entry_prices_anything(entry) or any(
value is not None and field.startswith("cache_") and is_server_derived_pricing_key(field)
for field, value in entry.items()
)
def _bedrock_pricing_info(
model: str | None,
invoked_model_id: str | None,
custom_llm_provider: str | None,
explicit_pricing: bool,
) -> ModelInfo | None:
if explicit_pricing or custom_llm_provider != "bedrock" or not model:
return None
registration: Final = next(
(
(key, entry)
for key, entry in runtime_model_cost_registrations(model, "bedrock")
if _bedrock_registration_has_prices(entry)
),
None,
)
profile: Final = _cost_map_model_info(invoked_model_id, "bedrock") if invoked_model_id else None
if profile is None or not _declares_rates(profile):
return _cost_map_model_info(registration[0], None) if registration is not None else None
if registration is None or registration[0] == profile["key"]:
return profile
return model_info_with_pricing_overrides(profile, registration[1])
def _declares_rates(model_info: ModelInfo | None) -> bool:
entry: Final = _raw_cost_map_entry(model_info["key"]) if model_info is not None else None
return entry is not None and _cost_map_entry_prices_anything(entry)
def _invoked_model_id_for_cost(
optional_params: object,
logging_obj: LitellmLoggingObject | None,
) -> str | None:
recorded: Final[object] = getattr(logging_obj, "model_call_details", None)
return _get_hidden_str_for_cost_calc(optional_params, "model_id") or _get_hidden_str_for_cost_calc(
recorded, "model_id"
)
def _select_model_name_for_cost_calc(
model: str | None,
completion_response: object | None,
@ -809,6 +873,7 @@ def _select_model_name_for_cost_calc(
custom_llm_provider: str | None = None,
router_model_id: str | None = None,
region_name: str | None = None,
invoked_model_id: str | None = None,
) -> str | None:
"""
1. If custom pricing is true, return received model name
@ -829,6 +894,13 @@ def _select_model_name_for_cost_calc(
hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None)
provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model")
explicit_pricing: Final = custom_pricing is True or base_model is not None
configured_model: Final = model or (
completion_response_model if isinstance(completion_response_model, str) else None
)
bedrock_model_info: Final = _bedrock_pricing_info(
configured_model, invoked_model_id, custom_llm_provider, explicit_pricing
)
bedrock_pricing_model: Final = bedrock_model_info["key"] if bedrock_model_info is not None else None
priced_from_response: Final = provider_response_model is not None or completion_response_model is not None
priced_region: Final = (
_get_hidden_str_for_cost_calc(hidden_params, "region_name") or region_name
@ -846,8 +918,8 @@ def _select_model_name_for_cost_calc(
else:
return_model = model
elif base_model is not None or provider_response_model is not None:
return_model = base_model if base_model is not None else provider_response_model
elif base_model is not None or bedrock_pricing_model is not None or provider_response_model is not None:
return_model = base_model if base_model is not None else bedrock_pricing_model or provider_response_model
elif completion_response_model is None and hidden_params is not None:
if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0:
@ -1426,6 +1498,10 @@ def completion_cost(
service_tier = _normalize_service_tier(service_tier)
explicit_pricing: Final = custom_pricing is True or base_model is not None
invoked_model_id: Final = _invoked_model_id_for_cost(
_LOGGED_LITELLM_PARAMS.validate_python(optional_params) if optional_params is not None else None,
litellm_logging_obj,
)
selected_model: Final = _select_model_name_for_cost_calc(
model=model,
completion_response=completion_response,
@ -1434,6 +1510,14 @@ def completion_cost(
base_model=base_model,
router_model_id=router_model_id,
region_name=region_name,
invoked_model_id=invoked_model_id,
)
bedrock_model_info: Final = _bedrock_pricing_info(
model or _get_response_model(completion_response),
invoked_model_id,
_get_provider_for_cost_calc(model, custom_llm_provider),
explicit_pricing,
)
potential_model_names: Final = [
@ -1808,6 +1892,7 @@ def completion_cost(
vertex_location=vertex_location,
response=completion_response,
custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id),
bedrock_model_info=bedrock_model_info,
)
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
@ -1896,6 +1981,7 @@ def completion_cost(
data_residency=data_residency,
vertex_location=vertex_location,
custom_cost_per_token=custom_cost_per_token,
model_info=bedrock_model_info,
)
_reasoning_cost = _token_type_breakdown.reasoning_cost
_cache_read_cost = _token_type_breakdown.cache_read_cost
@ -2137,7 +2223,16 @@ def pricing_entry_for_cost_calc(
custom_llm_provider=custom_llm_provider,
router_model_id=router_model_id,
region_name=region_name,
invoked_model_id=_invoked_model_id_for_cost(None, litellm_logging_obj),
)
bedrock_model_info: Final = _bedrock_pricing_info(
model or _get_response_model(completion_response),
_invoked_model_id_for_cost(None, litellm_logging_obj),
_get_provider_for_cost_calc(model, custom_llm_provider),
custom_pricing is True or base_model is not None,
)
if bedrock_model_info is not None:
return bedrock_model_info["key"], bedrock_model_info
candidates: Final = (selected_model, _get_response_model(completion_response), model)
resolved: Final = next(
(info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info),
@ -2571,6 +2666,8 @@ def batch_cost_calculator(
custom_llm_provider: str | None = None,
model_info: ModelInfo | None = None,
data_residency: str | None = None,
*,
vertex_location: str | None = None,
) -> tuple[float, float]:
"""
Calculate the cost of a batch job.
@ -2656,7 +2753,8 @@ def batch_cost_calculator(
total_prompt_cost *= uplift
total_completion_cost *= uplift
return total_prompt_cost, total_completion_cost
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
return total_prompt_cost * vertex_uplift, total_completion_cost * vertex_uplift
def _batch_prompt_cost(

View file

@ -1,9 +1,11 @@
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, replace
from typing import Final, TypeAlias
import httpx
from pydantic import ConfigDict, TypeAdapter, ValidationError
from pydantic_core import to_jsonable_python
from typing_extensions import assert_never
import litellm
@ -30,6 +32,7 @@ from litellm.types.decisions import (
OpenAIDecisionResponse,
UnsupportedDecisionsRequest,
)
from litellm.types.llms.openai import ChatCompletionUserMessage
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
@ -148,6 +151,30 @@ def _provider_ir_request(
return _UnsupportedSafetyIdentifier()
_LOGGED_FIELDS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _jsonable_fallback(value: object) -> object:
return dict(_LOGGED_FIELDS_ADAPTER.validate_python(value)) if isinstance(value, Mapping) else str(value)
def _logged_request_message(caller_request: object) -> ChatCompletionUserMessage:
caller_fields: Final = _LOGGED_FIELDS_ADAPTER.validate_python(
to_jsonable_python(caller_request, fallback=_jsonable_fallback, bytes_mode="base64")
)
sent_fields: Final = json.dumps(
dict((name, value) for name, value in caller_fields.items() if value not in (None, [], {}))
)
return ChatCompletionUserMessage(role="user", content=sent_fields)
def log_caller_request(kwargs: Mapping[str, object], caller_request: object) -> None:
"""Record every request field the caller sent as the logged `messages`, before any validation."""
logging_obj: Final = kwargs.get("litellm_logging_obj")
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.update_messages([_logged_request_message(caller_request)])
def prepare_call(
*,
model: str,

View file

@ -13,6 +13,7 @@ from litellm.decisions.call import (
SystemOneQuestions,
SystemOneRequestFields,
asend,
log_caller_request,
prepare_call,
send,
)
@ -59,6 +60,14 @@ def _prepare_decisions_call(
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> DecisionsCall:
log_caller_request(
kwargs,
{
"input": decision_input,
"questions": questions,
**({"state": kwargs["state"]} if "state" in kwargs else {}),
},
)
return prepare_call(
model=model,
fields=_request_fields(decision_input, questions, kwargs),

View file

@ -42,6 +42,7 @@ from litellm.constants import (
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
EMPTY_MAPPING,
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
PROVIDER_REQUEST_ID_HEADERS,
SENTRY_DENYLIST,
SENTRY_PII_DENYLIST,
@ -484,7 +485,7 @@ def _published_pricing(deployment_model: str | None) -> ModelInfo | None:
return None
def _resolve_vertex_location_for_cost(
def resolve_vertex_location_for_cost(
custom_llm_provider: str | None,
litellm_params: Mapping[str, object] | None,
optional_params: Mapping[str, object] | None,
@ -714,6 +715,7 @@ class Logging(LiteLLMLoggingBaseClass):
self._enqueue_deferred_logging: Callable[[], None] | None = None
self._async_success_scheduled: bool = False
self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
self._async_success_enqueued: bool = False
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
@ -1405,15 +1407,20 @@ class Logging(LiteLLMLoggingBaseClass):
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e
)
self.model_call_details["api_call_start_time"] = datetime.datetime.now()
# Set-once first provider-handoff instant. api_call_start_time
# is overwritten on every retry, so it can't measure one-time
# preprocessing; pinning the first attempt excludes retry loops
# + backoff. Logging object only — must NOT go into
# litellm_params["metadata"] (caller request metadata, typed
# Dict[str, str], echoed downstream; a datetime breaks it).
if self.model_call_details.get("first_api_call_start_time") is None:
self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"]
# A proxy-gate failure (auth, rate limit, client disconnect) fires pre_call for its callbacks
# without any provider handoff, so it must not stamp the handoff instants.
if LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL not in self.model_call_details:
self.model_call_details["api_call_start_time"] = datetime.datetime.now()
# Set-once first provider-handoff instant. api_call_start_time
# is overwritten on every retry, so it can't measure one-time
# preprocessing; pinning the first attempt excludes retry loops
# + backoff. Logging object only — must NOT go into
# litellm_params["metadata"] (caller request metadata, typed
# Dict[str, str], echoed downstream; a datetime breaks it).
if self.model_call_details.get("first_api_call_start_time") is None:
self.model_call_details["first_api_call_start_time"] = self.model_call_details[
"api_call_start_time"
]
# Input Integration Logging -> If you want to log the fact that an attempt to call the model was made
callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or [])
for callback in callbacks:
@ -1880,7 +1887,7 @@ class Logging(LiteLLMLoggingBaseClass):
router_model_id = self.get_router_model_id()
## RESPONSE COST ##
custom_pricing: Final = self._custom_pricing_for(priced_result)
custom_pricing: Final = self.custom_pricing_for(priced_result)
prompt = self._prompt_for_cost_calculation()
@ -1907,7 +1914,7 @@ class Logging(LiteLLMLoggingBaseClass):
if hasattr(self, "litellm_params") and self.litellm_params
else None
),
"vertex_location": _resolve_vertex_location_for_cost(
"vertex_location": resolve_vertex_location_for_cost(
custom_llm_provider=self.model_call_details.get("custom_llm_provider", None),
litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None),
optional_params=self.optional_params,
@ -2019,7 +2026,7 @@ class Logging(LiteLLMLoggingBaseClass):
model=model,
completion_response=result,
custom_llm_provider=custom_llm_provider,
custom_pricing=self._custom_pricing_for(result),
custom_pricing=self.custom_pricing_for(result),
base_model=_get_base_model_from_metadata(model_call_details=self.model_call_details),
router_model_id=router_model_id or self.get_router_model_id(),
region_name=_resolve_mantle_region_for_cost(
@ -2044,20 +2051,23 @@ class Logging(LiteLLMLoggingBaseClass):
usage=usage,
)
def _custom_pricing_for(self, result: object) -> bool:
litellm_params: Final = getattr(self, "litellm_params", None)
def spilled_over_ptu_for(self, result: object) -> bool:
result_hidden_params: Final = getattr(result, "_hidden_params", None) or MappingProxyType({})
additional_headers: Final = (
result_hidden_params.get("additional_headers")
if isinstance(result_hidden_params, dict)
else getattr(result_hidden_params, "additional_headers", None)
)
spilled_over: Final = is_spilled_over_ptu_request(
model_info=_deployment_model_info(litellm_params),
return is_spilled_over_ptu_request(
model_info=_deployment_model_info(getattr(self, "litellm_params", None)),
response_headers=self.model_call_details.get("response_headers"),
additional_headers=additional_headers,
)
return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params)
def custom_pricing_for(self, result: object) -> bool:
if self.spilled_over_ptu_for(result):
return False
return use_custom_pricing_for_model(litellm_params=getattr(self, "litellm_params", None))
def _prompt_for_cost_calculation(self) -> str:
"""
@ -2264,6 +2274,16 @@ class Logging(LiteLLMLoggingBaseClass):
except Exception:
return True
@property
def async_success_log_enqueued(self) -> bool:
"""True once the success log task was actually created, not merely claimed, so a later failure raised
around it (a disconnect cancelling a still-pending guardrail task) must not bill the prompt a second
time. A claimed log parked in the deferred slot never runs once a 499 interrupts, so it stays False."""
return self._async_success_enqueued
def mark_async_success_log_enqueued(self) -> None:
self._async_success_enqueued = True
def claim_async_success_log(self) -> bool:
"""One async success log per request. The innermost @client wrapper always exits first,
since the outer one is awaiting it, so it claims the log here and the outer wrapper gets
@ -3285,6 +3305,12 @@ class Logging(LiteLLMLoggingBaseClass):
model_name=self.get_deployment_model_for_cost(),
litellm_params=self.litellm_params,
model_info=self.get_router_deployment_model_info(),
vertex_location=resolve_vertex_location_for_cost(
custom_llm_provider=self.custom_llm_provider,
litellm_params=self.litellm_params if hasattr(self, "litellm_params") else None,
optional_params=None,
model=self.get_deployment_model_for_cost() or self.model,
),
)
result._hidden_params["response_cost"] = batch_result.cost
@ -5890,6 +5916,7 @@ class StandardLoggingPayloadSetup:
custom_llm_provider: str | None,
init_response_obj: object,
api_base: str | None = None,
router_model_id: str | None = None,
) -> StandardLoggingModelInformation:
model_cost_name: Final = _select_model_name_for_cost_calc(
model=base_model if custom_pricing else None,
@ -5897,12 +5924,23 @@ class StandardLoggingPayloadSetup:
base_model=base_model,
custom_pricing=custom_pricing,
)
pricing_model_name: Final = (
_select_model_name_for_cost_calc(
model=base_model if custom_pricing else None,
completion_response=init_response_obj,
base_model=base_model,
custom_pricing=custom_pricing,
router_model_id=router_model_id,
)
if router_model_id is not None
else model_cost_name
)
if model_cost_name is None:
model_cost_information = StandardLoggingModelInformation(model_map_key="", model_map_value=None)
else:
try:
_model_cost_information: Final = litellm.get_model_info(
model=model_cost_name,
model=pricing_model_name or model_cost_name,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
)
@ -6534,13 +6572,16 @@ def get_standard_logging_object_payload(
)
## Get model cost information ##
base_model = _get_base_model_from_metadata(model_call_details=kwargs)
# The router overrides completion_response.model to the model-group alias before
# this payload is built, so cost-map lookup via that alias always misses.
# Fall back to the actual deployment model set by the router in metadata.
if base_model is None:
base_model = metadata.get("deployment")
custom_pricing: Final = use_custom_pricing_for_model(litellm_params=litellm_params)
configured_base_model: Final = _get_base_model_from_metadata(model_call_details=kwargs)
spilled_over: Final = logging_obj.spilled_over_ptu_for(init_response_obj)
# The router overrides completion_response.model to the model-group alias before this payload is
# built, so cost-map lookup via that alias misses and the deployment model stands in for base_model.
# A PTU request Azure spilled onto pay-as-you-go capacity is billed at the served model's rates,
# which the deployment's zeroed PTU entry would misreport, so that fallback is skipped there.
base_model: Final = (
metadata.get("deployment") if configured_base_model is None and not spilled_over else configured_base_model
)
custom_pricing: Final = False if spilled_over else logging_obj.custom_pricing_for(init_response_obj)
raw_response_cost: Final = kwargs.get("response_cost")
llm_response_cost: Final[float] = raw_response_cost or 0.0
guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information"))
@ -6564,6 +6605,7 @@ def get_standard_logging_object_payload(
custom_llm_provider=kwargs.get("custom_llm_provider"),
init_response_obj=init_response_obj,
api_base=litellm_params.get("api_base"),
router_model_id=logging_obj.get_router_model_id(),
)
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(

View file

@ -75,6 +75,9 @@ _SERVICE_TIER_TO_COST_KEY_SUFFIX: Final[Mapping[str, str]] = MappingProxyType(
ServiceTier.ULTRAFAST.value: ServiceTier.ULTRAFAST.value,
}
)
SERVICE_TIER_COST_KEY_SUFFIXES: Final[tuple[str, ...]] = tuple(
sorted(frozenset(f"_{suffix}" for suffix in _SERVICE_TIER_TO_COST_KEY_SUFFIX.values()))
)
_INCLUSIVE_THRESHOLD_PROVIDERS: Final = frozenset({"xai"})
_BATCH_KEY_SUFFIX: Final = "_batches"
@ -272,7 +275,7 @@ def _parse_token_threshold(threshold: str) -> float:
def _parse_above_token_threshold(key: str) -> float:
return _parse_token_threshold(key.split("_above_")[1].split("_tokens")[0])
return _parse_token_threshold(key.rsplit("_above_", 1)[1].split("_tokens")[0])
def _prompt_exceeds_threshold(prompt_tokens: int, threshold: float, inclusive: bool) -> bool:
@ -665,20 +668,26 @@ def _get_token_base_cost(
# Exclude service_tier-specific variants (e.g. input_cost_per_token_above_200k_tokens_priority)
# so that the threshold detection loop only processes standard keys. The
# service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key.
tier_key_suffix: Final = _get_service_tier_cost_key("", service_tier)
threshold_prefixes: Final = tuple(f"{prefix}_above_" for prefix in _BATCH_RATE_PREFIXES)
threshold_keys: Final = [
k
for k in model_info
if k.startswith("input_cost_per_token_above_") and not k.endswith(_NON_STANDARD_THRESHOLD_SUFFIXES)
if k.startswith(threshold_prefixes)
and "_tokens" in k
and (
not k.endswith(_NON_STANDARD_THRESHOLD_SUFFIXES) or (tier_key_suffix != "" and k.endswith(tier_key_suffix))
)
]
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
threshold: float | None = None
for key in sorted(threshold_keys, key=_parse_above_token_threshold, reverse=True):
for key in sorted(threshold_keys, key=_parse_above_token_threshold):
value = model_info.get(key)
if value is not None:
try:
# Handle both formats: _above_128k_tokens and _above_128_tokens
threshold_str = key.split("_above_")[1].split("_tokens")[0]
threshold_str = key.rsplit("_above_", 1)[1].split("_tokens")[0]
threshold = _parse_above_token_threshold(key)
if _prompt_exceeds_threshold(usage.prompt_tokens, threshold, threshold_is_inclusive):
# Prefer a service_tier-specific above-threshold key when available,
@ -691,7 +700,7 @@ def _get_token_base_cost(
service_tier,
)
if service_tier
else key
else f"input_cost_per_token_above_{threshold_str}_tokens"
)
prompt_base_cost = cast(
float,
@ -747,8 +756,6 @@ def _get_token_base_cost(
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
break
except (IndexError, ValueError):
continue
except Exception:
@ -1064,6 +1071,25 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
)
def _modality_token_cost(
model_info: ModelInfo,
modality_cost_key: str,
tokens: int,
prompt_base_cost: float,
service_tier: str | None,
) -> float:
"""
Audio, image and video prompt tokens bill at their modality rate when the model has one.
Without one they are ordinary prompt tokens, so they bill at `prompt_base_cost`, which
`_get_token_base_cost` has already resolved for the service tier and long-context tier.
"""
tiered_key: Final = _get_service_tier_cost_key(modality_cost_key, service_tier)
cost_key: Final = tiered_key if model_info.get(tiered_key) is not None else modality_cost_key
if model_info.get(cost_key) is None:
return float(tokens) * prompt_base_cost
return calculate_cost_component(model_info, cost_key, tokens)
def _calculate_input_cost(
prompt_tokens_details: PromptTokensDetailsResult,
model_info: ModelInfo,
@ -1075,8 +1101,14 @@ def _calculate_input_cost(
) -> float:
"""
Calculates the input cost for a given model, prompt tokens, and completion tokens.
A model priced per character (`input_cost_per_character`) bills the reported character count
in place of its text tokens, so the same prompt is never charged under both units.
"""
prompt_cost = float(prompt_tokens_details["text_tokens"]) * prompt_base_cost
bills_characters: Final = (
prompt_tokens_details["character_count"] > 0 and model_info.get("input_cost_per_character") is not None
)
prompt_cost = 0.0 if bills_characters else float(prompt_tokens_details["text_tokens"]) * prompt_base_cost
### CACHE READ COST - Now uses tiered pricing
cache_hit_audio_tokens: Final = prompt_tokens_details["cache_hit_audio_tokens"]
@ -1094,28 +1126,37 @@ def _calculate_input_cost(
if prompt_tokens_details["audio_tokens"] and not (
prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None
):
audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier)
prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"])
prompt_cost += _modality_token_cost(
model_info,
"input_cost_per_audio_token",
prompt_tokens_details["audio_tokens"],
prompt_base_cost,
service_tier,
)
### IMAGE TOKEN COST
if prompt_tokens_details["image_tokens"] and not (
prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None
):
# For image token costs:
# First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
image_token_cost_key = "input_cost_per_image_token"
if model_info.get(image_token_cost_key) is None:
image_token_cost_key = "input_cost_per_token"
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
prompt_cost += _modality_token_cost(
model_info,
"input_cost_per_image_token",
prompt_tokens_details["image_tokens"],
prompt_base_cost,
service_tier,
)
### VIDEO TOKEN COST
if prompt_tokens_details["video_tokens"] and not (
prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None
):
video_token_cost_key = "input_cost_per_video_token"
if model_info.get(video_token_cost_key) is None:
video_token_cost_key = "input_cost_per_token"
prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"])
prompt_cost += _modality_token_cost(
model_info,
"input_cost_per_video_token",
prompt_tokens_details["video_tokens"],
prompt_base_cost,
service_tier,
)
### CACHE WRITING COST - Now uses tiered pricing
if (
@ -1629,20 +1670,23 @@ def get_billed_token_rates(
vertex_location: str | None = None,
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
model_info: ModelInfo | 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 = apply_provider_cache_read_default(
get_model_info(model=model, custom_llm_provider=custom_llm_provider),
resolved_model_info: Final = apply_provider_cache_read_default(
model_info
if model_info is not None
else get_model_info(model=model, custom_llm_provider=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,
model_info=resolved_model_info,
usage=usage,
custom_llm_provider=custom_llm_provider,
service_tier=service_tier,
@ -1661,6 +1705,7 @@ def get_token_type_cost_breakdown(
vertex_location: str | None = None,
current_time: datetime | None = None,
custom_cost_per_token: CostPerToken | None = None,
model_info: ModelInfo | None = None,
) -> TokenTypeCostBreakdown:
"""
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
@ -1685,6 +1730,7 @@ def get_token_type_cost_breakdown(
vertex_location=vertex_location,
current_time=current_time,
custom_cost_per_token=custom_cost_per_token,
model_info=model_info,
)
if rates is None:
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)

View file

@ -4,6 +4,7 @@ from functools import reduce
from typing import Final
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm.constants import LITELLM_DETAILED_TIMING
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers
@ -101,6 +102,32 @@ def response_timing_metrics(
return {"_response_ms": total_response_time_ms, "litellm_overhead_time_ms": overhead_ms}
_HIDDEN_PARAMS_DICT: Final = TypeAdapter(dict[str, object])
def _hidden_params_mapping(hidden_params: object) -> Mapping[str, object]:
try:
return _HIDDEN_PARAMS_DICT.validate_python(hidden_params)
except ValidationError:
return {}
def _existing_provider_response_model(result: object) -> str | None:
"""The pricing model a response already carries, like the decisions path sets, wins over the response model."""
hidden_params: Final[object] = getattr(result, "_hidden_params", None)
carried: Final[object] = _hidden_params_mapping(hidden_params).get("provider_response_model") or getattr(
hidden_params, "provider_response_model", None
)
return carried if isinstance(carried, str) and carried else None
def _provider_response_model(result: object) -> str | None:
"""The model the provider answered with, kept because the proxy later rewrites ``result.model`` to the alias the
client asked for and cost-map lookups for the standard logging payload would otherwise miss."""
model: Final = getattr(result, "model", None) if isinstance(result, ModelResponse) else None
return model if isinstance(model, str) and model else None
class ResponseMetadata:
"""
Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses
@ -133,6 +160,8 @@ class ResponseMetadata:
preserve_litellm_internal_headers=True,
),
"litellm_model_name": model,
"provider_response_model": _existing_provider_response_model(self.result)
or _provider_response_model(self.result),
}
self._update_hidden_params(new_params)

View file

@ -12,10 +12,15 @@ tell "user asked for this" from "nobody set it". This resolver answers that:
from __future__ import annotations
from typing import Final
from collections.abc import Mapping
from typing import Final, Literal, Protocol
from pydantic import TypeAdapter
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
_TIMEOUTS_BY_SERVICE_TIER: Final[TypeAdapter[Mapping[str, float] | None]] = TypeAdapter(Mapping[str, float] | None)
def get_configured_request_timeout() -> float | None:
"""Return the explicitly-configured ``litellm.request_timeout``, else ``None``."""
@ -27,3 +32,19 @@ def get_configured_request_timeout() -> float | None:
if timeout != float(DEFAULT_REQUEST_TIMEOUT_SECONDS):
return timeout
return None
class _DeploymentParams(Protocol):
def get(self, key: str, /) -> object: ...
def get_service_tier_timeout(
service_tier: object,
litellm_params: _DeploymentParams,
key: Literal["timeout_by_service_tier", "stream_timeout_by_service_tier"],
) -> float | None:
"""Return the deployment's ``key`` timeout for the request's ``service_tier``, else ``None``."""
if not isinstance(service_tier, str):
return None
timeouts: Final = _TIMEOUTS_BY_SERVICE_TIER.validate_python(litellm_params.get(key))
return None if timeouts is None else timeouts.get(service_tier)

View file

@ -103,6 +103,7 @@ from ..common_utils import (
process_anthropic_headers,
requires_native_compaction_beta,
strip_advisor_blocks_from_messages,
tools_marked_for_a_tool_config_point,
)
if TYPE_CHECKING:
@ -2056,6 +2057,19 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
# Remove internal LiteLLM parameters that should not be sent to Anthropic API
optional_params.pop("is_vertex_request", None)
optional_params.pop("client_metadata", None)
request_params: Final = cast( # cast-ok: optional_params is the untyped dict of the shared transform signature
dict[str, object], optional_params
)
marked_tools: Final = tools_marked_for_a_tool_config_point(
request_params.get("tools"), request_params.pop("cache_control_injection_points", None)
)
if marked_tools is not None:
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
request_litellm_params: Final = cast( # cast-ok: litellm_params is the untyped dict of the shared signature
Mapping[str, object], litellm_params
)
AnthropicCacheControlHook.record_gateway_injection(request_litellm_params, 1)
# ``top_k`` is a provider-specific kwarg that bypasses
# ``map_openai_params``; gate it here, the single boundary shared by
@ -2075,6 +2089,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
"model": model,
"messages": anthropic_messages,
**optional_params,
**({} if marked_tools is None else {"tools": list(marked_tools)}),
}
self._apply_output_config(data=data, model=model, optional_params=optional_params)

View file

@ -33,6 +33,7 @@ from litellm.types.llms.anthropic import (
ANTHROPIC_HOSTED_TOOLS,
ANTHROPIC_OAUTH_BETA_HEADER,
ANTHROPIC_OAUTH_TOKEN_PREFIX,
ANTHROPIC_TOOL_SEARCH_TOOL_TYPES,
AllAnthropicToolsValues,
AnthropicMcpServerTool,
AnthropicMessagesToolChoice,
@ -73,7 +74,7 @@ _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
_DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)")
_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/")
@ -115,6 +116,40 @@ def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None
)
def _tool_config_point_control(points: object) -> Mapping[object, object] | None:
point_fields: Final = tuple(_validated_claude_code_mapping(point) for point in validated_object_list(points) or ())
tool_config_point: Final = next(
(fields for fields in point_fields if fields is not None and fields.get("location") == "tool_config"), None
)
if tool_config_point is None:
return None
control: Final = tool_config_point.get("control")
return {"type": "ephemeral"} if control is None else _validated_claude_code_mapping(control)
def tools_marked_for_a_tool_config_point(tools: object, points: object) -> tuple[object, ...] | None:
"""Put the ``tool_config`` point's breakpoint on the last tool Anthropic caches.
Anthropic caches tools as one prefix ending at the marked tool, so the last one carries
the mark, skipping tool-search tools whose transform drops cache markers. A client that
already marked a tool keeps its own breakpoint. Returns ``None`` when nothing is marked.
"""
control: Final = _tool_config_point_control(points)
tool_list: Final = validated_object_list(tools) or ()
tool_fields: Final = tuple(_validated_claude_code_mapping(tool) for tool in tool_list)
if control is None or any(fields is not None and "cache_control" in fields for fields in tool_fields):
return None
cacheable: Final = tuple(
(position, fields)
for position, fields in enumerate(tool_fields)
if fields is not None and fields.get("type") not in ANTHROPIC_TOOL_SEARCH_TOOL_TYPES
)
if not cacheable:
return None
position, fields = cacheable[-1]
return (*tool_list[:position], {**fields, "cache_control": control}, *tool_list[position + 1 :])
def is_claude_code_user_agent(user_agent: str) -> bool:
"""Claude Code sends its API calls through the Anthropic SDK as `claude-cli/<version>` and its own
fetches, such as gateway model discovery, as `claude-code/<version>`"""
@ -128,9 +163,9 @@ def _validated_claude_code_mapping(value: object) -> dict[object, object] | None
return None
def _validated_claude_code_list(value: object) -> list[object] | None:
def validated_object_list(value: object) -> Sequence[object] | None:
try:
return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value)
return OBJECT_LIST_ADAPTER.validate_python(value)
except ValidationError:
return None
@ -157,7 +192,7 @@ def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None
def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None:
if isinstance(system, str):
return (system,)
blocks: Final = _validated_claude_code_list(system)
blocks: Final = validated_object_list(system)
if blocks is None:
return None
block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks)

View file

@ -12,6 +12,7 @@ from functools import partial
from typing import Any, Final, cast
import litellm
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.common_utils import (
@ -19,6 +20,7 @@ from litellm.llms.anthropic.common_utils import (
sanitize_tool_use_ids_in_anthropic_messages,
strip_empty_content_blocks_from_anthropic_messages,
strip_provider_specific_fields_from_anthropic_messages,
tools_marked_for_a_tool_config_point,
)
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
@ -636,7 +638,13 @@ def anthropic_messages_handler(
f"custom_llm_provider is required for Anthropic messages, passed in model={model}, custom_llm_provider={custom_llm_provider}"
)
local_vars.update(kwargs)
request_kwargs: Final = cast(dict[str, object], kwargs) # cast-ok: **kwargs of the handler is untyped
marked_tools: Final = tools_marked_for_a_tool_config_point(
local_vars.get("tools"), request_kwargs.pop("cache_control_injection_points", None)
)
if marked_tools is not None:
AnthropicCacheControlHook.record_gateway_injection(request_kwargs, 1)
local_vars.update({**request_kwargs, **({"tools": list(marked_tools)} if marked_tools is not None else {})})
anthropic_messages_optional_request_params: Final = (
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
params=local_vars,

View file

@ -3,6 +3,7 @@ from typing import Any, Final
import httpx
from litellm.anthropic_beta_headers_manager import is_beta_header_supported_by_bundled_config
from litellm.exceptions import AuthenticationError
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import verbose_logger
@ -611,6 +612,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
optional_params: dict,
custom_llm_provider: str = "anthropic",
messages: Sequence[object] = (),
beta_map_provider: str | None = None,
) -> dict:
"""
Auto-inject anthropic-beta headers based on features used.
@ -627,8 +629,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
optional_params: Optional parameters including tools, context_management, output_format, speed
custom_llm_provider: Provider name for looking up correct tool search header
messages: Request messages, scanned for per-message output_config
beta_map_provider: Provider whose beta map column gates safeguards, defaults to custom_llm_provider
"""
beta_values: Final[set] = set()
beta_values: Final[set[str]] = set()
existing_beta: Final = tuple(
piece.strip()
@ -678,6 +681,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
if optional_params.get("speed") == "fast":
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value)
if optional_params.get("safeguards") is not None and is_beta_header_supported_by_bundled_config(
ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value,
beta_map_provider or custom_llm_provider,
):
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value)
if _messages_carry_output_config(messages):
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value)

View file

@ -4,11 +4,12 @@ import asyncio
import json
import traceback
from collections import deque
from collections.abc import AsyncIterator, Iterator, Mapping
from collections.abc import AsyncGenerator, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final
from litellm import verbose_logger
from litellm._uuid import uuid
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
from litellm.litellm_core_utils.prompt_templates.common_utils import (
encrypted_reasoning_signature,
)
@ -25,6 +26,7 @@ from .transformation import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
from litellm.types.llms.openai import ResponseIncompleteEvent
class AnthropicResponsesStreamWrapper:
@ -391,8 +393,11 @@ class AnthropicResponsesStreamWrapper:
raise StopAsyncIteration
async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]:
"""Yield SSE-encoded bytes for each Anthropic event chunk."""
def async_anthropic_sse_wrapper(self) -> "AnthropicResponsesSSEByteStream":
"""SSE-encoded bytes for each Anthropic event chunk."""
return AnthropicResponsesSSEByteStream(self)
async def _sse_frames(self) -> AsyncGenerator[bytes, None]:
async for chunk in self:
if isinstance(chunk, dict):
event_type: str = str(chunk.get("type", "message"))
@ -400,3 +405,30 @@ class AnthropicResponsesStreamWrapper:
yield payload.encode()
else:
yield chunk
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
return partial_responses_incomplete_event_on_disconnect(self.responses_stream)
class AnthropicResponsesSSEByteStream:
"""
The byte stream the proxy relays for a default-route /v1/messages request.
Keeps the Responses API iterator reachable so a client disconnect can bill
the partial output the way native /v1/responses streams do.
"""
def __init__(self, wrapper: AnthropicResponsesStreamWrapper) -> None:
self._wrapper: Final = wrapper
self._frames: Final = wrapper._sse_frames() # pyright: ignore[reportPrivateUsage] # the wrapper's own byte stream
def __aiter__(self) -> "AnthropicResponsesSSEByteStream":
return self
async def __anext__(self) -> bytes:
return await self._frames.__anext__()
async def aclose(self) -> None:
await self._frames.aclose()
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
return self._wrapper.partial_response_on_disconnect()

View file

@ -1,12 +1,10 @@
"""Support for Azure OpenAI gpt-5 model family."""
from typing import Final
from typing import ClassVar, Final
import litellm
from litellm.exceptions import UnsupportedParamsError
from litellm.llms.openai.chat.gpt_5_transformation import (
OpenAIGPT5Config,
_get_effort_level,
is_gpt_reasoning_series_name,
)
from litellm.types.llms.openai import AllMessageValues
@ -15,9 +13,16 @@ from .gpt_transformation import AzureOpenAIConfig
class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
"""Azure specific handling for gpt-5 models."""
"""Azure specific handling for gpt-5 models.
reasoning_effort none joins xhigh as opt-in here: Azure rejected it on every gpt-5
deployment before gpt-5.1 (https://github.com/BerriAI/litellm/issues/16704), and a deployment
name the map never heard of must not forward it. See
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning
"""
GPT5_SERIES_ROUTE = "gpt5_series/"
OPT_IN_REASONING_EFFORTS: ClassVar[frozenset[str]] = frozenset(("xhigh", "none"))
@classmethod
def _model_map_lookup_name(cls, model: str) -> str:
@ -27,11 +32,18 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
model_prices_and_context_window.json; ``azure/gpt-5.1`` is. Overriding the shared
resolver rather than one lookup means the supports, explicitly-disabled and
default-effort answers all read the same entry.
An ``azure_ai/`` name reaches this config when ``get_llm_provider`` re-routes a Foundry
deployment on an OpenAI-v1 host to the azure provider with the prefix intact. It is the
same product the azure_ai config reads ``azure_ai/<model>`` for, so that row answers when
the map carries it, and the azure/ twin when it does not.
"""
if model.startswith(cls.GPT5_SERIES_ROUTE):
return "azure/" + model[len(cls.GPT5_SERIES_ROUTE) :]
if model.startswith("azure/"):
return model
if model.startswith("azure_ai/"):
return model if model in litellm.model_cost else "azure/" + model.removeprefix("azure_ai/")
return "azure/" + model
@classmethod
@ -75,34 +87,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
drop_params: bool,
api_version: str = "",
) -> dict:
reasoning_effort_value = non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort")
effective_effort: Final = _get_effort_level(reasoning_effort_value)
# gpt-5.1/5.2/5.4 support reasoning_effort='none', but other gpt-5 models don't
# See: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning
supports_none: Final = self._supports_reasoning_effort_level(model, "none")
if effective_effort == "none" and not supports_none:
if litellm.drop_params is True or (drop_params is not None and drop_params is True):
non_default_params = non_default_params.copy()
optional_params = optional_params.copy()
if _get_effort_level(non_default_params.get("reasoning_effort")) == "none":
non_default_params.pop("reasoning_effort")
if _get_effort_level(optional_params.get("reasoning_effort")) == "none":
optional_params.pop("reasoning_effort")
else:
raise UnsupportedParamsError(
status_code=400,
message=(
"Azure OpenAI does not support reasoning_effort='none' for this model. "
"Supported values are: 'low', 'medium', and 'high'. "
"To drop this parameter, set `litellm.drop_params=True` or for proxy:\n\n"
"`litellm_settings:\n drop_params: true`\n"
"Issue: https://github.com/BerriAI/litellm/issues/16704"
),
)
result: Final = OpenAIGPT5Config.map_openai_params(
return OpenAIGPT5Config.map_openai_params(
self,
non_default_params=non_default_params,
optional_params=optional_params,
@ -110,17 +95,6 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
drop_params=drop_params,
)
# Only drop reasoning_effort='none' for models that don't support it
result_effort: Final = _get_effort_level(result.get("reasoning_effort"))
if result_effort == "none" and not supports_none:
result.pop("reasoning_effort")
# Azure gpt-5.4+ with tools + reasoning_effort is now routed to the
# Responses API bridge (same as OpenAI), so we no longer need to drop
# reasoning_effort here. See: responses_api_bridge_check() in main.py.
return result
def transform_request(
self,
model: str,

View file

@ -30,10 +30,18 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.AZURE
@staticmethod
def _is_gpt_5_model(model: str) -> bool:
return AzureOpenAIGPT5Config.is_model_gpt_5_model(model)
@staticmethod
def _supports_reasoning_effort_none(model: str) -> bool:
return AzureOpenAIGPT5Config._supports_reasoning_effort_level(model, "none")
@staticmethod
def _effort_level_is_disabled(model: str, level: str) -> bool:
return AzureOpenAIGPT5Config.row_disables_reasoning_effort_level(model, level)
@staticmethod
def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
return AzureOpenAIGPT5Config.effort_resolves_to_none(model, effort)

View file

@ -3,6 +3,7 @@ from typing import Final
import httpx
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
from litellm.llms.azure_ai.chat.transformation import AzureAIGPT5Config
from litellm.llms.azure_ai.common_utils import (
AzureFoundryModelInfo,
api_key_header_for_base,
@ -27,6 +28,18 @@ class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig):
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.AZURE_AI
@staticmethod
def _supports_reasoning_effort_none(model: str) -> bool:
return AzureAIGPT5Config._supports_reasoning_effort_level(model, "none")
@staticmethod
def _effort_level_is_disabled(model: str, level: str) -> bool:
return AzureAIGPT5Config.row_disables_reasoning_effort_level(model, level)
@staticmethod
def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
return AzureAIGPT5Config.effort_resolves_to_none(model, effort)
def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict:
params: Final = litellm_params or GenericLiteLLMParams()
auth_headers: Final = get_azure_ai_auth_headers(

View file

@ -2659,13 +2659,13 @@ class AmazonConverseConfig(BaseConfig):
## HANDLE TOOL CALLS
_message: Final = Message(**chat_completion_message)
initial_finish_reason = map_finish_reason(completion_response["stopReason"])
mapped_finish_reason: Final = map_finish_reason(completion_response["stopReason"])
# 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 resolved_json_mode and not filtered_tools and tools:
initial_finish_reason = "stop"
initial_finish_reason: Final = (
"stop"
if resolved_json_mode and not filtered_tools and tools and mapped_finish_reason == "tool_calls"
else mapped_finish_reason
)
(
returned_message,

View file

@ -3,29 +3,103 @@ Bedrock Token Counter implementation using the CountTokens API.
"""
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Final
from pydantic import JsonValue
from pydantic import JsonValue, TypeAdapter
from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse
RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404})
_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
class _CountTokensReply(LiteLLMPydanticObjectBase):
input_tokens: int
@dataclass(frozen=True, slots=True)
class CountedTokens:
input_tokens: int
original_response: Mapping[str, JsonValue]
tokenizer_type: str
@dataclass(frozen=True, slots=True)
class CountTokensFailure:
status_code: int
message: str
tokenizer_type: str
CountTokensOutcome = CountedTokens | CountTokensFailure
def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool:
return (
isinstance(outcome, CountTokensFailure)
and outcome.status_code in _RUNTIME_FALLBACK_STATUSES
and is_mantle_claude_model(resolved_model)
)
class BedrockTokenCounter(BaseTokenCounter):
"""Token counter implementation for AWS Bedrock provider using the CountTokens API."""
"""Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens
on bedrock-mantle for the Claude models bedrock-runtime answers 400 or 404."""
def __init__(
self,
runtime_handler: BedrockCountTokensHandler | None = None,
mantle_handler: BedrockMantleCountTokensHandler | None = None,
client: AsyncHTTPHandler | None = None,
) -> None:
self._runtime_handler: Final = runtime_handler or BedrockCountTokensHandler()
self._mantle_handler: Final = mantle_handler or BedrockMantleCountTokensHandler()
self._client: Final = client
def should_use_token_counting_api(
self,
custom_llm_provider: str | None = None,
) -> bool:
"""
Returns True if we should use the Bedrock CountTokens API for token counting.
"""
return custom_llm_provider == LlmProviders.BEDROCK.value
async def _count_with(
self,
handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler,
tokenizer_type: str,
request_data: Mapping[str, object],
litellm_params: dict[str, object],
resolved_model: str,
) -> CountTokensOutcome:
try:
result: Final = await handler.handle_count_tokens_request(
request_data=request_data,
litellm_params=litellm_params,
resolved_model=resolved_model,
client=self._client,
)
reply: Final = _CountTokensReply.model_validate(result)
except BedrockError as e:
verbose_logger.debug(
"%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message
)
return CountTokensFailure(status_code=e.status_code, message=e.message, tokenizer_type=tokenizer_type)
except Exception as e:
verbose_logger.debug("Error calling %s CountTokens API: %s", tokenizer_type, e)
return CountTokensFailure(status_code=500, message=str(e), tokenizer_type=tokenizer_type)
return CountedTokens(input_tokens=reply.input_tokens, original_response=result, tokenizer_type=tokenizer_type)
async def count_tokens(
self,
model_to_use: str,
@ -38,81 +112,53 @@ class BedrockTokenCounter(BaseTokenCounter):
messages_format: TokenCountMessagesFormat = "anthropic",
output_format: JsonValue = None,
) -> TokenCountResponse | None:
"""
Count tokens using AWS Bedrock's CountTokens API.
This method calls the existing BedrockCountTokensHandler to make an API call
to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting.
Args:
model_to_use: The model identifier
messages: The messages to count tokens for
contents: Alternative content format (not used for Bedrock)
deployment: Deployment configuration containing litellm_params
request_model: The original request model name
Returns:
TokenCountResponse with token count, or None if counting fails
"""
if not messages:
return None
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
# Build request data in the format expected by BedrockCountTokensHandler
request_data: Final[dict[str, object]] = {
"model": model_to_use,
"messages": messages,
}
if tools:
request_data["tools"] = tools
if system:
request_data["system"] = system
# Get the resolved model (strip prefixes like bedrock/, converse/, etc.)
litellm_params: Final = (deployment or {}).get("litellm_params", {})
request_data: Final = MappingProxyType(
{
"model": model_to_use,
"messages": messages,
**MappingProxyType({name: value for name, value in (("tools", tools), ("system", system)) if value}),
}
)
resolved_model: Final = get_bedrock_base_model(model_to_use)
try:
handler: Final = BedrockCountTokensHandler()
result: Final = await handler.handle_count_tokens_request(
request_data=request_data,
litellm_params=litellm_params,
resolved_model=resolved_model,
runtime: Final = await self._count_with(
self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
)
outcome: Final = (
await self._count_with(
self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
)
# Transform response to TokenCountResponse
if result is not None:
if _runtime_rejected_claude_model(runtime, resolved_model)
else runtime
)
match outcome:
case CountedTokens():
return TokenCountResponse(
total_tokens=result.get("input_tokens", 0),
total_tokens=outcome.input_tokens,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
original_response=result,
tokenizer_type=outcome.tokenizer_type,
original_response=_COUNT_RESPONSE.validate_python(outcome.original_response),
)
except BedrockError as e:
verbose_logger.warning("Bedrock CountTokens API error: status=%s, message=%s", e.status_code, e.message)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
error=True,
error_message=e.message,
status_code=e.status_code,
)
except Exception as e:
verbose_logger.warning("Error calling Bedrock CountTokens API: %s", e)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
error=True,
error_message=str(e),
status_code=500,
)
return None
case CountTokensFailure():
verbose_logger.warning(
"%s CountTokens API error: status=%s, message=%s",
outcome.tokenizer_type,
outcome.status_code,
outcome.message,
)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type=outcome.tokenizer_type,
error=True,
error_message=outcome.message,
status_code=outcome.status_code,
)
case _:
assert_never(outcome)

View file

@ -4,6 +4,7 @@ AWS Bedrock CountTokens API handler.
Simplified handler leveraging existing LiteLLM Bedrock infrastructure.
"""
from collections.abc import Mapping
from typing import Any, Final
import httpx
@ -25,7 +26,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
async def handle_count_tokens_request(
self,
request_data: dict[str, Any],
request_data: Mapping[str, Any],
litellm_params: dict[str, Any],
resolved_model: str,
client: AsyncHTTPHandler | None = None,
@ -125,7 +126,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
raise
except httpx.HTTPStatusError as e:
# HTTP errors - preserve the actual status code
verbose_logger.error("HTTP error in CountTokens handler: %s", e)
verbose_logger.debug("HTTP error in CountTokens handler: %s", e)
raise BedrockError(
status_code=e.response.status_code,
message=e.response.text,

View file

@ -0,0 +1,113 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
import httpx
from pydantic import JsonValue, TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.llms.anthropic.count_tokens.transformation import (
COUNT_TOKEN_OPTION_NAMES,
AnthropicCountTokensConfig,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens"
MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01"
class MantleCountTokensRequest(TypedDict):
messages: ReadOnly[list[dict[str, JsonValue]]]
system: ReadOnly[NotRequired[JsonValue]]
tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]]
_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest)
_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
_COUNT_OPTIONS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mapping[str, JsonValue])
class BedrockMantleCountTokensHandler(BaseAWSLLM):
"""Counts tokens through Anthropic's count_tokens on the bedrock-mantle endpoint.
Claude models that Bedrock offers only through cross-region inference answer 400 on
bedrock-runtime's CountTokens; AWS documents Mantle's /anthropic/v1/messages/count_tokens
as the way to count them, with the base model id and the deployment's AWS credentials
"""
def __init__(self, anthropic_config: AnthropicCountTokensConfig | None = None) -> None:
super().__init__()
self._anthropic_config: Final = anthropic_config or AnthropicCountTokensConfig()
def get_mantle_count_tokens_endpoint(self, aws_region_name: str) -> str:
messages_url: Final = build_mantle_messages_url(
api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name
)
return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}"
async def handle_count_tokens_request(
self,
request_data: Mapping[str, object],
litellm_params: dict[str, object],
resolved_model: str,
client: AsyncHTTPHandler | None = None,
) -> dict[str, JsonValue]:
try:
request: Final = _COUNT_REQUEST.validate_python(request_data)
aws_region_name: Final = self._get_aws_region_name(
optional_params=litellm_params, model=resolved_model, model_id=None
)
endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name)
body: Final = self._anthropic_config.transform_request_to_count_tokens(
model=resolved_model,
messages=request["messages"],
tools=request.get("tools"),
system=request.get("system"),
optional_params=_COUNT_OPTIONS.validate_python(
MappingProxyType(
{name: value for name, value in request_data.items() if name in COUNT_TOKEN_OPTION_NAMES}
)
),
)
verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url)
api_key: Final = litellm_params.get("api_key")
signed_headers, signed_body = await run_aws_signing(
self._sign_request,
service_name="bedrock",
headers={ # mutable-ok: [LIT002] bearer signing may add the Authorization header
"Content-Type": "application/json",
"anthropic-version": MANTLE_ANTHROPIC_VERSION,
},
optional_params=litellm_params,
request_data=body,
api_base=endpoint_url,
model=resolved_model,
api_key=api_key if isinstance(api_key, str) else None,
)
async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
response: Final = await async_client.post(
endpoint_url, headers=signed_headers, data=signed_body, timeout=30.0
)
if response.status_code != 200:
raise BedrockError(
status_code=response.status_code,
message=response.text,
headers=response.headers,
response=response,
)
return _COUNT_RESPONSE.validate_json(response.content)
except BedrockError:
raise
except httpx.HTTPStatusError as e:
raise BedrockError(
status_code=e.response.status_code,
message=e.response.text,
headers=e.response.headers,
response=e.response,
)
except Exception as e:
raise BedrockError(status_code=500, message=f"bedrock-mantle count_tokens processing error: {e}")

View file

@ -45,6 +45,10 @@ def split_mantle_region_prefix(model: str) -> tuple[str | None, str]:
return None, model
def is_mantle_claude_model(model: str) -> bool:
return "claude" in model.lower()
def resolve_mantle_region(params: Mapping[str, object]) -> str:
region: Final = params.get("aws_region_name")
if isinstance(region, str) and region:

View file

@ -3,7 +3,7 @@ import json
import ssl
from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from functools import lru_cache, partial
from functools import lru_cache, partial, reduce
from types import MappingProxyType, ModuleType
from typing import (
TYPE_CHECKING,
@ -32,7 +32,7 @@ import litellm.litellm_core_utils
import litellm.types
import litellm.types.utils
from litellm._logging import _redact_string, verbose_logger
from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, update_headers_with_filtered_beta
from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
from litellm.files.types import FileContentStreamingResult
from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs
@ -46,6 +46,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
synthesize_subtitle_document,
)
from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields
from litellm.litellm_core_utils.realtime_errors import (
@ -2116,6 +2117,16 @@ class BaseLLMHTTPHandler:
merged_headers.update(extra_headers_from_kwargs)
if provider_specific_headers:
merged_headers.update(provider_specific_headers)
request_params: Final = reduce(
delete_nested_value,
litellm_params.get("additional_drop_params") or (),
anthropic_messages_optional_request_params,
)
if (
request_params.get("safeguards") is not None
and anthropic_messages_provider_config.should_filter_anthropic_beta_headers()
):
await aload_beta_headers_config()
(
headers,
api_base,
@ -2123,7 +2134,7 @@ class BaseLLMHTTPHandler:
headers=merged_headers or {},
model=model,
messages=messages,
optional_params=anthropic_messages_optional_request_params,
optional_params=request_params,
litellm_params=dict(litellm_params),
api_key=api_key,
api_base=api_base,
@ -2154,20 +2165,11 @@ class BaseLLMHTTPHandler:
custom_llm_provider=custom_llm_provider,
)
additional_drop_params: Final[list[str]] = litellm_params.get("additional_drop_params") or []
if additional_drop_params:
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
for path in additional_drop_params:
anthropic_messages_optional_request_params = delete_nested_value(
anthropic_messages_optional_request_params, path
)
# Prepare request body
request_body: Final = anthropic_messages_provider_config.transform_anthropic_messages_request(
model=model,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
anthropic_messages_optional_request_params=request_params,
litellm_params=litellm_params,
headers=headers,
)
@ -2273,7 +2275,7 @@ class BaseLLMHTTPHandler:
held_back_tool_names: Final = self._server_fulfilled_tools_in_request(
logging_obj=logging_obj,
tools=anthropic_messages_optional_request_params.get("tools"),
tools=request_params.get("tools"),
)
initial_response = AgenticAnthropicStreamingIterator(
completion_stream=completion_stream,
@ -2281,7 +2283,7 @@ class BaseLLMHTTPHandler:
model=model,
messages=messages,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
anthropic_messages_optional_request_params=request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs_for_agentic,
@ -2304,7 +2306,7 @@ class BaseLLMHTTPHandler:
model=model,
messages=messages,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
anthropic_messages_optional_request_params=request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs_for_agentic,

View file

@ -7,7 +7,7 @@ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequenc
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
import httpx
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
@ -21,6 +21,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
strip_name_from_message,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.base_utils import (
type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # base_utils helper returns an untyped dict
)
from litellm.types.llms.anthropic import AllAnthropicToolsValues
from litellm.types.llms.databricks import (
AllDatabricksContentValues,
@ -59,6 +62,10 @@ from ...anthropic.chat.transformation import (
from ...openai_like.chat.transformation import OpenAILikeChatConfig
from ..common_utils import DatabricksBase, DatabricksException
_RESPONSE_FORMAT_ADAPTER: ( # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object
Final[TypeAdapter[dict[str, object] | None]]
) = TypeAdapter(dict[str, object] | None)
def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool:
"""Databricks rejects assistant messages with neither content nor tool calls, e.g. a replayed
@ -68,6 +75,15 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool:
)
def _collapse_single_text_block(message: Mapping[str, object]) -> Mapping[str, object]:
content: Final = message.get("content")
match content:
case [{"type": "text", "text": str() as text}] if content == [{"type": "text", "text": text}]:
return {**message, "content": text}
case _:
return message
def _sanitize_empty_content(message_dict: dict[str, object]) -> None:
"""
Remove or filter content so empty text blocks are not sent.
@ -191,6 +207,14 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
def get_config(cls, *, model: str | None = None):
return super().get_config()
def get_json_schema_from_pydantic_object(
self,
response_format: ( # mutable-ok: matches BaseConfig override signature
type[BaseModel] | dict[str, object] | None
),
) -> dict[str, object] | None: # mutable-ok: BaseConfig contract returns a dict
return _RESPONSE_FORMAT_ADAPTER.validate_python(type_to_response_format_param(response_format=response_format))
def get_required_params(self) -> list[ProviderField]:
"""For a given provider, return it's required fields with a description"""
return [
@ -447,6 +471,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
Databricks does not support:
- 'name' in user message.
- litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages.
- non-Claude json_object "json" checks read only string content, so a single text block is sent as a string.
"""
new_messages = []
for idx, message in enumerate(messages):
@ -462,7 +487,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
_sanitize_empty_content(cast(dict[str, Any], _message))
if _is_bare_assistant_message(_message):
continue
new_messages.append(_message)
new_messages.append(_collapse_single_text_block(_message) if "claude" not in model else _message)
if "claude" not in model:
new_messages = _split_parallel_tool_calls(

View file

@ -1,7 +1,7 @@
"""Support for OpenAI gpt-5 model family."""
import re
from typing import Final
from typing import ClassVar, Final
import litellm
from litellm.utils import (
@ -67,6 +67,10 @@ def _get_effort_level(value: str | dict | None) -> str | None:
GPT_REASONING_SERIES_MARKERS: Final = ("gpt-5", "gpt-6")
def unsupported_reasoning_effort_message(param: str, level: str, model: str) -> str:
return f"{param}={level} is not supported for {model}. To drop unsupported params set `litellm.drop_params = True`"
def is_gpt_reasoning_series_name(model: str) -> bool:
normalized: Final = model.split("/")[-1]
return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) and not normalized.startswith("gpt-5-chat")
@ -82,6 +86,9 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
- Support for GPT-5-Codex models optimized for code generation.
"""
OPT_IN_REASONING_EFFORTS: ClassVar[frozenset[str]] = frozenset(("xhigh",))
ROUTE_SPECIFIC_REASONING_EFFORTS: ClassVar[frozenset[str]] = frozenset(("max",))
@classmethod
def is_model_gpt_5_model(cls, model: str) -> bool:
return is_gpt_reasoning_series_name(model)
@ -193,15 +200,10 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
return False
@classmethod
def _is_reasoning_effort_level_explicitly_disabled(cls, model: str, level: str) -> bool:
def is_reasoning_effort_level_explicitly_disabled(cls, model: str, level: str) -> bool:
"""Return True only when the model map explicitly sets the capability to False.
Unlike ``_supports_reasoning_effort_level`` (which requires an explicit True),
this method returns True only when ``supports_{level}_reasoning_effort`` is
explicitly set to ``False`` in the model map. A missing key is treated as
supported (i.e. this method returns False = not disabled).
Use this for opt-out checks where unknown models should be allowed through.
A missing key is not a disable. Unknown models pass through.
"""
return is_explicitly_disabled_factory(
model=cls._model_map_lookup_name(model),
@ -209,6 +211,39 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
key=f"supports_{level}_reasoning_effort",
)
@classmethod
def _is_reasoning_effort_level_explicitly_disabled(cls, model: str, level: str) -> bool:
return cls.is_reasoning_effort_level_explicitly_disabled(model, level)
@classmethod
def refuses_reasoning_effort_level(cls, model: str, level: str) -> bool:
"""The chat gate for any requested effort level, read off the model's map row alone.
An opt-in level (xhigh here; Azure adds none) is refused unless the row says true, since
the providers rejected it on most of the family when it arrived. Every other level is
refused only when the row says false: a model the map never heard of, or a row without the
flag, keeps forwarding what it was sent, and a level with no flag at all (medium, high) is
never refused. So a new model in the family needs its map row and no code change, the
row's supports_<level>_reasoning_effort flags being the whole contract.
A route-specific level (max) is never refused by either gate: the same model honors it
on Responses and refuses it on Chat Completions, so one row flag cannot gate it on
either route. Chat forwards it for the Responses bridge and a plain chat call gets the
provider's own 400, exactly as before the gate read every level.
"""
if level in cls.ROUTE_SPECIFIC_REASONING_EFFORTS:
return False
if level in cls.OPT_IN_REASONING_EFFORTS:
return not cls._supports_reasoning_effort_level(model, level)
return cls.is_reasoning_effort_level_explicitly_disabled(model, level)
@classmethod
def row_disables_reasoning_effort_level(cls, model: str, level: str) -> bool:
"""The Responses gate: opt-out for every level except the route-specific ones."""
if level in cls.ROUTE_SPECIFIC_REASONING_EFFORTS:
return False
return cls.is_reasoning_effort_level_explicitly_disabled(model, level)
def get_supported_openai_params(self, model: str) -> list:
if self.is_model_gpt_5_search_model(model):
return [
@ -269,12 +304,12 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
drop_params=drop_params,
)
# Get raw reasoning_effort and effective effort level for all guards.
# Use effective_effort (extracted string) for xhigh validation, "none" checks, and
# tool/sampling guards — dict inputs like {"effort": "none", "summary": "detailed"}
# must be treated as effort="none" to avoid incorrect tool-drop or sampling errors.
raw_reasoning_effort = non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort")
effective_effort: Final = _get_effort_level(raw_reasoning_effort)
# Dict inputs like {"effort": "none", "summary": "detailed"} are judged by their effort
# string so the gate, the sampling guard, and the temperature guard all see one level.
raw_reasoning_effort: Final = non_default_params.get("reasoning_effort") or optional_params.get(
"reasoning_effort"
)
requested_effort: Final = _get_effort_level(raw_reasoning_effort)
# Normalize dict reasoning_effort to string for Chat Completions API.
# Example: {"effort": "high", "summary": "detailed"} -> "high"
@ -286,31 +321,20 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
if "reasoning_effort" in optional_params:
optional_params["reasoning_effort"] = normalized
if effective_effort == "xhigh":
# xhigh is an opt-in capability: only allow if model explicitly supports it.
if not self._supports_reasoning_effort_level(model, effective_effort):
if litellm.drop_params or drop_params:
non_default_params.pop("reasoning_effort", None)
optional_params.pop("reasoning_effort", None)
else:
raise litellm.utils.UnsupportedParamsError(
message=(f"reasoning_effort={effective_effort} is not supported for this model."),
status_code=400,
)
elif effective_effort in ("minimal", "low"):
# minimal/low are opt-out: unknown models pass through; only block when
# the model map explicitly sets supports_{level}_reasoning_effort=false.
# Example: gpt-5.5-pro only accepts {medium, high, xhigh}, so it sets
# supports_low_reasoning_effort=false (and supports_minimal=false).
if self._is_reasoning_effort_level_explicitly_disabled(model, effective_effort):
if litellm.drop_params or drop_params:
non_default_params.pop("reasoning_effort", None)
optional_params.pop("reasoning_effort", None)
else:
raise litellm.utils.UnsupportedParamsError(
message=(f"reasoning_effort={effective_effort} is not supported for this model."),
status_code=400,
)
if requested_effort is not None and self.refuses_reasoning_effort_level(model, requested_effort):
if litellm.drop_params or drop_params:
non_default_params.pop("reasoning_effort", None)
optional_params.pop("reasoning_effort", None)
else:
raise litellm.utils.UnsupportedParamsError(
message=unsupported_reasoning_effort_message("reasoning_effort", requested_effort, model),
status_code=400,
)
# The level the provider will see: None once dropped, so the guards below judge by the model's default effort
effective_effort: Final = _get_effort_level(
non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort")
)
################################################################
# max_tokens is not supported for gpt-5 models on OpenAI API

View file

@ -22,7 +22,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
from litellm.llms.openai.chat.gpt_5_transformation import (
is_gpt_reasoning_series_name,
unsupported_reasoning_effort_message,
)
from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import *
@ -136,6 +139,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
return supports_none_reasoning_effort(model=model, custom_llm_provider=None)
@staticmethod
def _effort_level_is_disabled(model: str, level: str) -> bool:
"""Whether the model's cost-map row turns this effort level off.
The same row the chat gate reads, so a level the row refuses is refused on both wires
and a new model needs its row, not code. This surface stays opt-out for every level,
xhigh included, since it never refused a level the row omits; max is the exception, a
level the row cannot gate because the same model takes it here and refuses it on chat.
Azure overrides this so a bare deployment name reads the azure/ entry.
"""
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
return OpenAIGPT5Config.row_disables_reasoning_effort_level(model, level)
@staticmethod
def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
"""Whether this request's reasoning effort ends up as "none", the one condition
@ -242,7 +259,25 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
lookup_name: Final = self._model_map_lookup_name(model)
if self._is_gpt_5_model(model=lookup_name):
reasoning: Final = params.get("reasoning") or {}
effort: Final = reasoning.get("effort") if isinstance(reasoning, dict) else None
requested_effort: Final = reasoning.get("effort") if isinstance(reasoning, dict) else None
if (
isinstance(reasoning, dict)
and isinstance(requested_effort, str)
and self._effort_level_is_disabled(lookup_name, requested_effort)
):
if not (drop_params or litellm.drop_params):
raise litellm.UnsupportedParamsError(
message=unsupported_reasoning_effort_message("reasoning.effort", requested_effort, model),
status_code=400,
)
remaining: Final = reasoning.copy()
remaining.pop("effort", None)
if remaining:
params["reasoning"] = remaining
else:
params.pop("reasoning", None)
forwarded_reasoning: Final = params.get("reasoning")
effort: Final = forwarded_reasoning.get("effort") if isinstance(forwarded_reasoning, dict) else None
supports_none: Final = self._supports_reasoning_effort_none(model=lookup_name)
effort_is_none: Final = supports_none and self._effort_resolves_to_none(lookup_name, effort)

View file

@ -8,6 +8,7 @@ from litellm.llms.anthropic.experimental_pass_through.messages.transformation im
from litellm.llms.openai_like.json_loader import SimpleProviderConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01"
@ -57,6 +58,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig):
headers=normalized,
optional_params=optional_params,
messages=messages,
beta_map_provider=LlmProviders.OPENAI_LIKE.value,
)
return merged, api_base

View file

@ -54,6 +54,8 @@ class NativeVertexBatchCostCalculator(Protocol):
model: str,
custom_llm_provider: str | None = None,
model_info: ModelInfo | None = None,
*,
vertex_location: str | None = None,
) -> tuple[float, float]: ...
@ -90,6 +92,7 @@ def native_vertex_batch_row_stats(
model_info: ModelInfo | None,
calculate_usage: Callable[[GenerateContentResponseBody], Usage],
cost_calculator: NativeVertexBatchCostCalculator,
vertex_location: str | None = None,
) -> NativeVertexBatchRowStats | None:
"""
Usage and cost of one native Vertex predictions.jsonl row, a
@ -121,7 +124,11 @@ def native_vertex_batch_row_stats(
)
try:
prompt_cost, completion_cost = cost_calculator(
usage=usage, model=model, custom_llm_provider="vertex_ai", model_info=model_info
usage=usage,
model=model,
custom_llm_provider="vertex_ai",
model_info=model_info,
vertex_location=vertex_location,
)
except Exception as e: # noqa: BLE001 # one unpriceable row must not abort the batch's cost accounting
verbose_logger.warning(

View file

@ -1,14 +1,34 @@
import types
from typing import Final, Literal
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
from litellm.types.utils import EmbeddingResponse, Usage
from litellm.types.utils import EmbeddingResponse, PromptTokensDetailsWrapper, Usage
from .types import *
class _PredictMetadata(TypedDict, total=False):
billableCharacterCount: ReadOnly[int]
class _PredictResponseMetadata(TypedDict, total=False):
metadata: ReadOnly[_PredictMetadata]
_PREDICT_RESPONSE_METADATA: Final = TypeAdapter(_PredictResponseMetadata)
def billable_character_count(response: _PredictResponseMetadata) -> int | None:
"""
`metadata.billableCharacterCount` of a Vertex text embedding predict response, the count Google
bills per-character embedding models (text-embedding-005, text-multilingual-embedding-002) on.
"""
return response.get("metadata", {}).get("billableCharacterCount")
class VertexAITextEmbeddingConfig(BaseModel):
"""
Reference: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#TextEmbeddingInput
@ -248,7 +268,15 @@ class VertexAITextEmbeddingConfig(BaseModel):
model_response.object = "list"
model_response.data = embedding_response
model_response.model = model
usage: Final = Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens)
billable_characters: Final = billable_character_count(_PREDICT_RESPONSE_METADATA.validate_python(response))
usage: Final = Usage(
prompt_tokens=input_tokens,
completion_tokens=0,
total_tokens=input_tokens,
prompt_tokens_details=(
None if billable_characters is None else PromptTokensDetailsWrapper(character_count=billable_characters)
),
)
setattr(model_response, "usage", usage)
return model_response

View file

@ -1057,6 +1057,18 @@ def _resolve_openai_api_base(api_base: str | None) -> str:
)
def _chat_gate_refuses_reasoning_effort_none(model: str, custom_llm_provider: str) -> bool:
from litellm.llms.azure_ai.chat.transformation import AzureAIGPT5Config
match custom_llm_provider:
case "azure":
return litellm.AzureOpenAIGPT5Config.refuses_reasoning_effort_level(model, "none")
case "azure_ai":
return AzureAIGPT5Config.refuses_reasoning_effort_level(model, "none")
case _:
return OpenAIGPT5Config.refuses_reasoning_effort_level(model, "none")
def responses_api_bridge_check(
model: str,
custom_llm_provider: str,
@ -1129,10 +1141,13 @@ def responses_api_bridge_check(
)
for tool in (tools or ())
)
if isinstance(reasoning_effort, dict):
reasoning_active = reasoning_effort.get("effort") != "none" or reasoning_effort.get("summary") is not None
else:
reasoning_active = reasoning_effort != "none"
requested_effort: Final = reasoning_effort.get("effort") if isinstance(reasoning_effort, dict) else reasoning_effort
requested_summary: Final = reasoning_effort.get("summary") if isinstance(reasoning_effort, dict) else None
reasoning_active: Final = (
requested_effort != "none"
or requested_summary is not None
or _chat_gate_refuses_reasoning_effort_none(model, custom_llm_provider)
)
# The reasoning+tools constraint is enforced by the real OpenAI backend behind any api.openai.com
# host (the default URL or a PrivateLink hostname such as <region>.privatelink.api.openai.com) and
# by Azure OpenAI through the azure provider. Resolve the effective OpenAI base arg>global>env>default
@ -5707,6 +5722,7 @@ def completion(
optional_params=processed_non_default_params, # [IMPORTANT] - using processed_non_default_params ensures consistent params logged to langfuse for finetuning / eval datasets.
litellm_params=litellm_params,
custom_llm_provider=custom_llm_provider,
model_id=optional_params.get("model_id"),
)
if mock_response or mock_tool_calls or mock_timeout:
kwargs.pop("mock_timeout", None) # remove for any fallbacks triggered

File diff suppressed because it is too large Load diff

View file

@ -5,6 +5,8 @@ from typing import Final
from pydantic import TypeAdapter
from litellm.litellm_core_utils.request_timeout_resolver import get_service_tier_timeout
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS: Final = 600.0
_SECONDS: Final = TypeAdapter(float)
@ -48,21 +50,27 @@ def resolve_llm_passthrough_timeout(
Anthropic /v1/messages).
Non-streaming precedence: kwargs timeout/request_timeout -> litellm_params
timeout/request_timeout -> router_timeout -> general_settings.pass_through_request_timeout
-> 600s default.
timeout_by_service_tier[kwargs service_tier] -> litellm_params timeout/request_timeout -> router_timeout
-> general_settings.pass_through_request_timeout -> 600s default.
Streaming (``kwargs["stream"]`` truthy) resolves ``stream_timeout`` at every level before
any generic timeout, matching ``Router._get_stream_timeout`` on the completion route:
kwargs stream_timeout -> litellm_params stream_timeout -> router_stream_timeout, then the
non-streaming chain above.
kwargs stream_timeout -> litellm_params stream_timeout_by_service_tier[kwargs service_tier] ->
litellm_params stream_timeout -> router_stream_timeout, then the non-streaming chain above.
Only the first set value is validated as seconds, so a value in a lower-precedence
field never fails the call.
"""
request: Final = kwargs if kwargs is not None else _NO_PARAMS
deployment: Final = litellm_params if litellm_params is not None else _NO_PARAMS
service_tier: Final = request.get("service_tier")
stream_candidates: Final = (
(request.get("stream_timeout"), deployment.get("stream_timeout"), router_stream_timeout)
(
request.get("stream_timeout"),
get_service_tier_timeout(service_tier, deployment, "stream_timeout_by_service_tier"),
deployment.get("stream_timeout"),
router_stream_timeout,
)
if request.get("stream")
else ()
)
@ -70,6 +78,7 @@ def resolve_llm_passthrough_timeout(
*stream_candidates,
request.get("timeout"),
request.get("request_timeout"),
get_service_tier_timeout(service_tier, deployment, "timeout_by_service_tier"),
deployment.get("timeout"),
deployment.get("request_timeout"),
router_timeout,

View file

@ -111,6 +111,7 @@ from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _ar
from litellm.proxy.native_compaction import with_proxy_compaction_executor
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
from litellm.router import Router
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.router_utils.common_utils import resolve_model_group_alias
@ -418,6 +419,13 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
# A natural end-of-stream success event already fired and released the
# slot; do not bill again, and let the caller skip the slot release.
return True
partial_responses_event: Final = partial_responses_incomplete_event_on_disconnect(response)
if partial_responses_event is not None:
verbose_proxy_logger.debug(
"Billing partial streamed Responses API spend after client disconnect, litellm_call_id=%s",
logging_obj.litellm_call_id,
)
return await _dispatch_disconnect_billing_event(logging_obj, partial_responses_event)
chunks: Final[object] = getattr(response, "chunks", None)
if not isinstance(chunks, list) or not chunks:
return False
@ -447,6 +455,10 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
partial_usage: Final = getattr(partial_response, "usage", None)
if isinstance(partial_usage, Usage):
backfill_missing_cache_usage_fields(partial_usage)
return await _dispatch_disconnect_billing_event(logging_obj, partial_response)
async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool:
try:
await logging_obj.dispatch_success_handlers(
partial_response,

View file

@ -119,6 +119,23 @@ def coerce_numeric_form_fields(
}
def _parse_form_metadata(raw_metadata: str) -> object:
"""The JSON body path rejects a lone surrogate escape (``"\\ud800"``) with a
400 because no provider or database can encode it. ``json.loads`` accepts
the escape, so the multipart ``metadata`` field gets the same check."""
parsed: Final[object] = json.loads(raw_metadata)
try:
json.dumps(parsed, ensure_ascii=False).encode("utf-8")
except UnicodeEncodeError as e:
raise ProxyException(
message=f"Invalid form payload: metadata {e}",
type="invalid_request_error",
param="metadata",
code=status.HTTP_400_BAD_REQUEST,
) from e
return parsed
async def _read_request_body(request: Request | None) -> dict:
"""
Safely read the request body and parse it as JSON.
@ -158,9 +175,13 @@ async def _read_request_body(request: Request | None) -> dict:
param="request_body",
code=status.HTTP_400_BAD_REQUEST,
)
parsed_body = dict(form_data)
if "metadata" in parsed_body and isinstance(parsed_body["metadata"], str):
parsed_body["metadata"] = json.loads(parsed_body["metadata"])
form_body: Final = dict(form_data)
raw_metadata: Final = form_body.get("metadata")
parsed_body = (
{**form_body, "metadata": _parse_form_metadata(raw_metadata)}
if isinstance(raw_metadata, str)
else form_body
)
else:
# Read the request body
body: Final = await request.body()

View file

@ -88,6 +88,7 @@ if TYPE_CHECKING:
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
else:
PrismaClient = Any
ProxyLogging = Any
@ -97,6 +98,34 @@ RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, Call
_SPEND_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _turn_chains_through_the_db(payload: Mapping[str, object], llm_router: "Router | None") -> bool:
"""A Responses turn the chat-completions bridge replays from its spend log on the next `previous_response_id`.
Native Responses providers hold the conversation themselves, so their rows can wait for the batch write.
A row with no router deployment behind it is treated as bridged, which keeps today's immediate flush.
"""
if payload.get("call_type") not in RESPONSES_SESSION_CALL_TYPES:
return False
model_id: Final = payload.get("model_id")
deployment: Final = (
llm_router.get_deployment(model_id=model_id) if llm_router is not None and isinstance(model_id, str) else None
)
if deployment is None:
return True
from litellm.responses.main import provider_stores_the_conversation
litellm_params: Final = deployment.litellm_params
return not provider_stores_the_conversation(
litellm_params.model,
litellm_params.custom_llm_provider,
{
"use_chat_completions_api": litellm_params.use_chat_completions_api,
"model_info": deployment.model_info.model_dump(),
"api_base": litellm_params.api_base,
},
)
def _org_member_transaction_key(org_id: str, user_id: str) -> str:
return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}"
@ -527,12 +556,13 @@ class DBSpendUpdateWriter:
get_request_model_access_groups,
)
llm_router: Final = get_llm_router()
payload: Final = get_logging_payload(
kwargs=kwargs,
response_obj=completion_response,
start_time=start_time,
end_time=end_time,
llm_router=get_llm_router(),
llm_router=llm_router,
)
payload["spend"] = response_cost or 0.0
if isinstance(payload["startTime"], datetime):
@ -547,7 +577,10 @@ class DBSpendUpdateWriter:
payload["team_id"] = team_id
if not await self._record_spend_log(
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
payload=payload,
prisma_client=prisma_client,
disable_spend_logs=disable_spend_logs,
llm_router=llm_router,
):
return False
@ -608,14 +641,18 @@ class DBSpendUpdateWriter:
return True
async def _record_spend_log(
self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool
self,
payload: SpendLogsPayload,
prisma_client: "PrismaClient | None",
disable_spend_logs: bool,
llm_router: "Router | None" = None,
) -> bool:
if prisma_client is not None and _is_batch_cost_row(payload):
return await self._claim_batch_cost_spend_log(
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
)
if disable_spend_logs is False:
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client)
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client, llm_router=llm_router)
return True
async def _claim_batch_cost_spend_log(
@ -1462,6 +1499,7 @@ class DBSpendUpdateWriter:
payload: dict | SpendLogsPayload,
prisma_client: PrismaClient | None = None,
spend_logs_url: str | None = os.getenv("SPEND_LOGS_URL"),
llm_router: "Router | None" = None,
) -> PrismaClient | None:
verbose_proxy_logger.debug(
"Writing spend log to db - request_id: {}, spend: {}".format(
@ -1472,7 +1510,7 @@ class DBSpendUpdateWriter:
from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush
await enqueue_spend_logs(prisma_client, (payload,))
if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES:
if _turn_chains_through_the_db(_SPEND_METADATA_ADAPTER.validate_python(payload), llm_router):
request_spend_log_flush(prisma_client)
else:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")

View file

@ -1538,7 +1538,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
def _unmask_responses_api_completed_chunk(self, chunk: object, pii_tokens: dict[str, str]) -> None:
"""
Unmask PII tokens in-place for a ``response.completed`` Responses API event.
Unmask PII tokens in-place for a ``response.completed`` / ``response.incomplete`` Responses API event.
The chunk carries a ``response`` attribute (ResponsesAPIResponse) whose
``output`` list holds message items. Each item has a ``content`` list of
@ -1598,7 +1598,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
yield buffered_chunk
remaining_chunks = []
chunk_type = getattr(chunk, "type", None)
if chunk_type == "response.completed" and pii_tokens:
if chunk_type in ("response.completed", "response.incomplete") and pii_tokens:
self._unmask_responses_api_completed_chunk(chunk, pii_tokens)
saw_non_chat_chunk = True
yield chunk

View file

@ -6,8 +6,6 @@ from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from openai.types import Batch
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
@ -16,7 +14,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import BudgetConfig, StandardLoggingPayload
@ -121,17 +119,6 @@ def model_budget_start_time_cache_key(
return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}"
def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str:
return f"{spend_key}:batch:{batch_id}"
def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None:
"""A finished batch reports its whole cost on every poll, so its id is charged once per counter."""
if response_cost <= 0 or not is_batch_retrieve_call_type(call_type):
return None
return response_obj.id if isinstance(response_obj, Batch) else None
def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None:
"""Find the `model_max_budget` entry that governs `model`, or None."""
for candidate in _budget_model_candidates(model):

View file

@ -1,9 +1,14 @@
import asyncio
import traceback
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from collections.abc import Callable, Generator, Mapping, Sequence
from contextlib import contextmanager
from contextvars import ContextVar
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import batch_cost_is_final
@ -29,6 +34,7 @@ from litellm.proxy.db.db_spend_update_writer import (
debitable_model_access_groups,
get_llm_router,
)
from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.spend_tracking.spend_event import (
ObjectMapping,
@ -49,7 +55,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
get_request_model_access_groups,
should_store_prompts_and_responses_in_spend_logs,
)
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend
from litellm.types.utils import (
CallTypes,
LiteLLMBatch,
@ -61,6 +67,37 @@ from litellm.utils import get_end_user_id_for_cost_tracking
if TYPE_CHECKING:
from litellm.proxy.utils import ProxyLogging
class BatchAccountingOutcome:
def __init__(self, call_id: str) -> None:
self.call_id: Final = call_id
self.write_started = False
self.failed = False
@property
def retryable_failure(self) -> bool:
return self.failed and not self.write_started
_BATCH_ACCOUNTING: Final[ContextVar[BatchAccountingOutcome | None]] = ContextVar("batch_accounting", default=None)
_BATCH_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object])
@contextmanager
def track_batch_accounting(call_id: str) -> Generator[BatchAccountingOutcome, None, None]:
outcome: Final = BatchAccountingOutcome(call_id)
token: Final = _BATCH_ACCOUNTING.set(outcome)
try:
yield outcome
finally:
_BATCH_ACCOUNTING.reset(token)
def _batch_accounting_for(call_id: object) -> BatchAccountingOutcome | None:
outcome: Final = _BATCH_ACCOUNTING.get()
return outcome if outcome is not None and call_id == outcome.call_id else None
_UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
{
CallTypes.pass_through.value,
@ -91,6 +128,79 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter:
return proxy_logging_obj.db_spend_update_writer
def _billed_failed_attempts(request_data: Mapping[str, object]) -> tuple[BilledFailedAttempt, ...]:
"""The earlier billed attempts ``_failure_fields_to_lift`` stashed on the request, if any."""
attempts: Final = request_data.get("billed_failed_attempts")
if not isinstance(attempts, tuple):
return ()
lifted: Final = cast("tuple[object, ...]", attempts) # cast-ok: lifted by _failure_fields_to_lift, checked below
return tuple(attempt for attempt in lifted if isinstance(attempt, BilledFailedAttempt))
def _failure_spend_rows(
request_data: Mapping[str, object], final_cost: float, late_guardrails_snapshot: object
) -> tuple[tuple[float, Mapping[str, object]], ...]:
"""One ``(response_cost, kwargs)`` per spend row a failure writes: the final attempt under the
request id, then every earlier attempt the provider also billed under its own row."""
final_row: Final[Mapping[str, object]] = (
request_data
if late_guardrails_snapshot is None
else {**request_data, "standard_logging_object": late_guardrails_snapshot}
)
return (
(final_cost, final_row),
*(
(attempt.response_cost, _attempt_spend_row_kwargs(request_data, attempt, position))
for position, attempt in enumerate(_billed_failed_attempts(request_data))
),
)
def _attempt_spend_row_kwargs(
request_data: Mapping[str, object], attempt: BilledFailedAttempt, position: int
) -> Mapping[str, object]:
"""The final attempt's row kwargs re-keyed for one earlier billed attempt: its own
``request_id`` (``<id>-attempt-<position>``, since ``request_id`` is the spend-log primary key
and the retry count restarts after a fallback hop), its estimated usage and cost, its retry
count under ``attempted_retries`` and the deployment that billed it under ``model_info.id``."""
suffix: Final = f"-attempt-{position}"
raw_params: Final = cast("Mapping[str, object]", request_data.get("litellm_params") or {}) # cast-ok: request dict
litellm_params: Final = dict(raw_params)
metadata: Final = {
**cast("Mapping[str, object]", litellm_params.get("metadata") or {}), # cast-ok: request dict
"attempted_retries": attempt.attempted_retries,
}
metadata.pop("previous_models", None)
if attempt.deployment_id is not None:
metadata["model_info"] = {
**cast("Mapping[str, object]", metadata.get("model_info") or {}), # cast-ok: router dict
"id": attempt.deployment_id,
}
litellm_params["metadata"] = metadata
row: Final = {
**request_data,
"litellm_params": litellm_params,
"combined_usage_object": attempt.usage,
"response_cost": attempt.response_cost,
}
row.pop("billed_failed_attempts", None)
call_id: Final = request_data.get("litellm_call_id") or litellm_params.get("litellm_call_id")
if isinstance(call_id, str):
row["litellm_call_id"] = f"{call_id}{suffix}"
standard_logging_object: Final = request_data.get("standard_logging_object")
if isinstance(standard_logging_object, dict):
payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict
if isinstance(payload.get("id"), str):
row["standard_logging_object"] = {**payload, "id": f"{payload['id']}{suffix}"}
return row
def _str_keyed_fields(value: object) -> Mapping[str, object]:
if not isinstance(value, Mapping):
return {}
return cast("Mapping[str, object]", value) # cast-ok: payload dicts are str-keyed
class _ProxyDBLogger(CustomLogger):
def __init__(
self,
@ -258,20 +368,40 @@ class _ProxyDBLogger(CustomLogger):
existing_metadata.get("standard_logging_guardrail_information")
)
await self._spend_writer().update_database(
token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
response_cost=recovered_response_cost,
user_id=user_api_key_dict.user_id,
end_user_id=user_api_key_dict.end_user_id,
team_id=user_api_key_dict.team_id,
kwargs=request_data,
completion_response=original_exception,
await self._write_failure_spend_rows(
user_api_key_dict=user_api_key_dict,
original_exception=original_exception,
start_time=actual_start_time,
end_time=datetime.now(),
org_id=user_api_key_dict.org_id,
project_id=user_api_key_dict.project_id,
spend_rows=_failure_spend_rows(
cast("Mapping[str, object]", request_data), # cast-ok: request dict, copied per row below
recovered_response_cost,
None,
),
)
async def _write_failure_spend_rows(
self,
user_api_key_dict: UserAPIKeyAuth,
original_exception: Exception,
start_time: datetime,
spend_rows: Sequence[tuple[float, Mapping[str, object]]],
) -> None:
"""One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair."""
for row_cost, row_kwargs in spend_rows:
await self._spend_writer().update_database(
token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
response_cost=row_cost,
user_id=user_api_key_dict.user_id,
end_user_id=user_api_key_dict.end_user_id,
team_id=user_api_key_dict.team_id,
kwargs=dict(row_kwargs),
completion_response=original_exception,
start_time=start_time,
end_time=datetime.now(),
org_id=user_api_key_dict.org_id,
project_id=user_api_key_dict.project_id,
)
@log_db_metrics
async def _PROXY_track_cost_callback(
self,
@ -286,6 +416,11 @@ class _ProxyDBLogger(CustomLogger):
update_cache,
)
accounting: Final = (
_batch_accounting_for(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get("litellm_call_id"))
if _BATCH_ACCOUNTING.get() is not None
else None
)
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
try:
verbose_proxy_logger.debug(
@ -306,6 +441,8 @@ class _ProxyDBLogger(CustomLogger):
)
_write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
if accounting is not None and budget_reservation is not None:
accounting.write_started = True
if (
isinstance(completion_response, LiteLLMBatch)
and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
@ -329,8 +466,21 @@ class _ProxyDBLogger(CustomLogger):
key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None))
end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None)
sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
raw_cost: Final = _str_keyed_fields(
sl_object if sl_object is not None else kwargs # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload
).get("response_cost")
priced_cost: Final = raw_cost if isinstance(raw_cost, (int, float)) else None
response_cost: Final = (
sl_object.get("response_cost", None) if sl_object is not None else kwargs.get("response_cost", None)
await apply_search_free_tier(
kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload
sl_object=sl_object, # pyright: ignore[reportUnknownArgumentType] # payload carries untyped lists
response_cost=priced_cost,
completion_response=completion_response,
counter=proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache,
now=datetime.now(timezone.utc),
)
if litellm.apply_search_free_tier and sl_object is not None and priced_cost is not None
else priced_cost
)
tags: Final = _get_request_tags_for_cost_tracking(
sl_object=sl_object,
@ -448,6 +598,8 @@ class _ProxyDBLogger(CustomLogger):
f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
)
except Exception as e:
if accounting is not None:
accounting.failed = True
failing_model: Final = kwargs.get("model", "")
failing_call_type: Final = kwargs.get("call_type", "")
error_msg: Final = (
@ -698,7 +850,15 @@ async def _update_database_and_spend_counters(
budget_reservation=budget_reservation, response_cost=response_cost
)
try:
charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
writer: Final = proxy_logging_obj.db_spend_update_writer
accounting: Final = (
_batch_accounting_for(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get("litellm_call_id"))
if _BATCH_ACCOUNTING.get() is not None
else None
)
if accounting is not None:
accounting.write_started = True
charged: Final = await writer.update_database(
token=user_api_key,
response_cost=response_cost,
user_id=user_id,

View file

@ -0,0 +1,123 @@
from collections.abc import MutableMapping
from datetime import datetime
from types import MappingProxyType
from typing import Final, Protocol
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests_from_usage
from litellm.llms.gemini.cost_calculator import cost_per_web_search_request
from litellm.types.utils import CostBreakdown, ModelInfo, SearchFreeRequests, StandardLoggingPayload, Usage
_PERIOD_DURATION: Final = MappingProxyType({"day": "1d", "month": "1mo"})
_PERIOD_STAMP: Final = MappingProxyType({"day": "%Y-%m-%d", "month": "%Y-%m"})
_TTL_GRACE_SECONDS: Final = 3600
_COUNTED_KWARG: Final = "search_free_tier_counted"
_COST_BREAKDOWN: Final = TypeAdapter(CostBreakdown)
class SharedCounter(Protocol):
async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: ...
class _UsageCarrier(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
usage: Usage | None = None
def freed_units(allowance: int, units: int, total_after: int) -> int:
"""Units of this request that landed under the allowance, given the shared total after they were added."""
return max(0, min(units, allowance - (total_after - units)))
def search_free_tier_key(provider: str, allowance: SearchFreeRequests, now: datetime) -> str:
period_stamp: Final = now.strftime(_PERIOD_STAMP[allowance["period"]])
return f"litellm:search_free_tier:{provider}:{allowance['shared_group']}:{period_stamp}"
def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int:
reset: Final = get_next_standardized_reset_time(_PERIOD_DURATION[allowance["period"]], now, "UTC")
return int((reset - now).total_seconds()) + _TTL_GRACE_SECONDS
def _usage_of(completion_response: object) -> Usage | None:
try:
return _UsageCarrier.model_validate(completion_response, from_attributes=True).usage
except ValidationError:
return None
def _grounding_queries(usage: Usage) -> int:
details: Final = usage.prompt_tokens_details
from_details: Final = getattr(details, "web_search_requests", None) if details is not None else None
return from_details or get_web_search_requests_from_usage(usage) or 0
def _billable_units(usage: Usage, model_info: ModelInfo) -> int:
queries: Final = _grounding_queries(usage)
if queries == 0:
return 0
return queries if model_info.get("web_search_billing_unit") == "per_query" else 1
def _record_freed_cost(
kwargs: MutableMapping[str, object], # mutable-ok: writes the reduced cost where the spend row reads it
sl_object: StandardLoggingPayload,
original_cost: float,
freed_cost: float,
) -> float:
reduced_cost: Final = original_cost - freed_cost
breakdown: Final = sl_object["cost_breakdown"] or CostBreakdown(total_cost=original_cost)
sl_object["cost_breakdown"] = _COST_BREAKDOWN.validate_python( # rebind-ok: the spend row reads this payload
{
**breakdown,
"total_cost": reduced_cost,
"original_cost": breakdown.get("original_cost", original_cost),
"discount_amount": breakdown.get("discount_amount", 0.0) + freed_cost,
}
)
sl_object["response_cost"] = reduced_cost # rebind-ok: the spend row reads this payload
kwargs["response_cost"] = reduced_cost # rebind-ok: the spend row reads kwargs["response_cost"]
return reduced_cost
async def apply_search_free_tier(
kwargs: MutableMapping[str, object], # mutable-ok: writes the reduced cost where the spend row reads it
sl_object: StandardLoggingPayload,
response_cost: float,
completion_response: object,
counter: SharedCounter | None,
now: datetime,
) -> float:
"""Zero the grounding queries that landed under the model's shared free allowance, everything else is unchanged."""
model_info: Final = sl_object["model_map_information"]["model_map_value"]
allowance: Final = model_info.get("search_free_requests") if model_info is not None else None
usage: Final = _usage_of(completion_response)
if model_info is None or allowance is None or usage is None or sl_object["cache_hit"]:
return response_cost
units: Final = _billable_units(usage, model_info)
if units == 0 or kwargs.get(_COUNTED_KWARG):
return response_cost
if counter is None:
verbose_proxy_logger.warning(
"apply_search_free_tier is on but the proxy has no Redis, grounding is billed at full price"
)
return response_cost
provider: Final = sl_object["custom_llm_provider"] or model_info["litellm_provider"]
key: Final = search_free_tier_key(provider, allowance, now)
try:
total_after: Final = await counter.async_increment(key, units, ttl=search_free_tier_ttl(allowance, now))
except Exception: # noqa: BLE001 # any counter failure means full price, never a lost spend row
verbose_proxy_logger.warning(
"search free tier counter %s is unavailable, grounding is billed at full price", key
)
return response_cost
kwargs[_COUNTED_KWARG] = True # rebind-ok: the proxy runs this callback twice per request, count it once
freed: Final = freed_units(allowance["count"], units, int(total_after))
if freed == 0:
return response_cost
freed_cost: Final = min(response_cost, cost_per_web_search_request(usage, model_info) * freed / units)
return _record_freed_cost(kwargs, sl_object, response_cost, freed_cost)

View file

@ -358,8 +358,9 @@ _ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY: Final = "allow_client_mess
# ``litellm.model_cost`` map. Both effects belong to deployment configuration,
# not to user-supplied request bodies, so the proxy strips them before they
# reach the call path. Built from the Pydantic model so newly-added pricing
# fields are covered automatically.
_CLIENT_PRICING_CONTROL_FIELDS: Final = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
# fields are covered automatically. ``base_model`` selects which cost-map entry
# prices the row, so a caller naming a cheaper model lowers their own spend.
_CLIENT_PRICING_CONTROL_FIELDS: Final = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) | {"base_model"}
# ``model_info`` carries the same pricing fields when read by
# ``use_custom_pricing_for_model``; strip from metadata for the same reason.
# ``standard_logging_guardrail_information`` is proxy-written telemetry summed
@ -373,6 +374,7 @@ _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
{
"attempted_fallbacks",
"original_model_group",
"previous_models",
"request_retry_count",
CLIENT_OUTPUT_CEILING_METADATA_KEY,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,

View file

@ -46,11 +46,13 @@ from typing import (
overload,
)
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm import _custom_logger_compatible_callbacks_literal
from litellm.constants import (
DEFAULT_MODEL_CREATED_AT_TIME,
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
MAX_TEAM_LIST_LIMIT,
PROXY_REJECTED_BEFORE_ROUTING_KEY,
@ -143,7 +145,8 @@ from litellm.litellm_core_utils.core_helpers import (
independent_snapshot,
is_expected_client_error,
)
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.served_output_texts import (
@ -236,6 +239,7 @@ from litellm.repositories.verification_token_repository import (
VerificationTokenRepository,
)
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.deployment_pricing import pricing_models
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
@ -248,6 +252,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionRe
from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
from litellm.utils import (
_add_custom_logger_callback_to_specific_event, # pyright: ignore[reportPrivateUsage] # only string-to-logger helper
_get_base_model_from_metadata, # pyright: ignore[reportPrivateUsage] # main exposes it as get_base_model_from_metadata
)
if TYPE_CHECKING:
@ -939,20 +944,48 @@ def _failure_usage_to_lift(
model_call_details: Mapping[str, object],
request_body: Mapping[str, object],
dispatched: bool,
provider_billed: bool = False,
deployment_id: str | None = None,
base_model: str | None = None,
) -> tuple[object, object] | None:
"""A stream that broke mid-flight still billed the provider for the chunks
already delivered; the streaming handler stashes that recovered usage and
cost in model_call_details, so prefer it. Otherwise a request that was
dispatched to a provider and failed without upstream usage gets an
estimated input-side Usage with zero cost. The raw request body backfills
the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a
chat-completions provider) without filling optional_params. Returns the
estimated input-side Usage. When the failure is one the provider bills
anyway (``provider_billed``: a read timeout or a client disconnect, where
the provider had the prompt and kept working after LiteLLM hung up) that
usage is priced at the model's input rate, the way the success path prices it:
at the rate the router registered under ``deployment_id`` (``Logging.get_router_model_id``)
when the deployment declares one, else the cost map's rate for ``base_model`` (an Azure
deployment's ``model_info.base_model``) or the deployment model; any other post-dispatch
failure stays at zero cost. Returns the
(combined_usage_object, response_cost) pair to lift, or None."""
recovered_usage: Final = model_call_details.get("combined_usage_object")
if recovered_usage is not None:
return recovered_usage, model_call_details.get("response_cost")
if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL):
if not dispatched:
return None
if model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL) and not provider_billed:
# The proxy-only error handler stamps this flag on a client disconnect too, even though the
# request had already reached the provider, so a billed failure keeps its estimate.
return None
estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body)
if estimated_usage is None:
return None
if not provider_billed:
return estimated_usage, 0.0
return estimated_usage, _estimated_input_cost(
model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model
)
def _estimated_prompt_usage(
model_call_details: Mapping[str, object], request_body: Mapping[str, object]
) -> Usage | None:
"""The input-only Usage estimate for the prompt this request sent, or None when the call type has no
estimable input. The raw request body backfills the system prompt when the SDK bridges an endpoint
(e.g. /v1/messages on a chat-completions provider) without filling optional_params."""
if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES:
return None
optional_params: Final = model_call_details.get("optional_params")
@ -962,14 +995,205 @@ def _failure_usage_to_lift(
else None
)
system_input: Final = dispatched_system or request_body.get("system") or request_body.get("instructions")
estimated_usage: Final = _estimate_dispatched_failure_usage(
return _estimate_dispatched_failure_usage(
model=str(model_call_details.get("model") or ""),
request_input=model_call_details.get("messages"),
system_input=system_input,
)
if estimated_usage is None:
@dataclass(frozen=True)
class BilledFailedAttempt:
"""An earlier retry attempt the provider also billed, written as its own spend row."""
attempted_retries: int
deployment_id: str | None
usage: Usage
response_cost: float
_NO_RETRY_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
def _retry_metadata_bucket(litellm_params: Mapping[str, object] | None) -> Mapping[str, object]:
"""The request metadata bucket the router writes retry state into, ``litellm_metadata`` when the
route carries one (/messages, /responses) and ``metadata`` otherwise, same order as the deployment id."""
if litellm_params is None:
return _NO_RETRY_METADATA
for metadata_key in ("litellm_metadata", "metadata"):
metadata = litellm_params.get(metadata_key)
if isinstance(metadata, dict):
return cast("Mapping[str, object]", metadata) # cast-ok: untyped router dict, keys read with .get
return _NO_RETRY_METADATA
def _attempt_number(record: Mapping[str, object], index: int) -> int:
"""The attempt a retry record describes: its own ``attempted_retries`` count, else its position."""
recorded: Final = record.get("attempted_retries")
return recorded if isinstance(recorded, int) else index
def _retry_attempt_records(litellm_params: Mapping[str, object] | None) -> tuple[Mapping[str, object], ...]:
"""The router's ``RetryAttemptRecord`` breadcrumbs, one per failed attempt, read from the same bucket."""
if litellm_params is None:
return ()
for metadata_key in ("litellm_metadata", "metadata"):
metadata = litellm_params.get(metadata_key)
if not isinstance(metadata, dict):
continue
records = cast("Mapping[str, object]", metadata).get("previous_models") # cast-ok: untyped router dict
if isinstance(records, (list, tuple)):
return tuple(
cast("Mapping[str, object]", record) # cast-ok: RetryAttemptRecord TypedDict, read with .get
for record in cast("Sequence[object]", records) # cast-ok: untyped router tuple
if isinstance(record, dict)
)
return ()
def _record_deployment_id(record: Mapping[str, object]) -> str | None:
deployment_id: Final = record.get("deployment_id")
return deployment_id if isinstance(deployment_id, str) else None
def _attempt_identity(record: Mapping[str, object]) -> tuple[str | None, int | None]:
"""What tells one provider call from another across the request: the deployment that served it
and the retry count within its model group. Two records with the same identity describe the
same call (the fallback handler re-logs the attempt that exhausted the previous group)."""
attempted_retries: Final = record.get("attempted_retries")
return _record_deployment_id(record), attempted_retries if isinstance(attempted_retries, int) else None
def _timed_out_earlier_records(
litellm_params: Mapping[str, object] | None, final_identity: tuple[str | None, int | None]
) -> tuple[Mapping[str, object], ...]:
"""The ``Timeout`` breadcrumbs for calls other than the final attempt, one per provider call."""
by_identity: Final = {
_attempt_identity(record): record
for record in _retry_attempt_records(litellm_params)
if record.get("exception_type") == litellm.Timeout.__name__
}
return tuple(record for identity, record in by_identity.items() if identity != final_identity)
def _record_pricing(
model_call_details: Mapping[str, object], record: Mapping[str, object], base_model: str | None
) -> tuple[Mapping[str, object], str | None]:
"""The ``(model_call_details, base_model)`` to price one earlier attempt with: the record's own
deployment model and ``base_model`` when it served a different model than the final attempt
(the provider is then inferred from that model), else the final attempt's."""
record_model: Final = record.get("model")
if not isinstance(record_model, str) or record_model == model_call_details.get("model"):
return model_call_details, base_model
record_base_model: Final = record.get("base_model")
return (
{**model_call_details, "model": record_model, "custom_llm_provider": None},
record_base_model if isinstance(record_base_model, str) else None,
)
def _billed_earlier_attempt(
model_call_details: Mapping[str, object],
record: Mapping[str, object],
index: int,
estimated_usage: Usage,
base_model: str | None,
) -> BilledFailedAttempt:
pricing_details, pricing_base_model = _record_pricing(model_call_details, record, base_model)
return BilledFailedAttempt(
attempted_retries=_attempt_number(record, index),
deployment_id=_record_deployment_id(record),
usage=estimated_usage,
response_cost=_estimated_input_cost(
pricing_details,
estimated_usage.prompt_tokens,
deployment_id=_record_deployment_id(record),
base_model=pricing_base_model,
),
)
def _billed_earlier_attempts(
model_call_details: Mapping[str, object],
litellm_params: Mapping[str, object] | None,
estimated_usage: Usage,
base_model: str | None,
deployment_id: str | None,
) -> tuple[BilledFailedAttempt, ...]:
"""Every earlier retry attempt the provider also billed. The router records one
``RetryAttemptRecord`` per failed attempt in ``metadata.previous_models``, the final attempt
included, which the failure row itself bills. The final attempt is the record with the
request's own deployment and retry count (``attempted_retries`` restarts at 0 after a fallback
hop, so the count alone cannot tell the groups apart). An earlier attempt that ended in a
``Timeout`` had the same prompt accepted by its deployment, so it is billed like the final
attempt, at that deployment's input rate (else the cost map's rate for the record's own
``base_model`` or model, which a fallback hop can make a different model than the final one).
Only the exception class survives on the record, so an earlier connect timeout is billed too;
any other failure class is not."""
final_attempt: Final = _retry_metadata_bucket(litellm_params).get("attempted_retries")
final_identity: Final = (deployment_id, final_attempt if isinstance(final_attempt, int) else None)
return tuple(
_billed_earlier_attempt(model_call_details, record, index, estimated_usage, base_model)
for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity))
)
def _deployment_priced_input_cost(
model: str, custom_llm_provider: str | None, prompt_tokens: int, deployment_pricing: ModelInfo
) -> float:
"""Input price for the estimated prompt at the deployment's registered rates, read through
``deployment_pricing_model_info`` and priced by ``generic_cost_per_token`` the way the success
path prices them, so a long-context tier the deployment declares (``input_cost_per_token_above_200k_tokens``)
is honoured instead of flattened to the standard rate. An entry it cannot price keeps the row at $0."""
usage: Final = Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens)
try:
prompt_cost, _ = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider or str(deployment_pricing.get("litellm_provider") or ""),
model_info=deployment_pricing,
)
except Exception as exc: # noqa: BLE001 # a malformed registered entry must not drop the failure row
verbose_proxy_logger.debug("failure spend estimate: could not price deployment of %s: %s", model, exc)
return 0.0
return prompt_cost
def _priced_input_cost(pricing_model: str, custom_llm_provider: str | None, prompt_tokens: int) -> float | None:
try:
estimated_input_cost, _ = litellm.cost_per_token(
model=pricing_model,
custom_llm_provider=custom_llm_provider,
prompt_tokens=prompt_tokens,
)
except Exception as exc: # noqa: BLE001 # cost_per_token raises BadRequestError for an unmapped model and the row must still be written
verbose_proxy_logger.debug("failure spend estimate: could not price %s: %s", pricing_model, exc)
return None
return estimated_usage, 0.0
return estimated_input_cost
def _estimated_input_cost(
model_call_details: Mapping[str, object],
prompt_tokens: int,
deployment_id: str | None,
base_model: str | None,
) -> float:
"""Input price for the estimated prompt: the deployment's registered rates when it has any (the id
is ``Logging.get_router_model_id()``, never the flattened ``model_call_details`` copy a request body
key can overwrite), else the cost map's rate for ``base_model``, else the deployment model, the same
candidate order ``completion_cost`` walks (an Azure deployment name is not in the cost map, its
``base_model`` is; a ``base_model`` the cost map does not know falls through to the deployment
model). A model no candidate can price keeps the row at $0 rather than dropping the failure row."""
model: Final = str(model_call_details.get("model") or "")
custom_llm_provider: Final = str(model_call_details.get("custom_llm_provider") or "") or None
deployment_pricing: Final = deployment_pricing_model_info(model_id=deployment_id, deployment_model=model)
if deployment_pricing is not None:
return _deployment_priced_input_cost(model, custom_llm_provider, prompt_tokens, deployment_pricing)
pricing_models_to_try: Final = pricing_models(base_model, model)
candidate_costs: Final = (
_priced_input_cost(pricing_model, custom_llm_provider, prompt_tokens) for pricing_model in pricing_models_to_try
)
return next((cost for cost in candidate_costs if cost is not None), 0.0)
_EMPTY_LIFT: Final = MappingProxyType({})
@ -1085,7 +1309,38 @@ def _call_type_for_route(route: str | None) -> str | None:
_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException)
def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
_EXCEPTION_CHAIN_MAX_DEPTH: Final = 16
def _exception_chain(exception: BaseException) -> tuple[BaseException, ...]:
chain: Final[list[BaseException]] = [exception] # mutable-ok: built once, returned frozen
for _ in range(_EXCEPTION_CHAIN_MAX_DEPTH):
following = chain[-1].__cause__ or chain[-1].__context__
if following is None or following in chain:
break
chain.append(following)
return tuple(chain)
def _timed_out_reading_from_provider(exception: Exception) -> bool:
"""A read timeout means the request was sent and the provider was working on it; a connect
timeout (``httpx.ConnectTimeout``) never reached the provider, so it is not billed."""
return isinstance(exception, litellm.Timeout) and any(
isinstance(cause, httpx.ReadTimeout) for cause in _exception_chain(exception)
)
def _client_disconnected(exception: Exception) -> bool:
"""The 499 the proxy itself raises when the caller hung up while the provider was still working on the
request. A provider's own 499 arrives as a litellm exception and is a plain upstream failure."""
return isinstance(exception, HTTPException) and exception.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
def _provider_billed_the_failure(exception: Exception) -> bool:
return _timed_out_reading_from_provider(exception) or _client_disconnected(exception)
def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: bool = False) -> Mapping[str, object]:
"""Failure-path callbacks run after ``litellm_logging_obj`` is popped from
request_data (it is not serialisable), so the caller merges these fields
onto request_data first: the request start and first-handoff instants for
@ -1097,10 +1352,28 @@ def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str,
return _EMPTY_LIFT
_model_call_details: Final = getattr(_logging_obj, "model_call_details", {})
_first_handoff: Final = _model_call_details.get("first_api_call_start_time")
_litellm_params: Final = getattr(_logging_obj, "litellm_params", None)
_base_model: Final = _get_base_model_from_metadata(model_call_details=_model_call_details)
_deployment_id: Final = _logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None
_success_already_logged: Final = isinstance(_logging_obj, Logging) and _logging_obj.async_success_log_enqueued
_billed: Final = provider_billed and not _success_already_logged
_usage_to_lift: Final = _failure_usage_to_lift(
model_call_details=_model_call_details,
request_body=request_data,
dispatched=_first_handoff is not None,
provider_billed=_billed,
deployment_id=_deployment_id,
base_model=_base_model,
)
_prompt_estimate: Final = (
_estimated_prompt_usage(_model_call_details, request_data)
if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage")
else None
)
_earlier_attempts: Final = (
_billed_earlier_attempts(_model_call_details, _litellm_params, _prompt_estimate, _base_model, _deployment_id)
if _prompt_estimate is not None
else ()
)
_entries: Final = (
("start_time", _model_call_details.get("start_time")),
@ -1109,6 +1382,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str,
("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]),
("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)),
("standard_logging_object", _model_call_details.get("standard_logging_object")),
("billed_failed_attempts", _earlier_attempts or None),
)
return MappingProxyType({key: value for key, value in _entries if value is not None})
@ -3252,7 +3526,11 @@ class ProxyLogging:
original_exception=original_exception,
)
request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data))
request_data.update(
await offload_token_count(_failure_fields_to_lift)(
request_data, provider_billed=_provider_billed_the_failure(original_exception)
)
)
# Remove before callbacks iterate — not serialisable
request_data.pop("litellm_logging_obj", None)
@ -7790,7 +8068,11 @@ MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
def _is_transient_spend_log_write_error(e: Exception) -> bool:
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
return (
PrismaDBExceptionHandler.is_database_transport_error(e)
or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e)
or PrismaDBExceptionHandler.is_deadlock_error(e)
)
async def _create_spend_logs_with_poison_isolation(
@ -7824,7 +8106,7 @@ async def _create_spend_logs_with_poison_isolation(
await repo.table.create_many(data=rows, skip_duplicates=True)
return failure_budget
except Exception as e:
if not PrismaDBExceptionHandler.is_prisma_data_error(e):
if not (PrismaDBExceptionHandler.is_prisma_data_error(e) or isinstance(e, UnicodeEncodeError)):
raise
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
raise

View file

@ -14,7 +14,7 @@ from litellm.responses.litellm_completion_transformation.custom_tools import (
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
PART_UNION_TYPES,
@ -34,6 +34,7 @@ from litellm.types.llms.openai import (
ReasoningSummaryTextDoneEvent,
ResponseCompletedEvent,
ResponseCreatedEvent,
ResponseIncompleteEvent,
ResponseInProgressEvent,
ResponseInputParam,
ResponsesAPIOptionalRequestParams,
@ -912,7 +913,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
raise StopAsyncIteration
self.finished = self.is_stream_finished()
response_completed_event: Final = self._emit_response_completed_event(self.litellm_model_response)
response_completed_event: Final = self._emit_terminal_response_event(self.litellm_model_response)
if response_completed_event:
# Latch so wrappers (FallbackResponsesStreamWrapper) + proxy
# container-ownership hook can read completed_response.
@ -1258,7 +1259,46 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
)
return reasoning_aligned
def _emit_response_completed_event(self, litellm_model_response: ModelResponse) -> ResponseCompletedEvent | None:
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
"""
The response.incomplete event to bill when the client went away mid-stream.
Built from the chat completion chunks relayed so far, so the usage is what
stream_chunk_builder counts for them.
"""
if self.completed_response is not None or not self.collected_chat_completion_chunks:
return None
built: Final = stream_chunk_builder(
chunks=self.collected_chat_completion_chunks,
messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
input=self.request_input,
responses_api_request=dict(self.responses_api_request),
),
logging_obj=self.litellm_logging_obj,
)
if not isinstance(built, ModelResponse):
return None
responses_api_response: Final = (
LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
request_input=self.request_input,
chat_completion_response=built,
responses_api_request=self.responses_api_request,
)
)
partial: Final = self._with_encoded_response_id(
responses_api_response.model_copy(
update={
"id": self._cached_response_id or responses_api_response.id,
"status": "incomplete",
"output": list(self._output_with_streamed_item_ids(responses_api_response)),
}
)
)
stamp_responses_usage_cost(partial, self.litellm_logging_obj)
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
def _emit_terminal_response_event(
self, litellm_model_response: ModelResponse
) -> ResponseCompletedEvent | ResponseIncompleteEvent | None:
if litellm_model_response:
# Transform the response
responses_api_response: Final = (
@ -1278,6 +1318,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
# Encode the response ID to match non-streaming behavior
encoded_response: Final = self._with_encoded_response_id(responses_api_response)
if responses_api_response.status == "incomplete":
return ResponseIncompleteEvent(
type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
response=encoded_response,
)
return ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=encoded_response,

View file

@ -117,6 +117,44 @@ _INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_toke
)
class _EchoedResponsesRequestParams(TypedDict, total=False):
"""Responses API request fields that the response object echoes back."""
instructions: ReadOnly[str]
metadata: ReadOnly[dict[str, object]]
parallel_tool_calls: ReadOnly[bool]
temperature: ReadOnly[float]
tools: ReadOnly[list[dict[str, object]]]
top_p: ReadOnly[float]
max_output_tokens: ReadOnly[int]
previous_response_id: ReadOnly[str]
reasoning: ReadOnly[dict[str, object]]
text: ReadOnly[dict[str, object]]
truncation: ReadOnly[Literal["auto", "disabled"]]
user: ReadOnly[str]
store: ReadOnly[bool]
_ECHOED_PARAMS_ADAPTER: Final = TypeAdapter(_EchoedResponsesRequestParams)
def _is_echoable(name: str, value: object) -> bool:
try:
_ = _ECHOED_PARAMS_ADAPTER.validate_python({name: value})
except ValidationError:
return False
return True
def _echoable_request_params(request: Mapping[str, object]) -> _EchoedResponsesRequestParams:
echoable: Final = {
name: request[name]
for name in _EchoedResponsesRequestParams.__optional_keys__
if name in request and _is_echoable(name, request[name])
}
return _ECHOED_PARAMS_ADAPTER.validate_python(echoable)
@dataclass(frozen=True, slots=True)
class ResponsesToolChatForm:
chat_tools: tuple[ChatToolParam, ...]
@ -2294,7 +2332,7 @@ class LiteLLMCompletionResponsesConfig:
The corresponding responses API status value (one of ResponsesAPIStatus)
"""
if finish_reason is None:
return "completed"
return "incomplete"
# Map finish reasons to status
if finish_reason in ["stop", "tool_calls", "function_call"]:
@ -2309,11 +2347,18 @@ class LiteLLMCompletionResponsesConfig:
def _incomplete_details_for_finish_reason(
finish_reason: str | None,
existing: IncompleteDetails | None,
completion_tokens: int | None = None,
max_output_tokens: int | None = None,
) -> IncompleteDetails | None:
if existing is not None:
return existing
if finish_reason is None:
return None
hit_cap: Final = (
completion_tokens is not None
and max_output_tokens is not None
and completion_tokens >= max_output_tokens
)
return IncompleteDetails(reason="max_output_tokens") if hit_cap else None
reason: Final = _INCOMPLETE_REASON_BY_FINISH_REASON.get(finish_reason)
return IncompleteDetails(reason=reason) if reason is not None else None
@ -2433,9 +2478,13 @@ class LiteLLMCompletionResponsesConfig:
if choices and len(choices) > 0:
finish_reason = choices[0].finish_reason
echoed: Final = _echoable_request_params(responses_api_request)
chat_usage: Final = getattr(chat_completion_response, "usage", None)
incomplete_details: Final = LiteLLMCompletionResponsesConfig._incomplete_details_for_finish_reason(
finish_reason=finish_reason,
existing=getattr(chat_completion_response, "incomplete_details", None),
completion_tokens=getattr(chat_usage, "completion_tokens", None),
max_output_tokens=echoed.get("max_output_tokens"),
)
responses_api_response: Final[ResponsesAPIResponse] = ResponsesAPIResponse(
@ -2445,32 +2494,33 @@ class LiteLLMCompletionResponsesConfig:
object="response",
error=getattr(chat_completion_response, "error", None),
incomplete_details=incomplete_details,
instructions=getattr(chat_completion_response, "instructions", None),
metadata=getattr(chat_completion_response, "metadata", {}),
instructions=echoed.get("instructions"),
metadata=echoed.get("metadata") or {},
output=LiteLLMCompletionResponsesConfig._transform_chat_completion_choices_to_responses_output(
chat_completion_response=chat_completion_response,
choices=getattr(chat_completion_response, "choices", []),
responses_api_request=responses_api_request,
),
parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False),
temperature=getattr(chat_completion_response, "temperature", 0),
parallel_tool_calls=echoed.get("parallel_tool_calls", False),
temperature=echoed.get("temperature"),
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),
previous_response_id=getattr(chat_completion_response, "previous_response_id", None),
reasoning=None,
tools=echoed.get("tools") or [],
top_p=echoed.get("top_p"),
max_output_tokens=echoed.get("max_output_tokens"),
previous_response_id=echoed.get("previous_response_id"),
reasoning=echoed.get("reasoning"),
status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status(
finish_reason
),
text={},
truncation=getattr(chat_completion_response, "truncation", None),
text=echoed.get("text") or {},
truncation=echoed.get("truncation"),
usage=LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(
chat_completion_response=chat_completion_response
),
user=getattr(chat_completion_response, "user", None),
user=echoed.get("user"),
store=echoed.get("store"),
)
responses_api_response._hidden_params = getattr(chat_completion_response, "_hidden_params", {})

View file

@ -547,6 +547,22 @@ def _will_bridge_to_chat_completions(
)
def provider_stores_the_conversation(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool:
"""Whether the provider, not a LiteLLM spend log, holds this turn for a `previous_response_id` follow-up.
A bridged turn is replayed from its spend log, so the bridge needs that row in the DB, and a bridged
first MCP round runs as an internal call that writes no spend log at all.
"""
api_base: Final = _api_base_kwarg(kwargs)
return not _will_bridge_to_chat_completions(
model,
custom_llm_provider or litellm.get_llm_provider(model=model, api_base=api_base)[1],
bool(kwargs.get("use_chat_completions_api")),
kwargs.get("model_info"),
api_base,
)
@contextmanager
def _prompt_management_sees_a_provisional_message_list(
kwargs: dict[str, object], # mutable-ok: the signal is read and popped out of the caller's own kwargs

View file

@ -41,10 +41,12 @@ from litellm.types.integrations.custom_logger import converted_stream_requested
from litellm.types.llms.openai import (
PART_UNION_TYPES,
ResponseAPIUsage,
ResponseIncompleteEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
)
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
from litellm.types.utils import CallTypes
from litellm.utils import async_post_call_success_deployment_hook
@ -289,6 +291,7 @@ class BaseResponsesAPIStreamingIterator:
self._yielded_first_chunk = False
self._generated_content = ""
self._generated_tool_arguments = ""
self._created_response: ResponsesAPIResponse | None = None
self._completed_response_cached = False
self._completed_response_logged = False
self._completed_response_cache_hit: bool | None = None
@ -450,6 +453,10 @@ class BaseResponsesAPIStreamingIterator:
# Store the completed response (also for incomplete/failed so logging still fires)
_chunk_type: Final = getattr(openai_responses_api_chunk, "type", None)
openai_types = _get_openai_response_types()
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED:
_created: Final[object] = getattr(openai_responses_api_chunk, "response", None)
if isinstance(_created, ResponsesAPIResponse):
self._created_response = _created
if openai_responses_api_chunk and _chunk_type in (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
@ -481,7 +488,7 @@ class BaseResponsesAPIStreamingIterator:
else openai_responses_api_chunk.model_copy(update={"response": _billed_response})
)
self.completed_response = _terminal_chunk
_stamp_responses_usage_cost(_billed_response, self.logging_obj)
stamp_responses_usage_cost(_billed_response, self.logging_obj)
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
self._handle_logging_failed_response()
@ -874,6 +881,34 @@ class BaseResponsesAPIStreamingIterator:
except Exception:
pass
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
"""
The response.incomplete event to bill when the client went away mid-stream,
the same shape the success handler assembles for a natural end of stream.
Its response is the response.created payload (it carries the served
service_tier) plus the text and tool arguments relayed so far, with counted
usage since the provider never sent any. None when nothing was relayed yet
or the stream already reached its terminal event and billed itself.
"""
created: Final = self._created_response
generated: Final = self._generated_content + self._generated_tool_arguments
if self.completed_response is not None or created is None or not generated:
return None
usage: Final = _estimate_usage_safely(
self.model or "", self.request_data.get("input"), self.request_data, generated
)
if usage is None:
return None
partial: Final = created.model_copy(
update={
"status": "incomplete",
"usage": usage,
"output": _partial_output_items(self._generated_content),
}
)
stamp_responses_usage_cost(partial, self.logging_obj)
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
async def call_post_streaming_hooks_for_testing(
iterator: object, chunk: ResponsesAPIStreamingResponse
@ -1347,6 +1382,37 @@ def _add_text_like_part_events(
)
@runtime_checkable
class _BillsPartialStreamOnDisconnect(Protocol):
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: ...
def partial_responses_incomplete_event_on_disconnect(stream: object) -> ResponseIncompleteEvent | None:
"""Ask the stream behind a proxy or router wrapper for the partial response to bill, if it can build one."""
if not isinstance(stream, _BillsPartialStreamOnDisconnect):
return None
try:
partial: Final = stream.partial_response_on_disconnect()
except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown
verbose_logger.debug("Failed to build partial Responses API response for disconnect billing: %s", e)
return None
return partial if isinstance(partial, ResponseIncompleteEvent) else None
def _partial_output_items(generated_text: str) -> Sequence[GenericResponseOutputItem]:
if not generated_text:
return []
return [
GenericResponseOutputItem(
type="message",
id=f"msg_{uuid.uuid4()}",
status="incomplete",
role="assistant",
content=[OutputText(type="output_text", text=generated_text, annotations=[])],
)
]
def _billed_terminal_response(
response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None
) -> ResponsesAPIResponse | None:
@ -1453,7 +1519,7 @@ def _estimate_usage_safely(
return None
def _stamp_responses_usage_cost(
def stamp_responses_usage_cost(
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
) -> None:
if response_obj is None or logging_obj is None:
@ -1479,7 +1545,7 @@ def build_synthetic_response_events(
chunk_size: int,
) -> list[ResponsesAPIStreamingResponse]:
openai_types: Final = _get_openai_response_types()
_stamp_responses_usage_cost(transformed, logging_obj)
stamp_responses_usage_cost(transformed, logging_obj)
events: Final[list[ResponsesAPIStreamingResponse]] = [
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed),
@ -1696,6 +1762,8 @@ RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
"error",
]
_HISTORY_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"})
RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"})
_RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"})
@ -2479,7 +2547,7 @@ class ManagedResponsesWebSocketHandler:
@staticmethod
def _extract_response_id(completed_event: _MutableJsonObject) -> str | None:
"""
Pull the raw (decoded) response ID out of a ``response.completed`` event.
Pull the raw (decoded) response ID out of a terminal (``response.completed`` or ``response.incomplete``) event.
Returns *None* if the event doesn't contain a usable ID.
"""
resp_obj: Final = completed_event.get("response", {})
@ -2495,8 +2563,8 @@ class ManagedResponsesWebSocketHandler:
completed_event: _MutableJsonObject,
) -> list[dict[str, object]]:
"""
Convert the output items in a ``response.completed`` event into
Responses API message dicts suitable for the next turn's ``input``.
Convert output items in a terminal (``response.completed`` or ``response.incomplete``) event into Responses
API message dicts suitable for the next turn's ``input``.
"""
resp_obj: Final = completed_event.get("response", {})
if not isinstance(resp_obj, dict):
@ -2737,11 +2805,11 @@ class ManagedResponsesWebSocketHandler:
"""
Stream ``litellm.aresponses`` and forward every chunk over the WebSocket.
Captures the ``response.completed`` event type from the chunk object
directly (before serialization) to avoid a redundant JSON round-trip on
every chunk. Returns the completed event dict, or ``None``.
Captures the first terminal (``response.completed`` or ``response.incomplete``) event from the chunk object
directly (before serialization) to avoid a redundant JSON round-trip on every chunk. Returns the terminal
event dict, or ``None``.
"""
completed_event: _MutableJsonObject | None = None
terminal_event: _MutableJsonObject | None = None
stream_response: Final = await litellm.aresponses(model=model, **call_kwargs)
async for chunk in stream_response:
if chunk is None:
@ -2751,31 +2819,31 @@ class ManagedResponsesWebSocketHandler:
serialized = self._serialize_chunk(chunk)
if serialized is None:
continue
if chunk_type == "response.completed" and completed_event is None:
if chunk_type in _HISTORY_TERMINAL_EVENT_TYPES and terminal_event is None:
try:
completed_event = _load_json_object(serialized)
terminal_event = _load_json_object(serialized)
except Exception:
pass
try:
await self.websocket.send_text(serialized)
except Exception as send_exc:
verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc)
return completed_event # Client disconnected
return completed_event
return terminal_event # Client disconnected
return terminal_event
def _save_turn_history(
self,
completed_event: _MutableJsonObject | None,
terminal_event: _MutableJsonObject | None,
prior_history: list[dict[str, object]],
current_messages: list[dict[str, object]],
) -> None:
"""Store this turn in in-memory history for future previous_response_id lookups."""
if completed_event is None:
"""Store this terminal (``response.completed`` or ``response.incomplete``) event for future history lookups."""
if terminal_event is None:
return
new_response_id: Final = self._extract_response_id(completed_event)
new_response_id: Final = self._extract_response_id(terminal_event)
if not new_response_id:
return
output_msgs: Final = self._extract_output_messages(completed_event)
output_msgs: Final = self._extract_output_messages(terminal_event)
all_messages: Final = prior_history + current_messages + output_msgs
self._store_history(new_response_id, all_messages)
verbose_logger.debug(
@ -2801,7 +2869,8 @@ class ManagedResponsesWebSocketHandler:
2. Prepend those messages to the current ``input`` so the model has full
conversation context.
3. After the stream completes, extract the new response ID and output
messages from ``response.completed`` and store them in
messages from the terminal (``response.completed`` or
``response.incomplete``) event and store them in
``self._session_history`` for the next turn.
This in-memory approach avoids the async DB-write race condition that
@ -2855,13 +2924,13 @@ class ManagedResponsesWebSocketHandler:
call_kwargs.update(self.extra_kwargs)
try:
completed_event: Final = await self._stream_and_forward(model, call_kwargs)
terminal_event: Final = await self._stream_and_forward(model, call_kwargs)
except Exception as exc:
verbose_logger.exception("ManagedResponsesWS: error processing response.create: %s", exc)
await self._send_error(str(exc))
return
self._save_turn_history(completed_event, prior_history, current_messages)
self._save_turn_history(terminal_event, prior_history, current_messages)
# ------------------------------------------------------------------
# Main entry point

View file

@ -86,6 +86,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import (
is_registered_custom_provider,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES
from litellm.litellm_core_utils.ptu_pricing import (
PTU_COST_ATTRIBUTION_ENV_VAR,
declares_ptu,
@ -97,6 +98,7 @@ from litellm.litellm_core_utils.ptu_pricing import (
)
from litellm.litellm_core_utils.request_timeout_resolver import (
get_configured_request_timeout,
get_service_tier_timeout,
)
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.sensitive_data_masker import (
@ -198,6 +200,7 @@ from litellm.router_utils.cooldown_handlers import (
is_background_response_cost_poll_not_found,
is_caller_timeout_408,
)
from litellm.router_utils.deployment_pricing import deployment_pricing_models
from litellm.router_utils.fallback_event_handlers import (
MID_STREAM_FALLBACK_CONTROLS_KEY,
AttemptedFallbackTargets,
@ -317,12 +320,14 @@ from litellm.types.utils import (
StandardLoggingRoutingDecision,
Usage,
all_litellm_params,
is_server_derived_pricing_key,
shared_backend_model_info,
)
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.utils import (
CustomStreamWrapper,
EmbeddingResponse,
ModelInfoLookupMiss,
ModelResponse,
Rules,
function_setup,
@ -331,6 +336,8 @@ from litellm.utils import (
get_secret,
get_utc_datetime,
is_region_allowed,
lookup_model_info_for_registration,
model_info_with_pricing_overrides,
provider_rejectable_params,
set_live_deployment_replay,
)
@ -363,6 +370,7 @@ if TYPE_CHECKING:
from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
from litellm.types.llms.openai import (
ResponseAPIUsage,
ResponseIncompleteEvent,
ResponseInputParam,
ResponsesAPIResponse,
)
@ -470,6 +478,8 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({})
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
@ -629,6 +639,11 @@ class FallbackAwareAnthropicMessagesStream:
self._source_iterator = fallback_response
self.fallback_headers_adopted = True
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
return partial_responses_incomplete_event_on_disconnect(self._source_iterator)
def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream":
return self
@ -3242,6 +3257,7 @@ class Router:
from datetime import datetime
self._async_generator = async_generator
self._billing_source: object = source_iterator
# Mirror every attribute BaseResponsesAPIStreamingIterator.__init__
# would have set. The wrapper bypasses super().__init__ (it has no
# httpx.Response of its own and no provider config to drive), so
@ -3290,10 +3306,29 @@ class Router:
def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]:
prepared: Final = Router._prepare_fallback_hidden_params(fallback_response)
self._billing_source = fallback_response
self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} # mutable-ok: stream metadata
self.fallback_headers_adopted = True
return prepared
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
"""
The event to bill on a client disconnect. After a mid-stream hop the primary
stream's relayed tokens are folded into the fallback's event, the way
_combine_responses_fallback_usage does for a natural end of stream.
"""
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source)
if self._billing_source is source_iterator:
return current
primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator)
if current is None or primary is None:
return current or primary
if primary.response.usage is not None:
Router._combine_responses_fallback_usage(current, primary.response.usage)
return current
def __aiter__(self):
return self
@ -4033,6 +4068,7 @@ class Router:
"""Helper to get stream timeout from kwargs or deployment params"""
return (
kwargs.get("stream_timeout", None) # the params dynamically set by user
or get_service_tier_timeout(kwargs.get("service_tier"), data, "stream_timeout_by_service_tier")
or data.get("stream_timeout", None) # timeout set on litellm_params for this deployment
or self.stream_timeout # timeout set on router
or self.request_timeout # litellm_settings.request_timeout (per-attempt)
@ -4044,6 +4080,7 @@ class Router:
timeout: Final = (
kwargs.get("timeout", None) # the params dynamically set by user
or kwargs.get("request_timeout", None) # the params dynamically set by user
or get_service_tier_timeout(kwargs.get("service_tier"), data, "timeout_by_service_tier")
or data.get("timeout", None) # timeout set on litellm_params for this deployment
or data.get("request_timeout", None) # timeout set on litellm_params for this deployment
or self.request_timeout # litellm_settings.request_timeout (per-attempt)
@ -8299,17 +8336,24 @@ class Router:
When a retry or fallback happens, record which model group, deployment and attempt just failed and why,
and count it toward the request-wide num_retries_per_request cap
"""
from litellm.litellm_core_utils.get_litellm_params import _get_base_model_from_litellm_call_metadata
from litellm.types.router import RetryAttemptRecord
_metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var]
request_metadata: Final = cast("Mapping[str, object]", kwargs[_metadata_var]) # cast-ok: untyped kwargs
model_group: Final = kwargs.get("model")
model_info: Final = request_metadata.get("model_info")
deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None
base_model: Final = _get_base_model_from_litellm_call_metadata(
metadata=request_metadata
) or self._deployment_params_base_model(deployment_id if isinstance(deployment_id, str) else None)
deployment_model: Final = request_metadata.get("deployment")
attempted_retries: Final = request_metadata.get("attempted_retries")
attempt_record: Final[RetryAttemptRecord] = {
"model_group": model_group if isinstance(model_group, str) else None,
"deployment_id": deployment_id if isinstance(deployment_id, str) else None,
"model": deployment_model if isinstance(deployment_model, str) else None,
"base_model": base_model if isinstance(base_model, str) else None,
"exception_type": type(e).__name__,
"exception_string": str(e),
"attempted_retries": attempted_retries if type(attempted_retries) is int else None,
@ -8327,6 +8371,18 @@ class Router:
kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap
return kwargs
def _deployment_params_base_model(self, deployment_id: str | None) -> str | None:
"""The ``base_model`` a deployment declares under ``litellm_params`` (the Azure named-deployment
shape), the same key the success path's ``get_base_model_from_metadata`` reads."""
deployment: Final = self.get_deployment(model_id=deployment_id) if deployment_id is not None else None
if deployment is None:
return None
extra_params: Final = cast( # cast-ok: pydantic extras are untyped, base_model is isinstance-checked below
"Mapping[str, object]", deployment.litellm_params.model_extra or {}
)
base_model: Final = extra_params.get("base_model")
return base_model if isinstance(base_model, str) else None
def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int:
"""
Update deployment rpm for that minute
@ -8866,34 +8922,7 @@ class Router:
litellm_params=litellm_params,
model_info=_model_info,
)
for field in CustomPricingLiteLLMParams.model_fields:
if deployment.litellm_params.get(field) is not None:
_model_info[field] = deployment.litellm_params[field]
Router._inherit_builtin_base_rates_for_off_peak(
model_info=_model_info,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
if _model_info.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=_model_info,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_tiered_output_rate(
model_info=_model_info,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
## REGISTER MODEL INFO IN LITELLM MODEL COST MAP
Router._register_deployment_in_model_cost(
model_id=deployment.model_info.id,
model_info=_model_info,
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._register_deployment_pricing(deployment)
## Check if LLM Deployment is allowed for this deployment
if self.deployment_is_active_for_environment(deployment=deployment) is not True:
@ -9897,7 +9926,56 @@ class Router:
return (backend_key,)
@staticmethod
def _deployment_model_cost_payload(deployment: Deployment) -> dict: # mutable-ok: cost-map entry
def _inherit_builtin_service_tier_pricing(
model_info: dict, # mutable-ok: deployment cost-map entry filled in place
backend_model: str,
custom_llm_provider: str | None,
) -> None:
"""Inherit missing tier rates so a standalone entry does not fall back to custom standard rates."""
if ptu_terms(model_info) is not None and is_ptu_cost_attribution_enabled():
return
if all(model_info.get(field) is None for field in ("input_cost_per_token", "output_cost_per_token")):
return
try:
backend_info: Final = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider)
except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model
return
backend_entry: Final = litellm.model_cost.get(backend_info.get("key") or "")
if not isinstance(backend_entry, dict):
return
for field, backend_value in backend_entry.items():
if not field.endswith(SERVICE_TIER_COST_KEY_SUFFIXES):
continue
if model_info.get(field) is not None or backend_value is None:
continue
model_info[field] = copy.deepcopy(backend_value)
@staticmethod
def _cost_map_backend_model(deployment: Deployment) -> str:
model_info_base_model: Final = deployment.model_info.base_model
if isinstance(model_info_base_model, str) and model_info_base_model:
return model_info_base_model
params_base_model: Final = deployment.litellm_params.get("base_model")
if isinstance(params_base_model, str) and params_base_model:
return params_base_model
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True))
invoked_model: Final = params.get("model_id")
if (
isinstance(invoked_model, str)
and invoked_model
and (
deployment.litellm_params.custom_llm_provider == "bedrock"
or deployment.litellm_params.model.startswith("bedrock/")
)
):
profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}"
profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock")
if not isinstance(profile_info, ModelInfoLookupMiss):
return profile_model
return deployment.litellm_params.model
@staticmethod
def _deployment_model_cost_payload(deployment: Deployment) -> dict[str, object]: # mutable-ok: cost-map entry
"""The ``model_info`` a deployment contributes to ``litellm.model_cost``.
Custom pricing lives on ``litellm_params`` rather than ``model_info``, and
@ -9905,25 +9983,49 @@ class Router:
both are folded back in here. That keeps this reproducible from a
deployment alone, which is what lets a refresh rebuild the same entries.
"""
model_info: Final[dict] = deployment.model_info.model_dump(exclude_none=True) # mutable-ok: built in place
model_info: Final = dict(
_MODEL_INFO_ADAPTER.validate_python(deployment.model_info.model_dump(exclude_none=True))
)
for field in CustomPricingLiteLLMParams.model_fields:
field_value = deployment.litellm_params.get(field)
if field_value is not None:
model_info[field] = field_value
backend_model: Final = Router._cost_map_backend_model(deployment)
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True))
if (
params.get("model_id")
and backend_model != deployment.litellm_params.model
and (
deployment.litellm_params.custom_llm_provider == "bedrock"
or deployment.litellm_params.model.startswith("bedrock/")
)
):
profile_info: Final = lookup_model_info_for_registration(model=backend_model, custom_llm_provider="bedrock")
if not isinstance(profile_info, ModelInfoLookupMiss):
explicit_info: Final = _MODEL_INFO_ADAPTER.validate_python(model_info)
priced_info: Final = model_info_with_pricing_overrides(profile_info, explicit_info)
model_info.update(
{key: value for key, value in priced_info.items() if is_server_derived_pricing_key(key)}
)
Router._inherit_builtin_base_rates_for_off_peak(
model_info=model_info,
backend_model=deployment.litellm_params.model,
backend_model=backend_model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
if model_info.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=model_info,
backend_model=deployment.litellm_params.model,
backend_model=backend_model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_service_tier_pricing(
model_info=model_info,
backend_model=backend_model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_tiered_output_rate(
model_info=model_info,
backend_model=deployment.litellm_params.model,
backend_model=backend_model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
return model_info
@ -10325,19 +10427,9 @@ class Router:
)
for deployment in deployments
)
params: Final = tuple(deployment.get("litellm_params") or MappingProxyType({}) for deployment in deployments)
# base_model resolution mirrors get_router_model_info: unset or blank means the
# deployment's own model name is the cost-map key.
cost_map_keys: Final = tuple(
dict.fromkeys( # deduplicates while preserving config order
key
for key in (
model_info.get("base_model") or litellm_params.get("base_model") or litellm_params.get("model")
for model_info, litellm_params in zip(model_infos, params)
)
if isinstance(key, str) and key
)
)
raw_deployments: Final[tuple[object, ...]] = deployments
pricing_keys: Final = (deployment_pricing_models(deployment) for deployment in raw_deployments)
cost_map_keys: Final = tuple(dict.fromkeys(keys[0] for keys in pricing_keys if keys))
return DeploymentModelListingInfo(
cost_map_keys=cost_map_keys,
max_input_tokens=self._widest_configured_limit(model_infos, "max_input_tokens"),

View file

@ -27,6 +27,8 @@ from itertools import groupby
from types import MappingProxyType
from typing import Any, Final
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
@ -37,6 +39,7 @@ from litellm.litellm_core_utils.core_helpers import (
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.router_strategy.tag_based_routing import _get_tags_from_request_kwargs
from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once
from litellm.router_utils.cooldown_callbacks import (
_get_prometheus_logger_from_callbacks,
)
@ -46,6 +49,8 @@ from litellm.types.utils import BudgetConfig, GenericBudgetConfigType, StandardL
from litellm.types.utils import BudgetConfig as GenericBudgetInfo
DEFAULT_REDIS_SYNC_INTERVAL: Final = 1
_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object])
_BATCH_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float])
class _LiteLLMParamsDictView:
@ -500,14 +505,18 @@ class RouterBudgetLimiting(CustomLogger):
if standard_logging_payload is None:
raise ValueError("standard_logging_payload is required")
response_cost: Final[float] = standard_logging_payload.get("response_cost", 0)
batch_response, response_cost = _BATCH_CHARGE_VALUES.validate_python(
(response_obj, _CALLBACK_FIELDS.validate_python(standard_logging_payload).get("response_cost", 0))
)
model_id: Final[str] = str(standard_logging_payload.get("model_id", ""))
custom_llm_provider: Final[str] = kwargs.get("litellm_params", {}).get("custom_llm_provider", None)
if custom_llm_provider is None:
raise ValueError("custom_llm_provider is required")
custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider")
budget_config: Final = self._get_budget_config_for_provider(custom_llm_provider)
if budget_config:
batch_id: Final = batch_id_to_charge_once(kwargs.get("call_type"), batch_response, response_cost)
budget_config: Final = (
self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None
)
if custom_llm_provider is not None and budget_config is not None:
# increment spend for provider
spend_key: Final = f"provider_spend:{custom_llm_provider}:{budget_config.budget_duration}"
start_time_key: Final = f"provider_budget_start_time:{custom_llm_provider}"
@ -516,6 +525,7 @@ class RouterBudgetLimiting(CustomLogger):
spend_key=spend_key,
start_time_key=start_time_key,
response_cost=response_cost,
batch_id=batch_id,
)
deployment_budget_config: Final = self._get_budget_config_for_deployment(model_id)
@ -528,6 +538,7 @@ class RouterBudgetLimiting(CustomLogger):
spend_key=deployment_spend_key,
start_time_key=deployment_start_time_key,
response_cost=response_cost,
batch_id=batch_id,
)
request_tags: Final = _get_tags_from_request_kwargs(
@ -545,6 +556,7 @@ class RouterBudgetLimiting(CustomLogger):
spend_key=_tag_spend_key,
start_time_key=_tag_start_time_key,
response_cost=response_cost,
batch_id=batch_id,
)
async def _increment_spend_for_key(
@ -553,12 +565,15 @@ class RouterBudgetLimiting(CustomLogger):
spend_key: str,
start_time_key: str,
response_cost: float,
batch_id: str | None = None,
):
if budget_config.budget_duration is None:
return
current_time: Final = datetime.now(timezone.utc).timestamp()
ttl_seconds: Final = duration_in_seconds(budget_config.budget_duration)
if batch_id is not None and not await self._claim_batch_charge(spend_key, batch_id, ttl_seconds):
return
budget_start = await self._get_or_set_budget_start_time(
start_time_key=start_time_key,
@ -596,6 +611,17 @@ class RouterBudgetLimiting(CustomLogger):
verbose_router_logger.debug("Incremented spend for %s by %s", spend_key, response_cost)
async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool:
marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id)
redis_cache: Final = self.dual_cache.redis_cache
if redis_cache is not None:
polls: Final = await redis_cache.async_increment(key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True)
return polls == 1
local_polls: Final = await self.dual_cache.in_memory_cache.async_increment(
key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True
)
return local_polls == 1
async def periodic_sync_in_memory_spend_with_redis(self):
"""
Handler that triggers sync_in_memory_spend_with_redis every DEFAULT_REDIS_SYNC_INTERVAL seconds

View file

@ -1,5 +1,6 @@
#### What this does ####
# picks based on response time (for streaming, this is time to first token)
from collections.abc import Mapping
from datetime import datetime
from typing import Final
@ -9,6 +10,30 @@ from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed
def _deployment_cost_map_rates(deployment: object) -> tuple[float | None, float | None]:
"""Input and output per-token prices from the first cost-map row among the deployment's
pricing models that carries one; (None, None) when no row does, so the caller falls back to
its unmapped default."""
rows: Final = (str_keyed(_cost_map().get(model)) for model in deployment_pricing_models(deployment))
row: Final = next((row for row in rows if _declares_token_price(row)), None)
if row is None:
return None, None
return _as_rate(row.get("input_cost_per_token")), _as_rate(row.get("output_cost_per_token"))
def _cost_map() -> Mapping[str, object]:
return litellm.model_cost
def _declares_token_price(row: Mapping[str, object]) -> bool:
return row.get("input_cost_per_token") is not None or row.get("output_cost_per_token") is not None
def _as_rate(value: object) -> float | None:
return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None
class LowestCostLoggingHandler(CustomLogger):
@ -239,8 +264,8 @@ class LowestCostLoggingHandler(CustomLogger):
or _deployment.get("model_info", {}).get("rpm", None)
or float("inf")
)
item_litellm_model_name = _deployment.get("litellm_params", {}).get("model")
item_litellm_model_cost_map = litellm.model_cost.get(item_litellm_model_name, {})
deployment_view: object = _deployment
map_input_cost, map_output_cost = _deployment_cost_map_rates(deployment_view)
# check if user provided input_cost_per_token and output_cost_per_token in litellm_params
item_input_cost = None
@ -252,10 +277,10 @@ class LowestCostLoggingHandler(CustomLogger):
item_output_cost = _deployment.get("litellm_params", {}).get("output_cost_per_token")
if item_input_cost is None:
item_input_cost = item_litellm_model_cost_map.get("input_cost_per_token", 5.0)
item_input_cost = map_input_cost if map_input_cost is not None else 5.0
if item_output_cost is None:
item_output_cost = item_litellm_model_cost_map.get("output_cost_per_token", 5.0)
item_output_cost = map_output_cost if map_output_cost is not None else 5.0
# if litellm["model"] is not in model_cost map -> use item_cost = $10

View file

@ -3,6 +3,8 @@ import json
from os import PathLike
from typing import Final
from openai.types import Batch
from litellm._logging import verbose_logger
from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose
from litellm.types.utils import CallTypes
@ -189,3 +191,14 @@ def is_batch_retrieve_call_type(call_type: object) -> bool:
routing strategies' own state, must not be fed from it.
"""
return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES
def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str:
return f"{spend_key}:batch:{batch_id}"
def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None:
"""A finished batch reports its whole cost on every poll, so its id is charged once per counter."""
if response_cost <= 0 or not is_batch_retrieve_call_type(call_type):
return None
return response_obj.id if isinstance(response_obj, Batch) else None

View file

@ -0,0 +1,37 @@
from collections.abc import Iterator, Mapping
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter, ValidationError
_EMPTY: Final[Mapping[str, object]] = MappingProxyType({})
_STR_KEYED: Final = TypeAdapter(dict[str, object])
def pricing_models(base_model: object, model: object) -> tuple[str, ...]:
"""Cost-map keys to price a deployment under, ``model_info.base_model`` first, then the
deployment's own model, blanks and duplicates dropped. Every spend and routing surface that
prices a deployment walks this order, so a ``base_model`` the cost map does not know falls
through to the deployment model instead of billing or ranking at the unmapped default."""
return tuple(dict.fromkeys(_named_models(base_model, model)))
def deployment_pricing_models(deployment: object) -> tuple[str, ...]:
"""``pricing_models`` read off a raw deployment dict (``model_info.base_model``, then
``litellm_params.base_model``, then ``litellm_params.model``); malformed blocks count as unset."""
view: Final = str_keyed(deployment)
params: Final = str_keyed(view.get("litellm_params"))
info: Final = str_keyed(view.get("model_info"))
return pricing_models(info.get("base_model") or params.get("base_model"), params.get("model"))
def _named_models(*candidates: object) -> Iterator[str]:
return (candidate for candidate in candidates if isinstance(candidate, str) and candidate)
def str_keyed(value: object) -> Mapping[str, object]:
"""``value`` as a read-only str-keyed mapping, empty when it is not a dict."""
try:
return MappingProxyType(_STR_KEYED.validate_python(value))
except ValidationError:
return _EMPTY

View file

@ -1,20 +1,17 @@
"""Resolve which reasoning_effort values a deployment, and by intersection a model group, accepts.
An entry that states its levels outright in reasoning_effort_levels is read first and wins
whole, for a model whose set the per-level flags cannot express: Kimi K3 takes low, high and max,
and no flag can drop medium because medium has none. Every other entry answers through the
supports_*_reasoning_effort flags below, whose polarity mirrors how a request path reads that same
flag. medium and high are unconditional for a reasoning model. minimal and low are opt-out:
openai/chat/gpt_5_transformation.py refuses them only when the map says false. xhigh and max are
opt-in. none is opt-out everywhere except the azure gpt-5 family, whose config raises
UnsupportedParamsError without an explicit true.
An entry that states its levels outright in reasoning_effort_levels wins whole, for a model whose set
the per-level flags cannot express (Kimi K3 takes low, high and max, and no flag can drop medium).
Every other entry answers through its supports_*_reasoning_effort flags, read with the polarity the
chat gate uses: medium and high are unconditional, minimal and low are opt-out (refused only on an
explicit false), xhigh is opt-in, and none is opt-out except on the azure gpt-5 family, whose chat
config refuses it without an explicit true.
xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at
all: outside the gpt-6-astra rows every entry carrying supports_max_reasoning_effort is Claude-family,
and anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort
path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an
explicit flag is the only signal that the tier is a real one rather than litellm rounding the level
to a budget, and a missing flag costs advisory metadata rather than a rejected request.
The Responses gate reads none, minimal, low and xhigh as opt-out. Neither gate judges max, which
Chat Completions refuses and Responses honors on the gpt-6 family, so its flag is advisory and
advertised opt-in: outside the gpt-6 rows every entry carrying supports_max_reasoning_effort is
Claude-family, where anthropic/chat/transformation.py gates max on the output_config path, so an
explicit flag is the only signal that the tier is real.
A deployment the map describes with no effort flags at all resolves to None rather than to the
opt-out defaults. 689 of the map's 854 reasoning entries carry no flag, and the o-series, xai and

View file

@ -8,6 +8,7 @@ from litellm.decisions.call import (
SystemOneQuestions,
SystemOneRequestFields,
asend,
log_caller_request,
prepare_call,
reject_other_format,
send,
@ -29,10 +30,12 @@ def _prepare_systemone_call(
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> DecisionsCall:
fields: Final = SystemOneRequestFields(state=state, questions=questions)
log_caller_request(kwargs, fields)
reject_other_format(model=model, kwargs=kwargs, field_name="input", other_function="decisions")
return prepare_call(
model=model,
fields=SystemOneRequestFields(state=state, questions=questions),
fields=fields,
safety_identifier=safety_identifier,
api_key=api_key,
api_base=api_base,

View file

@ -22,7 +22,7 @@ class CacheControlMessageInjectionPoint(TypedDict):
class CacheControlToolConfigInjectionPoint(TypedDict):
"""Type for tool_config-level injection points (Bedrock)."""
"""Type for tool_config-level injection points (Anthropic and Bedrock)."""
location: Literal["tool_config"]
control: ChatCompletionCachedContent | None

View file

@ -72,6 +72,8 @@ class ProviderConnection:
request_timeout: float | None = None
force_timeout: float | None = None
stream_timeout: float | str | None = None
timeout_by_service_tier: Mapping[str, float] | None = None
stream_timeout_by_service_tier: Mapping[str, float] | None = None
max_retries: int | None = None
tenant_id: str | None = None
client_id: str | None = None

View file

@ -389,6 +389,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
otpm: int | None = None
timeout: float | str | httpx.Timeout | None = None # if str, pass in as os.environ/
stream_timeout: float | str | None = None # timeout when making stream=True calls, if str, pass in as os.environ/
# per service_tier overrides of timeout / stream_timeout, used when the request's service_tier matches a key
timeout_by_service_tier: Mapping[str, float] | None = None
stream_timeout_by_service_tier: Mapping[str, float] | None = None
max_retries: int | None = None
drop_params: bool | str | None = None
organization: str | None = None # for openai orgs
@ -573,6 +576,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
api_version: str | None
timeout: float | str | httpx.Timeout | None
stream_timeout: float | str | None
timeout_by_service_tier: ReadOnly[Mapping[str, float] | None]
stream_timeout_by_service_tier: ReadOnly[Mapping[str, float] | None]
max_retries: int | None
organization: list | str | None # for openai orgs
provider_affinity_header: ReadOnly[str | None]
@ -996,6 +1001,8 @@ class RouterModelGroupAliasItem(TypedDict):
class RetryAttemptRecord(TypedDict):
model_group: ReadOnly[str | None]
deployment_id: ReadOnly[str | None]
model: ReadOnly[str | None]
base_model: ReadOnly[str | None]
exception_type: ReadOnly[str]
exception_string: ReadOnly[str]
attempted_retries: ReadOnly[int | None]

View file

@ -223,6 +223,12 @@ class SearchContextCostPerQuery(TypedDict, total=False):
search_context_size_high: float
class SearchFreeRequests(TypedDict):
count: ReadOnly[int]
period: ReadOnly[Literal["day", "month"]]
shared_group: ReadOnly[str]
class AgenticLoopParams(TypedDict, total=False):
"""
Parameters passed to agentic loop hooks (e.g., WebSearch interception).
@ -284,6 +290,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
input_cost_per_token_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing
cache_creation_input_token_cost: float | None
cache_creation_input_token_cost_above_200k_tokens: float | None
cache_creation_input_token_cost_above_100k_tokens: ReadOnly[float | None]
cache_creation_input_token_cost_above_1hr_above_100k_tokens: ReadOnly[float | None]
cache_creation_input_token_cost_above_272k_tokens: float | None
cache_creation_input_token_cost_above_272k_tokens_priority: float | None
cache_creation_input_token_cost_above_272k_tokens_flex: float | None
@ -298,15 +306,18 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
cache_read_input_token_cost_priority: float | None # OpenAI priority service tier pricing
cache_read_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing
cache_read_input_token_cost_above_200k_tokens: float | None
cache_read_input_token_cost_above_100k_tokens: ReadOnly[float | None]
cache_read_input_token_cost_above_200k_tokens_priority: float | None
cache_read_input_token_cost_above_272k_tokens: float | None
cache_read_input_token_cost_above_272k_tokens_priority: float | None
cache_read_input_token_cost_above_272k_tokens_flex: float | None
cache_read_input_token_cost_above_512k_tokens: float | None
cache_read_input_token_cost_batches: ReadOnly[float | None]
cache_read_input_token_cost_above_100k_tokens_batches: ReadOnly[float | None]
cache_read_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None]
cache_creation_input_token_cost_batches: ReadOnly[float | None]
cache_creation_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None]
cache_creation_input_token_cost_above_100k_tokens_batches: ReadOnly[float | None]
cache_creation_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None]
# Smallest prefix this model will actually cache, whatever caching mechanism its provider uses.
# Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT.
@ -316,6 +327,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
input_cost_per_audio_token: float | None
input_cost_per_token_above_128k_tokens: float | None # only for vertex ai models
input_cost_per_token_above_200k_tokens: float | None # only for vertex ai gemini-2.5-pro models
input_cost_per_token_above_100k_tokens: ReadOnly[float | None]
input_cost_per_token_above_200k_tokens_priority: float | None
input_cost_per_token_above_272k_tokens: float | None # GPT-5.4/5.4-pro: prompts >272K priced at 2x input
input_cost_per_token_above_272k_tokens_priority: float | None
@ -333,8 +345,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
input_cost_per_second: float | None # for OpenAI Speech models
input_cost_per_token_batches: float | None
input_cost_per_video_token_batches: ReadOnly[float | None]
input_cost_per_token_above_100k_tokens_batches: ReadOnly[float | None]
input_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None]
output_cost_per_token_batches: float | None
output_cost_per_token_above_100k_tokens_batches: ReadOnly[float | None]
output_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None]
output_cost_per_token: Required[float | None]
output_cost_per_token_flex: float | None # OpenAI flex service tier pricing
@ -353,6 +367,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
output_cost_per_audio_token: float | None
output_cost_per_token_above_128k_tokens: float | None # only for vertex ai models
output_cost_per_token_above_200k_tokens: float | None # only for vertex ai gemini-2.5-pro models
output_cost_per_token_above_100k_tokens: ReadOnly[float | None]
output_cost_per_token_above_200k_tokens_priority: float | None
output_cost_per_token_above_272k_tokens: float | None # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output
output_cost_per_token_above_272k_tokens_priority: float | None
@ -390,6 +405,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
web_search_billing_unit: (
Literal["per_query", "per_prompt"] | None
) # "per_query" (Gemini 3.x) or "per_prompt" (Gemini 2.x)
search_free_requests: ReadOnly[SearchFreeRequests | None]
google_maps_grounding_cost_per_query: ReadOnly[float | None]
citation_cost_per_token: float | None # Cost per citation token for Perplexity
tiered_pricing: list[dict[str, Any]] | None # Tiered pricing structure for models like Dashscope
@ -3731,6 +3747,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
input_cost_per_token_priority: float | None = None
input_cost_per_token_ultrafast: float | None = None
cache_creation_input_token_cost_above_1hr: float | None = None
cache_creation_input_token_cost_above_100k_tokens: float | None = None
cache_creation_input_token_cost_above_1hr_above_100k_tokens: float | None = None
cache_creation_input_token_cost_above_200k_tokens: float | None = None
cache_creation_input_token_cost_above_272k_tokens: float | None = None
cache_creation_input_token_cost_above_272k_tokens_priority: float | None = None
@ -3742,13 +3760,16 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
cache_read_input_token_cost_flex: float | None = None
cache_read_input_token_cost_priority: float | None = None
cache_read_input_token_cost_ultrafast: float | None = None
cache_read_input_token_cost_above_100k_tokens: float | None = None
cache_read_input_token_cost_above_200k_tokens: float | None = None
cache_read_input_token_cost_above_200k_tokens_priority: float | None = None
cache_read_input_token_cost_above_272k_tokens_priority: float | None = None
cache_read_input_token_cost_above_272k_tokens_flex: float | None = None
cache_read_input_token_cost_batches: float | None = None
cache_read_input_token_cost_above_100k_tokens_batches: float | None = None
cache_read_input_token_cost_above_272k_tokens_batches: float | None = None
cache_creation_input_token_cost_batches: float | None = None
cache_creation_input_token_cost_above_100k_tokens_batches: float | None = None
cache_creation_input_token_cost_above_200k_tokens_batches: float | None = None
cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None
cache_read_input_audio_token_cost: float | None = None
@ -3757,10 +3778,12 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
input_cost_per_audio_token: float | None = None
input_cost_per_token_cache_hit: float | None = None
input_cost_per_token_above_128k_tokens: float | None = None
input_cost_per_token_above_100k_tokens: float | None = None
input_cost_per_token_above_200k_tokens: float | None = None
input_cost_per_token_above_200k_tokens_priority: float | None = None
input_cost_per_token_above_272k_tokens_priority: float | None = None
input_cost_per_token_above_272k_tokens_flex: float | None = None
input_cost_per_token_above_100k_tokens_batches: float | None = None
input_cost_per_token_above_272k_tokens_batches: float | None = None
input_cost_per_query: float | None = None
input_cost_per_image: float | None = None
@ -3781,10 +3804,12 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
output_cost_per_token_ultrafast: float | None = None
output_cost_per_audio_token: float | None = None
output_cost_per_token_above_128k_tokens: float | None = None
output_cost_per_token_above_100k_tokens: float | None = None
output_cost_per_token_above_200k_tokens: float | None = None
output_cost_per_token_above_200k_tokens_priority: float | None = None
output_cost_per_token_above_272k_tokens_priority: float | None = None
output_cost_per_token_above_272k_tokens_flex: float | None = None
output_cost_per_token_above_100k_tokens_batches: float | None = None
output_cost_per_token_above_272k_tokens_batches: float | None = None
output_cost_per_character_above_128k_tokens: float | None = None
output_cost_per_image: float | None = None
@ -3851,7 +3876,7 @@ def shared_backend_model_info(model_info: dict[str, Any]) -> dict[str, Any]:
return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS}
ABOVE_THRESHOLD_COST_KEY_PATTERN: Final = re.compile(r"_above_\d+k?_tokens$")
ABOVE_THRESHOLD_COST_KEY_PATTERN: Final = re.compile(r"_above_\d+k?_tokens(?:_batches)?$")
_PRICING_FIELD_EXEMPTIONS: Final[frozenset[str]] = frozenset({"output_vector_size"})

View file

@ -45,7 +45,7 @@ from httpx import Proxy
from httpx._utils import get_environment_proxies
from openai.lib import _parsing, _pydantic
from openai.types.chat.completion_create_params import ResponseFormat
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, with_config
import litellm
import litellm.litellm_core_utils
@ -954,10 +954,6 @@ def _ocr_document_summary(document: object) -> str:
# Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
_DECISION_CALL_TYPES: Final[frozenset[str]] = frozenset(
(CallTypes.decisions.value, CallTypes.adecisions.value, CallTypes.systemone.value, CallTypes.asystemone.value)
)
def function_setup(
original_function: str,
@ -1207,9 +1203,6 @@ def function_setup(
if isinstance(search_query, list)
else search_query
)
elif call_type in _DECISION_CALL_TYPES:
decisions_state: Final = args[1] if len(args) > 1 else kwargs.get("state") or kwargs.get("input") or ""
messages = decisions_state if isinstance(decisions_state, str) else json.dumps(decisions_state, default=str)
elif call_type in (CallTypes.image_edit.value, CallTypes.aimage_edit.value):
messages = args[1] if len(args) > 1 else kwargs.get("prompt")
elif call_type in (CallTypes.ocr.value, CallTypes.aocr.value):
@ -1375,6 +1368,7 @@ def _schedule_async_success_logging(
return
def _enqueue_async_logging() -> None:
logging_obj.mark_async_success_log_enqueued()
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
@ -3243,6 +3237,129 @@ def _get_builtin_model_info_for_registration(model: str) -> ModelInfo | None:
_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload
_runtime_registered_fields: Final[dict[str, frozenset[str]]] = {} # mutable-ok: fields each registration set
_runtime_registered_fields_owner: list[int] = [0] # mutable-ok: id of the model_cost map the fields describe
def _registered_fields_for_current_map() -> dict[str, frozenset[str]]:
if _runtime_registered_fields_owner[0] != id(litellm.model_cost):
_runtime_registered_fields.clear()
_runtime_registered_fields_owner[0] = id(litellm.model_cost)
return _runtime_registered_fields
_RATE_TABLE: Final = TypeAdapter(Mapping[str, object])
_RATE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
_TOKEN_COST_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token")
_DIRECTION_RATE_FIELDS: Final = MappingProxyType(
{
"input_cost_per_token": ("input_cost_per_token",),
"output_cost_per_token": ("output_cost_per_token", "output_cost_per_reasoning_token"),
}
)
def _is_overridden_rate(name: str, overridden_fields: tuple[str, ...]) -> bool:
is_variant: Final = name.startswith(tuple(f"{rate_field}_" for rate_field in overridden_fields))
return name in overridden_fields or (is_variant and "_cache_" not in name)
def _without_overridden_rates(rates: Mapping[str, object], overridden_fields: tuple[str, ...]) -> Mapping[str, object]:
return MappingProxyType(
{name: rate for name, rate in rates.items() if not _is_overridden_rate(name, overridden_fields)}
)
def _rate_table_with_flat_prices(
name: str, table: object, overridden_fields: tuple[str, ...], flat_prices: Mapping[str, object]
) -> object:
try:
if name == "tiered_pricing":
return [
{**_without_overridden_rates(tier, overridden_fields), **flat_prices}
for tier in _RATE_TIERS.validate_python(table)
]
if name == "off_peak_pricing":
return dict(_without_overridden_rates(_RATE_TABLE.validate_python(table), overridden_fields))
except ValidationError:
return table
return table
def _without_overridden_price_variants(
entry: Mapping[str, object], value: Mapping[str, object]
) -> Mapping[str, object]:
flat_prices: Final = MappingProxyType(
{cost_field: value[cost_field] for cost_field in _TOKEN_COST_FIELDS if value.get(cost_field) is not None}
)
overridden_fields: Final = tuple(
itertools.chain.from_iterable(_DIRECTION_RATE_FIELDS[cost_field] for cost_field in flat_prices)
)
if not overridden_fields:
return entry
return MappingProxyType(
{
name: _rate_table_with_flat_prices(name, field_value, overridden_fields, flat_prices)
for name, field_value in entry.items()
if not _is_overridden_rate(name, overridden_fields)
}
)
@with_config(ConfigDict(extra="allow"))
class _PricingModelInfo(ModelInfo):
pass
_PRICING_MODEL_INFO: Final = TypeAdapter(_PricingModelInfo)
def model_info_with_pricing_overrides(model_info: ModelInfo, registration: Mapping[str, object]) -> ModelInfo:
from litellm.types.utils import is_server_derived_pricing_key
overrides: Final = {
name: value for name, value in registration.items() if value is not None and is_server_derived_pricing_key(name)
}
return _PRICING_MODEL_INFO.validate_python(
_update_dictionary(dict(_without_overridden_price_variants(model_info, overrides)), overrides)
)
def runtime_model_cost_registrations(
model: str, custom_llm_provider: str
) -> tuple[tuple[str, Mapping[str, object]], ...]:
names: Final = _get_potential_model_names(model, custom_llm_provider)
candidates: Final = (
model,
names["combined_model_name"],
f"{custom_llm_provider}/{names['split_model']}",
names["split_model"],
names["stripped_model_name"],
)
registered: Final = (
(key, litellm.model_cost.get(key), _registered_fields_for_current_map().get(key, frozenset()))
for key in dict.fromkeys(candidates)
)
return tuple(
(key, MappingProxyType({name: value for name, value in entry.items() if name in fields}))
for key, entry, fields in registered
if isinstance(entry, Mapping) and fields
)
@dataclass(frozen=True, slots=True)
class ModelInfoLookupMiss:
tripped: bool
def lookup_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | ModelInfoLookupMiss:
try:
info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except ModelNotMappedError:
return ModelInfoLookupMiss(tripped=False)
except Exception:
return ModelInfoLookupMiss(tripped=True)
return ModelInfoLookupMiss(tripped=False) if is_generalized_model_info(info) else info
class _LiveDeploymentReplay:
@ -3284,6 +3401,7 @@ def reapply_runtime_model_cost_registrations() -> None:
have no such owner are recorded and replayed, and a registration describing
a single request opts out of even that.
"""
_runtime_registered_fields.clear()
if _LiveDeploymentReplay.callback is not None:
_LiveDeploymentReplay.callback()
if _runtime_registered_model_cost:
@ -3347,7 +3465,10 @@ def register_model(
if persist_across_reloads:
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost
for _registered_key, _registered_value in _registrations.items():
_runtime_registered_model_cost[_registered_key] = dict(_registered_value) # mutable-ok: caller-owned
_runtime_registered_model_cost[_registered_key] = _update_dictionary(
dict(_runtime_registered_model_cost.get(_registered_key, {})),
dict(_RATE_TABLE.validate_python(_registered_value)),
)
_skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
@ -3413,6 +3534,10 @@ def register_model(
## override / add new keys to the existing model cost dictionary
updated_dictionary = _update_dictionary(existing_model, value)
litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary)
registered_fields = _registered_fields_for_current_map()
registered_fields[str(model_cost_key)] = registered_fields.get(str(model_cost_key), frozenset()) | frozenset(
name for name, field_value in value.items() if field_value is not None
)
# Invalidate case-insensitive lookup map since model_cost was modified
_invalidate_model_cost_lowercase_map()
@ -5037,6 +5162,9 @@ def get_optional_params(
EXTRA_BODY_ROUTING_KEYS: Final = frozenset({"model"})
LITELLM_ONLY_PARAMS: Final = frozenset(("cache_control_injection_points",))
def add_provider_specific_params_to_optional_params(
optional_params: dict,
passed_params: dict,
@ -5052,8 +5180,9 @@ def add_provider_specific_params_to_optional_params(
# for openai, azure we should pass the extra/passed params within `extra_body` https://github.com/openai/openai-python/blob/ac33853ba10d13ac149b1fa3ca6dba7d613065c9/src/openai/resources/models.py#L46
if _should_drop_param(k="extra_body", additional_drop_params=additional_drop_params) is False:
extra_body: Final = dict(passed_params.pop("extra_body", None) or {})
withheld_params: Final = frozenset[str]() if custom_llm_provider == "litellm_proxy" else LITELLM_ONLY_PARAMS
for k in passed_params:
if k not in openai_params and passed_params[k] is not None:
if k not in openai_params and k not in withheld_params and passed_params[k] is not None:
extra_body[k] = passed_params[k]
if not isinstance(optional_params.get("extra_body"), dict):
optional_params["extra_body"] = {}
@ -6318,6 +6447,7 @@ def _get_model_info_helper(
supports_computer_use=_model_info.get("supports_computer_use", None),
search_context_cost_per_query=_model_info.get("search_context_cost_per_query", None),
web_search_billing_unit=_model_info.get("web_search_billing_unit", None),
search_free_requests=_model_info.get("search_free_requests", None),
google_maps_grounding_cost_per_query=_model_info.get("google_maps_grounding_cost_per_query", None),
tpm=_model_info.get("tpm", None),
rpm=_model_info.get("rpm", None),

File diff suppressed because it is too large Load diff

View file

@ -83,7 +83,12 @@
"minimum": 0,
"description": "USD per token written to the provider's prompt cache."
},
"cache_creation_input_token_cost_above_128k_tokens": {
"cache_creation_input_token_cost_above_100k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_100k_tokens_batches": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
@ -93,6 +98,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": {
"type": "number",
"minimum": 0,
@ -103,11 +113,6 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_256k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_272k_tokens": {
"type": "number",
"minimum": 0,
@ -118,6 +123,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_272k_tokens_decisions": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_272k_tokens_flex": {
"type": "number",
"minimum": 0,
@ -128,7 +138,7 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"cache_creation_input_token_cost_above_32k_tokens": {
"cache_creation_input_token_cost_above_272k_tokens_ultrafast": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
@ -137,6 +147,10 @@
"type": "number",
"minimum": 0
},
"cache_creation_input_token_cost_decisions": {
"type": "number",
"minimum": 0
},
"cache_creation_input_token_cost_flex": {
"type": "number",
"minimum": 0,
@ -147,6 +161,10 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"cache_creation_input_token_cost_ultrafast": {
"type": "number",
"minimum": 0
},
"cache_read_input_audio_token_cost": {
"type": "number",
"minimum": 0
@ -160,7 +178,12 @@
"minimum": 0,
"description": "USD per prompt token served from the provider's prompt cache."
},
"cache_read_input_token_cost_above_128k_tokens": {
"cache_read_input_token_cost_above_100k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_100k_tokens_batches": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
@ -170,16 +193,16 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_200k_tokens_batches": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_200k_tokens_priority": {
"type": "number",
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"cache_read_input_token_cost_above_256k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_272k_tokens": {
"type": "number",
"minimum": 0,
@ -190,6 +213,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_272k_tokens_decisions": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_272k_tokens_flex": {
"type": "number",
"minimum": 0,
@ -200,7 +228,7 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"cache_read_input_token_cost_above_32k_tokens": {
"cache_read_input_token_cost_above_272k_tokens_ultrafast": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
@ -210,10 +238,18 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_balanced": {
"type": "number",
"minimum": 0
},
"cache_read_input_token_cost_batches": {
"type": "number",
"minimum": 0
},
"cache_read_input_token_cost_decisions": {
"type": "number",
"minimum": 0
},
"cache_read_input_token_cost_flex": {
"type": "number",
"minimum": 0,
@ -224,6 +260,10 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"cache_read_input_token_cost_ultrafast": {
"type": "number",
"minimum": 0
},
"citation_cost_per_token": {
"type": "number",
"minimum": 0
@ -235,6 +275,10 @@
"comment": {
"type": "string"
},
"cost_per_second": {
"type": "number",
"minimum": 0
},
"default_reasoning_effort": {
"type": "string",
"description": "Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'.",
@ -253,6 +297,10 @@
"format": "date",
"pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$"
},
"explicit_cache_read_input_token_cost": {
"type": "number",
"minimum": 0
},
"gemini_audio_only_live": {
"type": "boolean"
},
@ -336,11 +384,25 @@
"type": "number",
"minimum": 0
},
"input_cost_per_text_message": {
"type": "number",
"minimum": 0
},
"input_cost_per_token": {
"type": "number",
"minimum": 0,
"description": "USD per prompt token."
},
"input_cost_per_token_above_100k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_100k_tokens_batches": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_128k_tokens": {
"type": "number",
"minimum": 0,
@ -351,16 +413,16 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_200k_tokens_batches": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_200k_tokens_priority": {
"type": "number",
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"input_cost_per_token_above_256k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_272k_tokens": {
"type": "number",
"minimum": 0,
@ -371,6 +433,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_272k_tokens_decisions": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_272k_tokens_flex": {
"type": "number",
"minimum": 0,
@ -381,7 +448,7 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"input_cost_per_token_above_32k_tokens": {
"input_cost_per_token_above_272k_tokens_ultrafast": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
@ -391,6 +458,10 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_balanced": {
"type": "number",
"minimum": 0
},
"input_cost_per_token_batches": {
"type": "number",
"minimum": 0,
@ -400,6 +471,10 @@
"type": "number",
"minimum": 0
},
"input_cost_per_token_decisions": {
"type": "number",
"minimum": 0
},
"input_cost_per_token_flex": {
"type": "number",
"minimum": 0,
@ -410,6 +485,10 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"input_cost_per_token_ultrafast": {
"type": "number",
"minimum": 0
},
"input_cost_per_video_per_second": {
"type": "number",
"minimum": 0
@ -474,6 +553,7 @@
"completion",
"embedding",
"evaluation",
"extraction",
"guardrail",
"image_edit",
"image_generation",
@ -656,6 +736,10 @@
"type": "number",
"minimum": 0
},
"output_cost_per_image_token_batches": {
"type": "number",
"minimum": 0
},
"output_cost_per_pixel": {
"type": "number",
"minimum": 0
@ -698,6 +782,16 @@
"minimum": 0,
"description": "USD per generated token."
},
"output_cost_per_token_above_100k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_100k_tokens_batches": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_128k_tokens": {
"type": "number",
"minimum": 0,
@ -708,16 +802,16 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_200k_tokens_batches": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_200k_tokens_priority": {
"type": "number",
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"output_cost_per_token_above_256k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_272k_tokens": {
"type": "number",
"minimum": 0,
@ -728,6 +822,11 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_272k_tokens_decisions": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_272k_tokens_flex": {
"type": "number",
"minimum": 0,
@ -738,7 +837,7 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"output_cost_per_token_above_32k_tokens": {
"output_cost_per_token_above_272k_tokens_ultrafast": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
@ -748,11 +847,19 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_balanced": {
"type": "number",
"minimum": 0
},
"output_cost_per_token_batches": {
"type": "number",
"minimum": 0,
"description": "USD per generated token via the provider's batch API."
},
"output_cost_per_token_decisions": {
"type": "number",
"minimum": 0
},
"output_cost_per_token_flex": {
"type": "number",
"minimum": 0,
@ -763,6 +870,10 @@
"minimum": 0,
"description": "Priority service-tier rate for the same-named base field."
},
"output_cost_per_token_ultrafast": {
"type": "number",
"minimum": 0
},
"output_cost_per_video_per_second": {
"type": "number",
"minimum": 0
@ -844,6 +955,33 @@
},
"additionalProperties": false
},
"search_free_requests": {
"type": "object",
"description": "Free Google Search grounding requests per period, shared across models in the same shared_group.",
"required": [
"count",
"period",
"shared_group"
],
"additionalProperties": false,
"properties": {
"count": {
"type": "integer",
"minimum": 0
},
"period": {
"type": "string",
"enum": [
"day",
"month"
]
},
"shared_group": {
"type": "string",
"pattern": "\\S"
}
}
},
"source": {
"type": "string",
"description": "URL of the provider pricing/model page this entry was taken from."
@ -918,6 +1056,24 @@
"supports_audio_output": {
"type": "boolean"
},
"supports_bedrock_converse_tool_result_images": {
"type": "boolean"
},
"supports_bedrock_native_structured_output": {
"type": "boolean"
},
"supports_bedrock_runtime_chat_completions_inline_reasoning": {
"type": "boolean"
},
"supports_bedrock_runtime_chat_completions_response_format": {
"type": "boolean"
},
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": {
"type": "boolean"
},
"supports_box_questions": {
"type": "boolean"
},
"supports_computer_use": {
"type": "boolean"
},
@ -990,6 +1146,9 @@
"supports_reasoning": {
"type": "boolean"
},
"supports_regex_lookaround": {
"type": "boolean"
},
"supports_response_schema": {
"type": "boolean"
},
@ -1088,6 +1247,31 @@
"minimum": 0,
"description": "Tokens added once per request to a local count estimate, for the provider's chat framing."
},
"token_count_model": {
"type": "object",
"required": [
"model"
],
"additionalProperties": false,
"properties": {
"model": {
"type": "string",
"pattern": "\\S"
},
"region": {
"type": "string",
"pattern": "\\S"
},
"litellm_provider": {
"type": "string",
"enum": [
"bedrock",
"bedrock_mantle"
]
}
},
"description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in."
},
"token_count_tokenizer": {
"type": "string",
"description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.",

View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.104.5"
version = "1.104.6"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -76,7 +76,7 @@ proxy = [
"httpx2>=2.5.0,<3",
"pydantic>=2.12.0,<3",
"litellm-proxy-extras==0.4.102.post1",
"litellm-enterprise==0.1.71",
"litellm-enterprise==0.1.71.post1",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",
@ -286,7 +286,7 @@ ci = [
"langchain-openai==1.1.14",
"langgraph>=1.2.4,<1.3.0",
"langgraph-prebuilt>=1.1.0,<1.3.0",
"claude-agent-sdk==0.1.44",
"claude-agent-sdk==0.2.165",
"google-cloud-speech==2.40.0",
]
healthcheck = [
@ -355,7 +355,7 @@ litellm-enterprise = { workspace = true }
members = ["enterprise", "litellm-proxy-extras"]
[tool.commitizen]
version = "1.104.5"
version = "1.104.6"
version_files = [
"pyproject.toml:^version",
]

View file

@ -87,6 +87,7 @@ ignored_function_names = [
"_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)
"_deployment_params_base_model", # Tested through log_retry in test_router_helper_utils.py
]

View file

@ -0,0 +1,61 @@
"""Live e2e: `/v1/messages/count_tokens` on a Bedrock Claude model that bedrock-runtime
cannot count.
Claude Opus 4.8 is offered only through cross-region inference, and bedrock-runtime's
CountTokens answers 400 for it. The proxy then has to count through bedrock-mantle's
Anthropic count_tokens, and the answer must sit within a few percent of what `/v1/messages`
bills as `usage.input_tokens`. The local tokenizer fallback undercounts these models by
about 40%, so this is the line that proves the real count is served. Both calls go through
the real Anthropic SDK, as a customer would use them.
"""
from __future__ import annotations
from typing import Final
import pytest
from anthropic.types import MessageParam
from e2e_config import unique_marker
from lifecycle import ResourceManager
from models import LiteLLMParamsBody
from proxy_client import ProxyClient
from sdk_clients import NO_PROXY_CACHE, SdkClients
pytestmark = pytest.mark.e2e
CROSS_REGION_ONLY_CLAUDE_BACKEND: Final = "bedrock/global.anthropic.claude-opus-4-8"
COUNT_TOLERANCE: Final = 0.05
def _register(proxy: ProxyClient, resources: ResourceManager, backend: str) -> tuple[str, str]:
model: Final = f"e2e-count-tokens-bedrock-{unique_marker()}"
model_id: Final = proxy.create_model(
model,
LiteLLMParamsBody(
model=backend,
aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
aws_region_name="os.environ/AWS_REGION",
),
)
resources.defer(lambda: proxy.delete_model(model_id))
return model, resources.key()
class TestBedrockMessagesCountTokens:
def test_count_matches_billed_input_tokens_for_a_model_bedrock_runtime_cannot_count(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register(proxy, resources, CROSS_REGION_ONLY_CLAUDE_BACKEND)
client: Final = sdk.anthropic(key)
prompt: Final = f"{unique_marker()} " + "The quick brown fox jumps over the lazy dog. " * 40
message: Final[MessageParam] = {"role": "user", "content": prompt}
counted: Final = client.messages.count_tokens(model=model, messages=[message])
answered: Final = client.messages.create(
model=model, max_tokens=1, messages=[message], extra_body=NO_PROXY_CACHE
)
billed: Final = answered.usage.input_tokens
assert billed, answered.usage
assert abs(counted.input_tokens - billed) <= billed * COUNT_TOLERANCE, (counted, answered.usage)

View file

@ -0,0 +1,349 @@
import asyncio
import json
import re
import threading
import uuid
from dataclasses import dataclass
from pathlib import Path
from queue import SimpleQueue
from types import MappingProxyType
from typing import Final
from urllib.parse import urlsplit
import httpx
import psutil
import pytest
import yaml
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.providers.test_anthropic_safeguards_beta_wire import (
ANTHROPIC_KEY,
ANTHROPIC_MODEL,
BETA,
OMITTED,
anthropic_peer,
marker_of,
messages_body,
stream_frames,
)
from pydantic import JsonValue, TypeAdapter
_CONFIG_MODEL: Final = "anthropic-safeguards-chaos"
_FOUNDRY_CONFIG_MODEL: Final = "azure-ai-safeguards-chaos"
_YAML_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_MODEL_LIST: Final = json.dumps(
{"object": "list", "data": [{"id": ANTHROPIC_MODEL, "object": "model", "owned_by": "anthropic"}]}
).encode()
@dataclass(frozen=True, slots=True)
class _Call:
stream: bool
marker: str
safeguards: bool = True
@dataclass(frozen=True, slots=True)
class _Served:
call: _Call
status: int
text: str
def _calls(count: int, stream: bool, safeguards: bool = True) -> tuple[_Call, ...]:
return tuple(_Call(stream=stream, marker=uuid.uuid4().hex, safeguards=safeguards) for _ in range(count))
def _streamed_reply(marker: str, abort_after: int | None = None, pause: float = 0) -> Reply:
return Reply(
content_type="text/event-stream",
chunks=stream_frames(marker),
abort_after=abort_after,
pause_between_chunks=pause,
)
async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served:
async with client.stream(
"POST",
"/v1/messages",
json=(
messages_body(model, call.marker, stream=call.stream)
if call.safeguards
else messages_body(model, call.marker, stream=call.stream, safeguards=OMITTED)
),
headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"},
) as response:
raw: Final = await response.aread()
return _Served(call=call, status=response.status_code, text=raw.decode())
async def _burst(
base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False
) -> tuple[_Served, ...]:
async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
results: Final = await asyncio.gather(
*(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors
)
for result in results:
assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
return tuple(result for result in results if isinstance(result, _Served))
def _assert_answered_with_its_own_marker(served: _Served) -> None:
assert served.status == 200, served.text
assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text
assert f"msg_{served.call.marker}" in served.text, served.text
if served.call.stream:
assert "event: message_stop" in served.text, served.text
def _assert_every_request_carried_the_beta(received: tuple[Request, ...], markers: frozenset[str]) -> None:
assert sorted(marker_of(request) for request in received) == sorted(markers)
missing: Final = [
frozenset(request.headers) for request in received if request.headers.get("anthropic-beta") != BETA
]
assert missing == [], missing
def _spend_success_ids(model: str, expected: frozenset[str]) -> None:
rows: Final = eventually(
lambda: read_rows(
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group=%s AND status=%s', (model, "success")
),
lambda found: expected <= {str(row["request_id"]) for row in found},
seconds=70,
)
identities: Final = [row["request_id"] for row in rows]
assert len(set(identities)) == len(identities), identities
async def test_burst_with_aborted_streams_keeps_the_beta_on_every_forwarded_request(gateway: Gateway) -> None:
calls: Final = tuple(_Call(stream=index % 2 == 0, marker=uuid.uuid4().hex) for index in range(30))
aborted: Final = frozenset(call.marker for index, call in enumerate(calls) if call.stream and index % 3 == 0)
def respond(request: Request) -> Reply:
marker: Final = marker_of(request)
if marker in aborted:
return _streamed_reply(marker, abort_after=0)
return anthropic_peer(request)
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
assert len(served) == 30
for item in served:
if item.call.marker in aborted:
assert item.status == 500, item.text
assert '"type":"api_error"' in item.text and "marker-" not in item.text, item.text
else:
_assert_answered_with_its_own_marker(item)
recovery: Final = _Call(stream=True, marker=uuid.uuid4().hex)
(recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,))
_assert_answered_with_its_own_marker(recovered)
_assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in (*calls, recovery)))
answered: Final = frozenset(f"msg_{call.marker}" for call in (*calls, recovery) if call.marker not in aborted)
_spend_success_ids(model, answered)
async def test_slow_streams_complete_with_the_beta_on_every_request(gateway: Gateway) -> None:
calls: Final = _calls(10, stream=True)
with (
wire_server(lambda request: _streamed_reply(marker_of(request), pause=0.3)) as wire,
gateway.scenario() as scenario,
):
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
assert len(served) == 10
for item in served:
_assert_answered_with_its_own_marker(item)
_assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in calls))
_spend_success_ids(model, frozenset(f"msg_{call.marker}" for call in calls))
def _anthropic_deployment(wire: Wire) -> dict[str, JsonValue]:
return {
"model_name": _CONFIG_MODEL,
"litellm_params": {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_base": wire.url, "api_key": ANTHROPIC_KEY},
}
def _foundry_deployment(wire: Wire) -> dict[str, JsonValue]:
return {
"model_name": _FOUNDRY_CONFIG_MODEL,
"litellm_params": {
"model": "azure_ai/claude-sonnet-4-6",
"api_base": wire.url,
"api_key": "synthetic-foundry-key",
},
}
def _chaos_config(tmp_path: Path, *deployments: dict[str, JsonValue]) -> Path:
base: Final = _YAML_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = {**base, "model_list": list(deployments)}
path: Final = tmp_path / "anthropic-safeguards-chaos.yaml"
path.write_text(yaml.safe_dump(config))
return path
def _open_upstream_connections(pid: int, upstream: str) -> int:
port: Final = urlsplit(upstream).port
return sum(
1
for connection in psutil.Process(pid).net_connections(kind="tcp")
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
)
@pytest.mark.skip(reason="Needs OwnedProxy.replacing_worker, which this stable line's integration harness does not have")
@pytest.mark.timeout(180)
async def test_worker_sigkill_mid_burst_leaves_the_sibling_sending_the_beta(gateway: Gateway, tmp_path: Path) -> None:
calls: Final = _calls(20, stream=False)
release: Final = threading.Event()
held_markers: Final[SimpleQueue[str]] = SimpleQueue()
def held(request: Request) -> Reply:
if (request.method, request.target) == ("GET", "/v1/models"):
return Reply(body=_MODEL_LIST)
held_markers.put(marker_of(request))
assert release.wait(timeout=60), "The burst was never released"
return anthropic_peer(request)
with wire_server(held) as wire:
path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire))
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
candidate: Final = owned.gateway
workers: Final = eventually(
lambda: tuple(int(found.group(1)) for found in _STARTED_WORKER.finditer(owned.log.read_text())),
lambda pids: len(pids) == 2,
seconds=30,
)
burst: Final = asyncio.create_task(
_burst(
str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True
)
)
await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60)
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers})
assert sum(held_by.values()) == 20, held_by
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
with owned.replacing_worker(victim_pid) as kill:
kill.freeze()
await asyncio.to_thread(kill)
release.set()
served: Final = await burst
assert held_by[survivor_pid] >= 10, held_by
assert len(served) == held_by[survivor_pid], (held_by, len(served))
for item in served:
_assert_answered_with_its_own_marker(item)
follow_up: Final = _Call(stream=False, marker=uuid.uuid4().hex)
(answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,))
_assert_answered_with_its_own_marker(answered)
received: Final = wire.drain()
posts: Final = tuple(request for request in received if request.method == "POST")
probes: Final[list[tuple[str, str]]] = [
(request.method, request.target) for request in received if request.method != "POST"
]
assert set(probes) <= {("GET", "/v1/models")}, probes
_assert_every_request_carried_the_beta(posts, frozenset(call.marker for call in (*calls, follow_up)))
def _beta_map_gets(received: tuple[Request, ...]) -> int:
return sum(1 for request in received if (request.method, request.target) == ("GET", "/beta.json"))
async def _chat(base_url: str, key: str, marker: str) -> httpx.Response:
async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
return await client.post(
"/v1/chat/completions",
json={
"model": _CONFIG_MODEL,
"max_tokens": 16,
"messages": [{"role": "user", "content": f"Question marker-{marker}"}],
},
headers={"Authorization": f"Bearer {key}"},
)
async def _answered_while_the_beta_map_is_held(
base_url: str, key: str, guarded: tuple[_Call, ...], chat_markers: tuple[str, ...]
) -> None:
served: Final = await _burst(base_url, key, _CONFIG_MODEL, guarded)
chats: Final = await asyncio.gather(*(_chat(base_url, key, marker) for marker in chat_markers))
for item in served:
_assert_answered_with_its_own_marker(item)
for marker, chat in zip(chat_markers, chats, strict=True):
assert chat.status_code == 200, chat.text
assert f"marker-{marker}" in chat.text, chat.text
@pytest.mark.timeout(240)
async def test_cold_workers_fetch_the_beta_map_only_for_filtered_safeguards_requests_off_the_shared_threads(
gateway: Gateway, tmp_path: Path
) -> None:
beta_map: Final = Path("litellm/anthropic_beta_headers_config.json").read_bytes()
release: Final = threading.Event()
def held_beta_map(_: Request) -> Reply:
assert release.wait(timeout=180), "The beta map was never released"
return Reply(body=beta_map)
def upstream(request: Request) -> Reply:
if (request.method, request.target) == ("GET", "/v1/models"):
return Reply(body=_MODEL_LIST)
return anthropic_peer(request)
plain: Final = _calls(8, stream=False, safeguards=False)
plain_foundry: Final = _calls(8, stream=False, safeguards=False)
direct: Final = _calls(8, stream=False)
foundry: Final = _calls(80, stream=False)
direct_while_held: Final = _calls(8, stream=False)
chat_markers: Final = tuple(uuid.uuid4().hex for _ in range(4))
with wire_server(held_beta_map) as beta_map_wire, wire_server(upstream) as wire:
path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire), _foundry_deployment(wire))
with owned_proxy_process(
gateway,
tmp_path,
{"LITELLM_ANTHROPIC_BETA_HEADERS_URL": f"{beta_map_wire.url}/beta.json"},
remove_environment=("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS",),
config=path,
workers=2,
) as owned:
base_url: Final = str(owned.gateway.client.base_url)
key: Final = owned.gateway.key
cold: Final = (
*await _burst(base_url, key, _CONFIG_MODEL, plain),
*await _burst(base_url, key, _FOUNDRY_CONFIG_MODEL, plain_foundry),
*await _burst(base_url, key, _CONFIG_MODEL, direct),
)
for served_cold in cold:
_assert_answered_with_its_own_marker(served_cold)
assert beta_map_wire.received.qsize() == 0, "a request that never filters betas fetched the beta map"
held: Final = asyncio.create_task(_burst(base_url, key, _FOUNDRY_CONFIG_MODEL, foundry))
try:
await asyncio.to_thread(eventually, beta_map_wire.received.qsize, lambda size: size >= 1, 60)
await asyncio.wait_for(
_answered_while_the_beta_map_is_held(base_url, key, direct_while_held, chat_markers), timeout=45
)
assert not held.done(), "Foundry safeguards requests were answered before the beta map loaded"
finally:
release.set()
for served_foundry in await held:
_assert_answered_with_its_own_marker(served_foundry)
received: Final = tuple(request for request in wire.drain() if request.method == "POST")
beta_by_marker: Final = MappingProxyType(
{marker_of(request): request.headers.get("anthropic-beta") for request in received}
)
assert {beta_by_marker[call.marker] for call in (*direct, *direct_while_held, *foundry)} == {BETA}, (
beta_by_marker
)
assert {beta_by_marker[call.marker] for call in (*plain, *plain_foundry)} == {None}, beta_by_marker
fetches: Final = _beta_map_gets(beta_map_wire.drain())
assert 1 <= fetches <= 2, fetches

View file

@ -0,0 +1,809 @@
import itertools
import json
import re
import time
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import anthropic
import httpx
import openai
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.vertex import service_account_json
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
BETA: Final = "dangerous-tool-use-2026-09-03"
SAFEGUARDS: Final[list[JsonValue]] = [
{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}
]
ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929"
ANTHROPIC_KEY: Final = "synthetic-anthropic-key"
_CLAUDE_CODE_BETAS: Final = (
"claude-code-20250219",
"interleaved-thinking-2025-05-14",
"fine-grained-tool-streaming-2025-05-14",
"context-management-2025-06-27",
)
_CLIENT_BETAS: Final = ",".join(_CLAUDE_CODE_BETAS)
_CLIENT_BETAS_WITH_DANGEROUS: Final = ",".join((*_CLAUDE_CODE_BETAS, BETA))
_MERGED_BETAS: Final = ",".join(sorted((*_CLAUDE_CODE_BETAS, BETA)))
_FOUNDRY_MODEL: Final = "azure_ai/claude-sonnet-4-6"
_FOUNDRY_KEY: Final = "synthetic-foundry-key"
_FOUNDRY_TARGET: Final = "/anthropic/v1/messages"
_FOUNDRY_ARRAY_ERROR: Final = "safeguards: Input should be a valid array"
_FOUNDRY_EXTRA_INPUT_ERROR: Final = "safeguards: Extra inputs are not permitted"
_ANTHROPIC_VERSION: Final = MappingProxyType({"anthropic-version": "2023-06-01"})
_MESSAGES_KEYS: Final = frozenset({"max_tokens", "messages", "model", "stream"})
_RESPONSES_KEYS: Final = frozenset({"include", "input", "max_output_tokens", "model"})
_SDK_EXTRA: Final[dict[str, JsonValue]] = {"safeguards": SAFEGUARDS, "cache": {"no-cache": True}, "num_retries": 0}
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
@dataclass(frozen=True, slots=True)
class _Omitted:
pass
OMITTED: Final = _Omitted()
@dataclass(frozen=True, slots=True)
class _Observed:
method: str
target: str
marker: str
headers: Mapping[str, str]
body: Mapping[str, JsonValue]
@property
def beta(self) -> str | None:
return self.headers.get("anthropic-beta")
@property
def keys(self) -> frozenset[str]:
return frozenset(self.body)
@property
def header_names(self) -> frozenset[str]:
return frozenset(self.headers)
def marker_of(request: Request) -> str:
found: Final = _MARKER.search(request.body.decode())
assert found is not None, request.body
return found.group(1)
def _message(marker: str) -> dict[str, JsonValue]:
return {
"id": f"msg_{marker}",
"type": "message",
"role": "assistant",
"model": ANTHROPIC_MODEL,
"content": [{"type": "text", "text": f"answer marker-{marker}"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 5},
}
def stream_frames(marker: str) -> tuple[bytes, ...]:
opening: Final[dict[str, JsonValue]] = {
**_message(marker),
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 12, "output_tokens": 0},
}
events: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = (
("message_start", {"message": opening}),
("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}),
("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": f"answer marker-{marker}"}}),
("content_block_stop", {"index": 0}),
("message_delta", {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 5}}),
("message_stop", {}),
)
return tuple(
f"event: {kind}\ndata: {json.dumps({'type': kind, **payload})}\n\n".encode() for kind, payload in events
)
def anthropic_peer(request: Request) -> Reply:
marker: Final = marker_of(request)
if _JSON_OBJECT.validate_json(request.body).get("stream") is True:
return Reply(content_type="text/event-stream", chunks=stream_frames(marker))
return Reply(body=json.dumps(_message(marker)).encode())
def _responses_peer(request: Request) -> Reply:
marker: Final = marker_of(request)
return Reply(
body=json.dumps(
{
"id": f"resp_{marker}",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"type": "message",
"id": f"msg_{marker}",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": f"answer marker-{marker}", "annotations": []}],
}
],
"usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17},
}
).encode()
)
def _error_peer(status: int, kind: str, message: str) -> Callable[[Request], Reply]:
body: Final = json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode()
return lambda request: Reply(status=status, body=body)
def messages_body(
model: str,
marker: str,
*,
safeguards: JsonValue | _Omitted = SAFEGUARDS,
stream: bool = False,
no_cache: bool = True,
) -> dict[str, JsonValue]:
base: Final[dict[str, JsonValue]] = {
"model": model,
"max_tokens": 16,
"messages": [{"role": "user", "content": f"Question marker-{marker}"}],
"num_retries": 0,
**({"stream": True} if stream else {}),
**({"cache": {"no-cache": True}} if no_cache else {}),
}
if isinstance(safeguards, _Omitted):
return base
return {**base, "safeguards": safeguards}
def _observe(request: Request) -> _Observed:
return _Observed(
request.method, request.target, marker_of(request), request.headers, _JSON_OBJECT.validate_json(request.body)
)
def _only(wire: Wire) -> _Observed:
(request,) = wire.drain()
return _observe(request)
def _post(gateway: Gateway, body: Mapping[str, JsonValue], *, beta: str | None = None) -> httpx.Response:
headers: Final = {**_ANTHROPIC_VERSION, **({"anthropic-beta": beta} if beta else {})}
return gateway.request("POST", "/v1/messages", body, headers=headers)
def _assert_answered(response: httpx.Response, marker: str) -> None:
assert response.status_code == 200, response.text
assert response.json()["id"] == f"msg_{marker}", response.text
def _assert_streamed(response: httpx.Response, marker: str) -> None:
assert response.status_code == 200, response.text
assert f"msg_{marker}" in response.text and f"answer marker-{marker}" in response.text, response.text
assert "event: message_stop" in response.text, response.text
def _spend_row(identity: str) -> None:
rows: Final = eventually(
lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)),
lambda found: len(found) == 1,
seconds=70,
)
assert rows[0]["status"] == "success", rows
def _cache_hit_spend_row(identity: str) -> None:
rows: Final = eventually(
lambda: read_rows(
'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s',
(f"{identity}\\_cache\\_hit%",),
),
lambda found: len(found) == 1,
seconds=70,
)
assert (rows[0]["status"], str(rows[0]["cache_hit"])) == ("success", "True"), rows
def _responses_spend_row(model: str) -> None:
rows: Final = eventually(
lambda: read_rows(
'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group = %s AND call_type = %s',
(model, "aresponses"),
),
lambda found: len(found) == 1,
seconds=70,
)
assert rows[0]["status"] == "success", rows
def _expect_messages(
observed: _Observed,
*,
target: str,
auth: tuple[str, str] | None,
beta: str | None,
safeguards: JsonValue | _Omitted = SAFEGUARDS,
keys: frozenset[str] = _MESSAGES_KEYS,
) -> None:
assert (observed.method, observed.target) == ("POST", target)
if auth is not None:
credential_matches: Final = observed.headers.get(auth[0]) == auth[1]
assert credential_matches, f"{auth[0]} differs from the deployment credential; headers {observed.header_names}"
assert observed.beta == beta, observed.header_names
if isinstance(safeguards, _Omitted):
assert observed.keys == keys
return
assert observed.keys == keys | {"safeguards"}
assert observed.body["safeguards"] == safeguards
def _expect_anthropic_version(observed: _Observed) -> None:
assert observed.headers.get("anthropic-version") == "2023-06-01", observed.header_names
@dataclass(frozen=True, slots=True)
class _ShapeCase:
sent: JsonValue | _Omitted
forwarded: JsonValue | _Omitted
client_beta: str | None
expected_beta: str | None
stream: bool = False
_ANTHROPIC_SHAPES: Final = MappingProxyType(
{
"D1_safeguards_without_client_header": _ShapeCase(SAFEGUARDS, SAFEGUARDS, None, BETA),
"D2_safeguards_with_claude_code_header": _ShapeCase(
SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, _MERGED_BETAS
),
"D3_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase(
SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS, _MERGED_BETAS
),
"D4_empty_safeguards": _ShapeCase([], [], None, BETA),
"D5_null_safeguards": _ShapeCase(None, OMITTED, None, None),
"D6_missing_safeguards": _ShapeCase(OMITTED, OMITTED, None, None),
}
)
_FOUNDRY_SHAPES: Final = MappingProxyType(
{
"D9_streamed_safeguards_with_claude_code_header": _ShapeCase(
SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, BETA, stream=True
),
"D11_empty_safeguards": _ShapeCase([], [], None, BETA),
"D12_null_safeguards": _ShapeCase(None, OMITTED, None, None),
"D13_missing_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase(
OMITTED, OMITTED, _CLIENT_BETAS, None
),
}
)
@pytest.mark.parametrize("row", tuple(_ANTHROPIC_SHAPES))
def test_anthropic_messages_derive_the_dangerous_tool_use_beta_from_safeguards(gateway: Gateway, row: str) -> None:
case: Final = _ANTHROPIC_SHAPES[row]
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
response: Final = _post(gateway, messages_body(model, marker, safeguards=case.sent), beta=case.client_beta)
_assert_answered(response, marker)
observed: Final = _only(wire)
_expect_messages(
observed,
target="/v1/messages",
auth=("x-api-key", ANTHROPIC_KEY),
beta=case.expected_beta,
safeguards=case.forwarded,
)
_expect_anthropic_version(observed)
_spend_row(f"msg_{marker}")
@dataclass(frozen=True, slots=True, kw_only=True)
class _DropCase:
deployment: Callable[[Gateway], Mapping[str, JsonValue]]
target: str
auth: tuple[str, str] | None
keys: frozenset[str] = _MESSAGES_KEYS
_VERTEX_PROJECT: Final = "scripted-project"
_VERTEX_LOCATION: Final = "us-east5"
_VERTEX_MODEL: Final = "claude-opus-5-5"
_DROP_CASES: Final = MappingProxyType(
{
"anthropic": _DropCase(
deployment=lambda _: {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_key": ANTHROPIC_KEY},
target="/v1/messages",
auth=("x-api-key", ANTHROPIC_KEY),
),
"azure_ai": _DropCase(
deployment=lambda _: {"model": _FOUNDRY_MODEL, "api_key": _FOUNDRY_KEY},
target=_FOUNDRY_TARGET,
auth=("x-api-key", _FOUNDRY_KEY),
),
"bedrock_claude_platform": _DropCase(
deployment=lambda _: {
"model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}",
"api_key": "synthetic-platform-key",
"workspace_id": "ws-synthetic",
"aws_region_name": "us-east-1",
},
target="/v1/messages",
auth=("x-api-key", "synthetic-platform-key"),
),
"vertex_ai": _DropCase(
deployment=lambda gateway: {
"model": f"vertex_ai/{_VERTEX_MODEL}",
"api_key": None,
"vertex_project": _VERTEX_PROJECT,
"vertex_location": _VERTEX_LOCATION,
"vertex_credentials": service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")),
},
target=(
f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}"
f"/publishers/anthropic/models/{_VERTEX_MODEL}:rawPredict"
),
auth=None,
keys=frozenset({"anthropic_version", "max_tokens", "messages", "stream"}),
),
}
)
@pytest.mark.parametrize("row", tuple(_DROP_CASES))
def test_a_deployment_dropping_safeguards_sends_neither_the_field_nor_the_beta(gateway: Gateway, row: str) -> None:
case: Final = _DROP_CASES[row]
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
api_base=wire.url, additional_drop_params=["safeguards"], **case.deployment(gateway)
)
response: Final = _post(gateway, messages_body(model, marker))
_assert_answered(response, marker)
_expect_messages(_only(wire), target=case.target, auth=case.auth, beta=None, safeguards=OMITTED, keys=case.keys)
_spend_row(f"msg_{marker}")
def test_a_deployment_dropping_context_management_sends_neither_the_field_nor_its_beta(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"anthropic/{ANTHROPIC_MODEL}",
api_base=wire.url,
api_key=ANTHROPIC_KEY,
additional_drop_params=["context_management"],
)
body: Final = {
**messages_body(model, marker, safeguards=OMITTED),
"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]},
}
response: Final = _post(gateway, body)
_assert_answered(response, marker)
_expect_messages(
_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None, safeguards=OMITTED
)
_spend_row(f"msg_{marker}")
@pytest.mark.parametrize("row", tuple(_FOUNDRY_SHAPES))
def test_azure_ai_messages_derive_the_beta_instead_of_forwarding_the_client_header(gateway: Gateway, row: str) -> None:
case: Final = _FOUNDRY_SHAPES[row]
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
response: Final = _post(
gateway, messages_body(model, marker, safeguards=case.sent, stream=case.stream), beta=case.client_beta
)
if case.stream:
_assert_streamed(response, marker)
else:
_assert_answered(response, marker)
observed: Final = _only(wire)
_expect_messages(
observed,
target=_FOUNDRY_TARGET,
auth=("x-api-key", _FOUNDRY_KEY),
beta=case.expected_beta,
safeguards=case.forwarded,
)
_expect_anthropic_version(observed)
_spend_row(f"msg_{marker}")
def test_anthropic_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
reply: Final = client.messages.create(
model=model,
max_tokens=16,
messages=[{"role": "user", "content": f"Question marker-{marker}"}],
extra_body=_SDK_EXTRA,
)
assert reply.id == f"msg_{marker}", reply
_expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA)
_spend_row(f"msg_{marker}")
async def test_anthropic_sdk_async_stream_safeguards_derive_the_beta_and_log_spend(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
client: Final = anthropic.AsyncAnthropic(
base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0
)
async with client.messages.with_streaming_response.create(
model=model,
max_tokens=16,
messages=[{"role": "user", "content": f"Question marker-{marker}"}],
stream=True,
extra_body=_SDK_EXTRA,
) as streamed:
pieces: Final = [
event.delta.text
async for event in await streamed.parse()
if event.type == "content_block_delta" and event.delta.type == "text_delta"
]
assert "".join(pieces) == f"answer marker-{marker}", pieces
_expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA)
_spend_row(f"msg_{marker}")
def test_azure_ai_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
reply: Final = client.messages.create(
model=model,
max_tokens=16,
messages=[{"role": "user", "content": f"Question marker-{marker}"}],
extra_body=_SDK_EXTRA,
)
assert reply.id == f"msg_{marker}", reply
_expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
_spend_row(f"msg_{marker}")
@dataclass(frozen=True, slots=True, kw_only=True)
class _ProviderCase:
parameters: Mapping[str, JsonValue]
auth: tuple[str, str] | None
expected_beta: str | None
target: str = "/v1/messages"
api_base_suffix: str = ""
model_info: Mapping[str, JsonValue] | None = None
client_beta: str | None = None
keys: frozenset[str] = _MESSAGES_KEYS
extra_headers: Mapping[str, str] = MappingProxyType({})
body_beta: list[str] | None = None
_PROVIDERS: Final = MappingProxyType(
{
"D14_deepseek_with_claude_code_header": _ProviderCase(
parameters={"model": "deepseek/deepseek-chat", "api_key": "synthetic-deepseek-key"},
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
target="/anthropic/v1/messages",
auth=("x-api-key", "synthetic-deepseek-key"),
expected_beta=None,
),
"D15_tencent_with_claude_code_header": _ProviderCase(
parameters={"model": "tencent/deepseek-v4-flash", "api_key": "synthetic-tencent-key"},
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
auth=("x-api-key", "synthetic-tencent-key"),
expected_beta=None,
),
"D16_minimax_with_claude_code_header": _ProviderCase(
parameters={"model": "minimax/MiniMax-M2", "api_key": "synthetic-minimax-key"},
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
auth=("x-api-key", "synthetic-minimax-key"),
expected_beta=None,
),
"D17_openai_compatible_native_passthrough": _ProviderCase(
parameters={"model": "openai/claude-sonnet-4-5", "api_key": "synthetic-openai-like-key"},
api_base_suffix="/v1",
model_info={"supported_endpoints": ["/v1/messages"]},
auth=("authorization", "Bearer synthetic-openai-like-key"),
expected_beta=None,
),
"D18_sail_with_claude_code_header": _ProviderCase(
parameters={"model": "sail/sail-claude", "api_key": "synthetic-sail-key"},
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
auth=("authorization", "Bearer synthetic-sail-key"),
expected_beta=None,
),
"D19_edenai": _ProviderCase(
parameters={"model": "edenai/claude-sonnet-4-5", "api_key": "synthetic-edenai-key"},
auth=("authorization", "Bearer synthetic-edenai-key"),
expected_beta=None,
),
"D21_bedrock_claude_platform": _ProviderCase(
parameters={
"model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}",
"api_key": "synthetic-platform-key",
"workspace_id": "ws-synthetic",
"aws_region_name": "us-east-1",
},
auth=("x-api-key", "synthetic-platform-key"),
expected_beta=BETA,
extra_headers={"anthropic-workspace-id": "ws-synthetic"},
),
"D22_bedrock_claude_platform_with_claude_code_header": _ProviderCase(
parameters={
"model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}",
"api_key": "synthetic-platform-key",
"workspace_id": "ws-synthetic",
"aws_region_name": "us-east-1",
},
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
auth=("x-api-key", "synthetic-platform-key"),
expected_beta=_MERGED_BETAS,
extra_headers={"anthropic-workspace-id": "ws-synthetic"},
),
"D23_bedrock_invoke": _ProviderCase(
parameters={
"model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
"api_key": "synthetic-bedrock-bearer",
"aws_region_name": "us-east-1",
},
target="/model/anthropic.claude-sonnet-4-5-20250929-v1:0/invoke",
auth=("authorization", "Bearer synthetic-bedrock-bearer"),
expected_beta=None,
keys=frozenset({"anthropic_beta", "anthropic_version", "max_tokens", "messages"}),
body_beta=[BETA],
),
"D24_bedrock_mantle": _ProviderCase(
parameters={
"model": "bedrock_mantle/anthropic.claude-sonnet-5-v1:0",
"api_key": "synthetic-mantle-bearer",
"aws_region_name": "us-east-1",
},
target="/anthropic/v1/messages",
auth=("authorization", "Bearer synthetic-mantle-bearer"),
expected_beta=BETA,
keys=frozenset({"max_tokens", "messages", "model"}),
),
}
)
@pytest.mark.parametrize("row", tuple(_PROVIDERS))
def test_anthropic_compatible_providers_receive_the_beta_only_where_the_proxy_derives_it(
gateway: Gateway, row: str
) -> None:
case: Final = _PROVIDERS[row]
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model_info=case.model_info, api_base=wire.url + case.api_base_suffix, **case.parameters
)
response: Final = _post(gateway, messages_body(model, marker), beta=case.client_beta)
_assert_answered(response, marker)
observed: Final = _only(wire)
_expect_messages(observed, target=case.target, auth=case.auth, beta=case.expected_beta, keys=case.keys)
for name, value in case.extra_headers.items():
assert observed.headers.get(name) == value, observed.header_names
assert observed.body.get("anthropic_beta") == case.body_beta, observed.body
_spend_row(f"msg_{marker}")
def test_github_copilot_messages_forward_safeguards_without_deriving_the_beta(
gateway: Gateway, tmp_path: Path
) -> None:
marker: Final = uuid.uuid4().hex
token_dir: Final = tmp_path / "copilot"
token_dir.mkdir()
with wire_server(anthropic_peer) as wire:
(token_dir / "api-key.json").write_text(
json.dumps({"token": "synthetic-copilot-token", "expires_at": 4102444800, "endpoints": {"api": wire.url}})
)
with (
owned_proxy(gateway, tmp_path, {"GITHUB_COPILOT_TOKEN_DIR": str(token_dir)}) as owned,
owned.scenario() as scenario,
):
model: Final = scenario.model(model="github_copilot/claude-sonnet-4-5", api_base=wire.url)
response: Final = _post(owned, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS)
_assert_answered(response, marker)
observed: Final = _only(wire)
_expect_messages(
observed, target="/v1/messages", auth=("authorization", "Bearer synthetic-copilot-token"), beta=None
)
assert observed.headers.get("openai-intent") == "messages-proxy", observed.header_names
_spend_row(f"msg_{marker}")
def test_non_native_deployments_bridge_to_the_responses_api_without_safeguards_or_the_beta(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(_responses_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-openai-key"
)
response: Final = _post(gateway, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS)
assert response.status_code == 200 and f"answer marker-{marker}" in response.text, response.text
observed: Final = _only(wire)
assert (observed.method, observed.target) == ("POST", "/v1/responses")
assert observed.headers.get("authorization") == "Bearer synthetic-openai-key", observed.header_names
assert observed.beta is None, observed.header_names
assert observed.keys == _RESPONSES_KEYS, observed.keys
_spend_row(str(_JSON_OBJECT.validate_json(response.content)["id"]))
def test_chat_completions_forward_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
client: Final = openai.OpenAI(
base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0
)
completion: Final = client.chat.completions.create(
model=model, messages=[{"role": "user", "content": f"Question marker-{marker}"}], extra_body=_SDK_EXTRA
)
assert completion.choices[0].message.content == f"answer marker-{marker}", completion
_expect_messages(
_only(wire),
target="/v1/messages",
auth=("x-api-key", ANTHROPIC_KEY),
beta=None,
keys=frozenset({"max_tokens", "messages", "model"}),
)
_spend_row(completion.id)
async def test_responses_stream_forwards_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
client: Final = openai.AsyncOpenAI(
base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0
)
async with client.responses.with_streaming_response.create(
model=model, input=f"Question marker-{marker}", stream=True, extra_body=_SDK_EXTRA
) as streamed:
events: Final = [event async for event in await streamed.parse()]
deltas: Final = [event.delta for event in events if event.type == "response.output_text.delta"]
assert "".join(deltas) == f"answer marker-{marker}", events
assert any(event.type == "response.completed" for event in events), events
_expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None)
_responses_spend_row(model)
_FOUNDRY_REJECTED_SHAPES: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"S1_int": 1, "S2_empty_string": "", "S3_five_kb_string": "x" * 5000}
)
@pytest.mark.parametrize("row", tuple(_FOUNDRY_REJECTED_SHAPES))
def test_foundry_rejecting_a_non_array_safeguards_value_reaches_the_caller_after_one_attempt(
gateway: Gateway, row: str
) -> None:
value: Final = _FOUNDRY_REJECTED_SHAPES[row]
marker: Final = uuid.uuid4().hex
with (
wire_server(_error_peer(400, "invalid_request_error", _FOUNDRY_ARRAY_ERROR)) as wire,
gateway.scenario() as scenario,
):
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
response: Final = _post(gateway, messages_body(model, marker, safeguards=value))
assert response.status_code == 400 and _FOUNDRY_ARRAY_ERROR in response.text, response.text
_expect_messages(
_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA, safeguards=value
)
def test_duplicate_safeguards_keys_resolve_to_the_last_value(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
raw: Final = (
f'{{"model": "{model}", "max_tokens": 16, "num_retries": 0, "cache": {{"no-cache": true}}, '
f'"messages": [{{"role": "user", "content": "Question marker-{marker}"}}], '
f'"safeguards": {json.dumps(SAFEGUARDS)}, "safeguards": null}}'
)
response: Final = gateway.client.post(
"/v1/messages",
content=raw.encode(),
headers={
"Authorization": f"Bearer {gateway.key}",
"content-type": "application/json",
**_ANTHROPIC_VERSION,
},
)
_assert_answered(response, marker)
_expect_messages(
_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=None, safeguards=OMITTED
)
_spend_row(f"msg_{marker}")
def test_the_same_safeguards_sent_twice_are_two_upstream_calls_with_two_spend_rows(gateway: Gateway) -> None:
markers: Final = (uuid.uuid4().hex, uuid.uuid4().hex)
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
responses: Final = tuple(_post(gateway, messages_body(model, marker)) for marker in markers)
for marker, response in zip(markers, responses, strict=True):
_assert_answered(response, marker)
received: Final = tuple(_observe(request) for request in wire.drain())
assert len(received) == 2, received
for observed in received:
_expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
assert sorted(item.marker for item in received) == sorted(markers)
for marker in markers:
_spend_row(f"msg_{marker}")
_UPSTREAM_ERRORS: Final = MappingProxyType(
{
"S6_unauthorized_deployment": (401, "authentication_error", "invalid x-api-key"),
"S7_foundry_extra_inputs_rejection": (400, "invalid_request_error", _FOUNDRY_EXTRA_INPUT_ERROR),
}
)
@pytest.mark.parametrize("row", tuple(_UPSTREAM_ERRORS))
def test_upstream_errors_reach_the_caller_with_the_beta_sent_once(gateway: Gateway, row: str) -> None:
status, kind, message = _UPSTREAM_ERRORS[row]
marker: Final = uuid.uuid4().hex
with wire_server(_error_peer(status, kind, message)) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
response: Final = _post(gateway, messages_body(model, marker))
assert response.status_code == status and message in response.text, response.text
_expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
def _post_noting_whether_upstream_was_called(
gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire
) -> tuple[httpx.Response, bool]:
before: Final = wire.received.qsize()
response: Final = _post(gateway, body)
assert response.status_code == 200, response.text
return response, wire.received.qsize() > before
def _until_served_from_cache(gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire) -> tuple[httpx.Response, int]:
deadline: Final = time.monotonic() + 20
for attempt in itertools.count(1):
response, called_upstream = _post_noting_whether_upstream_was_called(gateway, body, wire)
if not called_upstream:
return response, attempt
assert time.monotonic() < deadline, "the response cache never served the repeated request"
raise AssertionError("unreachable")
def test_the_response_cache_serves_the_repeated_request_without_an_upstream_call(gateway: Gateway) -> None:
marker: Final = uuid.uuid4().hex
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
body: Final = messages_body(model, marker, no_cache=False)
first: Final = _post(gateway, body)
_assert_answered(first, marker)
_expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
_spend_row(f"msg_{marker}")
hit, attempts = _until_served_from_cache(gateway, body, wire)
assert hit.json()["id"] == first.json()["id"], hit.text
assert hit.json()["content"] == first.json()["content"], hit.text
_cache_hit_spend_row(f"msg_{marker}")
misses: Final = tuple(_observe(request) for request in wire.drain())
assert len(misses) == attempts - 1, (attempts, misses)
for observed in misses:
_expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,186 @@
import json
import uuid
from collections.abc import Iterator, Mapping
from typing import Final
import pytest
from integration._support.wire import Reply, Request, wire_server
from pydantic import BaseModel, ConfigDict, Field, JsonValue
import litellm
_QWEN_BACKEND: Final = "databricks-qwen35-122b-a10b"
_CLAUDE_BACKEND: Final = "databricks-claude-haiku-4-5"
_API_KEY: Final = "synthetic-databricks-key"
_ORDER_JSON: Final = '{"buyer":{"name":"Ada"},"seller":{"name":"Bo"}}'
_MESSAGES: Final[tuple[Mapping[str, str], ...]] = ({"role": "user", "content": "Return the order as JSON."},)
class Person(BaseModel):
model_config = ConfigDict(frozen=True)
name: str
class Order(BaseModel):
model_config = ConfigDict(frozen=True)
buyer: Person
seller: Person
class _StrictObject(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
additional_properties: bool | None = Field(default=None, alias="additionalProperties")
class _OrderSchema(_StrictObject):
defs: Mapping[str, _StrictObject] = Field(alias="$defs")
class _JsonSchema(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
name: str
strict: bool | None = None
json_schema: dict[str, JsonValue] = Field(alias="schema")
class _ResponseFormat(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
type: str
json_schema: _JsonSchema
class _Function(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
name: str
parameters: dict[str, JsonValue]
class _Tool(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
function: _Function
class _Outbound(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
model: str
response_format: _ResponseFormat | None = None
tools: tuple[_Tool, ...] = ()
class _CallerMessage(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
content: str
class _CallerChoice(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
message: _CallerMessage
class _CallerCompletion(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
choices: tuple[_CallerChoice, ...]
def _completion(backend: str) -> Reply:
claude: Final = "claude" in backend
tool_call: Final[dict[str, JsonValue]] = {
"id": "json-tool-call",
"type": "function",
"function": {"name": "json_tool_call", "arguments": _ORDER_JSON},
}
message: Final[dict[str, JsonValue]] = (
{"role": "assistant", "content": None, "tool_calls": [tool_call]}
if claude
else {"role": "assistant", "content": _ORDER_JSON}
)
body: Final[dict[str, JsonValue]] = {
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": backend,
"choices": [{"index": 0, "message": message, "finish_reason": "tool_calls" if claude else "stop"}],
"usage": {"prompt_tokens": 9, "completion_tokens": 7, "total_tokens": 16},
}
return Reply(body=json.dumps(body).encode())
def _refs(node: JsonValue) -> Iterator[str]:
if isinstance(node, list):
for item in node:
yield from _refs(item)
return
if not isinstance(node, dict):
return
for key, value in node.items():
if key == "$ref" and isinstance(value, str):
yield value
else:
yield from _refs(value)
def _sent_schema(backend: str, received: tuple[Request, ...]) -> dict[str, JsonValue]:
assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
outbound: Final = _Outbound.model_validate_json(received[0].body)
assert outbound.model == backend
if "claude" in backend:
assert outbound.response_format is None
assert [tool.function.name for tool in outbound.tools] == ["json_tool_call"]
return outbound.tools[0].function.parameters
assert outbound.response_format is not None
assert (outbound.response_format.type, outbound.response_format.json_schema.name) == ("json_schema", "Order")
assert outbound.response_format.json_schema.strict is True
return outbound.response_format.json_schema.json_schema
def _assert_strict_order_schema(schema: dict[str, JsonValue]) -> None:
assert set(_refs(schema)) == {"#/$defs/Person"}, schema
parsed: Final = _OrderSchema.model_validate(schema)
assert (parsed.additional_properties, parsed.defs["Person"].additional_properties) == (False, False), schema
def _caller_order(response: litellm.ModelResponse) -> Order:
return Order.model_validate_json(
_CallerCompletion.model_validate_json(response.model_dump_json()).choices[0].message.content
)
_BACKENDS: Final = (
pytest.param(_QWEN_BACKEND, id="qwen"),
pytest.param(_CLAUDE_BACKEND, id="claude"),
)
_EXPECTED_ORDER: Final = Order(buyer=Person(name="Ada"), seller=Person(name="Bo"))
@pytest.mark.parametrize("backend", _BACKENDS)
def test_databricks_pydantic_response_format_sends_strict_defs_refs(backend: str) -> None:
with wire_server(lambda _request: _completion(backend)) as wire:
response: Final = litellm.completion(
model=f"databricks/{backend}",
api_base=wire.url,
api_key=_API_KEY,
messages=[dict(message) for message in _MESSAGES],
response_format=Order,
num_retries=0,
)
received: Final = wire.drain()
_assert_strict_order_schema(_sent_schema(backend, received))
assert isinstance(response, litellm.ModelResponse)
assert _caller_order(response) == _EXPECTED_ORDER
@pytest.mark.parametrize("backend", _BACKENDS)
async def test_databricks_pydantic_response_format_sends_strict_defs_refs_async(backend: str) -> None:
with wire_server(lambda _request: _completion(backend)) as wire:
response: Final = await litellm.acompletion(
model=f"databricks/{backend}",
api_base=wire.url,
api_key=_API_KEY,
messages=[dict(message) for message in _MESSAGES],
response_format=Order,
num_retries=0,
)
received: Final = wire.drain()
_assert_strict_order_schema(_sent_schema(backend, received))
assert isinstance(response, litellm.ModelResponse)
assert _caller_order(response) == _EXPECTED_ORDER

View file

@ -1,14 +1,20 @@
import json
import os
from datetime import datetime, timezone
from functools import partial
from hashlib import sha256
from itertools import chain
from typing import Final
from uuid import uuid4
import httpx
import pytest
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.upstream import delete_scenario, register_scenario
from integration._support.database import read_rows, write_rows
from integration._support.prisma_clients import connected_proxy_prisma
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
from pydantic import JsonValue
from pydantic import JsonValue, TypeAdapter
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
@ -59,17 +65,17 @@ def _accepting_routes() -> dict[str, JsonResponse | TextResponse]:
}
def _gone_at_provider_routes() -> RoutedResponse:
def _gone_at_provider_routes(*, retryable: bool = False) -> RoutedResponse:
return RoutedResponse(
content_type="application/x-routed",
routes={
**_accepting_routes(),
"GET /batches/batch-$REQUEST_ID": JsonResponse(
content_type="application/json",
status=404,
status=503 if retryable else 404,
body={
"error": {
"message": "No batch found with id 'batch-$REQUEST_ID'.",
"message": "Temporary failure" if retryable else "No batch found with id 'batch-$REQUEST_ID'.",
"type": "invalid_request_error",
"param": "id",
"code": "batch_not_found",
@ -80,7 +86,7 @@ def _gone_at_provider_routes() -> RoutedResponse:
)
def _completed_routes() -> RoutedResponse:
def _completed_routes(*, ready: bool = True) -> RoutedResponse:
output_line: Final = {
"id": "batch_req_1",
"custom_id": "r1",
@ -106,7 +112,8 @@ def _completed_routes() -> RoutedResponse:
routes={
**_accepting_routes(),
"GET /batches/batch-$REQUEST_ID": JsonResponse(
content_type="application/json", body=_batch("completed", "file-out-$REQUEST_ID")
content_type="application/json",
body=_batch("completed", "file-out-$REQUEST_ID") if ready else _batch("in_progress", None),
),
"GET /files/file-out-$REQUEST_ID/content": TextResponse(
content_type="application/jsonl", body=json.dumps(output_line, separators=(",", ":")) + "\n"
@ -115,7 +122,37 @@ def _completed_routes() -> RoutedResponse:
)
def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> str:
def _batch_was_polled(upstream: ScenarioHandle) -> bool:
response: Final = httpx.get(f"{upstream.control_url}/__observations", timeout=15, trust_env=False)
response.raise_for_status()
observations: Final = TypeAdapter(dict[str, list[dict[str, JsonValue]]]).validate_json(response.content)
return any(
item.get("method") == "GET" and string_value(item["path"]).startswith(f"/{upstream.scenario_id}/batches/")
for item in observations["requests"]
)
def _missing_output_routes() -> RoutedResponse:
return RoutedResponse(
content_type="application/x-routed",
routes={
**_completed_routes().routes,
"GET /files/file-out-$REQUEST_ID/content": JsonResponse(
content_type="application/json",
status=404,
body={
"error": {
"message": "No file found with id 'file-out-$REQUEST_ID'.",
"type": "invalid_request_error",
"code": "file_not_found",
}
},
),
},
)
def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> tuple[str, ScenarioHandle]:
scenario_id: Final = f"poll-{marker}-{sha256(os.urandom(16)).hexdigest()[:12]}"
handle: Final = register_scenario(scenario_id, routes)
scenario.cleanups.callback(delete_scenario, handle)
@ -133,7 +170,7 @@ def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse
},
)
scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
return string_value(created["model_name"])
return string_value(created["model_name"]), handle
def _submitted_batch_id(gateway: Gateway, key: str, model_name: str) -> str:
@ -173,20 +210,51 @@ def _managed_rows(batch_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]:
)
@pytest.mark.parametrize(
("retryable", "old_completed", "backlog"),
[(False, False, 0), (False, True, 0), (True, False, 0), (True, True, 0), (True, False, 1000)],
)
@pytest.mark.timeout(180)
@pytest.mark.covers("quota_management.spend_tracking.batch_costs.uncostable_rows_retire_so_newer_batches_are_costed")
def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_polling(gateway: Gateway) -> None:
def test_uncostable_batches_do_not_starve_billable_batches(
gateway: Gateway, retryable: bool, old_completed: bool, backlog: int
) -> None:
with gateway.scenario() as scenario:
key: Final = scenario.key()
owner: Final = f"batch-backlog-{uuid4().hex}"
if backlog:
_insert_scan_rows(owner, "old", backlog, datetime(2000, 1, 1, tzinfo=timezone.utc))
scenario.cleanups.callback(
write_rows, 'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,)
)
team_id: Final = scenario.team()
key: Final = scenario.key(team_id=team_id)
gone_batch_ids: Final = tuple(
_submitted_batch_id(
gateway, key, _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes())
gateway,
key,
_scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes(retryable=retryable))[0],
)
for index in range(MAX_OBJECTS_PER_POLL_CYCLE)
)
costable_batch_id: Final = _submitted_batch_id(
gateway, key, _scripted_deployment(scenario, "costable", _completed_routes())
missing_output_id: Final = _submitted_batch_id(
gateway, key, _scripted_deployment(scenario, "missing-output", _missing_output_routes())[0]
)
model, upstream = _scripted_deployment(scenario, "costable", _completed_routes(ready=False))
costable_batch_id: Final = _submitted_batch_id(gateway, key, model)
for batch_id in (*gone_batch_ids, missing_output_id, costable_batch_id):
scenario.cleanups.callback(
write_rows,
'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
(batch_id,),
)
if old_completed:
write_rows(
'UPDATE "LiteLLM_ManagedObjectTable" SET status = %s, '
"created_at = NOW() - INTERVAL '30 days' WHERE unified_object_id = %s",
("completed", costable_batch_id),
)
assert eventually(lambda: _batch_was_polled(upstream), bool, seconds=120)
register_scenario(upstream.scenario_id, _completed_routes(), control_url=upstream.control_url)
spend_rows: Final = eventually(
lambda: read_rows(
'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
@ -209,4 +277,134 @@ def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_pollin
}
]
assert _managed_rows((costable_batch_id,)) == [{"batch_processed": True}]
assert _managed_rows(gone_batch_ids) == [{"batch_processed": True}] * MAX_OBJECTS_PER_POLL_CYCLE
assert eventually(
lambda: _managed_rows((missing_output_id,)),
lambda rows: rows == [{"batch_processed": True}],
seconds=120,
) == [{"batch_processed": True}]
expected_blockers: Final = [{"batch_processed": not retryable}] * MAX_OBJECTS_PER_POLL_CYCLE
blockers: Final = eventually(
lambda: _managed_rows(gone_batch_ids), lambda rows: rows == expected_blockers, seconds=120
)
assert blockers == expected_blockers
expected_spend: Final = BATCH_COST_SHARE * (
PROMPT_TOKENS * INPUT_COST_PER_TOKEN + COMPLETION_TOKENS * OUTPUT_COST_PER_TOKEN
)
for table, column, identity in (
("LiteLLM_VerificationToken", "token", sha256(key.encode()).hexdigest()),
("LiteLLM_TeamTable", "team_id", team_id),
):
balance: Final = eventually(
partial(read_rows, f'SELECT spend FROM "{table}" WHERE {column} = %s', (identity,)),
lambda rows: len(rows) == 1 and rows[0]["spend"] == pytest.approx(expected_spend),
seconds=70,
)
assert balance == [{"spend": pytest.approx(expected_spend)}]
def _insert_scan_rows(owner: str, group: str, size: int, created_at: datetime) -> tuple[str, ...]:
prefix: Final = f"{owner}-{group}-"
write_rows(
'INSERT INTO "LiteLLM_ManagedObjectTable" '
"(id, unified_object_id, model_object_id, file_object, file_purpose, status, created_by, created_at, updated_at) "
"SELECT %s || lpad(n::text, 8, '0'), %s || n, %s || n, '{}'::jsonb, 'batch', 'completed', %s, "
"%s::timestamptz, NOW() FROM generate_series(0, %s::int - 1) n",
(prefix, prefix, prefix, owner, created_at.isoformat(), str(size)),
)
return tuple(f"{prefix}{index:08}" for index in range(size))
@pytest.mark.asyncio
@pytest.mark.parametrize("count", [2 * MAX_OBJECTS_PER_POLL_CYCLE, 2 * MAX_OBJECTS_PER_POLL_CYCLE + 1, 1000])
async def test_arrivals_progress_while_retry_scan_advances_and_wraps(count: int) -> None:
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
owner: Final = f"batch-scan-{uuid4().hex}"
try:
old_ids: Final = _insert_scan_rows(owner, "old", count, datetime(2000, 1, 1, tzinfo=timezone.utc))
async with connected_proxy_prisma() as prisma:
logging: Final = ProxyLogging(UserApiKeyCache())
router: Final = Router()
poller: Final = CheckBatchCost(logging, prisma, router)
where: Final = {"created_by": owner, "batch_processed": False}
first: Final = await poller._find_jobs(where)
assert tuple(row.id for row in first) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
recent_ids: Final = _insert_scan_rows(
owner, "recent", MAX_OBJECTS_PER_POLL_CYCLE + 1, datetime.now(timezone.utc)
)
async def poll_with_arrival(turn: int) -> tuple[str, ...]:
if turn % 4 == 0:
_insert_scan_rows(owner, f"arrival-{turn}", 1, datetime.now(timezone.utc))
page: Final = await poller._find_jobs(where)
assert len(page) <= MAX_OBJECTS_PER_POLL_CYCLE
write_rows(
'UPDATE "LiteLLM_ManagedObjectTable" SET batch_processed = true '
"WHERE created_by = %s AND id = ANY(%s) AND id LIKE %s",
(owner, [row.id for row in page], f"{owner}-arrival-%"),
)
return tuple(row.id for row in page)
pages: Final = tuple(
[await poll_with_arrival(turn) for turn in range(4 * (count // MAX_OBJECTS_PER_POLL_CYCLE + 2))]
)
selected: Final = tuple(chain.from_iterable(pages))
assert all(selected.count(identity) >= 2 for identity in recent_ids)
selected_old: Final = tuple(identity for identity in selected if identity in old_ids)
assert selected_old[: count - MAX_OBJECTS_PER_POLL_CYCLE] == old_ids[MAX_OBJECTS_PER_POLL_CYCLE:]
assert set(old_ids).issubset(selected_old)
latest_ids: Final = _insert_scan_rows(owner, "latest", 1, datetime.now(timezone.utc))
followups: Final = tuple([await poller._find_jobs(where) for _ in range(12)])
assert any(latest_ids[0] in tuple(row.id for row in page) for page in followups)
restarted: Final = CheckBatchCost(logging, prisma, router)
restart_page: Final = await restarted._find_jobs(where)
assert tuple(row.id for row in restart_page) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
finally:
write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,))
@pytest.mark.asyncio
@pytest.mark.parametrize("steady_arrivals", [False, True])
@pytest.mark.parametrize("backlog", [0, 1000])
async def test_recent_batches_are_revisited_without_scanning_old_backlog(steady_arrivals: bool, backlog: int) -> None:
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
owner: Final = f"batch-repeat-{uuid4().hex}"
try:
_insert_scan_rows(owner, "old", backlog, datetime(2000, 1, 1, tzinfo=timezone.utc))
async with connected_proxy_prisma() as prisma:
poller: Final = CheckBatchCost(ProxyLogging(UserApiKeyCache()), prisma, Router())
where: Final = {"created_by": owner, "batch_processed": False}
await poller._find_jobs(where)
target: Final = _insert_scan_rows(owner, "target", 1, datetime.now(timezone.utc))[0]
assert target in tuple(row.id for row in await poller._find_jobs(where))
for cycle in range(3):
if steady_arrivals:
_insert_scan_rows(owner, f"arrival-{cycle}", 1, datetime.now(timezone.utc))
pages: Final = tuple([await poller._find_jobs(where) for _ in range(4)])
assert any(target in tuple(row.id for row in page) for page in pages)
assert all(len(page) <= MAX_OBJECTS_PER_POLL_CYCLE for page in pages)
write_rows(
'UPDATE "LiteLLM_ManagedObjectTable" SET batch_processed = true WHERE created_by = %s AND id != %s',
(owner, target),
)
write_rows(
'UPDATE "LiteLLM_ManagedObjectTable" SET created_at = %s WHERE id = %s',
(datetime(2000, 1, 1, tzinfo=timezone.utc), target),
)
recovery: Final = tuple([await poller._find_jobs(where) for _ in range(4)])
assert any(target in tuple(row.id for row in page) for page in recovery)
finally:
write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,))

View file

@ -3,7 +3,7 @@ import pytest
import json
from io import BytesIO
from typing import Dict, List
from typing import Dict, List, Final
from litellm.router_utils.batch_utils import (
replace_model_in_jsonl,
_get_router_metadata_variable_name,
@ -333,3 +333,31 @@ def test_is_batch_retrieve_call_type_matches_only_batch_retrieves():
assert is_batch_retrieve_call_type(call_type.value) is False
assert is_batch_retrieve_call_type(None) is False
@pytest.mark.parametrize(
("call_type", "cost", "eligible"),
(("aretrieve_batch", 0.25, True), ("retrieve_batch", 0.25, True), ("acompletion", 0.25, False),
("aretrieve_batch", 0, False), ("aretrieve_batch", -0.25, False)),
)
def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None:
from openai.types import Batch
from litellm.router_utils.batch_utils import batch_id_to_charge_once
batch: Final = Batch(
id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions",
input_file_id="file-budget", object="batch", status="completed",
)
assert batch_id_to_charge_once(call_type, batch, cost) == (batch.id if eligible else None)
assert batch_id_to_charge_once(call_type, {"id": batch.id}, cost) is None
def test_batch_charge_markers_keep_counters_and_batches_independent() -> None:
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.router_utils.batch_utils import batch_charged_once_marker_key
cache: Final = InMemoryCache()
cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1)
assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "first-batch")) == 1
assert cache.get_cache(batch_charged_once_marker_key("deployment-spend", "first-batch")) is None
assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "second-batch")) is None

View file

@ -637,7 +637,12 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
"model": "gpt-3.5-turbo",
"api_key": "sk-must-not-be-recorded",
"messages": [{"role": "user", "content": "hi"}],
metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"},
metadata_key: {
"model_info": {"id": "deployment-1", "base_model": "azure/gpt-3.5-turbo"},
"deployment": "azure/my-35-deployment",
"attempted_retries": 2,
"user_api_key": "sk-proxy",
},
},
e=rate_limit_error,
)
@ -645,6 +650,8 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
{
"model_group": "gpt-3.5-turbo",
"deployment_id": "deployment-1",
"model": "azure/my-35-deployment",
"base_model": "azure/gpt-3.5-turbo",
"exception_type": "RateLimitError",
"exception_string": "litellm.RateLimitError: slow down",
"attempted_retries": 2,
@ -656,6 +663,38 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1
def test_log_retry_reads_base_model_from_the_deployments_litellm_params() -> None:
"""An Azure named deployment declares ``base_model`` under ``litellm_params``, not ``model_info``, and
that is the model the failure row is priced at, so the retry record has to carry it too"""
router = Router(
model_list=[
{
"model_name": "azure-group",
"litellm_params": {
"model": "azure/my-named-deployment",
"api_key": "sk-fake",
"api_base": "https://example.openai.azure.com",
"base_model": "azure/gpt-4o-mini",
},
"model_info": {"id": "dep-azure-named"},
}
]
)
new_kwargs = router.log_retry(
kwargs={
"model": "azure-group",
"metadata": {
"model_info": {"id": "dep-azure-named"},
"deployment": "azure/my-named-deployment",
"attempted_retries": 0,
},
},
e=litellm.Timeout(message="slow", llm_provider="azure", model="azure/my-named-deployment"),
)
record = new_kwargs["metadata"]["previous_models"][0]
assert (record["model"], record["base_model"]) == ("azure/my-named-deployment", "azure/gpt-4o-mini")
def test_update_usage(model_list):
"""Test if the '_update_usage' function is working correctly"""
router = Router(model_list=model_list)

View file

@ -1210,3 +1210,39 @@ class TestCoerceNumericFormFields:
numeric_fields=self.numeric_fields,
)
assert result == {"n": 3, "temperature": None, "image": buffer}
@pytest.mark.asyncio
async def test_form_data_metadata_with_lone_surrogate_escape_is_rejected_like_a_json_body():
"""``json.loads`` accepts ``"\\ud800"`` and yields a lone surrogate that no
provider or database can utf-8 encode. The JSON body path already answers
400, so the multipart ``metadata`` field must too, instead of letting the
value through to the spend log writers."""
mock_request = MagicMock()
mock_request.form = AsyncMock(
return_value=FormData({"model": "whisper-1", "metadata": '{"tags": ["poison-\\ud800-tag"]}'})
)
mock_request.headers = {"content-type": "multipart/form-data"}
mock_request.scope = {}
mock_request.state._cached_headers = None
with pytest.raises(ProxyException) as exc_info:
await _read_request_body(mock_request)
assert exc_info.value.code == "400"
assert "metadata" in exc_info.value.message
@pytest.mark.asyncio
async def test_form_data_metadata_keeps_non_ascii_and_valid_surrogate_pairs():
mock_request = MagicMock()
mock_request.form = AsyncMock(
return_value=FormData({"model": "whisper-1", "metadata": '{"tags": ["caf\\u00e9", "\\ud83d\\ude00"]}'})
)
mock_request.headers = {"content-type": "multipart/form-data"}
mock_request.scope = {}
mock_request.state._cached_headers = None
result = await _read_request_body(mock_request)
assert result["metadata"]["tags"] == ["café", "😀"]

View file

@ -3843,28 +3843,57 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
proxy_logging.failure_handler.assert_not_called()
def _responses_router() -> litellm.Router:
return litellm.Router(
model_list=[
{
"model_name": "native",
"litellm_params": {"model": "openai/gpt-5.1", "api_key": "sk-test"},
"model_info": {"id": "dep-native"},
},
{
"model_name": "bridged",
"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"},
"model_info": {"id": "dep-bridged"},
},
{
"model_name": "forced-bridge",
"litellm_params": {"model": "openai/gpt-5.1", "api_key": "sk-test", "use_chat_completions_api": True},
"model_info": {"id": "dep-forced-bridge"},
},
]
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call_type, expects_flush",
[("aresponses", True), ("responses", True), ("acompletion", False)],
"call_type, model_id, expects_flush",
[
("aresponses", "dep-bridged", True),
("responses", "dep-bridged", True),
("aresponses", "dep-forced-bridge", True),
("aresponses", None, True),
("aresponses", "dep-native", False),
("responses", "dep-native", False),
("acompletion", "dep-bridged", False),
],
)
async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back(
call_type: str, expects_flush: bool
async def test_insert_spend_log_asks_for_an_immediate_flush_only_on_rows_the_bridge_reads_back(
call_type: str, model_id: str | None, expects_flush: bool
):
"""
A `previous_response_id` chained straight off the previous turn reads the DB, so a
Responses row cannot sit in this worker's queue until the monitor's next poll.
A `previous_response_id` chained straight off a bridged turn reads that turn's row from the
DB, so it cannot sit in this worker's queue until the monitor's next poll. A native Responses
provider holds the conversation itself, so its row waits for the batch write like a chat row.
"""
db_writer = DBSpendUpdateWriter()
prisma = _tool_usage_prisma()
prisma.spend_log_flush_requested = asyncio.Event()
payload = {"request_id": "req-1", "call_type": call_type, "model_id": model_id}
await db_writer._insert_spend_log_to_db(
payload={"request_id": "req-1", "call_type": call_type},
prisma_client=prisma,
)
await db_writer._insert_spend_log_to_db(payload=payload, prisma_client=prisma, llm_router=_responses_router())
assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}]
assert prisma.spend_log_transactions == [payload]
assert prisma.spend_log_flush_requested.is_set() is expects_flush

View file

@ -2,12 +2,15 @@ import asyncio
import json
import logging
from datetime import datetime
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth
@ -27,7 +30,7 @@ from litellm.proxy.utils import ProxyUpdateSpend
from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event
from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, Usage
from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage
@pytest.mark.asyncio
@ -93,6 +96,65 @@ async def test_async_post_call_failure_hook():
assert metadata["original_key"] == "original_value"
@pytest.mark.asyncio
async def test_each_earlier_billed_attempt_gets_its_own_spend_row():
"""A timed-out request that the router retried writes one row per billed attempt: the final
attempt under the request id, every earlier Timeout attempt under ``<id>-attempt-<n>`` with its
own cost, attempt number and deployment."""
from litellm.proxy.utils import BilledFailedAttempt
logger = _ProxyDBLogger()
user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key", user_id="u1", team_id="t1")
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
request_data = {
"model": "gpt-4.1",
"litellm_call_id": "call-1",
"standard_logging_object": {"id": "call-1", "model_id": "dep-final"},
"metadata": {
"attempted_retries": 2,
"model_info": {"id": "dep-final", "input_cost_per_token": 0.5},
"previous_models": ({"exception_type": "Timeout"}, {"exception_type": "Timeout"}),
},
"combined_usage_object": usage,
"response_cost": 0.5,
"billed_failed_attempts": (
BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25),
BilledFailedAttempt(attempted_retries=0, deployment_id="dep-final", usage=usage, response_cost=0.5),
),
}
with patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database:
await logger.async_post_call_failure_hook(
request_data=request_data,
original_exception=litellm.Timeout("Request timed out.", model="gpt-4.1", llm_provider="openai"),
user_api_key_dict=user_api_key_dict,
)
rows = [call.kwargs for call in mock_update_database.call_args_list]
assert [row["response_cost"] for row in rows] == [0.5, 0.25, 0.5]
final_row, attempt_0, attempt_1 = (row["kwargs"] for row in rows)
assert final_row["litellm_call_id"] == "call-1"
assert final_row["litellm_params"]["metadata"]["attempted_retries"] == 2
assert attempt_0["litellm_call_id"] == "call-1-attempt-0"
assert attempt_0["standard_logging_object"]["id"] == "call-1-attempt-0"
assert attempt_0["litellm_params"]["metadata"]["attempted_retries"] == 0
assert attempt_0["litellm_params"]["metadata"]["model_info"]["id"] == "dep-a"
assert "previous_models" not in attempt_0["litellm_params"]["metadata"]
assert "billed_failed_attempts" not in attempt_0
assert attempt_0["combined_usage_object"] is usage
assert attempt_1["litellm_call_id"] == "call-1-attempt-1"
assert attempt_1["litellm_params"]["metadata"]["attempted_retries"] == 0
assert attempt_1["litellm_params"]["metadata"]["model_info"]["id"] == "dep-final"
assert final_row["litellm_params"]["metadata"]["model_info"]["id"] == "dep-final"
for row in rows:
assert row["token"] == "test_api_key"
assert row["user_id"] == "u1"
assert row["team_id"] == "t1"
@pytest.mark.asyncio
async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_metadata():
"""
@ -2712,3 +2774,165 @@ async def test_track_cost_callback_failure_alert_never_carries_request_metadata_
assert "headers" in failure_debug_lines[0]
else:
assert failure_debug_lines == []
_BATCH_CALL_START: Final = datetime(2026, 1, 1)
@pytest.mark.asyncio
async def test_batch_accounting_failures_stay_with_their_originating_call() -> None:
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
logger: Final = _ProxyDBLogger()
dependency: Final = SimpleNamespace(failed_tracking_alert=AsyncMock())
async def fail_before_accounting(call_id: str) -> None:
await logger.async_log_success_event(
kwargs={
"litellm_call_id": call_id,
"call_type": "aretrieve_batch",
"model": "batch-model",
"stream": False,
"response_cost": None,
},
response_obj=None,
start_time=_BATCH_CALL_START,
end_time=_BATCH_CALL_START,
)
async def independent_batch() -> bool:
with track_batch_accounting("independent") as outcome:
await fail_before_accounting("independent")
return outcome.retryable_failure
with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency):
with track_batch_accounting("outer") as outer:
await fail_before_accounting("unrelated")
assert outer.retryable_failure is False
with track_batch_accounting("inner") as inner:
await fail_before_accounting("inner")
assert inner.retryable_failure is True
assert outer.retryable_failure is False
independent, _ = await asyncio.gather(independent_batch(), fail_before_accounting("outer"))
assert independent is True
assert outer.retryable_failure is True
with pytest.raises(asyncio.CancelledError):
with track_batch_accounting("cancelled") as cancelled:
raise asyncio.CancelledError
await fail_before_accounting("cancelled")
assert cancelled.retryable_failure is False
@pytest.mark.asyncio
async def test_batch_accounting_does_not_retry_after_settling_a_budget_reservation() -> None:
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": ()}
dependency: Final = SimpleNamespace(failed_tracking_alert=AsyncMock())
with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency):
with track_batch_accounting("reserved") as outcome:
await _ProxyDBLogger().async_log_success_event(
kwargs={
"litellm_call_id": "reserved",
"call_type": "aretrieve_batch",
"model": "batch-model",
"stream": False,
"response_cost": None,
"litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}},
},
response_obj=None,
start_time=_BATCH_CALL_START,
end_time=_BATCH_CALL_START,
)
assert reservation["finalized"] is True
assert outcome.retryable_failure is False
class _SharedCounter:
def __init__(self) -> None:
self.totals: dict[str, float] = {}
async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float:
self.totals[key] = self.totals.get(key, 0) + value
return self.totals[key]
def _grounded_gemini_kwargs(response_cost: float) -> dict:
return {
"model": "vertex_ai/gemini-test",
"call_type": "acompletion",
"response_cost": response_cost,
"litellm_params": {"metadata": {"user_api_key": "sk-test", "user_api_key_user_id": "user-1"}},
"standard_logging_object": {
"response_cost": response_cost,
"request_tags": None,
"cache_hit": None,
"custom_llm_provider": "vertex_ai",
"cost_breakdown": {
"input_cost": 0.0,
"output_cost": 0.0,
"tool_usage_cost": 0.5,
"total_cost": response_cost,
},
"model_map_information": {
"model_map_key": "vertex_ai/gemini-test",
"model_map_value": {
"litellm_provider": "vertex_ai",
"mode": "chat",
"web_search_billing_unit": "per_query",
"search_context_cost_per_query": {"search_context_size_medium": 0.1},
"search_free_requests": {"count": 100, "period": "month", "shared_group": "test"},
},
},
},
"stream": False,
}
@pytest.mark.asyncio
@pytest.mark.parametrize("opted_in", [True, False])
async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when_opted_in(monkeypatch, opted_in):
monkeypatch.setattr(litellm, "apply_search_free_tier", opted_in, raising=False)
counter = _SharedCounter()
kwargs = _grounded_gemini_kwargs(response_cost=1.0)
response = ModelResponse(
usage=Usage(
prompt_tokens=10,
completion_tokens=5,
total_tokens=15,
prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=5),
)
)
with (
patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam
"litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock
),
patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam
"litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock
),
patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam
"litellm.proxy.proxy_server.proxy_logging_obj"
) as mock_proxy_logging,
):
mock_proxy_logging.failed_tracking_alert = AsyncMock()
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True)
mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
mock_proxy_logging.internal_usage_cache.dual_cache.redis_cache = counter
await _ProxyDBLogger()._PROXY_track_cost_callback(
kwargs=kwargs, completion_response=response, start_time=datetime.now(), end_time=datetime.now()
)
await asyncio.sleep(0)
mock_proxy_logging.failed_tracking_alert.assert_not_called()
mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once()
breakdown = kwargs["standard_logging_object"]["cost_breakdown"]
if opted_in:
assert kwargs["response_cost"] == pytest.approx(1.0 - 5 * 0.1)
assert breakdown["original_cost"] == 1.0 and breakdown["discount_amount"] == pytest.approx(5 * 0.1)
assert sum(counter.totals.values()) == 5
else:
assert kwargs["response_cost"] == 1.0 and "discount_amount" not in breakdown
assert counter.totals == {}

View file

@ -1249,6 +1249,32 @@ def test_resolve_llm_passthrough_timeout_validates_only_the_winning_value(
assert resolve_llm_passthrough_timeout(kwargs=kwargs, litellm_params=litellm_params) == expected
def test_resolve_llm_passthrough_timeout_uses_the_deployment_timeout_for_the_requests_service_tier():
litellm_params = {
"timeout": 610,
"stream_timeout": 30,
"timeout_by_service_tier": {"flex": 900},
"stream_timeout_by_service_tier": {"flex": 950},
}
assert resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex"}, litellm_params=litellm_params) == 900.0
assert resolve_llm_passthrough_timeout(kwargs={"service_tier": "default"}, litellm_params=litellm_params) == 610.0
assert (
resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex", "timeout": 20}, litellm_params=litellm_params)
== 20.0
)
assert (
resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex", "stream": True}, litellm_params=litellm_params)
== 950.0
)
assert (
resolve_llm_passthrough_timeout(
kwargs={"service_tier": "flex", "stream": True},
litellm_params={"timeout": 610, "stream_timeout": 30, "timeout_by_service_tier": {"flex": 900}},
)
== 30.0
)
def test_resolve_llm_passthrough_timeout_rejects_a_non_numeric_winner():
with pytest.raises(ValidationError):
resolve_llm_passthrough_timeout(kwargs={"timeout": httpx.Timeout(30.0)})

View file

@ -7637,6 +7637,118 @@ def _apply_stream_usage_tracking(
)
async def _start_partial_responses_stream(self, recorder: _RecordingSuccessLogger, *, relay_text: bool = True):
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
mocked: Final = await litellm.aresponses(
model="gpt-4o-mini",
input="tell me a story",
mock_response="The codename is AZURE-FALCON-42 and the story is long.",
stream=True,
api_key="test-key",
)
logging_obj: Final = mocked.logging_obj
logging_obj.dynamic_async_success_callbacks = [recorder]
stream: Final = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="gpt-4o-mini",
responses_api_provider_config=OpenAIResponsesAPIConfig(),
logging_obj=logging_obj,
request_data={"model": "gpt-4o-mini", "input": "tell me a story"},
custom_llm_provider="openai",
)
created: Final = {
"type": "response.created",
"sequence_number": 0,
"response": {
"id": "resp_partial",
"object": "response",
"created_at": 1700000000.0,
"model": "gpt-4o-mini",
"status": "in_progress",
"output": [],
"parallel_tool_calls": True,
"tool_choice": "auto",
"tools": [],
"service_tier": "priority",
},
}
delta: Final = {
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": "msg_partial",
"output_index": 0,
"content_index": 0,
"delta": "The codename is AZURE-FALCON-42",
}
for event in (created, delta) if relay_text else (created,):
stream._process_chunk(json.dumps(event))
return stream
async def _bill_and_wait(self, recorder: _RecordingSuccessLogger, streamed: object, logging_obj: object) -> bool:
billed: Final = await _bill_partial_streamed_spend_on_disconnect({"litellm_logging_obj": logging_obj}, streamed)
for _ in range(50):
if recorder.success_events:
break
await asyncio.sleep(0.1)
await asyncio.sleep(0.5)
return billed
@pytest.mark.asyncio
async def test_disconnect_bills_partial_native_responses_stream_with_counted_usage(self):
from litellm.types.llms.openai import ResponsesAPIResponse
recorder: Final = _RecordingSuccessLogger()
stream: Final = await self._start_partial_responses_stream(recorder)
billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj)
assert billed is True
assert len(recorder.success_events) == 1
partial_response: Final = recorder.success_events[0]["response_obj"]
assert isinstance(partial_response, ResponsesAPIResponse)
assert partial_response.status == "incomplete"
assert getattr(partial_response, "service_tier") == "priority"
assert partial_response.usage.output_tokens == litellm.token_counter(
model="gpt-4o-mini", text="The codename is AZURE-FALCON-42"
)
assert partial_response.usage.input_tokens > 0
@pytest.mark.asyncio
async def test_disconnect_bills_partial_default_route_messages_stream_through_the_responses_adapter(self):
from litellm.llms.anthropic.pass_through.responses_adapters.streaming_iterator import (
AnthropicResponsesStreamWrapper,
)
from litellm.router import FallbackAwareAnthropicMessagesStream
from litellm.types.llms.openai import ResponsesAPIResponse
async def _sse_frames() -> AsyncGenerator[bytes, None]:
yield b"event: message_start\n\n"
recorder: Final = _RecordingSuccessLogger()
stream: Final = await self._start_partial_responses_stream(recorder)
wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=stream, model="gpt-4o-mini")
streamed: Final = FallbackAwareAnthropicMessagesStream(_sse_frames(), wrapper.async_anthropic_sse_wrapper())
billed: Final = await self._bill_and_wait(recorder, streamed, stream.logging_obj)
assert billed is True
assert len(recorder.success_events) == 1
partial_response: Final = recorder.success_events[0]["response_obj"]
assert isinstance(partial_response, ResponsesAPIResponse)
assert partial_response.usage.output_tokens == litellm.token_counter(
model="gpt-4o-mini", text="The codename is AZURE-FALCON-42"
)
@pytest.mark.asyncio
async def test_disconnect_before_any_responses_output_bills_nothing(self):
recorder: Final = _RecordingSuccessLogger()
stream: Final = await self._start_partial_responses_stream(recorder, relay_text=False)
billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj)
assert billed is False
assert recorder.success_events == []
class TestApplyStreamUsageTracking:
def test_default_injects_usage_and_marks_strip_for_chat_completions(self):
data = {"stream": True, "model": "gpt-5.4-nano"}

View file

@ -7764,6 +7764,7 @@ _PLANTED_STAMPS = {
"attempted_fallbacks": 99,
"original_model_group": "spoofed-group",
"request_retry_count": -100,
"previous_models": [{"exception_type": "Timeout", "deployment_id": "spoofed", "attempted_retries": 101}],
"_client_output_ceiling": {"api_base": "https://attacker.example"},
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9,
"client_key": "client_value",
@ -7798,6 +7799,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
assert "original_model_group" not in updated["metadata"]
assert "_client_output_ceiling" not in updated["metadata"]
assert "request_retry_count" not in updated["metadata"]
assert "previous_models" not in updated["metadata"]
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"

View file

@ -58,9 +58,7 @@ class TestStripClientPricingOverrides:
# The strip set is built from the model so additions are picked up
# automatically — this test guards against the model and the strip
# set drifting apart if someone replaces the auto-derivation later.
assert _CLIENT_PRICING_CONTROL_FIELDS == frozenset(
CustomPricingLiteLLMParams.model_fields.keys()
)
assert _CLIENT_PRICING_CONTROL_FIELDS >= frozenset(CustomPricingLiteLLMParams.model_fields.keys())
# Sanity: the obvious top-level pricing fields are in the set.
for field in (
"input_cost_per_token",
@ -84,6 +82,22 @@ class TestStripClientPricingOverrides:
"messages": [{"role": "user", "content": "hi"}],
}
def test_root_base_model_dropped(self):
# ``base_model`` picks the cost-map entry a spend row is priced with
# (success rows via the standard logging payload, billed failure rows
# via the failure lift), so a caller naming a cheaper model lowers
# their own recorded spend.
data = {
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"base_model": "gpt-4o-mini",
}
_strip_client_pricing_overrides(data)
assert data == {
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
}
def test_metadata_model_info_dropped(self):
data = {
"model": "gpt-4",
@ -264,6 +278,7 @@ async def test_add_litellm_data_to_request_skips_strip_with_key_opt_in():
"model": "gpt-4",
"messages": [{"role": "user", "content": "hi"}],
"input_cost_per_token": 0.0001,
"base_model": "gpt-4o-mini",
"metadata": {"model_info": {"output_cost_per_token": 0.0002}},
}
@ -279,6 +294,7 @@ async def test_add_litellm_data_to_request_skips_strip_with_key_opt_in():
assert updated["input_cost_per_token"] == 0.0001
assert updated["metadata"]["model_info"] == {"output_cost_per_token": 0.0002}
assert updated["base_model"] == "gpt-4o-mini"
@pytest.mark.asyncio

View file

@ -2,6 +2,7 @@ import datetime as real_datetime
import smtplib
from typing import Final
import httpx
import pytest
from fastapi import HTTPException
@ -443,6 +444,47 @@ class TestPostCallFailureHookLiftsFirstApiCallStartTime:
assert "litellm_logging_obj" not in request_data
class TestPostCallFailureHookDoesNotInventAProviderHandoff:
"""A failure the proxy raised before any provider call (here an auth error) is logged through
_handle_logging_proxy_only_error, which fires pre_call for the callbacks. That must not stamp a
handoff instant onto the request, or every downstream reader (preprocessing latency, the
dispatched check behind the failure-row estimate) sees a provider call that never happened."""
@pytest.mark.asyncio
async def test_proxy_only_failure_lifts_no_first_api_call_start_time(self):
from datetime import datetime
from unittest.mock import AsyncMock, patch
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth
logging_obj = Logging(
model="gpt-4.1",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="acompletion",
start_time=datetime.now(),
litellm_call_id="pre-handoff-auth-failure",
function_id="pre-handoff",
)
request_data: Final[dict[str, object]] = {"litellm_logging_obj": logging_obj, "model": "gpt-4.1", "metadata": {}}
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.alert_types = []
with (
patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()),
patch.object(proxy_logging_obj, "_dispatch_proxy_only_failure_handlers", new=AsyncMock()),
):
await proxy_logging_obj.post_call_failure_hook(
request_data=request_data,
original_exception=Exception("Authentication Error, invalid key"),
user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"),
error_type=ProxyErrorTypes.auth_error,
)
assert request_data.get("first_api_call_start_time") is None
assert "api_call_start_time" not in logging_obj.model_call_details
class TestPostCallFailureHookLiftsRecoveredPartialSpend:
"""A stream that broke mid-flight still billed the provider for the chunks
already delivered. The streaming handler stashes that recovered usage and
@ -687,11 +729,13 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
"""A non-stream request that failed after dispatch (timeout, provider
error) consumed provider-billed input tokens but recovered no usage.
post_call_failure_hook must estimate the input side onto request_data so
the spend log's failure row records what was sent instead of zero, while
never charging spend for the failure (LIT-5690).
the spend log's failure row records what was sent instead of zero
(LIT-5690). On a timeout that estimate is priced at the model's input
rate since the provider bills the prompt it accepted (LIT-145); other
post-dispatch failures keep zero spend.
"""
async def _run(self, request_data):
async def _run(self, request_data, original_exception: Exception | None = None):
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import UserAPIKeyAuth
@ -701,17 +745,275 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()):
await proxy_logging_obj.post_call_failure_hook(
request_data=request_data,
original_exception=Exception("boom"),
original_exception=original_exception or Exception("boom"),
user_api_key_dict=UserAPIKeyAuth(),
)
def _timeout(self, model: str, cause: Exception) -> Exception:
try:
try:
raise cause
except Exception as inner:
raise litellm.Timeout("Request timed out.", model=model, llm_provider="openai") from inner
except litellm.Timeout as timed_out:
return timed_out
def _logging_obj(self, model_call_details):
logging_obj = MagicMock()
logging_obj.model_call_details = model_call_details
return logging_obj
@pytest.mark.asyncio
async def test_dispatched_failure_estimates_input_tokens_with_zero_cost(self):
async def test_timed_out_request_estimates_input_tokens_priced_at_input_rate(self):
from datetime import datetime
from litellm.types.utils import Usage
request_data = {
"litellm_logging_obj": self._logging_obj(
{
"first_api_call_start_time": datetime.now(),
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "count these input tokens please"}],
"call_type": "acompletion",
}
),
"metadata": {},
"response_cost": 123.0,
}
await self._run(
request_data, original_exception=self._timeout("gpt-3.5-turbo", httpx.ReadTimeout("no bytes in time"))
)
estimated = request_data["combined_usage_object"]
assert isinstance(estimated, Usage)
assert estimated.prompt_tokens > 0
assert estimated.completion_tokens == 0
assert estimated.total_tokens == estimated.prompt_tokens
expected_input_cost, _ = litellm.cost_per_token(model="gpt-3.5-turbo", prompt_tokens=estimated.prompt_tokens)
assert expected_input_cost > 0
assert request_data["response_cost"] == pytest.approx(expected_input_cost)
@pytest.mark.asyncio
async def test_client_disconnect_after_dispatch_is_priced_at_input_rate(self):
from datetime import datetime
from fastapi import HTTPException
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
from litellm.types.utils import Usage
request_data = {
"litellm_logging_obj": self._logging_obj(
{
"first_api_call_start_time": datetime.now(),
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "count these input tokens please"}],
"call_type": "acompletion",
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True,
}
),
"metadata": {},
}
await self._run(
request_data,
original_exception=HTTPException(
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
),
)
estimated = request_data["combined_usage_object"]
assert isinstance(estimated, Usage)
expected_input_cost, _ = litellm.cost_per_token(model="gpt-3.5-turbo", prompt_tokens=estimated.prompt_tokens)
assert expected_input_cost > 0
assert request_data["response_cost"] == pytest.approx(expected_input_cost)
@pytest.mark.asyncio
async def test_upstream_499_error_is_a_plain_failure_not_a_priced_disconnect(self):
"""A provider answering HTTP 499 surfaces as litellm.APIError with status_code 499. Only the
proxy's own disconnect HTTPException means the caller hung up mid-generation, so the upstream
error keeps the estimated usage at zero cost like any other post-dispatch provider error."""
from datetime import datetime
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
from litellm.types.utils import Usage
request_data = {
"litellm_logging_obj": self._logging_obj(
{
"first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 0),
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "count these input tokens please"}],
"call_type": "acompletion",
}
),
"metadata": {},
}
await self._run(
request_data,
original_exception=litellm.APIError(
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
message="client closed request",
llm_provider="openai",
model="gpt-3.5-turbo",
),
)
estimated = request_data["combined_usage_object"]
assert isinstance(estimated, Usage)
assert estimated.prompt_tokens > 0
assert request_data["response_cost"] == 0.0
@pytest.mark.asyncio
async def test_client_disconnect_after_the_call_already_completed_is_not_priced_again(self):
"""The LLM task finished and claimed its success log (which carries the real usage and cost)
while a during_call guardrail task was still pending, then the caller hung up. The 499 raised
for that disconnect must not price the prompt a second time."""
from datetime import datetime
from fastapi import HTTPException
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
from litellm.litellm_core_utils.litellm_logging import Logging
logging_obj = Logging(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "count these input tokens please"}],
stream=False,
call_type="acompletion",
start_time=datetime(2026, 1, 1, 12, 0, 0),
litellm_call_id="completed-before-disconnect",
function_id="fn",
)
logging_obj.model_call_details.update(
{
"first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 1),
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "count these input tokens please"}],
"call_type": "acompletion",
}
)
logging_obj.mark_async_success_log_enqueued()
request_data = {"litellm_logging_obj": logging_obj, "metadata": {}}
await self._run(
request_data,
original_exception=HTTPException(
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
),
)
assert request_data.get("response_cost", 0.0) == 0.0
@pytest.mark.asyncio
async def test_client_disconnect_with_a_claimed_but_deferred_success_log_is_still_priced(self):
"""With a post_call guardrail registered the success log is claimed but parked in the deferred slot,
and a 499 raised before the flush means it never runs. Nothing else bills that request, so the
failure row keeps the estimated input cost."""
from datetime import datetime
from fastapi import HTTPException
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
from litellm.litellm_core_utils.litellm_logging import Logging
logging_obj = Logging(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "count these input tokens please"}],
stream=False,
call_type="acompletion",
start_time=datetime(2026, 1, 1, 12, 0, 0),
litellm_call_id="deferred-success-never-flushed",
function_id="fn",
)
logging_obj.model_call_details.update(
{
"first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 1),
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "count these input tokens please"}],
"call_type": "acompletion",
}
)
logging_obj._defer_async_logging = True
logging_obj._enqueue_deferred_logging = lambda: None
request_data = {"litellm_logging_obj": logging_obj, "metadata": {}}
await self._run(
request_data,
original_exception=HTTPException(
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
),
)
assert request_data["response_cost"] > 0.0
@pytest.mark.asyncio
async def test_client_disconnect_before_provider_handoff_is_not_priced(self):
"""A caller that hung up while the request was still in pre-call hooks never reached the
provider, so there is no handoff instant on the logging object and the 499 row is not priced."""
from datetime import datetime
from unittest.mock import AsyncMock, patch
from fastapi import HTTPException
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy._types import UserAPIKeyAuth
logging_obj = Logging(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "count these input tokens please"}],
stream=False,
call_type="acompletion",
start_time=datetime.now(),
litellm_call_id="lit145-pre-handoff",
function_id="lit145",
)
assert logging_obj.model_call_details.get("first_api_call_start_time") is None
request_data = {"litellm_logging_obj": logging_obj, "model": "gpt-3.5-turbo", "metadata": {}}
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.alert_types = []
with (
patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()),
patch.object(proxy_logging_obj, "_dispatch_proxy_only_failure_handlers", new=AsyncMock()),
):
await proxy_logging_obj.post_call_failure_hook(
request_data=request_data,
original_exception=HTTPException(
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
),
user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"),
)
assert request_data.get("first_api_call_start_time") is None
assert "combined_usage_object" not in request_data
assert request_data.get("response_cost", 0.0) == 0.0
@pytest.mark.asyncio
async def test_connect_timeout_never_reached_the_provider_stays_zero_cost(self):
from datetime import datetime
from litellm.types.utils import Usage
request_data = {
"litellm_logging_obj": self._logging_obj(
{
"first_api_call_start_time": datetime.now(),
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "count these input tokens please"}],
"call_type": "acompletion",
}
),
"metadata": {},
}
await self._run(
request_data, original_exception=self._timeout("gpt-3.5-turbo", httpx.ConnectTimeout("no route"))
)
assert isinstance(request_data["combined_usage_object"], Usage)
assert request_data["response_cost"] == 0.0
@pytest.mark.asyncio
async def test_dispatched_non_timeout_failure_estimates_input_tokens_with_zero_cost(self):
from datetime import datetime
from litellm.types.utils import Usage
@ -733,8 +1035,6 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
estimated = request_data["combined_usage_object"]
assert isinstance(estimated, Usage)
assert estimated.prompt_tokens > 0
assert estimated.completion_tokens == 0
assert estimated.total_tokens == estimated.prompt_tokens
assert request_data["response_cost"] == 0.0
@pytest.mark.asyncio
@ -1901,7 +2201,355 @@ def test_usage_the_stream_already_recovered_beats_an_estimate():
assert lifted == (recovered, 0.25)
def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost():
def test_a_timed_out_request_on_an_unmapped_model_keeps_its_row_at_zero_cost():
from litellm.proxy.utils import _failure_usage_to_lift
lifted = _failure_usage_to_lift(
model_call_details={
"call_type": "acompletion",
"model": "my-unmapped-model",
"custom_llm_provider": "openai",
"messages": ONE_USER_MESSAGE,
},
request_body={},
dispatched=True,
provider_billed=True,
)
assert lifted is not None
usage, response_cost = lifted
assert usage.prompt_tokens > 0
assert response_cost == 0.0
def test_a_timed_out_request_is_priced_at_the_rate_the_router_registered_for_the_deployment(monkeypatch):
"""Router._create_deployment registers a deployment's own ``input_cost_per_token`` under its
``model_info.id``; the failure row prices with that entry through the same
``deployment_pricing_model_info`` the success path uses."""
from litellm.proxy.utils import _failure_usage_to_lift
monkeypatch.setitem(
litellm.model_cost,
"dep-lit145",
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 1.0},
)
lifted = _failure_usage_to_lift(
model_call_details={
"call_type": "acompletion",
"model": FAILURE_USAGE_MODEL,
"messages": ONE_USER_MESSAGE,
},
request_body={},
dispatched=True,
provider_billed=True,
deployment_id="dep-lit145",
)
assert lifted is not None
usage, response_cost = lifted
assert usage.prompt_tokens > 0
assert response_cost == pytest.approx(usage.prompt_tokens * 0.25)
@pytest.mark.parametrize(("prompt_tokens", "expected_rate"), [(1_000, 1e-7), (250_000, 5e-7)])
def test_a_timed_out_request_is_priced_at_the_deployments_long_context_tier(monkeypatch, prompt_tokens, expected_rate):
"""A deployment that declares ``input_cost_per_token_above_200k_tokens`` next to its standard rate
bills a 250k-token prompt at the tier on the success path; the failure row must not flatten it to
the standard rate."""
from litellm.proxy.utils import _estimated_input_cost
monkeypatch.setitem(
litellm.model_cost,
"dep-lit-f9-tiered",
{
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 1e-7,
"input_cost_per_token_above_200k_tokens": 5e-7,
"output_cost_per_token": 4e-7,
},
)
response_cost = _estimated_input_cost(
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
prompt_tokens=prompt_tokens,
deployment_id="dep-lit-f9-tiered",
base_model=None,
)
assert response_cost == pytest.approx(prompt_tokens * expected_rate)
def test_every_earlier_timed_out_attempt_is_billed_at_its_own_deployment_rate(monkeypatch):
"""The router leaves one RetryAttemptRecord per failed attempt in metadata.previous_models, the
final attempt included; the earlier Timeout ones are billed (each at the deployment that served
it), a 429 attempt is not, and the final attempt is left to the failure row itself."""
from litellm.proxy.utils import _billed_earlier_attempts
monkeypatch.setitem(
litellm.model_cost,
"dep-lit145-b",
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 1.0},
)
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
public_rate, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=24)
attempts = _billed_earlier_attempts(
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
litellm_params={
"metadata": {
"attempted_retries": 3,
"previous_models": (
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0},
{"deployment_id": "dep-lit145-b", "exception_type": "RateLimitError", "attempted_retries": 1},
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 2},
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 3},
),
}
},
estimated_usage=usage,
base_model=None,
deployment_id="dep-lit145-a",
)
assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [(0, "dep-lit145-a"), (2, "dep-lit145-b")]
assert attempts[0].response_cost == pytest.approx(public_rate)
assert attempts[0].response_cost > 0
assert attempts[1].response_cost == pytest.approx(24 * 0.25)
assert all(a.usage is usage for a in attempts)
def test_an_earlier_attempt_on_another_model_is_priced_at_that_models_rate_not_the_final_models(monkeypatch):
"""After a fallback hop the final attempt runs on another model group, so the records of the group
that timed out first name their own deployment model (and its ``base_model`` when the deployment
name is not in the cost map). Each is billed at that model's input rate, not at the final model's."""
from litellm.proxy.utils import _billed_earlier_attempts
for pricing_model, provider, input_rate in (
("lit145-final-model", "openai", 0.5),
("lit145-first-model", "openai", 0.02),
("azure/lit145-first-model", "azure", 0.02),
):
monkeypatch.setitem(
litellm.model_cost,
pricing_model,
{"litellm_provider": provider, "mode": "chat", "input_cost_per_token": input_rate, "output_cost_per_token": 1.0},
)
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
attempts = _billed_earlier_attempts(
model_call_details={"model": "openai/lit145-final-model", "custom_llm_provider": "openai"},
litellm_params={
"metadata": {
"attempted_retries": 1,
"model_info": {"id": "dep-lit145-final"},
"previous_models": (
{
"deployment_id": "dep-lit145-first",
"model": "openai/lit145-first-model",
"base_model": None,
"exception_type": "Timeout",
"attempted_retries": 0,
},
{
"deployment_id": "dep-lit145-first-azure",
"model": "azure/lit145-first-deployment",
"base_model": "azure/lit145-first-model",
"exception_type": "Timeout",
"attempted_retries": 1,
},
{
"deployment_id": "dep-lit145-final",
"model": "openai/lit145-final-model",
"base_model": None,
"exception_type": "Timeout",
"attempted_retries": 0,
},
{
"deployment_id": "dep-lit145-final",
"model": "openai/lit145-final-model",
"base_model": None,
"exception_type": "Timeout",
"attempted_retries": 1,
},
),
}
},
estimated_usage=usage,
base_model=None,
deployment_id="dep-lit145-final",
)
assert [(a.attempted_retries, a.deployment_id, a.response_cost) for a in attempts] == [
(0, "dep-lit145-first", pytest.approx(24 * 0.02)),
(1, "dep-lit145-first-azure", pytest.approx(24 * 0.02)),
(0, "dep-lit145-final", pytest.approx(24 * 0.5)),
]
def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_streams_recovered_usage():
"""A final attempt whose stream broke mid-flight lifts its recovered usage (prompt plus the chunks
delivered); the earlier timed-out attempts never produced output, so they are priced from the
input-only estimate and not from that recovered usage."""
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.proxy.utils import _estimated_prompt_usage, _failure_fields_to_lift
logging_obj = _dispatched_failure_logging_obj(
FAILURE_USAGE_MODEL,
get_litellm_params(
metadata={
"attempted_retries": 1,
"model_info": {"id": "dep-lit145-final"},
"previous_models": (
{"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0},
),
}
),
)
recovered = litellm.Usage(prompt_tokens=500, completion_tokens=40, total_tokens=540)
logging_obj.model_call_details["combined_usage_object"] = recovered
logging_obj.model_call_details["response_cost"] = 0.01
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
estimate = _estimated_prompt_usage(logging_obj.model_call_details, {})
assert lifted["combined_usage_object"] is recovered
assert [(a.usage.prompt_tokens, a.usage.completion_tokens) for a in lifted["billed_failed_attempts"]] == [
(estimate.prompt_tokens, 0)
]
assert estimate.prompt_tokens < recovered.prompt_tokens
def test_a_timed_out_attempt_from_the_group_a_fallback_left_is_still_billed():
"""After a fallback hop the router restarts ``attempted_retries`` at 0 for the new model group and
re-logs the attempt that exhausted the previous group, so the final attempt (dep-b, retry 1) must be
told from the previous group's attempt with the same count (dep-a, retry 1), and the re-logged
record must bill once."""
from litellm.proxy.utils import _billed_earlier_attempts
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
attempts = _billed_earlier_attempts(
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
litellm_params={
"metadata": {
"attempted_retries": 1,
"model_info": {"id": "dep-lit145-b"},
"previous_models": (
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0},
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1},
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1},
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 0},
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 1},
),
}
},
estimated_usage=usage,
base_model=None,
deployment_id="dep-lit145-b",
)
assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [
(0, "dep-lit145-a"),
(1, "dep-lit145-a"),
(0, "dep-lit145-b"),
]
assert all(a.response_cost > 0 for a in attempts)
def _dispatched_failure_logging_obj(model: str, litellm_params: dict) -> object:
"""A Logging object as the failure hook sees it: litellm_params in the real get_litellm_params()
shape (every key present, ``litellm_metadata`` None on chat completions) and a first-handoff
instant, so the lift treats the request as dispatched."""
from litellm.litellm_core_utils.litellm_logging import Logging
logging_obj = Logging(
model=model,
messages=ONE_USER_MESSAGE,
stream=False,
call_type="acompletion",
start_time=real_datetime.datetime.now(),
litellm_call_id="call-lit145",
function_id="fn-lit145",
)
logging_obj.update_environment_variables(model=model, user="", optional_params={}, litellm_params=litellm_params)
logging_obj.model_call_details["first_api_call_start_time"] = real_datetime.datetime.now()
logging_obj.model_call_details["messages"] = ONE_USER_MESSAGE
logging_obj.model_call_details["call_type"] = "acompletion"
return logging_obj
def test_the_deployment_rate_is_found_on_the_logging_objects_real_litellm_params(monkeypatch):
"""``get_litellm_params()`` always carries a ``litellm_metadata`` key (None on chat
completions) next to the router-stamped ``metadata.model_info.id``; the failure row must
still price at the deployment's own rate, like ``Logging.get_router_model_id`` on the
success path."""
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.proxy.utils import _failure_fields_to_lift
monkeypatch.setitem(
litellm.model_cost,
"dep-lit145-real-shape",
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 0.25},
)
litellm_params = get_litellm_params(metadata={"model_info": {"id": "dep-lit145-real-shape"}})
assert "litellm_metadata" in litellm_params and litellm_params["litellm_metadata"] is None
logging_obj = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, litellm_params)
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
assert lifted["response_cost"] == pytest.approx(lifted["combined_usage_object"].prompt_tokens * 0.25)
def test_an_azure_deployment_is_priced_at_its_base_model_like_the_success_path():
"""An Azure deployment name is not in the cost map; the success path prices it through
``model_info.base_model`` (``get_base_model_from_metadata``) and so must the failure row."""
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.proxy.utils import _failure_fields_to_lift
logging_obj = _dispatched_failure_logging_obj(
"azure/lit145-not-in-cost-map",
get_litellm_params(metadata={"model_info": {"base_model": "azure/gpt-4.1"}}),
)
logging_obj.model_call_details["custom_llm_provider"] = "azure"
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
expected_cost, _ = litellm.cost_per_token(
model="azure/gpt-4.1", prompt_tokens=lifted["combined_usage_object"].prompt_tokens
)
assert expected_cost > 0
assert lifted["response_cost"] == pytest.approx(expected_cost)
def test_a_litellm_params_key_in_the_request_body_cannot_zero_the_timeout_price():
"""``model_call_details`` is flattened with the request's optional_params, so a body key
named ``litellm_params`` can land there; the deployment rate must come from the logging
object's own ``litellm_params`` instead."""
from litellm.proxy.utils import _failure_usage_to_lift
lifted = _failure_usage_to_lift(
model_call_details={
"call_type": "acompletion",
"model": FAILURE_USAGE_MODEL,
"messages": ONE_USER_MESSAGE,
"litellm_params": {"input_cost_per_token": 0},
},
request_body={},
dispatched=True,
provider_billed=True,
deployment_id=None,
)
assert lifted is not None
usage, response_cost = lifted
expected_cost, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=usage.prompt_tokens)
assert response_cost == pytest.approx(expected_cost)
assert response_cost > 0
def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rate():
from litellm.proxy.utils import _failure_usage_to_lift
lifted = _failure_usage_to_lift(
@ -1912,13 +2560,16 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost():
},
request_body={},
dispatched=True,
provider_billed=True,
)
assert lifted is not None
usage, response_cost = lifted
assert usage.prompt_tokens > 0
assert usage.completion_tokens == 0
assert response_cost == 0.0
expected_input_cost, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=usage.prompt_tokens)
assert expected_input_cost > 0
assert response_cost == pytest.approx(expected_input_cost)
@pytest.mark.parametrize(

View file

@ -917,3 +917,53 @@ async def test_update_spend_logs_parks_failed_batch_in_redis_with_wire_safe_date
parked = await buffer.get_spend_logs_from_redis_buffer(limit=10)
assert mock_prisma_client.spend_log_transactions == []
assert [(row["request_id"], row["startTime"]) for row in parked] == [("a", started.isoformat())]
async def _create_many_while_postgres_is_restarting(**_: object) -> int:
from prisma.engine.utils import handle_response_errors
handle_response_errors(
MagicMock(),
[
{
"error": "Error in connector: Can't reach database server at `127.0.0.1:5432`",
"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server", "meta": None},
}
],
)
return 0
@pytest.mark.asyncio
async def test_update_spend_logs_retries_and_requeues_batch_when_postgres_restart_tears_down_prisma_engine(
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
"""While Postgres restarts the prisma engine answers the write with a P1001
payload whose ``meta`` is ``None``, which prisma's own response parser trips
over before it can build a typed error. That outage must be retried like
any other, and once retries exhaust the dequeued batch goes back to the head
of the queue so the rows land when Postgres is back.
"""
async def _fake_sleep(_: float) -> None:
return None
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(
side_effect=_create_many_while_postgres_is_restarting
)
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
with pytest.raises(AttributeError):
await ProxyUpdateSpend.update_spend_logs(
n_retry_times=2,
prisma_client=mock_prisma_client,
db_writer_client=None,
proxy_logging_obj=proxy_logging,
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
)
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 3

View file

@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable.
import json
import logging
from types import MappingProxyType
from typing import Final, Literal
import httpx
import pytest
@ -640,7 +641,7 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch):
monkeypatch.setattr(
bu,
"calculate_vertex_ai_batch_cost_and_usage",
lambda content, model, model_info=None: bu.BatchCostUsageResult(
lambda content, model, model_info=None, vertex_location=None: bu.BatchCostUsageResult(
cost=9.9,
usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
models=["gemini-2.0-flash-001"],
@ -671,7 +672,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch):
monkeypatch.setattr(
bu,
"calculate_vertex_ai_batch_cost_and_usage",
lambda content, model, model_info=None: pytest.fail("raw vertex path should not run"),
lambda content, model, model_info=None, vertex_location=None: pytest.fail("raw vertex path should not run"),
)
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai")
@ -1112,6 +1113,36 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk
assert result.failed_requests == 0
@pytest.mark.asyncio
async def test_handle_completed_vertex_batch_regional_location_costs_uplift_over_global(monkeypatch: pytest.MonkeyPatch) -> None:
import litellm.files.main as files_main
rows = [_vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5)]
async def fake_afile_content(**kw):
return type("R", (), {"content": _vertex_jsonl(rows)})()
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
global_result = await bu._handle_completed_batch(
_batch("gs://litellm-bucket/output/predictions.jsonl"),
custom_llm_provider="vertex_ai",
litellm_params={"vertex_project": "proj-1", "vertex_location": "global"},
vertex_location="global",
)
regional_result = await bu._handle_completed_batch(
_batch("gs://litellm-bucket/output/predictions.jsonl"),
custom_llm_provider="vertex_ai",
litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"},
vertex_location="us-central1",
)
uplift = litellm.model_cost["vertex_ai/gemini-3.6-flash"]["regional_endpoint_uplift_multiplier"]
assert global_result.cost > 0
assert regional_result.cost == pytest.approx(global_result.cost * uplift, rel=1e-9)
assert regional_result.usage == global_result.usage
@pytest.mark.asyncio
async def test_output_file_content_no_output_file_id_raises():
with pytest.raises(ValueError, match="Output file id is None"):
@ -1341,7 +1372,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch)
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
seen: dict = {}
def fake_vertex_calc(content, model, model_info=None):
def fake_vertex_calc(content, model, model_info=None, vertex_location=None):
seen["content"] = content
seen["model"] = model
return bu.BatchCostUsageResult(
@ -2417,3 +2448,30 @@ async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_
assert calls == []
assert (result.successful_requests, result.failed_requests) == (0, 1)
@pytest.mark.parametrize("provider", ("openai", "anthropic"))
@pytest.mark.parametrize("envelope", (1, "invalid", [1], None, {}))
@pytest.mark.parametrize("include_valid", (True, False))
def test_invalid_batch_envelopes_preserve_recoverable_cost(
provider: Literal["openai", "anthropic"], envelope: object, include_valid: bool
) -> None:
invalid: Final[dict[str, object]] = {"response" if provider == "openai" else "result": envelope}
valid: Final[dict[str, object]] = (
_success_row(model="batch-accounting-control", usage=_usage(100, 20))
if provider == "openai"
else {
"result": {
"type": "succeeded",
"message": {"model": "batch-accounting-control", "usage": {"input_tokens": 100, "output_tokens": 20}},
}
}
)
entries: Final = (valid, invalid, valid) if include_valid else (invalid,)
result: Final = bu._aggregate_batch_cost_usage_models(
entries, provider, model_info={"input_cost_per_token_batches": 0.001, "output_cost_per_token_batches": 0.002}
)
assert result.cost == pytest.approx(0.28 if include_valid else 0)
assert result.usage.prompt_tokens == (200 if include_valid else 0)
assert result.usage.completion_tokens == (40 if include_valid else 0)
assert (result.successful_requests, result.failed_requests) == (2 if include_valid else 0, 1)

View file

@ -0,0 +1,120 @@
import base64
import json
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Final
import pytest
import litellm
from litellm.decisions.main import _prepare_decisions_call
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.systemone.main import _prepare_systemone_call
QUESTION_TEXT: Final = "Is this about a billing problem?"
CONNECTION: Final = {
"model": "hosted_vllm/jev-fake",
"safety_identifier": None,
"api_key": "fake-provider-key",
"api_base": "http://127.0.0.1:4010",
"timeout": None,
"custom_llm_provider": None,
"extra_headers": None,
}
def _logging_obj(call_type: str) -> Logging:
return Logging(
model="hosted_vllm/jev-fake",
messages="placeholder",
stream=False,
call_type=call_type,
start_time=datetime.now(timezone.utc),
litellm_call_id="lit-9659",
function_id="lit-9659",
)
def _logged_request(logging_obj: Logging) -> Mapping[str, object]:
messages: Final = logging_obj.model_call_details["messages"]
assert len(messages) == 1 and messages[0]["role"] == "user"
return json.loads(messages[0]["content"])
def test_systemone_logs_every_request_field_the_caller_sent():
logging_obj: Final = _logging_obj("asystemone")
_prepare_systemone_call(
**CONNECTION,
state={"message": "My card was charged twice."},
questions={"billing": {"type": "noul", "instructions": QUESTION_TEXT}},
kwargs={"litellm_logging_obj": logging_obj},
)
logged: Final = _logged_request(logging_obj)
assert logged["state"] == {"message": "My card was charged twice."}
assert logged["questions"]["billing"]["instructions"] == QUESTION_TEXT
assert set(logged) == {"state", "questions"}
def test_decisions_logs_every_request_field_the_caller_sent():
logging_obj: Final = _logging_obj("adecisions")
_prepare_decisions_call(
**CONNECTION,
decision_input="My card was charged twice.",
questions=[{"type": "predicate", "instructions": QUESTION_TEXT}],
kwargs={"litellm_logging_obj": logging_obj},
)
logged: Final = _logged_request(logging_obj)
assert logged["input"] == "My card was charged twice."
assert logged["questions"][0]["instructions"] == QUESTION_TEXT
assert set(logged) == {"input", "questions"}
def test_request_is_logged_even_when_it_is_rejected_before_the_provider_call():
logging_obj: Final = _logging_obj("asystemone")
with pytest.raises(litellm.BadRequestError):
_prepare_systemone_call(
**CONNECTION,
state={"message": "My card was charged twice."},
questions={"billing": {"type": "noul", "instructions": QUESTION_TEXT}},
kwargs={"litellm_logging_obj": logging_obj, "input": "sent to the wrong function"},
)
logged: Final = _logged_request(logging_obj)
assert logged["questions"]["billing"]["instructions"] == QUESTION_TEXT
def test_decisions_request_is_logged_even_when_its_questions_fail_validation():
logging_obj: Final = _logging_obj("adecisions")
with pytest.raises(litellm.BadRequestError):
_prepare_decisions_call(
**CONNECTION,
decision_input="My card was charged twice.",
questions=[{"type": "not-a-question-type", "instructions": QUESTION_TEXT}],
kwargs={"litellm_logging_obj": logging_obj},
)
logged: Final = _logged_request(logging_obj)
assert logged["questions"][0]["instructions"] == QUESTION_TEXT
def test_non_utf8_bytes_in_the_request_are_logged_and_rejected_with_a_400():
logging_obj: Final = _logging_obj("adecisions")
raw_input: Final = b"\xff\xfe not utf-8"
with pytest.raises(litellm.BadRequestError):
_prepare_decisions_call(
**CONNECTION,
decision_input=raw_input,
questions=[{"type": "predicate", "instructions": QUESTION_TEXT}],
kwargs={"litellm_logging_obj": logging_obj},
)
logged: Final = _logged_request(logging_obj)
assert base64.urlsafe_b64decode(logged["input"]) == raw_input
assert logged["questions"][0]["instructions"] == QUESTION_TEXT

View file

@ -9,6 +9,8 @@ with deployment credentials, bypassing the managed files access-control hooks.
"""
import base64
import json
from typing import Final
import pytest
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
@ -367,13 +369,14 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti
# managed_files_obj.afile_content should NOT have been called
mock_managed_files_hook.afile_content.assert_not_called()
# Verify the DB update writes batch_processed, status, and file_object
mock_prisma.db.litellm_managedobjecttable.update.assert_called_once()
update_call_kwargs = (
mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs
claim_data: Final = next(
call.kwargs["data"]
for call in mock_prisma.db.litellm_managedobjecttable.update_many.call_args_list
if call.kwargs["where"].get("id") == mock_job.id
)
assert update_call_kwargs["data"]["batch_processed"] is True
assert update_call_kwargs["data"]["status"] == "complete"
assert (
"file_object" in update_call_kwargs["data"]
), "file_object must be written to DB so list_batches reads updated status"
assert claim_data["batch_processed"] is True
assert claim_data["status"] == "complete"
persisted_batch: Final = json.loads(claim_data["file_object"])
assert persisted_batch["status"] == "completed"
assert persisted_batch["usage"]["prompt_tokens"] == 10
assert persisted_batch["usage"]["completion_tokens"] == 5

View file

@ -99,6 +99,31 @@ def test_generic_cost_per_token_bills_cache_reads_at_input_rate_when_no_cache_re
assert completion_cost == pytest.approx(380 * 9.7e-7)
@pytest.mark.parametrize(
("model_info", "expected_prompt_cost"),
[
({"input_cost_per_token": 1e-7, "input_cost_per_character": 2.5e-8, "output_cost_per_token": 0.0}, 35 * 2.5e-8),
({"input_cost_per_token": 1.5e-7, "output_cost_per_token": 0.0}, 11 * 1.5e-7),
],
ids=["character_priced_model_bills_characters_instead_of_tokens", "token_priced_model_ignores_the_character_count"],
)
def test_generic_cost_per_token_bills_a_character_counted_prompt_under_one_unit(
model_info: ModelInfo, expected_prompt_cost: float
) -> None:
usage = Usage(
prompt_tokens=11,
completion_tokens=0,
total_tokens=11,
prompt_tokens_details=PromptTokensDetailsWrapper(character_count=35),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="embedding-fixture", usage=usage, custom_llm_provider="vertex_ai", model_info=model_info
)
assert (prompt_cost, completion_cost) == (pytest.approx(expected_prompt_cost), 0.0)
def test_apply_provider_cache_read_default_only_derives_a_rate_for_fireworks() -> None:
openai_info: ModelInfo = {"input_cost_per_token": 2e-6}
fireworks_info: ModelInfo = {"input_cost_per_token": 2e-6}
@ -450,6 +475,66 @@ def test_video_tokens_fallback_to_base_cost():
assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12)
def test_prompt_audio_and_video_tokens_without_modality_rates_bill_at_tiered_input_rate():
"""Prompt audio/video tokens with no modality rate are prompt tokens, so a >200k prompt bills them above-200k."""
from unittest.mock import patch
mock_model_info = {
"input_cost_per_token": 2e-6,
"input_cost_per_token_above_200k_tokens": 4e-6,
"output_cost_per_token": 1.2e-5,
"output_cost_per_token_above_200k_tokens": 1.8e-5,
}
usage = Usage(
prompt_tokens=816032,
completion_tokens=4559,
total_tokens=820591,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=9033, audio_tokens=14999, video_tokens=792000),
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=4559),
)
with patch(
"litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
return_value=mock_model_info,
):
prompt_cost, completion_cost = generic_cost_per_token(
model="test-model", usage=usage, custom_llm_provider="vertex_ai"
)
assert round(prompt_cost, 9) == round(816032 * 4e-6, 9)
assert round(completion_cost, 9) == round(4559 * 1.8e-5, 9)
def test_prompt_modality_rates_still_win_over_tiered_input_rate():
"""An explicit input_cost_per_audio_token / input_cost_per_video_token beats the fallback, even above 200k."""
from unittest.mock import patch
mock_model_info = {
"input_cost_per_token": 2e-6,
"input_cost_per_token_above_200k_tokens": 4e-6,
"input_cost_per_audio_token": 3e-6,
"input_cost_per_video_token": 1e-6,
"output_cost_per_token": 1.2e-5,
}
usage = Usage(
prompt_tokens=300000,
completion_tokens=10,
total_tokens=300010,
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100000, audio_tokens=100000, video_tokens=100000),
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10),
)
with patch(
"litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
return_value=mock_model_info,
):
prompt_cost, _ = generic_cost_per_token(model="test-model", usage=usage, custom_llm_provider="vertex_ai")
assert round(prompt_cost, 9) == round(100000 * 4e-6 + 100000 * 3e-6 + 100000 * 1e-6, 9)
def test_generic_cost_per_token_above_200k_tokens(_local_model_cost_map):
# gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing
model = "gemini-2.5-pro"
@ -3801,3 +3886,220 @@ def test_azure_gpt_5_6_alias_matches_sol_pricing(_local_model_cost_map, region_p
assert shared_cost_fields
for field in shared_cost_fields:
assert alias[field] == sol[field], field
# Per-token rates read 2026-10-07 from https://platform.claude.com/docs/en/about-claude/pricing (direct and
# azure_ai, which Microsoft bills at Anthropic's rates per
# https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models-billing) and from the
# AmazonBedrockFoundationModels price list at
# https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json (Bedrock)
@pytest.mark.parametrize(
("model", "custom_llm_provider", "prompt_tokens", "input_rate", "cache_read_rate", "output_rate"),
[
("claude-haiku-5-5", "anthropic", 100_000, 1e-07, 1e-08, 5e-07),
("claude-haiku-5-5", "anthropic", 100_001, 5e-07, 5e-08, 2.5e-06),
("azure_ai/claude-haiku-5-5", "azure_ai", 100_000, 1e-07, 1e-08, 5e-07),
("azure_ai/claude-haiku-5-5", "azure_ai", 100_001, 5e-07, 5e-08, 2.5e-06),
("global.anthropic.claude-haiku-5-5", "bedrock", 100_000, 1e-07, 1e-08, 5e-07),
("global.anthropic.claude-haiku-5-5", "bedrock", 100_001, 5e-07, 5e-08, 2.5e-06),
("us.anthropic.claude-haiku-5-5", "bedrock", 100_000, 1.1e-07, 1.1e-08, 5.5e-07),
("us.anthropic.claude-haiku-5-5", "bedrock", 100_001, 5.5e-07, 5.5e-08, 2.75e-06),
("us-gov.anthropic.claude-haiku-5-5", "bedrock", 100_000, 1.2e-07, 1.2e-08, 6e-07),
("us-gov.anthropic.claude-haiku-5-5", "bedrock", 100_001, 6e-07, 6e-08, 3e-06),
("bedrock_mantle/anthropic.claude-haiku-5-5", "bedrock_mantle", 100_001, 5.5e-07, 5.5e-08, 2.75e-06),
("bedrock/us-gov-west-1/anthropic.claude-haiku-5-5", "bedrock", 100_001, 6e-07, 6e-08, 3e-06),
("vertex_ai/claude-haiku-5-5", "vertex_ai", 100_000, 1e-07, 1e-08, 5e-07),
("vertex_ai/claude-haiku-5-5", "vertex_ai", 100_001, 5e-07, 5e-08, 2.5e-06),
],
)
def test_generic_cost_per_token_claude_haiku_5_5_prompt_length_tiers(
_local_model_cost_map: None,
model: str,
custom_llm_provider: str,
prompt_tokens: int,
input_rate: float,
cache_read_rate: float,
output_rate: float,
) -> None:
"""Claude Haiku 5.5 bills every token at 5x the base rates once the prompt is over 100,000 tokens."""
cached_tokens: Final = 10_000
completion_tokens: Final = 1_000
usage: Final = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens),
)
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider=custom_llm_provider,
)
assert prompt_cost == pytest.approx((prompt_tokens - cached_tokens) * input_rate + cached_tokens * cache_read_rate)
assert completion_cost == pytest.approx(completion_tokens * output_rate)
def test_vertex_regional_endpoint_uplift_scales_claude_haiku_5_5_over_100k_rates(
_local_model_cost_map: None,
) -> None:
"""Vertex regional endpoints bill 1.1x the global rate on all token types
(https://cloud.google.com/vertex-ai/generative-ai/pricing, 2026-10-07: regional
over-100K input is $0.55/MTok), so the uplift scales the over-100k rates too."""
cached_tokens: Final = 10_000
prompt_tokens: Final = 100_001
completion_tokens: Final = 1_000
usage: Final = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="vertex_ai/claude-haiku-5-5",
usage=usage,
custom_llm_provider="vertex_ai",
vertex_location="us-east5",
)
assert prompt_cost == pytest.approx((prompt_tokens - cached_tokens) * 5.5e-07 + cached_tokens * 5.5e-08)
assert completion_cost == pytest.approx(completion_tokens * 2.75e-06)
# Batch rates read 2026-10-07 from the Batch processing table at
# https://platform.claude.com/docs/en/about-claude/pricing: $0.05 / $0.25 per MTok input and $0.25 / $1.25 output,
# up to and over 100,000 prompt tokens
@pytest.mark.parametrize(
("prompt_tokens", "input_rate", "output_rate"),
[(100_000, 5e-08, 2.5e-07), (100_001, 2.5e-07, 1.25e-06)],
)
def test_batch_cost_calculator_claude_haiku_5_5_prompt_length_tiers(
_local_model_cost_map: None,
prompt_tokens: int,
input_rate: float,
output_rate: float,
) -> None:
from litellm.cost_calculator import batch_cost_calculator
completion_tokens: Final = 1_000
prompt_cost, completion_cost = batch_cost_calculator(
usage=Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
),
model="claude-haiku-5-5",
custom_llm_provider="anthropic",
)
assert prompt_cost == pytest.approx(prompt_tokens * input_rate)
assert completion_cost == pytest.approx(completion_tokens * output_rate)
@pytest.mark.parametrize(
("prompt_tokens", "expected"),
[
(100_000, (5e-08, 2.5e-07, 5e-09, 6.25e-08)),
(100_001, (2.5e-07, 1.25e-06, 2.5e-08, 3.125e-07)),
],
)
def test_get_batch_cost_rates_claude_haiku_5_5_prompt_length_tiers(
_local_model_cost_map: None,
prompt_tokens: int,
expected: tuple[float, float, float, float],
) -> None:
"""Cache write and cache read batch rates are 50% of the standard rates; Anthropic's batch table omits them."""
from litellm.litellm_core_utils.llm_cost_calc.utils import get_batch_cost_rates
rates: Final = get_batch_cost_rates(
litellm.get_model_info(model="claude-haiku-5-5", custom_llm_provider="anthropic"),
Usage(prompt_tokens=prompt_tokens, completion_tokens=1, total_tokens=prompt_tokens + 1),
"anthropic",
)
assert (rates.input, rates.output, rates.cache_read, rates.cache_creation) == expected
@pytest.mark.parametrize(
"field", ("output_cost_per_token", "cache_read_input_token_cost", "cache_creation_input_token_cost")
)
def test_flat_input_price_keeps_other_token_thresholds(field: str) -> None:
from typing import Final, cast
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost
info: Final = cast(
ModelInfo,
{
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
"cache_read_input_token_cost": 0.003,
"cache_creation_input_token_cost": 0.004,
f"{field}_above_300k_tokens": 0.005,
},
)
rates: Final = _get_token_base_cost(
model_info=info,
usage=Usage(prompt_tokens=300_001, completion_tokens=3, total_tokens=300_004),
)
assert rates[0] == 0.001
assert rates[1] == (0.005 if field == "output_cost_per_token" else 0.002)
assert rates[2] == (0.005 if field == "cache_creation_input_token_cost" else 0.004)
assert rates[4] == (0.005 if field == "cache_read_input_token_cost" else 0.003)
@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001, 512_000, 512_001))
@pytest.mark.parametrize("higher_rate", (0.0, 0.009))
def test_mixed_token_thresholds_preserve_each_crossed_rate(prompt_tokens: int, higher_rate: float) -> None:
from typing import Final, cast
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost
info: Final = cast(
ModelInfo,
{
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
"cache_read_input_token_cost": 0.003,
"cache_creation_input_token_cost": 0.004,
"input_cost_per_token_above_200k_tokens": 0.005,
"output_cost_per_token_above_200k_tokens": 0.006,
"cache_creation_input_token_cost_above_200k_tokens": 0.007,
"cache_read_input_token_cost_above_512k_tokens": higher_rate,
"output_cost_per_token_above_512k_tokens": higher_rate,
},
)
rates: Final = _get_token_base_cost(
model_info=info,
usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=3, total_tokens=prompt_tokens + 3),
)
assert rates[0] == (0.005 if prompt_tokens > 200_000 else 0.001)
assert rates[1] == (higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002)
assert rates[2] == (0.007 if prompt_tokens > 200_000 else 0.004)
assert rates[4] == (higher_rate if prompt_tokens > 512_000 else 0.003)
@pytest.mark.parametrize("rate", (0.0, 0.001))
def test_token_breakdown_uses_supplied_pricing_without_catalog_lookup(rate: float) -> None:
model_info: Final[ModelInfo] = {
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"cache_read_input_token_cost": rate,
"cache_creation_input_token_cost": rate,
}
usage: Final = Usage(
prompt_tokens=20, completion_tokens=3, total_tokens=23,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), cache_creation_input_tokens=5,
)
prompt_cost, output_cost = generic_cost_per_token(
model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info,
)
breakdown: Final = get_token_type_cost_breakdown(
model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info,
)
assert breakdown.cache_read_cost == pytest.approx(5 * rate)
assert breakdown.cache_creation_cost == pytest.approx(5 * rate)
assert breakdown.rates is not None
assert prompt_cost == pytest.approx(10 * breakdown.rates.input_cost_per_token + breakdown.cache_read_cost + breakdown.cache_creation_cost)
assert output_cost == pytest.approx(3 * breakdown.rates.output_cost_per_token)

View file

@ -363,15 +363,15 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model(
Regression for the provider-prefix fallback in _handle_web_search_cost. When the initial
get_model_info lookup fails for a "/"-containing model, the retry re-resolves model_info from
the prefix and must adopt that prefix's provider for routing. Otherwise an unrelated model
(here OpenRouter, which carries no web search pricing) is re-resolved but still routed through
(here Azure, which carries no web search pricing) is re-resolved but still routed through
the request's vertex_ai Gemini calculator, which charges its $0.035 per_prompt default for a
model that should cost nothing for web search.
"""
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
model = "openrouter/google/gemini-3.1-flash-lite"
model = "azure/gpt-4o"
model_info = litellm.get_model_info(model)
assert model_info["litellm_provider"] == "openrouter"
assert model_info["litellm_provider"] == "azure"
assert not model_info.get("search_context_cost_per_query")
usage = Usage(

View file

@ -17,6 +17,7 @@ import pytest
import litellm
import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod
import litellm.proxy.common_request_processing as common_request_processing_mod
from litellm.cost_calculator import _select_model_name_for_cost_calc
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
ResponseMetadata,
@ -651,3 +652,38 @@ def test_update_response_metadata_prices_per_second_deployment_from_its_stamped_
assert result._response_ms == pytest.approx(2000)
assert result._hidden_params["response_cost"] == pytest.approx((0.02 + 0.04) * 2)
def test_update_response_metadata_keeps_the_provider_model_after_the_proxy_rewrites_response_model():
start_time: Final = datetime.datetime(2026, 9, 21, 12, 0, 0)
logging_obj: Final = Logging(
model="azure/ptu-deployment",
messages=[{"role": "user", "content": "Hello"}],
stream=False,
call_type="completion",
start_time=start_time,
litellm_call_id="provider-model-response-metadata",
function_id="f",
)
logging_obj.update_environment_variables(
model="azure/ptu-deployment", litellm_params={}, optional_params={}, custom_llm_provider="azure"
)
result: Final = ModelResponse(
model="gpt-5.4-nano-2026-03-17",
usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18),
)
update_response_metadata(
result=result,
logging_obj=logging_obj,
model="azure/ptu-deployment",
kwargs={},
start_time=start_time,
end_time=start_time + datetime.timedelta(seconds=2),
)
result.model = "model-group-alias"
assert (
_select_model_name_for_cost_calc(model=None, completion_response=result, custom_llm_provider="azure")
== "azure/gpt-5.4-nano-2026-03-17"
)

View file

@ -970,19 +970,16 @@ def test_shipped_tool_search_rule_version_boundaries(shipped_cost_map, model, pr
def test_shipped_tool_search_rule_fills_mapped_claude_entries_without_flag(shipped_cost_map):
"""A mapped Claude 4.5+ entry with no supports_tool_search key gets it from the rule
on Anthropic direct, Vertex and Bedrock, a mapped pre-4.5 entry stays without one,
and Azure Foundry and reseller copies of the same model are not touched."""
on Anthropic direct, Azure Foundry, Vertex and Bedrock, a mapped pre-4.5 entry stays
without one, and reseller copies of the same model are not touched."""
for key, model, provider in (
("claude-opus-4-7", "claude-opus-4-7", "anthropic"),
("azure_ai/claude-opus-5", "claude-opus-5", "azure_ai"),
("vertex_ai/claude-opus-5", "claude-opus-5", "vertex_ai"),
):
assert "supports_tool_search" not in litellm.model_cost[key]
assert litellm.get_model_info(model, custom_llm_provider=provider)["supports_tool_search"] is True
assert "supports_tool_search" not in litellm.model_cost["azure_ai/claude-opus-5"]
azure_opus_5_info = litellm.get_model_info("claude-opus-5", custom_llm_provider="azure_ai")
assert azure_opus_5_info.get("supports_tool_search") is None
assert match_fill_missing_generalizations("claude-opus-5", "bedrock")["supports_tool_search"] is True
assert "supports_tool_search" not in match_fill_missing_generalizations("claude-opus-5", "azure_ai")
assert match_fill_missing_generalizations("claude-opus-5", "azure_ai")["supports_tool_search"] is True
assert match_fill_missing_generalizations("claude-opus-5", "perplexity") is None

View file

@ -212,7 +212,6 @@ def test_shipped_backup_marks_claude_4_6_plus_adaptive_not_4_0():
"anthropic.claude-opus-4-8",
"vertex_ai/claude-opus-4-6@default",
"us.anthropic.claude-sonnet-4-6",
"openrouter/anthropic/claude-opus-4.7",
"azure_ai/claude-opus-4-7",
]:
assert backup[adaptive]["supports_adaptive_thinking"] is True, adaptive

View file

@ -20,13 +20,14 @@ from openai._legacy_response import HttpxBinaryResponseContent
import litellm
from litellm._logging import session_id_var, trace_id_var
from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, SENTRY_DENYLIST, SENTRY_PII_DENYLIST
from litellm.cost_calculator import ocr_batch_cost
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.litellm_core_utils.litellm_logging import (
_extract_response_obj_and_hidden_params,
_get_status_fields,
get_standard_logging_object_payload,
set_callbacks,
)
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
@ -346,6 +347,20 @@ def test_get_masked_api_base(logging_obj):
assert type(masked_api_base) == str
@pytest.mark.parametrize("no_upstream_call", [False, True])
def test_pre_call_stamps_the_provider_handoff_only_when_a_provider_call_is_made(logging_obj, no_upstream_call):
"""The proxy fires pre_call for a failure it raised itself (auth, rate limit, client disconnect)
so callbacks see it; that call never reached a provider, so neither handoff instant may be set."""
logging_obj.update_environment_variables(litellm_params={}, optional_params={})
if no_upstream_call:
logging_obj.model_call_details[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True
logging_obj.pre_call(input="", api_key="")
assert ("api_call_start_time" in logging_obj.model_call_details) is not no_upstream_call
assert (logging_obj.model_call_details.get("first_api_call_start_time") is not None) is not no_upstream_call
def test_pre_call_tolerates_missing_api_base(logging_obj):
"""Presigned batch retrieves (Mistral, Bedrock) build their own URL and pass api_base=None
to pre_call; masking must not raise or the request's pre-call logging is silently lost."""
@ -1329,6 +1344,56 @@ class TestRetrieveBatchCostPassesModelIdentity:
assert captured["model_info"] is not None
assert captured["model_info"]["input_cost_per_token"] == 0.0
@pytest.mark.asyncio
async def test_forwards_vertex_location_of_the_deployment(self, monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.litellm_core_utils import litellm_logging as logging_module
captured: dict[str, object] = {}
from litellm.batches.batch_utils import BatchCostUsageResult
from litellm.types.utils import LiteLLMBatch, Usage
async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult:
captured.update(kwargs)
return BatchCostUsageResult(
cost=1.25,
usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800),
models=["gemini-3.5-flash"],
successful_requests=1,
failed_requests=0,
)
monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch)
obj = LitellmLogging(
model="vertex_ai/gemini-3.5-flash",
messages=[{"role": "user", "content": "Hey"}],
stream=False,
call_type="aretrieve_batch",
start_time=time.time(),
litellm_call_id="batch-call-2",
function_id="f",
)
obj.custom_llm_provider = "vertex_ai"
obj.litellm_params = {"vertex_project": "proj-1", "vertex_location": "us-central1"}
batch = LiteLLMBatch(
id="batch_vertex",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="file-in",
object="batch",
status="completed",
output_file_id="file-out",
)
with contextlib.suppress(Exception):
await obj._async_success_handler_body(result=batch, start_time=None, end_time=None)
assert captured, "handle_completed_batch was never called"
assert captured["vertex_location"] == "us-central1"
class TestRetrieveBatchPricesOnlyFinalBatches:
"""Regression (LIT-7048): retrieving a provider-id batch priced it on every poll.
@ -6794,10 +6859,10 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur
def _resolve(custom_llm_provider, litellm_params, optional_params, model):
from litellm.litellm_core_utils.litellm_logging import (
_resolve_vertex_location_for_cost,
resolve_vertex_location_for_cost,
)
return _resolve_vertex_location_for_cost(
return resolve_vertex_location_for_cost(
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
optional_params=optional_params,
@ -8336,6 +8401,7 @@ class TestAzurePTUSpilloverCost:
ROUTER_MODEL_ID: Final = "ptu-spill-router-model-id"
SERVED_MODEL: Final = "azure/spill-served-model-ptu"
PTU_DEPLOYMENT: Final = "azure/spill-ptu-deployment"
PTU_MODEL_INFO: Final = {
"id": ROUTER_MODEL_ID,
"team_id": "team-1",
@ -8371,7 +8437,9 @@ class TestAzurePTUSpilloverCost:
litellm.model_cost.pop(TestAzurePTUSpilloverCost.ROUTER_MODEL_ID, None)
litellm.model_cost.pop(TestAzurePTUSpilloverCost.SERVED_MODEL, None)
def _logging_obj(self, model_info: dict, *, flag: str, litellm_rate: float, monkeypatch) -> LitellmLogging:
def _logging_obj(
self, model_info: dict, *, flag: str, litellm_rate: float, monkeypatch, deployment: str | None = None
) -> LitellmLogging:
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", flag)
obj = LitellmLogging(
model=self.SERVED_MODEL,
@ -8388,7 +8456,7 @@ class TestAzurePTUSpilloverCost:
optional_params={},
litellm_params={
"api_base": "",
"metadata": {"model_info": model_info},
"metadata": {"model_info": model_info, "deployment": deployment},
"input_cost_per_token": litellm_rate,
"output_cost_per_token": litellm_rate,
},
@ -8458,6 +8526,84 @@ class TestAzurePTUSpilloverCost:
finally:
self._unregister_models()
def test_spillover_payload_reports_the_pay_as_you_go_entry_billing_used(self, monkeypatch) -> None:
self._register_models()
try:
obj = self._logging_obj(dict(self.PTU_MODEL_INFO), flag="True", litellm_rate=0.0, monkeypatch=monkeypatch)
response = self._response()
response._hidden_params["additional_headers"] = {"llm_provider-x-ms-is-spilled-over": "true"}
response_cost: Final = obj._response_cost_calculator(result=response)
payload: Final = get_standard_logging_object_payload(
kwargs={
"model": self.SERVED_MODEL,
"custom_llm_provider": "azure",
"litellm_params": obj.litellm_params,
"response_cost": response_cost,
},
init_response_obj=response,
start_time=datetime.datetime.now(),
end_time=datetime.datetime.now(),
logging_obj=obj,
status="success",
)
assert response_cost == pytest.approx(self.EXPECTED_SPILL_COST)
assert payload is not None
model_map_value: Final = payload["model_map_information"]["model_map_value"]
assert model_map_value is not None
assert model_map_value["input_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["input_cost_per_token"]
assert model_map_value["output_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["output_cost_per_token"]
assert response_cost == pytest.approx(
100 * model_map_value["input_cost_per_token"] + 50 * model_map_value["output_cost_per_token"]
)
finally:
self._unregister_models()
def test_spillover_without_base_model_reports_the_pay_as_you_go_rates_billing_used(self, monkeypatch) -> None:
self._register_models()
litellm.model_cost[self.PTU_DEPLOYMENT] = {"litellm_provider": "azure", "mode": "chat"}
try:
obj = self._logging_obj(
dict(self.PTU_MODEL_INFO),
flag="True",
litellm_rate=0.0,
monkeypatch=monkeypatch,
deployment=self.PTU_DEPLOYMENT,
)
response = self._response()
response._hidden_params["additional_headers"] = {"llm_provider-x-ms-is-spilled-over": "true"}
response._hidden_params["provider_response_model"] = response.model
response_cost: Final = obj._response_cost_calculator(result=response)
response.model = "spill-model-group-alias"
payload: Final = get_standard_logging_object_payload(
kwargs={
"model": self.SERVED_MODEL,
"custom_llm_provider": "azure",
"litellm_params": obj.litellm_params,
"response_cost": response_cost,
},
init_response_obj=response,
start_time=datetime.datetime.now(),
end_time=datetime.datetime.now(),
logging_obj=obj,
status="success",
)
assert response_cost == pytest.approx(self.EXPECTED_SPILL_COST)
assert payload is not None
model_map_information: Final = payload["model_map_information"]
assert model_map_information["model_map_key"] != self.PTU_DEPLOYMENT
model_map_value: Final = model_map_information["model_map_value"]
assert model_map_value is not None
assert response_cost == pytest.approx(
100 * model_map_value["input_cost_per_token"] + 50 * model_map_value["output_cost_per_token"]
)
finally:
litellm.model_cost.pop(self.PTU_DEPLOYMENT, None)
self._unregister_models()
def test_spillover_header_does_not_touch_non_ptu_custom_pricing(self, monkeypatch) -> None:
self._register_models()
custom_model_id: Final = "non-ptu-custom-router-model-id"
@ -8855,6 +9001,89 @@ async def test_prompt_management_with_unchanged_variables_replays_a_byte_identic
assert len(messages_n_plus_one) == len(messages_n) + 2
@pytest.mark.parametrize(
("input_cost_per_token", "output_cost_per_token"),
[(1e-06, 2e-06), (0.0, 0.0)],
ids=["custom_price", "explicit_zero_price"],
)
def test_standard_logging_payload_reports_the_deployment_entry_the_request_was_billed_from(
input_cost_per_token: float, output_cost_per_token: float
):
"""A deployment with its own price keeps the catalog model as model_map_key but is logged at its own rates."""
deployment_id: Final = "lit9660-priced-deployment"
deployment_model_info: Final = {
"id": deployment_id,
"input_cost_per_token": input_cost_per_token,
"output_cost_per_token": output_cost_per_token,
}
litellm.register_model(
model_cost={
deployment_id: {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": input_cost_per_token,
"output_cost_per_token": output_cost_per_token,
}
},
persist_across_reloads=False,
)
try:
logging_obj: Final = LitellmLogging(
model="gpt-5.4-nano",
messages=[{"role": "user", "content": "Hi"}],
stream=False,
call_type="completion",
start_time=time.time(),
litellm_call_id="lit9660",
function_id="fn",
)
litellm_params: Final = {
"metadata": {
"model_group": "lit9660-group",
"deployment": "openai/gpt-5.4-nano",
"model_info": deployment_model_info,
}
}
logging_obj.update_environment_variables(
model="gpt-5.4-nano", user="", optional_params={}, litellm_params=litellm_params, custom_llm_provider="openai"
)
response: Final = ModelResponse(
model="gpt-5.4-nano",
choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hi"))],
usage=litellm.Usage(prompt_tokens=10, completion_tokens=6, total_tokens=16),
)
response_cost: Final = logging_obj._response_cost_calculator(result=response)
payload: Final = get_standard_logging_object_payload(
kwargs={
"model": "gpt-5.4-nano",
"custom_llm_provider": "openai",
"litellm_params": litellm_params,
"response_cost": response_cost,
},
init_response_obj=response,
start_time=datetime.datetime.now(),
end_time=datetime.datetime.now(),
logging_obj=logging_obj,
status="success",
)
assert response_cost == 10 * input_cost_per_token + 6 * output_cost_per_token
assert payload is not None
assert payload["model_map_information"] == {
"model_map_key": "openai/gpt-5.4-nano",
"model_map_value": litellm.get_model_info(model=f"openai/{deployment_id}", custom_llm_provider="openai"),
}
model_map_value: Final = payload["model_map_information"]["model_map_value"]
assert model_map_value is not None
assert (model_map_value["input_cost_per_token"], model_map_value["output_cost_per_token"]) == (
input_cost_per_token,
output_cost_per_token,
)
finally:
litellm.model_cost.pop(deployment_id, None)
@pytest.mark.asyncio
async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch):
"""

View file

@ -29,7 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transfor
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import (
VertexAIAnthropicConfig,
)
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES
from litellm.types.utils import ServerToolUse, Usage
@ -6798,3 +6798,91 @@ def test_chat_dummy_tool_result_for_an_orphaned_tool_call_replays_a_byte_identic
_assert_prefix_stable(requests)
assert [m["role"] for m in requests[0]["messages"]] == ["user", "assistant", "user"]
assert requests[0]["messages"][2]["content"][0]["type"] == "tool_result"
def _make_transform_request(optional_params: dict, litellm_params: dict) -> dict:
return AnthropicConfig().transform_request(
model="claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "hi"}],
optional_params=optional_params,
litellm_params=litellm_params,
headers={},
)
def _two_function_tools() -> list[dict]:
return [
{
"type": "function",
"function": {"name": f"tool_{i}", "description": f"Tool {i}", "parameters": {"type": "object"}},
}
for i in range(2)
]
@pytest.mark.parametrize(
"control, expected_control",
[
(None, {"type": "ephemeral"}),
({"type": "ephemeral", "ttl": "1h"}, {"type": "ephemeral", "ttl": "1h"}),
],
)
def test_transform_request_tool_config_point_marks_last_tool_and_drops_the_key(control, expected_control):
"""LIT-9880: a tool_config injection point used to be spread into the Anthropic body, which
Anthropic rejected with 400 Extra inputs are not permitted."""
tools = _two_function_tools()
tools_before = copy.deepcopy(tools)
point = {"location": "tool_config", **({"control": control} if control is not None else {})}
litellm_params = {"metadata": {}, "model_info": {"id": "deployment-1"}}
data = _make_transform_request(
optional_params={"tools": tools, "cache_control_injection_points": [point]},
litellm_params=litellm_params,
)
assert "cache_control_injection_points" not in data
assert data["tools"][-1]["cache_control"] == expected_control
assert "cache_control" not in data["tools"][0]
assert tools == tools_before
assert litellm_params["metadata"]["litellm_gateway_injected_cache"] == "deployment-1"
def test_transform_request_tool_config_point_keeps_the_clients_tool_marker():
tools = _two_function_tools()
tools[0]["cache_control"] = {"type": "ephemeral"}
data = _make_transform_request(
optional_params={"tools": tools, "cache_control_injection_points": [{"location": "tool_config"}]},
litellm_params={},
)
assert "cache_control_injection_points" not in data
assert data["tools"][0]["cache_control"] == {"type": "ephemeral"}
assert "cache_control" not in data["tools"][1]
def test_transform_request_tool_config_point_skips_tool_search_tools():
tool_search_type = next(iter(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES))
tools = [*_two_function_tools(), {"type": tool_search_type, "name": "tool_search"}]
data = _make_transform_request(
optional_params={"tools": tools, "cache_control_injection_points": [{"location": "tool_config"}]},
litellm_params={},
)
assert data["tools"][1]["cache_control"] == {"type": "ephemeral"}
assert "cache_control" not in data["tools"][2]
assert "cache_control" not in data["tools"][0]
def test_transform_request_tool_config_point_without_tools_drops_the_key():
litellm_params = {"metadata": {}}
data = _make_transform_request(
optional_params={"cache_control_injection_points": [{"location": "tool_config"}]},
litellm_params=litellm_params,
)
assert "cache_control_injection_points" not in data
assert "tools" not in data
assert litellm_params["metadata"] == {}

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