Merge remote-tracking branch 'origin/main' into litellm_wildcard_discovery_provider_config

# Conflicts:
#	litellm/proxy/auth/model_checks.py
This commit is contained in:
mateo-berri 2026-10-30 00:54:13 -07:00
commit 683c4673b1
98 changed files with 5917 additions and 669 deletions

View file

@ -872,6 +872,11 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje
# ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params.
if "-to-" not in key and "fireworks-ai-default" not in key:
fireworks_ai_models.add(key)
elif value.get("litellm_provider") == "together_ai":
if _is_past_deprecation_date(value):
together_ai_models.discard(key)
elif not key.startswith("together-ai-"):
together_ai_models.add(key)
elif value.get("litellm_provider") == "baseten":
if _is_past_deprecation_date(value):
baseten_models.discard(key)

View file

@ -7,6 +7,7 @@ from typing import Final, Literal
from litellm.litellm_core_utils.env_utils import (
get_env_float,
get_env_float_in_range,
get_env_float_or_none,
get_env_int,
get_env_int_in_range,
@ -149,6 +150,9 @@ DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = get_env_float(
"DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5
) # default cooldown a deployment if 50% of requests fail in a given minute
DEFAULT_MAX_TOKENS: Final = get_env_int("DEFAULT_MAX_TOKENS", 4096)
NO_MAX_TOKENS_TPM_RESERVATION_FRACTION: Final = get_env_float_in_range(
"NO_MAX_TOKENS_TPM_RESERVATION_FRACTION", 0.25, 0.01, 1.0
)
DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3)
DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1)
DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5)
@ -1273,7 +1277,7 @@ empower_models: Final = set(
]
)
together_ai_models: Final[set] = set(
together_ai_models: Final[set[str]] = set(
[
# llama llms - chat
"togethercomputer/llama-2-70b-chat",

View file

@ -1544,6 +1544,12 @@ def _response_time_ms_for_cost(
return litellm_logging_obj.get_response_ms()
def _breakdown_total(litellm_logging_obj: LitellmLoggingObject | None) -> float | None:
breakdown: Final = None if litellm_logging_obj is None else litellm_logging_obj.cost_breakdown
total: Final[object] = None if breakdown is None else breakdown.get("total_cost")
return float(total) if isinstance(total, (int, float)) and not isinstance(total, bool) else None
def _adjusted_provider_billed_cost(
reported_cost: float,
custom_llm_provider: str | None,
@ -1561,9 +1567,9 @@ def _adjusted_provider_billed_cost(
)
_store_cost_breakdown_in_logging_obj(
litellm_logging_obj=litellm_logging_obj,
prompt_tokens_cost_usd_dollar=None,
completion_tokens_cost_usd_dollar=None,
cost_for_built_in_tools_cost_usd_dollar=None,
prompt_tokens_cost_usd_dollar=0.0,
completion_tokens_cost_usd_dollar=adjusted_cost,
cost_for_built_in_tools_cost_usd_dollar=0.0,
total_cost_usd_dollar=adjusted_cost,
original_cost=reported_cost,
discount_percent=discount_percent,
@ -2376,6 +2382,15 @@ def response_cost_calculator(
)
provider_response_cost: Final = get_response_cost_from_hidden_params(hidden_params)
if provider_response_cost is not None and not custom_pricing:
if _breakdown_total(litellm_logging_obj) != provider_response_cost:
_store_cost_breakdown_in_logging_obj(
litellm_logging_obj=litellm_logging_obj,
prompt_tokens_cost_usd_dollar=0.0,
completion_tokens_cost_usd_dollar=provider_response_cost,
cost_for_built_in_tools_cost_usd_dollar=0.0,
total_cost_usd_dollar=provider_response_cost,
original_cost=provider_response_cost,
)
return provider_response_cost
response_cost = completion_cost(

View file

@ -9,6 +9,7 @@
## LiteLLM versions of the OpenAI Exception Types
import datetime
import enum
from collections.abc import Mapping, Sequence
from typing import Any, Final
@ -115,6 +116,54 @@ def validate_rate_limit_type(value: object) -> str | None:
return None
def _upstream_body(response: object) -> bytes:
if not isinstance(response, httpx.Response):
return b""
try:
return response.content
except httpx.ResponseNotRead:
return b""
def _request_without_credentials(upstream: object, fallback_url: str) -> httpx.Request:
if not isinstance(upstream, httpx.Response):
return httpx.Request(method="POST", url=fallback_url)
try:
upstream_request: Final = upstream.request
except RuntimeError:
return httpx.Request(method="POST", url=fallback_url)
return httpx.Request(
method=upstream_request.method,
url=upstream_request.url.copy_with(userinfo=b"", query=None, fragment=None),
)
def _upstream_elapsed(upstream: object) -> datetime.timedelta | None:
if not isinstance(upstream, httpx.Response):
return None
try:
return upstream.elapsed
except RuntimeError:
return None
def error_response_without_credentials(status_code: int, upstream: object, fallback_url: str) -> httpx.Response:
response: Final = httpx.Response(
status_code=status_code,
stream=httpx.ByteStream(_upstream_body(upstream)),
request=_request_without_credentials(upstream, fallback_url),
)
_ = response.read()
response.headers = httpx.Headers(upstream.headers if isinstance(upstream, httpx.Response) else None)
elapsed: Final = _upstream_elapsed(upstream)
if elapsed is not None:
response.elapsed = elapsed
return response
_NO_UPSTREAM_REQUEST_URL: Final = "https://cloud.google.com/vertex-ai/"
_MINIMAL_ERROR_RESPONSE: httpx.Response | None = None
@ -519,20 +568,11 @@ class RateLimitError(openai.RateLimitError):
# headers stay reachable on `e.response.headers` for callers that
# explicitly want them; only the proxy-supplied `headers=` kwarg
# makes it onto `self.headers`.
_response_headers: Final = getattr(response, "headers", None) if response is not None else None
self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None
# Mirrors FastAPI HTTPException.detail so the same instance can be
# serialized through both the ProxyException and HTTPException paths.
self.detail = detail if detail is not None else self.message
self.response = httpx.Response(
status_code=429,
headers=_response_headers,
content=response.content if response is not None else None,
request=httpx.Request(
method="POST",
url=" https://cloud.google.com/vertex-ai/",
),
)
self.response = error_response_without_credentials(429, response, _NO_UPSTREAM_REQUEST_URL)
super().__init__(
self.message,
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
@ -757,15 +797,7 @@ class ServiceUnavailableError(openai.APIStatusError):
self.litellm_debug_info = litellm_debug_info
self.max_retries = max_retries
self.num_retries = num_retries
_response_headers: Final = getattr(response, "headers", None) if response is not None else None
self.response = httpx.Response(
status_code=self.status_code,
headers=_response_headers,
request=httpx.Request(
method="POST",
url=" https://cloud.google.com/vertex-ai/",
),
)
self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL)
super().__init__(
self.message, response=self.response, body=None
) # Call the base class constructor with the parameters it needs
@ -808,15 +840,7 @@ class BadGatewayError(openai.APIStatusError):
self.litellm_debug_info = litellm_debug_info
self.max_retries = max_retries
self.num_retries = num_retries
_response_headers: Final = getattr(response, "headers", None) if response is not None else None
self.response = httpx.Response(
status_code=self.status_code,
headers=_response_headers,
request=httpx.Request(
method="POST",
url=" https://cloud.google.com/vertex-ai/",
),
)
self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL)
super().__init__(
self.message,
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
@ -862,15 +886,7 @@ class InternalServerError(openai.InternalServerError):
self.litellm_debug_info = litellm_debug_info
self.max_retries = max_retries
self.num_retries = num_retries
_response_headers: Final = getattr(response, "headers", None) if response is not None else None
self.response = httpx.Response(
status_code=self.status_code,
headers=_response_headers,
request=httpx.Request(
method="POST",
url=" https://cloud.google.com/vertex-ai/",
),
)
self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL)
super().__init__(
self.message,
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
@ -908,6 +924,7 @@ class APIError(openai.APIError):
max_retries: int | None = None,
num_retries: int | None = None,
body: object | None = None,
response: httpx.Response | None = None,
):
self.status_code = status_code
self.message = f"litellm.APIError: {message}"
@ -916,6 +933,8 @@ class APIError(openai.APIError):
self.litellm_debug_info = litellm_debug_info
self.max_retries = max_retries
self.num_retries = num_retries
if response is not None:
self.response = response
if request is None:
request = httpx.Request(method="POST", url="https://api.openai.com/v1")
super().__init__(self.message, request=request, body=body)

View file

@ -23,6 +23,7 @@ from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.integrations.prompt_management_base import PromptManagementClient
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.litellm_core_utils.prompt_templates.common_utils import (
cache_control_as_prompt_cache_breakpoint,
with_prompt_cache_breakpoint,
@ -175,7 +176,10 @@ def _hosted_entry_flag(entry: Mapping[str, object], resolve_provider: Callable[[
def _hosted_openai_dialect_flag(
model: str, custom_llm_provider: str | None, resolve_provider: Callable[[str], str | None]
model: str,
custom_llm_provider: str | None,
resolve_provider: Callable[[str], str | None],
model_cost: Mapping[str, Mapping[str, object]] | None = None,
) -> bool | None:
"""
Explicit opt-in for an OpenAI-shaped deployment served by another provider.
@ -194,14 +198,19 @@ def _hosted_openai_dialect_flag(
"""
import litellm
exact_entry: Final = litellm.model_cost.get(model)
exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, dict) else None
cost_map: Final = (
cast(Mapping[str, Mapping[str, object]], litellm.model_cost) # cast-ok: the live map holds catalog entries
if model_cost is None
else model_cost
)
exact_entry: Final = cost_map.get(model)
exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, Mapping) else None
if custom_llm_provider is None and (exact_provider == "openai" or ("/" not in model and exact_provider is None)):
return None
provider: Final = custom_llm_provider or resolve_provider(model)
if provider is None or provider == "openai":
return None
if isinstance(exact_entry, dict) and exact_provider == provider:
if isinstance(exact_entry, Mapping) and exact_provider == provider:
return _hosted_entry_flag(exact_entry, lambda: provider)
from litellm.utils import get_potential_model_names
@ -214,8 +223,8 @@ def _hosted_openai_dialect_flag(
names["stripped_model_name"],
names["provider_prefixed_model_name"],
)
entries: Final = (litellm.model_cost.get(candidate) for candidate in candidates)
flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, dict))
entries: Final = (cost_map.get(candidate) for candidate in candidates)
flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, Mapping))
return next((flag for flag in flags if flag is not None), None)
@ -537,6 +546,18 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return False
return prompt_cache_options is not None or targets_openai_api(api_base)
@staticmethod
def catalog_targets_openai_prompt_cache_breakpoint(model: str, custom_llm_provider: str) -> bool:
return (
_hosted_openai_dialect_flag(
model,
custom_llm_provider,
AnthropicCacheControlHook._resolve_provider,
GetModelCostMap.loaded_model_cost_map(),
)
is True
)
@staticmethod
def _resolve_provider(model: str) -> str | None:
return _served_provider_for_model(model)

View file

@ -9,6 +9,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
import httpx
from pydantic import TypeAdapter
from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger
@ -218,6 +219,15 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None =
)
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuardrailParams]]:
if not isinstance(metadata, dict):
return []
return metadata.get("guardrails") or []
class CustomGuardrail(CustomLogger):
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
use_native_during_call_hook: ClassVar[bool] = False
@ -791,28 +801,22 @@ class CustomGuardrail(CustomLogger):
return True
raise
def get_guardrail_from_metadata(self, data: dict) -> list[str] | list[dict[str, DynamicGuardrailParams]]:
def get_guardrail_from_metadata(self, data: dict) -> list[str | dict[str, DynamicGuardrailParams]]:
"""
Returns the guardrail(s) to be run from the metadata or root
Returns the guardrail(s) to be run from the root, else from both metadata fields together
"""
if "guardrails" in data:
return data["guardrails"]
# Check both metadata locations. For regular endpoints move_guardrails_to_metadata
# writes to "metadata"; for thread/assistant endpoints it writes to
# "litellm_metadata". We check the one that actually contains the "guardrails"
# key so that a non-empty litellm_metadata without guardrails does not shadow
# the merged list stored in metadata (which would cause team guardrails to be
# silently skipped while default_on=True policy guardrails still fire).
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key) or {}
if isinstance(meta, dict) and "guardrails" in meta:
return meta.get("guardrails") or []
return []
request_fields: Final = _REQUEST_FIELDS.validate_python(data)
return [
*_guardrails_listed_in(request_fields.get("metadata")),
*_guardrails_listed_in(request_fields.get("litellm_metadata")),
]
def _guardrail_is_in_requested_guardrails(
self,
requested_guardrails: list[str] | list[dict[str, DynamicGuardrailParams]],
requested_guardrails: Sequence[str | dict[str, DynamicGuardrailParams]],
) -> bool:
for _guardrail in requested_guardrails:
if isinstance(_guardrail, dict):

View file

@ -398,14 +398,17 @@ class WebSearchInterceptionLogger(CustomLogger):
(t for t in tools if is_anthropic_native_web_search_tool(t)),
None,
)
emit_native_blocks: Final = native_tool is not None or (
kwargs is not None and bool(kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY))
)
outcome: Final = await self._short_circuit_search_outcome(query, kwargs=kwargs)
search_result_text: Final = WebSearchTransformation.search_outcome_text(outcome)
content: Final[list[dict[str, object]]] = []
if native_tool is not None:
if emit_native_blocks:
tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}"
tool_name: Final = native_tool.get("name") or "web_search"
tool_name: Final = (native_tool.get("name") if native_tool is not None else None) or "web_search"
content.append(
{
"type": "server_tool_use",
@ -415,10 +418,10 @@ class WebSearchInterceptionLogger(CustomLogger):
}
)
content.append(
WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome)
WebSearchTransformation.build_web_search_outcome_block(
tool_use_id=tool_use_id, outcome=outcome, with_snippets=False
)
)
# Keep the text block so non-native short-circuit callers (Claude Code,
# github_copilot, etc.) see the same payload they always have.
content.append({"type": "text", "text": search_result_text})
response: Final[dict[str, object]] = {
@ -435,7 +438,7 @@ class WebSearchInterceptionLogger(CustomLogger):
verbose_logger.debug(
"WebSearchInterception: Short-circuit search completed, returning synthetic response (%s chars, native_blocks=%s)",
len(search_result_text),
native_tool is not None,
emit_native_blocks,
)
return response

View file

@ -412,6 +412,7 @@ class WebSearchTransformation:
def build_web_search_tool_result_block(
tool_use_id: str,
search_response: SearchResponse | None,
with_snippets: bool = True,
) -> dict[str, object]:
"""
Build an Anthropic-native ``web_search_tool_result`` content block.
@ -445,6 +446,8 @@ class WebSearchTransformation:
``litellm.asearch()``. If None or empty, the block is still
emitted with an empty result list (signals "search ran, no
results" rather than "search did not run").
with_snippets: False when a text block in the same reply already
carries the snippets.
"""
items: Final[list[dict[str, object]]] = []
if search_response is not None:
@ -460,7 +463,7 @@ class WebSearchTransformation:
"title": title,
"page_age": page_age,
"encrypted_content": "",
"snippet": getattr(r, "snippet", "") or "",
**({"snippet": getattr(r, "snippet", "") or ""} if with_snippets else {}),
}
)
return {
@ -481,12 +484,15 @@ class WebSearchTransformation:
}
@staticmethod
def build_web_search_outcome_block(tool_use_id: str, outcome: SearchOutcome) -> dict[str, object]:
def build_web_search_outcome_block(
tool_use_id: str, outcome: SearchOutcome, with_snippets: bool = True
) -> dict[str, object]:
match outcome:
case SearchSucceeded(response=response):
return WebSearchTransformation.build_web_search_tool_result_block(
tool_use_id=tool_use_id,
search_response=response,
with_snippets=with_snippets,
)
case SearchFailed(error_code=error_code):
return WebSearchTransformation.build_web_search_tool_result_error_block(

View file

@ -1017,13 +1017,44 @@ def chunk_usage(chunk: object) -> object:
return getattr(chunk, "usage", None)
USAGE_COST_BILLED_PROVIDERS: Final = frozenset({"kyma"})
PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS: Final = frozenset({"xai", "perplexity", "opencode_go"})
UPSTREAM_USAGE_COST: Final = "upstream_usage_cost"
def _numeric(value: object) -> float | None:
return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None
def stamp_litellm_bill_on_usage(response: object, usage: object, bill: float) -> None:
"""Writes litellm's bill into ``usage.cost`` and the hidden ``response_cost`` cache, keeping the upstream's own
``usage.cost`` under ``upstream_usage_cost`` so a later pricing pass still bills the upstream's charge instead of
reading litellm's bill back as if the upstream had reported it."""
hidden: Final[object] = getattr(response, "_hidden_params", None)
if isinstance(hidden, dict):
hidden_params: Final = cast(dict[str, object], hidden) # cast-ok: hidden params ride on the response as a dict
upstream: Final = (
hidden_params[UPSTREAM_USAGE_COST]
if UPSTREAM_USAGE_COST in hidden_params
else _numeric(getattr(usage, "cost", None))
)
hidden_params.update({"response_cost": bill, UPSTREAM_USAGE_COST: upstream})
setattr(usage, "cost", bill)
def billed_usage_cost(response: object, custom_llm_provider: str | None) -> float | None:
if custom_llm_provider not in USAGE_COST_BILLED_PROVIDERS:
"""The charge an OpenAI-compatible upstream reports in ``usage.cost``, skipped for providers billed from their own
usage fields. Once litellm wrote its own bill into ``usage.cost`` (the stream builder and the Responses terminal
event do), the upstream's figure is read from ``upstream_usage_cost`` instead."""
if custom_llm_provider in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS:
return None
cost: Final = provider_reported_usage_cost(_usage_field(response, "usage"), custom_llm_provider)
hidden: Final = as_str_mapping(getattr(response, "_hidden_params", None))
cost: Final = (
_numeric(hidden[UPSTREAM_USAGE_COST])
if hidden is not None and UPSTREAM_USAGE_COST in hidden
else provider_reported_usage_cost(_usage_field(response, "usage"), custom_llm_provider)
)
return cost if cost is not None and math.isfinite(cost) and cost >= 0 else None

View file

@ -71,7 +71,15 @@ def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int)
Values outside the range fall back to the default and warn, so a misconfigured knob can
neither crash the caller nor silently change the meaning of what it computes.
"""
value: Final = get_env_int(env_var, default)
return _within_range(env_var, get_env_int(env_var, default), default, minimum, maximum)
def get_env_float_in_range(env_var: str, default: float, minimum: float, maximum: float) -> float:
"""Float counterpart of `get_env_int_in_range`."""
return _within_range(env_var, get_env_float(env_var, default), default, minimum, maximum)
def _within_range(env_var: str, value: _NumberT, default: _NumberT, minimum: _NumberT, maximum: _NumberT) -> _NumberT:
if minimum <= value <= maximum:
return value
logging.getLogger("LiteLLM").warning(

View file

@ -37,6 +37,7 @@ from ..exceptions import (
ServiceUnavailableError,
Timeout,
UnprocessableEntityError,
error_response_without_credentials,
)
@ -305,6 +306,16 @@ def _litellm_proxy_response(
)
def _without_credentials(response: httpx.Response | None) -> httpx.Response | None:
if not isinstance(response, httpx.Response):
return response
return error_response_without_credentials(response.status_code, response, "https://docs.litellm.ai/docs")
def _upstream_reply(response: httpx.Response | None) -> httpx.Response | None:
return response if response is not None and (response.headers or response.content) else None
def _map_openai_exception(
*,
model: str,
@ -315,7 +326,7 @@ def _map_openai_exception(
exception_provider: str,
extra_information: str,
) -> None:
response: Final = _litellm_proxy_response(original_exception, custom_llm_provider)
response: Final = _without_credentials(_litellm_proxy_response(original_exception, custom_llm_provider))
if (cyber_policy_body := _openai_cyber_policy_body(original_exception)) is not None:
raise ContentPolicyViolationError(
message=f"ContentPolicyViolationError: {exception_provider} - {original_exception}",
@ -378,6 +389,7 @@ def _map_openai_exception(
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=_upstream_reply(response),
)
elif (
("invalid_request_error" in error_str and "content_policy_violation" in error_str)
@ -449,6 +461,7 @@ def _map_openai_exception(
model=model,
llm_provider=custom_llm_provider,
body=getattr(original_exception, "body", None),
response=_upstream_reply(response),
)
elif "Request too large" in error_str:
raise RateLimitError(
@ -479,6 +492,7 @@ def _map_openai_exception(
model=model,
request=_request,
litellm_debug_info=extra_information,
response=_upstream_reply(response),
)
elif hasattr(original_exception, "status_code"):
if original_exception.status_code == 400:
@ -512,6 +526,7 @@ def _map_openai_exception(
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=_upstream_reply(response),
)
elif original_exception.status_code == 422:
raise BadRequestError(
@ -563,6 +578,7 @@ def _map_openai_exception(
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
exception_status_code=original_exception.status_code,
response=_upstream_reply(response),
)
else:
raise APIError(
@ -572,6 +588,7 @@ def _map_openai_exception(
model=model,
request=getattr(original_exception, "request", None),
litellm_debug_info=extra_information,
response=_upstream_reply(response),
)
else:
# if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors
@ -2019,6 +2036,7 @@ def _map_azure_exception(
exception_provider: str,
extra_information: str,
) -> None:
response: Final = _without_credentials(getattr(original_exception, "response", None))
message = get_error_message(error_obj=original_exception)
if message is None:
if hasattr(original_exception, "message"):
@ -2054,7 +2072,7 @@ def _map_azure_exception(
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif "DeploymentNotFound" in error_str:
raise NotFoundError(
@ -2062,7 +2080,7 @@ def _map_azure_exception(
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif azure_error_code == "content_policy_violation" or ExceptionCheckers.is_azure_content_policy_violation_error(
error_str
@ -2076,6 +2094,7 @@ def _map_azure_exception(
model=model,
extra_information=extra_information,
original_exception=original_exception,
response=response,
)
elif azure_error_code == "invalid_encrypted_content" or "could not be verified" in error_str:
helpful_message: Final = (
@ -2094,7 +2113,7 @@ def _map_azure_exception(
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
body=getattr(original_exception, "body", None),
)
elif getattr(original_exception, "status_code", None) in (None, 400) and (
@ -2106,7 +2125,7 @@ def _map_azure_exception(
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
body=getattr(original_exception, "body", None),
)
elif "invalid_request_error" in error_str and getattr(original_exception, "status_code", None) in (None, 400):
@ -2115,7 +2134,7 @@ def _map_azure_exception(
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
body=getattr(original_exception, "body", None),
)
elif "The api_key client option must be set either by passing api_key to the client or by setting" in error_str:
@ -2124,7 +2143,7 @@ def _map_azure_exception(
llm_provider=custom_llm_provider,
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif "Connection error" in error_str:
raise APIConnectionError(
@ -2140,7 +2159,7 @@ def _map_azure_exception(
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
body=getattr(original_exception, "body", None),
)
elif original_exception.status_code == 401:
@ -2149,7 +2168,7 @@ def _map_azure_exception(
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif original_exception.status_code == 404:
raise NotFoundError(
@ -2157,7 +2176,7 @@ def _map_azure_exception(
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif original_exception.status_code == 408:
raise Timeout(
@ -2165,6 +2184,7 @@ def _map_azure_exception(
model=model,
litellm_debug_info=extra_information,
llm_provider="azure",
response=_upstream_reply(response),
)
elif original_exception.status_code == 422:
raise BadRequestError(
@ -2172,7 +2192,7 @@ def _map_azure_exception(
model=model,
llm_provider="azure",
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif original_exception.status_code == 429:
raise RateLimitError(
@ -2180,7 +2200,7 @@ def _map_azure_exception(
model=model,
llm_provider="azure",
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif original_exception.status_code == 502:
raise BadGatewayError(
@ -2188,7 +2208,7 @@ def _map_azure_exception(
model=model,
llm_provider="azure",
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif original_exception.status_code == 503:
raise ServiceUnavailableError(
@ -2196,7 +2216,7 @@ def _map_azure_exception(
model=model,
llm_provider="azure",
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
)
elif original_exception.status_code == 504: # gateway timeout error
raise Timeout(
@ -2205,6 +2225,7 @@ def _map_azure_exception(
litellm_debug_info=extra_information,
llm_provider="azure",
exception_status_code=original_exception.status_code,
response=_upstream_reply(response),
)
else:
raise APIError(
@ -2214,6 +2235,7 @@ def _map_azure_exception(
litellm_debug_info=extra_information,
model=model,
request=httpx.Request(method="POST", url="https://openai.com/"),
response=_upstream_reply(response),
)
else:
# if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors
@ -2589,17 +2611,20 @@ def exception_type(
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=_upstream_reply(
_without_credentials(_litellm_proxy_response(mappable_exception, custom_llm_provider))
),
)
if (
custom_llm_provider == "litellm_proxy"
): # handle special case where calling litellm proxy + exception str contains error message
extract_and_raise_litellm_exception(
response=_litellm_proxy_response(mappable_exception, custom_llm_provider),
response=_without_credentials(_litellm_proxy_response(mappable_exception, custom_llm_provider)),
error_str=error_str,
model=model,
custom_llm_provider=custom_llm_provider,
body=getattr(original_exception, "body", None),
body=getattr(mappable_exception, "body", None),
)
if (
custom_llm_provider == "openai"

View file

@ -3,6 +3,8 @@ Helper functions for health check calls.
"""
import base64
import itertools
import re
from collections.abc import Awaitable, Callable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
@ -26,6 +28,9 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = (
)
ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16
WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT: Final = 2
WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT: Final = 3
BEDROCK_ANY_MODEL_ROUTES: Final = frozenset({"converse", "invoke"})
EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
{"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})}
)
@ -90,6 +95,62 @@ def _ocr_health_check_document(model: str, custom_llm_provider: str) -> Document
return native(model, custom_llm_provider)
def _provider_base_model(custom_llm_provider: str, model: str) -> str | None:
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
try:
model_info: Final = ProviderConfigManager.get_provider_model_info(
model=None, provider=LlmProviders(custom_llm_provider)
)
return None if model_info is None else model_info.get_base_model(model)
except (ValueError, NotImplementedError):
return None
def _is_route_segment(custom_llm_provider: str, segment: str, model_after_segment: str) -> bool:
if custom_llm_provider == "bedrock":
from litellm.llms.bedrock.common_utils import get_all_bedrock_regions
return segment in BEDROCK_ANY_MODEL_ROUTES or segment in get_all_bedrock_regions()
base_model: Final = _provider_base_model(custom_llm_provider, f"{segment}/{model_after_segment}")
return base_model is not None and not base_model.startswith(f"{segment}/")
def wildcard_health_check_models(wildcard: str, custom_llm_provider: str) -> tuple[str, ...]:
from litellm.litellm_core_utils.llm_request_utils import rank_chat_models_from_llm_provider
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
provider_prefix: Final = f"{custom_llm_provider}/"
segments: Final = wildcard.removeprefix(provider_prefix).split("/")
route_segments: Final = tuple(
itertools.takewhile(
lambda segment_index: _is_route_segment(
custom_llm_provider, segments[segment_index], "/".join(segments[segment_index + 1 :])
),
range(len(segments) - 1),
)
)
route_prefix: Final = "".join(f"{segments[segment_index]}/" for segment_index in route_segments)
model_regex: Final = re.compile(PatternMatchRouter().pattern_to_regex("/".join(segments[len(route_segments) :])))
ranked_model_ids: Final = tuple(
(ranked_model.model.removeprefix(provider_prefix), ranked_model)
for ranked_model in rank_chat_models_from_llm_provider(custom_llm_provider)
)
served_model_ids: Final = sorted(
(ranked for ranked in ranked_model_ids if model_regex.fullmatch(ranked[0])),
key=lambda ranked: (ranked[1].is_past_deprecation_date, ranked[1].route_rank, ranked[1].cost, ranked[0]),
)
served_probes: Final = itertools.islice(
dict.fromkeys(f"{provider_prefix}{route_prefix}{model_id}" for model_id, _ in served_model_ids),
WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT,
)
provider_probes: Final = (
f"{provider_prefix}{model_id}" for model_id, _ in ranked_model_ids[:WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT]
)
return tuple(dict.fromkeys((*served_probes, *provider_probes)))
class HealthCheckHelpers:
@staticmethod
async def ahealth_check_wildcard_models(
@ -99,23 +160,16 @@ class HealthCheckHelpers:
litellm_logging_obj: "Logging",
) -> dict:
from litellm import acompletion
from litellm.litellm_core_utils.llm_request_utils import (
pick_cheapest_chat_models_from_llm_provider,
)
# this is a wildcard model, we need to pick a random model from the provider
cheapest_models = pick_cheapest_chat_models_from_llm_provider(custom_llm_provider=custom_llm_provider, n=3)
if len(cheapest_models) == 0:
probe_models: Final = wildcard_health_check_models(wildcard=model, custom_llm_provider=custom_llm_provider)
if len(probe_models) == 0:
raise Exception(
f"Unable to health check wildcard model for provider {custom_llm_provider}. Add a model on your config.yaml or contribute here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json"
f"Unable to health check wildcard model {model} for provider {custom_llm_provider}: the cost map has "
"no chat model for this provider. Set model_info.health_check_model to a model this wildcard serves"
)
if len(cheapest_models) > 1:
fallback_models = cheapest_models[1:] # Pick the last 2 models from the shuffled list
else:
fallback_models = None
model_params["model"] = cheapest_models[0]
model_params["model"] = probe_models[0]
model_params["litellm_logging_obj"] = litellm_logging_obj
model_params["fallbacks"] = fallback_models
model_params["fallbacks"] = list(probe_models[1:]) or None
model_params["max_tokens"] = model_params.get("max_tokens", 16) # GPT-5 models require max_output_tokens >= 16
await acompletion(**model_params)
return {}

View file

@ -1,6 +1,11 @@
import datetime
from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from typing import Final, cast
from pydantic import TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider
@ -151,19 +156,52 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1):
Returns:
list[str]: A list of the n cheapest chat models.
"""
if custom_llm_provider not in litellm.models_by_provider:
return []
return [ranked_model.model for ranked_model in rank_chat_models_from_llm_provider(custom_llm_provider)[:n]]
@dataclass(frozen=True, slots=True)
class RankedChatModel:
model: str
route_rank: int
cost: float
is_past_deprecation_date: bool
def rank_chat_models_from_llm_provider(custom_llm_provider: str) -> tuple[RankedChatModel, ...]:
known_models: Final[Iterable[str]] = litellm.models_by_provider.get(custom_llm_provider, ())
ranked_rows: Final = (
(route_rank, cost, model)
for model, cost in _chat_row_costs(known_models, custom_llm_provider)
ranked_models: Final = (
RankedChatModel(
model=model, route_rank=route_rank, cost=cost, is_past_deprecation_date=is_past_deprecation_date
)
for model, cost, is_past_deprecation_date in _chat_row_costs(known_models, custom_llm_provider)
if (route_rank := _route_rank(model, custom_llm_provider)) is not None
)
return [model for _, _, model in sorted(ranked_rows, key=lambda row: row[:2])[:n]]
return tuple(sorted(ranked_models, key=lambda ranked_model: (ranked_model.route_rank, ranked_model.cost)))
def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float]]:
class _CostMapDeprecation(TypedDict):
deprecation_date: NotRequired[ReadOnly[object]]
_COST_MAP_DEPRECATION: Final = TypeAdapter(_CostMapDeprecation)
def _is_past_deprecation_date(cost_map_key: str) -> bool:
deprecation_date: Final = _COST_MAP_DEPRECATION.validate_python(
litellm.model_cost.get(cost_map_key) or {} # pyright: ignore[reportUnknownMemberType] # untyped cost map JSON
).get("deprecation_date")
if not isinstance(deprecation_date, str):
return False
try:
return (
datetime.date.fromisoformat(deprecation_date)
<= datetime.datetime.now(datetime.timezone.utc).astimezone().date()
)
except ValueError:
return False
def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float, bool]]:
for model in models:
try:
model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
@ -171,7 +209,11 @@ def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator
continue
if model_info.get("mode") != "chat":
continue
yield model, (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0)
yield (
model,
(model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0),
_is_past_deprecation_date(model_info["key"]),
)
def _route_rank(model: str, custom_llm_provider: str) -> int | None:

View file

@ -5878,38 +5878,36 @@ def function_call_prompt(
]
def response_schema_prompt(model: str, response_schema: dict, llm_provider: str) -> str:
def response_schema_prompt(model: str, response_schema: Mapping[str, object], llm_provider: str) -> str:
"""
Decides if a user-defined custom prompt or default needs to be used
Returns the prompt str that's passed to the model as a user message
"""
custom_prompt_details: dict | None = None
response_schema_as_message: Final = [{"role": "user", "content": f"{response_schema}"}]
if f"{model}/response_schema_prompt" in litellm.custom_prompt_dict:
custom_prompt_details = litellm.custom_prompt_dict[
f"{model}/response_schema_prompt"
] # allow user to define custom response schema prompt by model
elif "response_schema_prompt" in litellm.custom_prompt_dict:
custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"]
if custom_prompt_details is not None:
return render_custom_prompt_template(
template=custom_prompt_details, messages=response_schema_as_message, model=model, llm_provider=llm_provider
)
else:
prompts: Final[Mapping[str, Mapping[str, object]]] = litellm.custom_prompt_dict
model_prompt_key: Final = f"{model}/response_schema_prompt"
custom_prompt_details: Final = prompts.get(
model_prompt_key if model_prompt_key in prompts else "response_schema_prompt"
)
if custom_prompt_details is None:
return default_response_schema_prompt(response_schema=response_schema)
return render_custom_prompt_template(
template=custom_prompt_details,
messages=[{"role": "user", "content": f"{response_schema}"}],
model=model,
llm_provider=llm_provider,
)
def default_response_schema_prompt(response_schema: dict) -> str:
def default_response_schema_prompt(response_schema: Mapping[str, object]) -> str:
"""
Used if provider/model doesn't support 'response_schema' param.
This is the default prompt. Allow user to override this with a custom_prompt.
"""
prompt_str: Final = f"""Use this JSON schema:
```json
{response_schema}
```json
{json.dumps(response_schema, default=str)}
```"""
return prompt_str

View file

@ -13,9 +13,10 @@ from typing_extensions import ReadOnly, Required
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import (
USAGE_COST_BILLED_PROVIDERS,
PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS,
chunk_usage,
provider_reported_usage_cost,
stamp_litellm_bill_on_usage,
)
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping
@ -1380,9 +1381,6 @@ def stream_chunk_builder_text_completion(
return TextCompletionResponse(**response)
_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value}) | USAGE_COST_BILLED_PROVIDERS
def _logged_provider(logging_obj: Optional["Logging"]) -> object:
if logging_obj is None:
return None
@ -1390,10 +1388,6 @@ def _logged_provider(logging_obj: Optional["Logging"]) -> object:
return provider
def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool:
return _logged_provider(logging_obj) in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS
def _stream_builder_response_cost(
response: ModelResponse, logging_obj: Optional["Logging"], chunks: Sequence[object]
) -> float | None:
@ -1404,9 +1398,6 @@ def _stream_builder_response_cost(
for chunk in reversed(chunks)
)
return next((cost for cost in reported if cost is not None), None)
usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None)
if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj):
return float(usage_cost)
if logging_obj is not None:
return None
provider_hint: Final = response._hidden_params.get( # pyright: ignore[reportPrivateUsage] # no public accessor
@ -1453,13 +1444,15 @@ def _set_stream_builder_response_cost(
def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
if logging_obj is None:
return
if isinstance(getattr(usage, "cost", None), (int, float)):
if _logged_provider(logging_obj) in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS and isinstance(
getattr(usage, "cost", None), (int, float)
):
return
computed_cost: Final = logging_obj.response_cost_calculator(result=response)
if not isinstance(computed_cost, (int, float)) or computed_cost <= 0:
if isinstance(computed_cost, bool) or not isinstance(computed_cost, (int, float)) or computed_cost < 0:
return
setattr(usage, "cost", computed_cost)
_set_hidden_response_cost(response, float(computed_cost))
if computed_cost > 0 or isinstance(getattr(usage, "cost", None), (int, float)):
stamp_litellm_bill_on_usage(response, usage, float(computed_cost))
def _field_of(value: object, name: str) -> object:

View file

@ -17,6 +17,7 @@ import anyio
import anyio.lowlevel
import httpx
import tiktoken
from pydantic import TypeAdapter
from tiktoken.model import MODEL_TO_ENCODING
from typing_extensions import ParamSpec, TypeVar
@ -69,6 +70,7 @@ from litellm.types.llms.anthropic import (
AnthropicMessagesToolResultParam,
AnthropicMessagesToolUseParam,
AnthropicSearchResultParam,
CompactionBlock,
)
from litellm.types.llms.openai import (
AllMessageValues,
@ -937,6 +939,8 @@ def _count_messages(
default_token_count,
params.pdf_pricing,
)
elif key == "provider_specific_fields":
num_tokens += sum(map(params.count_function, _compaction_summaries(value)))
elif key == "search_results" and isinstance(value, list):
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_search_results_text,
@ -951,6 +955,25 @@ def _count_messages(
return num_tokens
_OBJECT_LIST: Final = TypeAdapter(list[object])
def _compaction_summaries(provider_specific_fields: object) -> tuple[str, ...]:
fields: Final = as_str_mapping(provider_specific_fields)
blocks: Final = fields.get("compaction_blocks") if fields is not None else None
block_fields: Final = (
tuple(map(as_str_mapping, _OBJECT_LIST.validate_python(blocks))) if isinstance(blocks, list) else ()
)
return tuple(
content
for block in block_fields
if block is not None
and block.get("type") == "compaction"
and isinstance(content := block.get("content"), str)
and content
)
def _count_extra(
count_function: TokenCounterFunction,
tools: list[ChatCompletionToolParam] | None,
@ -1648,6 +1671,11 @@ def _count_anthropic_content(
return tokens
_TEXT_FIELD_BY_BLOCK_TYPE: Final = MappingProxyType(
{"text": "text", "thinking": "thinking", "redacted_thinking": "thinking", "compaction": "content"}
)
def _count_content_list(
count_function: TokenCounterFunction,
content_list: str
@ -1659,6 +1687,7 @@ def _count_content_list(
| AnthropicMessagesImageParam
| AnthropicMessagesDocumentParam
| AnthropicSearchResultParam
| CompactionBlock
],
use_default_image_token_count: bool,
default_token_count: int | None,
@ -1670,8 +1699,10 @@ def _count_content_list(
for c in content_list:
if isinstance(c, str):
num_tokens += count_function(c)
elif c["type"] == "text":
num_tokens += count_function(str(c.get("text", "")))
elif c["type"] in _TEXT_FIELD_BY_BLOCK_TYPE:
block_text = str(c.get(_TEXT_FIELD_BY_BLOCK_TYPE[c["type"]]) or "")
if block_text:
num_tokens += count_function(block_text)
elif c["type"] == "image_url":
image_url = c.get("image_url")
num_tokens += _count_image_tokens(image_url, use_default_image_token_count)
@ -1704,12 +1735,6 @@ def _count_content_list(
default_token_count,
pdf_pricing,
)
elif c["type"] in ("thinking", "redacted_thinking"):
# Claude extended thinking content block
# Count the thinking text and skip the opaque blobs (signature, redacted data)
thinking_text = str(c.get("thinking", ""))
if thinking_text:
num_tokens += count_function(thinking_text)
elif c["type"] == "tool_reference":
# Anthropic tool-search reference block: a lightweight pointer to
# a deferred tool, e.g. {"type": "tool_reference", "tool_name": ...}.
@ -1727,7 +1752,7 @@ def _count_content_list(
f"Invalid content item type: {content_type}. "
f"Expected str or dict with 'type' field "
f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, "
f"tool_reference, search_result)."
f"compaction, tool_reference, search_result)."
)
return num_tokens
except Exception as e:

View file

@ -24,7 +24,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.messages.adapters.transformation import (
OPENAI_MAX_TOOL_NAME_LENGTH,
@ -61,6 +61,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
unappliable_request_rewrite,
)
from litellm.types.llms.anthropic import (
ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX,
AllAnthropicToolsValues,
AnthropicMessagesRequest,
)
@ -153,6 +154,7 @@ class _TextAttrBlock(Protocol):
text: str
@runtime_checkable
class _WritableMessage(Protocol):
@overload
def get(self, key: str, /) -> object | None: ...
@ -195,6 +197,140 @@ class _ToolCallShape:
arguments: str
@dataclass(frozen=True, slots=True)
class _WebSearchString:
block_idx: int
result_idx: int | None
field: str
value: str
def _streamed_web_search_result_blocks(response: ModelResponse) -> tuple[Mapping[str, object], ...]:
choice: Final = response.choices[0] if response.choices else None
fields: Final = as_str_mapping(choice.message.provider_specific_fields) if isinstance(choice, Choices) else None
blocks: Final = (
(as_str_mapping(item) for item in iter_items(fields.get("web_search_results")))
if fields is not None and isinstance(fields.get("web_search_results"), list)
else ()
)
return tuple(
deepcopy(block) for block in blocks if block is not None and block.get("type") == "web_search_tool_result"
)
def _web_search_results(block: Mapping[str, object]) -> Iterable[object]:
return iter_items(block.get("content")) if isinstance(block.get("content"), list) else ()
def _web_search_result_strings(block_idx: int, result_idx: int, result: object) -> tuple[_WebSearchString, ...]:
fields: Final = as_str_mapping(result)
if fields is None or fields.get("type") != "web_search_result":
return ()
values: Final = ((field, fields.get(field)) for field in ("title", "url"))
return tuple(
_WebSearchString(block_idx, result_idx, field, value)
for field, value in values
if isinstance(value, str) and value
)
def _is_server_tool_use(block: Mapping[str, object]) -> bool:
block_id: Final = block.get("id")
match block.get("type"):
case "server_tool_use":
return True
case "tool_use":
return isinstance(block_id, str) and block_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX)
case _:
return False
def _web_search_block_strings(
block_idx: int, block: Mapping[str, object], *, with_query: bool
) -> tuple[_WebSearchString, ...]:
"""The strings a client reads off a native web search block: the ``server_tool_use`` query and each
``web_search_tool_result`` result's title and url."""
match block.get("type"):
case "server_tool_use" | "tool_use" if with_query and _is_server_tool_use(block):
tool_input: Final = as_str_mapping(block.get("input"))
query: Final = None if tool_input is None else tool_input.get("query")
return (_WebSearchString(block_idx, None, "query", query),) if isinstance(query, str) and query else ()
case "web_search_tool_result":
return tuple(
chain.from_iterable(
_web_search_result_strings(block_idx, result_idx, result)
for result_idx, result in enumerate(_web_search_results(block))
)
)
case _:
return ()
def _rewritten_web_search_result(
result_idx: int, result: object, rewrites: Mapping[tuple[int | None, str], str]
) -> object:
fields: Final = as_str_mapping(result)
if fields is None:
return result
return {
**fields,
**{field: rewrites[(result_idx, field)] for field in ("title", "url") if (result_idx, field) in rewrites},
}
def _rewritten_web_search_block(
block: Mapping[str, object], rewrites: Sequence[tuple[_WebSearchString, str]]
) -> Mapping[str, object]:
by_location: Final = MappingProxyType({(target.result_idx, target.field): value for target, value in rewrites})
if _is_server_tool_use(block):
query: Final = by_location.get((None, "query"))
tool_input: Final = as_str_mapping(block.get("input")) or MappingProxyType({})
return block if query is None else {**block, "input": {**tool_input, "query": query}}
if not isinstance(block.get("content"), list):
return block
return {
**block,
"content": [
_rewritten_web_search_result(result_idx, result, by_location)
for result_idx, result in enumerate(_web_search_results(block))
],
}
def _write_back_web_search_block(content_block: object, rewrites: Sequence[tuple[_WebSearchString, str]]) -> None:
block: Final = as_str_mapping(content_block)
if block is None or not isinstance(content_block, _WritableMessage):
return
rewritten: Final = _rewritten_web_search_block(block, rewrites)
writable: Final = _as_writable(content_block)
for key in ("input", "content"):
if key in rewritten:
writable[key] = rewritten[key]
def _split_web_search_texts(
guardrailed: Sequence[str], web_search_count: int, reply_count: int
) -> tuple[Sequence[str], Sequence[str]]:
split_at: Final = web_search_count if len(guardrailed) == web_search_count + reply_count else 0
return guardrailed[:split_at], guardrailed[split_at:]
def _web_search_rewrites_by_block(
strings: Sequence[_WebSearchString], guardrailed: Sequence[object]
) -> Mapping[int, tuple[tuple[_WebSearchString, str], ...]]:
rewrites: Final = tuple(
(target, value)
for target, value in zip(strings, guardrailed)
if isinstance(value, str) and value != target.value
)
return MappingProxyType(
{
block_idx: tuple(rewrite for rewrite in rewrites if rewrite[0].block_idx == block_idx)
for block_idx in {target.block_idx for target, _ in rewrites}
}
)
def _is_client_tool_use(block: Mapping[str, object]) -> bool:
return (
block.get("type") == "tool_use"
@ -289,6 +425,7 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge
_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object])
_DECODED_SSE_DATA: Final = TypeAdapter(object)
_REBUILT_TOOL_CALL_BLOCK_TYPES: Final = frozenset({"tool_use", "server_tool_use"})
_RELEASED_TOOL_USE_STOP: Final = (
b"event: message_delta\n"
b'data: {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": null}, '
@ -315,6 +452,51 @@ def _write_back_tool_use(
block["name"] = shape.name
def _streamed_query_strings(
tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]], tool_calls: Sequence[_ToolCallShape]
) -> tuple[_WebSearchString, ...]:
if len(tool_call_blocks) != len(tool_calls):
return ()
return tuple(
chain.from_iterable(
_web_search_block_strings(
index, {**block, "input": _rewritten_tool_use_input(tool_call.arguments)}, with_query=True
)
for (index, block), tool_call in zip(tool_call_blocks, tool_calls)
)
)
def _query_rewrites(strings: Sequence[_WebSearchString], guardrailed: Sequence[object]) -> Mapping[int, str]:
return MappingProxyType(
{
target.block_idx: value
for target, value in zip(strings, guardrailed)
if isinstance(value, str) and value != target.value
}
)
def _with_query(tool_call: _ToolCallShape, query: str) -> _ToolCallShape:
arguments: Final = _rewritten_tool_use_input(tool_call.arguments)
if arguments is None:
return tool_call
return _ToolCallShape(name=tool_call.name, arguments=json.dumps({**arguments, "query": query}))
def _with_query_rewrites(
tool_calls: tuple[_ToolCallShape, ...],
tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]],
query_rewrites: Mapping[int, str],
) -> tuple[_ToolCallShape, ...]:
if not query_rewrites or len(tool_call_blocks) != len(tool_calls):
return tool_calls
return tuple(
_with_query(tool_call, query_rewrites[index]) if index in query_rewrites else tool_call
for (index, _), tool_call in zip(tool_call_blocks, tool_calls)
)
@dataclass(frozen=True, slots=True)
class _SSEFieldRewrite:
"""One field of one nested section of a buffered SSE event, rewritten."""
@ -402,7 +584,12 @@ class AnthropicMessagesHandler(BaseTranslation):
def post_call_hook_response(self, response: object) -> object:
if not isinstance(response, ModelResponse):
return response
return self.adapter.translate_openai_response_to_anthropic(response)
translated: Final = self.adapter.translate_openai_response_to_anthropic(response)
result_blocks: Final = _streamed_web_search_result_blocks(response)
if not result_blocks:
return translated
content: Final = translated.get("content") or ()
return {**translated, "content": [*result_blocks, *content]}
@staticmethod
def _build_streaming_usage_response(
@ -1273,9 +1460,17 @@ class AnthropicMessagesHandler(BaseTranslation):
task_mappings,
tool_calls_to_check,
)
block_fields: Final = tuple(map(self._output_block_fields, response_content))
web_search_strings: Final = tuple(
chain.from_iterable(
_web_search_block_strings(content_idx, fields[1], with_query=True)
for content_idx, fields in enumerate(block_fields)
if fields is not None
)
)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check or tool_calls_to_check:
if texts_to_check or tool_calls_to_check or web_search_strings:
request_data = self._prepare_request_data(
request_data,
response,
@ -1284,7 +1479,7 @@ class AnthropicMessagesHandler(BaseTranslation):
)
inputs: Final = self._build_guardrail_inputs(
texts_to_check,
[*(target.value for target in web_search_strings), *texts_to_check],
images_to_check,
tool_calls_to_check,
response,
@ -1297,14 +1492,19 @@ class AnthropicMessagesHandler(BaseTranslation):
logging_obj=litellm_logging_obj,
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
web_search_rewrites, reply_rewrites = _split_web_search_texts(
guardrailed_inputs.get("texts", []), len(web_search_strings), len(texts_to_check)
)
# Step 3: Map guardrail responses back to original response structure
await self._apply_guardrail_responses_to_output(
response=response,
responses=guardrailed_texts,
responses=list(reply_rewrites),
task_mappings=task_mappings,
)
self._write_back_web_search_strings(
response_content, _web_search_rewrites_by_block(web_search_strings, web_search_rewrites)
)
verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response)
@ -1350,8 +1550,14 @@ class AnthropicMessagesHandler(BaseTranslation):
string_so_far = first_choice.message.content
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ())
guardrail_inputs: Final = GenericGuardrailAPIInputs()
if string_so_far:
guardrail_inputs["texts"] = [string_so_far]
leading_texts: Final = (string_so_far,) if isinstance(string_so_far, str) and string_so_far else ()
tool_call_blocks: Final = self._streamed_tool_call_blocks(responses_so_far)
streamed_query_strings: Final = _streamed_query_strings(tool_call_blocks, pre_guardrail_tool_calls)
streamed_result_strings: Final = self._streamed_web_search_result_strings(responses_so_far)
streamed_web_search_strings: Final = (*streamed_result_strings, *streamed_query_strings)
texts_to_check: Final = [*(target.value for target in streamed_web_search_strings), *leading_texts]
if texts_to_check:
guardrail_inputs["texts"] = texts_to_check
if tool_calls_list:
guardrail_inputs["tool_calls"] = tool_calls_list
@ -1374,29 +1580,42 @@ class AnthropicMessagesHandler(BaseTranslation):
responses_so_far, request_data
)
raise
guardrailed_texts: Final = _guardrailed_inputs.get("texts")
web_search_rewrites, reply_rewrites = _split_web_search_texts(
_guardrailed_inputs.get("texts") or [], len(streamed_web_search_strings), len(leading_texts)
)
if (
deliver_ended_stream_rewrites
and isinstance(string_so_far, str)
and string_so_far
and guardrailed_texts
and guardrailed_texts[0] != string_so_far
and reply_rewrites
and reply_rewrites[0] != string_so_far
):
self._write_ended_stream_text_rewrite(
responses_so_far,
guardrailed_texts[0],
reply_rewrites[0],
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
if deliver_ended_stream_rewrites:
self._write_ended_stream_web_search_rewrites(
responses_so_far,
_web_search_rewrites_by_block(streamed_result_strings, web_search_rewrites),
)
if deliver_ended_stream_rewrites:
returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
self._write_ended_stream_tool_call_rewrites(
responses_so_far,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
post_guardrail_tool_calls=_tool_call_shapes(
returned_tool_calls
if isinstance(returned_tool_calls, list)
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
else tool_calls_list or ()
post_guardrail_tool_calls=_with_query_rewrites(
_tool_call_shapes(
returned_tool_calls
if isinstance(returned_tool_calls, list)
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
else tool_calls_list or ()
),
tool_call_blocks,
_query_rewrites(
streamed_query_strings, web_search_rewrites[len(streamed_result_strings) :]
),
),
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
@ -1451,6 +1670,14 @@ class AnthropicMessagesHandler(BaseTranslation):
request_data["litellm_metadata"] = user_metadata
return request_data
@staticmethod
def _write_back_web_search_strings(
response_content: Sequence[object],
rewrites_by_block: Mapping[int, tuple[tuple[_WebSearchString, str], ...]],
) -> None:
for block_idx, rewrites in rewrites_by_block.items():
_write_back_web_search_block(_content_block_at(response_content, block_idx), rewrites)
@staticmethod
def _get_response_content(response: object) -> Sequence[object]:
"""Extract content list from a dict or object response."""
@ -1573,30 +1800,22 @@ class AnthropicMessagesHandler(BaseTranslation):
) -> None:
"""Deliver ended-stream guardrail tool-call rewrites by rewriting the
buffered chunks in place: the rebuilt response lists tool calls in the
order of the stream's ``tool_use`` blocks, so the nth rewritten call lands
on the nth block, its first ``input_json_delta`` carrying the full rewritten
order of the stream's ``tool_use`` and ``server_tool_use`` blocks, so the nth
rewritten call lands on the nth block, its first ``input_json_delta`` carrying the full rewritten
arguments, every later one blanked, and ``content_block_start`` carrying the
rewritten name. Blocks that do not line up with the rebuilt tool calls make
the rewrite undeliverable, so the pipeline executor discards it and releases
the original chunks."""
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
block_indices: Final = tuple(
index
for item in responses_so_far
for event in cls._iter_sse_events(item)
if event.get("type") == "content_block_start"
and isinstance(block := event.get("content_block"), Mapping)
and block.get("type") == "tool_use"
and isinstance(index := event.get("index"), int)
)
block_indices: Final = tuple(index for index, _ in cls._streamed_tool_call_blocks(responses_so_far))
if len(block_indices) != len(post_guardrail_tool_calls):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(
guardrail_name,
f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried "
f"{len(block_indices)} tool_use blocks",
f"{len(block_indices)} tool_use or server_tool_use blocks",
)
rewrites_by_block: Final = MappingProxyType(
{
@ -1636,6 +1855,60 @@ class AnthropicMessagesHandler(BaseTranslation):
),
)
@classmethod
def _streamed_web_search_result_strings(cls, responses_so_far: Sequence[object]) -> tuple[_WebSearchString, ...]:
events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far)
result_blocks: Final = (
(index, block)
for event in events
if event.get("type") == "content_block_start"
and isinstance(index := event.get("index"), int)
and (block := as_str_mapping(event.get("content_block"))) is not None
)
return tuple(
chain.from_iterable(
_web_search_block_strings(index, block, with_query=False) for index, block in result_blocks
)
)
@classmethod
def _write_ended_stream_web_search_rewrites(
cls,
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
rewrites_by_block: Mapping[int, tuple[tuple[_WebSearchString, str], ...]],
) -> None:
"""Deliver ended-stream guardrail rewrites of web search result titles and urls into the
``content_block_start`` that carries each whole ``web_search_tool_result`` block."""
if not rewrites_by_block:
return
def rewrite_result_block(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
index: Final = event.get("index")
block: Final = as_str_mapping(event.get("content_block"))
rewrites: Final = rewrites_by_block.get(index) if isinstance(index, int) else None
if event.get("type") != "content_block_start" or block is None or rewrites is None:
return None
return _SSEFieldRewrite(
"content_block", "content", _rewritten_web_search_block(block, rewrites).get("content")
)
cls._rewrite_ended_stream_events(responses_so_far, rewrite_result_block)
@classmethod
def _streamed_tool_call_blocks(
cls, responses_so_far: Sequence[object]
) -> tuple[tuple[int, Mapping[str, object]], ...]:
events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far)
return tuple(
(index, block)
for event in events
if event.get("type") == "content_block_start"
and (block := as_str_mapping(event.get("content_block"))) is not None
and isinstance(block_type := block.get("type"), str)
and block_type in _REBUILT_TOOL_CALL_BLOCK_TYPES
and isinstance(index := event.get("index"), int)
)
@classmethod
def _streamed_tool_use_ids(cls, responses_so_far: Sequence[object]) -> Mapping[int, str]:
"""The id of each buffered ``tool_use`` block, by its content block index."""
@ -1738,8 +2011,10 @@ class AnthropicMessagesHandler(BaseTranslation):
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far)
streamed_text: Final = self.get_streaming_string_so_far(responses_so_far)
result_strings: Final = self._streamed_web_search_result_strings(responses_so_far) if stream_ended else ()
return StreamingScanKey(
texts=(self.get_streaming_string_so_far(responses_so_far),),
texts=(*(target.value for target in result_strings), streamed_text),
tool_calls=tool_use_fingerprints if stream_ended else (),
stream_ended=stream_ended,
tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended,

View file

@ -8,12 +8,13 @@ import re
from collections.abc import Container, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from itertools import chain
from types import MappingProxyType
from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar
from urllib.parse import quote, urlparse
import httpx
from pydantic import BeforeValidator, ConfigDict, Field, StrictBool, TypeAdapter, ValidationError
from pydantic import BeforeValidator, ConfigDict, Field, JsonValue, StrictBool, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
@ -24,6 +25,7 @@ from litellm.constants import (
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
)
from litellm.exceptions import UnsupportedParamsError
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_file_ids_from_messages,
is_encrypted_reasoning_block,
@ -46,6 +48,7 @@ from litellm.llms.anthropic.wif import (
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.anthropic import (
ANTHROPIC_BETA_HEADER_VALUES,
ANTHROPIC_HOSTED_TOOLS,
ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER,
ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,
@ -118,6 +121,7 @@ _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
_DECODED_JSON: Final = TypeAdapter(object)
_JSON_VALUE: Final = TypeAdapter(JsonValue)
_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/")
@ -177,24 +181,68 @@ def is_anthropic_messages_url(url: str) -> bool:
return parsed_url.hostname == "api.anthropic.com" or parsed_url.path.removesuffix("/").endswith("/v1/messages")
def _message_blocks(message: object) -> Sequence[object]:
fields: Final = as_str_mapping(message)
blocks: Final = validated_object_list(fields.get("content")) if fields is not None else None
return blocks or ()
def _is_signed_compaction_block(block: object) -> bool:
fields: Final = as_str_mapping(block)
if fields is None or fields.get("type") != "compaction":
return False
signature: Final = fields.get("signature")
return isinstance(signature, str) and bool(signature)
def carries_signed_compaction_block(messages: Sequence[object]) -> bool:
return any(map(_is_signed_compaction_block, chain.from_iterable(map(_message_blocks, messages))))
def requires_native_compaction_beta(
custom_llm_provider: str,
optional_params: Mapping[str, object],
messages: Sequence[object],
) -> bool:
return custom_llm_provider == "anthropic" and (
optional_params.get("compaction") is not None
or any(
isinstance(block, Mapping)
and block.get("type") == "compaction"
and isinstance(block.get("signature"), str)
and bool(block.get("signature"))
for message in messages
if isinstance(message, Mapping)
for content in (message.get("content"),)
if isinstance(content, (list, tuple))
for block in content
optional_params.get("compaction") is not None or carries_signed_compaction_block(messages)
)
def anthropic_context_management(context_management: JsonValue) -> JsonValue:
if not isinstance(context_management, list):
return context_management
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
edits: Final = AnthropicConfig.map_openai_context_management_to_anthropic(
[entry for entry in context_management if isinstance(entry, dict)]
)
return context_management if edits is None else _JSON_VALUE.validate_python(edits)
def _context_management_edits(spec: JsonValue) -> Sequence[object]:
mapping: Final = as_str_mapping(spec)
return (validated_object_list(mapping.get("edits")) if mapping is not None else None) or ()
def context_management_with_edits(context_management: JsonValue) -> JsonValue:
spec: Final = anthropic_context_management(context_management)
return spec if _context_management_edits(spec) else None
def context_management_betas(context_management: JsonValue) -> tuple[str, ...]:
edits: Final = _context_management_edits(anthropic_context_management(context_management))
edit_types: Final = tuple(edit.get("type") for edit in map(as_str_mapping, edits) if edit is not None)
return tuple(
beta
for beta, used in (
(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, "compact_20260112" in edit_types),
(
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
any(edit_type != "compact_20260112" for edit_type in edit_types),
),
)
if used
)
@ -1367,8 +1415,6 @@ class AnthropicModelInfo(BaseLLMModelInfo):
if is_vertex_request is True:
# Vertex AI requires web search beta header for web search to work
if web_search_tool_used:
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
headers["anthropic-beta"] = ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value
elif len(all_betas) > 0:
headers["anthropic-beta"] = ",".join(all_betas)

View file

@ -39,6 +39,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
tools: list[dict[str, JsonValue]] | None = None,
system: JsonValue = None,
optional_params: Mapping[str, JsonValue] | None = None,
context_management: JsonValue = None,
drop_params: bool = False,
) -> dict[str, JsonValue]:
"""
@ -70,6 +71,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
tools=tools,
system=system,
optional_params=optional_params,
context_management=context_management,
drop_params=drop_params,
)
@ -81,7 +83,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
verbose_logger.debug("Making request to: %s", endpoint_url)
# Get required headers
headers: Final = self.get_count_tokens_headers(auth_header)
headers: Final = self.get_count_tokens_headers(auth_header, messages, context_management)
# Use LiteLLM's async httpx client
async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC)

View file

@ -132,6 +132,7 @@ class AnthropicTokenCounter(BaseTokenCounter):
tools=tools,
system=system,
optional_params=count_options,
context_management=context_management.spec if context_management is not None else None,
drop_params=drops_unsupported_params(litellm_params.get("drop_params")),
)

View file

@ -27,11 +27,15 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import
)
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.anthropic.common_utils import (
carries_signed_compaction_block,
context_management_betas,
context_management_with_edits,
is_claude_model_target,
merge_anthropic_beta_headers,
tools_without_custom_field,
)
from litellm.llms.anthropic.wif import resolve_anthropic_base
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
@ -376,6 +380,7 @@ class AnthropicCountTokensConfig:
tools: list[dict[str, JsonValue]] | None = None,
system: JsonValue = None,
optional_params: Mapping[str, JsonValue] | None = None,
context_management: JsonValue = None,
drop_params: bool = False,
) -> dict[str, JsonValue]: # mutable-ok: provider transport requires JSON dictionaries
"""
@ -408,6 +413,7 @@ class AnthropicCountTokensConfig:
drop_custom=drop_params,
),
),
("context_management", context_management_with_edits(context_management)),
)
if value is not None
}
@ -419,18 +425,33 @@ class AnthropicCountTokensConfig:
)
)
def get_count_tokens_headers(self, auth_header: Mapping[str, str]) -> dict[str, str]:
def get_count_tokens_headers(
self,
auth_header: Mapping[str, str],
messages: Sequence[Mapping[str, JsonValue]],
context_management: JsonValue = None,
) -> dict[str, str]:
"""The count-tokens headers around a resolved Anthropic auth header
(``AnthropicModelInfo.get_auth_header``): x-api-key for a static key, an Authorization
bearer for ``ANTHROPIC_AUTH_TOKEN`` and for sk-ant-oat tokens, whose mandatory oauth beta
merges with the token-counting beta instead of replacing it."""
merges with the token-counting beta instead of replacing it. A history that carries a
signed compaction block also needs the compaction beta, or Anthropic rejects the block
as an unknown content type, and ``context_management`` edits need the betas inference sends."""
compaction_betas: Final = (
(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value,)
if carries_signed_compaction_block(messages)
else ()
)
request_betas: Final = (
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
*compaction_betas,
*context_management_betas(context_management),
)
return {
"Content-Type": "application/json",
"anthropic-version": "2023-06-01",
**auth_header,
"anthropic-beta": merge_anthropic_beta_headers(
auth_header.get("anthropic-beta"), ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
),
"anthropic-beta": merge_anthropic_beta_headers(auth_header.get("anthropic-beta"), request_betas),
}
def validate_request(

View file

@ -12,11 +12,20 @@ import json
from collections.abc import Mapping, Sequence
from typing import Final, cast
from pydantic import TypeAdapter
from litellm.llms.anthropic.messages.native.utils import anthropic_message_delta
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
_CITATIONS: Final = TypeAdapter(tuple[object, ...])
def _citations_of(block: Mapping[str, object]) -> tuple[object, ...]:
citations: Final = block.get("citations")
return _CITATIONS.validate_python(citations) if isinstance(citations, list) else ()
class FakeAnthropicMessagesStreamIterator:
"""
@ -47,16 +56,30 @@ class FakeAnthropicMessagesStreamIterator:
def _create_content_block_chunks(self, block_dict: Mapping[str, object], index: int) -> list[bytes]:
"""Build SSE chunks for a single content block."""
chunks: Final = []
chunks: Final = list[bytes]()
block_type: Final = block_dict.get("type")
if block_type == "text":
citations: Final = _citations_of(block_dict)
content_block_start = {
"type": "content_block_start",
"index": index,
"content_block": {"type": "text", "text": ""},
"content_block": {"type": "text", "text": "", "citations": ()}
if citations
else {"type": "text", "text": ""},
}
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
citation_deltas: Final = (
{
"type": "content_block_delta",
"index": index,
"delta": {"type": "citations_delta", "citation": citation},
}
for citation in citations
)
chunks.extend(
f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n".encode() for delta in citation_deltas
)
text: Final = block_dict.get("text", "")
content_block_delta = {
"type": "content_block_delta",
@ -89,36 +112,23 @@ class FakeAnthropicMessagesStreamIterator:
}
chunks.append(f"event: content_block_delta\ndata: {json.dumps(signature_delta)}\n\n".encode())
elif block_type == "redacted_thinking":
content_block_start = {
elif block_type in ("tool_use", "server_tool_use"):
empty_input: Final[Mapping[str, object]] = {}
tool_use_start: Final = {
"type": "content_block_start",
"index": index,
"content_block": {"type": "redacted_thinking"},
"content_block": {**block_dict, "input": empty_input},
}
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
elif block_type == "tool_use":
content_block_start = {
"type": "content_block_start",
"index": index,
"content_block": {
"type": "tool_use",
"id": block_dict.get("id"),
"name": block_dict.get("name"),
"input": {},
},
}
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
input_data: Final = block_dict.get("input", {})
content_block_delta = {
chunks.append(f"event: content_block_start\ndata: {json.dumps(tool_use_start)}\n\n".encode())
tool_use_delta: Final = {
"type": "content_block_delta",
"index": index,
"delta": {
"type": "input_json_delta",
"partial_json": json.dumps(input_data),
"partial_json": json.dumps(block_dict.get("input", {})),
},
}
chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode())
chunks.append(f"event: content_block_delta\ndata: {json.dumps(tool_use_delta)}\n\n".encode())
else:
passthrough_start: Final = {

View file

@ -43,6 +43,7 @@ from ...common_utils import (
OBJECT_LIST_ADAPTER,
AnthropicError,
AnthropicModelInfo,
context_management_betas,
is_claude_model_target,
merge_anthropic_beta_headers,
optionally_handle_anthropic_oauth,
@ -745,6 +746,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
)
if transformed_context_management is not None:
anthropic_messages_optional_request_params["context_management"] = transformed_context_management
elif drops_unsupported_params(litellm_params.drop_params):
anthropic_messages_optional_request_params.pop("context_management")
####### get required params for all anthropic messages requests ######
# Lazy %s: the f-string previously stringified the entire messages
@ -873,28 +876,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
if requires_native_compaction_beta(custom_llm_provider, optional_params, messages):
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value)
# Check for context management
context_management_param: Final = optional_params.get("context_management")
if context_management_param is not None:
# Check edits array for compact_20260112 type
edits: Final = context_management_param.get("edits", ())
has_compact = False
has_other = False
for edit in edits:
edit_type = edit.get("type", "")
if edit_type == "compact_20260112":
has_compact = True
else:
has_other = True
# Add compact header if any compact edits exist
if has_compact:
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value)
# Add context management header if any other edits exist
if has_other:
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value)
beta_values.update(context_management_betas(optional_params.get("context_management")))
# Check for structured outputs. Anthropic's newer request shape nests
# the schema under output_config.format; the older top-level

View file

@ -42,6 +42,7 @@ from ..openai.common_utils import (
build_output_token_limit_response,
is_output_token_limit_error,
model_dump_provider_response,
upstream_error_response,
)
from .common_utils import (
AzureOpenAIError,
@ -399,6 +400,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
raise AzureOpenAIError(
status_code=status_code,
message=str(e),
response=upstream_error_response(e),
headers=error_headers,
body=error_body,
)
@ -662,6 +664,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
raise AzureOpenAIError(
status_code=status_code,
message=message,
response=upstream_error_response(e),
headers=error_headers,
body=error_body,
)
@ -855,7 +858,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
error_text = error_response.text
raise AzureOpenAIError(status_code=status_code, message=error_text, headers=error_headers)
raise AzureOpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
)
async def make_async_azure_httpx_request(
self,
@ -1190,7 +1198,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
except Exception as e:
error_code: Final[object] = getattr(e, "status_code", None)
if error_code is not None:
raise AzureOpenAIError(status_code=error_code, message=str(e))
raise AzureOpenAIError(status_code=error_code, message=str(e), response=upstream_error_response(e))
else:
raise AzureOpenAIError(status_code=500, message=str(e))

View file

@ -8,6 +8,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.prompt_templates.factory import prompt_factory
from litellm.utils import CustomStreamWrapper, ModelResponse, TextCompletionResponse
from ...openai.common_utils import upstream_error_response
from ...openai.completion.transformation import OpenAITextCompletionConfig
from ..common_utils import AzureOpenAIError, BaseAzureLLM
@ -183,7 +184,12 @@ class AzureTextCompletion(BaseAzureLLM):
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers)
raise AzureOpenAIError(
status_code=status_code,
message=str(e),
headers=error_headers,
response=upstream_error_response(e),
)
async def acompletion(
self,
@ -244,7 +250,12 @@ class AzureTextCompletion(BaseAzureLLM):
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers)
raise AzureOpenAIError(
status_code=status_code,
message=str(e),
headers=error_headers,
response=upstream_error_response(e),
)
def streaming(
self,
@ -355,4 +366,9 @@ class AzureTextCompletion(BaseAzureLLM):
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers)
raise AzureOpenAIError(
status_code=status_code,
message=str(e),
headers=error_headers,
response=upstream_error_response(e),
)

View file

@ -1,5 +1,7 @@
from typing import Final
import httpx
from litellm.exceptions import ContentPolicyViolationError
@ -14,6 +16,7 @@ class AzureOpenAIExceptionMapping:
model: str,
extra_information: str,
original_exception: Exception,
response: httpx.Response | None,
) -> ContentPolicyViolationError:
"""
Create a content policy violation error
@ -39,7 +42,7 @@ class AzureOpenAIExceptionMapping:
llm_provider="azure",
model=model,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
response=response,
provider_specific_fields={
# Preserve legacy key for backward compatibility.
"innererror": inner_error,

View file

@ -6,15 +6,19 @@ from typing import TYPE_CHECKING, Final, cast
import httpx
from httpx import Response
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.litellm_core_utils.prompt_templates.common_utils import (
content_parts,
convert_content_list_to_str,
drop_tool_reference_parts_from_tool_messages,
filter_value_from_dict,
is_pdf_file_part,
prompt_cache_breakpoint_for_wire,
with_prompt_cache_breakpoint,
)
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.azure_ai.common_utils import (
@ -27,7 +31,7 @@ from litellm.llms.openai.common_utils import drop_params_from_unprocessable_enti
from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator, OpenAIConfig
from litellm.llms.xai.chat.transformation import XAIChatConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.openai import AllMessageValues, PromptCacheBreakpoint
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ModelResponse, ModelResponseStream, ProviderField
from litellm.utils import add_path_to_api_base, supports_tool_choice
@ -62,6 +66,72 @@ def _has_list_only_content_part(message: Mapping[str, object]) -> bool:
return any(_is_list_only_content_part(part) for part in content_parts(message))
_CONTENT_PARTS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
def _text_parts(message: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
content: Final = message.get("content")
if not isinstance(content, list):
return ()
try:
parts: Final = _CONTENT_PARTS_ADAPTER.validate_python(content)
except ValidationError:
return ()
return tuple(part for part in parts if part.get("type") == "text")
def _with_only_text_parts(message: AllMessageValues) -> AllMessageValues:
if _has_list_only_content_part(message) or not (text_parts := _text_parts(message)):
return message
text_only: Final = {**message, "content": list(text_parts)}
return cast(AllMessageValues, text_only) # cast-ok: the same message, text parts only
_FIELDS: Final = TypeAdapter(Mapping[str, object])
def _foundry_breakpoint(marker: object) -> PromptCacheBreakpoint | None:
if prompt_cache_breakpoint_for_wire(marker, drop_params=True) is None:
return None
return PromptCacheBreakpoint(mode="explicit")
def _part_with_foundry_breakpoint(part: object) -> object:
if not isinstance(part, Mapping):
return part
fields: Final = _FIELDS.validate_python(part)
if "prompt_cache_breakpoint" not in fields:
return fields
unmarked: Final = {key: value for key, value in fields.items() if key != "prompt_cache_breakpoint"}
return with_prompt_cache_breakpoint(unmarked, _foundry_breakpoint(fields["prompt_cache_breakpoint"]))
def _with_breakpoints_for_foundry(message: AllMessageValues, drop_params: bool) -> AllMessageValues:
renamed: Final = AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(_with_only_text_parts(message))
parts: Final = content_parts(renamed) if drop_params else ()
if not parts:
return renamed
foundry_message: Final = {**renamed, "content": [_part_with_foundry_breakpoint(part) for part in parts]}
return cast(AllMessageValues, foundry_message) # cast-ok: the same message, its breakpoints in Foundry's shape
def _deployment_reads_breakpoints(model: str, litellm_params: Mapping[str, object]) -> bool:
base_model: Final = litellm_params.get("base_model")
named_models: Final = (base_model, model) if isinstance(base_model, str) and base_model else (model,)
model_info: Final = litellm_params.get("model_info")
if not isinstance(model_info, Mapping):
return any(
AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(name, "azure_ai") for name in named_models
)
deployment_flag: Final = _FIELDS.validate_python(model_info).get("supports_prompt_cache_breakpoint")
if isinstance(deployment_flag, bool):
return deployment_flag
return any(
AnthropicCacheControlHook.catalog_targets_openai_prompt_cache_breakpoint(name, "azure_ai")
for name in named_models
)
class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator):
def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract
usage: Final = chunk.get("usage")
@ -251,6 +321,8 @@ class AzureAIStudioConfig(OpenAIConfig):
self,
messages: list[AllMessageValues],
model: str,
reads_breakpoints: bool | None = None,
drop_params: bool = False,
) -> list:
"""
- Azure AI Studio doesn't support content as a list. This handles:
@ -260,13 +332,31 @@ class AzureAIStudioConfig(OpenAIConfig):
Azure AI Foundry backends set additionalProperties=false and reject
these with "Extra inputs are not permitted", which breaks multi-turn
Anthropic-format clients that echo thinking blocks back as history.
2. Transforms list content to a string.
2. Transforms list content to a string, except on deployments that read OpenAI explicit
prompt cache breakpoints (the deployment's own model_info opt-in, else the cost-map entry
of its base_model or name, as loaded and without a sibling deployment's opt-in): only a
content part can carry the breakpoint (a message-level one is silently ignored), so those
keep their text parts, with each Anthropic `cache_control` mark renamed to a breakpoint the
way the openai provider renames it. Every other part is dropped, exactly as the string
conversion drops it, before the marks are renamed, so a message-level mark lands on the
last text part. Under drop_params each breakpoint goes out as `{"mode": "explicit"}`, the only
shape Foundry takes, and a malformed one is dropped.
3. If message contains an image, audio, or PDF file part, send as is (user-intended)
Operates on a deep copy so the caller's messages keep their thinking blocks
and provider metadata, which a fallback to another provider still needs.
"""
stripped_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages))
deployment_reads_breakpoints: Final = (
AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(model, "azure_ai")
if reads_breakpoints is None
else reads_breakpoints
)
copied_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages))
stripped_messages: Final = (
[_with_breakpoints_for_foundry(message, drop_params) for message in copied_messages]
if deployment_reads_breakpoints
else copied_messages
)
for message in stripped_messages:
message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy
for field in NON_OPENAI_SPEC_MESSAGE_FIELDS:
@ -275,6 +365,10 @@ class AzureAIStudioConfig(OpenAIConfig):
if _has_list_only_content_part(message):
continue
if deployment_reads_breakpoints and (text_parts := _text_parts(message)):
message_dict["content"] = list(text_parts)
continue
texts = convert_content_list_to_str(message=message)
if texts:
message["content"] = texts
@ -329,7 +423,16 @@ class AzureAIStudioConfig(OpenAIConfig):
if extra_body and isinstance(extra_body, dict):
optional_params.update(extra_body)
optional_params.pop("max_retries", None)
return super().transform_request(model, messages, optional_params, litellm_params, headers)
return {
"model": model,
"messages": self._transform_messages(
messages,
model,
_deployment_reads_breakpoints(model, _FIELDS.validate_python(litellm_params)),
bool(litellm_params.get("drop_params") or litellm.drop_params),
),
**optional_params,
}
def transform_response(
self,

View file

@ -26,6 +26,10 @@ class BigModelChatConfig(ZAIChatConfig):
def custom_llm_provider(self) -> str:
return "bigmodel"
@property
def prompts_unsupported_json_schema(self) -> bool:
return False
def get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]:
return get_bigmodel_url(api_base, "chat"), get_bigmodel_api_key(api_key)

View file

@ -29,6 +29,7 @@ if TYPE_CHECKING:
from aiohttp import ClientSession
import litellm
from litellm.exceptions import error_response_without_credentials
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object
from litellm.litellm_core_utils.token_counter import token_counter, usage_tokenizer_model
@ -148,6 +149,14 @@ class OpenAIError(BaseLLMException):
)
def upstream_error_response(error: Exception) -> httpx.Response | None:
if not isinstance(error, (openai.APIStatusError, httpx.HTTPStatusError)):
return None
return error_response_without_credentials(
error.response.status_code, error.response, "https://docs.litellm.ai/docs"
)
_DECODED_JSON: Final = TypeAdapter(object)

View file

@ -10,7 +10,7 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser
from litellm.types.utils import LlmProviders, ModelResponse, TextCompletionResponse
from litellm.utils import ProviderConfigManager
from ..common_utils import BaseOpenAILLM, OpenAIError
from ..common_utils import BaseOpenAILLM, OpenAIError, upstream_error_response
from ..workload_identity import build_async_openai_client, build_openai_client
from .transformation import OpenAITextCompletionConfig
@ -164,7 +164,12 @@ class OpenAITextCompletion(BaseLLM):
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
raise OpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
)
async def acompletion(
self,
@ -220,7 +225,12 @@ class OpenAITextCompletion(BaseLLM):
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
raise OpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
)
def streaming(
self,
@ -261,7 +271,12 @@ class OpenAITextCompletion(BaseLLM):
error_response: object = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
raise OpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
)
streamwrapper: Final = CustomStreamWrapper(
completion_stream=response,
model=model,

View file

@ -14,7 +14,7 @@ from litellm.utils import ProviderConfigManager
from ...base_llm.image_variations.transformation import BaseImageVariationConfig
from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj
from ..common_utils import OpenAIError
from ..common_utils import OpenAIError, upstream_error_response
from ..workload_identity import build_async_openai_client, build_openai_client
@ -129,7 +129,12 @@ class OpenAIImageVariationsHandler:
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
raise OpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
)
def image_variations(
self,
@ -245,4 +250,9 @@ class OpenAIImageVariationsHandler:
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
raise OpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
)

View file

@ -64,6 +64,7 @@ from .common_utils import (
is_openai_backed_api_base,
is_output_token_limit_error,
model_dump_provider_response,
upstream_error_response,
)
from .workload_identity import (
build_async_openai_client,
@ -983,6 +984,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
raise OpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
body=error_body,
)
@ -1110,6 +1112,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
raise OpenAIError(
status_code=status_code,
message=message,
response=upstream_error_response(e),
headers=error_headers,
body=exception_body,
)
@ -1282,6 +1285,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
raise OpenAIError(
status_code=getattr(e, "status_code", 500),
message=str(e),
response=upstream_error_response(e),
headers=error_headers,
body=exception_body,
)
@ -1409,7 +1413,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
raise OpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
)
def embedding(
self,
@ -1499,7 +1508,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
error_response: Final[object] = getattr(e, "response", None)
if error_headers is None and error_response:
error_headers = getattr(error_response, "headers", None)
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
raise OpenAIError(
status_code=status_code,
message=error_text,
response=upstream_error_response(e),
headers=error_headers,
)
async def aimage_generation(
self,
@ -1666,7 +1680,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
original_response=str(e),
)
if hasattr(e, "status_code"):
raise OpenAIError(status_code=getattr(e, "status_code", 500), message=str(e))
raise OpenAIError(
status_code=getattr(e, "status_code", 500),
message=str(e),
response=upstream_error_response(e),
)
else:
raise OpenAIError(status_code=500, message=str(e))

View file

@ -3,7 +3,7 @@ from typing import Final
from litellm.exceptions import UnsupportedParamsError
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, context_management_betas
from litellm.llms.anthropic.messages.native.transformation import (
AnthropicMessagesConfig,
messages_carry_output_config,
@ -171,28 +171,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
if existing_beta:
beta_values.update(b.strip() for b in existing_beta.split(","))
# Check for context management
context_management_param: Final = optional_params.get("context_management")
if context_management_param is not None:
# Check edits array for compact_20260112 type
edits: Final = context_management_param.get("edits", [])
has_compact = False
has_other = False
for edit in edits:
edit_type = edit.get("type", "")
if edit_type == "compact_20260112":
has_compact = True
else:
has_other = True
# Add compact header if any compact edits exist
if has_compact:
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value)
# Add context management header if any other edits exist
if has_other:
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value)
beta_values.update(context_management_betas(optional_params.get("context_management")))
# Check for web search tool
for tool in tools:

View file

@ -1,11 +1,43 @@
from collections.abc import Mapping
from typing import Final
from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam, ChatCompletionUserMessage
from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
ZAI_API_BASE: Final = "https://api.z.ai/api/paas/v4"
TOKEN_LIMIT_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens"})
def _is_json_object(
value: object,
) -> TypeIs[Mapping[str, object]]: # guard-ok: cost-map rows and request params are JSON
return isinstance(value, Mapping)
def _json_field(value: object, key: str) -> object:
return value.get(key) if _is_json_object(value) else None
def _exceeds_output_limit(model: str, max_tokens: object) -> bool:
info: Final = lookup_model_info_for_registration(model=model, custom_llm_provider="zai")
limit: Final = None if isinstance(info, ModelInfoLookupMiss) else info.get("max_output_tokens")
return isinstance(max_tokens, int) and isinstance(limit, int) and max_tokens > limit
def _schema_without_native_support(model: str, response_format: object) -> Mapping[str, object] | None:
if _json_field(response_format, "type") != "json_schema":
return None
if litellm.supports_response_schema(model=model, custom_llm_provider="zai"):
return None
schema: Final = _json_field(_json_field(response_format, "json_schema"), "schema")
return schema if _is_json_object(schema) else None
class ZAIChatConfig(OpenAIGPTConfig):
@ -13,6 +45,10 @@ class ZAIChatConfig(OpenAIGPTConfig):
def custom_llm_provider(self) -> str | None:
return "zai"
@property
def prompts_unsupported_json_schema(self) -> bool:
return True
def _get_openai_compatible_provider_info(
self, api_base: str | None, api_key: str | None
) -> tuple[str | None, str | None]:
@ -40,9 +76,11 @@ class ZAIChatConfig(OpenAIGPTConfig):
# GLM/ZAI supports cache_control, so return messages and tools unchanged
return messages, tools
def get_supported_openai_params(self, model: str) -> list:
def get_supported_openai_params(self, model: str) -> list[str]:
base_params: Final = [
"max_tokens",
"max_completion_tokens",
"response_format",
"stream",
"stream_options",
"temperature",
@ -52,8 +90,6 @@ class ZAIChatConfig(OpenAIGPTConfig):
"tool_choice",
]
import litellm
try:
if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
base_params.append("thinking")
@ -61,3 +97,53 @@ class ZAIChatConfig(OpenAIGPTConfig):
pass
return base_params
def map_openai_params(
self,
non_default_params: Mapping[str, object],
optional_params: Mapping[str, object],
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: request pipeline adds parameters to the returned body
supported_params: Final = self.get_supported_openai_params(model)
sent: Final = {param: value for param, value in non_default_params.items() if param in supported_params}
max_completion_tokens: Final = sent.get("max_completion_tokens")
falls_back_to_max_tokens: Final = max_completion_tokens is None or (
(litellm.drop_params or drop_params) and _exceeds_output_limit(model, max_completion_tokens)
)
max_tokens: Final = sent.get("max_tokens") if falls_back_to_max_tokens else max_completion_tokens
return {
**optional_params,
**{param: value for param, value in sent.items() if param not in TOKEN_LIMIT_PARAMS},
**({} if max_tokens is None else {"max_tokens": max_tokens}),
}
def transform_request(
self,
model: str,
messages: list[AllMessageValues], # mutable-ok: OpenAIGPTConfig.transform_request signature
optional_params: dict[str, object], # mutable-ok: OpenAIGPTConfig.transform_request signature
litellm_params: dict[str, object], # mutable-ok: OpenAIGPTConfig.transform_request signature
headers: dict[str, str], # mutable-ok: OpenAIGPTConfig.transform_request signature
) -> dict[str, object]: # mutable-ok: inherited contract returns the request body
schema: Final = (
_schema_without_native_support(model, optional_params.get("response_format"))
if self.prompts_unsupported_json_schema
else None
)
schema_message: Final = (
()
if schema is None
else (
ChatCompletionUserMessage(
role="user", content=response_schema_prompt(model=model, response_schema=schema, llm_provider="zai")
),
)
)
return super().transform_request(
model,
[*messages, *schema_message],
optional_params if schema is None else {**optional_params, "response_format": {"type": "json_object"}},
litellm_params,
headers,
)

View file

@ -6014,11 +6014,14 @@
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 1.25e-06,
"input_cost_per_token": 2.5e-06,
"input_cost_per_token_batches": 1.25e-06,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_batches": 5e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
@ -6313,16 +6316,19 @@
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 1.5e-07,
"input_cost_per_token_batches": 7.5e-08,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6e-07,
"output_cost_per_token_batches": 3e-07,
"search_context_cost_per_query": {
"search_context_size_high": 0.014,
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014
},
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
@ -11303,16 +11309,19 @@
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8e-06,
"output_cost_per_token_batches": 4e-06,
"search_context_cost_per_query": {
"search_context_size_high": 0.014,
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014
},
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
@ -11411,11 +11420,14 @@
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 1.1e-06,
"input_cost_per_token_batches": 5.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.4e-06,
"output_cost_per_token_batches": 2.2e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -11513,16 +11525,19 @@
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 2.75e-07,
"input_cost_per_token": 1.1e-06,
"input_cost_per_token_batches": 5.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.4e-06,
"output_cost_per_token_batches": 2.2e-06,
"search_context_cost_per_query": {
"search_context_size_high": 0.014,
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014
},
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
@ -54512,6 +54527,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true
},
"zai/glm-5.3-flash": {
@ -54526,10 +54542,26 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"source": "https://docs.z.ai/guides/overview/pricing",
"supports_vision": true
},
"zai/glm-5v-turbo": {
"cache_read_input_token_cost": 2.4e-07,
"input_cost_per_token": 1.2e-06,
"litellm_provider": "zai",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://vercel.com/ai-gateway/models/glm-5v-turbo",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"zai/glm-5.1": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 2.6e-07,
@ -81059,5 +81091,14 @@
"supports_response_schema": true,
"supports_tool_choice": false,
"supports_vision": true
},
"baseten/decision-models/baseten-joey-2b-latest": {
"input_cost_per_token": 1e-08,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
}
}

View file

@ -8,7 +8,7 @@ from typing import Final
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from fastapi.responses import JSONResponse
from pydantic import JsonValue, ValidationError
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
@ -54,6 +54,7 @@ from litellm.proxy.litellm_pre_call_utils import (
should_auto_drop_params_for_agentic_cli,
)
from litellm.router import Router
from litellm.types.utils import TokenCountResponse
router: Final = APIRouter()
@ -334,6 +335,15 @@ def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params:
}
_PROVIDER_COUNT_REPLY: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
def _provider_context_management(response: TokenCountResponse) -> Mapping[str, object]:
original: Final = _PROVIDER_COUNT_REPLY.validate_python(response.original_response)
context_management: Final = original.get("context_management") if original is not None else None
return {} if context_management is None else {"context_management": context_management}
@router.post(
"/v1/messages/count_tokens",
tags=["[beta] Anthropic Messages Token Counting"],
@ -413,7 +423,11 @@ async def count_tokens(
log_llm_api_exception(e, litellm_call_id)
return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
else:
return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)}
return {
"input_tokens": token_response.total_tokens,
**_provider_context_management(token_response),
**count_source_fields(token_response),
}
@router.post(

View file

@ -446,7 +446,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request,
from litellm.proxy.litellm_pre_call_utils import (
_trace_id_from_traceparent, # pyright: ignore[reportPrivateUsage] # reuse W3C traceparent parser
get_chain_id_from_headers,
metadata_variable_name_for_route,
metadata_variable_name_for_request,
)
if get_chain_id_from_headers(dict(request.headers)):
@ -460,7 +460,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request,
traceparent: Final[str | None] = request.headers.get("traceparent")
if isinstance(traceparent, str) and _trace_id_from_traceparent(traceparent):
return True
metadata_variable_name: Final = metadata_variable_name_for_route(get_request_route(request))
metadata_variable_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body)
selected_metadata: Final[object] = request_body.get(metadata_variable_name)
if isinstance(selected_metadata, dict):
selected_metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(selected_metadata)
@ -1367,11 +1367,15 @@ async def common_checks(
if not skip_all_budget_checks:
# Key metadata.tags are injected into request_body here so the tag budget
# check can read them; this mutation must run before the gathered checks.
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
if valid_token is not None:
for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(request, request_body):
LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=route)
LiteLLMProxyRequestSetup.merge_caller_body_tags(
request_data=tag_target,
metadata_variable_name=metadata_variable_name_for_request(route, tag_target),
)
LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(request_data=tag_target, user_api_key_dict=valid_token)
LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_body)

View file

@ -18,14 +18,21 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.responses.main import strip_responses_routing_prefix
from litellm.router import Router
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager, answer_from_cached_provider_listings, get_valid_models
from litellm.utils import (
ProviderConfigManager,
answer_from_cached_provider_listings,
get_valid_models,
potential_model_names_in_lookup_order,
)
LLM_PROVIDER_PREFIXES: Final = frozenset(llm_provider.value for llm_provider in LlmProviders)
_PROVIDERS_SENDING_THE_DEPLOYMENT_PATH_AS_THE_MODEL_ID: Final = frozenset({LlmProviders.TOGETHER_AI.value})
_P: Final = ParamSpec("_P")
_R: Final = TypeVar("_R")
@ -424,6 +431,25 @@ def _routes_back_to(
)
def _resolves_to_catalog_model(model: str, provider: str, provider_model_ids: frozenset[str]) -> bool:
provider_id_prefix: Final = f"{provider}/"
resolved_model: Final = (
f"{provider_id_prefix}{strip_responses_routing_prefix(model.removeprefix(provider_id_prefix))}"
)
return any(
name.removeprefix(provider_id_prefix) in provider_model_ids
for name in potential_model_names_in_lookup_order(resolved_model, provider)
)
def _routes_to_catalog_model(
routes_to_model: Callable[[str, str], bool], provider: str, provider_model_ids: frozenset[str]
) -> Callable[[str, str], bool]:
return lambda listed_id, model: (
routes_to_model(listed_id, model) and _resolves_to_catalog_model(model, provider, provider_model_ids)
)
def _ids_routing_to_catalog_models(
name_pattern: str,
deployment_pattern: str,
@ -437,8 +463,15 @@ def _ids_routing_to_catalog_models(
return ()
provider_id_prefix: Final = f"{provider}/"
provider_model_ids: Final = tuple(catalog_id.removeprefix(provider_id_prefix) for catalog_id in catalog_ids)
deployment_model_pattern_without_provider: Final = deployment_pattern.removeprefix(provider_id_prefix)
catalog_pattern: Final = _without_routing_segments(
deployment_pattern.removeprefix(provider_id_prefix), provider, name_pattern, provider_model_ids
deployment_model_pattern_without_provider, provider, name_pattern, provider_model_ids
)
routing_prefix: Final = deployment_model_pattern_without_provider.removesuffix(catalog_pattern)
reaches_catalog_model: Final = (
_routes_to_catalog_model(routes_to_model, provider, frozenset(provider_model_ids))
if routing_prefix and provider in _PROVIDERS_SENDING_THE_DEPLOYMENT_PATH_AS_THE_MODEL_ID
else routes_to_model
)
deployment_model_pattern: Final = f"{provider_id_prefix}{catalog_pattern}"
deployment_regex: Final = re.compile(PatternMatchRouter().pattern_to_regex(deployment_model_pattern))
@ -462,7 +495,7 @@ def _ids_routing_to_catalog_models(
for listed_id, catalog_id in listed_ids
if listed_id
and _routes_back_to(
listed_id, name_regex, (deployment_model_pattern, catalog_id), deployment_pattern, routes_to_model
listed_id, name_regex, (deployment_model_pattern, catalog_id), deployment_pattern, reaches_catalog_model
)
)
)

View file

@ -2401,9 +2401,10 @@ class ProxyBaseLLMRequestProcessing:
version=version,
proxy_config=proxy_config,
)
request_fields: Final = _REQUEST_FIELDS.validate_python(self.data)
metadata_variable_name: Final = get_metadata_variable_name(request, request_fields)
if logging_auth is not None:
metadata_variable_name: Final = get_metadata_variable_name(request)
request_metadata: Final = self.data.get(metadata_variable_name)
request_metadata: Final = request_fields.get(metadata_variable_name)
if isinstance(request_metadata, dict):
self.data[metadata_variable_name] = {
**TypeAdapter(dict[str, object]).validate_python(request_metadata),
@ -2431,12 +2432,11 @@ class ProxyBaseLLMRequestProcessing:
# Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved
if queue_time_seconds is not None:
_metadata_variable_name: Final = get_metadata_variable_name(request)
if _metadata_variable_name not in self.data:
self.data[_metadata_variable_name] = {}
if not isinstance(self.data[_metadata_variable_name], dict):
self.data[_metadata_variable_name] = {}
self.data[_metadata_variable_name]["queue_time_seconds"] = queue_time_seconds
if metadata_variable_name not in self.data:
self.data[metadata_variable_name] = {}
if not isinstance(self.data[metadata_variable_name], dict):
self.data[metadata_variable_name] = {}
self.data[metadata_variable_name]["queue_time_seconds"] = queue_time_seconds
if isinstance(model, str):
reject_url_valued_destination("model", model)

View file

@ -42,7 +42,11 @@ from litellm.caching.redis_batch import (
active_request_redis_batch,
)
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.constants import (
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
NO_MAX_TOKENS_TPM_RESERVATION_FRACTION,
)
from litellm.cost_calculator import get_usage_object
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens
@ -508,11 +512,6 @@ REDIS_NODE_HASHTAG_NAME: Final = "all_keys"
# *some* output budget; these define that fallback estimate.
DEFAULT_MAX_TOKENS_ESTIMATE: Final = 4096
DEFAULT_CHARS_PER_TOKEN: Final = 4
# Fraction of the available output budget reserved as the upfront floor when
# the request omits max_tokens. Applied to both DEFAULT_MAX_TOKENS_ESTIMATE
# (baseline floor) and to the smallest configured TPM limit (capped floor for
# small per-tenant TPM caps).
_TPM_FLOOR_FRACTION: Final = 4
# Both embeddings and the Responses API put their prompt in data["input"],
# but only embeddings have no output tokens. Every "is this an embedding"
# check on data["input"] must exclude these call types, or a Responses call
@ -1003,10 +1002,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
per-tenant cap can't be tripped by the floor alone. Returns the
baseline floor when no limit is provided.
"""
baseline: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION
baseline: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)
if min_configured_tpm_limit is None:
return baseline
return min(baseline, max(1, min_configured_tpm_limit // _TPM_FLOOR_FRACTION))
return min(baseline, max(1, int(min_configured_tpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)))
@staticmethod
def _is_embedding_request(data: object, call_type: str | None) -> bool:
@ -1138,7 +1137,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if call_type in RESPONSES_API_CALL_TYPES
else base_capped_floor
)
baseline_floor: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION
baseline_floor: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)
is_embedding: Final = PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type)
if (
capped_floor >= baseline_floor
@ -3313,7 +3312,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if key_agent_id:
return key_agent_id
metadata: Final = data.get("metadata") or {}
return metadata.get("agent_id")
litellm_metadata: Final = data.get("litellm_metadata") or {}
return metadata.get("agent_id") or litellm_metadata.get("agent_id")
def _get_session_id_from_data(self, data: dict) -> str | None:
"""Extract session_id from request metadata or litellm_session_id."""

View file

@ -37,7 +37,7 @@ from litellm.proxy.db.db_spend_update_writer import (
get_llm_router,
)
from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope
from litellm.proxy.spend_tracking.spend_event import (
@ -137,19 +137,34 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping)
def _proxy_stamped_used_client_oauth_token(
request_data: Mapping[str, object], request_route: str | None
) -> bool | None:
def _proxy_stamped_metadata(request_data: Mapping[str, object], request_route: str | None) -> Mapping[str, object]:
proxy_bucket: Final = (
get_metadata_variable_name_from_kwargs(request_data)
if request_route is None
else metadata_variable_name_for_route(request_route)
else metadata_variable_name_for_request(request_route, request_data)
)
proxy_metadata: Final = request_data.get(proxy_bucket)
stamped: Final = proxy_metadata.get("used_client_oauth_token") if isinstance(proxy_metadata, dict) else None
return _str_keyed_fields(request_data.get(proxy_bucket))
def _proxy_stamped_used_client_oauth_token(
request_data: Mapping[str, object], request_route: str | None
) -> bool | None:
stamped: Final = _proxy_stamped_metadata(request_data, request_route).get("used_client_oauth_token")
return stamped if isinstance(stamped, bool) else None
_PROXY_STAMPED_REQUEST_FIELDS: Final = ("requester_ip_address", "user_agent")
def _proxy_stamped_request_fields(
request_data: Mapping[str, object], request_route: str | None
) -> Mapping[str, object]:
if request_route is None:
return {}
proxy_metadata: Final = _proxy_stamped_metadata(request_data, request_route)
return {field: proxy_metadata[field] for field in _PROXY_STAMPED_REQUEST_FIELDS if field in proxy_metadata}
def _str_keyed_fields(value: object) -> Mapping[str, object]:
return value if is_str_keyed_mapping(value) else {}
@ -455,7 +470,13 @@ class _ProxyDBLogger(CustomLogger):
_metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_fields, request_route)
existing_metadata: Final[dict] = request_data.get("metadata", None) or {}
existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata})
existing_metadata.update(
{
**_routed_metadata_fallback(request_fields),
**_proxy_stamped_request_fields(request_fields, request_route),
**_metadata,
}
)
litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
litellm_metadata: Final[ObjectMapping] = (

View file

@ -379,6 +379,8 @@ LITELLM_METADATA_ROUTES: Final = (
"interactions",
)
CURSOR_CHAT_COMPLETIONS_ROUTE: Final = "/cursor/chat/completions"
LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset(
{
"mask_input",
@ -636,7 +638,7 @@ def _has_tag_budget(tag: "LiteLLM_TagTable | None") -> bool:
return tag.litellm_budget_table.max_budget is not None
def _normalized_metadata_slot(
def normalized_metadata_slot(
request_data: MutableMapping[str, object], metadata_variable_name: str
) -> dict[str, object]:
"""Return the request's metadata slot as a dict, normalising it in place first.
@ -653,6 +655,12 @@ def _normalized_metadata_slot(
return normalized
def _holds_metadata_object(value: object) -> bool:
if value is None or isinstance(value, dict):
return True
return isinstance(value, str) and isinstance(_DECODED_JSON.validate_python(safe_json_loads(value)), dict)
def _metadata_object_or_error(field: str, value: object) -> Mapping[str, object] | ProxyException:
if value is None:
return {}
@ -863,7 +871,7 @@ def _strip_router_reserved_metadata(
)
def get_metadata_variable_name(request: Request) -> str:
def get_metadata_variable_name(request: Request, data: Mapping[str, object] | None = None) -> str:
"""
Helper to return what the "metadata" field should be called in the request data
@ -874,7 +882,10 @@ def get_metadata_variable_name(request: Request) -> str:
# Inline imports — auth_utils/route_checks participate in a proxy import cycle.
from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415
return metadata_variable_name_for_route(get_request_route(request))
route: Final = get_request_route(request)
if data is None:
return metadata_variable_name_for_route(route)
return metadata_variable_name_for_request(route, data)
_get_metadata_variable_name: Final = get_metadata_variable_name
@ -890,6 +901,25 @@ def metadata_variable_name_for_route(route: str) -> Literal["metadata", "litellm
return "metadata"
def is_chat_completions_body(data: Mapping[str, object]) -> bool:
messages: Final = data.get("messages")
if isinstance(messages, list) and messages:
return True
return "messages" in data and "input" not in data
def is_cursor_responses_request(route: str, body: Mapping[str, object]) -> bool:
return CURSOR_CHAT_COMPLETIONS_ROUTE in route and not is_chat_completions_body(body)
def metadata_variable_name_for_request(
route: str, body: Mapping[str, object]
) -> Literal["metadata", "litellm_metadata"]:
if is_cursor_responses_request(route, body):
return "litellm_metadata"
return metadata_variable_name_for_route(route)
def _promoted_trace_control_fields(
requester_metadata: Mapping[str, object],
litellm_metadata: Mapping[str, object],
@ -2309,7 +2339,9 @@ class LiteLLMProxyRequestSetup:
``_tag_max_budget_check`` resolve to ``litellm_metadata``, leaving
header tags invisible to per-tag budget enforcement.
"""
if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES):
if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES) or is_cursor_responses_request(
route, request_data
):
request_data.setdefault("litellm_metadata", {})
@staticmethod
@ -2380,7 +2412,7 @@ class LiteLLMProxyRequestSetup:
return
tags: Final = _STRING_TAGS_ADAPTER.validate_python(raw_tags)
slot: Final = metadata_key if metadata_key is not None else get_metadata_variable_name_from_kwargs(request_data)
metadata: Final = _normalized_metadata_slot(request_data, slot)
metadata: Final = normalized_metadata_slot(request_data, slot)
existing: Final = metadata.get("tags")
if existing and not isinstance(existing, list):
return
@ -2393,6 +2425,34 @@ class LiteLLMProxyRequestSetup:
tags=tuple(tag for tag in tags if tag not in sent_tags)
)
@staticmethod
def merge_caller_body_tags(
request_data: MutableMapping[str, object], # mutable-ok: the request body the tag merges write into
metadata_variable_name: str | None = None,
) -> None:
slot: Final = (
metadata_variable_name
if metadata_variable_name is not None
else get_metadata_variable_name_from_kwargs(request_data)
)
if slot == "metadata" or not _holds_metadata_object(request_data.get(slot)):
return
raw_caller_metadata: Final = request_data.get("metadata")
caller_metadata: Final = (
_DECODED_JSON.validate_python(safe_json_loads(raw_caller_metadata))
if isinstance(raw_caller_metadata, str)
else raw_caller_metadata
)
caller_tags: Final = caller_metadata.get("tags") if isinstance(caller_metadata, dict) else None
if not isinstance(caller_tags, list) or not caller_tags:
return
metadata: Final = normalized_metadata_slot(request_data, slot)
existing_tags: Final = metadata.get("tags")
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
request_tags=existing_tags if isinstance(existing_tags, list) else None,
tags_to_add=caller_tags,
)
@staticmethod
def apply_key_tags_pre_auth(
request_data: dict,
@ -2408,7 +2468,7 @@ class LiteLLMProxyRequestSetup:
return
_metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data)
metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name)
metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name)
existing_tags: Final = metadata.get("tags")
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
@ -2460,7 +2520,7 @@ class LiteLLMProxyRequestSetup:
# from (litellm_metadata vs metadata) so the merged tags are visible
# to _tag_max_budget_check.
_metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data)
metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name)
metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name)
existing_tags: Final = metadata.get("tags")
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
@ -2637,7 +2697,7 @@ async def add_litellm_data_to_request(
}
safe_add_api_version_from_query_params(data, request)
_metadata_variable_name: Final = get_metadata_variable_name(request)
_metadata_variable_name: Final = get_metadata_variable_name(request, _REQUEST_FIELDS.validate_python(data))
if data.get(_metadata_variable_name, None) is None:
data[_metadata_variable_name] = {}
@ -2974,17 +3034,7 @@ async def add_litellm_data_to_request(
tags_to_add=tags,
)
_caller_body_metadata: Final = data.get("metadata") if _metadata_variable_name != "metadata" else None
_caller_body_tags: Final = (
_caller_body_metadata.get("tags")
if isinstance(_caller_body_metadata, dict) and isinstance(_caller_body_metadata.get("tags"), list)
else None
)
if _caller_body_tags:
data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup.merge_tags( # rebind-ok: matches file idiom
request_tags=data[_metadata_variable_name].get("tags"),
tags_to_add=_caller_body_tags,
)
LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data, metadata_variable_name=_metadata_variable_name)
# caller_tags: the tags the request's headers and body carry (x-litellm-tags, body "tags", and body
# "metadata.tags" / "litellm_metadata.tags" read before the mapped tags are added below),

View file

@ -98,15 +98,11 @@ def _count_change(sent: tuple[object, ...] | None, returned: tuple[object, ...]
return (len(sent), len(returned))
def _tool_call_mismatch_reason(
sent: tuple[tuple[object, object], ...] | None, returned: tuple[tuple[object, object], ...] | None
) -> str | None:
if sent == returned:
return None
def _tool_call_count_mismatch_reason(sent: Sequence[object] | None, returned: Sequence[object] | None) -> str | None:
sent_count: Final = len(sent or ())
returned_count: Final = len(returned or ())
if sent_count == returned_count:
return "the legacy hook changed a tool call's name or arguments, which this path cannot write back"
return None
return f"the legacy hook returned {returned_count} tool calls for a stream that carried {sent_count}"
@ -205,11 +201,10 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
non-streaming hooks, an exception it raises ends the stream through the executor's
fail/error classification, and the response it hands back, or the one it changed in place
and returned ``None`` for, is re-scanned by the same translation so its texts reach the
client through the translation's ended-stream write-back. A
replacement whose scanned texts do not line up with the originals, or whose tool calls
differ from them, is undeliverable, so the executor releases the original chunks. A stream
that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as
long as the hook left the tool calls alone."""
client through the translation's ended-stream write-back. Tool calls the hook rewrote go to the translation's
tool-call write-back like a unified guardrail's. A replacement whose scanned texts or tool
calls do not line up with the originals by count is undeliverable, so the executor releases
the original chunks."""
def __init__(
self,
@ -258,14 +253,19 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
guardrail_name,
f"the legacy hook returned {returned_text_count} texts for a stream that carried {sent_text_count}",
)
tool_call_mismatch: Final = _tool_call_mismatch_reason(
_tool_call_shapes(inputs.get("tool_calls")), _tool_call_shapes(rescanned.get("tool_calls"))
sent_tool_calls: Final = inputs.get("tool_calls")
rewritten_tool_calls: Final = rescanned.get("tool_calls")
tool_call_count_mismatch: Final = _tool_call_count_mismatch_reason(sent_tool_calls, rewritten_tool_calls)
if tool_call_count_mismatch is not None:
raise UndeliverableStreamRewrite(guardrail_name, tool_call_count_mismatch)
texts_update: Final[GenericGuardrailAPIInputs] = {"texts": rewritten} if rewritten else {}
tool_calls_update: Final[GenericGuardrailAPIInputs] = (
{"tool_calls": rewritten_tool_calls}
if rewritten_tool_calls is not None
and _rewrote(_tool_call_shapes(sent_tool_calls), _tool_call_shapes(rewritten_tool_calls))
else {}
)
if tool_call_mismatch is not None:
raise UndeliverableStreamRewrite(guardrail_name, tool_call_mismatch)
if not rewritten:
return inputs
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, **texts_update, **tool_calls_update}
return rewritten_inputs
async def _rescan(

View file

@ -79,6 +79,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry
from litellm.proxy.litellm_pre_call_utils import (
_CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list
_UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list
is_chat_completions_body,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
request_tags_from_request_data,
@ -213,13 +214,6 @@ def _normalize_tool_dialect(
return {**data, **{key: value for key, value in replaceable if key in data}}
def _is_chat_completions_body(data: Mapping[str, object]) -> bool:
messages: Final = data.get("messages")
if isinstance(messages, list) and messages:
return True
return "messages" in data and "input" not in data
_CURSOR_THINKING_SEPARATOR: Final = "-thinking-"
_CURSOR_FAST_SUFFIX: Final = "-fast"
_CURSOR_THINKING_LEVELS: Final[frozenset[str]] = frozenset(get_args(REASONING_EFFORT))
@ -261,7 +255,7 @@ def _resolve_cursor_model_variant(
resolved: Final = {**data, "model": variant.base_model}
if variant.reasoning_effort is None:
return resolved
if _is_chat_completions_body(data):
if is_chat_completions_body(data):
if "reasoning_effort" in data:
return resolved
return {**resolved, "reasoning_effort": variant.reasoning_effort}
@ -646,7 +640,7 @@ async def cursor_chat_completions(
raw_body: Final = await read_request_body(request=request)
if _is_chat_completions_body(raw_body):
if is_chat_completions_body(raw_body):
# Genuine chat completions body (Cursor sends these for models whose BYOK it
# already fixed); delegate so behavior matches /chat/completions exactly.
# Keyed on messages CONTENT, not key presence: Cursor can send a null or

View file

@ -250,7 +250,11 @@ from litellm.proxy.hooks.sensitive_data_routing import ( # noqa: F401, RUF100
PROXY_SensitiveDataRoutingHandler,
_PROXY_SensitiveDataRoutingHandler, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
)
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
add_guardrails_from_auth_metadata,
normalized_metadata_slot,
)
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.policy_engine.pipeline_executor import (
PipelineExecutor,
@ -3973,9 +3977,7 @@ class ProxyLogging:
)
request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller
request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP})
if "metadata" not in request_data:
request_data["metadata"] = {}
request_data["metadata"].update(user_api_key_logged_metadata)
normalized_metadata_slot(request_data, "metadata").update(user_api_key_logged_metadata)
if litellm_logging_obj is not None:
## UPDATE LOGGING INPUT

View file

@ -759,7 +759,7 @@ def _dispatched_model_name(
from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata
provider_model, _ = resolve_model_and_provider_for_metadata(model=model, custom_llm_provider=custom_llm_provider)
return _strip_responses_routing_prefix(provider_model)
return strip_responses_routing_prefix(provider_model)
def _will_bridge_to_chat_completions(
@ -1245,7 +1245,7 @@ def _pop_use_chat_completions_api_kw(kwargs: dict[str, object]) -> bool:
_RESPONSES_ROUTING_PREFIX: Final = "responses/"
def _strip_responses_routing_prefix(model: str) -> str:
def strip_responses_routing_prefix(model: str) -> str:
if not model.startswith(_RESPONSES_ROUTING_PREFIX):
return model
return model[len(_RESPONSES_ROUTING_PREFIX) :]
@ -1273,7 +1273,7 @@ def _resolve_model_provider_for_responses(
litellm_params.api_key = dynamic_api_key
if dynamic_api_base is not None:
litellm_params.api_base = dynamic_api_base
return _strip_responses_routing_prefix(provider_model), resolved_provider
return strip_responses_routing_prefix(provider_model), resolved_provider
def _apply_managed_file_id_mapping(
@ -2893,7 +2893,7 @@ async def _aresponses_websocket(
api_key=api_key,
litellm_params=litellm_params,
)
resolved_model: Final = _strip_responses_routing_prefix(provider_model)
resolved_model: Final = strip_responses_routing_prefix(provider_model)
litellm_params_dict["data_residency"] = infer_openai_data_residency(
_custom_llm_provider,

View file

@ -43,6 +43,7 @@ from litellm.litellm_core_utils.core_helpers import (
process_response_headers,
provider_reported_usage_cost,
set_response_cost_in_hidden_params,
stamp_litellm_bill_on_usage,
)
from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base
from litellm.litellm_core_utils.internal_call_metadata import (
@ -1869,7 +1870,7 @@ def stamp_responses_usage_cost(
if isinstance(bill, bool) or not isinstance(bill, (int, float)):
return
if bill > 0 or isinstance(getattr(usage_obj, "cost", None), (int, float)):
setattr(usage_obj, "cost", float(bill))
stamp_litellm_bill_on_usage(response_obj, usage_obj, float(bill))
def opening_item_payload(payload: Mapping[str, object]) -> Mapping[str, object]:

View file

@ -378,6 +378,7 @@ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check impor
EncryptedContentAffinityCheck,
)
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
async_io_token_refund_failure,
build_io_token_rate_limit_headers,
deployment_has_io_token_limits,
refund_stale_reservation_before_retry,
@ -9501,7 +9502,7 @@ class Router(DeprecatedRouterMembers):
raise original_exception
async def make_call(self, original_function: Any, *args, **kwargs):
async def make_call(self, original_function: Any, *args: object, **kwargs: object):
"""
Handler for making a call to the .completion()/.embeddings()/etc. functions.
"""
@ -9510,17 +9511,26 @@ class Router(DeprecatedRouterMembers):
endpoint_function: Final = (
generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function)
)
with endpoint_selection(
getattr(endpoint_function, "__name__", None),
getattr(endpoint_function, "__module__", None),
model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group),
):
response = original_function(*args, **kwargs)
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
response = await response
try:
with endpoint_selection(
getattr(endpoint_function, "__name__", None),
getattr(endpoint_function, "__module__", None),
model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group),
):
response = original_function(*args, **kwargs)
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
response = await response
except Exception:
with contextlib.suppress(Exception):
await async_io_token_refund_failure(
self.cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs)
)
raise
await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs)
## PROCESS RESPONSE HEADERS
response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs)
response = await self.set_response_headers(
response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs
)
return response

View file

@ -1,7 +1,8 @@
"""
Separate ITPM/OTPM (input/output tokens per minute) deployment rate limits.
- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM
- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM; a request
without max_tokens reserves the model's max output, capped at a fraction of the OTPM limit
- Post-call: reconcile ITPM to actual input tokens and OTPM to actual output tokens
- Cached prompt-read tokens are excluded from ITPM post-call accounting
@ -22,6 +23,7 @@ from litellm import token_counter
from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION
from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
@ -127,10 +129,15 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] |
return itpm_key, otpm_key
def _tokenizer_model(deployment: object) -> str:
def _served_model(deployment: object) -> str:
params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None
served: Final = params.get("model") if is_str_keyed_mapping(params) else None
return deployment_openai_model(deployment) or (served if isinstance(served, str) else "")
return served if isinstance(served, str) else ""
def _tokenizer_and_served_models(deployment: object) -> tuple[str, str]:
served: Final = _served_model(deployment)
return deployment_openai_model(deployment) or served, served
def _estimate_input_tokens(request_kwargs: dict[str, Any] | None, model: str = "") -> int:
@ -160,24 +167,61 @@ def _model_max_output_tokens(model_name: str) -> int | None:
return None
def _resolve_max_tokens(request_kwargs: dict[str, Any] | None, deployment: dict) -> int:
if request_kwargs:
# An explicit max_tokens=0 must be honored, not treated as absent and
# replaced by the model default.
explicit = request_kwargs.get("max_tokens")
if explicit is None:
explicit = request_kwargs.get("max_completion_tokens")
if explicit is None:
explicit = request_kwargs.get("max_output_tokens")
if explicit is not None:
return max(0, int(explicit))
_CHAT_OUTPUT_CAP_KEYS: Final = ("max_tokens", "max_completion_tokens")
_RESPONSES_OUTPUT_CAP_KEYS: Final = ("max_output_tokens",)
_OUTPUT_CAP_KEYS: Final = _CHAT_OUTPUT_CAP_KEYS + _RESPONSES_OUTPUT_CAP_KEYS
model_name: Final = (deployment.get("litellm_params") or {}).get("model")
if model_name:
model_max: Final = _model_max_output_tokens(model_name)
if model_max is not None:
return model_max
return 4096
def _explicit_max_tokens(
request_kwargs: Mapping[str, object] | None, cap_keys: tuple[str, ...] = _OUTPUT_CAP_KEYS
) -> int | None:
if not request_kwargs:
return None
explicit: Final = next(
(request_kwargs[key] for key in cap_keys if request_kwargs.get(key) is not None),
None,
)
match explicit:
case int() | float() | str():
return max(0, int(explicit))
case _:
return None
def _sent_deployment_cap_keys(request_kwargs: Mapping[str, object] | None) -> tuple[str, ...]:
is_responses_request: Final = (
request_kwargs is not None and "input" in request_kwargs and "messages" not in request_kwargs
)
return _RESPONSES_OUTPUT_CAP_KEYS if is_responses_request else _CHAT_OUTPUT_CAP_KEYS
def _deployment_output_cap(request_kwargs: Mapping[str, object] | None, deployment: Mapping[str, object]) -> int | None:
params: Final = deployment.get("litellm_params")
return (
_explicit_max_tokens(params, _sent_deployment_cap_keys(request_kwargs))
if is_str_keyed_mapping(params)
else None
)
def _default_otpm_reservation(model_name: str, otpm_limit: int) -> int:
model_max: Final = _model_max_output_tokens(model_name) if model_name else None
share_of_limit: Final = max(1, int(otpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION))
return min(4096 if model_max is None else model_max, share_of_limit)
def _otpm_reservation(
request_kwargs: Mapping[str, object] | None,
deployment: Mapping[str, object],
model_name: str,
otpm_limit: int,
) -> int:
from_request: Final = _explicit_max_tokens(request_kwargs)
if from_request is not None:
return from_request
default: Final = _default_otpm_reservation(model_name, otpm_limit)
from_deployment: Final = _deployment_output_cap(request_kwargs, deployment)
return default if from_deployment is None else min(from_deployment, default)
def _get_usage_tokens(usage: object) -> tuple[int, int, int]:
@ -210,6 +254,12 @@ def _extract_response_usage(response_obj: object) -> object:
return getattr(response_obj, "usage", None)
def _openai_shaped_usage(usage: object) -> object:
if isinstance(usage, dict) and litellm.AnthropicConfig.is_anthropic_usage_object(usage):
return litellm.AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None)
return usage
def _usage_is_present(usage: object) -> bool:
"""
True only if usage carries an actual input/output breakdown.
@ -238,7 +288,7 @@ def _resolve_reconcile_usage_tokens(
fields. When usage cannot be resolved, return ``usage_resolved=False`` so
callers keep the pre-call reservation instead of refunding it as zero usage.
"""
usage: Final = _extract_response_usage(response_obj)
usage: Final = _openai_shaped_usage(_extract_response_usage(response_obj))
if _usage_is_present(usage):
prompt_tokens, completion_tokens, cached_tokens = _get_usage_tokens(usage)
return max(0, prompt_tokens - cached_tokens), completion_tokens, True
@ -261,6 +311,21 @@ def _resolve_reconcile_usage_tokens(
return 0, 0, False
def _key_for_this_minute(reserved_key: str) -> str:
counter, _, _ = reserved_key.rpartition(":")
return f"{counter}:{get_utc_datetime().strftime('%H-%M')}"
def _reconcile_increments(
reservation: tuple[int, int, str | None, str | None], billable_input: int, completion_tokens: int
) -> tuple[tuple[str, int], ...]:
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
deltas: Final = ((itpm_key, billable_input - itpm_reserved), (otpm_key, completion_tokens - otpm_reserved))
return tuple(
(key if delta < 0 else _key_for_this_minute(key), delta) for key, delta in deltas if key is not None and delta
)
def _stash_reservation_in_metadata(
request_kwargs: dict[str, object] | None,
*,
@ -311,7 +376,9 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object,
return ()
top_level: Final = (kwargs.get("metadata"), kwargs.get("litellm_metadata"))
litellm_params: Final = _as_mutable_mapping(kwargs.get("litellm_params"))
from_params: Final = () if litellm_params is None else (litellm_params.get("metadata"),)
from_params: Final = (
() if litellm_params is None else (litellm_params.get("metadata"), litellm_params.get("litellm_metadata"))
)
standard_logging_object: Final = _as_mutable_mapping(kwargs.get("standard_logging_object"))
from_logging_object: Final = () if standard_logging_object is None else (standard_logging_object.get("metadata"),)
return top_level + from_params + from_logging_object
@ -433,9 +500,11 @@ def io_token_pre_call_check(
return deployment
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
_model: Final = _tokenizer_model(deployment)
_model, served_model = _tokenizer_and_served_models(deployment)
estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model)
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
otpm_to_reserve: Final = (
_otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0
)
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
@ -458,7 +527,7 @@ def io_token_pre_call_check(
)
if otpm_limit is not None:
otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit)
otpm_reserved = otpm_to_reserve
try:
_sync_increment_with_rollback(
dual_cache,
@ -497,9 +566,11 @@ async def async_io_token_pre_call_check(
return deployment
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
_model: Final = _tokenizer_model(deployment)
_model, served_model = _tokenizer_and_served_models(deployment)
estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model)
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
otpm_to_reserve: Final = (
_otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0
)
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
@ -523,7 +594,7 @@ async def async_io_token_pre_call_check(
)
if otpm_limit is not None:
otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit)
otpm_reserved = otpm_to_reserve
try:
await _increment_with_rollback(
dual_cache,
@ -564,7 +635,8 @@ def io_token_reconcile_success(
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
if itpm_key is None and otpm_key is None:
return
@ -572,23 +644,8 @@ def io_token_reconcile_success(
try:
if usage_resolved:
if itpm_key is not None:
itpm_delta: Final = billable_input - itpm_reserved
if itpm_delta != 0:
dual_cache.increment_cache(
key=itpm_key,
value=itpm_delta,
ttl=RoutingArgsTTL,
)
if otpm_key is not None:
otpm_delta: Final = completion_tokens - otpm_reserved
if otpm_delta != 0:
dual_cache.increment_cache(
key=otpm_key,
value=otpm_delta,
ttl=RoutingArgsTTL,
)
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
@ -618,39 +675,23 @@ async def async_io_token_reconcile_success(
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
if itpm_key is None and otpm_key is None:
return
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
# Reconcile against the exact key that held the reservation (which encodes
# the reservation's minute), not a key recomputed at response time. This
# tracks actual usage even when the pre-call estimate was 0, and avoids a
# minute-boundary mismatch for calls that span into the next minute. Always
# clear the stash afterwards (even if an increment throws) so a retry or a
# duplicate success event can't re-process it.
# A refund goes back to the minute that held the reservation, while usage above the reservation is charged
# to the minute the response finished in, so a call that runs past its minute still counts against a live
# window. Always clear the stash afterwards (even if an increment throws) so a retry or a duplicate success
# event can't re-process it.
try:
if usage_resolved:
if itpm_key is not None:
itpm_delta: Final = billable_input - itpm_reserved
if itpm_delta != 0:
await dual_cache.async_increment_cache(
key=itpm_key,
value=itpm_delta,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
if otpm_key is not None:
otpm_delta: Final = completion_tokens - otpm_reserved
if otpm_delta != 0:
await dual_cache.async_increment_cache(
key=otpm_key,
value=otpm_delta,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
await dual_cache.async_increment_cache(
key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span
)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",

View file

@ -26,7 +26,6 @@ from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
ITPM_RESERVED_KEY,
async_io_token_pre_call_check,
async_io_token_reconcile_success,
async_io_token_refund_failure,
@ -485,24 +484,12 @@ class ModelRateLimitingCheck(CustomLogger):
try:
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
# IO token reconciliation works purely from the cache keys stashed in
# kwargs/metadata, so it must run before the model_id guard below
# (which only the TPM-tracking path needs). Otherwise a request whose
# standard_logging_object lacks model_id would never return its
# reservation, leaving the counter elevated until the TTL expires.
slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {}
kwargs_metadata: Final = kwargs.get("metadata") or {}
if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata:
await async_io_token_reconcile_success(
self.dual_cache,
kwargs,
response_obj,
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
)
# Fall through: a deployment can also configure tpm/rpm alongside
# itpm/otpm, and that path's pre-call check reads the tpm_key
# counter tracked below, so it must still be incremented here.
await async_io_token_reconcile_success(
self.dual_cache,
kwargs,
response_obj,
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
)
if standard_logging_object is None:
return
@ -557,17 +544,7 @@ class ModelRateLimitingCheck(CustomLogger):
"""
try:
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {}
kwargs_metadata: Final = kwargs.get("metadata") or {}
if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata:
io_token_reconcile_success(
self.dual_cache,
kwargs,
response_obj,
)
# Fall through: a deployment can also configure tpm/rpm alongside
# itpm/otpm, and that path's pre-call check reads the tpm_key
# counter tracked below, so it must still be incremented here.
io_token_reconcile_success(self.dual_cache, kwargs, response_obj)
if standard_logging_object is None:
return

View file

@ -713,6 +713,23 @@ class AnthropicResponseContentBlockRedactedThinking(LiteLLMBaseModel):
data: str
class AnthropicResponseContentBlockServerToolUse(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
type: Literal["server_tool_use"]
id: str
name: str
input: Mapping[str, object]
class AnthropicResponseContentBlockWebSearchToolResult(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
type: Literal["web_search_tool_result"]
tool_use_id: str
content: Sequence[Mapping[str, object]] | Mapping[str, object]
class AnthropicResponseUsageBlock(LiteLLMBaseModel):
model_config = ConfigDict(extra="allow")
@ -744,6 +761,8 @@ class AnthropicResponse(LiteLLMBaseModel):
| AnthropicResponseContentBlockToolUse
| AnthropicResponseContentBlockThinking
| AnthropicResponseContentBlockRedactedThinking
| AnthropicResponseContentBlockServerToolUse
| AnthropicResponseContentBlockWebSearchToolResult
]
"""Content generated by the model."""

View file

@ -6710,6 +6710,26 @@ def _first_registered_match(
return next(matches, (None, None))
def _model_cost_lookup_order(
model: str, potential_model_names: PotentialModelNamesAndCustomLLMProvider
) -> tuple[str, ...]:
return (
potential_model_names["combined_model_name"],
model,
potential_model_names["region_free_combined_model_name"],
potential_model_names["split_model"],
potential_model_names["combined_stripped_model_name"],
potential_model_names["stripped_model_name"],
potential_model_names["provider_prefixed_model_name"],
)
def potential_model_names_in_lookup_order(model: str, custom_llm_provider: str) -> tuple[str, ...]:
return _model_cost_lookup_order(
model, get_potential_model_names(model=model, custom_llm_provider=custom_llm_provider)
)
def _get_model_info_from_generalization(
model: str,
potential_model_names: PotentialModelNamesAndCustomLLMProvider,
@ -6728,15 +6748,7 @@ def _get_model_info_from_generalization(
still had a priced exact name to try. O(number of rules); only call after the
exact lookups have missed.
"""
candidates: Final = (
potential_model_names["combined_model_name"],
model,
potential_model_names["region_free_combined_model_name"],
potential_model_names["split_model"],
potential_model_names["combined_stripped_model_name"],
potential_model_names["stripped_model_name"],
potential_model_names["provider_prefixed_model_name"],
)
candidates: Final = _model_cost_lookup_order(model, potential_model_names)
if any(get_model_cost_key(candidate) is not None for candidate in candidates):
return None
for candidate in candidates:
@ -6975,11 +6987,6 @@ def get_model_info_helper(
model=model, custom_llm_provider=custom_llm_provider or declared_authenticating_provider(model)
)
combined_model_name: Final = potential_model_names["combined_model_name"]
stripped_model_name: Final = potential_model_names["stripped_model_name"]
combined_stripped_model_name: Final = potential_model_names["combined_stripped_model_name"]
provider_prefixed_model_name: Final = potential_model_names["provider_prefixed_model_name"]
split_model: Final = potential_model_names["split_model"]
custom_llm_provider = potential_model_names["custom_llm_provider"]
model_cost_custom_llm_provider: Final = custom_llm_provider
#########################
@ -7046,16 +7053,9 @@ def get_model_info_helper(
and strip it. Tried last so no model that already resolves through 1-6 can change.
"""
lookup_order: Final = (
combined_model_name,
model,
potential_model_names["region_free_combined_model_name"],
split_model,
combined_stripped_model_name,
stripped_model_name,
provider_prefixed_model_name,
lookup: Final = _first_registered_match(
_model_cost_lookup_order(model, potential_model_names), model_cost_custom_llm_provider
)
lookup: Final = _first_registered_match(lookup_order, model_cost_custom_llm_provider)
key: str | None = lookup[0]
_model_info: dict[str, Any] | None = lookup[1]

View file

@ -6014,11 +6014,14 @@
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 1.25e-06,
"input_cost_per_token": 2.5e-06,
"input_cost_per_token_batches": 1.25e-06,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1e-05,
"output_cost_per_token_batches": 5e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
@ -6313,16 +6316,19 @@
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 1.5e-07,
"input_cost_per_token_batches": 7.5e-08,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6e-07,
"output_cost_per_token_batches": 3e-07,
"search_context_cost_per_query": {
"search_context_size_high": 0.014,
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014
},
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
@ -11303,16 +11309,19 @@
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8e-06,
"output_cost_per_token_batches": 4e-06,
"search_context_cost_per_query": {
"search_context_size_high": 0.014,
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014
},
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
@ -11411,11 +11420,14 @@
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 1.1e-06,
"input_cost_per_token_batches": 5.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.4e-06,
"output_cost_per_token_batches": 2.2e-06,
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
@ -11513,16 +11525,19 @@
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 2.75e-07,
"input_cost_per_token": 1.1e-06,
"input_cost_per_token_batches": 5.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.4e-06,
"output_cost_per_token_batches": 2.2e-06,
"search_context_cost_per_query": {
"search_context_size_high": 0.014,
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014
},
"source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
@ -54512,6 +54527,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true
},
"zai/glm-5.3-flash": {
@ -54526,10 +54542,26 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"source": "https://docs.z.ai/guides/overview/pricing",
"supports_vision": true
},
"zai/glm-5v-turbo": {
"cache_read_input_token_cost": 2.4e-07,
"input_cost_per_token": 1.2e-06,
"litellm_provider": "zai",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://vercel.com/ai-gateway/models/glm-5v-turbo",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"zai/glm-5.1": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 2.6e-07,
@ -81059,5 +81091,14 @@
"supports_response_schema": true,
"supports_tool_choice": false,
"supports_vision": true
},
"baseten/decision-models/baseten-joey-2b-latest": {
"input_cost_per_token": 1e-08,
"litellm_provider": "baseten",
"max_input_tokens": 16384,
"mode": "decisions",
"output_cost_per_token": 0.0,
"source": "https://inference.baseten.co/v1/models",
"supports_vision": true
}
}

View file

@ -2,6 +2,7 @@
PARITY pairs the streamed case with its plain twin where both bill the same row."""
from collections.abc import Mapping
from typing import Final, cast
from integration.cost_calculation.cost_tracking_case import (
@ -11878,7 +11879,215 @@ GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = GPT_5_6_TERRA_MESSAGES_BRIDGED_JS
}
)
# An OpenAI-compatible gateway that reports what it charged under usage.cost (#47533). The reported total is the bill,
# split across input and output by the ratio of the cost map rates (100 * 1e-6 : 20 * 3e-6), on plain and streamed
# replies alike. A deployment the operator priced keeps the operator's rates on both.
_GATEWAY_REPORTED_USAGE: Final = {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120, "cost": 0.037}
_GATEWAY_REPORTED_EXPECTED: Final = ExactExpected(
spend=0.037, input_cost=0.0, output_cost=0.037, prompt_tokens=100, completion_tokens=20
)
_GATEWAY_OPERATOR_PRICED: Final = Deployment(input_cost_per_token=2e-06, output_cost_per_token=6e-06)
_GATEWAY_OPERATOR_PRICED_EXPECTED: Final = ExactExpected(
spend=0.00032, input_cost=0.0002, output_cost=0.00012, prompt_tokens=100, completion_tokens=20
)
_GATEWAY_CHUNK: Final = {
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion.chunk",
"created": 1,
"model": "gateway-reported-cost-fixture",
}
GATEWAY_REPORTED_COST: Final = CostTrackingTestCase(
name="gateway-reported-cost",
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
model="gateway-reported-cost-fixture",
endpoint="/v1/chat/completions",
request={"model": "$MODEL", "messages": [{"role": "user", "content": "reported cost"}], "stream": False},
response=JsonResponse(
content_type="application/json",
body={
"id": "chatcmpl-$REQUEST_ID",
"object": "chat.completion",
"created": 1,
"model": "gateway-reported-cost-fixture",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": _GATEWAY_REPORTED_USAGE,
},
),
expected=_GATEWAY_REPORTED_EXPECTED,
)
GATEWAY_REPORTED_COST_STREAM: Final = CostTrackingTestCase(
name="gateway-reported-cost-stream",
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
model="gateway-reported-cost-fixture",
endpoint="/v1/chat/completions",
request={
"model": "$MODEL",
"messages": [{"role": "user", "content": "reported cost"}],
"stream": True,
"stream_options": {"include_usage": True},
},
response=SseResponse(
content_type="text/event-stream",
frames=sse_frames(
{**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]},
{**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
{**_GATEWAY_CHUNK, "choices": [], "usage": _GATEWAY_REPORTED_USAGE},
done=True,
),
),
expected=_GATEWAY_REPORTED_EXPECTED,
)
GATEWAY_REPORTED_COST_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST.model_copy(
update={
"name": "gateway-reported-cost-operator_priced",
"deployment": _GATEWAY_OPERATOR_PRICED,
"expected": _GATEWAY_OPERATOR_PRICED_EXPECTED,
}
)
GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_STREAM.model_copy(
update={
"name": "gateway-reported-cost-operator_priced-stream",
"deployment": _GATEWAY_OPERATOR_PRICED,
"expected": _GATEWAY_OPERATOR_PRICED_EXPECTED,
}
)
_GATEWAY_RESPONSES_USAGE: Final[Mapping[str, JsonValue]] = {"input_tokens": 100, "output_tokens": 20, "total_tokens": 120, "cost": 0.037}
_GATEWAY_RESPONSES_ITEM: Final[Mapping[str, JsonValue]] = {
"id": "msg_gateway",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
}
_GATEWAY_RESPONSES_BODY: Final[Mapping[str, JsonValue]] = {
"id": "resp_$REQUEST_ID",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gateway-reported-cost-fixture",
"output": [_GATEWAY_RESPONSES_ITEM],
"usage": _GATEWAY_RESPONSES_USAGE,
}
GATEWAY_REPORTED_COST_RESPONSES: Final = CostTrackingTestCase(
name="gateway-reported-cost-responses",
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
model="gateway-reported-cost-fixture",
endpoint="/v1/responses",
request={"model": "$MODEL", "input": "reported cost", "cache": {"no-cache": True}},
response=JsonResponse(content_type="application/json", body=_GATEWAY_RESPONSES_BODY),
expected=_GATEWAY_REPORTED_EXPECTED,
)
GATEWAY_REPORTED_COST_RESPONSES_STREAM: Final = GATEWAY_REPORTED_COST_RESPONSES.model_copy(
update={
"name": "gateway-reported-cost-responses-stream",
"request": {**GATEWAY_REPORTED_COST_RESPONSES.request, "stream": True},
"response": SseResponse(
content_type="text/event-stream",
frames=sse_frames(
{
"type": "response.created",
"sequence_number": 0,
"response": {**_GATEWAY_RESPONSES_BODY, "status": "in_progress", "output": [], "usage": None},
},
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": "msg_gateway",
"output_index": 0,
"content_index": 0,
"delta": "ok",
},
{"type": "response.completed", "sequence_number": 2, "response": _GATEWAY_RESPONSES_BODY},
),
),
}
)
GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_RESPONSES.model_copy(
update={
"name": "gateway-reported-cost-responses-operator_priced",
"deployment": _GATEWAY_OPERATOR_PRICED,
"expected": _GATEWAY_OPERATOR_PRICED_EXPECTED,
}
)
GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_RESPONSES_STREAM.model_copy(
update={
"name": "gateway-reported-cost-responses-operator_priced-stream",
"deployment": _GATEWAY_OPERATOR_PRICED,
"expected": _GATEWAY_OPERATOR_PRICED_EXPECTED,
}
)
_GATEWAY_ZERO_USAGE: Final = {**_GATEWAY_REPORTED_USAGE, "cost": 0}
GATEWAY_REPORTED_ZERO_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED.model_copy(
update={
"name": "gateway-reported-zero-operator_priced",
"response": JsonResponse(
content_type="application/json",
body={**GATEWAY_REPORTED_COST.response.body, "usage": _GATEWAY_ZERO_USAGE},
),
}
)
GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM.model_copy(
update={
"name": "gateway-reported-zero-operator_priced-stream",
"response": SseResponse(
content_type="text/event-stream",
frames=sse_frames(
{**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]},
{**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
{**_GATEWAY_CHUNK, "choices": [], "usage": _GATEWAY_ZERO_USAGE},
done=True,
),
),
}
)
_GATEWAY_OPENROUTER_OPERATOR_PRICED: Final = _GATEWAY_OPERATOR_PRICED.model_copy(
update={"model": "openrouter/vendor/gateway-reported-cost-fixture"}
)
GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED.model_copy(
update={
"name": "gateway-reported-cost-openrouter-operator_priced",
"deployment": _GATEWAY_OPENROUTER_OPERATOR_PRICED,
}
)
GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM: Final = (
GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM.model_copy(
update={
"name": "gateway-reported-cost-openrouter-operator_priced-stream",
"deployment": _GATEWAY_OPENROUTER_OPERATOR_PRICED,
}
)
)
CASES: Final[tuple[CostTrackingTestCase, ...]] = (
GATEWAY_REPORTED_COST,
GATEWAY_REPORTED_COST_STREAM,
GATEWAY_REPORTED_COST_OPERATOR_PRICED,
GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM,
GATEWAY_REPORTED_COST_RESPONSES,
GATEWAY_REPORTED_COST_RESPONSES_STREAM,
GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED,
GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM,
GATEWAY_REPORTED_ZERO_OPERATOR_PRICED,
GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM,
GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED,
GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM,
GPT_5_6_SLASH_ALIAS,
GPT_5_6_SLASH_ALIAS_STREAM,
GPT_6_SOL_RESPONSES_DEPLOYMENT_JSON,
@ -12120,6 +12329,22 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = (
)
PARITY: Final[tuple[StreamParityTestCase, ...]] = (
StreamParityTestCase(plain=GATEWAY_REPORTED_COST, streamed=GATEWAY_REPORTED_COST_STREAM),
StreamParityTestCase(
plain=GATEWAY_REPORTED_COST_OPERATOR_PRICED, streamed=GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM
),
StreamParityTestCase(
plain=GATEWAY_REPORTED_ZERO_OPERATOR_PRICED, streamed=GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM
),
StreamParityTestCase(
plain=GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED,
streamed=GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM,
),
StreamParityTestCase(plain=GATEWAY_REPORTED_COST_RESPONSES, streamed=GATEWAY_REPORTED_COST_RESPONSES_STREAM),
StreamParityTestCase(
plain=GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED,
streamed=GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM,
),
StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM),
StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM_INCOMPLETE),
StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM_UNVALIDATED),

View file

@ -26,6 +26,12 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = {
"input_cost_per_token_above_128k_tokens": 2.2e-07,
"output_cost_per_token_above_128k_tokens": 8.8e-07,
},
"gateway-reported-cost-fixture": {
"litellm_provider": "openai",
"mode": "chat",
"input_cost_per_token": 1e-06,
"output_cost_per_token": 3e-06,
},
"openrouter/anthropic/claude-sonnet-5": {
"litellm_provider": "openrouter",
"mode": "chat",

View file

@ -29,6 +29,24 @@ def sse_events(text: str) -> tuple[tuple[str, dict[str, object]], ...]:
)
def streamed_blocks(events: tuple[tuple[str, dict[str, object]], ...]) -> tuple[dict[str, object], ...]:
input_json: Final = tuple(
(event["index"], event["delta"]["partial_json"])
for name, event in events
if name == "content_block_delta" and event["delta"]["type"] == "input_json_delta"
)
return tuple(
{
**event["content_block"],
"input": json.loads("".join(part for index, part in input_json if index == event["index"])),
}
if any(index == event["index"] for index, _ in input_json)
else event["content_block"]
for name, event in events
if name == "content_block_start"
)
@pytest.mark.covers("other.provider_wire.bedrock.websearch_interception_streamed_capped_turn_ends_with_native_results")
def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_snippets_and_ordered_blocks(
gateway: Gateway, tmp_path: Path
@ -136,7 +154,7 @@ def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_sni
for earlier, earlier_event in events[:position]
if earlier == "content_block_start"
}, response.text
started: Final = tuple(event["content_block"] for name, event in events if name == "content_block_start")
started: Final = streamed_blocks(events)
search_ids: Final = tuple(block["id"] for block in started if block["type"] == "server_tool_use")
assert search_ids and all(search_id.startswith("srvtoolu_") for search_id in search_ids), response.text
assert started[-1] == {"type": "text", "text": ""}, response.text
@ -162,7 +180,12 @@ def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_sni
)
), response.text
assert (
"".join(event["delta"]["text"] for name, event in events if name == "content_block_delta") == "turn-1"
"".join(
event["delta"]["text"]
for name, event in events
if name == "content_block_delta" and event["delta"]["type"] == "text_delta"
)
== "turn-1"
), response.text
assert [event["delta"]["stop_reason"] for name, event in events if name == "message_delta"] == [
"end_turn"

View file

@ -3,11 +3,14 @@ from __future__ import annotations
import json
import os
import uuid
from collections.abc import Mapping
from hashlib import sha256
from importlib.resources import files
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, Scenario
from integration._support.client import Gateway, Scenario, eventually, list_value, object_value
from integration._support.responses_vendor import same_response
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.worker_readiness import wait_for_worker_readiness
@ -38,6 +41,7 @@ _CHAT_SUFFIXES: Final[tuple[tuple[str, str], ...]] = (("", ""),) + tuple(
(f"-thinking-{level}", level) for level in _CHAT_THINKING_LEVELS
)
_PROMPT: Final = "Inspect the requested change."
_REQUEST_TAG: Final = "cursor-integration"
_INPUT: Final = [{"role": "user", "content": [{"type": "input_text", "text": _PROMPT}]}]
_CUSTOM_TOOL: Final = {
"type": "custom",
@ -107,16 +111,30 @@ def _stream(identity: str) -> tuple[bytes, ...]:
) + (b"data: [DONE]\n\n",)
def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]:
def _chat_completion(stream: bool, identity: str) -> bytes | tuple[bytes, ...]:
if stream:
return (
b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{"content":"Patch inspected."},"finish_reason":null}]}\n\n',
b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n',
b"data: [DONE]\n\n",
)
return tuple(
f"data: {json.dumps(chunk)}\n\n".encode()
for chunk in (
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": _MODEL,
"choices": [{"index": 0, "delta": {"content": "Patch inspected."}, "finish_reason": None}],
},
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": _MODEL,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
},
)
) + (b"data: [DONE]\n\n",)
return json.dumps(
{
"id": "chat_cursor_1",
"id": identity,
"object": "chat.completion",
"created": 1,
"model": _MODEL,
@ -133,11 +151,12 @@ def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]:
).encode()
def _deployment(scenario: Scenario, wire: Wire) -> str:
def _deployment(scenario: Scenario, wire: Wire, **pricing: JsonValue) -> str:
model: Final = scenario.model(
model=f"openai/{_MODEL}",
api_key=_API_KEY,
api_base=f"{wire.url}/v1",
**pricing,
)
wait_for_worker_readiness(scenario.gateway, model, 1, _PROXY_WORKERS)
return model
@ -167,7 +186,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
assert request.method == "POST"
assert request.headers["authorization"] == f"Bearer {_API_KEY}", dict(request.headers)
if request.target == "/v1/chat/completions":
completion: Final = _chat_completion(stream)
completion: Final = _chat_completion(stream, raw_id)
return Reply(
content_type="text/event-stream" if stream else "application/json",
chunks=completion if isinstance(completion, tuple) else None,
@ -175,7 +194,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
)
assert request.target == "/v1/responses", request.target
body: Final = _JSON_OBJECT.validate_json(request.body)
assert {key: value for key, value in body.items() if key != "metadata"} == {
assert body == {
"model": _MODEL,
"input": _INPUT,
**reasoning,
@ -259,41 +278,69 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
assert [request.target for request in drain_contract_requests(wire)] == ["/v1/responses"]
def test_cursor_responses_body_carries_no_proxy_metadata(gateway: Gateway) -> None:
pytest.skip(
"BUG: /cursor/chat/completions forwards proxy-internal metadata (key hash, team, user, client ip, user agent) "
"to the provider in the Responses body"
)
@pytest.mark.parametrize(
"client_metadata", ({}, {"cursor_conversation": "integration"}), ids=("no_client_metadata", "client_metadata")
)
@pytest.mark.parametrize("stream", (False, True))
def test_cursor_responses_body_sends_only_client_metadata_and_bills_the_key(
gateway: Gateway,
stream: bool,
client_metadata: Mapping[str, JsonValue],
) -> None:
raw_id: Final = f"resp_cursor_metadata_{uuid.uuid4().hex}"
alias: Final = f"integration-cursor-{uuid.uuid4().hex}"
sent_metadata: Final = {"metadata": client_metadata} if client_metadata else {}
def respond(request: Request) -> Reply:
if request.method == "GET" and request.target == "/v1/models":
return model_discovery_reply(_MODEL)
assert request.method == "POST" and request.target == "/v1/responses", request.target
if stream:
return Reply(content_type="text/event-stream", chunks=_stream(raw_id))
return Reply(body=json.dumps(_response(raw_id)).encode())
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire)
key: Final = scenario.key(key_alias=alias)
hashed_key: Final = sha256(key.encode()).hexdigest()
response: Final = gateway.request(
"POST",
"/cursor/chat/completions",
{
"model": f"{model}-thinking-high",
"model": model,
"input": _INPUT,
"messages": [],
"tools": [_CUSTOM_TOOL],
"stream": stream,
**sent_metadata,
},
key=key,
)
assert response.status_code == 200, response.text
requests: Final = drain_contract_requests(wire)
assert len(requests) == 1, requests
body: Final = _JSON_OBJECT.validate_json(requests[0].body)
assert body == {
assert _JSON_OBJECT.validate_json(requests[0].body) == {
"model": _MODEL,
"input": _INPUT,
"reasoning": {"effort": "high"},
"tools": [_CUSTOM_TOOL],
"stream": stream,
**sent_metadata,
}, requests[0].body
(row,) = eventually(
lambda: gateway.rows(
'SELECT call_type, spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (hashed_key,)
),
lambda rows: len(rows) == 1,
seconds=70,
)
assert row["call_type"] == "aresponses", row
assert object_value(row["metadata"])["user_api_key_alias"] == alias, row
assert isinstance(row["spend"], float) and row["spend"] > 0, row
eventually(
lambda: gateway.rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed_key,)),
lambda tokens: len(tokens) == 1 and tokens[0]["spend"] == pytest.approx(row["spend"]),
seconds=70,
)
def test_cursor_stream_reports_usage_when_the_client_asks_for_it(gateway: Gateway) -> None:
@ -341,10 +388,12 @@ def test_cursor_messages_body_uses_chat_completions_wire(
suffix: str,
reasoning_effort: str,
) -> None:
chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}"
def respond(request: Request) -> Reply:
if request.method == "GET" and request.target == "/v1/models":
return model_discovery_reply(_MODEL)
completion: Final = _chat_completion(stream)
completion: Final = _chat_completion(stream, chat_id)
return Reply(
content_type="text/event-stream" if stream else "application/json",
chunks=completion if isinstance(completion, tuple) else None,
@ -367,6 +416,7 @@ def test_cursor_messages_body_uses_chat_completions_wire(
"stream": stream,
},
key=key,
headers={"x-litellm-tags": _REQUEST_TAG},
)
assert response.status_code == 200, response.text
if stream:
@ -393,7 +443,7 @@ def test_cursor_messages_body_uses_chat_completions_wire(
payload: Final = _JSON_OBJECT.validate_json(response.content)
identity: Final = payload.get("id")
created: Final = payload.get("created")
assert isinstance(identity, str) and same_response(identity, "chat_cursor_1"), response.text
assert isinstance(identity, str) and same_response(identity, chat_id), response.text
assert isinstance(created, int) and created > 0, response.text
assert payload == {
"id": identity,
@ -456,3 +506,72 @@ def test_cursor_messages_body_uses_chat_completions_wire(
**({"stream_options": {"include_usage": True}} if stream else {}),
**({"reasoning_effort": reasoning_effort} if reasoning_effort else {}),
}, requests[0].body
(row,) = eventually(
lambda: gateway.rows(
'SELECT call_type, team_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
(sha256(key.encode()).hexdigest(),),
),
lambda rows: len(rows) == 1,
seconds=70,
)
assert (row["call_type"], row["team_id"]) == ("acompletion", team), row
assert _REQUEST_TAG in list_value(row["request_tags"]), row
@pytest.mark.parametrize(
("body_shape", "metadata_as_json_string"), (("messages", False), ("input", False), ("input", True))
)
def test_cursor_body_keeps_client_and_key_tags_and_enforces_the_client_tag_budget(
gateway: Gateway, body_shape: str, metadata_as_json_string: bool
) -> None:
chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}"
response_id: Final = f"resp_cursor_tags_{uuid.uuid4().hex}"
budget_tag: Final = f"cursor-tag-budget-{uuid.uuid4().hex}"
key_tag: Final = f"cursor-key-tag-{uuid.uuid4().hex}"
client_metadata: Final = {"tags": [budget_tag]}
def respond(request: Request) -> Reply:
if request.method == "GET" and request.target == "/v1/models":
return model_discovery_reply(_MODEL)
if request.target == "/v1/responses":
return Reply(body=json.dumps(_response(response_id)).encode())
completion: Final = _chat_completion(False, chat_id)
return Reply(body=completion if isinstance(completion, bytes) else b"")
def delete_tag() -> None:
gateway.post("/tag/delete", {"name": budget_tag})
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _deployment(scenario, wire, input_cost_per_token=0.01, output_cost_per_token=0.01)
key: Final = scenario.key(models=[model], metadata={"tags": [key_tag]})
gateway.post("/tag/new", {"name": budget_tag, "max_budget": 0.0001})
scenario.cleanups.callback(delete_tag)
def tagged_chat(text: str) -> httpx.Response:
return gateway.request(
"POST",
"/cursor/chat/completions",
{
"model": model,
body_shape: [{"role": "user", "content": text}],
"metadata": json.dumps(client_metadata) if metadata_as_json_string else client_metadata,
},
key=key,
)
first: Final = tagged_chat("first tagged chat")
assert first.status_code == 200, first.text
(row,) = eventually(
lambda: gateway.rows(
'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
(sha256(key.encode()).hexdigest(),),
),
lambda rows: len(rows) == 1,
seconds=70,
)
assert {budget_tag, key_tag} <= set(list_value(row["request_tags"])), row
blocked: Final = eventually(
lambda: tagged_chat("tag budget probe"), lambda response: response.status_code != 200, seconds=70
)
assert blocked.status_code == 422, blocked.text
assert "budget" in blocked.text.lower(), blocked.text

View file

@ -153,15 +153,16 @@ def _foundry_deployment(rig: _Rig, scenario: Scenario, *, bare: bool) -> str:
return name
def _assert_foundry_chat_wire(received: Request, prompt: str, *, options: JsonValue | None) -> None:
def _assert_foundry_chat_wire(received: Request, user_content: JsonValue, *, options: JsonValue | None) -> None:
body: Final = body_of(received)
marked_system: Final = [{"type": "text", "text": SYSTEM, "prompt_cache_breakpoint": EXPLICIT}]
assert urlsplit(received.target).netloc == _FOUNDRY_HOST, received.target
assert urlsplit(received.target).path == _FOUNDRY_CHAT_PATH, received.target
assert "prompt_cache_breakpoint" not in received.body.decode(), received.body
assert body["model"] == _FOUNDRY_MODEL, received.body
assert body["messages"] == [{"role": "system", "content": SYSTEM}, {"role": "user", "content": prompt}], (
received.body
)
assert body["messages"] == [
{"role": "system", "content": marked_system},
{"role": "user", "content": user_content},
], received.body
assert body.get("prompt_cache_options") == options, received.body
@ -185,7 +186,7 @@ def test_b9_b12_foundry_gpt6_responses_reach_the_openai_dialect_on_the_foundry_h
@pytest.mark.timeout(_CELL_TIMEOUT)
def test_b10_foundry_gpt6_messages_bridge_to_chat_without_a_breakpoint(rig: _Rig) -> None:
def test_b10_foundry_gpt6_messages_bridge_to_chat_keeps_the_injected_breakpoint(rig: _Rig) -> None:
marker: Final = fresh_marker()
with rig.gateway.scenario() as scenario:
name: Final = _foundry_deployment(rig, scenario, bare=False)
@ -193,12 +194,12 @@ def test_b10_foundry_gpt6_messages_bridge_to_chat_without_a_breakpoint(rig: _Rig
rig.gateway, rig.wire, "messages", request_body("messages", name, prompt_text(marker))
)
assert_answered(outcome, marker)
_assert_foundry_chat_wire(received, prompt_text(marker), options=IMPLICIT)
_assert_foundry_chat_wire(received, [{"type": "text", "text": prompt_text(marker)}], options=IMPLICIT)
success_row(name, marker)
@pytest.mark.timeout(_CELL_TIMEOUT)
def test_b11_foundry_gpt6_chat_stays_on_the_plain_chat_wire(rig: _Rig) -> None:
def test_b11_foundry_gpt6_chat_stays_on_the_chat_wire_with_the_injected_breakpoint(rig: _Rig) -> None:
marker: Final = fresh_marker()
with rig.gateway.scenario() as scenario:
name: Final = _foundry_deployment(rig, scenario, bare=False)

View file

@ -0,0 +1,137 @@
from __future__ import annotations
import dataclasses
import threading
import uuid
from collections.abc import Callable, Iterator
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from queue import SimpleQueue
from typing import Final
import httpx
import pytest
from integration._support.client import JSON_OBJECT, gateway_from_environment
from integration._support.openai_wire import chat_reply
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.providers._cache_control_marks_support import owned_config
from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
GROUPS: Final = ("otpm-fraction-first", "otpm-fraction-retry")
OTPM_LIMIT: Final = 1000
FRACTION: Final = 0.5
DEFAULT_RESERVATION: Final = int(OTPM_LIMIT * FRACTION)
@dataclass(frozen=True, slots=True)
class Rig:
owned: OwnedProxy
upstream: Wire
gates: SimpleQueue[threading.Event]
@dataclass(frozen=True, slots=True)
class Admissions:
held_forwarded_max_tokens: bool
statuses: tuple[int, int, int]
one_minute: bool
def _reply_for(gates: SimpleQueue[threading.Event]) -> Callable[[Request], Reply]:
def reply(request: Request) -> Reply:
if is_model_info_probe(request):
return model_list_reply()
if JSON_OBJECT.validate_json(request.body).get("stream") is True:
return dataclasses.replace(
chat_reply("chatcmpl-held", OPENAI_MODEL, "held", stream=True), gate_before_first=gates.get(timeout=60)
)
return chat_reply("chatcmpl-capped", OPENAI_MODEL, "capped", stream=False)
return reply
def _config(directory: Path, upstream: Wire) -> Path:
return owned_config(
directory,
[
{
"model_name": group,
"litellm_params": {
"model": OPENAI_MODEL,
"api_base": upstream.url + "/v1",
"api_key": "synthetic-openai-key",
"otpm": OTPM_LIMIT,
},
"model_info": {"id": group},
}
for group in GROUPS
],
litellm_settings={"cache": False},
router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 0},
)
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("otpm-fraction")
gates: Final[SimpleQueue[threading.Event]] = SimpleQueue()
with (
gateway_from_environment() as gateway,
wire_server(_reply_for(gates)) as upstream,
owned_proxy_process(
gateway,
directory,
{"NO_MAX_TOKENS_TPM_RESERVATION_FRACTION": str(FRACTION)},
config=_config(directory, upstream),
) as owned,
):
yield Rig(owned, upstream, gates)
def _chat(rig: Rig, group: str, **params: int | bool) -> httpx.Response:
return rig.owned.gateway.request(
"POST",
"/v1/chat/completions",
{"model": group, "messages": [{"role": "user", "content": f"fraction {uuid.uuid4().hex}"}], **params},
)
def _await_forwarded_chat(upstream: Wire) -> Request:
request: Final = upstream.received.get(timeout=60)
return _await_forwarded_chat(upstream) if is_model_info_probe(request) else request
def _utc_minute() -> str:
return datetime.now(UTC).strftime("%Y-%m-%d %H:%M")
def _admissions_beside_a_held_stream(rig: Rig, group: str) -> Admissions:
release: Final = threading.Event()
rig.gates.put(release)
started: Final = _utc_minute()
with ThreadPoolExecutor(max_workers=1) as pool:
held: Final = pool.submit(_chat, rig, group, stream=True)
try:
forwarded: Final = JSON_OBJECT.validate_json(_await_forwarded_chat(rig.upstream).body)
over: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION + 1)
fits: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION)
ended: Final = _utc_minute()
finally:
release.set()
return Admissions(
held_forwarded_max_tokens="max_tokens" in forwarded,
statuses=(over.status_code, fits.status_code, held.result(timeout=60).status_code),
one_minute=started == ended,
)
def test_reservation_fraction_env_var_sets_what_a_request_without_max_tokens_holds(rig: Rig) -> None:
first: Final = _admissions_beside_a_held_stream(rig, GROUPS[0])
admissions: Final = first if first.one_minute else _admissions_beside_a_held_stream(rig, GROUPS[1])
assert admissions == Admissions(held_forwarded_max_tokens=False, statuses=(429, 200, 200), one_minute=True)

View file

@ -3699,6 +3699,20 @@ def test_get_guardrail_from_metadata() -> None:
) == [{"test-guardrail": {"extra_body": {"key": "value"}}}]
@pytest.mark.parametrize(
"caller_metadata",
[{"guardrails": []}, {"guardrails": ["not-a-guardrail"]}, {"guardrails": None}],
)
def test_caller_metadata_guardrails_never_hide_the_proxy_guardrails_in_litellm_metadata(caller_metadata: dict) -> None:
guardrail: Final = CustomGuardrail(
guardrail_name="key-guardrail", event_hook=GuardrailEventHooks.pre_call, default_on=False
)
request_data: Final = {"metadata": caller_metadata, "litellm_metadata": {"guardrails": ["key-guardrail"]}}
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is True
def test_guardrail_is_in_requested_guardrails() -> None:
guardrail: Final = CustomGuardrail(guardrail_name="test-guardrail")

View file

@ -5,7 +5,9 @@ Tests the short-circuit path that detects web-search-only /v1/messages requests
and executes the search directly without routing through the backend LLM.
"""
import datetime
import json
from collections.abc import Iterator, Mapping, Sequence
from typing import Final
from unittest.mock import AsyncMock, patch
@ -16,6 +18,7 @@ import litellm
from litellm.integrations.websearch_interception.handler import (
WebSearchInterceptionLogger,
)
from litellm.messages.dispatch import MessagesResult
# ---------------------------------------------------------------------------
# Detection tests
@ -73,6 +76,42 @@ class TestTryShortCircuitSearch:
assert "Result" in text_block["text"]
mock_search.assert_called_once_with("Search for Claude Code releases")
@pytest.mark.asyncio
async def test_emits_native_blocks_when_tools_already_converted(self, perplexity_search: respx.Route):
"""Pre-request/deployment hooks convert the native tool to the LiteLLM
standard shape before the short-circuit runs and record that in
WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY; the flag alone must be enough for
native blocks to be emitted."""
from litellm.integrations.websearch_interception.handler import (
WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY,
)
logger: Final = WebSearchInterceptionLogger(enabled_providers=["hosted_vllm"])
result: Final = await logger.try_short_circuit_search(
model="hosted_vllm/some-model",
messages=[{"role": "user", "content": "Search for something"}],
tools=[
{
"name": "litellm_web_search",
"description": "Search the web for information.",
"input_schema": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
}
],
custom_llm_provider="hosted_vllm",
kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True},
)
assert perplexity_search.call_count == 1
assert result is not None
assert [b["type"] for b in result["content"]] == ["server_tool_use", "web_search_tool_result", "text"]
assert result["content"][0]["name"] == "web_search"
assert [r["url"] for r in result["content"][1]["content"]] == ["https://example.com/releases"]
@pytest.mark.asyncio
async def test_does_not_short_circuit_mixed_tools(self):
"""Mix of web_search and other tools → NOT short-circuited"""
@ -468,6 +507,89 @@ class TestShortCircuitEntryPoint:
assert text_block["text"] == "results"
class TestShortCircuitThroughMessagesEntryPoint:
"""The pre-request hook renames the native tool before the short-circuit
runs, so only a request through anthropic_messages shows whether native
clients still get their server_tool_use and web_search_tool_result blocks."""
NATIVE_TOOL: Final = {"type": "web_search_20250305", "name": "web_search"}
@pytest.fixture(autouse=True)
def _intercept_web_search(self, perplexity_search: respx.Route) -> Iterator[respx.Route]:
with patch("litellm.callbacks", [WebSearchInterceptionLogger(enabled_providers=["bedrock"])]):
yield perplexity_search
@staticmethod
async def _send(*, tools: Sequence[Mapping[str, object]], stream: bool, **kwargs: object) -> MessagesResult:
return await litellm.anthropic_messages(
model="bedrock/converse/us.anthropic.claude-sonnet-5-5",
messages=[{"role": "user", "content": "latest python release"}],
tools=tools,
max_tokens=64,
stream=stream,
**kwargs,
)
@pytest.mark.asyncio
async def test_native_tool_gets_native_blocks(self, perplexity_search: respx.Route):
response: Final = await self._send(tools=[self.NATIVE_TOOL], stream=False)
assert perplexity_search.call_count == 1
assert [block["type"] for block in response["content"]] == [
"server_tool_use",
"web_search_tool_result",
"text",
]
server_tool_use, tool_result, text = response["content"]
assert server_tool_use["name"] == "web_search"
assert server_tool_use["input"] == {"query": "latest python release"}
assert tool_result["tool_use_id"] == server_tool_use["id"]
assert [result["url"] for result in tool_result["content"]] == ["https://example.com/releases"]
assert "Release notes" in text["text"]
assert "Release notes" not in json.dumps(tool_result)
@pytest.mark.asyncio
async def test_native_blocks_reply_can_be_costed(self):
from litellm.litellm_core_utils.litellm_logging import Logging
logging_obj: Final = Logging(
model="bedrock/converse/us.anthropic.claude-sonnet-5-5",
messages=[{"role": "user", "content": "latest python release"}],
stream=False,
call_type="anthropic_messages",
start_time=datetime.datetime(2026, 1, 1),
litellm_call_id="short-circuit-cost",
function_id="short-circuit-cost",
)
logging_obj.update_environment_variables(litellm_params={}, optional_params={})
response: Final = await self._send(tools=[self.NATIVE_TOOL], stream=False, litellm_logging_obj=logging_obj)
assert logging_obj.response_cost_calculator(result=response) == 0.0
@pytest.mark.asyncio
async def test_native_tool_streams_native_blocks(self):
stream: Final = await self._send(tools=[self.NATIVE_TOOL], stream=True)
body: Final = b"".join([chunk async for chunk in stream])
starts: Final = [
json.loads(line.removeprefix(b"data: "))["content_block"]
for line in body.splitlines()
if line.startswith(b"data: ") and b'"content_block_start"' in line
]
assert [block["type"] for block in starts] == ["server_tool_use", "web_search_tool_result", "text"]
assert starts[1]["tool_use_id"] == starts[0]["id"]
@pytest.mark.asyncio
async def test_litellm_standard_tool_stays_text_only(self):
from litellm.integrations.websearch_interception.tools import (
get_litellm_web_search_tool,
)
response: Final = await self._send(tools=[get_litellm_web_search_tool()], stream=False)
assert [block["type"] for block in response["content"]] == ["text"]
class TestShortCircuitSafeguardResults:
"""A short-circuited reply never reaches the backend, so the proxy answers Claude Code's safeguards itself."""

View file

@ -69,9 +69,7 @@ def test_preserve_upstream_non_openai_attributes_keeps_scalar_and_structured_val
scalar_attribute="scalar",
structured_attribute={"items": ["value"]},
)
model_response: Final = ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(content="hello"))]
)
model_response: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="hello"))])
preserve_upstream_non_openai_attributes(model_response, original_chunk)
@ -262,9 +260,7 @@ class TestMapFinishReasonAnthropic:
("content_filtered", "content_filter"),
],
)
def test_anthropic_finish_reasons(
self, provider_reason: str, expected: str
) -> None:
def test_anthropic_finish_reasons(self, provider_reason: str, expected: str) -> None:
assert map_finish_reason(provider_reason) == expected
def test_refusal(self):
@ -336,9 +332,7 @@ class TestMapFinishReasonZhipu:
class TestMapFinishReasonOpenAIPassthrough:
@pytest.mark.parametrize(
"reason", ["stop", "length", "tool_calls", "function_call", "content_filter"]
)
@pytest.mark.parametrize("reason", ["stop", "length", "tool_calls", "function_call", "content_filter"])
def test_openai_values_pass_through(self, reason):
assert map_finish_reason(reason) == reason
@ -349,9 +343,7 @@ class TestMapFinishReasonGenericError:
assert map_finish_reason("error") == "stop"
def test_lowercase_error_does_not_warn(self, mocker):
warn = mocker.patch(
"litellm.litellm_core_utils.core_helpers.verbose_logger.warning"
)
warn = mocker.patch("litellm.litellm_core_utils.core_helpers.verbose_logger.warning")
assert map_finish_reason("error") == "stop"
warn.assert_not_called()
@ -369,8 +361,7 @@ class TestFinishReasonMapOutputsAreValid:
"""Every value in _FINISH_REASON_MAP must be a valid OpenAI finish reason."""
for provider_reason, openai_reason in _FINISH_REASON_MAP.items():
assert openai_reason in VALID_OPENAI_FINISH_REASONS, (
f"Mapped value '{openai_reason}' (from '{provider_reason}') "
f"is not a valid OpenAI finish reason"
f"Mapped value '{openai_reason}' (from '{provider_reason}') is not a valid OpenAI finish reason"
)
@ -380,28 +371,18 @@ class TestRedactNestedMatchAndRegexKeys:
"assessments": [
{
"sensitiveInformationPolicy": {
"piiEntities": [
{"type": "NAME", "match": "secret-name", "action": "BLOCKED"}
]
},
"wordPolicy": {
"customWords": [{"match": "badword", "action": "BLOCKED"}]
"piiEntities": [{"type": "NAME", "match": "secret-name", "action": "BLOCKED"}]
},
"wordPolicy": {"customWords": [{"match": "badword", "action": "BLOCKED"}]},
}
],
"regex": "should-redact-key-named-regex",
}
out = redact_nested_match_and_regex_keys(payload)
assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
"match"
] == "[REDACTED]"
assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == (
"[REDACTED]"
)
assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]"
assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == ("[REDACTED]")
assert out["regex"] == "[REDACTED]"
assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][
0
]["match"] == "secret-name"
assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "secret-name"
def test_passes_through_none_and_str(self):
assert redact_nested_match_and_regex_keys(None) is None
@ -564,13 +545,17 @@ class TestIsExpectedClientError:
assert is_expected_client_error(over_budget) is True
litellm_limit = RateLimitError(
message="key over rpm", llm_provider="anthropic", model="claude-haiku-4-5",
message="key over rpm",
llm_provider="anthropic",
model="claude-haiku-4-5",
category=RateLimitErrorCategory.LITELLM_RATE_LIMIT,
)
assert is_expected_client_error(litellm_limit) is True
vendor_limit = RateLimitError(
message="rate limited upstream", llm_provider="anthropic", model="claude-haiku-4-5",
message="rate limited upstream",
llm_provider="anthropic",
model="claude-haiku-4-5",
category=RateLimitErrorCategory.VENDOR_RATE_LIMIT,
)
assert is_expected_client_error(vendor_limit) is False
@ -803,3 +788,34 @@ class TestReceiveStampSurvivesASecondMetadataBucket:
}
assert get_litellm_metadata_from_kwargs(kwargs)["litellm_received_at"] is own
def test_billed_usage_cost_bills_any_chat_upstream_without_its_own_reader() -> None:
from types import SimpleNamespace
from litellm.litellm_core_utils.core_helpers import billed_usage_cost
from litellm.types.utils import ModelResponse, Usage
usage = Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=0.037)
chat = ModelResponse(usage=usage)
assert billed_usage_cost(chat, "openai") == 0.037
assert billed_usage_cost(chat, "openrouter") == 0.037
assert billed_usage_cost(chat, "xai") is None
assert billed_usage_cost(SimpleNamespace(usage=usage), "openai") == 0.037
def test_billed_usage_cost_keeps_the_upstream_charge_after_litellm_wrote_its_bill_into_usage_cost() -> None:
from litellm.litellm_core_utils.core_helpers import billed_usage_cost, stamp_litellm_bill_on_usage
from litellm.types.utils import ModelResponse, Usage
reported = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=0.037))
stamp_litellm_bill_on_usage(reported, reported.usage, 0.0407)
stamp_litellm_bill_on_usage(reported, reported.usage, 0.0407)
assert reported.usage.cost == 0.0407
assert reported._hidden_params["response_cost"] == 0.0407
assert billed_usage_cost(reported, "openai") == 0.037
unreported = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120))
stamp_litellm_bill_on_usage(unreported, unreported.usage, 0.0015)
assert unreported.usage.cost == 0.0015
assert billed_usage_cost(unreported, "openai") is None

View file

@ -6,6 +6,7 @@ import pytest
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.env_utils import (
get_env_float,
get_env_float_in_range,
get_env_float_or_none,
get_env_int,
get_env_int_or_none,
@ -94,3 +95,12 @@ def test_unparsable_stdout_log_limit_does_not_break_logging(
assert "first line" in messages
assert "second line" in messages
assert len([message for message in messages if raw in message]) == 1
@pytest.mark.parametrize(("raw", "expected"), [("0.5", 0.5), ("1", 1.0), ("4", 0.25), ("0", 0.25), ("half", 0.25)])
def test_float_in_range_keeps_values_inside_the_range_and_falls_back_otherwise(
monkeypatch: pytest.MonkeyPatch, raw: str, expected: float
) -> None:
monkeypatch.setenv("ENV_UTILS_TEST_FRACTION", raw)
assert get_env_float_in_range("ENV_UTILS_TEST_FRACTION", 0.25, 0.01, 1.0) == expected

View file

@ -24,6 +24,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.proxy.common_utils.openai_error_payload import openai_error_type
from litellm.types.utils import LlmProviders
@ -54,8 +55,8 @@ def test_openai_cyber_policy_uses_the_structured_error_code(shape: str, code: st
assert caught.value.code == code
assert caught.value.body == body
assert caught.value.status_code == 400
assert caught.value.response is response
assert caught.value.request is request
assert caught.value.response.status_code == 400
assert caught.value.request.url == request.url
assert caught.value.response.headers["x-request-id"] == "cyber-policy-request"
@ -2940,3 +2941,80 @@ def test_transport_timeout_maps_to_timeout(
assert exc_info.value.status_code == 408
assert f"Request timed out. Error_str: {reason}" in str(exc_info.value)
def test_bedrock_server_error_keeps_the_aws_body_without_the_signed_request() -> None:
body: Final = '{"message":"Bedrock is unable to process your request."}'
signed_request: Final = httpx.Request(
"POST",
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-opus-5-5/converse",
headers={"authorization": "AWS4-HMAC-SHA256 Credential=AKIDSIGNED/20261029/us-east-1/bedrock/aws4_request"},
content=b'{"messages":[{"role":"user","content":[{"text":"hi"}]}]}',
)
upstream: Final = httpx.Response(
500, content=body, headers={"x-amzn-requestid": "req-aws"}, request=signed_request
)
with pytest.raises(litellm.ServiceUnavailableError) as raised:
exception_type(
model="bedrock/anthropic.claude-opus-5-5",
original_exception=BedrockError(status_code=500, message=body, response=upstream),
custom_llm_provider="bedrock",
)
assert raised.value.response.text == body
assert raised.value.response.headers["x-amzn-requestid"] == "req-aws"
assert "AKIDSIGNED" not in str(raised.value.request.headers)
assert raised.value.request.content == b""
@pytest.mark.parametrize("provider", ("gemini", "vertex_ai"))
def test_gemini_server_error_response_carries_the_upstream_error_body(provider: str) -> None:
body: Final = '{"error":{"code":500,"message":"Internal error encountered.","status":"INTERNAL"}}'
with pytest.raises(litellm.InternalServerError) as raised:
exception_type(
model="gemini-3.1-pro",
original_exception=VertexAIError(status_code=500, message=body, headers=None),
custom_llm_provider=provider,
)
assert raised.value.response.status_code == 500
assert raised.value.response.json() == json.loads(body)
@pytest.mark.parametrize(
("status_code", "message", "mapped_error"),
(
(500, "The server had an error processing your request. Sorry about that!", litellm.InternalServerError),
(520, "Web server is returning an unknown error", litellm.InternalServerError),
(504, "A timeout occurred", litellm.Timeout),
(500, "Mistral API raised a streaming error", litellm.APIError),
),
)
def test_openai_error_matched_by_its_message_keeps_the_upstream_body(
status_code: int, message: str, mapped_error: type[Exception]
) -> None:
body: Final = json.dumps({"error": {"message": message, "type": "server_error", "param": None, "code": None}})
sent_request: Final = httpx.Request(
"POST",
"https://api.openai.com/v1/chat/completions",
headers={"authorization": "Bearer sk-openai-secret"},
content=b'{"messages":[{"role":"user","content":"hi"}]}',
)
upstream: Final = httpx.Response(
status_code, content=body, headers={"x-request-id": "req-openai"}, request=sent_request
)
with pytest.raises(mapped_error) as raised:
exception_type(
model="gpt-5.6",
original_exception=openai.APIStatusError(message, response=upstream, body=json.loads(body)),
custom_llm_provider="openai",
)
error_response: Final = getattr(raised.value, "response")
assert error_response.text == body
assert error_response.headers["x-request-id"] == "req-openai"
assert "sk-openai-secret" not in str(error_response.request.headers)
assert error_response.request.content == b""

View file

@ -8,6 +8,7 @@ from types import MappingProxyType
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import responses
import respx
@ -19,6 +20,7 @@ from litellm.litellm_core_utils.health_check_helpers import (
HealthCheckHelpers,
default_health_check_mode,
native_health_check_mode,
wildcard_health_check_models,
)
from litellm.main import ahealth_check
from litellm.proxy._types import UserAPIKeyAuth
@ -1449,6 +1451,339 @@ async def test_groq_wildcard_health_check_uses_a_concrete_model(
assert response == {}
_WILDCARD_CATALOG: Final = MappingProxyType(
{
"databricks": (
"databricks/databricks-claude-fixture-d",
"databricks/databricks-claude-fixture-c",
"databricks/databricks-claude-fixture-b",
"databricks/databricks-claude-fixture-a",
"databricks/databricks-cheap-chat",
"databricks/databricks-claude-fixture-retired",
),
"bedrock": ("meta.llama-fixture", "anthropic.claude-fixture"),
"cloudflare": ("cloudflare/@cf/meta/fixture-chat",),
"vertex_ai": ("claude-fixture@default",),
}
)
_WILDCARD_CATALOG_COST: Final = MappingProxyType(
{
"databricks/databricks-claude-fixture-retired": 1e-10,
"databricks/databricks-cheap-chat": 1e-9,
"databricks/databricks-claude-fixture-a": 2e-9,
"databricks/databricks-claude-fixture-b": 3e-9,
"databricks/databricks-claude-fixture-c": 4e-9,
"databricks/databricks-claude-fixture-d": 5e-9,
"meta.llama-fixture": 1e-9,
"anthropic.claude-fixture": 2e-9,
"cloudflare/@cf/meta/fixture-chat": 1e-9,
"claude-fixture@default": 1e-9,
}
)
_DATABRICKS_FIXTURE_API_BASE: Final = "https://fixture.cloud.databricks.com/serving-endpoints"
@pytest.fixture
def wildcard_catalog(monkeypatch: pytest.MonkeyPatch) -> None:
providers: Final = {model: provider for provider, models in _WILDCARD_CATALOG.items() for model in models}
litellm.register_model(
{
model: {
"litellm_provider": providers[model],
"mode": "chat",
"max_tokens": 1024,
"input_cost_per_token": cost,
"output_cost_per_token": cost,
**({"deprecation_date": "2000-01-01"} if model.endswith("-retired") else {}),
}
for model, cost in _WILDCARD_CATALOG_COST.items()
}
)
for provider, models in _WILDCARD_CATALOG.items():
monkeypatch.setitem(litellm.models_by_provider, provider, list(models))
_DATABRICKS_PROVIDER_PROBES: Final = (
"databricks/databricks-claude-fixture-retired",
"databricks/databricks-cheap-chat",
"databricks/databricks-claude-fixture-a",
)
_BEDROCK_PROVIDER_PROBES: Final = ("bedrock/meta.llama-fixture", "bedrock/anthropic.claude-fixture")
_DATABRICKS_SERVED_CLAUDE_PROBES: Final = (
"databricks/databricks-claude-fixture-a",
"databricks/databricks-claude-fixture-b",
)
@pytest.mark.parametrize(
("custom_llm_provider", "wildcard", "expected"),
[
("databricks", "databricks-claude-*", (*_DATABRICKS_SERVED_CLAUDE_PROBES, *_DATABRICKS_PROVIDER_PROBES[:2])),
(
"databricks",
"databricks/databricks-claude-*",
(*_DATABRICKS_SERVED_CLAUDE_PROBES, *_DATABRICKS_PROVIDER_PROBES[:2]),
),
("databricks", "*", (*_DATABRICKS_PROVIDER_PROBES[1:], _DATABRICKS_PROVIDER_PROBES[0])),
("databricks", "databricks-claude-fixture-r*", _DATABRICKS_PROVIDER_PROBES),
("databricks", "databricks-llama-*", _DATABRICKS_PROVIDER_PROBES),
("databricks", "*-fixture", _DATABRICKS_PROVIDER_PROBES),
("databricks", "team-endpoints/*", _DATABRICKS_PROVIDER_PROBES),
("bedrock", "converse/anthropic.*", ("bedrock/converse/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)),
(
"bedrock",
"converse/*",
(
"bedrock/converse/meta.llama-fixture",
"bedrock/converse/anthropic.claude-fixture",
*_BEDROCK_PROVIDER_PROBES,
),
),
("bedrock", "invoke/anthropic.*", ("bedrock/invoke/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)),
("bedrock", "us-east-1/anthropic.*", ("bedrock/us-east-1/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)),
("bedrock", "agent/anthropic.*", _BEDROCK_PROVIDER_PROBES),
("bedrock", "nova/*", _BEDROCK_PROVIDER_PROBES),
("bedrock", "nova-2/*", _BEDROCK_PROVIDER_PROBES),
("bedrock", "mantle/*", _BEDROCK_PROVIDER_PROBES),
("bedrock", "openai/*", _BEDROCK_PROVIDER_PROBES),
("bedrock", "chat_completions/*", _BEDROCK_PROVIDER_PROBES),
("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/*", _BEDROCK_PROVIDER_PROBES),
(
"bedrock",
"converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/*",
_BEDROCK_PROVIDER_PROBES,
),
("cloudflare", "@cf/*", ("cloudflare/@cf/meta/fixture-chat",)),
("vertex_ai", "claude-*", ("vertex_ai/claude-fixture@default",)),
("vertex_ai", "meta/*", ("vertex_ai/claude-fixture@default",)),
("fixture-gateway", "team/*", ()),
],
)
@pytest.mark.usefixtures("wildcard_catalog")
def test_wildcard_health_check_probes_the_cheapest_served_chat_models_before_the_provider_cheapest(
custom_llm_provider: str, wildcard: str, expected: tuple[str, ...]
) -> None:
assert wildcard_health_check_models(wildcard=wildcard, custom_llm_provider=custom_llm_provider) == expected
def _fail_first_databricks_claude_probe(request: httpx.Request) -> httpx.Response:
model: Final = json.loads(request.content)["model"]
if model == "databricks-claude-fixture-a":
return httpx.Response(500, json={"error_code": "INTERNAL_ERROR", "message": "fixture outage"})
return httpx.Response(
200,
json={
"id": "chatcmpl-health",
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "2"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
)
@pytest.mark.asyncio
@pytest.mark.usefixtures("wildcard_catalog")
async def test_partial_wildcard_health_check_falls_back_within_the_wildcard_before_any_other_model(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock(
side_effect=_fail_first_databricks_claude_probe
)
response: Final = await ahealth_check(
{
"model": "databricks/databricks-claude-*",
"api_key": "fake-key",
"api_base": _DATABRICKS_FIXTURE_API_BASE,
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
}
)
probed_models: Final = [json.loads(call.request.content)["model"] for call in upstream.calls]
assert probed_models[0] == "databricks-claude-fixture-a"
assert probed_models[-1] == "databricks-claude-fixture-b"
assert all(model.startswith("databricks-claude-") for model in probed_models)
assert "databricks-claude-fixture-a" not in probed_models[1:]
assert response == {}
def _answer_only_the_retired_databricks_claude_model(request: httpx.Request) -> httpx.Response:
model: Final = json.loads(request.content)["model"]
if model != "databricks-claude-fixture-retired":
return httpx.Response(404, json={"error_code": "ENDPOINT_NOT_FOUND", "message": f"{model} does not exist"})
return _fail_first_databricks_claude_probe(request)
@pytest.mark.asyncio
@pytest.mark.usefixtures("wildcard_catalog")
async def test_partial_wildcard_health_check_stays_healthy_while_any_model_the_provider_probes_answers(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock(
side_effect=_answer_only_the_retired_databricks_claude_model
)
response: Final = await ahealth_check(
{
"model": "databricks/databricks-claude-*",
"api_key": "fake-key",
"api_base": _DATABRICKS_FIXTURE_API_BASE,
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
}
)
assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [
"databricks-claude-fixture-a",
"databricks-claude-fixture-b",
"databricks-claude-fixture-retired",
]
assert response == {}
@pytest.mark.asyncio
@pytest.mark.usefixtures("wildcard_catalog")
async def test_partial_wildcard_health_check_that_fails_on_every_probe_tries_its_own_models_then_the_provider_probes(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").respond(
401, json={"error_code": "UNAUTHENTICATED", "message": "fixture bad token"}
)
response: Final = await ahealth_check(
{
"model": "databricks/databricks-claude-*",
"api_key": "fake-key",
"api_base": _DATABRICKS_FIXTURE_API_BASE,
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
}
)
assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [
"databricks-claude-fixture-a",
"databricks-claude-fixture-b",
"databricks-claude-fixture-retired",
"databricks-cheap-chat",
]
assert "error" in response
_DATABRICKS_EQUAL_PRICE_MODELS: Final = (
"databricks/databricks-free-d",
"databricks/databricks-free-c",
"databricks/databricks-free-b",
"databricks/databricks-free-a",
)
def _register_equal_price_databricks_models(monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, ...]) -> None:
litellm.register_model(
{
model: {
"litellm_provider": "databricks",
"mode": "chat",
"max_tokens": 1024,
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
}
for model in listed_order
}
)
monkeypatch.setitem(litellm.models_by_provider, "databricks", list(listed_order))
@pytest.mark.parametrize("listed_order", [_DATABRICKS_EQUAL_PRICE_MODELS, _DATABRICKS_EQUAL_PRICE_MODELS[::-1]])
def test_partial_wildcard_health_check_orders_its_own_equal_price_models_by_id_whatever_the_listing_order(
monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, ...]
) -> None:
_register_equal_price_databricks_models(monkeypatch, listed_order)
assert wildcard_health_check_models(wildcard="databricks/databricks-free-*", custom_llm_provider="databricks")[
:2
] == ("databricks/databricks-free-a", "databricks/databricks-free-b")
def _answer_only_the_first_listed_databricks_free_model(request: httpx.Request) -> httpx.Response:
model: Final = json.loads(request.content)["model"]
if model != "databricks-free-d":
return httpx.Response(404, json={"error_code": "ENDPOINT_NOT_FOUND", "message": f"{model} does not exist"})
return _fail_first_databricks_claude_probe(request)
@pytest.mark.asyncio
async def test_wildcard_health_check_falls_back_to_equal_price_provider_models_in_the_provider_listing_order(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
_register_equal_price_databricks_models(monkeypatch, _DATABRICKS_EQUAL_PRICE_MODELS)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock(
side_effect=_answer_only_the_first_listed_databricks_free_model
)
response: Final = await ahealth_check(
{
"model": "databricks/databricks-claude-*",
"api_key": "fake-key",
"api_base": _DATABRICKS_FIXTURE_API_BASE,
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
}
)
assert [json.loads(call.request.content)["model"] for call in upstream.calls] == ["databricks-free-d"]
assert response == {}
@pytest.mark.asyncio
@pytest.mark.usefixtures("wildcard_catalog")
async def test_partial_wildcard_health_check_without_a_served_cost_map_model_probes_the_provider_cheapest_chat_model(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock(
side_effect=_fail_first_databricks_claude_probe
)
response: Final = await ahealth_check(
{
"model": "databricks/databricks-llama-*",
"api_key": "fake-key",
"api_base": _DATABRICKS_FIXTURE_API_BASE,
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
}
)
assert json.loads(upstream.calls[0].request.content)["model"] == "databricks-claude-fixture-retired"
assert response == {}
@pytest.mark.asyncio
async def test_partial_wildcard_health_check_on_a_provider_without_cost_map_chat_models_names_health_check_model(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
monkeypatch.setitem(litellm.models_by_provider, "databricks", [])
response: Final = await ahealth_check(
{
"model": "databricks/databricks-llama-*",
"api_key": "fake-key",
"api_base": _DATABRICKS_FIXTURE_API_BASE,
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
}
)
assert "model_info.health_check_model" in response["error"], response["error"]
assert respx_mock.calls.call_count == 0
@pytest.mark.asyncio
async def test_cohere_rerank_health_check_returns_a_successful_response(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
@ -1632,7 +1967,12 @@ async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_
monkeypatch.setitem(
litellm.model_cost,
"fake-together-chat",
{"input_cost_per_token": 1e-9, "output_cost_per_token": 1e-9, "mode": "chat", "litellm_provider": "together_ai"},
{
"input_cost_per_token": 1e-9,
"output_cost_per_token": 1e-9,
"mode": "chat",
"litellm_provider": "together_ai",
},
)
monkeypatch.setitem(litellm.models_by_provider, "together_ai", ("fake-together-chat",))
litellm.in_memory_llm_clients_cache.flush_cache()
@ -1647,4 +1987,3 @@ async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_
assert "error" not in result, result
assert json.loads(together_route.calls.last.request.content)["model"] == "fake-together-chat"
assert not other_hosts.called

View file

@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_request_utils import (
ensure_extra_body_is_safe,
flatten_form_field_values,
pick_cheapest_chat_models_from_llm_provider,
rank_chat_models_from_llm_provider,
serialize_multipart_form_fields,
)
@ -149,6 +150,24 @@ def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyP
"deepseek/fake-cheapest-embedding": {"input_cost_per_token": 1e-13, "mode": "embedding"},
"openai/fake-free-chat": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"},
"deepseek/fake-cheap-chat": {"input_cost_per_token": 1e-12, "output_cost_per_token": 1e-12, "mode": "chat"},
"deepseek/fake-retired-chat": {
"input_cost_per_token": 1e-14,
"output_cost_per_token": 1e-14,
"mode": "chat",
"deprecation_date": "2000-01-01",
},
"deepseek/fake-retiring-chat": {
"input_cost_per_token": 2e-12,
"output_cost_per_token": 2e-12,
"mode": "chat",
"deprecation_date": "2999-01-01",
},
"deepseek/fake-misdated-chat": {
"input_cost_per_token": 2.5e-12,
"output_cost_per_token": 2.5e-12,
"mode": "chat",
"deprecation_date": "next spring",
},
"fake-free-chat-without-provider": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"},
}
for name, row in rows.items():
@ -156,12 +175,73 @@ def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyP
monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, tuple(rows))
def test_pick_cheapest_chat_models_prefers_provider_routed_ids_and_skips_ids_routed_elsewhere(
def test_pick_cheapest_chat_models_prefers_provider_routed_ids_skips_ids_routed_elsewhere_and_keeps_retired_ones(
deepseek_rows_whose_cheapest_ids_route_elsewhere: None,
) -> None:
assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=10) == [
"deepseek/fake-retired-chat",
"deepseek/fake-cheap-chat",
"deepseek/fake-retiring-chat",
"deepseek/fake-misdated-chat",
"deepseek/fake-pricey-chat",
"fake-free-chat-without-provider",
]
assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-cheap-chat"]
assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-retired-chat"]
def test_rank_chat_models_flags_only_the_rows_whose_cost_map_deprecation_date_has_passed(
deepseek_rows_whose_cheapest_ids_route_elsewhere: None,
) -> None:
assert {
ranked_model.model: ranked_model.is_past_deprecation_date
for ranked_model in rank_chat_models_from_llm_provider(PICKER_PROVIDER)
} == {
"deepseek/fake-retired-chat": True,
"deepseek/fake-cheap-chat": False,
"deepseek/fake-retiring-chat": False,
"deepseek/fake-misdated-chat": False,
"deepseek/fake-pricey-chat": False,
"fake-free-chat-without-provider": False,
}
@pytest.mark.parametrize(
"listed_order",
[("deepseek/fake-tie-b", "deepseek/fake-tie-a"), ("deepseek/fake-tie-a", "deepseek/fake-tie-b")],
)
def test_pick_cheapest_chat_models_keeps_the_provider_listing_order_for_equal_price_models(
monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, str]
) -> None:
for name in listed_order:
monkeypatch.setitem(
litellm.model_cost,
name,
{
"input_cost_per_token": 1e-12,
"output_cost_per_token": 1e-12,
"mode": "chat",
"litellm_provider": PICKER_PROVIDER,
},
)
monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, listed_order)
assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=2) == list(listed_order)
def test_rank_chat_models_flags_a_retired_model_the_provider_lists_without_its_cost_map_prefix(
monkeypatch: pytest.MonkeyPatch,
) -> None:
rows: Final = {
"vertex_ai/fake-retired-vertex-chat": {"input_cost_per_token": 1e-14, "deprecation_date": "2000-01-01"},
"vertex_ai/fake-live-vertex-chat": {"input_cost_per_token": 1e-12, "deprecation_date": "2999-01-01"},
}
for name, row in rows.items():
monkeypatch.setitem(litellm.model_cost, name, {**row, "mode": "chat", "litellm_provider": "vertex_ai"})
monkeypatch.setitem(
litellm.models_by_provider, "vertex_ai", tuple(name.removeprefix("vertex_ai/") for name in rows)
)
assert [
(ranked_model.model, ranked_model.is_past_deprecation_date)
for ranked_model in rank_chat_models_from_llm_provider("vertex_ai")
] == [("fake-retired-vertex-chat", True), ("fake-live-vertex-chat", False)]

View file

@ -2659,7 +2659,7 @@ def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator():
) == pytest.approx(0.00503)
def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map():
def test_openai_compatible_streaming_reported_cost_is_billed_instead_of_the_cost_map():
import litellm
from litellm.cost_calculator import (
get_response_cost_from_hidden_params,
@ -2681,7 +2681,7 @@ def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map():
id="chatcmpl-openai-compatible",
model=model,
choices=[],
usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=3_144_000),
usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.037),
)
CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openai")
@ -2693,7 +2693,7 @@ def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map():
custom_llm_provider="openai",
call_type="completion",
optional_params={},
) == pytest.approx(2e-5)
) == pytest.approx(0.037)
def test_xai_streaming_reported_cost_still_takes_the_margin(monkeypatch):

View file

@ -1962,6 +1962,97 @@ def test_token_counter_with_redacted_thinking_content():
assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block)
def test_token_counter_with_compaction_block():
"""
A replayed on-demand compaction block (Anthropic's signed summary of the earlier turns) counts its
summary text like a text block and nothing for its signature. It used to raise, which turned every
/v1/messages/count_tokens call on a compacted history into a 500 once the provider count fell back here.
"""
model = "anthropic/claude-sonnet-5-5"
summary = "The user is building a recipe app in Python and asked for one-sentence class descriptions."
follow_up = {"role": "user", "content": [{"type": "text", "text": "Now do the same for Ingredient."}]}
signed_block = {"type": "compaction", "content": summary, "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40}
as_compaction = [{"role": "assistant", "content": [signed_block]}, follow_up]
unsigned = [{"role": "assistant", "content": [{"type": "compaction", "content": summary}]}, follow_up]
as_text = [{"role": "assistant", "content": [{"type": "text", "text": summary}]}, follow_up]
assert token_counter(model=model, messages=as_compaction) == token_counter(model=model, messages=as_text)
assert token_counter(model=model, messages=unsigned) == token_counter(model=model, messages=as_text)
def test_token_counter_counts_chat_compaction_blocks_like_their_summary_text():
"""
Chat clients replay a compaction block in the assistant message's provider_specific_fields, and inference
sends it to Anthropic as a content block, so the local count includes its summary text and skips its signature.
"""
model = "anthropic/claude-sonnet-5-5"
summary = "The user is building a recipe app in Python and asked for one-sentence class descriptions."
follow_up = {"role": "user", "content": "Now do the same for Ingredient."}
signed_block = {"type": "compaction", "content": summary, "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40}
replayed = [
{"role": "assistant", "content": None, "provider_specific_fields": {"compaction_blocks": [signed_block]}},
follow_up,
]
as_text = [{"role": "assistant", "content": summary}, follow_up]
assert token_counter(model=model, messages=replayed) == token_counter(model=model, messages=as_text)
class _BosEncoder:
def encode_batch_fast(self, texts: list[str]) -> list[list[int]]:
return [[0] * (len(text.split()) + 1) for text in texts]
def test_token_counter_charges_nothing_for_a_chat_compaction_block_with_no_summary():
"""
A tokenizer that prepends a BOS token charges one token for the empty string, so a replayed
compaction block with an empty summary must not reach it, like an empty thinking block.
"""
model = "anthropic/claude-sonnet-5-5"
tokenizer = {"type": "huggingface_tokenizer", "tokenizer": _BosEncoder()}
follow_up = {"role": "user", "content": "Now do the same for Ingredient."}
replayed = [
{
"role": "assistant",
"content": None,
"provider_specific_fields": {"compaction_blocks": [{"type": "compaction", "content": ""}]},
},
follow_up,
]
bare = [{"role": "assistant", "content": None}, follow_up]
assert token_counter(model=model, messages=replayed, custom_tokenizer=tokenizer) == token_counter(
model=model, messages=bare, custom_tokenizer=tokenizer
)
def test_count_content_list_skips_the_tokenizer_for_empty_block_text():
"""
An empty thinking block, a redacted_thinking block (no thinking text at all) and a compaction
block with no summary carry no tokens, even under a tokenizer that charges for the empty string
(the llama-2 tokenizer prepends its BOS token to every encode call). Counting them through the
tokenizer added one token per block on such models.
"""
from litellm.litellm_core_utils.token_counter import _count_content_list
def charges_for_empty_text(text: str) -> int:
return len(text.split()) + 1
empty_blocks = [
{"type": "thinking", "thinking": "", "signature": "EqcLCkYICxgCKkCrqu6lP..."},
{"type": "thinking", "thinking": None, "signature": "EqcLCkYICxgCKkCrqu6lP..."},
{"type": "redacted_thinking", "data": "EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIw"},
{"type": "compaction", "content": ""},
{"type": "text", "text": ""},
]
assert _count_content_list(charges_for_empty_text, empty_blocks, False, None) == 0
assert _count_content_list(charges_for_empty_text, [{"type": "thinking", "thinking": "two words"}], False, None) == 3
def test_token_counter_with_tool_reference_block():
"""
Regression test: a message containing an Anthropic tool-search

View file

@ -6,6 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic
"""
import json
import re
from typing import Any, Literal, Optional, Final
from unittest.mock import MagicMock, patch
@ -552,13 +553,25 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
assert chunks == original
@pytest.mark.asyncio
async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
async def test_deliver_ended_stream_tool_use_rewrite_lands_on_server_tool_use_and_tool_use_blocks(self):
handler = AnthropicMessagesHandler()
server_tool_use = [
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}),
(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}},
},
),
(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'},
},
),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
]
tool_use = self._ended_tool_use_sse_chunks()
@ -568,6 +581,26 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
+ [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]]
)
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', '{"fruit": "[MASKED]"}', "", ""]
raw = b"".join(chunks).decode()
assert '"id": "srvtoolu_1"' in raw and '"id": "toolu_1"' in raw
assert "persim" not in raw and '"query"' not in raw
@pytest.mark.asyncio
async def test_deliver_ended_stream_tool_use_rewrite_on_an_unindexed_tool_block_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = AnthropicMessagesHandler()
tool_use = self._ended_tool_use_sse_chunks()
chunks = tool_use[:1] + [tool_use[1].replace(b'"index": 0', b'"index": "0"')] + tool_use[2:]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
@ -576,6 +609,270 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
deliver_ended_stream_rewrites=True,
)
@staticmethod
def _python_masking_guardrail() -> CustomGuardrail:
class MaskPython(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional[LiteLLMLoggingObj] = None,
) -> GenericGuardrailAPIInputs:
for tool_call in inputs.get("tool_calls", []):
if isinstance(tool_call, ChatCompletionMessageToolCall):
tool_call.function.arguments = re.sub(
"python", "[MASKED]", tool_call.function.arguments, flags=re.I
)
return {
**inputs,
"texts": [re.sub("python", "[MASKED]", text, flags=re.I) for text in inputs.get("texts", [])],
}
return MaskPython(guardrail_name="test")
@staticmethod
def _python_texts_masking_guardrail() -> CustomGuardrail:
class MaskPythonTexts(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional[LiteLLMLoggingObj] = None,
) -> GenericGuardrailAPIInputs:
return {
**inputs,
"texts": [re.sub("python", "[MASKED]", text, flags=re.I) for text in inputs.get("texts", [])],
}
return MaskPythonTexts(guardrail_name="test")
@staticmethod
def _web_search_result_block() -> dict[str, object]:
return {
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_1",
"content": [
{
"type": "web_search_result",
"url": "https://docs.python.org/3.14/",
"title": "Python 3.14 docs",
"encrypted_content": "",
"page_age": None,
},
{
"type": "web_search_result",
"url": "https://peps.example/745",
"title": "PEP 745",
"encrypted_content": "",
"page_age": None,
},
],
}
@classmethod
def _web_search_stream_chunks(cls, server_tool_block_type: str = "server_tool_use") -> list[bytes]:
events = [
(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5-5",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
),
(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": server_tool_block_type,
"id": "srvtoolu_1",
"name": "web_search",
"input": {},
},
},
),
(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"query": "python 3.14 release date"}'},
},
),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
(
"content_block_start",
{"type": "content_block_start", "index": 1, "content_block": cls._web_search_result_block()},
),
("content_block_stop", {"type": "content_block_stop", "index": 1}),
(
"content_block_start",
{"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}},
),
(
"content_block_delta",
{
"type": "content_block_delta",
"index": 2,
"delta": {"type": "text_delta", "text": "Python 3.14 shipped"},
},
),
("content_block_stop", {"type": "content_block_stop", "index": 2}),
(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 2},
},
),
("message_stop", {"type": "message_stop"}),
]
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
@pytest.mark.parametrize("server_tool_block_type", ["server_tool_use", "tool_use"])
@pytest.mark.parametrize(
"guardrail_factory",
[_python_masking_guardrail, _python_texts_masking_guardrail],
ids=["texts_and_tool_calls", "texts_only"],
)
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_mask_the_web_search_query_results_and_text(
self, guardrail_factory, server_tool_block_type
):
chunks = self._web_search_stream_chunks(server_tool_block_type)
await AnthropicMessagesHandler().process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail_factory(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
raw = b"".join(chunks).decode()
assert "python" not in raw.lower()
result_block = next(
json.loads(line[len("data:") :])["content_block"]
for line in raw.split("\n")
if line.startswith("data:") and '"web_search_tool_result"' in line
)
assert [(result["title"], result["url"]) for result in result_block["content"]] == [
("[MASKED] 3.14 docs", "https://docs.[MASKED].org/3.14/"),
("PEP 745", "https://peps.example/745"),
]
assert '"id": "srvtoolu_1"' in raw and '"tool_use_id": "srvtoolu_1"' in raw
@staticmethod
def _last_text_redacting_guardrail() -> CustomGuardrail:
class RedactLastText(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": ["[REDACTED] " + inputs["texts"][-1]]}
return RedactLastText(guardrail_name="test")
@pytest.mark.asyncio
async def test_output_response_keeps_the_reply_as_the_last_text_a_guardrail_scans(self):
response = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5-5",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {"query": "python 3.14 release date"},
},
self._web_search_result_block(),
{"type": "text", "text": "Python 3.14 shipped"},
],
"stop_reason": "end_turn",
}
await AnthropicMessagesHandler().process_output_response(
response=response,
guardrail_to_apply=self._last_text_redacting_guardrail(),
litellm_logging_obj=MagicMock(),
)
assert response["content"][0]["input"] == {"query": "python 3.14 release date"}
assert response["content"][1] == self._web_search_result_block()
assert response["content"][2] == {"type": "text", "text": "[REDACTED] Python 3.14 shipped"}
@pytest.mark.asyncio
async def test_deliver_ended_stream_keeps_the_reply_as_the_last_text_a_guardrail_scans(self):
chunks = self._web_search_stream_chunks()
await AnthropicMessagesHandler().process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._last_text_redacting_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
raw = b"".join(chunks).decode()
assert "[REDACTED] Python 3.14 shipped" in raw
assert raw.count("[REDACTED]") == 1
assert json.dumps(self._web_search_result_block()) in raw
def test_ended_stream_scan_key_differs_from_a_mid_stream_key_by_the_web_search_results(self):
handler = AnthropicMessagesHandler()
chunks = self._web_search_stream_chunks()
mid_stream_key = handler.get_streaming_scan_key(chunks[:-2])
ended_key = handler.get_streaming_scan_key(chunks)
assert mid_stream_key.texts == ("Python 3.14 shipped",)
assert ended_key != mid_stream_key
assert ended_key.texts[-1] == "Python 3.14 shipped"
assert {"Python 3.14 docs", "https://docs.python.org/3.14/"} <= set(ended_key.texts)
@pytest.mark.asyncio
async def test_output_response_masks_the_web_search_query_results_and_text(self):
response = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5-5",
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_1",
"name": "web_search",
"input": {"query": "python 3.14 release date"},
},
self._web_search_result_block(),
{"type": "text", "text": "Python 3.14 shipped"},
],
"stop_reason": "end_turn",
}
await AnthropicMessagesHandler().process_output_response(
response=response,
guardrail_to_apply=self._python_masking_guardrail(),
litellm_logging_obj=MagicMock(),
)
assert "python" not in json.dumps(response).lower()
assert response["content"][0]["input"] == {"query": "[MASKED] 3.14 release date"}
assert [(result["title"], result["url"]) for result in response["content"][1]["content"]] == [
("[MASKED] 3.14 docs", "https://docs.[MASKED].org/3.14/"),
("PEP 745", "https://peps.example/745"),
]
assert response["content"][2] == {"type": "text", "text": "[MASKED] 3.14 shipped"}
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = AnthropicMessagesHandler()

View file

@ -24,7 +24,7 @@ from litellm.llms.tencent.messages.transformation import TencentAnthropicMessage
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import (
VertexAIPartnerModelsAnthropicMessagesConfig,
)
from litellm.types.llms.anthropic import AnthropicFinishReason
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, AnthropicFinishReason
from litellm.types.router import GenericLiteLLMParams
from tests.unit.llms.anthropic.test_anthropic_wif import (
ScriptedPoster,
@ -321,6 +321,66 @@ def _derived_beta_header(config: AnthropicMessagesConfig) -> dict:
)
_COMPACT_EDIT_BETA: Final = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value
_CLEARING_EDIT_BETA: Final = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
@pytest.mark.parametrize(
("context_management", "betas"),
(
pytest.param({"edits": [{"type": "compact_20260112"}]}, {_COMPACT_EDIT_BETA}, id="compact"),
pytest.param({"edits": [{"type": "clear_tool_uses_20250919"}]}, {_CLEARING_EDIT_BETA}, id="clearing"),
pytest.param(
{"edits": [{"type": "compact_20260112"}, {"type": "clear_thinking_20251015"}]},
{_COMPACT_EDIT_BETA, _CLEARING_EDIT_BETA},
id="both",
),
pytest.param(
[{"type": "compaction", "compact_threshold": 150000}], {_COMPACT_EDIT_BETA}, id="openai_list_form"
),
),
)
def test_context_management_edits_add_the_beta_each_edit_needs(
context_management: JsonValue, betas: set[str], local_beta_headers_config: None
) -> None:
headers: Final = AnthropicMessagesConfig()._update_headers_with_anthropic_beta(
headers={}, optional_params={"context_management": context_management}, messages=[]
)
assert set(headers["anthropic-beta"].split(",")) == betas
_UNMAPPABLE_LIST_FORM: Final = [{"type": "clear_tool_uses_20250919"}]
@pytest.mark.parametrize("config_class", [AnthropicMessagesConfig, VertexAIPartnerModelsAnthropicMessagesConfig])
@pytest.mark.parametrize(
("context_management", "drop_params", "forwarded"),
(
pytest.param(_UNMAPPABLE_LIST_FORM, True, None, id="unmappable_dropped_under_drop_params"),
pytest.param(_UNMAPPABLE_LIST_FORM, False, _UNMAPPABLE_LIST_FORM, id="unmappable_kept_for_the_provider_400"),
pytest.param(
[{"type": "compaction"}],
True,
{"edits": [{"type": "compact_20260112"}]},
id="mappable_still_mapped_under_drop_params",
),
),
)
def test_context_management_the_list_form_cannot_map_is_dropped_only_under_drop_params(
config_class: type[AnthropicMessagesConfig], context_management: JsonValue, drop_params: bool, forwarded: JsonValue
) -> None:
request: Final = config_class().transform_anthropic_messages_request(
model="claude-haiku-5-5",
messages=[{"role": "user", "content": "Hi"}],
anthropic_messages_optional_request_params={"max_tokens": 50, "context_management": context_management},
litellm_params=GenericLiteLLMParams(drop_params=drop_params),
headers={},
)
assert request.get("context_management") == forwarded
@pytest.mark.parametrize("config", _classifier_configs(), ids=type)
def test_routes_whose_platform_runs_the_classifier_keep_safeguards_and_the_beta(
config: AnthropicMessagesConfig, local_beta_headers_config: None

View file

@ -1,6 +1,10 @@
import json
from collections.abc import Iterable, Mapping
from typing import Final
from collections.abc import Iterable, Mapping, Sequence
from typing import Final, cast
import anthropic
import httpx
import pytest
from litellm.llms.anthropic.messages.native.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
@ -91,3 +95,140 @@ def test_message_delta_omits_absent_optional_fields():
"usage": {"input_tokens": 10, "output_tokens": 2},
}
assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None}
_WEB_SEARCH_RESULT_BLOCK: Final = {
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_01",
"content": [
{
"type": "web_search_result",
"url": "https://example.com/a",
"title": "Example A",
"encrypted_content": "abc123",
}
],
}
def _content_events(content: Sequence[Mapping[str, object]]) -> tuple[tuple[str, dict[str, object]], ...]:
response: Final = {
"id": "msg_ws",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": list(content),
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 5},
}
return tuple(
event for event in _decoded_events(FakeAnthropicMessagesStreamIterator(response=response))
if event[0].startswith("content_block_")
)
def test_server_tool_use_streams_empty_start_then_input_json_delta():
"""A real Anthropic stream opens server_tool_use with an empty input and
sends the query as an input_json_delta, which SDK accumulators read."""
caller: Final = {"type": "code_execution_20260120", "tool_id": "srvtoolu_00"}
events: Final = _content_events(
[
{
"type": "server_tool_use",
"id": "srvtoolu_01",
"name": "web_search",
"input": {"query": "example search"},
"caller": caller,
}
]
)
assert [event_type for event_type, _ in events] == [
"content_block_start",
"content_block_delta",
"content_block_stop",
]
assert events[0][1]["content_block"] == {
"type": "server_tool_use",
"id": "srvtoolu_01",
"name": "web_search",
"input": {},
"caller": caller,
}
delta: Final = events[1][1]["delta"]
assert isinstance(delta, dict)
assert delta["type"] == "input_json_delta"
assert json.loads(delta["partial_json"]) == {"query": "example search"}
def test_web_search_tool_result_rides_whole_in_content_block_start():
events: Final = _content_events([_WEB_SEARCH_RESULT_BLOCK])
assert [event_type for event_type, _ in events] == ["content_block_start", "content_block_stop"]
assert events[0][1]["content_block"] == _WEB_SEARCH_RESULT_BLOCK
def test_native_search_blocks_keep_indices_aligned_with_text():
events: Final = _content_events(
[
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "q"}},
_WEB_SEARCH_RESULT_BLOCK,
{"type": "text", "text": "Here is what I found."},
]
)
starts: Final = tuple(payload for event_type, payload in events if event_type == "content_block_start")
stops: Final = tuple(payload["index"] for event_type, payload in events if event_type == "content_block_stop")
assert [(start["index"], cast(dict[str, object], start["content_block"])["type"]) for start in starts] == [
(0, "server_tool_use"),
(1, "web_search_tool_result"),
(2, "text"),
]
assert stops == (0, 1, 2)
_CALLER: Final = {"type": "code_execution_20260120", "tool_id": "srvtoolu_00"}
_REPLAYED_BLOCKS: Final = (
{"type": "thinking", "thinking": "Look up the release page.", "signature": "sig_01"},
{"type": "redacted_thinking", "data": "EncryptedThinking=="},
{"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}, "caller": _CALLER},
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_fetch", "input": {"url": "https://example.com/a"}, "caller": _CALLER},
_WEB_SEARCH_RESULT_BLOCK,
{
"type": "text",
"text": "Example A is the newest release.",
"citations": [
{
"type": "web_search_result_location",
"cited_text": "Example A was released today.",
"url": "https://example.com/a",
"title": "Example A",
"encrypted_index": "idx_01",
}
],
},
)
def _sdk_rebuilt_content(block: Mapping[str, object]) -> list[dict[str, object]]:
response: Final = {
"id": "msg_replay",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5-5",
"content": [block],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 5},
}
body: Final = b"".join(FakeAnthropicMessagesStreamIterator(response=response))
transport: Final = httpx.MockTransport(
lambda _request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"})
)
client: Final = anthropic.Anthropic(api_key="sk-ant-test", http_client=httpx.Client(transport=transport))
with client.messages.stream(
model="claude-sonnet-5-5", max_tokens=64, messages=[{"role": "user", "content": "hi"}]
) as stream:
return [content_block.model_dump(exclude_none=True) for content_block in stream.get_final_message().content]
@pytest.mark.parametrize("block", _REPLAYED_BLOCKS, ids=lambda block: str(block["type"]))
def test_anthropic_sdk_rebuilds_every_block_field_from_the_stream(block: Mapping[str, object]):
assert _sdk_rebuilt_content(block) == [block]

View file

@ -26,7 +26,24 @@ from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAn
from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
AzureAIAnthropicCountTokensConfig,
)
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
COMPACTION_BETA = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value
COMPACTED_HISTORY = [
{
"role": "assistant",
"content": [
{
"type": "compaction",
"content": "The user is building a recipe app and asked for one-sentence class descriptions.",
"signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40,
}
],
},
{"role": "user", "content": "Now do the same for Ingredient."},
]
@pytest.mark.parametrize(
@ -890,3 +907,115 @@ async def test_counter_sends_the_structured_output_tool_choice_and_thinking_chat
inference["tool_choice"],
inference["tools"],
)
def test_headers_add_the_compaction_beta_only_for_a_compacted_history():
config = AnthropicCountTokensConfig()
auth_header = {"x-api-key": "sk-ant-api03-test-key"}
with_block = config.get_count_tokens_headers(auth_header, COMPACTED_HISTORY)["anthropic-beta"].split(",")
plain = config.get_count_tokens_headers(auth_header, [{"role": "user", "content": "hi"}])["anthropic-beta"]
assert COMPACTION_BETA in with_block
assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in with_block
assert plain == ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
@pytest.mark.asyncio
async def test_handler_sends_the_compaction_beta_with_a_compacted_history(httpx_transport_clients):
with respx.mock:
route = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 404})
)
result = await AnthropicCountTokensHandler().handle_count_tokens_request(
model="claude-sonnet-5-5",
messages=COMPACTED_HISTORY,
auth_header={"x-api-key": "sk-ant-api03-test-key"},
api_base="https://gateway.example",
)
assert result == {"input_tokens": 404}
betas = route.calls.last.request.headers["anthropic-beta"].split(",")
assert COMPACTION_BETA in betas
assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport_clients")
@pytest.mark.parametrize(
("context_management", "forwarded", "sent_beta", "unsent_beta"),
(
pytest.param(
{"edits": [{"type": "compact_20260112"}]},
{"edits": [{"type": "compact_20260112"}]},
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
id="compact",
),
pytest.param(
{"edits": [{"type": "clear_tool_uses_20250919"}]},
{"edits": [{"type": "clear_tool_uses_20250919"}]},
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
id="clearing",
),
pytest.param(
[{"type": "compaction", "compact_threshold": 150000}],
{"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]},
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
id="openai_list_form",
),
),
)
async def test_counter_forwards_context_management_with_the_beta_inference_sends(
context_management: JsonValue, forwarded: JsonValue, sent_beta: str, unsent_beta: str
) -> None:
"""Anthropic counts an unsigned compaction block only under the compact_20260112 edit that inference sends,
and a clearing edit changes the count, so the count carries the request's context_management as inference
sends it, the OpenAI list form mapped to edits, and its beta."""
with respx.mock:
route: Final = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 259})
)
result: Final = await AnthropicTokenCounter().count_tokens(
model_to_use="claude-opus-5-5",
messages=[{"role": "assistant", "content": [{"type": "compaction", "content": "Recipe app summary."}]}],
contents=None,
deployment={"litellm_params": {"api_key": "sk-ant-api03-test-key", "api_base": "https://gateway.example"}},
context_management=CountContextManagement(spec=context_management, drop_invalid_edits=False),
)
sent: Final = route.calls.last.request
betas: Final = sent.headers["anthropic-beta"].split(",")
assert result is not None and result.total_tokens == 259
assert TypeAdapter(dict[str, JsonValue]).validate_json(sent.content)["context_management"] == forwarded
assert (sent_beta in betas, unsent_beta in betas) == (True, False)
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport_clients")
@pytest.mark.parametrize("context_management", ({}, {"edits": []}, []), ids=("empty_object", "empty_edits", "empty_list"))
async def test_counter_leaves_out_context_management_without_edits(context_management: JsonValue) -> None:
"""Anthropic's count endpoint rejects a context_management with no edits, which inference accepts, so the count
leaves it out and adds no context-management beta."""
with respx.mock:
route: Final = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 11})
)
result: Final = await AnthropicTokenCounter().count_tokens(
model_to_use="claude-opus-5-5",
messages=[{"role": "user", "content": "Hello"}],
contents=None,
deployment={"litellm_params": {"api_key": "sk-ant-api03-test-key", "api_base": "https://gateway.example"}},
context_management=CountContextManagement(spec=context_management, drop_invalid_edits=False),
)
sent: Final = route.calls.last.request
betas: Final = frozenset(sent.headers.get("anthropic-beta", "").split(","))
assert result is not None and result.total_tokens == 11
assert "context_management" not in TypeAdapter(dict[str, JsonValue]).validate_json(sent.content)
assert not betas & {
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
}

View file

@ -48,7 +48,7 @@ FEDERATED_DEPLOYMENT = {
def count_tokens_headers_for(api_key: str) -> dict[str, str]:
auth_header = AnthropicModelInfo.get_auth_header(api_key=api_key)
assert auth_header is not None
return AnthropicCountTokensConfig().get_count_tokens_headers(auth_header)
return AnthropicCountTokensConfig().get_count_tokens_headers(auth_header, [{"role": "user", "content": "hi"}])
@pytest.fixture

View file

@ -1,3 +1,4 @@
import datetime
import json
from collections.abc import Awaitable, Callable, Mapping
from typing import Final
@ -12,6 +13,7 @@ import respx
import litellm
from litellm.exceptions import BadRequestError, ContentPolicyViolationError, ContextWindowExceededError
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
from litellm.llms.azure.common_utils import AzureOpenAIError
@respx.mock
@ -632,3 +634,202 @@ def test_azure_content_policy_wins_over_context_length_code(monkeypatch: pytest.
error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, policy_error))
assert type(error) is ContentPolicyViolationError
_AZURE_REFUSAL_BODY: Final = '{"error":{"message":"azure refused","type":"upstream_refusal","code":"refused"}}'
_AZURE_DEPLOYMENT_NOT_FOUND_BODY: Final = '{"error":{"code":"DeploymentNotFound","message":"no such deployment"}}'
_AZURE_INTERNAL_ERROR_BODY: Final = '{"error":{"code":"InternalServerError","message":"Internal server error"}}'
_AZURE_CONTENT_FILTER_BODY: Final = json.dumps(
{
"error": {
"message": "The response was filtered due to the prompt triggering Azure OpenAI's content management policy.",
"type": None,
"param": "prompt",
"code": "content_filter",
"status": 400,
"innererror": {
"code": "ResponsibleAIPolicyViolation",
"content_filter_result": {"violence": {"filtered": True, "severity": "high"}},
},
}
}
)
_AZURE_KEY: Final = "azure-provider-secret"
_AZURE_BASE: Final = "https://azure-upstream.openai.azure.com"
_AZURE_ARGS: Final = {"api_base": _AZURE_BASE, "api_version": "2024-10-21", "api_key": _AZURE_KEY, "max_retries": 0}
async def _azure_completion() -> None:
litellm.completion(model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], **_AZURE_ARGS)
async def _azure_completion_stream() -> None:
list(
litellm.completion(
model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **_AZURE_ARGS
)
)
async def _azure_acompletion() -> None:
await litellm.acompletion(model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], **_AZURE_ARGS)
async def _azure_acompletion_stream() -> None:
stream: Final = await litellm.acompletion(
model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **_AZURE_ARGS
)
async for _ in stream:
pass
async def _azure_embedding() -> None:
litellm.embedding(model="azure/text-embedding-3-large", input=["hi"], **_AZURE_ARGS)
async def _azure_aembedding() -> None:
await litellm.aembedding(model="azure/text-embedding-3-large", input=["hi"], **_AZURE_ARGS)
async def _azure_image_generation() -> None:
litellm.image_generation(model="azure/gpt-image-2", prompt="a cat", **_AZURE_ARGS)
async def _azure_text_completion() -> None:
litellm.text_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", **_AZURE_ARGS)
async def _azure_text_completion_stream() -> None:
list(litellm.text_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", stream=True, **_AZURE_ARGS))
async def _azure_atext_completion() -> None:
await litellm.atext_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", **_AZURE_ARGS)
async def _azure_atext_completion_stream() -> None:
stream: Final = await litellm.atext_completion(
model="azure_text/gpt-35-turbo-instruct", prompt="hi", stream=True, **_AZURE_ARGS
)
async for _ in stream:
pass
@pytest.mark.asyncio
@pytest.mark.parametrize(
"route",
(
_azure_completion,
_azure_completion_stream,
_azure_acompletion,
_azure_acompletion_stream,
_azure_embedding,
_azure_aembedding,
_azure_image_generation,
_azure_text_completion,
_azure_text_completion_stream,
_azure_atext_completion,
_azure_atext_completion_stream,
),
ids=lambda route: route.__name__.removeprefix("_azure_"),
)
@pytest.mark.parametrize(
("status_code", "body", "mapped_error"),
(
(400, _AZURE_REFUSAL_BODY, BadRequestError),
(400, _AZURE_CONTENT_FILTER_BODY, BadRequestError),
(403, _AZURE_REFUSAL_BODY, litellm.APIError),
(404, _AZURE_DEPLOYMENT_NOT_FOUND_BODY, litellm.NotFoundError),
(500, _AZURE_INTERNAL_ERROR_BODY, openai.APIError),
(502, _AZURE_REFUSAL_BODY, litellm.BadGatewayError),
(504, _AZURE_REFUSAL_BODY, litellm.Timeout),
),
)
@respx.mock
async def test_azure_error_keeps_the_upstream_body_without_the_api_key(
route: Callable[[], Awaitable[None]],
status_code: int,
body: str,
mapped_error: type[Exception],
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx.post(url__regex=r"https://azure-upstream\.openai\.azure\.com/.*").mock(
return_value=httpx.Response(status_code, content=body, headers={"content-type": "application/json"})
)
with pytest.raises(mapped_error) as raised:
await route()
error_response: Final = getattr(raised.value, "response")
assert error_response.status_code == status_code
assert error_response.text == body
assert error_response.json() == json.loads(body)
for exposed in (error_response.request, getattr(raised.value, "request")):
assert _AZURE_KEY not in str(exposed.headers)
assert exposed.content == b""
@pytest.mark.asyncio
@pytest.mark.parametrize(
"route",
(
_azure_completion,
_azure_completion_stream,
_azure_acompletion,
_azure_acompletion_stream,
_azure_embedding,
_azure_aembedding,
_azure_text_completion,
_azure_atext_completion,
_azure_atext_completion_stream,
),
ids=lambda route: route.__name__.removeprefix("_azure_"),
)
@respx.mock
async def test_azure_error_response_keeps_the_upstream_elapsed_time(
route: Callable[[], Awaitable[None]], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx.post(url__regex=r"https://azure-upstream\.openai\.azure\.com/.*").mock(
return_value=httpx.Response(400, content=_AZURE_REFUSAL_BODY, headers={"content-type": "application/json"})
)
with pytest.raises(BadRequestError) as raised:
await route()
assert raised.value.response.elapsed > datetime.timedelta(0)
def test_azure_client_side_missing_key_error_keeps_the_empty_placeholder_response() -> None:
missing_key: Final = openai.OpenAIError(
"The api_key client option must be set either by passing api_key to the client or by setting the "
"OPENAI_API_KEY environment variable"
)
with pytest.raises(litellm.AuthenticationError) as raised:
exception_type(model="azure/gpt-4.1-mini", original_exception=missing_key, custom_llm_provider="azure")
assert raised.value.response.status_code == 401
assert raised.value.response.text == ""
@pytest.mark.parametrize(
("status_code", "message", "mapped_error"),
(
(500, "azure_client is not an instance of AzureOpenAI or OpenAI", litellm.APIError),
(408, "the deployment did not answer in time", litellm.Timeout),
(504, "the deployment did not answer in time", litellm.Timeout),
),
)
def test_azure_error_litellm_raises_itself_has_no_upstream_response(
status_code: int, message: str, mapped_error: type[Exception]
) -> None:
with pytest.raises(mapped_error) as raised:
exception_type(
model="azure/gpt-4.1-mini",
original_exception=AzureOpenAIError(status_code=status_code, message=message),
custom_llm_provider="azure",
)
assert getattr(raised.value, "response", None) is None

View file

@ -1,6 +1,7 @@
import json
import traceback
from collections.abc import Callable
from types import MappingProxyType
from typing import Final
from unittest.mock import MagicMock, patch
@ -9,6 +10,7 @@ import respx
from httpx import Response
import litellm
from litellm import Router
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
from litellm.llms.azure_ai.azure_model_router.transformation import (
@ -880,6 +882,360 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages():
}
def test_azure_ai_keeps_prompt_cache_breakpoint_for_a_model_that_supports_it(_local_model_cost_map):
"""
Foundry GPT-6 deployments take OpenAI explicit prompt caching, and the breakpoint only
counts on a content part, so a system message carrying one must reach the request body as
the content-part list it arrived in while the Anthropic-only fields are still stripped.
"""
request: Final = AzureAIStudioConfig().transform_request(
model="gpt-6-astra",
messages=[
{
"role": "system",
"content": [
{
"type": "text",
"text": "stable instructions",
"prompt_cache_breakpoint": {"mode": "explicit"},
"cache_control": {"type": "ephemeral"},
}
],
},
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
],
optional_params={"prompt_cache_options": {"mode": "explicit"}},
litellm_params={},
headers={},
)
assert request["messages"][0]["content"] == [
{"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}}
]
assert request["messages"][1]["content"] == [{"type": "text", "text": "hi"}]
assert request["prompt_cache_options"] == {"mode": "explicit"}
def test_azure_ai_still_flattens_list_content_for_a_model_without_the_breakpoint_flag(_local_model_cost_map):
request: Final = AzureAIStudioConfig().transform_request(
model="gpt-4o",
messages=[
{
"role": "system",
"content": [
{"type": "text", "text": "stable ", "prompt_cache_breakpoint": {"mode": "explicit"}},
{"type": "text", "text": "instructions"},
],
},
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
],
optional_params={},
litellm_params={},
headers={},
)
assert request["messages"][0]["content"] == "stable instructions"
assert request["messages"][1]["content"] == "hi"
assert not _find_key_anywhere(request["messages"], "prompt_cache_breakpoint")
def test_azure_ai_drops_thinking_parts_from_history_while_keeping_the_breakpoint(_local_model_cost_map):
"""
A multi-turn client can echo an assistant turn back with thinking parts inside its content
list. Foundry rejects those parts, and the string conversion always dropped them, so a model
that keeps content parts must still drop them while the text parts and the breakpoint stay.
"""
request: Final = AzureAIStudioConfig().transform_request(
model="gpt-6-astra",
messages=[
{
"role": "system",
"content": [
{"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}}
],
},
{"role": "user", "content": "read the file"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "The user wants me to read a file.", "signature": "sig"},
{"type": "redacted_thinking", "data": "opaque"},
{"type": "text", "text": "Reading it now."},
],
},
{"role": "user", "content": [{"type": "text", "text": "thanks"}]},
],
optional_params={"prompt_cache_options": {"mode": "explicit"}},
litellm_params={},
headers={},
)
assert request["messages"][0]["content"] == [
{"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}}
]
assert request["messages"][2]["content"] == [{"type": "text", "text": "Reading it now."}]
assert request["messages"][3]["content"] == [{"type": "text", "text": "thanks"}]
assert not _find_key_anywhere(request["messages"], "thinking")
assert not _find_key_anywhere(request["messages"], "data")
def test_azure_ai_renames_a_cache_control_mark_to_a_breakpoint_for_a_model_that_reads_them(_local_model_cost_map):
request: Final = AzureAIStudioConfig().transform_request(
model="gpt-6-astra",
messages=[
{
"role": "system",
"content": [{"type": "text", "text": "stable instructions", "cache_control": {"type": "ephemeral"}}],
},
{"role": "user", "content": "hi", "cache_control": {"type": "ephemeral"}},
],
optional_params={"prompt_cache_options": {"mode": "explicit"}},
litellm_params={},
headers={},
)
assert request["messages"][0]["content"] == [
{"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}}
]
assert request["messages"][1]["content"] == [
{"type": "text", "text": "hi", "prompt_cache_breakpoint": {"mode": "explicit"}}
]
assert not _find_key_anywhere(request["messages"], "cache_control")
def test_azure_ai_drops_a_file_part_next_to_the_breakpoint_like_the_string_conversion(_local_model_cost_map):
messages: Final = [
{
"role": "user",
"content": [
{"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}},
{"type": "file", "file": {"file_id": "file-held-by-another-provider"}},
],
}
]
kept: Final = AzureAIStudioConfig().transform_request(
model="gpt-6-astra", messages=messages, optional_params={}, litellm_params={}, headers={}
)
flattened: Final = AzureAIStudioConfig().transform_request(
model="gpt-4o", messages=messages, optional_params={}, litellm_params={}, headers={}
)
assert kept["messages"][0]["content"] == [
{"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}}
]
assert flattened["messages"][0]["content"] == "summarize"
def test_azure_ai_sends_a_pdf_part_next_to_the_breakpoint_as_is(_local_model_cost_map):
content: Final = [
{"type": "file", "file": {"filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}},
{"type": "text", "text": "What is the codeword?", "prompt_cache_breakpoint": {"mode": "explicit"}},
]
request: Final = AzureAIStudioConfig().transform_request(
model="gpt-6-astra",
messages=[{"role": "user", "content": content}],
optional_params={},
litellm_params={},
headers={},
)
assert request["messages"] == [{"role": "user", "content": content}]
def test_azure_ai_puts_a_message_level_mark_on_the_last_text_part_when_a_file_part_follows(_local_model_cost_map):
request: Final = AzureAIStudioConfig().transform_request(
model="gpt-6-astra",
messages=[
{
"role": "user",
"cache_control": {"type": "ephemeral"},
"content": [
{"type": "text", "text": "summarize"},
{"type": "file", "file": {"file_id": "file-held-by-another-provider"}},
],
}
],
optional_params={"prompt_cache_options": {"mode": "explicit"}},
litellm_params={},
headers={},
)
assert request["messages"][0]["content"] == [
{"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}}
]
@pytest.mark.parametrize(
("registered", "keeps_parts"),
[
({"azure_ai/qa-foundry-alias": {"litellm_provider": "azure_ai", "supports_prompt_cache_breakpoint": True}}, True),
({"qa-foundry-alias": {"litellm_provider": "openai", "supports_prompt_cache_breakpoint": True}}, False),
],
ids=["azure_ai-entry-opts-in", "openai-entry-does-not-speak-for-azure_ai"],
)
def test_azure_ai_keeps_parts_only_when_the_azure_ai_entry_reads_breakpoints(
_local_model_cost_map, registered: dict[str, dict[str, object]], keeps_parts: bool
):
litellm.register_model(model_cost=registered, persist_across_reloads=False, add_to_provider_model_sets=False)
breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}
request: Final = AzureAIStudioConfig().transform_request(
model="qa-foundry-alias",
messages=[{"role": "system", "content": [breakpoint_part]}],
optional_params={},
litellm_params={"custom_llm_provider": "azure_ai", "drop_params": True},
headers={},
)
assert request["messages"][0]["content"] == ([breakpoint_part] if keeps_parts else "stable")
@pytest.mark.parametrize(
"registration_order", [("opted-in", "opted-out", "unflagged"), ("unflagged", "opted-out", "opted-in")]
)
def test_azure_ai_deployment_reads_only_its_own_breakpoint_flag_not_a_siblings_sharing_its_model_string(
_local_model_cost_map, respx_mock: respx.MockRouter, registration_order: tuple[str, ...]
) -> None:
breakpoint_flags: Final = MappingProxyType(
{
"opted-in": {"supports_prompt_cache_breakpoint": True},
"opted-out": {"supports_prompt_cache_breakpoint": False},
}
)
def deployment(host: str, model_info: dict[str, object]) -> dict[str, object]:
return {
"model_name": host,
"litellm_params": {
"model": "azure_ai/qa-shared-deployment",
"api_base": f"https://{host}.services.ai.azure.com/models",
"api_key": "azure-ai-test-key",
},
"model_info": model_info,
}
routes: Final = {
host: respx_mock.post(host=f"{host}.services.ai.azure.com", path="/models/chat/completions").mock(
return_value=Response(
200,
json={
"id": "chatcmpl-azure-ai",
"object": "chat.completion",
"created": 1700000000,
"model": "qa-shared-deployment",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
},
)
)
for host in registration_order
}
router: Final = Router(
model_list=[
deployment(host, {"id": f"{host}-id", **breakpoint_flags.get(host, {})}) for host in registration_order
]
)
marked_system: Final = {
"role": "system",
"content": [{"type": "text", "text": "stable", "cache_control": {"type": "ephemeral"}}],
}
for host in registration_order:
router.completion(model=host, messages=[marked_system, {"role": "user", "content": "hi"}])
def sent_system_content(host: str) -> object:
return json.loads(routes[host].calls.last.request.content)["messages"][0]["content"]
assert sent_system_content("opted-in") == [
{"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}
]
assert sent_system_content("opted-out") == "stable"
assert sent_system_content("unflagged") == "stable"
@pytest.mark.parametrize(
"litellm_params",
[
{"base_model": "azure_ai/gpt-6-astra"},
{"base_model": "azure_ai/gpt-6-astra", "model_info": {"id": "my-astra-deployment-id"}},
],
)
def test_azure_ai_keeps_parts_for_a_custom_deployment_name_whose_base_model_reads_breakpoints(
_local_model_cost_map, litellm_params: dict[str, object]
):
breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}
request: Final = AzureAIStudioConfig().transform_request(
model="my-astra-deployment",
messages=[{"role": "system", "content": [breakpoint_part]}],
optional_params={},
litellm_params=litellm_params,
headers={},
)
assert request["messages"][0]["content"] == [breakpoint_part]
@pytest.mark.parametrize(
"litellm_params",
[
{"base_model": "azure/cost-tracking-row"},
{"base_model": "azure/cost-tracking-row", "model_info": {"id": "astra-deployment-id"}},
],
)
def test_azure_ai_falls_back_to_the_model_name_when_the_base_model_row_does_not_read_breakpoints(
_local_model_cost_map, litellm_params: dict[str, object]
):
breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}
request: Final = AzureAIStudioConfig().transform_request(
model="gpt-6-astra",
messages=[{"role": "system", "content": [breakpoint_part]}],
optional_params={},
litellm_params=litellm_params,
headers={},
)
assert request["messages"][0]["content"] == [breakpoint_part]
@pytest.mark.parametrize(
"sent_marker, litellm_params, wire_part",
[
(
{"mode": "explicit", "ttl": "30m"},
{"drop_params": True},
{"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}},
),
({"type": "ephemeral"}, {"drop_params": True}, {"type": "text", "text": "stable"}),
(
{"mode": "explicit", "ttl": "30m"},
{},
{"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}},
),
],
)
def test_azure_ai_reshapes_breakpoints_for_foundry_only_under_drop_params(
_local_model_cost_map,
sent_marker: dict[str, object],
litellm_params: dict[str, object],
wire_part: dict[str, object],
):
request: Final = AzureAIStudioConfig().transform_request(
model="gpt-6-astra",
messages=[
{"role": "system", "content": [{"type": "text", "text": "stable", "prompt_cache_breakpoint": sent_marker}]}
],
optional_params={},
litellm_params=litellm_params,
headers={},
)
assert request["messages"][0]["content"] == [wire_part]
def test_azure_ai_reference_only_tool_result_is_sent_as_empty_text():
messages: Final = [
{"role": "user", "content": "load the WebFetch tool"},

View file

@ -197,6 +197,29 @@ def test_chat_does_not_send_the_openai_organization_to_bigmodel(monkeypatch: pyt
assert "openai-organization" not in route.calls.last.request.headers
@respx.mock
def test_chat_sends_json_schema_response_format_as_given() -> None:
route: Final = respx.post(f"{BASE}/api/paas/v4/chat/completions").respond(
json={
"id": "chat_test",
"object": "chat.completion",
"created": 1,
"model": "glm-4.6",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "{}"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
)
messages: Final = [{"role": "user", "content": "Ann likes 7"}]
response_format: Final = {
"type": "json_schema",
"json_schema": {"name": "profile", "schema": {"type": "object", "properties": {"n": {"type": "integer"}}}},
}
_ = litellm.completion(model="bigmodel/glm-4.6", messages=messages, response_format=response_format)
body: Final = json.loads(route.calls.last.request.content)
assert body["response_format"] == response_format
assert body["messages"] == messages
@pytest.mark.parametrize("is_async", (False, True))
@pytest.mark.asyncio
@respx.mock

View file

@ -93,24 +93,34 @@ async def test_acompletion_forwards_client_headers_to_provider(
assert request_headers["x-mycorp-llmcall-id"] == "abc-123"
_REJECTION_BODY: Final = '{"error":{"message":"bad prompt","type":"invalid_request_error"}}'
@pytest.fixture
def rejected_completions_endpoint():
return respx.post("https://api.openai.com/v1/completions").mock(
return_value=Response(
400,
json={"error": {"message": "bad prompt", "type": "invalid_request_error"}},
headers={"x-request-id": "req-1"},
content=_REJECTION_BODY,
headers={"x-request-id": "req-1", "content-type": "application/json"},
)
)
def _assert_carries_the_upstream_rejection(rejected: litellm.BadRequestError) -> None:
assert rejected.status_code == 400
assert rejected.litellm_response_headers["x-request-id"] == "req-1"
assert rejected.response.text == _REJECTION_BODY
assert rejected.response.json() == json.loads(_REJECTION_BODY)
assert rejected.response.headers["x-request-id"] == "req-1"
@respx.mock
def test_completion_surfaces_the_status_and_headers_of_a_rejected_request(rejected_completions_endpoint):
with pytest.raises(litellm.BadRequestError) as rejected:
text_completion(model="gpt-3.5-turbo-instruct", prompt="hello")
assert rejected.value.status_code == 400
assert rejected.value.litellm_response_headers["x-request-id"] == "req-1"
_assert_carries_the_upstream_rejection(rejected.value)
@respx.mock
@ -118,8 +128,7 @@ def test_streaming_completion_surfaces_the_status_and_headers_of_a_rejected_requ
with pytest.raises(litellm.BadRequestError) as rejected:
list(text_completion(model="gpt-3.5-turbo-instruct", prompt="hello", stream=True))
assert rejected.value.status_code == 400
assert rejected.value.litellm_response_headers["x-request-id"] == "req-1"
_assert_carries_the_upstream_rejection(rejected.value)
@respx.mock
@ -130,8 +139,7 @@ async def test_acompletion_surfaces_the_status_and_headers_of_a_rejected_request
with pytest.raises(litellm.BadRequestError) as rejected:
await atext_completion(model="gpt-3.5-turbo-instruct", prompt="hello")
assert rejected.value.status_code == 400
assert rejected.value.litellm_response_headers["x-request-id"] == "req-1"
_assert_carries_the_upstream_rejection(rejected.value)
@respx.mock

View file

@ -0,0 +1,51 @@
import json
from io import BytesIO
from typing import Final
import httpx
import pytest
import litellm
from litellm.llms.openai.common_utils import OpenAIError
_UPSTREAM_ERROR_BODY: Final = '{"error":{"message":"image is not a valid PNG","type":"invalid_request_error"}}'
def _upstream_rejection(_request: httpx.Request) -> httpx.Response:
return httpx.Response(
400,
content=_UPSTREAM_ERROR_BODY,
headers={"x-request-id": "req_variation_400", "content-type": "application/json"},
)
_VARIATION_KWARGS: Final = {
"custom_llm_provider": "openai",
"api_key": "transport-only",
"api_base": "https://upstream.example/v1",
"max_retries": 0,
}
async def _request_variation(is_async: bool) -> None:
if is_async:
await litellm.aimage_variation(image=BytesIO(b"image-bytes"), **_VARIATION_KWARGS)
return
litellm.image_variation(image=BytesIO(b"image-bytes"), **_VARIATION_KWARGS)
@pytest.mark.asyncio
@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
async def test_rejected_image_variation_carries_the_upstream_error_response(
is_async: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
transport: Final = httpx.MockTransport(_upstream_rejection)
monkeypatch.setattr(litellm, "client_session", httpx.Client(transport=transport))
monkeypatch.setattr(litellm, "aclient_session", httpx.AsyncClient(transport=transport))
with pytest.raises((litellm.BadRequestError, OpenAIError)) as raised:
await _request_variation(is_async)
assert raised.value.response.text == _UPSTREAM_ERROR_BODY
assert raised.value.response.json() == json.loads(_UPSTREAM_ERROR_BODY)
assert raised.value.response.headers["x-request-id"] == "req_variation_400"

View file

@ -1,8 +1,9 @@
import asyncio
import gzip
import importlib
import json
import os
from collections.abc import Callable, Mapping
from collections.abc import Awaitable, Callable, Mapping
from types import ModuleType
from typing import Final, Literal, Protocol
from unittest.mock import MagicMock, Mock, patch
@ -375,6 +376,277 @@ async def test_acompletion_preserves_public_errors_for_both_http_clients(
assert not http_client.is_closed
_UPSTREAM_ERROR_BODY: Final = (
'{"error":{"message":"upstream rejected the request","type":"invalid_request_error",'
'"param":"response_format.type","code":"invalid_value"}}'
)
_UPSTREAM_BASE: Final = "https://upstream.example/v1"
_UPSTREAM_TRANSPORT: Final = httpx.MockTransport(
lambda _request: httpx.Response(
400,
content=_UPSTREAM_ERROR_BODY,
headers={"x-request-id": "req_upstream_400", "x-upstream-trace": "trace-7", "content-type": "application/json"},
)
)
_MESSAGES: Final = [{"role": "user", "content": "hi"}]
def _rejecting_client() -> OpenAI:
return OpenAI(
api_key="transport-only",
base_url=_UPSTREAM_BASE,
max_retries=0,
http_client=httpx.Client(transport=_UPSTREAM_TRANSPORT),
)
def _rejecting_async_client() -> AsyncOpenAI:
return AsyncOpenAI(
api_key="transport-only",
base_url=_UPSTREAM_BASE,
max_retries=0,
http_client=httpx.AsyncClient(transport=_UPSTREAM_TRANSPORT),
)
def _upstream(provider: str, model: str) -> dict[str, object]:
return {"model": f"{provider}/{model}", "api_base": _UPSTREAM_BASE, "api_key": "transport-only", "num_retries": 0}
async def _completion(provider: str) -> None:
litellm.completion(messages=_MESSAGES, client=_rejecting_client(), **_upstream(provider, "gpt-5.6"))
async def _completion_stream(provider: str) -> None:
stream: Final = litellm.completion(
messages=_MESSAGES, client=_rejecting_client(), stream=True, **_upstream(provider, "gpt-5.6")
)
list(stream)
async def _acompletion(provider: str) -> None:
await litellm.acompletion(messages=_MESSAGES, client=_rejecting_async_client(), **_upstream(provider, "gpt-5.6"))
async def _acompletion_stream(provider: str) -> None:
stream: Final = await litellm.acompletion(
messages=_MESSAGES, client=_rejecting_async_client(), stream=True, **_upstream(provider, "gpt-5.6")
)
async for _ in stream:
pass
async def _embedding(provider: str) -> None:
litellm.embedding(input=["hi"], client=_rejecting_client(), **_upstream(provider, "text-embedding-3-large"))
async def _aembedding(provider: str) -> None:
await litellm.aembedding(
input=["hi"], client=_rejecting_async_client(), **_upstream(provider, "text-embedding-3-large")
)
async def _image_generation(provider: str) -> None:
litellm.image_generation(prompt="a cat", client=_rejecting_client(), **_upstream(provider, "gpt-image-2"))
@pytest.mark.asyncio
@pytest.mark.parametrize("provider", ("openai", "litellm_proxy"))
@pytest.mark.parametrize(
"route",
(_completion, _completion_stream, _acompletion, _acompletion_stream, _embedding, _aembedding, _image_generation),
ids=lambda route: route.__name__.lstrip("_"),
)
async def test_mapped_error_carries_the_upstream_error_response(
provider: str, route: Callable[[str], Awaitable[None]]
) -> None:
with pytest.raises(litellm.BadRequestError) as raised:
await route(provider)
assert raised.value.response.text == _UPSTREAM_ERROR_BODY
assert raised.value.response.json() == json.loads(_UPSTREAM_ERROR_BODY)
assert raised.value.response.headers["x-request-id"] == "req_upstream_400"
assert raised.value.response.headers["x-upstream-trace"] == "trace-7"
_UPSTREAM_SERVER_ERROR_BODY: Final = (
'{"error":{"message":"The server had an error while processing your request.","type":"server_error",'
'"param":null,"code":null}}'
)
_CHUNKED_UPSTREAM_HEADERS: Final = {
"content-type": "application/json",
"transfer-encoding": "chunked",
"x-request-id": "req_upstream_5xx",
}
def _chunked_upstream_client(status_code: int, content_encoding: Literal["gzip"] | None) -> httpx.AsyncClient:
body: Final = _UPSTREAM_SERVER_ERROR_BODY.encode()
encoding_headers: Final = {"content-encoding": content_encoding} if content_encoding else {}
return httpx.AsyncClient(
transport=httpx.MockTransport(
lambda _request: httpx.Response(
status_code,
headers={**_CHUNKED_UPSTREAM_HEADERS, **encoding_headers},
stream=httpx.ByteStream(gzip.compress(body) if content_encoding else body),
)
)
)
@pytest.mark.asyncio
@pytest.mark.parametrize("provider", ("openai", "litellm_proxy"))
@pytest.mark.parametrize("stream", (False, True))
@pytest.mark.parametrize(
("status_code", "mapped_error"),
(
(429, litellm.RateLimitError),
(500, litellm.InternalServerError),
(502, litellm.BadGatewayError),
(503, litellm.ServiceUnavailableError),
),
)
@pytest.mark.parametrize("content_encoding", (None, "gzip"))
async def test_retryable_error_keeps_the_decoded_upstream_body_and_its_headers(
provider: str,
stream: bool,
status_code: int,
mapped_error: type[Exception],
content_encoding: Literal["gzip"] | None,
) -> None:
client: Final = AsyncOpenAI(
api_key="transport-only",
base_url=_UPSTREAM_BASE,
max_retries=0,
http_client=_chunked_upstream_client(status_code, content_encoding),
)
with pytest.raises(mapped_error) as raised:
await litellm.acompletion(messages=_MESSAGES, client=client, stream=stream, **_upstream(provider, "gpt-5.6"))
error_response: Final = getattr(raised.value, "response")
assert error_response.status_code == status_code
assert error_response.text == _UPSTREAM_SERVER_ERROR_BODY
assert error_response.json() == json.loads(_UPSTREAM_SERVER_ERROR_BODY)
assert dict(error_response.headers) == {
**_CHUNKED_UPSTREAM_HEADERS,
**({"content-encoding": content_encoding} if content_encoding else {}),
}
_REFUSAL_BODY: Final = '{"error":{"message":"upstream refused","type":"upstream_refusal","param":null,"code":null}}'
_GATEWAY_TIMEOUT_BODY: Final = (
'{"error":{"message":"litellm.Timeout: APITimeoutError - Request timed out.","type":"invalid_request_error",'
'"param":null,"code":"408"}}'
)
_PROVIDER_KEY: Final = "sk-provider-secret"
_QUERY_SECRET: Final = "query-secret"
@pytest.mark.parametrize("provider", ("openai", "litellm_proxy"))
@pytest.mark.parametrize(
("status_code", "mapped_error", "body"),
(
(400, litellm.BadRequestError, _REFUSAL_BODY),
(403, litellm.APIError, _REFUSAL_BODY),
(408, litellm.Timeout, _REFUSAL_BODY),
(408, litellm.Timeout, _GATEWAY_TIMEOUT_BODY),
(409, litellm.APIError, _REFUSAL_BODY),
(413, litellm.APIError, _REFUSAL_BODY),
(429, litellm.RateLimitError, _REFUSAL_BODY),
(500, litellm.InternalServerError, _REFUSAL_BODY),
(504, litellm.Timeout, _REFUSAL_BODY),
),
)
def test_mapped_error_keeps_the_upstream_body_without_the_provider_credentials(
provider: str, status_code: int, mapped_error: type[Exception], body: str
) -> None:
client: Final = OpenAI(
api_key=_PROVIDER_KEY,
base_url=_UPSTREAM_BASE,
default_query={"key": _QUERY_SECRET},
max_retries=0,
http_client=httpx.Client(
transport=httpx.MockTransport(
lambda _request: httpx.Response(
status_code, content=body, headers={"content-type": "application/json"}
)
)
),
)
with pytest.raises(mapped_error) as raised:
litellm.completion(
model=f"{provider}/gpt-5.6",
messages=_MESSAGES,
client=client,
api_base=_UPSTREAM_BASE,
api_key=_PROVIDER_KEY,
num_retries=0,
)
error_response: Final = getattr(raised.value, "response")
assert error_response.status_code == status_code
assert error_response.text == body
assert error_response.json() == json.loads(body)
assert str(error_response.request.url) == f"{_UPSTREAM_BASE}/chat/completions"
for exposed in (error_response.request, getattr(raised.value, "request")):
assert _QUERY_SECRET not in str(exposed.url)
assert _PROVIDER_KEY not in str(exposed.headers)
assert exposed.content == b""
def _time_out(request: httpx.Request) -> httpx.Response:
raise httpx.ReadTimeout("no reply", request=request)
async def _sync_completion_timing_out(provider: str) -> None:
litellm.completion(
model=f"{provider}/gpt-5.6",
messages=_MESSAGES,
client=OpenAI(
api_key=_PROVIDER_KEY,
base_url=_UPSTREAM_BASE,
max_retries=0,
http_client=httpx.Client(transport=httpx.MockTransport(_time_out)),
),
api_base=_UPSTREAM_BASE,
api_key=_PROVIDER_KEY,
num_retries=0,
)
async def _async_stream_timing_out(provider: str) -> None:
stream: Final = await litellm.acompletion(
model=f"{provider}/gpt-5.6",
messages=_MESSAGES,
stream=True,
client=AsyncOpenAI(
api_key=_PROVIDER_KEY,
base_url=_UPSTREAM_BASE,
max_retries=0,
http_client=httpx.AsyncClient(transport=httpx.MockTransport(_time_out)),
),
api_base=_UPSTREAM_BASE,
api_key=_PROVIDER_KEY,
num_retries=0,
)
async for _ in stream:
pass
@pytest.mark.asyncio
@pytest.mark.parametrize("provider", ("openai", "litellm_proxy"))
@pytest.mark.parametrize("call", (_sync_completion_timing_out, _async_stream_timing_out))
async def test_timeout_before_any_reply_has_no_upstream_response(
provider: str, call: Callable[[str], Awaitable[None]]
) -> None:
with pytest.raises(litellm.Timeout) as raised:
await call(provider)
assert getattr(raised.value, "response", None) is None
@pytest.mark.asyncio
async def test_acompletion_streams_text_deltas_over_injected_transport(
http_backend: ModuleType, sdk_provider: str, monkeypatch: pytest.MonkeyPatch

View file

@ -254,7 +254,7 @@ def test_interaction_on_kyma_is_billed_at_the_cost_kyma_reports():
@pytest.mark.parametrize("model", ("openrouter/openai/gpt-4o-mini", "deepseek/deepseek-chat"))
def test_interaction_on_another_provider_keeps_its_upstream_cost_out_of_the_body(model: str):
def test_interaction_on_another_openai_compatible_provider_is_billed_at_its_reported_cost(model: str):
chat_usage: Final = {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7, "cost": 0.00721}
responses_body: Final = {
"id": "resp_upstream",
@ -278,7 +278,7 @@ def test_interaction_on_another_provider_keeps_its_upstream_cost_out_of_the_body
upstream.route(path__regex=r".*/responses$").respond(200, json=responses_body)
interaction: Final = litellm.interactions.create(model=model, input="ping", api_key="k")
assert "cost" not in interaction.model_dump(exclude_none=True)["usage"]
assert interaction.model_dump(exclude_none=True)["usage"]["cost"] == pytest.approx(0.00721)
def test_sdk_completion_cost_prices_kyma_at_the_cost_kyma_reports():
@ -410,8 +410,8 @@ def test_cost_headers_show_the_discount_and_margin_taken_from_the_cost_kyma_repo
assert float(headers["x-litellm-response-cost-discount-amount"] or "nan") == pytest.approx(0.004)
assert float(headers["x-litellm-response-cost-margin-amount"] or "nan") == pytest.approx(0.0004)
assert float(headers["x-litellm-response-cost"] or "nan") == pytest.approx(0.0044)
assert headers["x-litellm-response-cost-input"] is None
assert headers["x-litellm-response-cost-output"] is None
assert float(headers["x-litellm-response-cost-input"] or "nan") == 0.0
assert float(headers["x-litellm-response-cost-output"] or "nan") == pytest.approx(0.0044)
@pytest.mark.parametrize("thinking", ({"type": "enabled", "budget_tokens": 1024}, {"type": "adaptive"}))

View file

@ -666,7 +666,8 @@ async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_whe
*_agent_turns(2, first_round=3),
)
view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact))
with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")):
view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact))
assert isinstance(view, InferenceRequestView)
assert _tool_outputs_kept(view) == (True, True)

View file

@ -526,6 +526,29 @@ def test_both_compact_and_context_management_headers_added():
)
def test_openai_list_form_context_management_sends_the_compact_beta():
config: Final = VertexAIPartnerModelsAnthropicMessagesConfig()
with (
patch.object(config, "_ensure_access_token", return_value=("token", "test-project")),
patch.object(config, "get_complete_vertex_url", return_value="https://mock-url"),
):
updated_headers, _ = config.validate_anthropic_messages_environment(
headers={},
model="claude-vertex-ai-opus-4-6",
messages=[],
optional_params={"context_management": [{"type": "compaction", "compact_threshold": 150000}]},
litellm_params={
"vertex_ai_project": "test-project",
"vertex_ai_location": "us-central1",
"vertex_credentials": "{}",
},
api_base=None,
)
assert updated_headers["anthropic-beta"].split(",") == [ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value]
def test_validate_environment_always_refreshes_token_ignoring_stale_bearer():
"""Regression: stale Authorization in shared deployment extra_headers must not
skip token refresh on /v1/messages — _ensure_access_token is always called."""

View file

@ -2,13 +2,27 @@
Tests for Z.AI (Zhipu AI) provider - GLM models
"""
import json
import math
from collections.abc import Mapping
from typing import Final
import pytest
import respx
import litellm
from litellm import completion
from litellm.cost_calculator import cost_per_token
from litellm.llms.zai.chat.transformation import ZAIChatConfig
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
JSON_SCHEMA_RESPONSE_FORMAT: Final = {
"type": "json_schema",
"json_schema": {
"name": "profile",
"schema": {"type": "object", "properties": {"favorite_number": {"type": "integer"}}},
},
}
@pytest.fixture
@ -99,3 +113,186 @@ def test_zai_sync_completion(respx_mock, zai_response, monkeypatch):
assert response.choices[0].message.content == "Hello! How can I help you today?"
assert response.usage.total_tokens == 25
def test_zai_sends_response_format_and_maps_max_completion_tokens(
respx_mock: respx.MockRouter, zai_response: Mapping[str, object], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response)
completion(
model="zai/glm-5.3",
messages=[{"role": "user", "content": "Reply in JSON"}],
response_format={"type": "json_object"},
max_completion_tokens=64,
)
body: Final = json.loads(respx_mock.calls[0].request.content)
assert body["response_format"] == {"type": "json_object"}
assert body["max_tokens"] == 64
assert "max_completion_tokens" not in body
@pytest.mark.parametrize(
"token_limits",
[{"max_tokens": 300, "max_completion_tokens": 15}, {"max_completion_tokens": 15, "max_tokens": 300}],
)
def test_zai_max_completion_tokens_wins_over_max_tokens(token_limits: dict[str, int]) -> None:
mapped: Final = ZAIChatConfig().map_openai_params(token_limits, {}, "glm-5.3", drop_params=False)
assert mapped["max_tokens"] == 15
assert "max_completion_tokens" not in mapped
@pytest.mark.parametrize(
("over_limit_by", "drop_params", "forwarded"),
[(1, True, False), (1, False, True), (0, True, True)],
)
def test_zai_drop_params_drops_max_completion_tokens_above_model_limit(
local_model_cost_map: None,
respx_mock: respx.MockRouter,
zai_response: Mapping[str, object],
monkeypatch: pytest.MonkeyPatch,
over_limit_by: int,
drop_params: bool,
forwarded: bool,
) -> None:
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response)
requested: Final = litellm.model_cost["zai/glm-4.5-flash"]["max_output_tokens"] + over_limit_by
completion(
model="zai/glm-4.5-flash",
messages=[{"role": "user", "content": "Count"}],
max_completion_tokens=requested,
drop_params=drop_params,
)
body: Final = json.loads(respx_mock.calls[0].request.content)
assert body.get("max_tokens") == (requested if forwarded else None)
@pytest.mark.parametrize("with_over_limit_max_completion_tokens", [False, True])
def test_zai_drop_params_keeps_max_tokens_above_model_limit(
local_model_cost_map: None,
respx_mock: respx.MockRouter,
zai_response: Mapping[str, object],
monkeypatch: pytest.MonkeyPatch,
with_over_limit_max_completion_tokens: bool,
) -> None:
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response)
over_limit: Final = litellm.model_cost["zai/glm-4.5-flash"]["max_output_tokens"] + 1
over_limit_max_completion_tokens: Final = (
{"max_completion_tokens": over_limit + 1} if with_over_limit_max_completion_tokens else {}
)
completion(
model="zai/glm-4.5-flash",
messages=[{"role": "user", "content": "Count"}],
max_tokens=over_limit,
drop_params=True,
**over_limit_max_completion_tokens,
)
body: Final = json.loads(respx_mock.calls[0].request.content)
assert body["max_tokens"] == over_limit
@pytest.mark.parametrize("pass_through", [{}, {"allowed_openai_params": ["response_format"]}])
def test_zai_forwards_json_schema_to_models_that_enforce_it(
local_model_cost_map: None,
respx_mock: respx.MockRouter,
zai_response: Mapping[str, object],
monkeypatch: pytest.MonkeyPatch,
pass_through: dict[str, list[str]],
) -> None:
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response)
messages: Final = [{"role": "user", "content": "Ann likes 7"}]
completion(model="zai/glm-5.3", messages=messages, response_format=JSON_SCHEMA_RESPONSE_FORMAT, **pass_through)
body: Final = json.loads(respx_mock.calls[0].request.content)
assert body["response_format"] == JSON_SCHEMA_RESPONSE_FORMAT
assert body["messages"] == messages
@pytest.mark.parametrize("pass_through", [{}, {"allowed_openai_params": ["response_format"]}])
def test_zai_puts_json_schema_in_prompt_for_models_that_ignore_it(
local_model_cost_map: None,
respx_mock: respx.MockRouter,
zai_response: Mapping[str, object],
monkeypatch: pytest.MonkeyPatch,
pass_through: dict[str, list[str]],
) -> None:
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response)
completion(
model="zai/glm-4.6",
messages=[{"role": "user", "content": "Ann likes 7"}],
response_format=JSON_SCHEMA_RESPONSE_FORMAT,
**pass_through,
)
body: Final = json.loads(respx_mock.calls[0].request.content)
assert body["response_format"] == {"type": "json_object"}
assert body["messages"][0] == {"role": "user", "content": "Ann likes 7"}
assert body["messages"][-1]["role"] == "user"
prompt: Final = body["messages"][-1]["content"]
rendered_schema: Final = prompt[prompt.index("{") : prompt.rindex("}") + 1]
assert json.loads(rendered_schema) == JSON_SCHEMA_RESPONSE_FORMAT["json_schema"]["schema"]
def test_zai_glm_5v_turbo_bills_cached_input_at_cache_read_price(local_model_cost_map: None) -> None:
usage: Final = Usage(
prompt_tokens=1000,
completion_tokens=0,
total_tokens=1000,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=800),
)
prompt_cost, _ = cost_per_token(model="glm-5v-turbo", custom_llm_provider="zai", usage_object=usage)
entry: Final = litellm.model_cost["zai/glm-5v-turbo"]
assert math.isclose(
prompt_cost, 200 * entry["input_cost_per_token"] + 800 * entry["cache_read_input_token_cost"]
)
@pytest.mark.parametrize("model", ["glm-5v-turbo", "GLM-5V-Turbo", "GLM-5.3", "GLM-5.3-Flash"])
def test_zai_cost_resolves_case_insensitively(local_model_cost_map: None, model: str) -> None:
prompt_cost, completion_cost = cost_per_token(
model=model, custom_llm_provider="zai", prompt_tokens=1000, completion_tokens=1000
)
entry: Final = litellm.model_cost[f"zai/{model.lower()}"]
assert math.isclose(prompt_cost, 1000 * entry["input_cost_per_token"])
assert math.isclose(completion_cost, 1000 * entry["output_cost_per_token"])
def test_zai_glm_5v_turbo_accepts_thinking(
local_model_cost_map: None,
respx_mock: respx.MockRouter,
zai_response: Mapping[str, object],
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(json=zai_response)
completion(
model="zai/glm-5v-turbo",
messages=[{"role": "user", "content": "Hello"}],
thinking={"type": "disabled"},
)
body: Final = json.loads(respx_mock.calls[0].request.content)
assert body["thinking"] == {"type": "disabled"}

View file

@ -6,6 +6,7 @@ import json
import logging
import unittest
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -587,6 +588,48 @@ class TestErrorLogCarriesCallId:
assert call_id in record.getMessage()
class TestCountTokensCompactedHistory:
"""A history carrying Anthropic's signed compaction block used to fail /v1/messages/count_tokens once the count
ran on the local tokenizer, which rejected the block type."""
@pytest.mark.asyncio
async def test_count_tokens_counts_a_compaction_block_like_its_summary_text(self):
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import UserAPIKeyAuth
summary = "The user is building a recipe app and asked for one-sentence class descriptions."
follow_up = {"role": "user", "content": "Now do the same for Ingredient."}
compacted = {
"model": "claude-sonnet-5-5",
"messages": [
{
"role": "assistant",
"content": [{"type": "compaction", "content": summary, "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40}],
},
follow_up,
],
}
as_text = {
"model": "claude-sonnet-5-5",
"messages": [{"role": "assistant", "content": [{"type": "text", "text": summary}]}, follow_up],
}
request = MagicMock()
request.headers = {}
async def count(body):
with (
patch.object(ep, "read_request_body", new=AsyncMock(return_value=body)), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(proxy_server, "llm_router", None), # test-quality-ok: module global; with no router the count runs on the local tokenizer, the path that raised
):
return await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
counted = await count(compacted)
assert counted["input_tokens"] > 0
assert counted == await count(as_text)
def _claude_router() -> Router:
return Router(
model_list=[
@ -759,7 +802,10 @@ async def test_count_tokens_drops_malformed_tools_the_counted_deployment_drops_e
def _count_on_claude(
monkeypatch: pytest.MonkeyPatch, peer: Callable[[httpx.Request], httpx.Response], headers: Mapping[str, str]
monkeypatch: pytest.MonkeyPatch,
peer: Callable[[httpx.Request], httpx.Response],
headers: Mapping[str, str],
extra_body: Mapping[str, object] = MappingProxyType({}),
) -> httpx.Response:
from fastapi import FastAPI
@ -799,7 +845,7 @@ def _count_on_claude(
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth()
return TestClient(app).post(
"/v1/messages/count_tokens",
json={"model": "count-log-claude", "messages": [{"role": "user", "content": "Hello"}]},
json={"model": "count-log-claude", "messages": [{"role": "user", "content": "Hello"}], **extra_body},
headers=dict(headers),
)
@ -827,6 +873,24 @@ def test_count_tokens_logs_one_line_with_the_callers_call_id_model_and_count(
)
def test_count_tokens_returns_the_context_management_the_provider_applied(monkeypatch: pytest.MonkeyPatch) -> None:
applied: Final = {"original_input_tokens": 640}
response: Final = _count_on_claude(
monkeypatch,
lambda request: httpx.Response(
200, json={"input_tokens": 576, "context_management": applied}, request=request
),
{},
{"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}},
)
assert response.json() == {
"input_tokens": 576,
"context_management": applied,
"litellm_count_source": "provider",
}, response.text
def test_count_tokens_provider_rejection_logs_its_status_and_answers_the_same_error(
monkeypatch: pytest.MonkeyPatch, token_count_log_lines: Callable[[], tuple[dict[str, object], ...]]
) -> None:

View file

@ -1461,7 +1461,9 @@ def test_wildcard_lists_each_model_once_under_its_full_owner_id(provider: str):
import litellm
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
full_ids: Final = {f"{provider}/{model.removeprefix(f'{provider}/')}" for model in litellm.models_by_provider[provider]}
full_ids: Final = {
f"{provider}/{model.removeprefix(f'{provider}/')}" for model in litellm.models_by_provider[provider]
}
assert any(model.count("/") > 1 for model in full_ids)
assert sorted(get_known_models_from_wildcard(f"{provider}/*")) == sorted(full_ids)
@ -1776,12 +1778,57 @@ def test_partial_wildcard_lists_catalog_ids_stored_with_the_provider_prefix(
assert sorted(listed) == expected
@pytest.mark.parametrize(
("provider", "routing_segment"), [("bedrock", "converse"), ("bedrock", "us-east-1"), ("openai", "responses")]
)
def test_routing_segment_lists_the_same_catalog_ids_as_the_bare_provider_wildcard(
provider: str, routing_segment: str
) -> None:
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
from litellm.types.router import LiteLLM_Params
bare_listing: Final = get_known_models_from_wildcard(f"{provider}/*", LiteLLM_Params(model=f"{provider}/*"))
routed_listing: Final = get_known_models_from_wildcard(
f"{provider}/*", LiteLLM_Params(model=f"{provider}/{routing_segment}/*")
)
assert bare_listing
assert sorted(routed_listing) == sorted(bare_listing)
def test_provider_wildcard_lists_cost_map_models_but_not_pricing_tiers() -> None:
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
model_id: Final = "together_ai/partial-wildcard-org/model-t"
pricing_tier: Final = "together-ai-partial-wildcard-tier"
with _registered_chat_models({model_id: "together_ai", pricing_tier: "together_ai"}):
listed: Final = get_known_models_from_wildcard("together_ai/*")
assert model_id in listed
assert [listed_id for listed_id in listed if pricing_tier in listed_id] == []
def test_provider_wildcard_stops_listing_a_cost_map_model_once_it_is_retired() -> None:
import litellm
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
model_id: Final = "together_ai/partial-wildcard-org/model-retired"
with _registered_chat_models({model_id: "together_ai"}):
listed_while_served: Final = get_known_models_from_wildcard("together_ai/*")
litellm.add_known_models({model_id: {"litellm_provider": "together_ai", "deprecation_date": "2000-01-01"}})
listed_once_retired: Final = get_known_models_from_wildcard("together_ai/*")
assert model_id in listed_while_served
assert model_id not in listed_once_retired
_ROUTE_CF_ID: Final = "cloudflare/@cf/partial-route-org/model-a-instruct"
_ROUTE_CF_SUFFIXED_ID: Final = "cloudflare/@cf/partial-route-org/model-b-instruct-v2"
_ROUTE_HF_ID: Final = "cloudflare/@hf/partial-route-org/model-c"
_ROUTE_GEMINI_ID: Final = "gemini-partial-route-9-flash"
_ROUTE_OPENAI_ID: Final = "gpt-partial-route-9-mini"
_ROUTE_BEDROCK_ID: Final = "anthropic.partial-route-9-v1:0"
_ROUTE_TOGETHER_ID: Final = "together_ai/partial-route-org/model-t"
_ROUTE_CF_IDS: Final = {_ROUTE_CF_ID, _ROUTE_CF_SUFFIXED_ID}
@ -1860,6 +1907,29 @@ _ROUTE_CF_IDS: Final = {_ROUTE_CF_ID, _ROUTE_CF_SUFFIXED_ID}
set(),
id="same-name-namespace-the-catalog-lacks",
),
pytest.param(
"together/*", "together_ai/partial-route-missing-org/*", None, set(), id="alias-over-org-the-catalog-lacks"
),
pytest.param(
"openai/*",
"openai/chat_completions/*",
None,
{f"openai/chat_completions/{_ROUTE_OPENAI_ID}"},
id="chat-completions-bridge-route",
),
pytest.param(
"bedrock/*", "bedrock/llama/*", None, {f"bedrock/llama/{_ROUTE_BEDROCK_ID}"}, id="invoke-format-route"
),
pytest.param(
"together/*",
"together_ai/partial-route-org/*",
None,
{_ROUTE_TOGETHER_ID},
id="alias-over-org-a-cost-map-catalog-carries",
),
pytest.param(
"bedrock/*", "bedrock/us-east-1/*", None, {f"bedrock/us-east-1/{_ROUTE_BEDROCK_ID}"}, id="region-route"
),
pytest.param(
"*", "openai/*", "openai/gpt-partial-route-9*", set(), id="provider-prefixed-key-filter-over-catch-all"
),
@ -1905,6 +1975,7 @@ def test_every_listed_partial_wildcard_id_routes_to_the_catalog_model_it_names(
_ROUTE_GEMINI_ID: "gemini",
_ROUTE_OPENAI_ID: "openai",
_ROUTE_BEDROCK_ID: "bedrock",
_ROUTE_TOGETHER_ID: "together_ai",
}
):
v1_models: Final = get_complete_model_list(
@ -1939,6 +2010,20 @@ def test_every_listed_partial_wildcard_id_routes_to_the_catalog_model_it_names(
assert listed_model_info_names == sorted(reachable)
def test_inference_provider_route_lists_every_id_the_plain_provider_wildcard_lists() -> None:
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
from litellm.types.router import LiteLLM_Params
plain: Final = get_known_models_from_wildcard("huggingface/*", LiteLLM_Params(model="huggingface/*"))
routed: Final = get_known_models_from_wildcard(
"huggingface/*", LiteLLM_Params(model="huggingface/featherless-ai/*")
)
assert plain
assert [model for model in plain if "/" not in model.removeprefix("huggingface/")] == []
assert sorted(routed) == sorted(plain)
def test_catch_all_lists_no_id_a_more_specific_sibling_wildcard_routes_elsewhere() -> None:
from litellm import Router
from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list

View file

@ -2591,7 +2591,8 @@ async def test_agent_session_rate_limit_skipped_without_session_id():
@pytest.mark.asyncio
async def test_agent_rate_limit_from_metadata_agent_id():
@pytest.mark.parametrize("metadata_field", ("metadata", "litellm_metadata"))
async def test_agent_rate_limit_from_metadata_agent_id(metadata_field):
"""
Test that agent rate limits work when agent_id comes from
request metadata (header) rather than from the API key.
@ -2639,7 +2640,7 @@ async def test_agent_rate_limit_from_metadata_agent_id():
cache=local_cache,
data={
"model": "gpt-4",
"metadata": {"agent_id": _agent_id},
metadata_field: {"agent_id": _agent_id},
},
call_type="",
)

View file

@ -623,6 +623,56 @@ async def test_async_post_call_failure_hook_reads_used_client_oauth_token_from_t
assert json.loads(payload["metadata"])["used_client_oauth_token"] is expected
@pytest.mark.asyncio
@pytest.mark.parametrize(
"request_route, body",
[
(
"/cursor/chat/completions",
{
"input": "Hello",
"metadata": {"requester_ip_address": "203.0.113.9", "user_agent": "spoofed"},
"litellm_metadata": {"requester_ip_address": "127.0.0.1", "user_agent": "Cursor/1.0"},
},
),
(
"/v1/chat/completions",
{
"messages": [{"role": "user", "content": "Hello"}],
"metadata": {"requester_ip_address": "127.0.0.1", "user_agent": "Cursor/1.0"},
"litellm_metadata": {"requester_ip_address": "203.0.113.9", "user_agent": "spoofed"},
},
),
],
)
async def test_async_post_call_failure_hook_logs_the_proxy_stamped_ip_and_user_agent(request_route: str, body: dict):
request_data = {
"model": "gpt-4o",
"proxy_server_request": {"request_id": "test_request_id"},
**body,
}
with patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database:
await ProxyDBLogger().async_post_call_failure_hook(
request_data=request_data,
original_exception=Exception("invalid api key"),
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key", request_route=request_route),
)
payload = get_logging_payload(
kwargs=mock_update_database.call_args[1]["kwargs"],
response_obj={},
start_time=datetime.now(),
end_time=datetime.now(),
)
assert payload["requester_ip_address"] == "127.0.0.1"
assert json.loads(payload["metadata"])["user_agent"] == "Cursor/1.0"
@pytest.mark.asyncio
async def test_async_post_call_failure_hook_bills_guardrail_cost_on_blocked_request():
"""LIT-5651: a request blocked by a guardrail never reaches the LLM, but the

View file

@ -5,9 +5,12 @@ Uses mock guardrails to validate pipeline execution without external services.
"""
import copy
import json
import logging
import pickle
from typing import Literal
from collections.abc import Callable, Iterable, Mapping, Sequence
from itertools import chain
from typing import Final, Literal
from unittest.mock import MagicMock
import pytest
@ -15,6 +18,7 @@ import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
@ -1726,21 +1730,6 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up
assert chunks == [_chunk()]
@pytest.mark.asyncio
async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog):
masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}
guardrail = _LegacyHookGuardrail(
replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call])
)
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
_assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments")
assert chunks == [_chunk()]
@pytest.mark.asyncio
async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog):
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[]))
@ -1772,21 +1761,6 @@ async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alo
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog):
masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call]))
chunks = [_tool_only_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_legacy_streaming_step(
monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation()
)
_assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments")
assert chunks == [_tool_only_chunk()]
class _NoHooksGuardrail(CustomGuardrail):
pass
@ -1825,6 +1799,208 @@ async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_r
assert chunks == [_chunk()]
_AnthropicHookResponse = dict[str, list[dict[str, object]]]
class _AnthropicTextRewritingLegacyGuardrail(_LegacyHookGuardrail):
async def async_post_call_success_hook(
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: _AnthropicHookResponse
) -> None:
self.calls.append({"response": copy.deepcopy(response)})
for block in response["content"]:
if block["type"] == "text" and isinstance(text := block.get("text"), str):
block["text"] = "[REWRITTEN] " + text
def _anthropic_web_search_stream() -> list[bytes]:
result_block: Final = {
"type": "web_search_tool_result",
"tool_use_id": "srvtoolu_1",
"content": [
{"type": "web_search_result", "url": "https://example.test/a", "title": "A", "encrypted_content": ""}
],
}
events: Final = (
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "m",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"query": "a"}'},
},
{"type": "content_block_stop", "index": 0},
{"type": "content_block_start", "index": 1, "content_block": result_block},
{"type": "content_block_stop", "index": 1},
{"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 2, "delta": {"type": "text_delta", "text": "hello world"}},
{"type": "content_block_stop", "index": 2},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}},
{"type": "message_stop"},
)
return _anthropic_sse(events)
def _anthropic_sse(events: Iterable[Mapping[str, object]]) -> list[bytes]:
return [f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events]
def _anthropic_tool_only_stream() -> list[bytes]:
return _anthropic_sse(
(
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "m",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {}},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"ssn"'},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": ': "123"}'},
},
{"type": "content_block_stop", "index": 0},
{"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 2}},
{"type": "message_stop"},
)
)
def _streamed_events(chunks: Sequence[bytes]) -> list[dict[str, object]]:
return [
json.loads(line[len("data: ") :])
for line in b"".join(chunks).decode().splitlines()
if line.startswith("data: ")
]
def _streamed_block_deltas(chunks: Sequence[bytes], index: int, field: str) -> str:
deltas: Final = (
event["delta"]
for event in _streamed_events(chunks)
if event.get("index") == index and event["type"] == "content_block_delta"
)
return "".join(str(delta[field]) for delta in deltas if isinstance(delta, dict))
def _streamed_web_search_results(chunks: Sequence[bytes]) -> list[dict[str, str]]:
blocks: Final = (
event["content_block"] for event in _streamed_events(chunks) if event["type"] == "content_block_start"
)
results: Final = (
block["content"] for block in blocks if isinstance(block, dict) and block["type"] == "web_search_tool_result"
)
return list(chain.from_iterable(results))
@pytest.mark.asyncio
async def test_streaming_step_delivers_a_legacy_hook_mask_of_web_search_result_titles(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
chunks: Final = _anthropic_web_search_stream()
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result: Final = await _run_legacy_streaming_step(
monkeypatch, _AnthropicMaskingLegacyGuardrail(), chunks, translation=AnthropicMessagesHandler()
)
assert [step.outcome for step in result.step_results] == ["pass"]
assert [(item["title"], item["url"]) for item in _streamed_web_search_results(chunks)] == [
("[MASKED]", "https://example.test/a")
]
assert _streamed_block_deltas(chunks, 2, "text") == "[MASKED] world"
assert not any("discarded" in record.getMessage() for record in caplog.records)
class _AnthropicMaskingLegacyGuardrail(_LegacyHookGuardrail):
async def async_post_call_success_hook(
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: _AnthropicHookResponse
) -> None:
for block in response["content"]:
if block["type"] == "text" and isinstance(text := block.get("text"), str):
block["text"] = text.replace("hello", "[MASKED]")
if block["type"] == "tool_use" and isinstance(tool_input := block.get("input"), dict):
block["input"] = {key: "[MASKED]" for key in tool_input}
if block["type"] == "web_search_tool_result" and isinstance(results := block.get("content"), list):
block["content"] = [{**result, "title": "[MASKED]"} for result in results]
@pytest.mark.parametrize(
"stream, masked_key",
[(_anthropic_web_search_stream, "query"), (_anthropic_tool_only_stream, "ssn")],
ids=["web_search_query", "tool_only"],
)
@pytest.mark.asyncio
async def test_streaming_step_delivers_a_legacy_hook_mask_of_a_streamed_tool_input(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
stream: Callable[[], list[bytes]],
masked_key: str,
) -> None:
chunks: Final = stream()
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result: Final = await _run_legacy_streaming_step(
monkeypatch, _AnthropicMaskingLegacyGuardrail(), chunks, translation=AnthropicMessagesHandler()
)
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert json.loads(_streamed_block_deltas(chunks, 0, "partial_json")) == {masked_key: "[MASKED]"}
assert "hello" not in b"".join(chunks).decode()
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_delivers_a_legacy_text_rewrite_on_a_web_search_stream(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
guardrail: Final = _AnthropicTextRewritingLegacyGuardrail()
chunks: Final = _anthropic_web_search_stream()
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result: Final = await _run_legacy_streaming_step(
monkeypatch, guardrail, chunks, translation=AnthropicMessagesHandler()
)
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
raw: Final = b"".join(chunks).decode()
assert "[REWRITTEN] hello world" in raw
assert '"title": "A"' in raw and '"url": "https://example.test/a"' in raw
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch):
masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world"))

View file

@ -3938,7 +3938,7 @@ class TestCursorInputArmFlattening:
class TestChatCompletionsBodyDetection:
def test_routing_matrix(self):
from litellm.proxy.response_api_endpoints.endpoints import _is_chat_completions_body
from litellm.proxy.litellm_pre_call_utils import is_chat_completions_body as _is_chat_completions_body
assert _is_chat_completions_body({"messages": [{"role": "user", "content": "hi"}]}) is True
assert _is_chat_completions_body({"messages": [{"role": "user", "content": "hi"}], "input": []}) is True

View file

@ -2,6 +2,7 @@ import asyncio
import contextvars
import copy
import datetime
import gzip
import json
from collections.abc import AsyncIterable, Mapping
from types import MappingProxyType, SimpleNamespace
@ -13,6 +14,7 @@ from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stub
import anthropic
import httpx
import openai
import pytest
import respx
from fastapi import HTTPException, Request, Response, status
@ -9709,6 +9711,40 @@ async def test_saved_weights_override_caller_input_and_preserve_key_precedence(
assert logging.pre_call_hook.call_args.kwargs["data"].get("_router_weights") == expected
def _hook_data(**kwargs: object) -> object:
return kwargs["data"]
@pytest.mark.asyncio
@pytest.mark.parametrize("prompt, proxy_field, client_field", [
({"input": "hi"}, "litellm_metadata", "metadata"),
({"messages": [{"role": "user", "content": "hi"}]}, "metadata", "litellm_metadata"),
])
async def test_cursor_route_writes_proxy_fields_only_to_the_field_its_body_shape_picks(
prompt: dict[str, str | list[dict[str, str]]],
proxy_field: str,
client_field: str,
) -> None:
from litellm.proxy import proxy_server
logging: Final = MagicMock(spec=ProxyLogging)
logging.pre_call_hook = AsyncMock(side_effect=_hook_data)
processor: Final = ProxyBaseLLMRequestProcessing(data={"model": "gpt-x", "metadata": {"client": "cursor"}, **prompt})
data, _ = await processor.common_processing_pre_call_logic(
request=Request({"type": "http", "method": "POST", "path": "/cursor/chat/completions", "headers": []}),
general_settings={},
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash"),
proxy_logging_obj=logging,
proxy_config=proxy_server.ProxyConfig(),
route_type="acompletion",
logging_auth=ProxyUserAPIKeyAuth(api_key="logging-hash", key_alias="logging-alias"),
)
assert data[proxy_field]["user_api_key_alias"] == "logging-alias"
assert "queue_time_seconds" in data[proxy_field]
assert "user_api_key_alias" not in data.get(client_field, {})
assert "queue_time_seconds" not in data.get(client_field, {})
assert data["metadata"]["client"] == "cursor"
class TestPerRequestModelGroupAlias:
"""``router_settings.model_group_alias`` on a key or team has to be resolved
by the proxy: the Router resolves aliases from its own shared instance
@ -11675,6 +11711,67 @@ async def test_handle_llm_api_exception_forwards_litellm_response_headers_when_r
assert exc_info.value.headers["llm_provider-x-request-id"] == "req_openai_400"
_UPSTREAM_REJECTION_BODY: Final = b'{"error":{"message":"upstream said no","type":"upstream_rejection"}}'
def _upstream_rejection_wire(content_encoding: Literal["gzip"] | None) -> tuple[bytes, dict[str, str]]:
wire_body: Final = gzip.compress(_UPSTREAM_REJECTION_BODY) if content_encoding else _UPSTREAM_REJECTION_BODY
return wire_body, {
"x-request-id": "req_upstream",
"x-upstream-trace": "trace-7",
"content-type": "application/json",
"content-length": str(len(wire_body)),
**({"content-encoding": content_encoding} if content_encoding else {}),
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("status_code", "mapped_error"),
(
(400, litellm.BadRequestError),
(429, litellm.RateLimitError),
(503, litellm.ServiceUnavailableError),
),
)
@pytest.mark.parametrize("content_encoding", (None, "gzip"))
async def test_handle_llm_api_exception_forwards_exactly_the_upstream_headers_of_a_mapped_openai_error(
status_code: int, mapped_error: type[Exception], content_encoding: Literal["gzip"] | None
) -> None:
wire_body, upstream_headers = _upstream_rejection_wire(content_encoding)
transport: Final = httpx.MockTransport(
lambda _request: httpx.Response(status_code, content=wire_body, headers=upstream_headers)
)
with pytest.raises(mapped_error) as mapped:
await litellm.acompletion(
model="openai/gpt-5.6",
messages=[{"role": "user", "content": "hi"}],
num_retries=0,
client=openai.AsyncOpenAI(
api_key="transport-only",
base_url="https://upstream.example/v1",
max_retries=0,
http_client=httpx.AsyncClient(transport=transport),
),
)
proxy_logging_obj: Final = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
with pytest.raises(ProxyException) as forwarded:
await ProxyBaseLLMRequestProcessing(data={}).handle_llm_api_exception(
e=mapped.value,
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=proxy_logging_obj,
)
assert forwarded.value.code == str(status_code)
assert "upstream said no" in forwarded.value.message
assert {name: value for name, value in forwarded.value.headers.items() if name.startswith("llm_provider-")} == {
f"llm_provider-{name}": value for name, value in upstream_headers.items()
}
@pytest.mark.asyncio
async def test_handle_llm_api_exception_logs_bug_report_for_unmapped_error(
monkeypatch: pytest.MonkeyPatch,

View file

@ -11,6 +11,7 @@ import respx
import litellm
from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers
from litellm.litellm_core_utils.llm_request_utils import RankedChatModel
from litellm.proxy import health_check as hc_module
from litellm.proxy.health_check import (
_is_strategy_router_deployment,
@ -85,8 +86,10 @@ async def test_ahealth_check_wildcard_models_respects_max_tokens():
"""
with (
patch(
"litellm.litellm_core_utils.llm_request_utils.pick_cheapest_chat_models_from_llm_provider",
return_value=["gpt-4o-mini"],
"litellm.litellm_core_utils.llm_request_utils.rank_chat_models_from_llm_provider",
return_value=(
RankedChatModel(model="gpt-4o-mini", route_rank=0, cost=0.0, is_past_deprecation_date=False),
),
),
patch("litellm.acompletion", new_callable=AsyncMock),
):

View file

@ -121,6 +121,23 @@ class TestGetMetadataVariableName:
request = self._make_request("/chat/completions")
assert get_metadata_variable_name(request) == "metadata"
@pytest.mark.parametrize(
"body, expected",
[
({"messages": [{"role": "user", "content": "hi"}]}, "metadata"),
({"messages": []}, "metadata"),
({"input": [{"role": "user", "content": "hi"}]}, "litellm_metadata"),
({"input": "hi", "messages": []}, "litellm_metadata"),
],
)
def test_cursor_route_picks_the_metadata_field_by_body_shape(self, body, expected):
request = self._make_request("/cursor/chat/completions")
assert get_metadata_variable_name(request, body) == expected
def test_cursor_route_without_a_body_keeps_the_route_answer(self):
request = self._make_request("/cursor/chat/completions")
assert get_metadata_variable_name(request) == "metadata"
def test_returns_metadata_for_completions(self):
request = self._make_request("/v1/completions")
assert get_metadata_variable_name(request) == "metadata"
@ -681,6 +698,45 @@ async def test_add_litellm_data_to_request_rejects_non_object_metadata(field, va
assert field not in data
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/responses", "/cursor/chat/completions"])
@pytest.mark.parametrize("value, received_type", [("junk", "a string"), (5, "an integer"), ([1], "an array")])
async def test_invalid_litellm_metadata_answers_400_when_metadata_carries_tags(route, value, received_type):
from litellm.proxy.auth.auth_checks import common_checks
from litellm.proxy.utils import ProxyLogging
data = {"model": "gpt-4o", "input": "hi", "metadata": {"tags": ["client-tag"]}, "litellm_metadata": value}
request = _batches_request_mock()
await common_checks(
request_body=data,
team_object=None,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
valid_token=UserAPIKeyAuth(api_key="hashed-key"),
request=request,
)
with pytest.raises(ProxyException) as exc_info:
await add_litellm_data_to_request(
data=data,
request=request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert exc_info.value.code == "400"
assert exc_info.value.message == (
f"Invalid type for 'litellm_metadata': expected an object, but got {received_type} instead."
)
@pytest.mark.asyncio
async def test_add_litellm_data_to_request_removes_every_invalid_metadata_field_before_raising():
"""When both fields are invalid, the raise for the first must not leave the
@ -11098,3 +11154,15 @@ def test_mapped_baggage_tags_keep_non_string_client_tags_as_sent():
data: Final[dict[str, object]] = {"metadata": {"tags": [1, {"team": "x"}, "client-tag", "client-tag"]}}
LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data)
assert data["metadata"]["tags"] == [1, {"team": "x"}, "client-tag", "client-tag", "tenant:a"]
@pytest.mark.parametrize("caller_metadata", ({"tags": ["client-tag"]}, json.dumps({"tags": ["client-tag"]})))
def test_caller_tags_reach_litellm_metadata_whether_metadata_is_a_dict_or_a_json_string(caller_metadata):
data: Final[dict[str, object]] = {
"input": "hi",
"metadata": caller_metadata,
"litellm_metadata": {"tags": ["key-tag"]},
}
LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data)
assert set(data["litellm_metadata"]["tags"]) == {"key-tag", "client-tag"}
assert data["metadata"] == caller_metadata

View file

@ -4,6 +4,7 @@ and ``_handle_logging_proxy_only_error``."""
from __future__ import annotations
import asyncio
import json
from collections.abc import Callable
from datetime import datetime
from types import MappingProxyType
@ -1192,3 +1193,31 @@ async def test_handle_logging_proxy_only_path_flags_requests_that_died_before_re
blocked_by_guardrail.get(REJECTED_BEFORE_REQUEST_SETUP_KEY) is REJECTED_BEFORE_REQUEST_SETUP,
blocked_by_guardrail["litellm_call_id"],
) == (True, "client-supplied", False, "call-1")
@pytest.mark.asyncio
async def test_auth_failure_on_a_responses_body_with_a_json_string_metadata_is_logged_against_the_key(
proxy_logging, monkeypatch
):
from litellm.proxy import proxy_server
recorded: list[dict] = []
class _RecordingLogger(CustomLogger):
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
recorded.append(kwargs)
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(litellm, "callbacks", [_RecordingLogger()])
proxy_logging.alert_types = []
await proxy_logging.post_call_failure_hook(
request_data={"model": "internal-model", "input": "hi", "metadata": json.dumps({"tags": ["checkout"]})},
original_exception=litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-hashed", key_alias="checkout-key", request_route="/v1/responses"),
error_type=ProxyErrorTypes.auth_error,
route="/v1/responses",
)
assert len(recorded) == 1
assert recorded[0]["standard_logging_object"]["metadata"]["user_api_key_alias"] == "checkout-key"

View file

@ -19,7 +19,14 @@ from unittest.mock import patch
_ENV_HELPERS = frozenset(
{"get_env_int", "get_env_float", "get_env_int_in_range", "get_env_int_or_none", "get_env_float_or_none"}
{
"get_env_int",
"get_env_float",
"get_env_int_in_range",
"get_env_float_in_range",
"get_env_int_or_none",
"get_env_float_or_none",
}
)

View file

@ -13,6 +13,7 @@ from typing import Final
import httpx
import pytest
import litellm
from litellm.exceptions import (
BadRequestError,
ContentPolicyViolationError,
@ -418,3 +419,21 @@ class TestBedrockErrorHeaders:
headers = get_response_headers(dict(_response_headers))
assert headers.get("llm_provider-x-amzn-requestid") == "req-proxy-500"
@pytest.mark.parametrize(
"error_class",
(RateLimitError, ServiceUnavailableError, litellm.BadGatewayError, litellm.InternalServerError),
)
def test_error_built_from_an_unread_streaming_response_keeps_its_headers(error_class: type[Exception]) -> None:
unread: Final = httpx.Response(
500,
headers={"x-request-id": "req-unread"},
stream=httpx.ByteStream(b'{"error":"never read"}'),
)
error: Final = error_class(message="upstream failed", llm_provider="openai", model="gpt-5.6", response=unread)
error_response: Final = getattr(error, "response")
assert error_response.text == ""
assert dict(error_response.headers) == {"x-request-id": "req-unread"}

View file

@ -5443,6 +5443,46 @@ def test_stream_chunk_builder_keeps_provider_reported_usage_cost():
assert response._hidden_params["response_cost"] == pytest.approx(0.5)
def test_completion_cost_on_a_stamped_stream_still_bills_the_reported_charge():
usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "")
usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5)
chunks: Final = [
_stream_builder_text_chunk("gpt-4o", "Hello ", finish_reason="stop"),
usage_chunk,
]
response: Final = litellm.stream_chunk_builder(
chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj()
)
assert response is not None
assert litellm.completion_cost(completion_response=response) == pytest.approx(0.5)
def test_stream_chunk_builder_stamps_a_zero_operator_price_over_a_reported_cost():
logging_obj: Final = _stream_builder_logging_obj()
logging_obj.litellm_params = {
**logging_obj.litellm_params,
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
}
logging_obj.custom_cost_per_token = {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}
usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "")
usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5)
chunks: Final = [
_stream_builder_text_chunk("gpt-4o", "Hello ", finish_reason="stop"),
usage_chunk,
]
response: Final = litellm.stream_chunk_builder(
chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj
)
assert response is not None
assert getattr(response.usage, "cost", None) == 0.0
assert response._hidden_params["response_cost"] == 0.0
def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk():
from openai.types.completion_usage import CompletionUsage

View file

@ -3,6 +3,8 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits).
"""
import asyncio
from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime
from typing import Final
import pytest
@ -10,13 +12,16 @@ import pytest
import litellm
from litellm import Router
from litellm.caching.dual_cache import DualCache
from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
ITPM_CACHE_KEY,
ITPM_RESERVED_KEY,
OTPM_CACHE_KEY,
OTPM_RESERVED_KEY,
_explicit_max_tokens,
_reservation_value,
_resolve_max_tokens,
async_io_token_pre_call_check,
async_io_token_reconcile_success,
build_io_token_rate_limit_headers,
@ -30,7 +35,16 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
ModelRateLimitingCheck,
)
from litellm.types.router import RouterCacheEnum
from litellm.types.utils import ModelResponse, Usage
from litellm.utils import get_utc_datetime
MANTLE_OPUS = "bedrock_mantle/anthropic.claude-opus-4-7"
MANTLE_OPUS_MAX_OUTPUT = litellm.get_model_info(model=MANTLE_OPUS)["max_output_tokens"]
UNMAPPED_MODEL = "openai/litellm-io-unmapped-model"
UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT = 4096
AZURE_BASE_MODEL = "azure/gpt-5"
AZURE_BASE_MODEL_MAX_OUTPUT = litellm.get_model_info(model=AZURE_BASE_MODEL)["max_output_tokens"]
class TestIOTokenRateLimitHelpers:
@ -48,13 +62,23 @@ class TestIOTokenRateLimitHelpers:
# A real non-zero estimate is reserved as-is.
assert _reservation_value(42, 100) == 42
def test_resolve_max_tokens_respects_explicit_zero(self):
deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}}
# An explicit max_tokens=0 is honored, not replaced by the model default.
assert _resolve_max_tokens({"max_tokens": 0}, deployment) == 0
# max_completion_tokens is the fallback only when max_tokens is absent.
assert _resolve_max_tokens({"max_completion_tokens": 12}, deployment) == 12
assert _resolve_max_tokens({"max_output_tokens": 9}, deployment) == 9
@pytest.mark.parametrize(
"request_kwargs, expected",
[
({"max_tokens": 0}, 0),
({"max_completion_tokens": 12}, 12),
({"max_output_tokens": 9}, 9),
({"max_tokens": 5, "max_completion_tokens": 12}, 5),
({"max_tokens": None, "max_output_tokens": 9}, 9),
({"max_tokens": [100]}, None),
({"messages": []}, None),
(None, None),
],
)
def test_explicit_max_tokens_reads_every_output_cap_key(
self, request_kwargs: dict[str, object] | None, expected: int | None
):
assert _explicit_max_tokens(request_kwargs) == expected
def test_build_io_token_rate_limit_headers(self):
headers = build_io_token_rate_limit_headers(
@ -386,6 +410,244 @@ class TestModelRateLimitingCheckIOTokens:
normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY]
assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
@pytest.mark.parametrize(
"model, otpm_limit",
[
(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT // 2),
(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT - 1),
(UNMAPPED_MODEL, 1000),
(UNMAPPED_MODEL, 3),
],
)
async def test_request_without_max_tokens_is_admitted_below_model_max_output(
self, model: str, otpm_limit: int, use_sync_check: bool
):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": model, "otpm": otpm_limit},
"model_info": {"id": f"io-no-max-tokens-{otpm_limit}"},
"model_name": "no-max-tokens",
}
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
reserved: Final = request_kwargs["metadata"][OTPM_RESERVED_KEY]
assert isinstance(reserved, int)
assert 0 < reserved < otpm_limit
assert await dual_cache.async_get_cache(key=request_kwargs["metadata"][OTPM_CACHE_KEY]) == reserved
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model, otpm_limit",
[(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT), (UNMAPPED_MODEL, UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT)],
)
async def test_concurrent_requests_without_max_tokens_share_otpm_without_overshoot(
self, model: str, otpm_limit: int
):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": model, "otpm": otpm_limit},
"model_info": {"id": "io-no-max-tokens-race-id"},
"model_name": "no-max-tokens",
}
async def _attempt() -> dict[str, object] | None:
request_kwargs: Final = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
try:
await check.async_pre_call_check(deployment)
except litellm.RateLimitError:
return None
return request_kwargs["metadata"]
admitted: Final = [m for m in await asyncio.gather(*[_attempt() for _ in range(8)]) if m is not None]
otpm_key: Final = admitted[0][OTPM_CACHE_KEY]
assert len(admitted) == round(1 / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)
assert {m[OTPM_CACHE_KEY] for m in admitted} == {otpm_key}
assert await dual_cache.async_get_cache(key=otpm_key) == sum(m[OTPM_RESERVED_KEY] for m in admitted)
assert await dual_cache.async_get_cache(key=otpm_key) <= otpm_limit
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model, model_max_output",
[(MANTLE_OPUS, MANTLE_OPUS_MAX_OUTPUT), (UNMAPPED_MODEL, UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT)],
)
async def test_request_without_max_tokens_reserves_model_max_output_when_otpm_allows(
self, model: str, model_max_output: int
):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": model, "otpm": model_max_output * 10},
"model_info": {"id": "io-no-max-tokens-roomy-id"},
"model_name": "roomy",
}
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
await check.async_pre_call_check(deployment)
assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == model_max_output
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
async def test_base_model_does_not_change_the_default_otpm_reservation(self, use_sync_check: bool):
check: Final = ModelRateLimitingCheck(dual_cache=DualCache())
otpm_limit: Final = AZURE_BASE_MODEL_MAX_OUTPUT * 10
async def _reserved(model_info: dict[str, str]) -> int:
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
deployment: Final = {
"litellm_params": {"model": "azure/litellm-io-custom-deployment", "otpm": otpm_limit},
"model_info": model_info,
"model_name": "custom-name",
}
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
return request_kwargs["metadata"][OTPM_RESERVED_KEY]
with_base_model: Final = await _reserved({"id": "io-with-base-model", "base_model": AZURE_BASE_MODEL})
without_base_model: Final = await _reserved({"id": "io-without-base-model"})
assert with_base_model == without_base_model
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
async def test_unparseable_max_tokens_leaves_itpm_uncharged(self, use_sync_check: bool):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": MANTLE_OPUS, "itpm": 2000, "otpm": 2000},
"model_info": {"id": "io-bad-max-tokens-id"},
"model_name": "opus",
}
async def _pre_call(request_kwargs: dict[str, object]) -> None:
set_io_token_rate_limit_request_kwargs(request_kwargs)
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
await _pre_call({"messages": [{"role": "user", "content": "hello there"}], "max_tokens": "abc", "metadata": {}})
valid_kwargs: Final = {
"messages": [{"role": "user", "content": "hello there"}],
"max_tokens": 1,
"metadata": {},
}
await _pre_call(valid_kwargs)
itpm_counter: Final = await dual_cache.async_get_cache(key=valid_kwargs["metadata"][ITPM_CACHE_KEY])
assert itpm_counter == valid_kwargs["metadata"][ITPM_RESERVED_KEY]
@pytest.mark.asyncio
async def test_explicit_max_tokens_above_default_share_is_reserved_in_full(self):
dual_cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(dual_cache=dual_cache)
deployment: Final = {
"litellm_params": {"model": MANTLE_OPUS, "otpm": 100},
"model_info": {"id": "io-explicit-max-tokens-id"},
"model_name": "opus",
}
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 100, "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
await check.async_pre_call_check(deployment)
assert request_kwargs["metadata"][OTPM_RESERVED_KEY] == 100
set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": "hi"}], "metadata": {}})
with pytest.raises(litellm.RateLimitError):
await check.async_pre_call_check(deployment)
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
async def test_deployment_max_tokens_sets_the_reservation_unless_the_request_sets_one(self, use_sync_check: bool):
check: Final = ModelRateLimitingCheck(dual_cache=DualCache())
deployment: Final = {
"litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, "max_tokens": 100},
"model_info": {"id": "io-deployment-max-tokens-id"},
"model_name": "opus",
}
without_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
with_cap: Final = {"messages": [{"role": "user", "content": "hello"}], "max_tokens": 300, "metadata": {}}
for request_kwargs in (without_cap, with_cap):
set_io_token_rate_limit_request_kwargs(request_kwargs)
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
assert without_cap["metadata"][OTPM_RESERVED_KEY] == 100
assert with_cap["metadata"][OTPM_RESERVED_KEY] == 300
@pytest.mark.asyncio
async def test_deployment_cap_sets_the_reservation_only_through_the_param_its_api_sends(self):
check: Final = ModelRateLimitingCheck(dual_cache=DualCache())
chat: Final = {"messages": [{"role": "user", "content": "hello"}]}
responses: Final = {"input": "hello"}
async def reserved(deployment_cap: Mapping[str, int], request_shape: Mapping[str, object]) -> int:
request_kwargs: Final = {**request_shape, "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
await check.async_pre_call_check(
{
"litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, **deployment_cap},
"model_info": {"id": f"io-sent-cap-{sorted(deployment_cap)}-{sorted(request_shape)}"},
"model_name": "opus",
}
)
return request_kwargs["metadata"][OTPM_RESERVED_KEY]
chat_default: Final = await reserved({}, chat)
responses_default: Final = await reserved({}, responses)
assert chat_default == responses_default > 100
assert await reserved({"max_tokens": 100}, chat) == 100
assert await reserved({"max_completion_tokens": 100}, chat) == 100
assert await reserved({"max_output_tokens": 100}, chat) == chat_default
assert await reserved({"max_output_tokens": 100}, responses) == 100
assert await reserved({"max_tokens": 100}, responses) == responses_default
assert await reserved({"max_completion_tokens": 100}, responses) == responses_default
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_check", [False, True])
async def test_deployment_cap_above_the_default_reserves_the_default(self, use_sync_check: bool):
check: Final = ModelRateLimitingCheck(dual_cache=DualCache())
otpm_limit: Final = UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT * 5
async def reserved(deployment_cap: Mapping[str, int]) -> int:
request_kwargs: Final = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
set_io_token_rate_limit_request_kwargs(request_kwargs)
deployment: Final = {
"litellm_params": {"model": UNMAPPED_MODEL, "otpm": otpm_limit, **deployment_cap},
"model_info": {"id": f"io-deployment-cap-above-default-{deployment_cap}"},
"model_name": "unmapped",
}
if use_sync_check:
check.pre_call_check(deployment)
else:
await check.async_pre_call_check(deployment)
return request_kwargs["metadata"][OTPM_RESERVED_KEY]
default: Final = await reserved({})
assert await reserved({"max_tokens": default + 1}) == default
assert await reserved({"max_tokens": otpm_limit * 2}) == default
def test_sync_io_pre_call_reserves_and_reconciles(self):
from litellm.utils import get_utc_datetime
@ -415,9 +677,9 @@ class TestModelRateLimitingCheckIOTokens:
"standard_logging_object": {
"model_id": "io-sync-id",
"hidden_params": {"litellm_model_name": "bedrock_mantle/anthropic.claude-opus-4-7"},
"metadata": dict(request_kwargs["metadata"]),
"metadata": {},
},
"metadata": request_kwargs["metadata"],
"litellm_params": {"metadata": request_kwargs["metadata"]},
}
response = ModelResponse(
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
@ -709,6 +971,32 @@ class TestModelRateLimitingCheckIOTokens:
assert await dual_cache.async_get_cache(key=itpm_key) == 4
@pytest.mark.asyncio
@pytest.mark.parametrize("use_sync_reconcile", [False, True])
async def test_reconcile_charges_usage_above_the_reservation_to_the_minute_it_finished_in(
self, use_sync_reconcile: bool
):
dual_cache: Final = DualCache()
counter: Final = "global_router:io-late-id:bedrock_mantle/test:otpm"
reserved_minute_key: Final = f"{counter}:99-99"
await dual_cache.async_increment_cache(key=reserved_minute_key, value=10, ttl=60)
kwargs: Final = {"metadata": {ITPM_RESERVED_KEY: 0, OTPM_RESERVED_KEY: 10, OTPM_CACHE_KEY: reserved_minute_key}}
response: Final = ModelResponse(
choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}],
usage=Usage(prompt_tokens=4, completion_tokens=25, total_tokens=29),
)
minute_before: Final = get_utc_datetime().strftime("%H-%M")
if use_sync_reconcile:
io_token_reconcile_success(dual_cache, kwargs, response)
else:
await async_io_token_reconcile_success(dual_cache, kwargs, response)
finished_minutes: Final = {minute_before, get_utc_datetime().strftime("%H-%M")}
charged: Final = [await dual_cache.async_get_cache(key=f"{counter}:{minute}") for minute in finished_minutes]
assert await dual_cache.async_get_cache(key=reserved_minute_key) == 10
assert sum(value or 0 for value in charged) == 25 - 10
@pytest.mark.asyncio
async def test_reconcile_missing_usage_keeps_reservation(self):
dual_cache = DualCache()
@ -799,16 +1087,16 @@ class TestModelRateLimitingCheckIOTokens:
assert await dual_cache.async_get_cache(key=itpm_key) == 4
def test_sync_reconcile_anthropic_dict_usage(self):
def test_sync_reconcile_anthropic_dict_usage_counts_uncached_input_once(self):
dual_cache = DualCache()
itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00"
otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00"
dual_cache.set_cache(key=itpm_key, value=6, ttl=60)
dual_cache.set_cache(key=itpm_key, value=10, ttl=60)
dual_cache.set_cache(key=otpm_key, value=4, ttl=60)
kwargs = {
"metadata": {
ITPM_RESERVED_KEY: 6,
ITPM_RESERVED_KEY: 10,
OTPM_RESERVED_KEY: 4,
ITPM_CACHE_KEY: itpm_key,
OTPM_CACHE_KEY: otpm_key,
@ -816,12 +1104,17 @@ class TestModelRateLimitingCheckIOTokens:
}
response = {
"type": "message",
"usage": {"input_tokens": 3, "output_tokens": 2, "cache_read_input_tokens": 1},
"usage": {
"input_tokens": 3,
"output_tokens": 2,
"cache_read_input_tokens": 1,
"cache_creation_input_tokens": 4,
},
}
io_token_reconcile_success(dual_cache, kwargs, response)
assert dual_cache.get_cache(key=itpm_key) == 2
assert dual_cache.get_cache(key=itpm_key) == 3 + 4
assert dual_cache.get_cache(key=otpm_key) == 2
@pytest.mark.asyncio
@ -958,7 +1251,214 @@ class TestModelRateLimitingCheckIOTokens:
assert current == 0
ROUTER_IO_MODEL: Final = "anthropic/claude-haiku-5-5"
RouterCall = Callable[[Router, str | Exception], Awaitable[tuple[int, int]]]
def _router_with_io_limits() -> Router:
return Router(
model_list=[
{
"model_name": "limited",
"litellm_params": {"model": ROUTER_IO_MODEL, "api_key": "sk-test", "itpm": 50000, "otpm": 50000},
}
],
optional_pre_call_checks=["enforce_model_rate_limits"],
num_retries=1,
)
def _current_minute() -> str:
return get_utc_datetime().strftime("%H-%M")
def _counter(router: Router, template: RouterCacheEnum, minutes: frozenset[str]) -> int:
deployment: Final = router.model_list[0]
return sum(
router.cache.get_cache(
key=template.value.format(
id=deployment["model_info"]["id"],
model=deployment["litellm_params"]["model"],
current_minute=minute,
)
)
or 0
for minute in minutes
)
async def _io_counters_since(router: Router, first_minute: str) -> tuple[int, int]:
await asyncio.sleep(0)
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
minutes: Final = frozenset({first_minute, _current_minute()})
return _counter(router, RouterCacheEnum.ITPM, minutes), _counter(router, RouterCacheEnum.OTPM, minutes)
async def _sdk_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]:
response: Final = await router.acompletion(
model="limited",
messages=[{"role": "user", "content": "hello there"}],
max_tokens=100,
mock_response=mock_response,
)
return response.usage.prompt_tokens, response.usage.completion_tokens
async def _proxy_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]:
response: Final = await router.acompletion(
model="limited",
messages=[{"role": "user", "content": "hello there"}],
max_tokens=100,
metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_response,
)
return response.usage.prompt_tokens, response.usage.completion_tokens
async def _proxy_responses(router: Router, mock_response: str | Exception) -> tuple[int, int]:
response: Final = await router.aresponses(
model="limited",
input="hello there",
max_output_tokens=100,
litellm_metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_response,
)
return response.usage.input_tokens, response.usage.output_tokens
async def _proxy_messages(router: Router, mock_response: str | Exception) -> tuple[int, int]:
response: Final = await router.aanthropic_messages(
model="limited",
messages=[{"role": "user", "content": "hello there"}],
max_tokens=100,
litellm_metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_response,
)
return response["usage"]["input_tokens"], response["usage"]["output_tokens"]
ProxyCallSharingLogging = Callable[[Router, str | Exception, LiteLLMLoggingObject], Awaitable[object]]
PROXY_MESSAGES: Final = [{"role": "user", "content": "hello there"}]
def _logging_obj_shared_across_attempts(original_function: str) -> LiteLLMLoggingObject:
logging_obj, _ = litellm.utils.function_setup(
original_function,
litellm.utils.Rules(),
datetime.now(),
model="limited",
messages=PROXY_MESSAGES,
litellm_call_id="proxy-call",
)
return logging_obj
async def _proxy_chat_sharing_logging(
router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject
) -> object:
return await router.acompletion(
model="limited",
messages=PROXY_MESSAGES,
max_tokens=100,
metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_error,
litellm_call_id="proxy-call",
litellm_logging_obj=logging_obj,
)
async def _proxy_responses_sharing_logging(
router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject
) -> object:
return await router.aresponses(
model="limited",
input="hello there",
max_output_tokens=100,
litellm_metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_error,
litellm_call_id="proxy-call",
litellm_logging_obj=logging_obj,
)
async def _proxy_messages_sharing_logging(
router: Router, mock_error: str | Exception, logging_obj: LiteLLMLoggingObject
) -> object:
return await router.aanthropic_messages(
model="limited",
messages=PROXY_MESSAGES,
max_tokens=100,
litellm_metadata={"user_api_key_hash": "hashed-key"},
mock_response=mock_error,
litellm_call_id="proxy-call",
litellm_logging_obj=logging_obj,
)
class TestRouterIOTokenIntegration:
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call",
[_sdk_chat, _proxy_chat, _proxy_responses, _proxy_messages],
ids=["sdk_chat", "proxy_chat_metadata", "proxy_responses_litellm_metadata", "proxy_messages_litellm_metadata"],
)
async def test_success_reconciles_io_counters_to_real_usage(self, call: RouterCall):
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
usage: Final = await call(router, "hi")
assert await _io_counters_since(router, first_minute) == usage
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call", "mock_error"),
[
(
_proxy_responses,
litellm.InternalServerError(message="mock 500", llm_provider="anthropic", model=ROUTER_IO_MODEL),
),
(_proxy_messages, "litellm.InternalServerError"),
],
ids=["proxy_responses_litellm_metadata", "proxy_messages_litellm_metadata"],
)
async def test_retried_failure_refunds_each_io_reservation_once(
self, call: RouterCall, mock_error: str | Exception
):
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
with pytest.raises(litellm.InternalServerError):
await call(router, mock_error)
assert await _io_counters_since(router, first_minute) == (0, 0)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call", "original_function", "mock_error"),
[
(_proxy_chat_sharing_logging, "acompletion", "litellm.InternalServerError"),
(
_proxy_responses_sharing_logging,
"aresponses",
litellm.InternalServerError(message="mock 500", llm_provider="anthropic", model=ROUTER_IO_MODEL),
),
(_proxy_messages_sharing_logging, "anthropic_messages", "litellm.InternalServerError"),
],
ids=["proxy_chat", "proxy_responses", "proxy_messages"],
)
async def test_retried_failure_sharing_one_logging_object_refunds_every_attempt(
self, call: ProxyCallSharingLogging, original_function: str, mock_error: str | Exception
):
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
with pytest.raises(litellm.InternalServerError):
await call(router, mock_error, _logging_obj_shared_across_attempts(original_function))
assert await _io_counters_since(router, first_minute) == (0, 0)
@pytest.mark.asyncio
async def test_model_group_info_aggregates_io_limits(self):
router = Router(