mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
Merge remote-tracking branch 'origin/main' into litellm_wildcard_discovery_provider_config
# Conflicts: # litellm/proxy/auth/model_checks.py
This commit is contained in:
commit
683c4673b1
98 changed files with 5917 additions and 669 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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 {}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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] = (
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
0
tests/unit/llms/openai/image_variations/__init__.py
Normal file
0
tests/unit/llms/openai/image_variations/__init__.py
Normal 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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}))
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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="",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue