scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

Merge pull request #47736 from BerriAI/litellm_backport_p0_batch_translation_cost_stable_1_106_x

Browse source 
fix(batches)!: backport P0 batch-cost, pricing and translation fixes to stable/1.106.x for v1.106.2
This commit is contained in:
yuneng-jiang 2026-10-21 00:59:16 -07:00 • committed by GitHub
commit 500e19d09c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
69 changed files with 19672 additions and 23317 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
ci_cd/generate_model_prices_schema.py
View file

@ -62,6 +62,17 @@ OFF_PEAK_WINDOW: Final[JsonSchema] = {
}
OBJECT_KEYS: dict[str, JsonSchema] = {
"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"},
},
},
"off_peak_pricing": {
"type": "object",
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
@ -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.",
},
}

214
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
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
@ -31,7 +33,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",
@ -57,6 +64,15 @@ class _ManagedObjectRow(Protocol):
@property
def created_by(self) -> str | None: ...
@property
def created_at(self) -> datetime: ...
@property
def updated_at(self) -> datetime: ...
@property
def status(self) -> str | None: ...
@property
def file_object(self) -> object: ...
@ -122,6 +138,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:
@ -297,29 +318,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": {
@ -333,10 +334,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
@ -350,8 +400,8 @@ class CheckBatchCost:
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:
@ -359,12 +409,14 @@ class CheckBatchCost:
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
@ -378,10 +430,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(
@ -390,7 +448,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
@ -402,7 +462,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(
@ -436,13 +500,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
@ -457,7 +521,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"
@ -484,10 +548,12 @@ class CheckBatchCost:
"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"
)
@ -699,8 +765,6 @@ class CheckBatchCost:
@staticmethod
def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
import json
from litellm.types.utils import LiteLLMBatch
file_object: Final = job.file_object
@ -890,8 +954,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",
@ -922,27 +987,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:
@ -990,7 +1063,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,
@ -1003,8 +1076,6 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
self.batch_processed_support_confirmed = True
except Exception as query_err:
@ -1074,7 +1145,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
@ -1084,22 +1155,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 (

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

@ -24,7 +24,7 @@ from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import ValidationError
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, Unpack
import litellm
@ -347,6 +347,31 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
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 {}

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.76"
version = "0.1.76.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.76"
version = "0.1.76.post1"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

25
litellm-rust/crates/llms/src/anthropic/messages/transformation.rs
View file

@ -110,7 +110,10 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
}
fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
update_headers_with_anthropic_beta(headers, request)
merge_beta_headers(
update_headers_with_anthropic_beta(headers, request),
safeguards_betas(&request.params),
)
}
}
@ -198,6 +201,20 @@ fn uses_structured_output(params: &MessagesOptionalParams) -> bool {
.is_some_and(|config| config.format.is_some())
}
fn safeguards_betas(params: &MessagesOptionalParams) -> BetaSet {
carries_safeguards(params)
.then_some(AnthropicBeta::DangerousToolUse20260903)
.into_iter()
.collect()
}
fn carries_safeguards(params: &MessagesOptionalParams) -> bool {
params
.extra
.get("safeguards")
.is_some_and(|value| !value.is_null())
}
fn messages_carry_output_config(messages: &[Message]) -> bool {
messages
.iter()
@ -1052,6 +1069,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,

2
litellm/anthropic_beta_headers_config.json
View file

@ -229,7 +229,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",

45
litellm/anthropic_beta_headers_manager.py
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._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,

16
litellm/batches/batch_utils.py
View file

@ -777,11 +777,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[str, object]] = cast( # cast-ok: batch response bodies come from provider JSON
dict[str, object], batch_job_output_file.get("response", None) or {}
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
)
return _response.get("status_code", None) == 200

103
litellm/cost_calculator.py
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
@ -127,6 +127,7 @@ from litellm.types.utils import (
TranscriptionUsageTokensObject,
Usage,
VectorStoreSearchResponse,
is_server_derived_pricing_key,
)
from litellm.utils import (
CallTypes,
@ -139,6 +140,8 @@ from litellm.utils import (
TextCompletionResponse,
TranscriptionResponse,
cached_get_model_info_helper,
model_info_with_pricing_overrides,
runtime_model_cost_registrations,
token_counter,
)
@ -417,6 +420,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.
@ -704,6 +708,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(
@ -819,6 +832,60 @@ def _cost_map_entry_prices_anything(entry: Mapping[str, object]) -> bool:
)
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)
_LOGGED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
def _invoked_model_id_for_cost(
optional_params: object,
logging_obj: LitellmLoggingObject | None,
) -> str | None:
from litellm.llms.bedrock.common_utils import extract_model_name_from_bedrock_arn
recorded: Final[object] = getattr(logging_obj, "model_call_details", None)
invoked: Final = _get_hidden_str_for_cost_calc(optional_params, "model_id") or _get_hidden_str_for_cost_calc(
recorded, "model_id"
)
return (
extract_model_name_from_bedrock_arn(invoked) if invoked is not None and invoked.startswith("arn:") else invoked
)
def select_model_name_for_cost_calc(
model: str | None,
completion_response: object | None,
@ -827,6 +894,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
@ -847,6 +915,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
@ -864,8 +939,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:
@ -1472,6 +1547,10 @@ def completion_cost(
)
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,
@ -1480,8 +1559,15 @@ 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 = [
selected_model,
_get_response_model(completion_response),
@ -1844,6 +1930,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)
@ -1932,6 +2019,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
@ -2176,7 +2264,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),

15
litellm/litellm_core_utils/litellm_logging.py
View file

@ -6023,6 +6023,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,
@ -6030,12 +6031,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,
)
@ -6714,6 +6726,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(

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

@ -285,7 +285,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:
@ -681,10 +681,12 @@ def _get_token_base_cost(
# Optimization: collect threshold keys first to avoid sorting all model_info keys.
# Standard thresholds and thresholds suffixed for this request's service tier both count.
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_")
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))
)
@ -692,12 +694,12 @@ def _get_token_base_cost(
# 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,
@ -710,7 +712,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,
@ -766,8 +768,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:
@ -1654,20 +1654,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,
@ -1686,6 +1689,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
@ -1710,6 +1714,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)

11
litellm/llms/anthropic/pass_through/messages/transformation.py
View file

@ -4,6 +4,7 @@ from typing import Any, ClassVar, Final
import httpx
from litellm._logging import verbose_logger
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.llms.base_llm.anthropic_messages.transformation import (
@ -711,6 +712,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.
@ -727,8 +729,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()
@ -778,6 +781,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)

60
litellm/llms/azure/chat/gpt_5_transformation.py
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,

8
litellm/llms/azure/responses/transformation.py
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)

13
litellm/llms/azure_ai/responses/transformation.py
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(

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

@ -4,7 +4,7 @@ import json
import ssl
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, 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
from typing import (
TYPE_CHECKING,
@ -34,7 +34,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
@ -48,6 +48,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 (
@ -2223,6 +2224,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,
@ -2230,7 +2241,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,
@ -2261,20 +2272,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,
)
@ -2381,7 +2383,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,
@ -2389,7 +2391,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,
@ -2412,7 +2414,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,

29
litellm/llms/databricks/chat/transformation.py
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(

102
litellm/llms/openai/chat/gpt_5_transformation.py
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 (
@ -69,6 +69,10 @@ _get_effort_level = get_effort_level
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 "gpt-5-chat" not in normalized
@ -84,6 +88,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)
@ -203,15 +210,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),
@ -219,6 +221,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 [
@ -279,12 +314,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"
@ -296,31 +331,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

39
litellm/llms/openai/responses/transformation.py
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.responses.litellm_completion_transformation.reasoning_items import is_litellm_minted_reasoning_item
from litellm.secret_managers.main import get_secret_str
@ -141,6 +144,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
@ -247,7 +264,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)

2
litellm/llms/openai_like/messages/transformation.py
View file

@ -8,6 +8,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import (
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

28
litellm/main.py
View file

@ -1086,6 +1086,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,
@ -1135,7 +1147,9 @@ def responses_api_bridge_check(
# server-side), and Chat Completions rejects function tools whenever reasoning is
# on ("Function tools with reasoning_effort are not supported ... use
# /v1/responses or set reasoning_effort to 'none'"), so only an explicit
# ``"none"`` keeps the request chat-servable. Custom (grammar) tools are served
# ``"none"`` the chat gate forwards keeps the request chat-servable: a ``none``
# the model's map row refuses is dropped (or refused) before the call, which
# leaves the provider's default effort in force. Custom (grammar) tools are served
# natively by Chat Completions with reasoning on, so custom-only requests stay on
# chat and keep their native custom tool_call response shape.
# - The UNSET-effort arm only fires against endpoints known to enforce that
@ -1158,10 +1172,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
@ -5769,6 +5786,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

17299
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

15
litellm/proxy/hooks/model_max_budget_limiter.py
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._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
@ -17,7 +15,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
@ -122,17 +120,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):

56
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -1,9 +1,13 @@
import asyncio
import traceback
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Callable, Generator, Mapping, Sequence
from contextlib import contextmanager
from contextvars import ContextVar
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol, cast
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import batch_cost_is_final
@ -63,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,
@ -303,6 +338,11 @@ class _ProxyDBLogger(CustomLogger):
update_cache_read_keys,
)
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(
@ -323,6 +363,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
@ -473,6 +515,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 = (
@ -797,7 +841,15 @@ async def _update_database_and_spend_counters_in_batch(
from litellm.proxy.proxy_server import arm_update_cache_read
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,

84
litellm/router.py
View file

@ -342,12 +342,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,
@ -356,6 +358,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,
)
@ -497,6 +501,7 @@ 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])
_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
_RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None)
_ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
@ -9173,6 +9178,20 @@ class Router:
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
@ -9330,39 +9349,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_service_tier_pricing(
model_info=_model_info,
backend_model=Router._cost_map_backend_model(deployment),
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:
@ -10379,7 +10366,7 @@ class Router:
return (backend_key,)
@staticmethod
def _deployment_model_cost_payload(deployment: Deployment) -> dict: # mutable-ok: cost-map entry
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
@ -10387,30 +10374,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=Router._cost_map_backend_model(deployment),
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

28
litellm/router_strategy/budget_limiter.py
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._internal_context import with_service_target
from litellm._logging import verbose_router_logger
@ -38,6 +40,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,
)
@ -47,6 +50,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:
@ -502,10 +507,14 @@ 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 | None] = standard_logging_payload.get("custom_llm_provider")
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
)
@ -518,6 +527,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)
@ -530,6 +540,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(
@ -547,6 +558,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(
@ -555,12 +567,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,
@ -598,6 +613,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
@with_service_target("router_budgets")
async def periodic_sync_in_memory_spend_with_redis(self):
"""

13
litellm/router_utils/batch_utils.py
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
@ -192,3 +194,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

25
litellm/router_utils/reasoning_effort_capability.py
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

145
litellm/utils.py
View file

@ -45,7 +45,7 @@ from httpx import Proxy
from httpx._utils import get_environment_proxies
from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # OpenAI parser module is private
from openai.types.chat.completion_create_params import ResponseFormat
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, with_config
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config
import litellm
import litellm.litellm_core_utils
@ -3290,6 +3290,135 @@ def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: st
_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload
class _RegisteredPricingFields:
"""Explicitly registered rates per cost-map key, tied to the ``litellm.model_cost`` object they were
registered on so a replaced catalog starts without them, as the entries they describe do."""
cost_map: object = None
fields: dict[str, dict[str, object]] = {} # mutable-ok: rebuilt whenever litellm.model_cost is replaced
def _registered_pricing_fields() -> dict[str, dict[str, object]]:
if _RegisteredPricingFields.cost_map is not litellm.model_cost:
_RegisteredPricingFields.cost_map = litellm.model_cost
_RegisteredPricingFields.fields = {}
return _RegisteredPricingFields.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
_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"),
}
)
_RATE_TABLE: Final = TypeAdapter(Mapping[str, object])
_RATE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
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]:
"""A caller's flat ``input_cost_per_token`` / ``output_cost_per_token`` is that direction's price above every
token threshold, for every service tier and batch, in every ``tiered_pricing`` tier and off-peak window, so each
inherited rate that would stand in for it there, the reasoning-token rate included for output, is dropped, and
each tier carries the flat price instead."""
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_fields: Final = _registered_pricing_fields()
registered: Final = ((key, registered_fields.get(key)) for key in dict.fromkeys(candidates))
return tuple(
(key, MappingProxyType({name: value for name, value in entry.items() if value is not None}))
for key, entry in registered
if entry
)
class _LiveDeploymentReplay:
@ -3399,7 +3528,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)
_runtime_registered_model_cost[_registered_key] = update_dictionary(
dict(_runtime_registered_model_cost.get(_registered_key, {})),
_RATE_TABLE.validate_python(_registered_value),
)
_skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
@ -3467,6 +3599,15 @@ 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_pricing_fields().update(
{
name: {
**_registered_pricing_fields().get(name, {}),
**{field: field_value for field, field_value in value.items() if field_value is not None},
}
for name in dict.fromkeys((str(model_cost_key), _key_str))
}
)
# Invalidate case-insensitive lookup map since model_cost was modified
_invalidate_model_cost_lowercase_map()

17299
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

158
model_prices_and_context_window.schema.json
View file

@ -93,11 +93,6 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_128k_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": {
"type": "number",
"minimum": 0,
@ -118,16 +113,6 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_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_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,
@ -138,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,
@ -153,15 +143,14 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_above_32k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_creation_input_token_cost_batches": {
"type": "number",
"minimum": 0
},
"cache_creation_input_token_cost_decisions": {
"type": "number",
"minimum": 0
},
"cache_creation_input_token_cost_flex": {
"type": "number",
"minimum": 0,
@ -199,11 +188,6 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_128k_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_200k_tokens": {
"type": "number",
"minimum": 0,
@ -219,11 +203,6 @@
"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,
@ -234,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,
@ -249,11 +233,6 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"cache_read_input_token_cost_above_32k_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_512k_tokens": {
"type": "number",
"minimum": 0,
@ -267,6 +246,10 @@
"type": "number",
"minimum": 0
},
"cache_read_input_token_cost_decisions": {
"type": "number",
"minimum": 0
},
"cache_read_input_token_cost_flex": {
"type": "number",
"minimum": 0,
@ -314,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"
},
@ -397,6 +384,10 @@
"type": "number",
"minimum": 0
},
"input_cost_per_text_message": {
"type": "number",
"minimum": 0
},
"input_cost_per_token": {
"type": "number",
"minimum": 0,
@ -432,11 +423,6 @@
"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,
@ -447,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,
@ -462,11 +453,6 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_32k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"input_cost_per_token_above_512k_tokens": {
"type": "number",
"minimum": 0,
@ -485,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,
@ -563,6 +553,7 @@
"completion",
"embedding",
"evaluation",
"extraction",
"guardrail",
"image_edit",
"image_generation",
@ -821,11 +812,6 @@
"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,
@ -836,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,
@ -851,11 +842,6 @@
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_32k_tokens": {
"type": "number",
"minimum": 0,
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
},
"output_cost_per_token_above_512k_tokens": {
"type": "number",
"minimum": 0,
@ -870,6 +856,10 @@
"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,
@ -965,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."
@ -1039,6 +1056,12 @@
"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"
},
@ -1224,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.",

8
pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.106.1"
version = "1.106.2"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -78,7 +78,7 @@ proxy = [
"httpx2>=2.5.0,<3",
"pydantic>=2.12.0,<3",
"litellm-proxy-extras==0.4.108",
"litellm-enterprise==0.1.76",
"litellm-enterprise==0.1.76.post1",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",
@ -290,7 +290,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 = [
@ -373,7 +373,7 @@ litellm-enterprise = { workspace = true }
members = ["enterprise", "litellm-proxy-extras"]
[tool.commitizen]
version = "1.106.1"
version = "1.106.2"
version_files = [
"pyproject.toml:^version",
]

35
tests/integration/_support/prisma_clients.py Normal file
View file

@ -0,0 +1,35 @@
import os
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Final
from prisma import Prisma
from prisma.types import DatasourceOverride
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import PrismaClient, ProxyLogging
@asynccontextmanager
async def connected_prisma(database_url: str | None = None) -> AsyncIterator[Prisma]:
client: Final = Prisma(
datasource=DatasourceOverride(url=database_url)
) if database_url is not None else Prisma()
await client.connect()
try:
yield client
finally:
await client.disconnect()
@asynccontextmanager
async def connected_proxy_prisma(database_url: str | None = None) -> AsyncIterator[PrismaClient]:
client: Final = PrismaClient(
database_url=database_url or os.environ["DATABASE_URL"],
proxy_logging_obj=ProxyLogging(UserApiKeyCache()),
)
await client.connect()
try:
yield client
finally:
await client.disconnect()

349
tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py Normal file
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.messages_endpoint.providers.anthropic.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

809
tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py Normal file
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)

1324
tests/integration/providers/test_databricks_chat_wire.py
View file

File diff suppressed because it is too large Load diff

186
tests/integration/sdk/test_databricks_json_schema_refs_sdk.py Normal file
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

230
tests/integration/spend/test_batch_poll_starvation.py
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,))

28
tests/unit/batches/test_batch_utils.py
View file

@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable.
import asyncio, json, time
import logging
from types import MappingProxyType
from typing import Final, Literal
import httpx
import pytest
@ -3333,3 +3334,30 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
assert mock_batch._hidden_params["batch_successful_requests"] == 8
assert mock_batch._hidden_params["batch_failed_requests"] == 0
assert mock_batch.usage == expected_usage
@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)

21
tests/unit/enterprise/proxy/test_managed_files_access_check.py
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

82
tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -2432,6 +2432,63 @@ def test_get_cost_per_unit_falls_back_from_service_tier_key_to_base():
assert get_cost_per_unit(model_info_direct, "input_cost_per_token_priority") == 5e-6
@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)
def test_threshold_keys_exclude_service_tier_variants():
from typing import cast
@ -3965,3 +4022,28 @@ def test_get_batch_cost_rates_claude_haiku_5_5_prompt_length_tiers(
)
assert (rates.input, rates.output, rates.cache_read, rates.cache_creation) == expected
@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)

84
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -38,6 +38,7 @@ from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
_extract_response_obj_and_hidden_params,
_get_status_fields,
get_standard_logging_object_payload,
set_callbacks,
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@ -11563,3 +11564,86 @@ class CompletionCustomHandler(
except Exception:
print(f"Assertion Error: {traceback.format_exc()}")
self.errors.append(traceback.format_exc())
@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)

26
tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py
View file

@ -205,6 +205,23 @@ class TestTheNormalizedTierIsTheTierSent:
):
assert _reasoning_effort_sent(model, provider, effort) == expected
@pytest.mark.parametrize(
"model, provider",
[
("gpt-5.5", "azure_ai"),
("azure_ai/gpt-5.5", "azure_ai"),
("gpt-5.6", "azure"),
("us/gpt-5.6", "azure"),
],
)
def test_an_azure_hosted_deployment_without_the_flag_drops_to_the_tier_it_accepts(
self, local_model_cost_map, model, provider
):
"""Chat Completions answers ``max`` on gpt-5.6 with a 400 naming ``none`` through ``xhigh``
(OpenAI, 2026-10-01) and none of these Azure rows declares the level, so the adapter sends
the tier below instead of the rejected one."""
assert _reasoning_effort_sent(model, provider, "max") == "xhigh"
@pytest.mark.parametrize(
"model, provider",
[
@ -212,13 +229,16 @@ class TestTheNormalizedTierIsTheTierSent:
("azure_ai/gpt-6-astra", "azure_ai"),
("gpt-6-astra", "azure"),
("us/gpt-6-astra", "azure"),
("eu/gpt-6-astra", "azure"),
],
)
def test_an_azure_hosted_astra_deployment_drops_to_the_tier_it_accepts(
def test_an_azure_hosted_astra_deployment_drops_to_the_tier_its_chat_route_accepts(
self, local_model_cost_map, model, provider
):
"""The deployment answers ``max`` with a 400 naming ``none`` through ``xhigh``, so the rows
say so and the adapter sends the tier below instead of the rejected one."""
"""A Foundry gpt-6-astra answers a Chat Completions ``max`` with a 400 naming low through
xhigh while honoring the same level on the Responses route (2026-10-01). The row's flag is
route-blind, so it keeps max off and the adapter, which cannot know which route the call
takes, sends the tier below instead of the level chat rejects."""
assert _reasoning_effort_sent(model, provider, "max") == "xhigh"
def test_the_openai_hosted_twin_still_sends_max(self, local_model_cost_map):

145
tests/unit/llms/azure/chat/test_azure_gpt5_transformation.py
View file

@ -1,3 +1,6 @@
from collections.abc import Mapping
from typing import Final
import pytest
import litellm
@ -343,17 +346,22 @@ def test_azure_gpt_6_astra_takes_the_reasoning_series_request_shape():
model="gpt-6-astra",
custom_llm_provider="azure",
max_tokens=100,
reasoning_effort="max",
reasoning_effort="xhigh",
)
assert params["max_completion_tokens"] == 100
assert "max_tokens" not in params
assert params["reasoning_effort"] == "max"
assert params["reasoning_effort"] == "xhigh"
@pytest.mark.parametrize("model", ["azure/gpt-6-astra", "azure/us/gpt-6-astra"])
def test_azure_gpt6_astra_reasoning_effort_none_unlocks_temperature(config: AzureOpenAIGPT5Config, model: str):
"""Foundry's gpt-6-astra accepts reasoning_effort='none' and, only then, a non-default
temperature (verified live against a Foundry deployment), unlike OpenAI's gpt-6-astra."""
@pytest.mark.parametrize("model", ["azure/gpt-6-sol", "azure/us/gpt-6-sol"])
def test_azure_reasoning_effort_none_unlocks_temperature_off_the_azure_row(
config: AzureOpenAIGPT5Config, monkeypatch: pytest.MonkeyPatch, model: str
):
"""A deployment whose azure/ row takes none accepts a non-default temperature with it, read off
that row even when OpenAI's twin row says the model refuses none."""
monkeypatch.setitem(
litellm.model_cost, "gpt-6-sol", {**litellm.model_cost["gpt-6-sol"], "supports_none_reasoning_effort": False}
)
params = config.map_openai_params(
non_default_params={"temperature": 0.2, "reasoning_effort": "none"},
optional_params={},
@ -365,14 +373,127 @@ def test_azure_gpt6_astra_reasoning_effort_none_unlocks_temperature(config: Azur
assert params["reasoning_effort"] == "none"
@pytest.mark.parametrize("model", ["azure/gpt-6-astra", "azure/us/gpt-6-astra"])
def test_azure_gpt6_astra_rejects_reasoning_effort_minimal(config: AzureOpenAIGPT5Config, model: str):
"""Foundry's gpt-6-astra lists none, low, medium, high, xhigh and max but not minimal."""
with pytest.raises(litellm.utils.UnsupportedParamsError):
config.map_openai_params(
non_default_params={"reasoning_effort": "minimal"},
def _azure_chat_row_forwards(row: Mapping[str, object], level: str) -> bool:
"""Azure chat polarity: max is never gated, xhigh and none need an explicit true, every other
level only has to not be false."""
if level == "max":
return True
flag: Final = row.get(f"supports_{level}_reasoning_effort")
return flag is True if level in ("xhigh", "none") else flag is not False
@pytest.mark.parametrize("model", ["azure/gpt-6-astra", "azure/us/gpt-6-astra", "azure/eu/gpt-6-astra"])
@pytest.mark.parametrize("level", ["none", "minimal", "low", "medium", "high", "xhigh", "max"])
def test_azure_gpt6_astra_forwards_exactly_the_effort_levels_its_map_row_allows(
config: AzureOpenAIGPT5Config, model: str, level: str
):
"""The azure/ row is the whole contract for a Foundry deployment. A live Foundry gpt-6-astra
(model gpt-6-astra-2026-09-03) answered reasoning_effort none with a 400 naming low, medium,
high and xhigh on 2026-10-01, so its row turns none and minimal off. The row keeps max off too,
since chat refuses it, but no gate reads that flag: the Responses route honors max, so chat
forwards it for the bridge and a plain chat call gets the provider's own 400."""
forwarded: Final = _azure_chat_row_forwards(litellm.model_cost[model], level)
dropped: Final = config.map_openai_params(
non_default_params={"reasoning_effort": level},
optional_params={},
model=model,
drop_params=True,
api_version="2025-04-01-preview",
)
assert ("reasoning_effort" in dropped) is forwarded
if forwarded:
kept: Final = config.map_openai_params(
non_default_params={"reasoning_effort": level},
optional_params={},
model=model,
drop_params=False,
api_version="2025-04-01-preview",
)
assert kept["reasoning_effort"] == level
else:
with pytest.raises(litellm.utils.UnsupportedParamsError):
config.map_openai_params(
non_default_params={"reasoning_effort": level},
optional_params={},
model=model,
drop_params=False,
api_version="2025-04-01-preview",
)
@pytest.mark.parametrize("model", ["azure/gpt-6-astra", "azure/us/gpt-6-astra", "azure/eu/gpt-6-astra"])
def test_azure_gpt6_astra_row_turns_none_and_minimal_off(model: str):
"""Pinned on purpose so the row-derived test above cannot go vacuous: a Foundry gpt-6-astra
deployment refused none on chat and on Responses on 2026-10-01 (400 unsupported_value naming
low, medium, high, xhigh), and minimal the same way; max it refused on chat while honoring it
on Responses, so the route-blind flag stays off. When Azure adds a level, flip the row and
this assertion together."""
row: Final = litellm.model_cost[model]
assert row["supports_none_reasoning_effort"] is False
assert row["supports_minimal_reasoning_effort"] is False
assert row["supports_max_reasoning_effort"] is False
def test_azure_chat_gate_reads_the_foundry_row_when_azure_ai_prefix_survives_provider_remap(
config: AzureOpenAIGPT5Config, monkeypatch: pytest.MonkeyPatch
):
"""A Foundry deployment on an OpenAI-v1 host is re-routed to the azure provider with its
azure_ai/ prefix intact, and this config then gates it. The rows are made to disagree on low,
a level neither real row flags, so the prefixed name has to read azure_ai/gpt-6-sol while the
bare deployment name keeps reading azure/gpt-6-sol."""
monkeypatch.setitem(
litellm.model_cost,
"azure_ai/gpt-6-sol",
{**litellm.model_cost["azure_ai/gpt-6-sol"], "supports_low_reasoning_effort": False},
)
monkeypatch.setitem(
litellm.model_cost,
"azure/gpt-6-sol",
{**litellm.model_cost["azure/gpt-6-sol"], "supports_low_reasoning_effort": True},
)
prefixed: Final = config.map_openai_params(
non_default_params={"reasoning_effort": "low"},
optional_params={},
model="azure_ai/gpt-6-sol",
drop_params=True,
api_version="2025-04-01-preview",
)
assert "reasoning_effort" not in prefixed
bare: Final = config.map_openai_params(
non_default_params={"reasoning_effort": "low"},
optional_params={},
model="gpt-6-sol",
drop_params=True,
api_version="2025-04-01-preview",
)
assert bare["reasoning_effort"] == "low"
def test_azure_chat_gate_falls_back_to_the_azure_row_for_an_azure_ai_name_the_map_lacks(
config: AzureOpenAIGPT5Config, monkeypatch: pytest.MonkeyPatch
):
monkeypatch.delitem(litellm.model_cost, "azure_ai/gpt-6-sol")
monkeypatch.setitem(
litellm.model_cost,
"azure/gpt-6-sol",
{**litellm.model_cost["azure/gpt-6-sol"], "supports_low_reasoning_effort": False},
)
dropped: Final = config.map_openai_params(
non_default_params={"reasoning_effort": "low"},
optional_params={},
model="azure_ai/gpt-6-sol",
drop_params=True,
api_version="2025-04-01-preview",
)
assert "reasoning_effort" not in dropped
@pytest.mark.parametrize("key", ["azure/gpt-6.1-sol", "azure/gpt-6.1-sol-2026-09-29", "azure_ai/gpt-6.1-sol"])
def test_azure_gpt61_sol_rows_turn_none_off(key: str):
"""Pinned on purpose, as the astra row above: a Foundry gpt-6.1-sol deployment (model
gpt-6.1-sol-2026-09-29) refused none on chat and on Responses on 2026-10-02 (400 unsupported_value
naming low, medium, high, xhigh, and max on Responses), and minimal the same way, while these rows
still said none was on. When Azure adds a level, flip the row and this assertion together."""
row: Final = litellm.model_cost[key]
assert row["supports_none_reasoning_effort"] is False
assert row["supports_minimal_reasoning_effort"] is False

125
tests/unit/llms/azure/response/test_azure_transformation.py
View file

@ -656,16 +656,21 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
litellm.add_known_models(model_cost_map=litellm.model_cost)
def test_azure_responses_gpt6_astra_reasoning_effort_none_unlocks_temperature(local_model_cost_map: None):
"""Foundry's gpt-6-astra accepts reasoning.effort='none' with a non-default temperature
while OpenAI's gpt-6-astra does not, so the gate must read the azure/ cost-map entry
for the bare deployment name rather than OpenAI's."""
def test_azure_responses_reasoning_effort_none_unlocks_temperature_off_the_azure_row(
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
):
"""The gate reads the azure/ cost-map entry for the bare deployment name rather than OpenAI's
twin: with OpenAI's gpt-6-sol row turned off for none, azure/gpt-6-sol's row still lets none
through and, only then, a non-default temperature."""
monkeypatch.setitem(
litellm.model_cost, "gpt-6-sol", {**litellm.model_cost["gpt-6-sol"], "supports_none_reasoning_effort": False}
)
params = AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(
temperature=0.2,
reasoning={"effort": "none"},
),
model="gpt-6-astra",
model="gpt-6-sol",
drop_params=False,
)
assert params["temperature"] == 0.2
@ -684,6 +689,116 @@ def test_azure_responses_gpt6_astra_rejects_temperature_while_reasoning(local_mo
)
@pytest.mark.parametrize("model", ["gpt-6-astra", "gpt-6.1-sol"])
@pytest.mark.parametrize("level", ["none", "minimal", "low", "medium", "high", "xhigh", "max"])
def test_azure_responses_gpt6_effort_levels_follow_the_azure_row(local_model_cost_map: None, model: str, level: str):
"""A bare deployment name reads the azure/ row, opt-out for every level but max, so the none
and minimal Foundry refuses on this route (live 400s on 2026-10-01 for gpt-6-astra and 2026-10-02
for gpt-6.1-sol) are dropped or refused here while the max it honors goes through whatever the
route-blind flag says."""
forwarded = (
level == "max" or litellm.model_cost[f"azure/{model}"].get(f"supports_{level}_reasoning_effort") is not False
)
dropped = AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(reasoning={"effort": level}),
model=model,
drop_params=True,
)
assert ("reasoning" in dropped) is forwarded
if forwarded:
kept = AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(reasoning={"effort": level}),
model=model,
drop_params=False,
)
assert kept["reasoning"] == {"effort": level}
else:
with pytest.raises(litellm.UnsupportedParamsError):
AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(reasoning={"effort": level}),
model=model,
drop_params=False,
)
def test_azure_gpt6_astra_keeps_max_on_this_route_whatever_its_row_says(local_model_cost_map: None):
"""Pinned on purpose: the row-derived test above exempts max by hand, so this one has to fail
if the gate ever reads the flag. A Foundry gpt-6-astra deployment answered reasoning.effort max
through /v1/responses with 200, the effort echoed back and reasoning tokens billed, on
2026-10-01, while chat refuses the same level, as OpenAI's chat does. The row keeps max off for
the adapter and the picker, and this gate still forwards it even under drop_params."""
assert litellm.model_cost["azure/gpt-6-astra"]["supports_max_reasoning_effort"] is False
kept = AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(reasoning={"effort": "max"}),
model="gpt-6-astra",
drop_params=True,
)
assert kept["reasoning"] == {"effort": "max"}
@pytest.mark.parametrize("level", ["none", "minimal", "low", "medium", "high", "xhigh", "max"])
def test_azure_responses_effort_gate_fires_when_azure_ai_prefix_survives_provider_remap(
local_model_cost_map: None, level: str
):
"""litellm.responses(model="azure_ai/gpt-6.1-sol", api_base=<Foundry host>) is re-routed to the
azure provider with the prefix intact, so this config sees the prefixed name. The gate used to
skip every prefixed name as another provider's pass-through, and a live Foundry gpt-6.1-sol
answered the forwarded none and minimal with 400 on 2026-10-02 even under drop_params. The
prefixed name reads the azure_ai/ row and drops or refuses exactly what that row turns off."""
forwarded = (
level == "max"
or litellm.model_cost["azure_ai/gpt-6.1-sol"].get(f"supports_{level}_reasoning_effort") is not False
)
dropped = AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(
max_output_tokens=60, reasoning={"effort": level}
),
model="azure_ai/gpt-6.1-sol",
drop_params=True,
)
assert ("reasoning" in dropped) is forwarded
if forwarded:
return
with pytest.raises(litellm.UnsupportedParamsError):
AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(
max_output_tokens=60, reasoning={"effort": level}
),
model="azure_ai/gpt-6.1-sol",
drop_params=False,
)
def test_azure_responses_prefixed_name_reads_the_foundry_row_not_the_azure_twin(
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
):
"""The rows are made to disagree on low, a level neither real row flags, so the prefixed name
proving it reads azure_ai/gpt-6-sol cannot pass by reading azure/gpt-6-sol, which the bare
deployment name keeps reading."""
monkeypatch.setitem(
litellm.model_cost,
"azure_ai/gpt-6-sol",
{**litellm.model_cost["azure_ai/gpt-6-sol"], "supports_low_reasoning_effort": False},
)
monkeypatch.setitem(
litellm.model_cost,
"azure/gpt-6-sol",
{**litellm.model_cost["azure/gpt-6-sol"], "supports_low_reasoning_effort": True},
)
prefixed = AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(reasoning={"effort": "low"}),
model="azure_ai/gpt-6-sol",
drop_params=True,
)
assert "reasoning" not in prefixed
bare = AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=ResponsesAPIOptionalRequestParams(reasoning={"effort": "low"}),
model="gpt-6-sol",
drop_params=True,
)
assert bare["reasoning"] == {"effort": "low"}
def test_azure_responses_sends_the_deployment_name_when_azure_ai_prefix_survives_provider_remap():
request = AzureOpenAIResponsesAPIConfig().transform_responses_api_request(
model="azure_ai/gpt-5.4-nano",

17
tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py
View file

@ -148,17 +148,30 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url))
def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none(_local_model_cost_map):
def test_foundry_gpt_6_sol_keeps_sampling_params_when_reasoning_effort_is_none(_local_model_cost_map):
"""gpt-6-sol's Foundry row still says none, so the sampling params that ride on it stay. gpt-6-astra
used to be the model here until its row turned none off (a live Foundry deployment answered none
with a 400 naming low through xhigh on 2026-10-01); the test below pins that flip."""
optional_params = AzureAIStudioConfig().map_openai_params(
non_default_params={"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9},
optional_params={},
model="gpt-6-astra",
model="gpt-6-sol",
drop_params=False,
)
assert optional_params == {"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9}
def test_foundry_gpt_6_astra_refuses_reasoning_effort_none_before_the_call(_local_model_cost_map):
with pytest.raises(litellm.utils.UnsupportedParamsError):
AzureAIStudioConfig().map_openai_params(
non_default_params={"reasoning_effort": "none"},
optional_params={},
model="gpt-6-astra",
drop_params=False,
)
def test_azure_ai_grok_stop_parameter_handling():

81
tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py
View file

@ -7,13 +7,17 @@ sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
)
from typing import Final
from unittest.mock import patch
import httpx
import pytest
import litellm
from litellm.llms.azure_ai.anthropic.messages_transformation import (
AzureAnthropicMessagesConfig,
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.types.router import GenericLiteLLMParams
@ -470,3 +474,80 @@ def test_azure_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_fl
and info.get("supports_mid_conversation_system") is not True
]
assert missing == []
AUTO_MODE_SAFEGUARDS: Final = ({"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}},)
@pytest.mark.usefixtures("local_beta_headers_config")
@pytest.mark.parametrize(
("client_headers", "optional_params", "expected_beta"),
[
({}, {"safeguards": AUTO_MODE_SAFEGUARDS}, "dangerous-tool-use-2026-09-03"),
(
{"anthropic-beta": "context-1m-2025-08-07"},
{"safeguards": AUTO_MODE_SAFEGUARDS},
"context-1m-2025-08-07,dangerous-tool-use-2026-09-03",
),
({"anthropic-beta": "context-1m-2025-08-07"}, {}, "context-1m-2025-08-07"),
({}, {}, None),
],
)
def test_safeguards_request_carries_the_dangerous_tool_use_beta_without_the_client_header(
client_headers: dict[str, str], optional_params: dict[str, object], expected_beta: str | None
) -> None:
headers, _ = AzureAnthropicMessagesConfig().validate_anthropic_messages_environment(
headers=client_headers,
model="claude-sonnet-4-6",
messages=[{"role": "user", "content": "Run ls -la with the Bash tool"}],
optional_params=optional_params,
litellm_params={"api_key": "test-api-key"},
)
assert headers.get("anthropic-beta") == expected_beta, headers
def _foundry_answer(request: httpx.Request) -> httpx.Response:
return httpx.Response(
status_code=200,
headers={"content-type": "application/json"},
json={
"id": "msg_safeguards",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": [{"type": "text", "text": "hello from the gateway"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 12, "output_tokens": 6},
},
request=request,
)
@pytest.mark.asyncio
@pytest.mark.usefixtures("local_beta_headers_config")
async def test_outgoing_azure_request_sends_safeguards_with_the_dangerous_tool_use_beta() -> None:
received: Final[list[httpx.Request]] = []
def foundry_answers(request: httpx.Request) -> httpx.Response:
received.append(request)
return _foundry_answer(request)
upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers))
try:
await litellm.anthropic.messages.acreate(
max_tokens=256,
messages=[{"role": "user", "content": "Use the Bash tool to run: echo hello from the gateway"}],
model="azure_ai/claude-sonnet-4-6",
api_key="test-api-key",
api_base="https://test-resource.services.ai.azure.com",
safeguards=AUTO_MODE_SAFEGUARDS,
client=upstream,
)
finally:
await upstream.close()
assert len(received) == 1, received
sent: Final = received[0]
assert "dangerous-tool-use-2026-09-03" in sent.headers["anthropic-beta"].split(","), dict(sent.headers)
assert json.loads(sent.content)["safeguards"] == list(AUTO_MODE_SAFEGUARDS)

74
tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py
View file

@ -1,12 +1,15 @@
import json
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
from litellm.llms.azure_ai.responses.transformation import AzureAIResponsesAPIConfig
from litellm.responses.main import _will_bridge_to_chat_completions
from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager
@ -309,3 +312,74 @@ async def test_aresponses_serverless_host_stays_on_chat_bridge():
assert chat_route.called
assert not responses_route.called
assert chat_route.calls.last.request.headers["Authorization"] == "Bearer fake-key"
def _responses_params(**fields: object) -> ResponsesAPIOptionalRequestParams:
return ResponsesAPIOptionalRequestParams(**fields) # pyright: ignore[reportArgumentType] # TypedDict kwargs from a test helper
def test_foundry_responses_effort_gate_reads_the_azure_ai_row(
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
):
"""A Foundry deployment and its Azure OpenAI namesake are different products, so the Foundry
Responses route answers from the azure_ai/ row the Foundry chat route already reads, never from
the azure/ twin. The rows are made to disagree on low, a level neither real row flags, so a
leak from the wrong row cannot hide behind rows that happen to agree."""
monkeypatch.setitem(
litellm.model_cost,
"azure_ai/gpt-6-sol",
{**litellm.model_cost["azure_ai/gpt-6-sol"], "supports_low_reasoning_effort": False},
)
monkeypatch.setitem(
litellm.model_cost,
"azure/gpt-6-sol",
{**litellm.model_cost["azure/gpt-6-sol"], "supports_low_reasoning_effort": True},
)
dropped: Final = AzureAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=_responses_params(reasoning={"effort": "low"}),
model="gpt-6-sol",
drop_params=True,
)
assert "reasoning" not in dropped
with pytest.raises(litellm.UnsupportedParamsError):
AzureAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=_responses_params(reasoning={"effort": "low"}),
model="gpt-6-sol",
drop_params=False,
)
kept_by_azure_openai: Final = AzureOpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=_responses_params(reasoning={"effort": "low"}),
model="gpt-6-sol",
drop_params=True,
)
assert kept_by_azure_openai["reasoning"] == {"effort": "low"}
def test_foundry_responses_none_unlocks_temperature_off_the_azure_ai_row(
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
):
"""The sampling-param unlock that rides on none reads the same azure_ai/ row as the effort gate:
with the azure/ twin turned off for none, the Foundry row still lets none through and, only
then, a non-default temperature."""
monkeypatch.setitem(
litellm.model_cost,
"azure_ai/gpt-6-sol",
{**litellm.model_cost["azure_ai/gpt-6-sol"], "supports_none_reasoning_effort": True},
)
monkeypatch.setitem(
litellm.model_cost,
"azure/gpt-6-sol",
{**litellm.model_cost["azure/gpt-6-sol"], "supports_none_reasoning_effort": False},
)
params: Final = AzureAIResponsesAPIConfig().map_openai_params(
response_api_optional_params=_responses_params(temperature=0.2, reasoning={"effort": "none"}),
model="gpt-6-sol",
drop_params=False,
)
assert params["temperature"] == 0.2
assert params["reasoning"] == {"effort": "none"}

41
tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -302,3 +302,44 @@ async def test_async_converse_streaming_fails_at_the_request_timeout_not_the_ups
def test_sync_converse_streaming_fails_at_the_request_timeout_not_the_upstreams_pace() -> None:
with pytest.raises(litellm.Timeout):
litellm.completion(client=slow_upstream_sync_client(), **_converse_streaming_kwargs())
@pytest.mark.parametrize("stream", (False, True))
def test_converse_profile_override_records_cost_without_changing_response_model(
monkeypatch: pytest.MonkeyPatch, stream: bool
) -> None:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setenv("LITELLM_RUST", "0")
configuration.rust(False)
model: Final = "openai.gpt-5.6-sol"
profile: Final = f"us.{model}"
prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock")
def respond(request: httpx.Request) -> httpx.Response:
assert request.url.path == f"/model/{profile}/converse"
return httpx.Response(200, json=CONVERSE_RESPONSE)
with httpx.Client(transport=httpx.MockTransport(respond)) as transport:
response: Final = litellm.completion(
model=f"bedrock/converse/{model}",
model_id=profile,
messages=[{"role": "user", "content": "hi"}],
api_key="test-bearer",
aws_region_name="us-east-1",
client=HTTPHandler(client=transport),
stream=stream,
fake_stream=stream,
max_tokens=16,
)
result: Final = litellm.stream_chunk_builder(list(response)) if stream else response
assert result is not None
assert result.model == model
assert result.choices[0].message.content == "hi"
cost: Final = (
response.logging_obj.response_cost_calculator(result) if stream else result._hidden_params.get("response_cost")
)
assert cost == pytest.approx(
CONVERSE_RESPONSE["usage"]["inputTokens"] * prices["input_cost_per_token"]
+ CONVERSE_RESPONSE["usage"]["outputTokens"] * prices["output_cost_per_token"]
)

6
tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py
View file

@ -636,7 +636,11 @@ class TestAgentCoreSearchEdgeCases:
from litellm.search.cost_calculator import search_provider_cost_per_query
monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map())
assert search_provider_cost_per_query(model="agentcore/search", custom_llm_provider="agentcore") == (0.0, 0.0)
price_per_query = litellm.model_cost["agentcore/search"]["input_cost_per_query"]
assert price_per_query > 0
assert search_provider_cost_per_query(
model="agentcore/search", custom_llm_provider="agentcore", number_of_queries=3
) == (3 * price_per_query, 0.0)
def _transform_text(text: str) -> SearchResponse:

9
tests/unit/llms/bedrock/slow_upstream.py
View file

@ -7,6 +7,7 @@ gets the answer.
from __future__ import annotations
import json
from typing import Final
import httpx
@ -16,6 +17,14 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
STREAM_TIMEOUT_SECONDS: Final = 0.5
UPSTREAM_FIRST_BYTE_SECONDS: Final = 2.0
CONVERSE_ANSWER: Final = json.dumps(
{
"output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
"metrics": {"latencyMs": 1},
}
).encode()
_TIMEOUT_EXTENSION: Final = TypeAdapter(dict[str, float | None])

42
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -5121,3 +5121,45 @@ async def test_retrieve_file_rejects_unverified_bedrock_range_error(
)
finally:
sync_client.close()
@pytest.mark.asyncio
@pytest.mark.usefixtures("local_beta_headers_config")
async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_its_beta() -> None:
received: Final[list[httpx.Request]] = []
def anthropic_answers(request: httpx.Request) -> httpx.Response:
received.append(request)
return httpx.Response(
status_code=200,
json={
"id": "msg_dropped_safeguards",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 5, "output_tokens": 1},
},
request=request,
)
upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(anthropic_answers))
try:
await litellm.anthropic.messages.acreate(
max_tokens=16,
messages=[{"role": "user", "content": "hi"}],
model="anthropic/claude-sonnet-4-6",
api_key="sk-ant-test",
safeguards=({"type": "dangerous_tool_use"},),
additional_drop_params=["safeguards"],
client=upstream,
)
finally:
await upstream.close()
assert len(received) == 1, received
sent: Final = received[0]
assert "safeguards" not in json.loads(sent.content)
assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers)

202
tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py
View file

@ -1,3 +1,4 @@
import copy
import json
from collections.abc import Iterator
from typing import Final
@ -6,7 +7,9 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
from pydantic import TypeAdapter
from fastapi.testclient import TestClient
from pydantic import BaseModel, ConfigDict
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
@ -21,6 +24,7 @@ from litellm.llms.databricks.chat.transformation import (
DatabricksConfig,
_sanitize_empty_content,
)
from litellm.types.llms.openai import AllMessageValues
import asyncio
from unittest.mock import Mock
from litellm._version import version
@ -32,6 +36,25 @@ DATABRICKS_API_BASE: Final = "https://my.workspace.cloud.databricks.com/serving-
DATABRICKS_API_KEY: Final = "dapimykey"
DATABRICKS_CHAT_COMPLETIONS_URL: Final = f"{DATABRICKS_API_BASE}/chat/completions"
DATABRICKS_EMBEDDINGS_URL: Final = f"{DATABRICKS_API_BASE}/embeddings"
JSON_SCHEMA_RESPONSE_FORMAT: Final = {
"type": "json_schema",
"json_schema": {
"name": "P",
"strict": True,
"schema": {
"$defs": {
"Person": {
"type": "object",
"properties": {"name": {"type": "string"}},
"required": ["name"],
}
},
"type": "object",
"properties": {"p": {"$ref": "#/$defs/Person"}},
"required": ["p"],
},
},
}
@pytest.fixture()
@ -411,6 +434,119 @@ def test_transform_request_preserves_unity_model_service_name():
assert result["model"] == "system.ai.kimi-k3"
@pytest.mark.parametrize(
"model,messages,optional_params,expected_content",
[
pytest.param(
"databricks-meta-llama-3-3-70b-instruct",
[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
{"response_format": {"type": "json_object"}},
"Reply in JSON",
id="user",
),
pytest.param(
"databricks-meta-llama-3-3-70b-instruct",
[{"role": "system", "content": [{"type": "text", "text": "Reply in JSON"}]}],
{"response_format": {"type": "json_object"}},
"Reply in JSON",
id="system",
),
pytest.param(
"databricks-meta-llama-3-3-70b-instruct",
[{"role": "assistant", "content": [{"type": "text", "text": "Reply in JSON"}]}],
{"response_format": {"type": "json_object"}},
"Reply in JSON",
id="assistant",
),
pytest.param(
"databricks-meta-llama-3-3-70b-instruct",
[
{
"role": "tool",
"tool_call_id": "call_1",
"content": [{"type": "text", "text": "Reply in JSON"}],
}
],
{"response_format": {"type": "json_object"}},
"Reply in JSON",
id="tool",
),
pytest.param(
"databricks-meta-llama-3-3-70b-instruct",
[
{
"role": "user",
"content": [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}],
}
],
{"response_format": {"type": "json_object"}},
[{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}],
id="cache-control",
),
pytest.param(
"databricks-meta-llama-3-3-70b-instruct",
[{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}],
{"response_format": {"type": "json_object"}},
[{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}],
id="multiple-text-blocks",
),
pytest.param(
"databricks-meta-llama-3-3-70b-instruct",
[
{
"role": "user",
"content": [
{"type": "text", "text": "Reply in JSON"},
{"type": "image_url", "image_url": {"url": "https://example.com/image.png"}},
],
}
],
{"response_format": {"type": "json_object"}},
[
{"type": "text", "text": "Reply in JSON"},
{"type": "image_url", "image_url": {"url": "https://example.com/image.png"}},
],
id="text-and-image",
),
pytest.param(
"databricks-meta-llama-3-3-70b-instruct",
[{"role": "user", "content": "Reply in JSON"}],
{"response_format": {"type": "json_object"}},
"Reply in JSON",
id="plain-string",
),
pytest.param(
"databricks-claude-sonnet-5",
[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
{"response_format": {"type": "json_object"}},
[{"type": "text", "text": "Reply in JSON"}],
id="claude",
),
],
)
def test_transform_request_collapses_single_plain_text_block_for_non_claude_models(
model: str,
messages: list[AllMessageValues],
optional_params: dict[str, object],
expected_content: object,
) -> None:
config: Final[DatabricksConfig] = DatabricksConfig()
result: Final[dict[str, object]] = TypeAdapter(dict[str, object]).validate_python(
config.transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped
model=model,
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={},
)
)
outgoing_messages: Final[list[dict[str, object]]] = TypeAdapter(list[dict[str, object]]).validate_python(
result["messages"]
)
assert outgoing_messages == [{**messages[0], "content": expected_content}]
def test_transform_request_strips_thinking_blocks_and_reasoning_content():
"""Regression for LIT-6762: replaying an assistant turn that litellm decorated with
`thinking_blocks` / `reasoning_content` made Databricks 400 with
@ -2088,3 +2224,69 @@ def mock_chat_streaming_response_chunks() -> List[str]:
}
),
]
@pytest.mark.parametrize(
"model",
(
"databricks-meta-llama-3-3-70b-instruct",
"databricks-qwen35-122b-a10b",
"databricks-gpt-oss-120b",
"system.ai.deepseek-v4-1-flash",
),
)
def test_databricks_non_claude_json_schema_preserves_refs(model: str) -> None:
optional_params: Final = litellm.utils.get_optional_params(
model=model,
custom_llm_provider="databricks",
response_format=copy.deepcopy(JSON_SCHEMA_RESPONSE_FORMAT),
)
assert optional_params["response_format"] == JSON_SCHEMA_RESPONSE_FORMAT
def test_databricks_claude_json_schema_preserves_refs_in_tool_parameters() -> None:
optional_params: Final = litellm.utils.get_optional_params(
model="databricks-claude-haiku-4-5",
custom_llm_provider="databricks",
response_format=copy.deepcopy(JSON_SCHEMA_RESPONSE_FORMAT),
)
assert "response_format" not in optional_params
assert optional_params["tools"] == [
{
"type": "function",
"function": {
"name": "json_tool_call",
"parameters": JSON_SCHEMA_RESPONSE_FORMAT["json_schema"]["schema"],
},
}
]
def test_databricks_pydantic_json_schema_preserves_nested_refs() -> None:
class Person(BaseModel):
model_config = ConfigDict(frozen=True)
name: str
class Order(BaseModel):
model_config = ConfigDict(frozen=True)
person: Person
optional_params: Final = litellm.utils.get_optional_params(
model="databricks-meta-llama-3-3-70b-instruct",
custom_llm_provider="databricks",
response_format=Order,
)
schema: Final = optional_params["response_format"]["json_schema"]["schema"]
assert schema["properties"]["person"] == {"$ref": "#/$defs/Person"}
assert schema["$defs"]["Person"] == {
"additionalProperties": False,
"properties": {"name": {"title": "Name", "type": "string"}},
"required": ["name"],
"title": "Person",
"type": "object",
}

24
tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py
View file

@ -1,3 +1,4 @@
from typing import Final
from unittest.mock import MagicMock
import pytest
@ -184,6 +185,29 @@ def test_github_copilot_anthropic_messages_validate_environment_injects_context_
assert "context-management-2025-06-27" in validated_headers["anthropic-beta"]
@pytest.mark.usefixtures("local_beta_headers_config")
def test_github_copilot_safeguards_do_not_derive_the_dangerous_tool_use_beta() -> None:
config: Final = GithubCopilotAnthropicMessagesConfig()
config.authenticator = MagicMock()
config.authenticator.get_api_key.return_value = "gh.test-key"
config.authenticator.get_api_base.return_value = None
validated_headers, _ = config.validate_anthropic_messages_environment(
headers={},
model="github_copilot/claude-haiku-4.5",
messages=[{"role": "user", "content": "Hello"}],
optional_params={
"safeguards": ({"type": "dangerous_tool_use"},),
"context_management": {"edits": ({"type": "clear_tool_uses_20250919"},)},
},
litellm_params={},
api_key=None,
api_base=None,
)
assert validated_headers["anthropic-beta"] == "context-management-2025-06-27"
def test_github_copilot_anthropic_messages_validate_environment_auth_error():
"""Test error handling when authentication fails."""
config = GithubCopilotAnthropicMessagesConfig()

109
tests/unit/llms/openai/responses/test_openai_responses_transformation.py
View file

@ -2015,6 +2015,115 @@ class TestResponsesSurfaceSharesTheEffortRule:
)
assert ("top_p" in mapped) is top_p_survives
@pytest.mark.parametrize("level", ["none", "minimal", "low", "medium", "high", "xhigh", "max"])
def test_effort_levels_the_map_row_turns_off_are_dropped_or_refused(self, local_model_cost_map, level):
"""Same row as the chat gate, read opt-out for every level but max. gpt-6.1-sol's row turns
none and minimal off: OpenAI lists low, medium, high, xhigh and max for it
(https://developers.openai.com/api/docs/models/gpt-6.1-sol, via #43910), and a local proxy
got OpenAI's 400 for none and minimal through /v1/responses on 2026-09-30."""
forwarded = (
level == "max" or litellm.model_cost["gpt-6.1-sol"].get(f"supports_{level}_reasoning_effort") is not False
)
dropped = OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"reasoning": {"effort": level}},
model="gpt-6.1-sol",
drop_params=True,
)
assert ("reasoning" in dropped) is forwarded
if forwarded:
kept = OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"reasoning": {"effort": level}},
model="gpt-6.1-sol",
drop_params=False,
)
assert kept["reasoning"] == {"effort": level}
else:
with pytest.raises(litellm.UnsupportedParamsError):
OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"reasoning": {"effort": level}},
model="gpt-6.1-sol",
drop_params=False,
)
@pytest.mark.parametrize(
("flags", "level", "forwarded"),
[
({}, "none", True),
({"supports_none_reasoning_effort": False}, "none", False),
({"supports_none_reasoning_effort": True}, "none", True),
({"supports_minimal_reasoning_effort": False}, "minimal", False),
({"supports_max_reasoning_effort": False}, "max", True),
({}, "xhigh", True),
({"supports_xhigh_reasoning_effort": False}, "xhigh", False),
({"supports_none_reasoning_effort": False}, "medium", True),
],
)
def test_the_row_flags_alone_decide_whether_a_level_is_forwarded(
self, local_model_cost_map, monkeypatch, flags, level, forwarded
):
"""Same model name, rewritten row: this surface reads every level opt-out, xhigh included,
so the gate flips with the flags alone. max is the one level no flag can turn off here,
since the same model honors it on this route and refuses it on chat. The chat test module
pins the real row's facts."""
real_row = litellm.model_cost["gpt-6.1-sol"]
without_flags = {
field: value
for field, value in real_row.items()
if not (field.startswith("supports_") and field.endswith("_reasoning_effort"))
}
monkeypatch.setitem(litellm.model_cost, "gpt-6.1-sol", {**without_flags, **flags})
dropped = OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"reasoning": {"effort": level}},
model="gpt-6.1-sol",
drop_params=True,
)
assert ("reasoning" in dropped) is forwarded
if not forwarded:
with pytest.raises(litellm.UnsupportedParamsError):
OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"reasoning": {"effort": level}},
model="gpt-6.1-sol",
drop_params=False,
)
def test_dropping_a_refused_effort_keeps_the_rest_of_the_reasoning_object(self, local_model_cost_map):
refused = next(
level
for level in ("none", "minimal", "low")
if litellm.model_cost["gpt-6.1-sol"].get(f"supports_{level}_reasoning_effort") is False
)
mapped = OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"reasoning": {"effort": refused, "summary": "auto"}},
model="gpt-6.1-sol",
drop_params=True,
)
assert mapped["reasoning"] == {"summary": "auto"}
@pytest.mark.parametrize("level", ["none", "minimal", "low", "medium", "high", "xhigh", "max"])
def test_gpt5_3_codex_forwards_exactly_the_effort_levels_its_map_row_allows(self, local_model_cost_map, level):
"""The gate turns a stale false into a local refusal, so the Responses-only codex row has to
say what the route does. On 2026-10-01 /v1/responses answered gpt-5.3-codex none and xhigh
with 200 (effort echoed back) and minimal and max with 400 unsupported_value naming none,
low, medium, high and xhigh, so the row takes none and xhigh and turns minimal off."""
forwarded = (
level == "max" or litellm.model_cost["gpt-5.3-codex"].get(f"supports_{level}_reasoning_effort") is not False
)
mapped = OpenAIResponsesAPIConfig().map_openai_params(
response_api_optional_params={"reasoning": {"effort": level}},
model="gpt-5.3-codex",
drop_params=True,
)
assert ("reasoning" in mapped) is forwarded
def test_gpt5_3_codex_row_takes_none_and_xhigh(self, local_model_cost_map):
"""Pinned on purpose so the row-derived test above cannot go vacuous: the live answers of
2026-10-01 listed above. When OpenAI changes the set, flip the row and this assertion
together."""
row = litellm.model_cost["gpt-5.3-codex"]
assert row["supports_none_reasoning_effort"] is True
assert row["supports_xhigh_reasoning_effort"] is True
assert row["supports_minimal_reasoning_effort"] is False
def test_top_p_raises_without_drop_params(self, local_model_cost_map):
with pytest.raises(litellm.UnsupportedParamsError):
OpenAIResponsesAPIConfig().map_openai_params(

101
tests/unit/llms/openai/test_gpt5_transformation.py
View file

@ -1,4 +1,6 @@
import re
from collections.abc import Mapping
from typing import Final
import pytest
@ -431,6 +433,105 @@ def test_gpt5_drops_reasoning_effort_minimal_when_requested(config: OpenAIConfig
assert "reasoning_effort" not in params
def _chat_row_forwards(row: Mapping[str, object], level: str) -> bool:
"""Chat polarity: max is never gated, xhigh needs an explicit true, every other level only has
to not be false."""
if level == "max":
return True
flag: Final = row.get(f"supports_{level}_reasoning_effort")
return flag is True if level == "xhigh" else flag is not False
@pytest.mark.parametrize("level", ["none", "minimal", "low", "medium", "high", "xhigh", "max"])
def test_gpt6_1_sol_forwards_exactly_the_effort_levels_its_map_row_allows(config: OpenAIConfig, level: str):
"""The map row is the whole contract, so a new GPT-6 model is gated by its row with no code
change. gpt-6.1-sol's row turns none and minimal off: OpenAI lists low, medium, high, xhigh
and max for it (https://developers.openai.com/api/docs/models/gpt-6.1-sol, via #43910), and a
local proxy got OpenAI's 400 for none and minimal on both endpoints on 2026-09-30."""
forwarded: Final = _chat_row_forwards(litellm.model_cost["gpt-6.1-sol"], level)
dropped: Final = config.map_openai_params(
non_default_params={"reasoning_effort": level},
optional_params={},
model="gpt-6.1-sol",
drop_params=True,
)
assert ("reasoning_effort" in dropped) is forwarded
if forwarded:
kept: Final = config.map_openai_params(
non_default_params={"reasoning_effort": level},
optional_params={},
model="gpt-6.1-sol",
drop_params=False,
)
assert kept["reasoning_effort"] == level
else:
with pytest.raises(litellm.utils.UnsupportedParamsError):
config.map_openai_params(
non_default_params={"reasoning_effort": level},
optional_params={},
model="gpt-6.1-sol",
drop_params=False,
)
def test_gpt6_1_sol_row_turns_none_and_minimal_off():
"""Pinned on purpose so the row-derived test above cannot go vacuous: OpenAI lists low, medium,
high, xhigh and max for gpt-6.1-sol and neither none nor minimal
(https://developers.openai.com/api/docs/models/gpt-6.1-sol, read 2026-09-30 via #43910), and a
local proxy got OpenAI's 400 for both on 2026-09-30. When OpenAI adds a level, flip the row and
this assertion together."""
row: Final = litellm.model_cost["gpt-6.1-sol"]
assert row["supports_none_reasoning_effort"] is False
assert row["supports_minimal_reasoning_effort"] is False
def _row_without_effort_flags(row: Mapping[str, object]) -> dict[str, object]:
return {
field: value
for field, value in row.items()
if not (field.startswith("supports_") and field.endswith("_reasoning_effort"))
}
@pytest.mark.parametrize(
("flags", "level", "forwarded"),
[
({}, "none", True),
({"supports_none_reasoning_effort": False}, "none", False),
({"supports_none_reasoning_effort": True}, "none", True),
({"supports_minimal_reasoning_effort": False}, "minimal", False),
({"supports_max_reasoning_effort": False}, "max", True),
({}, "max", True),
({}, "xhigh", False),
({"supports_xhigh_reasoning_effort": True}, "xhigh", True),
({"supports_none_reasoning_effort": False}, "medium", True),
],
)
def test_the_row_flags_alone_decide_whether_a_chat_level_is_forwarded(
config: OpenAIConfig, monkeypatch: pytest.MonkeyPatch, flags: dict[str, bool], level: str, forwarded: bool
):
"""Same model name, rewritten row: the gate flips with the flags and nothing else, so breaking the
gate fails here while the pinned test above catches the real row losing a flag."""
monkeypatch.setitem(
litellm.model_cost, "gpt-6.1-sol", {**_row_without_effort_flags(litellm.model_cost["gpt-6.1-sol"]), **flags}
)
dropped: Final = config.map_openai_params(
non_default_params={"reasoning_effort": level},
optional_params={},
model="gpt-6.1-sol",
drop_params=True,
)
assert ("reasoning_effort" in dropped) is forwarded
if not forwarded:
with pytest.raises(litellm.utils.UnsupportedParamsError):
config.map_openai_params(
non_default_params={"reasoning_effort": level},
optional_params={},
model="gpt-6.1-sol",
drop_params=False,
)
def test_gpt5_minimal_dict_triggers_validation(config: OpenAIConfig):
"""Dict with effort='minimal' triggers minimal model-support validation."""
with pytest.raises(litellm.utils.UnsupportedParamsError):

23
tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py
View file

@ -516,3 +516,26 @@ def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(co
assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"}
assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"}
assert payload["messages"][2] == {"role": "user", "content": "a plain string message"}
@pytest.mark.usefixtures("local_beta_headers_config")
@pytest.mark.parametrize(
("client_headers", "expected_beta"),
[
({}, None),
({"anthropic-beta": "dangerous-tool-use-2026-09-03"}, "dangerous-tool-use-2026-09-03"),
],
)
def test_safeguards_reach_an_openai_like_host_with_only_the_client_betas(
config: OpenAILikeAnthropicMessagesConfig, client_headers: dict[str, str], expected_beta: str | None
) -> None:
headers, _ = config.validate_anthropic_messages_environment(
headers=client_headers,
model="some-model",
messages=[{"role": "user", "content": "hi"}],
optional_params={"safeguards": ({"type": "dangerous_tool_use"},)},
litellm_params={},
api_key="sk-test",
api_base="https://host/v1",
)
assert headers.get("anthropic-beta") == expected_beta, headers

726
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -9,12 +9,16 @@ ARN unified_object_id) batches with no managed unified id.
import asyncio
import json
from contextlib import contextmanager
from typing import TYPE_CHECKING
from copy import copy
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from litellm.batches.batch_utils import BatchCostUsageResult
@ -207,7 +211,7 @@ class TestCheckBatchCost:
find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
assert find_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
assert find_call[1]["order"] == {"created_at": "asc"}
assert find_call[1]["order"] == [{"created_at": "asc"}, {"id": "asc"}]
not_in = find_call[1]["where"]["status"]["not_in"]
assert "stale_expired" in not_in
# "complete"/"completed" are intentionally NOT excluded from the
@ -230,13 +234,13 @@ class TestCheckBatchCost:
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
# First find_many (primary query) raises with a schema error; second (fallback) returns empty
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
side_effect=[Exception("column batch_processed does not exist"), []]
side_effect=[Exception("column batch_processed does not exist"), [], []]
)
await check_batch_cost_instance.check_batch_cost()
calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
assert len(calls) == 2
assert len(calls) == 3
fallback_where = calls[1][1]["where"]
assert "batch_processed" not in fallback_where
assert "stale_expired" in fallback_where["status"]["not_in"]
@ -257,8 +261,7 @@ class TestCheckBatchCost:
await check_batch_cost_instance.check_batch_cost()
# Only one find_many call — the fallback directly, no primary query attempt
assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1
assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 2
fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"]
assert "batch_processed" not in fallback_where
@ -666,10 +669,16 @@ class TestCheckBatchCost:
assert check_batch_cost_instance._has_batch_processed_column is True
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
mock_response = MagicMock()
mock_response.status = "completed"
mock_response.output_file_id = "file-output-123"
mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
mock_response: Final = LiteLLMBatch(
id="batch-1",
status="completed",
output_file_id="file-output-123",
created_at=1,
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-input-1",
object="batch",
)
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
@ -728,10 +737,10 @@ class TestCheckBatchCost:
await check_batch_cost_instance.check_batch_cost()
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
"Expected update() to be called exactly once for the completed job"
)
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
assert update_data["batch_processed"] is True, (
"update() must include batch_processed=True when column is present"
)
@ -861,7 +870,7 @@ class TestCheckBatchCost:
mock_update_database.assert_awaited_once()
assert mock_update_database.call_args.kwargs["response_cost"] == 0.01
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
"the job must still be marked processed once cost tracking succeeds"
)
@ -963,10 +972,10 @@ class TestCheckBatchCost:
await check_batch_cost_instance.check_batch_cost()
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
f"Expected update() to be called exactly once for a {terminal_status} job"
)
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
assert update_data["status"] == terminal_status
assert update_data["batch_processed"] is True, (
"terminal-status update() must set batch_processed=True so polling stops"
@ -1069,9 +1078,13 @@ class TestCheckBatchCost:
assert owner.user_id == "user-1"
assert owner.team_id == "team-1"
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args
assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1"}
assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2
update_call = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
assert update_call.kwargs["where"] == {
"id": "job-terminal-mint-1",
"batch_processed": False,
"updated_at": mock_job.updated_at,
}
update_data = update_call.kwargs["data"]
assert update_data["status"] == terminal_status
assert update_data["batch_processed"] is True
@ -1136,10 +1149,10 @@ class TestCheckBatchCost:
) as mock_afile_content:
await check_batch_cost_instance.check_batch_cost()
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
"a completed batch with no output file must be marked processed exactly once"
)
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
assert update_data["status"] == completed_status
assert update_data["batch_processed"] is True, (
"completed-without-output update() must set batch_processed=True so polling stops"
@ -1283,11 +1296,16 @@ class TestCheckBatchCost:
assert check_batch_cost_instance._has_batch_processed_column is True
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
mock_response = MagicMock()
mock_response.status = terminal_status
mock_response.output_file_id = "file-output-123"
mock_response.error_file_id = None
mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}'
mock_response: Final = LiteLLMBatch(
id="batch-1",
status=terminal_status,
output_file_id="file-output-123",
created_at=1,
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-input-1",
object="batch",
)
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
@ -1350,8 +1368,8 @@ class TestCheckBatchCost:
f"{terminal_status} batch with an output file must fetch results and be billed"
)
mock_logging_obj.async_success_handler.assert_awaited_once()
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2
update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
assert update_data["batch_processed"] is True
assert update_data["status"] == terminal_status, (
f"billed {terminal_status} batch must keep its real terminal status in the DB"
@ -1523,10 +1541,10 @@ class TestCheckBatchCost:
assert mock_afile_content.await_count == 1
mock_calculate.assert_not_awaited()
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, (
"a terminal batch with a 404ing output file must be retired, not retried forever"
)
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
assert update_data["status"] == "failed"
assert update_data["batch_processed"] is True
@ -1845,13 +1863,16 @@ class TestUnmanagedVertexRouting:
router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash"
router.get_model_ids.return_value = ["deploy-1"]
mock_response = MagicMock()
mock_response.status = "completed"
mock_response.output_file_id = "gs://bucket/out/predictions.jsonl"
mock_response.error_file_id = None
mock_response.completed_at = None
mock_response.created_at = None
mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}'
mock_response: Final = LiteLLMBatch(
id="8823717160934178816",
status="completed",
output_file_id="gs://bucket/out/predictions.jsonl",
created_at=1,
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-input-1",
object="batch",
)
router.aretrieve_batch = AsyncMock(return_value=mock_response)
router.get_deployment_credentials_with_provider = MagicMock(
return_value={"vertex_project": "p", "vertex_location": "us-central1"}
@ -1919,8 +1940,8 @@ class TestUnmanagedVertexRouting:
mock_logging_obj.async_success_handler.assert_awaited_once()
assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01
assert prisma.db.litellm_managedobjecttable.update.call_count == 1
update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"]
assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2
update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
assert update_data["batch_processed"] is True
assert update_data["status"] == "complete"
@ -2063,13 +2084,16 @@ class TestUnmanagedBedrockRouting:
router.resolve_model_name_from_model_id.return_value = "claude-sonnet-4"
router.get_model_ids.return_value = ["deploy-1"]
mock_response = MagicMock()
mock_response.status = "completed"
mock_response.output_file_id = "s3://bucket/out/predictions.jsonl"
mock_response.error_file_id = None
mock_response.completed_at = None
mock_response.created_at = None
mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}'
mock_response: Final = LiteLLMBatch(
id=self._ARN,
status="completed",
output_file_id="s3://bucket/out/predictions.jsonl",
created_at=1,
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-input-1",
object="batch",
)
router.aretrieve_batch = AsyncMock(return_value=mock_response)
router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"})
@ -2133,8 +2157,8 @@ class TestUnmanagedBedrockRouting:
mock_logging_obj.async_success_handler.assert_awaited_once()
assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02
assert prisma.db.litellm_managedobjecttable.update.call_count == 1
update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"]
assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2
update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"]
assert update_data["batch_processed"] is True
assert update_data["status"] == "complete"
@ -2695,9 +2719,10 @@ class TestPollPageStarvation:
await self._instance(prisma, llm_router).check_batch_cost()
llm_router.aretrieve_batch.assert_not_awaited()
prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
call = prisma.db.litellm_managedobjecttable.update.call_args[1]
assert call["where"] == {"id": "job-no-model"}
prisma.db.litellm_managedobjecttable.update_many.assert_awaited()
call = prisma.db.litellm_managedobjecttable.update_many.call_args[1]
assert call["where"]["id"] == "job-no-model"
assert call["where"]["batch_processed"] is False
assert call["data"] == {"batch_processed": True}
@pytest.mark.asyncio
@ -2725,8 +2750,8 @@ class TestPollPageStarvation:
await self._instance(prisma, llm_router).check_batch_cost()
prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True}
prisma.db.litellm_managedobjecttable.update_many.assert_awaited()
assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"batch_processed": True}
@pytest.mark.asyncio
async def test_provider_404_with_deployment_gone_keeps_job(self):
@ -2785,24 +2810,7 @@ class TestPollPageStarvation:
await instance.check_batch_cost()
assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"}
@pytest.mark.asyncio
async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self):
"""A row already in a terminal status is never rewritten by the staleness sweep, so
it needs its own bound or it starves newer batches indefinitely."""
prisma = self._prisma([])
await self._instance(prisma, MagicMock()).check_batch_cost()
calls = prisma.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 2, "expected the staleness sweep plus the never-costed sweep"
where = calls[1][1]["where"]
assert where["file_purpose"] == "batch"
assert where["batch_processed"] is False
assert where["status"] == {"in": ["complete", "completed"]}
assert "created_at" in where
assert calls[1][1]["data"] == {"batch_processed": True}
assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"status": "stale_expired"}
@pytest.mark.asyncio
async def test_newer_batch_is_polled_once_dead_rows_are_retired(self):
@ -2840,7 +2848,11 @@ class TestPollPageStarvation:
await self._instance(prisma, llm_router).check_batch_cost()
retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list]
retired = [
call[1]["where"]["id"]
for call in prisma.db.litellm_managedobjecttable.update_many.call_args_list
if "id" in call[1]["where"]
]
assert retired == ["job-no-model", "job-gone"]
assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", (
"the newer healthy batch must still be polled in the same cycle"
@ -2888,7 +2900,8 @@ class _FakeManagedObjectRow:
self.team_id = None
self.api_key = None
self.request_tags = None
self.created_at = 1700000000
self.created_at = datetime(9999, 1, 1, tzinfo=timezone.utc)
self.updated_at = datetime(2000, 1, 1, tzinfo=timezone.utc)
self.file_object = json.dumps(
{
"id": "batch-456",
@ -2918,7 +2931,9 @@ class _FakeManagedObjectTable:
def _matches(self, where: dict) -> bool:
for key, value in where.items():
if key == "created_at":
return False
if "lt" in value and not self.row.created_at < value["lt"]:
return False
continue
if key == "status":
if self.row.status in value.get("not_in", []):
return False
@ -2931,7 +2946,9 @@ class _FakeManagedObjectTable:
async def _update_many(self, *, where: dict, data: dict) -> int:
if not self._matches(where):
return 0
if "batch_processed" in where:
if "updated_at" in where:
self.journal.append("finalize")
elif "batch_processed" in where:
self.journal.append("claim" if data.get("batch_processed") else "release")
for key, value in data.items():
setattr(self.row, key, value)
@ -2980,13 +2997,20 @@ class TestMultiPodBatchCostClaim:
@staticmethod
def _router():
response = MagicMock()
response.status = "completed"
response.output_file_id = _CLAIM_OUTPUT_FILE_ID
response.error_file_id = None
response.created_at = 1
response.completed_at = 2
response.model_dump_json.return_value = '{"id":"batch-456","status":"completed"}'
from litellm.types.utils import LiteLLMBatch
response: Final = LiteLLMBatch(
id="batch-456",
status="completed",
output_file_id=_CLAIM_OUTPUT_FILE_ID,
error_file_id=None,
created_at=1,
completed_at=2,
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-input-1",
object="batch",
)
deployment = MagicMock()
deployment.litellm_params.custom_llm_provider = "openai"
@ -3055,6 +3079,351 @@ class TestMultiPodBatchCostClaim:
):
yield logging_obj
@pytest.mark.asyncio
@pytest.mark.parametrize("failure_phase", ("before_write", "after_write", "external"))
@pytest.mark.parametrize("collector_enabled", (False, True))
async def test_builtin_accounting_retries_only_before_a_possible_charge(
self, failure_phase: str, collector_enabled: bool
) -> None:
import base64
from datetime import datetime
import httpx
import litellm
import respx
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.utils import PrismaClient
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.utils import SpecialEnums
class AccountingDependency:
def __init__(self) -> None:
self.failure_phase = failure_phase
self.writer: Final = DBSpendUpdateWriter()
self.failed_tracking_alert = AsyncMock()
self.slack_alerting_instance = MagicMock()
self.slack_alerting_instance.customer_spend_alert = AsyncMock(
side_effect=ConnectionError("Alert destination unavailable after charge")
if failure_phase == "after_write"
else None
)
@property
def db_spend_update_writer(self) -> DBSpendUpdateWriter:
if self.failure_phase == "before_write":
raise ConnectionError("Accounting dependency temporarily unavailable")
return self.writer
class ExternalCallback(CustomLogger):
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
if failure_phase == "external":
raise ConnectionError("External destination unavailable")
budget_cache: Final = DualCache()
RouterBudgetLimiting(
dual_cache=budget_cache,
provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}},
model_list=[
{
"model_name": "batch-alias",
"litellm_params": {"model": "openai/batch-test-model", "max_budget": 100, "budget_duration": "1d"},
"model_info": {"id": "model-123"},
}
],
)
dependency: Final = AccountingDependency()
spend_prisma: Final = PrismaClient(
database_url="postgresql://unused:unused@localhost:5432/unused", proxy_logging_obj=dependency
)
create_spend_rows: Final = AsyncMock(return_value=1)
input_rate: Final = 0.000002
litellm.register_model(
{
"batch-test-model": {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token_batches": input_rate,
"output_cost_per_token_batches": 0.000006,
}
}
)
content: Final = json.dumps(
{
"response": {
"status_code": 200,
"body": {
"model": "batch-test-model",
"usage": {"prompt_tokens": 100, "completion_tokens": 0, "total_tokens": 100},
},
},
"error": None,
}
).encode()
row: Final = _FakeManagedObjectRow()
original_status: Final = row.status
original_file_object: Final = row.file_object
row.unified_object_id = (
base64.urlsafe_b64encode(
SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode()
)
.decode()
.rstrip("=")
)
prisma: Final = self._prisma(row, [])
router: Final = self._router()
router.get_deployment.return_value = Deployment(
model_name="batch-alias",
litellm_params=LiteLLM_Params(model="openai/batch-test-model"),
model_info=ModelInfo(id="model-123"),
)
from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer
producer: Final = MagicMock(spec=SpendEventProducer) if collector_enabled else None
litellm.logging_callback_manager.add_litellm_async_success_callback(
ProxyDBLogger(spend_event_producer=producer)
)
litellm.logging_callback_manager.add_litellm_async_success_callback(ExternalCallback())
with (
respx.mock,
patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency),
patch("litellm.proxy.proxy_server.prisma_client", spend_prisma),
patch.object(
spend_prisma.writer_db._original_prisma,
"litellm_spendlogs",
MagicMock(create_many=create_spend_rows),
),
):
respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock(
return_value=httpx.Response(200, content=content)
)
worker: Final = self._instance(prisma, router)
await worker.check_batch_cost()
assert row.batch_processed is (failure_phase != "before_write")
if failure_phase == "before_write":
assert row.status == original_status
assert json.loads(row.file_object) == original_file_object
assert create_spend_rows.await_count == (0 if failure_phase == "before_write" else 1)
dependency.failure_phase = "recovered"
await worker.check_batch_cost()
await worker.check_batch_cost()
assert row.batch_processed is True
assert await budget_cache.async_get_cache("provider_spend:openai:1d") == pytest.approx(100 * input_rate)
assert await budget_cache.async_get_cache("deployment_spend:model-123:1d") == pytest.approx(100 * input_rate)
create_spend_rows.assert_awaited_once()
assert create_spend_rows.call_args.kwargs["data"][0]["spend"] == pytest.approx(100 * input_rate)
if producer is not None:
assert all(
json.loads(call.args[0])["call_type"] != "aretrieve_batch" for call in producer.publish.call_args_list
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("deployment_model", "successes", "failures", "expected_model"),
(
("openai/", 0, 2, "openai/"),
("", 0, 2, "model-123"),
("openai/", 2, 1, "batch-test-model"),
("openai/batch-test-model", 2, 1, "batch-test-model"),
("openai/batch-test-model", 2, 0, "batch-test-model"),
),
)
@pytest.mark.parametrize("failure_envelope", (None, 1, "invalid", [1]))
async def test_batch_accounting_preserves_models_and_mixed_results(
self, deployment_model: str, successes: int, failures: int, expected_model: str, failure_envelope: object
) -> None:
import base64
from datetime import datetime
import httpx
import litellm
import respx
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.utils import LiteLLMBatch, SpecialEnums
input_rate: Final = 0.000002
output_rate: Final = 0.000006
litellm.register_model(
{
"batch-test-model": {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token_batches": input_rate,
"output_cost_per_token_batches": output_rate,
}
}
)
success: Final = {
"response": {
"status_code": 200,
"body": {
"model": "batch-test-model",
"usage": {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120},
},
},
"error": None,
}
failure: Final = {"response": failure_envelope, "error": {"code": "batch_expired", "message": "expired"}}
rows: Final = (success,) * successes + (failure,) * failures
content: Final = b"\n".join(json.dumps(row).encode() for row in rows)
class BatchSpendRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.calls: tuple[dict[str, object], ...] = ()
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
if kwargs["call_type"] == "aretrieve_batch":
self.calls += (kwargs,)
callback: Final = BatchSpendRecorder()
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
row: Final = _FakeManagedObjectRow()
row.unified_object_id = (
base64.urlsafe_b64encode(
SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode()
)
.decode()
.rstrip("=")
)
prisma: Final = self._prisma(row, [])
router: Final = self._router()
router.get_deployment.return_value = Deployment(
model_name="batch-alias",
litellm_params=LiteLLM_Params(model=deployment_model, custom_llm_provider="openai"),
model_info=ModelInfo(id="model-123"),
)
router.aretrieve_batch.return_value = LiteLLMBatch(
id="batch-456",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="file-input-1",
output_file_id=_CLAIM_OUTPUT_FILE_ID,
object="batch",
status="completed",
)
router.get_deployment_credentials_with_provider.return_value = {
"api_key": "synthetic-key",
"custom_llm_provider": "openai",
}
with respx.mock:
respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock(
return_value=httpx.Response(200, content=content)
)
instance: Final = self._instance(prisma, router)
await instance.check_batch_cost()
await instance.check_batch_cost()
assert row.batch_processed is True
assert row.status == "complete"
router.aretrieve_batch.assert_awaited_once()
assert len(callback.calls) == 1
payload: Final = callback.calls[0]["standard_logging_object"]
assert payload["model"] == expected_model
assert payload["response_cost"] == pytest.approx(successes * (100 * input_rate + 20 * output_rate))
assert payload["prompt_tokens"] == successes * 100
assert payload["completion_tokens"] == successes * 20
assert payload["hidden_params"]["batch_successful_requests"] == successes
assert payload["hidden_params"]["batch_failed_requests"] == failures
assert payload["hidden_params"]["batch_models"] == (
["batch-test-model"] if deployment_model == "openai/batch-test-model" else ["batch-test-model"] * successes
)
@pytest.mark.asyncio
@pytest.mark.parametrize("error_source", ("batch", "output"))
async def test_a_404_for_another_identifier_preserves_recovery(self, error_source: str) -> None:
import base64
import httpx
import litellm
import respx
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import SpecialEnums
class SpendRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.costs: tuple[float, ...] = ()
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
if kwargs["call_type"] == "aretrieve_batch":
cost: Final = kwargs["response_cost"]
assert isinstance(cost, float)
self.costs += (cost,)
recorder: Final = SpendRecorder()
litellm.logging_callback_manager.add_litellm_async_success_callback(recorder)
litellm.register_model(
{
"batch-recovery-model": {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token_batches": 0.01,
"output_cost_per_token_batches": 0.02,
}
}
)
row: Final = _FakeManagedObjectRow()
row.unified_object_id = (
base64.urlsafe_b64encode(
SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode()
)
.decode()
.rstrip("=")
)
router: Final = self._router()
response: Final = router.aretrieve_batch.return_value
error: Final = litellm.NotFoundError(
message=f"No object found with id '{'batch-456' if error_source == 'batch' else _CLAIM_OUTPUT_FILE_ID}0'.",
model="gpt-4",
llm_provider="openai",
)
router.aretrieve_batch.side_effect = (error, response) if error_source == "batch" else (response, response)
prisma: Final = self._prisma(row, [])
output: Final = json.dumps(
{
"response": {
"status_code": 200,
"body": {
"model": "batch-recovery-model",
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
},
},
"error": None,
}
).encode()
with respx.mock:
route: Final = respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content")
route.mock(
side_effect=error if error_source == "output" else None,
return_value=None if error_source == "output" else httpx.Response(200, content=output),
)
await self._instance(prisma, router).check_batch_cost()
assert row.batch_processed is False
assert recorder.costs == ()
route.mock(return_value=httpx.Response(200, content=output))
await self._instance(prisma, router).check_batch_cost()
await self._instance(prisma, router).check_batch_cost()
assert row.batch_processed is True
assert row.status == "complete"
assert recorder.costs == (pytest.approx(10 * 0.01 + 5 * 0.02),)
@staticmethod
def _claim_calls(prisma) -> list:
return [
@ -3078,6 +3447,30 @@ class TestMultiPodBatchCostClaim:
with patch("litellm.proxy.proxy_server.scheduler", scheduler):
await guard._check_file_deletion_allowed(file_id)
@pytest.mark.asyncio
@pytest.mark.parametrize("status", ("completed", "failed", "expired", "cancelled"))
@pytest.mark.parametrize("metadata_text", ("before\x00after", r"\u0000"))
async def test_cost_claim_persists_terminal_metadata_without_a_later_write(
self, status: str, metadata_text: str
) -> None:
row: Final = _FakeManagedObjectRow()
journal: Final = []
prisma: Final = self._prisma(row, journal)
router: Final = self._router()
router.aretrieve_batch.return_value = router.aretrieve_batch.return_value.model_copy(
update={"status": status, "metadata": {"annotation": metadata_text}}
)
prisma.db.litellm_managedobjecttable.update.side_effect = RuntimeError("final write unavailable")
worker: Final = self._instance(prisma, router)
with self._billing_patches(journal) as logging_obj:
await worker.check_batch_cost()
await worker.check_batch_cost()
assert row.batch_processed is True
assert row.status == ("complete" if status == "completed" else status)
assert json.loads(row.file_object)["status"] == status
assert json.loads(row.file_object)["metadata"] == {"annotation": metadata_text.replace("\x00", "")}
assert logging_obj.async_success_handler.await_count == 1
@pytest.mark.asyncio
async def test_winning_pod_claims_the_row_between_fetching_and_billing(self):
"""The claim flips batch_processed false -> true after the results are in hand and
@ -3089,13 +3482,9 @@ class TestMultiPodBatchCostClaim:
with self._billing_patches(journal) as logging_obj:
await self._instance(prisma, self._router()).check_batch_cost()
assert journal == ["fetch", "claim", "bill", "finalize"]
assert self._claim_calls(prisma) == [
{
"where": {"id": "job-claim-1", "batch_processed": False},
"data": {"batch_processed": True},
}
]
assert journal == ["fetch", "claim", "bill"]
assert row.status == "complete"
assert json.loads(row.file_object)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5}
logging_obj.async_success_handler.assert_awaited_once()
assert row.batch_processed is True
@ -3115,30 +3504,28 @@ class TestMultiPodBatchCostClaim:
assert journal == ["fetch"]
logging_obj.async_success_handler.assert_not_awaited()
assert self._claim_calls(prisma) == [
{
"where": {"id": "job-claim-1", "batch_processed": False},
"data": {"batch_processed": True},
}
]
assert row.status == "in_progress"
assert json.loads(row.file_object)["status"] == "in_progress"
@pytest.mark.asyncio
async def test_a_failed_spend_log_write_releases_the_claim(self):
"""A transient failure while billing a claimed batch must hand the row back, or its
spend is silently lost instead of being retried on the next cycle."""
row = _FakeManagedObjectRow()
journal = []
prisma = self._prisma(row, journal)
@pytest.mark.parametrize("stored_as_json_string", (True, False))
async def test_an_unclassified_logging_failure_keeps_the_claim(self, stored_as_json_string: bool) -> None:
row: Final = _FakeManagedObjectRow()
row.file_object = row.file_object if stored_as_json_string else json.loads(row.file_object)
journal: Final = []
prisma: Final = self._prisma(row, journal)
with self._billing_patches(journal, bill_error=Exception("spend log write failed")):
with self._billing_patches(journal, bill_error=Exception("Accounting outcome unknown")):
await self._instance(prisma, self._router()).check_batch_cost()
assert journal == ["fetch", "claim", "bill", "release"]
assert row.batch_processed is False
assert self._claim_calls(prisma)[-1] == {
"where": {"id": "job-claim-1", "batch_processed": True},
"data": {"batch_processed": False},
}
assert journal == ["fetch", "claim", "bill"]
assert row.batch_processed is True
with self._billing_patches(journal) as logging_obj:
await self._instance(prisma, self._router()).check_batch_cost()
await self._instance(prisma, self._router()).check_batch_cost()
logging_obj.async_success_handler.assert_not_awaited()
@pytest.mark.asyncio
async def test_a_worker_interrupted_mid_costing_leaves_the_batch_billable(self):
@ -3170,7 +3557,7 @@ class TestMultiPodBatchCostClaim:
with self._billing_patches(survivor_journal) as survivor_logging:
await self._instance(survivor_prisma, self._router()).check_batch_cost()
assert survivor_journal == ["fetch", "claim", "bill", "finalize"]
assert survivor_journal == ["fetch", "claim", "bill"]
survivor_logging.async_success_handler.assert_awaited_once()
assert row.batch_processed is True
@ -3201,7 +3588,7 @@ class TestMultiPodBatchCostClaim:
finish_fetch.set()
await asyncio.wait_for(costing, timeout=5)
assert journal == ["fetch", "claim", "bill", "finalize"]
assert journal == ["fetch", "claim", "bill"]
assert row.batch_processed is True
await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID)
@ -3221,3 +3608,124 @@ class TestMultiPodBatchCostClaim:
assert self._claim_calls(prisma) == []
assert journal == ["fetch", "bill", "finalize"]
logging_obj.async_success_handler.assert_awaited_once()
@pytest.mark.asyncio
async def test_old_completed_batch_remains_available_for_costing() -> None:
row: Final = _FakeManagedObjectRow()
row.status = "completed"
row.created_at = datetime(2000, 1, 1, tzinfo=timezone.utc)
prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
await poller._cleanup_stale_managed_objects()
assert row.batch_processed is False
assert row.status == "completed"
@pytest.mark.asyncio
@pytest.mark.parametrize("progress", ["claimed", "refreshed"])
async def test_stale_retirement_preserves_concurrent_progress(progress: str) -> None:
row: Final = _FakeManagedObjectRow()
selected: Final = copy(row)
if progress == "claimed":
row.batch_processed = True
row.updated_at = datetime(2000, 1, 2, tzinfo=timezone.utc)
prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
await poller._retire_job(selected, "provider batch disappeared")
assert row.batch_processed is (progress == "claimed")
assert row.updated_at > selected.updated_at
@pytest.mark.asyncio
@pytest.mark.parametrize("progress", ["claimed", "refreshed"])
async def test_stale_terminal_response_preserves_concurrent_progress(progress: str) -> None:
from litellm.types.utils import LiteLLMBatch
row: Final = _FakeManagedObjectRow()
selected: Final = copy(row)
if progress == "claimed":
row.batch_processed = True
row.updated_at = datetime(2000, 1, 2, tzinfo=timezone.utc)
row.file_object = '{"revision":"fresh"}'
prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
response: Final = LiteLLMBatch(
id="batch-456",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="file-input-1",
object="batch",
status="failed",
)
await poller._finalize_unbilled_terminal_job(selected, response)
assert row.file_object == '{"revision":"fresh"}'
assert row.status == "in_progress"
assert row.batch_processed is (progress == "claimed")
@pytest.mark.asyncio
@pytest.mark.parametrize("legacy", [False, True])
async def test_permanently_lost_batch_retires_unchanged_row(legacy: bool) -> None:
row: Final = _FakeManagedObjectRow()
prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
poller._has_batch_processed_column = not legacy
await poller._retire_job(row, "provider batch disappeared")
assert row.batch_processed is (not legacy)
assert row.status == ("stale_expired" if legacy else "in_progress")
@pytest.mark.asyncio
async def test_failed_retirement_keeps_batch_available() -> None:
row: Final = _FakeManagedObjectRow()
prisma: Final = TestMultiPodBatchCostClaim._prisma(row, [])
prisma.db.litellm_managedobjecttable.update_many = AsyncMock(side_effect=RuntimeError("database unavailable"))
poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router())
await poller._retire_job(row, "provider batch disappeared")
assert row.batch_processed is False
assert row.status == "in_progress"
@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl"))
@pytest.mark.parametrize("other_id", ("{id}4", "prefix-{id}", "{id}/child", "{id}.backup", "{id}?generation=2"))
def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other_id: str) -> None:
from litellm.exceptions import NotFoundError
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
error: Final = NotFoundError(
message=f"No object found with id '{other_id.format(id=resource_id)}'.", model="test", llm_provider="openai"
)
assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is False
assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is False
@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl"))
@pytest.mark.parametrize("message", ("No object found with id '{id}'.", 'File "{id}" does not exist', "{id} not found"))
def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, message: str) -> None:
from litellm.exceptions import NotFoundError
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
error: Final = NotFoundError(message=message.format(id=resource_id), model="test", llm_provider="openai")
assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is True
assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True
@pytest.mark.parametrize("output_file_id", (None, ""))
def test_a_404_without_an_output_identity_remains_retryable(output_file_id: str | None) -> None:
from litellm.exceptions import NotFoundError
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
error: Final = NotFoundError(message="Output file not found", model="test", llm_provider="openai")
assert CheckBatchCost._is_output_file_gone_at_provider(error, output_file_id) is False

72
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -2,12 +2,15 @@ import asyncio
import json
import logging
from datetime import datetime, timedelta, timezone
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
@ -2916,3 +2919,72 @@ async def test_track_cost_callback_enqueue_emits_no_service_span(): # test-qual
assert update_database.await_count == 1, "the spend enqueue itself must still run"
assert success_hook.await_count == 0, [call.kwargs for call in success_hook.await_args_list]
@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=_CALL_START,
end_time=_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=_CALL_START,
end_time=_CALL_START,
)
assert reservation["finalized"] is True
assert outcome.retryable_failure is False

7
tests/unit/responses/test_responses_api_bridge_flag.py
View file

@ -356,11 +356,8 @@ class TestUseResponsesApiBridgeFlag:
assert upstream.call_count == 1
request_body: Final = json.loads(upstream.calls[0].request.read())
assert request_body["messages"] == [
{
"role": "system",
"content": [{"type": "text", "text": "You are terse."}, {"type": "text", "text": "Skills: none."}],
},
{"role": "user", "content": [{"type": "text", "text": "Hello"}]},
{"role": "system", "content": "You are terse.\n\nSkills: none."},
{"role": "user", "content": "Hello"},
]
assert "client_metadata" not in request_body
assert request_body["chat_template_kwargs"] == {"thinking": True}

89
tests/unit/router_strategy/test_budget_limiter.py
View file

@ -8,6 +8,7 @@ the callback raise before any spend was recorded, so those budgets never moved.
from typing import Final
import litellm
import pytest
from litellm.caching.caching import DualCache
@ -229,3 +230,91 @@ def cleanup_redis():
redis_client.delete(key)
except Exception as e:
print(f"Error cleaning up Redis: {str(e)}")
def _completed_budget_batch(batch_id: str):
from litellm.types.utils import LiteLLMBatch
return LiteLLMBatch(
id=batch_id,
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="file-budget",
object="batch",
status="completed",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("call_type", ("aretrieve_batch", "retrieve_batch"))
async def test_batch_polls_charge_each_router_budget_once(monkeypatch, call_type):
monkeypatch.setattr(litellm, "tag_budget_config", {"release": {"budget_limit": 100, "time_period": "1d"}})
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
cache: Final = DualCache()
limiter: Final = RouterBudgetLimiting(
dual_cache=cache,
provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}},
model_list=[
{
"model_name": "batch-budget",
"litellm_params": {"model": "openai/batch-budget", "max_budget": 100, "budget_duration": "1d"},
"model_info": {"id": "deployment-1"},
}
],
)
kwargs: Final = {
**_success_kwargs(provider_in_litellm_params=None, provider_in_payload="openai", call_type=call_type),
"metadata": {"tags": ["release"]},
}
response: Final = _completed_budget_batch("batch-first")
await asyncio.gather(*(limiter.async_log_success_event(kwargs, response, None, None) for _ in range(10)))
keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d", "tag_spend:release:1d")
assert tuple([await cache.async_get_cache(key) for key in keys]) == (0.25, 0.25, 0.25)
await limiter.async_log_success_event(kwargs, _completed_budget_batch("batch-second"), None, None)
await limiter.async_log_success_event({**kwargs, "call_type": "acompletion"}, response, None, None)
await limiter.async_log_success_event({**kwargs, "call_type": "acompletion"}, response, None, None)
assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0)
await limiter.async_log_success_event({**kwargs, "call_type": "_arealtime"}, response, None, None)
assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0)
@pytest.mark.asyncio
async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation():
from fakeredis import FakeServer, FakeRedis as FakeSyncRedis
from redis.exceptions import ConnectionError as RedisConnectionError
from fakeredis.aioredis import FakeRedis
from litellm.caching.redis_cache import RedisCache
server: Final = FakeServer()
client: Final = FakeRedis(server=server)
sync_client: Final = FakeSyncRedis(server=server)
redis_cache: Final = RedisCache(host="router-batch-fixture", connection_pool=sync_client.connection_pool)
litellm.in_memory_llm_clients_cache.set_cache(key=redis_cache._get_async_client_cache_key(), value=client)
first: Final = RouterBudgetLimiting(
dual_cache=DualCache(redis_cache=redis_cache),
provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}},
)
second: Final = RouterBudgetLimiting(
dual_cache=DualCache(redis_cache=redis_cache),
provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}},
)
kwargs: Final = _success_kwargs(
provider_in_litellm_params=None, provider_in_payload="openai", call_type="aretrieve_batch"
)
response: Final = _completed_budget_batch("batch-shared")
await asyncio.gather(
first.async_log_success_event(kwargs, response, None, None),
second.async_log_success_event(kwargs, response, None, None),
)
await first._push_in_memory_increments_to_redis()
await second._push_in_memory_increments_to_redis()
assert float(await client.get("provider_spend:openai:1d")) == 0.25
server.connected = False
with pytest.raises(RedisConnectionError):
await second.async_log_success_event(kwargs, _completed_budget_batch("batch-unconfirmed"), None, None)
server.connected = True
await second.async_log_success_event(kwargs, _completed_budget_batch("batch-unconfirmed"), None, None)
await second._push_in_memory_increments_to_redis()
assert float(await client.get("provider_spend:openai:1d")) == 0.5
await client.aclose()

30
tests/unit/router_utils/test_batch_utils.py
View file

@ -2,7 +2,7 @@ import asyncio
import importlib
import json
from io import BytesIO
from typing import Dict, List
from typing import Dict, Final, List
import pytest
@ -346,3 +346,31 @@ def setup_and_teardown():
yield
loop.close()
asyncio.set_event_loop(None)
@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

15
tests/unit/router_utils/test_reasoning_effort_capability.py
View file

@ -395,22 +395,25 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels:
[
("azure/gpt-6-astra", "azure"),
("azure/us/gpt-6-astra", "azure"),
("azure/eu/gpt-6-astra", "azure"),
("azure_ai/gpt-6-astra", "azure_ai"),
],
)
def test_an_azure_hosted_deployment_advertises_none_but_not_max(
def test_an_azure_hosted_deployment_advertises_the_levels_both_its_routes_accept(
self, local_model_cost_map, model, custom_llm_provider
):
"""Microsoft hosts the same model with a different level set than OpenAI does. Verified live
on both Azure routes: none returns 200 with zero reasoning tokens and unlocks temperature,
which OpenAI's API rejects, while max returns 400 unsupported_value naming none through
xhigh as the levels it does take."""
"""Verified live on 2026-10-01 against a Foundry gpt-6-astra deployment (model
gpt-6-astra-2026-09-03): none returns 400 unsupported_value on chat and on Responses, naming
low, medium, high and xhigh as the levels it takes; Responses also honors max with reasoning
tokens billed while chat refuses it, as OpenAI's chat does. The flag is route-blind, so the
row keeps max off and the advertised set stops at the level every route accepts.
Earlier evidence (2026-09-05, commit e79f3ec5205) had none accepted on chat; the deployment
no longer behaves that way."""
from litellm.utils import get_model_info_helper
model_info = dict(get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider))
assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == (
"none",
"low",
"medium",
"high",

622
tests/unit/test_cost_calculator.py
View file

@ -4,7 +4,10 @@ from collections.abc import Mapping
from types import MappingProxyType, SimpleNamespace
from typing import Final, Literal, cast
import json
import httpx
import respx
import pytest
from openai.types.completion_usage import CompletionUsage
from pydantic import BaseModel
@ -1024,7 +1027,7 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero():
{
"model_name": "semantic-ranker-default-004",
"litellm_params": {
"model": "vertex_ai/semantic-ranker-default-004",
"model": "vertex_ai/semantic-ranker-custom-test",
"vertex_project": "test-project",
"vertex_location": "us-east5",
},
@ -1033,7 +1036,7 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero():
]
)
router_model_id: Final = router.model_list[0]["model_info"]["id"]
assert litellm.model_cost["vertex_ai/semantic-ranker-default-004"].get("input_cost_per_query") is None
assert litellm.model_cost.get("vertex_ai/semantic-ranker-custom-test", {}).get("input_cost_per_query") is None
response: Final = RerankResponse(
id="vertex_ai_rerank_test",
@ -1043,7 +1046,7 @@ def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero():
cost: Final = completion_cost(
completion_response=response,
model="vertex_ai/semantic-ranker-default-004",
model="vertex_ai/semantic-ranker-custom-test",
custom_llm_provider="vertex_ai",
call_type="arerank",
custom_pricing=True,
@ -8044,3 +8047,616 @@ def test_batch_cost_calculator():
cost = completion_cost(**args)
assert cost > 0
@pytest.fixture
def isolated_pricing_registry(local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(litellm, "open_ai_chat_completion_models", set(litellm.open_ai_chat_completion_models))
monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider))
_CUSTOM_PRICING_MODEL: Final = "migration-custom-pricing-chat"
_CUSTOM_PRICING_RATES: Final = MappingProxyType(
{"input_cost_per_token": 0.0000008, "output_cost_per_token": 0.0000032}
)
_CUSTOM_PRICING_EXPECTED_COST: Final = 7 * 0.0000008 + 4 * 0.0000032
def _custom_pricing_chunk(choices: list[dict[str, object]], usage: dict[str, int] | None = None) -> str:
chunk: Final = {
"id": "chatcmpl-custom-pricing",
"object": "chat.completion.chunk",
"created": 1,
"model": _CUSTOM_PRICING_MODEL,
"choices": choices,
**({"usage": usage} if usage is not None else {}),
}
return f"data: {json.dumps(chunk)}\n\n"
def _custom_pricing_upstream(request: httpx.Request) -> httpx.Response:
if json.loads(request.content).get("stream") is True:
body: Final = "".join(
(
_custom_pricing_chunk(
[{"index": 0, "delta": {"role": "assistant", "content": "I am ready to help."}, "finish_reason": None}]
),
_custom_pricing_chunk([{"index": 0, "delta": {}, "finish_reason": "stop"}]),
_custom_pricing_chunk([], {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11}),
"data: [DONE]\n\n",
)
)
return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
return httpx.Response(
200,
json={
"id": "chatcmpl-custom-pricing",
"object": "chat.completion",
"created": 1,
"model": _CUSTOM_PRICING_MODEL,
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "I am ready to help."}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11},
},
)
@pytest.mark.parametrize("rate", (0.0, 0.001))
@pytest.mark.parametrize("cache_field", ("cache_read_input_token_cost", "cache_creation_input_token_cost"))
@pytest.mark.usefixtures("isolated_pricing_registry")
@respx.mock
def test_router_cache_only_pricing_preserves_uncached_token_charges(rate: float, cache_field: str) -> None:
litellm.register_model(
{_CUSTOM_PRICING_MODEL: {**_CUSTOM_PRICING_RATES, "litellm_provider": "openai", "mode": "chat"}},
persist_across_reloads=False,
)
respx.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_custom_pricing_upstream)
router: Final = litellm.Router(
model_list=[
{
"model_name": "cache-priced-deployment",
"litellm_params": {
"model": f"openai/{_CUSTOM_PRICING_MODEL}",
"api_key": "sk-test",
cache_field: rate,
},
"model_info": {"id": "cache-only-deployment"},
}
]
)
response: Final = router.completion(
model="cache-priced-deployment", messages=[{"role": "user", "content": "How are you?"}]
)
assert response._hidden_params["response_cost"] == pytest.approx(_CUSTOM_PRICING_EXPECTED_COST)
@pytest.mark.parametrize("profile_prefix", ("us", "global"))
@pytest.mark.parametrize(
"model_name",
("gpt-5.6-sol", "gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-sol", "gpt-6-luna", "gpt-6-astra", "gpt-6.1-sol"),
)
@pytest.mark.parametrize("arn", (False, True))
def test_bedrock_converse_cost_uses_invoked_profile(
_local_model_cost_map: None, model_name: str, profile_prefix: str, arn: bool
) -> None:
model: Final = f"openai.{model_name}"
profile: Final = f"{profile_prefix}.{model}"
invoked: Final = f"arn:aws:bedrock:us-east-1:123456789012:inference-profile/{profile}" if arn else profile
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock")
cost: Final = completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": invoked},
)
assert cost == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"])
assert response.model == model
@pytest.mark.parametrize("custom_pricing", (False, True))
def test_bedrock_converse_profile_preserves_explicit_base_and_custom_pricing(
_local_model_cost_map: None, custom_pricing: bool
) -> None:
model: Final = "openai.gpt-6-sol"
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
prices: Final = litellm.get_model_info(model, custom_llm_provider="bedrock")
cost: Final = completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
custom_pricing=custom_pricing,
base_model=None if custom_pricing else model,
)
assert cost == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"])
@pytest.mark.parametrize("stream", (False, True))
def test_bedrock_converse_cost_retains_profile_after_dispatch_consumes_optional_params(
_local_model_cost_map: None, stream: bool
) -> None:
from litellm.cost_calculator import pricing_entry_for_cost_calc
model: Final = "openai.gpt-5.6-sol"
profile: Final = f"us.{model}"
logging_obj: Final = Logging(
model=model,
messages=[{"role": "user", "content": "hi"}],
stream=stream,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="bedrock-profile-cost",
function_id="bedrock-profile-cost",
)
logging_obj.update_environment_variables(
litellm_params={}, optional_params={"model_id": profile}, custom_llm_provider="bedrock"
)
logging_obj.optional_params.pop("model_id")
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
prices: Final = litellm.get_model_info(profile, custom_llm_provider="bedrock")
assert logging_obj.response_cost_calculator(response) == pytest.approx(
10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
)
entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj)
assert entry is not None
assert entry[0] == prices["key"]
@pytest.mark.parametrize(
"invoked_model_id",
(None, "", 123, "unmapped-profile", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque"),
)
def test_bedrock_converse_cost_keeps_configured_rates_when_target_has_no_prices(
_local_model_cost_map: None, invoked_model_id: object
) -> None:
model: Final = "openai.gpt-6-sol"
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
expected: Final = completion_cost(completion_response=response, model=model, custom_llm_provider="bedrock")
assert (
completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
optional_params={"model_id": invoked_model_id},
)
== expected
)
def test_non_bedrock_cost_ignores_bedrock_model_id(_local_model_cost_map: None) -> None:
response: Final = ModelResponse(
model="gpt-6-sol", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
)
expected: Final = completion_cost(completion_response=response, custom_llm_provider="openai")
assert (
completion_cost(
completion_response=response,
custom_llm_provider="openai",
optional_params={"model_id": "us.openai.gpt-6-sol"},
)
== expected
)
@pytest.mark.parametrize("rate", (0.0, 0.001))
@pytest.mark.parametrize("persist_across_reloads", (True, False))
@pytest.mark.parametrize("routed_response", (True, False))
@pytest.mark.parametrize("metadata_prefix", ("bedrock/", "bedrock/converse/"))
def test_bedrock_converse_profile_preserves_registered_price_override(
_local_model_cost_map: None,
monkeypatch: pytest.MonkeyPatch,
rate: float,
persist_across_reloads: bool,
routed_response: bool,
metadata_prefix: str,
) -> None:
import litellm.utils as model_utils
from litellm.cost_calculator import pricing_entry_for_cost_calc
monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {})
catalog: Final = dict(litellm.model_cost)
model: Final = "openai.gpt-5.6-sol"
litellm.register_model(
{f"{metadata_prefix}{model}": {"supports_function_calling": True}},
persist_across_reloads=persist_across_reloads,
)
litellm.register_model(
{
model: {
"input_cost_per_token": rate,
"output_cost_per_token": rate,
"litellm_provider": "bedrock_converse",
"mode": "chat",
}
},
persist_across_reloads=persist_across_reloads,
)
litellm.register_model(
{model: {"supports_function_calling": True, "input_cost_per_token": None}},
persist_across_reloads=persist_across_reloads,
)
response: Final = ModelResponse(
model=f"bedrock/converse/{model}" if routed_response else model,
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="registered-pricing",
function_id="registered-pricing",
)
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
litellm_logging_obj=logging_obj,
) == pytest.approx(15 * rate)
assert logging_obj.billed_token_rates is not None
assert logging_obj.billed_token_rates.input_cost_per_token == rate
assert logging_obj.billed_token_rates.output_cost_per_token == rate
entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, None)
assert entry is not None
assert entry[0] == model
assert entry[1]["input_cost_per_token"] == rate
assert entry[1]["output_cost_per_token"] == rate
updated_rate: Final = rate + 0.002
litellm.register_model(
{model: {"input_cost_per_token": str(updated_rate)}}, persist_across_reloads=persist_across_reloads
)
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
) == pytest.approx(10 * updated_rate + 5 * rate)
monkeypatch.setattr(litellm, "model_cost", catalog)
model_utils._invalidate_model_cost_lowercase_map()
model_utils.reapply_runtime_model_cost_registrations()
profile_prices: Final = litellm.get_model_info(f"us.{model}", custom_llm_provider="bedrock")
expected_after_reload: Final = (
10 * updated_rate + 5 * rate
if persist_across_reloads
else 10 * profile_prices["input_cost_per_token"] + 5 * profile_prices["output_cost_per_token"]
)
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
) == pytest.approx(expected_after_reload)
def test_bedrock_converse_profile_ignores_registered_capabilities_without_prices(
_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
) -> None:
import litellm.utils as model_utils
monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {})
model: Final = "openai.gpt-6-sol"
litellm.register_model({model: {"supports_function_calling": True}})
response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
prices: Final = litellm.get_model_info(f"us.{model}", custom_llm_provider="bedrock")
assert completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
optional_params={"model_id": f"us.{model}"},
) == pytest.approx(10 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"])
@pytest.mark.parametrize("rate", (0.0, 0.001))
@pytest.mark.parametrize("cache_field", ("cache_read_input_token_cost", "cache_creation_input_token_cost"))
@pytest.mark.parametrize("has_base_rates", (False, True))
@pytest.mark.parametrize("persist_across_reloads", (False, True))
@pytest.mark.parametrize("profile_prefix", ("us", "global"))
@pytest.mark.parametrize("provider", (None, "bedrock"))
def test_bedrock_converse_profile_preserves_cache_only_registered_prices(
_local_model_cost_map: None,
monkeypatch: pytest.MonkeyPatch,
rate: float,
cache_field: str,
has_base_rates: bool,
persist_across_reloads: bool,
profile_prefix: str,
provider: str | None,
) -> None:
import litellm.utils as model_utils
from litellm.cost_calculator import pricing_entry_for_cost_calc
monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {})
model: Final = "openai.cache-overlay-test"
profile: Final = f"{profile_prefix}.{model}"
catalog: Final = {
**litellm.model_cost,
model: {
"litellm_provider": "bedrock",
"mode": "chat",
**({"input_cost_per_token": 0.01, "output_cost_per_token": 0.02} if has_base_rates else {}),
},
profile: {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.02,
"output_cost_per_token": 0.04,
"cache_read_input_token_cost": 0.005,
"cache_creation_input_token_cost": 0.025,
},
}
profile_before: Final = dict(catalog[profile])
monkeypatch.setattr(litellm, "model_cost", dict(catalog))
model_utils._invalidate_model_cost_lowercase_map()
litellm.register_model({model: {cache_field: rate}}, persist_across_reloads=persist_across_reloads)
response: Final = ModelResponse(
model=model,
usage=Usage(
prompt_tokens=15,
completion_tokens=3,
total_tokens=18,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=5 if cache_field == "cache_read_input_token_cost" else 0
),
cache_creation_input_tokens=5 if cache_field == "cache_creation_input_token_cost" else 0,
),
)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="cache-overlay",
function_id="cache-overlay",
)
logging_obj.model_call_details["model_id"] = profile
expected: Final = 10 * 0.02 + 3 * 0.04 + 5 * rate
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider=provider,
optional_params={"model_id": profile},
litellm_logging_obj=logging_obj,
) == pytest.approx(expected)
breakdown_key: Final = "cache_read_cost" if cache_field == "cache_read_input_token_cost" else "cache_creation_cost"
assert logging_obj.cost_breakdown is not None
assert logging_obj.cost_breakdown.get(breakdown_key, 0.0) == pytest.approx(5 * rate)
assert logging_obj.billed_token_rates is not None
assert (
logging_obj.billed_token_rates.cache_read_input_token_cost
if cache_field == "cache_read_input_token_cost"
else logging_obj.billed_token_rates.cache_creation_input_token_cost
) == rate
entry: Final = pricing_entry_for_cost_calc(
f"bedrock/converse/{model}", response, provider, False, None, None, None, logging_obj
)
assert entry is not None
assert entry[1]["input_cost_per_token"] == 0.02
assert entry[1]["output_cost_per_token"] == 0.04
assert entry[1][cache_field] == rate
assert litellm.model_cost[profile] == profile_before
monkeypatch.setattr(litellm, "model_cost", dict(catalog))
model_utils._invalidate_model_cost_lowercase_map()
model_utils.reapply_runtime_model_cost_registrations()
expected_after_reload: Final = (
10 * 0.02 + 3 * 0.04 + 5 * (rate if persist_across_reloads else catalog[profile][cache_field])
)
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider=provider,
optional_params={"model_id": profile},
litellm_logging_obj=logging_obj,
) == pytest.approx(expected_after_reload)
@pytest.mark.parametrize("profile", ("us.openai.missing-profile-test", "us.openai.unpriced-profile-test"))
def test_bedrock_cache_registration_with_unpriced_profile_keeps_registered_rates(
_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, profile: str
) -> None:
import litellm.utils as model_utils
model: Final = "openai.registered-cache-test"
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.model_cost,
model: {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
},
"us.openai.unpriced-profile-test": {"litellm_provider": "bedrock", "mode": "chat"},
},
)
model_utils._invalidate_model_cost_lowercase_map()
litellm.register_model({model: {"cache_read_input_token_cost": 0.001}}, persist_across_reloads=False)
response: Final = ModelResponse(
model=model,
usage=Usage(
prompt_tokens=15,
completion_tokens=3,
total_tokens=18,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5),
),
)
assert completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": profile},
) == pytest.approx(10 * 0.01 + 3 * 0.02 + 5 * 0.001)
@pytest.mark.parametrize("direction", ("input", "output"))
@pytest.mark.parametrize("rate", (0.0, 0.004))
@pytest.mark.parametrize("prompt_tokens", (100, 300_001))
def test_bedrock_partial_registration_overrides_only_its_profile_direction(
_local_model_cost_map: None,
monkeypatch: pytest.MonkeyPatch,
direction: str,
rate: float,
prompt_tokens: int,
) -> None:
from litellm.cost_calculator import pricing_entry_for_cost_calc
model: Final = "openai.partial-pricing-test"
profile: Final = f"us.{model}"
profile_entry: Final = {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"input_cost_per_token_above_300k_tokens": 0.03,
"output_cost_per_token_above_300k_tokens": 0.06,
"cache_read_input_token_cost": 0.001,
}
monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, profile: dict(profile_entry)})
litellm.register_model({model: {f"{direction}_cost_per_token": rate}}, persist_across_reloads=False)
response: Final = ModelResponse(
model=model,
usage=Usage(
prompt_tokens=prompt_tokens,
completion_tokens=3,
total_tokens=prompt_tokens + 3,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5),
),
)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="partial-pricing",
function_id="partial-pricing",
)
logging_obj.model_call_details["model_id"] = profile
input_rate: Final = rate if direction == "input" else (0.03 if prompt_tokens > 300_000 else 0.01)
output_rate: Final = rate if direction == "output" else (0.06 if prompt_tokens > 300_000 else 0.02)
assert completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
litellm_logging_obj=logging_obj,
) == pytest.approx((prompt_tokens - 5) * input_rate + 3 * output_rate + 5 * 0.001)
assert logging_obj.billed_token_rates is not None
assert logging_obj.billed_token_rates.input_cost_per_token == input_rate
assert logging_obj.billed_token_rates.output_cost_per_token == output_rate
assert logging_obj.cost_breakdown is not None
assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(5 * 0.001)
entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj)
assert entry is not None
assert entry[1][f"{direction}_cost_per_token"] == rate
assert litellm.model_cost[profile] == profile_entry
@pytest.mark.parametrize("prompt_tokens", (299_999, 300_000, 300_001))
@pytest.mark.parametrize("cache_field", ("cache_read_input_token_cost", "cache_creation_input_token_cost"))
@pytest.mark.parametrize("tier_override", (None, 0.0, 0.003))
def test_bedrock_registered_cache_prices_preserve_dynamic_threshold_rates(
_local_model_cost_map: None,
monkeypatch: pytest.MonkeyPatch,
prompt_tokens: int,
cache_field: str,
tier_override: float | None,
) -> None:
import litellm.utils as model_utils
from litellm.cost_calculator import pricing_entry_for_cost_calc
model: Final = "openai.threshold-overlay-test"
profile: Final = f"us.{model}"
tier_field: Final = f"{cache_field}_above_300k_tokens"
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.model_cost,
profile: {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"input_cost_per_token_above_300k_tokens": 0.03,
"output_cost_per_token_above_300k_tokens": 0.06,
tier_field: 0.007,
},
},
)
model_utils._invalidate_model_cost_lowercase_map()
litellm.register_model(
{
model: {
cache_field: 0.001,
**({tier_field: tier_override} if tier_override is not None else {}),
}
},
persist_across_reloads=False,
)
usage: Final = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=3,
total_tokens=prompt_tokens + 3,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=5 if cache_field == "cache_read_input_token_cost" else 0
),
cache_creation_input_tokens=5 if cache_field == "cache_creation_input_token_cost" else 0,
)
response: Final = ModelResponse(model=model, usage=usage)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="threshold-overlay",
function_id="threshold-overlay",
)
logging_obj.model_call_details["model_id"] = profile
expected_input_rate: Final = 0.03 if prompt_tokens > 300_000 else 0.01
expected_output_rate: Final = 0.06 if prompt_tokens > 300_000 else 0.02
expected_cache_rate: Final = (
(tier_override if tier_override is not None else 0.007) if prompt_tokens > 300_000 else 0.001
)
total: Final = completion_cost(
completion_response=response,
model=f"bedrock/converse/{model}",
custom_llm_provider="bedrock",
optional_params={"model_id": profile},
litellm_logging_obj=logging_obj,
)
assert total == pytest.approx(
(prompt_tokens - 5) * expected_input_rate + 3 * expected_output_rate + 5 * expected_cache_rate
)
breakdown_key: Final = "cache_read_cost" if cache_field == "cache_read_input_token_cost" else "cache_creation_cost"
assert logging_obj.cost_breakdown is not None
assert logging_obj.cost_breakdown.get(breakdown_key, 0.0) == pytest.approx(5 * expected_cache_rate)
assert logging_obj.billed_token_rates is not None
assert logging_obj.billed_token_rates.input_cost_per_token == expected_input_rate
assert logging_obj.billed_token_rates.output_cost_per_token == expected_output_rate
entry: Final = pricing_entry_for_cost_calc(model, response, "bedrock", False, None, None, None, logging_obj)
assert entry is not None
assert entry[1]["input_cost_per_token_above_300k_tokens"] == 0.03
assert entry[1][tier_field] == (tier_override if tier_override is not None else 0.007)

138
tests/unit/test_main.py
View file

@ -1063,6 +1063,50 @@ def test_responses_api_bridge_check_gpt_5_4_tools_with_reasoning_none_stays_chat
assert model_info.get("mode") != "responses"
@pytest.mark.parametrize(
"model, custom_llm_provider, reasoning_effort",
[
("gpt-6.1-sol", "openai", "none"),
("gpt-6.1-sol", "openai", {"effort": "none"}),
("gpt-6-astra", "azure", "none"),
],
)
def test_responses_api_bridge_check_tools_with_a_none_the_row_refuses_routes_to_responses(
local_model_cost_map: None, model: str, custom_llm_provider: str, reasoning_effort: object
) -> None:
"""The ``none`` escape hatch only holds when the chat gate forwards it. These rows refuse
``none``, so the gate drops it (or refuses the call) and the provider runs its default effort,
under which Chat Completions rejects function tools; the bridge has to fire the way it does
for an unset effort. Row-derived: the test follows the flag, never the vendor's answer."""
from litellm.main import responses_api_bridge_check
row_key: Final = model if custom_llm_provider == "openai" else f"{custom_llm_provider}/{model}"
assert litellm.model_cost[row_key]["supports_none_reasoning_effort"] is False
model_info, _ = responses_api_bridge_check(
model=model,
custom_llm_provider=custom_llm_provider,
tools=[{"type": "function", "function": {"name": "get_capital"}}],
reasoning_effort=reasoning_effort,
)
assert model_info.get("mode") == "responses"
def test_responses_api_bridge_check_a_none_the_row_refuses_stays_chat_without_tools(
local_model_cost_map: None,
) -> None:
"""Without function tools nothing forces the Responses route: the dropped ``none`` leaves a
plain chat call at the default effort, as for any other dropped level."""
from litellm.main import responses_api_bridge_check
model_info, _ = responses_api_bridge_check(
model="gpt-6.1-sol", custom_llm_provider="openai", reasoning_effort="none"
)
assert model_info.get("mode") != "responses"
def test_responses_api_bridge_check_reasoning_none_with_summary_still_routes_to_responses():
"""A reasoning summary is Responses-only regardless of effort value."""
from litellm.main import responses_api_bridge_check
@ -1439,11 +1483,15 @@ _FOUNDRY_FUNCTION_TOOL: Final = ({"type": "function", "function": {"name": "get_
pytest.param("gpt-6-astra", "https://myresource.openai.azure.com", None, id="gpt-6-azure-openai-host"),
pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "low", id="gpt-5.6-explicit-effort"),
pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, {"effort": "high"}, id="gpt-5.6-explicit-effort-dict"),
pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "none", id="gpt-6-row-refused-none"),
],
)
def test_responses_api_bridge_check_azure_ai_foundry_rejected_tools_route_to_responses(
model_name, api_base, reasoning_effort
):
"""The explicit ``none`` only keeps a Foundry gpt-6 request on chat while the deployment takes
it; the ``azure_ai/gpt-6-astra`` row says it does not, so the gate drops it and the call runs at
the default effort, under which Foundry rejects function tools from gpt-6 on."""
from litellm.main import responses_api_bridge_check
model_info, model = responses_api_bridge_check(
@ -1461,7 +1509,6 @@ def test_responses_api_bridge_check_azure_ai_foundry_rejected_tools_route_to_res
@pytest.mark.parametrize(
"model_name, api_base, reasoning_effort",
[
pytest.param("gpt-6-astra", _FOUNDRY_API_BASE, "none", id="explicit-none-stays-chat"),
pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, None, id="gpt-5.6-unset-effort-stays-chat"),
pytest.param("gpt-5.6-sol", _FOUNDRY_API_BASE, "none", id="gpt-5.6-explicit-none-stays-chat"),
pytest.param("gpt-5.5", _FOUNDRY_API_BASE, "high", id="gpt-5.5-explicit-effort-stays-chat"),
@ -6104,3 +6151,92 @@ async def test_async_responses_still_retries_provider_server_errors(monkeypatch)
)
assert result.status == "completed"
assert response.call_count == 2
@pytest.mark.parametrize("stream", (False, True))
@pytest.mark.parametrize("asynchronous", (False, True))
@pytest.mark.parametrize("profile,custom_price", ((None, False), ("us.openai.gpt-6-sol", False), (None, True)))
@pytest.mark.asyncio
async def test_converse_fallback_prices_the_current_attempt(
monkeypatch: pytest.MonkeyPatch, stream: bool, asynchronous: bool, profile: str | None, custom_price: bool
) -> None:
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.rust_bridge import configuration
from tests.unit.llms.bedrock.slow_upstream import CONVERSE_ANSWER
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
monkeypatch.setenv("LITELLM_RUST", "0")
configuration.rust(False)
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
failed_profile: Final = "us.amazon.nova-micro-v1:0"
model: Final = "openai.gpt-6-sol"
logging_obj: Final = LiteLLMLogging(
model="amazon.nova-micro-v1:0",
messages=[{"role": "user", "content": "hi"}],
stream=stream,
call_type="acompletion" if asynchronous else "completion",
start_time=datetime(2026, 1, 1),
litellm_call_id="bedrock-fallback-pricing",
function_id="bedrock-fallback-pricing",
)
def respond(request: httpx.Request) -> httpx.Response:
if failed_profile in request.url.path:
return httpx.Response(400, json={"message": "Input is too long for requested model."})
assert request.url.path == f"/model/{profile or model}/converse"
return httpx.Response(200, content=CONVERSE_ANSWER)
client: Final = (
AsyncHTTPHandler(transport=httpx.MockTransport(respond))
if asynchronous
else HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
)
kwargs: Final = {
"messages": [{"role": "user", "content": "hi"}],
"api_key": "test-bearer",
"aws_region_name": "us-east-1",
"client": client,
"litellm_logging_obj": logging_obj,
"num_retries": 0,
"stream": stream,
"fake_stream": stream,
}
if asynchronous:
with pytest.raises(litellm.BadRequestError):
await litellm.acompletion(model="bedrock/converse/amazon.nova-micro-v1:0", model_id=failed_profile, **kwargs)
else:
with pytest.raises(litellm.BadRequestError):
litellm.completion(model="bedrock/converse/amazon.nova-micro-v1:0", model_id=failed_profile, **kwargs)
overrides: Final = {
**({"model_id": profile} if profile is not None else {}),
**({"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} if custom_price else {}),
}
response: Final = (
await litellm.acompletion(model=f"bedrock/converse/{model}", **kwargs, **overrides)
if asynchronous
else litellm.completion(model=f"bedrock/converse/{model}", **kwargs, **overrides)
)
result: Final = (
litellm.stream_chunk_builder(
[chunk async for chunk in response] if asynchronous else list(response), logging_obj=response.logging_obj
)
if stream
else response
)
assert result is not None
prices: Final = litellm.get_model_info(profile or model, custom_llm_provider="bedrock")
expected: Final = (
0.001 + 0.002
if custom_price
else prices["input_cost_per_token"] + prices["output_cost_per_token"]
)
active_logging: Final = response.logging_obj if stream else logging_obj
assert active_logging.response_cost_calculator(result) == pytest.approx(expected)
active_logging.update_environment_variables(litellm_params={}, optional_params={})
assert active_logging.response_cost_calculator(result) == pytest.approx(expected)
if asynchronous:
await client.close()
else:
client.close()
configuration.reset_rust_configuration()

12
tests/unit/test_model_prices_schema.py
View file

@ -559,7 +559,17 @@ def test_deepseek_rows_bill_half_rate_outside_weekday_peak_hours(path: Path):
assert all(name in rows for name in DEEPSEEK_PRICED_ROWS)
PROVIDER_LABELS_WITHOUT_A_MODEL_SET: Final = frozenset({"sagemaker", "bedrock_converse"})
PROVIDER_LABELS_WITHOUT_A_MODEL_SET: Final = frozenset(
{
"sagemaker",
"bedrock_converse",
"coralbricks",
"inferenceapis",
"microsoft_365_copilot",
"openinfer",
"scaledown",
}
)
MODES_SERVED_OUTSIDE_THE_LLM_PROVIDER_REGISTRY: Final = frozenset({"search", "decisions", "evaluation"})
VERTEX_FAMILIES_A_VERTEX_WILDCARD_GRANT_DOES_NOT_LIST: Final = frozenset(
{

181
tests/unit/test_router_model_cost_isolation.py
View file

@ -796,6 +796,178 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override():
_restore_model_cost_entries(model_keys)
@pytest.mark.parametrize("service_tier", ("default", "priority", "flex"))
@pytest.mark.parametrize(
"field",
("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost", "input_cost_per_token_priority"),
)
@pytest.mark.parametrize("rate", (0.0, 0.004))
@pytest.mark.parametrize("explicit_base", (False, True))
@pytest.mark.parametrize("provider", (None, "bedrock"))
@pytest.mark.parametrize("registration_path", ("direct", "startup", "reload"))
def test_bedrock_deployment_inherits_invoked_profile_around_custom_prices(
monkeypatch: pytest.MonkeyPatch,
field: str,
rate: float,
explicit_base: bool,
provider: str | None,
registration_path: str,
service_tier: str,
) -> None:
from litellm.cost_calculator import completion_cost
from litellm.types.utils import ModelResponse, Usage, PromptTokensDetailsWrapper
model: Final = "openai.gpt-5.6-sol"
profile: Final = f"us.{model}"
base_model: Final = f"global.{model}"
profile_prices: Final = {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.01,
"output_cost_per_token": 0.02,
"cache_read_input_token_cost": 0.001,
"input_cost_per_token_priority": 0.05,
"output_cost_per_token_priority": 0.06,
"input_cost_per_token_flex": 0.007,
"output_cost_per_token_flex": 0.008,
}
base_prices: Final = {**profile_prices, "input_cost_per_token": 0.03, "output_cost_per_token": 0.04}
monkeypatch.setattr(
litellm, "model_cost", {**litellm.model_cost, profile: dict(profile_prices), base_model: dict(base_prices)}
)
deployment: Final = Deployment(
model_name="profile-priced",
litellm_params=LiteLLM_Params(
model=f"converse/{model}" if provider else f"bedrock/converse/{model}",
model_id=profile,
custom_llm_provider=provider,
),
model_info=ModelInfo(
**{
"id": "profile-priced-deployment",
field: rate,
**({"base_model": f"bedrock/{base_model}"} if explicit_base else {}),
}
),
)
if registration_path == "direct":
Router._register_deployment_pricing(deployment)
else:
router: Final = Router(model_list=[deployment.to_json(exclude_none=True)])
if registration_path == "reload":
router.replay_model_cost_registrations()
response: Final = ModelResponse(
model=model,
usage=Usage(
prompt_tokens=100,
completion_tokens=3,
total_tokens=103,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5),
),
)
catalog: Final = base_prices if explicit_base else profile_prices
input_field: Final = "input_cost_per_token" + (f"_{service_tier}" if service_tier != "default" else "")
output_field: Final = "output_cost_per_token" + (f"_{service_tier}" if service_tier != "default" else "")
input_rate: Final = rate if field == input_field else catalog[input_field]
output_rate: Final = rate if field == output_field else catalog[output_field]
cache_rate: Final = rate if field == "cache_read_input_token_cost" else 0.001
assert completion_cost(
model=model,
completion_response=response,
custom_llm_provider="bedrock",
custom_pricing=True,
router_model_id="profile-priced-deployment",
service_tier=service_tier,
) == pytest.approx(95 * input_rate + 3 * output_rate + 5 * cache_rate)
assert litellm.model_cost[profile] == profile_prices
@pytest.mark.parametrize("catalog_entry", (None, {"litellm_provider": 1, "mode": "chat"}))
@pytest.mark.parametrize("known_model", (False, True))
@pytest.mark.parametrize("explicit_base", (False, True))
@pytest.mark.parametrize("service_tier", ("default", "priority", "flex"))
@pytest.mark.parametrize("override", (None, "cache_read_input_token_cost", "input_cost_per_token_priority"))
@pytest.mark.parametrize("reload", (False, True))
def test_bedrock_deployment_keeps_custom_prices_when_profile_lookup_fails(
monkeypatch: pytest.MonkeyPatch,
catalog_entry: dict[str, object] | None,
known_model: bool,
explicit_base: bool,
service_tier: str,
override: str | None,
reload: bool,
) -> None:
from litellm.cost_calculator import completion_cost
from litellm.types.utils import ModelResponse, Usage, PromptTokensDetailsWrapper
model: Final = "anthropic.unmapped-profile-pricing-test"
profile: Final = (
"us.anthropic.invalid-profile-pricing-test"
if catalog_entry is not None
else "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/opaque-test"
)
configured_prices: Final = {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.03,
"output_cost_per_token": 0.04,
"cache_read_input_token_cost": 0.001,
"input_cost_per_token_priority": 0.05,
"output_cost_per_token_priority": 0.06,
"input_cost_per_token_flex": 0.007,
"output_cost_per_token_flex": 0.008,
}
monkeypatch.setattr(
litellm,
"model_cost",
{
**litellm.model_cost,
**({model: dict(configured_prices)} if known_model else {}),
**({profile: catalog_entry} if catalog_entry is not None else {}),
},
)
deployment: Final = Deployment(
model_name="custom-priced",
litellm_params=LiteLLM_Params(model=f"bedrock/converse/{model}", model_id=profile),
model_info=ModelInfo(
id="unmapped-profile-deployment",
input_cost_per_token=0.01,
output_cost_per_token=0.02,
**({override: 0.0} if override else {}),
**({"base_model": f"bedrock/{profile}"} if explicit_base else {}),
),
)
router: Final = Router(model_list=[deployment.to_json(exclude_none=True)])
if reload:
router.replay_model_cost_registrations()
inherit_configured: Final = known_model and not explicit_base
tier_rates: Final = inherit_configured and service_tier != "default"
input_rate: Final = configured_prices[f"input_cost_per_token_{service_tier}"] if tier_rates else 0.01
output_rate: Final = configured_prices[f"output_cost_per_token_{service_tier}"] if tier_rates else 0.02
effective_input: Final = 0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate
cache_rate: Final = 0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input)
assert completion_cost(
model=model,
completion_response=ModelResponse(
model=model,
usage=Usage(
prompt_tokens=100,
completion_tokens=3,
total_tokens=103,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=80),
),
),
custom_llm_provider="bedrock",
custom_pricing=True,
router_model_id="unmapped-profile-deployment",
service_tier=service_tier,
) == pytest.approx(20 * effective_input + 80 * cache_rate + 3 * output_rate)
if known_model:
assert {
key: value for key, value in litellm.model_cost[model].items() if key in configured_prices
} == configured_prices
def test_partial_custom_pricing_inherits_builtin_cache_pricing():
"""A deployment that overrides only input/output cost on a cache-supporting
model must still bill cache_read and cache_creation tokens. Before the
@ -2225,14 +2397,7 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None:
)
rebuilt = litellm.model_cost["priced-id"]
assert at_boot.items() <= rebuilt.items(), (
f"the rebuild changed or dropped a field the boot registration wrote: "
f"{ {k: (v, rebuilt.get(k)) for k, v in at_boot.items() if rebuilt.get(k) != v} }"
)
assert {field: rebuilt[field] for field in set(rebuilt) - set(at_boot)} == {
"db_model": False,
"member_auto_router": False,
}
assert at_boot == rebuilt
assert "member_auto_router" not in litellm.model_cost["gpt-4o"]
assert router.model_list
finally:

73
tests/unit/test_utils.py
View file

@ -833,6 +833,18 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"type": "object",
"properties": {
"supports_computer_use": {"type": "boolean"},
"cache_creation_input_token_cost_above_272k_tokens_decisions": {"type": "number"},
"cache_creation_input_token_cost_decisions": {"type": "number"},
"cache_read_input_token_cost_above_272k_tokens_decisions": {"type": "number"},
"cache_read_input_token_cost_decisions": {"type": "number"},
"explicit_cache_read_input_token_cost": {"type": "number"},
"input_cost_per_text_message": {"type": "number"},
"input_cost_per_token_above_272k_tokens_decisions": {"type": "number"},
"input_cost_per_token_decisions": {"type": "number"},
"output_cost_per_token_above_272k_tokens_decisions": {"type": "number"},
"output_cost_per_token_decisions": {"type": "number"},
"supports_bedrock_converse_tool_result_images": {"type": "boolean"},
"supports_bedrock_native_structured_output": {"type": "boolean"},
"cache_creation_input_audio_token_cost": {"type": "number"},
"cache_creation_input_token_cost": {"type": "number"},
"cache_creation_input_token_cost_above_1hr": {"type": "number"},
@ -959,6 +971,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"image_edit",
"embedding",
"evaluation",
"extraction",
"guardrail",
"image_generation",
"video_generation",
@ -1132,6 +1145,16 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"type": "object",
"additionalProperties": {"type": "number"},
},
"search_free_requests": {
"type": "object",
"properties": {
"count": {"type": "integer", "minimum": 1},
"period": {"type": "string", "enum": ["day", "month"]},
"shared_group": {"type": "string"},
},
"required": ["count", "period", "shared_group"],
"additionalProperties": False,
},
"search_context_cost_per_query": {
"type": "object",
"properties": {
@ -4384,7 +4407,11 @@ async def test_builtin_string_callback_registers_when_subclass_already_active(
@pytest.mark.usefixtures("isolated_openai_model_sets")
@pytest.mark.usefixtures("isolated_openai_model_sets")
def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypatch):
@pytest.mark.parametrize("rate", (0.0, 0.000123))
@pytest.mark.parametrize("updated_rate", (None, 0.0, 0.000456))
def test_reapply_runtime_registrations_replays_register_model_overrides(
monkeypatch: pytest.MonkeyPatch, rate: float, updated_rate: float | None
) -> None:
"""
register_model is the documented way to override pricing for a model. A
price-data reload swaps litellm.model_cost for a freshly fetched catalog,
@ -4415,11 +4442,17 @@ def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypa
"openai/gpt-4o": {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 0.000123,
"input_cost_per_token": rate,
}
}
)
litellm.register_model({"openai/gpt-4o": {"supports_function_calling": True, "input_cost_per_token": None}})
litellm.register_model(
{"openai/gpt-4o": {"input_cost_per_token": str(updated_rate) if updated_rate is not None else None}}
)
litellm.register_model({"openai/gpt-4o": {"supports_function_calling": False}})
litellm.model_cost = {
"openai/gpt-4o": {
"litellm_provider": "openai",
@ -4431,7 +4464,9 @@ def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypa
_invalidate_model_cost_lowercase_map()
reapply_runtime_model_cost_registrations()
assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == 0.000123
expected_rate: Final = rate if updated_rate is None else updated_rate
assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == expected_rate
assert litellm.model_cost["openai/gpt-4o"]["supports_function_calling"] is False
assert litellm.model_cost["openai/gpt-4o"]["max_input_tokens"] == 4242
finally:
litellm.model_cost = saved_model_cost
@ -8175,12 +8210,17 @@ def test_models_by_provider():
continue
elif k == "sample_spec":
continue
elif (
v["litellm_provider"] == "sagemaker"
or v["litellm_provider"] == "bedrock_converse"
elif v["litellm_provider"] in (
"sagemaker",
"bedrock_converse",
"coralbricks",
"inferenceapis",
"microsoft_365_copilot",
"openinfer",
"scaledown",
):
continue
elif v.get("mode") in ("search", "decisions", "evaluation"):
elif v.get("mode") in ("search", "decisions", "evaluation", "extraction"):
continue
else:
providers.add(v["litellm_provider"])
@ -9359,3 +9399,22 @@ def my_pre_call_rule(input: str):
if len(input) > 10:
return False
return True
@pytest.mark.parametrize(
"key", ("openai.test-converse", "bedrock/openai.test-converse", "bedrock/converse/openai.test-converse")
)
def test_runtime_model_cost_registrations_resolve_bedrock_routing(monkeypatch: pytest.MonkeyPatch, key: str) -> None:
from litellm import utils as model_utils
rates: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}
monkeypatch.setattr(
litellm, "model_cost", {**litellm.model_cost, key: {"litellm_provider": "bedrock_converse", "mode": "chat"}}
)
litellm.register_model({key: rates}, persist_across_reloads=False)
registrations: Final = dict(
model_utils.runtime_model_cost_registrations("bedrock/converse/openai.test-converse", "bedrock")
)
assert registrations[key] == rates
assert model_utils.runtime_model_cost_registrations("bedrock/converse/openai.other-converse", "bedrock") == ()

905
uv.lock generated
View file

File diff suppressed because it is too large Load diff