From 82ca10cc2de77536900647c70f135148fbb03933 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 14:01:50 +0000 Subject: [PATCH 01/64] feat(bedrock): apply bedrock_request_metadata_fields on the /bedrock pass-through routes (#47040) * feat(bedrock): apply bedrock_request_metadata_fields on the /bedrock pass-through routes * fix(bedrock): move Invoke body requestMetadata into the signed header and match Bedrock's ASCII whitespace rule * fix(bedrock): keep Invoke body requestMetadata in the body and answer bad metadata as AWS does Invoke requests whose body already carries requestMetadata keep the resolved pairs in the body instead of the 8,500-character signed header. Malformed caller metadata now answers 400 with the ValidationException error type the AWS SDK expects, and drop_params drops bad pairs on both the pass-through routes and /v1/chat/completions. The caller-pair dedupe is linear, and the chat validator treats \s as ASCII whitespace, as Bedrock does * fix(bedrock): match Bedrock's 400s and header limit on pass-through request metadata Invoke metadata too long for the 8,500-character header now goes in the body, Invoke requestMetadata null answers 400 like Bedrock, Converse type errors carry SerializationException, violation messages name the failing key, and the reserved user_api_key_ prefix is matched case-insensitively after trimming whitespace * fix(bedrock): trim the Invoke metadata header to Bedrock's limit and drop bad requestMetadata under drop_params without the setting * fix(bedrock): leave /bedrock pass-through bodies as sent when request metadata forwarding is off Global drop_params no longer rewrites requestMetadata on the pass-through routes unless bedrock_request_metadata_fields is set. Chat keeps dropping bad pairs under drop_params. Adds a tests/integration wire test for the four pass-through actions and removes an unused parameter from the Invoke header helper. * test(bedrock): drive the pass-through metadata wire test with boto3's bearer key env var * fix(bedrock): trim the pass-through invoke metadata header and stop Nova invoke from checking metadata it drops * test(bedrock): cover pass-through metadata header trimming, malformed caller metadata, and chat drop_params on the wire * fix(bedrock): match requestMetadata keys and values against the whole string so a trailing newline cannot pass the 256-character limit --- .../bedrock/chat/converse_transformation.py | 87 +---- .../amazon_nova_transformation.py | 11 +- litellm/llms/bedrock/files/transformation.py | 2 +- .../bedrock/passthrough/transformation.py | 106 +++++- litellm/llms/bedrock/request_metadata.py | 129 +++++-- .../llm_passthrough_endpoints.py | 22 +- ...drock_passthrough_request_metadata_wire.py | 316 ++++++++++++++++++ .../test_amazon_nova_transformation.py | 19 ++ .../chat/test_converse_transformation.py | 14 +- .../test_bedrock_files_transformation.py | 17 + ...test_bedrock_passthrough_transformation.py | 291 ++++++++++++++++ .../llms/bedrock/test_request_metadata.py | 111 +++++- .../test_llm_pass_through_endpoints.py | 42 +++ 13 files changed, 1043 insertions(+), 124 deletions(-) create mode 100644 tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 66ee920afd0..215db3a4ee4 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -66,7 +66,9 @@ from litellm.llms.bedrock.common_utils import ( ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, - bedrock_request_metadata_is_owned, + bedrock_request_metadata_is_rewritten, + bedrock_request_metadata_violation, + caller_request_metadata_violations_are_dropped, merge_bedrock_invoke_headers, resolve_bedrock_request_metadata, ) @@ -354,77 +356,15 @@ class AmazonConverseConfig(BaseConfig): and v is not None } - def _validate_request_metadata(self, metadata: dict) -> None: - """ - Validate requestMetadata according to AWS Bedrock Converse API constraints. - - Constraints: - - Maximum of 16 items - - Keys: 1-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{1,256} - - Values: 0-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{0,256} - """ - import re - - if not isinstance(metadata, dict): + def _validate_request_metadata(self, metadata: object) -> None: + violation: Final = bedrock_request_metadata_violation(metadata) + if violation is not None: raise litellm.exceptions.BadRequestError( - message="requestMetadata must be a dictionary", + message=violation.message, model="bedrock", llm_provider="bedrock", ) - if len(metadata) > 16: - raise litellm.exceptions.BadRequestError( - message="requestMetadata can contain a maximum of 16 items", - model="bedrock", - llm_provider="bedrock", - ) - - key_pattern: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$") - value_pattern: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$") - - for key, value in metadata.items(): - if not isinstance(key, str): - raise litellm.exceptions.BadRequestError( - message="requestMetadata keys must be strings", - model="bedrock", - llm_provider="bedrock", - ) - - if not isinstance(value, str): - raise litellm.exceptions.BadRequestError( - message="requestMetadata values must be strings", - model="bedrock", - llm_provider="bedrock", - ) - - if len(key) == 0 or len(key) > 256: - raise litellm.exceptions.BadRequestError( - message="requestMetadata key length must be 1-256 characters", - model="bedrock", - llm_provider="bedrock", - ) - - if len(value) > 256: - raise litellm.exceptions.BadRequestError( - message="requestMetadata value length must be 0-256 characters", - model="bedrock", - llm_provider="bedrock", - ) - - if not key_pattern.match(key): - raise litellm.exceptions.BadRequestError( - message=f"requestMetadata key '{key}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]", - model="bedrock", - llm_provider="bedrock", - ) - - if not value_pattern.match(value): - raise litellm.exceptions.BadRequestError( - message=f"requestMetadata value '{value}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]", - model="bedrock", - llm_provider="bedrock", - ) - @staticmethod def _is_openai_gpt_reasoning_model(model: str) -> bool: return re.search(r"openai\.gpt-\d", model) is not None @@ -1244,7 +1184,7 @@ class AmazonConverseConfig(BaseConfig): elif param == "context_management" and isinstance(value, (dict, list)): self._map_context_management_param(value, optional_params) if param == "requestMetadata": - self._map_request_metadata_param(value, optional_params) + self._map_request_metadata_param(value, optional_params, drop_params) if param == "service_tier" and isinstance(value, str): self._map_service_tier_param(value, optional_params) @@ -1277,9 +1217,10 @@ class AmazonConverseConfig(BaseConfig): return optional_params - def _map_request_metadata_param(self, value: object, optional_params: dict) -> None: + def _map_request_metadata_param(self, value: object, optional_params: dict, drop_params: bool) -> None: if value is not None and isinstance(value, dict): - self._validate_request_metadata(value) + if not caller_request_metadata_violations_are_dropped(drop_params): + self._validate_request_metadata(value) optional_params["requestMetadata"] = value def _map_context_management_param(self, value: dict | list, optional_params: dict) -> None: @@ -1737,7 +1678,7 @@ class AmazonConverseConfig(BaseConfig): # Extract requestMetadata before processing other parameters request_metadata: Final = inference_params.pop("requestMetadata", None) - if request_metadata is not None: + if request_metadata is not None and not caller_request_metadata_violations_are_dropped(drop_params): self._validate_request_metadata(request_metadata) output_config: OutputConfigBlock | None = inference_params.pop("outputConfig", None) @@ -2164,7 +2105,7 @@ class AmazonConverseConfig(BaseConfig): request_metadata: Final = resolve_bedrock_request_metadata( litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata") ) - if bedrock_request_metadata_is_owned(): + if bedrock_request_metadata_is_rewritten(litellm_params.get("drop_params") is True): _data.pop("requestMetadata", None) if request_metadata is not None: _data["requestMetadata"] = request_metadata @@ -2244,7 +2185,7 @@ class AmazonConverseConfig(BaseConfig): request_metadata: Final = resolve_bedrock_request_metadata( litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata") ) - if bedrock_request_metadata_is_owned(): + if bedrock_request_metadata_is_rewritten(litellm_params.get("drop_params") is True): _data.pop("requestMetadata", None) if request_metadata is not None: _data["requestMetadata"] = request_metadata diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py index b94435345d3..ce28e579dcf 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py @@ -95,7 +95,12 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): model_info: Mapping[str, object] | None = None, ) -> dict: return AmazonConverseConfig.map_openai_params( - self, non_default_params, optional_params, model, drop_params, model_info + self, + {key: value for key, value in non_default_params.items() if key != "requestMetadata"}, + optional_params, + model, + drop_params, + model_info, ) def transform_request( @@ -110,7 +115,9 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): self, model=model, messages=messages, - optional_params=_without_tool_config_injection_points(optional_params), + optional_params=_without_tool_config_injection_points( + {key: value for key, value in optional_params.items() if key != "requestMetadata"} + ), litellm_params=litellm_params, headers=headers, ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 4d40e0dec85..61e084fe544 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -1021,7 +1021,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): converse_config: Final = AmazonConverseConfig() mapped_params = converse_config.map_openai_params( - non_default_params=optional_params, + non_default_params={k: v for k, v in optional_params.items() if k != "requestMetadata"}, optional_params={}, model=model, drop_params=False, diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index fb0ecd5b825..36cd87aff06 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -10,6 +10,7 @@ from httpx import Response from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.constants import AWS_ERROR_TYPE_HEADER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage @@ -17,6 +18,16 @@ from litellm.types.utils import ModelResponseStream from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo +from ..request_metadata import ( + BEDROCK_REQUEST_METADATA_HEADER, + BedrockRequestMetadataViolation, + bedrock_request_metadata_headers, + bedrock_request_metadata_is_owned, + bedrock_request_metadata_violation, + caller_request_metadata_violations_are_dropped, + merge_bedrock_invoke_headers, + resolve_bedrock_request_metadata, +) if TYPE_CHECKING: from botocore.eventstream import EventStreamMessage @@ -28,10 +39,87 @@ if TYPE_CHECKING: BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"}) +_BEDROCK_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) +_BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"}) + + +def _bedrock_endpoint_action(endpoint: str) -> str: + return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] def is_bedrock_streaming_endpoint(endpoint: str) -> bool: - return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS + return _bedrock_endpoint_action(endpoint) in BEDROCK_STREAMING_ACTIONS + + +def _caller_request_metadata_violation( + action: str, body: Mapping[str, object] +) -> BedrockRequestMetadataViolation | None: + if action in _BEDROCK_INVOKE_ACTIONS and "requestMetadata" in body and body["requestMetadata"] is None: + return BedrockRequestMetadataViolation("type", "requestMetadata must be a dictionary") + return bedrock_request_metadata_violation(body.get("requestMetadata")) + + +def _aws_error_type(action: str, violation: BedrockRequestMetadataViolation) -> str: + if action in _BEDROCK_CONVERSE_ACTIONS and violation.kind == "type": + return "SerializationException" + return "ValidationException" + + +def _caller_request_metadata(action: str, body: Mapping[str, object], litellm_params: Mapping[str, object]) -> object: + violation: Final = ( + None + if caller_request_metadata_violations_are_dropped(litellm_params.get("drop_params") is True) + else _caller_request_metadata_violation(action, body) + ) + if violation is not None: + raise BedrockError( + status_code=400, + message=violation.message, + headers={AWS_ERROR_TYPE_HEADER: _aws_error_type(action, violation)}, + ) + return body.get("requestMetadata") + + +def _with_body_request_metadata( + action: str, body: Mapping[str, object], litellm_params: Mapping[str, object] +) -> Mapping[str, object]: + resolved: Final = resolve_bedrock_request_metadata( + litellm_params, caller_metadata=_caller_request_metadata(action, body, litellm_params) + ) + without_caller: Final = {key: value for key, value in body.items() if key != "requestMetadata"} + return without_caller if resolved is None else {**without_caller, "requestMetadata": resolved} + + +def _without_owned_request_metadata_header(headers: Mapping[str, str]) -> Mapping[str, str]: + if not bedrock_request_metadata_is_owned(): + return headers + return {name: value for name, value in headers.items() if name.lower() != BEDROCK_REQUEST_METADATA_HEADER.lower()} + + +def _with_request_metadata( + action: str | None, + headers: Mapping[str, str], + body: Mapping[str, object], + litellm_params: Mapping[str, object], +) -> tuple[Mapping[str, str], Mapping[str, object]]: + if action in _BEDROCK_CONVERSE_ACTIONS: + return headers, _with_body_request_metadata(action, body, litellm_params) + if action in _BEDROCK_INVOKE_ACTIONS and "requestMetadata" in body: + return _without_owned_request_metadata_header(headers), _with_body_request_metadata( + action, body, litellm_params + ) + if action in _BEDROCK_INVOKE_ACTIONS: + return _with_invoke_request_metadata(headers, body, litellm_params) + return headers, body + + +def _with_invoke_request_metadata( + headers: Mapping[str, str], + body: Mapping[str, object], + litellm_params: Mapping[str, object], +) -> tuple[Mapping[str, str], Mapping[str, object]]: + owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) + return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body _TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"}) @@ -284,18 +372,24 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD def sign_request( self, - headers: dict, - litellm_params: dict, - request_data: dict | None, + headers: dict[str, str], + litellm_params: dict[str, object], + request_data: dict[str, object] | None, api_base: str, model: str | None = None, ) -> tuple[dict, bytes | None]: optional_params: Final = litellm_params.copy() + signed_headers, signed_body = _with_request_metadata( + action=_bedrock_endpoint_action(api_base) if bedrock_request_metadata_is_owned() else None, + headers=headers, + body=request_data or {}, + litellm_params=litellm_params, + ) return self._sign_request( service_name="bedrock", - headers=headers, + headers=dict(signed_headers), optional_params=optional_params, - request_data=request_data or {}, + request_data=dict(signed_body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py index 1f4e5886508..71a2dc78578 100644 --- a/litellm/llms/bedrock/request_metadata.py +++ b/litellm/llms/bedrock/request_metadata.py @@ -8,17 +8,22 @@ sets ``litellm.bedrock_request_metadata_fields`` (``litellm_settings`` on the pr Two properties are load-bearing for that billing record and are asserted by the tests: proxy identity is resolved first so it can never be evicted by caller-supplied pairs, and the whole ``user_api_key_`` prefix is reserved so a caller cannot write a proxy-authoritative -looking key. Values that break Bedrock's constraints are dropped rather than sanitised or -rejected, because an operator flipping this setting on must not turn a working request into a -400 and a silently rewritten attribution key is worse than an absent one. +looking key. Identity and ``spend_logs_metadata`` values that break Bedrock's constraints are +dropped rather than sanitised or rejected, because an operator flipping this setting on must not +turn a working request into a 400 and a silently rewritten attribution key is worse than an +absent one. For the same reason an Invoke header that would pass Bedrock's 8,500-character limit +keeps only the leading pairs that fit. A caller's own malformed ``requestMetadata`` is rejected with +a 400 unless ``drop_params`` is set, in which case its bad pairs are dropped the same way, with or +without the setting. """ from __future__ import annotations import json import re -from collections.abc import Mapping -from typing import Final +from collections.abc import Mapping, Sized +from dataclasses import dataclass +from typing import Final, Literal import litellm @@ -26,15 +31,27 @@ BEDROCK_REQUEST_METADATA_HEADER: Final = "X-Amzn-Bedrock-Request-Metadata" BEDROCK_REQUEST_METADATA_MAX_PAIRS: Final = 16 BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX: Final = "user_api_key_" BEDROCK_REQUEST_METADATA_CLIENT_FIELD: Final = "spend_logs_metadata" +BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS: Final = 8500 _METADATA_PARAM_NAMES: Final[tuple[str, ...]] = ("metadata", "litellm_metadata") -_KEY_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$") -_VALUE_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$") +_ALLOWED_CHARACTERS: Final = r"[a-zA-Z0-9\s:_@$#=/+,.-]" +_KEY_PATTERN: Final = re.compile(rf"{_ALLOWED_CHARACTERS}{{1,256}}", re.ASCII) +_VALUE_PATTERN: Final = re.compile(rf"{_ALLOWED_CHARACTERS}{{0,256}}", re.ASCII) _OWNED_HEADER_NAMES: Final[frozenset[str]] = frozenset((BEDROCK_REQUEST_METADATA_HEADER.lower(),)) +@dataclass(frozen=True, slots=True) +class BedrockRequestMetadataViolation: + kind: Literal["type", "constraint"] + message: str + + def _is_forwardable(key: str, value: str) -> bool: - return _KEY_PATTERN.match(key) is not None and _VALUE_PATTERN.match(value) is not None + return _KEY_PATTERN.fullmatch(key) is not None and _VALUE_PATTERN.fullmatch(value) is not None + + +def _is_reserved(key: str) -> bool: + return key.strip().lower().startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) def _text_pairs(source: object) -> tuple[tuple[str, str], ...]: @@ -99,13 +116,16 @@ def _client_pairs( candidates: Final = tuple( (key, value) for key, value in (*_text_pairs(caller_metadata), *spend_logs_pairs) - if not key.startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) and _is_forwardable(key, value) + if not _is_reserved(key) and _is_forwardable(key, value) ) - return tuple( - pair - for index, pair in enumerate(candidates) - if pair[0] not in tuple(earlier for earlier, _ in candidates[:index]) - )[:budget] + first_index: Final = {key: index for index, (key, _) in reversed(tuple(enumerate(candidates)))} + return tuple(pair for index, pair in enumerate(candidates) if first_index[pair[0]] == index)[:budget] + + +def _forwardable_caller_pairs(caller_metadata: object) -> tuple[tuple[str, str], ...]: + return tuple(pair for pair in _text_pairs(caller_metadata) if _is_forwardable(*pair))[ + :BEDROCK_REQUEST_METADATA_MAX_PAIRS + ] def resolve_bedrock_request_metadata( @@ -113,17 +133,18 @@ def resolve_bedrock_request_metadata( caller_metadata: object = None, ) -> dict[str, str] | None: """ - Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when the feature is - off or nothing survives Bedrock's constraints. The result is a plain dict because it is - written straight onto the Converse body, which Bedrock types as ``dict[str, str]``. + Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when nothing survives + Bedrock's constraints. The result is a plain dict because it is written straight onto the + Converse body, which Bedrock types as ``dict[str, str]``. - ``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. It has already - been validated (and rejected with a 400) by the Converse transformation, so it is only - filtered here for the reserved identity prefix and the remaining slot budget. + ``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. Callers reject a + malformed one with a 400 unless ``drop_params`` is set, so here its pairs are only filtered for + Bedrock's rules, the reserved identity prefix, and the remaining slot budget. With the feature + off only Bedrock's rules apply, since no identity is written for a caller to forge. """ allowed_fields: Final = _allowed_fields() if not allowed_fields: - return None + return dict(_forwardable_caller_pairs(caller_metadata)) or None sources: Final = _metadata_sources(litellm_params) identity: Final = _identity_pairs(sources, allowed_fields) client: Final = _client_pairs( @@ -136,6 +157,36 @@ def resolve_bedrock_request_metadata( return resolved or None +def _pair_violation(key: str, value: str) -> BedrockRequestMetadataViolation | None: + if _KEY_PATTERN.fullmatch(key) is None: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata key {key!r} must be 1-256 characters of {_ALLOWED_CHARACTERS}" + ) + if _VALUE_PATTERN.fullmatch(value) is None: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata value for key {key!r} must be 0-256 characters of {_ALLOWED_CHARACTERS}" + ) + return None + + +def bedrock_request_metadata_violation(caller_metadata: object) -> BedrockRequestMetadataViolation | None: + """Why a caller's ``requestMetadata`` breaks Bedrock's rules, in the order Bedrock checks them: a + wrong JSON type before the pair count and character rules. ``None`` means absent or valid.""" + text_pairs: Final = _text_pairs(caller_metadata) + entry_count: Final = len(caller_metadata) if isinstance(caller_metadata, Sized) else 0 + if caller_metadata is None: + return None + if not isinstance(caller_metadata, Mapping): + return BedrockRequestMetadataViolation("type", "requestMetadata must be a dictionary") + if len(text_pairs) != entry_count: + return BedrockRequestMetadataViolation("type", "requestMetadata keys and values must be strings") + if len(text_pairs) > BEDROCK_REQUEST_METADATA_MAX_PAIRS: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata can contain a maximum of {BEDROCK_REQUEST_METADATA_MAX_PAIRS} items" + ) + return next((violation for key, value in text_pairs if (violation := _pair_violation(key, value))), None) + + def bedrock_request_metadata_is_owned() -> bool: """ Whether the proxy OWNS the request-metadata field and header name for this request. @@ -149,25 +200,51 @@ def bedrock_request_metadata_is_owned() -> bool: return bool(_allowed_fields()) +def caller_request_metadata_violations_are_dropped(drop_params: bool) -> bool: + return drop_params or litellm.drop_params is True + + +def bedrock_request_metadata_is_rewritten(drop_params: bool) -> bool: + return bedrock_request_metadata_is_owned() or caller_request_metadata_violations_are_dropped(drop_params) + + +def _header_value(pairs: Mapping[str, str]) -> str: + return json.dumps(pairs, separators=(",", ":")) + + +def _fits_header(pairs: Mapping[str, str]) -> bool: + return len(_header_value(pairs)) <= BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + + +def _longest_fitting_prefix(resolved: Mapping[str, str]) -> dict[str, str]: + pairs: Final = tuple(resolved.items()) + return next( + (dict(pairs[:count]) for count in range(len(pairs), 0, -1) if _fits_header(dict(pairs[:count]))), + {}, + ) + + def bedrock_request_metadata_headers( litellm_params: Mapping[str, object] | None, ) -> tuple[frozenset[str], tuple[tuple[str, str], ...]]: """ - The signed ``X-Amzn-Bedrock-Request-Metadata`` header for the Invoke paths, which have no - body field for request metadata. + The signed ``X-Amzn-Bedrock-Request-Metadata`` header for an Invoke request whose body carries + no ``requestMetadata`` of its own. Returns the header names the proxy OWNS and, separately, the pairs to send. Ownership is reported whenever forwarding is enabled, including when nothing resolves, because a caller can suppress the resolver (omit the allow-listed fields, or send values that all fail Bedrock's rules) and an owned-but-empty result must still evict the caller's header rather - than fall back to it. + than fall back to it. Only the leading pairs that fit Bedrock's header limit are sent, so + identity always survives and long ``spend_logs_metadata`` pairs drop off the end. """ if not bedrock_request_metadata_is_owned(): return frozenset(), () resolved: Final = resolve_bedrock_request_metadata(litellm_params) - if resolved is None: + fitting: Final = {} if resolved is None else _longest_fitting_prefix(resolved) + if not fitting: return _OWNED_HEADER_NAMES, () - return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, json.dumps(resolved, separators=(",", ":"))),) + return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, _header_value(fitting)),) def merge_bedrock_invoke_headers( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 464cfee270a..3aef5724bdc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1087,6 +1087,16 @@ def _guardrail_refusal_status(error: Exception) -> int | None: return None +def _aws_error_type(error: Exception, handled: ProxyException) -> str | None: + """The AWS error type to answer the handled error with: ValidationException for a guardrail block, else the type + the error carried, as LiteLLM's own Bedrock request checks set it. None when the shared handler changed the + guardrail block's status.""" + refusal_status: Final = _guardrail_refusal_status(error) + if refusal_status is not None: + return "ValidationException" if str(refusal_status) == handled.code else None + return next((value for name, value in handled.headers.items() if name.lower() == AWS_ERROR_TYPE_HEADER), None) + + def _bedrock_error_response(error: Exception, handled: HTTPException) -> Response | None: """Bedrock's own error response, for an error Bedrock answered the passthrough request with an AWS error type.""" @@ -1111,9 +1121,9 @@ async def bedrock_passthrough_error_response( user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLoggingType, ) -> Response: - """Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block or an error - Bedrock returned the way Bedrock does, so an AWS SDK reads its code and message. Anything else is raised as the - shared handler raises it.""" + """Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block, an error + Bedrock returned, or an error carrying an AWS error type the way Bedrock does, so an AWS SDK reads its code and + message. Anything else is raised as the shared handler raises it.""" try: await processor.handle_llm_api_exception( e=error, @@ -1126,9 +1136,9 @@ async def bedrock_passthrough_error_response( raise return response except ProxyException as handled: - refusal_status: Final = _guardrail_refusal_status(error) - if refusal_status is not None and str(refusal_status) == handled.code: - raise AwsErrorTypeProxyException(handled, "ValidationException") from handled + aws_error_type: Final = _aws_error_type(error, handled) + if aws_error_type is not None: + raise AwsErrorTypeProxyException(handled, aws_error_type) from handled raise diff --git a/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py b/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py new file mode 100644 index 00000000000..5fd12405ab2 --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py @@ -0,0 +1,316 @@ +import base64 +import json +import uuid +from urllib.parse import unquote +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import boto3 +import pytest +import yaml +from botocore.config import Config +from botocore.exceptions import ClientError +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_METADATA_HEADER: Final = "x-amzn-bedrock-request-metadata" +_HEADER_MAX_CHARS: Final = 8500 +_LONG_SPEND_LOGS_METADATA: Final = tuple((f"{index:02d}".ljust(256, "k"), "\n" * 256) for index in range(12)) +_CONVERSE_REPLY: Final[Mapping[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}, + "metrics": {"latencyMs": 3}, +} +_INVOKE_REPLY: Final[Mapping[str, JsonValue]] = { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, +} +_CONVERSE_STREAM_FRAMES: Final = ( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": "end_turn"}, "sc", "u"), + _aws_event_frame("metadata", {"usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}}, "sc", "u"), +) + + +def _bedrock_peer(request: Request) -> Reply: + if request.target.endswith("/converse-stream"): + return Reply(chunks=_CONVERSE_STREAM_FRAMES, content_type=_EVENT_STREAM) + if request.target.endswith("/invoke-with-response-stream"): + chunk: Final = {"bytes": base64.b64encode(json.dumps(_INVOKE_REPLY).encode()).decode()} + return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM) + if request.target.endswith("/converse"): + return Reply(body=json.dumps(_CONVERSE_REPLY).encode()) + return Reply(body=json.dumps(_INVOKE_REPLY).encode()) + + +def _messages(marker: str) -> Sequence[JsonValue]: + return [{"role": "user", "content": [{"text": marker}]}] + + +def _invoke_body(marker: str) -> Mapping[str, JsonValue]: + return { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + } + + +def _signed_header_names(request: Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _header_value(pairs: tuple[tuple[str, str], ...]) -> str: + return json.dumps(dict(pairs), separators=(",", ":")) + + +def _send_with_caller_metadata( + runtime: object, action: str, model: str, marker: str, caller_metadata: Mapping[str, JsonValue] +) -> None: + if action == "invoke": + getattr(runtime, "invoke_model")( + modelId=model, body=json.dumps({**_invoke_body(marker), "requestMetadata": caller_metadata}) + ) + return + getattr(runtime, action.replace("-", "_"))( + modelId=model, messages=_messages(marker), requestMetadata=caller_metadata + ) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + wire: Wire + strict_model: str + dropping_model: str + chat_model: str + key: str + identity: Mapping[str, str] + long_key: str + long_identity: tuple[tuple[str, str], ...] + + def runtime(self, key: str, monkeypatch: pytest.MonkeyPatch): # noqa: ANN201 # boto3 clients are untyped + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key) + return boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config( + read_timeout=case_timeout_seconds(), + retries={"total_max_attempts": 1}, + parameter_validation=False, + ), + ) + + def received(self, marker: str) -> tuple[Request, ...]: + return tuple(request for request in self.wire.drain() if marker.encode() in request.body) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + strict_model: Final = f"bedrock-metadata-{suffix}" + dropping_model: Final = f"bedrock-metadata-dropping-{suffix}" + chat_model: Final = f"bedrock-metadata-chat-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-request-metadata") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with wire_server(_bedrock_peer) as wire: + deployment: Final = { + "model": f"bedrock/{_MODEL_ID}", + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + } + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + {"model_name": strict_model, "litellm_params": deployment}, + {"model_name": dropping_model, "litellm_params": {**deployment, "drop_params": True}}, + { + "model_name": chat_model, + "litellm_params": {**deployment, "model": f"bedrock/converse/{_MODEL_ID}", "drop_params": True}, + }, + ], + "litellm_settings": { + **base_config.get("litellm_settings", {}), + "drop_params": False, + "bedrock_request_metadata_fields": [ + "user_api_key_alias", + "user_api_key_team_alias", + "spend_logs_metadata", + ], + }, + } + config_path: Final = directory / "bedrock-request-metadata.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(session_gateway, directory, {}, config=config_path, workers=2) as proxy, + proxy.scenario() as scenario, + ): + team_alias: Final = f"platform-{suffix}" + team_id: Final = scenario.team(team_alias=team_alias) + key_alias: Final = f"prod-key-{suffix}" + long_key_alias: Final = f"long-key-{suffix}" + yield Rig( + proxy=proxy, + wire=wire, + strict_model=strict_model, + dropping_model=dropping_model, + chat_model=chat_model, + key=scenario.key( + key_alias=key_alias, + team_id=team_id, + metadata={"spend_logs_metadata": {"agent": "triage-bot"}}, + ), + identity={ + "user_api_key_alias": key_alias, + "user_api_key_team_alias": team_alias, + "agent": "triage-bot", + }, + long_key=scenario.key( + key_alias=long_key_alias, + team_id=team_id, + metadata={"spend_logs_metadata": dict(_LONG_SPEND_LOGS_METADATA)}, + ), + long_identity=( + ("user_api_key_alias", long_key_alias), + ("user_api_key_team_alias", team_alias), + *_LONG_SPEND_LOGS_METADATA, + ), + ) + + +def test_boto3_passthrough_calls_carry_the_key_identity_in_signed_request_metadata( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = f"who-is-calling-{uuid.uuid4().hex}" + forged: Final = {"user_api_key_team_alias": "forged", "ticket": "T-42", "bad!": "x"} + runtime: Final = rig.runtime(rig.key, monkeypatch) + model: Final = rig.dropping_model + + converse: Final = runtime.converse(modelId=model, messages=_messages(marker), requestMetadata=forged) + converse_stream: Final = [ + next(iter(event)) + for event in runtime.converse_stream(modelId=model, messages=_messages(marker), requestMetadata=forged)[ + "stream" + ] + ] + invoke: Final = json.loads( + runtime.invoke_model(modelId=model, body=json.dumps(_invoke_body(marker)))["body"].read() + ) + invoke_stream: Final = [ + json.loads(event["chunk"]["bytes"]) + for event in runtime.invoke_model_with_response_stream(modelId=model, body=json.dumps(_invoke_body(marker)))[ + "body" + ] + ] + assert converse["output"] == _CONVERSE_REPLY["output"], converse + assert converse_stream == ["messageStart", "messageStop", "metadata"], converse_stream + assert invoke == _INVOKE_REPLY, invoke + assert invoke_stream == [_INVOKE_REPLY], invoke_stream + + received: Final = rig.received(marker) + assert [request.target for request in received] == [ + f"/model/{_MODEL_ID}/converse", + f"/model/{_MODEL_ID}/converse-stream", + f"/model/{_MODEL_ID}/invoke", + f"/model/{_MODEL_ID}/invoke-with-response-stream", + ], received + converse_requests: Final = received[:2] + invoke_requests: Final = received[2:] + assert [json.loads(request.body) for request in converse_requests] == [ + {"messages": _messages(marker), "requestMetadata": {**rig.identity, "ticket": "T-42"}}, + {"messages": _messages(marker), "requestMetadata": {**rig.identity, "ticket": "T-42"}}, + ], [request.body for request in converse_requests] + assert [_METADATA_HEADER in request.headers for request in converse_requests] == [False, False] + assert [json.loads(request.body) for request in invoke_requests] == [_invoke_body(marker), _invoke_body(marker)] + assert [json.loads(request.headers[_METADATA_HEADER]) for request in invoke_requests] == [ + rig.identity, + rig.identity, + ], [request.headers for request in invoke_requests] + assert all(_METADATA_HEADER in _signed_header_names(request) for request in invoke_requests), [ + request.headers["authorization"] for request in invoke_requests + ] + + +def test_invoke_metadata_too_long_for_the_header_signs_the_leading_pairs_that_fit( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = f"long-key-{uuid.uuid4().hex}" + runtime: Final = rig.runtime(rig.long_key, monkeypatch) + + invoke: Final = json.loads( + runtime.invoke_model(modelId=rig.strict_model, body=json.dumps(_invoke_body(marker)))["body"].read() + ) + + assert invoke == _INVOKE_REPLY, invoke + (request,) = rig.received(marker) + assert json.loads(request.body) == _invoke_body(marker) + header: Final = request.headers[_METADATA_HEADER] + sent_pairs: Final = len(json.loads(header)) + assert tuple(json.loads(header).items()) == rig.long_identity[:sent_pairs], header + assert 2 <= sent_pairs < len(rig.long_identity), header + assert len(header) <= _HEADER_MAX_CHARS < len(_header_value(rig.long_identity[: sent_pairs + 1])) + assert _METADATA_HEADER in _signed_header_names(request), request.headers["authorization"] + + +@pytest.mark.parametrize( + ("action", "caller_metadata", "error_code"), + [ + pytest.param("converse", {"ticket": 42}, "SerializationException", id="converse-number-value"), + pytest.param("converse-stream", {"bad!": "x"}, "ValidationException", id="converse-stream-bad-key"), + pytest.param("invoke", {"bad!": "x"}, "ValidationException", id="invoke-bad-key"), + ], +) +def test_malformed_caller_request_metadata_reaches_boto3_as_the_bedrock_error_without_an_upstream_call( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str, caller_metadata: Mapping[str, JsonValue], error_code: str +) -> None: + marker: Final = f"malformed-{uuid.uuid4().hex}" + runtime: Final = rig.runtime(rig.key, monkeypatch) + + with pytest.raises(ClientError) as raised: + _send_with_caller_metadata(runtime, action, rig.strict_model, marker, caller_metadata) + + assert raised.value.response["Error"]["Code"] == error_code, raised.value.response + assert raised.value.response["ResponseMetadata"]["HTTPStatusCode"] == 400, raised.value.response + assert rig.received(marker) == () + + +def test_chat_on_a_dropping_converse_deployment_sends_only_the_valid_caller_pairs(rig: Rig) -> None: + marker: Final = f"chat-{uuid.uuid4().hex}" + + response: Final = rig.proxy.client.post( + "/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": marker}], + "requestMetadata": {"ticket": "T-42", "bad!": "x"}, + }, + headers={"Authorization": f"Bearer {rig.key}"}, + timeout=case_timeout_seconds(), + ) + + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "ok", response.text + (request,) = rig.received(marker) + assert unquote(request.target) == f"/model/{_MODEL_ID}/converse", request.target + assert json.loads(request.body)["requestMetadata"] == {**rig.identity, "ticket": "T-42"}, request.body diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py index f0f0f9160fb..de4b135e6d0 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py @@ -115,3 +115,22 @@ def test_tool_config_injection_point_is_neither_placed_nor_credited(local_model_ assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["f"] assert "cachePoint" not in json.dumps(request) assert GATEWAY_INJECTED_CACHE_METADATA_KEY not in metadata + + +@pytest.mark.parametrize("value", ["b c", "bad!"], ids=["nbsp", "outside-bedrock-charset"]) +def test_caller_request_metadata_is_ignored_because_invoke_model_has_no_field_for_it(local_model_cost_map, value): + caller_metadata = {"ticket": value} + optional_params = AmazonInvokeNovaConfig().map_openai_params( + non_default_params={"max_tokens": 8, "requestMetadata": caller_metadata}, + optional_params={}, + model=MODEL, + drop_params=False, + ) + + request = _transform_request( + messages=[{"role": "user", "content": "hi"}], + optional_params={**optional_params, "requestMetadata": caller_metadata}, + ) + + assert "requestMetadata" not in request + assert request["inferenceConfig"] == {"maxTokens": 8} diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 0ed713f056e..6092934fc7a 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -4199,7 +4199,7 @@ def test_request_metadata_key_constraints(): # Test empty key invalid_metadata = {"": "value"} - with pytest.raises(Exception, match=r"(?i)key length|empty"): + with pytest.raises(Exception, match=r"key '' must be 1-256 characters"): config._transform_request( model="anthropic.claude-haiku-4-5-20251001-v1:0", messages=messages, @@ -4241,6 +4241,18 @@ def test_request_metadata_value_constraints(): ) +@pytest.mark.parametrize("metadata", [{"ticket": "a\u00a0b"}, {"a\u2003b": "ticket"}]) +def test_request_metadata_rejects_whitespace_outside_ascii(metadata): + with pytest.raises(litellm.exceptions.BadRequestError, match=r"must be [01]-256 characters of"): + AmazonConverseConfig()._transform_request( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hello!"}], + optional_params={"requestMetadata": metadata}, + litellm_params={}, + headers={}, + ) + + def test_request_metadata_character_pattern(): """Test character pattern validation for keys and values.""" config = AmazonConverseConfig() diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 425beefac8a..b050bb7625b 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1760,6 +1760,23 @@ class TestBedrockBatchNonChatEndpointRecords: assert "metadata" not in model_input assert "requestMetadata" not in model_input + def test_nova_record_request_metadata_stays_out_of_the_bedrock_body(self): + model_input = self._transform( + { + "custom_id": "4d", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": self.NOVA_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "requestMetadata": {"ticket": "b\u00a0c"}, + }, + } + ) + + assert "requestMetadata" not in model_input + assert model_input["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + @pytest.mark.parametrize( "body", [ diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 84d3160e609..da5c4699925 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -17,6 +17,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.passthrough.transformation import ( BedrockPassthroughConfig, is_bedrock_streaming_endpoint, @@ -898,3 +899,293 @@ def test_stream_collector_logs_nothing_for_an_unrecognized_endpoint(): collector.add(_event_frame("messageStart", {"role": "assistant"})) assert collector.build_logged_response(_converse_stream_logging_obj()) is None + + +REQUEST_METADATA_HEADER = "x-amzn-bedrock-request-metadata" +REQUEST_METADATA_FIELDS = ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"] +PROXY_IDENTITY = { + "user_api_key_alias": "prod-key", + "user_api_key_team_alias": "platform", + "spend_logs_metadata": {"agent": "triage-bot"}, +} +EXPECTED_REQUEST_METADATA = { + "user_api_key_alias": "prod-key", + "user_api_key_team_alias": "platform", + "agent": "triage-bot", +} +CONVERSE_BODY = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]} +INVOKE_BODY = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], +} + + +async def _sent_passthrough_request( + action: str, body: dict[str, object], identity: dict[str, object], **kwargs: object +) -> httpx.Request: + requests: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"ok": True}, request=request) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "bedrock-pt", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIAROUTERKEY", + "aws_secret_access_key": "router-secret", + "aws_region_name": "us-west-2", + }, + } + ] + ) + await router.allm_passthrough_route( + model="bedrock-pt", + method="POST", + endpoint=f"/model/bedrock-pt/{action}", + data=body, + client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), + litellm_metadata=identity, + **kwargs, + ) + return requests[0] + + +def _signed_header_names(request: httpx.Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "query_params"), + [("converse", None), ("converse-stream", None), ("converse/", None), ("converse", {"trace": "1"})], +) +async def test_converse_passthrough_signs_proxy_identity_into_request_metadata( + action: str, query_params: dict[str, str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**CONVERSE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged", "ticket": "T-42"}} + + request: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, request_query_params=query_params + ) + + assert json.loads(request.content) == { + **CONVERSE_BODY, + "requestMetadata": {**EXPECTED_REQUEST_METADATA, "ticket": "T-42"}, + } + assert caller_body["requestMetadata"] == {"user_api_key_team_alias": "forged", "ticket": "T-42"} + assert REQUEST_METADATA_HEADER not in request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_sends_proxy_identity_as_a_signed_header( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + PROXY_IDENTITY, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER}": '{"user_api_key_team_alias":"forged"}'}, + ) + + assert json.loads(request.headers[REQUEST_METADATA_HEADER]) == EXPECTED_REQUEST_METADATA + assert REQUEST_METADATA_HEADER in _signed_header_names(request) + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_keeps_caller_body_request_metadata_in_the_body( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**INVOKE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged", "ticket": "T-42"}} + + request: Final = await _sent_passthrough_request( + action, + caller_body, + PROXY_IDENTITY, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER.upper()}": '{"user_api_key_team_alias":"forged"}'}, + ) + + assert json.loads(request.content) == { + **INVOKE_BODY, + "requestMetadata": {**EXPECTED_REQUEST_METADATA, "ticket": "T-42"}, + } + assert REQUEST_METADATA_HEADER not in request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", [None, []]) +async def test_passthrough_requests_are_untouched_when_forwarding_is_off( + setting: list[str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", setting) + caller_body: Final = {**CONVERSE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged"}} + + converse: Final = await _sent_passthrough_request("converse", caller_body, PROXY_IDENTITY) + invoke: Final = await _sent_passthrough_request("invoke", INVOKE_BODY, PROXY_IDENTITY) + + assert json.loads(converse.content) == caller_body + assert REQUEST_METADATA_HEADER not in invoke.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("action", "body"), [("converse", CONVERSE_BODY), ("invoke", INVOKE_BODY)]) +async def test_passthrough_drops_reserved_caller_keys_when_nothing_resolves( + action: str, body: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**body, "requestMetadata": {"user_api_key_team_alias": "forged"}} + + request: Final = await _sent_passthrough_request(action, caller_body, {}) + + assert json.loads(request.content) == body + assert REQUEST_METADATA_HEADER not in request.headers + + +MALFORMED_REQUEST_METADATA: Final = ( + pytest.param( + {"user_api_key_team_alias": "forged", "bad!": "x"}, "key 'bad!' must be", "ValidationException", id="bad-key" + ), + pytest.param({"ticket": "a\u00a0b"}, "value for key 'ticket' must be", "ValidationException", id="bad-value"), + pytest.param("ticket=T-42", "must be a dictionary", "SerializationException", id="not-a-dictionary"), + pytest.param({f"key_{index}": "v" for index in range(17)}, "maximum of 16 items", "ValidationException", id="17"), + pytest.param({"ticket": 42}, "must be strings", "SerializationException", id="non-string-value"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["converse", "converse-stream"]) +@pytest.mark.parametrize(("caller_metadata", "violation", "converse_error_type"), MALFORMED_REQUEST_METADATA) +async def test_malformed_converse_request_metadata_is_rejected_with_the_error_type_bedrock_returns( + action: str, caller_metadata: object, violation: str, converse_error_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match=violation) as rejected: + await _sent_passthrough_request(action, {**CONVERSE_BODY, "requestMetadata": caller_metadata}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == converse_error_type + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +@pytest.mark.parametrize(("caller_metadata", "violation", "converse_error_type"), MALFORMED_REQUEST_METADATA) +async def test_malformed_invoke_request_metadata_is_rejected_as_a_bedrock_validation_error( + action: str, caller_metadata: object, violation: str, converse_error_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match=violation) as rejected: + await _sent_passthrough_request(action, {**INVOKE_BODY, "requestMetadata": caller_metadata}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == "ValidationException" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_null_request_metadata_is_rejected_like_bedrock_does( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match="must be a dictionary") as rejected: + await _sent_passthrough_request(action, {**INVOKE_BODY, "requestMetadata": None}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == "ValidationException" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["converse", "converse-stream"]) +async def test_converse_null_request_metadata_is_treated_as_absent_like_bedrock_does( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + request: Final = await _sent_passthrough_request(action, {**CONVERSE_BODY, "requestMetadata": None}, PROXY_IDENTITY) + + assert json.loads(request.content) == {**CONVERSE_BODY, "requestMetadata": EXPECTED_REQUEST_METADATA} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_metadata_too_long_for_the_header_sends_the_leading_pairs_that_fit( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + long_pairs: Final = {f"{index:02d}".ljust(256, "k"): "\n" * 256 for index in range(12)} + resolved: Final = tuple( + {"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform", **long_pairs}.items() + ) + + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + {**PROXY_IDENTITY, "spend_logs_metadata": long_pairs}, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER}": '{"user_api_key_team_alias":"forged"}'}, + ) + + header: Final = request.headers[REQUEST_METADATA_HEADER] + sent: Final = json.loads(header) + one_more_pair: Final = json.dumps(dict(resolved[: len(sent) + 1]), separators=(",", ":")) + assert sent == dict(resolved[: len(sent)]) + assert len(header) <= 8500 < len(one_more_pair) + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("action", "body"), [("converse", CONVERSE_BODY), ("invoke", INVOKE_BODY)]) +@pytest.mark.parametrize("drop_globally", [True, False]) +async def test_malformed_caller_request_metadata_is_dropped_when_params_are_dropped( + action: str, body: dict[str, object], drop_globally: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + caller_body: Final = {**body, "requestMetadata": {"user_api_key_team_alias": "forged", "bad!": "x"}} + + dropped: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, **({} if drop_globally else {"drop_params": True}) + ) + + assert json.loads(dropped.content) == {**body, "requestMetadata": EXPECTED_REQUEST_METADATA} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "body", "caller_metadata"), + [ + ("converse", CONVERSE_BODY, {"user_api_key_team_alias": "caller-team", "bad!": "x"}), + ("invoke", INVOKE_BODY, {"user_api_key_team_alias": "caller-team", "bad!": "x"}), + ("invoke", INVOKE_BODY, None), + ], +) +@pytest.mark.parametrize("drop_globally", [True, False]) +async def test_passthrough_body_is_sent_as_is_when_params_are_dropped_and_forwarding_is_off( + action: str, + body: dict[str, object], + caller_metadata: dict[str, str] | None, + drop_globally: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + caller_body: Final = {**body, "requestMetadata": caller_metadata} + + sent: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, **({} if drop_globally else {"drop_params": True}) + ) + + assert json.loads(sent.content) == caller_body + assert REQUEST_METADATA_HEADER not in sent.headers diff --git a/tests/unit/llms/bedrock/test_request_metadata.py b/tests/unit/llms/bedrock/test_request_metadata.py index 5a14bbea9f2..6b74ad90cff 100644 --- a/tests/unit/llms/bedrock/test_request_metadata.py +++ b/tests/unit/llms/bedrock/test_request_metadata.py @@ -1,8 +1,10 @@ import asyncio import json +from collections.abc import Callable, Mapping +from typing import Final import pytest - +from pydantic import TypeAdapter import litellm from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -18,6 +20,7 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran ) from litellm.llms.bedrock.request_metadata import ( BEDROCK_REQUEST_METADATA_HEADER, + BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS, BEDROCK_REQUEST_METADATA_MAX_PAIRS, resolve_bedrock_request_metadata, ) @@ -128,6 +131,20 @@ def test_identity_survives_a_caller_filling_every_slot( ) +def test_a_key_sent_twice_keeps_the_caller_value_and_uses_one_slot(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) + caller: Final = { + f"client_{index:02d}": "caller" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY) - 1) + } + spend_logs: Final = {**{key: "spend" for key in caller}, "extra": "spend"} + + resolved: Final = resolve_bedrock_request_metadata( + litellm_params("metadata", spend_logs_metadata=spend_logs, **IDENTITY), caller_metadata=caller + ) + + assert resolved == {**IDENTITY, **caller, "extra": "spend"} + + @pytest.mark.parametrize( "field_order", [ @@ -158,7 +175,14 @@ def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot( @pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"]) @pytest.mark.parametrize( "forged_key", - ["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"], + [ + "user_api_key_team_alias", + "user_api_key_org_alias", + "user_api_key_hash", + "USER_API_KEY_TEAM_ALIAS", + "User_Api_Key_Hash", + " user_api_key_user_id", + ], ) def test_caller_cannot_forge_or_shadow_a_reserved_identity_key( forged_key, @@ -181,14 +205,21 @@ def test_caller_cannot_forge_or_shadow_a_reserved_identity_key( assert "attacker-controlled" not in resolved.values() -def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize( + "team_alias", + ["O'Brien's team", "platform\u00a0team", "platform\u3000team", "x" * 256 + "\n"], + ids=["apostrophe", "no-break-space", "ideographic-space", "257-chars-ending-in-newline"], +) +def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds( + team_alias: str, monkeypatch: pytest.MonkeyPatch +) -> None: """A team alias with an apostrophe must not turn a working request into a 400 the moment - an operator flips the setting on.""" + an operator flips the setting on. Bedrock also rejects non-ASCII whitespace (checked live 2026-10-16).""" monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) params = litellm_params( "metadata", user_api_key_alias="prod-key", - user_api_key_team_alias="O'Brien's team", + user_api_key_team_alias=team_alias, user_api_key_user_email="x" * 300, ) @@ -198,13 +229,54 @@ def test_identity_violating_the_character_class_is_dropped_and_the_request_succe assert body["messages"] -def test_caller_supplied_violation_still_raises_bad_request(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize( + "caller_metadata", + [{"team": "O'Brien's team"}, {"k" * 256 + "\n": "v"}, {"team": "v" * 256 + "\n"}], + ids=["bad-character", "257-char-key-ending-in-newline", "257-char-value-ending-in-newline"], +) +def test_caller_supplied_violation_still_raises_bad_request( + caller_metadata: Mapping[str, str], monkeypatch: pytest.MonkeyPatch +) -> None: monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) with pytest.raises(litellm.exceptions.BadRequestError): - converse_body( - litellm_params("metadata", **IDENTITY), - {"requestMetadata": {"team": "O'Brien's team"}}, + converse_body(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_metadata}) + + +@pytest.mark.parametrize("driver", CONVERSE_DRIVERS) +@pytest.mark.parametrize("drop_globally", [True, False]) +@pytest.mark.parametrize(("fields", "forwarded_identity"), [(ALL_FIELDS, IDENTITY), (None, {})]) +def test_caller_violations_are_dropped_like_the_passthrough_routes_when_params_are_dropped( + driver: Callable[[Mapping[str, object], Mapping[str, object]], Mapping[str, object]], + drop_globally: bool, + fields: list[str] | None, + forwarded_identity: Mapping[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", fields) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + params: Final = {**litellm_params("metadata", **IDENTITY), **({} if drop_globally else {"drop_params": True})} + + body: Final = driver(params, {"requestMetadata": {"team": "O'Brien's team", "ticket": "T-42"}}) + + assert body["requestMetadata"] == {**forwarded_identity, "ticket": "T-42"} + + +@pytest.mark.parametrize("fields", [ALL_FIELDS, None]) +def test_mapped_caller_violations_are_dropped_when_params_are_dropped( + fields: list[str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", fields) + caller: Final = {"team": "O'Brien's team", "ticket": "T-42"} + + mapped: Final = AmazonConverseConfig().map_openai_params( + non_default_params={"requestMetadata": caller}, optional_params={}, model=MODEL, drop_params=True + ) + + assert mapped["requestMetadata"] == caller + with pytest.raises(litellm.exceptions.BadRequestError): + AmazonConverseConfig().map_openai_params( + non_default_params={"requestMetadata": caller}, optional_params={}, model=MODEL, drop_params=False ) @@ -320,6 +392,27 @@ def metadata_header_values(headers): return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()] +@pytest.mark.parametrize("driver", HEADER_DRIVERS) +def test_the_metadata_header_keeps_the_leading_pairs_that_fit_bedrocks_header_limit( + driver: Callable[[Mapping[str, str], Mapping[str, object]], Mapping[str, str]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) + long_pairs: Final = {f"{index:02d}".ljust(256, "k"): "\n" * 256 for index in range(12)} + params: Final = litellm_params("metadata", spend_logs_metadata=long_pairs, **IDENTITY) + resolved: Final = tuple((resolve_bedrock_request_metadata(params) or {}).items()) + + headers: Final = driver({}, params) + (header,) = (value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()) + sent: Final = TypeAdapter(dict[str, str]).validate_json(header) + one_more_pair: Final = json.dumps(dict(resolved[: len(sent) + 1]), separators=(",", ":")) + + assert len(header) <= BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + assert tuple(sent.items()) == resolved[: len(sent)] + assert IDENTITY.items() <= sent.items() + assert len(one_more_pair) > BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + + def test_converse_still_sets_the_bearer_authorization_header(monkeypatch: pytest.MonkeyPatch): """Converse owns the metadata header now, and that must not disturb the api_key path its validate_environment existed for. Closing the forgery hole cannot break authentication.""" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 21170abaeb1..ed89cc37544 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -33,6 +33,7 @@ from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException from litellm.proxy.utils import ProxyLogging from litellm.passthrough.main import raise_for_provider_status +from litellm.llms.bedrock.common_utils import BedrockError from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -2137,6 +2138,47 @@ class TestBedrockPassthroughAwsErrors: assert exc_info.value.aws_error_type == "ValidationException" assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) + + @pytest.mark.asyncio + async def test_bedrock_request_check_is_raised_with_the_aws_error_type_it_carries( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = BedrockError( + status_code=400, + message="requestMetadata can contain a maximum of 16 items", + headers={"x-amzn-errortype": "ValidationException"}, + ) + + with pytest.raises(AwsErrorTypeProxyException) as exc_info: + await self._answer(error) + + assert exc_info.value.aws_error_type == "ValidationException" + assert "requestMetadata can contain a maximum of 16 items" in exc_info.value.error.message + assert exc_info.value.error.code == "400" + + @pytest.mark.asyncio + async def test_bedrock_error_without_an_aws_error_type_keeps_the_proxy_shape( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(BedrockError(status_code=400, message="bad request")) + + assert exc_info.value.code == "400" + + @pytest.mark.asyncio + async def test_error_with_a_non_ascii_header_value_keeps_the_proxy_shape( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = HTTPException(status_code=403, detail="model not allowed", headers={"x-model-group": "modèle"}) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(error) + + assert (exc_info.value.code, exc_info.value.headers.get("x-model-group")) == ("403", "modèle") @pytest.mark.asyncio @pytest.mark.parametrize( ("error", "message", "code"), From c143954493baa810e2704917b9cdeb8fb3523718 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 15:41:24 +0000 Subject: [PATCH 02/64] fix(streaming): report the deployment's api_base when a stream fails mid-way (#47045) * fix(streaming): report the deployment's api_base when a stream fails mid-way * fix(responses): report the deployment's api_base when a native responses stream fails mid-way * fix(streaming): name the deployment's configured api_base in mid-stream errors and keep prompt text out Mid-stream errors on chat and native /v1/responses streams now prefer the router's configured api_base over the logged request URL, websocket turns keep the connection's api_base, and the mid-stream mapping no longer adds the prompt to the debug info. * fix(streaming): keep api_base credentials out of error debug info and leave websocket turn logging unchanged * fix(streaming): name the websocket connection's api_base on turn failures and strip scheme-less userinfo * fix(streaming): strip api_base userinfo before cutting the query so a raw ? or # in a password stays out Adds the audit's integration tests for the mid-stream alert API Base line and the SDK debug info. * fix(streaming): leave the API Base line out when an @ follows the host A raw / ? or # in a password and an @ in a query or path read the same in a malformed api_base, so either guess could show a secret or a wrong host. * fix(streaming): keep the API Base line when an @ sits in a provider path Vertex partner model ids like claude-opus-5-5@default put an @ in the URL path, so the fail-closed rule left out the line main printed. The line is now left out only when the text before the first / ? or # is not a host, and a :// counts as the scheme only after a valid scheme name. * test(streaming): cover a numeric password with a slash in the API Base line The audit found that a password like 1234/function-secret, which looks like a port, had no case. The line is left out on this branch and printed the raw secret on main --- .../exception_mapping_utils.py | 34 +- .../litellm_core_utils/streaming_handler.py | 24 +- litellm/responses/streaming_iterator.py | 25 +- .../test_stream_error_alert_api_base.py | 482 ++++++++++++++++++ .../sdk/test_stream_error_api_base_sdk.py | 69 +++ .../test_exception_mapping_utils.py | 98 ++++ .../test_streaming_handler.py | 139 ++++- .../test_streaming_iterator_error_events.py | 141 +++++ 8 files changed, 991 insertions(+), 21 deletions(-) create mode 100644 tests/integration/observability/test_stream_error_alert_api_base.py create mode 100644 tests/integration/sdk/test_stream_error_api_base_sdk.py diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index f8aadc4258d..13a53c55337 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -2413,6 +2413,35 @@ def _map_exception_by_status( ) +_NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) +_URL_SCHEME: Final = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*") +_HOST_WITH_OPTIONAL_PORT: Final = re.compile( + r"(?:(?:[A-Za-z0-9-]+\.)+[A-Za-z0-9-]+|localhost|\[[0-9A-Fa-f:.]+\])(?::\d+)?" +) + + +def with_deployment_api_base(litellm_params: object) -> Mapping[str, object]: + params: Final = as_str_mapping(litellm_params) or _NO_LITELLM_PARAMS + metadata: Final = as_str_mapping(params.get("metadata")) + deployment_api_base: Final = None if metadata is None else metadata.get("api_base") + if not isinstance(deployment_api_base, str) or not deployment_api_base: + return params + return MappingProxyType({**params, "api_base": deployment_api_base}) + + +def _api_base_without_credentials(api_base: str) -> str: + scheme, separator, _ = api_base.partition("://") + scheme_prefix: Final = scheme + separator if separator and _URL_SCHEME.fullmatch(scheme) else "" + authority_and_rest: Final = api_base[len(scheme_prefix) :] + authority: Final = re.split(r"[/?#]", authority_and_rest, maxsplit=1)[0] + after_authority: Final = authority_and_rest[len(authority) :] + host: Final = authority.rpartition("@")[2] + if "@" in after_authority and _HOST_WITH_OPTIONAL_PORT.fullmatch(host) is None: + return "" + path: Final = re.split(r"[?#]", after_authority, maxsplit=1)[0] + return f"{scheme_prefix}{host}{path}" + + def _is_guardrail_block(original_exception: Exception) -> bool: from litellm.integrations.custom_guardrail import is_guardrail_intervention @@ -2483,8 +2512,9 @@ def exception_type( if isinstance(custom_llm_provider, str) and len(custom_llm_provider) > 0: exception_provider = custom_llm_provider[0].upper() + custom_llm_provider[1:] + "Exception" - if _api_base: - extra_information += f"\nAPI Base: `{_api_base}`" + shown_api_base: Final = _api_base_without_credentials(_api_base) if _api_base else "" + if shown_api_base: + extra_information += f"\nAPI Base: `{shown_api_base}`" if messages and len(messages) > 0 and litellm.redact_messages_in_exceptions is False: extra_information += f"\nMessages: `{messages}`" diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 78fb82c98e7..78cfb58bedf 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -47,7 +47,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk from ..exceptions import OpenAIError from .core_helpers import map_finish_reason, process_response_headers -from .exception_mapping_utils import exception_type +from .exception_mapping_utils import exception_type, with_deployment_api_base from .llm_response_utils.get_api_base import get_api_base from .rules import Rules @@ -1579,11 +1579,7 @@ class CustomStreamWrapper: except Exception as e: traceback.format_exc() setattr(e, "message", str(e)) - raise exception_type( - model=self.model, - custom_llm_provider=self.custom_llm_provider, - original_exception=e, - ) + raise self._map_stream_exception(e) def set_logging_event_loop(self, loop): """ @@ -2370,6 +2366,14 @@ class CustomStreamWrapper: if not self._emitted_disqualifying_content and self._delta_disqualifies_continuation(delta): self._emitted_disqualifying_content = True + def _map_stream_exception(self, original_exception: Exception) -> Exception: + return exception_type( + model=self.model, + custom_llm_provider=self.custom_llm_provider, + original_exception=original_exception, + extra_kwargs=with_deployment_api_base(self.logging_obj.model_call_details.get("litellm_params")), + ) + def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn": """ Common error handling for both __next__ and __anext__. @@ -2414,13 +2418,7 @@ class CustomStreamWrapper: mapped_exception: Exception = e else: try: - mapped_exception = exception_type( - model=self.model, - custom_llm_provider=self.custom_llm_provider, - original_exception=e, - completion_kwargs={}, - extra_kwargs={}, - ) + mapped_exception = self._map_stream_exception(e) except Exception as mapping_error: mapped_exception = mapping_error self._restore_consumer_correlation_context() diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 89c32aed1fb..5897b0ac954 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -28,6 +28,7 @@ from litellm.constants import ( from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base @@ -264,7 +265,11 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: def _map_stream_error_to_exception( - error_obj: object, model: str, custom_llm_provider: str, headers: httpx.Headers | None = None + error_obj: object, + model: str, + custom_llm_provider: str, + headers: httpx.Headers | None = None, + logging_obj: LiteLLMLoggingObj | None = None, ) -> Exception: from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -277,13 +282,18 @@ def _map_stream_error_to_exception( body=error_body, headers=headers, ) + litellm_params: Final = ( + None + if logging_obj is None + else _typed_gets_litellm_params(logging_obj.model_call_details.get)("litellm_params", {}) + ) try: return litellm.exception_type( model=model, custom_llm_provider=custom_llm_provider, original_exception=provider_exception, completion_kwargs={}, - extra_kwargs={}, + extra_kwargs=with_deployment_api_base(litellm_params), ) except Exception as mapped_exception: return mapped_exception @@ -676,6 +686,7 @@ class BaseResponsesAPIStreamingIterator: self.model or "", self.custom_llm_provider or "", headers=provider_request_id_headers(self._raw_response_headers), + logging_obj=self.logging_obj, ) def _maybe_raise_for_error_event(self, result: object) -> None: @@ -1983,13 +1994,17 @@ class _ResponsesWsTurn: def with_event(self, event: Mapping[str, object]) -> _ResponsesWsTurn: return _ResponsesWsTurn(logging_obj=self.logging_obj, events=(*self.events, event)) - def failure(self, model: str, custom_llm_provider: str) -> Exception | None: + def failure( + self, model: str, custom_llm_provider: str, connection_logging_obj: LiteLLMLoggingObj + ) -> Exception | None: failed_event: Final = next( (event for event in self.events if event.get("type") in _RESPONSES_WS_FAILURE_EVENT_TYPES), None ) if failed_event is None: return None - return _map_stream_error_to_exception(_ws_event_error(failed_event), model, custom_llm_provider) + return _map_stream_error_to_exception( + _ws_event_error(failed_event), model, custom_llm_provider, logging_obj=connection_logging_obj + ) _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"}) @@ -2148,7 +2163,7 @@ class ResponsesWebSocketStreaming: def _log_turn(self, turn: _ResponsesWsTurn) -> None: events: Final = list(turn.events) - exception: Final = turn.failure(self.authorized_model or "", self.custom_llm_provider or "") + exception: Final = turn.failure(self.authorized_model or "", self.custom_llm_provider or "", self.logging_obj) if exception is None: asyncio.create_task(turn.logging_obj.dispatch_success_handlers(events, prefer_async_handlers=True)) return diff --git a/tests/integration/observability/test_stream_error_alert_api_base.py b/tests/integration/observability/test_stream_error_alert_api_base.py new file mode 100644 index 00000000000..449e378dfc7 --- /dev/null +++ b/tests/integration/observability/test_stream_error_alert_api_base.py @@ -0,0 +1,482 @@ +import asyncio +import hashlib +import json +import uuid +from collections.abc import Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import openai +import pytest +import yaml +from anthropic.types import RawContentBlockDeltaEvent, TextDelta +from integration._support.anthropic_sse import DROP_PAUSE, dropping_reply, stream_reply +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_stream import created, delta, failed, frame +from integration._support.responses_ws_peer import PROVIDER_MODEL, Script, events, responses_peer, session +from integration._support.tls import write_self_signed_cert +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types.chat import ChatCompletionChunk +from openai.types.responses import ResponseFailedEvent, ResponseTextDeltaEvent +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +AZURE_KEY: Final = "synthetic-azure-key" +API_VERSION: Final = "2025-04-01-preview" +AZURE_CHAT_BACKEND: Final = "azure/gpt-4o-mini" +PARTIAL_TEXT: Final = "partial answer" +ALERT_SECONDS: Final = 40 +NO_RETRY_NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0, "cache": {"no-cache": True}}) +EMPTY_MODEL_LIST: Final = Reply(body=b'{"object": "list", "data": []}') +VERTEX_PROJECT: Final = "alert-vertex-project" +VERTEX_LOCATION: Final = "us-east5" +VERTEX_PARTNER_BACKEND: Final = "claude-opus-5-5@default" +VERTEX_PARTNER_URL: Final = ( + f"https://{VERTEX_LOCATION}-aiplatform.googleapis.com/v1/projects/{VERTEX_PROJECT}" + f"/locations/{VERTEX_LOCATION}/publishers/anthropic/models/{VERTEX_PARTNER_BACKEND}" +) +VERTEX_OVERLOADED: Final = Reply( + status=500, body=b'{"error": {"code": 500, "message": "vertex partner overloaded", "status": "INTERNAL"}}' +) + +API_BASES: Final[Mapping[str, Callable[[str, str], str]]] = MappingProxyType( + { + "plain": lambda url, _: f"{url}/gateway", + "code": lambda url, secret: f"{url}/gateway?code={secret}", + "subscription-key": lambda url, secret: f"{url}/gateway?subscription-key={secret}", + "userinfo": lambda url, secret: url.replace("://", f"://alice:{secret}@", 1) + "/gateway", + } +) + + +@dataclass(frozen=True, slots=True) +class AlertingProxy: + gateway: Gateway + alerts: Wire + vertex: Wire + cert: tuple[Path, Path] + + +@dataclass(frozen=True, slots=True) +class Caller: + key: str + key_alias: str + team_alias: str + + +@dataclass(frozen=True, slots=True) +class Outcome: + text: str + failure: str + + +def accept_alert(request: Request) -> Reply: + return Reply(body=b"ok", content_type="text/plain") + + +def overloaded_vertex(request: Request) -> Reply: + return VERTEX_OVERLOADED + + +def alerting_config(directory: Path) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "general_settings": { + **object_value(base["general_settings"]), + "alerting": ["slack"], + "alert_types": ["llm_exceptions"], + }, + } + path: Final = directory / "alerting.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def alerting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[AlertingProxy]: + directory: Final = tmp_path_factory.mktemp("stream-error-alerts") + cert: Final = write_self_signed_cert(directory) + with ( + gateway_from_environment() as shared, + wire_server(accept_alert) as alerts, + wire_server(overloaded_vertex) as vertex, + ): + with owned_proxy( + shared, + directory, + { + "SLACK_WEBHOOK_URL": alerts.url, + "DEFAULT_FLUSH_INTERVAL_SECONDS": "1", + "SSL_CERT_FILE": str(cert[0]), + "VERTEXAI_API_BASE": vertex.url, + }, + config=alerting_config(directory), + ) as gateway: + yield AlertingProxy(gateway, alerts, vertex, cert) + + +def alert_naming(alerts: Wire, marker: str) -> str: + def matching() -> tuple[str, ...]: + texts: Final = (string_value(JSON_OBJECT.validate_json(request.body)["text"]) for request in alerts.drain()) + return tuple(text for text in texts if marker in text) + + return eventually(matching, lambda found: len(found) > 0, seconds=ALERT_SECONDS)[0] + + +def api_base_line(alert: str) -> str | None: + return next((line for line in alert.splitlines() if line.startswith("API Base:")), None) + + +def assert_names_deployment(alert: str, api_base: str, model_group: str, backend: str, caller: Caller) -> None: + assert api_base_line(alert) == f"API Base: `{api_base}`", alert + assert f"model_group: `{model_group}`" in alert, alert + assert f"deployment: `{backend}`" in alert, alert + assert f"Key Name: `{caller.key_alias}`\nTeam: `{caller.team_alias}`" in alert, alert + + +def keyed_caller(scenario: Scenario) -> Caller: + team_alias: Final = f"alert-team-{uuid.uuid4().hex[:12]}" + key_alias: Final = f"alert-key-{uuid.uuid4().hex[:12]}" + team: Final = scenario.team(team_alias=team_alias) + return Caller(scenario.key(key_alias=key_alias, team_id=team), key_alias, team_alias) + + +def azure_chat(scenario: Scenario, api_base: str, *, bridged: bool = False) -> str: + return scenario.model( + model=AZURE_CHAT_BACKEND, + api_base=api_base, + api_key=AZURE_KEY, + api_version=API_VERSION, + num_retries=0, + use_chat_completions_api=bridged, + ) + + +def provider_message(marker: str) -> str: + return f"upstream overloaded {marker}" + + +def chat_chunk(marker: str, content: Mapping[str, JsonValue]) -> bytes: + body: Final = { + "id": f"chatcmpl-{marker}", + "object": "chat.completion.chunk", + "created": 1, + "model": f"gpt-4o-mini-{marker}", + "choices": [{"index": 0, "delta": dict(content), "finish_reason": None}], + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def first_text(marker: str) -> bytes: + return chat_chunk(marker, {"role": "assistant", "content": PARTIAL_TEXT}) + + +def cut_after_text(marker: str) -> Reply: + return dropping_reply((first_text(marker), chat_chunk(marker, {"content": " never sent"})), abort_after=1) + + +def error_frame_after_text(marker: str) -> Reply: + error: Final = {"error": {"message": provider_message(marker), "type": "server_error", "code": "server_error"}} + return stream_reply((first_text(marker), b"data: " + json.dumps(error).encode() + b"\n\n"), pause=DROP_PAUSE) + + +def rejected_before_stream(marker: str) -> Reply: + return Reply( + status=500, body=json.dumps({"error": {"message": provider_message(marker), "type": "server_error"}}).encode() + ) + + +def chat_upstream(reply: Callable[[str], Reply], marker: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return EMPTY_MODEL_LIST + return reply(marker) + + return respond + + +def items_or_error(open_stream: Callable[[], Iterable[object]], error_type: type[Exception]) -> Iterator[object]: + try: + yield from open_stream() + except error_type as error: + yield error + + +def chunk_text(item: object) -> str: + if not isinstance(item, ChatCompletionChunk): + return "" + return "".join(choice.delta.content or "" for choice in item.choices) + + +def message_text(item: object) -> str: + if isinstance(item, RawContentBlockDeltaEvent) and isinstance(item.delta, TextDelta): + return item.delta.text + return "" + + +def responses_text(item: object) -> str: + return item.delta if isinstance(item, ResponseTextDeltaEvent) else "" + + +def failure_of(item: object) -> str: + if isinstance(item, Exception): + return str(item) + if isinstance(item, ResponseFailedEvent) and item.response.error is not None: + return item.response.error.message + return "" + + +def outcome(items: tuple[object, ...], text_of: Callable[[object], str]) -> Outcome: + return Outcome("".join(text_of(item) for item in items), "".join(failure_of(item) for item in items)) + + +def openai_client(gateway: Gateway, key: str) -> openai.OpenAI: + return openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=key, max_retries=0, timeout=60) + + +def chat_outcome(gateway: Gateway, key: str, model: str) -> Outcome: + with openai_client(gateway, key) as client: + items: Final = tuple( + items_or_error( + lambda: client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "stream an answer"}], + stream=True, + extra_body=dict(NO_RETRY_NO_CACHE), + ), + openai.APIError, + ) + ) + return outcome(items, chunk_text) + + +def messages_outcome(gateway: Gateway, key: str, model: str) -> Outcome: + with anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=key, max_retries=0, timeout=60) as client: + items: Final = tuple( + items_or_error( + lambda: client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": "stream an answer"}], + stream=True, + extra_body=dict(NO_RETRY_NO_CACHE), + ), + anthropic.APIError, + ) + ) + return outcome(items, message_text) + + +def responses_outcome(gateway: Gateway, key: str, model: str) -> Outcome: + with openai_client(gateway, key) as client: + items: Final = tuple( + items_or_error( + lambda: client.responses.create( + model=model, input="stream an answer", stream=True, extra_body=dict(NO_RETRY_NO_CACHE) + ), + openai.APIError, + ) + ) + return outcome(items, responses_text) + + +ENTRYPOINTS: Final[Mapping[str, Callable[[Gateway, str, str], Outcome]]] = MappingProxyType( + {"chat": chat_outcome, "messages": messages_outcome, "responses-bridge": responses_outcome} +) +FAILURES: Final[Mapping[str, Callable[[str], Reply]]] = MappingProxyType( + {"cut": cut_after_text, "error-frame": error_frame_after_text} +) + + +@pytest.mark.parametrize( + ("entrypoint", "failure"), + [("chat", "cut"), ("chat", "error-frame"), ("messages", "error-frame"), ("responses-bridge", "error-frame")], +) +def test_mid_stream_failure_alert_names_the_deployment(alerting: AlertingProxy, entrypoint: str, failure: str) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(chat_upstream(FAILURES[failure], marker)) as wire, alerting.gateway.scenario() as scenario: + model: Final = azure_chat(scenario, wire.url, bridged=entrypoint == "responses-bridge") + caller: Final = keyed_caller(scenario) + seen: Final = ENTRYPOINTS[entrypoint](alerting.gateway, caller.key, model) + assert seen.text == PARTIAL_TEXT and seen.failure, seen + assert failure == "cut" or provider_message(marker) in seen.failure, seen + assert_names_deployment(alert_naming(alerting.alerts, marker), wire.url, model, AZURE_CHAT_BACKEND, caller) + + +def test_native_responses_failed_event_alert_names_the_deployment(alerting: AlertingProxy) -> None: + marker: Final = uuid.uuid4().hex + identity: Final = f"resp_{marker}" + stream: Final = ( + frame(created(identity)), + frame(delta(identity, PARTIAL_TEXT)), + frame(failed(identity, "server_error", provider_message(marker))), + ) + + def respond(request: Request) -> Reply: + if request.method != "POST": + return EMPTY_MODEL_LIST + assert "/responses" in request.target, request.target + return stream_reply(stream) + + with wire_server(respond) as wire, alerting.gateway.scenario() as scenario: + model: Final = scenario.model(model=AZURE_CHAT_BACKEND, api_base=wire.url, api_key=AZURE_KEY, num_retries=0) + caller: Final = keyed_caller(scenario) + seen: Final = responses_outcome(alerting.gateway, caller.key, model) + assert seen.text == PARTIAL_TEXT and provider_message(marker) in seen.failure, seen + assert_names_deployment(alert_naming(alerting.alerts, marker), wire.url, model, AZURE_CHAT_BACKEND, caller) + + +def failing_turn(marker: str) -> Script: + def script( + response_id: str, + ) -> tuple[dict[str, JsonValue], ...]: # mutable-ok: responses_peer's Script returns dicts + opened, text_delta, completed_event = events(response_id, PARTIAL_TEXT) + response: Final = object_value(completed_event["response"]) + return ( + opened, + text_delta, + { + "type": "response.failed", + "response": { + **response, + "status": "failed", + "output": [], + "error": {"code": "server_error", "message": provider_message(marker)}, + }, + }, + ) + + return script + + +def test_responses_websocket_turn_failure_alert_names_the_deployment(alerting: AlertingProxy) -> None: + marker: Final = uuid.uuid4().hex + backend: Final = f"openai/{PROVIDER_MODEL}" + with responses_peer(alerting.cert, failing_turn(marker)) as peer, alerting.gateway.scenario() as scenario: + model: Final = scenario.model(model=backend, api_base=peer.url, num_retries=0) + caller: Final = keyed_caller(scenario) + (received,) = asyncio.run(session(str(alerting.gateway.client.base_url), caller.key, model, (marker,))) + assert [event["type"] for event in received] == [ + "response.created", + "response.output_text.delta", + "response.failed", + ], received + assert_names_deployment(alert_naming(alerting.alerts, marker), peer.url, model, backend, caller) + (row,) = eventually( + lambda: read_rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (hashlib.sha256(caller.key.encode()).hexdigest(),), + ), + lambda rows: len(rows) >= 1, + seconds=70, + ) + assert row["status"] == "failure", row + error_information: Final = object_value(object_value(row["metadata"])["error_information"]) + assert provider_message(marker) in string_value(error_information["error_message"]), error_information + + +@pytest.mark.parametrize( + ("phase", "placement"), + [ + ("pre-stream", "plain"), + ("pre-stream", "code"), + ("pre-stream", "subscription-key"), + ("pre-stream", "userinfo"), + ("mid-stream", "code"), + ], +) +def test_alert_api_base_line_keeps_credentials_out(alerting: AlertingProxy, phase: str, placement: str) -> None: + marker: Final = uuid.uuid4().hex + secret: Final = f"credential{uuid.uuid4().hex}" + pre_stream: Final = phase == "pre-stream" + reply: Final = rejected_before_stream if pre_stream else error_frame_after_text + with wire_server(chat_upstream(reply, marker)) as wire, alerting.gateway.scenario() as scenario: + model: Final = azure_chat(scenario, API_BASES[placement](wire.url, secret)) + caller: Final = keyed_caller(scenario) + seen: Final = chat_outcome(alerting.gateway, caller.key, model) + assert seen.failure and seen.text == ("" if pre_stream else PARTIAL_TEXT), seen + alert: Final = alert_naming(alerting.alerts, model if pre_stream else marker) + assert secret not in alert, alert + assert_names_deployment(alert, f"{wire.url}/gateway", model, AZURE_CHAT_BACKEND, caller) + + +def test_alert_api_base_line_keeps_a_malformed_password_out(alerting: AlertingProxy) -> None: + password_start: Final = f"credential{uuid.uuid4().hex[:8]}" + password_end: Final = uuid.uuid4().hex + with alerting.gateway.scenario() as scenario: + model: Final = azure_chat(scenario, f"http://alice:{password_start}?{password_end}@127.0.0.1:9/gateway") + caller: Final = keyed_caller(scenario) + seen: Final = chat_outcome(alerting.gateway, caller.key, model) + assert seen.failure and not seen.text, seen + line: Final = api_base_line(alert_naming(alerting.alerts, model)) or "" + assert not any(part in line for part in ("alice", password_start, password_end)), line + + +def test_alert_api_base_line_names_a_vertex_partner_model_version(alerting: AlertingProxy) -> None: + with alerting.gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"vertex_ai/{VERTEX_PARTNER_BACKEND}", + api_base=None, + api_key=None, + vertex_project=VERTEX_PROJECT, + vertex_location=VERTEX_LOCATION, + vertex_credentials=service_account_json(VERTEX_PROJECT, alerting.gateway.upstream_url.rstrip("/")), + num_retries=0, + ) + caller: Final = keyed_caller(scenario) + seen: Final = chat_outcome(alerting.gateway, caller.key, model) + assert seen.failure and not seen.text, seen + assert any( + f"/publishers/anthropic/models/{VERTEX_PARTNER_BACKEND}:" in request.target + for request in alerting.vertex.drain() + ) + line: Final = api_base_line(alert_naming(alerting.alerts, model)) or "" + assert line.startswith(f"API Base: `{VERTEX_PARTNER_URL}:"), line + + +def test_fallback_mid_stream_failure_alert_names_the_fallback_deployment(alerting: AlertingProxy) -> None: + marker: Final = uuid.uuid4().hex + primary_reply: Final = Reply(status=500, body=b'{"error": {"message": "primary is down", "type": "server_error"}}') + with ( + wire_server(chat_upstream(lambda _: primary_reply, marker)) as primary_wire, + wire_server(chat_upstream(cut_after_text, marker)) as backup_wire, + alerting.gateway.scenario() as scenario, + ): + primary: Final = azure_chat(scenario, primary_wire.url) + backup: Final = azure_chat(scenario, backup_wire.url) + caller: Final = keyed_caller(scenario) + with openai_client(alerting.gateway, caller.key) as client: + items: Final = tuple( + items_or_error( + lambda: client.chat.completions.create( + model=primary, + messages=[{"role": "user", "content": "stream an answer"}], + stream=True, + extra_body={**NO_RETRY_NO_CACHE, "fallbacks": [backup]}, + ), + openai.APIError, + ) + ) + seen: Final = outcome(items, chunk_text) + assert seen.text == PARTIAL_TEXT and seen.failure, seen + assert [request.method for request in primary_wire.drain()].count("POST") == 1 + assert [request.method for request in backup_wire.drain()].count("POST") == 1 + alert: Final = alert_naming(alerting.alerts, marker) + assert_names_deployment(alert, backup_wire.url, backup, AZURE_CHAT_BACKEND, caller) + assert primary_wire.url not in alert, alert diff --git a/tests/integration/sdk/test_stream_error_api_base_sdk.py b/tests/integration/sdk/test_stream_error_api_base_sdk.py new file mode 100644 index 00000000000..922e518c7bf --- /dev/null +++ b/tests/integration/sdk/test_stream_error_api_base_sdk.py @@ -0,0 +1,69 @@ +import json +from collections.abc import AsyncIterable, AsyncIterator +from typing import Final + +import openai +from integration._support.anthropic_sse import dropping_reply +from integration._support.wire import Reply, Request, wire_server + +import litellm +from litellm.exceptions import MidStreamFallbackError + +_PARTIAL_TEXT: Final = "partial answer" + + +def _chunk(content: str) -> bytes: + body: Final = { + "id": "chatcmpl-sdk-cut", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": content}, "finish_reason": None}], + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def _cut_after_text(request: Request) -> Reply: + assert request.method == "POST" and "/chat/completions" in request.target, request.target + return dropping_reply((_chunk(_PARTIAL_TEXT), _chunk(" never sent")), abort_after=1) + + +def _text(item: object) -> str: + if not isinstance(item, litellm.ModelResponseStream): + return "" + return "".join(choice.delta.content or "" for choice in item.choices) + + +async def _items_or_error(stream: AsyncIterable[object]) -> AsyncIterator[object]: + try: + async for item in stream: + yield item + except openai.APIError as error: + yield error + + +async def test_sdk_mid_stream_failure_debug_info_names_the_api_base_it_called() -> None: + with wire_server(_cut_after_text) as wire: + response: Final = await litellm.acompletion( + model="azure/gpt-4o-mini", + api_base=wire.url, + api_key="synthetic-azure-key", + api_version="2025-04-01-preview", + messages=[{"role": "user", "content": "stream an answer"}], + stream=True, + num_retries=0, + ) + assert isinstance(response, AsyncIterable), response + items: Final = [item async for item in _items_or_error(response)] + errors: Final = [item for item in items if isinstance(item, openai.APIError)] + assert len(errors) == 1 and items[-1] is errors[0], items + assert "".join(_text(item) for item in items) == _PARTIAL_TEXT, items + error: Final = errors[0] + mapped: Final = ( + error.original_exception + if isinstance(error, MidStreamFallbackError) and error.original_exception is not None + else error + ) + debug_info: Final = getattr(mapped, "litellm_debug_info", None) + assert isinstance(debug_info, str), mapped + assert f"\nAPI Base: `{wire.url}`" in debug_info, debug_info diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index b24348cfaa8..9eb99516e02 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -57,6 +57,104 @@ def test_openai_cyber_policy_uses_the_structured_error_code(shape: str, code: st assert caught.value.response.headers["x-request-id"] == "cyber-policy-request" + +def _debug_info_for_api_base(configured_api_base: str) -> str: + request: Final = httpx.Request("POST", "https://gateway.example/v1/chat/completions") + original: Final = OpenAIError( + status_code=500, message="upstream failed", request=request, response=httpx.Response(500, request=request) + ) + with pytest.raises(litellm.InternalServerError) as caught: + exception_type( + model="gpt-5.4-mini", + original_exception=original, + custom_llm_provider="openai", + completion_kwargs={}, + extra_kwargs={"api_base": configured_api_base}, + ) + debug_info: Final = caught.value.litellm_debug_info + assert isinstance(debug_info, str) + return debug_info + + +@pytest.mark.parametrize( + ("configured_api_base", "shown_api_base"), + [ + ("https://qa-user:function-secret@gateway.example/v1?code=function-secret", "https://gateway.example/v1"), + ("https://gateway.example/v1#function-secret", "https://gateway.example/v1"), + ("qa-user:function-secret@gateway.example:8080/v1", "gateway.example:8080/v1"), + ("https://gateway.example?code=function/secret", "https://gateway.example"), + ("https://gateway.example?user=alice@example.org&code=function-secret", "https://gateway.example"), + ("https://gateway.example?email=ops@other-host.example", "https://gateway.example"), + ("qa-user:function-secret@gateway.example/v1?next=https://other.example/y", "gateway.example/v1"), + ( + "https://qa-user:function-secret@gateway.example/models/claude-opus-5-5@default", + "https://gateway.example/models/claude-opus-5-5@default", + ), + ], + ids=[ + "userinfo_and_query", + "fragment", + "no_scheme", + "slash_in_query_without_path", + "at_sign_in_query_without_path", + "at_sign_names_other_host", + "no_scheme_with_url_in_query", + "userinfo_and_at_sign_in_path", + ], +) +def test_api_base_line_keeps_credentials_out(configured_api_base: str, shown_api_base: str) -> None: + debug_info: Final = _debug_info_for_api_base(configured_api_base) + assert f"API Base: `{shown_api_base}`" in debug_info + assert "function-secret" not in debug_info + assert "secret" not in debug_info.split("API Base:", 1)[1].split("\n", 1)[0] + + +@pytest.mark.parametrize( + "configured_api_base", + [ + "https://qa-user:qa?function-secret@gateway.example/v1", + "https://qa-user:qa#function-secret@gateway.example/v1", + "https://qa-user:qa/function-secret@gateway.example/v1", + "https://qa-user:qa/function?secret@gateway.example/v1", + "https://qa.user:qa/function-secret@gateway.example/v1", + "https://qa-user:1234/function-secret@gateway.example/v1", + "https://function-secret/more@gateway.example/v1", + "qa-user:qa://function-secret@gateway.example", + ], + ids=[ + "question_mark_in_password", + "hash_in_password", + "slash_in_password", + "slash_and_question_mark_in_password", + "dotted_username_and_slash_in_password", + "numeric_password_with_slash", + "slash_in_username", + "no_scheme_with_scheme_separator_in_password", + ], +) +def test_api_base_line_is_left_out_when_a_password_breaks_the_url(configured_api_base: str) -> None: + debug_info: Final = _debug_info_for_api_base(configured_api_base) + assert "API Base:" not in debug_info + assert not any(part in debug_info for part in ("qa-user", "qa.user", "function", "secret")) + + +@pytest.mark.parametrize( + "api_base", + [ + "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" + "/publishers/anthropic/models/claude-opus-5-5@default:rawPredict", + "us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1" + "/publishers/mistralai/models/mistral-large@2407:rawPredict", + "https://gateway.ai.cloudflare.com/v1/account/gateway/workers-ai/@cf/meta/llama-3.1-8b-instruct", + "http://localhost:4000/v1/models/claude@default", + "https://[::1]:8080/v1/models/claude@default", + ], + ids=["vertex_partner_model", "vertex_without_scheme", "workers_ai_path", "localhost", "ipv6"], +) +def test_api_base_line_keeps_an_at_sign_in_the_path(api_base: str) -> None: + assert f"API Base: `{api_base}`" in _debug_info_for_api_base(api_base) + + # Test cases for is_error_str_context_window_exceeded # Tuple format: (error_message, expected_result) context_window_test_cases = [ diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 93f700b1c2a..c07fe2d1c31 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -2,14 +2,17 @@ import asyncio import json import threading import time -from typing import Final, NoReturn, Optional +from typing import AsyncIterator, Final, Iterator, Literal, NoReturn, Optional from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import pytest import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import ( CustomStreamWrapper, @@ -1382,6 +1385,140 @@ async def test_async_streaming_read_timeout_triggers_midstream_fallback( assert isinstance(excinfo.value.original_exception, Exception) +_AZURE_DEPLOYMENT_API_BASE: Final = "https://mid-stream-debug.openai.azure.com" +_AZURE_DEPLOYMENT_REQUEST_URL: Final = ( + f"{_AZURE_DEPLOYMENT_API_BASE}/openai/deployments/gpt-5.4-mini/chat/completions?api-version=2025-04-01-preview" +) +_MID_STREAM_START_TIME: Final = 1_790_000_000.0 + + +def _azure_deployment_logging_obj(logged_api_base: str | None, routed_api_base: str | None) -> Logging: + logging_obj: Final = Logging( + model="gpt-5.4-mini", + messages=[{"role": "user", "content": "Count from 1 to 40"}], + stream=True, + call_type="acompletion", + start_time=_MID_STREAM_START_TIME, + litellm_call_id="mid-stream-debug-info", + function_id="mid-stream-debug-info", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-mini", + optional_params={}, + litellm_params=get_litellm_params( + api_base=logged_api_base, + custom_llm_provider="azure", + metadata={ + "model_group": "gpt-5.4-mini-azure", + "deployment": "azure/gpt-5.4-mini", + "api_base": routed_api_base, + }, + ), + ) + return logging_obj + + +def _content_then_failure(failure: Literal["transport_error", "unparseable_chunk"]) -> Iterator[object]: + yield ModelResponseStream( + model="gpt-5.4-mini-2026-03-17", + choices=[StreamingChoices(index=0, delta=Delta(content="1", role="assistant"))], + ) + if failure == "transport_error": + raise httpx.ReadError("Response payload is not completed") + yield object() + + +async def _as_async(chunks: Iterator[object]) -> AsyncIterator[object]: + for chunk in chunks: + yield chunk + + +async def _drain(wrapper: CustomStreamWrapper, sync_mode: bool) -> None: + if sync_mode: + for _ in wrapper: + pass + return + async for _ in wrapper: + pass + + +async def _mid_stream_debug_info( + logged_api_base: str | None, + routed_api_base: str | None, + failure: Literal["transport_error", "unparseable_chunk"], + sync_mode: bool, +) -> str: + chunks: Final = _content_then_failure(failure) + wrapper: Final = CustomStreamWrapper( + completion_stream=chunks if sync_mode else _as_async(chunks), + model="gpt-5.4-mini-2026-03-17", + logging_obj=_azure_deployment_logging_obj(logged_api_base, routed_api_base), + custom_llm_provider="azure", + ) + with pytest.raises(MidStreamFallbackError) as excinfo: + await _drain(wrapper, sync_mode) + assert excinfo.value.is_pre_first_chunk is False + debug_info: Final = getattr(excinfo.value.original_exception, "litellm_debug_info", None) + assert isinstance(debug_info, str) + return debug_info + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("failure", ["transport_error", "unparseable_chunk"]) +@pytest.mark.parametrize( + ("logged_api_base", "routed_api_base"), + [ + (_AZURE_DEPLOYMENT_REQUEST_URL, _AZURE_DEPLOYMENT_API_BASE), + (_AZURE_DEPLOYMENT_API_BASE, None), + (_AZURE_DEPLOYMENT_API_BASE, ""), + ], + ids=["routed", "sdk", "empty_routed"], +) +async def test_mid_stream_error_reports_the_deployment_it_was_sent_to( + logged_api_base: str, + routed_api_base: str | None, + failure: Literal["transport_error", "unparseable_chunk"], + sync_mode: bool, +): + debug_info: Final = await _mid_stream_debug_info(logged_api_base, routed_api_base, failure, sync_mode) + + assert f"API Base: `{_AZURE_DEPLOYMENT_API_BASE}`" in debug_info + assert "model_group: `gpt-5.4-mini-azure`" in debug_info + assert "deployment: `azure/gpt-5.4-mini`" in debug_info + assert "Count from 1 to 40" not in debug_info + assert "api.openai.com" not in debug_info + + +@pytest.mark.asyncio +async def test_mid_stream_error_leaves_out_an_api_base_it_cannot_resolve(): + debug_info: Final = await _mid_stream_debug_info(None, None, "transport_error", sync_mode=False) + + assert "model_group: `gpt-5.4-mini-azure`" in debug_info + assert "API Base:" not in debug_info + + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "routed_api_base", + [ + f"{_AZURE_DEPLOYMENT_API_BASE}?subscription-key=deployment-secret-value", + f"{_AZURE_DEPLOYMENT_API_BASE}?code=deployment-secret-value", + _AZURE_DEPLOYMENT_API_BASE.replace("https://", "https://qa-user:deployment-secret-value@"), + ], + ids=["subscription_key", "function_code", "userinfo"], +) +async def test_mid_stream_error_keeps_credentials_in_the_api_base_out(routed_api_base: str): + debug_info: Final = await _mid_stream_debug_info( + _AZURE_DEPLOYMENT_REQUEST_URL, routed_api_base, "transport_error", sync_mode=False + ) + + assert f"API Base: `{_AZURE_DEPLOYMENT_API_BASE}`" in debug_info + assert "deployment-secret" not in debug_info + + + def test_streaming_handler_with_created_time_propagation( initialized_custom_stream_wrapper: CustomStreamWrapper, logging_obj: Logging ): diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 001be7e1fbd..04d2b9ddf81 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -19,6 +19,7 @@ Also covers: ErrorEventError.param must accept dict payloads without raising a Pydantic ValidationError (previously typed as Optional[str]). """ +import asyncio import json from datetime import datetime from importlib import import_module @@ -31,15 +32,18 @@ import pytest import litellm from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( _ERROR_CODE_HTTP_STATUS, BaseResponsesAPIStreamingIterator, ResponsesAPIStreamingIterator, + ResponsesWebSocketStreaming, SyncResponsesAPIStreamingIterator, _status_code_for_error_fields, ) +from litellm.utils import ProviderConfigManager from litellm.types.llms.openai import ( ErrorEvent, ErrorEventError, @@ -658,3 +662,140 @@ def test_openai_top_level_error_event_keeps_the_provider_code_and_message( assert exc_info.value.status_code == status_code assert type(exc_info.value.original_exception) is exception_type assert f"upstream said {code}" in str(exc_info.value.original_exception) + + +_AZURE_RESPONSES_API_BASE: Final = "https://responses-mid-stream.openai.azure.com" +_AZURE_RESPONSES_REQUEST_URL: Final = f"{_AZURE_RESPONSES_API_BASE}/openai/responses?api-version=2025-04-01-preview" +_AZURE_RESPONSES_FAILED_EVENT: Final = { + "type": "response.failed", + "sequence_number": 2, + "response": { + "id": "resp_mid_stream", + "object": "response", + "created_at": 1790000000, + "status": "failed", + "model": "gpt-5.4-mini", + "output": [], + "error": {"code": "server_error", "message": "upstream failed mid-stream"}, + }, +} +_AZURE_RESPONSES_WS_URL: Final = "wss://responses-mid-stream.openai.azure.com/openai/v1/responses?api-version=preview" + + +def _azure_responses_logging_obj(logged_api_base: str, routed_api_base: str | None) -> LiteLLMLoggingObj: + logging_obj: Final = LiteLLMLoggingObj( + model="gpt-5.4-mini", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime(2026, 1, 1), + litellm_call_id="responses-mid-stream", + function_id="responses-mid-stream", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-mini", + optional_params={}, + litellm_params=get_litellm_params( + api_base=logged_api_base, + custom_llm_provider="azure", + metadata={ + "model_group": "gpt-5.4-mini-azure", + "deployment": "azure/gpt-5.4-mini", + "api_base": routed_api_base, + }, + ), + ) + return logging_obj + + +def _assert_names_the_azure_deployment(exception: BaseException, api_base: str = _AZURE_RESPONSES_API_BASE) -> None: + debug_info: Final = getattr(exception, "litellm_debug_info", None) + assert isinstance(debug_info, str) + assert f"API Base: `{api_base}`" in debug_info + assert "model_group: `gpt-5.4-mini-azure`" in debug_info + assert "deployment: `azure/gpt-5.4-mini`" in debug_info + assert "api.openai.com" not in debug_info + + +@pytest.mark.asyncio +async def test_response_failed_after_output_reports_the_deployment_it_was_sent_to(): + delta: Final = { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_mid_stream", + "output_index": 0, + "content_index": 0, + "delta": "1", + } + body: Final = "".join(f"data: {json.dumps(event)}\n\n" for event in (delta, _AZURE_RESPONSES_FAILED_EVENT)) + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body.encode()), + model="gpt-5.4-mini", + responses_api_provider_config=ProviderConfigManager.get_provider_responses_api_config( + provider="azure", model="gpt-5.4-mini" + ), + logging_obj=_azure_responses_logging_obj(_AZURE_RESPONSES_REQUEST_URL, _AZURE_RESPONSES_API_BASE), + custom_llm_provider="azure", + ) + stream: Final = aiter(iterator) + assert getattr(await anext(stream), "delta", None) == "1" + + with pytest.raises(MidStreamFallbackError) as exc_info: + await anext(stream) + assert exc_info.value.is_pre_first_chunk is False + _assert_names_the_azure_deployment(exc_info.value.original_exception) + + +class _ScriptedAzureSocket: + def __init__(self, events: tuple[str, ...]) -> None: + self._events = iter(events) + self._drained = asyncio.Event() + + async def recv(self, decode: bool = False) -> str: + event: Final = next(self._events, None) + if event is None: + self._drained.set() + raise ConnectionError("backend closed") + return event + + async def send(self, message: str) -> None: + return None + + async def close(self) -> None: + return None + + async def send_text(self, data: str) -> None: + return None + + async def receive_text(self) -> str: + await self._drained.wait() + raise ConnectionError("client closed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("routed_api_base", "reported_api_base"), + [ + (_AZURE_RESPONSES_API_BASE, _AZURE_RESPONSES_API_BASE), + (None, "wss://responses-mid-stream.openai.azure.com/openai/v1/responses"), + ], + ids=["router", "no_router"], +) +async def test_responses_websocket_turn_failure_reports_the_deployment_it_was_sent_to( + routed_api_base: str | None, reported_api_base: str +): + connection_logging_obj: Final = _azure_responses_logging_obj(_AZURE_RESPONSES_WS_URL, routed_api_base) + connection_logging_obj.pre_call(input=None, api_key="", additional_args={"api_base": _AZURE_RESPONSES_WS_URL}) + socket: Final = _ScriptedAzureSocket((json.dumps(_AZURE_RESPONSES_FAILED_EVENT),)) + streaming: Final = ResponsesWebSocketStreaming( + websocket=socket, + backend_ws=socket, + logging_obj=connection_logging_obj, + authorized_model="gpt-5.4-mini", + custom_llm_provider="azure", + ) + + failure: Final = await streaming.bidirectional_forward() + + assert failure is not None + _assert_names_the_azure_deployment(failure, reported_api_base) From b578dcf278fdf61bb66f3cf28611bd74c7530f8f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:29:56 -0700 Subject: [PATCH 03/64] test(translation): recapture anthropic claude bases after usage speed field (#47116) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_completions/bases/anthropic.py | 58 +++++++---- .../chat_completions/cross_model/anthropic.py | 96 ++++++++++++++++--- .../chat_completions/cross_model/case.py | 9 +- .../translation/messages/bases/anthropic.py | 37 +++++-- .../cross_model/test_messages_cross_model.py | 27 +++++- .../test_messages_service_tier_anthropic.py | 20 +++- .../translation/responses/bases/anthropic.py | 38 +++++--- .../cross_model/test_responses_cross_model.py | 15 ++- 8 files changed, 235 insertions(+), 65 deletions(-) diff --git a/tests/integration/translation/chat_completions/bases/anthropic.py b/tests/integration/translation/chat_completions/bases/anthropic.py index 16e1de53975..3a148c64859 100644 --- a/tests/integration/translation/chat_completions/bases/anthropic.py +++ b/tests/integration/translation/chat_completions/bases/anthropic.py @@ -42,7 +42,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011CfgDzBnuPYc7TdysZFRpo", + "id": "msg_011Cg7v1TRmiNAjM2RRxcxF8", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -58,6 +58,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -93,6 +94,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -123,7 +125,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfgDzH3AcaSNXtvurLgAG", + "id": "msg_011Cg7v28wK6VaszM4tL4D6n", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -139,6 +141,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -174,6 +177,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -223,6 +227,7 @@ CLAUDE_FABLE_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -272,6 +277,7 @@ CLAUDE_FABLE_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -321,6 +327,7 @@ CLAUDE_OPUS_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -370,6 +377,7 @@ CLAUDE_OPUS_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -400,7 +408,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-5", - "id": "msg_011CfgDzPkj2HxBq8qGJsGSi", + "id": "msg_011Cg7v2SdRUNPAvGScTF7au", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -417,6 +425,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -452,6 +461,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -482,7 +492,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfgDzV9vZCkuQJt15n3Z1", + "id": "msg_011Cg7v1hPQMJ1QhaC8HaSfQ", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -499,6 +509,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -534,6 +545,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -549,8 +561,16 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( "content-type": "application/json", }, expected_provider_request={**ANTHROPIC_MESSAGES_REQUEST, "model": "claude-sonnet-5-5"}, - mock_provider_response={**ANTHROPIC_MESSAGES_RESPONSE, "model": "claude-sonnet-5-5"}, - expected_litellm_response={**CHAT_COMPLETIONS_RESPONSE, "model": "anthropic/claude-sonnet-5-5"}, + mock_provider_response={ + **ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-sonnet-5-5", + "usage": {**object_value(ANTHROPIC_MESSAGES_RESPONSE["usage"]), "speed": "standard"}, + }, + expected_litellm_response={ + **CHAT_COMPLETIONS_RESPONSE, + "model": "anthropic/claude-sonnet-5-5", + "usage": {**object_value(CHAT_COMPLETIONS_RESPONSE["usage"]), "speed": "standard"}, + }, ) @@ -618,6 +638,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 149}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -641,7 +662,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( }, "id": "toolu_015ndrvU4VqqoMBzK35jt4TX", "type": "function", - }, + } ], "reasoning_content": "There's a mention of a chat_completions client I don't actually have access " "to, so I'll ignore that and stick to my available tools. I'll proceed by " @@ -658,7 +679,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "examine next.\n" "\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], "provider_specific_fields": { "citations": None, @@ -672,11 +693,11 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "deciding what to examine next.\n" "\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, - }, + } ], "usage": { "completion_tokens": 288, @@ -694,6 +715,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -757,6 +779,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 66}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -780,7 +803,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( }, "id": "toolu_01Ew2u9E9j2mkT6RbwSw13uR", "type": "function", - }, + } ], "reasoning_content": "I'll focus on reading the validation file first, around line 18, since that's " "what was specifically requested before looking at the tests.\n" @@ -793,7 +816,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "tests.\n" "\n", "signature": CLAUDE_OPUS_5_5_THINKING_SIGNATURE_0, - }, + } ], "provider_specific_fields": { "citations": None, @@ -805,11 +828,11 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "the tests.\n" "\n", "signature": CLAUDE_OPUS_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, - }, + } ], "usage": { "completion_tokens": 215, @@ -827,6 +850,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -869,6 +893,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -892,11 +917,11 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( }, "id": "toolu_01XMSjzRvHCE753Wx4dU2Vxo", "type": "function", - }, + } ], "provider_specific_fields": {"citations": None, "thinking_blocks": None}, }, - }, + } ], "usage": { "completion_tokens": 143, @@ -914,6 +939,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) diff --git a/tests/integration/translation/chat_completions/cross_model/anthropic.py b/tests/integration/translation/chat_completions/cross_model/anthropic.py index ccb543ccf51..a7a06d01361 100644 --- a/tests/integration/translation/chat_completions/cross_model/anthropic.py +++ b/tests/integration/translation/chat_completions/cross_model/anthropic.py @@ -2,10 +2,7 @@ from dataclasses import replace from typing import Final from integration._support.client import list_value, object_value -from integration.translation.case import ( - SSEFrame, - StreamingTranslationTestCase, -) +from integration.translation.case import SSEFrame, StreamingTranslationTestCase from integration.translation.chat_completions.cross_model.contracts import followup_request from integration.translation.chat_completions.bases.anthropic import ( CLAUDE_FABLE_5_1_TEST_CASE, @@ -20,9 +17,84 @@ from integration.translation.chat_completions.templates import ( coding_agent_stream_message, CODING_AGENT_RESPONSE, ) -from integration.translation.templates.anthropic import ( - ANTHROPIC_AGENT_REQUEST, - ANTHROPIC_AGENT_RESPONSE, +from integration.translation.templates.anthropic import ANTHROPIC_AGENT_REQUEST, ANTHROPIC_AGENT_RESPONSE + +CLAUDE_FABLE_5_AGENT_CROSS_MODEL_PROVIDER_USAGE: Final = { + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "input_tokens": 1184, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", +} +CLAUDE_FABLE_5_AGENT_CROSS_MODEL_CLIENT_USAGE: Final = { + "completion_tokens": 105, + "prompt_tokens": 1184, + "total_tokens": 1289, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1184, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "inference_geo": "global", + "service_tier": "standard", +} +CLAUDE_FABLE_5_AGENT_CROSS_MODEL_TEST_CASE: Final = replace( + CLAUDE_FABLE_5_TEST_CASE, + mock_provider_response={ + **CLAUDE_FABLE_5_TEST_CASE.mock_provider_response, + "usage": CLAUDE_FABLE_5_AGENT_CROSS_MODEL_PROVIDER_USAGE, + }, + expected_litellm_response={ + **CLAUDE_FABLE_5_TEST_CASE.expected_litellm_response, + "usage": CLAUDE_FABLE_5_AGENT_CROSS_MODEL_CLIENT_USAGE, + }, +) + +CLAUDE_OPUS_5_AGENT_CROSS_MODEL_PROVIDER_USAGE: Final = { + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "input_tokens": 1180, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", +} +CLAUDE_OPUS_5_AGENT_CROSS_MODEL_CLIENT_USAGE: Final = { + "completion_tokens": 105, + "prompt_tokens": 1180, + "total_tokens": 1285, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1180, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "inference_geo": "global", + "service_tier": "standard", +} +CLAUDE_OPUS_5_AGENT_CROSS_MODEL_TEST_CASE: Final = replace( + CLAUDE_OPUS_5_TEST_CASE, + mock_provider_response={ + **CLAUDE_OPUS_5_TEST_CASE.mock_provider_response, + "usage": CLAUDE_OPUS_5_AGENT_CROSS_MODEL_PROVIDER_USAGE, + }, + expected_litellm_response={ + **CLAUDE_OPUS_5_TEST_CASE.expected_litellm_response, + "usage": CLAUDE_OPUS_5_AGENT_CROSS_MODEL_CLIENT_USAGE, + }, ) CLAUDE_FABLE_5_1_AGENT_TEST_CASE: Final = replace( @@ -350,7 +422,7 @@ CLAUDE_FABLE_5_1_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_01324ezmdSmGSpmxGX55GvYF", "toolu_01MxHYehH7UYVY2A8WkRJZPu"), + ("toolu_01324ezmdSmGSpmxGX55GvYF", "toolu_01MxHYehH7UYVY2A8WkRJZPu") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -596,7 +668,7 @@ CLAUDE_FABLE_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_01UXBJZGdCaY4vxon1j4uapy", "toolu_0124RYbAPnwVoGtdzfPCNPc6"), + ("toolu_01UXBJZGdCaY4vxon1j4uapy", "toolu_0124RYbAPnwVoGtdzfPCNPc6") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -1494,7 +1566,7 @@ CLAUDE_HAIKU_5_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_0177HUb25ArysRAu8eF3u4DJ", "toolu_011gegGWvYoxPwxWY7LSK4Uk"), + ("toolu_0177HUb25ArysRAu8eF3u4DJ", "toolu_011gegGWvYoxPwxWY7LSK4Uk") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -1837,7 +1909,7 @@ CLAUDE_OPUS_5_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_013mPRCxrKKH6KzYGHKPukwv", "toolu_01FkXiaw4cTiXua3PoMd4wCu"), + ("toolu_013mPRCxrKKH6KzYGHKPukwv", "toolu_01FkXiaw4cTiXua3PoMd4wCu") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -2088,7 +2160,7 @@ CLAUDE_OPUS_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_017r3PoQsgodP2mgcgDnDLxU", "toolu_01T8ECEb4XwFyozSnCKUiuB2"), + ("toolu_017r3PoQsgodP2mgcgDnDLxU", "toolu_01T8ECEb4XwFyozSnCKUiuB2") ), expected_litellm_usage={ "prompt_tokens": 1180, diff --git a/tests/integration/translation/chat_completions/cross_model/case.py b/tests/integration/translation/chat_completions/cross_model/case.py index d0e842cfa67..07610b398e7 100644 --- a/tests/integration/translation/chat_completions/cross_model/case.py +++ b/tests/integration/translation/chat_completions/cross_model/case.py @@ -1,9 +1,6 @@ from typing import Final -from integration.translation.case import ( - StreamingTranslationTestCase, - TranslationTestCase, -) +from integration.translation.case import StreamingTranslationTestCase, TranslationTestCase from integration.translation.chat_completions.cross_model import ( anthropic, azure, @@ -105,7 +102,7 @@ MID: Final[tuple[AgentCases, ...]] = ( anthropic.CLAUDE_OPUS_5_5_AGENT_STREAMING_TEST_CASE, ), ( - anthropic.CLAUDE_OPUS_5_TEST_CASE, + anthropic.CLAUDE_OPUS_5_AGENT_CROSS_MODEL_TEST_CASE, anthropic.CLAUDE_OPUS_5_AGENT_FOLLOWUP_TEST_CASE, anthropic.CLAUDE_OPUS_5_AGENT_STREAMING_TEST_CASE, ), @@ -189,7 +186,7 @@ HIGH: Final[tuple[AgentCases, ...]] = ( anthropic.CLAUDE_FABLE_5_1_AGENT_STREAMING_TEST_CASE, ), ( - anthropic.CLAUDE_FABLE_5_TEST_CASE, + anthropic.CLAUDE_FABLE_5_AGENT_CROSS_MODEL_TEST_CASE, anthropic.CLAUDE_FABLE_5_AGENT_FOLLOWUP_TEST_CASE, anthropic.CLAUDE_FABLE_5_AGENT_STREAMING_TEST_CASE, ), diff --git a/tests/integration/translation/messages/bases/anthropic.py b/tests/integration/translation/messages/bases/anthropic.py index cc542499a07..be7f487f1b0 100644 --- a/tests/integration/translation/messages/bases/anthropic.py +++ b/tests/integration/translation/messages/bases/anthropic.py @@ -1,6 +1,7 @@ from collections.abc import Mapping from typing import Final +from integration._support.client import object_value from integration.translation.case import TranslationTestCase from integration.translation.messages.templates import MESSAGES_REQUEST, MESSAGES_RESPONSE from integration.translation.templates.anthropic import ANTHROPIC_MESSAGES_REQUEST, ANTHROPIC_MESSAGES_RESPONSE @@ -31,7 +32,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CffzUNHaEfzVxCh5hskBG", + "id": "msg_011Cg7v1w1xZexMrWpf5GU75", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -47,12 +48,13 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "anthropic/claude-sonnet-4-6", - "id": "msg_011CffzUNHaEfzVxCh5hskBG", + "id": "msg_011Cg7v1w1xZexMrWpf5GU75", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -68,6 +70,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -98,10 +101,10 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011CfgDz9Vj8wXvKzhoM4MQA", + "id": "msg_011Cg7v1Mt9Eid844XTV8AMc", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello."}], + "content": [{"type": "text", "text": "Hello!"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -114,15 +117,16 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "anthropic/claude-haiku-4-5", - "id": "msg_011CfgDz9Vj8wXvKzhoM4MQA", + "id": "msg_011Cg7v1Mt9Eid844XTV8AMc", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello."}], + "content": [{"type": "text", "text": "Hello!"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -135,6 +139,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -165,7 +170,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-5", - "id": "msg_011CfgDzLXmJmnHRTpeBoDmZ", + "id": "msg_011Cg7v2KRMCpM2FsiVygFAS", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -182,12 +187,13 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "anthropic/claude-sonnet-5", - "id": "msg_011CfgDzLXmJmnHRTpeBoDmZ", + "id": "msg_011Cg7v2KRMCpM2FsiVygFAS", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -204,6 +210,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -234,7 +241,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfgDzSFYYPC3aYE9rH98J", + "id": "msg_011Cg7v1ZPwvewtTtXqjdwzZ", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -251,12 +258,13 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "anthropic/claude-opus-4-8", - "id": "msg_011CfgDzSFYYPC3aYE9rH98J", + "id": "msg_011Cg7v1ZPwvewtTtXqjdwzZ", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -273,6 +281,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -377,11 +386,13 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_SONNET_5_5_RESPONSE_FIELDS, "model": "claude-sonnet-5-5", + "usage": {**object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_SONNET_5_5_RESPONSE_FIELDS, "model": "anthropic/claude-sonnet-5-5", + "usage": {**object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) @@ -460,11 +471,13 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_HAIKU_5_5_RESPONSE_FIELDS, "model": "claude-haiku-5-5", + "usage": {**object_value(CLAUDE_HAIKU_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_HAIKU_5_5_RESPONSE_FIELDS, "model": "anthropic/claude-haiku-5-5", + "usage": {**object_value(CLAUDE_HAIKU_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) @@ -555,11 +568,13 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_OPUS_5_5_RESPONSE_FIELDS, "model": "claude-opus-5-5", + "usage": {**object_value(CLAUDE_OPUS_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_OPUS_5_5_RESPONSE_FIELDS, "model": "anthropic/claude-opus-5-5", + "usage": {**object_value(CLAUDE_OPUS_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) @@ -651,10 +666,12 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_FABLE_5_1_RESPONSE_FIELDS, "model": "claude-fable-5-1", + "usage": {**object_value(CLAUDE_FABLE_5_1_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_FABLE_5_1_RESPONSE_FIELDS, "model": "anthropic/claude-fable-5-1", + "usage": {**object_value(CLAUDE_FABLE_5_1_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) diff --git a/tests/integration/translation/messages/cross_model/test_messages_cross_model.py b/tests/integration/translation/messages/cross_model/test_messages_cross_model.py index ae6fe70fe7c..e3bafc37993 100644 --- a/tests/integration/translation/messages/cross_model/test_messages_cross_model.py +++ b/tests/integration/translation/messages/cross_model/test_messages_cross_model.py @@ -3,11 +3,14 @@ from dataclasses import replace from typing import Final import pytest -from integration._support.client import JSON_OBJECT, Gateway +from integration._support.client import JSON_OBJECT, Gateway, object_value from integration._support.provider import SharedProvider from integration._support.wire import Reply from integration.translation.case import TranslationTestCase -from integration.translation.messages.bases.anthropic import CLAUDE_SONNET_5_5_TEST_CASE +from integration.translation.messages.bases.anthropic import ( + CLAUDE_SONNET_5_5_RESPONSE_FIELDS, + CLAUDE_SONNET_5_5_TEST_CASE, +) from integration.translation.messages.bases.azure_ai import ( CLAUDE_SONNET_5_5_TEST_CASE as AZURE_CLAUDE_SONNET_5_5_TEST_CASE, ) @@ -17,6 +20,18 @@ from integration.translation.response_scenarios import text_completion, two_mess from integration.translation.runner import assert_provider_request, assert_translation GPT_6_1_SOL_TWO_MESSAGES_BEFORE_TOOLS_TEST_CASE: Final = two_messages_before_tools(GPT_6_1_SOL_TEST_CASE) +CLAUDE_SONNET_5_5_CROSS_MODEL_NATIVE_TEST_CASE: Final = replace( + CLAUDE_SONNET_5_5_TEST_CASE, + scenario="cross_model_rich_request", + mock_provider_response={ + **CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, + "usage": object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), + }, + expected_litellm_response={ + **CLAUDE_SONNET_5_5_TEST_CASE.expected_litellm_response, + "usage": object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), + }, +) AZURE_CLAUDE_SONNET_5_5_CROSS_MODEL_TEST_CASE: Final = replace( AZURE_CLAUDE_SONNET_5_5_TEST_CASE, scenario="cross_model_rich_request", @@ -25,10 +40,14 @@ AZURE_CLAUDE_SONNET_5_5_CROSS_MODEL_TEST_CASE: Final = replace( **CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_request, "system": [{"type": "text", "text": CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_request["system"]}], }, - mock_provider_response=CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, + mock_provider_response={ + **CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, + "usage": object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), + }, expected_litellm_response={ **CLAUDE_SONNET_5_5_TEST_CASE.expected_litellm_response, "model": "azure_ai/claude-sonnet-5-5", + "usage": object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), }, ) @@ -44,7 +63,7 @@ def test_messages_keep_the_same_public_contract_on_native_and_bridged_routes( assert {key: value for key, value in case.litellm_request.items() if key != "model"} == { key: value for key, value in CLAUDE_SONNET_5_5_TEST_CASE.litellm_request.items() if key != "model" } - native: Final = assert_translation(CLAUDE_SONNET_5_5_TEST_CASE, gateway, provider) + native: Final = assert_translation(CLAUDE_SONNET_5_5_CROSS_MODEL_NATIVE_TEST_CASE, gateway, provider) compared: Final = assert_translation(case, gateway, provider) assert messages_structure(compared) == messages_structure(native) diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py index 5a3a760ffdc..6d7959b5de5 100644 --- a/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py +++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py @@ -8,6 +8,19 @@ from integration.translation.case import TranslationTestCase from integration.translation.messages.bases.anthropic import CLAUDE_SONNET_4_6_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE from integration.translation.runner import assert_translation +CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS: Final = { + "content": [{"type": "text", "text": "Hello!"}], + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 5, + "service_tier": "standard", + "inference_geo": "global", + }, +} + CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( CLAUDE_SONNET_4_6_TEST_CASE, scenario="service_tier", @@ -16,9 +29,14 @@ CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( **CLAUDE_SONNET_4_6_TEST_CASE.expected_provider_request, "service_tier": "standard_only", }, - mock_provider_response={**CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, "id": "msg_011Cg1RG36pCcTLE26htFFqC"}, + mock_provider_response={ + **CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, + **CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS, + "id": "msg_011Cg1RG36pCcTLE26htFFqC", + }, expected_litellm_response={ **CLAUDE_SONNET_4_6_TEST_CASE.expected_litellm_response, + **CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS, "id": "msg_011Cg1RG36pCcTLE26htFFqC", }, ) diff --git a/tests/integration/translation/responses/bases/anthropic.py b/tests/integration/translation/responses/bases/anthropic.py index e30c1f15699..8d5051acb92 100644 --- a/tests/integration/translation/responses/bases/anthropic.py +++ b/tests/integration/translation/responses/bases/anthropic.py @@ -35,7 +35,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011CfgDzBnuPYc7TdysZFRpo", + "id": "msg_011Cg7v1CTrRytLeZfaZojBR", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -51,6 +51,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -130,10 +131,10 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfgDzH3AcaSNXtvurLgAG", + "id": "msg_011Cg7v1kr1sYXyQA143nsda", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello!"}], + "content": [{"type": "text", "text": "Hello! How can I help you?"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -143,9 +144,10 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - "output_tokens": 5, + "output_tokens": 11, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -164,7 +166,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "id": ANY, "status": "completed", "role": "assistant", - "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], + "content": [{"type": "output_text", "text": "Hello! How can I help you?", "annotations": []}], "phase": None, } ], @@ -190,9 +192,9 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "video_tokens": None, "cache_write_tokens": 0, }, - "output_tokens": 5, - "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 0, "text_tokens": 5}, - "total_tokens": 23, + "output_tokens": 11, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 0, "text_tokens": 11}, + "total_tokens": 29, "cost": None, }, "user": None, @@ -225,7 +227,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-5", - "id": "msg_011CfgDzPkj2HxBq8qGJsGSi", + "id": "msg_011Cg7v215mkRTDBRoPcerLf", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -242,6 +244,7 @@ CLAUDE_SONNET_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -321,7 +324,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfgDzV9vZCkuQJt15n3Z1", + "id": "msg_011Cg7v1PgnYbGrw8eGpka6f", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -338,6 +341,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -462,6 +466,7 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 313}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={**RESPONSES_RESPONSE, "model": "anthropic/claude-sonnet-5-5"}, @@ -534,6 +539,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 67}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -553,7 +559,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "requested.\n" "\n", "annotations": [], - }, + } ], "phase": None, "encrypted_content": json.dumps( @@ -565,7 +571,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "requested.\n" "\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], separators=(",", ":"), ), @@ -605,7 +611,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "range near the top of the file since no specific lines were requested.\n" "\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, @@ -698,6 +704,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 130}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -719,7 +726,7 @@ CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( "zero, negative, and non-integer inputs before moving to the test file.\n" "\n", "annotations": [], - }, + } ], "phase": None, "encrypted_content": json.dumps( @@ -842,6 +849,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -860,7 +868,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "rejection, then I'll read the test file in a separate step since there's " "nothing to derive from it for the read.", "annotations": [], - }, + } ], "phase": None, }, diff --git a/tests/integration/translation/responses/cross_model/test_responses_cross_model.py b/tests/integration/translation/responses/cross_model/test_responses_cross_model.py index cb6f973d2df..a0b5778a200 100644 --- a/tests/integration/translation/responses/cross_model/test_responses_cross_model.py +++ b/tests/integration/translation/responses/cross_model/test_responses_cross_model.py @@ -22,12 +22,25 @@ from integration.translation.responses.bases.openai_responses import GPT_5_6_SOL from integration.translation.runner import assert_translation GPT_6_1_SOL_TWO_MESSAGES_BEFORE_TOOLS_TEST_CASE: Final = two_messages_before_tools(GPT_6_1_SOL_TEST_CASE) +CLAUDE_SONNET_5_5_CROSS_MODEL_PROVIDER_USAGE: Final = { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 509, + "output_tokens_details": {"thinking_tokens": 313}, + "service_tier": "standard", + "inference_geo": "global", +} AZURE_CLAUDE_SONNET_5_5_CROSS_MODEL_TEST_CASE: Final = replace( AZURE_CLAUDE_SONNET_5_5_TEST_CASE, scenario="cross_model_rich_request", litellm_request={**CLAUDE_SONNET_5_5_TEST_CASE.litellm_request, "model": "azure_ai/claude-sonnet-5-5"}, expected_provider_request=CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_request, - mock_provider_response=CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, + mock_provider_response={ + **CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, + "usage": CLAUDE_SONNET_5_5_CROSS_MODEL_PROVIDER_USAGE, + }, expected_litellm_response={ **CLAUDE_SONNET_5_5_TEST_CASE.expected_litellm_response, "model": "azure_ai/claude-sonnet-5-5", From 0c3087eb202145c330990d48e58fc2012779cbae Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:30:34 -0700 Subject: [PATCH 04/64] test(translation): recapture azure ai claude bases after usage speed field (#47117) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_completions/bases/azure_ai.py | 45 +++++++++++------- .../chat_completions/cross_model/azure_ai.py | 47 +++++++++++++++---- .../chat_completions/cross_model/case.py | 2 +- .../translation/messages/bases/azure_ai.py | 27 +++++++---- .../test_messages_service_tier_azure_ai.py | 20 +++++++- .../translation/responses/bases/azure_ai.py | 25 ++++++---- 6 files changed, 119 insertions(+), 47 deletions(-) diff --git a/tests/integration/translation/chat_completions/bases/azure_ai.py b/tests/integration/translation/chat_completions/bases/azure_ai.py index e3e8a84340b..89c75f3c597 100644 --- a/tests/integration/translation/chat_completions/bases/azure_ai.py +++ b/tests/integration/translation/chat_completions/bases/azure_ai.py @@ -43,7 +43,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011Cfjax6mBrNVhFGQopmVCp", + "id": "msg_011Cg7uyAiPY51HN4rSYzQW5", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -59,6 +59,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -94,6 +95,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -144,6 +146,7 @@ CLAUDE_FABLE_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -193,6 +196,7 @@ CLAUDE_FABLE_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -224,10 +228,10 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfjaxGKedCbjtb7HuAUBC", + "id": "msg_011Cg7uzQ1rtYRwCd9aeTyoF", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello!"}], + "content": [{"type": "text", "text": "Hello! How can I help you?"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -237,9 +241,10 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - "output_tokens": 5, + "output_tokens": 11, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -253,17 +258,17 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "finish_reason": "stop", "index": 0, "message": { - "content": "Hello!", + "content": "Hello! How can I help you?", "role": "assistant", "provider_specific_fields": {"citations": None, "thinking_blocks": None}, }, } ], "usage": { - "completion_tokens": 5, + "completion_tokens": 11, "prompt_tokens": 18, - "total_tokens": 23, - "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 5}, + "total_tokens": 29, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 11}, "prompt_tokens_details": { "cached_tokens": 0, "text_tokens": 18, @@ -275,6 +280,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -389,7 +395,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfjaxUiYCUsqMC4iqfc2U", + "id": "msg_011Cg7uyf9RNvQWkpJobt6JU", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -406,6 +412,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -441,6 +448,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -674,6 +682,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 167}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -697,7 +706,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( }, "id": "toolu_01M68cqREYA4WqXBUK8rqcEq", "type": "function", - }, + } ], "reasoning_content": "I need to read shop/validation.py next, keeping in mind edge cases like zero, " "negative, or non-integer quantities since the validation logic likely handles those. " @@ -709,7 +718,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "negative, or non-integer quantities since the validation logic likely handles those. " "I'll look at the function starting around line 18 to understand its scope.\n\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], "provider_specific_fields": { "citations": None, @@ -720,11 +729,11 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "negative, or non-integer quantities since the validation logic likely handles those. " "I'll look at the function starting around line 18 to understand its scope.\n\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, - }, + } ], "usage": { "completion_tokens": 324, @@ -742,6 +751,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "global", "service_tier": "standard", + "speed": "standard", }, }, ) @@ -763,10 +773,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "model": "claude-fable-5-1", "id": "msg_011Cg37NcfzdXibtC3dWhVWi", "content": [ - { - "type": "text", - "text": "I'll read the validation code and the test together since they are independent.", - }, + {"type": "text", "text": "I'll read the validation code and the test together since they are independent."}, { "type": "tool_use", "id": "toolu_01LxkeKf9WX9BVpiWrvfKUyc", @@ -791,6 +798,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, }, expected_litellm_response={ @@ -827,7 +835,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( ], "provider_specific_fields": {"citations": None, "thinking_blocks": None}, }, - }, + } ], "usage": { "completion_tokens": 216, @@ -845,6 +853,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "cache_read_input_tokens": 0, "inference_geo": "not_available", "service_tier": "standard", + "speed": "standard", }, }, ) diff --git a/tests/integration/translation/chat_completions/cross_model/azure_ai.py b/tests/integration/translation/chat_completions/cross_model/azure_ai.py index 72c14ef103c..15cbe4a3200 100644 --- a/tests/integration/translation/chat_completions/cross_model/azure_ai.py +++ b/tests/integration/translation/chat_completions/cross_model/azure_ai.py @@ -2,10 +2,7 @@ from dataclasses import replace from typing import Final from integration._support.client import list_value, object_value -from integration.translation.case import ( - SSEFrame, - StreamingTranslationTestCase, -) +from integration.translation.case import SSEFrame, StreamingTranslationTestCase from integration.translation.chat_completions.cross_model.contracts import followup_request from integration.translation.chat_completions.bases.azure_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, @@ -18,9 +15,41 @@ from integration.translation.chat_completions.templates import ( coding_agent_stream_message, CODING_AGENT_RESPONSE, ) -from integration.translation.templates.anthropic import ( - ANTHROPIC_AGENT_REQUEST, - ANTHROPIC_AGENT_RESPONSE, +from integration.translation.templates.anthropic import ANTHROPIC_AGENT_REQUEST, ANTHROPIC_AGENT_RESPONSE + +CLAUDE_FABLE_5_CROSS_MODEL_TEST_CASE: Final = replace( + CLAUDE_FABLE_5_TEST_CASE, + mock_provider_response={ + **CLAUDE_FABLE_5_TEST_CASE.mock_provider_response, + "usage": { + **object_value(ANTHROPIC_AGENT_RESPONSE["usage"]), + "input_tokens": 1184, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "not_available", + }, + }, + expected_litellm_response={ + **CLAUDE_FABLE_5_TEST_CASE.expected_litellm_response, + "usage": { + "completion_tokens": 105, + "prompt_tokens": 1184, + "total_tokens": 1289, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1184, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "inference_geo": "not_available", + "service_tier": "standard", + }, + }, ) CLAUDE_FABLE_5_1_AGENT_TEST_CASE: Final = replace( @@ -349,7 +378,7 @@ CLAUDE_FABLE_5_1_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_01PfMSQLj6cQzebc2vGbdsV8", "toolu_01SVc4RwmEaMiDK33hUEbTZW"), + ("toolu_01PfMSQLj6cQzebc2vGbdsV8", "toolu_01SVc4RwmEaMiDK33hUEbTZW") ), expected_litellm_usage={ "prompt_tokens": 1184, @@ -593,7 +622,7 @@ CLAUDE_FABLE_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( SSEFrame(data={"type": "message_stop"}, event="message_stop"), ), expected_litellm_message=coding_agent_stream_message( - ("toolu_01DprHn2GeJDAWFEyTqBzy2Z", "toolu_016kava9hqtKw71W4vBEedfW"), + ("toolu_01DprHn2GeJDAWFEyTqBzy2Z", "toolu_016kava9hqtKw71W4vBEedfW") ), expected_litellm_usage={ "prompt_tokens": 1184, diff --git a/tests/integration/translation/chat_completions/cross_model/case.py b/tests/integration/translation/chat_completions/cross_model/case.py index 07610b398e7..dc93bd3178f 100644 --- a/tests/integration/translation/chat_completions/cross_model/case.py +++ b/tests/integration/translation/chat_completions/cross_model/case.py @@ -196,7 +196,7 @@ HIGH: Final[tuple[AgentCases, ...]] = ( azure_ai.CLAUDE_FABLE_5_1_AGENT_STREAMING_TEST_CASE, ), ( - azure_ai.CLAUDE_FABLE_5_TEST_CASE, + azure_ai.CLAUDE_FABLE_5_CROSS_MODEL_TEST_CASE, azure_ai.CLAUDE_FABLE_5_AGENT_FOLLOWUP_TEST_CASE, azure_ai.CLAUDE_FABLE_5_AGENT_STREAMING_TEST_CASE, ), diff --git a/tests/integration/translation/messages/bases/azure_ai.py b/tests/integration/translation/messages/bases/azure_ai.py index 454a76400db..6de15226c48 100644 --- a/tests/integration/translation/messages/bases/azure_ai.py +++ b/tests/integration/translation/messages/bases/azure_ai.py @@ -1,6 +1,7 @@ from collections.abc import Mapping from typing import Final +from integration._support.client import object_value from integration.translation.case import TranslationTestCase from integration.translation.messages.templates import MESSAGES_REQUEST, MESSAGES_RESPONSE from integration.translation.templates.anthropic import ANTHROPIC_MESSAGES_REQUEST, ANTHROPIC_MESSAGES_RESPONSE @@ -31,10 +32,10 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011Cfjax2zxeHkddgerkEGRY", + "id": "msg_011Cg7vB62SdiZVj3M7xJXpr", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello."}], + "content": [{"type": "text", "text": "Hello!"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -47,15 +48,16 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "azure_ai/claude-haiku-4-5", - "id": "msg_011Cfjax2zxeHkddgerkEGRY", + "id": "msg_011Cg7vB62SdiZVj3M7xJXpr", "type": "message", "role": "assistant", - "content": [{"type": "text", "text": "Hello."}], + "content": [{"type": "text", "text": "Hello!"}], "container": None, "stop_reason": "end_turn", "stop_sequence": None, @@ -68,6 +70,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -98,7 +101,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfjaxAknNm3ZLkVxQRj8M", + "id": "msg_011Cg7vCCFGaPEhwGwoS6VPy", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -114,12 +117,13 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "azure_ai/claude-sonnet-4-6", - "id": "msg_011CfjaxAknNm3ZLkVxQRj8M", + "id": "msg_011Cg7vCCFGaPEhwGwoS6VPy", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -135,6 +139,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -234,7 +239,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfjaxPE8FBKSJ9KbtQzDE", + "id": "msg_011Cg7vBVpRK6kd1kkYXJXyb", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -251,12 +256,13 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, expected_litellm_response={ "model": "azure_ai/claude-opus-4-8", - "id": "msg_011CfjaxPE8FBKSJ9KbtQzDE", + "id": "msg_011Cg7vBVpRK6kd1kkYXJXyb", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -273,6 +279,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -483,11 +490,13 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_HAIKU_5_5_RESPONSE_FIELDS, "model": "claude-haiku-5-5", + "usage": {**object_value(CLAUDE_HAIKU_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_HAIKU_5_5_RESPONSE_FIELDS, "model": "azure_ai/claude-haiku-5-5", + "usage": {**object_value(CLAUDE_HAIKU_5_5_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) @@ -541,10 +550,12 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( **ANTHROPIC_MESSAGES_RESPONSE, **CLAUDE_FABLE_5_1_RESPONSE_FIELDS, "model": "claude-fable-5-1", + "usage": {**object_value(CLAUDE_FABLE_5_1_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, expected_litellm_response={ **MESSAGES_RESPONSE, **CLAUDE_FABLE_5_1_RESPONSE_FIELDS, "model": "azure_ai/claude-fable-5-1", + "usage": {**object_value(CLAUDE_FABLE_5_1_RESPONSE_FIELDS["usage"]), "speed": "standard"}, }, ) diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py index 13bf9dd888f..f987f9d1f2a 100644 --- a/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py +++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py @@ -8,6 +8,19 @@ from integration.translation.case import TranslationTestCase from integration.translation.messages.bases.azure_ai import CLAUDE_SONNET_4_6_TEST_CASE from integration.translation.runner import assert_translation +CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS: Final = { + "content": [{"type": "text", "text": "Hello!"}], + "usage": { + "input_tokens": 18, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 5, + "service_tier": "standard", + "inference_geo": "not_available", + }, +} + CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( CLAUDE_SONNET_4_6_TEST_CASE, scenario="service_tier", @@ -16,9 +29,14 @@ CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace( **CLAUDE_SONNET_4_6_TEST_CASE.expected_provider_request, "service_tier": "standard_only", }, - mock_provider_response={**CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, "id": "msg_011Cg1SnrstCdzZPwmTk2qQa"}, + mock_provider_response={ + **CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, + **CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS, + "id": "msg_011Cg1SnrstCdzZPwmTk2qQa", + }, expected_litellm_response={ **CLAUDE_SONNET_4_6_TEST_CASE.expected_litellm_response, + **CLAUDE_SONNET_4_6_SERVICE_TIER_RESPONSE_FIELDS, "id": "msg_011Cg1SnrstCdzZPwmTk2qQa", }, ) diff --git a/tests/integration/translation/responses/bases/azure_ai.py b/tests/integration/translation/responses/bases/azure_ai.py index 921fed14b45..dfcb454c482 100644 --- a/tests/integration/translation/responses/bases/azure_ai.py +++ b/tests/integration/translation/responses/bases/azure_ai.py @@ -31,7 +31,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-haiku-4-5-20251001", - "id": "msg_011Cfjax6mBrNVhFGQopmVCp", + "id": "msg_011Cg7vB393hLWW1Xxq6d6vm", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -47,6 +47,7 @@ CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -127,7 +128,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-sonnet-4-6", - "id": "msg_011CfjaxGKedCbjtb7HuAUBC", + "id": "msg_011Cg7vCAx8gxwjwBr2R1kxz", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], @@ -143,6 +144,7 @@ CLAUDE_SONNET_4_6_TEST_CASE: Final = TranslationTestCase( "output_tokens": 5, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, "diagnostics": None, }, @@ -320,7 +322,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( }, mock_provider_response={ "model": "claude-opus-4-8", - "id": "msg_011CfjaxUiYCUsqMC4iqfc2U", + "id": "msg_011Cg7vBUGeH2HFZFTDxVdBh", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello."}], @@ -337,6 +339,7 @@ CLAUDE_OPUS_4_8_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, "diagnostics": None, }, @@ -620,7 +623,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "suggesting a fix. Explain which file you need next, then call a read-only tool. " "Never change files. Use the responses client for this project review. Consider zero, " "negative and non-integer quantities before choosing the next file to inspect.", - }, + } ], }, mock_provider_response={ @@ -656,6 +659,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 73}, "service_tier": "standard", "inference_geo": "global", + "speed": "standard", }, }, expected_litellm_response={ @@ -673,7 +677,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "text": "I need to look at shop/validation.py to understand its contents, so I'll read " "through the first chunk of the file.\n\n", "annotations": [], - }, + } ], "phase": None, "encrypted_content": json.dumps( @@ -683,7 +687,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "thinking": "I need to look at shop/validation.py to understand its contents, so I'll read " "through the first chunk of the file.\n\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], separators=(",", ":"), ), @@ -699,7 +703,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "text": "I'll read `shop/validation.py`, starting with the `validate_quantity` function found " "in the search.", "annotations": [], - }, + } ], "phase": None, }, @@ -737,7 +741,7 @@ CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( "thinking": "I need to look at shop/validation.py to understand its contents, so I'll read " "through the first chunk of the file.\n\n", "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE_0, - }, + } ], }, }, @@ -764,7 +768,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "suggesting a fix. Explain which file you need next, then call a read-only tool. " "Never change files. Use the responses client for this project review. Consider zero, " "negative and non-integer quantities before choosing the next file to inspect.", - }, + } ], }, mock_provider_response={ @@ -801,6 +805,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "output_tokens_details": {"thinking_tokens": 0}, "service_tier": "standard", "inference_geo": "not_available", + "speed": "standard", }, }, expected_litellm_response={ @@ -818,7 +823,7 @@ CLAUDE_FABLE_5_1_TEST_CASE: Final = TranslationTestCase( "text": "I'll read the validation function; reading the test file at the same time is " "independent of it, so I'll fetch both now.", "annotations": [], - }, + } ], "phase": None, }, From 5dcf2c2062d57746bbf1f70ceb6294737b86bd6d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:36:40 -0700 Subject: [PATCH 05/64] test(translation): add sonnet 5.5 and gemini vertex rows to cross-model mid tier (#46896) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_completions/cross_model/README.md | 6 +- .../chat_completions/cross_model/case.py | 15 + .../chat_completions/cross_model/vertex_ai.py | 1215 ++++++++++++++++- 3 files changed, 1231 insertions(+), 5 deletions(-) diff --git a/tests/integration/translation/chat_completions/cross_model/README.md b/tests/integration/translation/chat_completions/cross_model/README.md index 0adce683f38..ee7c178da81 100644 --- a/tests/integration/translation/chat_completions/cross_model/README.md +++ b/tests/integration/translation/chat_completions/cross_model/README.md @@ -5,10 +5,10 @@ Every row sends `CODING_AGENT_REQUEST` to `/v1/chat/completions`, changing only | Tier | Model versions | | --- | --- | | Small | Haiku 4.5 and 5.5, GPT 5 Mini and 5.4 Mini, GPT 5.6 Luna and 6 Luna | -| Mid | Opus 5 and 5.5, GPT 5.6 Sol, 6 Sol and 6.1 Sol, GPT 5.6 Terra | +| Mid | Opus 5 and 5.5, Sonnet 5.5, Gemini 3.8 Flash and 3.1 Pro Preview, GPT 5.6 Sol, 6 Sol and 6.1 Sol, GPT 5.6 Terra | | High | Fable 5 and 5.1, Gemini 4 Argon, GPT 6 Astra | -The 45 rows cover Anthropic direct, OpenAI direct, Vertex, Azure OpenAI, Azure AI and Bedrock Converse where accessible. Vertex Haiku 5.5 was unavailable during capture and is not represented +The 48 rows cover Anthropic direct, OpenAI direct, Vertex, Azure OpenAI, Azure AI and Bedrock Converse where accessible. Vertex Haiku 5.5 was unavailable during capture and is not represented The request contains system instructions, conversation history, previous assistant tool calls and results, four function definitions, automatic parallel tools, low reasoning effort and an output limit. Each row exercises a tool-call response, raw streaming, SDK stream assembly and a follow-up that replays the returned assistant message and tool-call IDs @@ -33,7 +33,7 @@ Streaming cases store `expected_litellm_message` and `expected_litellm_usage`. M Usage requires nonnegative integer counts, positive prompt and completion counts, and `total_tokens == prompt_tokens + completion_tokens`. The complete normalized usage must equal the independently expected counters, including cached and reasoning tokens and derived text/cache details. Missing, changed or unexpected counters fail. Only null optional fields, the explicitly listed provider metadata/duplicate counters and computed `cost` are excluded. Thus a wrong prompt count with an adjusted total still fails -The raw stream check assembles tool names and argument fragments by index. The SDK check separately uses `OpenAI.chat.completions.stream()` without removing reasoning fields. All 45 SDK checks run, including the seven Haiku rows. [LIT-9653](https://linear.app/litellm-ai/issue/LIT-9653) was fixed in [#46611](https://github.com/BerriAI/litellm/pull/46611), so there are no Haiku skips +The raw stream check assembles tool names and argument fragments by index. The SDK check separately uses `OpenAI.chat.completions.stream()` without removing reasoning fields. All 48 SDK checks run, including the seven Haiku rows. [LIT-9653](https://linear.app/litellm-ai/issue/LIT-9653) was fixed in [#46611](https://github.com/BerriAI/litellm/pull/46611), so there are no Haiku skips ## Recorded upstreams diff --git a/tests/integration/translation/chat_completions/cross_model/case.py b/tests/integration/translation/chat_completions/cross_model/case.py index dc93bd3178f..fae659c5dca 100644 --- a/tests/integration/translation/chat_completions/cross_model/case.py +++ b/tests/integration/translation/chat_completions/cross_model/case.py @@ -177,6 +177,21 @@ MID: Final[tuple[AgentCases, ...]] = ( openai.GPT_6_1_SOL_AGENT_FOLLOWUP_TEST_CASE, openai.GPT_6_1_SOL_AGENT_STREAMING_TEST_CASE, ), + ( + vertex_ai.CLAUDE_SONNET_5_5_AGENT_TEST_CASE, + vertex_ai.CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEST_CASE, + vertex_ai.CLAUDE_SONNET_5_5_AGENT_STREAMING_TEST_CASE, + ), + ( + vertex_ai.GEMINI_3_8_FLASH_AGENT_TEST_CASE, + vertex_ai.GEMINI_3_8_FLASH_AGENT_FOLLOWUP_TEST_CASE, + vertex_ai.GEMINI_3_8_FLASH_AGENT_STREAMING_TEST_CASE, + ), + ( + vertex_ai.GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE, + vertex_ai.GEMINI_3_1_PRO_PREVIEW_AGENT_FOLLOWUP_TEST_CASE, + vertex_ai.GEMINI_3_1_PRO_PREVIEW_AGENT_STREAMING_TEST_CASE, + ), ) HIGH: Final[tuple[AgentCases, ...]] = ( diff --git a/tests/integration/translation/chat_completions/cross_model/vertex_ai.py b/tests/integration/translation/chat_completions/cross_model/vertex_ai.py index 42f363b1920..0dbb4b277c5 100644 --- a/tests/integration/translation/chat_completions/cross_model/vertex_ai.py +++ b/tests/integration/translation/chat_completions/cross_model/vertex_ai.py @@ -6,19 +6,22 @@ from integration.translation.case import ( SSEFrame, StreamingTranslationTestCase, ) -from integration.translation.chat_completions.cross_model.contracts import followup_request from integration.translation.chat_completions.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_FABLE_5_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_OPUS_5_TEST_CASE, + CLAUDE_SONNET_5_5_TEST_CASE, + GEMINI_3_1_PRO_PREVIEW_TEST_CASE, + GEMINI_3_8_FLASH_TEST_CASE, GEMINI_4_ARGON_TEST_CASE, ) +from integration.translation.chat_completions.cross_model.contracts import followup_request from integration.translation.chat_completions.templates import ( CODING_AGENT_REQUEST, - coding_agent_stream_message, CODING_AGENT_RESPONSE, + coding_agent_stream_message, ) from integration.translation.templates.anthropic import ( ANTHROPIC_AGENT_REQUEST, @@ -1721,6 +1724,1214 @@ CLAUDE_OPUS_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( ) +CLAUDE_SONNET_5_5_AGENT_TEST_CASE: Final = replace( + CLAUDE_SONNET_5_5_TEST_CASE, + scenario="agent_tools", + litellm_request={**CODING_AGENT_REQUEST, "model": "vertex_ai/claude-sonnet-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-sonnet-5-5@default:rawPredict", + expected_provider_headers={ + "authorization": "Bearer scripted-token", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "x-api-key": "scripted-token", + }, + expected_provider_request={ + **ANTHROPIC_AGENT_REQUEST, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "low"}, + "anthropic_version": "vertex-2023-10-16", + }, + mock_provider_response={ + **ANTHROPIC_AGENT_RESPONSE, + "model": "claude-sonnet-5-5", + "id": "msg_vrtx_011Cg63qBoZy88hTvxUdtSEr", + "content": [ + { + "type": "tool_use", + "id": "toolu_vrtx_01AtmBHyfEtHxHSJs8imPFWJ", + "name": "search_code", + "input": {"query": "retries_remaining"}, + }, + { + "type": "tool_use", + "id": "toolu_vrtx_019seNAPMm8gmCrNPf8Fwzvf", + "name": "run_tests", + "input": {"path": "tests/test_retry.py"}, + }, + ], + "stop_reason": "tool_use", + "usage": { + "input_tokens": 1184, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "model": "vertex_ai/claude-sonnet-5-5@default", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "function": {"arguments": '{"query": "retries_remaining"}', "name": "search_code"}, + "id": "toolu_vrtx_01AtmBHyfEtHxHSJs8imPFWJ", + "type": "function", + }, + { + "index": 1, + "function": {"arguments": '{"path": "tests/test_retry.py"}', "name": "run_tests"}, + "id": "toolu_vrtx_019seNAPMm8gmCrNPf8Fwzvf", + "type": "function", + }, + ], + "provider_specific_fields": {"citations": None, "thinking_blocks": None}, + "content": None, + }, + } + ], + "usage": { + "completion_tokens": 105, + "prompt_tokens": 1184, + "total_tokens": 1289, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1184, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, + }, +) + + +CLAUDE_SONNET_5_5_AGENT_FOLLOWUP_TEST_CASE: Final = replace( + CLAUDE_SONNET_5_5_AGENT_TEST_CASE, + scenario="agent_followup", + litellm_request=followup_request( + CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_request, CLAUDE_SONNET_5_5_AGENT_TEST_CASE.expected_litellm_response + ), + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-sonnet-5-5@default:rawPredict", + expected_provider_request={ + **ANTHROPIC_AGENT_REQUEST, + "messages": [ + *list_value(ANTHROPIC_AGENT_REQUEST["messages"]), + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_vrtx_01AtmBHyfEtHxHSJs8imPFWJ", + "name": "search_code", + "input": {"query": "retries_remaining"}, + }, + { + "type": "tool_use", + "id": "toolu_vrtx_019seNAPMm8gmCrNPf8Fwzvf", + "name": "run_tests", + "input": {"path": "tests/test_retry.py"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_vrtx_01AtmBHyfEtHxHSJs8imPFWJ", + "content": '{"matches": [{"path": "src/worker.py", "line": 12, "text": "remaining = ' + 'retries_remaining(3, attempts)"}]}', + }, + { + "type": "tool_result", + "tool_use_id": "toolu_vrtx_019seNAPMm8gmCrNPf8Fwzvf", + "content": '{"exit_code": 1, "output": "FAILED tests/test_retry.py::test_retries_remaining - assert ' + '4 == 2"}', + }, + { + "type": "text", + "text": "Stop after these diagnostics. Do not call more tools or change files. Summarize the " + "bug, the caller found and the failed test in two short sentences.", + }, + ], + }, + ], + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "low"}, + "anthropic_version": "vertex-2023-10-16", + }, + mock_provider_response={ + **ANTHROPIC_AGENT_RESPONSE, + "model": "claude-sonnet-5-5", + "id": "msg_vrtx_011Cg63wSU3MbRaa6kwyvRLk", + "content": [ + { + "type": "text", + "text": "The bug is that `retries_remaining` in `src/retry.py` returns `limit + attempts` " + "instead of `max(0, limit - attempts)`, so it overcounts and can't be clamped at zero. " + "The one other caller is `src/worker.py` line 12 (`retries_remaining(3, attempts)`), and " + "`tests/test_retry.py::test_retries_remaining` fails with `assert 4 == 2`.", + } + ], + "stop_reason": "end_turn", + "usage": { + "input_tokens": 1492, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 138, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "model": "vertex_ai/claude-sonnet-5-5@default", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "The bug is that `retries_remaining` in `src/retry.py` returns `limit + attempts` " + "instead of `max(0, limit - attempts)`, so it overcounts and can't be clamped at zero. " + "The one other caller is `src/worker.py` line 12 (`retries_remaining(3, attempts)`), and " + "`tests/test_retry.py::test_retries_remaining` fails with `assert 4 == 2`.", + "role": "assistant", + "provider_specific_fields": {"citations": None, "thinking_blocks": None}, + }, + } + ], + "usage": { + "completion_tokens": 138, + "prompt_tokens": 1492, + "total_tokens": 1630, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 138}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1492, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, + }, +) + +CLAUDE_SONNET_5_5_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="agent_streaming", + litellm_endpoint=CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_endpoint, + litellm_request={ + **CLAUDE_SONNET_5_5_AGENT_TEST_CASE.litellm_request, + "stream": True, + "stream_options": {"include_usage": True}, + }, + expected_provider_headers=CLAUDE_SONNET_5_5_AGENT_TEST_CASE.expected_provider_headers, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-sonnet-5-5@default:streamRawPredict?alt=sse", + expected_provider_request={ + **ANTHROPIC_AGENT_REQUEST, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "low"}, + "anthropic_version": "vertex-2023-10-16", + "stream": True, + }, + mock_provider_frames=( + SSEFrame( + data={ + "type": "message_start", + "message": { + "model": "claude-sonnet-5-5", + "id": "msg_vrtx_011Cg63qMXiZfvJr7NuvT9B1", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1184, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 16, + }, + }, + }, + event="message_start", + ), + SSEFrame(data={"type": "ping"}, event="ping"), + SSEFrame( + data={ + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "toolu_vrtx_01EJMuteinVLQX3UQ8bxJEps", + "name": "search_code", + "input": {}, + }, + }, + event="content_block_start", + ), + SSEFrame( + data={"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query'}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '": "ret'}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": "ries_re"}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": "maining"}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '"}'}, + }, + event="content_block_delta", + ), + SSEFrame(data={"type": "content_block_stop", "index": 0}, event="content_block_stop"), + SSEFrame( + data={ + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "tool_use", + "id": "toolu_vrtx_015Mb3cdfYTHCA9QowGSX2RP", + "name": "run_tests", + "input": {}, + }, + }, + event="content_block_start", + ), + SSEFrame( + data={"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": ""}}, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": '{"path"'}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": ': "test'}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": "s/test_"}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": "retry.p"}, + }, + event="content_block_delta", + ), + SSEFrame( + data={ + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": 'y"}'}, + }, + event="content_block_delta", + ), + SSEFrame(data={"type": "content_block_stop", "index": 1}, event="content_block_stop"), + SSEFrame( + data={ + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1184, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 105, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + event="message_delta", + ), + SSEFrame(data={"type": "message_stop"}, event="message_stop"), + ), + expected_litellm_message=coding_agent_stream_message( + ("toolu_vrtx_01EJMuteinVLQX3UQ8bxJEps", "toolu_vrtx_015Mb3cdfYTHCA9QowGSX2RP") + ), + expected_litellm_usage={ + "prompt_tokens": 1184, + "completion_tokens": 105, + "total_tokens": 1289, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 1184, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, + }, +) + + +GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE: Final = replace( + GEMINI_3_1_PRO_PREVIEW_TEST_CASE, + scenario="agent_tools", + litellm_request={**CODING_AGENT_REQUEST, "model": "vertex_ai/gemini-3.1-pro-preview"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:generateContent", + expected_provider_headers={"content-type": "application/json", "authorization": "Bearer scripted-token"}, + expected_provider_request={**VERTEX_AI_AGENT_REQUEST}, + mock_provider_response={ + **VERTEX_AI_AGENT_RESPONSE, + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "**My Thought Process on Addressing the `retries_remaining` Issue**\n" + "\n" + "Alright, so my immediate priority is to tackle the `retries_remaining` function. The " + "core of the problem, as I see it, lies in its current implementation where it seems to " + "be adding `limit` and `attempts` rather than performing the correct subtraction to " + "determine remaining retries. This, of course, is a fundamental logic error. " + "Furthermore, I've identified a critical oversight: the function doesn't account for " + "scenarios where the number of `attempts` might actually exceed the `limit`. This is a " + "potential edge case that could lead to unexpected behavior and incorrect results, so it " + "absolutely needs to be addressed.\n" + "\n" + "My proposed solution is a straightforward fix: `return max(0, limit - attempts)`. This " + "ensures that we're correctly calculating the difference between the allowed `limit` and " + "the current `attempts`. Crucially, by wrapping this subtraction within `max(0, ...)`, " + "I'm enforcing a floor of zero. This elegantly handles the situation where `attempts` " + "might be greater than or equal to `limit`, ensuring that the function will never return " + "a negative number of remaining retries, which would be nonsensical in this context.\n" + "\n" + "With that fix in mind, my plan is to proceed with the original request: I will first " + "execute the tests, specifically `tests/test_retry.py`, and run them concurrently for " + "efficiency. Immediately after, I'll perform a thorough search for all call sites of the " + "`retries_remaining` function. This two-pronged approach will allow me to verify the " + "correctness of my proposed fix and understand the broader impact of this change across " + "the codebase.\n" + "\n", + "thought": True, + }, + { + "functionCall": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_2861166", + }, + "thoughtSignature": "synthetic-signature", + }, + { + "functionCall": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_2861167", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 404, + "candidatesTokenCount": 41, + "totalTokenCount": 551, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 404}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 41}], + "thoughtsTokenCount": 106, + }, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:10:08.649083Z", + "responseId": "QDDSavvOJ6-q_9MP5byA4Qo", + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "id": "QDDSavvOJ6-q_9MP5byA4Qo", + "model": "vertex_ai/gemini-3.1-pro-preview", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "provider_specific_fields": {"thought_signature": "synthetic-signature"}, + "function": {"arguments": '{"path": "tests/test_retry.py"}', "name": "run_tests"}, + "id": "call_2861166__thought__synthetic-signature", + "type": "function", + }, + { + "index": 1, + "function": {"arguments": '{"query": "retries_remaining"}', "name": "search_code"}, + "id": "call_2861167", + "type": "function", + }, + ], + "images": [], + "reasoning_content": "**My Thought Process on Addressing the `retries_remaining` Issue**\n" + "\n" + "Alright, so my immediate priority is to tackle the `retries_remaining` function. The " + "core of the problem, as I see it, lies in its current implementation where it seems to " + "be adding `limit` and `attempts` rather than performing the correct subtraction to " + "determine remaining retries. This, of course, is a fundamental logic error. " + "Furthermore, I've identified a critical oversight: the function doesn't account for " + "scenarios where the number of `attempts` might actually exceed the `limit`. This is a " + "potential edge case that could lead to unexpected behavior and incorrect results, so it " + "absolutely needs to be addressed.\n" + "\n" + "My proposed solution is a straightforward fix: `return max(0, limit - attempts)`. This " + "ensures that we're correctly calculating the difference between the allowed `limit` and " + "the current `attempts`. Crucially, by wrapping this subtraction within `max(0, ...)`, " + "I'm enforcing a floor of zero. This elegantly handles the situation where `attempts` " + "might be greater than or equal to `limit`, ensuring that the function will never return " + "a negative number of remaining retries, which would be nonsensical in this context.\n" + "\n" + "With that fix in mind, my plan is to proceed with the original request: I will first " + "execute the tests, specifically `tests/test_retry.py`, and run them concurrently for " + "efficiency. Immediately after, I'll perform a thorough search for all call sites of the " + "`retries_remaining` function. This two-pronged approach will allow me to verify the " + "correctness of my proposed fix and understand the broader impact of this change across " + "the codebase.\n" + "\n", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "**My Thought Process on Addressing the `retries_remaining` Issue**\n" + "\n" + "Alright, so my immediate priority is to tackle the `retries_remaining` function. The " + "core of the problem, as I see it, lies in its current implementation where it seems to " + "be adding `limit` and `attempts` rather than performing the correct subtraction to " + "determine remaining retries. This, of course, is a fundamental logic error. " + "Furthermore, I've identified a critical oversight: the function doesn't account for " + "scenarios where the number of `attempts` might actually exceed the `limit`. This is a " + "potential edge case that could lead to unexpected behavior and incorrect results, so it " + "absolutely needs to be addressed.\n" + "\n" + "My proposed solution is a straightforward fix: `return max(0, limit - attempts)`. This " + "ensures that we're correctly calculating the difference between the allowed `limit` and " + "the current `attempts`. Crucially, by wrapping this subtraction within `max(0, ...)`, " + "I'm enforcing a floor of zero. This elegantly handles the situation where `attempts` " + "might be greater than or equal to `limit`, ensuring that the function will never return " + "a negative number of remaining retries, which would be nonsensical in this context.\n" + "\n" + "With that fix in mind, my plan is to proceed with the original request: I will first " + "execute the tests, specifically `tests/test_retry.py`, and run them concurrently for " + "efficiency. Immediately after, I'll perform a thorough search for all call sites of the " + "`retries_remaining` function. This two-pronged approach will allow me to verify the " + "correctness of my proposed fix and understand the broader impact of this change across " + "the codebase.\n" + "\n", + } + ], + "provider_specific_fields": {"thought_signatures": ["synthetic-signature"]}, + "content": None, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + } + ], + "usage": { + "completion_tokens": 147, + "prompt_tokens": 404, + "total_tokens": 551, + "completion_tokens_details": {"reasoning_tokens": 106, "text_tokens": 41}, + "prompt_tokens_details": {"text_tokens": 404}, + }, + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + }, +) + + +GEMINI_3_1_PRO_PREVIEW_AGENT_FOLLOWUP_TEST_CASE: Final = replace( + GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE, + scenario="agent_followup", + litellm_request=followup_request( + GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.litellm_request, + GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.expected_litellm_response, + ), + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:generateContent", + expected_provider_request={ + **VERTEX_AI_AGENT_REQUEST, + "contents": [ + *list_value(VERTEX_AI_AGENT_REQUEST["contents"]), + { + "role": "model", + "parts": [ + { + "thought": True, + "text": "**My Thought Process on Addressing the `retries_remaining` Issue**\n" + "\n" + "Alright, so my immediate priority is to tackle the `retries_remaining` function. The " + "core of the problem, as I see it, lies in its current implementation where it seems to " + "be adding `limit` and `attempts` rather than performing the correct subtraction to " + "determine remaining retries. This, of course, is a fundamental logic error. " + "Furthermore, I've identified a critical oversight: the function doesn't account for " + "scenarios where the number of `attempts` might actually exceed the `limit`. This is a " + "potential edge case that could lead to unexpected behavior and incorrect results, so it " + "absolutely needs to be addressed.\n" + "\n" + "My proposed solution is a straightforward fix: `return max(0, limit - attempts)`. This " + "ensures that we're correctly calculating the difference between the allowed `limit` and " + "the current `attempts`. Crucially, by wrapping this subtraction within `max(0, ...)`, " + "I'm enforcing a floor of zero. This elegantly handles the situation where `attempts` " + "might be greater than or equal to `limit`, ensuring that the function will never return " + "a negative number of remaining retries, which would be nonsensical in this context.\n" + "\n" + "With that fix in mind, my plan is to proceed with the original request: I will first " + "execute the tests, specifically `tests/test_retry.py`, and run them concurrently for " + "efficiency. Immediately after, I'll perform a thorough search for all call sites of the " + "`retries_remaining` function. This two-pronged approach will allow me to verify the " + "correctness of my proposed fix and understand the broader impact of this change across " + "the codebase.\n" + "\n", + }, + { + "function_call": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_2861166", + }, + "thoughtSignature": "synthetic-signature", + }, + { + "function_call": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_2861167", + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "name": "run_tests", + "response": { + "exit_code": 1, + "output": "FAILED tests/test_retry.py::test_retries_remaining - assert 4 == 2", + }, + "id": "call_2861166", + } + }, + { + "function_response": { + "name": "search_code", + "response": { + "matches": [ + { + "path": "src/worker.py", + "line": 12, + "text": "remaining = retries_remaining(3, attempts)", + } + ] + }, + "id": "call_2861167", + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "text": "Stop after these diagnostics. Do not call more tools or change files. Summarize the " + "bug, the caller found and the failed test in two short sentences." + } + ], + }, + ], + }, + mock_provider_response={ + **VERTEX_AI_AGENT_RESPONSE, + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "**My Debugging Session**\n" + "\n" + "It appears I've stumbled upon a rather straightforward arithmetic error within the " + "`retries_remaining` function. My current understanding is that it's incorrectly adding " + "the `limit` and `attempts` parameters instead of performing a subtraction, and " + "consequently, it's not ensuring the result is capped at zero. I've identified that " + "`src/worker.py` is a caller of this function, and specifically, the " + "`tests/test_retry.py::test_retries_remaining` test is failing because, for instance, " + "`retries_remaining(3, 1)` is returning 4 when it should be returning 2.\n" + "\n", + "thought": True, + }, + { + "text": "The bug in `retries_remaining` is that it adds `limit` and `attempts` instead of " + "calculating `max(0, limit - attempts)`. This causes `tests/test_retry.py` to fail, " + "which was verified, and the function is called by `src/worker.py`.", + "thoughtSignature": "synthetic-signature", + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 525, + "candidatesTokenCount": 64, + "totalTokenCount": 706, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 525}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 64}], + "thoughtsTokenCount": 117, + }, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:11:28.565345Z", + "responseId": "kDDSauHAIq6d4_UPv9qnMQ", + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "id": "kDDSauHAIq6d4_UPv9qnMQ", + "model": "vertex_ai/gemini-3.1-pro-preview", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "The bug in `retries_remaining` is that it adds `limit` and `attempts` instead of " + "calculating `max(0, limit - attempts)`. This causes `tests/test_retry.py` to fail, " + "which was verified, and the function is called by `src/worker.py`.", + "role": "assistant", + "images": [], + "reasoning_content": "**My Debugging Session**\n" + "\n" + "It appears I've stumbled upon a rather straightforward arithmetic error within the " + "`retries_remaining` function. My current understanding is that it's incorrectly adding " + "the `limit` and `attempts` parameters instead of performing a subtraction, and " + "consequently, it's not ensuring the result is capped at zero. I've identified that " + "`src/worker.py` is a caller of this function, and specifically, the " + "`tests/test_retry.py::test_retries_remaining` test is failing because, for instance, " + "`retries_remaining(3, 1)` is returning 4 when it should be returning 2.\n" + "\n", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "**My Debugging Session**\n" + "\n" + "It appears I've stumbled upon a rather straightforward arithmetic error within the " + "`retries_remaining` function. My current understanding is that it's incorrectly adding " + "the `limit` and `attempts` parameters instead of performing a subtraction, and " + "consequently, it's not ensuring the result is capped at zero. I've identified that " + "`src/worker.py` is a caller of this function, and specifically, the " + "`tests/test_retry.py::test_retries_remaining` test is failing because, for instance, " + "`retries_remaining(3, 1)` is returning 4 when it should be returning 2.\n" + "\n", + } + ], + "provider_specific_fields": {"thought_signatures": ["synthetic-signature"]}, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + } + ], + "usage": { + "completion_tokens": 181, + "prompt_tokens": 525, + "total_tokens": 706, + "completion_tokens_details": {"reasoning_tokens": 117, "text_tokens": 64}, + "prompt_tokens_details": {"text_tokens": 525}, + }, + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + }, +) + +GEMINI_3_1_PRO_PREVIEW_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="agent_streaming", + litellm_endpoint=GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.litellm_endpoint, + litellm_request={ + **GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.litellm_request, + "stream": True, + "stream_options": {"include_usage": True}, + }, + expected_provider_headers=GEMINI_3_1_PRO_PREVIEW_AGENT_TEST_CASE.expected_provider_headers, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:streamGenerateContent?alt=sse", + expected_provider_request={**VERTEX_AI_AGENT_REQUEST}, + mock_provider_frames=( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_1016270", + }, + "thoughtSignature": "synthetic-signature", + } + ], + } + } + ], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:10:13.305463Z", + "responseId": "RTDSarfSEuXG4_UP_a-tuQo", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_1016276", + } + } + ], + } + } + ], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:10:13.305463Z", + "responseId": "RTDSarfSEuXG4_UP_a-tuQo", + } + ), + SSEFrame( + data={ + "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 404, + "candidatesTokenCount": 41, + "totalTokenCount": 473, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 404}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 41}], + "thoughtsTokenCount": 28, + }, + "modelVersion": "gemini-3.1-pro-preview", + "createTime": "2026-10-16T14:10:13.305463Z", + "responseId": "RTDSarfSEuXG4_UP_a-tuQo", + } + ), + ), + expected_litellm_message={ + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1016270__thought__synthetic-signature", + "type": "function", + "function": {"name": "run_tests", "arguments": '{"path": "tests/test_retry.py"}'}, + }, + { + "id": "call_1016276", + "type": "function", + "function": {"name": "search_code", "arguments": '{"query": "retries_remaining"}'}, + }, + ], + "reasoning_content": None, + "thinking_blocks": [], + "reasoning_items": [], + }, + expected_litellm_usage={ + "prompt_tokens": 404, + "completion_tokens": 69, + "total_tokens": 473, + "prompt_tokens_details": {"text_tokens": 404}, + "completion_tokens_details": {"reasoning_tokens": 28, "text_tokens": 41}, + }, +) + + +GEMINI_3_8_FLASH_AGENT_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="agent_tools", + litellm_request={**CODING_AGENT_REQUEST, "model": "vertex_ai/gemini-3.8-flash"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.8-flash:generateContent", + expected_provider_headers={"content-type": "application/json", "authorization": "Bearer scripted-token"}, + expected_provider_request={**VERTEX_AI_AGENT_REQUEST}, + mock_provider_response={ + **VERTEX_AI_AGENT_RESPONSE, + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_1969624", + }, + "thoughtSignature": "synthetic-signature", + }, + { + "functionCall": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_1969625", + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 404, + "candidatesTokenCount": 41, + "totalTokenCount": 445, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 404}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 41}], + }, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:10:02.505616Z", + "responseId": "OjDSapDuHpCG9LsPreTxQA", + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "id": "OjDSapDuHpCG9LsPreTxQA", + "model": "vertex_ai/gemini-3.8-flash", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "provider_specific_fields": {"thought_signature": "synthetic-signature"}, + "function": {"arguments": '{"query": "retries_remaining"}', "name": "search_code"}, + "id": "call_1969624__thought__synthetic-signature", + "type": "function", + }, + { + "index": 1, + "function": {"arguments": '{"path": "tests/test_retry.py"}', "name": "run_tests"}, + "id": "call_1969625", + "type": "function", + }, + ], + "images": [], + "thinking_blocks": [], + "provider_specific_fields": {"thought_signatures": ["synthetic-signature"]}, + "content": None, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + } + ], + "usage": { + "completion_tokens": 41, + "prompt_tokens": 404, + "total_tokens": 445, + "completion_tokens_details": {"text_tokens": 41}, + "prompt_tokens_details": {"text_tokens": 404}, + }, + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + }, +) + + +GEMINI_3_8_FLASH_AGENT_FOLLOWUP_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_AGENT_TEST_CASE, + scenario="agent_followup", + litellm_request=followup_request( + GEMINI_3_8_FLASH_AGENT_TEST_CASE.litellm_request, GEMINI_3_8_FLASH_AGENT_TEST_CASE.expected_litellm_response + ), + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.8-flash:generateContent", + expected_provider_request={ + **VERTEX_AI_AGENT_REQUEST, + "contents": [ + *list_value(VERTEX_AI_AGENT_REQUEST["contents"]), + { + "role": "model", + "parts": [ + { + "function_call": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_1969624", + }, + "thoughtSignature": "synthetic-signature", + }, + { + "function_call": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_1969625", + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "name": "search_code", + "response": { + "matches": [ + { + "path": "src/worker.py", + "line": 12, + "text": "remaining = retries_remaining(3, attempts)", + } + ] + }, + "id": "call_1969624", + } + }, + { + "function_response": { + "name": "run_tests", + "response": { + "exit_code": 1, + "output": "FAILED tests/test_retry.py::test_retries_remaining - assert 4 == 2", + }, + "id": "call_1969625", + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "text": "Stop after these diagnostics. Do not call more tools or change files. Summarize the " + "bug, the caller found and the failed test in two short sentences." + } + ], + }, + ], + }, + mock_provider_response={ + **VERTEX_AI_AGENT_RESPONSE, + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "`retries_remaining` erroneously adds `limit` and `attempts` instead of calculating " + "`max(0, limit - attempts)`, causing `assert retries_remaining(3, 1) == 2` to fail with " + "`assert 4 == 2`. Aside from its unit test, `src/worker.py` (line 12) was found as the " + "only other caller in the repository.", + "thoughtSignature": "synthetic-signature", + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 525, + "candidatesTokenCount": 85, + "totalTokenCount": 610, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 525}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 85}], + }, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:11:25.313560Z", + "responseId": "jTDSatiRE9ap4_UPuPC7wQo", + }, + expected_litellm_response={ + **CODING_AGENT_RESPONSE, + "id": "jTDSatiRE9ap4_UPuPC7wQo", + "model": "vertex_ai/gemini-3.8-flash", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "`retries_remaining` erroneously adds `limit` and `attempts` instead of calculating " + "`max(0, limit - attempts)`, causing `assert retries_remaining(3, 1) == 2` to fail with " + "`assert 4 == 2`. Aside from its unit test, `src/worker.py` (line 12) was found as the " + "only other caller in the repository.", + "role": "assistant", + "images": [], + "thinking_blocks": [], + "provider_specific_fields": {"thought_signatures": ["synthetic-signature"]}, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + } + ], + "usage": { + "completion_tokens": 85, + "prompt_tokens": 525, + "total_tokens": 610, + "completion_tokens_details": {"text_tokens": 85}, + "prompt_tokens_details": {"text_tokens": 525}, + }, + "vertex_ai_grounding_metadata": [], + "vertex_ai_url_context_metadata": [], + "vertex_ai_safety_results": [], + "vertex_ai_citation_metadata": [], + }, +) + +GEMINI_3_8_FLASH_AGENT_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="agent_streaming", + litellm_endpoint=GEMINI_3_8_FLASH_AGENT_TEST_CASE.litellm_endpoint, + litellm_request={ + **GEMINI_3_8_FLASH_AGENT_TEST_CASE.litellm_request, + "stream": True, + "stream_options": {"include_usage": True}, + }, + expected_provider_headers=GEMINI_3_8_FLASH_AGENT_TEST_CASE.expected_provider_headers, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-3.8-flash:streamGenerateContent?alt=sse", + expected_provider_request={**VERTEX_AI_AGENT_REQUEST}, + mock_provider_frames=( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "search_code", + "args": {"query": "retries_remaining"}, + "id": "call_302633", + }, + "thoughtSignature": "synthetic-signature", + } + ], + } + } + ], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:10:05.862754Z", + "responseId": "PTDSaqLUNOXG4_UP_a-tuQo", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "run_tests", + "args": {"path": "tests/test_retry.py"}, + "id": "call_302642", + } + } + ], + } + } + ], + "usageMetadata": {"trafficType": "ON_DEMAND"}, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:10:05.862754Z", + "responseId": "PTDSaqLUNOXG4_UP_a-tuQo", + } + ), + SSEFrame( + data={ + "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 404, + "candidatesTokenCount": 41, + "totalTokenCount": 445, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 404}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 41}], + }, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-16T14:10:05.862754Z", + "responseId": "PTDSaqLUNOXG4_UP_a-tuQo", + } + ), + ), + expected_litellm_message=coding_agent_stream_message(("call_302633__thought__synthetic-signature", "call_302642")), + expected_litellm_usage={ + "prompt_tokens": 404, + "completion_tokens": 41, + "total_tokens": 445, + "prompt_tokens_details": {"text_tokens": 404}, + "completion_tokens_details": {"text_tokens": 41}, + }, +) + + GEMINI_4_ARGON_AGENT_TEST_CASE: Final = replace( GEMINI_4_ARGON_TEST_CASE, scenario="agent_tools", From eb5295b20c94ccb0ac7cd59df666280a404c24a7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:46:25 +0000 Subject: [PATCH 06/64] fix(anthropic): add the dangerous-tool-use beta when a /v1/messages request carries safeguards (#43955) * fix(anthropic): add the dangerous-tool-use beta when a /v1/messages request carries safeguards * test(azure_ai): pin the outgoing Foundry request keeps the dangerous-tool-use beta next to safeguards * test(azure_ai): inject the Foundry client and scope the beta cache reset to the test * test(azure_ai): own and close the mock transport client, drop the docstring entry * test(integration): audit the safeguards beta derivation across anthropic-compatible providers * fix(github_copilot): keep safeguards from deriving the dangerous-tool-use beta * test(integration): kill the chaos victim worker through the harness helper * fix(anthropic): derive the dangerous-tool-use beta only where the beta map supports it and honor additional_drop_params The safeguards beta is now added only for providers whose beta map column accepts it, so github_copilot, Eden AI and the OpenAI-compatible JSON providers keep their base wire. additional_drop_params now runs before the provider headers are built, so a dropped param no longer drives its derived beta. * fix(anthropic): load the beta map off the event loop and pin a dropped context_management The first /v1/messages request that needs the remote beta map now loads it in a worker thread, so a slow fetch no longer stalls every other request on the worker. An integration cell pins that a deployment dropping context_management sends neither the field nor its derived beta. * fix(anthropic): load the beta map only for safeguards requests, once per worker * fix(anthropic): gate safeguards on the bundled beta map and load the remote map on its own thread --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/anthropic_beta_headers_manager.py | 45 +- .../pass_through/messages/transformation.py | 11 +- litellm/llms/custom_httpx/llm_http_handler.py | 34 +- .../openai_like/messages/transformation.py | 2 + .../test_anthropic_safeguards_beta_chaos.py | 348 ++++++++ .../test_anthropic_safeguards_beta_wire.py | 809 ++++++++++++++++++ .../anthropic/test_anthropic_wire.py | 5 +- ...azure_anthropic_messages_transformation.py | 81 ++ .../custom_httpx/test_llm_http_handler.py | 42 + ..._github_copilot_messages_transformation.py | 25 + ..._like_anthropic_messages_transformation.py | 23 + 11 files changed, 1401 insertions(+), 24 deletions(-) create mode 100644 tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py create mode 100644 tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 8b7fa1c9347..96c38b69104 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -22,17 +22,28 @@ Environment Variables: - LITELLM_ANTHROPIC_BETA_HEADERS_URL: Custom URL for remote config (optional) """ +import asyncio import json import os +import threading +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from functools import cache from importlib.resources import files -from typing import Final +from typing import Final, TypeAlias import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger # Cache for the loaded configuration _BETA_HEADERS_CONFIG: dict | None = None +_BETA_HEADERS_CONFIG_LOCK: Final = threading.Lock() +_BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-beta-headers-config") + +BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]] +_BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) class GetAnthropicBetaHeadersConfig: @@ -174,13 +185,27 @@ def _load_beta_headers_config() -> dict: if _BETA_HEADERS_CONFIG is not None: return _BETA_HEADERS_CONFIG - # Get the URL from environment or use default - from litellm import anthropic_beta_headers_url + with _BETA_HEADERS_CONFIG_LOCK: + if _BETA_HEADERS_CONFIG is not None: + return _BETA_HEADERS_CONFIG - _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url) - verbose_logger.debug("Loaded and cached beta headers config") + # Get the URL from environment or use default + from litellm import anthropic_beta_headers_url - return _BETA_HEADERS_CONFIG + _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url) + verbose_logger.debug("Loaded and cached beta headers config") + + return _BETA_HEADERS_CONFIG + + +async def aload_beta_headers_config() -> None: + if _BETA_HEADERS_CONFIG is None: + await asyncio.get_running_loop().run_in_executor(_BETA_HEADERS_CONFIG_LOADER, _load_beta_headers_config) + + +@cache +def _bundled_beta_headers_config() -> BetaHeadersConfig: + return _BETA_HEADERS_CONFIG_ADAPTER.validate_python(GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()) def reload_beta_headers_config() -> dict: @@ -288,6 +313,14 @@ def is_beta_header_supported( return beta_header in provider_mapping and provider_mapping[beta_header] is not None +def is_beta_header_supported_by_bundled_config(beta_header: str, provider: str) -> bool: + config: Final = _bundled_beta_headers_config() + aliases: Final = config.get("provider_aliases", {}) + canonical_provider: Final = (aliases.get(provider) if isinstance(aliases, Mapping) else None) or provider + provider_mapping: Final = config.get(canonical_provider, {}) + return isinstance(provider_mapping, Mapping) and provider_mapping.get(beta_header) is not None + + def get_provider_beta_header( anthropic_beta_header: str, provider: str, diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index 640736d6463..ad45485fa56 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -4,6 +4,7 @@ from typing import ClassVar, Final import httpx from litellm._logging import verbose_logger +from litellm.anthropic_beta_headers_manager import is_beta_header_supported_by_bundled_config from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -708,6 +709,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params: dict, custom_llm_provider: str = "anthropic", messages: Sequence[object] = (), + beta_map_provider: str | None = None, ) -> dict: """ Auto-inject anthropic-beta headers based on features used. @@ -724,8 +726,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params: Optional parameters including tools, context_management, output_format, speed custom_llm_provider: Provider name for looking up correct tool search header messages: Request messages, scanned for per-message output_config + beta_map_provider: Provider whose beta map column gates safeguards, defaults to custom_llm_provider """ - beta_values: Final[set] = set() + beta_values: Final[set[str]] = set() existing_beta: Final = tuple( piece.strip() @@ -775,6 +778,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): if optional_params.get("speed") == "fast": beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value) + if optional_params.get("safeguards") is not None and is_beta_header_supported_by_bundled_config( + ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value, + beta_map_provider or custom_llm_provider, + ): + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value) + if messages_carry_output_config(messages): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index fffcdd67150..4ff4c43428f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4,7 +4,7 @@ import json import ssl from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager -from functools import lru_cache, partial +from functools import lru_cache, partial, reduce from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -35,7 +35,7 @@ import litellm.litellm_core_utils import litellm.types import litellm.types.utils from litellm._logging import redact_string, verbose_logger -from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta +from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, update_headers_with_filtered_beta from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs @@ -49,6 +49,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( synthesize_subtitle_document, ) from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params +from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.realtime_errors import ( @@ -2304,6 +2305,16 @@ class BaseLLMHTTPHandler: merged_headers.update(extra_headers_from_kwargs) if provider_specific_headers: merged_headers.update(provider_specific_headers) + request_params: Final = reduce( + delete_nested_value, + litellm_params.get("additional_drop_params") or (), + anthropic_messages_optional_request_params, + ) + if ( + request_params.get("safeguards") is not None + and anthropic_messages_provider_config.should_filter_anthropic_beta_headers() + ): + await aload_beta_headers_config() ( headers, api_base, @@ -2311,7 +2322,7 @@ class BaseLLMHTTPHandler: headers=merged_headers or {}, model=model, messages=messages, - optional_params=anthropic_messages_optional_request_params, + optional_params=request_params, litellm_params=dict(litellm_params), api_key=api_key, api_base=api_base, @@ -2342,20 +2353,11 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, ) - additional_drop_params: Final[list[str]] = litellm_params.get("additional_drop_params") or [] - if additional_drop_params: - from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value - - for path in additional_drop_params: - anthropic_messages_optional_request_params = delete_nested_value( - anthropic_messages_optional_request_params, path - ) - # Prepare request body request_body: Final = anthropic_messages_provider_config.transform_anthropic_messages_request( model=model, messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -2462,7 +2464,7 @@ class BaseLLMHTTPHandler: held_back_tool_names: Final = self._server_fulfilled_tools_in_request( logging_obj=logging_obj, - tools=anthropic_messages_optional_request_params.get("tools"), + tools=request_params.get("tools"), ) initial_response = AgenticAnthropicStreamingIterator( completion_stream=completion_stream, @@ -2470,7 +2472,7 @@ class BaseLLMHTTPHandler: model=model, messages=messages, anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs=kwargs_for_agentic, @@ -2493,7 +2495,7 @@ class BaseLLMHTTPHandler: model=model, messages=messages, anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs=kwargs_for_agentic, diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 2e9a300e2fd..94f9792da35 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -8,6 +8,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -57,6 +58,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): headers=normalized, optional_params=optional_params, messages=messages, + beta_map_provider=LlmProviders.OPENAI_LIKE.value, ) return merged, api_base diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py new file mode 100644 index 00000000000..e2f23b7558f --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py @@ -0,0 +1,348 @@ +import asyncio +import json +import re +import threading +import uuid +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.messages_endpoint.providers.anthropic.test_anthropic_safeguards_beta_wire import ( + ANTHROPIC_KEY, + ANTHROPIC_MODEL, + BETA, + OMITTED, + anthropic_peer, + marker_of, + messages_body, + stream_frames, +) +from pydantic import JsonValue, TypeAdapter + +_CONFIG_MODEL: Final = "anthropic-safeguards-chaos" +_FOUNDRY_CONFIG_MODEL: Final = "azure-ai-safeguards-chaos" +_YAML_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_MODEL_LIST: Final = json.dumps( + {"object": "list", "data": [{"id": ANTHROPIC_MODEL, "object": "model", "owned_by": "anthropic"}]} +).encode() + + +@dataclass(frozen=True, slots=True) +class _Call: + stream: bool + marker: str + safeguards: bool = True + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _calls(count: int, stream: bool, safeguards: bool = True) -> tuple[_Call, ...]: + return tuple(_Call(stream=stream, marker=uuid.uuid4().hex, safeguards=safeguards) for _ in range(count)) + + +def _streamed_reply(marker: str, abort_after: int | None = None, pause: float = 0) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=stream_frames(marker), + abort_after=abort_after, + pause_between_chunks=pause, + ) + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + async with client.stream( + "POST", + "/v1/messages", + json=( + messages_body(model, call.marker, stream=call.stream) + if call.safeguards + else messages_body(model, call.marker, stream=call.stream, safeguards=OMITTED) + ), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return _Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _assert_answered_with_its_own_marker(served: _Served) -> None: + assert served.status == 200, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + assert f"msg_{served.call.marker}" in served.text, served.text + if served.call.stream: + assert "event: message_stop" in served.text, served.text + + +def _assert_every_request_carried_the_beta(received: tuple[Request, ...], markers: frozenset[str]) -> None: + assert sorted(marker_of(request) for request in received) == sorted(markers) + missing: Final = [ + frozenset(request.headers) for request in received if request.headers.get("anthropic-beta") != BETA + ] + assert missing == [], missing + + +def _spend_success_ids(model: str, expected: frozenset[str]) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group=%s AND status=%s', (model, "success") + ), + lambda found: expected <= {str(row["request_id"]) for row in found}, + seconds=70, + ) + identities: Final = [row["request_id"] for row in rows] + assert len(set(identities)) == len(identities), identities + + +async def test_burst_with_aborted_streams_keeps_the_beta_on_every_forwarded_request(gateway: Gateway) -> None: + calls: Final = tuple(_Call(stream=index % 2 == 0, marker=uuid.uuid4().hex) for index in range(30)) + aborted: Final = frozenset(call.marker for index, call in enumerate(calls) if call.stream and index % 3 == 0) + + def respond(request: Request) -> Reply: + marker: Final = marker_of(request) + if marker in aborted: + return _streamed_reply(marker, abort_after=0) + return anthropic_peer(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 30 + for item in served: + if item.call.marker in aborted: + assert item.status == 500, item.text + assert '"type":"api_error"' in item.text and "marker-" not in item.text, item.text + else: + _assert_answered_with_its_own_marker(item) + recovery: Final = _Call(stream=True, marker=uuid.uuid4().hex) + (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,)) + _assert_answered_with_its_own_marker(recovered) + _assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in (*calls, recovery))) + answered: Final = frozenset(f"msg_{call.marker}" for call in (*calls, recovery) if call.marker not in aborted) + _spend_success_ids(model, answered) + + +async def test_slow_streams_complete_with_the_beta_on_every_request(gateway: Gateway) -> None: + calls: Final = _calls(10, stream=True) + with ( + wire_server(lambda request: _streamed_reply(marker_of(request), pause=0.3)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 10 + for item in served: + _assert_answered_with_its_own_marker(item) + _assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in calls)) + _spend_success_ids(model, frozenset(f"msg_{call.marker}" for call in calls)) + + +def _anthropic_deployment(wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": _CONFIG_MODEL, + "litellm_params": {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_base": wire.url, "api_key": ANTHROPIC_KEY}, + } + + +def _foundry_deployment(wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": _FOUNDRY_CONFIG_MODEL, + "litellm_params": { + "model": "azure_ai/claude-sonnet-4-6", + "api_base": wire.url, + "api_key": "synthetic-foundry-key", + }, + } + + +def _chaos_config(tmp_path: Path, *deployments: dict[str, JsonValue]) -> Path: + base: Final = _YAML_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = {**base, "model_list": list(deployments)} + path: Final = tmp_path / "anthropic-safeguards-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(180) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_sending_the_beta(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = _calls(20, stream=False) + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + if (request.method, request.target) == ("GET", "/v1/models"): + return Reply(body=_MODEL_LIST) + held_markers.put(marker_of(request)) + assert release.wait(timeout=60), "The burst was never released" + return anthropic_peer(request) + + with wire_server(held) as wire: + path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire)) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(found.group(1)) for found in _STARTED_WORKER.finditer(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst( + str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True + ) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) + follow_up: Final = _Call(stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) + _assert_answered_with_its_own_marker(answered) + received: Final = wire.drain() + posts: Final = tuple(request for request in received if request.method == "POST") + probes: Final[list[tuple[str, str]]] = [ + (request.method, request.target) for request in received if request.method != "POST" + ] + assert set(probes) <= {("GET", "/v1/models")}, probes + _assert_every_request_carried_the_beta(posts, frozenset(call.marker for call in (*calls, follow_up))) + + +def _beta_map_gets(received: tuple[Request, ...]) -> int: + return sum(1 for request in received if (request.method, request.target) == ("GET", "/beta.json")) + + +async def _chat(base_url: str, key: str, marker: str) -> httpx.Response: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + return await client.post( + "/v1/chat/completions", + json={ + "model": _CONFIG_MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Question marker-{marker}"}], + }, + headers={"Authorization": f"Bearer {key}"}, + ) + + +async def _answered_while_the_beta_map_is_held( + base_url: str, key: str, guarded: tuple[_Call, ...], chat_markers: tuple[str, ...] +) -> None: + served: Final = await _burst(base_url, key, _CONFIG_MODEL, guarded) + chats: Final = await asyncio.gather(*(_chat(base_url, key, marker) for marker in chat_markers)) + for item in served: + _assert_answered_with_its_own_marker(item) + for marker, chat in zip(chat_markers, chats, strict=True): + assert chat.status_code == 200, chat.text + assert f"marker-{marker}" in chat.text, chat.text + + +@pytest.mark.timeout(240) +async def test_cold_workers_fetch_the_beta_map_only_for_filtered_safeguards_requests_off_the_shared_threads( + gateway: Gateway, tmp_path: Path +) -> None: + beta_map: Final = Path("litellm/anthropic_beta_headers_config.json").read_bytes() + release: Final = threading.Event() + + def held_beta_map(_: Request) -> Reply: + assert release.wait(timeout=180), "The beta map was never released" + return Reply(body=beta_map) + + def upstream(request: Request) -> Reply: + if (request.method, request.target) == ("GET", "/v1/models"): + return Reply(body=_MODEL_LIST) + return anthropic_peer(request) + + plain: Final = _calls(8, stream=False, safeguards=False) + plain_foundry: Final = _calls(8, stream=False, safeguards=False) + direct: Final = _calls(8, stream=False) + foundry: Final = _calls(80, stream=False) + direct_while_held: Final = _calls(8, stream=False) + chat_markers: Final = tuple(uuid.uuid4().hex for _ in range(4)) + with wire_server(held_beta_map) as beta_map_wire, wire_server(upstream) as wire: + path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire), _foundry_deployment(wire)) + with owned_proxy_process( + gateway, + tmp_path, + {"LITELLM_ANTHROPIC_BETA_HEADERS_URL": f"{beta_map_wire.url}/beta.json"}, + remove_environment=("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS",), + config=path, + workers=2, + ) as owned: + base_url: Final = str(owned.gateway.client.base_url) + key: Final = owned.gateway.key + cold: Final = ( + *await _burst(base_url, key, _CONFIG_MODEL, plain), + *await _burst(base_url, key, _FOUNDRY_CONFIG_MODEL, plain_foundry), + *await _burst(base_url, key, _CONFIG_MODEL, direct), + ) + for served_cold in cold: + _assert_answered_with_its_own_marker(served_cold) + assert beta_map_wire.received.qsize() == 0, "a request that never filters betas fetched the beta map" + + held: Final = asyncio.create_task(_burst(base_url, key, _FOUNDRY_CONFIG_MODEL, foundry)) + try: + await asyncio.to_thread(eventually, beta_map_wire.received.qsize, lambda size: size >= 1, 60) + await asyncio.wait_for( + _answered_while_the_beta_map_is_held(base_url, key, direct_while_held, chat_markers), timeout=45 + ) + assert not held.done(), "Foundry safeguards requests were answered before the beta map loaded" + finally: + release.set() + for served_foundry in await held: + _assert_answered_with_its_own_marker(served_foundry) + + received: Final = tuple(request for request in wire.drain() if request.method == "POST") + beta_by_marker: Final = MappingProxyType( + {marker_of(request): request.headers.get("anthropic-beta") for request in received} + ) + assert {beta_by_marker[call.marker] for call in (*direct, *direct_while_held, *foundry)} == {BETA}, ( + beta_by_marker + ) + assert {beta_by_marker[call.marker] for call in (*plain, *plain_foundry)} == {None}, beta_by_marker + fetches: Final = _beta_map_gets(beta_map_wire.drain()) + assert 1 <= fetches <= 2, fetches diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py new file mode 100644 index 00000000000..0c41efffe73 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py @@ -0,0 +1,809 @@ +import itertools +import json +import re +import time +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +BETA: Final = "dangerous-tool-use-2026-09-03" +SAFEGUARDS: Final[list[JsonValue]] = [ + {"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}} +] +ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929" +ANTHROPIC_KEY: Final = "synthetic-anthropic-key" +_CLAUDE_CODE_BETAS: Final = ( + "claude-code-20250219", + "interleaved-thinking-2025-05-14", + "fine-grained-tool-streaming-2025-05-14", + "context-management-2025-06-27", +) +_CLIENT_BETAS: Final = ",".join(_CLAUDE_CODE_BETAS) +_CLIENT_BETAS_WITH_DANGEROUS: Final = ",".join((*_CLAUDE_CODE_BETAS, BETA)) +_MERGED_BETAS: Final = ",".join(sorted((*_CLAUDE_CODE_BETAS, BETA))) +_FOUNDRY_MODEL: Final = "azure_ai/claude-sonnet-4-6" +_FOUNDRY_KEY: Final = "synthetic-foundry-key" +_FOUNDRY_TARGET: Final = "/anthropic/v1/messages" +_FOUNDRY_ARRAY_ERROR: Final = "safeguards: Input should be a valid array" +_FOUNDRY_EXTRA_INPUT_ERROR: Final = "safeguards: Extra inputs are not permitted" +_ANTHROPIC_VERSION: Final = MappingProxyType({"anthropic-version": "2023-06-01"}) +_MESSAGES_KEYS: Final = frozenset({"max_tokens", "messages", "model", "stream"}) +_RESPONSES_KEYS: Final = frozenset({"include", "input", "max_output_tokens", "model"}) +_SDK_EXTRA: Final[dict[str, JsonValue]] = {"safeguards": SAFEGUARDS, "cache": {"no-cache": True}, "num_retries": 0} +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Omitted: + pass + + +OMITTED: Final = _Omitted() + + +@dataclass(frozen=True, slots=True) +class _Observed: + method: str + target: str + marker: str + headers: Mapping[str, str] + body: Mapping[str, JsonValue] + + @property + def beta(self) -> str | None: + return self.headers.get("anthropic-beta") + + @property + def keys(self) -> frozenset[str]: + return frozenset(self.body) + + @property + def header_names(self) -> frozenset[str]: + return frozenset(self.headers) + + +def marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body + return found.group(1) + + +def _message(marker: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": f"answer marker-{marker}"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 5}, + } + + +def stream_frames(marker: str) -> tuple[bytes, ...]: + opening: Final[dict[str, JsonValue]] = { + **_message(marker), + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 12, "output_tokens": 0}, + } + events: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ("message_start", {"message": opening}), + ("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": f"answer marker-{marker}"}}), + ("content_block_stop", {"index": 0}), + ("message_delta", {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 5}}), + ("message_stop", {}), + ) + return tuple( + f"event: {kind}\ndata: {json.dumps({'type': kind, **payload})}\n\n".encode() for kind, payload in events + ) + + +def anthropic_peer(request: Request) -> Reply: + marker: Final = marker_of(request) + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply(content_type="text/event-stream", chunks=stream_frames(marker)) + return Reply(body=json.dumps(_message(marker)).encode()) + + +def _responses_peer(request: Request) -> Reply: + marker: Final = marker_of(request) + return Reply( + body=json.dumps( + { + "id": f"resp_{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": f"msg_{marker}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": f"answer marker-{marker}", "annotations": []}], + } + ], + "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}, + } + ).encode() + ) + + +def _error_peer(status: int, kind: str, message: str) -> Callable[[Request], Reply]: + body: Final = json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode() + return lambda request: Reply(status=status, body=body) + + +def messages_body( + model: str, + marker: str, + *, + safeguards: JsonValue | _Omitted = SAFEGUARDS, + stream: bool = False, + no_cache: bool = True, +) -> dict[str, JsonValue]: + base: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Question marker-{marker}"}], + "num_retries": 0, + **({"stream": True} if stream else {}), + **({"cache": {"no-cache": True}} if no_cache else {}), + } + if isinstance(safeguards, _Omitted): + return base + return {**base, "safeguards": safeguards} + + +def _observe(request: Request) -> _Observed: + return _Observed( + request.method, request.target, marker_of(request), request.headers, _JSON_OBJECT.validate_json(request.body) + ) + + +def _only(wire: Wire) -> _Observed: + (request,) = wire.drain() + return _observe(request) + + +def _post(gateway: Gateway, body: Mapping[str, JsonValue], *, beta: str | None = None) -> httpx.Response: + headers: Final = {**_ANTHROPIC_VERSION, **({"anthropic-beta": beta} if beta else {})} + return gateway.request("POST", "/v1/messages", body, headers=headers) + + +def _assert_answered(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + assert response.json()["id"] == f"msg_{marker}", response.text + + +def _assert_streamed(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + assert f"msg_{marker}" in response.text and f"answer marker-{marker}" in response.text, response.text + assert "event: message_stop" in response.text, response.text + + +def _spend_row(identity: str) -> None: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda found: len(found) == 1, + seconds=70, + ) + assert rows[0]["status"] == "success", rows + + +def _cache_hit_spend_row(identity: str) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', + (f"{identity}\\_cache\\_hit%",), + ), + lambda found: len(found) == 1, + seconds=70, + ) + assert (rows[0]["status"], str(rows[0]["cache_hit"])) == ("success", "True"), rows + + +def _responses_spend_row(model: str) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group = %s AND call_type = %s', + (model, "aresponses"), + ), + lambda found: len(found) == 1, + seconds=70, + ) + assert rows[0]["status"] == "success", rows + + +def _expect_messages( + observed: _Observed, + *, + target: str, + auth: tuple[str, str] | None, + beta: str | None, + safeguards: JsonValue | _Omitted = SAFEGUARDS, + keys: frozenset[str] = _MESSAGES_KEYS, +) -> None: + assert (observed.method, observed.target) == ("POST", target) + if auth is not None: + credential_matches: Final = observed.headers.get(auth[0]) == auth[1] + assert credential_matches, f"{auth[0]} differs from the deployment credential; headers {observed.header_names}" + assert observed.beta == beta, observed.header_names + if isinstance(safeguards, _Omitted): + assert observed.keys == keys + return + assert observed.keys == keys | {"safeguards"} + assert observed.body["safeguards"] == safeguards + + +def _expect_anthropic_version(observed: _Observed) -> None: + assert observed.headers.get("anthropic-version") == "2023-06-01", observed.header_names + + +@dataclass(frozen=True, slots=True) +class _ShapeCase: + sent: JsonValue | _Omitted + forwarded: JsonValue | _Omitted + client_beta: str | None + expected_beta: str | None + stream: bool = False + + +_ANTHROPIC_SHAPES: Final = MappingProxyType( + { + "D1_safeguards_without_client_header": _ShapeCase(SAFEGUARDS, SAFEGUARDS, None, BETA), + "D2_safeguards_with_claude_code_header": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, _MERGED_BETAS + ), + "D3_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS, _MERGED_BETAS + ), + "D4_empty_safeguards": _ShapeCase([], [], None, BETA), + "D5_null_safeguards": _ShapeCase(None, OMITTED, None, None), + "D6_missing_safeguards": _ShapeCase(OMITTED, OMITTED, None, None), + } +) +_FOUNDRY_SHAPES: Final = MappingProxyType( + { + "D9_streamed_safeguards_with_claude_code_header": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, BETA, stream=True + ), + "D11_empty_safeguards": _ShapeCase([], [], None, BETA), + "D12_null_safeguards": _ShapeCase(None, OMITTED, None, None), + "D13_missing_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase( + OMITTED, OMITTED, _CLIENT_BETAS, None + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_ANTHROPIC_SHAPES)) +def test_anthropic_messages_derive_the_dangerous_tool_use_beta_from_safeguards(gateway: Gateway, row: str) -> None: + case: Final = _ANTHROPIC_SHAPES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + response: Final = _post(gateway, messages_body(model, marker, safeguards=case.sent), beta=case.client_beta) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + beta=case.expected_beta, + safeguards=case.forwarded, + ) + _expect_anthropic_version(observed) + _spend_row(f"msg_{marker}") + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _DropCase: + deployment: Callable[[Gateway], Mapping[str, JsonValue]] + target: str + auth: tuple[str, str] | None + keys: frozenset[str] = _MESSAGES_KEYS + + +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_MODEL: Final = "claude-opus-5-5" + +_DROP_CASES: Final = MappingProxyType( + { + "anthropic": _DropCase( + deployment=lambda _: {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_key": ANTHROPIC_KEY}, + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + ), + "azure_ai": _DropCase( + deployment=lambda _: {"model": _FOUNDRY_MODEL, "api_key": _FOUNDRY_KEY}, + target=_FOUNDRY_TARGET, + auth=("x-api-key", _FOUNDRY_KEY), + ), + "bedrock_claude_platform": _DropCase( + deployment=lambda _: { + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + target="/v1/messages", + auth=("x-api-key", "synthetic-platform-key"), + ), + "vertex_ai": _DropCase( + deployment=lambda gateway: { + "model": f"vertex_ai/{_VERTEX_MODEL}", + "api_key": None, + "vertex_project": _VERTEX_PROJECT, + "vertex_location": _VERTEX_LOCATION, + "vertex_credentials": service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")), + }, + target=( + f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}" + f"/publishers/anthropic/models/{_VERTEX_MODEL}:rawPredict" + ), + auth=None, + keys=frozenset({"anthropic_version", "max_tokens", "messages", "stream"}), + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_DROP_CASES)) +def test_a_deployment_dropping_safeguards_sends_neither_the_field_nor_the_beta(gateway: Gateway, row: str) -> None: + case: Final = _DROP_CASES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + api_base=wire.url, additional_drop_params=["safeguards"], **case.deployment(gateway) + ) + response: Final = _post(gateway, messages_body(model, marker)) + _assert_answered(response, marker) + _expect_messages(_only(wire), target=case.target, auth=case.auth, beta=None, safeguards=OMITTED, keys=case.keys) + _spend_row(f"msg_{marker}") + + +def test_a_deployment_dropping_context_management_sends_neither_the_field_nor_its_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", + api_base=wire.url, + api_key=ANTHROPIC_KEY, + additional_drop_params=["context_management"], + ) + body: Final = { + **messages_body(model, marker, safeguards=OMITTED), + "context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}, + } + response: Final = _post(gateway, body) + _assert_answered(response, marker) + _expect_messages( + _only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None, safeguards=OMITTED + ) + _spend_row(f"msg_{marker}") + + +@pytest.mark.parametrize("row", tuple(_FOUNDRY_SHAPES)) +def test_azure_ai_messages_derive_the_beta_instead_of_forwarding_the_client_header(gateway: Gateway, row: str) -> None: + case: Final = _FOUNDRY_SHAPES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post( + gateway, messages_body(model, marker, safeguards=case.sent, stream=case.stream), beta=case.client_beta + ) + if case.stream: + _assert_streamed(response, marker) + else: + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, + target=_FOUNDRY_TARGET, + auth=("x-api-key", _FOUNDRY_KEY), + beta=case.expected_beta, + safeguards=case.forwarded, + ) + _expect_anthropic_version(observed) + _spend_row(f"msg_{marker}") + + +def test_anthropic_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + reply: Final = client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + extra_body=_SDK_EXTRA, + ) + assert reply.id == f"msg_{marker}", reply + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +async def test_anthropic_sdk_async_stream_safeguards_derive_the_beta_and_log_spend(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0 + ) + async with client.messages.with_streaming_response.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + stream=True, + extra_body=_SDK_EXTRA, + ) as streamed: + pieces: Final = [ + event.delta.text + async for event in await streamed.parse() + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ] + assert "".join(pieces) == f"answer marker-{marker}", pieces + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +def test_azure_ai_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + reply: Final = client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + extra_body=_SDK_EXTRA, + ) + assert reply.id == f"msg_{marker}", reply + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ProviderCase: + parameters: Mapping[str, JsonValue] + auth: tuple[str, str] | None + expected_beta: str | None + target: str = "/v1/messages" + api_base_suffix: str = "" + model_info: Mapping[str, JsonValue] | None = None + client_beta: str | None = None + keys: frozenset[str] = _MESSAGES_KEYS + extra_headers: Mapping[str, str] = MappingProxyType({}) + body_beta: list[str] | None = None + + +_PROVIDERS: Final = MappingProxyType( + { + "D14_deepseek_with_claude_code_header": _ProviderCase( + parameters={"model": "deepseek/deepseek-chat", "api_key": "synthetic-deepseek-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + target="/anthropic/v1/messages", + auth=("x-api-key", "synthetic-deepseek-key"), + expected_beta=None, + ), + "D15_tencent_with_claude_code_header": _ProviderCase( + parameters={"model": "tencent/deepseek-v4-flash", "api_key": "synthetic-tencent-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-tencent-key"), + expected_beta=None, + ), + "D16_minimax_with_claude_code_header": _ProviderCase( + parameters={"model": "minimax/MiniMax-M2", "api_key": "synthetic-minimax-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-minimax-key"), + expected_beta=None, + ), + "D17_openai_compatible_native_passthrough": _ProviderCase( + parameters={"model": "openai/claude-sonnet-4-5", "api_key": "synthetic-openai-like-key"}, + api_base_suffix="/v1", + model_info={"supported_endpoints": ["/v1/messages"]}, + auth=("authorization", "Bearer synthetic-openai-like-key"), + expected_beta=None, + ), + "D18_sail_with_claude_code_header": _ProviderCase( + parameters={"model": "sail/sail-claude", "api_key": "synthetic-sail-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("authorization", "Bearer synthetic-sail-key"), + expected_beta=None, + ), + "D19_edenai": _ProviderCase( + parameters={"model": "edenai/claude-sonnet-4-5", "api_key": "synthetic-edenai-key"}, + auth=("authorization", "Bearer synthetic-edenai-key"), + expected_beta=None, + ), + "D21_bedrock_claude_platform": _ProviderCase( + parameters={ + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + auth=("x-api-key", "synthetic-platform-key"), + expected_beta=BETA, + extra_headers={"anthropic-workspace-id": "ws-synthetic"}, + ), + "D22_bedrock_claude_platform_with_claude_code_header": _ProviderCase( + parameters={ + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-platform-key"), + expected_beta=_MERGED_BETAS, + extra_headers={"anthropic-workspace-id": "ws-synthetic"}, + ), + "D23_bedrock_invoke": _ProviderCase( + parameters={ + "model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + "api_key": "synthetic-bedrock-bearer", + "aws_region_name": "us-east-1", + }, + target="/model/anthropic.claude-sonnet-4-5-20250929-v1:0/invoke", + auth=("authorization", "Bearer synthetic-bedrock-bearer"), + expected_beta=None, + keys=frozenset({"anthropic_beta", "anthropic_version", "max_tokens", "messages"}), + body_beta=[BETA], + ), + "D24_bedrock_mantle": _ProviderCase( + parameters={ + "model": "bedrock_mantle/anthropic.claude-sonnet-5-v1:0", + "api_key": "synthetic-mantle-bearer", + "aws_region_name": "us-east-1", + }, + target="/anthropic/v1/messages", + auth=("authorization", "Bearer synthetic-mantle-bearer"), + expected_beta=BETA, + keys=frozenset({"max_tokens", "messages", "model"}), + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_PROVIDERS)) +def test_anthropic_compatible_providers_receive_the_beta_only_where_the_proxy_derives_it( + gateway: Gateway, row: str +) -> None: + case: Final = _PROVIDERS[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model_info=case.model_info, api_base=wire.url + case.api_base_suffix, **case.parameters + ) + response: Final = _post(gateway, messages_body(model, marker), beta=case.client_beta) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages(observed, target=case.target, auth=case.auth, beta=case.expected_beta, keys=case.keys) + for name, value in case.extra_headers.items(): + assert observed.headers.get(name) == value, observed.header_names + assert observed.body.get("anthropic_beta") == case.body_beta, observed.body + _spend_row(f"msg_{marker}") + + +def test_github_copilot_messages_forward_safeguards_without_deriving_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + token_dir: Final = tmp_path / "copilot" + token_dir.mkdir() + with wire_server(anthropic_peer) as wire: + (token_dir / "api-key.json").write_text( + json.dumps({"token": "synthetic-copilot-token", "expires_at": 4102444800, "endpoints": {"api": wire.url}}) + ) + with ( + owned_proxy(gateway, tmp_path, {"GITHUB_COPILOT_TOKEN_DIR": str(token_dir)}) as owned, + owned.scenario() as scenario, + ): + model: Final = scenario.model(model="github_copilot/claude-sonnet-4-5", api_base=wire.url) + response: Final = _post(owned, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, target="/v1/messages", auth=("authorization", "Bearer synthetic-copilot-token"), beta=None + ) + assert observed.headers.get("openai-intent") == "messages-proxy", observed.header_names + _spend_row(f"msg_{marker}") + + +def test_non_native_deployments_bridge_to_the_responses_api_without_safeguards_or_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-openai-key" + ) + response: Final = _post(gateway, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS) + assert response.status_code == 200 and f"answer marker-{marker}" in response.text, response.text + observed: Final = _only(wire) + assert (observed.method, observed.target) == ("POST", "/v1/responses") + assert observed.headers.get("authorization") == "Bearer synthetic-openai-key", observed.header_names + assert observed.beta is None, observed.header_names + assert observed.keys == _RESPONSES_KEYS, observed.keys + _spend_row(str(_JSON_OBJECT.validate_json(response.content)["id"])) + + +def test_chat_completions_forward_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = openai.OpenAI( + base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0 + ) + completion: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": f"Question marker-{marker}"}], extra_body=_SDK_EXTRA + ) + assert completion.choices[0].message.content == f"answer marker-{marker}", completion + _expect_messages( + _only(wire), + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + beta=None, + keys=frozenset({"max_tokens", "messages", "model"}), + ) + _spend_row(completion.id) + + +async def test_responses_stream_forwards_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = openai.AsyncOpenAI( + base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0 + ) + async with client.responses.with_streaming_response.create( + model=model, input=f"Question marker-{marker}", stream=True, extra_body=_SDK_EXTRA + ) as streamed: + events: Final = [event async for event in await streamed.parse()] + deltas: Final = [event.delta for event in events if event.type == "response.output_text.delta"] + assert "".join(deltas) == f"answer marker-{marker}", events + assert any(event.type == "response.completed" for event in events), events + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None) + _responses_spend_row(model) + + +_FOUNDRY_REJECTED_SHAPES: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"S1_int": 1, "S2_empty_string": "", "S3_five_kb_string": "x" * 5000} +) + + +@pytest.mark.parametrize("row", tuple(_FOUNDRY_REJECTED_SHAPES)) +def test_foundry_rejecting_a_non_array_safeguards_value_reaches_the_caller_after_one_attempt( + gateway: Gateway, row: str +) -> None: + value: Final = _FOUNDRY_REJECTED_SHAPES[row] + marker: Final = uuid.uuid4().hex + with ( + wire_server(_error_peer(400, "invalid_request_error", _FOUNDRY_ARRAY_ERROR)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post(gateway, messages_body(model, marker, safeguards=value)) + assert response.status_code == 400 and _FOUNDRY_ARRAY_ERROR in response.text, response.text + _expect_messages( + _only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA, safeguards=value + ) + + +def test_duplicate_safeguards_keys_resolve_to_the_last_value(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + raw: Final = ( + f'{{"model": "{model}", "max_tokens": 16, "num_retries": 0, "cache": {{"no-cache": true}}, ' + f'"messages": [{{"role": "user", "content": "Question marker-{marker}"}}], ' + f'"safeguards": {json.dumps(SAFEGUARDS)}, "safeguards": null}}' + ) + response: Final = gateway.client.post( + "/v1/messages", + content=raw.encode(), + headers={ + "Authorization": f"Bearer {gateway.key}", + "content-type": "application/json", + **_ANTHROPIC_VERSION, + }, + ) + _assert_answered(response, marker) + _expect_messages( + _only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=None, safeguards=OMITTED + ) + _spend_row(f"msg_{marker}") + + +def test_the_same_safeguards_sent_twice_are_two_upstream_calls_with_two_spend_rows(gateway: Gateway) -> None: + markers: Final = (uuid.uuid4().hex, uuid.uuid4().hex) + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + responses: Final = tuple(_post(gateway, messages_body(model, marker)) for marker in markers) + for marker, response in zip(markers, responses, strict=True): + _assert_answered(response, marker) + received: Final = tuple(_observe(request) for request in wire.drain()) + assert len(received) == 2, received + for observed in received: + _expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + assert sorted(item.marker for item in received) == sorted(markers) + for marker in markers: + _spend_row(f"msg_{marker}") + + +_UPSTREAM_ERRORS: Final = MappingProxyType( + { + "S6_unauthorized_deployment": (401, "authentication_error", "invalid x-api-key"), + "S7_foundry_extra_inputs_rejection": (400, "invalid_request_error", _FOUNDRY_EXTRA_INPUT_ERROR), + } +) + + +@pytest.mark.parametrize("row", tuple(_UPSTREAM_ERRORS)) +def test_upstream_errors_reach_the_caller_with_the_beta_sent_once(gateway: Gateway, row: str) -> None: + status, kind, message = _UPSTREAM_ERRORS[row] + marker: Final = uuid.uuid4().hex + with wire_server(_error_peer(status, kind, message)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post(gateway, messages_body(model, marker)) + assert response.status_code == status and message in response.text, response.text + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + + +def _post_noting_whether_upstream_was_called( + gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire +) -> tuple[httpx.Response, bool]: + before: Final = wire.received.qsize() + response: Final = _post(gateway, body) + assert response.status_code == 200, response.text + return response, wire.received.qsize() > before + + +def _until_served_from_cache(gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire) -> tuple[httpx.Response, int]: + deadline: Final = time.monotonic() + 20 + for attempt in itertools.count(1): + response, called_upstream = _post_noting_whether_upstream_was_called(gateway, body, wire) + if not called_upstream: + return response, attempt + assert time.monotonic() < deadline, "the response cache never served the repeated request" + raise AssertionError("unreachable") + + +def test_the_response_cache_serves_the_repeated_request_without_an_upstream_call(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + body: Final = messages_body(model, marker, no_cache=False) + first: Final = _post(gateway, body) + _assert_answered(first, marker) + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + hit, attempts = _until_served_from_cache(gateway, body, wire) + assert hit.json()["id"] == first.json()["id"], hit.text + assert hit.json()["content"] == first.json()["content"], hit.text + _cache_hit_spend_row(f"msg_{marker}") + misses: Final = tuple(_observe(request) for request in wire.drain()) + assert len(misses) == attempts - 1, (attempts, misses) + for observed in misses: + _expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py index dcfd1d19ac2..a448a8cd2da 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py @@ -10,6 +10,7 @@ from anthropic.types import Message from integration._support import claude_code as cc from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows +from integration._support.safeguards_rig import DANGEROUS_TOOL_USE_BETA from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -280,7 +281,9 @@ def test_anthropic_messages_preserves_cli_beta_safeguards_and_result(gateway: Ga def respond(request: Request) -> Reply: assert request.method == "POST" and request.target == "/v1/messages", request.target assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers - assert sorted(request.headers["anthropic-beta"].split(",")) == sorted(beta.split(",")), request.headers + assert sorted(request.headers["anthropic-beta"].split(",")) == sorted( + (*beta.split(","), DANGEROUS_TOOL_USE_BETA) + ), request.headers expected_body: Final = { "model": "claude-sonnet-4-6", "max_tokens": 32, diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 866bdf173a9..176dcfc86c9 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -7,13 +7,17 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) +from typing import Final from unittest.mock import patch +import httpx import pytest +import litellm from litellm.llms.azure_ai.anthropic.messages_transformation import ( AzureAnthropicMessagesConfig, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.router import GenericLiteLLMParams @@ -470,3 +474,80 @@ def test_azure_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_fl and info.get("supports_mid_conversation_system") is not True ] assert missing == [] + + +AUTO_MODE_SAFEGUARDS: Final = ({"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}},) + + +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("client_headers", "optional_params", "expected_beta"), + [ + ({}, {"safeguards": AUTO_MODE_SAFEGUARDS}, "dangerous-tool-use-2026-09-03"), + ( + {"anthropic-beta": "context-1m-2025-08-07"}, + {"safeguards": AUTO_MODE_SAFEGUARDS}, + "context-1m-2025-08-07,dangerous-tool-use-2026-09-03", + ), + ({"anthropic-beta": "context-1m-2025-08-07"}, {}, "context-1m-2025-08-07"), + ({}, {}, None), + ], +) +def test_safeguards_request_carries_the_dangerous_tool_use_beta_without_the_client_header( + client_headers: dict[str, str], optional_params: dict[str, object], expected_beta: str | None +) -> None: + headers, _ = AzureAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers=client_headers, + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Run ls -la with the Bash tool"}], + optional_params=optional_params, + litellm_params={"api_key": "test-api-key"}, + ) + assert headers.get("anthropic-beta") == expected_beta, headers + + +def _foundry_answer(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + json={ + "id": "msg_safeguards", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello from the gateway"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + request=request, + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_outgoing_azure_request_sends_safeguards_with_the_dangerous_tool_use_beta() -> None: + received: Final[list[httpx.Request]] = [] + + def foundry_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _foundry_answer(request) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=256, + messages=[{"role": "user", "content": "Use the Bash tool to run: echo hello from the gateway"}], + model="azure_ai/claude-sonnet-4-6", + api_key="test-api-key", + api_base="https://test-resource.services.ai.azure.com", + safeguards=AUTO_MODE_SAFEGUARDS, + client=upstream, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + sent: Final = received[0] + assert "dangerous-tool-use-2026-09-03" in sent.headers["anthropic-beta"].split(","), dict(sent.headers) + assert json.loads(sent.content)["safeguards"] == list(AUTO_MODE_SAFEGUARDS) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 50ed70e09f5..23c142b61ab 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5313,3 +5313,45 @@ async def test_retrieve_file_rejects_unverified_bedrock_range_error( ) finally: sync_client.close() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_its_beta() -> None: + received: Final[list[httpx.Request]] = [] + + def anthropic_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response( + status_code=200, + json={ + "id": "msg_dropped_safeguards", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(anthropic_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="anthropic/claude-sonnet-4-6", + api_key="sk-ant-test", + safeguards=({"type": "dangerous_tool_use"},), + additional_drop_params=["safeguards"], + client=upstream, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + sent: Final = received[0] + assert "safeguards" not in json.loads(sent.content) + assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) diff --git a/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py index 5cbb9a0a7b4..46062f16c9a 100644 --- a/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py +++ b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import MagicMock import httpx @@ -9,6 +10,10 @@ from litellm.llms.github_copilot.common_utils import GetAPIKeyError from litellm.llms.github_copilot.messages.transformation import ( GithubCopilotAnthropicMessagesConfig, ) +from litellm.llms.github_copilot.per_user_auth import ( + GITHUB_COPILOT_USER_SESSION_KWARG_KEY, + GithubCopilotUserSession, +) def test_github_copilot_anthropic_messages_get_complete_url(): @@ -185,6 +190,26 @@ def test_github_copilot_anthropic_messages_validate_environment_injects_context_ assert "context-management-2025-06-27" in validated_headers["anthropic-beta"] +@pytest.mark.usefixtures("local_beta_headers_config") +def test_github_copilot_safeguards_do_not_derive_the_dangerous_tool_use_beta() -> None: + session: Final = GithubCopilotUserSession(token="gh.test-key", api_base="https://api.githubcopilot.com") + + validated_headers, _ = GithubCopilotAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="github_copilot/claude-haiku-4.5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={ + "safeguards": ({"type": "dangerous_tool_use"},), + "context_management": {"edits": ({"type": "clear_tool_uses_20250919"},)}, + }, + litellm_params={GITHUB_COPILOT_USER_SESSION_KWARG_KEY: session}, + api_key=None, + api_base=None, + ) + + assert validated_headers["anthropic-beta"] == "context-management-2025-06-27" + + def test_github_copilot_anthropic_messages_validate_environment_auth_error(): """Test error handling when authentication fails.""" config = GithubCopilotAnthropicMessagesConfig() diff --git a/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 9167853bf64..1d5ef50ca3a 100644 --- a/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -516,3 +516,26 @@ def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(co assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"} assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"} assert payload["messages"][2] == {"role": "user", "content": "a plain string message"} + + +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("client_headers", "expected_beta"), + [ + ({}, None), + ({"anthropic-beta": "dangerous-tool-use-2026-09-03"}, "dangerous-tool-use-2026-09-03"), + ], +) +def test_safeguards_reach_an_openai_like_host_with_only_the_client_betas( + config: OpenAILikeAnthropicMessagesConfig, client_headers: dict[str, str], expected_beta: str | None +) -> None: + headers, _ = config.validate_anthropic_messages_environment( + headers=client_headers, + model="some-model", + messages=[{"role": "user", "content": "hi"}], + optional_params={"safeguards": ({"type": "dangerous_tool_use"},)}, + litellm_params={}, + api_key="sk-test", + api_base="https://host/v1", + ) + assert headers.get("anthropic-beta") == expected_beta, headers From 5b042c0db6c4f629e5daccfe9dbdeafbc90265cd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:48:05 +0000 Subject: [PATCH 07/64] fix(rust): add the dangerous-tool-use beta when an anthropic /v1/messages request carries safeguards (#46102) * fix(rust): add the dangerous-tool-use beta when an anthropic /v1/messages request carries safeguards * test(rust): check the dangerous-tool-use beta reaches the wire with safeguards * fix(rust): derive the dangerous-tool-use beta only for hosts with a beta map column DeepSeek and the other hosts without a column in anthropic_beta_headers_config.json keep forwarding safeguards and client betas as sent. Adds the wire-level audit tests. * fix(rust): add the safeguards beta on Azure AI and Claude Platform hosts Azure AI and Claude Platform Messages transforms now add the dangerous-tool-use beta for a non-null safeguards field, through the Azure AI and Anthropic beta-map columns. The safeguards wire and chaos tests get unique file names so they do not collide with the Python fix's tests, the Python-route test asserts the fixed behavior behind a BUG skip, and the DeepSeek test checks that no beta is added * fix(rust): keep Claude Platform out of the safeguards beta until a route can prove it Claude Platform has no beta-map column and no QA account, so it goes back to sending safeguards with no derived beta, like DeepSeek. Also release held chaos-test requests before the thread pool waits on them, and mark the no-column test as rstest. * fix(rust): derive the safeguards beta for Claude Platform and keep the null-cell test off Databricks Claude Platform on AWS answers 400 to safeguards without dangerous-tool-use-2026-09-03 and 200 with it, so its Messages transform derives the beta through the Anthropic column again. The null-cell case now uses only Bedrock Converse, since Databricks support for the beta is changing in a sibling PR. * fix(rust): drop a null safeguards key instead of sending it without its beta Anthropic rejects any safeguards key, null included, when the dangerous-tool-use beta is missing, so an explicit "safeguards": null on the Rust gateway still answered 400. The shared Messages transform now drops a null safeguards key on every host, as the proxy already drops None params before calling Rust, so it answers 200 with no beta. The scripted peer in the integration wire tests now rejects the key on presence, matching the real API. * test: cover Bedrock Mantle in the null safeguards case and match the vendor's non-list wording The Rust null-safeguards regression test now also runs for bedrock_mantle, whose beta is keyed on the field being present. The scripted Anthropic peer answers a non-list safeguards value with the real API's "valid array" wording. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../tests/messages/request.rs | 71 +++ .../src/anthropic/messages/transformation.rs | 91 +++- .../src/azure_ai/messages/transformation.rs | 12 +- .../messages/transformation.rs | 21 +- .../src/deepseek/messages/transformation.rs | 23 + .../src/vertex_ai/messages/transformation.rs | 19 +- .../_anthropic_safeguards_beta_support.py | 285 ++++++++++++ ...st_rust_anthropic_safeguards_beta_chaos.py | 158 +++++++ ...est_rust_anthropic_safeguards_beta_wire.py | 408 ++++++++++++++++++ 9 files changed, 1062 insertions(+), 26 deletions(-) create mode 100644 tests/integration/providers/_anthropic_safeguards_beta_support.py create mode 100644 tests/integration/providers/test_rust_anthropic_safeguards_beta_chaos.py create mode 100644 tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index aa8cb05b170..93bbbb071ba 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -343,6 +343,77 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header( assert_eq!(sent, expected); } +#[rstest] +#[case::anthropic("anthropic", &[AnthropicBeta::DangerousToolUse20260903])] +#[case::deepseek("deepseek", &[])] +#[case::minimax("minimax", &[])] +#[tokio::test] +async fn safeguards_add_the_dangerous_tool_use_beta_only_where_the_beta_map_has_a_column( + call: MessagesCall, + #[case] provider: &str, + #[case] derived: &[AnthropicBeta], +) { + let upstream = upstream([message_response()]).await; + let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]); + + run_message(with_fields( + MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]), + ..call + }, + json!({"safeguards": safeguards}), + )) + .await; + + let request = only_request(&upstream).await; + let expected: BetaSet = derived + .iter() + .cloned() + .chain([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())]) + .collect(); + assert_eq!(sent_betas(&request), expected); + assert_eq!( + request.body_json::().unwrap()["safeguards"], + safeguards + ); +} + +#[rstest] +#[tokio::test] +async fn null_safeguards_are_dropped_instead_of_sent_without_their_beta( + call: MessagesCall, + #[values("anthropic", "azure_ai", "bedrock_mantle", "deepseek", "minimax")] provider: &str, +) { + let upstream = upstream([message_response()]).await; + + run_message(with_fields( + MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]), + ..call + }, + json!({"safeguards": null}), + )) + .await; + + let request = only_request(&upstream).await; + assert_eq!( + sent_betas(&request), + BetaSet::from_iter([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())]) + ); + assert!( + !request + .body_json::>() + .unwrap() + .contains_key("safeguards") + ); +} + #[rstest] #[tokio::test] async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) { diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index aa926662b7a..bcd1e08ad86 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -109,7 +109,10 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { } fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { - update_headers_with_anthropic_beta(headers, request) + merge_beta_headers( + headers, + provider_feature_betas(request, BetaProvider::Anthropic), + ) } } @@ -136,12 +139,20 @@ pub(crate) fn transform_messages_request( messages: strip_encrypted_reasoning_blocks(messages), params: MessagesOptionalParams { context_management, + extra: without_null_safeguards(request.params.extra), ..request.params }, ..request }) } +fn without_null_safeguards(extra: Map) -> Map { + extra + .into_iter() + .filter(|(key, value)| key != "safeguards" || !value.is_null()) + .collect() +} + pub(crate) fn update_headers_with_anthropic_beta( headers: Headers, request: &MessagesRequest, @@ -149,12 +160,34 @@ pub(crate) fn update_headers_with_anthropic_beta( merge_beta_headers(headers, feature_betas(request)) } -/// The betas a request needs on a host other than the first-party API: the Anthropic set -/// filtered and renamed through the host's `BetaProvider` policy. +pub(crate) fn update_headers_with_safeguards_beta( + headers: Headers, + request: &MessagesRequest, + provider: BetaProvider, +) -> Headers { + merge_beta_headers( + headers, + feature_betas(request).union(on_provider(mapped_host_betas(&request.params), provider)), + ) +} + +/// The betas a request needs on a host with a column in the beta map, filtered and renamed +/// through that column. pub(crate) fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet { - feature_betas(request) - .iter() - .filter_map(|beta| beta.on(provider)) + on_provider( + feature_betas(request).union(mapped_host_betas(&request.params)), + provider, + ) +} + +fn on_provider(betas: BetaSet, provider: BetaProvider) -> BetaSet { + betas.iter().filter_map(|beta| beta.on(provider)).collect() +} + +fn mapped_host_betas(params: &MessagesOptionalParams) -> BetaSet { + carries_safeguards(params) + .then_some(AnthropicBeta::DangerousToolUse20260903) + .into_iter() .collect() } @@ -206,6 +239,13 @@ fn uses_structured_output(params: &MessagesOptionalParams) -> bool { .is_some_and(|config| config.format.is_some()) } +fn carries_safeguards(params: &MessagesOptionalParams) -> bool { + params + .extra + .get("safeguards") + .is_some_and(|value| !value.is_null()) +} + fn messages_carry_output_config(messages: &[Message]) -> bool { messages .iter() @@ -1025,10 +1065,44 @@ mod tests { json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}), &["per-turn-control-2026-07-01"] )] + #[case::safeguards( + json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}), + &[] + )] fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { assert_eq!(feature_betas(&request(fields)), betas(expected)); } + #[rstest] + #[case::anthropic(BetaProvider::Anthropic, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])] + #[case::vertex_ai(BetaProvider::VertexAi, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])] + #[case::empty_list(BetaProvider::Anthropic, json!([]), &["dangerous-tool-use-2026-09-03"])] + #[case::not_a_list(BetaProvider::Anthropic, json!(1), &["dangerous-tool-use-2026-09-03"])] + #[case::null(BetaProvider::Anthropic, Value::Null, &[])] + #[case::bedrock_converse_has_a_null_cell(BetaProvider::BedrockConverse, json!([{"type": "dangerous_tool_use"}]), &[])] + fn safeguards_add_the_dangerous_tool_use_beta_where_the_beta_map_has_a_cell( + #[case] provider: BetaProvider, + #[case] safeguards: Value, + #[case] expected: &[&str], + ) { + assert_eq!( + provider_feature_betas(&request(json!({"safeguards": safeguards})), provider), + betas(expected) + ); + } + + #[rstest] + fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() { + let client = headers(&[("anthropic-beta", "claude-code-20250219")]); + assert_eq!( + update_headers_with_anthropic_beta( + client.clone(), + &request(json!({"safeguards": [{"type": "dangerous_tool_use"}]})) + ), + client + ); + } + #[rstest] #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}), &[("x-api-key", "k"), ("anthropic-version", "2023-06-01")])] #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}), &[("Anthropic-Beta", " , "), ("x-api-key", "k")])] @@ -1067,6 +1141,11 @@ mod tests { json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), &[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")], )] + #[case::safeguards_add_the_dangerous_tool_use_beta_to_the_client_betas( + &[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")], + json!({"safeguards": [{"type": "dangerous_tool_use"}]}), + &[("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14")], + )] fn request_headers_merge_the_feature_betas( #[case] input: &[(&str, &str)], #[case] fields: Value, diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index 9eaaa657f6f..3bec142ecf8 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -2,7 +2,7 @@ use litellm_auth::SecretValue; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_llms_types::{ formats::messages::MessagesRequest, - providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH}, + providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; @@ -12,7 +12,7 @@ use crate::{ common_utils::API_KEY_PLACEMENT, messages::{ handler::shape_anthropic_messages_request, - transformation::{transform_messages_request, update_headers_with_anthropic_beta}, + transformation::{transform_messages_request, update_headers_with_safeguards_beta}, }, }, azure_ai::common_utils::{ @@ -104,7 +104,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { } fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { - update_headers_with_anthropic_beta(headers, request) + update_headers_with_safeguards_beta(headers, request, BetaProvider::AzureAi) } } @@ -544,6 +544,12 @@ mod tests { &[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")] )] #[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])] + #[case::safeguards_need_the_dangerous_tool_use_beta( + json!({"safeguards": [{"type": "dangerous_tool_use"}]}), + &[("anthropic-beta", "claude-code-20250219")], + &[("x-api-key", "k"), ("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03")] + )] + #[case::null_safeguards_need_no_beta(json!({"safeguards": null}), &[], &[("x-api-key", "k")])] fn request_headers_carry_the_anthropic_feature_betas( #[case] fields: serde_json::Value, #[case] forwarded: &[(&str, &str)], diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 1b99eca975f..e2e64e54664 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -6,7 +6,7 @@ use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{header_value, with_header, without_headers}; use litellm_llms_types::{ formats::messages::MessagesRequest, - providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH}, + providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; use serde_json::Value; @@ -15,7 +15,7 @@ use crate::{ Error, ErrorDetail, anthropic::messages::{ handler::shape_anthropic_messages_request, - transformation::{transform_messages_request, update_headers_with_anthropic_beta}, + transformation::{transform_messages_request, update_headers_with_safeguards_beta}, }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, @@ -174,7 +174,7 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig { Some(workspace) => with_header(headers, WORKSPACE_HEADER, workspace.into()), None => headers, }; - update_headers_with_anthropic_beta(headers, request) + update_headers_with_safeguards_beta(headers, request, BetaProvider::Anthropic) } fn wire_body(&self, body: Value) -> Value { @@ -356,6 +356,21 @@ mod tests { ); } + #[rstest] + #[case::safeguards(json!([{"type": "dangerous_tool_use"}]), Some("dangerous-tool-use-2026-09-03"))] + #[case::null_safeguards(Value::Null, None)] + fn safeguards_need_the_dangerous_tool_use_beta( + #[case] safeguards: Value, + #[case] expected: Option<&str>, + ) { + let request = serde_json::from_value(json!({"model": "claude-test", "max_tokens": 16, "messages": [], "safeguards": safeguards})).unwrap(); + let headers = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG.request_headers(vec![], &request); + assert_eq!( + header_value(&headers, "anthropic-beta").map(|value| value.to_string()), + expected.map(str::to_string) + ); + } + #[rstest] #[case::default(json!({}), false)] #[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)] diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 9359184f40e..fefedbbbf5e 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -451,4 +451,27 @@ mod tests { json!({"type": "enabled", "budget_tokens": 1024}) ); } + + #[rstest] + #[case::without_client_betas(&[])] + #[case::with_client_betas(&[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")])] + fn safeguards_and_client_betas_reach_deepseek_as_sent(#[case] client: &[(&str, &str)]) { + let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]); + let body = json!({ + "model": "deepseek-v4-pro", + "max_tokens": 100, + "messages": [{"role": "user", "content": "hi"}], + "safeguards": safeguards + }); + let headers: Headers = client + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + assert_eq!( + DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG + .request_headers(headers.clone(), &request_from(body.clone())), + headers + ); + assert_eq!(transformed(body)["safeguards"], safeguards); + } } diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index e6a3661a81f..0e665304661 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -213,23 +213,14 @@ fn has_web_search_tool(tools: Option<&[Recognized]>) -> bool { } /// Python's Vertex `validate_anthropic_messages_environment` beta set: the shared feature -/// betas under Vertex's policy, plus the web search and safeguards betas Vertex needs spelled -/// out. +/// betas under Vertex's policy, plus the web search beta Vertex needs spelled out. fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet { let params = &request.params; provider_feature_betas(request, BetaProvider::VertexAi).union( - [ - has_web_search_tool(params.tools.as_deref()) - .then_some(AnthropicBeta::WebSearch20250305), - params - .extra - .get("safeguards") - .is_some_and(|value| !value.is_null()) - .then_some(AnthropicBeta::DangerousToolUse20260903), - ] - .into_iter() - .flatten() - .collect(), + has_web_search_tool(params.tools.as_deref()) + .then_some(AnthropicBeta::WebSearch20250305) + .into_iter() + .collect(), ) } diff --git a/tests/integration/providers/_anthropic_safeguards_beta_support.py b/tests/integration/providers/_anthropic_safeguards_beta_support.py new file mode 100644 index 00000000000..0307ea11724 --- /dev/null +++ b/tests/integration/providers/_anthropic_safeguards_beta_support.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import json +import re +import threading +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import yaml +from integration._support.anthropic_sse import ( + CONTENT_BLOCK_START, + error_body, + message_start, + sse, + text_delta, +) +from integration._support.client import Gateway, object_value +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue, TypeAdapter + +DANGEROUS_TOOL_USE: Final = "dangerous-tool-use-2026-09-03" +WEB_SEARCH: Final = "web-search-2025-03-05" +CLIENT_BETAS: Final = "interleaved-thinking-2025-05-14,claude-code-20250219,interleaved-thinking-2025-05-14" +MERGED_BETAS: Final = f"claude-code-20250219,{DANGEROUS_TOOL_USE},interleaved-thinking-2025-05-14" +BACKEND: Final = "claude-sonnet-4-6" +ANTHROPIC_DEPLOYMENT: Final = "safeguarded-claude" +VERTEX_DEPLOYMENT: Final = "safeguarded-vertex-claude" +DEEPSEEK_DEPLOYMENT: Final = "safeguarded-deepseek" +ANTHROPIC_KEY: Final = "synthetic-anthropic-key" +DEEPSEEK_KEY: Final = "synthetic-deepseek-key" +PROJECT: Final = "scripted-project" +LOCATION: Final = "us-east5" +_VERTEX_MODEL_PATH: Final = f"/v1/projects/{PROJECT}/locations/{LOCATION}/publishers/anthropic/models/{BACKEND}" +VERTEX_TARGET: Final = f"{_VERTEX_MODEL_PATH}:rawPredict" +VERTEX_STREAM_TARGET: Final = f"{_VERTEX_MODEL_PATH}:streamRawPredict?alt=sse" +ANTHROPIC_TARGET: Final = "/v1/messages" +DEEPSEEK_TARGET: Final = "/anthropic/v1/messages" +SCRIPTED_TARGETS: Final = frozenset({ANTHROPIC_TARGET, DEEPSEEK_TARGET, VERTEX_TARGET, VERTEX_STREAM_TARGET}) +REPLY_TEXT: Final = "safeguards beta control" +EXTRA_INPUTS_MESSAGE: Final = "safeguards: Extra inputs are not permitted" +NOT_A_LIST_MESSAGE: Final = "safeguards: Input should be a valid array" +OVERLOADED_MESSAGE: Final = "Overloaded" +SAFEGUARDS: Final[tuple[JsonValue, ...]] = ( + { + "type": "dangerous_tool_use", + "classifier_context": {"v": 1, "permission_mode": "auto"}, + }, +) +SAFEGUARD_RESULTS: Final[tuple[JsonValue, ...]] = ( + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}, +) +BASH_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "name": "Bash", + "description": "Run a shell command", + "input_schema": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + } +) +WEB_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "web_search_20250305", "name": "web_search", "max_uses": 1} +) +RUST_ROUTE_HEADER: Final = "x-litellm-rust" +_MARKER: Final = re.compile(rb"marker-[0-9a-f]{32}") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_STARTED_WORKER: Final = "Started server process [" + + +@dataclass(frozen=True, slots=True) +class Delivery: + target: str + beta: str | None + safeguards: JsonValue + marker: str + authorization: str | None + api_key: str | None + + +def new_marker() -> str: + return f"marker-{uuid.uuid4().hex}" + + +def marker_of(request: Request) -> str: + found: Final = tuple(match.group(0) for match in _MARKER.finditer(request.body)) + return found[-1].decode() if found else "unmarked" + + +def body_of(request: Request) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(request.body) + + +def json_object(text: str) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(text) + + +def betas_of(request: Request) -> frozenset[str]: + return frozenset(part.strip() for part in request.headers.get("anthropic-beta", "").split(",") if part.strip()) + + +def started_workers(log: Path) -> tuple[int, ...]: + return tuple( + int(line.rsplit("[", 1)[1].strip(" ]\r")) for line in log.read_text().splitlines() if _STARTED_WORKER in line + ) + + +def _streams(request: Request) -> bool: + return request.target == VERTEX_STREAM_TARGET or body_of(request).get("stream") is True + + +def _results(safeguards: JsonValue) -> list[JsonValue] | None: + if safeguards is None: + return None + return list(SAFEGUARD_RESULTS) if safeguards else [] + + +def _message(identity: str, results: list[JsonValue] | None) -> dict[str, JsonValue]: + verdict: Final[dict[str, JsonValue]] = {} if results is None else {"safeguard_results": results} + return { + "id": identity, + "type": "message", + "role": "assistant", + "model": BACKEND, + "content": [{"type": "text", "text": REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + **verdict, + } + + +def _stream(identity: str, results: list[JsonValue] | None) -> tuple[bytes, ...]: + verdict: Final[dict[str, JsonValue]] = {} if results is None else {"safeguard_results": results} + return ( + message_start(identity, BACKEND), + CONTENT_BLOCK_START, + text_delta(REPLY_TEXT), + sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, **verdict}, + "usage": {"output_tokens": 3}, + }, + ), + sse("message_stop", {"type": "message_stop"}), + ) + + +def anthropic_peer(request: Request) -> Reply: + if request.target not in SCRIPTED_TARGETS: + return Reply( + status=404, + body=json.dumps({"error": f"unscripted target {request.target}"}).encode(), + ) + if request.target == DEEPSEEK_TARGET: + return Reply(body=json.dumps(_message(f"msg_{marker_of(request)}", None)).encode()) + body: Final = body_of(request) + safeguards: Final = body.get("safeguards") + if "safeguards" in body and DANGEROUS_TOOL_USE not in betas_of(request): + return Reply(status=400, body=error_body(400, EXTRA_INPUTS_MESSAGE)) + if safeguards is not None and not isinstance(safeguards, list): + return Reply(status=400, body=error_body(400, NOT_A_LIST_MESSAGE)) + identity: Final = f"msg_{marker_of(request)}" + results: Final = _results(safeguards) + if _streams(request): + return Reply(content_type="text/event-stream", chunks=_stream(identity, results)) + return Reply(body=json.dumps(_message(identity, results)).encode()) + + +def outage_peer(outage: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if outage.is_set(): + return Reply(status=529, body=error_body(529, OVERLOADED_MESSAGE)) + return anthropic_peer(request) + + return respond + + +def holding_peer(held: SimpleQueue[str], release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + held.put(marker_of(request)) + assert release.wait(timeout=60), "Held request was never released" + return anthropic_peer(request) + + return respond + + +def deliveries(wire: Wire) -> tuple[Delivery, ...]: + return tuple( + Delivery( + request.target, + request.headers.get("anthropic-beta"), + body_of(request).get("safeguards"), + marker_of(request), + request.headers.get("authorization"), + request.headers.get("x-api-key"), + ) + for request in wire.drain() + ) + + +def one_delivery(wire: Wire, marker: str) -> Delivery: + (sent,) = deliveries(wire) + assert sent.marker == marker, sent + return sent + + +def rust_routed(headers: Mapping[str, str]) -> bool: + return headers.get(RUST_ROUTE_HEADER) == "true" + + +def messages_body(model: str, marker: str, *, stream: bool = False) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": marker}], + "tools": [dict(BASH_TOOL)], + } + + +def safeguarded(body: Mapping[str, JsonValue], safeguards: JsonValue) -> dict[str, JsonValue]: + return {**body, "safeguards": safeguards} + + +def chat_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}]} + + +def responses_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "input": marker} + + +def owned_config(directory: Path, gateway: Gateway, api_base: str) -> Path: + config: Final = _JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + merged: Final = { + **config, + "model_list": [ + { + "model_name": ANTHROPIC_DEPLOYMENT, + "litellm_params": { + "model": f"anthropic/{BACKEND}", + "api_base": api_base, + "api_key": ANTHROPIC_KEY, + }, + }, + { + "model_name": VERTEX_DEPLOYMENT, + "litellm_params": { + "model": f"vertex_ai/{BACKEND}", + "api_base": api_base, + "vertex_project": PROJECT, + "vertex_location": LOCATION, + "vertex_credentials": service_account_json(PROJECT, gateway.upstream_url.rstrip("/")), + }, + }, + { + "model_name": DEEPSEEK_DEPLOYMENT, + "litellm_params": { + "model": "deepseek/deepseek-chat", + "api_base": api_base, + "api_key": DEEPSEEK_KEY, + }, + }, + ], + "router_settings": { + **object_value(config["router_settings"]), + "num_retries": 0, + }, + } + path: Final = directory / f"safeguards-beta-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path diff --git a/tests/integration/providers/test_rust_anthropic_safeguards_beta_chaos.py b/tests/integration/providers/test_rust_anthropic_safeguards_beta_chaos.py new file mode 100644 index 00000000000..f90d68c1a63 --- /dev/null +++ b/tests/integration/providers/test_rust_anthropic_safeguards_beta_chaos.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpcore +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import wire_server +from integration.providers._anthropic_safeguards_beta_support import ( + ANTHROPIC_DEPLOYMENT, + DANGEROUS_TOOL_USE, + SAFEGUARDS, + deliveries, + holding_peer, + messages_body, + new_marker, + owned_config, + rust_routed, + safeguarded, + started_workers, +) +from pydantic import TypeAdapter + +_CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_RESTART_TIMEOUT: Final = int(4 * graceful_stop_seconds() + 240) +_BURST: Final = 12 +_SECOND_WAVE: Final = 6 +_CLIENT_ADDRESS: Final = TypeAdapter(tuple[str, int]) + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=120, trust_env=False)) for _ in range(count) + ) + + +def _post(client: httpx.Client, key: str, marker: str) -> int: + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), list(SAFEGUARDS)) + response: Final = client.post("/v1/messages", json=body, headers={"Authorization": f"Bearer {key}"}) + assert response.status_code != 200 or rust_routed(response.headers), dict(response.headers) + return response.status_code + + +def _post_or_dropped(client: httpx.Client, key: str, marker: str) -> int | None: + try: + return _post(client, key, marker) + except httpx.TransportError: + return None + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + stream: Final[object] = response.extensions["network_stream"] # pyright: ignore[reportAny] # httpx types extensions as Any + assert isinstance(stream, httpcore.NetworkStream), stream + port: Final = _CLIENT_ADDRESS.validate_python(stream.get_extra_info("client_addr"))[1] + response.read() + assert response.status_code == 200, response.text + return port + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_worker_sigkill_mid_safeguards_burst_leaves_the_replacement_adding_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(holding_peer(held, release))) + config: Final = owned_config(tmp_path, gateway, wire.url) + owned: Final = stack.enter_context( + owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) + ) + proxy_url: Final = owned.gateway.client.base_url + workers: Final = eventually(lambda: started_workers(owned.log), lambda pids: len(pids) == 2, seconds=30) + clients: Final = _clients(stack, str(proxy_url), _BURST) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + stack.callback(release.set) + ports: Final = tuple(_local_port(client) for client in clients) + markers: Final = tuple(new_marker() for _ in clients) + futures: Final = tuple( + pool.submit(_post_or_dropped, client, gateway.key, marker) + for client, marker in zip(clients, markers, strict=True) + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else 200), ( + port, + result, + shares, + ) + second_wave: Final = tuple(new_marker() for _ in range(_SECOND_WAVE)) + for marker in second_wave: + assert _post(owned.gateway.client, gateway.key, marker) == 200 + delivered: Final = deliveries(wire) + assert tuple(sent.beta for sent in delivered) == (DANGEROUS_TOOL_USE,) * (len(clients) + len(second_wave)), ( + delivered + ) + seen: Final = tuple(sum(sent.marker == marker for sent in delivered) for marker in markers + second_wave) + assert seen == (1,) * len(seen), seen + assert owned.process.poll() is None + + +@pytest.mark.timeout(_RESTART_TIMEOUT) +def test_proxy_shutdown_mid_safeguards_burst_then_a_fresh_boot_keeps_adding_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(holding_peer(held, release))) + config: Final = owned_config(tmp_path, gateway, wire.url) + markers: Final = tuple(new_marker() for _ in range(_BURST)) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(markers))) + stack.callback(release.set) + with owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) as first: + clients: Final = _clients(stack, str(first.gateway.client.base_url), len(markers)) + futures: Final = tuple( + pool.submit(_post_or_dropped, client, gateway.key, marker) + for client, marker in zip(clients, markers, strict=True) + ) + eventually(held.qsize, lambda size: size == len(markers), seconds=30) + first.begin_shutdown() + release.set() + results: Final = tuple(future.result(timeout=120) for future in futures) + during: Final = deliveries(wire) + assert all(result in (None, 200) for result in results), results + assert tuple(sent.beta for sent in during) == (DANGEROUS_TOOL_USE,) * len(markers), during + assert sorted(sent.marker for sent in during) == sorted(markers), during + with owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) as second: + after: Final = tuple(new_marker() for _ in range(_SECOND_WAVE)) + statuses: Final = tuple(_post(second.gateway.client, gateway.key, marker) for marker in after) + assert statuses == (200,) * len(after), statuses + delivered: Final = deliveries(wire) + assert tuple(sent.beta for sent in delivered) == (DANGEROUS_TOOL_USE,) * len(after), delivered + assert sorted(sent.marker for sent in delivered) == sorted(after), delivered diff --git a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py new file mode 100644 index 00000000000..6ed7e34a583 --- /dev/null +++ b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py @@ -0,0 +1,408 @@ +from __future__ import annotations + +import json +import threading +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Final + +import anthropic +import httpx +import pytest +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, +) +from integration._support.database import read_rows +from integration._support.process import ( + OwnedProxy, + graceful_stop_seconds, + owned_proxy_process, +) +from integration._support.wire import Wire, wire_server +from integration.providers._anthropic_safeguards_beta_support import ( + ANTHROPIC_DEPLOYMENT, + ANTHROPIC_KEY, + ANTHROPIC_TARGET, + BACKEND, + BASH_TOOL, + CLIENT_BETAS, + DANGEROUS_TOOL_USE, + DEEPSEEK_DEPLOYMENT, + DEEPSEEK_KEY, + DEEPSEEK_TARGET, + MERGED_BETAS, + NOT_A_LIST_MESSAGE, + OVERLOADED_MESSAGE, + REPLY_TEXT, + SAFEGUARD_RESULTS, + SAFEGUARDS, + VERTEX_DEPLOYMENT, + VERTEX_STREAM_TARGET, + VERTEX_TARGET, + WEB_SEARCH, + WEB_SEARCH_TOOL, + anthropic_peer, + chat_body, + deliveries, + json_object, + messages_body, + new_marker, + one_delivery, + outage_peer, + owned_config, + responses_body, + rust_routed, + safeguarded, + started_workers, +) +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.xdist_group("anthropic_safeguards_beta") + +_CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_BURST: Final = 24 +_RECOVERY: Final = 6 +_LARGE_ENTRY: Final[JsonValue] = [ + { + "type": "dangerous_tool_use", + "classifier_context": {"v": 1, "permission_mode": "auto", "note": "x" * 5120}, + } +] + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + wire: Wire + outage: threading.Event + + @property + def gateway(self) -> Gateway: + return self.owned.gateway + + +@dataclass(frozen=True, slots=True) +class _Answer: + path: str + marker: str + status: int + text: str + rust: bool + + +@pytest.fixture(scope="module") +def shared_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("anthropic-safeguards-beta") + outage: Final = threading.Event() + with ( + gateway_from_environment() as environment, + wire_server(outage_peer(outage)) as wire, + ): + config: Final = owned_config(directory, environment, wire.url) + with owned_proxy_process(environment, directory, {"LITELLM_RUST": "true"}, config=config, workers=2) as owned: + eventually( + lambda: started_workers(owned.log), + lambda pids: len(pids) == 2, + seconds=60, + ) + wire.drain() + yield _Rig(owned, wire, outage) + + +@pytest.fixture +def rig(shared_rig: _Rig) -> _Rig: + shared_rig.outage.clear() + shared_rig.wire.drain() + return shared_rig + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', + (call_id,), + ) + + +def _post(gateway: Gateway, body: Mapping[str, JsonValue], *, beta: str | None = None) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + body, + headers={} if beta is None else {"anthropic-beta": beta}, + ) + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_sdk_messages_with_safeguards_go_upstream_with_the_dangerous_tool_use_beta( + rig: _Rig, +) -> None: + marker: Final = new_marker() + with httpx.Client(trust_env=False, timeout=60) as http: + client: Final = anthropic.Anthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=http, + ) + raw: Final = client.messages.with_raw_response.create( + model=ANTHROPIC_DEPLOYMENT, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + extra_body={"safeguards": list(SAFEGUARDS), "tools": [dict(BASH_TOOL)]}, + ) + message: Final = raw.parse() + assert rust_routed(raw.headers), dict(raw.headers) + received: Final = json_object(message.model_dump_json()) + assert message.id == f"msg_{marker}", received + assert received.get("safeguard_results") == list(SAFEGUARD_RESULTS), received + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.api_key, sent.beta) == ( + ANTHROPIC_TARGET, + ANTHROPIC_KEY, + DANGEROUS_TOOL_USE, + ), sent + assert sent.safeguards == list(SAFEGUARDS), sent + assert "x-litellm-call-id" in raw.headers, dict(raw.headers) + call_id: Final = raw.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: _spend_rows(call_id), lambda found: len(found) == 1, seconds=70) + assert len(rows) == 1, rows + + +@pytest.mark.timeout(_CELL_TIMEOUT) +async def test_sdk_streamed_messages_with_safeguards_go_upstream_with_the_beta( + rig: _Rig, +) -> None: + marker: Final = new_marker() + async with httpx.AsyncClient(trust_env=False, timeout=60) as http: + client: Final = anthropic.AsyncAnthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=http, + ) + async with client.messages.stream( + model=ANTHROPIC_DEPLOYMENT, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + extra_body={"safeguards": list(SAFEGUARDS), "tools": [dict(BASH_TOOL)]}, + ) as stream: + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert rust_routed(stream.response.headers), dict(stream.response.headers) + assert (final.id, text) == (f"msg_{marker}", REPLY_TEXT), (final, text) + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.beta, sent.safeguards) == ( + ANTHROPIC_TARGET, + DANGEROUS_TOOL_USE, + list(SAFEGUARDS), + ), sent + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_raw_messages_merge_the_safeguards_beta_into_the_client_betas(rig: _Rig, stream: bool) -> None: + marker: Final = new_marker() + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker, stream=stream), list(SAFEGUARDS)) + response: Final = _post(rig.gateway, body, beta=CLIENT_BETAS) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.beta) == (ANTHROPIC_TARGET, MERGED_BETAS), sent + + +@pytest.mark.parametrize("extra", [{"safeguards": None}, {}], ids=["null", "absent"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_null_or_absent_safeguards_add_no_beta(rig: _Rig, extra: Mapping[str, JsonValue]) -> None: + marker: Final = new_marker() + response: Final = _post(rig.gateway, {**messages_body(ANTHROPIC_DEPLOYMENT, marker), **extra}) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" not in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.beta, sent.safeguards) == (None, None), sent + + +@pytest.mark.parametrize( + ("safeguards", "status", "fragment"), + [ + (7, 400, NOT_A_LIST_MESSAGE), + ("dangerous_tool_use", 400, NOT_A_LIST_MESSAGE), + ({}, 400, NOT_A_LIST_MESSAGE), + ([], 200, "safeguard_results"), + (_LARGE_ENTRY, 200, "safeguard_results"), + ], + ids=["int", "string", "object", "empty_list", "5kb_entry"], +) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_malformed_safeguards_still_add_the_beta_and_relay_the_vendor_verdict( + rig: _Rig, safeguards: JsonValue, status: int, fragment: str +) -> None: + marker: Final = new_marker() + response: Final = _post( + rig.gateway, + safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), safeguards), + ) + assert (response.status_code, fragment in response.text) == (status, True), response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.beta, sent.safeguards) == (DANGEROUS_TOOL_USE, safeguards), sent + control: Final = new_marker() + after: Final = _post(rig.gateway, messages_body(ANTHROPIC_DEPLOYMENT, control)) + assert after.status_code == 200, after.text + assert rust_routed(after.headers), dict(after.headers) + assert one_delivery(rig.wire, control).beta is None + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize("last", ["null", "list"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_duplicated_safeguards_keys_resolve_to_the_last_value(rig: _Rig, last: str) -> None: + marker: Final = new_marker() + opening: Final = json.dumps(messages_body(ANTHROPIC_DEPLOYMENT, marker))[:-1] + null_field: Final = '"safeguards": null' + list_field: Final = f'"safeguards": {json.dumps(list(SAFEGUARDS))}' + ordered: Final = (null_field, list_field) if last == "list" else (list_field, null_field) + raw: Final = f"{opening}, {ordered[0]}, {ordered[1]}}}".encode() + response: Final = rig.gateway.client.post( + "/v1/messages", + content=raw, + headers={ + "Authorization": f"Bearer {rig.gateway.key}", + "content-type": "application/json", + }, + ) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert ("safeguard_results" in response.text) == (last == "list"), response.text + sent: Final = one_delivery(rig.wire, marker) + assert sent.beta == (DANGEROUS_TOOL_USE if last == "list" else None), sent + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_safeguards_request_with_a_bad_key_never_reaches_the_peer(rig: _Rig) -> None: + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, new_marker()), list(SAFEGUARDS)) + response: Final = rig.gateway.request("POST", "/v1/messages", body, key="sk-not-a-key") + assert response.status_code == 401, response.text + assert deliveries(rig.wire) == () + + +@pytest.mark.parametrize("web_search", [False, True], ids=["bash_only", "with_web_search"]) +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_vertex_messages_with_safeguards_keep_the_beta_from_the_shared_helper( + rig: _Rig, stream: bool, web_search: bool +) -> None: + marker: Final = new_marker() + tools: Final[list[JsonValue]] = [ + dict(BASH_TOOL), + *((dict(WEB_SEARCH_TOOL),) if web_search else ()), + ] + body: Final = safeguarded(messages_body(VERTEX_DEPLOYMENT, marker, stream=stream), list(SAFEGUARDS)) + response: Final = _post(rig.gateway, {**body, "tools": tools}) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert sent.target == (VERTEX_STREAM_TARGET if stream else VERTEX_TARGET), sent + assert sent.authorization == "Bearer scripted-token", sent + assert sent.beta == (f"{DANGEROUS_TOOL_USE},{WEB_SEARCH}" if web_search else DANGEROUS_TOOL_USE), sent + + +def _beta_set(header: str | None) -> frozenset[str]: + return frozenset(part.strip() for part in (header or "").split(",") if part.strip()) + + +@pytest.mark.parametrize("beta", [None, CLIENT_BETAS], ids=["no-client-betas", "client-betas"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_deepseek_messages_forward_safeguards_without_adding_a_beta(rig: _Rig, beta: str | None) -> None: + marker: Final = new_marker() + response: Final = _post( + rig.gateway, + safeguarded(messages_body(DEEPSEEK_DEPLOYMENT, marker), list(SAFEGUARDS)), + beta=beta, + ) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.api_key, sent.safeguards) == (DEEPSEEK_TARGET, DEEPSEEK_KEY, list(SAFEGUARDS)), sent + assert _beta_set(sent.beta) <= _beta_set(beta) - {DANGEROUS_TOOL_USE}, sent + + +def test_python_route_adds_the_beta_for_safeguards(gateway: Gateway, scenario: Scenario) -> None: + pytest.skip("BUG: Python /v1/messages sends safeguards without the dangerous-tool-use beta, Anthropic answers 400") + with wire_server(anthropic_peer) as wire: + deployment: Final = scenario.model(model=f"anthropic/{BACKEND}", api_base=wire.url, api_key=ANTHROPIC_KEY) + marker: Final = new_marker() + response: Final = _post(gateway, safeguarded(messages_body(deployment, marker), list(SAFEGUARDS))) + assert response.status_code == 200, response.text + assert "safeguard_results" in response.text, response.text + (sent,) = deliveries(wire) + assert (sent.marker, sent.safeguards) == (marker, list(SAFEGUARDS)), sent + assert DANGEROUS_TOOL_USE in _beta_set(sent.beta), sent + assert not rust_routed(response.headers), dict(response.headers) + + +def _burst_request(index: int, marker: str) -> tuple[str, dict[str, JsonValue]]: + if index % 4 == 0: + return "/v1/messages", safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), list(SAFEGUARDS)) + if index % 4 == 1: + return "/v1/messages", safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker, stream=True), list(SAFEGUARDS)) + if index % 4 == 2: + return "/v1/chat/completions", chat_body(ANTHROPIC_DEPLOYMENT, marker) + return "/v1/responses", responses_body(ANTHROPIC_DEPLOYMENT, marker) + + +def _answer(gateway: Gateway, path: str, marker: str, body: Mapping[str, JsonValue]) -> _Answer: + response: Final = gateway.request("POST", path, body) + return _Answer(path, marker, response.status_code, response.text, rust_routed(response.headers)) + + +def _wave(pool: ThreadPoolExecutor, gateway: Gateway, markers: tuple[str, ...], offset: int) -> tuple[_Answer, ...]: + planned: Final = tuple((marker, *_burst_request(offset + index, marker)) for index, marker in enumerate(markers)) + submitted: Final = tuple(pool.submit(_answer, gateway, path, marker, body) for marker, path, body in planned) + return tuple(future.result(timeout=120) for future in submitted) + + +def _assert_served(answer: _Answer) -> None: + assert answer.status == 200, answer + assert REPLY_TEXT in answer.text, answer + assert answer.rust == (answer.path == "/v1/messages"), answer + if answer.path == "/v1/messages": + assert f"msg_{answer.marker}" in answer.text and "safeguard_results" in answer.text, answer + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_peer_outage_mid_safeguards_burst_relays_each_verdict_once_and_recovers( + rig: _Rig, +) -> None: + healthy: Final = tuple(new_marker() for _ in range(_BURST // 2)) + outage: Final = tuple(new_marker() for _ in range(_BURST // 2)) + recovery: Final = tuple(new_marker() for _ in range(_RECOVERY)) + with ThreadPoolExecutor(max_workers=8) as pool: + before: Final = _wave(pool, rig.gateway, healthy, 0) + rig.outage.set() + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + during: Final = _wave(pool, rig.gateway, outage, len(healthy)) + rig.outage.clear() + after: Final = _wave(pool, rig.gateway, recovery, len(healthy) + len(outage)) + for answer in (*before, *after): + _assert_served(answer) + for answer in during: + assert answer.status >= 500 and OVERLOADED_MESSAGE in answer.text, answer + markers: Final = (*healthy, *outage, *recovery) + safeguarded_markers: Final = frozenset( + marker for index, marker in enumerate(markers) if _burst_request(index, marker)[0] == "/v1/messages" + ) + delivered: Final = deliveries(rig.wire) + seen: Final = tuple(sum(sent.marker == marker for sent in delivered) for marker in markers) + assert seen == (1,) * len(seen), seen + expected: Final = { + marker: (True, DANGEROUS_TOOL_USE) if marker in safeguarded_markers else (False, None) for marker in markers + } + for sent in delivered: + assert (sent.safeguards is not None, sent.beta) == expected[sent.marker], sent From 22e17b69faa887cdcf1a639e5e2c1076a65ec0aa Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 09:48:07 -0700 Subject: [PATCH 08/64] test(e2e): event_loop_stalls charges a hold to the longest gc pass that fits it, not to the sum of every pass since the last hold (#47089) * test(e2e): event_loop_stalls charges a hold to the longest gc pass that fits it, not to the sum of every pass since the last hold * test(e2e): event_loop_stalls lists the pass-through post-stream logging step seen once at 0.105 s --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/perf/cases/event_loop_stalls.py | 1 + tests/e2e/perf/setup/event_loop.py | 12 +++++------- tests/e2e_harness/perf/setup/test_event_loop.py | 10 ++++++++++ 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/e2e/perf/cases/event_loop_stalls.py b/tests/e2e/perf/cases/event_loop_stalls.py index 66c99686151..04424de2e8d 100644 --- a/tests/e2e/perf/cases/event_loop_stalls.py +++ b/tests/e2e/perf/cases/event_loop_stalls.py @@ -43,6 +43,7 @@ CASE: Final = replace( KnownStall("RedisBatch.flush", max_count=2), KnownStall("run_coroutine_job", max_count=2), KnownStall("user_api_key_auth.py", max_count=2), + KnownStall("route_streaming_logging_to_handler", max_count=2), ), ), ) diff --git a/tests/e2e/perf/setup/event_loop.py b/tests/e2e/perf/setup/event_loop.py index 00d05c79f73..dee45c374d9 100644 --- a/tests/e2e/perf/setup/event_loop.py +++ b/tests/e2e/perf/setup/event_loop.py @@ -115,13 +115,11 @@ def _parse_event(line: str) -> Stall | GcPass | None: def _with_passes(stall: Stall, passes: tuple[GcPass, ...]) -> Stall: - if not passes: + fitting: Final = tuple(gc_pass for gc_pass in passes if gc_pass.seconds <= stall.seconds) + if not fitting: return stall - return replace( - stall, - gc_seconds=sum(gc_pass.seconds for gc_pass in passes), - gc_generation=max(gc_pass.generation for gc_pass in passes), - ) + longest: Final = max(fitting, key=lambda gc_pass: gc_pass.seconds) + return replace(stall, gc_seconds=longest.seconds, gc_generation=longest.generation) def parse_stalls(lines: Iterable[str]) -> tuple[Stall, ...]: @@ -137,7 +135,7 @@ def parse_stalls(lines: Iterable[str]) -> tuple[Stall, ...]: def _reported(spec: EventLoopStalls, stall: Stall) -> Stall: """A hold the collector accounts for, what is left under the threshold, is the collector's, not the task's.""" - if stall.gc_generation is None or stall.gc_seconds > stall.seconds: + if stall.gc_generation is None: return stall if stall.seconds - stall.gc_seconds >= spec.max_hold_seconds: return stall diff --git a/tests/e2e_harness/perf/setup/test_event_loop.py b/tests/e2e_harness/perf/setup/test_event_loop.py index 4728003db7f..57868305716 100644 --- a/tests/e2e_harness/perf/setup/test_event_loop.py +++ b/tests/e2e_harness/perf/setup/test_event_loop.py @@ -78,6 +78,16 @@ def test_a_hold_well_past_its_collector_pass_stays_on_the_task() -> None: assert [(row.coroutine, row.offending) for row in rows] == [("TaskHandle._run_coro()", True)] +def test_a_hold_is_charged_to_the_longest_collector_pass_that_fits_it_not_to_their_sum() -> None: + thread_side_pass: Final = _GC_PASS.replace("0.711", "0.627") + rows: Final = summarize( + _SPEC, parse_stalls((thread_side_pass, _GC_PASS.replace("0.711", "0.605"), _STREAM.replace("0.412", "0.614"))) + ) + assert [(row.coroutine, row.where, row.max_seconds, row.offending) for row in rows] == [ + ("garbage collection", "gen=2", 0.614, False) + ] + + def test_a_collector_pass_explains_only_the_next_hold_and_only_one_at_least_as_long() -> None: rows: Final = summarize(_SPEC, parse_stalls((_GC_PASS, _STARTUP, _HANDLER))) assert sorted((row.coroutine, row.offending) for row in rows) == [ From 46ae68021eba19645f74042c03238061ff3d6b5d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:48:37 +0000 Subject: [PATCH 09/64] fix(anthropic): carry safeguard_results and stop_details in synthesized message_delta events (#46096) * fix(anthropic): carry safeguard_results and stop_details in synthesized message_delta events When the proxy renders a complete Anthropic Messages response as an SSE stream for a client that asked to stream, the message_delta event carried only stop_reason and stop_sequence. A real Anthropic stream puts safeguard_results and stop_details inside message_delta.delta, so a client such as Claude Code auto mode lost the server-side classifier verdict and a refusal lost its stop_details. Both stream builders now copy those keys from the source response through one shared helper when they are present. * fix(guardrails): keep safeguard_results and stop_details when re-emitting a rewritten Anthropic stream A post-call guardrail that rewrites an Anthropic stream (presidio, bedrock, tool permission, model armor) now carries the upstream message_delta safeguard_results and stop_details into the re-emitted stream. The router fallback message_start no longer leaks those final fields early * fix(guardrails): keep stop_details and safeguard_results in rewritten Anthropic streams without explanations A guardrail that rewrites a /v1/messages stream now carries stop_details (type and category) and the dangerous_tool_use verdicts for the tool calls it still delivers, with every explanation dropped so a masked stream never echoes the unmasked text through those fields. * fix(guardrails): copy only upstream stop_details and forward-compatible safeguard_results into rewritten Anthropic streams A rewritten stream now takes stop_details and safeguard_results from the upstream message_delta alone, so a refusal the upstream sent without stop_details is never given one by the OpenAI round trip. Unknown safeguard types, statuses, reasons and extra fields pass through, a verdict follows a tool call only when its input is unchanged, and tool_permission keeps the explanations since it rewrites no text. * fix(guardrails): apply the safeguard_results and stop_details rules to policy pipeline stream rewrites * fix(router): deliver a refused stream built from a finished agentic turn instead of falling back Websearch, MCP and advisor streams replay a complete reply whose tools already ran, so a refusal fallback would run every tool again. Also give the recovered response's message_start a null stop_details as the API sends it, drop safeguard_results a masking rewrite cannot read, and add the audit cells for synthesized streams. * test(router): annotate the finished-turn refusal stream test locals with Final --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../chat/guardrail_translation/handler.py | 62 +++ .../messages/fake_stream_iterator.py | 15 +- .../messages/streaming_iterator.py | 26 +- .../anthropic/pass_through/messages/utils.py | 91 +++- .../anthropic/pass_through/stream_assembly.py | 25 +- litellm/proxy/guardrails/anthropic_sse.py | 83 +++- .../guardrail_hooks/bedrock_guardrails.py | 2 +- .../model_armor/model_armor.py | 2 +- .../guardrails/guardrail_hooks/presidio.py | 2 +- .../guardrail_hooks/tool_permission.py | 2 +- litellm/router.py | 10 +- ...test_safeguard_results_mcp_auto_execute.py | 51 +- ...t_safeguard_results_synthesized_streams.py | 470 ++++++++++++++++++ .../test_presidio_streaming_output.py | 93 +++- .../test_anthropic_guardrail_handler.py | 95 +++- .../messages/test_fake_stream_iterator.py | 93 ++++ .../messages/test_streaming_iterator.py | 93 ++++ .../pass_through/test_stream_assembly.py | 36 ++ .../guardrail_hooks/test_tool_permission.py | 49 ++ .../proxy/guardrails/test_anthropic_sse.py | 316 ++++++++++++ .../test_fallback_event_handlers.py | 37 ++ 21 files changed, 1621 insertions(+), 32 deletions(-) create mode 100644 tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py create mode 100644 tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py create mode 100644 tests/unit/proxy/guardrails/test_anthropic_sse.py diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 6b6974becac..7654b56ad9a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -30,6 +30,7 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) +from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details from litellm.llms.anthropic.pass_through.stream_assembly import ( build_complete_streaming_response, build_usage_only_response_from_chunks, @@ -1456,6 +1457,9 @@ class AnthropicMessagesHandler(BaseTranslation): return _SSEFieldRewrite("delta", "text", next(replacements)) cls._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta) + cls._redact_ended_stream_message_delta( + responses_so_far, kept_tool_use_ids=frozenset(cls._streamed_tool_use_ids(responses_so_far).values()) + ) @classmethod def _write_ended_stream_tool_call_rewrites( @@ -1523,6 +1527,64 @@ class AnthropicMessagesHandler(BaseTranslation): return None cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use) + streamed_tool_use_ids: Final = cls._streamed_tool_use_ids(responses_so_far) + cls._redact_ended_stream_message_delta( + responses_so_far, + kept_tool_use_ids=frozenset( + tool_use_id for index, tool_use_id in streamed_tool_use_ids.items() if index not in rewrites_by_block + ), + ) + + @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.""" + events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far) + tool_use_blocks: Final = ( + (event.get("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 block.get("type") == "tool_use" + ) + return MappingProxyType( + { + index: tool_use_id + for index, block in tool_use_blocks + if isinstance(index, int) and isinstance(tool_use_id := block.get("id"), str) + } + ) + + @classmethod + def _redact_ended_stream_message_delta( + cls, + responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place + *, + kept_tool_use_ids: frozenset[str], + ) -> None: + """Hold the buffered ``message_delta`` to the rules a re-emitted rewrite follows: a verdict survives + only for a tool call delivered unchanged, and no free-text ``explanation`` describes the text the + guardrail rewrote.""" + + def rewrite_stop_details(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None + if delta is None or "stop_details" not in delta: + return None + return _SSEFieldRewrite( + "delta", "stop_details", redacted_stop_details(delta["stop_details"], keep_explanations=False) + ) + + def rewrite_safeguard_results(event: Mapping[str, object]) -> _SSEFieldRewrite | None: + delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None + if delta is None or "safeguard_results" not in delta: + return None + return _SSEFieldRewrite( + "delta", + "safeguard_results", + filtered_safeguard_results(delta["safeguard_results"], kept_tool_use_ids, keep_explanations=False), + ) + + cls._rewrite_ended_stream_events(responses_so_far, rewrite_stop_details) + cls._rewrite_ended_stream_events(responses_so_far, rewrite_safeguard_results) @staticmethod def _rewrite_ended_stream_events( diff --git a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py index 5930739d445..8d1b358c8f1 100644 --- a/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/fake_stream_iterator.py @@ -12,6 +12,7 @@ import json from collections.abc import Mapping from typing import Any, Final, cast +from litellm.llms.anthropic.pass_through.messages.utils import anthropic_message_delta from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -34,8 +35,13 @@ class FakeAnthropicMessagesStreamIterator: - message_stop """ - def __init__(self, response: AnthropicMessagesResponse): + def __init__( + self, + response: AnthropicMessagesResponse, + message_delta_fields: Mapping[str, object] | None = None, + ): self.response = response + self.message_delta_fields = message_delta_fields self.chunks = self._create_streaming_chunks() self.current_index = 0 @@ -175,14 +181,9 @@ class FakeAnthropicMessagesStreamIterator: cache_read_input_tokens: Final = usage.get("cache_read_input_tokens") if cache_read_input_tokens is not None: delta_usage["cache_read_input_tokens"] = cache_read_input_tokens - safeguard_results: Final = self.response.get("safeguard_results") message_delta: Final = { "type": "message_delta", - "delta": { - "stop_reason": self.response.get("stop_reason"), - "stop_sequence": self.response.get("stop_sequence"), - **({"safeguard_results": safeguard_results} if safeguard_results is not None else {}), - }, + "delta": anthropic_message_delta(response_dict, self.message_delta_fields), "usage": delta_usage, } chunks.append(f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode()) diff --git a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py index e65fe596da7..44db193ea83 100644 --- a/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/messages/streaming_iterator.py @@ -18,7 +18,11 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP -from litellm.llms.anthropic.pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE +from litellm.llms.anthropic.pass_through.messages.utils import ( + INCOMPLETE_STREAM_ERROR_MESSAGE, + MESSAGE_DELTA_OPTIONAL_KEYS, + anthropic_message_delta, +) from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -280,30 +284,32 @@ def anthropic_messages_response_as_sse_events(response: AnthropicMessagesRespons content_events: Final = ( event for index, block in enumerate(content_blocks) for event in _anthropic_content_block_events(index, block) ) - # A real message_start always carries a null stop_reason/stop_sequence and - # a zero output_tokens - those are only known once generation finishes, so - # copying the completed response's final values here would let a client - # treat the message as already finished, or double-count output tokens. + # A real message_start always carries a null stop_reason/stop_sequence/stop_details, + # a zero output_tokens and no safeguard_results - those are only known once + # generation finishes, so copying the completed response's final values here + # would let a client treat the message as already finished, or double-count + # output tokens. message_start_usage: Final = { **(response.get("usage") or {}), "output_tokens": 0, } + message_start_fields: Final = { + key: value for key, value in response.items() if key not in MESSAGE_DELTA_OPTIONAL_KEYS + } message_start_payload: Final = { "type": "message_start", "message": { - **response, + **message_start_fields, "content": [], "stop_reason": None, "stop_sequence": None, + "stop_details": None, "usage": message_start_usage, }, } message_delta_payload: Final = { "type": "message_delta", - "delta": { - "stop_reason": response.get("stop_reason"), - "stop_sequence": response.get("stop_sequence"), - }, + "delta": anthropic_message_delta(response), "usage": response.get("usage") or {}, } return ( diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py index d8b96c1b7a0..d91a5e2fb5c 100644 --- a/litellm/llms/anthropic/pass_through/messages/utils.py +++ b/litellm/llms/anthropic/pass_through/messages/utils.py @@ -2,7 +2,7 @@ from collections.abc import Iterable, Mapping, Sequence from functools import lru_cache from typing import TYPE_CHECKING, Final, cast, get_type_hints -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import as_str_mapping @@ -49,6 +49,95 @@ def refusal_stop_details(explanation: str | None) -> AnthropicStopDetails: return AnthropicStopDetails(type="refusal", category=None, explanation=explanation) +MESSAGE_DELTA_OPTIONAL_KEYS: Final = ("stop_details", "safeguard_results") + + +def anthropic_message_delta( + response: Mapping[str, object], optional_fields: Mapping[str, object] | None = None +) -> Mapping[str, object]: + """ + The ``delta`` of the ``message_delta`` event synthesized from a complete + Anthropic Messages response. A real stream carries ``stop_details`` and + ``safeguard_results`` there, so a client that asked to stream reads them + from this event and nowhere else. ``optional_fields`` replaces the + response's own when it was rebuilt from a stream that carried them. + """ + optional_source: Final = response if optional_fields is None else optional_fields + present_optional: Final = { + key: value for key in MESSAGE_DELTA_OPTIONAL_KEYS if (value := optional_source.get(key)) is not None + } + return { + "stop_reason": response.get("stop_reason"), + "stop_sequence": response.get("stop_sequence"), + **present_optional, + } + + +_SAFEGUARD_RESULT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def redacted_stop_details(raw: object, keep_explanations: bool) -> object: + """``stop_details`` as a rewritten stream may deliver it: its free-text ``explanation`` can quote + the text a guardrail rewrote, so it is nulled unless the rewrite changed no text.""" + stop_details: Final = as_str_mapping(raw) + if keep_explanations or stop_details is None or "explanation" not in stop_details: + return raw + return {**stop_details, "explanation": None} + + +def _without_explanation(raw: object) -> object: + mapping: Final = as_str_mapping(raw) + if mapping is None: + return raw + return {key: value for key, value in mapping.items() if key != "explanation"} + + +def _status_without_explanations(raw: object) -> object: + status: Final = as_str_mapping(_without_explanation(raw)) + if status is None: + return raw + tool_uses: Final = as_str_mapping(status.get("tool_uses")) + if tool_uses is None: + return status + return { + **status, + "tool_uses": {tool_use_id: _without_explanation(verdict) for tool_use_id, verdict in tool_uses.items()}, + } + + +def _result_without_explanations(raw: object) -> object: + result: Final = as_str_mapping(_without_explanation(raw)) + if result is None: + return raw + if "status" not in result: + return result + return {**result, "status": _status_without_explanations(result["status"])} + + +def _kept_tool_use_verdicts(raw: object, kept_tool_use_ids: frozenset[str]) -> object: + result: Final = as_str_mapping(raw) + status: Final = as_str_mapping(result.get("status")) if result is not None else None + tool_uses: Final = as_str_mapping(status.get("tool_uses")) if status is not None else None + if result is None or status is None or tool_uses is None: + return raw + kept_tool_uses: Final = { + tool_use_id: verdict for tool_use_id, verdict in tool_uses.items() if tool_use_id in kept_tool_use_ids + } + return {**result, "status": {**status, "tool_uses": kept_tool_uses}} + + +def filtered_safeguard_results(raw: object, kept_tool_use_ids: frozenset[str], keep_explanations: bool) -> object: + """``safeguard_results`` as a rewritten stream may deliver them: a ``tool_uses`` verdict survives only + for a tool call the client gets with the input the classifier judged, and the free-text ``explanation`` of a + result, its status and each verdict is dropped unless the rewrite changed no text.""" + try: + results: Final = _SAFEGUARD_RESULT_ITEMS.validate_python(raw) + except ValidationError: + return raw if keep_explanations else None + kept: Final = tuple(_kept_tool_use_verdicts(result, kept_tool_use_ids) for result in results) + return kept if keep_explanations else tuple(_result_without_explanations(result) for result in kept) + + def _mapping_field(container: object, key: str) -> object | None: """One key of a raw provider payload, or None when the payload is not a mapping.""" if not isinstance(container, Mapping): diff --git a/litellm/llms/anthropic/pass_through/stream_assembly.py b/litellm/llms/anthropic/pass_through/stream_assembly.py index 89961bbe255..90cb9f4c335 100644 --- a/litellm/llms/anthropic/pass_through/stream_assembly.py +++ b/litellm/llms/anthropic/pass_through/stream_assembly.py @@ -1,16 +1,18 @@ import codecs import json from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import as_str_mapping, map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.handler import ( ModelResponseIterator as AnthropicModelResponseIterator, ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.pass_through.messages.utils import MESSAGE_DELTA_OPTIONAL_KEYS from litellm.types.utils import ( Choices, Message, @@ -447,6 +449,27 @@ def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None: ) +def anthropic_message_delta_fields(all_chunks: Sequence[object]) -> Mapping[str, object]: + """The ``stop_details`` and ``safeguard_results`` of the stream's ``message_delta``. + + The assembled OpenAI shape has no slot for them, and a client that asked to stream reads them + from that event and nowhere else (a Claude Code auto-mode session takes its classifier verdict + from ``safeguard_results``), so a caller re-emitting a rewritten response merges them back in. + """ + sse_stream: Final = _joined_sse_stream(all_chunks) + if sse_stream is None: + return MappingProxyType({}) + message_delta: Final = next( + (event for event in reversed(_parsed_sse_events(sse_stream)) if event.get("type") == "message_delta"), None + ) + delta: Final = as_str_mapping(message_delta.get("delta")) if message_delta is not None else None + if delta is None: + return MappingProxyType({}) + return MappingProxyType( + {key: value for key in MESSAGE_DELTA_OPTIONAL_KEYS if (value := delta.get(key)) is not None} + ) + + def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool: """Whether raw SSE frames are Anthropic Messages events. diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index c3a6f259f95..2baeabdce7b 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -8,8 +8,13 @@ hook scan such a stream, and re-emit it when the guardrail rewrote the response. from __future__ import annotations import json +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details +from litellm.llms.anthropic.pass_through.stream_assembly import anthropic_message_delta_fields from litellm.llms.anthropic.pass_through.stream_assembly import ( assemble_anthropic_sse_stream as assemble_anthropic_sse_stream, # noqa: PLC0414 # public re-export ) @@ -50,7 +55,62 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( # noqa: E402 ) -def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]: +def _tool_use_inputs(content: Sequence[object]) -> Mapping[str, object]: + blocks: Final = (block for raw in content if (block := as_str_mapping(raw)) is not None) + return MappingProxyType( + { + tool_use_id: block.get("input") + for block in blocks + if block.get("type") == "tool_use" and isinstance(tool_use_id := block.get("id"), str) + } + ) + + +def _unchanged_tool_use_ids(source_content: Sequence[object], delivered_content: Sequence[object]) -> frozenset[str]: + source_inputs: Final = _tool_use_inputs(source_content) + return frozenset( + tool_use_id + for tool_use_id, tool_input in _tool_use_inputs(delivered_content).items() + if tool_use_id in source_inputs and source_inputs[tool_use_id] == tool_input + ) + + +def _rewritten_message_delta_fields( + source_chunks: Sequence[object], unchanged_tool_use_ids: frozenset[str], keep_explanations: bool +) -> Mapping[str, object]: + source_fields: Final = anthropic_message_delta_fields(source_chunks) + stop_details: Final = source_fields.get("stop_details") + safeguard_results: Final = source_fields.get("safeguard_results") + return MappingProxyType( + { + **( + {"stop_details": redacted_stop_details(stop_details, keep_explanations)} + if stop_details is not None + else {} + ), + **( + { + "safeguard_results": filtered_safeguard_results( + safeguard_results, unchanged_tool_use_ids, keep_explanations + ) + } + if safeguard_results is not None + else {} + ), + } + ) + + +def anthropic_sse_chunks_from_response( + assembled: ModelResponse, source_chunks: Sequence[object] = (), *, keep_explanations: bool = False +) -> tuple[bytes, ...]: + """Re-emit a rewritten response as the SSE events of the ``source_chunks`` stream it replaces. + + ``message_delta`` carries the source stream's ``stop_details`` and ``safeguard_results`` and + nothing the OpenAI round trip made up, with a verdict kept only for a tool call delivered with + the input the classifier judged. Their free-text ``explanation`` was written about the text the + guardrail rewrote, so it is dropped unless the caller rewrites no text (``keep_explanations``). + """ from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, ) @@ -58,7 +118,22 @@ def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, FakeAnthropicMessagesStreamIterator, ) - anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( - response=assembled + adapter: Final = LiteLLMAnthropicMessagesAdapter() + anthropic_response: Final = adapter.translate_openai_response_to_anthropic(response=assembled) + source_response: Final = assemble_anthropic_sse_stream(source_chunks) + source_content: Final = ( + adapter.translate_openai_response_to_anthropic(response=source_response).get("content") + if source_response is not None + else None + ) + unchanged_tool_use_ids: Final = _unchanged_tool_use_ids( + source_content or (), anthropic_response.get("content") or () + ) + return tuple( + FakeAnthropicMessagesStreamIterator( + response=anthropic_response, + message_delta_fields=_rewritten_message_delta_fields( + source_chunks, unchanged_tool_use_ids, keep_explanations + ), + ).chunks ) - return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index b92020c89be..0672ccf6a66 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -3058,7 +3058,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################################### if raw_sse: for sse_chunk in ( - anthropic_sse_chunks_from_response(assembled_model_response) + anthropic_sse_chunks_from_response(assembled_model_response, all_chunks) if model_response_text(assembled_model_response) != pre_guardrail_text else all_chunks ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index e590e5c2e43..b9871b1e276 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1263,7 +1263,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Return sanitized stream if surface is _StreamSurface.ANTHROPIC_MESSAGES: - for sse_chunk in anthropic_sse_chunks_from_response(assembled_response): + for sse_chunk in anthropic_sse_chunks_from_response(assembled_response, all_chunks): yield sse_chunk return mock_response: Final = MockResponseIterator(model_response=assembled_response) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 01a9259203a..f21aabc7f33 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1532,7 +1532,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): await self._process_response_for_pii(response=assembled, request_data=request_data, mode="mask") if model_response_text(assembled) == original_text: return chunks - return anthropic_sse_chunks_from_response(assembled) + return anthropic_sse_chunks_from_response(assembled, chunks) @staticmethod def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str]) -> bytes: diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index ebd50c8cd39..fa92a0c90e0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -971,7 +971,7 @@ class ToolPermissionGuardrail(CustomGuardrail): return self._modify_response_with_permission_errors(anthropic_response, anthropic_denials) - for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response): + for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response, all_chunks, keep_explanations=True): yield sse_chunk def _check_assembled_stream( diff --git a/litellm/router.py b/litellm/router.py index 99adc6c5fa6..1b3857dfb71 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6055,6 +6055,9 @@ class Router: including pings and provider error frames, is forwarded live and verbatim instead. """ + from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( aclose_if_supported, aclose_superseded_attempt, @@ -6066,6 +6069,7 @@ class Router: ) source_iterator: Final = response + replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator) async def close_source() -> None: with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): @@ -6126,7 +6130,11 @@ class Router: if not has_generated_content and error_event is None else None ) - if refusal_stop_details is not None and self._refusal_fallback_available(model, initial_kwargs): + if ( + refusal_stop_details is not None + and not replays_a_finished_agentic_turn + and self._refusal_fallback_available(model, initial_kwargs) + ): refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) raise MidStreamFallbackError( message=refusal_error.message, diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py index c9da24484e2..57ab6e0e745 100644 --- a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_mcp_auto_execute.py @@ -4,28 +4,35 @@ with no verdicts.""" from __future__ import annotations +import json from collections.abc import Callable, Iterator, Mapping +from itertools import chain from typing import Final import pytest -from integration._support.client import object_value +from integration._support.anthropic_refusal import DOCS_STOP_DETAILS +from integration._support.client import list_value, object_value from integration._support.mcp import McpPeer, mcp_peer, tool_calls from integration._support.process import graceful_stop_seconds from integration._support.safeguards_rig import ( BASH_TOOL, CHAT_MODEL, + CLAUDE_MODEL, Responder, Rig, auto_mode_body, available, chat_backend, json_body, + native_message, nonce, request_json, safeguards_rig, + sse_events, stream_results, stream_tool_use_ids, tool_id, + tool_use_block, tool_use_ids, ) from integration._support.wire import Reply, Request @@ -115,3 +122,45 @@ def test_auto_executed_mcp_tool_costs_no_classifier_call(rig: Rig, peer: McpPeer received: Final = rig.received(request_nonce) assert rig.classifier_calls(received) == (), [request.target for request in received] assert len(rig.backend_calls(received)) == 2, [request.target for request in received] + + +def carries_a_tool_result_block(request: Request) -> bool: + messages: Final = list_value(request_json(request).get("messages", [])) + contents: Final = tuple(object_value(message).get("content") for message in messages) + blocks: Final = chain.from_iterable(content for content in contents if isinstance(content, list)) + return any(object_value(block).get("type") == "tool_result" for block in blocks) + + +def claude_mcp_tool_then_refusal(request_nonce: str) -> Responder: + identity: Final = tool_id(request_nonce) + tool_turn: Final = native_message( + identity, + content=(tool_use_block(f"{identity}m", name=f"{MCP_ALIAS}-add", arguments={"a": 2, "b": 3}),), + model="claude-sonnet-4-6", + ) + refusal_turn: Final = { + **native_message(f"{identity}r", content=(), stop_reason="refusal", model="claude-sonnet-4-6"), + "stop_details": DOCS_STOP_DETAILS, + } + + def respond(request: Request) -> Reply: + turn: Final = refusal_turn if carries_a_tool_result_block(request) else tool_turn + return Reply(body=json.dumps(turn).encode()) + + return respond + + +def test_auto_executed_mcp_loop_streams_the_final_turn_refusal_stop_details(rig: Rig, peer: McpPeer) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, claude_mcp_tool_then_refusal(request_nonce)) + peer.drain() + status, raw = rig.stream(auto_mode_body(CLAUDE_MODEL, request_nonce, tools=[dict(AUTO_EXECUTED), dict(BASH_TOOL)])) + assert status == 200, raw + assert executed_tools(peer) == ("add",) + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "message_delta" + ) + assert len(deltas) == 1, raw + assert deltas[0]["stop_reason"] == "refusal", raw + assert deltas[0].get("stop_details") == DOCS_STOP_DETAILS, raw + assert len(rig.backend_calls(rig.received(request_nonce))) == 2 diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py new file mode 100644 index 00000000000..99a6dd18d48 --- /dev/null +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_synthesized_streams.py @@ -0,0 +1,470 @@ +"""A /v1/messages stream LiteLLM builds from a complete reply (web search interception, a fallback's answer, a guardrail +rewrite) carries the ``stop_details`` and ``safeguard_results`` a real Anthropic stream puts on its ``message_delta``.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Iterator, Mapping, Sequence +from typing import Final +from urllib.parse import urlsplit + +import pytest +from integration._support.client import eventually, object_value +from integration._support.process import graceful_stop_seconds +from integration._support.safeguards_rig import ( + BASH_TOOL, + CLAUDE_MODEL, + CLAUDE_TARGET, + TAVILY_TARGET, + Responder, + Rig, + auto_mode_body, + available, + deployment, + native_message, + native_stream_frames, + nonce, + request_json, + safeguards_rig, + sse_events, + stream_results, + stream_tool_use_ids, + text_block, + tool_id, + tool_use_block, + unavailable, +) +from integration._support.wire import Reply, Request +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +CLAUDE: Final = "anthropic/claude-sonnet-4-6" +PRIMARY_MODEL: Final = "synthesized-primary" +FALLBACK_MODEL: Final = "synthesized-fallback" +REFUSED_MODEL: Final = "synthesized-refused" +MOCK_MODEL: Final = "synthesized-mock" +PIPELINE_MODEL: Final = "synthesized-pipeline" +AZURE_REFUSED_MODEL: Final = "synthesized-azure-refused" +CACHED_FALLBACK_MODEL: Final = "synthesized-cached-fallback" +PRIMARY_TARGET: Final = "/primary/v1/messages" +FALLBACK_TARGET: Final = "/fallback/v1/messages" +REFUSED_TARGET: Final = "/refused/v1/messages" +PIPELINE_TARGET: Final = "/pipeline/v1/messages" +AZURE_REFUSED_TARGET: Final = "/azure/anthropic/v1/messages" +CACHED_FALLBACK_TARGET: Final = "/cached-fallback/v1/messages" +MOCK_ANSWER: Final = "mocked fallback answer" +TOOL_PERMISSION: Final = "synthesized-tool-permission" +PIPELINE_REWRITE: Final = "synthesized-pipeline-rewrite" +SECRET: Final = "SECRET" +MASKED: Final = "[MASKED]" +REWRITE_CODE: Final = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " if input_type != 'response':\n" + " return allow()\n" + f" texts = [text.replace({SECRET!r}, {MASKED!r}) for text in inputs.get('texts') or []]\n" + " calls = [\n" + " {'id': call['id'], 'type': 'function', 'function': {'name': call['function']['name'],\n" + f" 'arguments': call['function']['arguments'].replace({SECRET!r}, {MASKED!r})}}}}\n" + " for call in inputs.get('tool_calls') or []\n" + " ]\n" + " return modify(texts=texts, tool_calls=calls)\n" +) +WEB_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = {"type": "web_search_20250305", "name": "web_search", "max_uses": 1} +READ_TOOL: Final[Mapping[str, JsonValue]] = { + "name": "Read", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, +} +NO_CACHE: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}} +REFUSAL_STOP_DETAILS: Final[Mapping[str, JsonValue]] = { + "type": "refusal", + "category": "cyber", + "explanation": "This request was declined because it could enable cyber harm.", +} +_FRAME_BOUNDARY: Final = re.compile(rb"(?<=\n\n)") + +LITELLM_SETTINGS: Final[Mapping[str, JsonValue]] = { + "callbacks": ["websearch_interception"], + "websearch_interception_params": {"enabled_providers": ["anthropic"], "search_tool_name": "tavily-search"}, +} + + +def extra_models(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return ( + deployment(PRIMARY_MODEL, CLAUDE, f"{wire_url}/primary"), + deployment(FALLBACK_MODEL, CLAUDE, f"{wire_url}/fallback"), + deployment(REFUSED_MODEL, CLAUDE, f"{wire_url}/refused"), + deployment(MOCK_MODEL, CLAUDE, f"{wire_url}/mock", mock_response=MOCK_ANSWER), + deployment(PIPELINE_MODEL, CLAUDE, f"{wire_url}/pipeline"), + deployment(AZURE_REFUSED_MODEL, "azure_ai/claude-sonnet-4-6", f"{wire_url}/azure/anthropic"), + deployment(CACHED_FALLBACK_MODEL, CLAUDE, f"{wire_url}/cached-fallback"), + ) + + +def top_level(wire_url: str) -> Mapping[str, JsonValue]: + tavily: Final[Mapping[str, JsonValue]] = { + "search_tool_name": "tavily-search", + "litellm_params": { + "search_provider": "tavily", + "api_key": "synthetic-tavily-key", + "api_base": f"{wire_url}/tavily", + }, + } + tool_permission: Final[Mapping[str, JsonValue]] = { + "guardrail_name": TOOL_PERMISSION, + "litellm_params": { + "guardrail": "tool_permission", + "mode": "post_call", + "default_on": False, + "default_action": "allow", + "on_disallowed_action": "rewrite", + "rules": [{"id": "deny_bash", "tool_name": "^Bash$", "decision": "deny"}], + }, + } + pipeline_rewrite: Final[Mapping[str, JsonValue]] = { + "guardrail_name": PIPELINE_REWRITE, + "litellm_params": { + "guardrail": "custom_code", + "mode": "post_call", + "default_on": False, + "custom_code": REWRITE_CODE, + }, + } + masking_policy: Final[Mapping[str, JsonValue]] = { + "guardrails": {"add": [PIPELINE_REWRITE]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": PIPELINE_REWRITE, "on_pass": "allow", "on_fail": "block"}], + }, + } + return { + "search_tools": [dict(tavily)], + "guardrails": [dict(tool_permission), dict(pipeline_rewrite)], + "policies": {"synthesized-masking": dict(masking_policy)}, + "policy_attachments": [{"policy": "synthesized-masking", "models": [PIPELINE_MODEL]}], + "router_settings": { + "disable_cooldowns": True, + "content_policy_fallbacks": [ + {PRIMARY_MODEL: [FALLBACK_MODEL]}, + {REFUSED_MODEL: [MOCK_MODEL]}, + {AZURE_REFUSED_MODEL: [CACHED_FALLBACK_MODEL]}, + ], + }, + } + + +@pytest.fixture(scope="module") +def synthesizing(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("synthesized-streams") + with safeguards_rig( + directory, litellm_settings=LITELLM_SETTINGS, top_level=top_level, extra_models=extra_models + ) as rig: + yield rig + + +def claude_backend(message: Mapping[str, JsonValue], delta_fields: Mapping[str, JsonValue]) -> Responder: + def respond(request: Request) -> Reply: + if request_json(request).get("stream") is True: + return Reply( + content_type="text/event-stream", chunks=native_stream_frames(message, delta_extra=delta_fields) + ) + return Reply(body=json.dumps({**message, **delta_fields}).encode()) + + return respond + + +def by_target(backends: Mapping[str, Responder]) -> Responder: + def respond(request: Request) -> Reply: + return backends[urlsplit(request.target).path](request) + + return respond + + +def refusal(identity: str) -> Mapping[str, JsonValue]: + return native_message(identity, content=(), stop_reason="refusal", model="claude-sonnet-4-6") + + +def answer(identity: str, text: str) -> Mapping[str, JsonValue]: + return native_message(identity, content=(text_block(text),), stop_reason="end_turn", model="claude-sonnet-4-6") + + +def refusal_fields() -> Mapping[str, JsonValue]: + return {"stop_details": dict(REFUSAL_STOP_DETAILS), "safeguard_results": list(available({}))} + + +def search_body(model: str, request_nonce: str) -> Mapping[str, JsonValue]: + return auto_mode_body(model, request_nonce, tools=[dict(WEB_SEARCH_TOOL)], prompt="search the web for", **NO_CACHE) + + +def message_delta(raw: bytes) -> Mapping[str, JsonValue]: + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "message_delta" + ) + assert len(deltas) == 1, raw + return deltas[0] + + +def message_start(raw: bytes) -> Mapping[str, JsonValue]: + starts: Final = tuple( + object_value(event["message"]) for event in sse_events(raw) if event.get("type") == "message_start" + ) + assert len(starts) == 1, raw + return starts[0] + + +def streamed_text(raw: bytes) -> str: + deltas: Final = tuple( + object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "content_block_delta" + ) + return "".join(str(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + + +def delta_frame(raw: bytes) -> bytes: + frames: Final = tuple(frame for frame in _FRAME_BOUNDARY.split(raw) if b'"message_delta"' in frame) + assert len(frames) == 1, raw + return frames[0] + + +def targets(rig: Rig, request_nonce: str) -> tuple[str, ...]: + return tuple(urlsplit(request.target).path for request in rig.received(request_nonce)) + + +def test_intercepted_refusal_stream_delivers_the_upstream_stop_details(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + synthesizing.register(request_nonce, claude_backend(refusal(tool_id(request_nonce)), refusal_fields())) + status, raw = synthesizing.stream(search_body(CLAUDE_MODEL, request_nonce)) + assert status == 200, raw + assert message_delta(raw) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": dict(REFUSAL_STOP_DETAILS), + "safeguard_results": list(available({})), + }, raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_intercepted_refusal_stream_with_a_content_policy_fallback_delivers_the_refusal_without_falling_back( + synthesizing: Rig, +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + synthesizing.register( + request_nonce, + by_target( + { + PRIMARY_TARGET: claude_backend(refusal(identity), refusal_fields()), + FALLBACK_TARGET: claude_backend(answer(f"{identity}f", f"fallback answer {request_nonce}"), {}), + } + ), + ) + status, raw = synthesizing.stream(search_body(PRIMARY_MODEL, request_nonce)) + assert status == 200, raw + assert message_delta(raw) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": dict(REFUSAL_STOP_DETAILS), + "safeguard_results": list(available({})), + }, raw + assert targets(synthesizing, request_nonce) == (PRIMARY_TARGET,) + + +def test_tool_permission_rewrite_keeps_only_the_delivered_tool_verdict_with_its_explanation(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + denied: Final = tool_id(request_nonce) + allowed: Final = f"{denied}r" + allowed_verdict: Final[Mapping[str, JsonValue]] = { + "type": "evaluated", + "outcome": "flagged", + "explanation": f"reads a secret {allowed}", + } + verdicts: Final[Mapping[str, JsonValue]] = { + denied: {"type": "evaluated", "outcome": "flagged", "explanation": f"destructive {denied}"}, + allowed: dict(allowed_verdict), + } + message: Final = native_message( + denied, + content=( + text_block("Cleaning up."), + tool_use_block(denied), + tool_use_block(allowed, name="Read", arguments={"path": "/work/repo/.env"}), + ), + model="claude-sonnet-4-6", + ) + synthesizing.register(request_nonce, claude_backend(message, {"safeguard_results": list(available(verdicts))})) + body: Final = auto_mode_body( + CLAUDE_MODEL, request_nonce, tools=[dict(BASH_TOOL), dict(READ_TOOL)], guardrails=[TOOL_PERMISSION], **NO_CACHE + ) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert stream_tool_use_ids(raw) == (allowed,), raw + assert stream_results(raw) == available({allowed: dict(allowed_verdict)}), raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_intercepted_answer_stream_keeps_its_safeguard_results(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + text: Final = f"found it {request_nonce}" + results: Final = list(available({})) + synthesizing.register( + request_nonce, + claude_backend(answer(tool_id(request_nonce), text), {"stop_details": None, "safeguard_results": results}), + ) + status, raw = synthesizing.stream(search_body(CLAUDE_MODEL, request_nonce)) + assert status == 200, raw + assert streamed_text(raw) == text, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": results}, raw + received: Final = targets(synthesizing, request_nonce) + assert received == (CLAUDE_TARGET,) and TAVILY_TARGET not in received, received + + +def test_native_refusal_stream_without_a_guardrail_keeps_the_upstream_message_delta_bytes(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + message: Final = native_message( + identity, content=(text_block("Partial."),), stop_reason="refusal", model="claude-sonnet-4-6" + ) + fields: Final = {**refusal_fields(), "safeguard_results": list(unavailable((identity,), "error"))} + sent: Final = native_stream_frames(message, delta_extra=fields) + synthesizing.register(request_nonce, claude_backend(message, fields)) + status, raw = synthesizing.stream(auto_mode_body(CLAUDE_MODEL, request_nonce, **NO_CACHE)) + assert status == 200, raw + assert delta_frame(raw) == delta_frame(b"".join(sent)), raw + assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,) + + +def test_native_refusal_falling_back_to_a_mock_reply_renders_a_fresh_message_lifecycle(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + synthesizing.register(request_nonce, claude_backend(refusal(tool_id(request_nonce)), refusal_fields())) + status, raw = synthesizing.stream(auto_mode_body(REFUSED_MODEL, request_nonce, **NO_CACHE)) + assert status == 200, raw + started: Final = message_start(raw) + assert started["stop_reason"] is None, raw + assert started.get("stop_details") is None and "safeguard_results" not in started, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None}, raw + assert streamed_text(raw) == MOCK_ANSWER, raw + assert b'"refusal"' not in raw, raw + assert targets(synthesizing, request_nonce) == (REFUSED_TARGET,) + + +def test_native_refusal_falling_back_to_a_cached_intercepted_reply_streams_its_fields_in_message_delta( + synthesizing: Rig, +) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + text: Final = f"cached answer {request_nonce}" + results: Final = list(available({})) + synthesizing.register( + request_nonce, + by_target( + { + AZURE_REFUSED_TARGET: claude_backend(refusal(identity), refusal_fields()), + CACHED_FALLBACK_TARGET: claude_backend( + answer(f"{identity}c", text), {"stop_details": None, "safeguard_results": results} + ), + } + ), + ) + cacheable: Final = auto_mode_body( + CACHED_FALLBACK_MODEL, request_nonce, tools=[dict(WEB_SEARCH_TOOL)], prompt="search the web for" + ) + + def prime() -> tuple[str, ...]: + primed: Final = synthesizing.post({**cacheable, "stream": False}) + assert primed.status_code == 200, primed.text + return targets(synthesizing, request_nonce) + + first_prime: Final = prime() + assert first_prime == (CACHED_FALLBACK_TARGET,), first_prime + eventually(prime, lambda received: received == ()) + status, raw = synthesizing.stream({**cacheable, "model": AZURE_REFUSED_MODEL}) + assert status == 200, raw + assert streamed_text(raw) == text, raw + started: Final = message_start(raw) + assert started["stop_details"] is None and "safeguard_results" not in started, raw + assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": results}, raw + assert targets(synthesizing, request_nonce) == (AZURE_REFUSED_TARGET,) + + +def flagged_verdict(explanation: str) -> Mapping[str, JsonValue]: + return {"type": "evaluated", "outcome": "flagged", "explanation": explanation} + + +FLAGGED_WITHOUT_EXPLANATION: Final[Mapping[str, JsonValue]] = {"type": "evaluated", "outcome": "flagged"} + + +def masked_refusal(identity: str) -> tuple[Mapping[str, JsonValue], Mapping[str, JsonValue], Mapping[str, JsonValue]]: + message: Final = native_message( + identity, content=(text_block(f"Partial {SECRET}."),), stop_reason="refusal", model="claude-sonnet-4-6" + ) + expected: Final[Mapping[str, JsonValue]] = { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {**REFUSAL_STOP_DETAILS, "explanation": None}, + "safeguard_results": list(available({})), + } + return message, refusal_fields(), expected + + +def masked_tool_turn( + identity: str, +) -> tuple[Mapping[str, JsonValue], Mapping[str, JsonValue], Mapping[str, JsonValue]]: + message: Final = native_message( + identity, + content=(text_block(f"Reading {SECRET}."), tool_use_block(identity, name="Read", arguments={"path": "/work"})), + model="claude-sonnet-4-6", + ) + sent: Final[Mapping[str, JsonValue]] = { + "safeguard_results": list(available({identity: dict(flagged_verdict(f"reads {identity}"))})) + } + expected: Final[Mapping[str, JsonValue]] = { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": list(available({identity: dict(FLAGGED_WITHOUT_EXPLANATION)})), + } + return message, sent, expected + + +MASKED_STREAMS: Final = {"refusal_stop_details": masked_refusal, "verdict_explanation": masked_tool_turn} + + +@pytest.mark.parametrize("shape", tuple(MASKED_STREAMS)) +def test_pipeline_text_rewrite_drops_explanations_from_the_upstream_message_delta( + synthesizing: Rig, shape: str +) -> None: + request_nonce: Final = nonce() + message, sent, expected = MASKED_STREAMS[shape](tool_id(request_nonce)) + synthesizing.register(request_nonce, claude_backend(message, sent)) + body: Final = auto_mode_body(PIPELINE_MODEL, request_nonce, tools=[dict(READ_TOOL)], **NO_CACHE) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert MASKED in streamed_text(raw) and SECRET not in streamed_text(raw), raw + assert message_delta(raw) == expected, raw + assert targets(synthesizing, request_nonce) == (PIPELINE_TARGET,) + + +def test_pipeline_tool_call_rewrite_keeps_only_the_untouched_tool_verdict(synthesizing: Rig) -> None: + request_nonce: Final = nonce() + rewritten: Final = tool_id(request_nonce) + untouched: Final = f"{rewritten}r" + message: Final = native_message( + rewritten, + content=( + tool_use_block(rewritten, arguments={"command": f"cat {SECRET}"}), + tool_use_block(untouched, name="Read", arguments={"path": "/work"}), + ), + model="claude-sonnet-4-6", + ) + verdicts: Final[Mapping[str, JsonValue]] = { + rewritten: dict(flagged_verdict(f"reads a file {rewritten}")), + untouched: dict(flagged_verdict(f"reads a directory {untouched}")), + } + synthesizing.register(request_nonce, claude_backend(message, {"safeguard_results": list(available(verdicts))})) + body: Final = auto_mode_body(PIPELINE_MODEL, request_nonce, tools=[dict(BASH_TOOL), dict(READ_TOOL)], **NO_CACHE) + status, raw = synthesizing.stream(body) + assert status == 200, raw + assert stream_tool_use_ids(raw) == (rewritten, untouched), raw + assert SECRET.encode() not in raw and MASKED.encode() in raw, raw + assert stream_results(raw) == available({untouched: dict(FLAGGED_WITHOUT_EXPLANATION)}), raw + assert targets(synthesizing, request_nonce) == (PIPELINE_TARGET,) diff --git a/tests/integration/observability/test_presidio_streaming_output.py b/tests/integration/observability/test_presidio_streaming_output.py index d784fecd146..d69866ec43d 100644 --- a/tests/integration/observability/test_presidio_streaming_output.py +++ b/tests/integration/observability/test_presidio_streaming_output.py @@ -2,7 +2,7 @@ import json import re import threading import uuid -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager from dataclasses import dataclass @@ -10,6 +10,7 @@ from pathlib import Path from typing import Final import httpx +import pytest import yaml from integration._support.client import Gateway, eventually from integration._support.process import OwnedProxy, owned_proxy_process @@ -457,6 +458,96 @@ def test_anthropic_messages_per_request_guardrails_selects_masking(gateway: Gate assert len(rig.analyzer.drain()) == 1 +READ_TOOL_USE_ID: Final = "toolu_presidio_read" +READ_INPUT: Final = {"path": "/work/report.txt"} + + +def anthropic_tool_turn_stream(identity: str, text: str, delta: Mapping[str, object]) -> tuple[bytes, ...]: + opening, text_start, text_delta, text_stop, _, closing = anthropic_stream(identity, text) + tool_use: Final = {"type": "tool_use", "id": READ_TOOL_USE_ID, "name": "Read", "input": {}} + return ( + opening, + text_start, + text_delta, + text_stop, + anthropic_frame("content_block_start", {"type": "content_block_start", "index": 1, "content_block": tool_use}), + anthropic_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": json.dumps(READ_INPUT)}, + }, + ), + anthropic_frame("content_block_stop", {"type": "content_block_stop", "index": 1}), + anthropic_frame("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 4}}), + closing, + ) + + +def read_verdicts(verdict: Mapping[str, object]) -> Sequence[Mapping[str, object]]: + return [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {READ_TOOL_USE_ID: verdict}}}] + + +def anthropic_deltas(received: Received, event_type: str) -> tuple[Mapping[str, object], ...]: + events: Final = tuple( + json.loads(line.removeprefix("data: ")) for line in received.text.split("\n") if line.startswith("data: ") + ) + return tuple(event["delta"] for event in events if event.get("type") == event_type) + + +def anthropic_message_delta(received: Received) -> Mapping[str, object]: + deltas: Final = anthropic_deltas(received, "message_delta") + assert len(deltas) == 1, received.text + return deltas[0] + + +def anthropic_text_deltas(received: Received) -> str: + deltas: Final = anthropic_deltas(received, "content_block_delta") + return "".join(str(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + + +@pytest.mark.parametrize( + ("delta", "field", "kept"), + [ + ( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": read_verdicts( + {"type": "evaluated", "outcome": "flagged", "explanation": f"{PERSON} asked for the report"} + ), + }, + "safeguard_results", + read_verdicts({"type": "evaluated", "outcome": "flagged"}), + ), + ( + { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": f"declined for {PERSON}"}, + }, + "stop_details", + {"type": "refusal", "category": "cyber", "explanation": None}, + ), + ], + ids=["verdict_explanation", "stop_details_explanation"], +) +def test_anthropic_messages_masked_stream_keeps_upstream_safeguard_fields_without_explanations( + gateway: Gateway, tmp_path: Path, delta: Mapping[str, object], field: str, kept: object +) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + provider: Final = anthropic_provider(anthropic_tool_turn_stream(identity, f"{PERSON} designed it.", delta)) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert received.status == 200, received.text + assert anthropic_text_deltas(received) == f"{MASK} designed it.", received.text + assert anthropic_message_id(received) == identity, received.text + assert anthropic_message_delta(received).get(field) == kept, received.text + assert PERSON not in received.text, received.text + assert len(rig.upstream.drain()) == 1 + + def openai_provider(identity: str) -> Callable[[Request], Reply]: def provider(request: Request) -> Reply: assert request.target == "/v1/chat/completions", request.target diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 892736b5043..0aca60351bf 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -6,7 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic """ import json -from typing import Any, Literal, Optional +from typing import Any, Literal, Optional, Final from unittest.mock import MagicMock, patch import pytest @@ -18,7 +18,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import Streami from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs class MockPassThroughGuardrail(CustomGuardrail): @@ -383,6 +383,97 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert '"name": "lookup_fruit"' not in raw assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"} + @staticmethod + def _message_delta(chunks: list[bytes]) -> dict[str, object]: + lines: Final = b"".join(chunks).decode().split("\n") + events: Final = (json.loads(line[len("data:") :]) for line in lines if line.startswith("data:")) + return next(event["delta"] for event in events if event.get("type") == "message_delta") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_text_rewrite_drops_the_upstream_explanations(self): + delta_fields: Final = { + "stop_details": {"type": "refusal", "category": "cyber", "explanation": "the reply says world"}, + "safeguard_results": [ + {"type": "prompt_injection", "status": {"type": "available", "explanation": "quotes world"}} + ], + } + source: Final[list[bytes]] = self._ended_sse_chunks() + chunks: Final = [ + chunk.replace( + b'"stop_sequence": null}', b'"stop_sequence": null, ' + json.dumps(delta_fields)[1:-1].encode() + b"}" + ) + for chunk in source + ] + + await AnthropicMessagesHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert self._delta_texts(chunks) == ["hello [MASKED]", ""] + assert self._message_delta(chunks) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": None}, + "safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}], + } + + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_use_rewrite_drops_the_verdict_of_the_rewritten_call_only(self): + class MaskPersimmon(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) and "persimmon" in tool_call.function.arguments: + tool_call.function.arguments = '{"fruit": "[MASKED]"}' + return inputs + + first_call: Final[list[bytes]] = self._ended_tool_use_sse_chunks() + second_call: Final = [ + chunk.replace(b'"index": 0', b'"index": 1') + .replace(b"toolu_1", b"toolu_2") + .replace(b"persim", b"pear") + .replace(b'mon\\"}', b'\\"}') + for chunk in first_call[1:6] + ] + verdicts: Final = { + "toolu_1": {"type": "evaluated", "outcome": "flagged", "explanation": "looks up persimmon"}, + "toolu_2": {"type": "evaluated", "outcome": "not_flagged", "explanation": "looks up pear"}, + } + safeguard_results: Final = [ + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": verdicts}} + ] + message_delta: Final = first_call[6].replace( + b'"stop_sequence": null}', + b'"stop_sequence": null, "safeguard_results": ' + json.dumps(safeguard_results).encode() + b"}", + ) + chunks: Final = [*first_call[:6], *second_call, message_delta, first_call[7]] + + await AnthropicMessagesHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MaskPersimmon(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert self._partial_jsons(chunks)[:3] == ['{"fruit": "[MASKED]"}', "", ""] + assert self._message_delta(chunks)["safeguard_results"] == [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"toolu_2": {"type": "evaluated", "outcome": "not_flagged"}}, + }, + } + ] + @pytest.mark.asyncio async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self): handler = AnthropicMessagesHandler() diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py b/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py new file mode 100644 index 00000000000..476cea2833c --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/messages/test_fake_stream_iterator.py @@ -0,0 +1,93 @@ +import json +from collections.abc import Iterable, Mapping +from typing import Final + +from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, +) + +_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] + + +def _decoded_event(chunk: object) -> tuple[str, dict[str, object]]: + assert isinstance(chunk, bytes) + event_line, data_line = chunk.decode().split("\n")[:2] + assert event_line.startswith("event: ") + return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: ")) + + +def _decoded_events(chunks: Iterable[object]) -> tuple[tuple[str, dict[str, object]], ...]: + return tuple(_decoded_event(chunk) for chunk in chunks) + + +def _message_delta(response: Mapping[str, object]) -> object: + decoded: Final = _decoded_events(FakeAnthropicMessagesStreamIterator(response=response)) + assert tuple(event_type for event_type, _ in decoded)[-2:] == ("message_delta", "message_stop") + deltas: Final = tuple(payload["delta"] for event_type, payload in decoded if event_type == "message_delta") + assert len(deltas) == 1 + return deltas[0] + + +def test_message_delta_carries_safeguard_results(): + """Websearch interception and the MCP loop render a complete response as + this stream for a client that asked to stream (the guardrail path does too, + after merging the upstream fields back in, which + tests/unit/proxy/guardrails/test_anthropic_sse.py covers); a Claude Code + auto-mode client reads the classifier verdict from message_delta.delta, + where a real Anthropic stream puts it.""" + response: Final = { + "id": "msg_1", + "model": "claude-sonnet-4-6", + "role": "assistant", + "type": "message", + "content": [{"type": "text", "text": "I won't run that."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + assert _message_delta(response) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + + +def test_message_delta_carries_refusal_stop_details(): + stop_details: Final = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response: Final = { + "id": "msg_2", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "usage": {"input_tokens": 10, "output_tokens": 0}, + } + assert _message_delta(response) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + + +def test_message_delta_omits_absent_optional_fields(): + response: Final = { + "id": "msg_3", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": {"input_tokens": 10, "output_tokens": 2}, + } + assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None} diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py index 7e983be835b..b995c6dba9b 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py @@ -1306,6 +1306,33 @@ def test_anthropic_messages_streaming_response_reports_withheld_output_of_its_st ) +def test_anthropic_messages_response_as_sse_events_keeps_message_delta_fields_out_of_message_start(): + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}] + stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response = { + "id": "msg_1", + "model": "claude-haiku", + "role": "assistant", + "type": "message", + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "safeguard_results": safeguard_results, + "content": [{"type": "text", "text": "no"}], + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + message_start = decoded[0][1]["message"] + assert message_start["stop_details"] is None + assert "safeguard_results" not in message_start + assert decoded[-2][1]["delta"] == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + "safeguard_results": safeguard_results, + } + + def test_anthropic_messages_response_as_sse_events_no_content_blocks(): response = {"id": "msg_4", "content": [], "stop_reason": "end_turn"} decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) @@ -1530,6 +1557,72 @@ async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_co assert deferred_fired.is_set() +_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] + + +def _sse_events_message_delta(response: Mapping[str, object]) -> object: + decoded: Final = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + deltas: Final = tuple(payload["delta"] for event_type, payload in decoded if event_type == "message_delta") + assert len(deltas) == 1 + return deltas[0] + + +def test_anthropic_messages_response_as_sse_events_message_delta_carries_safeguard_results(): + """A Claude Code auto-mode request's classifier verdict rides in + message_delta.delta on a real Anthropic stream; a synthesized stream that + drops it leaves the client without the verdict.""" + response = { + "id": "msg_7", + "content": [{"type": "text", "text": "I won't run that."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + assert _sse_events_message_delta(response) == { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS, + } + + +def test_anthropic_messages_response_as_sse_events_message_delta_carries_refusal_stop_details(): + stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"} + response = { + "id": "msg_8", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + assert _sse_events_message_delta(response) == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": stop_details, + } + + +def test_anthropic_messages_response_as_sse_events_message_delta_omits_absent_optional_fields(): + response = { + "id": "msg_9", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + } + assert _sse_events_message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None} + + class _OneVerdictClassifier: def __init__(self, reply: str) -> None: self._reply = reply diff --git a/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py b/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py index 33f95ccc9ff..dbd49064e4b 100644 --- a/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py +++ b/tests/unit/llms/anthropic/pass_through/test_stream_assembly.py @@ -4,6 +4,7 @@ from typing import Final import pytest from litellm.llms.anthropic.pass_through.stream_assembly import ( + anthropic_message_delta_fields, assemble_anthropic_sse_stream, is_anthropic_sse_stream, is_raw_sse_stream, @@ -55,6 +56,35 @@ ANTHROPIC_STREAM: Final = ( _frame("message_stop", {"type": "message_stop"}), ) +SAFEGUARD_RESULTS: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"} + }, + }, + } +] +SAFEGUARDED_STREAM: Final = ( + *ANTHROPIC_STREAM[:5], + _frame( + "message_delta", + { + "type": "message_delta", + "delta": { + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "safeguard_results": SAFEGUARD_RESULTS, + }, + "usage": {"output_tokens": 2}, + }, + ), + _frame("message_stop", {"type": "message_stop"}), +) + GOOGLE_STREAM: Final = (b'data: {"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}\n\n',) ANTHROPIC_ERROR: Final = _frame("error", {"type": "error", "error": {"type": "overloaded_error", "message": "busy"}}) CHAT_ERROR: Final = b'data: {"error": {"message": "blocked by guardrail"}}\n\n' @@ -108,6 +138,12 @@ def test_assemble_anthropic_sse_stream_restores_upstream_identity_only_when_aske assert (restored.id, restored.model) == (MESSAGE_ID, MODEL) +def test_anthropic_message_delta_fields_reads_the_fields_the_openai_shape_drops(): + assert dict(anthropic_message_delta_fields(ANTHROPIC_STREAM)) == {} + assert dict(anthropic_message_delta_fields(SAFEGUARDED_STREAM)) == {"safeguard_results": SAFEGUARD_RESULTS} + assert dict(anthropic_message_delta_fields(GOOGLE_STREAM)) == {} + + @pytest.mark.parametrize( "chunks", [GOOGLE_STREAM, INVALID_UTF8, (ANTHROPIC_ERROR,), ()], diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 4e19c7ed328..2ee4c8c72c0 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -5,6 +5,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) import json import logging import re +from typing import Final from unittest.mock import patch import pytest @@ -1379,6 +1380,54 @@ class TestToolPermissionGuardrailAnthropicMessages: assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id" assert message_start["model"] == "claude-sonnet-4-5", "the rewritten stream reports the model the upstream served" + @pytest.mark.asyncio + async def test_rewrite_mode_keeps_the_allowed_tool_use_verdict_with_its_explanation(self): + """Claude Code auto mode reads the classifier verdict from message_delta. The rewrite changes + no text the explanation could repeat, so the allowed call keeps its verdict whole.""" + default_on_rewriting: Final = ToolPermissionGuardrail( + guardrail_name="anthropic-rewrite-default-on", + rules=self.rules, + default_action="deny", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + bash_verdict: Final = {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"} + safeguard_results: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"tu_ok": bash_verdict, "tu_bad": {"type": "evaluated", "outcome": "not_flagged"}}, + }, + } + ] + bash_chunks: Final = self._sse_chunks("Bash", "tu_ok") + read_chunks: Final = self._sse_chunks("Read", "tu_bad") + read_tool_use: Final = [chunk.replace(b'"index": 1', b'"index": 2') for chunk in read_chunks[4:7]] + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 5}, + } + chunks: Final = [ + *bash_chunks[:7], + *read_tool_use, + f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode(), + bash_chunks[-1], + ] + + body: Final = b"".join(await self._drain(default_on_rewriting, chunks)).decode() + delta: Final = next( + json.loads(line[6:])["delta"] + for line in body.splitlines() + if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_delta" + ) + assert "Permission denied" in body + assert delta["safeguard_results"] == [ + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {"tu_ok": bash_verdict}}} + ] + @pytest.mark.asyncio async def test_message_start_without_a_dict_message_fails_closed(self): """Malformed SSE must not be forwarded unscanned. diff --git a/tests/unit/proxy/guardrails/test_anthropic_sse.py b/tests/unit/proxy/guardrails/test_anthropic_sse.py new file mode 100644 index 00000000000..95f596aab11 --- /dev/null +++ b/tests/unit/proxy/guardrails/test_anthropic_sse.py @@ -0,0 +1,316 @@ +import json +from collections.abc import Callable +from itertools import chain +from typing import Final + +import pytest + +from litellm.proxy.guardrails.anthropic_sse import ( + anthropic_sse_chunks_from_response, + assemble_anthropic_sse_stream, +) +from litellm.types.utils import Choices, ModelResponse + + +def _frame(event: str, data: dict[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +def _decoded_event(chunk: bytes) -> tuple[str, dict[str, object]]: + event_line, data_line = chunk.decode().split("\n")[:2] + assert event_line.startswith("event: ") + return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: ")) + + +def _decode(chunks: tuple[bytes, ...]) -> tuple[tuple[str, dict[str, object]], ...]: + return tuple(_decoded_event(chunk) for chunk in chunks) + + +def _rewritten_delta( + upstream: tuple[bytes, ...], rewrite: Callable[[ModelResponse], None], keep_explanations: bool = False +) -> object: + assembled: Final = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + rewrite(assembled) + rewritten: Final = anthropic_sse_chunks_from_response(assembled, upstream, keep_explanations=keep_explanations) + return next(payload["delta"] for event, payload in _decode(rewritten) if event == "message_delta") + + +def _set_text(text: str) -> Callable[[ModelResponse], None]: + def rewrite(assembled: ModelResponse) -> None: + choice: Final = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = text + + return rewrite + + +def _tool_use_frames(index: int, tool_use_id: str, name: str) -> tuple[bytes, ...]: + return ( + _frame( + "content_block_start", + { + "type": "content_block_start", + "index": index, + "content_block": {"type": "tool_use", "id": tool_use_id, "name": name, "input": {}}, + }, + ), + _frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": index, + "delta": {"type": "input_json_delta", "partial_json": '{"command": "ls"}'}, + }, + ), + _frame("content_block_stop", {"type": "content_block_stop", "index": index}), + ) + + +def _upstream_stream( + delta: dict[str, object], text: str = "call 555-0100", tool_uses: tuple[tuple[str, str], ...] = () +) -> tuple[bytes, ...]: + tool_use_frames = tuple( + chain.from_iterable( + _tool_use_frames(index, tool_use_id, name) for index, (tool_use_id, name) in enumerate(tool_uses, start=1) + ) + ) + return ( + _frame( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_upstream_1", + "type": "message", + "role": "assistant", + "model": "claude-upstream", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + }, + ), + _frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ), + _frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + *tool_use_frames, + _frame("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 2}}), + _frame("message_stop", {"type": "message_stop"}), + ) + + +def _safeguard_results(tool_uses: dict[str, dict[str, str]]) -> list[dict[str, object]]: + return [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_uses}}] + + +SAFEGUARD_RESULTS: Final = _safeguard_results({}) +REFUSAL_STOP_DETAILS: Final = {"type": "refusal", "category": "cyber", "explanation": None} + + +@pytest.mark.parametrize( + ("stop_reason", "delta_fields"), + [ + ("end_turn", {"safeguard_results": SAFEGUARD_RESULTS}), + ("refusal", {"stop_details": REFUSAL_STOP_DETAILS}), + ], + ids=["safeguard_results", "refusal_stop_details"], +) +def test_rewritten_stream_keeps_the_upstream_message_delta_fields(stop_reason: str, delta_fields: dict[str, object]): + """A guardrail that rewrites a streamed response assembles it into the OpenAI shape, which has no + slot for ``safeguard_results`` or ``stop_details``, so the re-emitted ``message_delta`` used to + drop the classifier verdict a Claude Code auto-mode session reads from it.""" + upstream_delta: dict[str, object] = {"stop_reason": stop_reason, "stop_sequence": None, **delta_fields} + upstream = _upstream_stream(upstream_delta) + assembled = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = "call " + + decoded = _decode(anthropic_sse_chunks_from_response(assembled, upstream)) + payloads = {event: payload for event, payload in decoded} + + assert payloads["content_block_delta"]["delta"] == {"type": "text_delta", "text": "call "} + assert payloads["message_delta"]["delta"] == upstream_delta + assert not set(delta_fields) & set(payloads["message_start"]["message"]) + + +SSN: Final = "123-45-6789" + + +@pytest.mark.parametrize( + ("upstream_delta", "tool_uses", "expected_fields"), + [ + ( + { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"type": "refusal", "category": "cyber", "explanation": f"Request involves SSN {SSN}"}, + }, + (), + {"stop_details": {"type": "refusal", "category": "cyber", "explanation": None}}, + ), + ( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": f"leaks SSN {SSN}"}} + ), + }, + (("toolu_bash", "Bash"),), + {"safeguard_results": _safeguard_results({"toolu_bash": {"type": "evaluated", "outcome": "flagged"}})}, + ), + ( + { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": [ + { + "type": "prompt_injection", + "status": {"type": "available", "explanation": f"quotes SSN {SSN}"}, + "explanation": f"saw SSN {SSN}", + } + ], + }, + (), + {"safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}]}, + ), + ], + ids=["stop_details_explanation", "verdict_explanation", "unknown_shape_explanation"], +) +def test_masked_stream_drops_the_unmasked_upstream_explanations( + upstream_delta: dict[str, object], tool_uses: tuple[tuple[str, str], ...], expected_fields: dict[str, object] +): + upstream = _upstream_stream(upstream_delta, text=f"my SSN is {SSN}", tool_uses=tool_uses) + assembled = assemble_anthropic_sse_stream(upstream, restore_identity=True) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.content = "my SSN is " + + rewritten = anthropic_sse_chunks_from_response(assembled, upstream) + message_delta = next(payload for event, payload in _decode(rewritten) if event == "message_delta") + + assert not [chunk for chunk in rewritten if SSN.encode() in chunk] + assert {key: message_delta["delta"][key] for key in expected_fields} == expected_fields + + +def test_masked_stream_drops_safeguard_results_it_cannot_read(): + upstream_delta: Final = { + "stop_reason": "end_turn", + "stop_sequence": None, + "safeguard_results": {"type": "dangerous_tool_use", "explanation": f"saw SSN {SSN}"}, + } + + assert _rewritten_delta(_upstream_stream(upstream_delta), _set_text("my SSN is ")) == { + "stop_reason": "end_turn", + "stop_sequence": None, + } + + +def test_rewritten_stream_keeps_only_the_verdicts_of_tool_uses_it_still_delivers(): + kept_verdict = {"type": "evaluated", "outcome": "not_flagged"} + upstream = _upstream_stream( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_glob": {"type": "evaluated", "outcome": "flagged"}, "toolu_bash": kept_verdict} + ), + }, + tool_uses=(("toolu_glob", "Glob"), ("toolu_bash", "Bash")), + ) + assembled = assemble_anthropic_sse_stream(upstream) + assert assembled is not None + choice = assembled.choices[0] + assert isinstance(choice, Choices) + choice.message.tool_calls = [call for call in choice.message.tool_calls or [] if call.id == "toolu_bash"] + choice.message.content = "Permission denied: Tool 'Glob' denied" + + decoded = _decode(anthropic_sse_chunks_from_response(assembled, upstream)) + message_delta = next(payload for event, payload in decoded if event == "message_delta") + delivered_tool_use_ids = [ + payload["content_block"]["id"] + for event, payload in decoded + if event == "content_block_start" and payload["content_block"]["type"] == "tool_use" + ] + + assert delivered_tool_use_ids == ["toolu_bash"] + assert message_delta["delta"]["safeguard_results"] == _safeguard_results({"toolu_bash": kept_verdict}) + + +def test_rewritten_refusal_gets_no_stop_details_the_upstream_did_not_send(): + """The OpenAI round trip sets ``stop_details`` on every refusal, and the repo reads a refusal + carrying one as a safeguard refusal, so a plain refusal must not come back as one.""" + upstream_delta: Final = {"stop_reason": "refusal", "stop_sequence": None} + + assert _rewritten_delta(_upstream_stream(upstream_delta), _set_text("call ")) == upstream_delta + + +def test_rewritten_stream_keeps_safeguard_shapes_it_does_not_know(): + upstream_delta: Final = { + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": {"type": "pause_turn", "category": None, "explanation": None, "policy_url": "https://x.test"}, + "safeguard_results": [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": {"toolu_bash": {"type": "unavailable", "reason": "overloaded"}}, + }, + "classifier_version": "2026-10", + }, + {"type": "prompt_injection", "status": {"type": "skipped"}}, + ], + } + upstream: Final = _upstream_stream(upstream_delta, tool_uses=(("toolu_bash", "Bash"),)) + + assert _rewritten_delta(upstream, _set_text("call ")) == upstream_delta + + +def test_rewrite_that_changes_no_text_keeps_the_upstream_explanations(): + upstream_delta: Final = { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results( + {"toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"}} + ), + } + upstream: Final = _upstream_stream(upstream_delta, tool_uses=(("toolu_bash", "Bash"),)) + + assert _rewritten_delta(upstream, _set_text("Permission denied"), keep_explanations=True) == upstream_delta + + +def test_rewritten_stream_drops_the_verdict_of_a_tool_use_whose_input_the_guardrail_changed(): + """A verdict scored on the upstream input must not vouch for the input the client runs.""" + verdict: Final = {"type": "evaluated", "outcome": "not_flagged"} + upstream: Final = _upstream_stream( + { + "stop_reason": "tool_use", + "stop_sequence": None, + "safeguard_results": _safeguard_results({"toolu_masked": verdict, "toolu_untouched": verdict}), + }, + tool_uses=(("toolu_masked", "Bash"), ("toolu_untouched", "Bash")), + ) + + def mask_first_tool_input(assembled: ModelResponse) -> None: + choice: Final = assembled.choices[0] + assert isinstance(choice, Choices) + tool_calls: Final = choice.message.tool_calls or [] + assert [call.id for call in tool_calls] == ["toolu_masked", "toolu_untouched"] + tool_calls[0].function.arguments = '{"command": ""}' + + rewritten_delta: Final = _rewritten_delta(upstream, mask_first_tool_input) + + assert isinstance(rewritten_delta, dict) + assert rewritten_delta["safeguard_results"] == _safeguard_results({"toolu_untouched": verdict}) diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 56184fa299a..637fdfa6873 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -14,6 +14,7 @@ import respx import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import FakeAnthropicMessagesStreamIterator from litellm.litellm_core_utils import get_llm_provider_logic from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure from litellm.router_utils.fallback_event_handlers import ( @@ -32,6 +33,7 @@ from litellm.router_utils.fallback_event_handlers import ( is_bad_request_without_fallback, log_failure_fallback_event, log_success_fallback_event, + mid_stream_fallback_hop_kwargs, mid_stream_fallback_snapshot_kwargs, mid_stream_retry_kwargs, record_pre_routing_selection, @@ -1865,6 +1867,41 @@ async def test_streaming_refusal_without_fallback_row_passes_frames_through(): assert len(fake.calls) == 1 +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +@pytest.mark.asyncio +async def test_refused_stream_built_from_a_finished_agentic_turn_is_delivered_without_a_fallback(): + """Websearch, MCP and advisor loops run their tools before streaming the final reply, so a + fallback would run every tool call again for a reply the client then gets from another model.""" + fake: Final = FakeAnthropicUpstream() + router: Final = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream: Final = await router._aanthropic_messages_streaming_iterator( + response=FakeAnthropicMessagesStreamIterator(REFUSAL_RESPONSE), + initial_kwargs=mid_stream_fallback_hop_kwargs( + model="fable-tier", + original_generic_function=litellm.anthropic_messages, + controls=None, + kwargs={ + "model": "fable-tier", + "max_tokens": 16, + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + ), + ) + events: Final = [ + json.loads(line[len(b"data: ") :]) + for line in (await _collect(stream)).splitlines() + if line.startswith(b"data: ") + ] + + assert fake.calls == [] + message_delta: Final = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_details"] == REFUSAL_RESPONSE["stop_details"] + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.asyncio async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata(): From 415b84232f72d4a53663b347e80f8f1e206048ce Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 10:01:48 -0700 Subject: [PATCH 10/64] test(e2e): event_loop_stalls caps the request-handler bucket at 2, twice the gc-aware maximum (#47128) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/perf/cases/event_loop_stalls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/perf/cases/event_loop_stalls.py b/tests/e2e/perf/cases/event_loop_stalls.py index 04424de2e8d..0d67a98257f 100644 --- a/tests/e2e/perf/cases/event_loop_stalls.py +++ b/tests/e2e/perf/cases/event_loop_stalls.py @@ -32,7 +32,7 @@ CASE: Final = replace( allowlist=( KnownStall("garbage collection", max_count=4), KnownStall("LifespanOn.main", max_count=4), - KnownStall("RequestResponseCycle.run_asgi", max_count=6), + KnownStall("RequestResponseCycle.run_asgi", max_count=2), KnownStall("with_proxy_compaction_executor", max_count=2), KnownStall("Logging.async_success_handler", max_count=2), KnownStall("Logging.dispatch_success_handlers", max_count=2), From 0349cff097dfe1f02c2d67f62fdb2f6f8fa555b1 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 10:10:42 -0700 Subject: [PATCH 11/64] fix(backend): retain Prometheus metrics during startup (#47129) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- backend/routes/allowlist.py | 1 + tests/unit/proxy/test_component_allowlists.py | 23 ++++++++++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 60d695653b5..2a7522cf34e 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -165,6 +165,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset( BACKEND_MOUNT_PATHS: frozenset[str] = frozenset( { "/admin", + "/metrics", "/swagger", # API documentation static assets belong to the backend "/mcp", # lazily-mounted MCP sub-app serves on the backend component } diff --git a/tests/unit/proxy/test_component_allowlists.py b/tests/unit/proxy/test_component_allowlists.py index 42da00fcb6a..eb058d01770 100644 --- a/tests/unit/proxy/test_component_allowlists.py +++ b/tests/unit/proxy/test_component_allowlists.py @@ -57,7 +57,7 @@ _PRE_EXISTING_ENV = {key: os.environ.get(key) for key in _THROWAWAY_ENV} for _key, _value in _THROWAWAY_ENV.items(): os.environ.setdefault(_key, _value) -from prometheus_client import make_asgi_app +from prometheus_client import CollectorRegistry, Counter, make_asgi_app # gateway/ and backend/ live at the repo root, not inside litellm/. _REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) @@ -71,6 +71,7 @@ from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app +from litellm.integrations.prometheus_metrics_endpoint import make_metrics_asgi_app from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter for _key, _previous in _PRE_EXISTING_ENV.items(): @@ -355,6 +356,26 @@ def test_gateway_trim_keeps_metrics_mount(): "/metrics Mount must survive the gateway route trim" +@pytest.mark.parametrize("enabled", (False, True)) +def test_backend_serves_metrics_after_startup_when_enabled(enabled: bool) -> None: + registry: Final = CollectorRegistry() + failures: Final = Counter("litellm_postgres_failed_requests", "Failed database requests", registry=registry) + failures.inc(2) + + @asynccontextmanager + async def lifespan(application: Starlette) -> AsyncGenerator[None, None]: + if enabled: + application.mount("/metrics", make_metrics_asgi_app(registry)) + yield + + application: Final = FastAPI(lifespan=partial(_backend_lifespan, lifespan=lifespan)) + with TestClient(application) as client: + response: Final = client.get("/metrics") + assert response.status_code == (200 if enabled else 404), response.text + if enabled: + assert "litellm_postgres_failed_requests_total 2.0\n" in response.text + + def test_gateway_drops_ui_and_swagger_mounts(): """UI static and swagger Mounts must still be trimmed from the gateway.""" for path in ("/ui", "/_next", "/litellm-asset-prefix/_next", "/swagger"): From fa716ec7c646a57466fac32404618af3c50bf2dc Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:15:05 +0000 Subject: [PATCH 12/64] test(translation): allow anthropic usage speed in cross-model contract (#47130) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_completions/cross_model/anthropic.py | 78 ------------------- .../chat_completions/cross_model/azure_ai.py | 35 --------- .../chat_completions/cross_model/case.py | 6 +- .../chat_completions/cross_model/contracts.py | 4 +- .../cross_model/test_messages_cross_model.py | 27 +------ .../translation/response_contracts.py | 1 + .../cross_model/test_responses_cross_model.py | 15 +--- 7 files changed, 12 insertions(+), 154 deletions(-) diff --git a/tests/integration/translation/chat_completions/cross_model/anthropic.py b/tests/integration/translation/chat_completions/cross_model/anthropic.py index a7a06d01361..3df7da17ef6 100644 --- a/tests/integration/translation/chat_completions/cross_model/anthropic.py +++ b/tests/integration/translation/chat_completions/cross_model/anthropic.py @@ -19,84 +19,6 @@ from integration.translation.chat_completions.templates import ( ) from integration.translation.templates.anthropic import ANTHROPIC_AGENT_REQUEST, ANTHROPIC_AGENT_RESPONSE -CLAUDE_FABLE_5_AGENT_CROSS_MODEL_PROVIDER_USAGE: Final = { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - "input_tokens": 1184, - "output_tokens": 105, - "output_tokens_details": {"thinking_tokens": 0}, - "service_tier": "standard", - "inference_geo": "global", -} -CLAUDE_FABLE_5_AGENT_CROSS_MODEL_CLIENT_USAGE: Final = { - "completion_tokens": 105, - "prompt_tokens": 1184, - "total_tokens": 1289, - "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, - "prompt_tokens_details": { - "cached_tokens": 0, - "text_tokens": 1184, - "cache_write_tokens": 0, - "cache_creation_tokens": 0, - "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - }, - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "inference_geo": "global", - "service_tier": "standard", -} -CLAUDE_FABLE_5_AGENT_CROSS_MODEL_TEST_CASE: Final = replace( - CLAUDE_FABLE_5_TEST_CASE, - mock_provider_response={ - **CLAUDE_FABLE_5_TEST_CASE.mock_provider_response, - "usage": CLAUDE_FABLE_5_AGENT_CROSS_MODEL_PROVIDER_USAGE, - }, - expected_litellm_response={ - **CLAUDE_FABLE_5_TEST_CASE.expected_litellm_response, - "usage": CLAUDE_FABLE_5_AGENT_CROSS_MODEL_CLIENT_USAGE, - }, -) - -CLAUDE_OPUS_5_AGENT_CROSS_MODEL_PROVIDER_USAGE: Final = { - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - "input_tokens": 1180, - "output_tokens": 105, - "output_tokens_details": {"thinking_tokens": 0}, - "service_tier": "standard", - "inference_geo": "global", -} -CLAUDE_OPUS_5_AGENT_CROSS_MODEL_CLIENT_USAGE: Final = { - "completion_tokens": 105, - "prompt_tokens": 1180, - "total_tokens": 1285, - "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, - "prompt_tokens_details": { - "cached_tokens": 0, - "text_tokens": 1180, - "cache_write_tokens": 0, - "cache_creation_tokens": 0, - "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - }, - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "inference_geo": "global", - "service_tier": "standard", -} -CLAUDE_OPUS_5_AGENT_CROSS_MODEL_TEST_CASE: Final = replace( - CLAUDE_OPUS_5_TEST_CASE, - mock_provider_response={ - **CLAUDE_OPUS_5_TEST_CASE.mock_provider_response, - "usage": CLAUDE_OPUS_5_AGENT_CROSS_MODEL_PROVIDER_USAGE, - }, - expected_litellm_response={ - **CLAUDE_OPUS_5_TEST_CASE.expected_litellm_response, - "usage": CLAUDE_OPUS_5_AGENT_CROSS_MODEL_CLIENT_USAGE, - }, -) - CLAUDE_FABLE_5_1_AGENT_TEST_CASE: Final = replace( CLAUDE_FABLE_5_1_TEST_CASE, scenario="agent_tools", diff --git a/tests/integration/translation/chat_completions/cross_model/azure_ai.py b/tests/integration/translation/chat_completions/cross_model/azure_ai.py index 15cbe4a3200..73d0ffbb8df 100644 --- a/tests/integration/translation/chat_completions/cross_model/azure_ai.py +++ b/tests/integration/translation/chat_completions/cross_model/azure_ai.py @@ -17,41 +17,6 @@ from integration.translation.chat_completions.templates import ( ) from integration.translation.templates.anthropic import ANTHROPIC_AGENT_REQUEST, ANTHROPIC_AGENT_RESPONSE -CLAUDE_FABLE_5_CROSS_MODEL_TEST_CASE: Final = replace( - CLAUDE_FABLE_5_TEST_CASE, - mock_provider_response={ - **CLAUDE_FABLE_5_TEST_CASE.mock_provider_response, - "usage": { - **object_value(ANTHROPIC_AGENT_RESPONSE["usage"]), - "input_tokens": 1184, - "output_tokens": 105, - "output_tokens_details": {"thinking_tokens": 0}, - "service_tier": "standard", - "inference_geo": "not_available", - }, - }, - expected_litellm_response={ - **CLAUDE_FABLE_5_TEST_CASE.expected_litellm_response, - "usage": { - "completion_tokens": 105, - "prompt_tokens": 1184, - "total_tokens": 1289, - "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 105}, - "prompt_tokens_details": { - "cached_tokens": 0, - "text_tokens": 1184, - "cache_write_tokens": 0, - "cache_creation_tokens": 0, - "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, - }, - "cache_creation_input_tokens": 0, - "cache_read_input_tokens": 0, - "inference_geo": "not_available", - "service_tier": "standard", - }, - }, -) - CLAUDE_FABLE_5_1_AGENT_TEST_CASE: Final = replace( CLAUDE_FABLE_5_1_TEST_CASE, scenario="agent_tools", diff --git a/tests/integration/translation/chat_completions/cross_model/case.py b/tests/integration/translation/chat_completions/cross_model/case.py index fae659c5dca..d81e6a269f6 100644 --- a/tests/integration/translation/chat_completions/cross_model/case.py +++ b/tests/integration/translation/chat_completions/cross_model/case.py @@ -102,7 +102,7 @@ MID: Final[tuple[AgentCases, ...]] = ( anthropic.CLAUDE_OPUS_5_5_AGENT_STREAMING_TEST_CASE, ), ( - anthropic.CLAUDE_OPUS_5_AGENT_CROSS_MODEL_TEST_CASE, + anthropic.CLAUDE_OPUS_5_TEST_CASE, anthropic.CLAUDE_OPUS_5_AGENT_FOLLOWUP_TEST_CASE, anthropic.CLAUDE_OPUS_5_AGENT_STREAMING_TEST_CASE, ), @@ -201,7 +201,7 @@ HIGH: Final[tuple[AgentCases, ...]] = ( anthropic.CLAUDE_FABLE_5_1_AGENT_STREAMING_TEST_CASE, ), ( - anthropic.CLAUDE_FABLE_5_AGENT_CROSS_MODEL_TEST_CASE, + anthropic.CLAUDE_FABLE_5_TEST_CASE, anthropic.CLAUDE_FABLE_5_AGENT_FOLLOWUP_TEST_CASE, anthropic.CLAUDE_FABLE_5_AGENT_STREAMING_TEST_CASE, ), @@ -211,7 +211,7 @@ HIGH: Final[tuple[AgentCases, ...]] = ( azure_ai.CLAUDE_FABLE_5_1_AGENT_STREAMING_TEST_CASE, ), ( - azure_ai.CLAUDE_FABLE_5_CROSS_MODEL_TEST_CASE, + azure_ai.CLAUDE_FABLE_5_TEST_CASE, azure_ai.CLAUDE_FABLE_5_AGENT_FOLLOWUP_TEST_CASE, azure_ai.CLAUDE_FABLE_5_AGENT_STREAMING_TEST_CASE, ), diff --git a/tests/integration/translation/chat_completions/cross_model/contracts.py b/tests/integration/translation/chat_completions/cross_model/contracts.py index 4fbf6a880b5..a15a43985f1 100644 --- a/tests/integration/translation/chat_completions/cross_model/contracts.py +++ b/tests/integration/translation/chat_completions/cross_model/contracts.py @@ -121,7 +121,9 @@ PROVIDER_EXTENSIONS: Final = MappingProxyType( ("provider_specific_fields", "reasoning_content", "thinking_blocks", "reasoning_items", "images") ), "choices[].message.tool_calls[]": frozenset(("index", "caller", "provider_specific_fields")), - "usage": frozenset(("cache_creation_input_tokens", "cache_read_input_tokens", "inference_geo", "service_tier")), + "usage": frozenset( + ("cache_creation_input_tokens", "cache_read_input_tokens", "inference_geo", "service_tier", "speed") + ), } ) diff --git a/tests/integration/translation/messages/cross_model/test_messages_cross_model.py b/tests/integration/translation/messages/cross_model/test_messages_cross_model.py index e3bafc37993..ae6fe70fe7c 100644 --- a/tests/integration/translation/messages/cross_model/test_messages_cross_model.py +++ b/tests/integration/translation/messages/cross_model/test_messages_cross_model.py @@ -3,14 +3,11 @@ from dataclasses import replace from typing import Final import pytest -from integration._support.client import JSON_OBJECT, Gateway, object_value +from integration._support.client import JSON_OBJECT, Gateway from integration._support.provider import SharedProvider from integration._support.wire import Reply from integration.translation.case import TranslationTestCase -from integration.translation.messages.bases.anthropic import ( - CLAUDE_SONNET_5_5_RESPONSE_FIELDS, - CLAUDE_SONNET_5_5_TEST_CASE, -) +from integration.translation.messages.bases.anthropic import CLAUDE_SONNET_5_5_TEST_CASE from integration.translation.messages.bases.azure_ai import ( CLAUDE_SONNET_5_5_TEST_CASE as AZURE_CLAUDE_SONNET_5_5_TEST_CASE, ) @@ -20,18 +17,6 @@ from integration.translation.response_scenarios import text_completion, two_mess from integration.translation.runner import assert_provider_request, assert_translation GPT_6_1_SOL_TWO_MESSAGES_BEFORE_TOOLS_TEST_CASE: Final = two_messages_before_tools(GPT_6_1_SOL_TEST_CASE) -CLAUDE_SONNET_5_5_CROSS_MODEL_NATIVE_TEST_CASE: Final = replace( - CLAUDE_SONNET_5_5_TEST_CASE, - scenario="cross_model_rich_request", - mock_provider_response={ - **CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, - "usage": object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), - }, - expected_litellm_response={ - **CLAUDE_SONNET_5_5_TEST_CASE.expected_litellm_response, - "usage": object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), - }, -) AZURE_CLAUDE_SONNET_5_5_CROSS_MODEL_TEST_CASE: Final = replace( AZURE_CLAUDE_SONNET_5_5_TEST_CASE, scenario="cross_model_rich_request", @@ -40,14 +25,10 @@ AZURE_CLAUDE_SONNET_5_5_CROSS_MODEL_TEST_CASE: Final = replace( **CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_request, "system": [{"type": "text", "text": CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_request["system"]}], }, - mock_provider_response={ - **CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, - "usage": object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), - }, + mock_provider_response=CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, expected_litellm_response={ **CLAUDE_SONNET_5_5_TEST_CASE.expected_litellm_response, "model": "azure_ai/claude-sonnet-5-5", - "usage": object_value(CLAUDE_SONNET_5_5_RESPONSE_FIELDS["usage"]), }, ) @@ -63,7 +44,7 @@ def test_messages_keep_the_same_public_contract_on_native_and_bridged_routes( assert {key: value for key, value in case.litellm_request.items() if key != "model"} == { key: value for key, value in CLAUDE_SONNET_5_5_TEST_CASE.litellm_request.items() if key != "model" } - native: Final = assert_translation(CLAUDE_SONNET_5_5_CROSS_MODEL_NATIVE_TEST_CASE, gateway, provider) + native: Final = assert_translation(CLAUDE_SONNET_5_5_TEST_CASE, gateway, provider) compared: Final = assert_translation(case, gateway, provider) assert messages_structure(compared) == messages_structure(native) diff --git a/tests/integration/translation/response_contracts.py b/tests/integration/translation/response_contracts.py index c3182d269ff..6af3f9d681d 100644 --- a/tests/integration/translation/response_contracts.py +++ b/tests/integration/translation/response_contracts.py @@ -197,6 +197,7 @@ MESSAGES_OPTIONAL: Final = MappingProxyType( "server_tool_use", "service_tier", "inference_geo", + "speed", ) ), } diff --git a/tests/integration/translation/responses/cross_model/test_responses_cross_model.py b/tests/integration/translation/responses/cross_model/test_responses_cross_model.py index a0b5778a200..cb6f973d2df 100644 --- a/tests/integration/translation/responses/cross_model/test_responses_cross_model.py +++ b/tests/integration/translation/responses/cross_model/test_responses_cross_model.py @@ -22,25 +22,12 @@ from integration.translation.responses.bases.openai_responses import GPT_5_6_SOL from integration.translation.runner import assert_translation GPT_6_1_SOL_TWO_MESSAGES_BEFORE_TOOLS_TEST_CASE: Final = two_messages_before_tools(GPT_6_1_SOL_TEST_CASE) -CLAUDE_SONNET_5_5_CROSS_MODEL_PROVIDER_USAGE: Final = { - "input_tokens": 4, - "cache_creation_input_tokens": 873, - "cache_read_input_tokens": 0, - "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, - "output_tokens": 509, - "output_tokens_details": {"thinking_tokens": 313}, - "service_tier": "standard", - "inference_geo": "global", -} AZURE_CLAUDE_SONNET_5_5_CROSS_MODEL_TEST_CASE: Final = replace( AZURE_CLAUDE_SONNET_5_5_TEST_CASE, scenario="cross_model_rich_request", litellm_request={**CLAUDE_SONNET_5_5_TEST_CASE.litellm_request, "model": "azure_ai/claude-sonnet-5-5"}, expected_provider_request=CLAUDE_SONNET_5_5_TEST_CASE.expected_provider_request, - mock_provider_response={ - **CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, - "usage": CLAUDE_SONNET_5_5_CROSS_MODEL_PROVIDER_USAGE, - }, + mock_provider_response=CLAUDE_SONNET_5_5_TEST_CASE.mock_provider_response, expected_litellm_response={ **CLAUDE_SONNET_5_5_TEST_CASE.expected_litellm_response, "model": "azure_ai/claude-sonnet-5-5", From 49b5740be8c40880b565b3b07ede3122fcead128 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 17 Oct 2026 10:23:21 -0700 Subject: [PATCH 13/64] feat(terraform)!: default the component images to the module release tag (#46960) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/litellm/aws/.terraform.lock.hcl | 2 + terraform/litellm/aws/README.md | 21 +++- terraform/litellm/aws/ecs.tf | 10 +- .../aws/examples/default/.terraform.lock.hcl | 2 + .../litellm/aws/examples/default/main.tf | 3 + .../examples/default/terraform.tfvars.example | 15 ++- .../litellm/aws/examples/default/variables.tf | 12 ++ .../litellm/aws/examples/default/versions.tf | 2 +- terraform/litellm/aws/locals.tf | 5 + terraform/litellm/aws/migrations.tf | 2 +- .../aws/tests/byo_infrastructure.tftest.hcl | 1 + .../litellm/aws/tests/collector.tftest.hcl | 3 +- .../aws/tests/connection_pool.tftest.hcl | 1 + terraform/litellm/aws/tests/images.tftest.hcl | 104 ++++++++++++++++++ .../aws/tests/metrics_sidecar.tftest.hcl | 1 + .../aws/tests/workload_autoscaling.tftest.hcl | 1 + terraform/litellm/aws/variables.tf | 37 +++++-- terraform/litellm/aws/versions.tf | 2 +- terraform/litellm/gcp/.terraform.lock.hcl | 3 + terraform/litellm/gcp/README.md | 12 +- .../gcp/examples/default/.terraform.lock.hcl | 3 + .../gcp/examples/default/deploystack.json | 2 +- .../litellm/gcp/examples/default/main.tf | 2 +- .../examples/default/terraform.tfvars.example | 11 +- .../litellm/gcp/examples/default/variables.tf | 2 +- .../litellm/gcp/examples/default/versions.tf | 2 +- .../litellm/gcp/tests/collector.tftest.hcl | 1 + .../gcp/tests/connection_pool.tftest.hcl | 1 + .../litellm/gcp/tests/deps_only.tftest.hcl | 15 +++ .../gcp/tests/metrics_sidecar.tftest.hcl | 1 + terraform/litellm/gcp/variables.tf | 10 +- terraform/litellm/gcp/versions.tf | 2 +- 32 files changed, 245 insertions(+), 46 deletions(-) create mode 100644 terraform/litellm/aws/tests/images.tftest.hcl diff --git a/terraform/litellm/aws/.terraform.lock.hcl b/terraform/litellm/aws/.terraform.lock.hcl index 30b2a194c4c..79097043254 100644 --- a/terraform/litellm/aws/.terraform.lock.hcl +++ b/terraform/litellm/aws/.terraform.lock.hcl @@ -6,6 +6,7 @@ provider "registry.terraform.io/hashicorp/aws" { constraints = "~> 5.60" hashes = [ "h1:Ijt7pOlB7Tr7maGQIqtsLFbl7pSMIj06TVdkoSBcYOw=", + "h1:edXOJWE4ORX8Fm+dpVpICzMZJat4AX0VRCAy/xkcOc0=", "zh:054b8dd49f0549c9a7cc27d159e45327b7b65cf404da5e5a20da154b90b8a644", "zh:0b97bf8d5e03d15d83cc40b0530a1f84b459354939ba6f135a0086c20ebbe6b2", "zh:1589a2266af699cbd5d80737a0fe02e54ec9cf2ca54e7e00ac51c7359056f274", @@ -28,6 +29,7 @@ provider "registry.terraform.io/hashicorp/random" { version = "3.8.1" constraints = "~> 3.6" hashes = [ + "h1:Eexl06+6J+s75uD46+WnZtpJZYRVUMB0AiuPBifK6Jc=", "h1:u8AKlWVDTH5r9YLSeswoVEjiY72Rt4/ch7U+61ZDkiQ=", "zh:08dd03b918c7b55713026037c5400c48af5b9f468f483463321bd18e17b907b4", "zh:0eee654a5542dc1d41920bbf2419032d6f0d5625b03bd81339e5b33394a3e0ae", diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 6fcbdc2f500..bedb6f4e98b 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -419,8 +419,9 @@ To run *many* tenants from a single config, call the module with ```hcl module "litellm" { for_each = toset(["acme", "globex"]) - source = "github.com/BerriAI/litellm//terraform/litellm/aws?ref=" - tenant = each.key + source = "github.com/BerriAI/litellm//terraform/litellm/aws?ref=" + image_tag = "" + tenant = each.key env = "prod" region = "us-west-2" azs = ["us-west-2a", "us-west-2b"] @@ -485,6 +486,8 @@ provider "aws" { module "litellm" { source = "github.com/BerriAI/litellm//terraform/litellm/aws?ref=" + image_tag = "" # a git ref ships no image default, a registry version does + region = "us-west-2" tenant = "acme" env = "prod" @@ -499,12 +502,20 @@ top — set org-wide tags there, per-deployment tags via the `tags` input. ## Image pulls -The defaults pull from `ghcr.io/berriai/litellm-:v1.86.0-dev`, +The defaults pull `ghcr.io/berriai/litellm-:${image_tag}`, which is anonymous-readable. There are four images: `litellm-gateway`, `litellm-backend`, `litellm-ui`, and `litellm-migrations` (slim image used only by the one-off migration task — runs `prisma migrate deploy` against -the writer DB and exits). Bump them together when bumping LiteLLM. To pull -from a private registry: +the writer DB and exits). A published module version (Terraform Registry +`BerriAI/litellm/aws`) ships with `image_tag` stamped to its own release +tag, so bumping the module version bumps all four images with it. A git +ref of this monorepo has no published images of its own, so there +`image_tag` is required. Set `image_tag` to pin a different LiteLLM +release, or +`image_registry` to pull the same four images from a mirror. Set +`gateway_image` / `backend_image` / `ui_image` / `migrations_image` only +for a per-component override (custom build, different tag). To pull from a +private registry: - **ECR (same account)**: the execution role already has `AmazonECSTaskExecutionRolePolicy`, which grants ECR pull for repos in diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 2b235c2bad5..412e6c423f5 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -229,7 +229,7 @@ locals { gateway_metrics_container = local.metrics_enabled ? [ { name = "metrics" - image = var.gateway_image + image = local.gateway_image essential = false entryPoint = ["python", "-m", "litellm.proxy.prometheus_metrics_server"] command = ["--port", tostring(var.gateway_metrics_port)] @@ -306,7 +306,7 @@ locals { collector_container = var.collector_enabled ? [{ name = "collector" - image = var.gateway_image + image = local.gateway_image essential = false cpu = var.collector_cpu memory = var.collector_memory @@ -389,7 +389,7 @@ resource "aws_ecs_task_definition" "gateway" { merge( { name = "gateway" - image = var.gateway_image + image = local.gateway_image essential = true portMappings = [{ containerPort = 4000, protocol = "tcp" }] @@ -499,7 +499,7 @@ resource "aws_ecs_task_definition" "backend" { merge( { name = "backend" - image = var.backend_image + image = local.backend_image essential = true portMappings = [{ containerPort = 4001, protocol = "tcp" }] @@ -591,7 +591,7 @@ resource "aws_ecs_task_definition" "ui" { container_definitions = jsonencode([ { name = "ui" - image = var.ui_image + image = local.ui_image essential = true portMappings = [{ containerPort = 3000, protocol = "tcp" }] diff --git a/terraform/litellm/aws/examples/default/.terraform.lock.hcl b/terraform/litellm/aws/examples/default/.terraform.lock.hcl index 4a059b2b268..ce530fa767b 100644 --- a/terraform/litellm/aws/examples/default/.terraform.lock.hcl +++ b/terraform/litellm/aws/examples/default/.terraform.lock.hcl @@ -6,6 +6,7 @@ provider "registry.terraform.io/hashicorp/aws" { constraints = "~> 5.60" hashes = [ "h1:Ijt7pOlB7Tr7maGQIqtsLFbl7pSMIj06TVdkoSBcYOw=", + "h1:edXOJWE4ORX8Fm+dpVpICzMZJat4AX0VRCAy/xkcOc0=", "zh:054b8dd49f0549c9a7cc27d159e45327b7b65cf404da5e5a20da154b90b8a644", "zh:0b97bf8d5e03d15d83cc40b0530a1f84b459354939ba6f135a0086c20ebbe6b2", "zh:1589a2266af699cbd5d80737a0fe02e54ec9cf2ca54e7e00ac51c7359056f274", @@ -29,6 +30,7 @@ provider "registry.terraform.io/hashicorp/random" { constraints = "~> 3.6" hashes = [ "h1:OO+IuvQJSPmWdN8AyyIEvPJbLvDQpgX/zbktoa9KsJE=", + "h1:UlBuNVuCGJ39tTv2c5gz2NRZnQbXfbIWbTzWcth5o74=", "zh:161ad0bd9a75768c82f53fb6e7172a9d8be2d4889b012645a34795031aaf1bf1", "zh:19dc9a5b17729725ccfc4f45b0500af0ee5bc6b6b160c7adb8f2bf617d2c80ea", "zh:269eda8fe42daa7974d5a34d166c3ba9defe80cde86c01e4dadcfdf2e1f05e5f", diff --git a/terraform/litellm/aws/examples/default/main.tf b/terraform/litellm/aws/examples/default/main.tf index 115003f7ddd..5e7a65f907f 100644 --- a/terraform/litellm/aws/examples/default/main.tf +++ b/terraform/litellm/aws/examples/default/main.tf @@ -38,6 +38,9 @@ module "litellm" { litellm_license = var.litellm_license ui_password = var.ui_password + image_registry = var.image_registry + image_tag = var.image_tag != "" ? var.image_tag : null + acm_certificate_arn = var.acm_certificate_arn allow_plaintext_alb = var.allow_plaintext_alb s3_force_destroy = var.s3_force_destroy diff --git a/terraform/litellm/aws/examples/default/terraform.tfvars.example b/terraform/litellm/aws/examples/default/terraform.tfvars.example index 880ecf56555..dd392d9ecc5 100644 --- a/terraform/litellm/aws/examples/default/terraform.tfvars.example +++ b/terraform/litellm/aws/examples/default/terraform.tfvars.example @@ -58,12 +58,15 @@ env = "stage" # s3_force_destroy = false # skip_final_snapshot = false -# Component images and per-task sizing/autoscaling are NOT exposed as -# variables in this example (it keeps the curated surface small). They -# default to working public GHCR images. To pin images or tune -# CPU/memory/workers/autoscaling, set those inputs directly on the -# `module "litellm"` block in main.tf — the full list is in -# ../../variables.tf — or call the module from your own root config. +# Images. image_tag is applied to all four litellm-* images. A published +# module version ships with it stamped to its own release tag, a checkout of +# the monorepo has no default and needs it set. image_registry can point at +# an ECR pull-through mirror. Per-component *_image overrides and per-task +# sizing/autoscaling are NOT exposed here. Set them directly on the +# `module "litellm"` block in main.tf (full list in ../../variables.tf) or +# call the module from your own root config. +# image_registry = "ghcr.io/berriai" +# image_tag = "" # ---------- proxy_config (mirrors helm gateway.config.proxy_config) ---------- # proxy_config = { diff --git a/terraform/litellm/aws/examples/default/variables.tf b/terraform/litellm/aws/examples/default/variables.tf index f8140266fca..27c6c95e553 100644 --- a/terraform/litellm/aws/examples/default/variables.tf +++ b/terraform/litellm/aws/examples/default/variables.tf @@ -105,6 +105,18 @@ variable "ui_password" { } # TLS — provide an ACM cert for production, or opt into HTTP-only for dev. +variable "image_registry" { + description = "Registry path prefix; images composed as /litellm-:." + type = string + default = "ghcr.io/berriai" +} + +variable "image_tag" { + description = "Tag applied to all four litellm-* images. Bump in lockstep." + type = string + default = "" +} + variable "acm_certificate_arn" { description = "ACM cert ARN for the ALB HTTPS listener. Empty → no TLS." type = string diff --git a/terraform/litellm/aws/examples/default/versions.tf b/terraform/litellm/aws/examples/default/versions.tf index 73b88e91dce..06ebe5bd2c5 100644 --- a/terraform/litellm/aws/examples/default/versions.tf +++ b/terraform/litellm/aws/examples/default/versions.tf @@ -1,5 +1,5 @@ terraform { - required_version = ">= 1.6.0" + required_version = ">= 1.9.0" required_providers { aws = { diff --git a/terraform/litellm/aws/locals.tf b/terraform/litellm/aws/locals.tf index b89bd486d02..c06b171d6bd 100644 --- a/terraform/litellm/aws/locals.tf +++ b/terraform/litellm/aws/locals.tf @@ -11,6 +11,11 @@ locals { # the stack can reference local.name. name = "${var.tenant}-litellm-${var.env}" + gateway_image = var.gateway_image != "" ? var.gateway_image : "${var.image_registry}/litellm-gateway:${var.image_tag}" + backend_image = var.backend_image != "" ? var.backend_image : "${var.image_registry}/litellm-backend:${var.image_tag}" + ui_image = var.ui_image != "" ? var.ui_image : "${var.image_registry}/litellm-ui:${var.image_tag}" + migrations_image = var.migrations_image != "" ? var.migrations_image : "${var.image_registry}/litellm-migrations:${var.image_tag}" + # This is a reusable module — it declares no `provider` block, so the AWS # provider's `default_tags` is the caller's concern, not ours. To keep the # same per-resource tagging the stack had when it owned the provider, the diff --git a/terraform/litellm/aws/migrations.tf b/terraform/litellm/aws/migrations.tf index e924b29eba0..57c33f69a8b 100644 --- a/terraform/litellm/aws/migrations.tf +++ b/terraform/litellm/aws/migrations.tf @@ -28,7 +28,7 @@ resource "aws_ecs_task_definition" "migrations" { container_definitions = jsonencode([{ name = "migrations" - image = var.migrations_image + image = local.migrations_image essential = true # No entryPoint/command override — the image's ENTRYPOINT runs run.py. diff --git a/terraform/litellm/aws/tests/byo_infrastructure.tftest.hcl b/terraform/litellm/aws/tests/byo_infrastructure.tftest.hcl index 5a619bc98b1..702ba37184c 100644 --- a/terraform/litellm/aws/tests/byo_infrastructure.tftest.hcl +++ b/terraform/litellm/aws/tests/byo_infrastructure.tftest.hcl @@ -14,6 +14,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" diff --git a/terraform/litellm/aws/tests/collector.tftest.hcl b/terraform/litellm/aws/tests/collector.tftest.hcl index 1465130232f..082f17b2c08 100644 --- a/terraform/litellm/aws/tests/collector.tftest.hcl +++ b/terraform/litellm/aws/tests/collector.tftest.hcl @@ -13,6 +13,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" @@ -66,7 +67,7 @@ run "enabled_adds_a_sidecar_that_shares_the_gateway_transport" { assert { condition = ( - local.collector_container[0].image == var.gateway_image && + local.collector_container[0].image == local.gateway_image && local.collector_container[0].entryPoint == ["sh", "-c"] && local.collector_container[0].command == ["exec python -m litellm.proxy.collector"] && local.collector_container[0].essential == false && diff --git a/terraform/litellm/aws/tests/connection_pool.tftest.hcl b/terraform/litellm/aws/tests/connection_pool.tftest.hcl index 7408b1179ed..c618bd11768 100644 --- a/terraform/litellm/aws/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/aws/tests/connection_pool.tftest.hcl @@ -12,6 +12,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" diff --git a/terraform/litellm/aws/tests/images.tftest.hcl b/terraform/litellm/aws/tests/images.tftest.hcl new file mode 100644 index 00000000000..2d03a46b7a3 --- /dev/null +++ b/terraform/litellm/aws/tests/images.tftest.hcl @@ -0,0 +1,104 @@ +mock_provider "aws" { + mock_data "aws_iam_policy_document" { + defaults = { + json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}" + } + } +} +mock_provider "random" {} + +variables { + image_tag = "v9.9.9" + region = "us-east-1" + tenant = "acme" + env = "test" + allow_plaintext_alb = true + azs = ["us-east-1a", "us-east-1b"] +} + +run "every_component_follows_image_registry_and_image_tag" { + command = plan + + variables { + image_registry = "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror" + image_tag = "v9.9.9" + } + + assert { + condition = ( + local.gateway_image == "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror/litellm-gateway:v9.9.9" && + local.backend_image == "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror/litellm-backend:v9.9.9" && + local.ui_image == "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror/litellm-ui:v9.9.9" && + local.migrations_image == "123456789012.dkr.ecr.us-east-1.amazonaws.com/mirror/litellm-migrations:v9.9.9" + ) + error_message = "All four images must be composed as /litellm-:." + } +} + +run "an_empty_image_tag_fails_closed_unless_every_image_is_overridden" { + command = plan + + variables { + image_tag = "" + } + + expect_failures = [var.image_tag] +} + +run "a_null_image_tag_takes_the_module_default" { + command = plan + + variables { + image_tag = null + } + + expect_failures = [var.image_tag] +} + +run "four_explicit_uris_need_no_image_tag" { + command = plan + + variables { + image_tag = "" + gateway_image = "example.com/custom/gateway:nightly" + backend_image = "example.com/custom/backend:nightly" + ui_image = "example.com/custom/ui:nightly" + migrations_image = "example.com/custom/migrations:nightly" + } + + assert { + condition = local.migrations_image == "example.com/custom/migrations:nightly" + error_message = "Explicit per-component URIs must satisfy the module without an image_tag." + } +} + +run "composed_defaults_pull_from_ghcr_at_the_image_tag" { + command = plan + + assert { + condition = alltrue([ + for image in [local.gateway_image, local.backend_image, local.ui_image, local.migrations_image] : + startswith(image, "ghcr.io/berriai/litellm-") && endswith(image, ":${var.image_tag}") + ]) + error_message = "With no overrides every component must come from ghcr.io/berriai at image_tag." + } +} + +run "a_per_component_uri_wins_over_the_composed_image" { + command = plan + + variables { + image_tag = "v9.9.9" + gateway_image = "example.com/custom/gateway:nightly" + } + + assert { + condition = ( + local.gateway_image == "example.com/custom/gateway:nightly" && + local.backend_image == "ghcr.io/berriai/litellm-backend:v9.9.9" && + local.ui_image == "ghcr.io/berriai/litellm-ui:v9.9.9" && + local.migrations_image == "ghcr.io/berriai/litellm-migrations:v9.9.9" + ) + error_message = "A per-component override must replace only that component's image." + } +} diff --git a/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl b/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl index de839d7e5be..b330d4b91c3 100644 --- a/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl +++ b/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl @@ -11,6 +11,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" diff --git a/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl b/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl index 94e281faf94..ea6c343767c 100644 --- a/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl +++ b/terraform/litellm/aws/tests/workload_autoscaling.tftest.hcl @@ -11,6 +11,7 @@ mock_provider "aws" { mock_provider "random" {} variables { + image_tag = "v9.9.9" region = "us-east-1" tenant = "acme" env = "test" diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 580a0cc657a..9b7a279a907 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -135,38 +135,53 @@ variable "azs" { # ---------- Component images ---------- # -# Defaults pin the four componentized images at the same release tag on -# GHCR. Override on a per-component basis in tfvars when bumping; bump them -# together when bumping the LiteLLM release. +variable "image_registry" { + description = "Registry path prefix used to compose the four LiteLLM image URIs as `/litellm-:`. Per-component overrides (`gateway_image`, `backend_image`, `ui_image`, `migrations_image`) bypass this entirely when set." + type = string + default = "ghcr.io/berriai" +} + +variable "image_tag" { + description = "Tag applied to all four litellm-* images when composed from `image_registry`. A published module version ships with this stamped to its own release tag, which matches the `v`-prefixed tag published to GHCR for the same LiteLLM release. Set it to pin a different LiteLLM release. When sourcing the module from a git ref of the monorepo it has no default and must be set." + type = string + default = "" + nullable = false + + validation { + condition = var.image_tag != "" || (var.gateway_image != "" && var.backend_image != "" && var.ui_image != "" && var.migrations_image != "") + error_message = "image_tag is empty. Published module versions (registry.terraform.io/modules/BerriAI/litellm) ship with it stamped to their own release. A git ref of the monorepo has no published images of its own, so set image_tag to a published LiteLLM release tag or set all four *_image URIs." + } +} variable "gateway_image" { - description = "Container image for the gateway (data plane, port 4000). Tag must match a tag actually published to GHCR — the split images use the `v`-prefixed semver convention." + description = "Full image URI for the gateway (data plane, port 4000). Empty (default) composes from `image_registry` + `image_tag`." type = string - default = "ghcr.io/berriai/litellm-gateway:v1.86.0-dev" + default = "" } variable "backend_image" { - description = "Container image for the backend (management API, port 4001)." + description = "Full image URI for the backend (management API, port 4001). Empty (default) composes from `image_registry` + `image_tag`." type = string - default = "ghcr.io/berriai/litellm-backend:v1.86.0-dev" + default = "" } variable "ui_image" { - description = "Container image for the UI (nginx static export, port 3000)." + description = "Full image URI for the UI (nginx static export, port 3000). Empty (default) composes from `image_registry` + `image_tag`." type = string - default = "ghcr.io/berriai/litellm-ui:v1.86.0-dev" + default = "" } variable "migrations_image" { description = <<-EOT - Container image for the one-off prisma migration task. Built from + Full image URI for the one-off prisma migration task. Empty (default) + composes from `image_registry` + `image_tag`. Built from `migrations/Dockerfile` — slim image whose ENTRYPOINT runs `python3 /app/run.py` (assembles DATABASE_URL from DATABASE_* env vars via DatabaseURLSettings, then runs `prisma migrate deploy`). Should track the same release tag as gateway/backend/ui. EOT type = string - default = "ghcr.io/berriai/litellm-migrations:v1.86.0-dev" + default = "" } # ---------- Service sizing ---------- diff --git a/terraform/litellm/aws/versions.tf b/terraform/litellm/aws/versions.tf index 73b88e91dce..06ebe5bd2c5 100644 --- a/terraform/litellm/aws/versions.tf +++ b/terraform/litellm/aws/versions.tf @@ -1,5 +1,5 @@ terraform { - required_version = ">= 1.6.0" + required_version = ">= 1.9.0" required_providers { aws = { diff --git a/terraform/litellm/gcp/.terraform.lock.hcl b/terraform/litellm/gcp/.terraform.lock.hcl index 013391b53d4..ad8b77ce422 100644 --- a/terraform/litellm/gcp/.terraform.lock.hcl +++ b/terraform/litellm/gcp/.terraform.lock.hcl @@ -6,6 +6,7 @@ provider "registry.terraform.io/hashicorp/google" { constraints = "~> 6.10" hashes = [ "h1:79CwMTsp3Ud1nOl5hFS5mxQHyT0fGVye7pqpU0PPlHI=", + "h1:faTJQOetP9/RYuHwA3r2SWnuYoyzQNm4tUWZrZggcgY=", "zh:1f3513fcfcbf7ca53d667a168c5067a4dd91a4d4cccd19743e248ff31065503c", "zh:3da7db8fc2c51a77dd958ea8baaa05c29cd7f829bd8941c26e2ea9cb3aadc1e5", "zh:3e09ac3f6ca8111cbb659d38c251771829f4347ab159a12db195e211c76068bb", @@ -26,6 +27,7 @@ provider "registry.terraform.io/hashicorp/google-beta" { constraints = "~> 6.10" hashes = [ "h1:P2GiUJM1frlPtBViwKn1A9V2dVBdGuWcX80w9TdH8ZE=", + "h1:uxh4ME3hhSzVjmiWgA1IQqYqg25MV6FMVArHyA6Ki5o=", "zh:18b442bd0a05321d39dda1e9e3f1bdede4e61bc2ac62cc7a67037a3864f75101", "zh:2e387c51455862828bec923a3ec81abf63a4d998da470cf00e09003bda53d668", "zh:3942e708fa84ebe54996086f4b1398cb747fe19cbcd0be07ace528291fb35dee", @@ -45,6 +47,7 @@ provider "registry.terraform.io/hashicorp/random" { version = "3.8.1" constraints = "~> 3.6" hashes = [ + "h1:Eexl06+6J+s75uD46+WnZtpJZYRVUMB0AiuPBifK6Jc=", "h1:u8AKlWVDTH5r9YLSeswoVEjiY72Rt4/ch7U+61ZDkiQ=", "zh:08dd03b918c7b55713026037c5400c48af5b9f468f483463321bd18e17b907b4", "zh:0eee654a5542dc1d41920bbf2419032d6f0d5625b03bd81339e5b33394a3e0ae", diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index c90af26ebf2..3227dcac179 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -54,11 +54,15 @@ Then point the stack at it via `image_registry`: ```hcl image_registry = "us-central1-docker.pkg.dev/my-gcp-project/litellm/berriai" -image_tag = "v1.86.0-dev" ``` -The four `litellm-:${image_tag}` URIs are composed from those -two vars. Set `gateway_image` / `backend_image` / `ui_image` / +The four `litellm-:${image_tag}` URIs are composed from +`image_registry` and `image_tag`. A published module version (Terraform +Registry `BerriAI/litellm/google`) ships with `image_tag` stamped to its +own release tag, so bumping the module version bumps all four images with +it. A git ref of this monorepo has no published images of its own, so +there `image_tag` is required. Set `image_tag` to pin a different LiteLLM +release. Set `gateway_image` / `backend_image` / `ui_image` / `migrations_image` only if you need a per-component override (custom build, different tag). @@ -485,6 +489,8 @@ provider "google-beta" { module "litellm" { source = "github.com/BerriAI/litellm//terraform/litellm/gcp?ref=" + image_tag = "" # a git ref ships no image default, a registry version does + project = "my-gcp-project" region = "us-central1" tenant = "acme" diff --git a/terraform/litellm/gcp/examples/default/.terraform.lock.hcl b/terraform/litellm/gcp/examples/default/.terraform.lock.hcl index e6285567315..55568afb834 100644 --- a/terraform/litellm/gcp/examples/default/.terraform.lock.hcl +++ b/terraform/litellm/gcp/examples/default/.terraform.lock.hcl @@ -6,6 +6,7 @@ provider "registry.terraform.io/hashicorp/google" { constraints = "~> 6.10" hashes = [ "h1:79CwMTsp3Ud1nOl5hFS5mxQHyT0fGVye7pqpU0PPlHI=", + "h1:faTJQOetP9/RYuHwA3r2SWnuYoyzQNm4tUWZrZggcgY=", "zh:1f3513fcfcbf7ca53d667a168c5067a4dd91a4d4cccd19743e248ff31065503c", "zh:3da7db8fc2c51a77dd958ea8baaa05c29cd7f829bd8941c26e2ea9cb3aadc1e5", "zh:3e09ac3f6ca8111cbb659d38c251771829f4347ab159a12db195e211c76068bb", @@ -26,6 +27,7 @@ provider "registry.terraform.io/hashicorp/google-beta" { constraints = "~> 6.10" hashes = [ "h1:P2GiUJM1frlPtBViwKn1A9V2dVBdGuWcX80w9TdH8ZE=", + "h1:uxh4ME3hhSzVjmiWgA1IQqYqg25MV6FMVArHyA6Ki5o=", "zh:18b442bd0a05321d39dda1e9e3f1bdede4e61bc2ac62cc7a67037a3864f75101", "zh:2e387c51455862828bec923a3ec81abf63a4d998da470cf00e09003bda53d668", "zh:3942e708fa84ebe54996086f4b1398cb747fe19cbcd0be07ace528291fb35dee", @@ -46,6 +48,7 @@ provider "registry.terraform.io/hashicorp/random" { constraints = "~> 3.6" hashes = [ "h1:OO+IuvQJSPmWdN8AyyIEvPJbLvDQpgX/zbktoa9KsJE=", + "h1:UlBuNVuCGJ39tTv2c5gz2NRZnQbXfbIWbTzWcth5o74=", "zh:161ad0bd9a75768c82f53fb6e7172a9d8be2d4889b012645a34795031aaf1bf1", "zh:19dc9a5b17729725ccfc4f45b0500af0ee5bc6b6b160c7adb8f2bf617d2c80ea", "zh:269eda8fe42daa7974d5a34d166c3ba9defe80cde86c01e4dadcfdf2e1f05e5f", diff --git a/terraform/litellm/gcp/examples/default/deploystack.json b/terraform/litellm/gcp/examples/default/deploystack.json index 47d1fd914ce..732d4a187b1 100644 --- a/terraform/litellm/gcp/examples/default/deploystack.json +++ b/terraform/litellm/gcp/examples/default/deploystack.json @@ -25,7 +25,7 @@ { "name": "image_tag", "description": "Tag for the four litellm-* images (gateway, backend, ui, migrations). Bump together when bumping LiteLLM", - "default": "v1.86.0-dev" + "default": "" }, { "name": "image_registry", diff --git a/terraform/litellm/gcp/examples/default/main.tf b/terraform/litellm/gcp/examples/default/main.tf index 782d4ee4b65..9b18a6df863 100644 --- a/terraform/litellm/gcp/examples/default/main.tf +++ b/terraform/litellm/gcp/examples/default/main.tf @@ -41,7 +41,7 @@ module "litellm" { ui_password = var.ui_password image_registry = var.image_registry - image_tag = var.image_tag + image_tag = var.image_tag != "" ? var.image_tag : null lb_domains = var.lb_domains allow_plaintext_lb = var.allow_plaintext_lb diff --git a/terraform/litellm/gcp/examples/default/terraform.tfvars.example b/terraform/litellm/gcp/examples/default/terraform.tfvars.example index ec6206734e2..6456b525d08 100644 --- a/terraform/litellm/gcp/examples/default/terraform.tfvars.example +++ b/terraform/litellm/gcp/examples/default/terraform.tfvars.example @@ -38,12 +38,13 @@ env = "stage" # Images. Cloud Run rejects ghcr.io, so a real deploy must point # image_registry at an Artifact Registry remote repo (see README "Image -# pulls"); image_tag is applied to all four litellm-* images. Per-component -# *_image overrides are NOT exposed here — set them directly on the -# `module "litellm"` block in main.tf (see ../../variables.tf) if you need -# to mix-and-match versions. +# pulls"); image_tag is applied to all four litellm-* images. A published +# module version ships with it stamped to its own release tag, a checkout of +# the monorepo has no default and needs it set. Per-component *_image overrides are +# NOT exposed here. Set them directly on the `module "litellm"` block in +# main.tf (see ../../variables.tf) if you need to mix-and-match versions. # image_registry = "us-central1-docker.pkg.dev/my-gcp-project/litellm/berriai" -# image_tag = "v1.86.0-dev" +# image_tag = "" # ---------- proxy_config (mirrors helm gateway.config.proxy_config) ---------- # proxy_config = { diff --git a/terraform/litellm/gcp/examples/default/variables.tf b/terraform/litellm/gcp/examples/default/variables.tf index 08b78346df3..dcf84572fd7 100644 --- a/terraform/litellm/gcp/examples/default/variables.tf +++ b/terraform/litellm/gcp/examples/default/variables.tf @@ -85,7 +85,7 @@ variable "image_registry" { variable "image_tag" { description = "Tag applied to all four litellm-* images. Bump in lockstep." type = string - default = "v1.86.0-dev" + default = "" } # TLS — provide DNS names for a managed cert, or opt into HTTP-only for dev. diff --git a/terraform/litellm/gcp/examples/default/versions.tf b/terraform/litellm/gcp/examples/default/versions.tf index a630c59afd0..55214b5c51a 100644 --- a/terraform/litellm/gcp/examples/default/versions.tf +++ b/terraform/litellm/gcp/examples/default/versions.tf @@ -1,5 +1,5 @@ terraform { - required_version = ">= 1.6.0" + required_version = ">= 1.9.0" required_providers { google = { diff --git a/terraform/litellm/gcp/tests/collector.tftest.hcl b/terraform/litellm/gcp/tests/collector.tftest.hcl index 7a5ea781acb..f16d02d9c7e 100644 --- a/terraform/litellm/gcp/tests/collector.tftest.hcl +++ b/terraform/litellm/gcp/tests/collector.tftest.hcl @@ -7,6 +7,7 @@ mock_provider "google-beta" {} mock_provider "random" {} variables { + image_tag = "v9.9.9" project_id = "acme-test" region = "us-central1" tenant = "acme" diff --git a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl index 999e4f0ff95..c0ad842ff95 100644 --- a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl @@ -18,6 +18,7 @@ mock_provider "google-beta" {} mock_provider "random" {} variables { + image_tag = "v9.9.9" project_id = "test-project" tenant = "tenant" env = "test" diff --git a/terraform/litellm/gcp/tests/deps_only.tftest.hcl b/terraform/litellm/gcp/tests/deps_only.tftest.hcl index 610c49d5b52..11815af9051 100644 --- a/terraform/litellm/gcp/tests/deps_only.tftest.hcl +++ b/terraform/litellm/gcp/tests/deps_only.tftest.hcl @@ -14,6 +14,7 @@ mock_provider "google-beta" {} mock_provider "random" {} variables { + image_tag = "v9.9.9" project_id = "test-project" tenant = "tenant" env = "test" @@ -112,6 +113,20 @@ run "deps_only_creates_no_runtime" { } } +run "deps_only_needs_no_image_tag" { + command = plan + + variables { + create_runtime = false + image_tag = "" + } + + assert { + condition = length(google_cloud_run_v2_service.gateway) == 0 + error_message = "deps-only mode must plan without an image tag since it deploys no images" + } +} + run "existing_network_attaches_data_stores" { command = plan diff --git a/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl b/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl index 7de62065d23..fb5446b4c56 100644 --- a/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl +++ b/terraform/litellm/gcp/tests/metrics_sidecar.tftest.hcl @@ -14,6 +14,7 @@ mock_provider "google-beta" {} mock_provider "random" {} variables { + image_tag = "v9.9.9" project_id = "test-project" tenant = "tenant" env = "test" diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index 412f919ab89..6161e893e02 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -144,9 +144,15 @@ variable "image_registry" { } variable "image_tag" { - description = "Tag applied to all four litellm-* images when composed from `image_registry`. Bump in lockstep when bumping LiteLLM. Must match a tag actually published to GHCR — the split images use the `v`-prefixed semver convention (e.g. `v1.86.0-dev`)." + description = "Tag applied to all four litellm-* images when composed from `image_registry`. A published module version ships with this stamped to its own release tag, which matches the `v`-prefixed tag published to GHCR for the same LiteLLM release. Set it to pin a different LiteLLM release. When sourcing the module from a git ref of the monorepo it has no default and must be set." type = string - default = "v1.86.0-dev" + default = "" + nullable = false + + validation { + condition = !var.create_runtime || var.image_tag != "" || (var.gateway_image != "" && var.backend_image != "" && var.ui_image != "" && var.migrations_image != "") + error_message = "image_tag is empty. Published module versions (registry.terraform.io/modules/BerriAI/litellm) ship with it stamped to their own release. A git ref of the monorepo has no published images of its own, so set image_tag to a published LiteLLM release tag or set all four *_image URIs." + } } variable "gateway_image" { diff --git a/terraform/litellm/gcp/versions.tf b/terraform/litellm/gcp/versions.tf index a630c59afd0..55214b5c51a 100644 --- a/terraform/litellm/gcp/versions.tf +++ b/terraform/litellm/gcp/versions.tf @@ -1,5 +1,5 @@ terraform { - required_version = ">= 1.6.0" + required_version = ">= 1.9.0" required_providers { google = { From 8b1859c50d337b9b8410dca35e58e293c87dd1a5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 17 Oct 2026 10:25:44 -0700 Subject: [PATCH 14/64] fix(mcp)!: cap the ID-JAG credential by its signed resource claim and refuse a resource the gateway does not serve (#47094) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/gateway_dcr_flow.py | 3 +- .../_experimental/mcp_server/id_jag_grant.py | 44 +++++++++-- .../mcp_server/test_gateway_dcr_flow.py | 1 + .../mcp_server/test_id_jag_grant.py | 75 +++++++++++++++++++ 4 files changed, 117 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 90289fd4c20..70e4bac01b5 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -248,6 +248,7 @@ class SubjectTokenRefusal(LiteLLMBaseModel): "invalid_grant", "invalid_client", "invalid_scope", + "invalid_target", "temporarily_unavailable", ] description: str = Field(min_length=1) @@ -1527,7 +1528,7 @@ def _subject_token_refusal_response(refusal: SubjectTokenRefusal) -> Response: return _oauth_error(503, refusal.error, refusal.description) case "invalid_client": return _oauth_error(401, refusal.error, refusal.description) - case "unsupported_grant_type" | "invalid_request" | "invalid_grant" | "invalid_scope": + case "unsupported_grant_type" | "invalid_request" | "invalid_grant" | "invalid_scope" | "invalid_target": return _oauth_error(400, refusal.error, refusal.description) case _: assert_never(refusal.error) diff --git a/litellm/proxy/_experimental/mcp_server/id_jag_grant.py b/litellm/proxy/_experimental/mcp_server/id_jag_grant.py index 155d5280f5a..f2177df5a26 100644 --- a/litellm/proxy/_experimental/mcp_server/id_jag_grant.py +++ b/litellm/proxy/_experimental/mcp_server/id_jag_grant.py @@ -23,6 +23,7 @@ from urllib.parse import unquote_plus import jwt from fastapi import Request +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.models.user import LiteLLM_UserTable @@ -48,6 +49,7 @@ _INACTIVE_SUBJECT: Final = frozenset(("no_active_key", "unresolvable")) _REPLAY_KEY_PREFIX: Final = "mcp_gateway_dcr:id_jag:" SCOPE_MCP: Final = "litellm:mcp" SCOPE_LLM: Final = "litellm:llm" +_RESOURCE_CLAIM_ADAPTER: Final = TypeAdapter(list[str]) SCOPE_ROUTE_GROUPS: Final = MappingProxyType( {SCOPE_MCP: ("mcp_inference_routes",), SCOPE_LLM: ("llm_api_routes_without_mcp",)} ) @@ -65,9 +67,9 @@ class GrantedScope: def granted_scope(claim: object, *, mcp_resource: bool) -> GrantedScope | SubjectTokenRefusal: """The routes the assertion buys. With ``aud`` the bare origin the ``scope`` claim decides and - no claim is the full credential. With ``aud`` the MCP resource (``{base}/mcp``) the identity - provider named that resource, so the credential is ceilinged to the MCP routes whatever the - claim says: ``litellm:llm`` cannot be granted and is refused.""" + no claim is the full credential. With the MCP resource (``{base}/mcp``) as ``aud`` or as every + ``resource``, the identity provider named that resource, so the credential is ceilinged to the + MCP routes whatever the claim says: ``litellm:llm`` cannot be granted and is refused.""" if claim is not None and not isinstance(claim, str): return SubjectTokenRefusal(error="invalid_scope", description=_SCOPE_DESCRIPTION) requested: Final = frozenset(claim.split()) if claim else frozenset() @@ -90,6 +92,32 @@ def names_mcp_resource(aud: object, base_url: str | None) -> bool: return base_url is not None and f"{base_url}/mcp" in audiences and base_url not in audiences +def _resource_claim_values(claim: object) -> tuple[str, ...] | None: + if isinstance(claim, str): + return (claim,) + if not isinstance(claim, list): + return None + try: + values: Final = _RESOURCE_CLAIM_ADAPTER.validate_python(claim, strict=True) + except ValidationError: + return None + return tuple(values) if values else None + + +def resource_names_mcp(claims: Mapping[str, object], base_url: str | None) -> bool | SubjectTokenRefusal: + """The signed RFC 8707 ``resource`` claim (a URI or an array of URIs) restricts the credential like an + ``aud`` of ``{base}/mcp`` does.""" + if "resource" not in claims: + return False + resources: Final = _resource_claim_values(claims["resource"]) + if resources is None: + return SubjectTokenRefusal(error="invalid_target", description=_RESOURCE_DESCRIPTION) + served_resources: Final = frozenset(gateway_audiences(base_url)) + if any(resource not in served_resources for resource in resources): + return SubjectTokenRefusal(error="invalid_target", description=_RESOURCE_DESCRIPTION) + return frozenset(resources) == frozenset((f"{base_url}/mcp",)) + + def _route_groups(scopes: Sequence[str]) -> Iterator[str]: for scope in scopes: yield from SCOPE_ROUTE_GROUPS[scope] @@ -98,8 +126,10 @@ def _route_groups(scopes: Sequence[str]) -> Iterator[str]: _SCOPE_DESCRIPTION: Final = ( f"scope must be {SCOPE_MCP} (MCP gateway routes), {SCOPE_LLM} (LLM API routes), or both, space separated" ) +_RESOURCE_DESCRIPTION: Final = "the assertion's resource claim must name this gateway's origin or its /mcp resource" _MCP_RESOURCE_DESCRIPTION: Final = ( - f"the /mcp audience only grants MCP routes ({SCOPE_MCP}); use the bare origin as audience for LLM routes" + f"the /mcp audience or resource only grants MCP routes ({SCOPE_MCP}); use the bare origin as audience and " + "resource for LLM routes" ) @@ -450,8 +480,12 @@ async def identity_from_id_jag( if isinstance(marker, SubjectTokenRefusal): return marker replay_key, replay_ttl_seconds = marker + mcp_resource: Final = resource_names_mcp(verified, prerequisites.base_url) + if isinstance(mcp_resource, SubjectTokenRefusal): + return mcp_resource scope: Final = granted_scope( - verified.get("scope"), mcp_resource=names_mcp_resource(verified.get("aud"), prerequisites.base_url) + verified.get("scope"), + mcp_resource=mcp_resource or names_mcp_resource(verified.get("aud"), prerequisites.base_url), ) if isinstance(scope, SubjectTokenRefusal): return scope diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index ef95df23ff8..530667fdfdd 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -3384,6 +3384,7 @@ async def test_jwt_bearer_grant_claims_the_jti_before_it_mints(): ("invalid_client", 401), ("invalid_grant", 400), ("invalid_request", 400), + ("invalid_target", 400), ("unsupported_grant_type", 400), ("temporarily_unavailable", 503), ], diff --git a/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py b/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py index fdfa0544ad1..f4fa51d57ba 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_id_jag_grant.py @@ -33,6 +33,7 @@ from litellm.proxy._experimental.mcp_server.id_jag_grant import ( jwt_bearer_available, names_mcp_resource, proxy_user_lookup, + resource_names_mcp, select_team, subject_resolver, ) @@ -87,6 +88,16 @@ def _assertion(typ: str = ID_JAG_TYP, **claims: object) -> str: ) +def _unscoped_assertion(**claims: object) -> str: + unscoped_claims: Final = {key: value for key, value in _claims(**claims).items() if key != "scope"} + return jwt.encode( + unscoped_claims, + "unused-by-the-fake-verifier-but-long-enough-for-hs256", + algorithm="HS256", + headers={"typ": ID_JAG_TYP}, + ) + + class _Verifier: """Stands in for the injected ``DecodeAssertion`` port; records each call as a new tuple.""" @@ -714,6 +725,40 @@ def test_the_mcp_resource_audience_is_the_mcp_path_without_the_bare_origin(aud, assert names_mcp_resource(aud, None) is False +@pytest.mark.parametrize( + "claims, expected", + [ + ({}, False), + ({"resource": f"{GATEWAY}/mcp"}, True), + ({"resource": [f"{GATEWAY}/mcp"]}, True), + ({"resource": GATEWAY}, False), + ({"resource": [GATEWAY, f"{GATEWAY}/mcp"]}, False), + ({"resource": "https://other.example"}, None), + ({"resource": [f"{GATEWAY}/mcp", "https://other.example"]}, None), + ({"resource": f"{GATEWAY}/mcp/"}, None), + ({"resource": None}, None), + ({"resource": 7}, None), + ({"resource": []}, None), + ({"resource": [7]}, None), + ], +) +def test_the_signed_resource_claim_must_name_only_gateway_resources( + claims: Mapping[str, object], expected: bool | None +) -> None: + result: Final = resource_names_mcp(claims, GATEWAY) + if isinstance(result, bool): + assert expected is result + return + assert expected is None + assert result.error == "invalid_target" + + +def test_the_signed_resource_claim_is_refused_without_a_gateway_base_url() -> None: + result: Final = resource_names_mcp({"resource": f"{GATEWAY}/mcp"}, None) + assert isinstance(result, SubjectTokenRefusal) + assert result.error == "invalid_target" + + @pytest.mark.asyncio async def test_redeeming_for_the_mcp_resource_narrows_the_identity_and_refuses_llm_scope(): """End to end through identity_from_id_jag: the same verified claims, with ``aud`` the MCP @@ -736,3 +781,33 @@ async def test_redeeming_for_the_mcp_resource_narrows_the_identity_and_refuses_l scoped = await _identity(_redemption(_assertion(aud=AUDIENCES[1], scope="litellm:llm")), _Verifier(), _Resolver()) assert isinstance(scoped, AssertionIdentity) assert scoped.scope == "litellm:llm" + + +@pytest.mark.asyncio +async def test_redeeming_for_the_signed_resource_caps_the_identity_before_subject_lookup() -> None: + mcp_claim: Final = _unscoped_assertion(aud=GATEWAY, resource=f"{GATEWAY}/mcp") + narrowed: Final = await _identity(_redemption(mcp_claim), _Verifier(), _Resolver()) + assert isinstance(narrowed, AssertionIdentity) + assert (narrowed.scope, narrowed.allowed_routes) == ("litellm:mcp", ("mcp_inference_routes",)) + + llm_claim: Final = _assertion(aud=GATEWAY, resource=f"{GATEWAY}/mcp", scope="litellm:llm") + refused_scope: Final = await _identity(_redemption(llm_claim), _Verifier(), _Resolver()) + assert isinstance(refused_scope, SubjectTokenRefusal) + assert refused_scope.error == "invalid_scope" + + other_resource: Final = _unscoped_assertion(aud=GATEWAY, resource="https://other.example") + resolver: Final = _Resolver() + refused_resource: Final = await _identity(_redemption(other_resource), _Verifier(), resolver) + assert isinstance(refused_resource, SubjectTokenRefusal) + assert refused_resource.error == "invalid_target" + assert resolver.calls == () + + origin_resource: Final = _unscoped_assertion(aud=GATEWAY, resource=GATEWAY) + full_resource: Final = await _identity(_redemption(origin_resource), _Verifier(), _Resolver()) + assert isinstance(full_resource, AssertionIdentity) + assert (full_resource.scope, full_resource.allowed_routes) == (None, None) + + no_resource: Final = _unscoped_assertion(aud=GATEWAY) + full_control: Final = await _identity(_redemption(no_resource), _Verifier(), _Resolver()) + assert isinstance(full_control, AssertionIdentity) + assert (full_control.scope, full_control.allowed_routes) == (None, None) From 3c19b4149b4cb3858278b6927bebd7092a0cefa8 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 10:44:49 -0700 Subject: [PATCH 15/64] fix(cost-map): fall back to the bundled map when urlsplit rejects LITELLM_MODEL_COST_MAP_URL (#47133) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/get_model_cost_map.py | 22 ++++++++++++++----- .../test_get_model_cost_map.py | 19 ++++++++++++++++ 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 6e3769260cc..b68db22183a 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -324,10 +324,16 @@ def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _Fetc return ModelCostMapReloaded(model_cost_map=model_cost_map, revision=git_blob_id(content), etag=etag) -def _attempt_fetch_s3(url: str, timeout: int) -> _FetchAttemptOutcome: +def _split_url(url: str) -> urllib.parse.SplitResult | ModelCostMapReloadUnavailable: + try: + return urllib.parse.urlsplit(url) + except ValueError as error: + return ModelCostMapReloadUnavailable(reason=f"invalid URL {url}: {error}") + + +def _attempt_fetch_s3(parts: urllib.parse.SplitResult, url: str, timeout: int) -> _FetchAttemptOutcome: from botocore.exceptions import BotoCoreError, ClientError - parts: Final = urllib.parse.urlsplit(url) bucket: Final = parts.netloc key: Final = parts.path.lstrip("/") if not bucket or not key: @@ -343,9 +349,11 @@ def _attempt_fetch_s3(url: str, timeout: int) -> _FetchAttemptOutcome: async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: - parts: Final = urllib.parse.urlsplit(url) + parts: Final = _split_url(url) + if isinstance(parts, ModelCostMapReloadUnavailable): + return parts if parts.scheme == "s3": - return await asyncio.to_thread(_attempt_fetch_s3, url, timeout) + return await asyncio.to_thread(_attempt_fetch_s3, parts, url, timeout) try: response: Final = await client.get(url, timeout=timeout) except (httpx.HTTPError, httpx.InvalidURL) as e: @@ -354,9 +362,11 @@ async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _Fe def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: - parts: Final = urllib.parse.urlsplit(url) + parts: Final = _split_url(url) + if isinstance(parts, ModelCostMapReloadUnavailable): + return parts if parts.scheme == "s3": - return _attempt_fetch_s3(url, timeout) + return _attempt_fetch_s3(parts, url, timeout) try: response: Final = client.get(url, timeout=timeout) except (httpx.HTTPError, httpx.InvalidURL) as e: diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 0256ecf24ed..eca9bbdc170 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -686,6 +686,25 @@ def test_boot_load_does_not_retry_non_retryable_failure(): assert source["fallback_reason"] is not None +def test_boot_load_falls_back_when_urlsplit_rejects_the_url(): + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map( + url="http://[bad/model_prices.json", + sleep=sleeper, + rng=random.Random(0), + client=client, + ) + assert calls["count"] == 0 + assert sleeper.waits == [] + assert _retry_threads() == [] + assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert "invalid URL" in (source["fallback_reason"] or "") + + def test_boot_load_respects_local_env_override(monkeypatch): """LITELLM_LOCAL_MODEL_COST_MAP=True still short-circuits to the backup with zero HTTP.""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") From 656c63c0b437b8d0f28db3f2599cc4a91c87d1bf Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:53:59 +0000 Subject: [PATCH 16/64] test(translation): add claude-haiku-5-5 vertex ai cases (#47136) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/proxy_config.yaml | 7 + .../chat_completions/bases/vertex_ai.py | 148 +++++++++++++++++ .../test_chat_completions_basic_vertex_ai.py | 2 + .../translation/messages/bases/vertex_ai.py | 101 ++++++++++++ .../basic/test_messages_basic_vertex_ai.py | 2 + .../translation/responses/bases/vertex_ai.py | 155 ++++++++++++++++++ .../basic/test_responses_basic_vertex_ai.py | 2 + 7 files changed, 417 insertions(+) diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index 37bcc4683dd..4e2cab56d27 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -473,6 +473,13 @@ model_list: vertex_project: synthetic-vertex-project vertex_location: global vertex_credentials: '{"type": "external_account", "audience": "synthetic-vertex-audience", "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": "http://127.0.0.1:8190/_oauth/token", "credential_source": {"url": "http://127.0.0.1:8190/health"}}' + - model_name: vertex_ai/claude-haiku-5-5@default + litellm_params: + model: vertex_ai/claude-haiku-5-5@default + api_base: http://127.0.0.1:8191 + vertex_project: synthetic-vertex-project + vertex_location: global + vertex_credentials: '{"type": "external_account", "audience": "synthetic-vertex-audience", "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": "http://127.0.0.1:8190/_oauth/token", "credential_source": {"url": "http://127.0.0.1:8190/health"}}' - model_name: vertex_ai/claude-opus-5-5@default litellm_params: model: vertex_ai/claude-opus-5-5@default diff --git a/tests/integration/translation/chat_completions/bases/vertex_ai.py b/tests/integration/translation/chat_completions/bases/vertex_ai.py index d1022da78c7..447566b0ee3 100644 --- a/tests/integration/translation/chat_completions/bases/vertex_ai.py +++ b/tests/integration/translation/chat_completions/bases/vertex_ai.py @@ -577,6 +577,154 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_THINKING_SIGNATURE: Final = ( + "CAQS8gcKEggSEAIYAjgBQgh0aGlua2luZxIMKxM+8GdyGfxkJrU6Ggw3AMSc5G/SaulUVJkiMJgPNRytmwmfF4bj7WugkKlRxBr+SRwcL9iQ" + "+S5uvmIi/D8thAXX+oVg67diUo3rxiqNB0KrfSB/DXD50OomDcy90ai3a+RMcwIy9wEMzOdOIFqQ24tUMm1YorcQLQpresEnXA9ItiMAzP+P" + "8NMJ1br7O3JHYLIi/12UzVr36AKPrMJQk86VBn5fVxo0xcphyl071Y30m3wO0INGtU+d37IXSrYCJzQdnUSKxekBK28fxsA4tmxx1BE/1Hfl" + "UTsq36+CE2MatE0ymL2pQA+povoJ1iVlB/PBZvkOvpRcQubSwm8RqJrBfmse8baU2UJot0VNNC7OprQ3wYx1j3BGmXOQCGkvCn6jd2lPLjLS" + "F1CjbWQs7X18gB69y2sUosm3aI+MnQNHkUqe8JtIaii42uu3Xe/JheeRTRrcsuAdMZJ76ID9wRZnf3/QCHaMyrYwHXf5SsoUjX3Ofs5MuQT0" + "ISf4T7q+QGf51Bxb4jdfF879h061ZSejzknCHI3fV3SwuldjIIV+FD50d63iLqM/D8GOjaGfDrYugkkdBjCNEQkTpn9NQmqalsrJWOdQQrhc" + "L8TDZG4HBasd/kGr7EnlEbLKrGaWbAWYo4vssoUE7Q57wPumuMrYh4OKInksQTnB0OV9OIZEiTgb9fyaxJBCUiypuQDGMlAP8nb0TQF1vxjR" + "vsXV/qnN0ovR6/7raePN4NM4nMEXESRQyC22Ml88z+WZhe2VG5CLH8tFNapisH+msyqtGyDa+SEie/e6xcDfI1fLAZwEHM34/YjW+C1E1PPd" + "RsGBE9Ms6JUaPrcITwImVKMLGH1pj+TsS+qK7JTGjH8FKwR1Eds/dJTFDEQaAYcmjXWxDR5HBknx7xv90K2uoKt4KmAgulWCQv/LaEAt27lm" + "aGP+4PQ9tpb/Otsub8spmcXKfNqdjWTDH3aEF5MsJ5eSddn92pFV9JAo4oFcj1fUTujfwumGCYDADjLBmyvQao9S42zpO5idmYrHWHt7h7m6" + "BoU5b4uKVS4w4ELEl8y+Oq1uiQLSB8e0cxLarWfCiiTL6v6aJUDEvUTQU+O3XiOxNMCkc00J7e1TnNnnoFqs0sqhV0mMT346vLHocBMVn2iE" + "ZSwFb9OWlHO1e9Y7ccd8XL7KjXtoZeuBTrdYlxrf3g6QHVZg3OOM3OVSOWawFU0tBZUUgH6rYH4oHh8JqwGmbTI1mtAtUX89KrHzGWgUs619" + "SPbFptVh8Vfb3zr0D0OzDqjzv1KCNW/tZcJnCaL9Zl6tAt2v5FiBTOjIbBgB" +) + +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/chat/completions", + litellm_request={**CHAT_COMPLETIONS_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={ + "authorization": "Bearer scripted-token", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "x-api-key": "scripted-token", + }, + expected_provider_request=VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8AfAmbYAQCTSseeQSZJ", + "content": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually in my " + "available toolset — I only have search_files and read_file. I'll note that discrepancy and proceed by " + "explaining which file I need, then calling read_file on shop/validation.py, probably starting around the " + "first 60 lines to get context near line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` next, starting at the `validate_quantity` definition on line 18 and " + "including the surrounding context. The project instructions mention a chat_completions client, but none " + "is available in my tools, so I'm using the read-only file tools." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_017jjdVAoN8bus5ng5NkzafQ", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 876, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 876, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 348, + "output_tokens_details": {"thinking_tokens": 168}, + }, + }, + expected_litellm_response={ + **CHAT_COMPLETIONS_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": ( + "I'll read `shop/validation.py` next, starting at the `validate_quantity` definition on line 18 and " + "including the surrounding context. The project instructions mention a chat_completions client, but " + "none is available in my tools, so I'm using the read-only file tools." + ), + "role": "assistant", + "tool_calls": [ + { + "index": 2, + "function": { + "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 60}', + "name": "read_file", + }, + "id": "toolu_vrtx_017jjdVAoN8bus5ng5NkzafQ", + "type": "function", + } + ], + "reasoning_content": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually in my " + "available toolset — I only have search_files and read_file. I'll note that discrepancy and proceed " + "by explaining which file I need, then calling read_file on shop/validation.py, probably starting " + "around the first 60 lines to get context near line 18.\n\n" + ), + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually " + "in my available toolset — I only have search_files and read_file. I'll note that discrepancy " + "and proceed by explaining which file I need, then calling read_file on shop/validation.py, " + "probably starting around the first 60 lines to get context near line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + } + ], + "provider_specific_fields": { + "citations": None, + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not " + "actually in my available toolset — I only have search_files and read_file. I'll note " + "that discrepancy and proceed by explaining which file I need, then calling read_file on " + "shop/validation.py, probably starting around the first 60 lines to get context near " + "line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + } + ], + }, + }, + } + ], + "usage": { + "completion_tokens": 348, + "prompt_tokens": 880, + "total_tokens": 1228, + "completion_tokens_details": {"reasoning_tokens": 168, "text_tokens": 180}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 4, + "cache_write_tokens": 876, + "cache_creation_tokens": 876, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 876, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 876, + "cache_read_input_tokens": 0, + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/chat/completions", diff --git a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py index 71cbd84543e..7b7f5a7876c 100644 --- a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py +++ b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.chat_completions.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, diff --git a/tests/integration/translation/messages/bases/vertex_ai.py b/tests/integration/translation/messages/bases/vertex_ai.py index 6f506b39d93..120b764a63c 100644 --- a/tests/integration/translation/messages/bases/vertex_ai.py +++ b/tests/integration/translation/messages/bases/vertex_ai.py @@ -10,6 +10,7 @@ from integration.translation.templates.vertex_ai import ( ) GEMINI_3_5_FLASH_THOUGHT_SIGNATURE: Final = "AY89a1+kriNs9YsPegJQxgjTWsMlxsqPU/UkHDyz4EK+3FRrE52ygqPIHTvkblgKKiO5Fn8R9D3iu//i+YnuuZ+e7W0yxUoFScas89b5rn3vZHsihnUA7VTe3b7beth1IftteoTGZ3F8OSLrSIP8HKn4xTIx35eeNEt5oiD3K5rTQrdklqO277HvTDhu/LuZhZ0pzaWciGQEc61epIBBpC2ygu8V4L8TekUVZnz6qzSzxipdiQStQ8dZou43334hMnxOR6JjJCGxqt/1K+C0vXX6YP6f+Bph5VhPSnvy2Z1uApi3TUxAcsDDyVewIr+ha2G6DPhN1Oi5uZZNhmbvH01Uq5OOquF1uEICq2m/35LhJIbDNM9N7oI1Z6L94P9xdugeFXM2S7wdJfc8sdg09vEzNeA++GK25zRuB4hU7XyaG4XnWjmu2kswi4sf+EyrO/K0P2tb/ZCJu7kJTtWTiD1wgSIl1eqS8ORKKsXDN2hIbFrtsNR50WIvVzd6" +CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE: Final = "CAQS4wYKEggSEAIYAjgBQgh0aGlua2luZxIMgzBnOw4aYVVqPKRPGgyphQhlJ0pt7tgasWIiMOWM4ATud/TWad7qo/Eh9ZSXhK3IcgnwdgqVeXog9D1q/y69qnyrHdN3a8RYUrAgtyr+BYkQvRW/cF0vH0l86vIM95N77usMXdZKRX3I8XW68KpQedYwFoR12sS9q7A24CcN5UMyQljcF76GhmAFnvpk4yEs7pk8b7ZRSXzBsU6wJfMYCoGXC9r3hZN8m4sBHLJGJRv5ekzEVGVvGCYooE1fBZ8ik/j2DCNOyVdDs4/7xzkqRo7dgFsNJI9Nm7+HgwEdvb1rORObJWooeYb8tfofTFVO43qvNdvM2p2d1Dcf6lsQmvE5fgDFc0t75AuvSZUrgDNe2bMXd2USxxL9YtJ1I9ae96hZAnyOqYZlyiuZf5lZtQQSRouVIdZvFUNonOHLy/Uyw/BGF2bO1/5rIFJBDJC8C6BpbyfDQPpal4YkALulDdgeJWI6dmJYUiK+A+3zcVgQh7at5JPbG5n0tQgO2IFkayqLjCf67AaGS6Xq8J97F+FSwppKvS0TvRbTEvIK1iO3HZZRwcRIcFBn5Llvv6R9JFknXEtohUqvfTnZfuPQ104vxqm91G2rFK1R6C1md6Ze3pJWD0wgkgZyQ7PIKsAmBF1qNoniFi5GVRJx6XHSzr4ZAPDDO1aYABIRl5OWilrkQP3xedj5pcAjJqFuto3cXlhQSNUrocOpEV5IIXxSJc5g5VwYAybpMU4FVImzkCJGJaHLWQKRo0k4IcvXLxBPNXZhop2CkEeazZskF4de0jn1dJa3sQ9rT+7fNVEWPvP+lkBqu/lfObBPVh5jGeouDrRcJvyeYKNZ0NQBdznM5rv4rN8HczeYg1WaYuNOMBvkqT5Zg/UhLHhHV7NWBSa0zt9DPevU/S3OBR1E0YIzCtPFTa0Ia4gGTX40hEc+supCA3h4HpehqMqETGe++Ah+U10HRhlb6lVfRhkU0ypdPqtpodgw/OdUOy7f3/sc31xknudQMoEQqv48F4gR34ee18Op+w1MGn0FAzAFE95D2BuJObRqUhZ6N6nWtlWOqJjCPR+0OrwR3JlxWh19Vxps/ASL6HaJKjClVG2ffojvmVnY8aHjEfI8ifbHE6EYAQ==" CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", @@ -486,6 +487,106 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/messages", + litellm_request={**MESSAGES_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={"authorization": "Bearer scripted-token", "content-type": "application/json"}, + expected_provider_request={ + **VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + "messages": MESSAGES_REQUEST["messages"], + "stream": False, + "system": [ + { + "type": "text", + "text": ( + "You help maintain a synthetic Python project. Read the relevant files before suggesting a fix. " + "Explain which file you need next, then call a read-only tool. Never change files. Use the messages " + "client for this project review. Consider zero, negative and non-integer quantities before choosing " + "the next file to inspect." + ), + }, + ], + "tools": MESSAGES_REQUEST["tools"], + }, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8BfjWDUi4ZYQxpy6aVM", + "content": [ + { + "type": "thinking", + "thinking": ( + "I found validate_quantity defined around line 18, so I'll read the validation.py file from the start " + "through roughly line 60 to capture the full context, keeping in mind I'll need to check how it handles " + "zero, negative, and non-integer cases.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` first. The search put `validate_quantity` at line 18, so I'll start " + "there and include the lines around it to see how zero, negative, and non-integer values are handled." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01YSY1kAV4ayCLVdBJkSv4dt", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 292, + "output_tokens_details": {"thinking_tokens": 129}, + }, + }, + expected_litellm_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "id": "msg_vrtx_011Cg8BfjWDUi4ZYQxpy6aVM", + "content": [ + { + "type": "thinking", + "thinking": ( + "I found validate_quantity defined around line 18, so I'll read the validation.py file from the start " + "through roughly line 60 to capture the full context, keeping in mind I'll need to check how it handles " + "zero, negative, and non-integer cases.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` first. The search put `validate_quantity` at line 18, so I'll start " + "there and include the lines around it to see how zero, negative, and non-integer values are handled." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01YSY1kAV4ayCLVdBJkSv4dt", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 292, + "output_tokens_details": {"thinking_tokens": 129}, + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", diff --git a/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py b/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py index a5b01dfd2e4..4f74a648d8e 100644 --- a/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py +++ b/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.messages.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, diff --git a/tests/integration/translation/responses/bases/vertex_ai.py b/tests/integration/translation/responses/bases/vertex_ai.py index 734ebf13613..59637e380ce 100644 --- a/tests/integration/translation/responses/bases/vertex_ai.py +++ b/tests/integration/translation/responses/bases/vertex_ai.py @@ -12,6 +12,7 @@ from integration.translation.templates.vertex_ai import ( ) GEMINI_3_5_FLASH_THOUGHT_SIGNATURE: Final = "AY89a1+kriNs9YsPegJQxgjTWsMlxsqPU/UkHDyz4EK+3FRrE52ygqPIHTvkblgKKiO5Fn8R9D3iu//i+YnuuZ+e7W0yxUoFScas89b5rn3vZHsihnUA7VTe3b7beth1IftteoTGZ3F8OSLrSIP8HKn4xTIx35eeNEt5oiD3K5rTQrdklqO277HvTDhu/LuZhZ0pzaWciGQEc61epIBBpC2ygu8V4L8TekUVZnz6qzSzxipdiQStQ8dZou43334hMnxOR6JjJCGxqt/1K+C0vXX6YP6f+Bph5VhPSnvy2Z1uApi3TUxAcsDDyVewIr+ha2G6DPhN1Oi5uZZNhmbvH01Uq5OOquF1uEICq2m/35LhJIbDNM9N7oI1Z6L94P9xdugeFXM2S7wdJfc8sdg09vEzNeA++GK25zRuB4hU7XyaG4XnWjmu2kswi4sf+EyrO/K0P2tb/ZCJu7kJTtWTiD1wgSIl1eqS8ORKKsXDN2hIbFrtsNR50WIvVzd6" +CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE: Final = "CAQS0gYKEggSEAIYAjgBQgh0aGlua2luZxIMcbnLkNvj2D6NoBl7Ggw+Qgax2/68TAR1maEiMH6XEeKEqR1ewAiuj01fMNuD1KI7pCTXUsWWr6H7nSzP5BiRGXRrDCIxPYu2Xbvx7yrtBbbyId4ReZzx21wZ2+d0vkkpZM82PXY72wQRJOHLR4JNLRjiR2Jn+FrZ8mBo2mofFrlTSuioOVHc1hheiH4InRK57w2X9xaTXa+mwoWW5epPieHD26+NNJirF9vB832fvqDVOd+f8N0ZDUXUut8zl/Jxx+0BMOktMqyfehIOGumQmyv49E3JHVQGrx+Dd3KZjaQXrsRFf1K3AJ7p5XElR+cYis40W7EL3DE9EAtriZuDEt0fTct79fzZIZgHd6Dnua/rGRnjNjaaZKxGjnhfkVHwsZW4HqTNgwWQiD3lt2ncoLfsfd2IWwMEI5uWpN3GtvRkQvfzidmgGmx3dI8g0nBvLPuq0jzVIuuoKIKqzT7A5ae3Gs9eTlSDr2wOGfdyNf6Q2TLw0w7EegI4wYSExKhFkx0EytCGHJcbN0WClXXmHsRtrXnP7sxCMMY7B8TIpAoN7SXyNjJizOGgzoYL7PNUgqp5nVz/C9vTvudhSWbjLITdxFJ5hearL2rMhryfHBKvRGYC9RMy/EjjtINNwYpxdVX29mUE+8KT+41HymtJ1Ewjtoo8V/7XQnQk3cB7X2RFHcR89hAM/pSsB14ofMFoDk+jJitS+sCDZI0Xh7/ySU0o/OI4l+oD/yctZTHx/mIOiWRtnAOrveNT1ms0/2+4tcO0wi3pRrOB2s5ugYc1I3+GcRwILEIM3a4ABd5Dw2nW2fgUhfSTTGURm4Z91ciih8pKUS1ZPatFTOJ1ehKeUCXflkqY2PHnftiNB/8GzlomsQ2NoRw8UXD00oMmBvSvZpPN7khFLxWHtxLjleKkeKCzN7QZkOqvjAnUwy9TacT2ZvSXIxrvaep7dOyrlMEta/1KsYEcWuBxLr87Va4txXnmoV4RqGgaauNcR2wOGyRFJppxKAgYbmMnxams7cv/7IY2RVSgSQf2k4hFQYzRwrHngjapuZzgNS7UCMjvJJs4VdIypVOATsZdPzJToMgWaiEoRZ3zApQDgYhhGAE=" CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", @@ -790,6 +791,160 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/responses", + litellm_request={**RESPONSES_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={ + "authorization": "Bearer scripted-token", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "x-api-key": "scripted-token", + }, + expected_provider_request={ + **VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + "system": [ + { + "type": "text", + "text": ( + "You help maintain a synthetic Python project. Read the relevant files before suggesting a fix. " + "Explain which file you need next, then call a read-only tool. Never change files. Use the responses " + "client for this project review. Consider zero, negative and non-integer quantities before choosing " + "the next file to inspect." + ), + }, + ], + }, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8BhBrWmiX5ZtMxpM9zA", + "content": [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer values " + "before moving to the next file—I'll grab the first 60 lines since I don't know the full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read the validation function in `shop/validation.py`, starting around the match on line 18." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 255, + "output_tokens_details": {"thinking_tokens": 126}, + }, + }, + expected_litellm_response={ + **RESPONSES_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "output": [ + { + "type": "reasoning", + "id": ANY, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "annotations": [], + }, + ], + "phase": None, + "encrypted_content": json.dumps( + [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + ], + separators=(",", ":"), + ), + }, + { + "type": "message", + "id": ANY, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "I'll read the validation function in `shop/validation.py`, starting around the match on line 18.", + "annotations": [], + }, + ], + "phase": None, + }, + { + "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 60}', + "call_id": "toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "name": "read_file", + "type": "function_call", + "id": "fc_toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "namespace": None, + "status": "completed", + }, + ], + "usage": { + "input_tokens": 877, + "input_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "cached_tokens_details": None, + "image_tokens": None, + "text_tokens": 4, + "video_tokens": None, + "cache_write_tokens": 873, + }, + "output_tokens": 255, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 126, "text_tokens": 129}, + "total_tokens": 1132, + "cost": None, + }, + "provider_specific_fields": { + "citations": None, + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + ], + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/responses", diff --git a/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py b/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py index 67c3cf4f7a9..bd9118482d0 100644 --- a/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py +++ b/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.responses.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, From dd8b2853a0a839ef89392bdca025cd36ff3d41ae Mon Sep 17 00:00:00 2001 From: litellm-stampbot Date: Sat, 17 Oct 2026 10:59:53 -0700 Subject: [PATCH 17/64] feat(deepinfra): add DeepSeek-V4.1-Flash and DeepSeek-V4.1-Flash-1017 to the cost map (#47132) * feat(deepinfra): add DeepSeek-V4.1-Flash and DeepSeek-V4.1-Flash-1017 to the cost map * fix(deepinfra): price DeepSeek-V4.1-Flash rows at the current 30% promo rate --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 34 +++++++++++++++++++ model_prices_and_context_window.json | 34 +++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 58338f185c7..8545a5b568e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -65561,6 +65561,40 @@ "supports_vision": false, "source": "https://deepinfra.com/pricing" }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash" + }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash-1017": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash-1017" + }, "deepinfra/openai/gpt-oss-120b-Ultra": { "max_tokens": 131072, "max_input_tokens": 131072, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 58338f185c7..8545a5b568e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -65561,6 +65561,40 @@ "supports_vision": false, "source": "https://deepinfra.com/pricing" }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash" + }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash-1017": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash-1017" + }, "deepinfra/openai/gpt-oss-120b-Ultra": { "max_tokens": 131072, "max_input_tokens": 131072, From 571ef3e2aeafd79f1cf16ba9addc17cf484543a1 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 11:00:51 -0700 Subject: [PATCH 18/64] fix(responses): run pre_call guardrails on every Responses WebSocket turn and bill each turn its own guardrail cost (#47013) * fix(responses): run pre_call guardrails on every Responses WebSocket turn and bill each turn its own guardrail cost * test(responses): keep the websocket unit test file on main's formatting * fix(responses): type the websocket frame guard for the lint gates * fix(responses): pass the websocket frame guard as a named parameter and mark the guardrail dict edits * fix(responses): bill a guardrail-rejected Responses WebSocket frame as a failed request with its guardrail cost * fix(responses): type the recorded guardrail records as Mapping for the lint gate * fix(responses): bill a rejected websocket frame through the spend logger only so the connection keeps its rate-limit slot * fix(responses): guard nested response.create frames, strip forged user_api_key_* frame metadata, copy guarded tools back and give each rejected frame its own spend row id * fix(responses): mark the spend logger's bare-dict failure hook for the type gate * fix(responses): reject response.create frames that carry input at both levels and keep billing of a rejected frame non-blocking * fix(responses): name the non-blocking billing catch for the strict lint gate * fix(responses): keep the admitted guardrail list on every websocket frame, a frame cannot replace it with its own guardrails list * test(responses): wait for the key spend total before asserting it in the websocket guardrail cost test * fix(responses): copy the admitted metadata shallowly per websocket frame, the OTEL parent span cannot be deep-copied * fix(responses): reject a WebSocket turn a pre_call guardrail rerouted to another model * fix(responses): bill a rerouted WS turn and carry Presidio pii_tokens from the frame guard to the socket unmasker * fix(responses): default pii_tokens through default_factory, Python 3.13 rejects a mappingproxy dataclass default * fix(responses): run the per-frame guardrail walk on managed Responses WebSockets too * fix(responses): unmask managed Responses WebSocket replies with the frame guard's Presidio pii_tokens * fix(responses): strip client guardrail cost records from WS frames so a blocked frame bills the real scan * fix(responses): forward a WS frame without the fields a guardrail deleted --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/llm_http_handler.py | 64 +-- .../proxy/response_api_endpoints/endpoints.py | 153 ++++++- litellm/responses/main.py | 3 + litellm/responses/streaming_iterator.py | 280 +++++++++--- ...est_responses_websocket_guardrail_spend.py | 168 +++++++ .../response_api_endpoints/test_endpoints.py | 426 ++++++++++++++++-- .../test_responses_websocket_all_providers.py | 264 +++++++++++ 7 files changed, 1225 insertions(+), 133 deletions(-) create mode 100644 tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4ff4c43428f..1bccef7dd9e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -216,6 +216,7 @@ if TYPE_CHECKING: FakeAnthropicMessagesStreamIterator, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard from litellm.types.google_genai.main import GenerateContentResponse from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -6867,6 +6868,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs: object, ) -> Exception | None: """ @@ -6882,6 +6884,35 @@ class BaseLLMHTTPHandler: """ _ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks() + _ws_guardrail_callbacks: list = [] + _ws_output_guardrail_callbacks: list = [] + try: + import litellm as _litellm + + # Use duck-typing so any guardrail that exposes the PII + # masking interface works, not just _OPTIONAL_PresidioPIIMasking. + # This avoids a layering violation (SDK importing from proxy). + _ws_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if callable(getattr(cb, "check_pii", None)) + and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) + and callable(getattr(cb, "_unmask_pii_text", None)) + and getattr(cb, "output_parse_pii", False) + ] + _ws_output_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if callable(getattr(cb, "check_pii", None)) + and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) + and getattr(cb, "apply_to_output", False) + ] + except Exception as _guardrail_exc: + verbose_logger.warning( + "Responses WebSocket: failed to collect guardrail callbacks — PII masking will be skipped. Error: %s", + _guardrail_exc, + ) + if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket(): from litellm.responses.streaming_iterator import ( ManagedResponsesWebSocketHandler, @@ -6899,6 +6930,8 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, + guardrail_callbacks=_ws_guardrail_callbacks, # pyright: ignore[reportUnknownArgumentType] # duck-typed callbacks from litellm.callbacks + frame_guard=frame_guard, **kwargs, ) await handler.run() @@ -6980,36 +7013,6 @@ class BaseLLMHTTPHandler: if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata - _ws_guardrail_callbacks: list = [] - _ws_output_guardrail_callbacks: list = [] - try: - import litellm as _litellm - - # Use duck-typing so any guardrail that exposes the PII - # masking interface works, not just _OPTIONAL_PresidioPIIMasking. - # This avoids a layering violation (SDK importing from proxy). - _ws_guardrail_callbacks = [ - cb - for cb in _litellm.callbacks - if callable(getattr(cb, "check_pii", None)) - and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) - and callable(getattr(cb, "_unmask_pii_text", None)) - and getattr(cb, "output_parse_pii", False) - ] - _ws_output_guardrail_callbacks = [ - cb - for cb in _litellm.callbacks - if callable(getattr(cb, "check_pii", None)) - and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) - and getattr(cb, "apply_to_output", False) - ] - except Exception as _guardrail_exc: - verbose_logger.warning( - "Responses WebSocket: failed to collect guardrail " - "callbacks — PII masking will be skipped. Error: %s", - _guardrail_exc, - ) - streaming: Final = ResponsesWebSocketStreaming( websocket=websocket, backend_ws=cast(ClientConnection, backend_ws), @@ -7023,6 +7026,7 @@ class BaseLLMHTTPHandler: authorized_model=model, custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, + frame_guard=frame_guard, ) return await streaming.bidirectional_forward() diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 207e966ef28..aebcab75601 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,6 +1,7 @@ import asyncio import contextlib import json +import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum @@ -28,10 +29,11 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError from starlette.websockets import WebSocket, WebSocketDisconnect from typing_extensions import ReadOnly, TypedDict +import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.tool_fields import first_validation_error @@ -58,11 +60,17 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, safe_set_request_parsed_body, ) +from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id +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 +) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -1810,6 +1818,143 @@ async def _enforce_responses_ws_first_frame_model_auth( ) +_WS_FRAME: Final = TypeAdapter(dict[str, object]) +_GUARDRAIL_INFORMATION: Final = TypeAdapter(list[dict[str, object]]) +_PII_TOKENS: Final = TypeAdapter(dict[str, str]) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" +_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") +_CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") +# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it +_CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") + + +class _PreCallGuardrailWalk(Protocol): + async def pre_call_hook( + self, + *, + user_api_key_dict: UserAPIKeyAuth, + data: dict[str, object], # mutable-ok: ProxyLogging.pre_call_hook hands the guardrails a dict they edit + call_type: Literal["aresponses"], + guardrails_only: bool, + ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns the guardrails' own dict + + +def _responses_ws_frame_guard( + admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: _PreCallGuardrailWalk +) -> ResponsesWsFrameGuard: + """Socket admission skips guardrails (the first frame's ``input`` is not part of the admitted request), so every + ``response.create`` frame gets the pre_call guardrail walk the HTTP ``/v1/responses`` path runs on its body.""" + metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted) + admitted_bucket: Final = admitted.get(metadata_key) + bucket: Final = { + key: value + for key, value in ( + _WS_FRAME.validate_python(admitted_bucket) if isinstance(admitted_bucket, dict) else {} + ).items() + if key != _GUARDRAIL_INFORMATION_KEY + } + model: Final = admitted.get("model") + + # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation + rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) + + async def guard(frame: str) -> GuardedResponsesWsFrame: + frame_obj: Final = _WS_FRAME.validate_json(frame) + nested: Final = frame_obj.get("response") + if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS): + raise ValueError( + "response.create carries input, instructions or tools both at the top level and under response" + ) + body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj + frame_data: Final = { + **{ + key: _without_untrusted_metadata(key, value) + for key, value in body.items() + if key != "type" and key not in _CLIENT_GUARDRAIL_SELECTORS + }, + "model": model, + "litellm_call_id": str(uuid.uuid4()), + # shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied + metadata_key: {**bucket}, + } + try: + guarded_data: Final = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True + ) + except Exception as e: + await _bill_rejected_frame(rejected_frame_auth, e, frame_data) + raise + if guarded_data.get("model") != model: + rerouted: Final = ValueError( + f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket" + f" is bound to {model!r}, send this turn over HTTP /v1/responses" + ) + await _bill_rejected_frame(rejected_frame_auth, rerouted, frame_data) + raise rerouted + # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them + guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data} + kept: Final = {key: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS} + rewritten: Final = ( + {**frame_obj, "response": {**kept, **guarded_fields}} + if isinstance(nested, dict) + else {**kept, **guarded_fields} + ) + return GuardedResponsesWsFrame( + frame=json.dumps(rewritten), + guardrail_information=_recorded_guardrail_information(guarded_data.get(metadata_key)), + pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + ) + + return guard + + +def _without_untrusted_metadata(key: str, value: object) -> object: + """Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge.""" + if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict): + return value + return { + name: item + for name, item in _WS_FRAME.validate_python(value).items() + if not name.startswith("user_api_key_") + and name not in _UNTRUSTED_METADATA_CONTROL_FIELDS + and name not in _CLIENT_PRICING_METADATA_FIELDS + and name not in ("headers", "pii_tokens", *_CLIENT_GUARDRAIL_SELECTORS) + } + + +async def _bill_rejected_frame( + user_api_key_dict: UserAPIKeyAuth, + exception: Exception, + frame_data: dict[str, object], # mutable-ok: the spend logger stamps the failure row on the request dict +) -> None: + """Only the spend logger sees a rejected frame; the other failure hooks (the rate limiter's slot release, + alerting) belong to the live connection, which stays open. Billing never replaces the guardrail's own error.""" + for callback in litellm.callbacks: + if isinstance(callback, ProxyDBLogger): + try: + await callback.async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict + request_data=frame_data, original_exception=exception, user_api_key_dict=user_api_key_dict + ) + except Exception: # noqa: BLE001 # billing never replaces the guardrail's own error + verbose_proxy_logger.exception("[Non-Blocking] Responses WebSocket: billing a rejected frame failed") + + +def _recorded_pii_tokens(metadata: object) -> Mapping[str, str]: + """Presidio stores the tokens it masked a frame with in ``metadata.pii_tokens``; the socket unmasks with them.""" + recorded: Final = _WS_FRAME.validate_python(metadata).get("pii_tokens") if isinstance(metadata, dict) else None + return MappingProxyType(_PII_TOKENS.validate_python(recorded)) if isinstance(recorded, dict) else _NO_PII_TOKENS + + +_NO_PII_TOKENS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def _recorded_guardrail_information(bucket: object) -> tuple[Mapping[str, object], ...]: + recorded: Final = ( + _WS_FRAME.validate_python(bucket).get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None + ) + return tuple(_GUARDRAIL_INFORMATION.validate_python(recorded)) if recorded else () + + async def _responses_websocket_session( websocket: WebSocket, model: str | None, @@ -1885,6 +2030,7 @@ async def _responses_websocket_session( user_api_base=user_api_base, model=resolved_model, route_type="_aresponses_websocket", + skip_guardrails=True, ) except Exception as e: verbose_proxy_logger.exception("Responses WebSocket pre-call error") @@ -1906,7 +2052,10 @@ async def _responses_websocket_session( return routed_data: Final = dict( - data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message) + data, + user_api_key_dict=user_api_key_dict, + responses_ws_frame_guard=_responses_ws_frame_guard(data, user_api_key_dict, proxy_logging_obj), + **_routing_hints_from_first_ws_frame(first_message), ) # Phase 2: route to upstream provider try: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index b7385c3b3d5..d32d48c51ca 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -62,6 +62,7 @@ from litellm.types.llms.openai import ( # Handle ResponseText import with fallback if TYPE_CHECKING: + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard from litellm.types.llms.openai import ResponseText else: ResponseText = str # Fallback for ResponseText import @@ -2598,6 +2599,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs, ) -> Exception | None: """ @@ -2690,5 +2692,6 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + frame_guard=responses_ws_frame_guard, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 5897b0ac954..145189616af 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -7,12 +7,23 @@ import time import traceback import uuid from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime from functools import lru_cache from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, cast, overload, runtime_checkable +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NoReturn, + Protocol, + TypeAlias, + cast, + overload, + runtime_checkable, +) import httpx from openai._streaming import SSEDecoder @@ -1932,6 +1943,73 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) _PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +_PII_DELTA_EVENT_TYPES: Final = frozenset( + { + "response.output_text.delta", + "response.reasoning_summary_text.delta", + "response.refusal.delta", + "response.function_call_arguments.delta", + } +) + + +def _unmask_responses_event(response_str: str, pii_tokens: Mapping[str, str], unmask_pii_text: _UnmasksPiiText) -> str: + """Put the original PII back into one client-bound event: ``response.completed`` output text, or a delta.""" + try: + evt_obj: Final = _load_json_object(response_str) + except (json.JSONDecodeError, TypeError): + return response_str + event_type: Final = evt_obj.get("type") + + if event_type == "response.completed": + modified = False + response_obj: Final = evt_obj.get("response") + if not _is_json_object(response_obj): + return response_str + output_items: Final = response_obj.get("output") + for output_item in output_items if _is_json_array(output_items) else []: + if not _is_json_object(output_item): + continue + content = output_item.get("content") + if not _is_json_array(content): + continue + for content_block in content: + if not _is_json_object(content_block): + continue + text = content_block.get("text") + if isinstance(text, str): + unmasked = unmask_pii_text(text, pii_tokens) + if unmasked != text: + content_block["text"] = unmasked + modified = True + return json.dumps(evt_obj) if modified else response_str + + if event_type in _PII_DELTA_EVENT_TYPES: + delta: Final = evt_obj.get("delta") + if isinstance(delta, str): + unmasked = unmask_pii_text(delta, pii_tokens) + if unmasked != delta: + evt_obj["delta"] = unmasked + return json.dumps(evt_obj) + + return response_str + + +@dataclass(frozen=True, slots=True) +class GuardedResponsesWsFrame: + """A ``response.create`` frame after the proxy's pre_call guardrails ran on it: the frame to forward (guardrails + may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, and the + Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply.""" + + frame: str + guardrail_information: tuple[Mapping[str, object], ...] = () + pii_tokens: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + + +ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] _RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES @@ -1964,6 +2042,30 @@ def _user_texts_in_input_item(item: object) -> tuple[str, ...]: ) +def _set_turn_guardrail_information( + turn_logging_obj: LiteLLMLoggingObj, guardrail_information: tuple[Mapping[str, object], ...] +) -> None: + """A turn bills the guardrails that ran on its own ``response.create`` frame, never the records the + connection-level logging object carries from an earlier frame.""" + litellm_params: Final = _json_object_or_none( + turn_logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if litellm_params is None: + return + for key in _PER_TURN_METADATA_KEYS: + _replace_guardrail_information(litellm_params.get(key), guardrail_information) + + +def _replace_guardrail_information(bucket: object, guardrail_information: tuple[Mapping[str, object], ...]) -> None: + if not _is_json_object(bucket): + return + bucket.pop(_GUARDRAIL_INFORMATION_KEY, None) + if guardrail_information: + bucket[_GUARDRAIL_INFORMATION_KEY] = [ # rebind-ok: the turn copy's own metadata bucket + dict(entry) for entry in guardrail_information + ] + + def _logged_user_messages(client_frame: object) -> tuple[Mapping[str, object], ...]: frame: Final = _json_object_or_none(client_frame) if frame is None or frame.get("type") != "response.create": @@ -2092,11 +2194,13 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + frame_guard: ResponsesWsFrameGuard | None = None, ): self.websocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict + self.frame_guard: ResponsesWsFrameGuard | None = frame_guard self.request_data: dict[str, object] = request_data or {} litellm_metadata: Final = self.request_data.get("litellm_metadata") self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {} @@ -2112,12 +2216,15 @@ class ResponsesWebSocketStreaming: self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults - def _start_turn(self, client_frame: str | None) -> None: + def _start_turn( + self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = () + ) -> None: start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time turn_logging_obj: Final = self.logging_obj.copy_for_new_call( litellm_call_id=str(uuid.uuid4()), start_time=start_time ) self._detach_budget_reservation() + _set_turn_guardrail_information(turn_logging_obj, guardrail_information) turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) turn_logging_obj.pre_call(input=client_frame, api_key="") self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) @@ -2136,10 +2243,46 @@ class ResponsesWebSocketStreaming: for key, value in litellm_params.items() } - def _record_client_frame(self, client_frame: str) -> None: + def _record_client_frame( + self, client_frame: str, guardrail_information: tuple[Mapping[str, object], ...] = () + ) -> None: frame: Final = _json_object_or_none(client_frame) if frame is not None and frame.get("type") == "response.create": - self._start_turn(client_frame) + self._start_turn(client_frame, guardrail_information) + + async def _guard_frame(self, message: str) -> GuardedResponsesWsFrame | None: + """Run the proxy's pre_call guardrails on one ``response.create`` frame. + + Returns None when a guardrail rejects the frame: the client gets an ``error`` event and the + frame is dropped while the connection stays open for the next turn. + """ + frame: Final = _json_object_or_none(message) + if self.frame_guard is None or frame is None or frame.get("type") != "response.create": + return GuardedResponsesWsFrame(frame=message) + try: + return await self.frame_guard(message) + except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event + verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e) + await self._send_error_event("invalid_request_error", str(e)) + return None + + async def _forward_guarded_frame(self, message: str) -> None: + guarded: Final = await self._guard_frame(message) + if guarded is None: + return + if guarded.pii_tokens: + self._remember_pii_tokens(guarded.pii_tokens) + masked: Final = await self._mask_response_create(guarded.frame) + self._record_client_frame(masked, guarded.guardrail_information) + await self.backend_ws.send(masked) + + def _remember_pii_tokens(self, pii_tokens: Mapping[str, str]) -> None: + metadata: Final = self.request_data.get("metadata") + known: Final = metadata.get("pii_tokens") if _is_json_object(metadata) else None + self.request_data["metadata"] = { + **(metadata if _is_json_object(metadata) else {}), + "pii_tokens": {**(known if _is_str_mapping(known) else {}), **pii_tokens}, + } def _record_backend_event(self, event: str) -> None: event_obj: Final = _json_object_or_none(event) @@ -2383,14 +2526,7 @@ class ResponsesWebSocketStreaming: return json.dumps(msg_obj) if modified else message # Delta event types whose ``delta`` field may contain PII tokens. - _DELTA_EVENT_TYPES = frozenset( - { - "response.output_text.delta", - "response.reasoning_summary_text.delta", - "response.refusal.delta", - "response.function_call_arguments.delta", - } - ) + _DELTA_EVENT_TYPES = _PII_DELTA_EVENT_TYPES # Terminal events that carry the full output text or tool-call arguments # already delivered by ``response.completed``. Suppressed when output masking @@ -2430,47 +2566,9 @@ class ResponsesWebSocketStreaming: if not pii_tokens: return response_str - try: - evt_obj: Final = _load_json_object(response_str) - except (json.JSONDecodeError, TypeError): - return response_str - cb: Final = self.guardrail_callbacks[0] unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR) - event_type: Final = evt_obj.get("type") - - if event_type == "response.completed": - modified = False - response_obj: Final = evt_obj.get("response") - if not _is_json_object(response_obj): - return response_str - output_items: Final = response_obj.get("output") - for output_item in output_items if _is_json_array(output_items) else []: - if not _is_json_object(output_item): - continue - content = output_item.get("content") - if not _is_json_array(content): - continue - for content_block in content: - if not _is_json_object(content_block): - continue - text = content_block.get("text") - if isinstance(text, str): - unmasked = unmask_pii_text(text, pii_tokens) - if unmasked != text: - content_block["text"] = unmasked - modified = True - return json.dumps(evt_obj) if modified else response_str - - if event_type in self._DELTA_EVENT_TYPES: - delta: Final = evt_obj.get("delta") - if isinstance(delta, str): - unmasked = unmask_pii_text(delta, pii_tokens) - if unmasked != delta: - evt_obj["delta"] = unmasked - return json.dumps(evt_obj) - - return response_str + return _unmask_responses_event(response_str, pii_tokens, unmask_pii_text) async def _mask_response_completed(self, response_str: str) -> str: """ @@ -2564,38 +2662,29 @@ class ResponsesWebSocketStreaming: self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message ) except RateLimitError as e: - try: - await self.websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "rate_limit_exceeded", - "message": str(e), - }, - } - ) - ) - except Exception: # noqa: BLE001, S110 # client may already be gone - pass + await self._send_error_event("rate_limit_exceeded", str(e)) return False return True + async def _send_error_event(self, error_type: str, message: str) -> None: + try: + await self.websocket.send_text( + json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) + ) + except Exception: # noqa: BLE001, S110 # client may already be gone + pass + async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): - masked_first: Final = await self._mask_response_create(self.first_message) - self._record_client_frame(masked_first) - await self.backend_ws.send(masked_first) + await self._forward_guarded_frame(self.first_message) while True: message = await self.websocket.receive_text() if not await self._enforce_or_reject_frame(message): continue - masked = await self._mask_response_create(message) - self._record_client_frame(masked) - await self.backend_ws.send(masked) + await self._forward_guarded_frame(message) except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) @@ -2672,10 +2761,15 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, + frame_guard: ResponsesWsFrameGuard | None = None, + guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, **kwargs: object, ) -> None: self.websocket = websocket self.model = model + self.frame_guard = frame_guard + self.guardrail_callbacks: tuple[PresidioGuardrailCallback, ...] = tuple(guardrail_callbacks or ()) + self._pii_tokens: Mapping[str, str] = {} self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.litellm_metadata: Mapping[str, object] = litellm_metadata or {} @@ -2983,6 +3077,27 @@ class ManagedResponsesWebSocketHandler: if self.litellm_metadata: call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) + async def _guard_frame(self, raw_message: str) -> GuardedResponsesWsFrame | None: + """Same per-frame pre_call guardrail walk as the native socket: a rejected frame errors and is not sent.""" + if self.frame_guard is None: + return GuardedResponsesWsFrame(frame=raw_message) + try: + return await self.frame_guard(raw_message) + except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event + verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e) + await self._send_error(str(e), error_type="invalid_request_error") + return None + + @staticmethod + def _metadata_with_turn_guardrails( + existing: object, guardrail_information: tuple[Mapping[str, object], ...] + ) -> Mapping[str, object]: + """Each turn is its own ``litellm.aresponses`` call, so its guardrail records ride on that call's metadata.""" + return { + **(existing if _is_json_object(existing) else {}), + _GUARDRAIL_INFORMATION_KEY: [dict(record) for record in guardrail_information], + } + @staticmethod def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None: """Update proxy_server_request body so spend logs record the full request.""" @@ -3026,12 +3141,19 @@ class ManagedResponsesWebSocketHandler: except Exception: pass try: - await self.websocket.send_text(serialized) + await self.websocket.send_text(self._unmask_for_client(serialized)) except Exception as send_exc: verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc) return terminal_event # Client disconnected return terminal_event + def _unmask_for_client(self, serialized: str) -> str: + """The provider and the turn history keep Presidio's tokens, the client gets the original text back.""" + if not self._pii_tokens or not self.guardrail_callbacks: + return serialized + unmask_pii_text: Final[_UnmasksPiiText] = getattr(self.guardrail_callbacks[0], _UNMASK_PII_TEXT_ATTR) + return _unmask_responses_event(serialized, self._pii_tokens, unmask_pii_text) + def _save_turn_history( self, terminal_event: _MutableJsonObject | None, @@ -3101,7 +3223,13 @@ class ManagedResponsesWebSocketHandler: await self._send_error(str(e), error_type="rate_limit_exceeded") return - call_kwargs: Final = self._build_base_call_kwargs(msg_obj) + guarded: Final = await self._guard_frame(raw_message) + if guarded is None: + return + if guarded.pii_tokens: + self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} + + call_kwargs: Final = self._build_base_call_kwargs(_load_json_object(guarded.frame)) call_kwargs["stream"] = True # A frame that repeats the connection's public alias (model_group) must @@ -3121,6 +3249,10 @@ class ManagedResponsesWebSocketHandler: self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history) self._inject_credentials(call_kwargs, model=model) + if guarded.guardrail_information: + call_kwargs["litellm_metadata"] = self._metadata_with_turn_guardrails( + call_kwargs.get("litellm_metadata"), guarded.guardrail_information + ) self._update_proxy_request(call_kwargs, requested_model or self.model_group or model) call_kwargs.update(self.extra_kwargs) diff --git a/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py new file mode 100644 index 00000000000..8af70d39186 --- /dev/null +++ b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py @@ -0,0 +1,168 @@ +"""A native Responses WebSocket runs the proxy's pre_call guardrails on every ``response.create`` frame, bills each +turn's spend row with the guardrail cost of that turn alone, and bills a blocked frame's guardrail scan as a failed +turn (LIT-9744)""" + +import asyncio +import hashlib +import json +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import websockets +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, create, responses_peer, turn +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_INPUT_COST_PER_TOKEN: Final = 0.002 +_OUTPUT_COST_PER_TOKEN: Final = 0.004 +_COST_PER_WORD: Final = 0.1 +_MARKER: Final = " [guarded]" +_BLOCKED_WORD: Final = "forbidden" +_GUARDRAIL_MODULE: Final = f""" +from litellm.integrations.custom_guardrail import CustomGuardrail + + +class TurnGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + blocks = [block for item in data.get("input") or [] for block in item.get("content") or []] + words = sum(len(block["text"].split()) for block in blocks) + blocked = any({_BLOCKED_WORD!r} in block["text"] for block in blocks) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={{"words": words}}, + request_data=data, + guardrail_status="guardrail_intervened" if blocked else "success", + tracing_detail={{"guardrail_cost": round({_COST_PER_WORD} * words, 10)}}, + ) + if blocked: + raise ValueError("turn guardrail blocked the request") + for block in blocks: + block["text"] = block["text"] + {_MARKER!r} + return data +""" + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + cert: tuple[Path, Path] + + +def _config(directory: Path) -> Path: + (directory / "turn_guardrail.py").write_text(_GUARDRAIL_MODULE) + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["guardrails"] = [ + { + "guardrail_name": "turn-guardrail", + "litellm_params": {"guardrail": "turn_guardrail.TurnGuardrail", "mode": "pre_call", "default_on": True}, + } + ] + path: Final = directory / "responses-ws-guardrail.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("responses-ws-guardrail") + cert: Final = write_self_signed_cert(directory) + with ( + gateway_from_environment() as shared, + owned_proxy(shared, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_config(directory)) as owned, + ): + yield _Rig(owned, cert) + + +def _nested(frame: str) -> str: + """The ``{"type": "response.create", "response": {...}}`` wire shape the WS layer accepts next to the flat one.""" + flat: Final = json.loads(frame) + return json.dumps({"type": flat["type"], "response": {key: value for key, value in flat.items() if key != "type"}}) + + +async def _session( + proxy_url: str, key: str, model: str, frames: tuple[str, ...] +) -> tuple[tuple[dict[str, JsonValue], ...], ...]: + proxy: Final = proxy_url.rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/responses?model={model}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=10, + ) as connection: + return tuple([await turn(connection, frame) for frame in frames]) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT spend, status FROM "LiteLLM_SpendLogs" WHERE api_key = %s ORDER BY "startTime"', + (hashlib.sha256(key.encode()).hexdigest(),), + ) + + +def _key_spend(key: str) -> float: + (row,) = read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashlib.sha256(key.encode()).hexdigest(),) + ) + return float(row["spend"]) + + +def _llm_cost() -> float: + return USAGE["input_tokens"] * _INPUT_COST_PER_TOKEN + USAGE["output_tokens"] * _OUTPUT_COST_PER_TOKEN + + +def test_every_turn_runs_the_pre_call_guardrail_and_bills_only_its_own_guardrail_cost(rig: _Rig) -> None: + with responses_peer(rig.cert) as peer, rig.gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + input_cost_per_token=_INPUT_COST_PER_TOKEN, + output_cost_per_token=_OUTPUT_COST_PER_TOKEN, + ) + key: Final = scenario.key(models=[model]) + tag: Final = uuid.uuid4().hex + texts: Final = ( + f"one {tag}", + f"one two {tag}", + f"one two three {tag}", + f"{_BLOCKED_WORD} {tag}", + f"one {tag}", + f"{_BLOCKED_WORD} again {tag}", + ) + frames: Final = tuple(create(model, text) for text in texts[:-1]) + (_nested(create(model, texts[-1])),) + turns: Final = asyncio.run(_session(str(rig.gateway.client.base_url), key, model, frames)) + + assert [events[-1]["type"] for events in turns] == [ + "response.completed", + "response.completed", + "response.completed", + "error", + "response.completed", + "error", + ], [[event["type"] for event in events] for events in turns] + assert "blocked" in str(turns[3][-1]["error"]), turns[3] + + forwarded: Final = tuple(peer.frames.get_nowait() for _ in range(peer.frames.qsize())) + assert [frame["input"][0]["content"][0]["text"] for frame in forwarded] == [ + f"{text}{_MARKER}" for text in texts if _BLOCKED_WORD not in text + ], forwarded + + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 6, seconds=70) + assert [row["status"] for row in rows] == ["success", "success", "success", "failure", "success", "failure"], ( + rows + ) + expected: Final = [ + (0.0 if _BLOCKED_WORD in text else _llm_cost()) + _COST_PER_WORD * len(text.split()) for text in texts + ] + assert [float(row["spend"]) for row in rows] == pytest.approx(expected), rows + total: Final = sum(expected) + assert eventually( + lambda: _key_spend(key), lambda spend: abs(spend - total) < 1e-6, seconds=70 + ) == pytest.approx(total) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index a7b955f5b84..52e95901a16 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -480,7 +480,12 @@ def _advertised_cost_split(headers: Mapping[str, str]) -> list[tuple[float, floa @pytest.mark.asyncio @pytest.mark.parametrize( "background, stored, status, expected_billed_events", - [(True, True, "completed", 1), (True, True, "in_progress", 0), (True, False, "completed", 0), (False, True, "completed", 0)], + [ + (True, True, "completed", 1), + (True, True, "in_progress", 0), + (True, False, "completed", 0), + (False, True, "completed", 0), + ], ) async def test_reading_a_finished_background_response_bills_it( monkeypatch: pytest.MonkeyPatch, background: bool, stored: bool, status: str, expected_billed_events: int @@ -531,8 +536,12 @@ async def test_reading_a_finished_background_response_bills_it( prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")), ): await _bill_finished_background_response( @@ -662,9 +671,15 @@ async def _read_twice_on_two_pods( router.aget_responses = AsyncMock(side_effect=AssertionError("billing must not read the provider again")) with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")), ): await asyncio.gather( @@ -702,7 +717,9 @@ async def test_a_finished_response_the_create_could_not_store_is_billed_once_und monkeypatch, table, response_id, - UserAPIKeyAuth(api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42"), + UserAPIKeyAuth( + api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42" + ), ) (billed,) = payloads @@ -746,7 +763,9 @@ async def test_a_finished_read_never_bills_a_response_already_billed_or_billed_b ) table: Final = _UniqueManagedObjectTable(stored) - payloads: Final = await _read_twice_on_two_pods(monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader")) + payloads: Final = await _read_twice_on_two_pods( + monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader") + ) assert payloads == () assert table.rows == stored @@ -878,9 +897,15 @@ async def test_a_finished_read_or_cancel_bills_usage_only_the_configured_deploym "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("the configured deployment answers billing reads")), ): await getattr(endpoints, endpoint_name)( @@ -941,9 +966,15 @@ async def test_deleting_a_finished_unbilled_background_response_bills_it_before_ "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): result: Final = await delete_response( response_id=stored.response_id, @@ -989,7 +1020,9 @@ async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(ca caller: Final = UserAPIKeyAuth(api_key="sk-caller", user_role=caller_role) async def pipeline(**_: object) -> DeleteResponseResult: - await security.async_pre_call_hook(caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses") + await security.async_pre_call_hook( + caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses" + ) return DeleteResponseResult(id=stored.response_id, object="response", deleted=True) processor: Final = MagicMock() @@ -1002,9 +1035,15 @@ async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(ca "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): with contextlib.suppress(HTTPException): await delete_response( @@ -1058,9 +1097,15 @@ def _capture_spend_events(monkeypatch: pytest.MonkeyPatch) -> _SpendEvents: @contextlib.contextmanager def _proxy_over(prisma_client: MagicMock, proxy_logging: MagicMock, router: MagicMock): with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): yield @@ -1211,7 +1256,9 @@ async def test_a_streamed_background_create_is_billed_wherever_it_sets_backgroun "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", return_value=False, ), - patch("litellm.proxy.proxy_server.select_data_generator", provider_events), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.select_data_generator", provider_events + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): await responses_api( request=_request_for("POST", "/v1/responses", body), @@ -1329,7 +1376,12 @@ async def test_background_create_the_provider_confirms_is_stored_under_its_provi } request = Request( - {"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"content-type", b"application/json")]}, + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"content-type", b"application/json")], + }, receive, ) @@ -1918,9 +1970,327 @@ class TestResponsesWSFirstFrameModelAuth: assert routed["input"] == replayed_input assert routed["previous_response_id"] == "resp_prev" assert processor.common_processing_pre_call_logic.await_args.kwargs["model"] == "gpt-4o-mini" + assert processor.common_processing_pre_call_logic.await_args.kwargs["skip_guardrails"] is True + assert callable(routed["responses_ws_frame_guard"]) assert mock_route_request.await_args.kwargs["route_type"] == "_aresponses_websocket" ws.close.assert_not_awaited() + @pytest.mark.asyncio + async def test_frame_guard_runs_the_guardrail_walk_on_the_frame_body_and_returns_its_records(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + assert call_type == "aresponses" and guardrails_only is True + assert "type" not in data and data["model"] == "gpt-4o-mini" + assert data["litellm_metadata"]["user_api_key_hash"] == "hash" + assert "standard_logging_guardrail_information" not in data["litellm_metadata"] + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_name": "turn"}] + return {**data, "input": data["input"] + " [guarded]"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + admitted = { + "model": "gpt-4o-mini", + "litellm_metadata": { + "user_api_key_hash": "hash", + "standard_logging_guardrail_information": [{"guardrail_name": "admission"}], + }, + } + guard = _responses_ws_frame_guard(admitted, MagicMock(), proxy_logging_obj) + + first = await guard(json.dumps({"type": "response.create", "model": "alias", "input": "hi"})) + second = await guard(json.dumps({"type": "response.create", "input": "again"})) + + assert json.loads(first.frame) == {"type": "response.create", "model": "alias", "input": "hi [guarded]"} + assert first.guardrail_information == ({"guardrail_name": "turn"},) + assert json.loads(second.frame)["input"] == "again [guarded]" + assert admitted["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_name": "admission"} + ] + + @pytest.mark.asyncio + async def test_frame_guard_bills_a_rejected_frame_as_a_failed_request(self): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + raise ValueError("blocked") + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + proxy_logging_obj.post_call_failure_hook = AsyncMock() + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + rate_limiter = MagicMock() + rate_limiter.async_post_call_failure_hook = AsyncMock() + user_api_key_dict = MagicMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, user_api_key_dict, proxy_logging_obj + ) + + with ( + patch.object(litellm, "callbacks", [rate_limiter, spend_logger]), + pytest.raises(ValueError, match="blocked"), + ): + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + rate_limiter.async_post_call_failure_hook.assert_not_awaited() + spend_logger.async_post_call_failure_hook.assert_awaited_once() + failure = spend_logger.async_post_call_failure_hook.await_args.kwargs + user_api_key_dict.model_copy.assert_called_once_with(update={"budget_reservation": None}) + assert failure["user_api_key_dict"] is user_api_key_dict.model_copy.return_value + assert str(failure["original_exception"]) == "blocked" + assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_cost": 0.5} + ] + assert failure["request_data"]["input"] == "hi" + assert len(failure["request_data"]["litellm_call_id"]) == 36 + + @pytest.mark.asyncio + async def test_frame_guard_guards_the_nested_response_body_and_copies_back_tools(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return {**data, "input": data["input"] + " [guarded]", "tools": data["tools"][:1]} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + tools = [{"type": "function", "name": "a"}, {"type": "function", "name": "b"}] + + guarded = await guard(json.dumps({"type": "response.create", "response": {"input": "hi", "tools": tools}})) + + assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["input"] == "hi" + assert json.loads(guarded.frame) == { + "type": "response.create", + "response": {"input": "hi [guarded]", "tools": tools[:1]}, + } + + @pytest.mark.asyncio + async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + metadata = { + "user_api_key_metadata": {"disable_global_guardrails": True}, + "disable_global_guardrails": True, + "headers": {"x-litellm-tags": "forged"}, + "tag": "kept", + } + + await guard(json.dumps({"type": "response.create", "input": "hi", "metadata": metadata})) + + assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["metadata"] == {"tag": "kept"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "input": "forbidden", "guardrails": []}, + {"type": "response.create", "input": "forbidden", "metadata": {"guardrails": []}}, + {"type": "response.create", "response": {"input": "forbidden", "guardrails": [], "policies": []}}, + ], + ) + async def test_frame_guard_keeps_the_admitted_guardrail_list_when_a_frame_tries_to_replace_it(self, frame): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + MagicMock(), + proxy_logging_obj, + ) + + await guard(json.dumps(frame)) + + scanned = proxy_logging_obj.pre_call_hook.await_args.kwargs["data"] + assert "guardrails" not in scanned and "policies" not in scanned + assert "guardrails" not in (scanned.get("metadata") or {}) + assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + + @pytest.mark.asyncio + async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self): + import threading + + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + span = threading.RLock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"litellm_parent_otel_span": span}}, + MagicMock(), + proxy_logging_obj, + ) + + guarded = await guard(json.dumps({"type": "response.create", "input": "hi"})) + + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + assert ( + proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["litellm_metadata"]["litellm_parent_otel_span"] + is span + ) + + @pytest.mark.asyncio + async def test_frame_guard_rejects_a_turn_a_guardrail_rerouted_to_another_model(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + return {**data, "model": "onprem-llama"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with ( + patch.object(litellm, "callbacks", [spend_logger]), + pytest.raises(ValueError, match="rerouted this turn to 'onprem-llama'"), + ): + await guard(json.dumps({"type": "response.create", "input": "my SSN is 123-45-6789"})) + + failure = spend_logger.async_post_call_failure_hook.await_args.kwargs + assert "rerouted this turn to 'onprem-llama'" in str(failure["original_exception"]) + assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_cost": 0.5} + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_frame_guard_drops_a_client_s_guardrail_cost_records_so_a_blocked_frame_bills_the_real_scan( + self, nested: bool + ): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + raise ValueError("blocked") + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + body = { + "input": "forbidden", + "metadata": {"trace": "t1", "standard_logging_guardrail_information": [], "model_info": {"x": 1}}, + } + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with patch.object(litellm, "callbacks", [spend_logger]), pytest.raises(ValueError, match="blocked"): + await guard(json.dumps(frame)) + + billed = spend_logger.async_post_call_failure_hook.await_args.kwargs["request_data"] + assert billed["metadata"] == {"trace": "t1"} + assert billed["litellm_metadata"]["standard_logging_guardrail_information"] == [{"guardrail_cost": 0.5}] + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_frame_guard_forwards_a_frame_without_the_instructions_a_guardrail_deleted(self, nested: bool): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return {key: value for key, value in data.items() if key != "instructions"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + body = {"instructions": "SECRET system prompt", "input": "hi", "store": False} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with patch.object(litellm, "callbacks", []): + guarded = await guard(json.dumps(frame)) + + forwarded = json.loads(guarded.frame) + assert forwarded["type"] == "response.create" + sent_body = forwarded["response"] if nested else {k: v for k, v in forwarded.items() if k != "type"} + assert sent_body == {"input": "hi", "store": False} + + @pytest.mark.asyncio + async def test_frame_guard_hands_back_the_pii_tokens_presidio_masked_with_and_drops_forged_ones(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + assert data["metadata"] == {"trace": "t1"}, "a client cannot seed pii_tokens" + data["metadata"]["pii_tokens"] = {"": "alice@example.com"} + return {**data, "input": "email please"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "input": "email alice@example.com please", + "metadata": {"trace": "t1", "pii_tokens": {"": "mallory@example.com"}}, + } + ) + ) + + assert json.loads(guarded.frame)["input"] == "email please" + assert dict(guarded.pii_tokens) == {"": "alice@example.com"} + + @pytest.mark.asyncio + async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): + await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + + proxy_logging_obj.pre_call_hook.assert_not_awaited() + + @pytest.mark.asyncio + async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=ValueError("blocked")) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock(side_effect=RuntimeError("db down")) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", [spend_logger]), pytest.raises(ValueError, match="blocked"): + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + spend_logger.async_post_call_failure_hook.assert_awaited_once() + @pytest.mark.asyncio @pytest.mark.parametrize("provider_rejected", [True, False]) async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool): @@ -3614,16 +3984,18 @@ class TestResponsesInputTokens: model_list=[ { "model_name": "gpt-5.2", - "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test", "additional_drop_params": ["tools"]}, + "litellm_params": { + "model": "openai/gpt-5.2", + "api_key": "sk-test", + "additional_drop_params": ["tools"], + }, }, {"model_name": "backup", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test"}}, ], fallbacks=[{"gpt-5.2": ["backup"]}], ) - response, counter = self._post_input_tokens( - {"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router - ) + response, counter = self._post_input_tokens({"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router) assert response.status_code == 200, response.text assert counter.await_args.kwargs["request"].tools is None diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 2b78aa98b04..b38808fb57c 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -285,6 +285,151 @@ class TestManagedWebSocketHandlerIntegration: assert captured["model"] == "bedrock_mantle/openai.gpt-5.5" + @pytest.mark.asyncio + async def test_managed_socket_runs_the_frame_guard_on_every_turn(self, monkeypatch): + import json + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ManagedResponsesWebSocketHandler + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + if "forbidden" in body["input"]: + raise ValueError("blocked by guardrail") + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": body["input"].upper()}), + guardrail_information=({"guardrail_name": "g", "guardrail_cost": 0.5},), + ) + + async def empty_stream(): + return + yield {} + + aresponses: Final = AsyncMock(side_effect=lambda **kwargs: empty_stream()) + monkeypatch.setattr(litellm, "aresponses", aresponses) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-sonnet-4-5", + logging_obj=Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + litellm_metadata={"user_api_key_hash": "k1"}, + frame_guard=guard, + ) + + await handler._process_response_create(json.dumps({"type": "response.create", "input": "say forbidden"})) + await handler._process_response_create(json.dumps({"type": "response.create", "input": "hello"})) + + aresponses.assert_awaited_once() + error_event: Final = json.loads(websocket.send_text.await_args_list[0].args[0]) + assert error_event["error"] == {"type": "invalid_request_error", "message": "blocked by guardrail"} + sent: Final = aresponses.await_args.kwargs + assert sent["input"] == "HELLO" + assert sent["litellm_metadata"]["user_api_key_hash"] == "k1" + assert sent["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_name": "g", "guardrail_cost": 0.5} + ] + + @pytest.mark.asyncio + async def test_managed_socket_unmasks_the_reply_with_the_frame_guard_s_pii_tokens(self, monkeypatch): + import json + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ManagedResponsesWebSocketHandler + + presidio = _FakeWSGuardrail() + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + masked = await presidio.check_pii(body["input"], True, None, walk := {}) + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": masked}), pii_tokens=walk["metadata"]["pii_tokens"] + ) + + provider_saw: dict = {} + + async def fake_aresponses(**kwargs): + provider_saw.update(kwargs) + + async def stream(): + yield {"type": "response.output_text.delta", "delta": "Sent to "} + yield { + "type": "response.completed", + "response": { + "id": "resp_1", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "Sent to "}]} + ], + }, + } + + return stream() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-sonnet-4-5", + logging_obj=Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + frame_guard=guard, + guardrail_callbacks=[presidio], + ) + + await handler._process_response_create( + json.dumps({"type": "response.create", "input": "email alice@example.com please"}) + ) + + assert provider_saw["input"] == "email please" + sent: Final = [json.loads(call.args[0]) for call in websocket.send_text.await_args_list] + assert sent[0]["delta"] == "Sent to alice@example.com" + assert sent[1]["response"]["output"][0]["content"][0]["text"] == "Sent to alice@example.com" + + @pytest.mark.asyncio + async def test_http_handler_hands_the_frame_guard_to_a_managed_socket(self): + import json + from unittest.mock import AsyncMock, MagicMock + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + guard: Final = AsyncMock(side_effect=ValueError("blocked by guardrail")) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=ConnectionError("client left")) + + await BaseLLMHTTPHandler().async_responses_websocket( + model="anthropic/claude-sonnet-4-5", + websocket=websocket, + logging_obj=MagicMock(), + responses_api_provider_config=None, + first_message=json.dumps({"type": "response.create", "input": "say forbidden"}), + frame_guard=guard, + ) + + guard.assert_awaited_once() + error_event: Final = json.loads(websocket.send_text.await_args.args[0]) + assert error_event["error"]["message"] == "blocked by guardrail" + @pytest.mark.parametrize("terminal_type", ["response.completed", "response.incomplete"]) @pytest.mark.asyncio async def test_truncated_turn_is_kept_as_history_for_the_next_turn(self, monkeypatch, terminal_type: str): @@ -1319,6 +1464,97 @@ class TestWebSocketProjectQuotaEnforcement: mock_websocket.send_text.assert_called_once() assert "rate_limit_exceeded" in mock_websocket.send_text.call_args[0][0] + @pytest.mark.asyncio + async def test_native_handler_runs_the_frame_guard_on_every_turn_and_bills_each_turn_its_own_guardrail_records( + self, + ): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWebSocketStreaming + + def connection_logging(): + logging_obj = MagicMock() + logging_obj.litellm_params = { + "metadata": { + "standard_logging_guardrail_information": [{"guardrail_name": "admission", "guardrail_cost": 9.0}] + } + } + logging_obj.model_call_details = {} + logging_obj.copy_for_new_call.side_effect = lambda **_: connection_logging() + return logging_obj + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": body["input"] + " [guarded]"}), + guardrail_information=({"guardrail_name": "turn", "guardrail_cost": 0.1 * len(body["input"])},), + ) + + frames = [json.dumps({"type": "response.create", "input": text}) for text in ("hi", "hello")] + mock_backend_ws = MagicMock() + mock_backend_ws.send = AsyncMock() + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_websocket.receive_text = AsyncMock(side_effect=[frames[1], RuntimeError("client closed")]) + + handler = ResponsesWebSocketStreaming( + websocket=mock_websocket, + backend_ws=mock_backend_ws, + logging_obj=connection_logging(), + first_message=frames[0], + frame_guard=guard, + ) + await handler.client_to_backend() + + sent = [json.loads(call.args[0]) for call in mock_backend_ws.send.call_args_list] + assert [frame["input"] for frame in sent] == ["hi [guarded]", "hello [guarded]"] + assert [ + turn.logging_obj.litellm_params["metadata"]["standard_logging_guardrail_information"] + for turn in handler._open_turns + ] == [ + [{"guardrail_name": "turn", "guardrail_cost": pytest.approx(0.2)}], + [{"guardrail_name": "turn", "guardrail_cost": pytest.approx(0.5)}], + ] + + @pytest.mark.asyncio + async def test_native_handler_drops_a_frame_the_guard_rejects_and_keeps_the_socket_open(self): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWebSocketStreaming + + async def guard(frame: str) -> GuardedResponsesWsFrame: + if "forbidden" in frame: + raise ValueError("turn guardrail blocked the request") + return GuardedResponsesWsFrame(frame=frame) + + blocked = json.dumps({"type": "response.create", "input": "forbidden"}) + allowed = json.dumps({"type": "response.create", "input": "fine"}) + mock_backend_ws = MagicMock() + mock_backend_ws.send = AsyncMock() + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_websocket.receive_text = AsyncMock(side_effect=[allowed, RuntimeError("client closed")]) + logging_obj = MagicMock() + logging_obj.copy_for_new_call.return_value = logging_obj + + handler = ResponsesWebSocketStreaming( + websocket=mock_websocket, + backend_ws=mock_backend_ws, + logging_obj=logging_obj, + first_message=blocked, + frame_guard=guard, + ) + await handler.client_to_backend() + + error_event = json.loads(mock_websocket.send_text.call_args[0][0]) + assert error_event["type"] == "error" + assert error_event["error"] == { + "type": "invalid_request_error", + "message": "turn guardrail blocked the request", + } + assert [json.loads(call.args[0])["input"] for call in mock_backend_ws.send.call_args_list] == ["fine"] + assert len(handler._open_turns) == 1 + @pytest.mark.asyncio async def test_native_handler_forwards_frame_allowed_by_quota_callback(self): from unittest.mock import AsyncMock, MagicMock @@ -1905,6 +2141,34 @@ def _make_streaming(**kwargs): class TestNativeWebSocketGuardrailMasking: """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming.""" + @pytest.mark.asyncio + async def test_a_reply_to_a_frame_the_guard_masked_is_unmasked_with_the_guard_s_tokens(self): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame + + guardrail = _FakeWSGuardrail() + + async def guard(frame: str) -> GuardedResponsesWsFrame: + masked = await guardrail.check_pii(json.loads(frame)["input"], True, None, walk := {}) + return GuardedResponsesWsFrame( + frame=json.dumps({"type": "response.create", "input": masked}), + pii_tokens=walk["metadata"]["pii_tokens"], + ) + + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + handler = _make_streaming( + backend_ws=backend_ws, request_data={}, guardrail_callbacks=[guardrail], frame_guard=guard + ) + handler.logging_obj.copy_for_new_call.return_value = handler.logging_obj + + await handler._forward_guarded_frame(json.dumps({"type": "response.create", "input": "email alice@example.com"})) + + assert json.loads(backend_ws.send.await_args.args[0])["input"] == "email " + reply = json.dumps({"type": "response.output_text.delta", "delta": "Sent to "}) + assert json.loads(handler._unmask_response_event(reply))["delta"] == "Sent to alice@example.com" + @pytest.mark.asyncio async def test_mask_response_create_flat_string_input(self): guardrail = _FakeWSGuardrail() From 1436a8851afcd7752ed1489421f601a0b161eb18 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 11:02:20 -0700 Subject: [PATCH 19/64] perf(proxy): gc.freeze() the startup heap once the proxy is up so gen-2 passes stop re-walking it (#47137) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27bfd03d7f8..1e16eb626fe 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1,6 +1,7 @@ import asyncio import copy import enum +import gc import importlib import inspect import io @@ -1746,6 +1747,8 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState async with AsyncExitStack() as admin_mcp_stack: try: await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app)) + gc.freeze() + verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count()) yield state finally: if model_info_scheduler is not None and model_info_scheduler.running: From 836f32661119e9c288ff297374121db9444a0663 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 11:11:16 -0700 Subject: [PATCH 20/64] test(integration): remove BUG skips for 2 tests that pass on main (#47142) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/management/test_scim_admin_group.py | 4 ---- .../providers/test_rust_anthropic_safeguards_beta_wire.py | 1 - 2 files changed, 5 deletions(-) diff --git a/tests/integration/management/test_scim_admin_group.py b/tests/integration/management/test_scim_admin_group.py index 1c0cf42ad13..b7bf3cfe968 100644 --- a/tests/integration/management/test_scim_admin_group.py +++ b/tests/integration/management/test_scim_admin_group.py @@ -661,10 +661,6 @@ def test_scim_admin_group_put_members_promotes_and_demotes(gateway: Gateway, tmp @pytest.mark.timeout(300) def test_scim_group_put_returns_the_updated_members(gateway: Gateway, tmp_path: Path) -> None: - pytest.skip( - "BUG: PUT /scim/v2/Groups/{id} returns the members list from before the update, so the 200 body omits " - "the members it just added" - ) group_name: Final = f"scim-admin-put-response-{uuid.uuid4().hex}" with _scim_admin_group_fixture(gateway, tmp_path, group_name) as fixture: expected_member: Final = ScimGroupMember( diff --git a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py index 6ed7e34a583..75c5efc3aa0 100644 --- a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py +++ b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py @@ -334,7 +334,6 @@ def test_deepseek_messages_forward_safeguards_without_adding_a_beta(rig: _Rig, b def test_python_route_adds_the_beta_for_safeguards(gateway: Gateway, scenario: Scenario) -> None: - pytest.skip("BUG: Python /v1/messages sends safeguards without the dangerous-tool-use beta, Anthropic answers 400") with wire_server(anthropic_peer) as wire: deployment: Final = scenario.model(model=f"anthropic/{BACKEND}", api_base=wire.url, api_key=ANTHROPIC_KEY) marker: Final = new_marker() From cd20ecace18848e474fff822d0295c173b81beb7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 11:11:46 -0700 Subject: [PATCH 21/64] test(cost): add tier 1 cost tracking cells and sync usage shapes for 2026-10-17 (#47141) * test(cost): add tier 1 cost tracking cells and sync usage shapes for 2026-10-17 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): restore tiered cost cases and revert out of scope usage edits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(cost): add remaining tier 1 endpoint cells for 2026-10-17 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cost_calculation/cases/anthropic.py | 45 ++ .../cost_calculation/cases/azure.py | 148 +++++++ .../cases/bedrock_converse.py | 125 ++++++ .../cost_calculation/cases/bedrock_mantle.py | 259 ++++++++++- .../cost_calculation/cases/gemini.py | 166 ++++++- .../cost_calculation/cases/misc.py | 112 ++++- .../cost_calculation/cases/openai.py | 360 ++++++++++++++++ .../cost_calculation/cases/vertex_ai.py | 404 +++++++++++++++++- .../integration/cost_calculation/cost_map.py | 20 + 9 files changed, 1601 insertions(+), 38 deletions(-) diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index a0c81d8f06d..5ad4afa9530 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -6899,6 +6899,35 @@ CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_T } ) +def _responses_bridged_stream_case( + sibling: CostTrackingTestCase, stream_sibling: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return sibling.model_copy( + update={ + "name": name, + "request": {**sibling.request, "stream": True}, + "response": stream_sibling.response, + } + ) + + +CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_MESSAGES_STREAM, + "claude-fable-5-1-responses_bridged_stream", +) +CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + "claude-haiku-5-5-responses_bridged_stream", +) +CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_MESSAGES_STREAM, + "claude-opus-5-5-responses_bridged_stream", +) +CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_SONNET_5_5_MESSAGES_STREAM, + "claude-sonnet-5-5-responses_bridged_stream", +) + CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K: Final = with_usage( CLAUDE_HAIKU_5_5_INPUT_TEXT, "claude-haiku-5-5-tiered_input_above_100k", @@ -6994,6 +7023,10 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_SONNET_5_RESPONSES_BRIDGED_JSON, CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_FABLE_5_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, @@ -7119,6 +7152,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( + StreamParityTestCase( + plain=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM + ), StreamParityTestCase(plain=CLAUDE_HAIKU_4_5_INPUT_TEXT, streamed=CLAUDE_HAIKU_4_5_STREAM), StreamParityTestCase( plain=CLAUDE_HAIKU_4_5_RESPONSE_MODEL_OVERRIDE, streamed=CLAUDE_HAIKU_4_5_STREAM_RESPONSE_MODEL_OVERRIDE diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index 1360f65b802..d9f17c12bda 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -21,6 +21,8 @@ from integration.cost_calculation.cost_tracking_case import ( ) from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from integration.cost_calculation.cases.openai import _responses_json_body, _responses_stream_case +from pydantic import JsonValue def _response_for_model( @@ -47,6 +49,37 @@ def _stream_for_model(response: SseResponse, model: str) -> SseResponse: ) +def _azure_responses_json_case( + base: CostTrackingTestCase, name: str, response_model: str, content_filters: JsonValue +) -> CostTrackingTestCase: + response_case: Final = bridge_pairs((base,), ("/v1/responses",))[0].bridged + return response_case.model_copy( + update={ + "name": name, + "response": JsonResponse( + content_type="application/json", + body=_responses_json_body( + response_model, + include_content_filters=True, + content_filters=content_filters, + ), + ), + } + ) + + +def _azure_messages_stream_case( + messages_json: CostTrackingTestCase, stream_case: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return messages_json.model_copy( + update={ + "name": name, + "request": {**messages_json.request, "stream": True}, + "response": stream_case.response, + } + ) + + # azure/gpt-5.4-mini AZURE_GPT_5_4_MINI_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-5.4-mini-input_text", @@ -5812,6 +5845,85 @@ AZURE_GPT_6_ASTRA_CACHE_WRITE_5M: Final = with_response_usage( ExactExpected(spend=0.22375, input_cost=0.2, output_cost=0.02375, prompt_tokens=12928, completion_tokens=380), ) +AZURE_GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K: Final = with_usage( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, + "azure/gpt-5.6-terra-tiered_input_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + }, + ExactExpected( + spend=1.3281576, input_cost=1.32, output_cost=0.0081576, prompt_tokens=300000, completion_tokens=412 + ), +) +AZURE_GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K: Final = with_usage( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, + "azure/gpt-5.6-terra-tiered_cache_read_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 200000}, + }, + ExactExpected( + spend=0.5361576, input_cost=0.528, output_cost=0.0081576, prompt_tokens=300000, completion_tokens=412 + ), +) + + +AZURE_GPT_6_1_SOL_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_1_SOL_INPUT_TEXT, "azure-gpt-6.1-sol-responses_json", "gpt-6.1-sol", None +) +AZURE_GPT_6_LUNA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_LUNA_INPUT_TEXT, "azure-gpt-6-luna-responses_json", "gpt-6-luna", None +) +AZURE_GPT_6_ASTRA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_ASTRA_INPUT_TEXT, "azure-gpt-6-astra-responses_json", "gpt-6-astra", [] +) +AZURE_GPT_5_6_TERRA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, "azure-gpt-5.6-terra-responses_json", "gpt-5.6-terra", None +) +AZURE_GPT_6_1_SOL_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_1_SOL_RESPONSES_JSON, + "azure-gpt-6.1-sol-responses_stream", + "gpt-6.1-sol", + 14, + include_content_filters=True, +) +AZURE_GPT_6_LUNA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_LUNA_RESPONSES_JSON, + "azure-gpt-6-luna-responses_stream", + "gpt-6-luna", + 25, + include_content_filters=True, +) +AZURE_GPT_6_ASTRA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_ASTRA_RESPONSES_JSON, + "azure-gpt-6-astra-responses_stream", + "gpt-6-astra", + 16, + include_content_filters=True, + content_filters=[], +) +AZURE_GPT_5_6_TERRA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_5_6_TERRA_RESPONSES_JSON, + "azure-gpt-5.6-terra-responses_stream", + "gpt-5.6-terra", + 15, + include_content_filters=True, +) +AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_1_SOL_STREAM, "azure-gpt-6.1-sol-messages_bridged_stream" +) +AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_LUNA_STREAM, "azure-gpt-6-luna-messages_bridged_stream" +) +AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_ASTRA_STREAM, "azure-gpt-6-astra-messages_bridged_stream" +) +AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, AZURE_GPT_5_6_TERRA_STREAM, "azure-gpt-5.6-terra-messages_bridged_stream" +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_SOL_MESSAGES_BRIDGED_JSON, @@ -5903,6 +6015,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_LUNA_CACHE_WRITE_5M, AZURE_GPT_6_ASTRA_CACHE_WRITE_5M, AZURE_GPT_5_6_TERRA_INPUT_TEXT, + AZURE_GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K, + AZURE_GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_5_6_TERRA_STREAM, AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX, AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, @@ -5911,6 +6025,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_5_6_TERRA_CACHE_WRITE_5M, AZURE_GPT_5_6_TERRA_REASONING, AZURE_GPT_5_6_TERRA_TOOL_CALL, + AZURE_GPT_6_1_SOL_RESPONSES_JSON, + AZURE_GPT_6_1_SOL_RESPONSES_STREAM, + AZURE_GPT_6_LUNA_RESPONSES_JSON, + AZURE_GPT_6_LUNA_RESPONSES_STREAM, + AZURE_GPT_6_ASTRA_RESPONSES_JSON, + AZURE_GPT_6_ASTRA_RESPONSES_STREAM, + AZURE_GPT_5_6_TERRA_RESPONSES_JSON, + AZURE_GPT_5_6_TERRA_RESPONSES_STREAM, + AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -5924,6 +6050,22 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_INPUT_TEXT, streamed=AZURE_GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_INPUT_TEXT, streamed=AZURE_GPT_6_1_SOL_STREAM), StreamParityTestCase(plain=AZURE_GPT_5_6_TERRA_INPUT_TEXT, streamed=AZURE_GPT_5_6_TERRA_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_RESPONSES_JSON, streamed=AZURE_GPT_6_1_SOL_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_LUNA_RESPONSES_JSON, streamed=AZURE_GPT_6_LUNA_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_RESPONSES_JSON, streamed=AZURE_GPT_6_ASTRA_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_5_6_TERRA_RESPONSES_JSON, streamed=AZURE_GPT_5_6_TERRA_RESPONSES_STREAM), + StreamParityTestCase( + plain=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -5938,4 +6080,10 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=AZURE_GPT_5_6_TERRA_INPUT_TEXT, bridged=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON ), + BridgeCostParityTestCase(native=AZURE_GPT_6_1_SOL_STREAM, bridged=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=AZURE_GPT_6_LUNA_STREAM, bridged=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=AZURE_GPT_6_ASTRA_STREAM, bridged=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase( + native=AZURE_GPT_5_6_TERRA_STREAM, bridged=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) diff --git a/tests/integration/cost_calculation/cases/bedrock_converse.py b/tests/integration/cost_calculation/cases/bedrock_converse.py index f78bdfecfa3..858db06cafc 100644 --- a/tests/integration/cost_calculation/cases/bedrock_converse.py +++ b/tests/integration/cost_calculation/cases/bedrock_converse.py @@ -6088,6 +6088,71 @@ US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_CACHE_WRITE_1H_ABOVE_100K: Final = with_usa ), ) +def _us_anthropic_messages_bridge_pairs( + input_case: CostTrackingTestCase, stream_case: CostTrackingTestCase +) -> tuple[BridgeCostParityTestCase, ...]: + return tuple( + BridgeCostParityTestCase( + native=pair.native, + bridged=pair.bridged.model_copy( + update={ + "name": ( + f"us.anthropic." + f"{pair.native.name.removeprefix('us.anthropic.').removesuffix('-input_text').removesuffix('-stream')}" + f"-messages_bridged_{'stream' if pair.native.request.get('stream') is True else 'json'}" + ) + } + ), + ) + for pair in bridge_pairs((input_case, stream_case), ("/v1/messages",)) + ) + + +def _us_anthropic_responses_stream_case( + json_case: CostTrackingTestCase, stream_case: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return json_case.model_copy( + update={ + "name": name, + "request": {**json_case.request, "stream": True}, + "response": stream_case.response, + } + ) + + +US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + "us.anthropic.claude-fable-5-1-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, + "us.anthropic.claude-haiku-5-5-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM, + "us.anthropic.claude-opus-5-5-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM, + "us.anthropic.claude-sonnet-5-5-responses_bridged_stream", +) +_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM +) +_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM +) +_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM +) +_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_SONNET_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_SONNET_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, @@ -6210,6 +6275,14 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING, US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND, US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS, + US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + *tuple(pair.bridged for pair in _US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -6235,6 +6308,38 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS[1].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -6267,4 +6372,24 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + ), + *_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS, + *_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS, + *_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS, + *_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS, ) diff --git a/tests/integration/cost_calculation/cases/bedrock_mantle.py b/tests/integration/cost_calculation/cases/bedrock_mantle.py index 2fa940d66f4..3cd9150033a 100644 --- a/tests/integration/cost_calculation/cases/bedrock_mantle.py +++ b/tests/integration/cost_calculation/cases/bedrock_mantle.py @@ -18,7 +18,7 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_OPUS_5_5_THINKING, CLAUDE_OPUS_5_5_TOOL_CALL, ) -from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, Deployment, @@ -535,14 +535,40 @@ _MANTLE_GPT_CHAT_REQUEST: Final = { "messages": [{"role": "user", "content": "cost tracking fixture"}], "max_completion_tokens": 64, } +_MANTLE_GPT_TOOL_REQUEST: Final = { + **_MANTLE_GPT_RESPONSES_REQUEST, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ], + "tool_choice": "required", +} +_MANTLE_GPT_CACHE_WRITE_REQUEST: Final = { + **_MANTLE_GPT_RESPONSES_REQUEST, + "input": "cacheable pricing-harness prefix " * 1000, + "prompt_cache_key": "cost-cases-20261017", + "max_output_tokens": 64, +} def _mantle_gpt_usage( - input_tokens: int, output_tokens: int, reasoning_tokens: int, cached_tokens: int = 0 + input_tokens: int, + output_tokens: int, + reasoning_tokens: int, + cached_tokens: int = 0, + cache_write_tokens: int = 0, ) -> dict[str, JsonValue]: return { "input_tokens": input_tokens, - "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": cached_tokens}, + "input_tokens_details": {"cache_write_tokens": cache_write_tokens, "cached_tokens": cached_tokens}, "output_tokens": output_tokens, "output_tokens_details": {"reasoning_tokens": reasoning_tokens}, "total_tokens": input_tokens + output_tokens, @@ -628,6 +654,20 @@ def _mantle_gpt_cases( response_expected, _MANTLE_GPT_CHAT_REQUEST, ) + chat_stream_case: Final = chat_case.model_copy( + update={ + "name": f"bedrock_mantle-{provider_model}-chat_stream", + "request": {**chat_case.request, "stream": True}, + "response": stream_case.response, + } + ) + messages_pairs: Final = bridge_pairs((response_case, stream_case), ("/v1/messages",)) + messages_json: Final = messages_pairs[0].bridged.model_copy( + update={"name": f"bedrock_mantle-{provider_model}-messages_bridged_json"} + ) + messages_stream: Final = messages_pairs[1].bridged.model_copy( + update={"name": f"bedrock_mantle-{provider_model}-messages_bridged_stream"} + ) cache_case: Final = _mantle_gpt_json_case( BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, provider_model, @@ -660,7 +700,73 @@ def _mantle_gpt_cases( tiered_boundary_expected, {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, ) - return (response_case, stream_case, reasoning_case, chat_case, cache_case, tiered_input, tiered_cache, tiered_boundary) + return ( + response_case, + stream_case, + reasoning_case, + chat_case, + cache_case, + tiered_input, + tiered_cache, + tiered_boundary, + chat_stream_case, + messages_json, + messages_stream, + ) + + +def _mantle_gpt_audit3_cases( + provider_model: str, + tool_input_tokens: int, + tool_output_tokens: int, + tool_reasoning_tokens: int, + tool_expected: ExactExpected, + cache_output_tokens: int, + cache_reasoning_tokens: int, + cache_expected: ExactExpected, +) -> tuple[CostTrackingTestCase, ...]: + tool_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tool_call", + _mantle_gpt_usage(tool_input_tokens, tool_output_tokens, tool_reasoning_tokens), + tool_expected, + _MANTLE_GPT_TOOL_REQUEST, + ) + if not isinstance(tool_case.response, JsonResponse): + raise TypeError(f"{tool_case.name}: expected a JSON response fixture") + tool_response: Final = tool_case.response.model_copy( + update={ + "body": { + **tool_case.response.body, + "output": [ + { + "arguments": '{"city":"Paris"}', + "call_id": "call_$REQUEST_ID", + "id": "fc_$REQUEST_ID", + "name": "get_weather", + "status": "completed", + "type": "function_call", + } + ], + } + } + ) + tool_case_with_response: Final = tool_case.model_copy(update={"response": tool_response}) + cache_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "cache_write_5m", + _mantle_gpt_usage( + 1667, + cache_output_tokens, + cache_reasoning_tokens, + cache_write_tokens=1665, + ), + cache_expected, + _MANTLE_GPT_CACHE_WRITE_REQUEST, + ) + return (tool_case_with_response, cache_case) def _mantle_gpt_bridge(case: CostTrackingTestCase) -> BridgeCostParityTestCase: @@ -723,6 +829,83 @@ _MANTLE_GPT_56_TERRA_CASES: Final = _mantle_gpt_cases( ExactExpected(spend=0.84537376, input_cost=0.83776, output_cost=0.00761376, prompt_tokens=272000, completion_tokens=412), ) +_MANTLE_GPT_AUDIT3_CASES: Final = ( + *_mantle_gpt_audit3_cases( + "openai.gpt-6.1-sol", + 50, + 20, + 0, + ExactExpected(spend=0.000363, input_cost=0.000121, output_cost=0.000242, prompt_tokens=50, completion_tokens=20), + 28, + 0, + ExactExpected( + spend=0.005381765, + input_cost=0.005042965, + output_cost=0.0003388, + prompt_tokens=1667, + completion_tokens=28, + cache_creation_cost=0.005038125, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-6-luna", + 51, + 18, + 0, + ExactExpected( + spend=0.000018612, input_cost=0.000006732, output_cost=0.00001188, prompt_tokens=51, completion_tokens=18 + ), + 64, + 40, + ExactExpected( + spend=0.000317229, + input_cost=0.000274989, + output_cost=0.00004224, + prompt_tokens=1667, + completion_tokens=64, + cache_creation_cost=0.000274725, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-6-astra", + 51, + 18, + 0, + ExactExpected( + spend=0.0020163, input_cost=0.0007293, output_cost=0.001287, prompt_tokens=51, completion_tokens=18 + ), + 27, + 0, + ExactExpected( + spend=0.031720975, + input_cost=0.029790475, + output_cost=0.0019305, + prompt_tokens=1667, + completion_tokens=27, + cache_creation_cost=0.029761875, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-5.6-terra", + 51, + 28, + 8, + ExactExpected( + spend=0.00067452, input_cost=0.00015708, output_cost=0.00051744, prompt_tokens=51, completion_tokens=28 + ), + 28, + 0, + ExactExpected( + spend=0.00693385, + input_cost=0.00641641, + output_cost=0.00051744, + prompt_tokens=1667, + completion_tokens=28, + cache_creation_cost=0.00641025, + ), + ), +) + _MANTLE_CLAUDE_OPUS_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5" _ADAPTIVE_USAGE: Final = { "input_tokens": 24, @@ -737,7 +920,25 @@ _ADAPTIVE_EXPECTED: Final = ExactExpected( spend=0.0007524, input_cost=0.0001584, output_cost=0.000594, prompt_tokens=24, completion_tokens=18 ) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = with_usage( + +def _mantle_without_usage_fields(case: CostTrackingTestCase) -> CostTrackingTestCase: + if not isinstance(case.response, JsonResponse): + raise TypeError(f"{case.name}: expected a JSON response fixture") + usage: Final = case.response.body.get("usage") + if not isinstance(usage, dict): + raise TypeError(f"{case.name}: expected a usage object") + response: Final = case.response.model_copy( + update={ + "body": { + **case.response.body, + "usage": {key: value for key, value in usage.items() if key not in {"inference_geo", "speed"}}, + } + } + ) + return case.model_copy(update={"response": response}) + + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_INPUT_TEXT, "bedrock_mantle-anthropic.claude-opus-5-5-input_text", _ADAPTIVE_USAGE, @@ -751,7 +952,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = with_usage( "allowed_openai_params": ["reasoning_effort"], }, } -) +)) _MANTLE_CLAUDE_STREAM_RESPONSE: Final = SseResponse( content_type="text/event-stream", @@ -841,7 +1042,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CLAUDE_OPUS_5_5_MESSAGES } ) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = with_usage( +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_CACHE_READ, "bedrock_mantle-anthropic.claude-opus-5-5-cache_read", { @@ -856,7 +1057,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = with_usage( ExactExpected( spend=0.00308319, input_cost=0.00153219, output_cost=0.001551, prompt_tokens=4225, completion_tokens=47 ), -).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL})) BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = with_usage( CLAUDE_OPUS_5_5_CACHE_WRITE_5M, @@ -892,12 +1093,12 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = with_usage( ), ).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING: Final = with_usage( +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_THINKING, "bedrock_mantle-anthropic.claude-opus-5-5-thinking", _ADAPTIVE_USAGE, _ADAPTIVE_EXPECTED, -).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL})) BEDROCK_MANTLE_CLAUDE_OPUS_5_5_TOOL_CALL: Final = with_usage( CLAUDE_OPUS_5_5_TOOL_CALL, @@ -921,6 +1122,21 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON: Final = BEDROCK_MANTLE_CLAUDE_ } ) +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON: Final = BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-responses_json", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, + } +) +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM: Final = BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-responses_stream", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line", "stream": True}, + } +) + BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( _mantle_gpt_bridge(_MANTLE_GPT_61_SOL_CASES[0]), _mantle_gpt_bridge(_MANTLE_GPT_6_LUNA_CASES[0]), @@ -930,6 +1146,14 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( native=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, bridged=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, ), + *tuple( + BridgeCostParityTestCase(native=cases[0], bridged=cases[9]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), + *tuple( + BridgeCostParityTestCase(native=cases[1], bridged=cases[10]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), ) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( @@ -945,11 +1169,14 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( *_MANTLE_GPT_6_LUNA_CASES, *_MANTLE_GPT_6_ASTRA_CASES, *_MANTLE_GPT_56_TERRA_CASES, + *_MANTLE_GPT_AUDIT3_CASES, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H, @@ -972,4 +1199,16 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, ), + StreamParityTestCase( + plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON, + streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM, + ), + *tuple( + StreamParityTestCase(plain=cases[3], streamed=cases[8]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), + *tuple( + StreamParityTestCase(plain=cases[9], streamed=cases[10]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), ) diff --git a/tests/integration/cost_calculation/cases/gemini.py b/tests/integration/cost_calculation/cases/gemini.py index bd64f4aeb11..ceca816c300 100644 --- a/tests/integration/cost_calculation/cases/gemini.py +++ b/tests/integration/cost_calculation/cases/gemini.py @@ -2,8 +2,10 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" +import json from typing import Final, cast +from pydantic import JsonValue from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, @@ -17,26 +19,74 @@ from integration.cost_calculation.cost_tracking_case import ( from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames -def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: +def _gemini_case_response( + response: JsonResponse | SseResponse, + model_version: str, + usage_fields: dict[str, JsonValue] | None = None, +) -> JsonResponse | SseResponse: match response: case JsonResponse(): - return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + usage: Final = response.body.get("usageMetadata") + body: Final = {**response.body, "modelVersion": model_version} + return response.model_copy( + update={ + "body": { + **body, + **( + { + "usageMetadata": { + **cast(dict[str, JsonValue], usage), + **usage_fields, + } + } + if usage_fields is not None and isinstance(usage, dict) + else {} + ), + } + } + ) case SseResponse(): return response.model_copy( update={ - "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + "frames": tuple( + _gemini_case_frame(frame, model_version, usage_fields) for frame in response.frames + ) } ) +def _gemini_case_frame( + frame: str, model_version: str, usage_fields: dict[str, JsonValue] | None +) -> str: + replaced: Final = frame.replace("gemini-3.8-flash", model_version) + if usage_fields is None: + return replaced + prefix, separator, payload = replaced.partition("data: ") + if not separator or payload == "[DONE]": + return replaced + event: Final = cast(dict[str, JsonValue], json.loads(payload)) + usage: Final = event.get("usageMetadata") + if not isinstance(usage, dict): + return replaced + return f"{prefix}data: {json.dumps({**event, 'usageMetadata': {**usage, **usage_fields}})}" + + def _gemini_case( - sibling: CostTrackingTestCase, name: str, model: str, model_version: str, expected: ExactExpected + sibling: CostTrackingTestCase, + name: str, + model: str, + model_version: str, + expected: ExactExpected, + *, + usage_fields: dict[str, JsonValue] | None = None, ) -> CostTrackingTestCase: return sibling.model_copy( update={ "name": name, "model": model, - "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), model_version), + "response": _gemini_case_response( + cast(JsonResponse | SseResponse, sibling.response), model_version, usage_fields + ), "expected": expected, } ) @@ -1540,6 +1590,7 @@ GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1594,6 +1645,7 @@ GEMINI_GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( "totalTokenCount": 13308, "cachedContentTokenCount": 12288, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1653,6 +1705,7 @@ GEMINI_GEMINI_3_8_FLASH_REASONING: Final = CostTrackingTestCase( "thoughtsTokenCount": 3480, "totalTokenCount": 5280, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1240}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1934,6 +1987,7 @@ GEMINI_GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -3275,19 +3329,46 @@ GEMINI_GEMINI_3_1_PRO_CACHE_CREATION: Final = CostTrackingTestCase( GEMINI_GEMINI_4_ARGON_INPUT_TEXT: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-4-argon-input_text", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}]}, ) GEMINI_GEMINI_4_ARGON_STREAM: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-4-argon-stream", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}]}, ) GEMINI_GEMINI_4_ARGON_CACHE_READ: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-4-argon-cache_read", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00693968, input_cost=0.00275968, output_cost=0.00418, prompt_tokens=12928, completion_tokens=380) + "gemini-4-argon", ExactExpected(spend=0.00693968, input_cost=0.00275968, output_cost=0.00418, prompt_tokens=12928, completion_tokens=380), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}]}, ) GEMINI_GEMINI_4_ARGON_REASONING: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-4-argon-reasoning", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040) + "gemini-4-argon", ExactExpected(spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}]}, +) +GEMINI_GEMINI_4_ARGON_SERVICE_TIER_FLEX: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, + "gemini-gemini-4-argon-service_tier_flex", + "gemini/gemini-4-argon", + "gemini-4-argon", + ExactExpected(spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412), +) +GEMINI_GEMINI_4_ARGON_SERVICE_TIER_PRIORITY: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, + "gemini-gemini-4-argon-service_tier_priority", + "gemini/gemini-4-argon", + "gemini-4-argon", + ExactExpected(spend=0.015444, input_cost=0.0072864, output_cost=0.0081576, prompt_tokens=1840, completion_tokens=412), +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_SERVICE_TIER_PRIORITY: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, + "gemini-gemini-3.1-pro-preview-service_tier_priority", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=0.02018016, input_cost=0.0086112, output_cost=0.01156896, prompt_tokens=1840, completion_tokens=412 + ), ) GEMINI_GEMINI_4_ARGON_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-4-argon-tool_call", "gemini/gemini-4-argon", @@ -3296,28 +3377,60 @@ GEMINI_GEMINI_4_ARGON_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-3.1-pro-preview-input_text", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-3.1-pro-preview-stream", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-3.1-pro-preview-cache_read", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.01078688, input_cost=0.00485888, output_cost=0.005928, prompt_tokens=12928, completion_tokens=380) + ExactExpected(spend=0.01078688, input_cost=0.00485888, output_cost=0.005928, prompt_tokens=12928, completion_tokens=380), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-3.1-pro-preview-reasoning", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.066248, input_cost=0.003224, output_cost=0.063024, prompt_tokens=1240, completion_tokens=4040) + ExactExpected(spend=0.066248, input_cost=0.003224, output_cost=0.063024, prompt_tokens=1240, completion_tokens=4040), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-3.1-pro-preview-tool_call", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) ) +GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K: Final = _gemini_case( + GEMINI_GEMINI_3_1_PRO_TIERED_INPUT_ABOVE_200K, + "gemini-gemini-3.1-pro-preview-tiered_input_above_200k", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=1.079468, input_cost=1.06496, output_cost=0.014508, prompt_tokens=204800, completion_tokens=620 + ), + usage_fields={"serviceTier": "standard"}, +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K: Final = _gemini_case( + GEMINI_GEMINI_3_1_PRO_TIERED_CACHE_READ_ABOVE_200K, + "gemini-gemini-3.1-pro-preview-tiered_cache_read_above_200k", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=0.13742976, input_cost=0.12619776, output_cost=0.011232, prompt_tokens=205824, completion_tokens=480 + ), + usage_fields={"serviceTier": "standard"}, +) +GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_GEMINI_4_ARGON_INPUT_TEXT, GEMINI_GEMINI_4_ARGON_STREAM), + ("/v1/responses", "/v1/messages"), +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM), + ("/v1/responses", "/v1/messages"), +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_1_PRO_INPUT_TEXT, @@ -3379,12 +3492,19 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_4_ARGON_STREAM, GEMINI_GEMINI_4_ARGON_CACHE_READ, GEMINI_GEMINI_4_ARGON_REASONING, + GEMINI_GEMINI_4_ARGON_SERVICE_TIER_FLEX, + GEMINI_GEMINI_4_ARGON_SERVICE_TIER_PRIORITY, GEMINI_GEMINI_4_ARGON_TOOL_CALL, GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM, GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K, + GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K, GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING, GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, + GEMINI_GEMINI_3_1_PRO_PREVIEW_SERVICE_TIER_PRIORITY, + *(pair.bridged for pair in GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES), + *(pair.bridged for pair in GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3410,6 +3530,22 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM ), + StreamParityTestCase( + plain=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[3].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[3].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3422,6 +3558,6 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( ), ("/v1/responses", "/v1/messages"), ), - *bridge_pairs((GEMINI_GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), - *bridge_pairs((GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), + *GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES, + *GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES, ) diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index 1a8b2fcb66b..760f77f91cc 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -19,13 +19,17 @@ from integration.cost_calculation.cost_tracking_case import ( RecountRates, SseResponse, WavUpload, + with_usage, ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_CACHE_READ, CLAUDE_FABLE_5_1_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, CLAUDE_FABLE_5_1_STREAM, CLAUDE_FABLE_5_1_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, @@ -33,12 +37,30 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, CLAUDE_HAIKU_5_5_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_HAIKU_5_5_STREAM, CLAUDE_HAIKU_5_5_THINKING, CLAUDE_HAIKU_5_5_TOOL_CALL, ) +def _azure_ai_claude_messages_case( + response_case: CostTrackingTestCase, + request_case: CostTrackingTestCase, + name: str, + model: str, +) -> CostTrackingTestCase: + return response_case.model_copy( + update={ + "name": name, + "model": model, + "endpoint": "/v1/messages", + "request": request_case.request, + } + ) + + def _azure_ai_claude_usage( input_tokens: int, output_tokens: int, @@ -2027,7 +2049,66 @@ AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = _azure_ai_claude_shape_case( ), ) - +AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + "azure_ai/claude-fable-5-1-messages_input_text", + "azure_ai/claude-fable-5-1", +) +AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, + "azure_ai/claude-fable-5-1-messages_stream", + "azure_ai/claude-fable-5-1", +) +AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES: Final = bridge_pairs( + (AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM), + ("/v1/responses",), +) +AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + "azure_ai/claude-haiku-5-5-messages_input_text", + "azure_ai/claude-haiku-5-5", +) +AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_HAIKU_5_5_STREAM, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + "azure_ai/claude-haiku-5-5-messages_stream", + "azure_ai/claude-haiku-5-5", +) +AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES: Final = bridge_pairs( + (AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K: Final = with_usage( + AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + "azure_ai/claude-haiku-5-5-tiered_input_above_100k", + { + "input_tokens": 150000, + "output_tokens": 480, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + ExactExpected( + spend=0.14478, input_cost=0.1425, output_cost=0.00228, prompt_tokens=150000, completion_tokens=480 + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_CACHE_READ_ABOVE_100K: Final = with_usage( + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + "azure_ai/claude-haiku-5-5-tiered_cache_read_above_100k", + { + "input_tokens": 4096, + "output_tokens": 480, + "cache_read_input_tokens": 110592, + "cache_creation_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + ExactExpected( + spend=0.01667744, input_cost=0.01439744, output_cost=0.00228, prompt_tokens=114688, completion_tokens=480 + ), +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( PERPLEXITY_SONAR_NEXT_OFF_PEAK, NOVA_NEXT_TRANSCRIPTIONS_PER_SECOND, @@ -2088,13 +2169,21 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, AZURE_AI_CLAUDE_FABLE_5_1_THINKING, AZURE_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + *(pair.bridged for pair in AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES), AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, AZURE_AI_CLAUDE_HAIKU_5_5_STREAM, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K, + AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_CACHE_READ_ABOVE_100K, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, AZURE_AI_CLAUDE_HAIKU_5_5_THINKING, AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL, + AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -2105,4 +2194,25 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=AZURE_AI_CLAUDE_HAIKU_5_5_STREAM ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + streamed=AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES[0].bridged, + streamed=AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES[1].bridged, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + streamed=AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES[0].bridged, + streamed=AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES[1].bridged, + ), +) + +BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( + *AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES, + *AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES, ) diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index 85ac05dd69d..f610fd9591a 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -34,6 +34,256 @@ def _stream_for_model(response: SseResponse, model: str) -> SseResponse: return response.model_copy(update={"frames": tuple(frame.replace("gpt-6-luna", model) for frame in response.frames)}) +def _responses_json_body( + model: str, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> dict[str, JsonValue]: + reasoning_item: Final[dict[str, JsonValue]] = { + "content": [], + "encrypted_content": "$ENCRYPTED_CONTENT", + "id": "rs_$REQUEST_ID", + "summary": [], + "type": "reasoning", + } + message_item: Final[dict[str, JsonValue]] = { + "content": [ + { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + } + ], + "id": "msg_$REQUEST_ID", + "phase": "final_answer", + "role": "assistant", + "status": "completed", + "type": "message", + } + output: Final = [reasoning_item, message_item] if include_reasoning else [message_item] + body: Final[dict[str, JsonValue]] = { + "access_programs": [], + "background": False, + "completed_at": 1791855453, + "created_at": 1791855452, + "error": None, + "frequency_penalty": None, + "id": "resp_$REQUEST_ID", + "incomplete_details": None, + "instructions": None, + "max_output_tokens": None, + "max_tool_calls": None, + "metadata": {}, + "moderation": None, + "model": model, + "object": "response", + "output": output, + "parallel_tool_calls": True, + "presence_penalty": None, + "previous_response_id": None, + "prompt_cache_key": None, + "prompt_cache_retention": None, + "reasoning": {"effort": "medium"}, + "safety_identifier": None, + "service_tier": "default", + "status": "completed", + "store": True, + "temperature": 1.0, + "text": None, + "tool_choice": "auto", + "tool_usage": None, + "tools": [], + "top_logprobs": None, + "top_p": 1.0, + "truncation": "disabled", + "usage": { + "input_tokens": 1840, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2252, + }, + "user": None, + } + return {**body, "content_filters": content_filters} if include_content_filters else body + + +def _responses_stream_response( + model: str, + delta_count: int, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> SseResponse: + response_body: Final = _responses_json_body( + model, + include_content_filters=include_content_filters, + content_filters=content_filters, + include_reasoning=include_reasoning, + ) + in_progress_body: Final = { + **response_body, + "output": [], + "status": "in_progress", + "usage": None, + } + reasoning_item: Final = { + "content": [], + "encrypted_content": "$ENCRYPTED_CONTENT", + "id": "rs_$REQUEST_ID", + "summary": [], + "type": "reasoning", + } + message_item: Final = { + "content": [], + "id": "msg_$REQUEST_ID", + "role": "assistant", + "status": "in_progress", + "type": "message", + } + delta_events: Final = tuple( + { + "content_index": 0, + "delta": "scripted response" if index == 0 else "", + "item_id": "msg_$REQUEST_ID", + "logprobs": [], + "obfuscation": None, + "output_index": 1 if include_reasoning else 0, + "sequence_number": (6 if include_reasoning else 4) + index, + "type": "response.output_text.delta", + } + for index in range(delta_count) + ) + delta_start: Final = 6 if include_reasoning else 4 + message_output_index: Final = 1 if include_reasoning else 0 + reasoning_events: Final = ( + ( + { + "item": reasoning_item, + "output_index": 0, + "sequence_number": 2, + "type": "response.output_item.added", + }, + { + "item": reasoning_item, + "output_index": 0, + "sequence_number": 3, + "type": "response.output_item.done", + }, + ) + if include_reasoning + else () + ) + completed_message_item: Final = { + "content": [ + { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + } + ], + "id": "msg_$REQUEST_ID", + "phase": "final_answer", + "role": "assistant", + "status": "completed", + "type": "message", + } + return SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "response": in_progress_body, + "sequence_number": 0, + "type": "response.created", + }, + { + "response": in_progress_body, + "sequence_number": 1, + "type": "response.in_progress", + }, + *reasoning_events, + { + "item": message_item, + "output_index": message_output_index, + "sequence_number": 4 if include_reasoning else 2, + "type": "response.output_item.added", + }, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "output_index": message_output_index, + "part": {"annotations": [], "logprobs": [], "text": "", "type": "output_text"}, + "sequence_number": 5 if include_reasoning else 3, + "type": "response.content_part.added", + }, + *delta_events, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "logprobs": [], + "output_index": message_output_index, + "sequence_number": delta_start + delta_count, + "text": "scripted response", + "type": "response.output_text.done", + }, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "output_index": message_output_index, + "part": { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + }, + "sequence_number": delta_start + delta_count + 1, + "type": "response.content_part.done", + }, + { + "item": completed_message_item, + "output_index": message_output_index, + "sequence_number": delta_start + delta_count + 2, + "type": "response.output_item.done", + }, + { + "response": response_body, + "sequence_number": delta_start + delta_count + 3, + "type": "response.completed", + }, + ), + ) + + +def _responses_stream_case( + base: CostTrackingTestCase, + name: str, + response_model: str, + delta_count: int, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> CostTrackingTestCase: + return base.model_copy( + update={ + "name": name, + "request": {**base.request, "stream": True}, + "response": _responses_stream_response( + response_model, + delta_count, + include_content_filters=include_content_filters, + content_filters=content_filters, + include_reasoning=include_reasoning, + ), + } + ) + + # gpt-5.3-codex GPT_5_3_CODEX_INPUT_TEXT: Final = CostTrackingTestCase( name="gpt-5.3-codex-input_text", @@ -11356,6 +11606,31 @@ GPT_5_6_TERRA_TOOL_CALL: Final = GPT_6_LUNA_TOOL_CALL.model_copy( } ) +GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K: Final = with_usage( + GPT_5_6_TERRA_INPUT_TEXT, + "gpt-5.6-terra-tiered_input_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + }, + ExactExpected( + spend=1.5696408, input_cost=1.56, output_cost=0.0096408, prompt_tokens=300000, completion_tokens=412 + ), +) +GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K: Final = with_usage( + GPT_5_6_TERRA_INPUT_TEXT, + "gpt-5.6-terra-tiered_cache_read_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 200000}, + }, + ExactExpected( + spend=0.6336408, input_cost=0.624, output_cost=0.0096408, prompt_tokens=300000, completion_tokens=412 + ), +) + GPT_5_6_SLASH_ALIAS: Final = GPT_5_6_INPUT_TEXT.model_copy( update={"name": "gpt-5.6-slash_alias", "alias": "openai/custom/gpt-5.6-$MARKER"} @@ -11446,6 +11721,63 @@ GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS: Final = CostTrackingTestCase( ), ) +GPT_6_1_SOL_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_1_SOL_RESPONSES_NATIVE_JSON, + "gpt-6.1-sol-responses_native_stream", + "gpt-6.1-sol", + 1, + include_reasoning=False, +) +GPT_6_LUNA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_LUNA_RESPONSES_NATIVE_JSON, + "gpt-6-luna-responses_native_stream", + "gpt-6-luna", + 45, +) +GPT_6_ASTRA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_ASTRA_RESPONSES_NATIVE_JSON, + "gpt-6-astra-responses_native_stream", + "gpt-6-astra", + 36, +) +GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, + "gpt-5.6-terra-responses_native_stream", + "gpt-5.6-terra", + 95, +) +GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON: Final = bridge_pairs( + (GPT_5_6_TERRA_RESPONSES_NATIVE_JSON,), ("/v1/messages",) +)[0].bridged.model_copy(update={"name": "gpt-5.6-terra-messages_bridged_json"}) +GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM: Final = GPT_6_1_SOL_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6.1-sol-messages_bridged_stream", + "request": {**GPT_6_1_SOL_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_1_SOL_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_6_LUNA_MESSAGES_BRIDGED_STREAM: Final = GPT_6_LUNA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6-luna-messages_bridged_stream", + "request": {**GPT_6_LUNA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_LUNA_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM: Final = GPT_6_ASTRA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6-astra-messages_bridged_stream", + "request": {**GPT_6_ASTRA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_ASTRA_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-messages_bridged_stream", + "request": {**GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM.response, + } +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_SLASH_ALIAS, GPT_5_6_SLASH_ALIAS_STREAM, @@ -11659,6 +11991,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_6_LUNA_CACHE_WRITE_5M, GPT_6_ASTRA_CACHE_WRITE_5M, GPT_5_6_TERRA_INPUT_TEXT, + GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K, + GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K, GPT_5_6_TERRA_STREAM, GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON, @@ -11669,6 +12003,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_TERRA_SERVICE_TIER_FLEX, GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, GPT_5_6_TERRA_TOOL_CALL, + GPT_6_1_SOL_RESPONSES_NATIVE_STREAM, + GPT_6_LUNA_RESPONSES_NATIVE_STREAM, + GPT_6_ASTRA_RESPONSES_NATIVE_STREAM, + GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM, + GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, + GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM, + GPT_6_LUNA_MESSAGES_BRIDGED_STREAM, + GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM, + GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -11698,6 +12041,16 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=GPT_6_ASTRA_INPUT_TEXT, streamed=GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=GPT_6_1_SOL_INPUT_TEXT, streamed=GPT_6_1_SOL_STREAM), StreamParityTestCase(plain=GPT_5_6_TERRA_INPUT_TEXT, streamed=GPT_5_6_TERRA_STREAM), + StreamParityTestCase(plain=GPT_6_1_SOL_RESPONSES_NATIVE_JSON, streamed=GPT_6_1_SOL_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_LUNA_RESPONSES_NATIVE_JSON, streamed=GPT_6_LUNA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_ASTRA_RESPONSES_NATIVE_JSON, streamed=GPT_6_ASTRA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, streamed=GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, streamed=GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase(plain=GPT_6_LUNA_MESSAGES_BRIDGED_JSON, streamed=GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase(plain=GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, streamed=GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase( + plain=GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, streamed=GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -11726,6 +12079,13 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_NATIVE_JSON, bridged=GPT_6_LUNA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_NATIVE_JSON, bridged=GPT_6_ASTRA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_NATIVE_JSON, bridged=GPT_6_1_SOL_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, bridged=GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_NATIVE_STREAM, bridged=GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_NATIVE_STREAM, bridged=GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_NATIVE_STREAM, bridged=GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase( + native=GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM, bridged=GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), BridgeCostParityTestCase(native=GPT_6_SOL_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_SOL_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_LUNA_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_ASTRA_CHAT_BRIDGED_JSON), diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 761078ef593..0e95eb865bf 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -2,6 +2,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" +import json from typing import Final, cast from pydantic import JsonValue @@ -11,13 +12,19 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, CLAUDE_FABLE_5_1_STREAM, CLAUDE_FABLE_5_1_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_OPUS_5_5_CACHE_READ, CLAUDE_OPUS_5_5_CACHE_WRITE_1H, CLAUDE_OPUS_5_5_CACHE_WRITE_5M, CLAUDE_OPUS_5_5_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_STREAM, CLAUDE_OPUS_5_5_STREAM, CLAUDE_OPUS_5_5_THINKING, CLAUDE_OPUS_5_5_TOOL_CALL, @@ -25,6 +32,8 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_SONNET_5_5_CACHE_WRITE_1H, CLAUDE_SONNET_5_5_CACHE_WRITE_5M, CLAUDE_SONNET_5_5_INPUT_TEXT, + CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_SONNET_5_5_MESSAGES_STREAM, CLAUDE_SONNET_5_5_STREAM, CLAUDE_SONNET_5_5_THINKING, CLAUDE_SONNET_5_5_TOOL_CALL, @@ -41,26 +50,72 @@ from integration.cost_calculation.cost_tracking_case import ( from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames -def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: +def _gemini_case_response( + response: JsonResponse | SseResponse, + model_version: str, + usage_fields: dict[str, JsonValue] | None = None, +) -> JsonResponse | SseResponse: match response: case JsonResponse(): - return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + usage: Final = response.body.get("usageMetadata") + body: Final = {**response.body, "modelVersion": model_version} + return response.model_copy( + update={ + "body": { + **body, + **( + { + "usageMetadata": { + **cast(dict[str, JsonValue], usage), + **usage_fields, + } + } + if usage_fields is not None and isinstance(usage, dict) + else {} + ), + } + } + ) case SseResponse(): return response.model_copy( update={ - "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + "frames": tuple( + _gemini_case_frame(frame, model_version, usage_fields) for frame in response.frames + ) } ) +def _gemini_case_frame(frame: str, model_version: str, usage_fields: dict[str, JsonValue] | None) -> str: + replaced: Final = frame.replace("gemini-3.8-flash", model_version) + if usage_fields is None: + return replaced + prefix, separator, payload = replaced.partition("data: ") + if not separator or payload == "[DONE]": + return replaced + event: Final = cast(dict[str, JsonValue], json.loads(payload)) + usage: Final = event.get("usageMetadata") + if not isinstance(usage, dict): + return replaced + return f"{prefix}data: {json.dumps({**event, 'usageMetadata': {**usage, **usage_fields}})}" + + def _gemini_case( - sibling: CostTrackingTestCase, name: str, model: str, expected: ExactExpected + sibling: CostTrackingTestCase, + name: str, + model: str, + expected: ExactExpected, + *, + usage_fields: dict[str, JsonValue] | None = None, + model_version: str = "gemini-3.1-pro-preview", ) -> CostTrackingTestCase: return sibling.model_copy( update={ "name": name, "model": model, - "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), "gemini-3.1-pro-preview"), + "response": _gemini_case_response( + cast(JsonResponse | SseResponse, sibling.response), model_version, usage_fields + ), "expected": expected, } ) @@ -97,6 +152,7 @@ def _vertex_claude_shape_case( expected: ExactExpected, response_content: list[dict[str, JsonValue]] | None = None, adaptive_thinking: bool = False, + thinking_effort: str = "medium", ) -> CostTrackingTestCase: response: Final = cast(JsonResponse, sibling.response) request: Final = ( @@ -107,7 +163,7 @@ def _vertex_claude_shape_case( if key not in {"reasoning_effort", "allowed_openai_params"} }, "thinking": {"type": "adaptive"}, - "output_config": {"effort": "medium"}, + "output_config": {"effort": thinking_effort}, "allowed_openai_params": ["thinking", "output_config"], } if adaptive_thinking @@ -1716,6 +1772,8 @@ GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -1770,6 +1828,8 @@ GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( "totalTokenCount": 13308, "cachedContentTokenCount": 12288, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -2117,6 +2177,8 @@ GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -3122,28 +3184,258 @@ VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( ), ) +VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT, + "vertex_ai/claude-haiku-5-5-input_text", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=1840, output_tokens=412), + ExactExpected(spend=0.00507, input_cost=0.002392, output_cost=0.002678, prompt_tokens=1840, completion_tokens=412), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM: Final = VERTEX_AI_CLAUDE_SONNET_5_5_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-stream", + "model": "vertex_ai/claude-haiku-5-5", + "response": _vertex_claude_stream_response("claude-haiku-5-5", "end_turn", ("scripted answer",)), + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_READ: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ, + "vertex_ai/claude-haiku-5-5-cache_read", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=14, output_tokens=6, cache_read_input_tokens=27003), + ExactExpected( + spend=0.00356759, + input_cost=0.00352859, + output_cost=0.000039, + prompt_tokens=27017, + completion_tokens=6, + cache_read_cost=0.00351039, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + "vertex_ai/claude-haiku-5-5-cache_write_5m", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.043950075, + input_cost=0.043911075, + output_cost=0.000039, + prompt_tokens=27025, + completion_tokens=6, + cache_creation_cost=0.043892875, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + "vertex_ai/claude-haiku-5-5-cache_write_1h", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.0702858, + input_cost=0.0702468, + output_cost=0.000039, + prompt_tokens=27025, + completion_tokens=6, + cache_creation_cost=0.0702286, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_THINKING: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_THINKING, + "vertex_ai/claude-haiku-5-5-thinking", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=18, output_tokens=144, thinking_tokens=37), + ExactExpected( + spend=0.0009594, input_cost=0.0000234, output_cost=0.000936, prompt_tokens=18, completion_tokens=144 + ), + adaptive_thinking=True, + thinking_effort="low", +) +VERTEX_AI_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_TOOL_CALL, + "vertex_ai/claude-haiku-5-5-tool_call", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.0008255, input_cost=0.0005005, output_cost=0.000325, prompt_tokens=385, completion_tokens=50 + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-messages_input_text", + "model": "vertex_ai/claude-haiku-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM: Final = CLAUDE_HAIKU_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-messages_stream", + "model": "vertex_ai/claude-haiku-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) + +VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-1-messages_input_text", + "model": "vertex_ai/claude-fable-5-1", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_1_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-1-messages_stream", + "model": "vertex_ai/claude-fable-5-1", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_1_STREAM.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-opus-5-5-messages_input_text", + "model": "vertex_ai/claude-opus-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CLAUDE_OPUS_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-opus-5-5-messages_stream", + "model": "vertex_ai/claude-opus-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_OPUS_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-sonnet-5-5-messages_input_text", + "model": "vertex_ai/claude-sonnet-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM: Final = CLAUDE_SONNET_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-sonnet-5-5-messages_stream", + "model": "vertex_ai/claude-sonnet-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_SONNET_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) + GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( GEMINI_3_8_FLASH_INPUT_TEXT, "vertex_ai/gemini-3.1-pro-preview-input_text", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( GEMINI_3_8_FLASH_STREAM, "vertex_ai/gemini-3.1-pro-preview-stream", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( GEMINI_3_8_FLASH_CACHE_READ, "vertex_ai/gemini-3.1-pro-preview-cache_read", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380) + ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( GEMINI_3_1_PRO_REASONING, "vertex_ai/gemini-3.1-pro-preview-reasoning", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040) + ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}], + "trafficType": "ON_DEMAND", + }, +) +GEMINI_3_8_FLASH_REASONING: Final = _gemini_case( + GEMINI_3_1_PRO_REASONING, + "vertex_ai/gemini-3.8-flash-reasoning", + "vertex_ai/gemini-3.8-flash", + ExactExpected( + spend=0.01144, + input_cost=0.0006448, + output_cost=0.0107952, + prompt_tokens=1240, + completion_tokens=4040, + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}], + "trafficType": "ON_DEMAND", + }, + model_version="gemini-3.8-flash", ) GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( GEMINI_3_8_FLASH_TOOL_CALL, "vertex_ai/gemini-3.1-pro-preview-tool_call", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) ) +GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K: Final = _gemini_case( + GEMINI_3_1_PRO_TIERED_INPUT_ABOVE_200K, + "vertex_ai/gemini-3.1-pro-preview-tiered_input_above_200k", + "vertex_ai/gemini-3.1-pro-preview", + ExactExpected( + spend=0.913396, input_cost=0.90112, output_cost=0.012276, prompt_tokens=204800, completion_tokens=620 + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 620}], + "trafficType": "ON_DEMAND", + }, +) +GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K: Final = _gemini_case( + GEMINI_3_1_PRO_TIERED_CACHE_READ_ABOVE_200K, + "vertex_ai/gemini-3.1-pro-preview-tiered_cache_read_above_200k", + "vertex_ai/gemini-3.1-pro-preview", + ExactExpected( + spend=0.11628672, input_cost=0.10678272, output_cost=0.009504, prompt_tokens=205824, completion_tokens=480 + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 480}], + "trafficType": "ON_DEMAND", + "cacheTokensDetails": [{"modality": "TEXT", "tokenCount": 201728}], + }, +) + GEMINI_4_ARGON_INPUT_TEXT: Final = GEMINI_3_8_FLASH_INPUT_TEXT.model_copy( update={ "name": "vertex_ai/gemini-4-argon-input_text", @@ -3301,6 +3593,15 @@ GEMINI_4_ARGON_TOOL_CALL: Final = GEMINI_3_8_FLASH_TOOL_CALL.model_copy( } ) +GEMINI_4_ARGON_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_4_ARGON_INPUT_TEXT, GEMINI_4_ARGON_STREAM), + ("/v1/responses", "/v1/messages"), +) +GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_3_1_PRO_PREVIEW_STREAM), + ("/v1/responses", "/v1/messages"), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_INPUT_TEXT, GEMINI_3_1_PRO_CACHE_READ, @@ -3326,6 +3627,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_STREAM_TOOL_CALL, GEMINI_3_1_PRO_STREAM_FULL_USAGE, GEMINI_3_8_FLASH_INPUT_TEXT, + GEMINI_3_8_FLASH_REASONING, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, @@ -3367,9 +3669,30 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, VERTEX_AI_CLAUDE_FABLE_5_1_THINKING, VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_HAIKU_5_5_THINKING, + VERTEX_AI_CLAUDE_HAIKU_5_5_TOOL_CALL, + VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS), GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_3_1_PRO_PREVIEW_STREAM, GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K, + GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K, GEMINI_3_1_PRO_PREVIEW_REASONING, GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, GEMINI_4_ARGON_INPUT_TEXT, @@ -3377,6 +3700,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_4_ARGON_CACHE_READ, GEMINI_4_ARGON_REASONING, GEMINI_4_ARGON_TOOL_CALL, + *(pair.bridged for pair in GEMINI_4_ARGON_ROUTE_BRIDGES), + *(pair.bridged for pair in GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3397,8 +3722,59 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( ), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS[1].bridged, + ), StreamParityTestCase(plain=GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_3_1_PRO_PREVIEW_STREAM), StreamParityTestCase(plain=GEMINI_4_ARGON_INPUT_TEXT, streamed=GEMINI_4_ARGON_STREAM), + StreamParityTestCase( + plain=GEMINI_4_ARGON_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_4_ARGON_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_4_ARGON_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_4_ARGON_ROUTE_BRIDGES[3].bridged, + ), + StreamParityTestCase( + plain=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[3].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3406,6 +3782,10 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( (GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_STREAM), ("/v1/responses", "/v1/messages"), ), - *bridge_pairs((GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), - *bridge_pairs((GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), + *GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES, + *GEMINI_4_ARGON_ROUTE_BRIDGES, + *VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS, ) diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 2ef42249a7e..ba7be7647da 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -889,6 +889,26 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "max_output_tokens": 128000, "max_tokens": 128000, }, + "vertex_ai/claude-haiku-5-5": { + "litellm_provider": "vertex_ai-anthropic_models", + "mode": "chat", + "input_cost_per_token": 1.3e-06, + "input_cost_per_token_above_100k_tokens": 6.5e-06, + "output_cost_per_token": 6.5e-06, + "output_cost_per_token_above_100k_tokens": 3.25e-05, + "cache_creation_input_token_cost": 1.625e-06, + "cache_creation_input_token_cost_above_1hr": 2.6e-06, + "cache_creation_input_token_cost_above_100k_tokens": 8.125e-06, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.3e-05, + "cache_read_input_token_cost": 1.3e-07, + "cache_read_input_token_cost_above_100k_tokens": 6.5e-07, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + }, "groq/qwen/qwen3.8-27b": { "litellm_provider": "groq", "mode": "chat", From 223800f916378661314e8936fc6c6998a27fe16f Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 17 Oct 2026 11:21:03 -0700 Subject: [PATCH 22/64] fix(presidio): stop overlapping entity spans from corrupting numbered PII tokens (#47024) * fix(presidio): resolve overlapping entity spans in numbered tokens anonymization (#42130) Mirror of BerriAI/litellm#42133 by @PINYOPATTANAWASANPORN, squashed onto main with the unit test moved to tests/unit after the test tree relocation * fix(presidio): label an overlapping run by its outermost detection, add integration coverage A run of overlapping detections now takes its entity type from the highest-scoring detection that no other detection in the run contains, matching Presidio's own anonymizer, so a phone number with a date-shaped fragment inside stays PHONE_NUMBER. The resolver returns typed detections instead of Any. * fix(presidio): drop the redundant comment in the numbered token loop * fix(presidio): bump a numbered token another part of the request already holds Each message and content part numbers its tokens from 1 into one request-wide pii_tokens map. Collapsing overlapping detections made an overlapping value's token well formed, so it collided with the same type in another part and both restored as the later value. A token that maps to a different value now moves to the next free number, and a token for the same value is reused * fix(presidio): give a clashing token a spare number instead of scanning for one Scanning upward from the clashing number probed every token an earlier part had claimed, so two parts with thousands of values of one type took seconds of synchronous work. No token in the map is numbered above the map's size, so numbers above both that size and the part's run count are always free * test(presidio): make the anonymizer double reject detections the analyzer never returned Pick the token in one conditional expression instead of reassigning it, and have the scripted anonymizer answer 500 when the proxy forwards spans that differ from the analyzer's, so the default-scope cells check what was sent --------- Co-authored-by: PINYO PATTANAWASANPORN Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../guardrails/guardrail_hooks/presidio.py | 94 +++++-- .../test_presidio_overlapping_spans.py | 251 ++++++++++++++++++ .../guardrail_hooks/test_presidio.py | 240 +++++++++++++++++ 3 files changed, 562 insertions(+), 23 deletions(-) create mode 100644 tests/integration/observability/test_presidio_overlapping_spans.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index f21aabc7f33..cadc9ca689b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -16,6 +16,7 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaita from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime +from itertools import accumulate, count, pairwise from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast import aiohttp @@ -182,6 +183,59 @@ async def _coalesce_first_sse_frame(stream: AsyncIterator[object]) -> AsyncGener yield chunk +@dataclass(frozen=True, slots=True) +class _Detection: + entity_type: str + start: int + end: int + score: float + + +def _strictly_contains(outer: _Detection, inner: _Detection) -> bool: + return ( + outer.start <= inner.start and inner.end <= outer.end and (outer.start, outer.end) != (inner.start, inner.end) + ) + + +def _collapse_run(run: Sequence[_Detection]) -> _Detection: + label: Final = max( + (detection for detection in run if not any(_strictly_contains(other, detection) for other in run)), + key=lambda detection: (detection.score, detection.end - detection.start, detection.entity_type), + ) + return _Detection(label.entity_type, run[0].start, max(detection.end for detection in run), label.score) + + +def _to_detection(item: PresidioAnalyzeResponseItem) -> _Detection: + return _Detection( + str(item.get("entity_type")), + int(item.get("start") or 0), + int(item.get("end") or 0), + float(item.get("score") or 0.0), + ) + + +def _resolve_overlapping_spans(analyze_results: Sequence[PresidioAnalyzeResponseItem]) -> tuple[_Detection, ...]: + """ + Collapse each run of overlapping detections into one detection covering the whole run. + + The run takes its entity type from its highest-scoring detection that no other + detection in the run strictly contains, the same containment rule Presidio's + anonymizer applies, so a phone number with a date-shaped fragment inside it stays + a PHONE_NUMBER. Runs come back sorted by start and never overlap. + """ + detections: Final = sorted( + (detection for detection in map(_to_detection, analyze_results) if detection.start < detection.end), + key=lambda detection: detection.start, + ) + run_ends: Final = tuple(accumulate((detection.end for detection in detections), max)) + boundaries: Final = ( + 0, + *(index for index in range(1, len(detections)) if detections[index].start >= run_ends[index - 1]), + len(detections), + ) + return tuple(_collapse_run(detections[first:last]) for first, last in pairwise(boundaries) if first < last) + + class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -822,31 +876,25 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data["metadata"] = {} if "pii_tokens" not in request_data["metadata"]: request_data["metadata"]["pii_tokens"] = {} - pii_tokens: Final = request_data["metadata"]["pii_tokens"] + pii_tokens: Final = cast("dict[str, str]", request_data["metadata"]["pii_tokens"]) # cast-ok: this hook's map - # Assign sequence numbers in forward (left-to-right) order so - # that is the first entity in the text, etc. - sorted_forward: Final = sorted(analyze_results, key=lambda x: x["start"]) - seq_map: Final = {} - for idx, ar in enumerate(sorted_forward, start=1): - seq_map[(ar["start"], ar["end"])] = idx + runs: Final = _resolve_overlapping_spans(analyze_results) + # Every message and content part numbers its runs from 1 into this one map, so a + # token another part already holds for a different value takes a spare number. + # No token in the map is numbered above its size, so spares start above both. + spare_numbers: Final = count(max(len(pii_tokens), len(runs)) + 1) - # Apply replacements in reverse order by start position so - # that replacing later spans first does not shift earlier - # coordinates in the original text. - for ar in reversed(sorted_forward): - start = ar["start"] - end = ar["end"] - entity_type = ar["entity_type"] - replacement = f"<{entity_type}>" - seq = seq_map[(start, end)] - if replacement.endswith(">"): - replacement = f"{replacement[:-1]}_{seq}>" - else: - replacement = f"{replacement}_{seq}" - pii_tokens[replacement] = text[start:end] - new_text = new_text[:start] + replacement + new_text[end:] - masked_entity_count[entity_type] = masked_entity_count.get(entity_type, 0) + 1 + for seq, run in reversed(tuple(enumerate(runs, start=1))): + original = text[run.start : run.end] + numbered = f"<{run.entity_type}_{seq}>" + replacement = ( + numbered + if pii_tokens.get(numbered, original) == original + else f"<{run.entity_type}_{next(spare_numbers)}>" + ) + pii_tokens[replacement] = original + new_text = new_text[: run.start] + replacement + new_text[run.end :] + masked_entity_count[run.entity_type] = masked_entity_count.get(run.entity_type, 0) + 1 return new_text async def anonymize_text( diff --git a/tests/integration/observability/test_presidio_overlapping_spans.py b/tests/integration/observability/test_presidio_overlapping_spans.py new file mode 100644 index 00000000000..a2d92b15330 --- /dev/null +++ b/tests/integration/observability/test_presidio_overlapping_spans.py @@ -0,0 +1,251 @@ +import json +import re +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +MODEL: Final = "gpt-6.1-sol" +PHONE: Final = "4251234567" +EMAIL: Final = "john.smith@example.com" +INTERNATIONAL_PHONE: Final = "+33 1 42 68 53 00" +US_PHONE: Final = "+1 (206) 555-0199" +SAME_SPAN: Final = f"Repeat exactly: my number is {PHONE}" +NESTED: Final = f"contact {EMAIL} today" +CONTAINED_FRAGMENT: Final = f"her phone is {INTERNATIONAL_PHONE}." +SEPARATE_PARTS: Final = (f"My number is {PHONE}.", f"Also call {US_PHONE}.") +# Spans, scores and anonymized texts recorded from presidio-analyzer and presidio-anonymizer 2.2.364 on 2026-10-16 +DETECTIONS: Final = MappingProxyType( + { + PHONE: (("PHONE_NUMBER", 0, 10, 0.75), ("US_BANK_NUMBER", 0, 10, 0.05), ("US_DRIVER_LICENSE", 0, 10, 0.01)), + EMAIL: (("EMAIL_ADDRESS", 0, 22, 1.0), ("URL", 0, 7, 0.5), ("URL", 11, 22, 0.5)), + "today": (("DATE_TIME", 0, 5, 0.85),), + INTERNATIONAL_PHONE: (("DATE_TIME", 12, 17, 0.85), ("PHONE_NUMBER", 0, 17, 0.75)), + US_PHONE: (("PHONE_NUMBER", 0, 17, 0.4),), + } +) +ANONYMIZED: Final = MappingProxyType( + { + SAME_SPAN: "Repeat exactly: my number is ", + NESTED: "contact ", + CONTAINED_FRAGMENT: "her phone is .", + SEPARATE_PARTS[0]: "My number is .", + SEPARATE_PARTS[1]: "Also call .", + } +) +CASES: Final = ( + pytest.param(SAME_SPAN, "Repeat exactly: my number is ", id="same_span"), + pytest.param(NESTED, "contact ", id="nested"), + pytest.param(CONTAINED_FRAGMENT, "her phone is .", id="container_outranks_higher_scored_fragment"), +) +STREAM: Final = pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + + +def _detections(text: str, value: str) -> Iterator[dict[str, str | int | float]]: + for match in re.finditer(rf"(? list[dict[str, str | int | float]]: + findings: Final = chain.from_iterable(_detections(text, value) for value in DETECTIONS) + return sorted(findings, key=lambda item: -item["score"]) + + +def _spans(results: list[dict[str, str | int | float]]) -> list[tuple[str | int | float, ...]]: + return sorted(tuple(result[key] for key in ("entity_type", "start", "end", "score")) for result in results) + + +def _analyze(request: Request) -> Reply: + assert request.target == "/analyze", request.target + return Reply(body=json.dumps(_findings(json.loads(request.body)["text"])).encode()) + + +def _anonymize(request: Request) -> Reply: + assert request.target == "/anonymize", request.target + body: Final = json.loads(request.body) + text: Final = body["text"] + if text not in ANONYMIZED: + return Reply(status=500, body=json.dumps({"error": f"no recorded anonymization for {text!r}"}).encode()) + if _spans(body["analyzer_results"]) != _spans(_findings(text)): + return Reply( + status=500, body=json.dumps({"error": f"analyzer_results differ from the detections for {text!r}"}).encode() + ) + return Reply(body=json.dumps({"text": ANONYMIZED[text], "items": []}).encode()) + + +def _texts(messages: list[dict]) -> tuple[str, ...]: + return tuple( + chain.from_iterable( + (message["content"],) + if isinstance(message["content"], str) + else (part["text"] for part in message["content"]) + for message in messages + ) + ) + + +def _chunk(text: str, finish: str | None) -> bytes: + delta: Final = {"role": "assistant", "content": text} if finish is None else {} + payload: Final = { + "id": "chatcmpl-echo", + "object": "chat.completion.chunk", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _echo(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + body: Final = json.loads(request.body) + text: Final = " | ".join(_texts(body["messages"])) + if body.get("stream"): + pieces: Final = tuple(text[index : index + 6] for index in range(0, len(text), 6)) + return Reply( + content_type="text/event-stream", + chunks=(*(_chunk(piece, None) for piece in pieces), _chunk("", "stop"), b"data: [DONE]\n\n"), + ) + reply: Final = { + "id": "chatcmpl-echo", + "object": "chat.completion", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 9, "total_tokens": 18}, + } + return Reply(body=json.dumps(reply).encode()) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + provider: Wire + + +def _config(root: Path, provider: Wire, analyzer: Wire, anonymizer: Wire, scope: Literal["input"] | None) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"]["cache"] = False + config["model_list"].append( + { + "model_name": MODEL, + "litellm_params": {"model": f"openai/{MODEL}", "api_base": provider.url + "/v1", "api_key": "synthetic"}, + } + ) + guardrail: Final = { + "guardrail": "presidio", + "mode": "pre_call", + "output_parse_pii": True, + "default_on": True, + "presidio_analyzer_api_base": analyzer.url + "/", + "presidio_anonymizer_api_base": anonymizer.url + "/", + } + scoped: Final = guardrail if scope is None else {**guardrail, "presidio_filter_scope": scope} + config["guardrails"] = [{"guardrail_name": "presidio-pii", "litellm_params": scoped}] + path: Final = root / "presidio-overlapping-spans.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@contextmanager +def _rig(root: Path, scope: Literal["input"] | None) -> Iterator[Rig]: + with ( + gateway_from_environment() as gateway, + wire_server(_echo) as provider, + wire_server(_analyze) as analyzer, + wire_server(_anonymize) as anonymizer, + owned_proxy_process( + gateway, root, {}, config=_config(root, provider, analyzer, anonymizer, scope), workers=2 + ) as owned, + ): + yield Rig(owned.gateway, provider) + + +@pytest.fixture(scope="module") +def default_scope_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory.mktemp("presidio-default-scope"), None) as rig: + yield rig + + +@pytest.fixture(scope="module") +def input_scope_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory.mktemp("presidio-input-scope"), "input") as rig: + yield rig + + +def _send(rig: Rig, messages: list[dict], stream: bool) -> tuple[int, str]: + rig.provider.drain() + body: Final = {"model": MODEL, "messages": messages, "stream": stream} + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + if not stream: + response: Final = rig.proxy.client.post("/v1/chat/completions", json=body, headers=headers) + if response.status_code != 200: + return response.status_code, response.text + return response.status_code, response.json()["choices"][0]["message"]["content"] + with rig.proxy.client.stream("POST", "/v1/chat/completions", json=body, headers=headers) as streamed: + if streamed.status_code != 200: + return streamed.status_code, streamed.read().decode() + frames: Final = tuple( + json.loads(line.removeprefix("data: ")) + for line in streamed.iter_lines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + return 200, "".join(frame["choices"][0]["delta"].get("content") or "" for frame in frames if frame["choices"]) + + +def _forwarded(rig: Rig) -> tuple[str, ...]: + received: Final = rig.provider.drain() + assert len(received) == 1, received + return _texts(json.loads(received[0].body)["messages"]) + + +@STREAM +@pytest.mark.parametrize(("text", "forwarded"), CASES) +def test_overlapping_detections_reach_the_provider_as_one_token_and_come_back_unmasked( + input_scope_rig: Rig, text: str, forwarded: str, stream: bool +) -> None: + status, replied = _send(input_scope_rig, [{"role": "user", "content": text}], stream) + assert status == 200, replied + assert _forwarded(input_scope_rig) == (forwarded,) + assert replied == text + + +@STREAM +@pytest.mark.parametrize(("text", "forwarded"), CASES) +def test_default_scope_masks_the_cleanly_restored_reply_like_presidio_masks_the_original( + default_scope_rig: Rig, text: str, forwarded: str, stream: bool +) -> None: + status, replied = _send(default_scope_rig, [{"role": "user", "content": text}], stream) + assert status == 200, replied + assert _forwarded(default_scope_rig) == (forwarded,) + assert replied == ANONYMIZED[text] + + +@STREAM +@pytest.mark.parametrize( + "messages", + [ + pytest.param( + [{"role": "user", "content": [{"type": "text", "text": part} for part in SEPARATE_PARTS]}], + id="content_parts", + ), + pytest.param([{"role": "user", "content": part} for part in SEPARATE_PARTS], id="messages"), + ], +) +def test_one_type_in_separate_parts_gets_distinct_tokens_and_each_part_comes_back_as_itself( + input_scope_rig: Rig, messages: list[dict], stream: bool +) -> None: + status, replied = _send(input_scope_rig, messages, stream) + assert status == 200, replied + assert _forwarded(input_scope_rig) == ("My number is .", "Also call .") + assert replied == " | ".join(SEPARATE_PARTS) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py index 7d0c81af4c2..bd1e9ae4068 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -3789,6 +3789,246 @@ async def test_output_parse_pii_numbered_tokens_across_chunks(): assert pii_tokens[""] == CHUNK_MARKER_TWO +@pytest.mark.parametrize( + ("text", "analyze_results", "expected_text", "expected_tokens", "expected_counts"), + [ + pytest.param( + "Repeat exactly: my number is 4251234567", + [ + {"entity_type": "PHONE_NUMBER", "start": 29, "end": 39, "score": 0.75}, + {"entity_type": "US_BANK_NUMBER", "start": 29, "end": 39, "score": 0.05}, + {"entity_type": "US_DRIVER_LICENSE", "start": 29, "end": 39, "score": 0.01}, + ], + "Repeat exactly: my number is ", + {"": "4251234567"}, + {"PHONE_NUMBER": 1}, + id="same-span", + ), + pytest.param( + "contact john.smith@example.com today", + [ + {"entity_type": "EMAIL_ADDRESS", "start": 8, "end": 30, "score": 1.0}, + {"entity_type": "DATE_TIME", "start": 31, "end": 36, "score": 0.85}, + {"entity_type": "URL", "start": 8, "end": 15, "score": 0.5}, + {"entity_type": "URL", "start": 19, "end": 30, "score": 0.5}, + ], + "contact ", + {"": "john.smith@example.com", "": "today"}, + {"EMAIL_ADDRESS": 1, "DATE_TIME": 1}, + id="nested", + ), + pytest.param( + "her phone is +33 1 42 68 53 00.", + [ + {"entity_type": "DATE_TIME", "start": 25, "end": 30, "score": 0.85}, + {"entity_type": "PHONE_NUMBER", "start": 13, "end": 30, "score": 0.75}, + ], + "her phone is .", + {"": "+33 1 42 68 53 00"}, + {"PHONE_NUMBER": 1}, + id="container-outranks-higher-scored-fragment", + ), + pytest.param( + "ids ab0123456789 and cd0123456789", + [ + {"entity_type": "AAA", "start": 4, "end": 16, "score": 0.5}, + {"entity_type": "BBB", "start": 21, "end": 29, "score": 0.5}, + {"entity_type": "CCC", "start": 25, "end": 33, "score": 0.9}, + ], + "ids and ", + {"": "ab0123456789", "": "cd0123456789"}, + {"AAA": 1, "CCC": 1}, + id="partial-overlap-in-second-run", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 3, "end": 15, "score": 0.5}, + {"entity_type": "BBB", "start": 3, "end": 8, "score": 0.9}, + ], + "id ", + {"": "ab0123456789"}, + {"AAA": 1}, + id="container-sharing-start-outranks-fragment", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 3, "end": 10, "score": 0.9}, + {"entity_type": "BBB", "start": 7, "end": 15, "score": 0.5}, + ], + "id ", + {"": "ab0123456789"}, + {"AAA": 1}, + id="earlier-higher-scored-detection-labels-the-whole-run", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 6, "end": 15, "score": 0.5}, + {"entity_type": "BBB", "start": 3, "end": 8, "score": 0.5}, + ], + "id ", + {"": "ab0123456789"}, + {"AAA": 1}, + id="score-tie-goes-to-the-longer-detection", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 3, "end": 9, "score": 0.5}, + {"entity_type": "BBB", "start": 9, "end": 15, "score": 0.5}, + ], + "id ", + {"": "ab0123", "": "456789"}, + {"AAA": 1, "BBB": 1}, + id="touching-detections-stay-separate", + ), + pytest.param( + "id ab0123456789", + [ + {"entity_type": "AAA", "start": 3, "end": 15, "score": 0.5}, + {"entity_type": "BBB", "start": 3, "end": 15, "score": 0.5}, + ], + "id ", + {"": "ab0123456789"}, + {"BBB": 1}, + id="identical-spans-with-equal-scores-label-the-same-in-any-order", + ), + pytest.param( + "My name is John Smith, phone +1 415 555 2671, please check the order.", + [ + {"entity_type": "EMPTY", "start": 25, "end": 25, "score": 0.99}, + {"entity_type": "INVERTED", "start": 40, "end": 30, "score": 0.99}, + {"entity_type": "PERSON", "start": 11, "end": 21, "score": 0.9}, + {"entity_type": "PHONE_NUMBER", "start": 29, "end": 44, "score": 0.85}, + {"entity_type": "UK_NHS", "start": 32, "end": 44, "score": 0.6}, + ], + "My name is , phone , please check the order.", + {"": "John Smith", "": "+1 415 555 2671"}, + {"PERSON": 1, "PHONE_NUMBER": 1}, + id="empty-and-inverted-spans-are-ignored", + ), + pytest.param( + "Code ABC-12345-XYZ is valid.", + [ + {"entity_type": "GENERIC_ID", "start": 5, "end": 14, "score": 0.5}, + {"entity_type": "SECRET_KEY", "start": 9, "end": 18, "score": 0.95}, + ], + "Code is valid.", + {"": "ABC-12345-XYZ"}, + {"SECRET_KEY": 1}, + id="later-higher-scored-detection-labels-the-whole-run", + ), + ], +) +@pytest.mark.parametrize("order", ["analyzer", "reversed"]) +def test_numbered_tokens_collapse_overlapping_detections_into_one_token_per_run( + text: str, + analyze_results: list[PresidioAnalyzeResponseItem], + expected_text: str, + expected_tokens: dict[str, str], + expected_counts: dict[str, int], + order: Literal["analyzer", "reversed"], +) -> None: + guardrail: Final = OPTIONAL_PresidioPIIMasking(mock_testing=True) + request_data: Final = {"metadata": {}} + masked_entity_count: Final = {} + ordered: Final = analyze_results if order == "analyzer" else analyze_results[::-1] + + masked: Final = guardrail._finalize_presidio_anonymize_numbered_tokens( + text, ordered, request_data, masked_entity_count + ) + + assert masked == expected_text + assert request_data["metadata"]["pii_tokens"] == expected_tokens + assert masked_entity_count == expected_counts + assert OPTIONAL_PresidioPIIMasking._unmask_pii_text(masked, expected_tokens) == text + + +@pytest.mark.parametrize( + ("parts", "expected_masked", "expected_tokens"), + [ + pytest.param( + ( + ( + "My number is 4251234567.", + [ + {"entity_type": "PHONE_NUMBER", "start": 13, "end": 23, "score": 0.75}, + {"entity_type": "US_BANK_NUMBER", "start": 13, "end": 23, "score": 0.05}, + {"entity_type": "US_DRIVER_LICENSE", "start": 13, "end": 23, "score": 0.01}, + ], + ), + ( + "Also call +1 (206) 555-0199.", + [{"entity_type": "PHONE_NUMBER", "start": 10, "end": 27, "score": 0.4}], + ), + ), + ("My number is .", "Also call ."), + {"": "4251234567", "": "+1 (206) 555-0199"}, + id="overlapping-run-then-plain-detection", + ), + pytest.param( + ( + ("Call John Smith.", [{"entity_type": "PERSON", "start": 5, "end": 15, "score": 0.85}]), + ("Again John Smith.", [{"entity_type": "PERSON", "start": 6, "end": 16, "score": 0.85}]), + ), + ("Call .", "Again ."), + {"": "John Smith"}, + id="same-value-keeps-its-token", + ), + pytest.param( + ( + ("Call John Smith.", [{"entity_type": "PERSON", "start": 5, "end": 15, "score": 0.85}]), + ( + "Ann Lee and Bob Ray.", + [ + {"entity_type": "PERSON", "start": 0, "end": 7, "score": 0.85}, + {"entity_type": "PERSON", "start": 12, "end": 19, "score": 0.85}, + ], + ), + ), + ("Call .", " and ."), + {"": "John Smith", "": "Bob Ray", "": "Ann Lee"}, + id="bump-skips-every-claimed-number", + ), + pytest.param( + ( + ( + "Ann Lee and Bob Ray.", + [ + {"entity_type": "PERSON", "start": 0, "end": 7, "score": 0.85}, + {"entity_type": "PERSON", "start": 12, "end": 19, "score": 0.85}, + ], + ), + ("Call John Smith.", [{"entity_type": "PERSON", "start": 5, "end": 15, "score": 0.85}]), + ), + (" and .", "Call ."), + {"": "Ann Lee", "": "Bob Ray", "": "John Smith"}, + id="spare-number-starts-above-every-earlier-part", + ), + ], +) +def test_numbered_tokens_stay_unique_across_parts_of_one_request( + parts: tuple[tuple[str, list[PresidioAnalyzeResponseItem]], ...], + expected_masked: tuple[str, ...], + expected_tokens: dict[str, str], +) -> None: + guardrail: Final = OPTIONAL_PresidioPIIMasking(mock_testing=True) + request_data: Final = {"metadata": {}} + + masked: Final = tuple( + guardrail._finalize_presidio_anonymize_numbered_tokens(text, results, request_data, {}) + for text, results in parts + ) + + assert masked == expected_masked + assert request_data["metadata"]["pii_tokens"] == expected_tokens + assert OPTIONAL_PresidioPIIMasking._unmask_pii_text(" | ".join(masked), expected_tokens) == " | ".join( + text for text, _ in parts + ) + + @pytest.mark.asyncio async def test_analyze_text_chunked_failure_stays_fail_closed(): """If one chunk still fails, the chunked path raises exactly like a single From 8226b2be9ad26b297839f2b84bfb418493173f26 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 18:25:32 +0000 Subject: [PATCH 23/64] fix(otel): report cache hits at zero cost with a cache-hit flag in V1 and V2 spans (#43604) * fix(litellm): report zero cost and cache-hit attrs on cache-served spans * fix(litellm): stamp cache-hit marker on HiddenParams carriers and restore probe side effects gauntlet-mini on ace57d25f verified two live gaps. TextCompletionResponse and ResponsesAPIResponse carry _hidden_params as a HiddenParams object, so the dict-gated marker write skipped them and the served hit still stamped the full price. The marker write is now container-agnostic, which it already was for the partial-embedding stamp, and a regression test exercises the carrier end to end. _uncached_response_cost also restores billed_token_rates, zero_cost_warned, and the response_cost_failure_debug_information key alongside cost_breakdown, so the probe leaves no pricing state on the call it reads. * test(caching): cover cache-hit marker on HiddenParams carriers through the caching handler * fix(logging): keep the uncached price probe's restores to the state it can leak * fix(otel): emit V2 LLM-call span for cache hits with explicit cache attributes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): read cache payload fields defensively without casts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(otel): clarify pyright suppression reason on cache payload reads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): build the cache-hit view from typing_extensions.TypedDict so import works on Python 3.10 and 3.11 * fix(otel): report cache hits at zero cost everywhere and mark misses Cache hits now carry a zero cost breakdown and a zero hidden response cost, the /v1/messages hit cost header follows chat and responses hits, every span carries litellm.cache_hit, V2 hit spans drop server.address, an async hit keeps the provider-prefixed model name, and the V2 attribute budget counts preset keys * fix(otel): send a 0 cost header on cache hits and keep hits out of deployment TPM Cache hits now send x-litellm-response-cost: 0.0 on chat, responses and /v1/messages, while a fresh call whose cost is 0 still sends no header. The async hit path stamps the same model name as main again. Stamping the provider-prefixed name made a hit's tokens count toward the deployment's TPM limit, so repeated cached requests hit 429s under usage-based routing. Drop comments that restated fields and annotate new test values as Final. * fix(otel): keep the served pricing basis on a cache hit's cost breakdown A hit's breakdown now keeps service_tier, data_residency, vertex_location, and the discount and margin rates the served call recorded, with every amount at 0, so savings reports still read the basis. The hit header test now drives the live non-streaming header path. * fix(otel): send zero breakdown headers on a cache hit and give hit spans a fresh call's cost keys A hit's x-litellm-response-cost-original, -input, and -output headers now read 0.0 instead of going missing or carrying the full price. A hit's cost breakdown now comes from pricing the response as fresh, so every route exports the same keys, pricing basis, and rates with every amount at 0. * test(integration): assert zero cost and the cache-hit flag on cache-served OTel V1 and V2 spans * test: wait for the cache entry before the cache-hit request and build cost headers without a route_request double * test: annotate the new cache-hit unit test locals with Final --------- Co-authored-by: shrey kharbanda Co-authored-by: jesus Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/caching/caching_handler.py | 8 +- litellm/integrations/opentelemetry.py | 18 +- litellm/integrations/otel/emitter.py | 32 +- litellm/integrations/otel/logger.py | 8 +- litellm/integrations/otel/mappers/genai.py | 2 + litellm/integrations/otel/model/metadata.py | 5 + litellm/integrations/otel/model/payloads.py | 7 +- litellm/integrations/otel/model/semconv.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 82 +++- litellm/proxy/common_request_processing.py | 48 ++- .../observability/test_otel_cache_hit_cost.py | 370 ++++++++++++++++++ tests/unit/caching/test_caching_handler.py | 88 +++++ .../otel/test_otel_v2_components.py | 13 + .../integrations/otel/test_otel_v2_emitter.py | 65 ++- .../integrations/otel/test_otel_v2_logger.py | 43 ++ tests/unit/integrations/test_opentelemetry.py | 69 ++++ .../test_litellm_logging.py | 310 +++++++++++++++ .../proxy/test_common_request_processing.py | 96 +++++ 18 files changed, 1205 insertions(+), 61 deletions(-) create mode 100644 tests/integration/observability/test_otel_cache_hit_cost.py diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 063faa70f2f..7089ff78835 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -346,6 +346,7 @@ class LLMCachingHandler: if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): verbose_logger.debug("Checking Async Cache") + logging_obj.reset_caching_details() cached_result = await self._retrieve_from_cache( call_type=call_type, kwargs=kwargs, @@ -357,7 +358,7 @@ class LLMCachingHandler: verbose_logger.debug("Cache Hit!") cache_hit: Final = True end_time: Final = datetime.datetime.now() - model, custom_llm_provider, _, _ = litellm.get_llm_provider( + provider_model, custom_llm_provider, _, _ = litellm.get_llm_provider( model=model, custom_llm_provider=kwargs.get("custom_llm_provider", None), api_base=kwargs.get("api_base", None), @@ -366,7 +367,7 @@ class LLMCachingHandler: cache_duration_ms: Final = (cache_check_end_time - cache_check_start_time) * 1000 self._update_litellm_logging_obj_environment( logging_obj=logging_obj, - model=model, + model=provider_model, kwargs=kwargs, cached_result=cached_result, is_async=True, @@ -381,7 +382,7 @@ class LLMCachingHandler: call_type=call_type, kwargs=kwargs, logging_obj=logging_obj, - model=model, + model=provider_model, custom_llm_provider=kwargs.get("custom_llm_provider", None), args=args, ) @@ -467,6 +468,7 @@ class LLMCachingHandler: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) print_verbose("Checking Sync Cache") + logging_obj.reset_caching_details() with response_cache_phase("get"): cached_result = litellm.cache.get_cache(**new_kwargs) if is_response_without_output(cached_result): diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 90383ddb60d..fb0723289ee 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -6,9 +6,10 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, cast from pydantic import ConfigDict, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -106,6 +107,14 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i # Cap on credential-scoped providers held at once; each one owns an exporter thread. _MAX_DYNAMIC_TRACER_PROVIDERS: Final = 256 + +class _CacheHitView(TypedDict, total=False): + cache_hit: ReadOnly[object] + saved_cache_cost: ReadOnly[object] + + +_CACHE_HIT_VIEW: Final = TypeAdapter(_CacheHitView) + # Dedicated so a slow exporter shutdown cannot starve the shared logging executor. _PROVIDER_SHUTDOWN_EXECUTOR: Final = ThreadPoolExecutor(max_workers=4, thread_name_prefix="OtelProviderShutdown") @@ -2434,6 +2443,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): litellm_params=litellm_params, ) # Cost breakdown tracking + cache_view: Final = _CACHE_HIT_VIEW.validate_python(standard_logging_payload) + cache_hit: Final = cache_view.get("cache_hit") is True cost_breakdown: Final[CostBreakdown | None] = standard_logging_payload.get("cost_breakdown") if cost_breakdown: for key, value in cost_breakdown.items(): @@ -2443,6 +2454,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): key=f"gen_ai.cost.{key}", value=value, ) + self.safe_set_attribute(span=span, key=LiteLLM.CACHE_HIT, value=cache_hit) + if cache_hit: + saved_cache_cost: Final = cache_view.get("saved_cache_cost") + if isinstance(saved_cache_cost, (int, float)): + self.safe_set_attribute(span=span, key=LiteLLM.SAVED_CACHE_COST, value=saved_cache_cost) ############################################# ########## LLM Request Attributes ########### ############################################# diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py index a4cb0a84541..8779153959f 100644 --- a/litellm/integrations/otel/emitter.py +++ b/litellm/integrations/otel/emitter.py @@ -100,23 +100,19 @@ def _carried_keys( return frozenset(key for key, _value in attributes) -def attribute_budget(span: Span, reserved: int, overwrites: frozenset[str] = frozenset()) -> int | None: - """How many mapped attributes fit on ``span`` next to what it already carries and ``reserved`` more. +def carried_keys(span: Span) -> frozenset[str]: + """The attribute keys ``span`` already holds, empty for a span that exposes none.""" + if not isinstance(span, ReadableSpan): + return frozenset() + return _carried_keys(span.attributes or ()) - ``overwrites`` are the mapped keys already present on ``span``: setting one - replaces the value in place and consumes no slot against the limit, so only - the genuinely new pre-existing keys reduce the budget. Counting the - overwritten ones too reserves slots the fit can never spend and sheds - indexed message attributes for nothing. - """ + +def attribute_budget(span: Span, reserved: int) -> int | None: + """How many new attribute keys fit on ``span`` next to what it already carries and ``reserved`` more.""" limit: Final = span_attribute_limit(span) if limit is None: return None - if not isinstance(span, ReadableSpan): - return limit - reserved - existing_keys: Final = _carried_keys(span.attributes or ()) - fresh: Final = len(existing_keys - overwrites) if overwrites else len(existing_keys) - return limit - fresh - reserved + return limit - len(carried_keys(span)) - reserved def stamp_error( @@ -303,8 +299,14 @@ class SpanEmitter: ) stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES reserved: Final = len(stamped_later.keys() - mapped.keys()) - for key, value in fit_indexed_messages(mapped, attribute_budget(span, reserved, frozenset(mapped))).items(): - span.set_attribute(key, value) + overwrites: Final = carried_keys(span) & mapped.keys() + fitted: Final = fit_indexed_messages( + MappingProxyType({key: value for key, value in mapped.items() if key not in overwrites}), + attribute_budget(span, reserved), + ) + for key, value in mapped.items(): + if key in overwrites or key in fitted: + span.set_attribute(key, value) if error: stamped: Final = stamp_error(span, error) if stamped is not None and self._event_recorder is not None and role is SpanRole.LLM_CALL: diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index 9eb9cb40caf..73448221acb 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -538,8 +538,12 @@ class OpenTelemetryV2(CustomLogger): # logger is a success/failure callback only, so ``pre_call`` never reaches it # and no carrier exists. The payload plus the request-level provider-handoff # stamp (``upstream_started``) is the affirmative signal of a real call; a - # gate rejection carries ``is_no_upstream_call`` and gets no span. - if carrier is None and (call.is_no_upstream_call or not call.upstream_started or call.payload is None): + # gate rejection carries ``is_no_upstream_call`` and gets no span. A cache + # hit never hands off to a provider, so ``cache_hit`` substitutes for the + # upstream stamp, since the payload proves the call completed. + if carrier is None and ( + call.is_no_upstream_call or call.payload is None or (not call.upstream_started and not call.cache_hit) + ): return None try: return self._finish_carrier(carrier, call, start_time, end_time) diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 0ac8b47e239..4e07d090edd 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -91,6 +91,8 @@ class GenAIMapper: f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: d.cost.margin_fixed_amount, f"{LiteLLM.COST_PREFIX}margin_percent": lambda d: d.cost.margin_percent, f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: d.cost.margin_total_amount, + LiteLLM.CACHE_HIT: lambda d: d.cache_hit, + LiteLLM.SAVED_CACHE_COST: lambda d: d.saved_cache_cost, LiteLLM.REQUEST_STREAMING: lambda d: d.is_streaming, LiteLLM.REQUEST_ROUTE: lambda d: d.request_route, LiteLLM.REQUEST_PURPOSE: lambda d: d.request_purpose, diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index cbbb1a9f6d4..a3151c24519 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -249,6 +249,10 @@ class LLMCallEvent: # ``InternalCallOrigin`` such as ``autorouter_classifier``), ``None`` for the # caller's own provider attempt. purpose: str | None + # True for a response served from the litellm response cache. A cache hit + # never hands off to a provider, so ``upstream_started`` stays False even + # though the call completed and is worth a span. + cache_hit: bool # A best-effort ``"{operation} {model}"`` name known at ``pre_call`` time. The # span is renamed from the typed payload at close (``finish_span``); this only # needs to be reasonable for a span that never gets closed (a leak). @@ -273,6 +277,7 @@ class LLMCallEvent: upstream_started=kwargs.get("api_call_start_time") is not None, upstream_start_seconds=_epoch_seconds(kwargs.get("api_call_start_time")), purpose=internal_call_origin(payload, kwargs), + cache_hit=bool(payload.get("cache_hit")) if payload else False, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown provisional_span_name=f"{operation.value} {model}".strip(), time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs), trace=trace, diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 101e6285946..5bca8e01aec 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -441,6 +441,8 @@ class LLMCallSpanData: embedding_output: EmbeddingOutput | None = None promoted_metadata: Mapping[str, str] = field(default_factory=lambda: _EMPTY_METADATA) routing_attributes: Mapping[str, RoutingAttributeValue] = field(default_factory=lambda: MappingProxyType({})) + cache_hit: bool | None = None + saved_cache_cost: float | None = None @classmethod def from_standard_logging_payload( @@ -456,6 +458,7 @@ class LLMCallSpanData: metadata_keys: tuple[str, ...] = (), ) -> LLMCallSpanData: params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {}) + cache_hit: Final = as_bool(payload.get("cache_hit")) is True # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown # The single parse of the request's metadata — the request-vs-provider # model split, the response model, api base, and identity all come from # here rather than being re-derived from the raw payload dicts. @@ -490,7 +493,7 @@ class LLMCallSpanData: error=_parse_error(payload), response_cost=as_float(payload.get("response_cost")), cost=LLMCost.from_breakdown(cast("Mapping[str, object] | None", payload.get("cost_breakdown"))), - server=ServerInfo.from_api_base(context.api_base), + server=None if cache_hit else ServerInfo.from_api_base(context.api_base), identity=context.identity, promoted_metadata=allowlisted_metadata(context.identity.metadata, metadata_keys), is_streaming=as_bool(payload.get("stream")), @@ -507,6 +510,8 @@ class LLMCallSpanData: session_id=session_id or None, embedding_output=embedding_output if capture_content else None, routing_attributes=routing_decision_attributes(decision), + cache_hit=cache_hit, + saved_cache_cost=as_float(payload.get("saved_cache_cost")) if cache_hit else None, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown ) diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 38cf7266593..b7f955a75ed 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -340,6 +340,8 @@ class LiteLLM: # semconv key for an MCP server's *name* (the convention uses ``server.address`` # for its network location), so it lives under the vendor namespace. MCP_SERVER_NAME: Final = "litellm.mcp.server.name" + CACHE_HIT: Final = "litellm.cache_hit" + SAVED_CACHE_COST: Final = "litellm.saved_cache_cost" class Metric: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 438b38f3f09..457a17ab8fe 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -13,6 +13,7 @@ import sys import time import traceback from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime as dt_object from functools import lru_cache from types import MappingProxyType, TracebackType @@ -558,6 +559,12 @@ def mask_api_base_credentials(api_base: str) -> str: return api_base[:key_end] + "*" * 5 + api_base[-4:] +@dataclass(frozen=True, slots=True) +class UncachedPrice: + cost: float | None + cost_breakdown: CostBreakdown | None + + def _timestamp_seconds(moment: object) -> float | None: if isinstance(moment, datetime.datetime): return moment.timestamp() @@ -811,6 +818,10 @@ class Logging(LiteLLMLoggingBaseClass): new_call.sync_streaming_chunks = [] return new_call + def reset_caching_details(self) -> None: + """Forget an earlier cache lookup's result, so a call that looks up again reports only its latest.""" + self.caching_details = None + def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" self.response_timing_metrics = dict(timing_metrics) @@ -1890,6 +1901,7 @@ class Logging(LiteLLMLoggingBaseClass): cache_hit: bool | None = None, litellm_model_name: str | None = None, router_model_id: str | None = None, + ignore_stamped_cost: bool = False, ) -> float | None: """ Calculate response cost using result + logging object variables. @@ -1925,7 +1937,9 @@ class Logging(LiteLLMLoggingBaseClass): ): hidden_params: Final = result_hidden_params if ( - "response_cost" in hidden_params and hidden_params["response_cost"] is not None + not ignore_stamped_cost + and "response_cost" in hidden_params + and hidden_params["response_cost"] is not None ): # use cost if already calculated self._record_zero_cost_diagnostic( priced_result, @@ -2051,6 +2065,20 @@ class Logging(LiteLLMLoggingBaseClass): _response_cost_calculator = response_cost_calculator + def uncached_price(self, result: object) -> UncachedPrice: + previous_cost_breakdown: Final = self.cost_breakdown + previous_billed_token_rates: Final = self.billed_token_rates + debug_key: Final = "response_cost_failure_debug_information" + had_failure_debug_info: Final = debug_key in self.model_call_details + try: + cost: Final = self.response_cost_calculator(result=result, cache_hit=False, ignore_stamped_cost=True) + return UncachedPrice(cost=cost, cost_breakdown=self.cost_breakdown) + finally: + self.cost_breakdown = previous_cost_breakdown + self.billed_token_rates = previous_billed_token_rates + if not had_failure_debug_info: + self.model_call_details.pop(debug_key, None) + def _record_zero_cost_diagnostic( self, result: object, @@ -6909,6 +6937,36 @@ def _autorouter_savings_for_payload( return None +_CACHE_HIT_ZEROED_AMOUNTS: Final = MappingProxyType( + { + "cache_read_cost": CostBreakdown(cache_read_cost=0.0), + "cache_creation_cost": CostBreakdown(cache_creation_cost=0.0), + "reasoning_cost": CostBreakdown(reasoning_cost=0.0), + "original_cost": CostBreakdown(original_cost=0.0), + "discount_amount": CostBreakdown(discount_amount=0.0), + "margin_fixed_amount": CostBreakdown(margin_fixed_amount=0.0), + "margin_total_amount": CostBreakdown(margin_total_amount=0.0), + } +) + + +def _cache_hit_cost_breakdown(served: CostBreakdown | None) -> CostBreakdown: + """Every amount at 0, keeping the pricing basis and rates a fresh call would have had, which savings reports read.""" + zero_costs: Final = CostBreakdown(input_cost=0.0, output_cost=0.0, total_cost=0.0, tool_usage_cost=0.0) + if served is None: + return zero_costs + zeroed_additional_costs: Final = ( + CostBreakdown(additional_costs=dict.fromkeys(served["additional_costs"], 0.0)) + if "additional_costs" in served + else CostBreakdown() + ) + return functools.reduce( + lambda breakdown, zeroed: breakdown | zeroed, + (_CACHE_HIT_ZEROED_AMOUNTS[key] for key in served if key in _CACHE_HIT_ZEROED_AMOUNTS), + served | zero_costs | zeroed_additional_costs, + ) + + def get_standard_logging_object_payload( kwargs: dict | None, init_response_obj: object, @@ -6992,17 +7050,10 @@ def get_standard_logging_object_payload( "user", None ) # maintain backwards compatibility with old request body check - saved_cache_cost: float = 0.0 if cache_hit is True or served_from_cache(init_response_obj): id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id - if cache_hit is True: - saved_cache_cost = ( - logging_obj.response_cost_calculator( - result=init_response_obj, - cache_hit=False, - ) - or 0.0 - ) + uncached_price: Final = logging_obj.uncached_price(result=init_response_obj) if cache_hit is True else None + saved_cache_cost: Final = (uncached_price.cost if uncached_price is not None else None) or 0.0 ## Get model cost information ## base_model = get_base_model_from_metadata(model_call_details=kwargs) @@ -7019,7 +7070,7 @@ def get_standard_logging_object_payload( # clean up litellm hidden params clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params) - if clean_hidden_params["response_cost"] is None and raw_response_cost is not None: + if cache_hit is True or (clean_hidden_params["response_cost"] is None and raw_response_cost is not None): clean_hidden_params["response_cost"] = llm_response_cost if clean_hidden_params["litellm_overhead_time_ms"] is None and status == "success": # /v1/messages dict results and the bridge stream wrappers keep it on the logging object; @@ -7081,7 +7132,14 @@ def get_standard_logging_object_payload( ): model_name = response_model_name - request_cost_breakdown: Final = cost_breakdown_with_guardrail(logging_obj.cost_breakdown, guardrail_cost) + request_cost_breakdown: Final = cost_breakdown_with_guardrail( + ( + _cache_hit_cost_breakdown(uncached_price.cost_breakdown) + if uncached_price is not None + else logging_obj.cost_breakdown + ), + guardrail_cost, + ) captured_baseline: Final = logging_obj.baseline_observation autorouter_savings: Final = ( None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 74ee87c772b..e52f11641b0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1663,19 +1663,26 @@ def _totals_to_zero(response_cost: float | str | None) -> bool: return False +def _is_cache_hit(litellm_logging_obj: LiteLLMLoggingObj) -> bool: + caching_details: Final = litellm_logging_obj.caching_details + return caching_details is not None and caching_details.get("cache_hit") is True + + def _get_cost_breakdown_from_logging_obj( litellm_logging_obj: LiteLLMLoggingObj | None, response_cost: float | str | None = None, ) -> CostBreakdownHeaderValues: """Extract discount, margin, and per-component cost information from logging object's cost breakdown. - A non-inference call that priced at zero never records a breakdown, so its components are - reported as zero here. Any such call that did price normally (retrieving a background response, - and the cost poller's read of one) reports the breakdown it stored, or nothing at all when the - breakdown has not landed yet. + A cache hit reports every component as zero. A non-inference call that priced at zero never + records a breakdown, so its components are reported as zero too. Any such call that did price + normally (retrieving a background response, and the cost poller's read of one) reports the + breakdown it stored, or nothing at all when the breakdown has not landed yet. """ if not litellm_logging_obj or not hasattr(litellm_logging_obj, "cost_breakdown"): return CostBreakdownHeaderValues() + if _is_cache_hit(litellm_logging_obj): + return _ZERO_COST_BREAKDOWN cost_breakdown: Final = litellm_logging_obj.cost_breakdown if not cost_breakdown: @@ -1705,6 +1712,13 @@ def _header_value(value: float | None) -> str | None: return str(value) if value is not None else None +def _response_cost_header_value(hidden_params: Mapping[str, object]) -> float | str: + if hidden_params.get("cache_hit") is True: + return 0.0 + response_cost: Final = hidden_params.get("response_cost") + return response_cost if isinstance(response_cost, (int, float, str)) and response_cost else "" + + def cost_headers( response_cost: float | str | None, litellm_logging_obj: LiteLLMLoggingObj | None ) -> Mapping[str, str | None]: @@ -2084,18 +2098,12 @@ class ProxyBaseLLMRequestProcessing: Build LiteLLM proxy response headers for routes that call the LLM directly (e.g. Google native :generateContent) instead of base_process_llm_request. """ - if isinstance(response, dict): - hidden_params = get_hidden_params_dict(response) - else: - hidden_params = getattr(response, "_hidden_params", None) or {} - if not isinstance(hidden_params, dict): - hidden_params = {} - + hidden_params: Final = get_hidden_params_dict(response) model_id: Final = ProxyBaseLLMRequestProcessing._get_model_id_from_response(hidden_params, request_data) cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" - response_cost: Final = hidden_params.get("response_cost", None) or "" + response_cost: Final = _response_cost_header_value(hidden_params) fastest_response_batch_completion: Final = hidden_params.get("fastest_response_batch_completion", None) additional_headers: Final = hidden_params.get("additional_headers", {}) or {} @@ -2640,7 +2648,7 @@ class ProxyBaseLLMRequestProcessing: cache_key=hidden_params.get("cache_key") or "", api_base=hidden_params.get("api_base") or "", version=version, - response_cost=hidden_params.get("response_cost") or "", + response_cost=_response_cost_header_value(hidden_params), model_region=getattr(user_api_key_dict, "allowed_model_region", ""), fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"), request_data=self.data, @@ -2685,13 +2693,17 @@ class ProxyBaseLLMRequestProcessing: hold the attribute at all, and Google :generateContent carries ``_hidden_params`` but no synchronously-populated ``response_cost``. In both cases the cost is read back from the logging object instead, recomputing from - the same calculator only when it has not been stored yet. + the same calculator only when it has not been stored yet. A response served + from the cache costs 0, which the calculator cannot tell from a fresh call, and any + other zero or unknown cost comes back empty so no cost header is sent. """ + if _is_cache_hit(logging_obj): + return 0.0 stored_cost: Final = logging_obj.model_call_details.get("response_cost") if isinstance(stored_cost, (int, float)): - return float(stored_cost) + return float(stored_cost) or "" recomputed_cost: Final = logging_obj.response_cost_calculator(result=response) - return recomputed_cost if isinstance(recomputed_cost, (int, float)) else "" + return recomputed_cost or "" if isinstance(recomputed_cost, (int, float)) else "" def _debug_log_request_payload(self) -> None: """Log request payload at DEBUG level, truncating if too large.""" @@ -2900,7 +2912,7 @@ class ProxyBaseLLMRequestProcessing: cache_key, api_base, response_cost = ( hidden_params.get("cache_key", None) or "", hidden_params.get("api_base", None) or "", - hidden_params.get("response_cost", None) or "", + _response_cost_header_value(hidden_params), ) fastest_response_batch_completion, additional_headers = ( hidden_params.get("fastest_response_batch_completion", None), @@ -3170,7 +3182,7 @@ class ProxyBaseLLMRequestProcessing: recover_response_cost: Final = not response_cost and hidden_params.get("response_cost") is None computed_cost_for_headers: Final = ( - self._response_cost_from_logging_obj(response=response, logging_obj=logging_obj) or "" + self._response_cost_from_logging_obj(response=response, logging_obj=logging_obj) if recover_response_cost else response_cost ) diff --git a/tests/integration/observability/test_otel_cache_hit_cost.py b/tests/integration/observability/test_otel_cache_hit_cost.py new file mode 100644 index 00000000000..76c21b529ef --- /dev/null +++ b/tests/integration/observability/test_otel_cache_hit_cost.py @@ -0,0 +1,370 @@ +import json +import os +import re +import uuid +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, list_value, object_value +from integration._support.database import read_rows +from integration._support.otlp_sink import Span, owned_sinks, recorded_spans +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +INPUT_PRICE: Final = 0.001 +OUTPUT_PRICE: Final = 0.002 +PROMPT_TOKENS: Final = 7 +COMPLETION_TOKENS: Final = 2 +CHAT_PRICE: Final = PROMPT_TOKENS * INPUT_PRICE + COMPLETION_TOKENS * OUTPUT_PRICE +EMBEDDING_PRICE: Final = PROMPT_TOKENS * INPUT_PRICE +CLIENT: Final = 3 +COST_HEADERS: Final = ("x-litellm-response-cost", "x-litellm-response-cost-input", "x-litellm-response-cost-output") +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +MARKER: Final = re.compile(rb"cache-hit-cost-[0-9a-f]{32}") + + +def _sse(frames: tuple[Mapping[str, JsonValue], ...]) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=(*(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames), b"data: [DONE]\n\n"), + ) + + +def _answer(marker: str) -> str: + return "answer to " + marker + + +def _embedding_value(marker: str) -> float: + return float(int(marker[-12:], 16)) + + +def _chat_reply(stream: bool, marker: str) -> Reply: + identity: Final = "chatcmpl-" + uuid.uuid4().hex + message: Final = {"role": "assistant", "content": _answer(marker)} + usage: Final = { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + } + if not stream: + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": message, "finish_reason": "stop"}], + "usage": usage, + } + ).encode() + ) + chunk: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"} + return _sse( + ( + {**chunk, "choices": [{"index": 0, "delta": message}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**chunk, "choices": [], "usage": usage}, + ) + ) + + +def _responses_reply(marker: str) -> Reply: + identity: Final = "resp_" + uuid.uuid4().hex + return Reply( + body=json.dumps( + { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": "msg_" + identity, + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": _answer(marker), "annotations": []}], + } + ], + "usage": { + "input_tokens": PROMPT_TOKENS, + "output_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + } + ).encode() + ) + + +def _embedding_reply(marker: str) -> Reply: + return Reply( + body=json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [_embedding_value(marker), 0.5, 0.75]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": PROMPT_TOKENS, "total_tokens": PROMPT_TOKENS}, + } + ).encode() + ) + + +def _provider(request: Request) -> Reply: + found: Final = MARKER.search(request.body) + assert found is not None, request.body + marker: Final = found.group().decode() + if request.target.endswith("/embeddings"): + return _embedding_reply(marker) + if request.target.endswith("/responses"): + return _responses_reply(marker) + return _chat_reply(JSON_OBJECT.validate_json(request.body).get("stream") is True, marker) + + +@dataclass(frozen=True, slots=True) +class Route: + path: str + stream: bool + price: float + body: Callable[[str, str], Mapping[str, JsonValue]] + cached: Callable[[str], str] = _answer + + +ROUTES: Final[Mapping[str, Route]] = MappingProxyType( + { + "chat": Route( + "/v1/chat/completions", + False, + CHAT_PRICE, + lambda model, marker: {"model": model, "messages": [{"role": "user", "content": marker}]}, + ), + "chat_stream": Route( + "/v1/chat/completions", + True, + CHAT_PRICE, + lambda model, marker: { + "model": model, + "messages": [{"role": "user", "content": marker}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + ), + "messages": Route( + "/v1/messages", + False, + CHAT_PRICE, + lambda model, marker: {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + ), + "responses": Route("/v1/responses", False, CHAT_PRICE, lambda model, marker: {"model": model, "input": marker}), + "embeddings": Route( + "/v1/embeddings", + False, + EMBEDDING_PRICE, + lambda model, marker: {"model": model, "input": marker}, + lambda marker: repr(_embedding_value(marker)), + ), + } +) + + +@dataclass(frozen=True, slots=True) +class Served: + call_id: str + headers: httpx.Headers + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + sink: str + chat_model: str + embedding_model: str + upstream: Wire + otel_v2: bool + + def send(self, route: str, marker: str) -> Served: + chosen: Final = ROUTES[route] + model: Final = self.embedding_model if route == "embeddings" else self.chat_model + with self.proxy.client.stream( + "POST", + chosen.path, + json=dict(chosen.body(model, marker)), + headers={"Authorization": f"Bearer {self.proxy.key}"}, + ) as response: + body: Final = response.read() + assert response.status_code == 200, body + return Served(response.headers["x-litellm-call-id"], response.headers) + + def upstream_calls(self, marker: str) -> int: + return sum(1 for request in self.upstream.drain() if marker.encode() in request.body) + + def cost_span(self, call_id: str) -> Span: + def matching() -> tuple[Span, ...]: + return tuple(span for span in recorded_spans(self.sink)[1] if self._is_cost_span(span, call_id)) + + found: Final = eventually(matching, lambda spans: len(spans) >= 1, seconds=30) + assert len(found) == 1, found + return found[0] + + def _is_cost_span(self, span: Span, call_id: str) -> bool: + if span["attributes"].get("litellm.call_id") != call_id: + return False + if self.otel_v2: + return span["kind"] == CLIENT and "gen_ai.operation.name" in span["attributes"] + return span["name"] == "litellm_request" + + def total_cost(self, span: Span) -> float: + return _number(span["attributes"].get("litellm.cost.total" if self.otel_v2 else "gen_ai.cost.total_cost")) + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float, str)) and not isinstance(value, bool), value + return float(value) + + +def _config(provider: Wire, sink: str, chat_model: str, embedding_model: str) -> dict[str, JsonValue]: + config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + priced: Final = {"input_cost_per_token": INPUT_PRICE, "output_cost_per_token": OUTPUT_PRICE} + return { + **config, + "model_list": [ + *list_value(config.get("model_list") or []), + { + "model_name": chat_model, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": provider.url + "/v1", + **priced, + }, + }, + { + "model_name": embedding_model, + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "integration-provider-key", + "api_base": provider.url + "/v1", + **priced, + }, + }, + ], + "litellm_settings": {**object_value(config["litellm_settings"]), "callbacks": ["otel"]}, + "callback_settings": {"otel": {"exporter": "http/json", "endpoint": sink, "use_simple_processor": True}}, + } + + +@pytest.fixture(scope="module") +def provider() -> Iterator[Wire]: + with wire_server(_provider) as wire: + yield wire + + +def _rig(provider: Wire, directory: Path, otel_v2: bool) -> Iterator[Rig]: + chat_model: Final = "cachecost-chat-" + uuid.uuid4().hex[:12] + embedding_model: Final = "cachecost-embed-" + uuid.uuid4().hex[:12] + with owned_sinks(directory / "sinks") as sinks: + path: Final = directory / f"cache-hit-cost-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(_config(provider, sinks.operator, chat_model, embedding_model))) + v1_exporter: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": sinks.operator} + exporter: Final = {"LITELLM_OTEL_V2": "1"} if otel_v2 else v1_exporter + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + directory, + exporter, + config=path, + remove_environment=() if otel_v2 else ("LITELLM_OTEL_V2",), + workers=2, + ) as owned, + ): + yield Rig(owned.gateway, sinks.operator, chat_model, embedding_model, provider, otel_v2) + + +@pytest.fixture(scope="module") +def v1(provider: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + yield from _rig(provider, tmp_path_factory.mktemp("cache-hit-cost-v1"), otel_v2=False) + + +@pytest.fixture(scope="module") +def v2(provider: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + yield from _rig(provider, tmp_path_factory.mktemp("cache-hit-cost-v2"), otel_v2=True) + + +def _spend_row(call_id: str) -> Mapping[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, cache_hit, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,) + ), + lambda found: len(found) == 1, + seconds=70, + ) + return rows[0] + + +def _breakdown_total(row: Mapping[str, JsonValue]) -> float: + return _number(object_value(object_value(row["metadata"])["cost_breakdown"])["total_cost"]) + + +def _cache_holds(needle: str) -> bool: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as store: + return any(needle.encode() in _bytes(store.get(key)) for key in store.scan_iter(count=1000, _type="STRING")) + + +def _bytes(value: object) -> bytes: + return value if isinstance(value, bytes) else b"" + + +def _miss_then_hit(rig: Rig, route: str) -> tuple[Served, Served]: + marker: Final = "cache-hit-cost-" + uuid.uuid4().hex + miss: Final = rig.send(route, marker) + eventually(lambda: _cache_holds(ROUTES[route].cached(marker)), bool, seconds=30) + hit: Final = rig.send(route, marker) + assert rig.upstream_calls(marker) == 1 + return miss, hit + + +@pytest.mark.parametrize("route", tuple(ROUTES)) +def test_a_v1_cache_hit_exports_zero_cost_and_answers_zero_cost_headers(v1: Rig, route: str) -> None: + price: Final = ROUTES[route].price + miss, hit = _miss_then_hit(v1, route) + + hit_span: Final = v1.cost_span(hit.call_id) + assert v1.total_cost(hit_span) == 0, hit_span + assert hit_span["attributes"].get("litellm.cache_hit") is True, hit_span + assert _number(hit_span["attributes"].get("litellm.saved_cache_cost")) == pytest.approx(price), hit_span + + assert {name: hit.headers.get(name) for name in COST_HEADERS} == dict.fromkeys(COST_HEADERS, "0.0"), hit.headers + + hit_row: Final = _spend_row(hit.call_id) + assert (_number(hit_row["spend"]), hit_row["cache_hit"]) == (0, "True"), hit_row + assert _breakdown_total(hit_row) == 0, hit_row + assert _number(_spend_row(miss.call_id)["spend"]) == pytest.approx(price) + + miss_span: Final = v1.cost_span(miss.call_id) + assert v1.total_cost(miss_span) == pytest.approx(price), miss_span + assert miss_span["attributes"].get("litellm.cache_hit") is False, miss_span + + +@pytest.mark.parametrize("route", tuple(ROUTES)) +def test_a_v2_cache_hit_exports_an_llm_span_at_zero_cost(v2: Rig, route: str) -> None: + price: Final = ROUTES[route].price + miss, hit = _miss_then_hit(v2, route) + + hit_span: Final = v2.cost_span(hit.call_id) + assert v2.total_cost(hit_span) == 0, hit_span + assert hit_span["attributes"].get("litellm.cache_hit") is True, hit_span + assert _number(hit_span["attributes"].get("litellm.saved_cache_cost")) == pytest.approx(price), hit_span + + miss_span: Final = v2.cost_span(miss.call_id) + assert v2.total_cost(miss_span) == pytest.approx(price), miss_span + assert miss_span["attributes"].get("litellm.cache_hit") is False, miss_span diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 7d6b1e15fc6..388e4c42ea0 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -15,6 +15,7 @@ from pydantic import TypeAdapter import litellm from litellm import aembedding, completion, embedding, aresponses, responses +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.caching.caching import Cache from litellm.responses.streaming_iterator import CachedResponsesAPIStreamingIterator @@ -303,6 +304,93 @@ async def test_cache_hit_reports_zero_response_cost(copilot_response_cache: Cach assert second._hidden_params["response_cost"] == 0 +@pytest.mark.asyncio +async def test_async_cache_hits_never_use_up_a_deployment_tpm_limit(copilot_response_cache: Cache) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "tpm-limited", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + "mock_response": "cached answer", + "tpm": 100, + }, + } + ], + routing_strategy="usage-based-routing-v2", + num_retries=0, + ) + messages: Final = [{"role": "user", "content": "tpm limit and cache hits"}] + + async def call_and_settle() -> litellm.ModelResponse: + response: Final = await router.acompletion(model="tpm-limited", messages=messages, caching=True) + await asyncio.gather(*_PENDING_CACHE_WRITES) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + return response + + first: Final = await call_and_settle() + hits: Final = [await call_and_settle() for _ in range(8)] + + assert first.usage.total_tokens * len(hits) > 100 + assert all(hit._hidden_params["cache_hit"] is True for hit in hits) + + +def _logging_obj_left_with_a_cache_hit(messages: list[dict[str, str]]) -> LiteLLMLogging: + logging_obj: Final = LiteLLMLogging( + litellm_call_id=str(uuid.uuid4()), + call_type=CallTypes.acompletion.value, + model="gpt-4o-mini", + messages=messages, + function_id=str(uuid.uuid4()), + stream=False, + start_time=datetime.now(), + ) + logging_obj.caching_details = {"cache_hit": True, "cache_duration_ms": 1.0} + return logging_obj + + +@pytest.mark.asyncio +async def test_async_cache_miss_clears_the_hit_an_earlier_lookup_left_on_the_call( + copilot_response_cache: Cache, +) -> None: + messages: Final = [{"role": "user", "content": f"never cached {uuid.uuid4()}"}] + logging_obj: Final = _logging_obj_left_with_a_cache_hit(messages) + handler: Final = LLMCachingHandler( + original_function=litellm.acompletion, request_kwargs={}, start_time=datetime.now() + ) + + response: Final = await handler.async_get_cache( + model="gpt-4o-mini", + original_function=litellm.acompletion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.acompletion.value, + kwargs={"model": "gpt-4o-mini", "messages": messages, "caching": True}, + ) + + assert response is not None and response.cached_result is None + assert logging_obj.caching_details is None + + +def test_sync_cache_miss_clears_the_hit_an_earlier_lookup_left_on_the_call(copilot_response_cache: Cache) -> None: + messages: Final = [{"role": "user", "content": f"never cached {uuid.uuid4()}"}] + logging_obj: Final = _logging_obj_left_with_a_cache_hit(messages) + handler: Final = LLMCachingHandler(original_function=completion, request_kwargs={}, start_time=datetime.now()) + + response: Final = handler.sync_get_cache( + model="gpt-4o-mini", + original_function=completion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.completion.value, + kwargs={"model": "gpt-4o-mini", "messages": messages, "caching": True}, + ) + + assert response.cached_result is None + assert logging_obj.caching_details is None + + chat_completion_response = litellm.ModelResponse( id=str(uuid.uuid4()), choices=[ diff --git a/tests/unit/integrations/otel/test_otel_v2_components.py b/tests/unit/integrations/otel/test_otel_v2_components.py index bc230a8bcb2..3f4389765d3 100644 --- a/tests/unit/integrations/otel/test_otel_v2_components.py +++ b/tests/unit/integrations/otel/test_otel_v2_components.py @@ -2,6 +2,7 @@ baggage helpers, metrics, the typed coercion helpers, mapper branches, span-name builders, and the registry validator's failure paths. Needs the OTel SDK.""" +from typing import Final import contextlib import json import threading @@ -358,6 +359,18 @@ def test_genai_mapper_cost_breakdown_absent(): assert not any(k.startswith(LiteLLM.COST_PREFIX) and k != f"{LiteLLM.COST_PREFIX}total" for k in attrs) +def test_genai_mapper_cache_hit_attrs(): + from litellm.integrations.otel.model.semconv import LiteLLM + + attrs: Final = GenAIMapper().map(replace(_full_llm_call(), cache_hit=True, saved_cache_cost=0.002)) + assert attrs[LiteLLM.CACHE_HIT] is True + assert attrs[LiteLLM.SAVED_CACHE_COST] == 0.002 + + uncached: Final = GenAIMapper().map(_full_llm_call()) + assert LiteLLM.CACHE_HIT not in uncached + assert LiteLLM.SAVED_CACHE_COST not in uncached + + def test_llm_cost_from_breakdown_maps_costbreakdown_keys(): cost = LLMCost.from_breakdown( { diff --git a/tests/unit/integrations/otel/test_otel_v2_emitter.py b/tests/unit/integrations/otel/test_otel_v2_emitter.py index 949eb97bccd..0905e9e7c30 100644 --- a/tests/unit/integrations/otel/test_otel_v2_emitter.py +++ b/tests/unit/integrations/otel/test_otel_v2_emitter.py @@ -2,6 +2,7 @@ legacy dual-emit, hierarchy, error status, and idempotency. Needs the OTel SDK.""" import json +from typing import Final import pytest @@ -642,25 +643,71 @@ def test_preset_message_key_the_fit_sheds_still_never_overflows_the_span(): assert len(finished.attributes) <= SpanLimits().max_span_attributes -def test_attribute_budget_counts_only_keys_the_fit_does_not_overwrite(): +def test_preset_keys_the_mappers_overwrite_shed_no_message_attribute(monkeypatch): """Pre-set attributes the mapped set overwrites consume no slot against the span limit. A boundary-opened LLM-call span already carries a few stamped attributes, - most of which the mappers re-emit under the same keys. Charging those - against the budget reserves slots the fit can never spend and sheds - indexed message attributes for nothing. + most of which the mappers re-emit under the same keys, so a span sized to + exactly its mapped attributes keeps every indexed message. """ + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", "1000") + full: Final = _conversation_span(["genai", "openinference"], _conversation_payload(6)).attributes + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", str(len(full))) + cfg: Final = OpenTelemetryV2Config( + exporter="in_memory", + legacy_compat=False, + mapper_names=["genai", "openinference"], + capture_message_content="span_only", + ) + provider, exporter = providers.in_memory_provider(cfg) + engine: Final = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg) + span: Final = engine.start_span(SpanRole.LLM_CALL, "chat gpt-4o") + span.set_attribute("gen_ai.request.model", "gpt-4o") + span.set_attribute("llm.input_messages.1.message.role", "stale-role") + engine.finish_span( + SpanRole.LLM_CALL, + span, + LLMCallSpanData.from_standard_logging_payload(_conversation_payload(6), capture_content=True), + ) + (finished,) = exporter.get_finished_spans() + assert finished.dropped_attributes == 0 + assert dict(finished.attributes) == dict(full) + + +def test_a_preset_message_key_the_fit_sheds_never_overflows_an_exact_limit(monkeypatch): + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", "1000") + full: Final = _conversation_span(["genai", "openinference"], _conversation_payload(6)).attributes + limit: Final = len(full) - 2 + monkeypatch.setenv("OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT", str(limit)) + cfg: Final = OpenTelemetryV2Config( + exporter="in_memory", + legacy_compat=False, + mapper_names=["genai", "openinference"], + capture_message_content="span_only", + ) + provider, exporter = providers.in_memory_provider(cfg) + engine: Final = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg) + span: Final = engine.start_span(SpanRole.LLM_CALL, "chat gpt-4o") + span.set_attribute("llm.input_messages.1.message.role", "stale-role") + engine.finish_span( + SpanRole.LLM_CALL, + span, + LLMCallSpanData.from_standard_logging_payload(_conversation_payload(6), capture_content=True), + ) + (finished,) = exporter.get_finished_spans() + assert finished.dropped_attributes == 0 + assert len(finished.attributes) <= limit + assert finished.attributes.get("llm.input_messages.1.message.role") != "stale-role" + + +def test_attribute_budget_charges_every_key_the_span_already_carries(): cfg = OpenTelemetryV2Config(exporter="in_memory") provider, _exporter = providers.in_memory_provider(cfg) span = providers.get_tracer(provider, "litellm-test").start_span("s") span.set_attribute("gen_ai.request.model", "gpt-4o") span.set_attribute("litellm.metadata.user_api_key_alias", "edge-key") try: - assert attribute_budget(span, 0) == SpanLimits().max_span_attributes - 2 - assert ( - attribute_budget(span, 0, frozenset({"gen_ai.request.model"})) - == SpanLimits().max_span_attributes - 1 - ) + assert attribute_budget(span, 1) == SpanLimits().max_span_attributes - 3 finally: span.end() diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index 6a9a963948e..10e12981ec1 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -494,6 +494,49 @@ def test_no_span_when_pre_call_never_ran(): assert exporter.get_finished_spans() == () # no phantom LLM span +def test_cache_hit_without_pre_call_emits_zero_cost_span(): + """A response served from the litellm cache never hands off to a provider, + so ``pre_call`` never runs and ``api_call_start_time`` is never stamped, + but the completed call's payload still warrants a span, carrying the cache + attributes rather than being dropped like a gate rejection.""" + logger, exporter = _logger() + payload: Final = _payload( + cache_hit=True, + response_cost=0.0, + saved_cache_cost=0.002, + cost_breakdown={"input_cost": 0.0, "output_cost": 0.0, "total_cost": 0.0, "tool_usage_cost": 0.0}, + ) + asyncio.run( + logger.async_log_success_event(_kwargs(payload=payload), None, None, None) + ) + (span,) = exporter.get_finished_spans() + assert span.attributes[LiteLLM.CACHE_HIT] is True + assert span.attributes[LiteLLM.SAVED_CACHE_COST] == 0.002 + assert { + name: span.attributes[f"{LiteLLM.COST_PREFIX}{name}"] for name in ("total", "input", "output", "tool_usage") + } == {"total": 0.0, "input": 0.0, "output": 0.0, "tool_usage": 0.0} + assert "server.address" not in span.attributes + + +def test_cache_miss_span_reports_cache_hit_false(): + logger, exporter = _logger() + _emit_llm(logger, _kwargs(payload=_payload(cache_hit=False, saved_cache_cost=0.0))) + (span,) = exporter.get_finished_spans() + assert span.attributes[LiteLLM.CACHE_HIT] is False + assert LiteLLM.SAVED_CACHE_COST not in span.attributes + + +def test_cache_miss_without_pre_call_still_emits_no_span(): + """The cache-hit relaxation is not a general loosening: a success callback + with no ``pre_call`` carrier and no cache hit still emits nothing.""" + logger, exporter = _logger() + payload: Final = _payload(cache_hit=False, response_cost=0.002) + asyncio.run( + logger.async_log_success_event(_kwargs(payload=payload), None, None, None) + ) + assert exporter.get_finished_spans() == () + + def test_real_llm_failure_still_emitted(): """A genuine LLM failure: ``pre_call`` ran (the call was attempted), so the CLIENT span is opened at the boundary and closed ERROR.""" diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 9ba09085ce8..b40c1a96500 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -478,6 +478,75 @@ class TestOpenTelemetryCostBreakdown(unittest.TestCase): assert ("gen_ai.cost.tool_usage_cost", 0.0001) not in call_args_list assert ("gen_ai.cost.original_cost", 0.004) not in call_args_list + def test_cache_hit_span_carries_cache_indicator_and_saved_cost(self): + """ + A cache-hit span reports litellm.cache_hit and litellm.saved_cache_cost so cost + attribution can tell a cache-served response apart from a billed call. + """ + otel: Final = OpenTelemetry() + mock_span: Final = MagicMock() + + kwargs: Final = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + "cache_hit": True, + "saved_cache_cost": 0.003, + "cost_breakdown": {"input_cost": 0.0, "output_cost": 0.0, "total_cost": 0.0, "tool_usage_cost": 0.0}, + }, + } + + response_obj: Final = { + "id": "test-response-id", + "model": "gpt-4", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + + otel.set_attributes(span=mock_span, kwargs=kwargs, response_obj=response_obj) + + mock_span.set_attribute.assert_any_call("litellm.cache_hit", True) + mock_span.set_attribute.assert_any_call("litellm.saved_cache_cost", 0.003) + mock_span.set_attribute.assert_any_call("gen_ai.cost.total_cost", 0.0) + mock_span.set_attribute.assert_any_call("gen_ai.cost.input_cost", 0.0) + mock_span.set_attribute.assert_any_call("gen_ai.cost.output_cost", 0.0) + + def test_non_cache_hit_span_reports_cache_hit_false(self): + otel: Final = OpenTelemetry() + mock_span: Final = MagicMock() + + kwargs: Final = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + "cache_hit": False, + }, + } + + response_obj: Final = { + "id": "test-response-id", + "model": "gpt-4", + "choices": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + + otel.set_attributes(span=mock_span, kwargs=kwargs, response_obj=response_obj) + + call_args_list: Final = [call[0] for call in mock_span.set_attribute.call_args_list] + assert ("litellm.cache_hit", False) in call_args_list + assert all(key != "litellm.saved_cache_cost" for key, _value in call_args_list) + assert all(not key.startswith("gen_ai.cost.") for key, _value in call_args_list) + class TestOpenTelemetryProviderInitialization(unittest.TestCase): """Test suite for verifying provider initialization respects existing providers""" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f314c45dae2..3d164b59bfe 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -40,6 +40,7 @@ from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, _extract_response_obj_and_hidden_params, _get_status_fields, + get_standard_logging_object_payload, set_callbacks, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -52,6 +53,7 @@ from litellm.proxy.utils import InternalUsageCache from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, + CostBreakdown, ImageResponse, LiteLLMRealtimeStreamLoggingObject, ModelResponse, @@ -7400,6 +7402,75 @@ def test_payload_guardrail_cost_merges_into_existing_cost_breakdown(logging_obj) assert logging_obj.cost_breakdown["total_cost"] == pytest.approx(0.0000429) +def test_cache_hit_payload_keeps_zero_components_beside_guardrail_cost(logging_obj): + logging_obj.set_cost_breakdown( + input_cost=0.00003, + output_cost=0.0000129, + total_cost=0.0000429, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) + payload: Final = _build_success_payload(logging_obj, {**_guardrail_kwargs(response_cost=0.0), "cache_hit": True}) + + assert payload is not None + assert payload["response_cost"] == pytest.approx(0.0003) + assert payload["cost_breakdown"] == { + "input_cost": 0.0, + "output_cost": 0.0, + "tool_usage_cost": 0.0, + "service_tier": None, + "data_residency": None, + "vertex_location": None, + "guardrail_cost": pytest.approx(0.0003), + "total_cost": pytest.approx(0.0003), + } + + +def test_cache_hit_payload_keeps_the_served_pricing_basis_with_every_amount_at_zero(logging_obj): + logging_obj.set_cost_breakdown( + input_cost=0.00003, + output_cost=0.0000129, + total_cost=0.0000474, + cost_for_built_in_tools_cost_usd_dollar=0.0000011, + additional_costs={"azure_model_router_flat_cost": 0.00014}, + original_cost=0.00005, + discount_percent=0.1, + discount_amount=0.000005, + margin_percent=0.05, + margin_fixed_amount=0.000001, + margin_total_amount=0.0000035, + cache_read_cost=0.000002, + cache_creation_cost=0.000003, + reasoning_cost=0.000004, + service_tier="priority", + data_residency="eu", + vertex_location="europe-west4", + ) + kwargs: Final = {"litellm_call_id": "hit-basis-call", "model": "gpt-4o", "messages": [], "cache_hit": True} + + payload: Final = _build_success_payload(logging_obj, kwargs) + + assert payload is not None + assert payload["cost_breakdown"] == { + "input_cost": 0.0, + "output_cost": 0.0, + "total_cost": 0.0, + "tool_usage_cost": 0.0, + "additional_costs": {"azure_model_router_flat_cost": 0.0}, + "original_cost": 0.0, + "discount_percent": 0.1, + "discount_amount": 0.0, + "margin_percent": 0.05, + "margin_fixed_amount": 0.0, + "margin_total_amount": 0.0, + "cache_read_cost": 0.0, + "cache_creation_cost": 0.0, + "reasoning_cost": 0.0, + "service_tier": "priority", + "data_residency": "eu", + "vertex_location": "europe-west4", + } + + def test_payload_without_guardrail_cost_is_unchanged(logging_obj): kwargs = { "litellm_call_id": "no-guardrail-call", @@ -10187,6 +10258,245 @@ def test_signoz_dispatch_requires_an_endpoint(monkeypatch): is_otel_v2_enabled.cache_clear() +def _lit8679_priced_logging_obj(deployment_id: str, litellm_call_id: str) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model="gpt-5.4-nano", + messages=[{"role": "user", "content": "Hi"}], + stream=False, + call_type="completion", + start_time=time.time(), + litellm_call_id=litellm_call_id, + function_id="fn", + ) + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + user="", + optional_params={}, + litellm_params={"metadata": {"model_group": "lit8679-group", "model_info": {"id": deployment_id}}}, + custom_llm_provider="openai", + ) + return logging_obj + + +def test_uncached_price_recomputes_price_past_a_zero_stamp(): + """A probe for the price a cache-served response would have had ignores the 0.0 stamped on it.""" + deployment_id: Final = "lit8679-uncached-cost-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-uncached") + oracle_logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-oracle") + cached_response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + cached_response._hidden_params = {"cache_hit": True, "cache_key": "lit8679-key", "response_cost": 0.0} + fresh_response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + cost: Final = logging_obj.uncached_price(result=cached_response).cost + uncached_price: Final = oracle_logging_obj._response_cost_calculator(result=fresh_response) + + assert cost == uncached_price + assert cost is not None and cost > 0.0 + assert logging_obj.cost_breakdown is None + finally: + litellm.model_cost.pop(deployment_id, None) + + +def test_uncached_price_leaves_no_pricing_state_on_the_call(): + """The uncached price probe leaves no cost breakdown or billed rates behind on the cache-hit call.""" + deployment_id: Final = "lit8679-probe-state-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-probe-state") + cached_response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + cached_response._hidden_params = {"cache_hit": True, "cache_key": "lit8679-key", "response_cost": 0.0} + assert (logging_obj.cost_breakdown, logging_obj.billed_token_rates) == (None, None) + + cost: Final = logging_obj.uncached_price(result=cached_response).cost + + assert cost is not None and cost > 0.0 + assert (logging_obj.cost_breakdown, logging_obj.billed_token_rates) == (None, None) + finally: + litellm.model_cost.pop(deployment_id, None) + + +def _assert_zeroed_copy(hit_breakdown: CostBreakdown | None, fresh_breakdown: CostBreakdown) -> None: + basis_and_rates: Final = ("service_tier", "data_residency", "vertex_location", "discount_percent", "margin_percent") + assert hit_breakdown is not None and hit_breakdown.keys() == fresh_breakdown.keys() + assert all(hit_breakdown[key] == fresh_breakdown[key] for key in basis_and_rates if key in hit_breakdown) + assert all(value == 0.0 for key, value in hit_breakdown.items() if key not in basis_and_rates) + + +def test_payload_saved_cache_cost_keeps_uncached_price_when_result_stamped_zero(): + """A cache-hit payload still reports what the call would have cost, not the 0.0 stamped on the served response.""" + deployment_id: Final = "lit8679-saved-cost-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-payload") + oracle_logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-payload-oracle") + response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + response._hidden_params = {"cache_hit": True, "cache_key": "lit8679-key", "response_cost": 0.0} + fresh_response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + logging_obj.model_call_details["cache_hit"] = True + now: Final = datetime.datetime.now() + payload: Final = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-5.4-nano", + "messages": [], + "cache_hit": True, + "litellm_params": logging_obj.litellm_params, + }, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + uncached_price: Final = oracle_logging_obj._response_cost_calculator(result=fresh_response) + + assert payload is not None + assert payload["saved_cache_cost"] == uncached_price + assert payload["saved_cache_cost"] > 0.0 + assert payload["response_cost"] == 0.0 + assert logging_obj.cost_breakdown is None + assert oracle_logging_obj.cost_breakdown is not None + _assert_zeroed_copy(payload["cost_breakdown"], oracle_logging_obj.cost_breakdown) + finally: + litellm.model_cost.pop(deployment_id, None) + + +def test_cache_hit_payload_reports_zero_cost_though_the_call_priced_the_response_as_fresh(): + """The proxy prices a /v1/messages body for its cost header before the success callback learns it was a cache + hit, and a replayed stream's assembled response carries its original price, so the call holds full-price numbers + that the hit's payload must not export.""" + deployment_id: Final = "lit8679-stale-breakdown-deployment" + litellm.register_model( + model_cost={ + deployment_id: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + try: + logging_obj: Final = _lit8679_priced_logging_obj(deployment_id, "lit8679-stale-breakdown") + response: Final = ModelResponse( + model="gpt-5.4-nano", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + header_price: Final = logging_obj.response_cost_calculator(result=response) + response._hidden_params["response_cost"] = header_price + logging_obj.model_call_details["cache_hit"] = True + now: Final = datetime.datetime.now() + payload: Final = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-5.4-nano", + "messages": [], + "cache_hit": True, + "response_cost": 0.0, + "litellm_params": logging_obj.litellm_params, + }, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert header_price is not None and header_price > 0.0 + assert logging_obj.cost_breakdown is not None + assert payload is not None + assert payload["response_cost"] == 0.0 + assert payload["hidden_params"]["response_cost"] == 0.0 + _assert_zeroed_copy(payload["cost_breakdown"], logging_obj.cost_breakdown) + assert payload["saved_cache_cost"] == header_price + finally: + litellm.model_cost.pop(deployment_id, None) + + +def test_uncached_price_leaves_no_failure_debug_info_for_unpriced_model(): + """A probe that cannot price the model returns None without stamping the pricing failure on the cache-hit call.""" + logging_obj: Final = LitellmLogging( + model="lit8679-model-with-no-price", + messages=[{"role": "user", "content": "Hi"}], + stream=False, + call_type="completion", + start_time=time.time(), + litellm_call_id="lit8679-unpriced-probe", + function_id="fn", + ) + logging_obj.update_environment_variables( + model="lit8679-model-with-no-price", + user="", + optional_params={}, + litellm_params={"metadata": {}}, + custom_llm_provider="openai", + ) + cached_response: Final = ModelResponse( + model="lit8679-model-with-no-price", + choices=[litellm.Choices(message=litellm.Message(role="assistant", content="cached"))], + usage=litellm.Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + cached_response._hidden_params = {"cache_hit": True, "cache_key": "lit8679-key", "response_cost": 0.0} + assert "response_cost_failure_debug_information" not in logging_obj.model_call_details + + cost: Final = logging_obj.uncached_price(result=cached_response).cost + + assert cost is None + assert "response_cost_failure_debug_information" not in logging_obj.model_call_details + + def test_enterprise_alerting_loggers_resolves_real_classes(): from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 6cc4e22a26a..679e1bfd85a 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic +from litellm.types.utils import CachingDetails from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -11637,3 +11638,98 @@ class TestResponseCostFromLoggingObjForBridgedMessages: assert logging_obj.cost_breakdown is not None assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(expected_output) assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) + + +class TestResponseCostFromLoggingObjForCacheHits: + BODY: Final = MappingProxyType( + { + "id": "msg_cached", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "cached answer"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1200, "output_tokens": 400}, + } + ) + + @staticmethod + def _logging_obj(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = Logging( + model="claude-haiku-5-5", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id=str(uuid.uuid4()), + function_id="cache-hit-header", + ) + logging_obj.optional_params = {} + return logging_obj + + def test_a_v1_messages_cache_hit_costs_nothing_though_its_usage_prices_as_a_fresh_call(self): + fresh_cost: Final = ProxyBaseLLMRequestProcessing._response_cost_from_logging_obj( + response=dict(self.BODY), logging_obj=self._logging_obj() + ) + hit: Final = self._logging_obj() + hit.caching_details = {"cache_hit": True, "cache_duration_ms": 2.0} + + hit_cost: Final = ProxyBaseLLMRequestProcessing._response_cost_from_logging_obj( + response=dict(self.BODY), logging_obj=hit + ) + + assert fresh_cost is not None and fresh_cost > 0 + assert hit_cost == 0.0 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("hidden_params", "caching_details", "expected_header"), + ( + ( + {"cache_hit": True, "response_cost": 0.0}, + CachingDetails(cache_hit=True, cache_duration_ms=2.0), + "0.0", + ), + ({"response_cost": 0.0}, None, None), + ), + ) + async def test_only_a_cache_hit_sends_zero_cost_headers( + self, + hidden_params: dict[str, object], + caching_details: CachingDetails | None, + expected_header: str | None, + ): + response: Final = litellm.ModelResponse() + response._hidden_params = hidden_params + logging_obj: Final = self._logging_obj() + logging_obj.caching_details = caching_details + + headers: Final = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data={"model": "gpt-5.4-mini"}, + request=Request({"type": "http", "method": "POST", "headers": []}), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + logging_obj=logging_obj, + version=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + cost_headers: Final = ( + "x-litellm-response-cost", + "x-litellm-response-cost-original", + "x-litellm-response-cost-input", + "x-litellm-response-cost-output", + ) + assert {name: headers.get(name) for name in cost_headers} == dict.fromkeys(cost_headers, expected_header) + + def test_a_fresh_zero_cost_v1_messages_call_still_sends_no_cost_header(self): + logging_obj: Final = self._logging_obj() + logging_obj.model_call_details["response_cost"] = 0.0 + + cost: Final = ProxyBaseLLMRequestProcessing._response_cost_from_logging_obj( + response=dict(self.BODY), logging_obj=logging_obj + ) + + assert cost == "" From 03e3e71efa13966e6177f7f4156e5e4e72a04107 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 17 Oct 2026 11:31:37 -0700 Subject: [PATCH 24/64] fix(proxy): let proxy_admin_viewer store its own MCP credentials (#46994) Co-authored-by: michelligabriele Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/route_checks.py | 25 ++++++++---- tests/unit/proxy/auth/test_route_checks.py | 45 ++++++++++++++++++++++ 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 8dab693cd58..f5e8304c83c 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -72,14 +72,22 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES: Final = ("/regenerate", "/reset_spe _LLM_API_ROUTE_GROUPS: Final = frozenset(("llm_api_routes", "llm_api_routes_without_mcp")) _AUTH_ENFORCED_PASS_THROUGH_ROUTE_GROUPS: Final = frozenset(("openai_routes", *_LLM_API_ROUTE_GROUPS)) -# Method-scoped because the generic credential CRUD handlers share these paths: -# the route check cannot admit PATCH/DELETE on a name like "user_connections" -# or "*/user_connection*" without also opening the admin CRUD handlers. -_INTERNAL_USER_CREDENTIAL_CONNECTION_ROUTES: Final = ( +# Writes that only touch the caller's own credential row (BYOK connections and +# per-user MCP credentials). Method-scoped because the generic credential CRUD +# handlers share these paths: the route check cannot admit PATCH/DELETE on a +# name like "user_connections" without also opening the admin CRUD handlers. +_SELF_SCOPED_CREDENTIAL_ROUTES: Final = ( ("GET", "/credentials/user_connections"), ("POST", "/credentials/{credential_name:path}/user_connection/start"), ("POST", "/credentials/{credential_name:path}/user_connection/poll"), ("DELETE", "/credentials/{credential_name:path}/user_connection"), + ("POST", "/v1/mcp/server/oauth/{server_id}/register"), + ("POST", "/v1/mcp/server/{server_id}/oauth-user-credential"), + ("DELETE", "/v1/mcp/server/{server_id}/oauth-user-credential"), + ("POST", "/v1/mcp/server/{server_id}/user-credential"), + ("DELETE", "/v1/mcp/server/{server_id}/user-credential"), + ("POST", "/v1/mcp/server/{server_id}/user-env-vars"), + ("DELETE", "/v1/mcp/server/{server_id}/user-env-vars"), ) @@ -349,7 +357,7 @@ class RouteChecks: _user_role == LitellmUserRoles.INTERNAL_USER.value and ( RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value) - or RouteChecks.is_internal_user_credential_connection_route(route=route, request=request) + or RouteChecks.is_self_scoped_credential_route(route=route, request=request) ) or user_is_org_admin(request_data=request_data, user_object=user_obj) and RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value) @@ -702,13 +710,13 @@ class RouteChecks: _get_request_method = get_request_method @staticmethod - def is_internal_user_credential_connection_route(route: str, request: Request | None) -> bool: + def is_self_scoped_credential_route(route: str, request: Request | None) -> bool: method: Final = RouteChecks.get_request_method(request) if method is None: return False return any( method == expected_method and RouteChecks.check_route_access(route=route, allowed_routes=(pattern,)) - for expected_method, pattern in _INTERNAL_USER_CREDENTIAL_CONNECTION_ROUTES + for expected_method, pattern in _SELF_SCOPED_CREDENTIAL_ROUTES ) @staticmethod @@ -1033,6 +1041,9 @@ class RouteChecks: if route == "/session/logout": return + if RouteChecks.is_self_scoped_credential_route(route=route, request=request): + return + # Hard-block known write routes regardless of HTTP method (defensive # — these are POSTs in practice, but pinning them here protects # against future GET-shaped writes). diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index 21958e88008..7bdcb8afd37 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -4782,6 +4782,51 @@ def test_view_only_user_denied_credential_connection_routes(method: str, route: _internal_user_check(route, method, role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) +_SELF_SCOPED_CREDENTIAL_WRITES = ( + ("POST", "/credentials/x/user_connection/start"), + ("POST", "/credentials/x/user_connection/poll"), + ("DELETE", "/credentials/x/user_connection"), + ("POST", "/v1/mcp/server/oauth/srv-1/register"), + ("POST", "/v1/mcp/server/srv-1/oauth-user-credential"), + ("DELETE", "/v1/mcp/server/srv-1/oauth-user-credential"), + ("POST", "/v1/mcp/server/srv-1/user-credential"), + ("DELETE", "/v1/mcp/server/srv-1/user-credential"), + ("POST", "/v1/mcp/server/srv-1/user-env-vars"), + ("DELETE", "/v1/mcp/server/srv-1/user-env-vars"), +) + + +@pytest.mark.parametrize(("method", "route"), _SELF_SCOPED_CREDENTIAL_WRITES) +def test_proxy_admin_viewer_allowed_self_scoped_credential_writes(method: str, route: str) -> None: + """Regression for LIT-9320: these handlers only write the caller's own credential + row, so a read-only admin can use them like /user/password/change.""" + assert _internal_user_check(route, method, role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) is None + + +@pytest.mark.parametrize( + ("method", "route"), + [ + ("POST", "/v1/mcp/server"), + ("PUT", "/v1/mcp/server"), + ("DELETE", "/v1/mcp/server/srv-1"), + ("POST", "/v1/mcp/server/import"), + ("POST", "/v1/mcp/server/srv-1/pin"), + ("PATCH", "/v1/mcp/server/srv-1/oauth-user-credential"), + ("PUT", "/v1/mcp/server/srv-1/user-credential"), + ("PATCH", "/v1/mcp/server/srv-1/user-env-vars"), + ("POST", "/v1/mcp/server/srv-1/user-credential/rotate"), + ("POST", "/v1/mcp/server/oauth/srv-1/token"), + ("DELETE", "/v1/mcp/server/oauth/srv-1/register"), + ("PATCH", "/credentials/x/user_connection"), + ("DELETE", "/credentials/user_connections"), + ("POST", "/key/generate"), + ], +) +def test_proxy_admin_viewer_denied_mcp_admin_and_sibling_writes(method: str, route: str) -> None: + with pytest.raises(HTTPException, match="user not allowed to access this route"): + _internal_user_check(route, method, role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + + def test_user_connection_route_registration_order_wins_over_generic_credential_routes() -> None: """The :path placeholders in the generic CRUD routes could swallow the user_connection paths; the concrete routes must register (and match) first.""" From b5037f28d172f6f6375ba3d08dcdf2c10deb8d2d Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 11:39:00 -0700 Subject: [PATCH 25/64] test(perf): qualify database upgrades with shared 5 RPS traffic (#46802) * test(perf): qualify database upgrades with shared 5 RPS traffic * test(perf): share regular latency and memory limits for upgrades * fix(perf): escape index format placeholder in upgrade worker * fix(perf): handle uninitialized multiprocess failure counters * test(perf): match persisted upgrade requests by call ID --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: kerry --- tests/e2e/perf/setup/locustfile.py | 25 +- tests/e2e/perf/setup/run.py | 111 ++++++-- tests/e2e/perf/setup/upgrade.py | 209 +++++++++++++++ tests/e2e/perf/setup/upgrade_worker.py | 146 +++++++++++ tests/e2e/perf/setup/wire.py | 30 ++- tests/e2e/perf/test_database_upgrade.py | 85 +++++++ tests/e2e_harness/perf/test_run.py | 66 +++++ tests/e2e_harness/perf/test_upgrade.py | 325 ++++++++++++++++++++++++ 8 files changed, 959 insertions(+), 38 deletions(-) create mode 100644 tests/e2e/perf/setup/upgrade.py create mode 100644 tests/e2e/perf/setup/upgrade_worker.py create mode 100644 tests/e2e/perf/test_database_upgrade.py create mode 100644 tests/e2e_harness/perf/test_run.py create mode 100644 tests/e2e_harness/perf/test_upgrade.py diff --git a/tests/e2e/perf/setup/locustfile.py b/tests/e2e/perf/setup/locustfile.py index 1609617bd73..ccc7cf698f8 100644 --- a/tests/e2e/perf/setup/locustfile.py +++ b/tests/e2e/perf/setup/locustfile.py @@ -19,9 +19,10 @@ from typing import Final, cast import gevent from locust import FastHttpUser, task from locust.contrib.fasthttp import FastResponse, ResponseContextManager -from setup.wire import Endpoint, WireRecord, WireRequest, request_body from pydantic import TypeAdapter +from setup.wire import Endpoint, PlanProgress, WireRecord, WireRequest, request_body, response_id + _PLAN_ADAPTER: Final = TypeAdapter(list[WireRequest]) _JSON: Final = TypeAdapter[object](object) _PLAN: Final = _PLAN_ADAPTER.validate_json(Path(os.environ["PERF_PLAN_FILE"]).read_bytes()) @@ -33,7 +34,8 @@ _INTERVAL_S: Final = _USERS / _RPS _JITTER: Final = 0.25 _RNG: Final = random.Random(_PLAN[0].marker if _PLAN else "empty-plan") _RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8") -_written = 0 +_PROGRESS: Final = PlanProgress(len(_PLAN)) +_STOP_FILE: Final = os.environ.get("PERF_STOP_FILE") BLOCK_REASON: Final = "perf-test-block" _TOKEN_MAX_AGE_S: Final = 240.0 _tokens: dict[str, tuple[str, float]] = {} @@ -155,10 +157,9 @@ def _judge(item: WireRequest, status: int, body_text: str) -> tuple[bool, str | def _record(record: WireRecord) -> None: - global _written _ = _RECORDS.write(record.model_dump_json() + "\n") _RECORDS.flush() - _written += 1 + _PROGRESS.completed += 1 def _authorization(user: FastHttpUser, item: WireRequest) -> str: @@ -203,7 +204,7 @@ class PlanUser(FastHttpUser): a random point of its first interval and jitters every later interval by up to `_JITTER` either way, so the users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate.""" - network_timeout = 120.0 + network_timeout = float(os.environ.get("PERF_REQUEST_TIMEOUT", "120")) connection_timeout = 30.0 def on_start(self) -> None: @@ -221,9 +222,10 @@ class PlanUser(FastHttpUser): @task def send(self) -> None: - item: Final = next(_NEXT, None) + take: Final = _PROGRESS.take(stop=_STOP_FILE is not None and Path(_STOP_FILE).exists()) + item: Final = next(_NEXT, None) if take else None if item is None: - if _written >= len(_PLAN): + if _PROGRESS.drained: _quit(self) return started_at: Final = time.time() @@ -237,6 +239,7 @@ class PlanUser(FastHttpUser): error: str | None = None header_s = 0.0 response: FastResponse | None = None + identifier: str | None = None # rebind-ok: the ID arrives in a JSON body or the first SSE event try: headers: Final = {"Authorization": _authorization(self, item), "Content-Type": "application/json"} context: ResponseContextManager @@ -265,6 +268,7 @@ class PlanUser(FastHttpUser): buffer += line events, buffer = _events(buffer) for event in events: + identifier = identifier or response_id(event) if first_token_s is None and token_in_event(item.endpoint, event): first_token_s = now text += delta_of_event(item.endpoint, event) @@ -272,7 +276,9 @@ class PlanUser(FastHttpUser): body_text = context.text or "" first_byte_s = time.perf_counter() - started if status == 200: - text = text_of_json(item.endpoint, _as_dict(_JSON.validate_json(body_text))) + payload: Final = _as_dict(_JSON.validate_json(body_text)) + identifier = response_id(payload) + text = text_of_json(item.endpoint, payload) ok, error = _judge(item, status, body_text) if ok: context.success() @@ -302,7 +308,8 @@ class PlanUser(FastHttpUser): mock_delay_s=float(delay_header) if delay_header else None, call_id=_header(response, "x-litellm-call-id") if response is not None else None, error=error, + response_id=identifier, ) ) - if _written >= len(_PLAN): + if _PROGRESS.drained: _quit(self) diff --git a/tests/e2e/perf/setup/run.py b/tests/e2e/perf/setup/run.py index 9ba9b45104d..272c64754c4 100644 --- a/tests/e2e/perf/setup/run.py +++ b/tests/e2e/perf/setup/run.py @@ -9,9 +9,11 @@ import shutil import subprocess import sys import tempfile +import time +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path -from typing import Final +from typing import BinaryIO, Final from e2e_config import PROXY_BASE_URL from pydantic import TypeAdapter @@ -73,7 +75,59 @@ def _keep_records_artifact(records_path: Path) -> None: _ = shutil.copyfile(records_path, Path(artifact_dir) / "perf-records.jsonl") -def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: CaseSetup) -> RunOutcome: +class RecordStream: + def __init__(self, source: BinaryIO) -> None: + self.source: Final = source + self.partial = b"" + + def read(self) -> tuple[WireRecord, ...]: + chunks: Final = (self.partial + self.source.read()).split(b"\n") + self.partial = chunks[-1] + return tuple(_RECORD_ADAPTER.validate_json(line) for line in chunks[:-1] if line) + + +def run_until( + command: list[str], + env: dict[str, str], + records_path: Path, + stop_path: Path, + stop_when: Callable[[tuple[WireRecord, ...]], bool], + timeout: int, +) -> subprocess.CompletedProcess[str]: + with records_path.open("rb") as records, records_path.with_suffix(".log").open("w+") as output: + with subprocess.Popen(command, env=env, stdout=output, stderr=subprocess.STDOUT, text=True) as process: + deadline: Final = time.monotonic() + timeout + reader: Final = RecordStream(records) + try: + while process.poll() is None: + if stop_when(reader.read()): + stop_path.touch() + if time.monotonic() >= deadline: + raise subprocess.TimeoutExpired(command, timeout) + try: + _ = process.wait(timeout=1) + except subprocess.TimeoutExpired: + pass + finally: + stop_path.touch() + try: + _ = process.wait(timeout=35) + except subprocess.TimeoutExpired: + process.kill() + _ = process.wait(timeout=10) + _keep_records_artifact(records_path) + _ = output.seek(0) + return subprocess.CompletedProcess(command, process.returncode, output.read(), "") + + +def run_plan( + case: PerformanceTestCase, + plan: tuple[PlannedRequest, ...], + setup: CaseSetup, + *, + stop_when: Callable[[tuple[WireRecord, ...]], bool] | None = None, + request_timeout: float = 120.0, +) -> RunOutcome: """Replay `plan` with Locust in a child process and collect its records. The run-time ceiling is well past the planned duration so Locust's own clock never cuts a slow run short; the locustfile quits itself once the last record is written.""" @@ -82,6 +136,7 @@ def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: with tempfile.TemporaryDirectory(prefix="e2e-perf-") as directory: plan_path: Final = Path(directory) / "plan.json" records_path: Final = Path(directory) / "records.jsonl" + stop_path: Final = Path(directory) / "stop" _ = plan_path.write_bytes(_PLAN_ADAPTER.dump_json(wire)) records_path.touch() env: Final = { @@ -91,33 +146,33 @@ def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: "PERF_RECORDS_FILE": str(records_path), "PERF_RPS": str(case.traffic.rps), "PERF_USERS": str(case.traffic.concurrency), + "PERF_REQUEST_TIMEOUT": str(request_timeout), + **({"PERF_STOP_FILE": str(stop_path)} if stop_when is not None else {}), } - completed: Final = subprocess.run( - [ - sys.executable, - "-P", - "-m", - "locust", - "--headless", - "--only-summary", - "--locustfile", - str(_LOCUSTFILE), - "--host", - PROXY_BASE_URL, - "--users", - str(case.traffic.concurrency), - "--spawn-rate", - str(case.traffic.concurrency), - "--run-time", - f"{ceiling_s}s", - "--exit-code-on-error", - "0", - ], - env=env, - capture_output=True, - text=True, - timeout=ceiling_s + 120, - check=False, + command: Final = [ + sys.executable, + "-P", + "-m", + "locust", + "--headless", + "--only-summary", + "--locustfile", + str(_LOCUSTFILE), + "--host", + PROXY_BASE_URL, + "--users", + str(case.traffic.concurrency), + "--spawn-rate", + str(case.traffic.concurrency), + "--run-time", + f"{ceiling_s}s", + "--exit-code-on-error", + "0", + ] + completed: Final = ( + run_until(command, env, records_path, stop_path, stop_when, ceiling_s + 120) + if stop_when is not None + else subprocess.run(command, env=env, capture_output=True, text=True, timeout=ceiling_s + 120, check=False) ) records: Final = tuple( _RECORD_ADAPTER.validate_json(line) diff --git a/tests/e2e/perf/setup/upgrade.py b/tests/e2e/perf/setup/upgrade.py new file mode 100644 index 00000000000..37e5c7bc904 --- /dev/null +++ b/tests/e2e/perf/setup/upgrade.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +import math +import time +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import psycopg +from cases.small_bodies_5rps import CASE +from e2e_config import MASTER_KEY +from e2e_http import send_text_external +from prometheus_client.parser import text_string_to_metric_families +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter + +from setup.performance_test_case import PerformanceTestCase +from setup.upgrade_worker import UpgradeState, fingerprints, missing_indexes, read_state, release_indexes +from setup.wire import WireRecord + +MAX_TRAFFIC_SECONDS: Final = 6 * 3600 +_ROWS: Final = TypeAdapter(tuple[tuple[str, str | None], ...]) +_ACTIVITY: Final = TypeAdapter(tuple[list[int], int]) +_COUNT: Final = TypeAdapter(tuple[int]) +_SERVICES: Final = ("postgres", "batch_write_to_db") + + +def upgrade_case() -> PerformanceTestCase: + return replace( + CASE, + name="database_upgrade_5rps", + traffic=replace(CASE.traffic, total_requests=int(CASE.traffic.rps * MAX_TRAFFIC_SECONDS)), + limits=tuple( + limit + for limit in CASE.limits + if limit.metric + in { + "unexpected_status_count", + "gateway_added_seconds", + "time_to_first_token_seconds", + "memory_growth_megabytes", + } + ), + ) + + +class MetricsHeaders(BaseModel): + model_config = ConfigDict(frozen=True) + authorization: str = Field(default=f"Bearer {MASTER_KEY}", alias="Authorization") + + +def postgres_failures(text: str) -> float: + families: Final = tuple(text_string_to_metric_families(text)) + names: Final = frozenset(family.name for family in families) + required: Final = frozenset(f"litellm_{service}_failed_requests" for service in _SERVICES) + samples: Final = tuple(sample for family in families for sample in family.samples) + has_database_activity: Final = any( + sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0 + for sample in samples + ) + assert required <= names or has_database_activity, f"Missing database service counters: {required - names}" + failures: Final = tuple(sample.value for sample in samples if sample.name in {f"{name}_total" for name in required}) + assert all(math.isfinite(value) and value >= 0 for value in failures), "Invalid database failure counter" + return sum(failures) + + +def assert_database_metrics(urls: tuple[str, ...]) -> None: + assert urls, "E2E_DATABASE_METRICS_URLS must name every gateway/backend /metrics endpoint" + for url in urls: + assert_database_metrics_url(url) + + +def assert_database_metrics_url(url: str) -> None: + result: Final = send_text_external("GET", url, headers=MetricsHeaders(), timeout=10) + assert result.status_code == 200, f"Database metrics unavailable at {url}: {result.status_code}" + assert postgres_failures(result.body) == 0, f"Failed PostgreSQL operations reported at {url}" + + +def persisted(connection: psycopg.Connection[tuple[object, ...]], records: tuple[WireRecord, ...]) -> frozenset[str]: + ids: Final = tuple(record.call_id for record in records if record.call_id) + if not ids: + return frozenset() + rows: Final = _ROWS.validate_python( + connection.execute( + 'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=ANY(%s)', (list(ids),) + ).fetchall() + ) + return matched_response_ids(records, rows) + + +def matched_response_ids(records: tuple[WireRecord, ...], rows: tuple[tuple[str, str | None], ...]) -> frozenset[str]: + call_ids: Final = frozenset(call_id for _, call_id in rows if call_id) + return frozenset(record.response_id for record in records if record.response_id and record.call_id in call_ids) + + +def activity(connection: psycopg.Connection[tuple[object, ...]]) -> tuple[list[int], int]: + return _ACTIVITY.validate_python( + connection.execute("""SELECT + ARRAY(SELECT pid FROM pg_stat_progress_create_index WHERE datname=current_database()), + (SELECT count(*) FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock') + """).fetchone() + ) + + +def assert_finished(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None: + assert state.phase == "complete", f"Upgrade did not finish: {state.phase}" + assert not missing_indexes(connection, state.required_indexes), "Required target indexes are missing or invalid" + (missing,) = _COUNT.validate_python( + connection.execute( + """ + SELECT count(*) FROM unnest(%s::text[]) name WHERE NOT EXISTS ( + SELECT 1 FROM "_prisma_migrations" WHERE migration_name=name + AND finished_at IS NOT NULL AND rolled_back_at IS NULL)""", + (list(state.required_migrations),), + ).fetchone() + ) + assert missing == 0, f"{missing} target migrations were not applied" + (unfinished,) = _COUNT.validate_python( + connection.execute(""" + SELECT count(*) FROM "_prisma_migrations" WHERE finished_at IS NULL AND rolled_back_at IS NULL + """).fetchone() + ) + assert unfinished == 0, f"{unfinished} unfinished migrations" + assert fingerprints(connection, tuple(state.baseline_rows)) == state.baseline_rows, ( + "Baseline rows changed or disappeared" + ) + + +class UpgradeMonitor: + def __init__( + self, + *, + state: Callable[[], UpgradeState], + check_metrics: Callable[[], None], + read_activity: Callable[[], tuple[list[int], int]], + read_persisted: Callable[[tuple[WireRecord, ...]], frozenset[str]], + start_indexes: Callable[[], None], + clock: Callable[[], float] = time.monotonic, + wall_clock: Callable[[], float] = time.time, + ) -> None: + self.state: Final = state + self.check_metrics: Final = check_metrics + self.read_activity: Final = read_activity + self.read_persisted: Final = read_persisted + self.start_indexes: Final = start_indexes + self.clock: Final = clock + self.wall_clock: Final = wall_clock + self.started: Final = clock() + self.status_limit: Final = next(limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count") + self.pending: tuple[WireRecord, ...] = () + self.processed = 0 + self.successes = 0 + self.persisted_count = 0 + self.overlap_writes = 0 + self.max_lock_waiters = 0 + self.previous_builds: frozenset[int] = frozenset() + self.build_seen_at: float | None = None + self.released = False + self.completed_at: float | None = None + + def poll(self, records: tuple[WireRecord, ...]) -> bool: + self.check_metrics() + current: Final = self.state() + assert current.phase != "failed", "Target migration/index worker failed" + assert self.clock() - self.started < MAX_TRAFFIC_SECONDS, "Upgrade exceeded its six-hour operational timeout" + successful: Final = tuple(record for record in records if record.status == 200) + self.processed += len(records) + self.successes += len(successful) + assert self.processed - self.successes <= self.status_limit, "Too many unexpected non-200 responses" + assert all(record.ok and record.text for record in successful), "Failed or empty inference response" + assert all(record.response_id and record.call_id for record in successful), ( + "Response missing spend-log identifiers" + ) + builds, lock_waiters = self.read_activity() + if not self.previous_builds.intersection(builds): + self.build_seen_at = self.wall_clock() if builds else None + self.max_lock_waiters = max(self.max_lock_waiters, lock_waiters) + self.pending += successful + found: Final = self.read_persisted(self.pending) + self.persisted_count += len(found) + still_building, _ = self.read_activity() + since: Final = self.build_seen_at + if since is not None and frozenset(builds).intersection(still_building): + self.overlap_writes += sum( + 1 for record in self.pending if record.response_id in found and record.started_at >= since + ) + self.previous_builds = frozenset(builds) + self.pending = tuple(record for record in self.pending if record.response_id not in found) + if self.successes and not self.released: + self.start_indexes() + self.released = True + if current.phase == "complete" and self.completed_at is None: + self.completed_at = self.clock() + return self.completed_at is not None and self.clock() - self.completed_at >= 60 and self.successes >= 300 + + def assert_complete(self, state: UpgradeState) -> None: + assert self.successes >= 300, f"Only {self.successes} requests completed" + assert not self.pending, f"{len(self.pending)} successful requests have no matching spend log" + assert self.persisted_count == self.successes, "Spend-log count does not match successful requests" + assert not state.pending_indexes or self.overlap_writes > 0, "No persisted requests observed during index work" + + +def monitor_for(connection: psycopg.Connection[tuple[object, ...]], urls: tuple[str, ...]) -> UpgradeMonitor: + return UpgradeMonitor( + state=lambda: read_state(connection), + check_metrics=lambda: assert_database_metrics(urls), + read_activity=lambda: activity(connection), + read_persisted=lambda records: persisted(connection, records), + start_indexes=lambda: release_indexes(connection), + ) diff --git a/tests/e2e/perf/setup/upgrade_worker.py b/tests/e2e/perf/setup/upgrade_worker.py new file mode 100644 index 00000000000..e69e841814f --- /dev/null +++ b/tests/e2e/perf/setup/upgrade_worker.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import json +import os +import sys +import time +from pathlib import Path +from typing import Final, Literal + +import psycopg +from pydantic import BaseModel, ConfigDict, TypeAdapter + + +class UpgradeState(BaseModel): + model_config = ConfigDict(frozen=True) + + phase: Literal["schema", "schema_ready", "indexes", "complete", "failed"] + started_at: float + schema_seconds: float = 0 + index_seconds: float = 0 + required_indexes: tuple[str, ...] + required_migrations: tuple[str, ...] + pending_indexes: tuple[str, ...] + baseline_rows: dict[str, str] + + +_STATE: Final = TypeAdapter(tuple[UpgradeState]) +_BOOL: Final = TypeAdapter(tuple[bool]) +_STRINGS: Final = TypeAdapter(tuple[tuple[str], ...]) +_FINGERPRINTS: Final = TypeAdapter(tuple[tuple[str, str], ...]) +_FINGERPRINT: Final = """md5(ROW(request_id, api_key, spend, total_tokens, prompt_tokens, + completion_tokens, model, "startTime", "endTime")::text)""" + + +def connect() -> psycopg.Connection[tuple[object, ...]]: + return psycopg.connect( + os.environ["DATABASE_URL"], + autocommit=True, + connect_timeout=10, + options="-c statement_timeout=10000 -c timezone=UTC", + ) + + +def fingerprints(connection: psycopg.Connection[tuple[object, ...]], ids: tuple[str, ...]) -> dict[str, str]: + rows: Final = connection.execute( + f'SELECT request_id, {_FINGERPRINT} FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', (list(ids),) + ).fetchall() + return dict(_FINGERPRINTS.validate_python(rows)) + + +def missing_indexes(connection: psycopg.Connection[tuple[object, ...]], names: tuple[str, ...]) -> tuple[str, ...]: + rows: Final = connection.execute( + """SELECT name FROM unnest(%s::text[]) AS name WHERE NOT EXISTS ( + SELECT 1 FROM pg_index WHERE indexrelid = to_regclass(format('public.%%I', name)) + AND indisvalid AND indisready)""", + (list(names),), + ).fetchall() + return tuple(row[0] for row in _STRINGS.validate_python(rows)) + + +def read_state(connection: psycopg.Connection[tuple[object, ...]]) -> UpgradeState: + (state,) = _STATE.validate_python( + connection.execute('SELECT state FROM "_LiteLLMPerfUpgrade" WHERE id=1').fetchone() + ) + return state + + +def write_state(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None: + _ = connection.execute( + """INSERT INTO "_LiteLLMPerfUpgrade" (id, state) VALUES (1, %s::jsonb) + ON CONFLICT (id) DO UPDATE SET state = EXCLUDED.state""", + (state.model_dump_json(),), + ) + + +def release_indexes(connection: psycopg.Connection[tuple[object, ...]]) -> None: + _ = connection.execute('UPDATE "_LiteLLMPerfUpgrade" SET traffic_ready=true WHERE id=1') + + +def wait_for_traffic(connection: psycopg.Connection[tuple[object, ...]], deadline: float) -> bool: + while time.monotonic() < deadline: + (ready,) = _BOOL.validate_python( + connection.execute('SELECT traffic_ready FROM "_LiteLLMPerfUpgrade" WHERE id=1').fetchone() + ) + if ready: + return True + time.sleep(1) + return False + + +def main() -> int: + assert os.environ.get("E2E_DATABASE_UPGRADE") == "1", "Only run against a disposable restored baseline" + assert os.environ.get("DATABASE_PERF_CONFIRM_DISPOSABLE") == "1", "Disposable database confirmation required" + from litellm_proxy_extras import request_log_indexes + from litellm_proxy_extras.utils import ProxyExtrasDBManager + + with connect() as connection: + _ = connection.execute("""CREATE TABLE "_LiteLLMPerfUpgrade" ( + id integer PRIMARY KEY CHECK (id=1), state jsonb NOT NULL, traffic_ready boolean NOT NULL DEFAULT false)""") + sample: Final = connection.execute(""" + (SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id LIMIT 32) + UNION (SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id DESC LIMIT 32)""").fetchall() + ids: Final = tuple(row[0] for row in _STRINGS.validate_python(sample)) + assert ids, "Restored baseline has no spend logs" + names: Final = tuple(index.name for index in request_log_indexes.REQUEST_LOG_INDEXES) + migrations: Final = Path(request_log_indexes.__file__).parent / "migrations" + initial: Final = UpgradeState( + phase="schema", + started_at=time.time(), + required_indexes=names, + required_migrations=tuple(sorted(path.parent.name for path in migrations.glob("*/migration.sql"))), + pending_indexes=missing_indexes(connection, names), + baseline_rows=fingerprints(connection, ids), + ) + assert initial.required_migrations and initial.required_indexes, "Target schema requirements are missing" + write_state(connection, initial) + try: + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True), ( + "Schema migration failed" + ) + ready: Final = initial.model_copy( + update={ + "phase": "schema_ready", + "schema_seconds": time.time() - initial.started_at, + "pending_indexes": missing_indexes(connection, names), + } + ) + write_state(connection, ready) + assert wait_for_traffic(connection, time.monotonic() + 1800), "Traffic did not start within 30 minutes" + write_state(connection, ready.model_copy(update={"phase": "indexes"})) + index_start: Final = time.monotonic() + assert ProxyExtrasDBManager.build_request_log_indexes(), "Request-log index build failed" + complete: Final = ready.model_copy( + update={"phase": "complete", "index_seconds": time.monotonic() - index_start} + ) + write_state(connection, complete) + _ = sys.stdout.write(complete.model_dump_json() + "\n") + return 0 + except Exception as exc: + write_state(connection, read_state(connection).model_copy(update={"phase": "failed"})) + _ = sys.stderr.write(json.dumps({"error": type(exc).__name__}) + "\n") + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/e2e/perf/setup/wire.py b/tests/e2e/perf/setup/wire.py index 07dc5d170ae..ced79d6970c 100644 --- a/tests/e2e/perf/setup/wire.py +++ b/tests/e2e/perf/setup/wire.py @@ -9,12 +9,13 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final, Literal -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from setup.corpus.session import Endpoint as CorpusEndpoint from setup.corpus.session import request_body as corpus_body Endpoint = Literal["chat_completions", "responses", "messages"] +_PAYLOAD: Final = TypeAdapter(dict[str, object]) class TokenSource(BaseModel): @@ -70,6 +71,33 @@ class WireRecord(BaseModel): mock_delay_s: float | None call_id: str | None error: str | None + response_id: str | None = None + + +def response_id(payload: Mapping[str, object]) -> str | None: + nested: Final = payload.get("response") or payload.get("message") + source: Final = _PAYLOAD.validate_python(nested) if isinstance(nested, Mapping) else payload + identifier: Final = source.get("id") + return identifier if isinstance(identifier, str) else None + + +class PlanProgress: + def __init__(self, total: int) -> None: + self.total: Final = total + self.issued = 0 + self.completed = 0 + self.stopped = False + + def take(self, *, stop: bool) -> bool: + self.stopped = self.stopped or stop + if self.stopped or self.issued >= self.total: + return False + self.issued += 1 + return True + + @property + def drained(self) -> bool: + return (self.stopped or self.issued == self.total) and self.completed == self.issued _CORPUS_ENDPOINT: Final[Mapping[Endpoint, CorpusEndpoint]] = MappingProxyType( diff --git a/tests/e2e/perf/test_database_upgrade.py b/tests/e2e/perf/test_database_upgrade.py new file mode 100644 index 00000000000..f165b9b738f --- /dev/null +++ b/tests/e2e/perf/test_database_upgrade.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path +from typing import Final + +import pytest +from idp import Keycloak +from lifecycle import ResourceManager +from memory_readings import read_rss_everywhere +from proxy_client import ProxyClient +from setup.datastore import read_db_counters, read_redis_counters +from setup.observe import Observation, evaluate, read_mock_requests, render, reset_mock, set_mock_latency +from setup.plan import build_plan +from setup.podcpu import gateway_cpu_sampler +from setup.provision import MOCK_PROVIDERS, build_setup, mock_url +from setup.run import run_plan +from setup.upgrade import assert_finished, monitor_for, upgrade_case +from setup.upgrade_worker import connect, read_state + +pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] + + +@pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only") +def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: ResourceManager) -> None: + case: Final = upgrade_case() + urls: Final = tuple(mock_url(provider) for provider in MOCK_PROVIDERS) + for url in urls: + _ = reset_mock(url) + _ = set_mock_latency(url, case.mock_latency) + setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources) + metrics: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip()) + with connect() as connection: + assert read_state(connection).phase == "schema_ready", "Expected the target schema ready, index work waiting" + rss_before: Final = read_rss_everywhere(proxy) + db_before: Final = read_db_counters() + redis_before: Final = read_redis_counters() + monitor: Final = monitor_for(connection, metrics) + with gateway_cpu_sampler() as cpu: + outcome: Final = run_plan(case, build_plan(case), setup, stop_when=monitor.poll, request_timeout=30) + assert outcome.locust_exit == 0, outcome.locust_output + remaining: Final = outcome.records[monitor.processed :] + _ = monitor.poll(remaining) + deadline: Final = time.monotonic() + 120 + while monitor.pending and time.monotonic() < deadline: + time.sleep(1) + _ = monitor.poll(()) + state: Final = read_state(connection) + _ = sys.stdout.write( + f"successful_requests={monitor.successes} persisted_requests={monitor.persisted_count} " + f"unmatched_requests={len(monitor.pending)}\n" + ) + monitor.assert_complete(state) + assert_finished(connection, state) + rss_after: Final = read_rss_everywhere(proxy) + assert {r.worker for r in rss_before.readings} == {r.worker for r in rss_after.readings}, ( + "Gateway workers restarted" + ) + observed: Final = Observation( + records=outcome.records, + mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)), + guardrail_calls=(), + rss_before=rss_before, + rss_after=rss_after, + wall_seconds=outcome.wall_seconds, + spend_logged_share=1, + db_before=db_before, + db_after=read_db_counters(), + redis_before=redis_before, + redis_after=read_redis_counters(), + gateway_cpu_seconds=cpu.seconds, + ) + rows, exceeded = evaluate(case, observed) + _ = sys.stdout.write(render(case, observed, rows) + "\n") + _ = sys.stdout.write( + f"schema_seconds={state.schema_seconds:.3f} index_seconds={state.index_seconds:.3f} " + f"overlap_writes={monitor.overlap_writes} max_lock_waiters={monitor.max_lock_waiters} " + f"pending_indexes={state.pending_indexes}\n" + ) + artifact: Final = os.environ.get("E2E_ARTIFACT_DIR") + if artifact: + _ = (Path(artifact) / "database-upgrade.json").write_text(state.model_dump_json(), encoding="utf-8") + assert not exceeded, "\n".join(exceeded) diff --git a/tests/e2e_harness/perf/test_run.py b/tests/e2e_harness/perf/test_run.py new file mode 100644 index 00000000000..23f62c9c560 --- /dev/null +++ b/tests/e2e_harness/perf/test_run.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import os +import sys +from pathlib import Path +from typing import Final + +import pytest +from setup.run import run_until +from setup.wire import WireRecord +from test_upgrade import record + + +class StopAfterRecord: + def __init__(self, *, partial: Path, resume: Path, fail: bool = False) -> None: + self.partial: Final = partial + self.resume: Final = resume + self.fail: Final = fail + self.received: tuple[WireRecord, ...] = () + + def __call__(self, records: tuple[WireRecord, ...]) -> bool: + self.received += records + if self.partial.exists(): + self.resume.touch() + assert not (self.fail and records), "Database failed" + return bool(records) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_controlled_process_preserves_partial_lines_and_drains_on_stop_or_failure(tmp_path: Path, fail: bool) -> None: + records: Final = tmp_path / "records.jsonl" + stop: Final = tmp_path / "stop" + partial: Final = tmp_path / "partial" + resume: Final = tmp_path / "resume" + records.touch() + encoded: Final = record().model_dump_json() + "\n" + program: Final = f""" +import time +from pathlib import Path +with Path({str(records)!r}).open("w") as f: + text = {encoded!r} + f.write(text[:10]) + f.flush() + Path({str(partial)!r}).touch() + while not Path({str(resume)!r}).exists(): + time.sleep(0.01) + f.write(text[10:]) + f.flush() + while not Path({str(stop)!r}).exists(): + time.sleep(0.01) + f.write({record(1).model_dump_json()!r} + "\\n") + f.flush() +""" + callback: Final = StopAfterRecord(partial=partial, resume=resume, fail=fail) + command: Final = [sys.executable, "-c", program] + if fail: + with pytest.raises(AssertionError, match="Database failed"): + _ = run_until(command, dict(os.environ), records, stop, callback, timeout=10) + else: + result: Final = run_until(command, dict(os.environ), records, stop, callback, timeout=10) + assert result.returncode == 0, result.stdout + assert callback.received == (record(),) + assert tuple(WireRecord.model_validate_json(line) for line in records.read_text().splitlines()) == ( + record(), + record(1), + ) diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py new file mode 100644 index 00000000000..0904b7a7b5a --- /dev/null +++ b/tests/e2e_harness/perf/test_upgrade.py @@ -0,0 +1,325 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final + +import pytest +from cases.small_bodies_5rps import CASE +from memory_readings import RssCapture, RssReading +from setup.observe import Observation, evaluate +from setup.upgrade import MAX_TRAFFIC_SECONDS, UpgradeMonitor, matched_response_ids, postgres_failures, upgrade_case +from setup.upgrade_worker import UpgradeState +from setup.wire import PlanProgress, WireRecord, response_id + + +def record(seq: int = 0, *, started_at: float = 100) -> WireRecord: + return WireRecord( + seq=seq, + marker=f"request-{seq}", + caller="key", + endpoint="chat_completions", + stream=False, + blocked=False, + started_at=started_at, + status=200, + header_s=1, + first_byte_s=1, + first_token_s=None, + total_s=1, + text="answer", + ok=True, + mock_delay_s=1, + call_id=f"call-{seq}", + error=None, + response_id=f"response-{seq}", + ) + + +def test_spend_logs_match_call_ids_when_response_ids_are_transformed() -> None: + records: Final = (record(0), record(1), record(2)) + rows: Final = (("provider-response-id", "call-0"), ("response-1", "unrelated-call")) + assert matched_response_ids(records, rows) == frozenset({"response-0"}) + + +def state(phase: str = "schema_ready", *, pending: bool = True) -> UpgradeState: + return UpgradeState.model_validate( + { + "phase": phase, + "started_at": 0, + "required_indexes": ["index"], + "required_migrations": ["migration"], + "pending_indexes": ["index"] if pending else [], + "baseline_rows": {"seed-row": "hash"}, + } + ) + + +class Probe: + def __init__(self) -> None: + self.now = 0.0 + self.current = state() + self.builds: list[int] = [] + self.found: frozenset[str] = frozenset() + self.releases = 0 + self.metric_checks = 0 + self.errors = False + + def check_metrics(self) -> None: + self.metric_checks += 1 + assert not self.errors, "Failed PostgreSQL operation" + + def persisted(self, records: tuple[WireRecord, ...]) -> frozenset[str]: + return frozenset(row.response_id for row in records if row.response_id in self.found and row.response_id) + + def release(self) -> None: + self.releases += 1 + + def monitor(self) -> UpgradeMonitor: + return UpgradeMonitor( + state=lambda: self.current, + check_metrics=self.check_metrics, + read_activity=lambda: (self.builds, 2), + read_persisted=self.persisted, + start_indexes=self.release, + clock=lambda: self.now, + wall_clock=lambda: self.now, + ) + + +def test_stop_drains_issued_requests_without_issuing_the_rest() -> None: + progress: Final = PlanProgress(100) + assert progress.take(stop=False) + assert progress.take(stop=False) + progress.completed = 1 + assert not progress.take(stop=True) + assert not progress.drained + assert not progress.take(stop=False) + progress.completed = 2 + assert progress.drained + assert progress.issued == 2 + + +def test_finite_plan_still_waits_for_its_last_response() -> None: + progress: Final = PlanProgress(1) + assert progress.take(stop=False) + assert not progress.take(stop=False) + assert not progress.drained + progress.completed = 1 + assert progress.drained + + +@pytest.mark.parametrize( + "payload,expected", + [ + ({"id": "chat-id"}, "chat-id"), + ({"type": "response.created", "response": {"id": "response-id"}}, "response-id"), + ({"type": "message_start", "message": {"id": "message-id"}}, "message-id"), + ({"type": "response.output_item.added", "item": {"id": "item-id"}}, None), + ({"type": "content_block_delta", "delta": {"text": "hello"}}, None), + ], +) +def test_response_id_uses_response_not_output_item(payload: Mapping[str, object], expected: str | None) -> None: + assert response_id(payload) == expected + + +def metrics(*, failure: str = "0", batch: str = "0") -> str: + return f"""# TYPE litellm_postgres_failed_requests_total counter +litellm_postgres_failed_requests_total{{function_name="query",error_class="Timeout"}} {failure} +# TYPE litellm_batch_write_to_db_failed_requests_total counter +litellm_batch_write_to_db_failed_requests_total{{function_name="flush",error_class="DatabaseError"}} {batch} +""" + + +def test_failed_sql_and_async_batch_writes_are_both_counted() -> None: + assert postgres_failures(metrics(failure="2", batch="3")) == 5 + assert postgres_failures(metrics()) == 0 + + +def test_no_error_label_series_is_a_valid_zero() -> None: + assert ( + postgres_failures("""# TYPE litellm_postgres_failed_requests_total counter +# TYPE litellm_batch_write_to_db_failed_requests_total counter +""") + == 0 + ) + + +@pytest.mark.parametrize("service", [None, "postgres", "batch_write_to_db"]) +def test_multiprocess_metrics_omit_failure_families_until_a_failure(service: str | None) -> None: + total: Final = """# TYPE litellm_postgres_total_requests_total counter +litellm_postgres_total_requests_total{postgres="postgres"} 1 +""" + failure: Final = ( + f"# TYPE litellm_{service}_failed_requests_total counter\n" + f'litellm_{service}_failed_requests_total{{function_name="query",error_class="Timeout"}} 1\n' + if service is not None + else "" + ) + assert postgres_failures(total + failure) == (1 if service is not None else 0) + + +@pytest.mark.parametrize("total", ["0", "-1", "NaN", "+Inf"]) +def test_missing_failure_families_require_observed_database_activity(total: str) -> None: + with pytest.raises(AssertionError): + _ = postgres_failures( + "# TYPE litellm_postgres_total_requests_total counter\n" + f'litellm_postgres_total_requests_total{{postgres="postgres"}} {total}\n' + ) + + +@pytest.mark.parametrize("body", ["", "# TYPE unrelated counter\n", metrics(failure="NaN"), metrics(batch="-1")]) +def test_missing_or_invalid_counters_do_not_pass_as_zero(body: str) -> None: + with pytest.raises(AssertionError): + _ = postgres_failures(body) + + +def test_index_completion_stops_after_the_post_upgrade_window() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + assert not monitor.poll((record(),)) + assert probe.releases == 1 + probe.builds = [17] + probe.now = 101 + assert not monitor.poll(()) + probe.found = frozenset({"response-0", "response-1"}) + probe.now = 103 + assert not monitor.poll((record(1, started_at=102),)) + assert monitor.overlap_writes == 1 + assert monitor.persisted_count == 2 + assert not monitor.pending + assert monitor.max_lock_waiters == 2 + probe.current = state("complete") + assert not monitor.poll(()) + probe.now = 162 + assert not monitor.poll(()) + probe.now = 163 + assert not monitor.poll(()) + probe.found = frozenset(f"response-{i}" for i in range(300)) + assert monitor.poll(tuple(record(i) for i in range(2, 300))) + assert probe.releases == 1 + assert probe.metric_checks == 7 + + +def test_http_200_without_spend_row_cannot_complete() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + records: Final = tuple(record(i) for i in range(300)) + _ = monitor.poll(records) + with pytest.raises(AssertionError, match="no matching spend log"): + monitor.assert_complete(state("complete", pending=False)) + + +def test_no_pending_index_work_still_requires_300_persisted_requests() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + records: Final = tuple(record(i) for i in range(300)) + probe.found = frozenset(f"response-{i}" for i in range(300)) + _ = monitor.poll(records) + monitor.assert_complete(state("complete", pending=False)) + assert monitor.successes == monitor.persisted_count == 300 + with pytest.raises(AssertionError, match="during index work"): + monitor.assert_complete(state("complete", pending=True)) + + +def test_postgres_failure_fails_even_when_inference_succeeded() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + probe.errors = True + with pytest.raises(AssertionError, match="PostgreSQL"): + _ = monitor.poll((record(),)) + + +def test_failed_migration_and_operational_timeout_are_not_success() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + probe.current = state("failed") + with pytest.raises(AssertionError, match="worker failed"): + _ = monitor.poll(()) + probe.current = state() + probe.now = MAX_TRAFFIC_SECONDS + with pytest.raises(AssertionError, match="operational timeout"): + _ = monitor.poll(()) + + +def test_upgrade_case_uses_5rps_with_the_same_request_and_memory_limits() -> None: + case: Final = upgrade_case() + assert case.traffic.rps == 5 + assert not {"cpu_seconds_per_request", "postgres_calls_per_request", "redis_calls_per_request"}.intersection( + limit.metric for limit in case.limits + ) + assert case.limits == tuple( + limit + for limit in CASE.limits + if limit.metric + in { + "unexpected_status_count", + "gateway_added_seconds", + "time_to_first_token_seconds", + "memory_growth_megabytes", + } + ) + + +def failed_record(seq: int) -> WireRecord: + return record(seq).model_copy(update={"status": 503, "ok": False, "text": "", "response_id": None, "call_id": None}) + + +def test_three_non_200s_do_not_start_indexes_or_require_spend_rows() -> None: + probe: Final = Probe() + monitor: Final = probe.monitor() + assert not monitor.poll(tuple(failed_record(i) for i in range(3))) + assert probe.releases == 0 + assert monitor.processed == 3 + assert monitor.successes == 0 + assert not monitor.pending + probe.found = frozenset(f"response-{i}" for i in range(3, 303)) + assert not monitor.poll(tuple(record(i) for i in range(3, 303))) + monitor.assert_complete(state("complete", pending=False)) + assert probe.releases == 1 + assert monitor.processed == 303 + assert monitor.successes == monitor.persisted_count == 300 + + +def test_fourth_non_200_fails_across_poll_batches() -> None: + monitor: Final = Probe().monitor() + assert not monitor.poll(tuple(failed_record(i) for i in range(3))) + with pytest.raises(AssertionError, match="Too many unexpected non-200"): + _ = monitor.poll((failed_record(3),)) + + +@pytest.mark.parametrize( + "failed", [record().model_copy(update={"text": ""}), record().model_copy(update={"ok": False})] +) +def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None: + with pytest.raises(AssertionError, match="Failed or empty inference"): + _ = Probe().monitor().poll((failed,)) + + +@pytest.mark.parametrize( + "failures,overhead,ttft,growth,expected", + [ + (3, 0.29, 3.7, 100, ()), + (4, 0.29, 3.7, 100, ("unexpected_status_count",)), + (3, 0.30, 3.7, 100, ("gateway_added_seconds percentile 50",)), + (3, 0.29, 3.8, 100, ("time_to_first_token_seconds percentile 99",)), + (3, 0.29, 3.7, 101, ("memory_growth_megabytes",)), + ], +) +def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples( + failures: int, overhead: float, ttft: float, growth: float, expected: tuple[str, ...] +) -> None: + nonstream: Final = tuple( + record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100) + ) + stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft}) + observed: Final = Observation( + records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))), + mock_requests=(), + guardrail_calls=(), + rss_before=RssCapture((RssReading("gateway", "host", 1, 1000),), ()), + rss_after=RssCapture((RssReading("gateway", "host", 1, 1000 + growth),), ()), + wall_seconds=60, + ) + rows, exceeded = evaluate(upgrade_case(), observed) + assert tuple(row.label for row in rows if row.over or row.missing) == expected + assert len(exceeded) == len(expected) From 9e9dccb9d4e906c3bfe380d3711f8ba5d2b26479 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 19:10:04 +0000 Subject: [PATCH 26/64] fix(register_model)!: keep bedrock geo deployments priced across price data reloads (#41523) * fix(register_model): keep bedrock geo deployments priced across price data reloads A request-time pricing registration inserting into litellm.model_cost while a price data reload walked it could abort the reload (dictionary changed size during iteration in add_known_models), or trip the lookup map rebuild inside get_model_info so the deployment replay fell through to the unmapped path and wrote a raw bedrock/au.anthropic. entry with no token pricing, a $0 shadow re-rolled on every later reload. Walk snapshots of the map, bind the case-insensitive lookup map locally, and let the unmapped path inherit the whole built-in entry of a provider get_model_info would accept. * fix(register_model): keep synthesized zero costs out of inherited built-in pricing An unmapped alias that inherits its built-in entry now takes only the token prices the catalogue really carries, so a synthesized zero never lands next to the caller's custom price. The router regression test reports the writer's failures instead of swallowing them, and the local names carry Final. * fix(register_model): match inherited built-ins by the key provider prefix when the value names none * fix(utils): resolve a tripped lookup from the key's own catalog row before its region sibling * fix(register_model): keep custom flat prices above inherited tier rates * fix(register_model): keep the other direction's tier rates when a custom flat price covers one * fix(register_model): keep a catalog row's tiers when its lookup trips, and snapshot request-path cost map walks A tripped get_model_info on a key that is itself a catalog row now inherits that row whole, so a custom flat price bills a long prompt at the row's above-threshold rate on both lookup paths instead of flipping until the next reload. The fal_ai, fireworks, ollama and gpt-5 request-path walks over litellm.model_cost iterate a snapshot so a concurrent reload cannot raise mid-request. * fix(register_model): price route-prefixed geo keys on a tripped lookup and replace cost map entries atomically * fix(register_model): move every alias of a cost map entry to its registered price * perf(register_model): read alias names from the shared entry instead of scanning the cost map Each registration scanned every cost map entry to find the names sharing it, about 1 ms on the 4,189-entry map, and a proxy request to a custom-priced deployment registers twice. Aliased entries now carry the names they are served under, so a registration reads them directly, and an alias the caller pointed at its own entry keeps that entry. The tests this PR added lose their docstrings * fix(utils): keep capability flags and write through on a tripped lookup in register_model A key the lookup cannot map now seeds its registered entry from the capability rules its unregistered lookup resolves to, so registering a region alias no longer turns off supports_system_messages and similar flags. A lookup that trips mid-walk on a key naming a catalog row writes into that row instead of leaving a shadow key that goes stale. Adds regression tests for both and for the request-path snapshot walks. * fix(utils): retry a tripped registration lookup and rebuild a registration a reload races A lookup a concurrent write trips is retried instead of writing into the catalog row its key names, so a per-request price never reprices a shared row. The inherited entry leaves the matched row's provider behind, so a provider-less azure_ai key stays priced, and a registration whose map a reload swapped mid-build is rebuilt from the new map. * fix(utils): write each registration into the map it was built from and rebuild it when a reload replaced that map * test(integration): price data reload keeps every registered deployment priced --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 2 +- .../litellm_core_utils/get_model_cost_map.py | 36 +- litellm/llms/fal_ai/cost_calculator.py | 2 +- .../llms/fireworks_ai/chat/transformation.py | 2 +- litellm/llms/ollama/common_utils.py | 2 +- .../llms/openai/chat/gpt_5_transformation.py | 4 +- litellm/utils.py | 396 ++++++++---- .../pricing/test_price_data_reload_prices.py | 433 +++++++++++++ .../test_register_model_custom_pricing.py | 608 +++++++++++++++++- .../unit/test_router_model_cost_isolation.py | 75 +++ tests/unit/test_utils.py | 119 +++- 11 files changed, 1546 insertions(+), 133 deletions(-) create mode 100644 tests/integration/pricing/test_price_data_reload_prices.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 959c9fb406d..9cc8965b628 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -742,7 +742,7 @@ def is_openai_finetune_model(key: str) -> bool: def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None: - for key, value in model_cost_map.items(): + for key, value in tuple(model_cost_map.items()): if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): open_ai_chat_completion_models.add(key) elif value.get("litellm_provider") == "text-completion-openai": diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index b68db22183a..073e3a3455e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -557,22 +557,39 @@ def get_model_cost_map_loaded_at() -> "datetime | None": return _cost_map_source_info.loaded_at -def _expand_model_aliases(model_cost: dict) -> dict: +class AliasedCostEntry(dict[str, object]): + __slots__ = ("names",) + + def __init__(self, entry: Mapping[str, object], names: tuple[str, ...]) -> None: + super().__init__(entry) + self.names: tuple[str, ...] = names + + +def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, ...]]: + return MappingProxyType( + { + owner: (owner, *(alias for alias, alias_owner in alias_owners.items() if alias_owner == owner)) + for owner in dict.fromkeys(alias_owners.values()) + } + ) + + +def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, dict[str, object]]: """ Expand ``aliases`` lists in model cost entries into top-level entries. - Each alias gets a reference to the **same** dict object as the canonical - entry (zero memory overhead). The ``aliases`` key is removed from the - entry so downstream code never sees it. + Each alias gets a reference to the **same** ``AliasedCostEntry`` as the + canonical entry, which carries every name it is served under. The + ``aliases`` key is removed from the entry so downstream code never sees it. If an alias collides with an existing canonical entry the alias is skipped and a warning is logged. """ - aliases_to_add: Final[dict[str, dict]] = {} + alias_owners: Final[dict[str, str]] = {} keys_with_aliases: Final[list[str]] = [] for model_name, model_info in model_cost.items(): - aliases: list | None = model_info.get("aliases") + aliases: object = model_info.get("aliases") if aliases is None: continue keys_with_aliases.append(model_name) @@ -594,7 +611,7 @@ def _expand_model_aliases(model_cost: dict) -> dict: model_name, ) continue - if alias in aliases_to_add: + if alias in alias_owners: verbose_logger.warning( "LiteLLM model alias conflict: alias '%s' (from '%s') " "was already claimed by another entry — skipping.", @@ -602,13 +619,14 @@ def _expand_model_aliases(model_cost: dict) -> dict: model_name, ) continue - aliases_to_add[alias] = model_info # same dict reference + alias_owners[alias] = model_name # Remove the ``aliases`` key from entries so it doesn't pollute model info for key in keys_with_aliases: model_cost[key].pop("aliases", None) - model_cost.update(aliases_to_add) + for owner, names in _alias_groups(alias_owners).items(): + model_cost.update(dict.fromkeys(names, AliasedCostEntry(model_cost[owner], names))) return model_cost diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index aa3d7942174..2acf1c69594 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -87,7 +87,7 @@ def _keyed_rows(model: str, quality: str) -> tuple[tuple[int, int, float], ...]: suffix: Final = f"/{model}" return tuple( (width, height, float(raw_cost)) - for key in litellm.model_cost + for key in (*litellm.model_cost,) if isinstance(key, str) and key.startswith(prefix) and key.endswith(suffix) for size in (key[len(prefix) : -len(suffix)],) for dimensions in (_parse_keyed_dimensions(size),) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 8b6946ee562..14ae59d4f85 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -544,7 +544,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): return cached[2] index: Final[list[tuple[str, dict]]] = [] - for key, model_info in model_cost.items(): + for key, model_info in (*model_cost.items(),): if not key.startswith("fireworks_ai/"): continue if not isinstance(model_info, dict): diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index b5328b1ba4d..284a6ffea62 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -197,7 +197,7 @@ class OllamaModelInfo(BaseLLMModelInfo): "ollama/" + stripped_model, "ollama_chat/" + stripped_model, } - model_cost_keys: Final = {key.lower() for key in model_cost} + model_cost_keys: Final = {key.lower() for key in (*model_cost,)} return any(name.lower() in model_cost_keys for name in potential_model_names) @staticmethod diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 94c5c5d016d..5bf19844287 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -28,7 +28,9 @@ def _catalogue_declares_default_effort() -> bool: the one litellm gave before it existed. Scanning costs ~80us on the largest published map and only on the fallback path, which is noise beside the request it precedes. """ - return any(isinstance(entry, dict) and "default_reasoning_effort" in entry for entry in litellm.model_cost.values()) + return any( + isinstance(entry, dict) and "default_reasoning_effort" in entry for entry in (*litellm.model_cost.values(),) + ) def _normalize_reasoning_effort_for_chat_completion( diff --git a/litellm/utils.py b/litellm/utils.py index fc3acd3065d..b37827a3b13 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -44,7 +44,7 @@ from httpx import Proxy from httpx._utils import get_environment_proxies from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # OpenAI parser module is private from openai.types.chat.completion_create_params import ResponseFormat -from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, with_config +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config import litellm import litellm.litellm_core_utils @@ -89,6 +89,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens @@ -3303,46 +3304,176 @@ _BEDROCK_REGION_PREFIXES: Final = ( "ap-northeast-1.", ) -_CACHE_PRICING_FIELDS: Final = ( - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost", - "cache_read_input_token_cost_above_200k_tokens", +_TOKEN_COST_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token") + + +def _builtin_cost_map() -> dict[str, Mapping[str, object]]: + return cast(dict[str, Mapping[str, object]], litellm.model_cost) + + +def _replace_cost_map_entry(cost_map: dict[str, Mapping[str, object]], key: str, entry: Mapping[str, object]) -> None: + current: Final = cost_map.get(key) + merged: Final = {**(current or {}), **entry} + if not isinstance(current, AliasedCostEntry): + cost_map.update({key: merged}) + return + sharing_names: Final = tuple(name for name in current.names if cost_map.get(name) is current) + cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, current.names))) + + +def _without_synthesized_costs( + model_info: Mapping[str, object], key: str, value: Mapping[str, object] +) -> Mapping[str, object]: + """``_get_model_info_helper`` synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` = 0 when the raw + entry has neither (#30198). Persisting those zeros would turn a name-priced entry into a free one and make + ``_is_cost_explicitly_configured`` disable budget enforcement on the next re-registration.""" + builtin_cost_map: Final = _builtin_cost_map() + raw_fields: Final = frozenset(builtin_cost_map.get(str(model_info["key"]), builtin_cost_map.get(key, {}))) + return MappingProxyType( + { + name: field_value + for name, field_value in model_info.items() + if name not in _TOKEN_COST_FIELDS or name in raw_fields or name in value + } + ) + + +_DIRECTION_RATE_FIELDS: Final = MappingProxyType( + { + "input_cost_per_token": ("input_cost_per_token",), + "output_cost_per_token": ("output_cost_per_token", "output_cost_per_reasoning_token"), + } ) +_RATE_TABLE: Final = TypeAdapter(Mapping[str, object]) +_RATE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) -def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, object] | None: +def _is_overridden_rate(name: str, overridden_fields: tuple[str, ...]) -> bool: + is_variant: Final = name.startswith(tuple(f"{rate_field}_" for rate_field in overridden_fields)) + return name in overridden_fields or (is_variant and "_cache_" not in name) + + +def _without_overridden_rates(rates: Mapping[str, object], overridden_fields: tuple[str, ...]) -> Mapping[str, object]: + return MappingProxyType( + {name: rate for name, rate in rates.items() if not _is_overridden_rate(name, overridden_fields)} + ) + + +def _rate_table_with_flat_prices( + name: str, table: object, overridden_fields: tuple[str, ...], flat_prices: Mapping[str, object] +) -> object: + try: + if name == "tiered_pricing": + return [ + {**_without_overridden_rates(tier, overridden_fields), **flat_prices} + for tier in _RATE_TIERS.validate_python(table) + ] + if name == "off_peak_pricing": + return dict(_without_overridden_rates(_RATE_TABLE.validate_python(table), overridden_fields)) + except ValidationError: + return table + return table + + +def _without_overridden_price_variants( + entry: Mapping[str, object], value: Mapping[str, object] +) -> Mapping[str, object]: + """A caller's flat ``input_cost_per_token`` / ``output_cost_per_token`` is that direction's price above every + token threshold, for every service tier and batch, in every ``tiered_pricing`` tier and off-peak window, so each + inherited rate that would stand in for it there, the reasoning-token rate included for output, is dropped, and + each tier carries the flat price instead.""" + flat_prices: Final = MappingProxyType( + {cost_field: value[cost_field] for cost_field in _TOKEN_COST_FIELDS if value.get(cost_field) is not None} + ) + overridden_fields: Final = tuple( + itertools.chain.from_iterable(_DIRECTION_RATE_FIELDS[cost_field] for cost_field in flat_prices) + ) + if not overridden_fields: + return entry + return MappingProxyType( + { + name: _rate_table_with_flat_prices(name, field_value, overridden_fields, flat_prices) + for name, field_value in entry.items() + if not _is_overridden_rate(name, overridden_fields) + } + ) + + +@dataclass(frozen=True, slots=True) +class _BuiltinCostMatch: + key: str + entry: Mapping[str, object] + names_the_row: bool + + +@dataclass(frozen=True, slots=True) +class _LookupMiss: + tripped: bool + + +@dataclass(frozen=True, slots=True) +class _FallbackRegistration: + entry: Mapping[str, object] + matched_builtin: bool + + +def _fallback_registration(key: str, provider: str, value: Mapping[str, object]) -> _FallbackRegistration: + """The key keeps its own entry, seeded from the capability rules an unregistered lookup resolves it to, with the + matched built-in entry on top. The matched row's provider stays behind, since the key's own lookup can reject it + (an ``azure_ai/`` key never resolves to an ``openai`` entry), and the key keeps the provider its caller gave.""" + match: Final = _resolve_builtin_model_cost_entry(key=key, provider=provider) + generalized: Final[Mapping[str, object]] = match_capability_generalizations(key) or {} + if match is None: + return _FallbackRegistration(entry=generalized, matched_builtin=False) + inherited: Final = match.entry if match.names_the_row else _without_overridden_price_variants(match.entry, value) + own_fields: Final = {name: field_value for name, field_value in inherited.items() if name != "litellm_provider"} + return _FallbackRegistration(entry={**generalized, **own_fields}, matched_builtin=True) + + +def _resolve_builtin_model_cost_entry(key: str, provider: str) -> _BuiltinCostMatch | None: """Best-effort lookup of a built-in ``model_cost`` entry for a custom key whose shape ``get_model_info`` cannot resolve (repeated provider prefixes like ``bedrock/bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region - aliases). + aliases), or whose lookup tripped. - Returns a copy of the matching entry so the caller can inherit its defaults - (most importantly cache pricing) without mutating the shared built-in. - Returns ``None`` when no safe match exists. + The key itself is tried first, then the key behind its provider prefixes + and bedrock route, so a catalog row whose lookup tripped inherits from + itself rather than from a region-stripped sibling priced at the base rate. + ``names_the_row`` is set when the match is the row the key names, which + ``get_model_info`` resolves it to, so the caller keeps that row's rates as + an untripped registration does. Returns ``None`` when no candidate exists + under a provider ``get_model_info`` would accept for ``provider``, or, when + the value names none, for the provider prefix the key itself carries, + which is what ``get_model_info`` derives when it reads the key back. """ - candidates: Final[list[str]] = [] + from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix + segments: Final = key.split("/") - idx = 0 - while idx < len(segments) - 1 and segments[idx] in LlmProvidersSet: - idx += 1 - candidates.append("/".join(segments[idx:])) - - base: Final = candidates[-1] if candidates else key - for region_prefix in _BEDROCK_REGION_PREFIXES: - if base.startswith(region_prefix): - candidates.append(base[len(region_prefix) :]) - - if provider: - stripped: Final = _strip_model_name(model=base, custom_llm_provider=provider) - if stripped != base: - candidates.append(stripped) + match_provider: Final = provider or (segments[0] if len(segments) > 1 and segments[0] in LlmProvidersSet else "") + provider_prefix_count: Final = next( + (idx for idx, segment in enumerate(segments[:-1]) if segment not in LlmProvidersSet), len(segments) - 1 + ) + prefix_stripped: Final = tuple("/".join(segments[idx:]) for idx in range(1, provider_prefix_count + 1)) + prefix_free: Final = prefix_stripped[-1] if prefix_stripped else key + base: Final = ( + strip_bedrock_routing_prefix(prefix_free) if match_provider in ("bedrock", "bedrock_converse") else prefix_free + ) + region_stripped: Final = tuple( + base[len(region_prefix) :] for region_prefix in _BEDROCK_REGION_PREFIXES if base.startswith(region_prefix) + ) + stripped: Final = _strip_model_name(model=base, custom_llm_provider=match_provider) if match_provider else base + row_names: Final = (key, *prefix_stripped, base, stripped) + candidates: Final = (key, *prefix_stripped, base, *region_stripped, stripped) + builtin_cost_map: Final = _builtin_cost_map() for candidate in candidates: - entry = litellm.model_cost.get(candidate) - if entry is not None and entry.get("litellm_provider") is not None: - return dict(entry) + entry = builtin_cost_map.get(candidate) + if ( + entry is not None + and entry.get("litellm_provider") is not None + and _check_provider_match(dict(entry), match_provider) + ): + return _BuiltinCostMatch(key=candidate, entry=dict(entry), names_the_row=candidate in row_names) return None @@ -3360,19 +3491,115 @@ def is_generalized_model_info(model_info: ModelInfo) -> bool: return key not in litellm.model_cost and match_capability_generalizations(key) is not None -def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | None: - """Resolve ``model`` to its built-in cost-map entry for registration merging. +_REGISTRATION_ATTEMPTS: Final = 3 - Returns ``None`` when the lookup raises or when it resolved via a - fallback-generalization capability rule. A rule-derived entry carries no - pricing, so treating it as a hit would skip the built-in cache-pricing - inheritance for prefix-mangled keys. - """ + +def _lookup_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | _LookupMiss: try: info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except ModelNotMappedError: + return _LookupMiss(tripped=False) except Exception: - return None - return None if is_generalized_model_info(info) else info + return _LookupMiss(tripped=True) + return _LookupMiss(tripped=False) if is_generalized_model_info(info) else info + + +def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | _LookupMiss: + """Resolve ``model`` to its built-in cost-map entry for registration merging. + + Misses when the lookup cannot map the model or resolved it via a + fallback-generalization capability rule, since a rule-derived entry carries no + pricing. A lookup that raised anything else, such as a concurrent registration + changing the map mid-walk, is retried, so a transient failure registers exactly + where the untripped lookup would. + """ + lookups: Final = (_lookup_for_registration(model, custom_llm_provider) for _ in range(_REGISTRATION_ATTEMPTS)) + return next( + (lookup for lookup in lookups if not (isinstance(lookup, _LookupMiss) and lookup.tripped)), + _LookupMiss(tripped=True), + ) + + +@dataclass(frozen=True, slots=True) +class _RegistrationWrite: + key: str + entry: Mapping[str, object] + + +def _registration_base( + key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None +) -> _RegistrationWrite: + provider: Final = str(value.get("litellm_provider") or "") + if provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO or any( + key.startswith(f"{p}/") for p in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO + ): + return _RegistrationWrite(key=key, entry=dict(_builtin_cost_map().get(key, {}))) + builtin_model_info: Final = _get_builtin_model_info_for_registration( + model=key, custom_llm_provider=custom_llm_provider + ) + if not isinstance(builtin_model_info, _LookupMiss): + return _RegistrationWrite( + key=builtin_model_info["key"], + entry=dict(_without_synthesized_costs(builtin_model_info, key=key, value=value)), + ) + fallback: Final = _fallback_registration(key=key, provider=provider, value=value) + if not fallback.matched_builtin and ( + value.get("cache_creation_input_token_cost") is None + and value.get("cache_read_input_token_cost") is None + and value.get("tiered_pricing") is None + and (value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None) + ): + verbose_logger.warning( + "register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info", + warning_display_name or key, + ) + return _RegistrationWrite(key=key, entry=dict(fallback.entry)) + + +def _registration_write( + key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None +) -> _RegistrationWrite: + base: Final = _registration_base(key, value, custom_llm_provider, warning_display_name) + # ``get_model_info`` returns ``litellm_provider: None`` when the + # provider is unknown (e.g. custom deployments registered via + # ``Router.add_deployment``). Persisting that None into + # ``litellm.model_cost`` causes ``_check_provider_match`` to drop + # custom pricing on subsequent cost lookups. + existing_model: Final = { + name: field_value + for name, field_value in base.entry.items() + if name != "litellm_provider" or field_value is not None + } + return _RegistrationWrite(key=base.key, entry=update_dictionary(existing_model, value)) + + +def _registration_attempt( + key: str, + value: Mapping[str, object], + custom_llm_provider: str | None, + warning_display_name: str | None, + final_attempt: bool, +) -> _RegistrationWrite | None: + source: Final = _builtin_cost_map() + write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name) + _replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry) + return write if final_attempt or litellm.model_cost is source else None + + +def _register_on_current_map( + key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None +) -> _RegistrationWrite: + """A price data reload that swaps the map while the entry is being built or written would otherwise get the old + map's prices written over its own, or drop the entry with the old map, so each attempt writes into the map it was + built from and the entry is rebuilt when a reload replaced that map before the write finished. The last attempt + writes into the current map, so the entry is never lost.""" + attempts: Final = ( + _registration_attempt( + key, value, custom_llm_provider, warning_display_name, final_attempt=attempt == _REGISTRATION_ATTEMPTS - 1 + ) + for attempt in range(_REGISTRATION_ATTEMPTS) + ) + return next(write for write in attempts if write is not None) _runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload @@ -3492,79 +3719,15 @@ def register_model( for _registered_key, _registered_value in _registrations.items(): _runtime_registered_model_cost[_registered_key] = dict(_registered_value) - _skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO - for key, value in loaded_model_cost.items(): - ## get model info ## - provider = value.get("litellm_provider", "") - _key_str = str(key) - if provider in _skip_get_model_info_providers or any( - _key_str.startswith(f"{p}/") for p in _skip_get_model_info_providers - ): - existing_model = litellm.model_cost.get(key, {}) - model_cost_key = key - else: - builtin_model_info = _get_builtin_model_info_for_registration( - model=_key_str, custom_llm_provider=custom_llm_provider - ) - if builtin_model_info is not None: - existing_model = cast(dict, builtin_model_info) - model_cost_key = existing_model["key"] - else: - # An exact entry ends the lookup ladder before the capability rules are - # consulted, so seed from them: otherwise registering an unmapped model - # shadows the very defaults it would have resolved to unregistered. - existing_model = dict(match_capability_generalizations(_key_str) or {}) - model_cost_key = key - builtin_entry = _resolve_builtin_model_cost_entry(key=_key_str, provider=provider) - if builtin_entry is not None: - for field in _CACHE_PRICING_FIELDS: - if value.get(field) is None and builtin_entry.get(field) is not None: - existing_model[field] = builtin_entry[field] - elif ( - value.get("cache_creation_input_token_cost") is None - and value.get("cache_read_input_token_cost") is None - and value.get("tiered_pricing") is None - and ( - value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None - ) - ): - verbose_logger.warning( - "register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info", - warning_display_name or key, - ) - # ``get_model_info`` returns ``litellm_provider: None`` when the - # provider is unknown (e.g. custom deployments registered via - # ``Router.add_deployment``). Persisting that None into - # ``litellm.model_cost`` causes ``_check_provider_match`` to drop - # custom pricing on subsequent cost lookups. - if existing_model.get("litellm_provider") is None: - existing_model.pop("litellm_provider", None) - # Same pattern for cost fields (#30198): ``_get_model_info_helper`` - # synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` - # = 0 when they are absent from the raw entry. Writing those zeros - # back flips a sparse entry from "no cost keys" (priced via name) - # to "cost keys = 0" (free), which makes - # ``_is_cost_explicitly_configured`` return True and silently - # disables budget enforcement on the next re-registration. - _raw_entry = litellm.model_cost.get(model_cost_key) - if _raw_entry is None: - _raw_entry = litellm.model_cost.get(key) - if _raw_entry is None: - _raw_entry = {} - for _cost_field in ("input_cost_per_token", "output_cost_per_token"): - if _cost_field not in _raw_entry and _cost_field not in value: - existing_model.pop(_cost_field, None) - ## override / add new keys to the existing model cost dictionary - updated_dictionary = update_dictionary(existing_model, value) - litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary) + write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name) # Invalidate case-insensitive lookup map since model_cost was modified _invalidate_model_cost_lowercase_map() - verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", model_cost_key, model_cost_key) + verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", write.key, write.key) if add_to_provider_model_sets and value.get("litellm_provider") == "openrouter": - litellm.openrouter_models.add(_key_str.split("/", 1)[-1]) + litellm.openrouter_models.add(str(key).split("/", 1)[-1]) if add_to_provider_model_sets: litellm.add_known_models(loaded_model_cost) return model_cost @@ -5732,7 +5895,7 @@ def _rebuild_model_cost_lowercase_map() -> dict[str, str]: The rebuilt map (guaranteed to be not None). """ global _model_cost_lowercase_map - _model_cost_lowercase_map = {k.lower(): k for k in litellm.model_cost} + _model_cost_lowercase_map = {k.lower(): k for k in tuple(litellm.model_cost)} return _model_cost_lowercase_map @@ -5747,9 +5910,7 @@ def _handle_stale_map_entry_rebuild( Returns: The matched key if found after rebuild, None otherwise. """ - global _model_cost_lowercase_map - _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() - matched_key: Final = _model_cost_lowercase_map.get(potential_key_lower) + matched_key: Final = _rebuild_model_cost_lowercase_map().get(potential_key_lower) if matched_key is not None and matched_key in litellm.model_cost: return matched_key return None @@ -5767,7 +5928,7 @@ def _handle_new_key_with_scan( The matched key if found, None otherwise. """ global _model_cost_lowercase_map - for key in litellm.model_cost: + for key in tuple(_builtin_cost_map()): if key.lower() == potential_key_lower: _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() return key @@ -5789,18 +5950,19 @@ def get_model_cost_key(potential_key: str) -> str | None: called and confirmed not to cause performance issues, add it to the allowed_helpers list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py """ - global _model_cost_lowercase_map - # Exact match (O(1)) if potential_key in litellm.model_cost: return potential_key - # Case-insensitive lookup via map (O(1)) - if _model_cost_lowercase_map is None: - _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() + # Case-insensitive lookup via map (O(1)); bound locally because another thread's + # registration can invalidate the module-level map at any point + cached_lowercase_map: Final = _model_cost_lowercase_map + lowercase_map: Final = ( + cached_lowercase_map if cached_lowercase_map is not None else _rebuild_model_cost_lowercase_map() + ) potential_key_lower: Final = potential_key.lower() - matched_key = _model_cost_lowercase_map.get(potential_key_lower) + matched_key = lowercase_map.get(potential_key_lower) # Verify key exists (O(1) - handles model_cost.pop() case) if matched_key is not None and matched_key in litellm.model_cost: diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py new file mode 100644 index 00000000000..bb24130aa6c --- /dev/null +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -0,0 +1,433 @@ +import itertools +import json +import re +import threading +import uuid +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_ROW: Final = "anthropic.claude-audit-reload-v1:0" +_GEO: Final = f"au.{_ROW}" +_FINETUNE: Final = "ft:gpt-audit-reload-2025-01-01" +_CLAUDE: Final = "claude-audit-reload-sonnet" +_MINI: Final = "gpt-audit-reload-mini" +_USAGE: Final = re.compile(r"bill (\d+) (\d+) (\d+)") +_CUSTOM: Final = MappingProxyType({"input_cost_per_token": 1.1e-06, "output_cost_per_token": 2.3e-06}) +_COST_MAP_LIMITS: Final = MappingProxyType({"MODEL_COST_MAP_MIN_MODEL_COUNT": "1", "MODEL_COST_MAP_MAX_SHRINK_RATIO": "0"}) +_PUBLISHED_SCALE: Final = 2.0 +_CATALOG: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + _ROW: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + }, + _GEO: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + }, + _FINETUNE: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 3.75e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 1.875e-06, + }, + _CLAUDE: { + "litellm_provider": "anthropic", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + }, + _MINI: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2e-06, + }, + } +) + + +@dataclass(frozen=True, slots=True) +class _Usage: + prompt: int + cached: int + completion: int + + def text(self) -> str: + return f"bill {self.prompt} {self.cached} {self.completion}" + + +@dataclass(frozen=True, slots=True) +class _Rates: + input: float + output: float + cache_read: float + + def cost(self, usage: _Usage) -> float: + return ( + (usage.prompt - usage.cached) * self.input + usage.cached * self.cache_read + usage.completion * self.output + ) + + +@dataclass(frozen=True, slots=True) +class _Deployment: + name: str + model: str + prices: Mapping[str, float] + usage: _Usage + rates: Callable[[float], _Rates | None] + + +@dataclass(frozen=True, slots=True) +class _Charge: + request_id: str + expected: float + + +def _catalog(scale: float) -> dict[str, dict[str, JsonValue]]: + return { + model: {name: value * scale if isinstance(value, float) else value for name, value in row.items()} + for model, row in _CATALOG.items() + } + + +def _rate(model: str, field: str, scale: float) -> float: + rate: Final = _CATALOG[model][field] + assert isinstance(rate, float), rate + return rate * scale + + +def _catalog_rates(model: str) -> Callable[[float], _Rates]: + return lambda scale: _Rates( + input=_rate(model, "input_cost_per_token", scale), + output=_rate(model, "output_cost_per_token", scale), + cache_read=_rate(model, "cache_read_input_token_cost", scale) + if "cache_read_input_token_cost" in _CATALOG[model] + else _rate(model, "input_cost_per_token", scale), + ) + + +def _custom_rates(_scale: float) -> _Rates: + return _Rates( + input=_CUSTOM["input_cost_per_token"], + output=_CUSTOM["output_cost_per_token"], + cache_read=_CUSTOM["input_cost_per_token"], + ) + + +def _unpriced(_scale: float) -> None: + return None + + +def _expected_cost(rates: _Rates | None, usage: _Usage) -> float: + return 0.0 if rates is None else rates.cost(usage) + + +def _expected_listing(rates: _Rates | None) -> tuple[object, object]: + if rates is None: + return None, None + return pytest.approx(rates.input), pytest.approx(rates.output) + + +def _connection(model: str, wire_url: str) -> dict[str, JsonValue]: + if model.startswith("bedrock/"): + return {"api_base": wire_url, "api_key": "synthetic-bedrock-bearer", "aws_region_name": "ap-northeast-1"} + if model.startswith("azure/"): + return {"api_base": wire_url, "api_key": "synthetic-azure-key", "api_version": "2025-01-01-preview"} + return {"api_base": f"{wire_url}/v1", "api_key": "synthetic-openai-key"} + + +def _config(directory: Path, deployments: tuple[_Deployment, ...], wire_url: str) -> Path: + config: Final = directory / "price_data_reload_config.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": deployment.name, + "litellm_params": { + "model": deployment.model, + **_connection(deployment.model, wire_url), + **deployment.prices, + }, + } + for deployment in deployments + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + "router_settings": {"disable_cooldowns": True}, + } + ) + ) + return config + + +def _usage_in(request: Request) -> _Usage: + found: Final = _USAGE.search(request.body.decode()) + assert found is not None, request + return _Usage(prompt=int(found.group(1)), cached=int(found.group(2)), completion=int(found.group(3))) + + +def _reply(request: Request) -> Reply: + usage: Final = _usage_in(request) + if "/converse" in request.target: + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "priced"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": usage.prompt, + "outputTokens": usage.completion, + "totalTokens": usage.prompt + usage.completion, + }, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + if "/invoke" in request.target: + return Reply( + body=json.dumps( + { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": [{"type": "text", "text": "priced"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": usage.prompt, "output_tokens": usage.completion}, + } + ).encode() + ) + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "scripted", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "priced"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": usage.prompt, + "completion_tokens": usage.completion, + "total_tokens": usage.prompt + usage.completion, + "prompt_tokens_details": {"cached_tokens": usage.cached}, + }, + } + ).encode() + ) + + +def _responder(published: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/model_prices.json": + return Reply(body=json.dumps(_catalog(_PUBLISHED_SCALE if published.is_set() else 1.0)).encode()) + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + return _reply(request) + + return respond + + +def _listed_rates(proxy: Gateway, name: str) -> tuple[JsonValue, JsonValue]: + entries: Final = proxy.get("/model/info")["data"] + assert isinstance(entries, list), entries + info: Final = next( + object_value(object_value(entry)["model_info"]) + for entry in entries + if object_value(entry)["model_name"] == name + ) + return info.get("input_cost_per_token"), info.get("output_cost_per_token") + + +def _billed(client: httpx.Client, key: str, deployment: _Deployment) -> tuple[str, float | None]: + response: Final = client.post( + "/v1/chat/completions", + json={"model": deployment.name, "messages": [{"role": "user", "content": deployment.usage.text()}]}, + headers={"Authorization": f"Bearer {key}"}, + ) + assert response.status_code == 200, response.text + reported: Final = response.headers.get("x-litellm-response-cost") + return string_value(object_value(response.json())["id"]), None if reported is None else float(reported) + + +def _charged_as_listed(proxy: Gateway, deployment: _Deployment, scale: float) -> _Charge: + rates: Final = deployment.rates(scale) + assert _listed_rates(proxy, deployment.name) == _expected_listing(rates), deployment.name + request_id, charged = _billed(proxy.client, proxy.key, deployment) + expected: Final = _expected_cost(rates, deployment.usage) + assert charged == (None if rates is None else pytest.approx(expected, rel=1e-6)), deployment.name + return _Charge(request_id=request_id, expected=expected) + + +def _spend(request_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + spend: Final = rows[0]["spend"] + assert isinstance(spend, (int, float)), rows + return float(spend) + + +def _proxy_environment(wire_url: str) -> dict[str, str]: + return {**_COST_MAP_LIMITS, "LITELLM_MODEL_COST_MAP_URL": f"{wire_url}/model_prices.json"} + + +def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments: tuple[_Deployment, ...]) -> None: + published: Final = threading.Event() + with ( + wire_server(_responder(published)) as wire, + owned_proxy( + gateway, + tmp_path, + _proxy_environment(wire.url), + config=_config(tmp_path, deployments, wire.url), + remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",), + ) as proxy, + ): + before: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in deployments) + published.set() + reload: Final = proxy.request("POST", "/reload/model_cost_map") + assert reload.status_code == 200, reload.text + after: Final = tuple(_charged_as_listed(proxy, deployment, _PUBLISHED_SCALE) for deployment in deployments) + for charge in (*before, *after): + assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_price_data_reload_reprices_catalog_deployments_and_keeps_custom_prices( + gateway: Gateway, tmp_path: Path +) -> None: + _assert_priced_across_a_reload( + gateway, + tmp_path, + ( + _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)), + _Deployment("mini-custom", f"openai/{_MINI}", _CUSTOM, _Usage(1000, 0, 500), _custom_rates), + ), + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize( + "deployment", + ( + _Deployment( + "region-unpriced", f"bedrock/ap-northeast-1.{_ROW}", {}, _Usage(1000, 0, 500), _catalog_rates(_ROW) + ), + _Deployment( + "finetune-unpriced", + f"openai/{_FINETUNE}:org:suffix:y", + {}, + _Usage(2000, 1000, 1000), + _catalog_rates(_FINETUNE), + ), + _Deployment("finetune-custom", f"openai/{_FINETUNE}:org:suffix:z", _CUSTOM, _Usage(2000, 0, 1000), _custom_rates), + _Deployment("cross-provider-unpriced", f"azure/{_CLAUDE}", {}, _Usage(2000, 1000, 1000), _unpriced), + _Deployment( + "region-custom-above-threshold", + f"bedrock/ap-northeast-1.{_ROW}", + _CUSTOM, + _Usage(250_000, 0, 1000), + _custom_rates, + ), + ), + ids=lambda deployment: deployment.name, +) +def test_deployment_bills_the_catalog_row_its_registration_resolves_to_across_a_reload( + gateway: Gateway, tmp_path: Path, deployment: _Deployment +) -> None: + _assert_priced_across_a_reload(gateway, tmp_path, (deployment,)) + + +def _storm_deployment(index: int) -> _Deployment: + prices: Final = MappingProxyType( + {"input_cost_per_token": (index + 1) * 1e-07, "output_cost_per_token": (index + 1) * 3e-07} + ) + return _Deployment( + f"storm-{index}", + f"openai/gpt-audit-reload-storm-{index}", + prices, + _Usage(100, 0, 50), + lambda _scale: _Rates( + input=prices["input_cost_per_token"], + output=prices["output_cost_per_token"], + cache_read=prices["input_cost_per_token"], + ), + ) + + +def _storm_traffic( + base_url: str, key: str, deployments: tuple[_Deployment, ...], offset: int, finished: threading.Event +) -> tuple[tuple[float, float], ...]: + rotation: Final = ( + deployments[(offset + turn) % len(deployments)] for turn, _ in enumerate(iter(finished.is_set, True)) + ) + with httpx.Client(base_url=base_url, timeout=60, trust_env=False) as client: + return tuple( + (_billed(client, key, deployment)[1], _expected_cost(deployment.rates(1.0), deployment.usage)) + for deployment in rotation + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_price_data_reloads_under_custom_priced_traffic_keep_every_deployment_priced( + gateway: Gateway, tmp_path: Path +) -> None: + storm: Final = tuple(_storm_deployment(index) for index in range(24)) + geo: Final = _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)) + finished: Final = threading.Event() + with ( + wire_server(_responder(threading.Event())) as wire, + owned_proxy( + gateway, + tmp_path, + _proxy_environment(wire.url), + config=_config(tmp_path, (geo, *storm), wire.url), + remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",), + ) as proxy, + ThreadPoolExecutor(max_workers=8) as pool, + ): + traffic: Final = tuple( + pool.submit(_storm_traffic, str(proxy.client.base_url), proxy.key, storm, offset, finished) + for offset in range(8) + ) + reloads: Final = tuple(proxy.request("POST", "/reload/model_cost_map") for _ in range(40)) + finished.set() + charges: Final = tuple(itertools.chain.from_iterable(calls.result() for calls in traffic)) + assert tuple(reload.status_code for reload in reloads) == (200,) * 40, tuple( + reload.text for reload in reloads if reload.status_code != 200 + ) + assert len(charges) >= 8 + assert tuple(charged for charged, _ in charges) == pytest.approx(tuple(expected for _, expected in charges)) + settled: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in (geo, *storm)) + for charge in settled: + assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6) diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index 02c90e0d97b..35fbd28eae7 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -10,13 +10,16 @@ calculations for DB-sourced models with prompt caching pricing. """ import copy +import itertools import os -from typing import Final +from collections.abc import Callable, Mapping +from typing import Final, NoReturn import pytest import litellm +from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases, adopt_model_cost_map from litellm.main import _build_custom_pricing_entry from litellm.utils import _invalidate_model_cost_lowercase_map @@ -796,6 +799,609 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key(): _invalidate_model_cost_lowercase_map() +@pytest.mark.parametrize("route", ["", "converse/", "invoke/"]) +def test_register_model_prices_a_geo_alias_from_its_builtin_when_the_lookup_fails( + monkeypatch: pytest.MonkeyPatch, route: str +) -> None: + from litellm.types.utils import ModelResponse, Usage + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + canonical_key: Final = "au.anthropic.claude-sonnet-5" + alias: Final = f"bedrock/{route}{canonical_key}" + canonical: Final = litellm.model_cost[canonical_key] + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration) + litellm.register_model( + {alias: {"litellm_provider": "bedrock", "mode": "chat", "supports_vision": True}}, + persist_across_reloads=False, + ) + + usage: Final = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) + alias_cost: Final = litellm.completion_cost( + completion_response=ModelResponse(model=alias, usage=usage), model=alias, custom_llm_provider="bedrock" + ) + canonical_cost: Final = litellm.completion_cost( + completion_response=ModelResponse(model=canonical_key, usage=usage), + model=canonical_key, + custom_llm_provider="bedrock", + ) + try: + assert alias_cost == canonical_cost > 0 + assert litellm.get_model_info(alias)["input_cost_per_token"] == canonical["input_cost_per_token"] + assert litellm.get_model_info(alias)["supports_vision"] is True + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None: + geo_key: Final = "au.anthropic.lit-5853-geo-model" + geo_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "cache_read_input_token_cost": 3.3e-07, + } + sibling_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + } + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + geo_key: dict(geo_row), + "anthropic.lit-5853-geo-model": sibling_row, + }, + ) + _invalidate_model_cost_lowercase_map() + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration) + litellm.register_model({geo_key: {"mode": "chat"}}, persist_across_reloads=False) + + try: + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + assert litellm.model_cost[geo_key][field] == geo_row[field] + assert litellm.get_model_info(geo_key)[field] == geo_row[field] + assert litellm.model_cost[geo_key]["litellm_provider"] == geo_row["litellm_provider"] + finally: + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize("route", ["", "converse/", "invoke/"]) +def test_register_model_leaves_no_shadow_key_when_the_lookup_trips_once( + monkeypatch: pytest.MonkeyPatch, route: str +) -> None: + catalog_key: Final = "au.anthropic.lit-5853-geo-model" + alias: Final = f"bedrock/{route}{catalog_key}" + repriced_input_cost: Final = 9e-06 + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + catalog_key: { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + }, + }, + ) + _invalidate_model_cost_lowercase_map() + real_lookup: Final = litellm.utils.get_model_info + lookups: Final = itertools.count() + + def lookup_tripped_once_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> object: + if next(lookups) == 0: + raise RuntimeError("dictionary changed size during iteration") + return real_lookup(model=model, custom_llm_provider=custom_llm_provider) + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_once_by_a_concurrent_registration) + litellm.register_model({alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False) + litellm.register_model({catalog_key: {"input_cost_per_token": repriced_input_cost}}, persist_across_reloads=False) + + try: + assert litellm.get_model_info(alias)["key"] == catalog_key + assert litellm.get_model_info(alias)["input_cost_per_token"] == repriced_input_cost + finally: + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize("key_prefix", ["", "bedrock/", "bedrock/converse/"]) +def test_register_model_bills_a_catalog_rows_tier_rates_the_same_when_the_lookup_fails( + monkeypatch: pytest.MonkeyPatch, key_prefix: str +) -> None: + from litellm.types.utils import ModelResponse, Usage + + catalog_key: Final = "au.anthropic.lit-5853-tiered-geo-model" + registered_key: Final = f"{key_prefix}{catalog_key}" + catalog_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + "input_cost_per_token_above_200k_tokens": 6.6e-06, + "output_cost_per_token_above_200k_tokens": 2.475e-05, + } + builtin_cost_map: Final = litellm.get_model_cost_map(url="") + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + def long_prompt_cost_after_registering_with(lookup: Callable[..., object]) -> float: + monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)}) + _invalidate_model_cost_lowercase_map() + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup) + litellm.register_model( + { + registered_key: { + "litellm_provider": "bedrock", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + return litellm.completion_cost( + completion_response=ModelResponse( + model=registered_key, + usage=Usage(prompt_tokens=250_000, completion_tokens=1_000, total_tokens=251_000), + ), + model=registered_key, + custom_llm_provider="bedrock", + ) + + try: + assert long_prompt_cost_after_registering_with(lookup_tripped_by_a_concurrent_registration) == pytest.approx( + long_prompt_cost_after_registering_with(litellm.utils.get_model_info) + ) + finally: + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize("lookup_trips", [False, True]) +def test_register_model_prices_a_provider_less_azure_ai_key_without_repricing_the_openai_row( + monkeypatch: pytest.MonkeyPatch, lookup_trips: bool +) -> None: + openai_key: Final = "lit-5853-openai-model" + azure_ai_key: Final = f"azure_ai/{openai_key}" + openai_row: Final = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 5e-06, + "output_cost_per_token": 1.5e-05, + } + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://lit-5853.openai.azure.com") + monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), openai_key: dict(openai_row)}) + monkeypatch.setattr( + litellm, "open_ai_chat_completion_models", {*litellm.open_ai_chat_completion_models, openai_key} + ) + _invalidate_model_cost_lowercase_map() + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + with pytest.MonkeyPatch.context() as lookup_patch: + if lookup_trips: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration) + litellm.register_model( + {azure_ai_key: {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}}, + persist_across_reloads=False, + ) + + try: + assert litellm.model_cost[openai_key] == openai_row + assert litellm.cost_per_token(model=azure_ai_key, prompt_tokens=1000, completion_tokens=1000) == pytest.approx( + (1e-03, 2e-03) + ) + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monkeypatch: pytest.MonkeyPatch) -> None: + catalog_key: Final = "au.anthropic.lit-5853-reloaded-model" + deployment_key: Final = f"bedrock/{catalog_key}" + catalog_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + } + reloaded_input_cost: Final = 6.6e-06 + builtin_cost_map: Final = litellm.get_model_cost_map(url="") + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)}) + _invalidate_model_cost_lowercase_map() + reloaded_cost_map: Final = { + **builtin_cost_map, + catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost}, + } + real_lookup: Final = litellm.utils.get_model_info + + def lookup_racing_a_price_data_reload(model: str, custom_llm_provider: str | None = None) -> object: + info: Final = real_lookup(model=model, custom_llm_provider=custom_llm_provider) + if litellm.model_cost is not reloaded_cost_map: + adopt_model_cost_map(reloaded_cost_map) + return info + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_racing_a_price_data_reload) + litellm.register_model({deployment_key: {"litellm_provider": "bedrock"}}, persist_across_reloads=False) + + try: + assert litellm.model_cost[catalog_key]["input_cost_per_token"] == reloaded_input_cost + assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_its_price_when_a_reload_adopts_a_map_as_it_writes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + catalog_key: Final = "au.anthropic.lit-5853-reloaded-as-written-model" + deployment_key: Final = f"bedrock/{catalog_key}" + catalog_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3.3e-06, + "output_cost_per_token": 1.65e-05, + } + reloaded_input_cost: Final = 6.6e-06 + registered_output_cost: Final = 2e-05 + builtin_cost_map: Final = litellm.get_model_cost_map(url="") + reloaded_cost_map: Final = { + **builtin_cost_map, + catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost}, + } + + class CostMapAReloadReplacesAsItIsWritten(dict[str, dict[str, object]]): + def update(self, other: Mapping[str, dict[str, object]], /) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # registration writes pass one mapping + if litellm.model_cost is self: + adopt_model_cost_map(reloaded_cost_map) + super().update(other) + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + CostMapAReloadReplacesAsItIsWritten({**builtin_cost_map, catalog_key: dict(catalog_row)}), + ) + _invalidate_model_cost_lowercase_map() + + litellm.register_model( + {deployment_key: {"litellm_provider": "bedrock", "output_cost_per_token": registered_output_cost}}, + persist_across_reloads=False, + ) + + try: + assert litellm.model_cost is reloaded_cost_map + assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost + assert litellm.get_model_info(deployment_key)["output_cost_per_token"] == registered_output_cost + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_leaves_the_entry_a_lookup_is_walking_intact(monkeypatch: pytest.MonkeyPatch) -> None: + catalog_key: Final = "lit-5853-live-entry-model" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + catalog_key: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 2e-06}, + }, + ) + _invalidate_model_cost_lowercase_map() + walk: Final = iter(litellm.model_cost[catalog_key].items()) + first_field: Final = next(walk) + + try: + litellm.register_model( + {catalog_key: {"litellm_provider": "openai", "input_cost_per_token": 1e-06, "supports_vision": True}}, + persist_across_reloads=False, + ) + assert dict((first_field, *walk)) == { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 2e-06, + } + assert litellm.get_model_info(catalog_key)["input_cost_per_token"] == 1e-06 + assert litellm.get_model_info(catalog_key)["supports_vision"] is True + finally: + _invalidate_model_cost_lowercase_map() + + +_CANONICAL_KEY: Final = "lit-5853-canonical-model" +_ALIAS_KEY: Final = "lit-5853-alias-model" + + +def _install_aliased_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + _expand_model_aliases( + { + **litellm.get_model_cost_map(url=""), + _CANONICAL_KEY: { + "aliases": [_ALIAS_KEY], + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 2e-06, + "output_cost_per_token": 8e-06, + }, + } + ), + ) + _invalidate_model_cost_lowercase_map() + + +def _token_prices(name: str) -> tuple[float | None, float | None]: + info: Final = litellm.get_model_info(name) + return info["input_cost_per_token"], info["output_cost_per_token"] + + +@pytest.mark.parametrize("registered_key", [_CANONICAL_KEY, _ALIAS_KEY]) +def test_register_model_prices_every_alias_of_the_registered_entry( + monkeypatch: pytest.MonkeyPatch, registered_key: str +) -> None: + _install_aliased_cost_map(monkeypatch) + + try: + litellm.register_model( + {registered_key: {"litellm_provider": "openai", "input_cost_per_token": 1e-06}}, + persist_across_reloads=False, + ) + litellm.register_model( + {registered_key: {"litellm_provider": "openai", "output_cost_per_token": 4e-06}}, + persist_across_reloads=False, + ) + assert {name: _token_prices(name) for name in (_CANONICAL_KEY, _ALIAS_KEY)} == { + _CANONICAL_KEY: (1e-06, 4e-06), + _ALIAS_KEY: (1e-06, 4e-06), + } + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_leaves_an_alias_repointed_at_its_own_entry(monkeypatch: pytest.MonkeyPatch) -> None: + _install_aliased_cost_map(monkeypatch) + litellm.model_cost[_ALIAS_KEY] = { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 5e-06, + "output_cost_per_token": 6e-06, + } + + try: + litellm.register_model( + {_CANONICAL_KEY: {"litellm_provider": "openai", "input_cost_per_token": 1e-06}}, + persist_across_reloads=False, + ) + assert {name: _token_prices(name) for name in (_CANONICAL_KEY, _ALIAS_KEY)} == { + _CANONICAL_KEY: (1e-06, 8e-06), + _ALIAS_KEY: (5e-06, 6e-06), + } + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_inherits_builtin_token_pricing_for_unmapped_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + builtin_key: Final = "us.anthropic.claude-sonnet-4-6" + registered_key: Final = f"bedrock/bedrock/bedrock/{builtin_key}" + builtin: Final = litellm.model_cost[builtin_key] + + try: + litellm.register_model( + {registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False + ) + info: Final = litellm.get_model_info(registered_key) + assert info["input_cost_per_token"] == builtin["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] == builtin["output_cost_per_token"] > 0 + assert info["cache_read_input_token_cost"] == builtin["cache_read_input_token_cost"] + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_the_capabilities_an_unregistered_region_alias_resolves_to( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.fallback_generalizations import ( + get_fallback_generalization_rules, + set_fallback_generalizations, + ) + + sibling_key: Final = "anthropic.lit-5853-flagless-model" + region_alias: Final = f"bedrock/ap-northeast-1.{sibling_key}" + sibling_row: Final = { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + } + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), sibling_key: sibling_row}) + _invalidate_model_cost_lowercase_map() + previous_rules: Final = list(get_fallback_generalization_rules()) + set_fallback_generalizations( + [ + { + "name": "lit-5853-flagless", + "pattern": "lit-5853-flagless", + "model_info": {"supports_system_messages": True, "supports_parallel_function_calling": True}, + } + ] + ) + + try: + unregistered: Final = litellm.get_model_info(region_alias) + litellm.register_model( + {region_alias: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False + ) + registered: Final = litellm.get_model_info(region_alias) + assert registered["supports_system_messages"] is unregistered["supports_system_messages"] is True + assert ( + registered["supports_parallel_function_calling"] + is unregistered["supports_parallel_function_calling"] + is True + ) + assert registered["input_cost_per_token"] == sibling_row["input_cost_per_token"] + finally: + set_fallback_generalizations(previous_rules) + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize( + ("builtin_prices", "custom_prices", "usage", "expected_cost"), + [ + pytest.param( + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + }, + {"input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}, + {"prompt_tokens": 250_000, "completion_tokens": 1_000}, + 250_000 * 1e-06 + 1_000 * 2e-06, + id="above-200k-tier", + ), + pytest.param( + { + "tiered_pricing": [ + {"input_cost_per_token": 5e-08, "output_cost_per_token": 4e-07, "range": [0, 256_000]}, + {"input_cost_per_token": 2.5e-07, "output_cost_per_token": 2e-06, "range": [256_000, 1_000_000]}, + ] + }, + {"input_cost_per_token": 1e-06}, + {"prompt_tokens": 300_000, "completion_tokens": 1_000}, + 300_000 * 1e-06 + 1_000 * 2e-06, + id="tiered-pricing-input-only", + ), + pytest.param( + {"input_cost_per_token": 3e-06, "output_cost_per_token": 1.5e-05, "output_cost_per_reasoning_token": 3e-05}, + {"output_cost_per_token": 2e-06}, + {"prompt_tokens": 100, "completion_tokens": 1_000, "reasoning_tokens": 600}, + 100 * 3e-06 + 1_000 * 2e-06, + id="reasoning-rate-output-only", + ), + pytest.param( + { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "off_peak_pricing": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 5e-06, + "hours_utc": "00:00-00:00", + }, + }, + {"output_cost_per_token": 2e-06}, + {"prompt_tokens": 100, "completion_tokens": 1_000}, + 100 * 1e-06 + 1_000 * 2e-06, + id="off-peak-output-only", + ), + ], +) +def test_register_model_keeps_custom_flat_prices_over_inherited_rate_variants( + monkeypatch: pytest.MonkeyPatch, + builtin_prices: dict[str, object], + custom_prices: dict[str, float], + usage: dict[str, int], + expected_cost: float, +) -> None: + from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, Usage + + builtin_key: Final = "anthropic.lit-5853-tiered-model-v1:0" + alias: Final = f"bedrock/ap-northeast-1.{builtin_key}" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + builtin_key: {"litellm_provider": "bedrock", "mode": "chat", **builtin_prices}, + }, + ) + _invalidate_model_cost_lowercase_map() + + try: + litellm.register_model( + {alias: {"litellm_provider": "bedrock", **custom_prices}}, + persist_across_reloads=False, + ) + response_usage: Final = Usage( + prompt_tokens=usage["prompt_tokens"], + completion_tokens=usage["completion_tokens"], + total_tokens=usage["prompt_tokens"] + usage["completion_tokens"], + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=usage.get("reasoning_tokens", 0)), + ) + cost: Final = litellm.completion_cost( + completion_response=ModelResponse(model=alias, usage=response_usage), + model=alias, + custom_llm_provider="bedrock", + ) + assert cost == pytest.approx(expected_cost) + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_never_inherits_pricing_across_providers(monkeypatch: pytest.MonkeyPatch) -> None: + openai_key: Final = "lit-5853-openai-only-model" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.get_model_cost_map(url=""), + openai_key: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1e-05, + "cache_read_input_token_cost": 1.25e-06, + }, + }, + ) + _invalidate_model_cost_lowercase_map() + registered_key: Final = f"bedrock/{openai_key}" + + try: + litellm.register_model( + {registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False + ) + registered: Final = litellm.model_cost[registered_key] + assert "cache_read_input_token_cost" not in registered + assert "input_cost_per_token" not in registered + assert "output_cost_per_token" not in registered + finally: + _invalidate_model_cost_lowercase_map() + + def test_update_dictionary_merges_nested_dicts_without_aliasing(): """A nested dict must be merged copy-on-write: the pre-existing nested dict object stays untouched, and the caller's incoming nested dict is never diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d8625ff9da3..b353dc86d5f 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -9,9 +9,14 @@ should still use the built-in pricing. import asyncio import copy +import itertools import logging import os +import queue import re +import sys +import threading +from types import MappingProxyType from typing import Final from unittest.mock import Mock, patch @@ -3241,6 +3246,76 @@ def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info( assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price +def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + catalog: Final = litellm.get_model_cost_map(url="") + monkeypatch.setattr(litellm, "model_cost", dict(catalog)) + _invalidate_model_cost_lowercase_map() + backends: Final = tuple( + f"bedrock/au.anthropic.{model}" + for model in ("claude-sonnet-5", "claude-opus-5", "claude-haiku-4-5-20251001-v1:0", "claude-sonnet-4-6") + ) + canonical_prices: Final = MappingProxyType( + {backend: catalog[backend.removeprefix("bedrock/")]["input_cost_per_token"] for backend in backends} + ) + assert all(price > 0 for price in canonical_prices.values()) + router: Final = Router( + model_list=[ + { + "model_name": backend.rsplit(".", 1)[-1], + "litellm_params": {"model": backend, "aws_region_name": "ap-southeast-2"}, + "model_info": {"id": f"lit-5853-dep-{index}"}, + } + for index, backend in enumerate(backends) + ] + ) + stop: Final = threading.Event() + writer_outcome: Final[queue.SimpleQueue[int | Exception]] = queue.SimpleQueue() + + def register_until_stopped() -> None: + try: + for i in itertools.count(): + if stop.is_set(): + writer_outcome.put(i) + return + litellm.register_model( + { + f"lit-5853-writer-{i}": { + "litellm_provider": "lit-5853-writer", + "mode": "chat", + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 1e-08, + } + }, + persist_across_reloads=False, + ) + except Exception as failure: + writer_outcome.put(failure) + + writer: Final = threading.Thread(target=register_until_stopped, daemon=True) + saved_switch_interval: Final = sys.getswitchinterval() + sys.setswitchinterval(1e-4) + writer.start() + try: + for _ in range(10): + _simulate_price_data_reload_with_provider_sets(monkeypatch, dict(catalog)) + finally: + stop.set() + writer.join() + sys.setswitchinterval(saved_switch_interval) + + registrations_landed: Final = writer_outcome.get(timeout=5) + try: + assert isinstance(registrations_landed, int) and registrations_landed > 0, registrations_landed + assert tuple(backend for backend in backends if backend in litellm.model_cost) == () + for backend in backends: + assert litellm.get_model_info(backend)["input_cost_per_token"] == canonical_prices[backend] + assert tuple(router.get_model_names()) == tuple(backend.rsplit(".", 1)[-1] for backend in backends) + finally: + _invalidate_model_cost_lowercase_map() + + _COLLIDING_BUILTIN_KEY: Final = "baseten/zai-org/GLM-5.2" _COLLIDING_SHARED_OPENAI_KEY: Final = "openai/zai-org/GLM-5.2" _COLLIDING_MODEL_INFO: Final = { diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 87f9d58e247..a8e90fded83 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone from pathlib import Path, PurePath -from typing import Final, cast +from typing import Final, SupportsIndex, cast from unittest import mock from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -4887,6 +4887,104 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey _invalidate_model_cost_lowercase_map() +def test_model_cost_walks_survive_a_registration_landing_mid_walk(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.utils import _invalidate_model_cost_lowercase_map, _rebuild_model_cost_lowercase_map + + class _KeyRegisteredMidWalk(str): + def lower(self) -> str: + litellm.model_cost.setdefault("lit-5853-registered-during-the-rebuild", {"mode": "chat"}) + return super().lower() + + class _EntryRegisteredMidWalk(dict): + def get(self, field: str, default: object = None) -> object: + litellm.model_cost.setdefault("lit-5853-registered-during-add-known-models", {"mode": "chat"}) + return super().get(field, default) + + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + litellm.model_cost[_KeyRegisteredMidWalk("LIT-5853-WALKED-LAST")] = _EntryRegisteredMidWalk(mode="chat") + try: + assert _rebuild_model_cost_lowercase_map()["lit-5853-walked-last"] == "LIT-5853-WALKED-LAST" + litellm.add_known_models(model_cost_map=litellm.model_cost) + assert "lit-5853-registered-during-the-rebuild" in litellm.model_cost + assert "lit-5853-registered-during-add-known-models" in litellm.model_cost + finally: + _invalidate_model_cost_lowercase_map() + + +class _KeyWhoseLoweringLosesTheRace(str): + """Lowering the key is the one Python-level step between the lookup map's rebuild and + its read, so another thread's registration invalidating the map lands right there.""" + + def lower(self) -> str: + from litellm.utils import _invalidate_model_cost_lowercase_map + + _invalidate_model_cost_lowercase_map() + return super().lower() + + +def test_case_insensitive_model_cost_lookup_survives_an_invalidation_mid_lookup() -> None: + from litellm.utils import _get_model_cost_key, _invalidate_model_cost_lowercase_map + + _invalidate_model_cost_lowercase_map() + try: + assert _get_model_cost_key(_KeyWhoseLoweringLosesTheRace("GPT-4O")) == "gpt-4o" + finally: + _invalidate_model_cost_lowercase_map() + + +def _register_a_model_mid_walk() -> None: + litellm.model_cost.setdefault("lit-5853-registered-mid-walk", {"mode": "chat"}) + + +class _KeyReadDuringARegistration(str): + def startswith( + self, + prefix: str | tuple[str, ...], + start: SupportsIndex | None = None, + end: SupportsIndex | None = None, + ) -> bool: + _register_a_model_mid_walk() + return super().startswith(prefix, start, end) + + def lower(self) -> str: + _register_a_model_mid_walk() + return super().lower() + + +class _EntryReadDuringARegistration(dict[str, object]): + def __contains__(self, key: object) -> bool: + _register_a_model_mid_walk() + return super().__contains__(key) + + +@pytest.mark.parametrize("walk", ["fal_ai", "fireworks_ai", "ollama", "gpt-5", "case-insensitive scan"]) +def test_request_path_cost_map_walks_survive_a_registration_mid_walk(monkeypatch: pytest.MonkeyPatch, walk: str) -> None: + from litellm.llms.fal_ai.cost_calculator import _keyed_rows + from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig + from litellm.llms.ollama.common_utils import OllamaModelInfo + from litellm.llms.openai.chat.gpt_5_transformation import _catalogue_declares_default_effort + from litellm.utils import _handle_new_key_with_scan, _invalidate_model_cost_lowercase_map + + walks: Final[Mapping[str, Callable[[], object]]] = { + "fal_ai": lambda: _keyed_rows("lit-5853-model", "standard"), + "fireworks_ai": FireworksAIConfig._get_fireworks_index, + "ollama": lambda: OllamaModelInfo._is_static_ollama_model("ollama/lit-5853-model"), + "gpt-5": _catalogue_declares_default_effort, + "case-insensitive scan": lambda: _handle_new_key_with_scan("lit-5853-model"), + } + monkeypatch.setattr( + litellm, + "model_cost", + {_KeyReadDuringARegistration("lit-5853-catalog-model"): _EntryReadDuringARegistration(mode="chat")}, + ) + + try: + walks[walk]() + assert "lit-5853-registered-mid-walk" in litellm.model_cost + finally: + _invalidate_model_cost_lowercase_map() + + class _JsonCapture(logging.Handler): def __init__(self): super().__init__() @@ -7091,6 +7189,25 @@ def test_get_model_info_gemini(monkeypatch): assert info.get("rpm") is not None, f"{model} does not have rpm" +def test_register_model_without_a_provider_never_inherits_another_providers_entry(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.utils import _invalidate_model_cost_lowercase_map + + key: Final = "azure/claude-sonnet-4-6" + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + litellm.model_cost.pop(key, None) + try: + litellm.register_model( + {key: {"input_cost_per_token": 1e-07, "output_cost_per_token": 4e-07}}, persist_across_reloads=False + ) + assert litellm.model_cost[key].get("litellm_provider") is None + info: Final = litellm.get_model_info(key) + assert info["litellm_provider"] == "azure" + assert info["input_cost_per_token"] == 1e-07 + assert info["output_cost_per_token"] == 4e-07 + finally: + _invalidate_model_cost_lowercase_map() + + @pytest.mark.parametrize( "model, litellm_params, expects_native_config", [ From 7049a5ff4d0619efa0f45b72b36467195a14cfd5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 17 Oct 2026 12:14:28 -0700 Subject: [PATCH 27/64] fix(mcp)!: register the gateway's own DCR client for bridge servers instead of relaying the caller's loopback redirect (#46940) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 75 +++- .../mcp_management_endpoints.py | 5 +- .../mcp_server/test_discoverable_endpoints.py | 327 +++++++++++++++++- .../test_mcp_management_endpoints.py | 91 +++++ 4 files changed, 481 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 449adaa73e9..3b2c2e09df5 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -862,12 +862,17 @@ def _raise_unless_oauth2_discovery_server( def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool: """True when a DCR-bridge server relays client registration to the upstream authorization - server instead of short-circuiting to an admin-configured OAuth client. In the relay arm the + server instead of short-circuiting to the server's own OAuth client. In the relay arm the upstream holds each client's own registration, so the authorize and token relays pass the client's ``client_id`` and ``redirect_uri`` through verbatim and the authorization code returns directly to the client's redirect URI without transiting the gateway. Gateway-side redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit - arm, where the upstream only knows the gateway's own callback.""" + arm, where the upstream only knows the gateway's own callback. + + Authorize and token decide per request from the server's current client: once the gateway + holds its shared client (registered at ``/register`` under ``PROXY_BASE_URL``, or entered by an + admin) every caller runs the gateway-callback arm, while a server that still has none keeps + relaying, so clients registered before the shared client existed keep their own binding.""" return bool( mcp_server.is_dcr_bridge and mcp_server.effective_registration_url @@ -880,6 +885,27 @@ def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool: ) +def _register_relays_bridge_client(mcp_server: MCPServer) -> bool: + """``/register`` relays the caller's own registration only when the gateway cannot keep a shared + client: with ``PROXY_BASE_URL`` set it registers one client bound to its own ``/callback`` and + keeps it as the server's client, so an upstream that restricts redirect domains (Atlassian) + never sees a laptop loopback redirect. Without a database the shared client would live only in + this process, so a database-free gateway keeps relaying as well.""" + return _dcr_bridge_relays_client_registration(mcp_server) and ( + configured_proxy_base_url() is None or not _shared_client_store_available() + ) + + +def _shared_client_store_available() -> bool: + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # avoids circular import + + try: + get_prisma_client_or_throw("Database not connected. Cannot keep the shared MCP OAuth client.") + except HTTPException: + return False + return True + + def _require_s256_pkce( code_challenge: str | None, code_challenge_method: str | None, @@ -1028,6 +1054,17 @@ async def _managed_authorization_scope( return selected_scope, scope_user_id +async def _hydrate_bridge_shared_client(mcp_server: MCPServer, request: Request) -> None: + """Whichever worker served ``/register`` persisted the bridge's shared gateway client; a worker + whose catalog snapshot predates that write still holds no ``client_id``, so the gateway-callback + arms load the persisted client before presenting one upstream. Callers that carry their own + client (a gateway-minted ephemeral client sealed into the state or the code) skip this, since the + upstream bound their code to that client.""" + if not mcp_server.is_dcr_bridge or mcp_server.client_id: + return + await _reuse_persisted_dcr_client_if_available(mcp_server, f"{get_request_base_url(request)}/callback") + + async def authorize_with_server( request: Request, mcp_server: MCPServer, @@ -1069,6 +1106,8 @@ async def authorize_with_server( # calling this for its enforcement side effect, then falls through to the gateway # /callback flow below, which reads the original code_challenge names. bridge_challenge, bridge_method = _require_s256_pkce(code_challenge, code_challenge_method) + if ephemeral_dcr_client is None: + await _hydrate_bridge_shared_client(resolved_server, request) # A gateway-minted ephemeral client is registered against {base}/callback, so its # flow must run the short-circuit arm; the relay arm is only for clients that # registered themselves through the front door and hold their own redirect binding. @@ -1233,6 +1272,7 @@ async def exchange_token_with_server( refresh_token: str | None = None, scope: str | None = None, client_token_endpoint_auth_method: MCPTokenEndpointAuthMethod | None = None, + sealed_client: bool = False, ): raise_if_not_oauth2(mcp_server) if grant_type not in ("authorization_code", "refresh_token"): @@ -1282,11 +1322,14 @@ async def exchange_token_with_server( # recovered from a sealed code) must authenticate the way its own registration was granted, # not the way the server row is configured; callers that carry no method keep the row's method # as before. - resolved_client_id: Final = resolved_server.client_id if resolved_server.client_id else client_id - resolved_client_secret: Final = resolved_server.client_secret if resolved_server.client_id else client_secret + if not sealed_client: + await _hydrate_bridge_shared_client(resolved_server, request) + server_client_wins: Final = bool(resolved_server.client_id) and not sealed_client + resolved_client_id: Final = resolved_server.client_id if server_client_wins else client_id + resolved_client_secret: Final = resolved_server.client_secret if server_client_wins else client_secret resolved_auth_method: Final = ( resolved_server.token_endpoint_auth_method - if resolved_server.client_id + if server_client_wins else (client_token_endpoint_auth_method or resolved_server.token_endpoint_auth_method) ) try: @@ -1830,10 +1873,13 @@ async def _persist_dcr_client_registration( empty and otherwise adopts whatever a concurrent registration persisted first. The client-forwarded token modes (``true_passthrough`` / ``oauth_delegate``) are skipped - unconditionally: the caller holds the upstream token and the gateway must hold no OAuth - client identity for these servers. Persisting here would stamp ``oauth2_flow`` and a - ``client_id`` onto a server whose mode promises the gateway stores nothing, making a - fresh pass-through server read as gateway-authorized. + unless the server is a DCR bridge: the caller holds the upstream token and the gateway holds + no OAuth client identity for these servers, so persisting would stamp ``oauth2_flow`` and a + ``client_id`` onto a server whose mode promises the gateway stores nothing. A DCR bridge is + the exception by design: its gateway-hosted sign-in registers one shared client bound to the + gateway ``/callback`` (the same app an admin may enter by hand), and that client must survive + for every later caller's authorize and token exchange. ``oauth2_flow`` is still never stamped + on a client-forwarded row. ``redirect_uris`` records what the client is bound to so a later origin change can be detected as a positive mismatch and trigger re-registration instead of stranding the @@ -1842,7 +1888,7 @@ async def _persist_dcr_client_registration( ``update_mcp_server`` merges credential blobs: a re-registered public client must not inherit the previous client's secret or auth method. """ - if mcp_server.is_client_forwarded_token: + if mcp_server.is_client_forwarded_token and not mcp_server.is_dcr_bridge: return "skipped" try: @@ -1910,20 +1956,23 @@ async def _persist_dcr_client_registration( ): await _refresh_persisted_dcr_server(stored) return "reused" + stamped_flow: Final[Literal["authorization_code"] | None] = ( + None if mcp_server.is_client_forwarded_token else "authorization_code" + ) updated_row: Final = await update_mcp_server( prisma_client=prisma_client, data=( UpdateMCPServerRequest( server_id=mcp_server.server_id, credentials=credentials, - oauth2_flow="authorization_code", + oauth2_flow=stamped_flow, token_url=mcp_server.token_url, ) if mcp_server.token_url else UpdateMCPServerRequest( server_id=mcp_server.server_id, credentials=credentials, - oauth2_flow="authorization_code", + oauth2_flow=stamped_flow, ) ), touched_by="mcp_oauth_dcr", @@ -2211,7 +2260,7 @@ async def register_client_with_server( ) return dummy_return - bridge_relay: Final = _dcr_bridge_relays_client_registration(resolved_server) + bridge_relay: Final = _register_relays_bridge_client(resolved_server) if bridge_relay and not client_redirect_uris: raise HTTPException( status_code=400, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d014ff6cf07..a014b9cf8a5 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2517,7 +2517,9 @@ if MCP_AVAILABLE: resolved_redirect_uri: Final = f"{get_request_base_url(request)}/callback" if sealed_code else redirect_uri caller_client_id: Final = sealed_code.client_id if sealed_code else client_id caller_client_secret: Final = sealed_code.client_secret if sealed_code else client_secret - resolved_client_id: Final = mcp_server.client_id or caller_client_id or "" + resolved_client_id: Final = ( + caller_client_id if sealed_code else (mcp_server.client_id or caller_client_id or "") + ) if not resolved_client_id: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -2541,6 +2543,7 @@ if MCP_AVAILABLE: refresh_token=refresh_token, scope=scope, client_token_endpoint_auth_method=sealed_code.token_endpoint_auth_method if sealed_code else None, + sealed_client=sealed_code is not None, ) @router.post( diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 5fdc7669496..a43dacaf100 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -5,11 +5,12 @@ import json import time from base64 import urlsafe_b64encode from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import HTTPException, Request +from fastapi import HTTPException, Request, Response +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import LiteLLM_MCPServerTable from litellm.types.mcp import MCPAuth @@ -26,6 +27,14 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +class _BridgeRegistrationPayload(TypedDict, total=False): + client_name: ReadOnly[str] + grant_types: ReadOnly[list[str]] + response_types: ReadOnly[list[str]] + token_endpoint_auth_method: ReadOnly[str] + redirect_uris: ReadOnly[list[str]] + + def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None): credential = {"type": "oauth2", "access_token": access_token} if refresh_token is not None: @@ -5849,7 +5858,12 @@ def test_oauth_authorization_server_404_for_non_bridge_client_forwarded_server() assert exc.value.status_code == 404 -async def _bridge_register_response(server, request_payload, replace_stale_client=False): +async def _bridge_register_response( + server: "MCPServer", + request_payload: _BridgeRegistrationPayload, + replace_stale_client: bool = False, + database_connected: bool = True, +) -> tuple[Response, MagicMock, AsyncMock]: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client_with_server, ) @@ -5879,6 +5893,10 @@ async def _bridge_register_response(server, request_payload, replace_stale_clien new_callable=AsyncMock, return_value=False, ), + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + **({"return_value": MagicMock()} if database_connected else {"side_effect": HTTPException(status_code=500)}), + ), ): response = await register_client_with_server( request=_bridge_mock_request(), @@ -6033,6 +6051,309 @@ async def test_register_bridge_relay_never_persists(): mock_persist.assert_not_called() +_LAPTOP_LOOPBACK_REDIRECT: Final = "http://127.0.0.1:33418/" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type_value", ["oauth_delegate", "true_passthrough"]) +async def test_register_bridge_with_proxy_base_url_registers_gateway_client_not_laptop_loopback( + monkeypatch: pytest.MonkeyPatch, auth_type_value: Literal["oauth_delegate", "true_passthrough"] +) -> None: + """A laptop MCP client (VS Code) registers with its own loopback redirect. With PROXY_BASE_URL set + the gateway must not relay that registration upstream: it registers its one shared client bound + to the gateway /callback, keeps it, and hands the caller a client bound to the caller's loopback. + Relaying the loopback is the customer failure on upstreams that restrict redirect domains.""" + import json + + from litellm.types.mcp import MCPAuth + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + response, mock_async_client, mock_persist = await _bridge_register_response( + _bridge_server(auth_type=MCPAuth(auth_type_value)), + {"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]}, + ) + + posted: Final = mock_async_client.post.call_args.kwargs["json"] + assert posted["redirect_uris"] == ["https://litellm.example.com/callback"], posted + assert posted["client_name"] == "LiteLLM Proxy", posted + mock_persist.assert_awaited_once() + + payload: Final = json.loads(response.body.decode("utf-8")) + assert payload["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT], payload + + +@pytest.mark.asyncio +async def test_register_bridge_without_database_still_relays_client_registration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A database-free gateway cannot keep the shared client past this process, so even with + PROXY_BASE_URL set it relays the caller's own registration instead of stranding the next authorize.""" + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + _, mock_async_client, mock_persist = await _bridge_register_response( + _bridge_server(), + {"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]}, + database_connected=False, + ) + + assert mock_async_client.post.call_args.kwargs["json"]["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT] + mock_persist.assert_not_called() + + +@pytest.mark.asyncio +async def test_register_bridge_without_proxy_base_url_still_relays_client_registration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Without PROXY_BASE_URL the gateway cannot keep a shared client, so the relay arm remains the + fallback and the caller's own redirect goes upstream as before.""" + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + _, mock_async_client, mock_persist = await _bridge_register_response( + _bridge_server(), + {"client_name": "Visual Studio Code", "redirect_uris": [_LAPTOP_LOOPBACK_REDIRECT]}, + ) + + assert mock_async_client.post.call_args.kwargs["json"]["redirect_uris"] == [_LAPTOP_LOOPBACK_REDIRECT] + mock_persist.assert_not_called() + + +def _persisted_shared_gateway_client( + server: "MCPServer", client_id: str = "shared-gateway-client" +) -> LiteLLM_MCPServerTable: + request: Final = _bridge_mock_request() + credentials: Final = { + "client_id": client_id, + "client_secret": None, + "token_endpoint_auth_method": None, + "redirect_uris": [f"{request.base_url.rstrip('/')}/callback"], + "dcr_issuer": server.issuer, + "dcr_server_url": server.url, + } + return LiteLLM_MCPServerTable.model_validate({**server.model_dump(exclude_none=True), "credentials": credentials}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("shared_client_persisted", "expected_redirect_uri", "expected_client_id"), + [ + (True, "https://litellm.example.com/callback", "shared-gateway-client"), + (False, _LAPTOP_LOOPBACK_REDIRECT, "dcr-client-123"), + ], +) +async def test_authorize_clientless_bridge_loads_shared_client_before_choosing_the_arm( + monkeypatch: pytest.MonkeyPatch, + shared_client_persisted: bool, + expected_redirect_uri: str, + expected_client_id: str, +) -> None: + """This worker's catalog snapshot has no client_id. When another worker already persisted the shared + client, authorize must load it and send it with the gateway /callback upstream, never the caller's + own client_id or loopback. When no shared client exists yet (a client registered before the gateway + kept one), the caller keeps its own upstream binding through the relay.""" + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding") + server: Final = _bridge_server(url="https://mcp.example.com/mcp") + persisted: Final = _persisted_shared_gateway_client(server) if shared_client_persisted else None + prisma_client: Final = MagicMock() + mock_get_mcp_server: Final = AsyncMock(return_value=persisted) + mock_update_server: Final = AsyncMock() + with ( + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object( + global_mcp_server_manager, + "update_server", + new=mock_update_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper", + return_value="mocked_encrypted_state", + ), + ): + response: Final = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="dcr-client-123", + redirect_uri=_LAPTOP_LOOPBACK_REDIRECT, + state="s", + code_challenge="chal", + code_challenge_method="S256", + ) + + location: Final = response.headers["location"] + assert location.startswith("https://provider.com/oauth/authorize"), location + query: Final = parse_qs(urlparse(location).query) + assert query["redirect_uri"] == [expected_redirect_uri], query + assert query["client_id"] == [expected_client_id], query + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("shared_client_persisted", "expected_redirect_uri", "expected_client_id"), + [ + (True, "https://litellm.example.com/callback", "shared-gateway-client"), + (False, _LAPTOP_LOOPBACK_REDIRECT, "dcr-client-123"), + ], +) +async def test_token_clientless_bridge_loads_shared_client_before_choosing_the_arm( + monkeypatch: pytest.MonkeyPatch, + shared_client_persisted: bool, + expected_redirect_uri: str, + expected_client_id: str, +) -> None: + """The token leg mirrors authorize: with the shared client persisted it redeems the code with that + client and the gateway /callback, and with none yet it relays the caller's own client and redirect.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding") + server: Final = _bridge_server(url="https://mcp.example.com/mcp") + persisted: Final = _persisted_shared_gateway_client(server) if shared_client_persisted else None + prisma_client: Final = MagicMock() + mock_get_mcp_server: Final = AsyncMock(return_value=persisted) + mock_update_server: Final = AsyncMock() + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma_client), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object( + global_mcp_server_manager, + "update_server", + new=mock_update_server, + ), + ): + data: Final = await _bridge_token_post_data(server, redirect_uri=_LAPTOP_LOOPBACK_REDIRECT) + + assert data["redirect_uri"] == expected_redirect_uri, data + assert data["client_id"] == expected_client_id, data + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("dcr_bridge", "expected"), + [(True, "persisted"), (False, "skipped")], +) +async def test_persist_dcr_client_for_client_forwarded_server_only_when_it_is_a_bridge( + dcr_bridge: bool, + expected: Literal["persisted", "skipped"], +) -> None: + """A DCR bridge keeps the shared gateway client on the server row (without stamping oauth2_flow on + a client-forwarded row), while a client-forwarded server without the bridge still stores nothing.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + + server: Final = _bridge_server(auth_type=MCPAuth.oauth_delegate, dcr_bridge=dcr_bridge) + prisma: Final = MagicMock() + prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) + mock_update: Final = AsyncMock(return_value=MagicMock()) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=prisma), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=mock_update), + patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False), + ): + result: Final = await _persist_dcr_client_registration( + server, + {"client_id": "shared-gateway-client", "token_endpoint_auth_method": "none"}, + "https://litellm.example.com/callback", + replace_stale_client=False, + ) + + assert result == expected + if expected == "skipped": + mock_update.assert_not_awaited() + return + written: Final = mock_update.await_args.kwargs["data"] + assert written.oauth2_flow is None, written + assert written.credentials["client_id"] == "shared-gateway-client", written + assert server.client_id == "shared-gateway-client" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("server_has_shared_client", [False, True]) +async def test_token_sealed_ephemeral_client_is_not_replaced_by_persisted_gateway_client( + monkeypatch: pytest.MonkeyPatch, + server_has_shared_client: bool, +) -> None: + """The token leg must use the sealed ephemeral client even if this worker's catalog has a shared client.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import exchange_token_with_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-registration-binding") + server: Final = _bridge_server( + url="https://mcp.example.com/mcp", + client_id="shared-gateway-client" if server_has_shared_client else None, + client_secret="shared-secret" if server_has_shared_client else None, + token_endpoint_auth_method="client_secret_basic" if server_has_shared_client else None, + ) + fake_http_response: Final = MagicMock() + fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client: Final = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + prisma_client: Final = MagicMock() + mock_get_mcp_server: Final = AsyncMock( + return_value=_persisted_shared_gateway_client(server, client_id="persisted-gateway-client") + ) + mock_update_server: Final = AsyncMock() + with ( + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object( + global_mcp_server_manager, + "update_server", + new=mock_update_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + ): + await exchange_token_with_server( + request=_bridge_mock_request(), + mcp_server=server, + grant_type="authorization_code", + code="upstream-code", + redirect_uri="https://litellm.example.com/callback", + client_id="ephemeral-client", + client_secret="ephemeral-secret" if server_has_shared_client else None, + code_verifier="verifier", + client_token_endpoint_auth_method="client_secret_post" if server_has_shared_client else None, + sealed_client=True, + ) + + data: Final = fake_http_client.post.call_args.kwargs["data"] + assert data["client_id"] == "ephemeral-client", data + assert data["redirect_uri"] == "https://litellm.example.com/callback", data + if server_has_shared_client: + headers: Final = fake_http_client.post.call_args.kwargs["headers"] + assert data["client_secret"] == "ephemeral-secret", data + assert "Authorization" not in headers, headers + + _BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef" diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 0d59c655835..65c4a8faeb7 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -3545,6 +3545,95 @@ class TestTemporaryMCPSessionEndpoints: assert exchange_mock.await_args.kwargs["redirect_uri"] == "https://litellm.example.com/callback" assert exchange_mock.await_args.kwargs["client_token_endpoint_auth_method"] == "client_secret_basic" + @pytest.mark.asyncio + async def test_mcp_token_prefers_sealed_client_when_server_has_shared_client( + self: "TestTemporaryMCPSessionEndpoints", + ) -> None: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + seal_passthrough_authorization_code, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_token, + ) + + request: Final = MagicMock() + request.base_url = "https://litellm.example.com/" + request.headers = {} + server: Final = generate_mock_mcp_server_config_record(server_id="server-1").model_copy( + update={ + "auth_type": MCPAuth.true_passthrough, + "client_id": "shared-gateway-client", + "client_secret": "shared-secret", + "token_endpoint_auth_method": "client_secret_basic", + "token_url": "https://provider.example/oauth/token", + } + ) + admin_auth: Final = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN) + fake_http_response: Final = MagicMock() + fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client: Final = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + prisma_client: Final = MagicMock() + mock_get_mcp_server: Final = AsyncMock(return_value=None) + mock_update_server: Final = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.master_key", "sk-lit4581-test-master-key"), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._oauth_server_operation", + return_value=nullcontext(server), + ), + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=prisma_client, + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=mock_get_mcp_server, + ), + patch.object( + global_mcp_server_manager, + "update_server", + new=mock_update_server, + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ), + ): + sealed: Final = seal_passthrough_authorization_code( + upstream_code="up-code", + client_id="ephemeral-client", + client_secret="ephemeral-secret", + mcp_server_id="server-1", + token_endpoint_auth_method="client_secret_post", + ) + result: Final = await mcp_token( + request=request, + server_id="server-1", + user_api_key_dict=admin_auth, + grant_type="authorization_code", + code=sealed, + redirect_uri="https://example.com/callback", + client_id=None, + client_secret=None, + code_verifier="verifier", + refresh_token=None, + scope=None, + ) + + data: Final = fake_http_client.post.call_args.kwargs["data"] + headers: Final = fake_http_client.post.call_args.kwargs["headers"] + assert data["client_id"] == "ephemeral-client", data + assert data["client_secret"] == "ephemeral-secret", data + assert data["code"] == "up-code", data + assert data["redirect_uri"] == "https://litellm.example.com/callback", data + assert "Authorization" not in headers, headers + assert result.status_code == 200 + assert result.body == b'{"access_token":"tok","token_type":"Bearer"}' + @pytest.mark.asyncio async def test_mcp_token_refresh_grant_never_opens_sealed_code(self): """The minted client is unrecoverable outside the single authorization_code flow by @@ -3841,6 +3930,7 @@ class TestTemporaryMCPSessionEndpoints: refresh_token=None, scope=None, client_token_endpoint_auth_method=None, + sealed_client=False, ) @pytest.mark.asyncio @@ -3895,6 +3985,7 @@ class TestTemporaryMCPSessionEndpoints: refresh_token="rt-123", scope=None, client_token_endpoint_auth_method=None, + sealed_client=False, ) @pytest.mark.asyncio From d81472f196323ec677a3f359a984492e37e46b4f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 19:16:17 +0000 Subject: [PATCH 28/64] fix(model_prices): registry audit 2026-10-17, Perplexity decider v1.2 rows, Azure Mistral and Vertex Gemini deprecation dates (#47114) * fix(model_prices): registry audit 2026-10-17, Azure Mistral and Vertex Gemini deprecation dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(model_prices): add Perplexity pplx-decider v1.2 and flash v1.2 rows with latest aliases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 39 +++++++++++++++++++ model_prices_and_context_window.json | 39 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8545a5b568e..6ebb54c0035 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12877,6 +12877,7 @@ "supports_prompt_caching": true }, "azure_ai/ministral-3b": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 4e-08, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12940,6 +12941,7 @@ "supports_vision": true }, "azure_ai/mistral-medium-2505": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 4e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12977,6 +12979,7 @@ "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -29699,6 +29702,7 @@ }, "vertex_ai/gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -29806,6 +29810,7 @@ "supports_native_streaming": true }, "vertex_ai/gemini-3.8-live": { + "deprecation_date": "2027-09-25", "input_cost_per_audio_token": 3e-06, "input_cost_per_image_token": 1e-06, "input_cost_per_token": 7.5e-07, @@ -31960,6 +31965,7 @@ }, "gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -45425,6 +45431,38 @@ "output_cost_per_token": 0.0, "source": "https://docs.perplexity.ai/docs/getting-started/pricing" }, + "perplexity/pplx-decider-v1.2": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-latest": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-flash-v1.2": { + "input_cost_per_token": 1e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-flash-latest": { + "input_cost_per_token": 1e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, "perplexity/sonar": { "input_cost_per_token": 1e-06, "litellm_provider": "perplexity", @@ -75357,6 +75395,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-4-argon": { + "deprecation_date": "2027-10-13", "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8545a5b568e..6ebb54c0035 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12877,6 +12877,7 @@ "supports_prompt_caching": true }, "azure_ai/ministral-3b": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 4e-08, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -12940,6 +12941,7 @@ "supports_vision": true }, "azure_ai/mistral-medium-2505": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 4e-07, "litellm_provider": "azure_ai", "max_input_tokens": 131072, @@ -12977,6 +12979,7 @@ "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { + "deprecation_date": "2026-12-09", "input_cost_per_token": 1e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, @@ -29699,6 +29702,7 @@ }, "vertex_ai/gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -29806,6 +29810,7 @@ "supports_native_streaming": true }, "vertex_ai/gemini-3.8-live": { + "deprecation_date": "2027-09-25", "input_cost_per_audio_token": 3e-06, "input_cost_per_image_token": 1e-06, "input_cost_per_token": 7.5e-07, @@ -31960,6 +31965,7 @@ }, "gemini-3.8-flash": { "prompt_cache_min_tokens": 4096, + "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -45425,6 +45431,38 @@ "output_cost_per_token": 0.0, "source": "https://docs.perplexity.ai/docs/getting-started/pricing" }, + "perplexity/pplx-decider-v1.2": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-latest": { + "input_cost_per_token": 2e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-flash-v1.2": { + "input_cost_per_token": 1e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, + "perplexity/pplx-decider-flash-latest": { + "input_cost_per_token": 1e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 262144, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "source": "https://docs.perplexity.ai/docs/decisions/quickstart" + }, "perplexity/sonar": { "input_cost_per_token": 1e-06, "litellm_provider": "perplexity", @@ -75357,6 +75395,7 @@ "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-4-argon": { + "deprecation_date": "2027-10-13", "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, From 4b635f18e4a15476ebcd5d2a2eceab0c0856a6ac Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 17 Oct 2026 12:16:46 -0700 Subject: [PATCH 29/64] fix(mcp): bound paginated tools/list by LITELLM_MCP_TOOL_LISTING_TIMEOUT (#46513) Preserve explicit per-server timeouts while bounding paginated listing Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 15 +++- .../mcp_server/test_mcp_server_manager.py | 68 +++++++++++++++++++ 2 files changed, 82 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 02799b90566..96cf6ee11aa 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4531,7 +4531,7 @@ class MCPServerManager: ) else: page: Final = ( - await client.list_tools_page(params) + await self._fetch_tools_page_with_timeout(client, server, params) if params is not None else ListToolsResult(tools=await self._fetch_tools_with_timeout(client, server.name)) ) @@ -5621,6 +5621,19 @@ class MCPServerManager: ) raise_classified_list_failure(e, server_name) + async def _fetch_tools_page_with_timeout( + self, client: MCPClient, server: MCPServer, params: PaginatedRequestParams + ) -> ListToolsResult: + deadline: Final = ( + MCP_TOOL_LISTING_TIMEOUT if server.timeout is None else max(server.timeout, MCP_TOOL_LISTING_TIMEOUT) + ) + try: + with anyio.fail_after(deadline): + return await client.list_tools_page(params) + except TimeoutError as e: + verbose_logger.warning("Timeout while listing tools from %s", server.name) + raise MCPServerListError(ServerListFault(tag="timeout"), server.name) from e + _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 def assign_unique_short_prefix( diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 691da35b8c4..d47b2f73a48 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -20056,6 +20056,74 @@ async def test_disabled_stdio_tools_are_empty_initially_and_reject_continuation( assert page.next_cursor is None +@pytest.mark.asyncio +async def test_aggregate_bounds_each_paginated_server_by_its_own_listing_deadline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from mcp.types import JSONRPCRequest, PaginatedRequestParams + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import catalog, operations + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR, get_server_prefix + + monkeypatch.setenv("LITELLM_SALT_KEY", "paginated-listing-deadline-test") + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(manager_module, "MCP_TOOL_LISTING_TIMEOUT", 0.0) + extended: Final = MCPServer( + server_id="extended", + name="extended", + transport=MCPTransport.http, + url="https://extended.example/mcp", + timeout=60.0, + ) + global_deadline: Final = MCPServer( + server_id="global", name="global", transport=MCPTransport.http, url="https://global.example/mcp" + ) + servers: Final = (extended, global_deadline) + manager: Final = MCPServerManager() + manager.registry = {server.server_id: server for server in servers} + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + context: Final = operations.prepare_context(UserAPIKeyAuth(api_key="owned-caller", user_id="alice")) + + async def respond(request: httpx2.Request) -> httpx2.Response: + if request.method == "DELETE": + return httpx2.Response(200) + payload: Final = _JSONRPC_ADAPTER.validate_json(request.content) + if not isinstance(payload, JSONRPCRequest): + return httpx2.Response(202) + result: Final = { + "initialize": { + "protocolVersion": "2025-03-26", + "serverInfo": {"name": "test", "version": "1"}, + "capabilities": {"tools": {}}, + }, + "server/discover": { + "supportedVersions": ["2026-07-28"], + "capabilities": {"tools": {}}, + "ttlMs": 60000, + "cacheScope": "private", + "resultType": "complete", + }, + "tools/list": { + "tools": [{"name": "echo", "inputSchema": {"type": "object"}}], + "ttlMs": 60000, + "cacheScope": "private", + "resultType": "complete", + }, + }[payload.method] + return httpx2.Response(200, json={"jsonrpc": "2.0", "id": payload.id, "result": result}) + + with _mcp_upstream(respond): + listing: Final = await catalog.aggregate_gateway_tools( + context, PaginatedRequestParams(), servers, {}, record_listing=True, enforce_rate_limits=False + ) + + assert [tool.name for tool in listing.tools] == [f"{get_server_prefix(extended)}{MCP_TOOL_PREFIX_SEPARATOR}echo"] + assert listing.outcomes["extended"].tag == "ok" + assert listing.outcomes["global"].tag == "timeout" + + @pytest.mark.asyncio async def test_failed_aggregate_continuation_preserves_only_delivered_tool_metadata(monkeypatch): from mcp.shared.exceptions import MCPError From 84d91bbe8742aff639c3cd009acdce4d3c8efc47 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 19:17:18 +0000 Subject: [PATCH 30/64] fix(vertex): match standalone 403 status tokens instead of any substring (#47134) * fix(vertex): match standalone 403 status tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex): keep text 403 fallback for wrapped 500 errors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex): keep BadRequestError for explicit 403 text matches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex): let structured body codes override the 403 text fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../exception_mapping_utils.py | 13 +- .../test_google_genai_vertex_routes_wire.py | 2 +- .../test_exception_mapping_utils.py | 126 ++++++++++++++++++ 3 files changed, 138 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 13a53c55337..c52bc534058 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1196,6 +1196,9 @@ def _map_sagemaker_exception( ) +_VERTEX_403_STATUS_CODE_PATTERN: Final = re.compile(r"(? None: + status_code: Final[object] = getattr(original_exception, "status_code", None) + body_code: Final[int | None] = _get_body_error_code(error_str) if "Vertex AI API has not been used in project" in error_str or "Unable to find your project" in error_str: raise BadRequestError( message=f"litellm.BadRequestError: {custom_llm_provider}Exception - {error_str}", @@ -1252,7 +1257,11 @@ def _map_vertex_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, ) - elif "403" in error_str: + elif ( + (not isinstance(status_code, int) or status_code in (403, 500)) + and body_code in (None, 403) + and _VERTEX_403_STATUS_CODE_PATTERN.search(error_str) + ): raise BadRequestError( message=f"{custom_llm_provider.capitalize()}Exception BadRequestError - {error_str}", model=model, @@ -1306,7 +1315,7 @@ def _map_vertex_exception( elif ( isinstance(getattr(original_exception, "status_code", None), int) and 500 <= original_exception.status_code < 600 - and _get_body_error_code(error_str) == 429 + and body_code == 429 ): # upstream gateway wraps a 429 inside a 5xx envelope # e.g. HTTP 500/503 with {"error":{"code":429,...}}. diff --git a/tests/integration/providers/test_google_genai_vertex_routes_wire.py b/tests/integration/providers/test_google_genai_vertex_routes_wire.py index d2c2f3bafb2..d7c2aa72c74 100644 --- a/tests/integration/providers/test_google_genai_vertex_routes_wire.py +++ b/tests/integration/providers/test_google_genai_vertex_routes_wire.py @@ -388,7 +388,7 @@ def test_adapter_route_upstream_error_reaches_the_caller_with_its_status( gateway: Gateway, status: int, stream: bool ) -> None: route: Final = _ROUTES["model_garden_publisher"] - message: Final = f"scripted vertex rejection {uuid.uuid4().hex}" + message: Final = f"scripted vertex rejection {uuid.uuid4().hex}f4039" with wire_server(_vertex_peer(status=status, error=message)) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire, route, num_retries=0) key: Final = scenario.key(models=[model]) diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 9eb99516e02..731e842090a 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import ( exception_type, extract_and_raise_litellm_exception, ) +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 @@ -578,6 +579,131 @@ def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_lim ) +def test_vertex_status_code_429_ignores_403_token_in_message() -> None: + original_exception: Final = BaseLLMException(status_code=429, message="scripted vertex rejection f4039") + + with pytest.raises(litellm.RateLimitError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 429 + + +def test_vertex_status_code_403_maps_to_403() -> None: + original_exception: Final = BaseLLMException(status_code=403, message="permission denied") + + with pytest.raises(litellm.PermissionDeniedError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + +def test_vertex_status_code_403_with_text_token_maps_to_bad_request() -> None: + original_exception: Final = BaseLLMException( + status_code=403, + message='{"error":{"code":403,"message":"Permission denied"}}', + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + +def test_vertex_text_403_without_status_code_still_maps_to_bad_request() -> None: + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=Exception("403 Forbidden: permission denied"), + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + +def test_vertex_status_substring_without_status_code_does_not_map_to_403() -> None: + with pytest.raises(litellm.APIConnectionError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=Exception("scripted vertex rejection f4039a"), + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code != 403 + + +def test_vertex_wrapped_500_with_403_body_maps_to_bad_request() -> None: + original_exception: Final = BaseLLMException( + status_code=500, + message='Error code: 403 - {"error": {"code": 403, "message": "Permission denied"}}', + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + +def test_vertex_wrapped_500_with_embedded_403_text_does_not_map_to_403() -> None: + original_exception: Final = BaseLLMException(status_code=500, message="scripted vertex rejection f4039") + + with pytest.raises(litellm.InternalServerError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 500 + + +def test_vertex_wrapped_500_with_body_code_429_and_403_token_maps_to_rate_limit() -> None: + original_exception: Final = BaseLLMException( + status_code=500, + message='{"error":{"code":429,"message":"upstream unavailable (request id: 403-xyz)"}}', + ) + + with pytest.raises(litellm.RateLimitError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 429 + + +def test_vertex_wrapped_500_with_body_code_403_maps_to_bad_request() -> None: + original_exception: Final = BaseLLMException( + status_code=500, + message='{"error":{"code":403,"message":"Permission denied"}}', + ) + + with pytest.raises(litellm.BadRequestError) as excinfo: + exception_type( + model="gemini/gemini-2.5-flash", + original_exception=original_exception, + custom_llm_provider="vertex_ai", + ) + + assert excinfo.value.status_code == 403 + + class TestGetBodyErrorCode: """Unit tests for _get_body_error_code helper.""" From f6ae25939e86b0036d9d1c2023ac9b8e74091eb2 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 17 Oct 2026 12:22:41 -0700 Subject: [PATCH 31/64] feat(guardrails): add num_retries to generic_guardrail_api for checks lost after connect (#46965) * feat(guardrails): retry generic_guardrail_api checks lost after connect with num_retries * fix(guardrails): close each guardrail retry connection before the next attempt and cap num_retries at 5 A failed retry's fresh client stayed open until every later retry finished, so a check that kept timing out held one idle client per retry. Each retry now closes its own client as soon as it ends, and num_retries is capped at 5 so the retry chain stays shallow * refactor(guardrails): type the generic_guardrail_api inbound header sanitizer as a mapping Keeps the type check gate at its merge-base count after the num_retries constructor params added unknown kwargs errors in the Qohash subclass * chore(guardrails): regenerate OpenAPI snapshot and UI types for num_retries, allow the retry client The guardrail retry opens its own client on purpose to get a new connection and closes it after the attempt, so the per-request async client check allowlists generic_guardrail_api.py --- litellm/proxy/_lazy_openapi_snapshot.json | 28 +++ .../generic_guardrail_api/__init__.py | 1 + .../generic_guardrail_api.py | 82 ++++++-- litellm/types/guardrails.py | 22 ++- .../guardrail_hooks/generic_guardrail_api.py | 3 + .../ensure_async_clients_test.py | 3 + .../test_generic_guardrail_num_retries.py | 179 ++++++++++++++++++ .../providers/_idle_upstream_support.py | 11 +- .../test_idle_upstream_guardrail_cut_wire.py | 134 +++++++++++++ .../test_generic_guardrail_api.py | 153 ++++++++++++++- tests/unit/types/test_guardrails.py | 14 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 + 12 files changed, 615 insertions(+), 25 deletions(-) create mode 100644 tests/integration/observability/test_generic_guardrail_num_retries.py create mode 100644 tests/integration/providers/test_idle_upstream_guardrail_cut_wire.py diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 961d32ab664..7198d74909c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -14412,6 +14412,20 @@ "description": "Optional field if guardrail requires a 'model' parameter", "title": "Model" }, + "num_retries": { + "anyOf": [ + { + "maximum": 5.0, + "minimum": 0.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "description": "How many more times to send a guardrail check whose connection timed out waiting for the answer or was dropped after the request went out. Each retry opens a fresh connection and gets the full timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', other guardrails ignore it.", + "title": "Num Retries" + }, "on_sensitive_data": { "anyOf": [ { @@ -17001,6 +17015,20 @@ "description": "If True, logs violations without blocking. Defaults to False if not provided", "title": "Monitor Mode" }, + "num_retries": { + "anyOf": [ + { + "maximum": 5.0, + "minimum": 0.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "description": "How many more times to send a guardrail check whose connection timed out waiting for the answer or was dropped after the request went out. Each retry opens a fresh connection and gets the full timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', other guardrails ignore it.", + "title": "Num Retries" + }, "on_disallowed_action": { "default": "block", "description": "Choose whether disallowed tools block the request or get rewritten out of the payload", diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index 242354ac417..44705c4d166 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -44,6 +44,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" litellm_params, optional_params, "streaming_buffer_until_moderated" ), timeout=litellm_params.timeout, + num_retries=litellm_params.num_retries, ) litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 93b08831d84..199da7df30d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -22,6 +22,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) @@ -63,6 +64,21 @@ _HEADER_VALUE_ALLOWLIST: Final = frozenset( # Placeholder for headers that exist but are not on the allowlist (we don't expose their value). _HEADER_PRESENT_PLACEHOLDER: Final = "[present]" +_AFTER_CONNECT_FAILURES: Final = ( + httpx.ReadTimeout, + httpx.WriteTimeout, + httpx.ReadError, + httpx.WriteError, + httpx.RemoteProtocolError, +) + + +def _after_connect_failure(error: Timeout | httpx.TransportError) -> httpx.TransportError | None: + """AsyncHTTPHandler re-raises connect and read timeouts as the same litellm Timeout, so the phase comes from + the httpx error it was raised while handling.""" + origin: Final = error if isinstance(error, httpx.TransportError) else error.__cause__ or error.__context__ + return origin if isinstance(origin, _AFTER_CONNECT_FAILURES) else None + def _header_value_allowed( header_name: str, @@ -81,7 +97,7 @@ def _header_value_allowed( def _sanitize_inbound_headers( - headers: object, + headers: Mapping[object, object] | None, extra_allowlist: set[str] | None = None, ) -> dict[str, str] | None: """ @@ -254,9 +270,18 @@ class GenericGuardrailAPI(CustomGuardrail): streaming_sampling_rate: int | None = None, streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None, streaming_buffer_until_moderated: bool | None = None, + num_retries: int | None = None, + async_handler: AsyncHTTPHandler | None = None, + fresh_handler_factory: Callable[[], AsyncHTTPHandler] | None = None, **kwargs, ): - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + self.fresh_handler_factory: Callable[[], AsyncHTTPHandler] = fresh_handler_factory or ( + lambda: AsyncHTTPHandler(timeout=self.async_handler.timeout) + ) + self.num_retries: int = num_retries or 0 self.headers = headers or {} self.extra_headers = extra_headers or [] @@ -384,12 +409,9 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs.update(inputs) return return_inputs - def _build_request_headers(self) -> dict: + def _build_request_headers(self) -> Mapping[str, str]: """Build HTTP headers for the guardrail API request.""" - headers: Final = {"Content-Type": "application/json"} - if self.headers: - headers.update(self.headers) - return headers + return {"Content-Type": "application/json", **self.headers} def _request_body( self, @@ -463,6 +485,35 @@ class GenericGuardrailAPI(CustomGuardrail): verbose_proxy_logger.error("Generic Guardrail API: failed to make request: %s", str(error)) raise Exception(f"Generic Guardrail API failed: {error}") + async def _post_check( + self, + payload: Mapping[str, object], + headers: Mapping[str, str], + retries_left: int, + fresh: bool = False, + ) -> httpx.Response: + handler: Final = self.fresh_handler_factory() if fresh else self.async_handler + try: + return await handler.post( + url=self.api_base, json=dict(payload), headers=dict(headers), timeout=self.timeout + ) + except (Timeout, httpx.TransportError) as error: + failure: Final = _after_connect_failure(error) + if retries_left < 1 or failure is None: + raise + verbose_proxy_logger.warning( + "Generic Guardrail API check timed out or lost its connection, retrying on a fresh connection. " + "guardrail_name=%s api_base=%s error_type=%s retries_left=%d", + getattr(self, "guardrail_name", None), + self.api_base, + type(failure).__name__, + retries_left - 1, + ) + finally: + if fresh: + await handler.close() + return await self._post_check(payload, headers, retries_left - 1, fresh=True) + @log_guardrail_information async def apply_guardrail( self, @@ -508,12 +559,8 @@ class GenericGuardrailAPI(CustomGuardrail): request_body: Final = request_data.get("body") or {} # Merge additional provider specific params from config and dynamic params - additional_params: Final = {**self.additional_provider_specific_params} - - # Get dynamic params from request if available - dynamic_params: Final = self.get_guardrail_dynamic_request_body_params(request_body) - if dynamic_params: - additional_params.update(dynamic_params) + dynamic_params: Final[Mapping[str, object]] = self.get_guardrail_dynamic_request_body_params(request_body) or {} + additional_params: Final = {**self.additional_provider_specific_params, **dynamic_params} # Extract user API key metadata user_metadata: Final = self._extract_user_api_key_metadata(request_data) @@ -546,11 +593,10 @@ class GenericGuardrailAPI(CustomGuardrail): # Make the API request # Use mode="json" to ensure all iterables are converted to lists - response: Final = await self.async_handler.post( - url=self.api_base, - json=dict(self._request_body(guardrail_request, structured_messages)), - headers=headers, - timeout=self.timeout, + response: Final = await self._post_check( + self._request_body(guardrail_request, structured_messages), + headers, + self.num_retries, ) response.raise_for_status() diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 33db8cbc6f0..ff063d2e233 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -2,7 +2,7 @@ from collections.abc import Mapping from datetime import datetime from enum import Enum from types import MappingProxyType -from typing import Final, Literal, cast +from typing import Annotated, Final, Literal, TypeAlias, cast from pydantic import ConfigDict, Field, field_validator, model_validator from pydantic.fields import FieldInfo @@ -1012,6 +1012,22 @@ def _default_if_blank(field: FieldInfo | None, value: object) -> object: return field.get_default(call_default_factory=True) +GuardrailNumRetries: TypeAlias = Annotated[ + int | None, + Field( + ge=0, + le=5, + description=( + "How many more times to send a guardrail check whose connection timed out waiting for the answer " + "or was dropped after the request went out. Each retry opens a fresh connection and gets the full " + "timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 " + "sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', " + "other guardrails ignore it." + ), + ), +] + + class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails api_key: str | None = Field(default=None, description="API key for the guardrail service") api_base: str | None = Field(default=None, description="Base URL for the guardrail service API") @@ -1220,6 +1236,8 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up description="Python-like code containing the apply_guardrail function for custom guardrail logic", ) + num_retries: GuardrailNumRetries = None + timeout: float | None = Field( default=None, description=( @@ -1437,7 +1455,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o super().__init__(**kwargs) # pyright: ignore[reportArgumentType] # pydantic validates each field's type at runtime - def __contains__(self, key) -> bool: + def __contains__(self, key: str) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index d34a3a13838..d38e54986b9 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -4,6 +4,7 @@ from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows h from pydantic import ConfigDict, Field, SerializerFunctionWrapHandler, model_serializer from typing_extensions import TypedDict +from litellm.types.guardrails import GuardrailNumRetries from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( AllMessageValues, @@ -70,6 +71,8 @@ class GenericGuardrailAPIOptionalParams(LiteLLMBaseModel): ), ) + num_retries: GuardrailNumRetries = None + forward_messages_as_sent: bool | None = Field( default=False, description=( diff --git a/tests/code_coverage_tests/ensure_async_clients_test.py b/tests/code_coverage_tests/ensure_async_clients_test.py index 5a534199cbc..93ae2e2837b 100644 --- a/tests/code_coverage_tests/ensure_async_clients_test.py +++ b/tests/code_coverage_tests/ensure_async_clients_test.py @@ -5,6 +5,9 @@ ALLOWED_FILES = [ # Lens data traffic owns one pool per app lifespan, isolated from model traffic and closed on shutdown. "../../litellm/tracing/remote.py", "./litellm/tracing/remote.py", + # Only a guardrail check retry opens a client, on purpose to get a new connection, and closes it right after. + "../../litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py", + "./litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py", # local files "../../litellm/__init__.py", "../../litellm/llms/custom_httpx/http_handler.py", diff --git a/tests/integration/observability/test_generic_guardrail_num_retries.py b/tests/integration/observability/test_generic_guardrail_num_retries.py new file mode 100644 index 00000000000..df761d25968 --- /dev/null +++ b/tests/integration/observability/test_generic_guardrail_num_retries.py @@ -0,0 +1,179 @@ +"""generic_guardrail_api `num_retries` resends a check that timed out waiting for its answer, through a real proxy. + +Each guardrail points at its own owned stub. A stub that holds a check never answers it before the 1s guardrail +timeout, and holding lasts until teardown so a resent check can only be answered by a later arrival. +""" + +import itertools +import json +import threading +from collections.abc import Callable, Iterator, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +MODEL: Final = "guarded-chat" +UPSTREAM_MODEL: Final = "openai/gpt-6.1-sol" +GUARDRAIL_TIMEOUT_SECONDS: Final = 1 +HOLDS_FIRST_RETRIED: Final = "holds-first-check-retried-once" +HOLDS_FIRST_UNSET: Final = "holds-first-check-num-retries-unset" +HOLDS_ALL_RETRIED_TWICE: Final = "holds-every-check-retried-twice" +UNAVAILABLE_RETRIED_TWICE: Final = "unavailable-retried-twice" +ANSWER: Final = "upstream answer" +NO_ACTION: Final = Reply(body=json.dumps({"action": "NONE"}).encode()) +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class Rig: + owned: OwnedProxy + upstream: Wire + guardrails: Mapping[str, Wire] + + +def _holds_first(release: threading.Event) -> Callable[[Request], Reply]: + arrivals: Final = itertools.count() + + def respond(_request: Request) -> Reply: + if next(arrivals) == 0: + release.wait() + return NO_ACTION + + return respond + + +def _holds_all(release: threading.Event) -> Callable[[Request], Reply]: + def respond(_request: Request) -> Reply: + release.wait() + return NO_ACTION + + return respond + + +def _unavailable(_request: Request) -> Reply: + return Reply(status=503, body=b'{"error": "guardrail unavailable"}') + + +def _upstream(_request: Request) -> Reply: + return chat_reply("chatcmpl-num-retries", "gpt-6.1-sol", ANSWER, stream=False) + + +def _guardrail(name: str, wire: Wire, num_retries: int | None) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": wire.url, + "api_key": "synthetic-guardrail-key", + "timeout": GUARDRAIL_TIMEOUT_SECONDS, + **({} if num_retries is None else {"num_retries": num_retries}), + }, + } + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("generic-guardrail-num-retries") + release: Final = threading.Event() + with ExitStack() as stack: + upstream: Final = stack.enter_context(wire_server(_upstream)) + guardrails: Final = MappingProxyType( + { + HOLDS_FIRST_RETRIED: stack.enter_context(wire_server(_holds_first(release), keep_alive=True)), + HOLDS_FIRST_UNSET: stack.enter_context(wire_server(_holds_first(release), keep_alive=True)), + HOLDS_ALL_RETRIED_TWICE: stack.enter_context(wire_server(_holds_all(release), keep_alive=True)), + UNAVAILABLE_RETRIED_TWICE: stack.enter_context(wire_server(_unavailable, keep_alive=True)), + } + ) + stack.callback(release.set) + retries: Final = MappingProxyType( + {HOLDS_FIRST_RETRIED: 1, HOLDS_FIRST_UNSET: None, HOLDS_ALL_RETRIED_TWICE: 2, UNAVAILABLE_RETRIED_TWICE: 2} + ) + base: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + path: Final = directory / "generic-guardrail-num-retries.yaml" + path.write_text( + yaml.safe_dump( + { + **base, + "model_list": [ + { + "model_name": MODEL, + "litellm_params": { + "model": UPSTREAM_MODEL, + "api_base": f"{upstream.url}/v1", + "api_key": "synthetic-openai-key", + }, + } + ], + "guardrails": [_guardrail(name, wire, retries[name]) for name, wire in guardrails.items()], + } + ) + ) + owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, {}, config=path)) + yield Rig(owned, upstream, guardrails) + + +def _chat(rig: Rig, guardrail: str) -> httpx.Response: + return rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": MODEL, "messages": [{"role": "user", "content": "is this allowed?"}], "guardrails": [guardrail]}, + ) + + +def _checks(wire: Wire) -> tuple[bytes, ...]: + return tuple(request.body for request in wire.drain()) + + +def test_check_held_past_the_timeout_is_resent_on_a_fresh_connection_and_the_request_is_served(rig: Rig) -> None: + response: Final = _chat(rig, HOLDS_FIRST_RETRIED) + + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == ANSWER + stub: Final = rig.guardrails[HOLDS_FIRST_RETRIED] + checks: Final = _checks(stub) + assert len(checks) == 2, checks + assert checks[0] == checks[1] + assert stub.connections() == 2 + assert len(rig.upstream.drain()) == 1 + + +def test_check_held_past_the_timeout_is_sent_once_when_num_retries_is_unset(rig: Rig) -> None: + response: Final = _chat(rig, HOLDS_FIRST_UNSET) + + assert response.status_code != 200, response.text + assert "Generic Guardrail API failed" in response.text + assert len(_checks(rig.guardrails[HOLDS_FIRST_UNSET])) == 1 + assert rig.upstream.drain() == () + + +def test_num_retries_bounds_how_many_times_a_held_check_is_sent(rig: Rig) -> None: + response: Final = _chat(rig, HOLDS_ALL_RETRIED_TWICE) + + assert response.status_code != 200, response.text + assert "Generic Guardrail API failed" in response.text + assert len(_checks(rig.guardrails[HOLDS_ALL_RETRIED_TWICE])) == 3 + assert rig.upstream.drain() == () + + +def test_error_status_from_the_guardrail_is_not_resent(rig: Rig) -> None: + response: Final = _chat(rig, UNAVAILABLE_RETRIED_TWICE) + + assert response.status_code != 200, response.text + assert "Generic Guardrail API failed" in response.text + assert len(_checks(rig.guardrails[UNAVAILABLE_RETRIED_TWICE])) == 1 + assert rig.upstream.drain() == () diff --git a/tests/integration/providers/_idle_upstream_support.py b/tests/integration/providers/_idle_upstream_support.py index ca808abc008..5fccd9804a6 100644 --- a/tests/integration/providers/_idle_upstream_support.py +++ b/tests/integration/providers/_idle_upstream_support.py @@ -294,11 +294,11 @@ def chat_httpx(target: ProxyTarget, model: str, marker: str) -> Answer: ) -def chat_with_guardrail(target: ProxyTarget, model: str, marker: str) -> Answer: +def chat_with_guardrail(target: ProxyTarget, model: str, marker: str, guardrail: str = GUARDRAIL) -> Answer: return _post( target, "/v1/chat/completions", - {"model": model, "messages": [{"role": "user", "content": marker}], "guardrails": [GUARDRAIL]}, + {"model": model, "messages": [{"role": "user", "content": marker}], "guardrails": [guardrail]}, _chat_body, ) @@ -517,18 +517,21 @@ def pass_through_deployment(_model: str, url: str, _steady: str) -> ConfigPart: return ConfigPart(pass_through=({"path": PASS_THROUGH, "target": f"{url}/passthrough"},)) -def guardrail_deployment(model: str, url: str, steady: str) -> ConfigPart: +def guardrail_deployment( + model: str, url: str, steady: str, guardrail: str = GUARDRAIL, num_retries: int | None = None +) -> ConfigPart: return ConfigPart( models=openai_deployment(model, steady, steady).models, guardrails=( { - "guardrail_name": GUARDRAIL, + "guardrail_name": guardrail, "litellm_params": { "guardrail": "generic_guardrail_api", "mode": "pre_call", "default_on": False, "api_base": url, "api_key": PROVIDER_KEY, + **({} if num_retries is None else {"num_retries": num_retries}), }, }, ), diff --git a/tests/integration/providers/test_idle_upstream_guardrail_cut_wire.py b/tests/integration/providers/test_idle_upstream_guardrail_cut_wire.py new file mode 100644 index 00000000000..6b12e61aa4e --- /dev/null +++ b/tests/integration/providers/test_idle_upstream_guardrail_cut_wire.py @@ -0,0 +1,134 @@ +import threading +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from functools import partial +from itertools import chain +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.idle_gateway import Arrival +from integration._support.process import graceful_stop_seconds +from integration._support.wire import Reply, Request +from integration.providers._idle_upstream_support import ( + REUSE_IDLE_SECONDS, + ProxyTarget, + Row, + Sent, + UpstreamProxy, + answer, + answered, + chat_with_guardrail, + expected_answers, + guardrail_deployment, + landed, + on_each_worker, + send, + spend_statuses, + stay_idle, + upstream_proxy, +) + +CUT_AFTER_SECONDS: Final = 3.0 +POOLED_PER_WORKER: Final = 2 +RETRIED: Final = "guardrail-cut-num-retries-1" +UNSET: Final = "guardrail-cut-num-retries-unset" + +pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 180 + REUSE_IDLE_SECONDS)) + + +@dataclass(frozen=True, slots=True) +class PairedWarmup: + """While warming, each check waits for a second one so every worker pools two connections to the guardrail.""" + + warming: threading.Event = field(default_factory=threading.Event) + pair: threading.Barrier = field(default_factory=lambda: threading.Barrier(POOLED_PER_WORKER)) + + def respond(self, request: Request) -> Reply: + if self.warming.is_set(): + self.pair.wait(timeout=30) + return answer(request) + + +@dataclass(frozen=True, slots=True) +class StaleGuardrails: + proxy: UpstreamProxy + pooled: Mapping[str, frozenset[int]] + + +WARMUPS: Final = MappingProxyType({RETRIED: PairedWarmup(), UNSET: PairedWarmup()}) +ROWS: Final = ( + Row( + RETRIED, + partial(guardrail_deployment, guardrail=RETRIED, num_retries=1), + partial(chat_with_guardrail, guardrail=RETRIED), + respond=WARMUPS[RETRIED].respond, + ), + Row( + UNSET, + partial(guardrail_deployment, guardrail=UNSET), + partial(chat_with_guardrail, guardrail=UNSET), + respond=WARMUPS[UNSET].respond, + ), +) + + +def _pair(row: Row, target: ProxyTarget) -> tuple[Sent, ...]: + with ThreadPoolExecutor(max_workers=POOLED_PER_WORKER) as pool: + return tuple(pool.map(send, (row,) * POOLED_PER_WORKER, (target,) * POOLED_PER_WORKER)) + + +def _warm(proxy: UpstreamProxy, row: Row) -> frozenset[int]: + warmup: Final = WARMUPS[row.name] + warmup.warming.set() + sents: Final = on_each_worker(proxy.owned, proxy.workers, lambda: _pair(row, proxy.target)) + warmup.warming.clear() + assert answered(sents) == expected_answers(row, sents), sents + return frozenset(arrival.connection for arrival in proxy.gateways[row.name].drain()) + + +@pytest.fixture(scope="module") +def stale(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StaleGuardrails]: + directory: Final = tmp_path_factory.mktemp("idle-upstream-guardrail-cut") + with upstream_proxy(session_gateway, directory, ROWS, {}, drop_after=CUT_AFTER_SECONDS) as proxy: + pooled: Final = MappingProxyType({row.name: _warm(proxy, row) for row in ROWS}) + assert {name: len(connections) for name, connections in pooled.items()} == { + row.name: 2 * POOLED_PER_WORKER for row in ROWS + }, pooled + stay_idle(tuple(proxy.gateways.values()), REUSE_IDLE_SECONDS) + yield StaleGuardrails(proxy, pooled) + + +def _cut_once(sents: tuple[Sent, ...]) -> tuple[tuple[str, bool], ...]: + return tuple(sorted((sent.marker, True) for sent in sents)) + + +def _cut_once_then_landed(sents: tuple[Sent, ...]) -> tuple[tuple[str, bool], ...]: + return tuple(sorted(chain.from_iterable(((sent.marker, False), (sent.marker, True)) for sent in sents))) + + +def _landed_on(arrivals: tuple[Arrival, ...]) -> frozenset[int]: + return frozenset(arrival.connection for arrival in arrivals if not arrival.dropped) + + +def test_retried_check_cut_on_a_stale_pooled_connection_lands_on_a_connection_that_was_never_pooled( + stale: StaleGuardrails, +) -> None: + sents, arrivals = stale.proxy.probe(RETRIED) + + assert answered(sents) == expected_answers(stale.proxy.rows[RETRIED], sents), (sents, arrivals) + assert landed(arrivals) == _cut_once_then_landed(sents), arrivals + assert _landed_on(arrivals).isdisjoint(stale.pooled[RETRIED]), (arrivals, stale.pooled) + assert spend_statuses(sents) == ("success",) * len(sents), sents + + +def test_check_cut_on_a_stale_pooled_connection_is_not_resent_when_num_retries_is_unset( + stale: StaleGuardrails, +) -> None: + sents, arrivals = stale.proxy.probe(UNSET) + + assert [sent.answer.status == 200 for sent in sents] == [False] * len(sents), sents + assert all("Generic Guardrail API failed" in sent.answer.text for sent in sents), sents + assert landed(arrivals) == _cut_once(sents), arrivals diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index f50da1bff64..8308a7ba386 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -8,6 +8,9 @@ specifically focusing on metadata extraction and passing. import json import logging import os +from collections.abc import Callable +from functools import partial +from queue import SimpleQueue from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -26,7 +29,7 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( _HEADER_PRESENT_PLACEHOLDER, ) -from litellm.types.utils import Choices, Message +from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message @pytest.fixture @@ -2470,3 +2473,151 @@ class TestFailOnError: request_data={}, input_type="response", ) + + +Outcome = Exception | Callable[[], httpx.Response] + +REWRITTEN: Final = partial(httpx.Response, 200, json={"action": "GUARDRAIL_INTERVENED", "texts": ["redacted"]}) + + +def _scripted(*outcomes: Outcome) -> tuple[httpx.MockTransport, SimpleQueue[bytes]]: + """Each check gets the next outcome; the last one repeats once the script runs out.""" + sent: Final[SimpleQueue[bytes]] = SimpleQueue() + script: Final = iter(outcomes) + + def answer(request: httpx.Request) -> httpx.Response: + sent.put(request.read()) + outcome: Final = next(script, outcomes[-1]) + if isinstance(outcome, Exception): + raise outcome + return outcome() + + return httpx.MockTransport(answer), sent + + +def _guardrail_on( + transport: httpx.MockTransport, num_retries: int | None, fail_on_error: bool = True +) -> GenericGuardrailAPI: + return GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="retrying-guardrail", + fail_on_error=fail_on_error, + num_retries=num_retries, + timeout=5, + async_handler=AsyncHTTPHandler(transport=transport), + fresh_handler_factory=lambda: AsyncHTTPHandler(transport=transport), + ) + + +async def _check(guardrail: GenericGuardrailAPI) -> GenericGuardrailAPIInputs: + return await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type="request") + + +async def _attempts_until_it_raises(outcome: Outcome, num_retries: int | None) -> int: + transport, sent = _scripted(outcome) + with pytest.raises(Exception, match=r"Generic Guardrail API failed|policy violation"): + await _check(_guardrail_on(transport, num_retries)) + return sent.qsize() + + +class _OpenCloseLoggingHandler(AsyncHTTPHandler): + def __init__(self, transport: httpx.MockTransport, events: SimpleQueue[str]) -> None: + super().__init__(transport=transport) + self._events: Final = events + events.put("open") + + async def close(self) -> None: + self._events.put("close") + await super().close() + + +class TestNumRetries: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("outcomes", "fresh_connections"), + [ + ((httpx.ReadTimeout("no answer before the read timeout"),), 3), + ((httpx.ReadTimeout("no answer"), httpx.ReadTimeout("no answer"), REWRITTEN), 2), + ], + ids=["every_retry_fails", "second_retry_succeeds"], + ) + async def test_each_retry_connection_closes_before_the_next_attempt( + self, outcomes: tuple[Outcome, ...], fresh_connections: int + ) -> None: + transport, _ = _scripted(*outcomes) + events: Final[SimpleQueue[str]] = SimpleQueue() + guardrail: Final = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="retrying-guardrail", + fail_on_error=False, + num_retries=3, + timeout=5, + async_handler=AsyncHTTPHandler(transport=transport), + fresh_handler_factory=lambda: _OpenCloseLoggingHandler(transport, events), + ) + + await _check(guardrail) + + assert [events.get_nowait() for _ in range(events.qsize())] == ["open", "close"] * fresh_connections + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "failure", + [ + httpx.ReadTimeout("no answer before the read timeout"), + httpx.WriteTimeout("request body stalled"), + httpx.ReadError("connection reset by peer"), + httpx.WriteError("broken pipe"), + httpx.RemoteProtocolError("Server disconnected without sending a response."), + ], + ids=["read_timeout", "write_timeout", "read_error", "write_error", "remote_protocol_error"], + ) + async def test_check_lost_after_connect_is_resent_and_the_retry_verdict_applies(self, failure: Exception) -> None: + transport, sent = _scripted(failure, REWRITTEN) + + result: Final = await _check(_guardrail_on(transport, num_retries=1)) + + bodies: Final = [sent.get_nowait() for _ in range(sent.qsize())] + assert len(bodies) == 2 + assert bodies[0] == bodies[1] + assert result == {"texts": ["redacted"]} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "outcome", + [ + httpx.ConnectTimeout("no TCP handshake before the timeout"), + httpx.ConnectError("connection refused"), + httpx.PoolTimeout("no free pooled connection"), + partial(httpx.Response, 503, json={"error": "unavailable"}), + partial(httpx.Response, 500, json={"error": "boom"}), + partial(httpx.Response, 200, json={"action": "BLOCKED", "blocked_reason": "policy violation"}), + ], + ids=["connect_timeout", "connect_error", "pool_timeout", "http_503", "http_500", "blocked"], + ) + async def test_num_retries_does_not_resend_connect_failures_error_statuses_or_blocks( + self, outcome: Outcome + ) -> None: + assert await _attempts_until_it_raises(outcome, num_retries=3) == await _attempts_until_it_raises( + outcome, num_retries=None + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("num_retries", [None, 0]) + async def test_unset_num_retries_sends_a_timed_out_check_once(self, num_retries: int | None) -> None: + transport, sent = _scripted(httpx.ReadTimeout("no answer before the read timeout"), REWRITTEN) + + with pytest.raises(Exception, match="Generic Guardrail API failed"): + await _check(_guardrail_on(transport, num_retries)) + + assert sent.qsize() == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("num_retries", [1, 3]) + async def test_num_retries_bounds_the_attempts_before_fail_on_error_applies(self, num_retries: int) -> None: + transport, sent = _scripted(httpx.ReadTimeout("no answer before the read timeout")) + + result: Final = await _check(_guardrail_on(transport, num_retries, fail_on_error=False)) + + assert sent.qsize() == num_retries + 1 + assert result == {"texts": ["hi"]} diff --git a/tests/unit/types/test_guardrails.py b/tests/unit/types/test_guardrails.py index 5c652a10e16..97bab8295f4 100644 --- a/tests/unit/types/test_guardrails.py +++ b/tests/unit/types/test_guardrails.py @@ -7,6 +7,7 @@ from typing import Any, List, Optional, Tuple from unittest.mock import MagicMock, patch import pytest +from pydantic import ValidationError import litellm import litellm.litellm_core_utils @@ -178,3 +179,16 @@ def test_litellm_params_treats_blank_values_as_unset(blank: str) -> None: assert with_blank_values.model_dump() == omitted.model_dump() assert BaseLitellmParams(fail_on_error=blank).model_dump() == BaseLitellmParams().model_dump() + + +@pytest.mark.parametrize(("configured", "loaded"), [("2", 2), (2, 2), ("0", 0), (5, 5), ("", None)]) +def test_num_retries_loads_the_dashboard_string_and_blank_means_unset(configured: object, loaded: int | None) -> None: + assert ( + LitellmParams(guardrail="generic_guardrail_api", mode="pre_call", num_retries=configured).num_retries == loaded + ) + + +@pytest.mark.parametrize("configured", [-1, 6, "1.5", "twice"]) +def test_num_retries_rejects_values_that_are_not_a_retry_count(configured: object) -> None: + with pytest.raises(ValidationError, match="num_retries"): + LitellmParams(guardrail="generic_guardrail_api", mode="pre_call", num_retries=configured) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7fae3437f22..e5c1b9f2c26 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27052,6 +27052,11 @@ export interface components { * @description Optional field if guardrail requires a 'model' parameter */ model?: string | null; + /** + * Num Retries + * @description How many more times to send a guardrail check whose connection timed out waiting for the answer or was dropped after the request went out. Each retry opens a fresh connection and gets the full timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', other guardrails ignore it. + */ + num_retries?: number | null; /** * On Sensitive Data * @description Action to take when sensitive data is detected. 'block' raises an exception (default behavior). 'route' reroutes the request to the model specified in sensitive_data_route_to_model. @@ -36638,6 +36643,11 @@ export interface components { * @description If True, logs violations without blocking. Defaults to False if not provided */ monitor_mode?: boolean | null; + /** + * Num Retries + * @description How many more times to send a guardrail check whose connection timed out waiting for the answer or was dropped after the request went out. Each retry opens a fresh connection and gets the full timeout. Connect failures, HTTP error statuses and BLOCKED verdicts are not retried. Unset or 0 sends each check once, and the most allowed is 5. Implemented by guardrail='generic_guardrail_api', other guardrails ignore it. + */ + num_retries?: number | null; /** * On Disallowed Action * @description Choose whether disallowed tools block the request or get rewritten out of the payload From a79dae764f0b55baf27abf8e59fef4b5d0404ec8 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Sat, 17 Oct 2026 14:32:03 -0500 Subject: [PATCH 32/64] fix(proxy): route team key vector store calls to a team model of the store's provider (#46894) * test(vector-stores): reproduce team key search routed to a team model of the wrong provider A team key's model-less vector store search picks the team's first deployment, so a langgraph team model turns a Bedrock store search into "Vector store search is not supported for langgraph". Cover map_team_model picking by provider and a registered Bedrock store search reaching Bedrock with the registry's credentials. * fix(router): match model-less team deployment to the requested provider map_team_model now takes an optional custom_llm_provider. With no model given, it returns the first team deployment serving that provider, or None when none does, instead of the team's first deployment. Without a provider it behaves as before, so BYOK teams calling raw vector store IDs keep their team key. * fix(proxy): pass the vector store provider when mapping team models Vector store search, retrieve, update and delete pass the store's custom_llm_provider to map_team_model, so a team key's model-less call lands on a team deployment of that provider, or on the direct path with the store's own credentials when the team has none. Other model-less routes keep today's behavior. * fix(proxy): type the vector store provider passed to map_team_model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(vector-stores): cover team key vector store search provider routing end to end Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): ruff format route_llm_request.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): type vector store routing tests and mock Bedrock transport Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): pin model-less vector store calls when a public name mixes providers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): keep team provider routing in map_team_model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/route_llm_request.py | 18 +- litellm/router.py | 58 ++++-- ...t_vector_store_team_model_provider_wire.py | 192 ++++++++++++++++++ tests/unit/proxy/test_route_llm_request.py | 92 +++++++++ tests/unit/test_router/test_router.py | 99 +++++++++ 5 files changed, 444 insertions(+), 15 deletions(-) create mode 100644 tests/integration/providers/test_vector_store_team_model_provider_wire.py diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 6f93bcd82dc..8bab7215b17 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -813,7 +813,23 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr route=ROUTE_ENDPOINT_MAPPING.get(route_type, route_type), param="model", ) - team_model_name: Final = llm_router.map_team_model(data["model"], team_id) if team_id is not None else None + custom_llm_provider: Final[object] = ( + JSON_OBJECT_ADAPTER.validate_python(data).get("custom_llm_provider") + if route_type + in ( + "avector_store_search", + "avector_store_retrieve", + "avector_store_update", + "avector_store_delete", + ) + else None + ) + team_model_provider: Final[str | None] = custom_llm_provider if isinstance(custom_llm_provider, str) else None + team_model_name: Final = ( + llm_router.map_team_model(data["model"], team_id, custom_llm_provider=team_model_provider) + if team_id is not None + else None + ) if team_model_name is not None: data["model"] = team_model_name return getattr(llm_router, f"{route_type}")(**data) diff --git a/litellm/router.py b/litellm/router.py index 1b3857dfb71..8d77fd0dac5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12565,26 +12565,56 @@ class Router: members: Final = routing_group.models if routing_group is not None else () return served_model if served_model == target or served_model in members else model_name - def map_team_model(self, team_model_name: str | None, team_id: str) -> str | None: + def map_team_model( + self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None + ) -> str | None: """ - Check if team_model_name resolves to team-specific deployments. + Resolve explicit team models or select a deployment for model-less calls. - Returns the public model name (unchanged) so the router can find all - sibling deployments via team_id filtering, instead of collapsing to a - single internal model_name. - - When team_model_name is None (e.g. vector store / file endpoints that - don't include a model in their request), returns the first matching - team deployment's team_public_model_name so the router can inject BYOK - credentials from the team-scoped deployment. - - Returns: - - str: the team_model_name if team deployments exist for this team - - None: if no team-specific model is found + Model-less calls with no provider use the first team deployment. With a provider, they + use the first team deployment of that provider, pinned by ID when its public name also + covers another provider. """ models: Final = self.get_model_list(model_name=team_model_name, team_id=team_id) if not models: return None + if team_model_name is None and custom_llm_provider is not None: + + def team_deployment_public_name(deployment: DeploymentTypedDict) -> str | None: + model_info: Final = _configured_model_info(deployment) + if model_info.get("team_id") != team_id: + return None + team_public_model_name: Final = model_info.get("team_public_model_name") + return ( + team_public_model_name + if isinstance(team_public_model_name, str) and team_public_model_name + else deployment["model_name"] + ) + + matching_deployment: Final = next( + ( + model + for model in models + if team_deployment_public_name(model) is not None + and provider_for_generic_call(model["litellm_params"]) == custom_llm_provider + ), + None, + ) + if matching_deployment is None: + return None + public_name: Final = team_deployment_public_name(matching_deployment) + if public_name is None: + return None + same_public_name_deployments: Final = tuple( + candidate for candidate in models if team_deployment_public_name(candidate) == public_name + ) + if all( + provider_for_generic_call(candidate["litellm_params"]) == custom_llm_provider + for candidate in same_public_name_deployments + ): + return public_name + deployment_id: Final = _configured_model_info(matching_deployment).get("id") + return deployment_id if isinstance(deployment_id, str) else None for model in models: if model.get("model_info", {}).get("team_id") == team_id: if team_model_name is None: diff --git a/tests/integration/providers/test_vector_store_team_model_provider_wire.py b/tests/integration/providers/test_vector_store_team_model_provider_wire.py new file mode 100644 index 00000000000..5b986a64ff6 --- /dev/null +++ b/tests/integration/providers/test_vector_store_team_model_provider_wire.py @@ -0,0 +1,192 @@ +import json +import uuid +from collections.abc import Callable +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +QUERY: Final = "team provider vector search" +TEAM_ACCESS_KEY: Final = "AKIAINTEGRATIONTEAM01" +STORE_ACCESS_KEY: Final = "AKIASTOREKEY00000001" +_RETRIEVE_RESPONSE: Final = json.dumps( + { + "retrievalResults": [ + { + "content": {"text": "matching Bedrock content"}, + "score": 0.87, + "metadata": { + "x-amz-bedrock-kb-source-uri": "s3://synthetic-bucket/vector-store.txt", + "x-amz-bedrock-kb-chunk-id": "chunk-1", + }, + } + ] + } +).encode() +_OPENAI_RESULT: Final[dict[str, JsonValue]] = { + "file_id": "file-search-wire", + "filename": "search.txt", + "score": 0.91, + "attributes": {"tenant": "a"}, + "content": [{"type": "text", "text": "matching OpenAI content"}], +} +_BEDROCK_RESULT: Final[dict[str, JsonValue]] = { + "score": 0.87, + "content": [{"text": "matching Bedrock content", "type": "text"}], + "file_id": "s3://synthetic-bucket/vector-store.txt", + "filename": "vector-store.txt", + "attributes": { + "x-amz-bedrock-kb-source-uri": "s3://synthetic-bucket/vector-store.txt", + "x-amz-bedrock-kb-chunk-id": "chunk-1", + }, +} + + +def _bedrock_retrieve_peer(knowledge_base_id: str, access_key: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == f"/knowledgebases/{knowledge_base_id}/retrieve", request.target + assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={access_key}/") + assert JSON_OBJECT.validate_json(request.body) == { + "retrievalQuery": {"text": QUERY}, + "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 3}}, + } + return Reply(body=_RETRIEVE_RESPONSE) + + return respond + + +def _openai_search_peer(store_id: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == f"/v1/vector_stores/{store_id}/search", request.target + assert request.headers["authorization"] == "Bearer team-openai-key" + assert JSON_OBJECT.validate_json(request.body)["query"] == QUERY + return Reply( + body=json.dumps( + { + "object": "vector_store.search_results.page", + "search_query": QUERY, + "data": [_OPENAI_RESULT], + "has_more": False, + "next_page": None, + } + ).encode() + ) + + return respond + + +def test_bedrock_vector_store_search_uses_bedrock_team_model_credentials(gateway: Gateway) -> None: + knowledge_base_id: Final = f"KB{uuid.uuid4().hex[:8].upper()}" + with wire_server(_bedrock_retrieve_peer(knowledge_base_id, TEAM_ACCESS_KEY)) as wire, gateway.scenario() as scenario: + team: Final = scenario.team() + scenario.model( + model="langgraph/repro-agent", + api_base="http://127.0.0.1:1", + model_info={"team_id": team}, + ) + scenario.model( + model="bedrock/us.anthropic.claude-sonnet-5-5", + api_base=wire.url, + aws_access_key_id=TEAM_ACCESS_KEY, + aws_secret_access_key="synthetic-team-secret-key", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + model_info={"team_id": team}, + ) + team_key: Final = scenario.key(team_id=team) + gateway.post( + "/vector_store/new", + { + "vector_store_id": knowledge_base_id, + "custom_llm_provider": "bedrock", + "litellm_params": { + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": wire.url, + }, + }, + ) + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": knowledge_base_id}) + + response: Final = gateway.request( + "POST", + f"/v1/vector_stores/{knowledge_base_id}/search", + {"query": QUERY, "max_num_results": 3}, + key=team_key, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["data"] == [_BEDROCK_RESULT], response.text + assert len(wire.drain()) == 1 + + +def test_bedrock_vector_store_search_uses_store_credentials_without_matching_team_model( + gateway: Gateway, +) -> None: + knowledge_base_id: Final = f"KB{uuid.uuid4().hex[:8].upper()}" + with wire_server(_bedrock_retrieve_peer(knowledge_base_id, STORE_ACCESS_KEY)) as wire, gateway.scenario() as scenario: + team: Final = scenario.team() + scenario.model( + model="langgraph/repro-agent", + api_base="http://127.0.0.1:1", + model_info={"team_id": team}, + ) + team_key: Final = scenario.key(team_id=team) + gateway.post( + "/vector_store/new", + { + "vector_store_id": knowledge_base_id, + "custom_llm_provider": "bedrock", + "litellm_params": { + "aws_access_key_id": STORE_ACCESS_KEY, + "aws_secret_access_key": "synthetic-store-secret-key", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": wire.url, + }, + }, + ) + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": knowledge_base_id}) + + response: Final = gateway.request( + "POST", + f"/v1/vector_stores/{knowledge_base_id}/search", + {"query": QUERY, "max_num_results": 3}, + key=team_key, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["data"] == [_BEDROCK_RESULT], response.text + assert len(wire.drain()) == 1 + + +def test_openai_vector_store_search_without_store_key_uses_team_byok_key(gateway: Gateway) -> None: + store_id: Final = f"vs_team_{uuid.uuid4().hex}" + with wire_server(_openai_search_peer(store_id)) as wire, gateway.scenario() as scenario: + team: Final = scenario.team() + scenario.model( + model="openai/gpt-4o-mini", + api_key="team-openai-key", + api_base=f"{wire.url}/v1", + model_info={"team_id": team}, + ) + team_key: Final = scenario.key(team_id=team) + gateway.post( + "/vector_store/new", + { + "vector_store_id": store_id, + "custom_llm_provider": "openai", + "litellm_params": {}, + }, + ) + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store_id}) + + response: Final = gateway.request( + "POST", + f"/v1/vector_stores/{store_id}/search", + {"query": QUERY}, + key=team_key, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["data"] == [_OPENAI_RESULT], response.text + assert len(wire.drain()) == 1 diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 67f29220cb6..4e4e2f23950 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,4 +1,7 @@ +import httpx +import litellm import pytest +import respx from types import SimpleNamespace @@ -1890,3 +1893,92 @@ async def test_null_user_config_is_treated_as_absent() -> None: ) assert response.choices[0].message.content == "from the main router" + + +def _registered_bedrock_kb_search(team_id: str) -> dict[str, object]: + """The request a team key's /v1/vector_stores/{id}/search carries once the registry row is merged in.""" + return { + "model": None, + "vector_store_id": "KB12345678", + "query": "litellm proxy", + "custom_llm_provider": "bedrock", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIAREGISTRYKEY00000", + "aws_secret_access_key": "registry-secret", + "metadata": {"user_api_key_team_id": team_id}, + } + + +def _team_a_router(*team_models: tuple[str, dict[str, object]]) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": f"internal-team-a-{public_name}", + "litellm_params": litellm_params, + "model_info": { + "id": f"team-a-{public_name}", + "team_id": "team-a", + "team_public_model_name": public_name, + }, + } + for public_name, litellm_params in team_models + ] + ) + + +_LANGGRAPH_AGENT_PARAMS: Final[dict[str, object]] = { + "model": "langgraph/repro-agent", + "api_base": "http://127.0.0.1:2024", +} +_BEDROCK_CLAUDE_PARAMS: Final[dict[str, object]] = { + "model": "bedrock/anthropic.claude-sonnet-4-6", + "aws_region_name": "us-east-1", +} + + +@pytest.mark.parametrize( + "team_models", + [ + pytest.param( + [("a-langgraph-agent", _LANGGRAPH_AGENT_PARAMS), ("a-bedrock-claude", _BEDROCK_CLAUDE_PARAMS)], + id="langgraph-first", + ), + pytest.param( + [("a-bedrock-claude", _BEDROCK_CLAUDE_PARAMS), ("a-langgraph-agent", _LANGGRAPH_AGENT_PARAMS)], + id="bedrock-first", + ), + pytest.param([("a-langgraph-agent", _LANGGRAPH_AGENT_PARAMS)], id="no-bedrock-team-model"), + ], +) +@pytest.mark.asyncio +async def test_team_key_search_of_registered_bedrock_store_reaches_bedrock_with_registry_credentials( + monkeypatch: pytest.MonkeyPatch, + team_models: list[tuple[str, dict[str, object]]], +) -> None: + """A team key's model-less search of a registered Bedrock store was routed through the team's first + deployment, so a langgraph team model turned it into "Vector store search is not supported for langgraph".""" + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + with respx.mock(assert_all_called=True) as respx_mock: + route: Final = respx_mock.post( + "https://bedrock-agent-runtime.us-east-1.amazonaws.com/knowledgebases/KB12345678/retrieve" + ).mock(return_value=httpx.Response(200, json={"retrievalResults": []})) + response: Final = await ( + await route_request( + _registered_bedrock_kb_search("team-a"), + _team_a_router(*team_models), + None, + "avector_store_search", + ) + ) + + request: Final = route.calls.last.request + posted_url: Final = str(request.url) + signed_auth: Final = request.headers["Authorization"] + assert response["object"] == "vector_store.search_results.page", response + assert route.call_count == 1 + assert posted_url == "https://bedrock-agent-runtime.us-east-1.amazonaws.com/knowledgebases/KB12345678/retrieve", ( + posted_url + ) + assert "Credential=AKIAREGISTRYKEY00000/" in signed_auth, signed_auth diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 50a9d98385d..bd3cf738af8 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -23912,6 +23912,105 @@ def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name result = router.map_team_model(team_model_name="team-model", team_id="team-1") assert result == "team-model", f"Expected public name 'team-model', got {result}" + +def _team_deployment( + public_name: str, + litellm_model: str, + team_id: str = "team-a", + deployment_id: str | None = None, +) -> DeploymentTypedDict: + return { + "model_name": f"internal-{team_id}-{public_name}", + "litellm_params": {"model": litellm_model}, + "model_info": { + "id": deployment_id or f"{team_id}-{public_name}", + "team_id": team_id, + "team_public_model_name": public_name, + }, + } + + +_LANGGRAPH_AGENT: Final[DeploymentTypedDict] = _team_deployment("a-langgraph-agent", "langgraph/repro-agent") +_BEDROCK_CLAUDE: Final[DeploymentTypedDict] = _team_deployment( + "a-bedrock-claude", "bedrock/anthropic.claude-sonnet-4-6" +) +_OPENAI_GPT: Final[DeploymentTypedDict] = _team_deployment("a-gpt", "openai/gpt-4o-mini") + + +@pytest.mark.parametrize( + "team_models, custom_llm_provider, expected", + [ + pytest.param([_LANGGRAPH_AGENT, _BEDROCK_CLAUDE], "bedrock", "a-bedrock-claude", id="match-is-second"), + pytest.param([_BEDROCK_CLAUDE, _LANGGRAPH_AGENT], "bedrock", "a-bedrock-claude", id="match-is-first"), + pytest.param([_LANGGRAPH_AGENT, _OPENAI_GPT], "bedrock", None, id="no-match-returns-none"), + pytest.param([_LANGGRAPH_AGENT, _BEDROCK_CLAUDE], None, "a-langgraph-agent", id="no-provider-keeps-first"), + pytest.param( + [_team_deployment("b-bedrock-claude", "bedrock/anthropic.claude-sonnet-4-6", team_id="team-b")], + "bedrock", + None, + id="other-teams-match-ignored", + ), + pytest.param( + [ + _team_deployment("shared", "langgraph/repro-agent", deployment_id="shared-langgraph"), + _team_deployment( + "shared", + "bedrock/anthropic.claude-sonnet-4-6", + deployment_id="shared-bedrock", + ), + ], + "bedrock", + "shared-bedrock", + id="mixed-public-name-langgraph-first", + ), + pytest.param( + [ + _team_deployment( + "shared", + "bedrock/anthropic.claude-sonnet-4-6", + deployment_id="shared-bedrock", + ), + _team_deployment("shared", "langgraph/repro-agent", deployment_id="shared-langgraph"), + ], + "bedrock", + "shared-bedrock", + id="mixed-public-name-bedrock-first", + ), + pytest.param( + [ + _team_deployment("shared", "bedrock/anthropic.claude-sonnet-4-6", deployment_id="bedrock-one"), + _team_deployment("shared", "bedrock/anthropic.claude-sonnet-4-6", deployment_id="bedrock-two"), + ], + "bedrock", + "shared", + id="same-provider-public-name-keeps-load-balancing", + ), + ], +) +def test_map_team_model_without_model_picks_deployment_serving_the_requested_provider( + team_models: list[DeploymentTypedDict], custom_llm_provider: str | None, expected: str | None +) -> None: + """A model-less vector store call routed through the team's first deployment, so a langgraph + team model made a Bedrock vector store search fail with "not supported for langgraph".""" + router = Router(model_list=team_models) + + result = router.map_team_model(None, "team-a", custom_llm_provider=custom_llm_provider) + + assert result == expected, ( + f"team models {[m['litellm_params']['model'] for m in team_models]}, provider {custom_llm_provider}" + ) + + +def test_map_team_model_with_explicit_model_ignores_requested_provider() -> None: + router = Router(model_list=[_LANGGRAPH_AGENT, _BEDROCK_CLAUDE]) + + result = router.map_team_model("a-langgraph-agent", "team-a", custom_llm_provider="bedrock") + + assert result == "a-langgraph-agent", ( + "an explicit model is the caller's choice and must not be re-picked by provider" + ) + + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestPreCallChecksOptimization: """ From 4143e0869eff3db190a5e8f73069bfe2c9428303 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sat, 17 Oct 2026 14:49:31 -0500 Subject: [PATCH 33/64] fix(batches): handle empty models when logging batch cost (#47155) Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 3 +- .../common_utils/test_check_batch_cost.py | 118 +++++++++++++++++- 2 files changed, 119 insertions(+), 2 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 734092c4feb..94c3674945d 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -755,8 +755,9 @@ class CheckBatchCost: failed_requests=output_file_result.failed_requests + error_file_failed_requests, ) ) + logging_model: Final = batch_result.models[0] if batch_result.models else litellm_model_name or model_id logging_obj = LiteLLMLogging( - model=batch_result.models[0], + model=logging_model, messages=[{"role": "user", "content": ""}], stream=False, call_type="aretrieve_batch", diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index a5ee7fe23bb..c974e55c0c9 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -11,7 +11,7 @@ import json import logging from contextlib import contextmanager from hashlib import sha256 -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -3078,6 +3078,122 @@ class TestMultiPodBatchCostClaim: ): yield logging_obj + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("deployment_model", "successes", "failures", "expected_model"), + ( + ("openai/", 0, 2, "openai/"), + ("", 0, 2, "model-123"), + ("openai/", 2, 1, "batch-test-model"), + ("openai/batch-test-model", 2, 1, "batch-test-model"), + ("openai/batch-test-model", 2, 0, "batch-test-model"), + ), + ) + async def test_batch_accounting_preserves_models_and_mixed_results( + self, deployment_model: str, successes: int, failures: int, expected_model: str + ) -> None: + import base64 + from datetime import datetime + + import httpx + import litellm + import respx + + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + from litellm.types.utils import LiteLLMBatch, SpecialEnums + + input_rate: Final = 0.000002 + output_rate: Final = 0.000006 + litellm.register_model( + { + "batch-test-model": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token_batches": input_rate, + "output_cost_per_token_batches": output_rate, + } + } + ) + success: Final = { + "response": { + "status_code": 200, + "body": { + "model": "batch-test-model", + "usage": {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120}, + }, + }, + "error": None, + } + failure: Final = {"response": None, "error": {"code": "batch_expired", "message": "expired"}} + rows: Final = (success,) * successes + (failure,) * failures + content: Final = b"\n".join(json.dumps(row).encode() for row in rows) + + class BatchSpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.calls: tuple[dict[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if kwargs["call_type"] == "aretrieve_batch": + self.calls += (kwargs,) + + callback: Final = BatchSpendRecorder() + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + row: Final = _FakeManagedObjectRow() + row.unified_object_id = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode() + ) + .decode() + .rstrip("=") + ) + prisma: Final = self._prisma(row, []) + router: Final = self._router() + router.get_deployment.return_value = Deployment( + model_name="batch-alias", + litellm_params=LiteLLM_Params(model=deployment_model, custom_llm_provider="openai"), + model_info=ModelInfo(id="model-123"), + ) + router.aretrieve_batch.return_value = LiteLLMBatch( + id="batch-456", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + output_file_id=_CLAIM_OUTPUT_FILE_ID, + object="batch", + status="completed", + ) + router.get_deployment_credentials_with_provider.return_value = { + "api_key": "synthetic-key", + "custom_llm_provider": "openai", + } + with respx.mock: + respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock( + return_value=httpx.Response(200, content=content) + ) + instance: Final = self._instance(prisma, router) + await instance.check_batch_cost() + await instance.check_batch_cost() + + assert row.batch_processed is True + assert row.status == "complete" + router.aretrieve_batch.assert_awaited_once() + assert len(callback.calls) == 1 + payload: Final = callback.calls[0]["standard_logging_object"] + assert payload["model"] == expected_model + assert payload["response_cost"] == pytest.approx(successes * (100 * input_rate + 20 * output_rate)) + assert payload["prompt_tokens"] == successes * 100 + assert payload["completion_tokens"] == successes * 20 + assert payload["hidden_params"]["batch_successful_requests"] == successes + assert payload["hidden_params"]["batch_failed_requests"] == failures + assert payload["hidden_params"]["batch_models"] == ( + ["batch-test-model"] if deployment_model == "openai/batch-test-model" else ["batch-test-model"] * successes + ) + @staticmethod def _claim_calls(prisma) -> list: return [ From 4846cfabc7f4b8cf6a83a960b7d1123eebef4315 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 12:51:47 -0700 Subject: [PATCH 34/64] fix(vector-stores): keep proxy context out of provider vector store metadata (#47021) * fix(vector-stores): keep proxy context out of provider vector store metadata The proxy merged its request context (user_api_key_* fields, endpoint, user_agent, requester_ip_address, litellm_api_version) and the router's deployment info into the vector store request's metadata. For vector stores that field is also OpenAI's own metadata, which is capped at 16 keys, so a store created with a team key was full on creation and every later metadata update returned 400. Vector store routes now keep that context in litellm_metadata, like batches, files and responses, and the router does the same for avector_store_create. Provider-side metadata holds only what the caller sent. Adds db_scripts/strip_vector_store_internal_metadata.py, an opt-in script that removes the proxy-written keys from stores created before this change. It is a dry run unless --apply is passed. * fix(vector-stores): page the metadata cleanup script without recursion and tolerate non-string values Listing every store recursed once per page, which overflows the interpreter stack on large projects. Pages are now walked iteratively. Metadata values are no longer required to be strings, and a malformed response exits 2 with a fixed message instead of a traceback that could echo metadata values. * fix(vector-stores): read proxy agent id, tags and spend_logs_metadata from litellm_metadata Vector store routes now keep proxy context in litellm_metadata, but the v3 limiter only read agent_id from metadata, and logging let provider metadata tags or spend_logs_metadata replace the proxy's values. A string tags value also crashed get_request_tags. The limiter now prefers the proxy-resolved agent_id in litellm_metadata. Logging combines list tags from both buckets (metadata first, deduplicated) and merges dict spend_logs_metadata with litellm_metadata winning, ignoring values of the wrong type. * fix(vector-stores): narrow the logging fallback to mistyped caller values The tag union and spend_logs_metadata merge changed logged tags and spend_logs_metadata on every route that has both metadata buckets, deduped tags quadratically, and stopped background interaction polling from keeping string tags as sent. Caller values now win exactly as before, and the proxy's litellm_metadata tags or spend_logs_metadata only replace a caller value of the wrong type. The limiter reads agent ids through a pydantic model instead of a cast and ignores non-string ids. * fix(vector-stores): log the proxy's full tag list on litellm_metadata routes On routes that keep proxy context in litellm_metadata, the proxy builds litellm_metadata.tags from the caller's tags plus key, team and header tags, but get_request_tags preferred the caller's metadata.tags. A caller sending metadata.tags on a vector store call dropped the key's tags from request_tags, which tag spend tracking reads. get_request_tags now prefers a non-empty litellm_metadata tag list, and falls back to metadata.tags as before. * fix(vector-stores): strip LiteLLM-written keys at the provider boundary instead of moving metadata buckets Moving vector store proxy context into litellm_metadata broke shared readers (agent limits, request tags, merged logging metadata), and fixing them widened the change to every litellm_metadata route. This reverts the bucket move, the router change and the reader fixes, so proxy context stays in metadata exactly as on main. Vector store create and update now drop the keys LiteLLM itself writes (user_api_key*, requester and deployment keys, agent ids) from the metadata sent to the provider, in the shared optional-param helper both paths use. The cleanup script reuses the same key list. * fix(vector-stores): pass non-string metadata keys through the provider metadata filter * fix(vector-stores): strip LiteLLM-written metadata keys only on proxy and router calls * fix(vector-stores): send the provider exactly the metadata the caller sent through the proxy * fix(vector-stores): snapshot caller metadata from the raw request body and keep it out of logged request bodies * fix(vector-stores): copy the parsed-body metadata when the raw request body cannot be reparsed --- .github/workflows/test-unit.yml | 1 + .../strip_vector_store_internal_metadata.py | 201 ++++++++++++++++++ litellm/constants.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 3 +- .../proxy/vector_store_endpoints/endpoints.py | 9 + litellm/vector_stores/utils.py | 91 +++++++- tests/unit/db_scripts/__init__.py | 0 ...st_strip_vector_store_internal_metadata.py | 198 +++++++++++++++++ .../test_vector_store_endpoints.py | 157 ++++++++++++++ .../unit/test_router/test_router_endpoints.py | 37 ++++ tests/unit/vector_stores/test_main.py | 90 +++++++- 11 files changed, 782 insertions(+), 6 deletions(-) create mode 100644 db_scripts/strip_vector_store_internal_metadata.py create mode 100644 tests/unit/db_scripts/__init__.py create mode 100644 tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 1b486d78f04..cb1b595f3f0 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -658,6 +658,7 @@ jobs: - shard: proxy-extras test-path: >- tests/unit/litellm_proxy_extras + tests/unit/db_scripts python-version: "3.12" workers: 2 reruns: 0 diff --git a/db_scripts/strip_vector_store_internal_metadata.py b/db_scripts/strip_vector_store_internal_metadata.py new file mode 100644 index 00000000000..34ffa85dcd1 --- /dev/null +++ b/db_scripts/strip_vector_store_internal_metadata.py @@ -0,0 +1,201 @@ +import argparse +import os +import sys +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from itertools import accumulate, repeat, takewhile +from typing import Final, Literal, TextIO + +import httpx +from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError + +from litellm.vector_stores.utils import is_litellm_written_metadata_key + +DEFAULT_API_BASE: Final = "https://api.openai.com/v1" +LIST_PAGE_SIZE: Final = "100" + + +class VectorStore(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + metadata: dict[str, JsonValue] | None = None + + +class VectorStorePage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[VectorStore, ...] + has_more: bool = False + + +class OpenAIErrorBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + message: str = "" + + +class OpenAIErrorResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + error: OpenAIErrorBody + + +@dataclass(frozen=True, slots=True) +class Options: + api_key: str + api_base: str + vector_store_ids: tuple[str, ...] + extra_keys: frozenset[str] + apply: bool + + +@dataclass(frozen=True, slots=True) +class StoreResult: + store_id: str + matched_keys: tuple[str, ...] + status: Literal["dry_run", "updated", "failed"] + error: str = "" + + +def is_proxy_written_key(key: str, extra_keys: frozenset[str]) -> bool: + return is_litellm_written_metadata_key(key) or key in extra_keys + + +def matching_keys(store: VectorStore, extra_keys: frozenset[str]) -> tuple[str, ...]: + return tuple(sorted(key for key in (store.metadata or {}) if is_proxy_written_key(key, extra_keys))) + + +def fetch_page(client: httpx.Client, after: str | None) -> VectorStorePage: + params: Final = {"limit": LIST_PAGE_SIZE} if after is None else {"limit": LIST_PAGE_SIZE, "after": after} + return VectorStorePage.model_validate_json(client.get("/vector_stores", params=params).raise_for_status().content) + + +def fetch_next_page(client: httpx.Client, page: VectorStorePage | None) -> VectorStorePage | None: + if page is None or not page.has_more or not page.data: + return None + return fetch_page(client, after=page.data[-1].id) + + +def list_all_stores(client: httpx.Client) -> Iterator[VectorStore]: + pages: Final = accumulate( + repeat(None), lambda page, _: fetch_next_page(client, page), initial=fetch_page(client, after=None) + ) + for page in takewhile(lambda page: page is not None, pages): + yield from page.data + + +def retrieve_stores(client: httpx.Client, store_ids: tuple[str, ...]) -> Iterator[VectorStore]: + for store_id in store_ids: + yield VectorStore.model_validate_json(client.get(f"/vector_stores/{store_id}").raise_for_status().content) + + +def describe_http_error(response: httpx.Response) -> str: + try: + message: Final = OpenAIErrorResponse.model_validate_json(response.content).error.message + except ValidationError: + return f"HTTP {response.status_code}" + return f"HTTP {response.status_code}: {message}" + + +def delete_keys(client: httpx.Client, store_id: str, keys: tuple[str, ...]) -> StoreResult: + try: + response: Final = client.post(f"/vector_stores/{store_id}", json={"metadata": dict.fromkeys(keys)}) + except httpx.TransportError as error: + return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=type(error).__name__) + if response.is_success: + return StoreResult(store_id=store_id, matched_keys=keys, status="updated") + return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=describe_http_error(response)) + + +def clean_store(client: httpx.Client, store: VectorStore, options: Options) -> StoreResult | None: + keys: Final = matching_keys(store, options.extra_keys) + if not keys: + return None + if not options.apply: + return StoreResult(store_id=store.id, matched_keys=keys, status="dry_run") + return delete_keys(client, store.id, keys) + + +def format_result(result: StoreResult) -> str: + key_list: Final = ", ".join(result.matched_keys) + count: Final = len(result.matched_keys) + match result.status: + case "dry_run": + return f"{result.store_id}: would delete {count} keys: {key_list}" + case "updated": + return f"{result.store_id}: deleted {count} keys: {key_list}" + case "failed": + return f"{result.store_id}: failed to delete {count} keys ({result.error}): {key_list}" + + +def clean_and_report( + client: httpx.Client, stores: Iterator[VectorStore], options: Options, out: TextIO +) -> Iterator[StoreResult | None]: + for store in stores: + result: Final = clean_store(client, store, options) + if result is not None: + print(format_result(result), file=out) + yield result + + +def run(options: Options, client: httpx.Client, out: TextIO) -> int: + stores: Final = ( + retrieve_stores(client, options.vector_store_ids) if options.vector_store_ids else list_all_stores(client) + ) + scanned: Final = tuple(clean_and_report(client, stores, options, out)) + results: Final = tuple(result for result in scanned if result is not None) + updated: Final = sum(1 for result in results if result.status == "updated") + failed: Final = sum(1 for result in results if result.status == "failed") + mode: Final = "applied" if options.apply else "dry run, pass --apply to delete" + print( + f"scanned={len(scanned)} matched={len(results)} updated={updated} failed={failed} ({mode})", + file=out, + ) + return 1 if failed else 0 + + +def parse_options(argv: Sequence[str] | None) -> Options: + parser: Final = argparse.ArgumentParser( + description="Remove LiteLLM proxy-written keys from OpenAI vector store metadata. Dry run by default." + ) + parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"), help="defaults to $OPENAI_API_KEY") + parser.add_argument("--api-base", default=DEFAULT_API_BASE) + parser.add_argument("--vector-store-id", action="append", default=[], help="repeatable, defaults to every store") + parser.add_argument("--extra-key", action="append", default=[], help="repeatable, another key name to delete") + parser.add_argument("--apply", action="store_true", help="delete the matching keys instead of listing them") + args: Final = parser.parse_args(argv) + if not args.api_key: + parser.error("--api-key or OPENAI_API_KEY is required") + return Options( + api_key=args.api_key, + api_base=args.api_base.rstrip("/"), + vector_store_ids=tuple(args.vector_store_id), + extra_keys=frozenset(args.extra_key), + apply=args.apply, + ) + + +def main(argv: Sequence[str] | None = None, transport: httpx.BaseTransport | None = None) -> int: + options: Final = parse_options(argv) + with httpx.Client( + base_url=options.api_base, + headers={"Authorization": f"Bearer {options.api_key}"}, + timeout=60.0, + transport=transport, + ) as client: + try: + return run(options, client, sys.stdout) + except httpx.HTTPStatusError as error: + sys.stderr.write(f"could not read vector stores: {describe_http_error(error.response)}\n") + return 2 + except httpx.TransportError as error: + sys.stderr.write(f"could not reach {options.api_base}: {type(error).__name__}\n") + return 2 + except ValidationError: + sys.stderr.write(f"unexpected vector store response shape from {options.api_base}\n") + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/litellm/constants.py b/litellm/constants.py index 77f2961d6a1..94ff02a8c26 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1631,6 +1631,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" +VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata" ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3ad5caade1b..29fe636f5e4 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -41,6 +41,7 @@ from litellm.constants import ( SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, TYPESAFE_DANGER_ZONE_HEADER_PREFIX, + VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) from litellm.integrations.custom_guardrail import without_server_streaming_classification @@ -2160,7 +2161,7 @@ def refresh_proxy_server_request_body_snapshot( if not isinstance(proxy_server_request, dict): return _body_snapshot_exclude: Final = ( - frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj"}) + frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG}) | _TRANSPORT_ONLY_CREDENTIAL_KEYS | _CALLBACK_CREDENTIAL_KEYS ) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 22721dc26bc..1b54dddcf4f 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -9,12 +9,14 @@ from typing import ( from fastapi import APIRouter, Depends, HTTPException, Request, Response +from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body from litellm.proxy.utils import jsonify_object from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, @@ -23,6 +25,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( ) from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse +from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry router: Final = APIRouter() @@ -213,6 +216,9 @@ async def vector_store_create( ) data: Final = await read_request_body(request=request) + data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot( + await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data) + ) # Check for target_model_names parameter target_model_names: Final = data.pop("target_model_names", None) @@ -448,6 +454,9 @@ async def vector_store_update( data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line if "vector_store_id" not in data: data["vector_store_id"] = vector_store_id + data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot( + await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data) + ) data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict diff --git a/litellm/vector_stores/utils.py b/litellm/vector_stores/utils.py index a2b5a6732a8..eb61491853f 100644 --- a/litellm/vector_stores/utils.py +++ b/litellm/vector_stores/utils.py @@ -1,11 +1,92 @@ +import copy +from collections.abc import Container, Mapping from typing import Final, cast, get_type_hints +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm.constants import ROUTING_REQUEST_TAGS_METADATA_KEY, VECTOR_STORE_CALLER_METADATA_KWARG from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.types.vector_stores import ( VectorStoreCreateOptionalRequestParams, VectorStoreSearchOptionalRequestParams, ) +LITELLM_WRITTEN_METADATA_KEYS: Final = frozenset( + { + "requester_ip_address", + "user_agent", + "endpoint", + "litellm_api_version", + "deployment", + "deployment_model_name", + "api_base", + "agent_id", + "actor_agent_id", + "target_agent_id", + "billing_agent_id", + "agent_execution_mode", + "verified_human_user_id", + "litellm_caller_session_id", + } +) + + +def is_litellm_written_metadata_key(key: str) -> bool: + return key == "user_api_key" or key.startswith("user_api_key_") or key in LITELLM_WRITTEN_METADATA_KEYS + + +def _is_callers_key(key: object) -> bool: + return not isinstance(key, str) or not is_litellm_written_metadata_key(key) + + +def _sent_by_proxy_or_router(params: Mapping[str, object], metadata: Container[object]) -> bool: + return params.get("proxy_server_request") is not None or ROUTING_REQUEST_TAGS_METADATA_KEY in metadata + + +def _callers_metadata(metadata: Mapping[object, object]) -> Mapping[object, object]: + return {key: value for key, value in metadata.items() if _is_callers_key(key)} + + +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) + + +class VectorStoreWriteBody(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + metadata: object = None + + +def _json_object(raw: str | bytes) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_json(raw) + except ValidationError: + return None + + +def _decoded_metadata(metadata: object) -> object: + if not isinstance(metadata, str): + return metadata + decoded: Final = _json_object(metadata) + return metadata if decoded is None else decoded + + +def callers_metadata_snapshot(raw_json_body: bytes | None, parsed_body: VectorStoreWriteBody) -> object: + raw_body: Final = None if raw_json_body is None else _json_object(raw_json_body) + metadata: Final = copy.deepcopy(parsed_body.metadata) if raw_body is None else raw_body.get("metadata") + return _decoded_metadata(metadata) + + +def _metadata_for_provider(params: Mapping[str, object]) -> object: + if VECTOR_STORE_CALLER_METADATA_KWARG in params: + return params[VECTOR_STORE_CALLER_METADATA_KWARG] + metadata: Final = params.get("metadata") + if not isinstance(metadata, Mapping): + return metadata + request_metadata: Final = cast(Mapping[object, object], metadata) # cast-ok: request metadata is a JSON object + if not _sent_by_proxy_or_router(params, request_metadata): + return request_metadata + return _callers_metadata(request_metadata) + class VectorStoreRequestUtils: """Helper utils for constructing Vector Store search requests""" @@ -48,7 +129,11 @@ class VectorStoreRequestUtils: Returns: VectorStoreCreateOptionalRequestParams instance with only the valid parameters """ - valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() + valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() - {"metadata"} filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None} - - return cast(VectorStoreCreateOptionalRequestParams, filtered_params) + metadata: Final = _metadata_for_provider(params) + if metadata is None: + return cast(VectorStoreCreateOptionalRequestParams, filtered_params) + return cast( # cast-ok: same params as above plus the metadata the provider should store + VectorStoreCreateOptionalRequestParams, {**filtered_params, "metadata": metadata} + ) diff --git a/tests/unit/db_scripts/__init__.py b/tests/unit/db_scripts/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py b/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py new file mode 100644 index 00000000000..05b055f75a1 --- /dev/null +++ b/tests/unit/db_scripts/test_strip_vector_store_internal_metadata.py @@ -0,0 +1,198 @@ +import importlib.util +import json +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import httpx +import pytest + +_SCRIPT: Final = Path(__file__).resolve().parents[3] / "db_scripts" / "strip_vector_store_internal_metadata.py" +_spec: Final = importlib.util.spec_from_file_location("strip_vector_store_internal_metadata", _SCRIPT) +strip_script: Final = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = strip_script +_spec.loader.exec_module(strip_script) + +POLLUTED: Final = { + "id": "vs_1", + "metadata": { + "user_api_key": "hash-value-1", + "user_api_key_hash": "hash-value-1", + "user_api_key_user_email": "member@example.com", + "endpoint": "http://proxy/v1/vector_stores", + "deployment": "openai/gpt-5", + "team": "llmproxy", + "trace_id": "trace-123", + }, +} +CLEAN: Final = {"id": "vs_2", "metadata": {"team": "llmproxy"}} +NO_METADATA: Final = {"id": "vs_3", "metadata": None} + + +def _fake_openai( + stores: tuple[Mapping[str, object], ...], page_size: int = 100, post_status: int = 200 +) -> tuple[httpx.MockTransport, list[httpx.Request]]: + seen: Final[list[httpx.Request]] = [] + by_id: Final = {store["id"]: store for store in stores} + + def handle(request: httpx.Request) -> httpx.Response: + seen.append(request) + path: Final = request.url.path.removeprefix("/v1") + if request.method == "POST": + if post_status != 200: + return httpx.Response(post_status, json={"error": {"message": "upstream rejected the update"}}) + return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")]) + if path == "/vector_stores": + after: Final = request.url.params.get("after") + start: Final = 0 if after is None else [store["id"] for store in stores].index(after) + 1 + page: Final = stores[start : start + page_size] + return httpx.Response(200, json={"data": list(page), "has_more": start + page_size < len(stores)}) + return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")]) + + return httpx.MockTransport(handle), seen + + +def _run(args: list[str], transport: httpx.MockTransport) -> int: + return strip_script.main(["--api-key", "sk-test", *args], transport=transport) + + +def _posts(seen: list[httpx.Request]) -> list[tuple[str, dict[str, object]]]: + return [(request.url.path, json.loads(request.content)) for request in seen if request.method == "POST"] + + +def test_dry_run_lists_matching_key_names_without_values_and_sends_no_update(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN)) + + exit_code: Final = _run([], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 0 + assert ( + "vs_1: would delete 5 keys: deployment, endpoint, user_api_key, user_api_key_hash, user_api_key_user_email" + in out + ) + assert "scanned=2 matched=1 updated=0 failed=0" in out + assert "hash-value-1" not in out + assert "member@example.com" not in out + assert _posts(seen) == [] + assert {request.headers["authorization"] for request in seen} == {"Bearer sk-test"} + + +def test_apply_nulls_exactly_the_proxy_written_keys_and_leaves_caller_keys_alone(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA)) + + exit_code: Final = _run(["--apply"], transport) + + assert exit_code == 0 + assert _posts(seen) == [ + ( + "/v1/vector_stores/vs_1", + { + "metadata": { + "deployment": None, + "endpoint": None, + "user_api_key": None, + "user_api_key_hash": None, + "user_api_key_user_email": None, + } + }, + ) + ] + assert "scanned=3 matched=1 updated=1 failed=0" in capsys.readouterr().out + + +def test_extra_key_is_deleted_alongside_the_built_in_keys(): + transport, seen = _fake_openai((POLLUTED,)) + + _run(["--apply", "--extra-key", "trace_id"], transport) + + assert _posts(seen)[0][1]["metadata"]["trace_id"] is None + assert "team" not in _posts(seen)[0][1]["metadata"] + + +def test_listing_follows_has_more_with_the_last_store_id(capsys): + transport, seen = _fake_openai((CLEAN, POLLUTED, NO_METADATA), page_size=2) + + _run([], transport) + + assert [request.url.params.get("after") for request in seen] == [None, "vs_1"] + assert "scanned=3 matched=1" in capsys.readouterr().out + + +def test_listing_walks_more_pages_than_the_interpreter_recursion_limit(capsys): + stores: Final = tuple({"id": f"vs_{index}", "metadata": None} for index in range(sys.getrecursionlimit() + 50)) + transport, _ = _fake_openai(stores, page_size=1) + + assert _run([], transport) == 0 + assert f"scanned={len(stores)} matched=0" in capsys.readouterr().out + + +def test_non_string_metadata_values_are_scanned_without_printing_them(capsys): + transport, _ = _fake_openai(({"id": "vs_9", "metadata": {"user_api_key": "secret-value", "count": 7}},)) + + exit_code: Final = _run([], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 0 + assert "vs_9: would delete 1 keys: user_api_key" in out + assert "secret-value" not in out + + +def test_an_unexpected_response_shape_exits_non_zero_without_echoing_metadata(capsys): + def malformed(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": [{"metadata": {"user_api_key": "secret-value"}}], "has_more": False}) + + exit_code: Final = _run([], httpx.MockTransport(malformed)) + + captured: Final = capsys.readouterr() + assert exit_code == 2 + assert "unexpected vector store response shape" in captured.err + assert "secret-value" not in captured.out + captured.err + + +def test_vector_store_id_retrieves_only_the_named_stores_without_listing(capsys): + transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA)) + + _run(["--vector-store-id", "vs_2", "--vector-store-id", "vs_1"], transport) + + assert [request.url.path for request in seen] == ["/v1/vector_stores/vs_2", "/v1/vector_stores/vs_1"] + assert "scanned=2 matched=1" in capsys.readouterr().out + + +def test_a_failed_update_is_reported_and_exits_non_zero(capsys): + transport, _ = _fake_openai((POLLUTED,), post_status=400) + + exit_code: Final = _run(["--apply"], transport) + + out: Final = capsys.readouterr().out + assert exit_code == 1 + assert "vs_1: failed to delete 5 keys (HTTP 400: upstream rejected the update)" in out + assert "updated=0 failed=1" in out + + +def test_a_listing_error_exits_non_zero_without_updating_anything(): + seen: Final[list[httpx.Request]] = [] + + def reject(request: httpx.Request) -> httpx.Response: + seen.append(request) + return httpx.Response(401, json={"error": {"message": "bad key"}}) + + assert _run(["--apply"], httpx.MockTransport(reject)) == 2 + assert [request.method for request in seen] == ["GET"] + + +def test_an_unreachable_api_exits_non_zero(): + def unreachable(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused", request=request) + + assert _run([], httpx.MockTransport(unreachable)) == 2 + + +def test_an_api_key_is_required(monkeypatch): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + + with pytest.raises(SystemExit) as exit_info: + strip_script.main([]) + + assert exit_info.value.code == 2 diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 2f70c1f5188..48642a9b50d 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -6,6 +6,7 @@ import pytest from fastapi import HTTPException, Request import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) @@ -3256,3 +3257,159 @@ class TestConfigOwnedVectorStores: assert response["status"] == "success", response prisma.db.litellm_managedvectorstorestable.delete.assert_awaited_once_with(where={"vector_store_id": self.DB_ID}) assert registry.get_litellm_managed_vector_store_from_registry(self.DB_ID) is None + + +CALLERS_METADATA = {"team": "llmproxy", "endpoint": "ingest", "trace_id": "callers-trace"} + + +class _PreCallBodies(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_bodies = [] + + def log_pre_api_call(self, model, messages, kwargs): + self.logged_bodies.append(kwargs["litellm_params"]["proxy_server_request"]["body"]) + + +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize("model", [None, "gpt-5"]) +@pytest.mark.parametrize( + ("metadata_field", "expected_metadata"), + [ + ({"metadata": CALLERS_METADATA}, CALLERS_METADATA), + ({"metadata": {}}, {}), + ({"metadata": json.dumps(CALLERS_METADATA)}, CALLERS_METADATA), + ({"metadata": None}, None), + ({}, None), + ], +) +def test_proxy_vector_store_writes_send_the_provider_exactly_the_callers_metadata( + monkeypatch, operation, model, metadata_field, expected_metadata +): + import httpx + import respx + from fastapi.testclient import TestClient + + from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setattr(proxy_server, "master_key", "sk-master-test") + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}] + ), + ) + pre_call_bodies = _PreCallBodies() + monkeypatch.setattr(litellm, "callbacks", [pre_call_bodies]) + sent_bodies = [] + + def echo(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + sent_bodies.append(body) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1" + routing = {} if model is None else {"model": model} + with respx.mock(assert_all_called=False) as upstream: + upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo) + response = TestClient(app).post( + path, + json={"name": "x", **routing, **metadata_field}, + headers={ + "Authorization": "Bearer sk-master-test", + "x-litellm-tags": "header-tag", + "x-litellm-trace-id": "header-trace", + "x-litellm-session-id": "header-session", + }, + ) + + assert response.status_code == 200, response.text + assert len(sent_bodies) == 1 + assert sent_bodies[0].get("metadata") == expected_metadata + assert len(pre_call_bodies.logged_bodies) == 1 + assert "name" in pre_call_bodies.logged_bodies[0] + assert VECTOR_STORE_CALLER_METADATA_KWARG not in pre_call_bodies.logged_bodies[0] + + +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize( + "request_body", + [ + {"data": {"name": "x", "metadata": json.dumps(CALLERS_METADATA)}}, + { + "content": json.dumps({"name": "x", "metadata": CALLERS_METADATA})[:-1] + + ', "nested": ' + + "[" * 300 + + "]" * 300 + + "}", + "headers": {"Content-Type": "application/json"}, + }, + ], + ids=["form", "json-too-deep-to-reparse"], +) +def test_proxy_vector_store_writes_send_only_the_callers_metadata_when_the_raw_body_cannot_be_reparsed( + monkeypatch, operation, request_body +): + import httpx + import respx + from fastapi.testclient import TestClient + + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setattr(proxy_server, "master_key", "sk-master-test") + monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=[])) + sent_bodies = [] + + def echo(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + sent_bodies.append(body) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1" + headers = { + "Authorization": "Bearer sk-master-test", + "x-litellm-trace-id": "header-trace", + **request_body.get("headers", {}), + } + body_kwargs = {key: value for key, value in request_body.items() if key != "headers"} + with respx.mock(assert_all_called=False) as upstream: + upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo) + response = TestClient(app).post(path, headers=headers, **body_kwargs) + + assert response.status_code == 200, response.text + assert len(sent_bodies) == 1 + assert sent_bodies[0].get("metadata") == CALLERS_METADATA diff --git a/tests/unit/test_router/test_router_endpoints.py b/tests/unit/test_router/test_router_endpoints.py index 74a1bfefeee..815b32b33f7 100644 --- a/tests/unit/test_router/test_router_endpoints.py +++ b/tests/unit/test_router/test_router_endpoints.py @@ -1,10 +1,14 @@ import asyncio +import json import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import litellm import pytest from litellm import Router +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @pytest.fixture @@ -573,6 +577,39 @@ async def test_init_vector_store_api_endpoints(): mock_original_function.assert_called_once_with(vector_store_id="test-store", custom_llm_provider="openai") +@pytest.mark.asyncio +async def test_avector_store_create_sends_only_the_callers_metadata_to_the_provider(): + def echo_vector_store(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + return httpx.Response( + 200, + json={ + "id": "vs_echo", + "object": "vector_store", + "created_at": 0, + "name": body["name"], + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + router: Final = Router( + model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}] + ) + + response: Final = await router.avector_store_create( + model="gpt-5", + name="x", + metadata={"team": "llmproxy"}, + client=AsyncHTTPHandler(transport=httpx.MockTransport(echo_vector_store)), + ) + + assert response["id"] == "vs_echo" + assert response["metadata"] == {"team": "llmproxy"} + + def test_apply_default_settings(): """ Test the apply_default_settings method. diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index 1c968126c42..25a378270b1 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -8,16 +8,19 @@ model_dump() it (the #19550 serialization trap). """ import json +from collections.abc import Mapping +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest import litellm.vector_stores.main as vector_stores_main from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) -from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.vector_stores.main import search +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.vector_stores.main import acreate, aupdate, search MOCK_SEARCH_RESPONSE = { "object": "vector_store.search_results.page", @@ -114,3 +117,86 @@ def test_search_forwards_top_level_user_context_to_bedrock_retrieve(): posted = json.loads(client.post.call_args.kwargs["data"]) assert posted["userContext"] == {"userId": "alice@example.com"} assert posted["retrievalQuery"] == {"text": "q"} + + +PROXY_STAMPED_METADATA: Final = { + "team": "llmproxy", + "user_api_key": "hashed-key", + "user_api_key_hash": "hashed-key", + "user_api_key_team_id": "team-1", + "requester_ip_address": "127.0.0.1", + "user_agent": "curl/8.7.1", + "endpoint": "http://localhost:4000/v1/vector_stores", + "litellm_api_version": "1.0.0", + "deployment": "openai/gpt-5", + "deployment_model_name": "gpt-5", + "api_base": "https://api.openai.com/v1", + "agent_id": "agent-1", + "actor_agent_id": "agent-2", + "target_agent_id": "agent-3", + "billing_agent_id": "agent-4", + "agent_execution_mode": "autonomous", + "verified_human_user_id": "user-1", + "litellm_caller_session_id": "session-1", + "user_api_keyboard": "qwerty", + 1: "a", +} + + +def _copy(metadata: Mapping[object, object] | None) -> dict[object, object] | None: + return None if metadata is None else dict(metadata) + + +async def _metadata_sent_to_provider( + operation: str, metadata: Mapping[object, object] | None, **kwargs: object +) -> list[object]: + sent_metadata: Final[list[object]] = [] + + def echo_vector_store(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + sent_metadata.append(body.get("metadata")) + return httpx.Response( + 200, + json={ + "id": "vs_1", + "object": "vector_store", + "created_at": 0, + "name": "x", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", + "metadata": body.get("metadata"), + }, + ) + + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(echo_vector_store)) + if operation == "create": + await acreate(name="x", metadata=_copy(metadata), api_key="sk-test", client=client, **kwargs) + else: + await aupdate(vector_store_id="vs_1", metadata=_copy(metadata), api_key="sk-test", client=client, **kwargs) + return sent_metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["create", "update"]) +async def test_proxy_vector_store_writes_send_only_the_callers_metadata_to_the_provider(operation: str): + sent_metadata: Final = await _metadata_sent_to_provider( + operation, + PROXY_STAMPED_METADATA, + proxy_server_request={"url": "http://localhost:4000/v1/vector_stores", "method": "POST", "body": {}}, + ) + + assert sent_metadata == [{"team": "llmproxy", "user_api_keyboard": "qwerty", "1": "a"}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["create", "update"]) +@pytest.mark.parametrize( + "callers_metadata", [{"team": "llmproxy", "endpoint": "ingest", "deployment": "prod", "agent_id": "bot-7"}, None] +) +async def test_direct_sdk_vector_store_writes_send_the_callers_metadata_unchanged( + operation: str, callers_metadata: Mapping[object, object] | None +): + sent_metadata: Final = await _metadata_sent_to_provider(operation, callers_metadata) + + assert sent_metadata == [callers_metadata] From bbd446b476526864e6f0abeb3a6cb6728e9772ed Mon Sep 17 00:00:00 2001 From: alberto Date: Sat, 17 Oct 2026 20:52:10 +0100 Subject: [PATCH 35/64] fix(proxy): report requested model in streamed Responses response.model (#46662) * chore: Add failing test case Signed-off-by: Alberto Perdomo * fix: Restamp nested response.model in streamed Responses events Signed-off-by: Alberto Perdomo * chore: Add tests Signed-off-by: Alberto Perdomo * chore: Run format and lint, and fix Signed-off-by: Alberto Perdomo * fix: Type nested Responses restamp to satisfy basedpyright gate Signed-off-by: Alberto Perdomo * fix(proxy): restamp the nested model on streamed response.queued events Background streams send response.queued as a GenericEvent, so its nested response.model still carried the backend name. Also drop the new comments * fix(proxy): restamp a nested response kept as a dict on streamed lifecycle events * test(proxy): drive typed events through the nested restamp no-op test * fix(proxy): restamp the nested model on guardrail block frames in streamed responses * test(proxy): assert a cache-polled background response reports the requested model * fix(proxy): restamp nested model on dict events and on every SSE data line --------- Signed-off-by: Alberto Perdomo Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 105 ++++- ...est_background_response_billing_proxies.py | 4 +- .../proxy/test_response_model_sanitization.py | 415 +++++++++++++++++- 3 files changed, 500 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1e16eb626fe..3ed3c57a96c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -954,7 +954,14 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionSystemMessage, ChatCompletionToolParam, + GenericEvent, HttpxBinaryResponseContent, + ResponseCompletedEvent, + ResponseCreatedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponseInProgressEvent, + ResponsesAPIResponse, ) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -9922,7 +9929,7 @@ def _restamp_streaming_chunk_model( # Note: This warning is intentionally verbose. A mismatch is a useful signal that an # internal provider/deployment identifier is leaking into the public API, and helps # maintainers/operators catch regressions while preserving OpenAI-compatible output. - if not target_model or not isinstance(chunk, (BaseModel, dict)): + if not target_model or not isinstance(chunk, (BaseModel, dict, bytes)): return chunk, model_mismatch_logged # For Azure Model Router, preserve the actual model used in each chunk @@ -9934,9 +9941,21 @@ def _restamp_streaming_chunk_model( if not fallback_was_attempted and request_data.get("fastest_response", False): return chunk, model_mismatch_logged - downstream_model: Final = chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None) + if isinstance(chunk, bytes): + if request_data.get("_litellm_raw_sse_stream"): + return chunk, model_mismatch_logged + return _restamp_nested_response_model_in_sse_frame(chunk, target_model), model_mismatch_logged + + typed_chunk: Final[BaseModel | Mapping[str, object]] = chunk + restamped_chunk: Final = _restamp_nested_response_model(typed_chunk, target_model) + + downstream_model: Final = ( + restamped_chunk.get("model") + if isinstance(restamped_chunk, Mapping) + else getattr(restamped_chunk, "model", None) + ) if downstream_model == target_model: - return chunk, model_mismatch_logged + return restamped_chunk, model_mismatch_logged if not model_mismatch_logged and downstream_model != target_model: verbose_proxy_logger.debug( @@ -9954,22 +9973,92 @@ def _restamp_streaming_chunk_model( logging_obj: Final = request_data.get("litellm_logging_obj") if isinstance(logging_obj, LiteLLMLoggingObj): logging_obj.client_facing_stream_model = target_model - if isinstance(chunk, dict): - return {**chunk, "model": target_model}, model_mismatch_logged + if isinstance(restamped_chunk, Mapping): + return {**restamped_chunk, "model": target_model}, model_mismatch_logged try: - return chunk.model_copy(update={"model": target_model}), model_mismatch_logged + return restamped_chunk.model_copy(update={"model": target_model}), model_mismatch_logged except Exception as e: verbose_proxy_logger.error( "litellm_call_id=%s: failed to override chunk.model=%r on chunk_type=%s. error=%s", request_data.get("litellm_call_id"), target_model, - type(chunk), + type(restamped_chunk), str(e), exc_info=True, ) - return chunk, model_mismatch_logged + return restamped_chunk, model_mismatch_logged + + +def _nested_response_of(chunk: BaseModel) -> object: + if isinstance( + chunk, + ( + ResponseCreatedEvent, + ResponseInProgressEvent, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ), + ): + return chunk.response + if isinstance(chunk, GenericEvent): + extra_fields: Final[Mapping[str, object]] = chunk.model_extra or {} + return extra_fields.get("response") + return None + + +def _restamp_nested_response_model( + chunk: BaseModel | Mapping[str, object], target_model: str +) -> BaseModel | Mapping[str, object]: + if not isinstance(chunk, BaseModel): + restamped_event: Final = _event_with_restamped_nested_model(chunk, target_model) + return chunk if restamped_event is None else restamped_event + nested_response: Final = _nested_response_of(chunk) + if isinstance(nested_response, ResponsesAPIResponse): + if nested_response.model in (None, target_model): + return chunk + return chunk.model_copy(update={"response": nested_response.model_copy(update={"model": target_model})}) + untyped_response: Final = as_str_mapping(nested_response) + if untyped_response is None or untyped_response.get("model") in (None, target_model): + return chunk + return chunk.model_copy(update={"response": {**untyped_response, "model": target_model}}) + + +def _event_with_restamped_nested_model(event: Mapping[str, object], target_model: str) -> Mapping[str, object] | None: + nested_response: Final = as_str_mapping(event.get("response")) + if nested_response is None or nested_response.get("model") in (None, target_model): + return None + return {**event, "response": {**nested_response, "model": target_model}} + + +_SSE_DATA_FIELD: Final = b"data:" +_SSE_NESTED_RESPONSE_MARKER: Final = b'"response"' +_SSE_EVENT_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def _restamp_nested_response_model_in_sse_frame(frame: bytes, target_model: str) -> bytes: + if _SSE_NESTED_RESPONSE_MARKER not in frame: + return frame + return b"".join( + _restamp_nested_response_model_in_sse_line(line, target_model) for line in frame.splitlines(keepends=True) + ) + + +def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) -> bytes: + if not line.startswith(_SSE_DATA_FIELD): + return line + try: + event: Final = _SSE_EVENT_OBJECT.validate_json(line.removeprefix(_SSE_DATA_FIELD).strip()) + except ValidationError: + return line + restamped_event: Final = _event_with_restamped_nested_model(event, target_model) + if restamped_event is None: + return line + serialized_event: Final = json.dumps(restamped_event, separators=(",", ":"), ensure_ascii=False) + line_ending: Final = line[len(line.rstrip(b"\r\n")) :] + return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending def _fast_serialize_simple_model_response_stream( diff --git a/tests/integration/spend/test_background_response_billing_proxies.py b/tests/integration/spend/test_background_response_billing_proxies.py index e83f2629f33..4b4185f5da2 100644 --- a/tests/integration/spend/test_background_response_billing_proxies.py +++ b/tests/integration/spend/test_background_response_billing_proxies.py @@ -549,7 +549,7 @@ def _cache_polled_create(proxy: Gateway, key: str, model: str, create: Mapping[s @pytest.mark.timeout(OWNED_PROXY_SECONDS + 120) @pytest.mark.parametrize("create", CACHE_POLLED_CREATES.keys()) -def test_a_background_response_the_proxy_polls_through_its_cache_bills_its_creator_once( +def test_a_background_response_polled_through_the_cache_reports_its_alias_and_bills_its_creator_once( gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, create: str ) -> None: body, provider_sees_background = CACHE_POLLED_CREATES[create] @@ -564,6 +564,8 @@ def test_a_background_response_the_proxy_polls_through_its_cache_bills_its_creat seconds=SPEND_RETRY_SECONDS, ) assert finished.text == "pong" + polled: Final = object_value(proxy.request("GET", f"/v1/responses/{polling_id}", key=key).json()) + assert polled["model"] == handle.scenario_id assert_charged_exactly(lambda: key_spend(key), PRICE) assert [read_response(proxy, key, polling_id).text for _ in range(READS)] == ["pong"] * READS assert_charged_exactly(lambda: key_spend(key), PRICE) diff --git a/tests/unit/proxy/test_response_model_sanitization.py b/tests/unit/proxy/test_response_model_sanitization.py index ba66d27609b..f52736453c1 100644 --- a/tests/unit/proxy/test_response_model_sanitization.py +++ b/tests/unit/proxy/test_response_model_sanitization.py @@ -68,6 +68,19 @@ def _decode_sse_chunk(chunk) -> str: return chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk +def _make_response_created_event(model: str | None): + from litellm.types.llms.openai import ( + ResponseCreatedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + return ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, + response=ResponsesAPIResponse(id="resp_test", created_at=0, output=[], model=model), + ) + + def test_restamp_streaming_chunk_skips_matching_model(): from litellm.proxy.proxy_server import _restamp_streaming_chunk_model @@ -117,19 +130,13 @@ def test_fast_serialize_returns_none_when_model_field_is_missing(): # produce a serialized result via the slow-path fallback, and it must # not contain ``"model": null``. serialized = _serialize_streaming_chunk(chunk) - payload_str = ( - serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized - ) + payload_str = serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized assert '"model": null' not in payload_str assert '"model":null' not in payload_str - assert json.loads(payload_str) == json.loads( - chunk.model_dump_json(exclude_none=True, exclude_unset=True) - ) + assert json.loads(payload_str) == json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True)) -def test_proxy_chat_completion_does_not_return_provider_prefixed_model( - tmp_path, monkeypatch -): +def test_proxy_chat_completion_does_not_return_provider_prefixed_model(tmp_path, monkeypatch): """ Regression test: @@ -159,15 +166,11 @@ def test_proxy_chat_completion_does_not_return_provider_prefixed_model( monkeypatch.setattr( proxy_server.llm_router, # type: ignore[arg-type] "acompletion", - AsyncMock( - return_value=_make_minimal_chat_completion_response(model=internal_model) - ), + AsyncMock(return_value=_make_minimal_chat_completion_response(model=internal_model)), ) # Also no-op proxy logging hooks to keep this test focused and deterministic. - monkeypatch.setattr( - proxy_server.proxy_logging_obj, "during_call_hook", AsyncMock(return_value=None) - ) + monkeypatch.setattr(proxy_server.proxy_logging_obj, "during_call_hook", AsyncMock(return_value=None)) monkeypatch.setattr( proxy_server.proxy_logging_obj, "update_request_status", @@ -482,3 +485,385 @@ async def test_proxy_streaming_fastest_response_preserves_winning_model(monkeypa payload = json.loads(first[len("data: ") :].strip()) assert payload["model"] == winning_model assert payload["model"] != comma_separated_models + + +def test_restamp_streaming_chunk_rewrites_nested_response_model(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_created_event("hosted_vllm/some-model") + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response.model == "client-model" + assert chunk.response.model == "hosted_vllm/some-model" + + +def _make_response_completed_event(model: str): + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse(id="resp_test", created_at=0, output=[], model=model), + ) + + +def test_restamp_streaming_chunk_rewrites_nested_model_on_completed_event(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_completed_event("hosted_vllm/some-model") + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response.model == "client-model" + assert chunk.response.model == "hosted_vllm/some-model", "the priced original must not change" + + +@pytest.mark.parametrize("event_type", ["response.in_progress", "response.failed", "response.incomplete"]) +def test_restamp_streaming_chunk_rewrites_nested_model_on_other_lifecycle_events(event_type): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + payload = { + "type": event_type, + "sequence_number": 1, + "response": {"id": "resp_test", "object": "response", "created_at": 0, "output": [], "model": "gpt-backend"}, + } + chunk = OpenAIResponsesAPIConfig.get_event_model_class(event_type).model_validate(payload) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response.model == "client-model" + assert chunk.response.model == "gpt-backend", "the priced original must not change" + + +def test_restamp_streaming_chunk_leaves_nested_none_model_untouched(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_created_event(None) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response.model is None + + +@pytest.mark.parametrize("event_type", ["response.output_text.delta", "response.custom_event"]) +def test_restamp_streaming_chunk_does_not_add_nested_response_to_other_events(event_type): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = OpenAIResponsesAPIConfig.get_event_model_class(event_type).model_validate( + { + "type": event_type, + "sequence_number": 3, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "hi", + } + ) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert "response" not in result.model_dump() + assert result.delta == "hi" + + +def test_restamp_streaming_chunk_keeps_nested_model_when_raw_name_requested(): + from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_created_event("hosted_vllm/some-model") + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: True}}, + model_mismatch_logged=False, + ) + + assert result.response.model == "hosted_vllm/some-model" + + +def test_restamp_streaming_chunk_uses_fallback_model_for_nested_response(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = _make_response_created_event("hosted_vllm/some-model") + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + fallback_was_attempted=True, + fallback_model_from_metadata="fallback-public", + ) + + assert result.response.model == "fallback-public" + + +@pytest.mark.asyncio +async def test_proxy_responses_stream_reports_requested_model_in_nested_response(monkeypatch): + """ + End-to-end regression for /v1/responses streaming: + + The nested `response.model` on lifecycle events reaching the wire must report the + client-requested model, not the provider-prefixed deployment name. + """ + client_model = "vllm-model" + internal_model = f"hosted_vllm/{client_model}" + + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + async def _iterator_hook( + user_api_key_dict: UserAPIKeyAuth, + response: AsyncGenerator, + request_data: dict, + ): + yield _make_response_created_event(model=internal_model) + + monkeypatch.setattr( + proxy_server.proxy_logging_obj, + "async_post_call_streaming_iterator_hook", + _iterator_hook, + ) + monkeypatch.setattr( + proxy_server.proxy_logging_obj, + "async_post_call_streaming_hook", + AsyncMock(side_effect=lambda **kwargs: kwargs["response"]), + ) + for name in ("has_streaming_callbacks", "needs_iterator_wrap", "needs_per_chunk_streaming_hook"): + monkeypatch.setattr(proxy_server.proxy_logging_obj, name, MagicMock(return_value=True)) + + gen = proxy_server.async_data_generator( + response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key=MASTER_KEY), + request_data={"model": client_model}, + responses_stream_errors=True, + ) + + payloads = [] + async for item in gen: + text = _decode_sse_chunk(item) + for line in text.splitlines(): + if line.startswith("data: ") and line.strip() != "data: [DONE]": + payloads.append(json.loads(line[len("data: ") :].strip())) + + created = next(p for p in payloads if p.get("type") == "response.created") + assert created["response"]["model"] == client_model + assert not created["response"]["model"].startswith("hosted_vllm/") + + +def test_restamp_streaming_chunk_rewrites_nested_model_on_queued_event(): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + queued_payload = { + "type": "response.queued", + "sequence_number": 0, + "response": {"id": "resp_test", "object": "response", "status": "queued", "model": "gpt-backend-2026-01-01"}, + } + chunk = OpenAIResponsesAPIConfig.get_event_model_class("response.queued").model_validate(queued_payload) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.model_dump()["response"] == {**queued_payload["response"], "model": "client-model"} + assert chunk.model_dump()["response"]["model"] == "gpt-backend-2026-01-01", "the priced original must not change" + + +def test_restamp_streaming_chunk_rewrites_nested_model_when_lifecycle_event_kept_raw_response(): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + partial_response = {"id": "resp_test", "object": "response", "status": "in_progress", "model": "gpt-backend-2026-01-01"} + chunk = OpenAIResponsesAPIConfig().transform_streaming_response( + model="gpt-backend", + parsed_chunk={"type": "response.created", "sequence_number": 0, "response": partial_response}, + logging_obj=MagicMock(), + ) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result.response == {**partial_response, "model": "client-model"} + assert chunk.response["model"] == "gpt-backend-2026-01-01", "the priced original must not change" + + +def test_restamp_streaming_chunk_leaves_generic_event_without_nested_response_untouched(): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + shell_payload = {"type": "response.shell_call.in_progress", "sequence_number": 3, "item_id": "sh_1"} + chunk = OpenAIResponsesAPIConfig.get_event_model_class(shell_payload["type"]).model_validate(shell_payload) + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert "response" not in result.model_dump() + + +@pytest.mark.parametrize("stream_started", [False, True]) +def test_restamp_streaming_chunk_rewrites_nested_model_in_guardrail_block_frames(stream_started): + from litellm.exceptions import ModifyResponseException + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + sent_before_block = ( + [{"type": "response.created", "response": {"id": "resp_live", "model": "gpt-backend-2026-01-01"}}] + if stream_started + else [] + ) + block = ModifyResponseException( + message="Blocked by policy.", + model="gpt-backend-2026-01-01", + request_data={}, + original_response=sent_before_block, + ) + frames = OpenAIResponsesHandler().build_block_sse_chunks( + block, stream_started=stream_started, responses_so_far=sent_before_block + ) + client_frames = [ + _restamp_streaming_chunk_model( + chunk=frame, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + )[0] + for frame in frames + ] + events = [json.loads(frame.decode().removeprefix("data: ")) for frame in client_frames] + + assert all(frame.endswith(b"\n\n") for frame in client_frames) + assert events[-1]["type"] == "response.completed" + assert events[-1]["response"]["output"][0]["content"][0]["text"] == "Blocked by policy." + assert {event["response"]["model"] for event in events if "response" in event} == {"client-model"} + assert [frame for frame, event in zip(client_frames, events) if "response" not in event] == [ + frame for frame in frames if b'"response"' not in frame + ] + + +def test_restamp_streaming_chunk_rewrites_nested_model_on_dict_lifecycle_events(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + chunk = {"type": "response.completed", "response": {"id": "resp_test", "model": "gpt-backend-2026-01-01"}} + result, _ = _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + assert result["response"] == {"id": "resp_test", "model": "client-model"} + assert chunk["response"]["model"] == "gpt-backend-2026-01-01", "the priced original must not change" + + +_BACKEND_COMPLETED_EVENT = b'{"type":"response.completed","response":{"id":"resp_1","model":"gpt-backend-2026-01-01"}}' + + +@pytest.mark.parametrize( + "frame", + [ + pytest.param(b"event: response.completed\ndata: " + _BACKEND_COMPLETED_EVENT + b"\n\n", id="event_line_first"), + pytest.param(b"data:" + _BACKEND_COMPLETED_EVENT + b"\n\n", id="no_space_after_data"), + pytest.param( + b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\ndata: " + _BACKEND_COMPLETED_EVENT + b"\r\n\r\n", + id="coalesced_frames", + ), + ], +) +def test_restamp_streaming_chunk_rewrites_nested_model_on_every_sse_data_line(frame): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + result, _ = _restamp_streaming_chunk_model( + chunk=frame, + requested_model_from_client="client-model", + request_data={"litellm_call_id": "test-call-id"}, + model_mismatch_logged=False, + ) + + def is_data_line(line: bytes) -> bool: + return line.startswith(b"data:") + + result_lines = result.splitlines(keepends=True) + frame_lines = frame.splitlines(keepends=True) + data_events = [json.loads(line.removeprefix(b"data:")) for line in result_lines if is_data_line(line)] + assert len(data_events) == len([line for line in frame_lines if is_data_line(line)]) + assert {event["response"]["model"] for event in data_events} == {"client-model"} + assert [line for line in result_lines if not is_data_line(line)] == [ + line for line in frame_lines if not is_data_line(line) + ] + assert [line[len(line.rstrip(b"\r\n")) :] for line in result_lines] == [ + line[len(line.rstrip(b"\r\n")) :] for line in frame_lines + ] + + +_SKIPPED_RESTAMPS = { + "azure_model_router": ("azure_ai/model_router", {}), + "fastest_response": ("openai/gpt-4o,gemini/gemini-2.5-flash", {"fastest_response": True}), +} + + +@pytest.mark.parametrize("skip", _SKIPPED_RESTAMPS.keys()) +def test_restamp_streaming_chunk_keeps_nested_model_where_the_top_level_model_is_kept(skip): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + requested_model, request_data = _SKIPPED_RESTAMPS[skip] + frame = b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\n" + results = [ + _restamp_streaming_chunk_model( + chunk=chunk, + requested_model_from_client=requested_model, + request_data=request_data, + model_mismatch_logged=False, + )[0] + for chunk in (_make_response_created_event("gpt-backend-2026-01-01"), frame) + ] + + assert results[0].response.model == "gpt-backend-2026-01-01" + assert results[1] == frame + + +def test_restamp_streaming_chunk_leaves_raw_sse_passthrough_frames_untouched(): + from litellm.proxy.proxy_server import _restamp_streaming_chunk_model + + frame = b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\n" + result, _ = _restamp_streaming_chunk_model( + chunk=frame, + requested_model_from_client="client-model", + request_data={"_litellm_raw_sse_stream": True}, + model_mismatch_logged=False, + ) + + assert result == frame From 0f007c05fcd2fceef3b6063d31a578616a1b2af7 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 13:24:05 -0700 Subject: [PATCH 36/64] fix(logging): run legacy s3 callback for streamed /v1/messages calls (#47156) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- .../test_litellm_logging.py | 37 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 457a17ab8fe..99ebf479b20 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3059,7 +3059,7 @@ class Logging(LiteLLMLoggingBaseClass): complete_streaming_response: ( ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None ) = None - if "complete_streaming_response" in self.model_call_details: + if "complete_streaming_response" in self.model_call_details and not for_async_request: return # break out of this. complete_streaming_response = self._get_assembled_streaming_response( result=result, diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 3d164b59bfe..c7450c852c0 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -12596,3 +12596,40 @@ def test_a_websocket_turn_row_is_keyed_by_the_response_id_or_else_its_call_id( ) assert payload is not None assert payload["id"] == expected_row_id + + +@pytest.mark.parametrize("for_async_request, expected_s3_writes", [(True, 1), (False, 0)]) +def test_legacy_s3_logs_a_streamed_anthropic_messages_call_handed_off_from_the_async_side( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch, for_async_request: bool, expected_s3_writes: int +) -> None: + """LIT-9834: the /v1/messages stream handler stores ``complete_streaming_response`` before any logger runs. + + The one hand-off from the async side (``for_async_request=True``) must still reach the legacy ``s3`` writer, + while the chunk-by-chunk sync path (``for_async_request=False``) keeps treating a stored response as logged.""" + complete_response: Final = ModelResponse( + id="msg_fake_stream_001", + model="claude-sonnet-4-5", + choices=[ + { + "message": {"role": "assistant", "content": "Hello from the fake upstream"}, + "finish_reason": "stop", + "index": 0, + } + ], + usage={"prompt_tokens": 8, "completion_tokens": 6, "total_tokens": 14}, + ) + logging_obj.stream = True + logging_obj.call_type = "anthropic_messages" + logging_obj.model_call_details["litellm_params"] = {"aanthropic_messages": True} + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + logging_obj.model_call_details["complete_streaming_response"] = complete_response + s3_writer: Final = MagicMock() + monkeypatch.setattr(litellm_logging, "s3Logger", s3_writer) + monkeypatch.setattr(litellm, "success_callback", ["s3"]) + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + + logging_obj.success_handler(complete_response, start_time=now, end_time=now, for_async_request=for_async_request) + + assert s3_writer.log_event.call_count == expected_s3_writes + if expected_s3_writes: + assert s3_writer.log_event.call_args.kwargs["response_obj"] is complete_response From 90fccd4d8758d120bd94b1e8527e0657ddaa6a2c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 13:25:29 -0700 Subject: [PATCH 37/64] test(e2e): get_text_external accepts JSON next to text so a pod-log error reads as the Kubernetes Status, not a 406 (#47150) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/e2e_http.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 90a9e9367eb..2aea32909ce 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -507,11 +507,12 @@ def get_text_external( ca_bundle: str | None = None, ) -> str | NetworkError: """GET an absolute URL outside the proxy whose body is plain text (e.g. a pod log), not JSON. - A non-2xx status is reported as a NetworkError so the caller can fail loudly.""" + A non-2xx status is reported as a NetworkError so the caller can fail loudly. JSON is accepted too so an API + that answers an error as a JSON document (the Kubernetes Status on a pod log) returns that error instead of a 406.""" try: resp = requests.get( url, - headers={"Accept": "text/plain", **(_headers(headers) if headers is not None else {})}, + headers={"Accept": "text/plain, application/json", **(_headers(headers) if headers is not None else {})}, timeout=timeout, verify=ca_bundle if ca_bundle is not None else True, ) From 4273996985b31a1fdb575631532172c4ba0a759d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 13:26:09 -0700 Subject: [PATCH 38/64] fix(spend_logs): persist user email in spend log metadata so it shows in UI request logs (#41150) * fix(spend_logs): persist user email in spend log metadata so it shows in UI request logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_logs): update payload expectations for user email metadata Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(gcs_pubsub): add user_api_key_user_email to the spend log fixture Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_logs): drop untyped metadata key args from the debug log to keep the type gate flat Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_logs): pass metadata keys to the debug log as a lazy %-style argument Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(spend_logs): store user email in spend logs only behind store_user_email_in_spend_logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * revert(spend_logs): store user email in spend logs by default, drop store_user_email_in_spend_logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_logs): store an empty user email as null, narrow the expected payload test diff Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: jesus Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + .../spend_tracking/spend_tracking_utils.py | 12 ++++-- .../gcs_pubsub/spend_logs_payload.json | 2 +- .../test_spend_management_endpoints.py | 2 +- .../test_spend_tracking_utils.py | 38 +++++++++++++++++++ .../test_session_handler_with_cold_storage.py | 12 +++--- .../request/RequestLogsTableColumns.test.tsx | 35 +++++++++++++++++ .../logs/request/RequestLogsTableColumns.tsx | 5 ++- 8 files changed, 95 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2c98cfeb33d..dc7383569dc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3667,6 +3667,7 @@ class SpendLogsMetadata(TypedDict): user_api_key_project_alias: str | None user_api_key_org_id: str | None user_api_key_user_id: str | None + user_api_key_user_email: ReadOnly[str | None] user_api_key_team_alias: str | None spend_logs_metadata: dict | None # special param to log k,v pairs to spendlogs for a call requester_ip_address: str | None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index c7ca08fb513..d1303b012c8 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -179,10 +179,16 @@ _STAMPED_METADATA_KEYS: Final = frozenset( "autorouter_baseline_observation", "used_client_oauth_token", "litellm_roi_estimator", + "user_api_key_user_email", ) ) +def _stored_user_email(metadata: Mapping[str, object]) -> str | None: + email: Final = metadata.get("user_api_key_user_email") + return email if isinstance(email, str) and email.strip() else None + + def _get_spend_logs_metadata( metadata: dict | None, applied_guardrails: list[str] | None = None, @@ -214,6 +220,7 @@ def _get_spend_logs_metadata( user_api_key_project_alias=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, @@ -251,9 +258,7 @@ def _get_spend_logs_metadata( azure_spillover=azure_spillover, used_client_oauth_token=used_client_oauth_token, ) - verbose_proxy_logger.debug( - "getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys())) - ) + verbose_proxy_logger.debug("getting payload for SpendLogs, available keys in metadata: %s", [*metadata]) # Filter the metadata dictionary to include only the specified keys clean_metadata: Final = SpendLogsMetadata( @@ -267,6 +272,7 @@ def _get_spend_logs_metadata( azure_spillover=azure_spillover, used_client_oauth_token=used_client_oauth_token, litellm_roi_estimator=metadata.get("litellm_roi_estimator") is True, + user_api_key_user_email=_stored_user_email(metadata), ) _raw_key: Final = clean_metadata.get("user_api_key") _trusted_hash: Final = metadata.get("user_api_key_hash") diff --git a/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json index f7ebf357d83..c6cf8ee7ee4 100644 --- a/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json +++ b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index f6b36c184ec..47cfa2df117 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3725,7 +3725,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "used_client_oauth_token": null, "litellm_roi_estimator": false, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"user_api_key_user_email": null, "actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "used_client_oauth_token": null, "litellm_roi_estimator": false, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 65dcce97b90..262357fe1bd 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -526,6 +526,7 @@ def _make_standard_logging_payload_with_usage_object(usage_object: dict) -> Stan user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, @@ -1957,6 +1958,7 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, @@ -2090,6 +2092,7 @@ def test_get_logging_payload_handles_missing_overhead_gracefully(): user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, @@ -2677,6 +2680,7 @@ def test_get_logging_payload_includes_retry_info_in_spend_logs_metadata(): user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, @@ -2772,6 +2776,7 @@ def test_get_logging_payload_handles_missing_retry_info_gracefully(): user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, @@ -5243,6 +5248,37 @@ def test_get_spend_logs_metadata_sibling_fields_preserved(): assert meta["user_api_key_team_id"] == "team-123" +@pytest.mark.parametrize( + ("user_email", "stored_email"), + [ + ("sso-user@example.com", "sso-user@example.com"), + ("", None), + (" ", None), + (None, None), + ], +) +def test_get_logging_payload_stores_user_email_in_spend_log_metadata(user_email, stored_email): + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_params": { + "metadata": { + "user_api_key_user_email": user_email, + "user_api_key_user_id": "okta|abc123", + } + }, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-user-email", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + metadata = json.loads(payload["metadata"]) + + assert metadata["user_api_key_user_email"] == stored_email + assert metadata["user_api_key_user_id"] == "okta|abc123" + + def test_redact_logged_api_key_partial_sha256_is_hashed(): partial_hex = "a" * 63 result = _redact_logged_api_key(partial_hex) @@ -5342,6 +5378,7 @@ def test_get_logging_payload_includes_fallback_info_in_spend_logs_metadata(): user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, @@ -5437,6 +5474,7 @@ def test_get_logging_payload_handles_missing_fallback_info_gracefully(): user_api_key_team_id=None, user_api_key_org_id=None, user_api_key_user_id=None, + user_api_key_user_email=None, user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py index f2fbcda59a4..1565820cbbd 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py @@ -10,20 +10,18 @@ Tests for the changes to integrate cold storage handling across different compon """ import json -from datetime import datetime, timezone -from typing import Optional -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest from litellm.integrations.s3_v2 import S3Logger -from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload +from litellm.proxy._types import SpendLogsMetadata from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler from litellm.proxy.spend_tracking.spend_tracking_utils import _get_spend_logs_metadata from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) -from litellm.types.utils import StandardLoggingMetadata, StandardLoggingPayload +from litellm.types.utils import StandardLoggingMetadata class TestColdStorageObjectKeyIntegration: @@ -57,7 +55,9 @@ class TestColdStorageObjectKeyIntegration: """ # Create a SpendLogsMetadata instance with cold_storage_object_key metadata = SpendLogsMetadata( - user_api_key="test_key", cold_storage_object_key="test/path/to/object.json" + user_api_key="test_key", + user_api_key_user_email=None, + cold_storage_object_key="test/path/to/object.json", ) # Verify the field can be set and accessed diff --git a/ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.test.tsx index 6d4a6e789b4..2f2246a3878 100644 --- a/ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.test.tsx @@ -147,6 +147,41 @@ describe("Internal User column", () => { expect(await screen.findByText("alice@example.com (106514937785257944828)")).toBeInTheDocument(); }); + it("shows the email stored in the log metadata when the user list has no entry for that user", async () => { + const user = userEvent.setup(); + renderRows( + [ + logEntry({ + request_id: "req-sso-user", + user: "okta|abc123", + metadata: { user_api_key_user_email: "sso-user@example.com" }, + }), + ], + deps, + ); + + const emailCell = screen.getByText("sso-user@example.com"); + expect(screen.queryByText("okta|abc123")).not.toBeInTheDocument(); + await user.hover(emailCell); + expect(await screen.findByText("sso-user@example.com (okta|abc123)")).toBeInTheDocument(); + }); + + it("prefers the user list email over the one stored in the log metadata", () => { + renderRows( + [ + logEntry({ + request_id: "req-renamed-user", + user: "106514937785257944828", + metadata: { user_api_key_user_email: "old@example.com" }, + }), + ], + deps, + ); + + expect(screen.getByText("alice@example.com")).toBeInTheDocument(); + expect(screen.queryByText("old@example.com")).not.toBeInTheDocument(); + }); + it("falls back to the raw id when no email is known for the user", () => { renderRows([logEntry({ request_id: "req-unknown-user", user: "unknown-user-id" })], deps); diff --git a/ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.tsx index 5d5918e4657..b3d44361d0a 100644 --- a/ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.tsx @@ -365,7 +365,10 @@ export const getRequestLogsTableColumns = ({ cell: ({ row }) => ( ), }, From 603ba6fe00be0f6e7f7f33e29d9191a34936f96d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 13:51:31 -0700 Subject: [PATCH 39/64] test(integration): run the bedrock mantle prompt cache and output limit wire tests on their own proxy (#47152) * test(integration): run the bedrock mantle prompt cache and output limit wire tests on their own proxy These two files register deployments with their own model settings on the gateway. They now boot a module-scoped proxy on a scratch database, so their setup cannot affect other tests that share the worker's proxy * test(integration): seed the isolated gateway with the same router settings as the worker gateway --- tests/integration/_support/process.py | 20 +++++++++++++++++-- ...mantle_gpt_prompt_cache_breakpoint_wire.py | 11 ++++++++++ .../test_bedrock_mantle_output_limit_wire.py | 13 ++++++++++-- 3 files changed, 40 insertions(+), 4 deletions(-) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 00b9dca5184..5cbe1309075 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -17,7 +17,7 @@ from ipaddress import IPv4Address from itertools import accumulate from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import TYPE_CHECKING, Final from urllib.parse import urlsplit import httpx @@ -30,9 +30,13 @@ from cryptography.x509.oid import NameOID from pydantic import TypeAdapter from tests.integration._support.budget import import_seconds, readiness_for, readiness_seconds -from tests.integration._support.client import GATEWAY_LIMITS, Gateway +from tests.integration._support.client import GATEWAY_LIMITS, Gateway, gateway_from_environment +from tests.integration._support.database import scratch_database from tests.integration._support.ports import allocate_port +if TYPE_CHECKING: + import pytest + DB_PUSH: Final = ("--use_prisma_db_push",) MIGRATE_DEPLOY: Final = () LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",) @@ -379,6 +383,18 @@ def owned_proxy( yield owned.gateway +@contextmanager +def isolated_gateway(directory: Path, environment: "pytest.MonkeyPatch") -> Iterator[Gateway]: + with gateway_from_environment() as shared, scratch_database() as database: + environment.setenv("DATABASE_URL", database) + environment.setenv("INTEGRATION_PROXY_DATABASE_URL", database) + environment.delenv("INTEGRATION_PROXY_READ_REPLICA_URL", raising=False) + environment.delenv("DATABASE_URL_READ_REPLICA", raising=False) + with owned_proxy(shared, directory, {}) as owned: + owned.post("/config/update", {"router_settings": {"num_retries": 0}}) + yield owned + + def _stop(process: subprocess.Popen[bytes]) -> None: try: root_stopped: Final = stop_root_process(process) diff --git a/tests/integration/providers/test_bedrock_mantle_gpt_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_bedrock_mantle_gpt_prompt_cache_breakpoint_wire.py index 2af41c2b683..bc46cf5cde1 100644 --- a/tests/integration/providers/test_bedrock_mantle_gpt_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_gpt_prompt_cache_breakpoint_wire.py @@ -1,10 +1,12 @@ import asyncio import json +from collections.abc import Iterator from typing import Final from urllib.parse import urlsplit import pytest from integration._support.client import Gateway, eventually +from integration._support.process import isolated_gateway from integration._support.wire import Request, wire_server from integration.providers._mantle_gpt_prompt_cache_support import ( ANTHROPIC_PATH, @@ -68,6 +70,15 @@ from pydantic import JsonValue pytestmark = pytest.mark.timeout(600) +@pytest.fixture(scope="module") +def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with ( + pytest.MonkeyPatch.context() as environment, + isolated_gateway(tmp_path_factory.mktemp("gateway"), environment) as owned, + ): + yield owned + + @pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) @pytest.mark.parametrize("endpoint", ENDPOINTS) def test_a_configured_system_point_reaches_mantle_as_a_breakpoint_over_httpx( diff --git a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py index 343d5253900..9b07cce2f3a 100644 --- a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py @@ -5,7 +5,7 @@ import os import re import threading import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from pathlib import Path from queue import SimpleQueue @@ -21,7 +21,7 @@ import pytest import yaml from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.process import graceful_stop_seconds, isolated_gateway, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -318,6 +318,15 @@ def _spend_statuses(model: str, expected: int) -> list[JsonValue]: return [row["status"] for row in rows] +@pytest.fixture(scope="module") +def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with ( + pytest.MonkeyPatch.context() as environment, + isolated_gateway(tmp_path_factory.mktemp("gateway"), environment) as owned, + ): + yield owned + + def _proxy_url(gateway: Gateway) -> str: return str(gateway.client.base_url).rstrip("/") From 519906e6c209b51d17815fe1dead78537ff0c619 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 13:51:38 -0700 Subject: [PATCH 40/64] test(e2e): enable prometheus_system on the redis chaos proxy config (#47160) #46484 moved the litellm_redis_circuit_breaker_* metrics behind service_callback prometheus_system, so the chaos proxy, which only set callbacks prometheus, stopped exposing them and the breaker assertion read 0 --- tests/e2e/gateway/redis_chaos_ci_config.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/e2e/gateway/redis_chaos_ci_config.yml b/tests/e2e/gateway/redis_chaos_ci_config.yml index f7a71c50a71..1b098966a57 100644 --- a/tests/e2e/gateway/redis_chaos_ci_config.yml +++ b/tests/e2e/gateway/redis_chaos_ci_config.yml @@ -5,6 +5,7 @@ general_settings: litellm_settings: callbacks: ["prometheus"] + service_callback: ["prometheus_system"] require_auth_for_metrics_endpoint: false enable_redis_auth_cache: true cache: true From 0431e757167c17344e44f90e72c1a823f09a4ee5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 13:51:53 -0700 Subject: [PATCH 41/64] fix(tests): wait for every thread of a killed worker before treating it as gone (#47148) WorkerKill returned once the victim's main thread was a zombie. On Linux the process keeps its sockets open until its last thread exits, so a pooled keep-alive connection to the killed worker still looked idle and the next request through it got ECONNRESET. A zombie now counts as gone only once num_threads() is at most 1. Where psutil cannot read threads of a zombie (macOS) it raises ZombieProcess, a NoSuchProcess, so the old behavior holds --- tests/integration/_support/process.py | 2 +- .../unit/integration_support/test_process.py | 43 +++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 5cbe1309075..d5ce6f916a2 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -115,7 +115,7 @@ def _is_serving_worker(process: psutil.Process) -> bool: def _is_gone(process: psutil.Process) -> bool: try: - return process.status() == psutil.STATUS_ZOMBIE + return process.status() == psutil.STATUS_ZOMBIE and process.num_threads() <= 1 except psutil.NoSuchProcess: return True diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 4b053fb5a97..7741de93a80 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -2,6 +2,8 @@ from __future__ import annotations import importlib import os +import select +import socket import subprocess import sys import threading @@ -26,6 +28,7 @@ from tests.integration._support.process import ( OwnedProxy, WorkerHealthcheck, WorkerKill, + _is_gone, healthcheck_seconds, kill_stray_process, owned_proxy_process, @@ -469,6 +472,46 @@ def test_a_frozen_victim_stays_stopped_until_the_kill() -> None: assert victim.status() == psutil.STATUS_ZOMBIE +_LEADER_EXITS_FIRST: Final = """import ctypes +import socket +import threading + +listener = socket.create_server(("127.0.0.1", 0)) + + +def hold() -> None: + connection, _ = listener.accept() + connection.recv(1) + + +threading.Thread(target=hold).start() +print(listener.getsockname()[1], flush=True) +ctypes.CDLL(None).pthread_exit(None) +""" + + +def _closed_by_peer(connection: socket.socket) -> bool: + return bool(select.select((connection,), (), (), 0)[0]) + + +@pytest.mark.skipif(sys.platform != "linux", reason="only Linux keeps a zombie thread-group leader's other threads") +def test_a_process_is_not_gone_while_a_thread_outliving_its_leader_still_holds_its_sockets() -> None: + with subprocess.Popen((sys.executable, "-I", "-c", _LEADER_EXITS_FIRST), stdout=subprocess.PIPE) as child: + try: + assert child.stdout is not None + port: Final = int(child.stdout.readline()) + with socket.create_connection(("127.0.0.1", port), timeout=5) as connection: + leader: Final = psutil.Process(child.pid) + eventually(leader.status, lambda status: status == psutil.STATUS_ZOMBIE, seconds=10) + assert not _closed_by_peer(connection) + assert not _is_gone(leader) + connection.sendall(b"x") + eventually(lambda: _is_gone(leader), bool, seconds=10) + assert _closed_by_peer(connection) + finally: + child.kill() + + def _kill_after_freezing_the_supervisor(owned: OwnedProxy, worker: int, supervisor: psutil.Process) -> None: with owned.replacing_worker(worker) as kill: supervisor.suspend() From 1552b5b69f712d2d4fc557914dbb6f96d823c027 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 13:52:01 -0700 Subject: [PATCH 42/64] fix(proxy): warn at boot when a config environment_variables key replaces a value already set in the process environment (#47165) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 17 ++++++++++--- .../proxy/proxy_server/test_proxy_config.py | 24 +++++++++++++++++++ 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3ed3c57a96c..dba4b987725 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6535,10 +6535,13 @@ class ProxyConfig: environment_variables: Final = config.get("environment_variables", None) if environment_variables: for key, value in environment_variables.items(): + if not isinstance(key, str): + verbose_proxy_logger.warning("Skipping an environment_variables key that is not a string") + continue if key in self._BLOCKED_ENV_KEYS: verbose_proxy_logger.warning("Skipping blocked environment variable key: %s", key) continue - if isinstance(key, str) and is_mcp_stdio_flag_key(key): + if is_mcp_stdio_flag_key(key): verbose_proxy_logger.warning( "Ignoring %s set in the config file. Set it in the proxy's environment instead", MCP_STDIO_ENABLED_ENV_VAR, @@ -6554,7 +6557,7 @@ class ProxyConfig: if isinstance(value, str) and value.startswith("os.environ/"): resolved_secret_string: str | None = get_secret_str(secret_name=value) if resolved_secret_string is not None: - os.environ[key] = resolved_secret_string + self._set_config_environment_variable(key=key, value=resolved_secret_string) else: ######################################################### # handles this scenario: @@ -6563,13 +6566,21 @@ class ProxyConfig: # ARIZE_ENDPOINT: https://otlp.arize.com/v1 # ``` ######################################################### - os.environ[key] = str(value) + self._set_config_environment_variable(key=key, value=str(value)) # check if litellm_license in general_settings if "LITELLM_LICENSE" in environment_variables: _license_check.license_str = os.getenv("LITELLM_LICENSE", None) premium_user = _license_check.is_premium() + def _set_config_environment_variable(self, key: str, value: str) -> None: + if key in os.environ and os.environ[key] != value: + verbose_proxy_logger.warning( + "environment_variables.%s from the config file replaces the value already set in the process environment", + key, + ) + os.environ[key] = value + def _warn_on_misplaced_jwt_keys(self, config: dict) -> tuple[str, ...]: misplaced_jwt_keys = tuple(key for key in ("enable_jwt_auth", "litellm_jwtauth") if key in config) if not misplaced_jwt_keys: diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 51c84618a74..8837fa095e8 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -2381,6 +2381,30 @@ def test_ProxyConfig__load_environment_variables_sets_env(monkeypatch): assert result == {"TEST_LOAD_ENV_X": "hello", "set": True, "len": 1} +def test_ProxyConfig__load_environment_variables_warns_when_the_config_replaces_a_process_value( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + monkeypatch.setenv("TEST_LOAD_ENV_OVERRIDDEN", "from-process") + monkeypatch.setenv("TEST_LOAD_ENV_SAME", "same-value") + monkeypatch.delenv("TEST_LOAD_ENV_FRESH", raising=False) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyConfig()._load_environment_variables( + { + "environment_variables": { + "TEST_LOAD_ENV_OVERRIDDEN": "from-config", + "TEST_LOAD_ENV_SAME": "same-value", + "TEST_LOAD_ENV_FRESH": "fresh-value", + } + } + ) + warnings = [record.getMessage() for record in caplog.records if record.levelno == logging.WARNING] + assert os.environ["TEST_LOAD_ENV_OVERRIDDEN"] == "from-config" + assert len(warnings) == 1 + assert "TEST_LOAD_ENV_OVERRIDDEN" in warnings[0] + assert "from-process" not in warnings[0] and "from-config" not in warnings[0] + assert "TEST_LOAD_ENV_SAME" not in warnings[0] and "TEST_LOAD_ENV_FRESH" not in warnings[0] + + def test_ProxyConfig__load_environment_variables_blocks_dangerous_keys(monkeypatch): original_path = os.environ.get("PATH", "") pc = ProxyConfig() From 77c2ac18bd88ee89ae256fb627d4bc14aa5bb492 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 14:00:03 -0700 Subject: [PATCH 43/64] test(caching): assert sync cache-hit logging without depending on thread order (#47093) Sync success callbacks run on a thread pool, so the miss and the two hits can reach the logger in any order. Compare counts instead of the list order. --- tests/unit/caching/test_caching_handler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index 388e4c42ea0..51754c80dd6 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -42,6 +42,7 @@ from litellm.types.utils import ( Embedding, ) from litellm.types.llms.openai import ResponsesAPIResponse +from collections import Counter from collections.abc import Awaitable, Callable, Mapping from datetime import timedelta, datetime from typing import Final, cast @@ -104,7 +105,7 @@ def test_sync_cache_hits_reach_sync_custom_logger( call(**request_kwargs) assert all_logged.wait(timeout=10) - assert logged_cache_hits == [None, True, True] + assert Counter(logged_cache_hits) == Counter({None: 1, True: 2}), logged_cache_hits @pytest.fixture From 159f185ff45e624e6507bc24a795e7102b5b84b1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 14:00:48 -0700 Subject: [PATCH 44/64] test(integration): report uncollected shard files as collection errors instead of crashing xdist (#47075) The integration conftest raised UsageError from pytest_collection_modifyitems when a shard-selected file collected nothing. Under xdist that ran on the workers after the unfiltered item list was already published, so the worker ended its session with tests still queued and the controller hit xdist's crashitem assertion, naming an unrelated queued test and hiding the real ImportError. A selected file that collects nothing now produces a failed collection report instead. Files whose own collection already failed keep pytest's real error. Serial runs stop with "Interrupted: N errors during collection", xdist runs report the error and exit 1, and both still fail the shard. --- tests/integration/conftest.py | 15 ++- .../database/test_harness_lifecycle.py | 13 ++- .../unit/integration_support/test_sharding.py | 103 ++++++++++++++++++ 3 files changed, 126 insertions(+), 5 deletions(-) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index e82430c3add..92fed7200e1 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -38,6 +38,7 @@ REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]() WORKER_INVENTORY: Final = TypeAdapter(tuple[str, ...]) SHARD_PATH: Final[TypeAdapter[Path | None]] = TypeAdapter(Path | None) SHARD_FILES: Final = pytest.StashKey[frozenset[str] | None]() +FAILED_COLLECTIONS: Final = pytest.StashKey[frozenset[str]]() GROUP_SUFFIX: Final = re.compile(r"@integration_fixture_[0-9a-f]{64}$") INTEGER_OPTION: Final = TypeAdapter(int) STRING_ENVIRONMENT: Final = TypeAdapter(dict[str, str]) @@ -104,6 +105,12 @@ class IntegrationReportPlugin: def pytest_runtest_logreport(self, report: pytest.TestReport) -> None: self.config.stash[REPORTS].append(report) + def pytest_collectreport(self, report: pytest.CollectReport) -> None: + if report.failed: + self.config.stash[FAILED_COLLECTIONS] = self.config.stash.get(FAILED_COLLECTIONS, frozenset()) | { + report.nodeid.split("::", 1)[0] + } + @pytest.hookimpl(optionalhook=True) def pytest_xdist_node_collection_finished(self, node: object, ids: Sequence[str]) -> None: self.config.stash[COLLECTED] = tuple(_case_identity(nodeid) for nodeid in ids if _owned(nodeid)) @@ -224,8 +231,12 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item raise pytest.UsageError("Parallel integration runs require CircleCI's shard file selection") if files is not None: available: Final = frozenset(item.nodeid.split("::", 1)[0] for item in owned) - if files - available: - raise pytest.UsageError(f"Integration shard selected uncollected files: {sorted(files - available)}") + for file in sorted(files - available - config.stash.get(FAILED_COLLECTIONS, frozenset())): + config.hook.pytest_collectreport( + report=pytest.CollectReport( + file, "failed", f"Integration shard selected {file}, which has no tests", [] + ) + ) selected: Final = tuple(item for item in items if item not in owned or item.nodeid.split("::", 1)[0] in files) deselected: Final = tuple(item for item in owned if item.nodeid.split("::", 1)[0] not in files) config.hook.pytest_deselected(items=deselected) diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index 60bd7ae72a3..d600a92c5a2 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -784,9 +784,16 @@ def test_worker_crash_keeps_the_original_junit_identity_and_an_incomplete_execut assert cases[0].find("error") is not None -@pytest.mark.parametrize("assignment", ("absent", "missing", "foreign")) +@pytest.mark.parametrize( + ("assignment", "exit_code"), + ( + ("absent", pytest.ExitCode.USAGE_ERROR), + ("missing", pytest.ExitCode.USAGE_ERROR), + ("foreign", pytest.ExitCode.INTERRUPTED), + ), +) def test_invalid_native_selections_fail_before_any_case_runs( - tmp_path: Path, assignment: Literal["absent", "missing", "foreign"] + tmp_path: Path, assignment: Literal["absent", "missing", "foreign"], exit_code: pytest.ExitCode ) -> None: root: Final = Path(__file__).resolve().parents[3] events: Final = tmp_path / "executed.txt" @@ -823,7 +830,7 @@ def test_invalid_native_selections_fail_before_any_case_runs( timeout=15, check=False, ) - assert result.returncode == pytest.ExitCode.USAGE_ERROR, result.stdout + result.stderr + assert result.returncode == exit_code, result.stdout + result.stderr assert not events.exists() diff --git a/tests/unit/integration_support/test_sharding.py b/tests/unit/integration_support/test_sharding.py index 2317279f2b2..0ffd7f3572a 100644 --- a/tests/unit/integration_support/test_sharding.py +++ b/tests/unit/integration_support/test_sharding.py @@ -1,5 +1,8 @@ from __future__ import annotations +import os +import subprocess +import sys from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -13,6 +16,21 @@ from tests.integration.conftest import COLLECTED, IntegrationReportPlugin, _shar from tests.integration.conftest import INVENTORY as WORKER_INVENTORY INVENTORY: Final = ("test_owned.py::test_first", "test_owned.py::test_skipped", "test_owned.py::test_last") +INTEGRATION_PACKAGE: Final = Path(__file__).resolve().parents[2] / "integration" +HEALTHY: Final = "tests/integration/providers/test_healthy.py" +BROKEN: Final = "tests/integration/providers/test_broken.py" +RENAMED: Final = "tests/integration/providers/test_renamed.py" +CLASS_ERROR: Final = "tests/integration/providers/test_class_error.py" +MODULES: Final = MappingProxyType( + { + HEALTHY: "def test_one():\n pass\n\n\ndef test_two():\n pass\n", + BROKEN: "from tests.integration._support.process import _removed_helper\n\n\ndef test_never():\n pass\n", + CLASS_ERROR: ( + "import pytest\n\n\nclass TestBroken:\n" + " @pytest.mark.parametrize('unused', [1])\n def test_case(self):\n pass\n" + ), + } +) def _execution( @@ -111,3 +129,88 @@ def test_worker_results_retain_the_full_inventory_and_reject_disagreement() -> N incomplete: Final = WorkerResult(MappingProxyType({"integration_inventory": INVENTORY[:1]})) with pytest.raises(pytest.UsageError, match="inventories disagree"): plugin.pytest_testnodedown(incomplete, None) + + +XDIST: Final = ("-n", "2", "--dist=loadgroup") +SERIAL: Final = ("-p", "no:xdist") + + +def _run_shard( + root: Path, selected: tuple[str, ...], distribution: tuple[str, ...] = XDIST +) -> subprocess.CompletedProcess[str]: + package: Final = root / "tests" / "integration" + (package / "providers").mkdir(parents=True) + (root / "pytest.ini").write_text("[pytest]\n") + (root / "tests" / "__init__.py").write_text("") + for name in ("__init__.py", "conftest.py", "run.py", "_support"): + (package / name).symlink_to(INTEGRATION_PACKAGE / name) + for file in selected: + if file in MODULES: + (root / file).write_text(MODULES[file]) + (root / "shard.txt").write_text("".join(f"{file}\n" for file in selected)) + return subprocess.run( + ( + sys.executable, + "-P", + "-m", + "pytest", + "tests/integration/providers", + "-p", + "no:cacheprovider", + *distribution, + "--integration-shard-count=2", + "--integration-shard-index=0", + "--integration-shard-files", + "shard.txt", + ), + cwd=root, + env={ + "PATH": os.environ.get("PATH", ""), + "HOME": str(root), + "PYTHONPATH": os.pathsep.join((str(root), str(root / "tests"))), + }, + capture_output=True, + text=True, + timeout=180, + check=False, + ) + + +@pytest.mark.parametrize( + ("distribution", "exit_code", "outcome"), + ( + (XDIST, pytest.ExitCode.TESTS_FAILED, "2 passed, 1 error"), + (SERIAL, pytest.ExitCode.INTERRUPTED, "Interrupted: 1 error during collection"), + ), + ids=("xdist", "serial"), +) +def test_a_broken_import_in_a_shard_is_a_collection_error_not_a_worker_crash( + tmp_path: Path, distribution: tuple[str, ...], exit_code: pytest.ExitCode, outcome: str +) -> None: + result: Final = _run_shard(tmp_path, (HEALTHY, BROKEN), distribution) + assert result.returncode == exit_code, result.stdout + result.stderr + assert outcome in result.stdout + assert result.stdout.count(f"ERROR collecting {BROKEN}") == 1 + assert "cannot import name '_removed_helper'" in result.stdout + assert "INTERNALERROR" not in result.stdout + result.stderr + + +def test_a_shard_file_that_collects_no_tests_still_fails_as_a_collection_error(tmp_path: Path) -> None: + result: Final = _run_shard(tmp_path, (HEALTHY, RENAMED)) + assert result.returncode == pytest.ExitCode.TESTS_FAILED, result.stdout + result.stderr + assert result.stdout.count(f"ERROR collecting {RENAMED}") == 1 + assert "2 passed" in result.stdout + assert "INTERNALERROR" not in result.stdout + result.stderr + + +def test_a_class_level_collection_error_is_not_also_reported_as_a_file_without_tests(tmp_path: Path) -> None: + result: Final = _run_shard(tmp_path, (HEALTHY, CLASS_ERROR), SERIAL) + assert result.returncode == pytest.ExitCode.INTERRUPTED, result.stdout + result.stderr + assert f"ERROR {CLASS_ERROR}::TestBroken" in result.stdout + assert "which has no tests" not in result.stdout + + +def test_a_healthy_shard_runs_every_selected_test(tmp_path: Path) -> None: + result: Final = _run_shard(tmp_path, (HEALTHY,)) + assert result.returncode == pytest.ExitCode.OK, result.stdout + result.stderr + assert "2 passed" in result.stdout From e0514d92b83fbb98e69dedfc6637b885c035fae2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 14:05:42 -0700 Subject: [PATCH 45/64] test(e2e): hold failing-request RSS to a loose ceiling over two measured batches (#47170) The 48 MB budget on the third batch kept failing on the x86 release gate with no leak behind it. Isolated rc/1.107.0 and v1.105.1 stacks both grew about 60 MB on the third batch with gc-tracked objects flat, and both levelled off to a couple of MB per batch after twenty batches. Measure batches 3-4 against a 200 MB per-worker ceiling instead, which a native leak of 400 KB per failed request still trips. The stored-request size check stays the deterministic catch for the v1.100.0 leak --- tests/e2e/AGENTS.md | 2 +- tests/e2e/e2e_config.py | 2 +- .../e2e/router/test_reliability_memory_e2e.py | 83 +++++++++++-------- 3 files changed, 52 insertions(+), 35 deletions(-) diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index d267059cdc7..62ebfd12fe5 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -43,7 +43,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below), plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection -- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) +- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) - `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 8fc7e50d95c..efbc95966c4 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -237,7 +237,7 @@ MEMORY_TRANSCRIPT_TURNS = int(os.environ.get("E2E_MEMORY_TRANSCRIPT_TURNS", "40" MEMORY_CONCURRENCY = int(os.environ.get("E2E_MEMORY_CONCURRENCY", "4")) MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES", "15")) MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1")) -MEMORY_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_RSS_BUDGET_MB", "48")) +MEMORY_RSS_CEILING_MB = float(os.environ.get("E2E_MEMORY_RSS_CEILING_MB", "200")) MEMORY_IDLE_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_IDLE_RSS_BUDGET_MB", "768")) MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64")) diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/router/test_reliability_memory_e2e.py index 196784b5345..fb114092303 100644 --- a/tests/e2e/router/test_reliability_memory_e2e.py +++ b/tests/e2e/router/test_reliability_memory_e2e.py @@ -1,5 +1,5 @@ """Live e2e: a few hundred requests that fail before any provider answers must not -grow the proxy's resident memory past a fixed budget once the proxy is warm. +grow the proxy's resident memory past a loose ceiling once the proxy is warm. The regression this guards shipped in v1.100.0: every retry breadcrumb copied the whole request and the copies nested into one router-global list, so a proxy under @@ -13,26 +13,27 @@ deployment that fails a handful of times in a row is benched (a bad-credential 4 included), the router answers "No deployments available" without retrying, and the retry loop that leaks stops running; cooldown_time 0 keeps it running. -Three identical batches run back to back. The first two are the warmup that grows -the proxy's caches and allocator arenas to their steady state, the third is the one -the budget applies to, so a healthy proxy shows the third batch adding roughly -nothing while a leaking one adds a fixed amount per request. One warmup batch is not -enough: each batch's in-flight retries peak well above the live heap, the allocator -keeps that high-water mark, and on a cold worker the second batch still set a new one -(about 50 MB of RSS with the live Python heap flat to within 1 MB), where the later -batches each settled under 25 MB. RSS is read through /debug/memory/summary on every +Four identical batches run back to back. The first two are the warmup that grows +the proxy's caches and allocator arenas, the last two are the measured window, and +each worker's RSS may grow at most a loose ceiling over that window. RSS cannot be +held to a tight per-batch budget: on the x86 release gate (rc/1.107.0 and v1.105.1 +on isolated stacks, 28 batches each) the allocator kept setting new high-water marks +for the first several batches, about 60 MB on the third batch of a fresh stack and +more after the rest of the suite, with the gc-tracked object count flat to 0.1%, and +only settled to a couple of MB per batch after twenty. So the ceiling catches a gross +leak and nothing slower. RSS is read through /debug/memory/summary on every configured replica; a burst of failing calls leaves a transient bulge of garbage that gc reclaims within seconds, so each checkpoint samples until no new worker has answered for a settle window and keeps the lowest reading per worker. A checkpoint runs before the warmup and after every batch, and each batch's settled growth per worker rides out as a repeated `rss_batch_growth` JUnit property as soon as that -batch settles, so a run that fails keeps the curve up to the batch it failed in. The budget is judged per worker -(by replica address, hostname and pid, since pods in their own pid namespaces report -the same pids) so each worker is compared with itself, and the checkpoints around the -measured batch must see the same workers: a single load-balanced address reaches the -workers behind it one answer at a time, and a worker that answered only one -checkpoint would otherwise drop out of the comparison, which is where a leaking -worker could hide. +batch settles, so a run that fails keeps the curve up to the batch it failed in. The +ceiling is judged per worker (by replica address, hostname and pid, since pods in +their own pid namespaces report the same pids) so each worker is compared with +itself, and the checkpoints around the measured window must see the same workers: a +single load-balanced address reaches the workers behind it one answer at a time, and +a worker that answered only one checkpoint would otherwise drop out of the +comparison, which is where a leaking worker could hide. RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS @@ -79,7 +80,7 @@ from e2e_config import ( MEMORY_IDLE_RSS_BUDGET_MB, MEMORY_REQUESTS_PER_PHASE, MEMORY_RETRIES_PER_REQUEST, - MEMORY_RSS_BUDGET_MB, + MEMORY_RSS_CEILING_MB, MEMORY_RSS_SAMPLE_INTERVAL_SECONDS, MEMORY_RSS_SETTLE_SAMPLES, MEMORY_STORED_REQUEST_BUDGET_KB, @@ -99,7 +100,9 @@ pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] DEPLOYMENTS_PER_GROUP: Final = 2 RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES WARMUP_BATCHES: Final = 2 -MEASURED_BATCH: Final = WARMUP_BATCHES + 1 +MEASURED_BATCHES: Final = 2 +FIRST_MEASURED_BATCH: Final = WARMUP_BATCHES + 1 +LAST_BATCH: Final = WARMUP_BATCHES + MEASURED_BATCHES @dataclass(frozen=True, slots=True) @@ -207,7 +210,7 @@ def _record_batch_growth( ) -def _settled_rss_around_measured_batch( +def _settled_rss_after_batches( proxy: ProxyClient, key: str, model: str, @@ -216,21 +219,25 @@ def _settled_rss_around_measured_batch( record_property: Callable[[str, object], None], before: Mapping[WorkerKey, RssReading], batch: int, -) -> tuple[Mapping[WorkerKey, RssReading], Mapping[WorkerKey, RssReading]]: + last_batch: int, +) -> Mapping[WorkerKey, RssReading]: calls: Final = _fail_many(proxy, key, model, override) _assert_every_call_failed_through_fallback(calls, fallback) after: Final = _settled_rss_per_worker(proxy) _record_batch_growth(record_property, batch, before, after) - if batch == MEASURED_BATCH: - return before, after - return _settled_rss_around_measured_batch(proxy, key, model, override, fallback, record_property, after, batch + 1) + if batch == last_batch: + return after + return _settled_rss_after_batches( + proxy, key, model, override, fallback, record_property, after, batch + 1, last_batch + ) def _heaviest_worker_growth( warm: Mapping[WorkerKey, RssReading], after: Mapping[WorkerKey, RssReading] ) -> WorkerGrowth: assert warm.keys() == after.keys(), ( - f"the workers answering /debug/memory/summary changed around measured batch {MEASURED_BATCH}, so not every " + f"the workers answering /debug/memory/summary changed around measured batches {FIRST_MEASURED_BATCH}-" + f"{LAST_BATCH}, so not every " f"worker can be compared with itself: gone after it {sorted(warm.keys() - after.keys())} (a worker " f"that died or was restarted under failing traffic, which is what an OOM kill looks like), first seen " f"after it {sorted(after.keys() - warm.keys())} (the warm window never reached them, so they have no " @@ -332,7 +339,7 @@ class TestReliabilityMemory: f"request into the stored snapshot the way the v1.100.0 ones did" ) - warm, after = _settled_rss_around_measured_batch( + warm: Final = _settled_rss_after_batches( client.proxy, scoped_key, primary, @@ -341,16 +348,26 @@ class TestReliabilityMemory: record_property, _settled_rss_per_worker(client.proxy), 1, + WARMUP_BATCHES, + ) + after: Final = _settled_rss_after_batches( + client.proxy, + scoped_key, + primary, + override, + fallback, + record_property, + warm, + FIRST_MEASURED_BATCH, + LAST_BATCH, ) heaviest: Final = _heaviest_worker_growth(warm, after) - assert heaviest.growth_mb <= MEMORY_RSS_BUDGET_MB, ( - f"proxy RSS grew {heaviest.growth_mb:.1f} MB over measured batch {MEASURED_BATCH} of " - f"{MEMORY_REQUESTS_PER_PHASE} failing requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) " - f"after {WARMUP_BATCHES} identical warmup batches, past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: " - f"worker pid {heaviest.warm.worker_pid} on " + assert heaviest.growth_mb <= MEMORY_RSS_CEILING_MB, ( + f"proxy RSS grew {heaviest.growth_mb:.1f} MB over measured batches {FIRST_MEASURED_BATCH}-{LAST_BATCH} " + f"({MEASURED_BATCHES * MEMORY_REQUESTS_PER_PHASE} failing requests) after {WARMUP_BATCHES} identical " + f"warmup batches, past the {MEMORY_RSS_CEILING_MB:.0f} MB ceiling: worker pid {heaviest.warm.worker_pid} on " f"{heaviest.warm.hostname or 'an unnamed host'} behind {heaviest.warm.replica} settled at " - f"{heaviest.warm.ram_usage_mb:.1f} MB warm and " - f"{heaviest.after.ram_usage_mb:.1f} MB after; failing requests are leaking memory the way the " - f"v1.100.0 retry breadcrumbs did" + f"{heaviest.warm.ram_usage_mb:.1f} MB warm and {heaviest.after.ram_usage_mb:.1f} MB after; failing " + f"requests are leaking memory the way the v1.100.0 retry breadcrumbs did" ) From ef1d3959a06289eae04b89c94426e7c93ebbb9a5 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:15:38 +0000 Subject: [PATCH 46/64] fix(count_tokens)!: answer 400 for Vertex partner models without a count API and OpenAI-shaped provider errors on /v1/responses/input_tokens (#47163) * fix(count_tokens): answer 400 for Vertex partner models without a count API and OpenAI-shaped provider errors on /v1/responses/input_tokens * fix(count_tokens): answer 400 for Vertex Mistral counts and type the input_tokens error tests Vertex has no count-tokens route for Mistral or codestral models, so they now get the same 400 as other partner models without a count API instead of a 404 that points operators at IAM. * fix(count_tokens): unwrap prefixed provider errors on /v1/responses/input_tokens * fix(count_tokens): say LiteLLM, not Vertex, cannot count a Vertex partner model * fix(count_tokens): count a Vertex partner model with its configured custom_tokenizer and unwrap prefixed flat provider errors * fix(count_tokens): keep the 400 for contents-only Vertex partner counts with a custom_tokenizer * fix(count_tokens): pass through an OpenAI provider error that omits param or code * fix(count_tokens): keep the status-based type for Anthropic count errors on the Responses route * test(count_tokens): cover the Vertex partner count gate and Responses count error shapes on a real proxy * test(count_tokens): add regression tests for the Vertex partner count gate and Responses count error shapes --- litellm/llms/vertex_ai/common_utils.py | 18 ++ .../count_tokens/handler.py | 36 ++- .../common_utils/openai_error_payload.py | 4 + .../proxy/response_api_endpoints/endpoints.py | 97 +++++++- .../test_openai_count_tokens_tools_wire.py | 25 ++ .../test_vertex_partner_count_tokens_wire.py | 177 ++++++++++++++- .../vertex_ai/test_vertex_ai_common_utils.py | 29 +++ .../response_api_endpoints/test_endpoints.py | 213 ++++++++++++++++++ tests/unit/proxy/test_custom_tokenizer_bug.py | 60 ++++- 9 files changed, 627 insertions(+), 32 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 27ca4acf05d..3c9a8429882 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1278,6 +1278,7 @@ class VertexAITokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIError as PartnerVertexAIError, ) @@ -1293,6 +1294,23 @@ class VertexAITokenCounter(BaseTokenCounter): ): return None + without_count_route: Final = ( + VertexAIPartnerModels.is_vertex_partner_model(model_to_use) and count_tokens_publisher(model_to_use) is None + ) + model_info: Final[Mapping[str, object]] = deployment.get("model_info") or {} + if without_count_route and messages and model_info.get("custom_tokenizer") is not None: + return None + if without_count_route: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="vertex_ai_partner_models", + error=True, + error_message=f"LiteLLM does not support counting tokens for Vertex AI model {model_to_use}.", + status_code=400, + ) + # Check if this is a partner model (Claude, Mistral, etc.) if VertexAIPartnerModels.is_vertex_partner_model(model_to_use): # Use partner models token counter diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 285380f5fc7..7100c9efde3 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -1,5 +1,5 @@ """ -Token counter for Vertex AI Partner Models (Anthropic Claude, Mistral, etc.) +Token counter for Vertex AI Partner Models (Anthropic Claude, Meta Llama) This handler provides token counting for partner models hosted on Vertex AI. Unlike Gemini models which use Google's token counting API, partner models use @@ -12,36 +12,30 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +_COUNT_TOKENS_PUBLISHERS: Final = ( + ("claude", "anthropic"), + ("llama", "meta"), + ("meta/", "meta"), +) + + +def count_tokens_publisher(model: str) -> str | None: + return next((publisher for marker, publisher in _COUNT_TOKENS_PUBLISHERS if marker in model), None) + class VertexAIPartnerModelsTokenCounter(VertexBase): """ Token counter for Vertex AI Partner Models. - Handles token counting for models like Claude (Anthropic), Mistral, etc. + Handles token counting for models like Claude (Anthropic) and Llama (Meta) that are available through Vertex AI's partner model program. """ def _get_publisher_for_model(self, model: str) -> str: - """ - Determine the publisher name for the given model. - - Args: - model: The model name (e.g., "claude-3-5-sonnet-20241022") - - Returns: - Publisher name to use in the Vertex AI endpoint URL - - Raises: - ValueError: If the model is not a recognized partner model - """ - if "claude" in model: - return "anthropic" - elif "mistral" in model or "codestral" in model: - return "mistralai" - elif "llama" in model or "meta/" in model: - return "meta" - else: + publisher: Final = count_tokens_publisher(model) + if publisher is None: raise ValueError(f"Unknown partner model: {model}") + return publisher def _build_count_tokens_endpoint( self, diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index f092f9637f8..2f01dbe11f7 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -42,6 +42,10 @@ def openai_error_type(exc: object, status_code: int) -> str: carried: Final = attribute_of(exc, "type") if isinstance(carried, str) and carried != STRINGIFIED_NONE: return carried + return openai_error_type_for_status(status_code) + + +def openai_error_type_for_status(status_code: int) -> str: mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: return mapped diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index aebcab75601..3c727a6414c 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -25,12 +25,13 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from openai.types.responses import ResponseItemList from openai.types.responses.response_create_params import ResponseInputParam -from pydantic import ConfigDict, TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config from starlette.websockets import WebSocket, WebSocketDisconnect -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger +from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs @@ -54,12 +55,18 @@ from litellm.proxy.common_request_processing import ( create_response, resolve_counted_tool_fields, ) +from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_set_request_parsed_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, + error_status_code, + openai_error_type_for_status, +) from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id from litellm.proxy.litellm_pre_call_utils import ( @@ -1411,6 +1418,85 @@ def _responses_invalid_request_response(message: str, param: str | None, code: s return JSONResponse(status_code=400, content=body) +class _ProviderErrorMessage(TypedDict): + message: ReadOnly[str] + + +class _NestedProviderError(TypedDict): + error: ReadOnly[_ProviderErrorMessage] + + +@with_config(ConfigDict(extra="ignore")) +class _ProviderOpenAIErrorDetail(TypedDict): + message: ReadOnly[str] + type: ReadOnly[str] + param: NotRequired[ReadOnly[str | None]] + code: NotRequired[ReadOnly[str | None]] + + +@with_config(ConfigDict(extra="forbid")) +class _ProviderOpenAIError(TypedDict): + error: ReadOnly[_ProviderOpenAIErrorDetail] + + +_PROVIDER_COUNT_ERROR_TYPE: Final = "token_counting_error" +_PROVIDER_OPENAI_ERROR: Final = TypeAdapter(_ProviderOpenAIError) +_NESTED_PROVIDER_ERROR: Final = TypeAdapter(_NestedProviderError) + + +def _provider_error_json(raw_message: str) -> str: + start: Final = raw_message.find("{") + return raw_message if start == -1 else raw_message[start:] + + +def _provider_error_message(raw_message: str) -> str: + error_json: Final = _provider_error_json(raw_message) + try: + return _NESTED_PROVIDER_ERROR.validate_json(error_json)["error"]["message"] + except ValidationError: + flat_message: Final = AnthropicExceptionMapping.extract_error_message(error_json) + return raw_message if flat_message == error_json else flat_message + + +def _provider_count_error_body(raw_message: str, status_code: int) -> _ResponsesApiErrorBody: + try: + provider_error: Final = _PROVIDER_OPENAI_ERROR.validate_json(_provider_error_json(raw_message))["error"] + except ValidationError: + return { + "error": { + "message": _provider_error_message(raw_message), + "type": openai_error_type_for_status(status_code), + "param": None, + "code": None, + } + } + return { + "error": { + "message": provider_error["message"], + "type": provider_error["type"], + "param": provider_error.get("param"), + "code": provider_error.get("code"), + } + } + + +def _provider_count_error_response(exc: ProxyException, request: Request) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does + general_settings_view, + ) + + status_code: Final = error_status_code(exc, 500) + _close_dangling_otel_server_span(request, status_code, exc=exc) + body: Final = _provider_count_error_body(exc.message, status_code) + call_id: Final = error_body_call_id(general_settings_view(), exc.headers.get(LITELLM_CALL_ID_HEADER)) + return JSONResponse( + status_code=status_code, + content={"error": with_call_id({**body["error"]}, call_id)}, + headers=exc.headers, + ) + + def _missing_responses_param_response(param: str) -> JSONResponse: return _responses_invalid_request_response( message=f"Missing required parameter: '{param}'.", @@ -1515,7 +1601,12 @@ async def responses_input_tokens( message=f"Invalid request for token counting: {e}", param=None, code=None ) - token_response: Final = await token_counter(request=token_request, call_endpoint=True) + try: + token_response: Final = await token_counter(request=token_request, call_endpoint=True) + except ProxyException as e: + if e.type != _PROVIDER_COUNT_ERROR_TYPE: + raise + return _provider_count_error_response(e, request) result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 2b23ccdf4f2..16a5b8685ad 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -391,3 +391,28 @@ def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: O ) assert response.status_code == 200, response.text assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English." + + +@pytest.mark.parametrize( + "error", + [ + pytest.param( + {"message": "Invalid input", "type": "invalid_request_error", "param": "input", "code": "invalid_value"}, + id="full", + ), + pytest.param({"message": "Invalid input", "type": "invalid_request_error"}, id="bare"), + ], +) +def test_responses_input_tokens_passes_an_openai_count_error_through( + counting_proxy: OwnedProxy, error: dict[str, JsonValue] +) -> None: + def rejecting(_request: Request) -> Reply: + return Reply(status=400, body=json.dumps({"error": error}).encode()) + + gateway: Final = counting_proxy.gateway + with wire_server(rejecting) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {"model": model, "input": "hello"}) + assert response.status_code == 400, response.text + assert response.json() == {"error": {"param": None, "code": None, **error}}, response.text + assert _count_body(wire)["input"] == [{"role": "user", "content": "hello"}] diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 4cf9a3159fe..8d520a038af 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -5,8 +5,10 @@ import uuid from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack +from dataclasses import dataclass from pathlib import Path from queue import SimpleQueue +from types import MappingProxyType from typing import Final import httpx @@ -14,11 +16,12 @@ import psutil import pytest import yaml from integration._support.client import Gateway, Scenario, eventually -from integration._support.process import owned_proxy_process +from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter +from tokenizers import Regex, Tokenizer, models, pre_tokenizers _BACKEND: Final = "claude-sonnet-4-6" _PROJECT: Final = "scripted-project" @@ -34,9 +37,19 @@ _REPLY_TEXT: Final = "scripted reply" _OWNED_MODEL: Final = "partner-claude" _OWNED_UNREACHABLE_MODEL: Final = "partner-claude-unreachable" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_GLM_BACKEND: Final = "zai-org/glm-5.2-maas" +_MISTRAL_BACKEND: Final = "mistral-medium-3" +_LLAMA_BACKEND: Final = "meta/llama-4-maverick-17b-128e-instruct-maas" +_LLAMA_COUNT_TARGET: Final = ( + f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/meta/models/count-tokens:rawPredict" +) +_COUNTED_TEXT: Final = "Count this message" +_TOKENIZER_REPO: Final = "scripted/char-tokenizer" +_TOKENIZER_TARGET: Final = f"/{_TOKENIZER_REPO}/resolve/main/tokenizer.json" +_TOKENIZER_HEADERS: Final = MappingProxyType({"X-Repo-Commit": "0" * 40, "ETag": '"scripted-char-tokenizer"'}) _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this message"}] +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": _COUNTED_TEXT}] _SYSTEM: Final = "You are a terse assistant that answers in one sentence" _SYSTEM_BLOCKS: Final[list[JsonValue]] = [ {"type": "text", "text": "You are a terse assistant"}, @@ -114,16 +127,16 @@ def _counted(_request: Request) -> Reply: return Reply(body=json.dumps({"input_tokens": _PEER_COUNT}).encode()) -def _rejected(status: int) -> Reply: +def _rejected(status: int, error_type: str = "invalid_request_error") -> Reply: return Reply( status=status, - body=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": _REJECTION}}).encode(), + body=json.dumps({"type": "error", "error": {"type": error_type, "message": _REJECTION}}).encode(), ) -def _rejecting(status: int) -> Callable[[Request], Reply]: +def _rejecting(status: int, error_type: str = "invalid_request_error") -> Callable[[Request], Reply]: def count(_request: Request) -> Reply: - return _rejected(status) + return _rejected(status, error_type) return count @@ -223,6 +236,78 @@ def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: return total +@dataclass(frozen=True, slots=True) +class _CountRoute: + path: Callable[[str], str] + body: Callable[[str], dict[str, JsonValue]] + rejection: Callable[[str], dict[str, JsonValue]] + count_field: str + params: Mapping[str, str] = MappingProxyType({}) + + def send(self, gateway: Gateway, model: str) -> httpx.Response: + return gateway.request("POST", self.path(model), self.body(model), params=self.params) + + +_MESSAGES_ROUTE: Final = _CountRoute( + path=lambda _model: "/v1/messages/count_tokens", + body=_bare, + rejection=lambda message: {"type": "error", "error": {"type": "invalid_request_error", "message": message}}, + count_field="input_tokens", +) +_RESPONSES_ROUTE: Final = _CountRoute( + path=lambda _model: "/v1/responses/input_tokens", + body=lambda model: {"model": model, "input": _COUNTED_TEXT}, + rejection=lambda message: { + "error": {"message": message, "type": "invalid_request_error", "param": None, "code": None} + }, + count_field="input_tokens", +) +_GOOGLE_ROUTE: Final = _CountRoute( + path=lambda model: f"/v1beta/models/{model}:countTokens", + body=lambda _model: {"contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]}, + rejection=lambda message: {"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}}, + count_field="totalTokens", +) +_UTILS_ROUTE: Final = _CountRoute( + path=lambda _model: "/utils/token_counter", + body=_bare, + rejection=lambda message: { + "error": {"message": message, "type": "token_counting_error", "param": "model", "code": "400"} + }, + count_field="total_tokens", + params=MappingProxyType({"call_endpoint": "true"}), +) +_COUNT_ROUTES: Final = (_MESSAGES_ROUTE, _RESPONSES_ROUTE, _GOOGLE_ROUTE, _UTILS_ROUTE) + + +def _error_message(payload: Mapping[str, JsonValue]) -> str: + error: Final = payload.get("error") + assert isinstance(error, dict), payload + message: Final = error.get("message") + assert isinstance(message, str), payload + return message + + +def _route_count(gateway: Gateway, route: _CountRoute, model: str) -> tuple[int, JsonValue]: + response: Final = route.send(gateway, model) + return response.status_code, _JSON_OBJECT.validate_json(response.content).get(route.count_field) + + +def _char_tokenizer() -> bytes: + tokenizer: Final = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Split(Regex("."), behavior="isolated") + return tokenizer.to_str().encode() + + +def _tokenizer_hub(tokenizer: bytes) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == _TOKENIZER_TARGET: + return Reply(body=tokenizer, headers=_TOKENIZER_HEADERS) + return Reply(status=404, body=json.dumps({"error": f"unscripted target {request.target}"}).encode()) + + return respond + + def _closed_port_url() -> str: return f"http://127.0.0.1:{allocate_port()}" @@ -371,11 +456,89 @@ def test_responses_input_tokens_surfaces_the_peer_error(gateway: Gateway) -> Non {"model": model, "input": "Count this message", "instructions": "Be terse", "tools": _RESPONSES_TOOLS}, ) assert response.status_code == 400, response.text - assert "input_tokens" not in response.text, response.text + assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text (sent,) = _counted_bodies(wire) assert sent["tools"] == _RESPONSES_TOOLS, sent +def test_responses_input_tokens_types_a_peer_error_by_its_status(gateway: Gateway) -> None: + with wire_server(_peer(_rejecting(404, "not_found_error"))) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url) + response: Final = _RESPONSES_ROUTE.send(gateway, model) + assert response.status_code == 404, response.text + assert response.json() == _RESPONSES_ROUTE.rejection(_REJECTION), response.text + assert _counted_bodies(wire) == (_PEER_BARE,) + + +@pytest.mark.parametrize( + ("backend", "route"), + [ + pytest.param(_GLM_BACKEND, _MESSAGES_ROUTE, id="glm-messages"), + pytest.param(_GLM_BACKEND, _RESPONSES_ROUTE, id="glm-responses"), + pytest.param(_GLM_BACKEND, _GOOGLE_ROUTE, id="glm-google"), + pytest.param(_GLM_BACKEND, _UTILS_ROUTE, id="glm-utils"), + pytest.param(_MISTRAL_BACKEND, _MESSAGES_ROUTE, id="mistral-messages"), + ], +) +def test_partner_without_a_count_route_answers_400_without_calling_vertex( + gateway: Gateway, backend: str, route: _CountRoute +) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{backend}") + response: Final = route.send(gateway, model) + assert response.status_code == 400, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + message: Final = _error_message(payload) + assert backend in message, message + assert payload == route.rejection(message), response.text + assert wire.drain() == () + + +def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> None: + with wire_server(_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}") + response: Final = _count(gateway, _bare(model)) + assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + (sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET) + assert sent["messages"] == _MESSAGES, sent + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_custom_tokenizer_partner_counts_locally_on_every_count_route(gateway: Gateway, tmp_path: Path) -> None: + environment: Final = {"HF_HOME": str(tmp_path / "hf-home")} + with ( + wire_server(_tokenizer_hub(_char_tokenizer())) as hub, + wire_server(_peer()) as vertex, + owned_proxy_process(gateway, tmp_path, {**environment, "HF_ENDPOINT": hub.url}) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = _deployment( + owned.gateway, + scenario, + vertex.url, + model=f"vertex_ai/{_GLM_BACKEND}", + model_info={"custom_tokenizer": {"identifier": _TOKENIZER_REPO, "revision": "main", "auth_token": None}}, + ) + local: Final = _payload( + owned.gateway.request("POST", "/utils/token_counter", _bare(model), params={"call_endpoint": "false"}) + ) + assert local["tokenizer_type"] == "huggingface_tokenizer", local + total: Final = local["total_tokens"] + assert isinstance(total, int) and total >= len(_COUNTED_TEXT), local + counts: Final = tuple(_route_count(owned.gateway, route, model) for route in _COUNT_ROUTES) + assert counts == ((200, total),) * len(_COUNT_ROUTES), counts + contents_only: Final = owned.gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "contents": [{"role": "user", "parts": [{"text": _COUNTED_TEXT}]}]}, + params={"call_endpoint": "true"}, + ) + assert contents_only.status_code == 400, contents_only.text + assert _GLM_BACKEND in _error_message(_JSON_OBJECT.validate_json(contents_only.content)), contents_only.text + assert vertex.drain() == () + assert {request.target for request in hub.drain()} >= {_TOKENIZER_TARGET} + + def test_gemini_count_tokens_route_reaches_the_partner_peer(gateway: Gateway) -> None: with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 97c0a657ca4..20dba332c08 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -51,6 +51,35 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request( assert "text blocks" in (result.error_message or "") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + ( + "zai-org/glm-5.2-maas", + "deepseek-ai/deepseek-v3.2-maas", + "qwen/qwen3-coder-480b-a35b-instruct-maas", + "mistral-medium-3", + "codestral-2501", + "openai/gpt-oss-120b-maas", + "google/gemma-4-26b-a4b-it-maas", + ), +) +async def test_vertex_partner_model_without_an_inference_matching_count_answers_400_naming_the_model(model: str) -> None: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use=model, + messages=[{"role": "user", "content": "Count this"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "global"}}, + request_model=f"vertex_ai/{model}", + messages_format="anthropic", + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert result.error_message == f"LiteLLM does not support counting tokens for Vertex AI model {model}." + + class _VertexCountTokensPeer: def __init__(self, status_code: int) -> None: self.status_code: Final = status_code diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 52e95901a16..fc329eaec5f 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -4029,6 +4029,219 @@ class TestResponsesInputTokens: assert response.status_code == 429, response.text assert response.json()["error"]["message"] == "rate limited" + @pytest.mark.parametrize( + "provider_body, status_code, expected_error", + [ + pytest.param( + { + "error": { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": "missing_required_parameter", + } + }, + 400, + { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": "missing_required_parameter", + }, + id="openai_error_passes_through", + ), + pytest.param( + { + "error": { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + } + }, + 400, + { + "message": "Missing required parameter: 'tools[0].name'.", + "type": "invalid_request_error", + "param": "tools[0].name", + "code": None, + }, + id="openai_error_without_code_keeps_type_and_param", + ), + pytest.param( + { + "type": "error", + "error": {"type": "invalid_request_error", "message": "tools.0: Input tag 'function' is invalid"}, + "request_id": "req_vrtx_1", + }, + 400, + { + "message": "tools.0: Input tag 'function' is invalid", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="anthropic_error_message_unwrapped", + ), + pytest.param( + { + "type": "error", + "error": {"type": "not_found_error", "message": "model: claude-typo"}, + "request_id": "req_vrtx_2", + }, + 404, + {"message": "model: claude-typo", "type": "invalid_request_error", "param": None, "code": None}, + id="anthropic_error_type_follows_the_status", + ), + pytest.param( + {"error": {"code": 404, "message": "Publisher model not found.", "status": "NOT_FOUND"}}, + 404, + {"message": "Publisher model not found.", "type": "invalid_request_error", "param": None, "code": None}, + id="google_error_message_unwrapped", + ), + pytest.param( + "litellm.APIError: Google Gen AI Studio API error: 404 - " + + json.dumps( + {"error": {"code": 404, "message": "models/gemini-typo is not found.", "status": "NOT_FOUND"}} + ), + 404, + { + "message": "models/gemini-typo is not found.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="prefixed_gemini_error_message_unwrapped", + ), + pytest.param( + "litellm.BadRequestError: AzureException - " + + json.dumps( + { + "error": { + "message": "Unrecognized request argument supplied: foo", + "type": "invalid_request_error", + "param": "foo", + "code": "unknown_parameter", + } + } + ), + 400, + { + "message": "Unrecognized request argument supplied: foo", + "type": "invalid_request_error", + "param": "foo", + "code": "unknown_parameter", + }, + id="prefixed_openai_error_passes_through", + ), + pytest.param( + {"message": "The provided model doesn't support counting tokens."}, + 400, + { + "message": "The provided model doesn't support counting tokens.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="bedrock_error_message_unwrapped", + ), + pytest.param( + "BedrockException - " + json.dumps({"message": "The provided model doesn't support counting tokens."}), + 400, + { + "message": "The provided model doesn't support counting tokens.", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="prefixed_bedrock_error_message_unwrapped", + ), + pytest.param( + "tool schema {unclosed is invalid", + 400, + { + "message": "tool schema {unclosed is invalid", + "type": "invalid_request_error", + "param": None, + "code": None, + }, + id="plain_text_with_a_brace_kept_whole", + ), + pytest.param( + "upstream count service unavailable", + 503, + { + "message": "upstream count service unavailable", + "type": "internal_server_error", + "param": None, + "code": None, + }, + id="plain_text_error", + ), + ], + ) + def test_provider_count_error_answers_in_openai_error_shape( + self, + provider_body: str | Mapping[str, object], + status_code: int, + expected_error: Mapping[str, str | None], + ) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message=provider_body if isinstance(provider_body, str) else json.dumps(provider_body), + type="token_counting_error", + param="model", + code=str(status_code), + ) + ) + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == status_code, response.text + assert response.json() == {"error": expected_error} + + def test_non_provider_count_error_keeps_its_own_type(self) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message="Token counting is disabled and no provider API result available", + type="token_counting_disabled", + param="model", + code=503, + ) + ) + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == 503, response.text + assert response.json()["error"]["type"] == "token_counting_disabled" + + def test_provider_count_error_names_the_call_id_when_opted_in(self) -> None: + from litellm.proxy._types import ProxyException + + failing_counter: Final = AsyncMock( + side_effect=ProxyException( + message="LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.", + type="token_counting_error", + param="model", + code="400", + headers={"x-litellm-call-id": "call-9564"}, + ) + ) + with patch("litellm.proxy.proxy_server.general_settings", {"include_call_id_in_error_body": True}): + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas.", + "type": "invalid_request_error", + "param": None, + "code": None, + "litellm_call_id": "call-9564", + } + } + class TestGuardrailBlockedResponsesShape: """A pre_call block raises ModifyResponseException before any provider call. diff --git a/tests/unit/proxy/test_custom_tokenizer_bug.py b/tests/unit/proxy/test_custom_tokenizer_bug.py index 23b6b302202..29625ea1a32 100644 --- a/tests/unit/proxy/test_custom_tokenizer_bug.py +++ b/tests/unit/proxy/test_custom_tokenizer_bug.py @@ -9,6 +9,7 @@ The HuggingFace download boundary (Tokenizer.from_pretrained) is mocked so these stay hermetic unit tests; the proxy's extraction-and-selection path runs for real. """ +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -17,7 +18,7 @@ import litellm import litellm.proxy.proxy_server import litellm.utils from litellm import Router -from litellm.proxy._types import TokenCountRequest +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.proxy_server import token_counter @@ -102,3 +103,60 @@ async def test_model_without_custom_tokenizer_uses_default(monkeypatch): assert response.tokenizer_type == "openai_tokenizer" assert response.model_used == "gpt-4" assert response.total_tokens > 0 + + + +def _vertex_glm_router_with_its_own_tokenizer() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-glm", + "litellm_params": { + "model": "vertex_ai/zai-org/glm-5.2-maas", + "vertex_project": "test-project", + "vertex_location": "global", + }, + "model_info": { + "custom_tokenizer": {"identifier": "zai-org/GLM-5.2", "revision": "main", "auth_token": None}, + }, + } + ] + ) + + +@pytest.mark.asyncio +async def test_custom_tokenizer_counts_a_vertex_partner_model_without_a_count_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) + + with patch.object(litellm.utils, "tokenizer_dispatch") as mock_tokenizer_cls: + mock_tokenizer_cls.from_pretrained.return_value = _fake_hf_tokenizer(9) + + response: Final = await token_counter( + request=TokenCountRequest(model="vertex-glm", messages=[{"role": "user", "content": "Return exactly OK."}]), + call_endpoint=True, + ) + + mock_tokenizer_cls.from_pretrained.assert_called_once_with("zai-org/GLM-5.2", revision="main", token=None) + assert response.error is not True + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens >= 9 + + +@pytest.mark.asyncio +async def test_contents_only_count_of_a_vertex_partner_model_without_a_count_route_answers_400( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _vertex_glm_router_with_its_own_tokenizer()) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest( + model="vertex-glm", contents=[{"role": "user", "parts": [{"text": "Return exactly OK."}]}] + ), + call_endpoint=True, + ) + + assert raised.value.code == "400" + assert raised.value.message == "LiteLLM does not support counting tokens for Vertex AI model zai-org/glm-5.2-maas." From 1e3641a52a6ce0566988bfded602a128d242838f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 14:16:21 -0700 Subject: [PATCH 47/64] fix(router): read an application profile's base model from metadata on a cache hit (#47144) * fix(router): read an application profile's base model from metadata on a cache hit The 1h prompt-cache pin check read base_model only from the top level of litellm_params. On a response-cache hit the params it sees carry the base model only under metadata.model_info, so an application-profile ARN resolved to itself and the pin was written for 5 minutes while the wire carried a 1h TTL. Read it through get_base_model_from_metadata, which also checks metadata.model_info * test(router): pin top-level base_model precedence over metadata model_info --- .../prompt_caching_deployment_check.py | 10 ++++- .../test_prompt_caching_deployment_check.py | 42 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index b0da3ed264a..b2f0c658161 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -21,7 +21,11 @@ from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload -from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt +from litellm.utils import ( + get_base_model_from_metadata, + get_prompt_cache_min_tokens, + is_prompt_caching_valid_prompt, +) from ..prompt_caching_cache import PromptCachingCache @@ -61,7 +65,9 @@ def _honors_one_hour_cache_ttl( ) -> bool: if standard_logging_object["custom_llm_provider"] == "bedrock": deployment_model: Final = standard_logging_object["model"] - capability_model: Final = bedrock_capability_model(deployment_model, litellm_params.get("base_model")) + capability_model: Final = bedrock_capability_model( + deployment_model, get_base_model_from_metadata({"litellm_params": litellm_params}) + ) return bedrock_cache_point_ttl_supported(capability_model, deployment_model) model_map_value: Final = standard_logging_object["model_map_information"]["model_map_value"] return model_map_value is not None and (model_map_value.get("cache_creation_input_token_cost_above_1hr") or 0) > 0 diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index e5565a3df16..5f018ab5dda 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -143,6 +143,48 @@ def test_bedrock_application_profile_uses_its_base_model_for_one_hour_ttl() -> N ) is True +def test_bedrock_application_profile_cache_hit_reads_base_model_from_metadata_model_info() -> None: + standard_logging_object: Final = cast( + StandardLoggingPayload, + { + "custom_llm_provider": "bedrock", + "model": "bedrock/converse/arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/pinprofile", + }, + ) + assert ( + _honors_one_hour_cache_ttl( + standard_logging_object, + {"metadata": {"model_info": {"base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}}}, + ) + is True + ) + + +def test_bedrock_application_profile_top_level_base_model_wins_over_metadata_model_info() -> None: + standard_logging_object: Final = cast( + StandardLoggingPayload, + { + "custom_llm_provider": "bedrock", + "model": "bedrock/converse/arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/pinprofile", + }, + ) + top_level: Final = "bedrock/anthropic.claude-3-haiku-20240307-v1:0" + in_metadata: Final = "bedrock/us.anthropic.claude-sonnet-5-5" + top_level_alone: Final = _honors_one_hour_cache_ttl(standard_logging_object, {"base_model": top_level}) + metadata_alone: Final = _honors_one_hour_cache_ttl( + standard_logging_object, {"metadata": {"model_info": {"base_model": in_metadata}}} + ) + assert top_level_alone != metadata_alone + + assert ( + _honors_one_hour_cache_ttl( + standard_logging_object, + {"base_model": top_level, "metadata": {"model_info": {"base_model": in_metadata}}}, + ) + is top_level_alone + ) + + @pytest.mark.asyncio async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minimum(): """ From b3c0f3457df5e8158842b758aff346cc53ab998c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 17 Oct 2026 14:17:46 -0700 Subject: [PATCH 48/64] fix: handle missing choice indices in streamed output recording (#46947) * fix: handle missing choice indices in streamed output recording * test: verify recorded partial output on client disconnect * test: correlate streaming regressions with persisted request IDs * test: build progress-stream configs without mutation --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../litellm_core_utils/served_output_texts.py | 4 +- .../observability/test_guardrail_effects.py | 279 +++++++++++++++++- .../test_served_output_texts.py | 27 ++ 3 files changed, 307 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/served_output_texts.py b/litellm/litellm_core_utils/served_output_texts.py index ce9a568ea72..e23db5200a3 100644 --- a/litellm/litellm_core_utils/served_output_texts.py +++ b/litellm/litellm_core_utils/served_output_texts.py @@ -126,13 +126,13 @@ def _chat_stream_choice_text(chunks: Sequence[ModelResponseStream], index: int) content for chunk in chunks for choice in chunk.choices - if choice.index == index and isinstance(content := choice.delta.content, str) + if (choice.index or 0) == index and isinstance(content := choice.delta.content, str) ) return "".join(contents) if contents else None def _chat_stream_texts(chunks: Sequence[ModelResponseStream]) -> ServedTexts | None: - choice_count: Final = max((choice.index + 1 for chunk in chunks for choice in chunk.choices), default=0) + choice_count: Final = max(((choice.index or 0) + 1 for chunk in chunks for choice in chunk.choices), default=0) texts: Final = tuple(_chat_stream_choice_text(chunks, index) for index in range(choice_count)) return texts if any(text is not None for text in texts) else None diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index f267cff04f3..4e0c3d45cdf 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -8,6 +8,7 @@ from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from contextlib import contextmanager from dataclasses import dataclass +from itertools import chain from pathlib import Path from types import MappingProxyType from typing import Final @@ -22,7 +23,8 @@ from integration._support.mcp import mcp_peer, register_mcp, tool_names from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port -from openai import AsyncOpenAI, OpenAI +from openai import APIError, AsyncOpenAI, OpenAI +from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, JsonValue @@ -2858,6 +2860,281 @@ def test_a_fully_read_stream_is_scanned_exactly_once( assert len(rig.scans.response_scans(rig.secret())) == 1, rig.scans.seen +def _pre_call_progress_config(tmp_path: Path, identity: str, policy_url: str, *, store_prompts: bool) -> Path: + path: Final = _post_call_config(tmp_path, identity, policy_url, {**_END_OF_STREAM_ONLY, "mode": "pre_call"}, True) + if store_prompts: + config: Final = yaml.safe_load(path.read_text()) + settings: Final = {**config["general_settings"], "store_prompts_in_spend_logs": True} + path.write_text(yaml.safe_dump({**config, "general_settings": settings})) + return path + + +def _openai_chat_chunks(candidate: Gateway, model: str, prompt: str) -> tuple[ChatCompletionChunk, ...]: + with OpenAI( + base_url=str(candidate.client.base_url) + "/v1", api_key=candidate.key, max_retries=0, timeout=30 + ) as client: + return tuple( + client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + stream=True, + extra_headers={"x-litellm-call-id": prompt}, + ) + ) + + +def _async_openai_chat_chunks(candidate: Gateway, model: str, prompt: str) -> tuple[ChatCompletionChunk, ...]: + async def consume() -> tuple[ChatCompletionChunk, ...]: + async with AsyncOpenAI( + base_url=str(candidate.client.base_url) + "/v1", api_key=candidate.key, max_retries=0, timeout=30 + ) as client: + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + stream=True, + extra_headers={"x-litellm-call-id": prompt}, + ) + return tuple([chunk async for chunk in stream]) + + return asyncio.run(consume()) + + +_PROGRESS_STREAM_READERS: Final = ( + pytest.param(_openai_chat_chunks, id="sync-openai-sdk"), + pytest.param(_async_openai_chat_chunks, id="async-openai-sdk"), +) + + +@pytest.mark.parametrize("read_stream", _PROGRESS_STREAM_READERS) +@pytest.mark.parametrize("index_fields", ({}, {"index": None}), ids=("omitted-index", "null-index")) +@pytest.mark.parametrize("progress_content", (None, "progress "), ids=("null-content", "text-content")) +def test_progress_chunk_without_choice_index_streams_and_logs_served_answer( + gateway: Gateway, + tmp_path: Path, + record_property: Callable[[str, object], None], + index_fields: dict[str, JsonValue], + progress_content: str | None, + read_stream: Callable[[Gateway, str, str], tuple[ChatCompletionChunk, ...]], +) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + record_property("litellm_call_id", identity) + answer: Final = "synthetic answer " + identity + reasoning: Final = "backend progress " + identity + expected: Final = (progress_content or "") + answer + frames: Final = ( + _chat_frame( + identity, + ( + { + **index_fields, + "delta": {"reasoning_content": reasoning, "content": progress_content}, + "finish_reason": None, + }, + ), + ), + _chat_frame(identity, (_chat_choice(0, {"content": answer}),)), + _chat_frame(identity, (_chat_choice(0, {}, "stop"),)), + b"data: [DONE]\n\n", + ) + + def provider(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + body: Final = json.loads(request.body) + assert body["stream"] is True, body + assert body["messages"] == [{"role": "user", "content": identity}], body + return Reply(content_type="text/event-stream", chunks=frames) + + with wire_server(_allowing_guardrail) as policy, wire_server(provider) as upstream: + path: Final = _pre_call_progress_config(tmp_path, identity, policy.url, store_prompts=True) + with ( + owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + model: Final = scenario.model(model_info={"id": identity}, api_base=upstream.url + "/v1") + eventually(lambda: owned.serves_only_on_every_worker(model, identity), bool, seconds=70) + chunks: Final = read_stream(candidate, model, identity) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + actual: Final = "".join(choice.delta.content or "" for choice in choices) + assert actual == expected, chunks + assert tuple(choice.index for choice in choices if choice.finish_reason == "stop") == (0,), chunks + assert reasoning in json.dumps([chunk.model_dump() for chunk in chunks]), chunks + assert len(upstream.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, response, metadata FROM "LiteLLM_SpendLogs" ' + "WHERE model_group=%s AND (request_id=%s OR metadata->>'litellm_call_id'=%s)", + (model, identity, identity), + ), + lambda values: len(values) == 1, + seconds=70, + ) + record_property("spend_request_id", rows[0]["request_id"]) + assert rows[0]["status"] == "success", rows + logged_choices: Final = object_value(rows[0]["response"])["choices"] + assert isinstance(logged_choices, list), rows + assert tuple(object_value(object_value(choice)["message"])["content"] for choice in logged_choices) == ( + expected, + ), rows + entries: Final = object_value(rows[0]["metadata"])["guardrail_information"] + assert isinstance(entries, list), rows + assert tuple( + (entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"]) + for entry in map(object_value, entries) + ) == ((identity, "pre_call", "success"),), rows + scans: Final = tuple(json.loads(request.body) for request in policy.drain()) + assert len(scans) == 1 and scans[0]["input_type"] == "request", scans + assert scans[0]["texts"] == [identity], scans + + +@pytest.mark.parametrize("read_stream", _PROGRESS_STREAM_READERS) +def test_progress_chunk_without_choice_index_preserves_upstream_stream_error( + gateway: Gateway, + tmp_path: Path, + record_property: Callable[[str, object], None], + read_stream: Callable[[Gateway, str, str], tuple[ChatCompletionChunk, ...]], +) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + record_property("litellm_call_id", identity) + message: Final = "synthetic upstream failure marker " + identity + frames: Final = ( + _chat_frame( + identity, + ({"delta": {"reasoning_content": "backend progress", "content": None}, "finish_reason": None},), + ), + b"data: " + + json.dumps({"error": {"message": message, "type": "server_error", "code": "upstream_failed"}}).encode() + + b"\n\n", + b"data: [DONE]\n\n", + ) + + def provider(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + body: Final = json.loads(request.body) + assert body["stream"] is True, body + assert body["messages"] == [{"role": "user", "content": identity}], body + return Reply(content_type="text/event-stream", chunks=frames) + + with wire_server(_allowing_guardrail) as policy, wire_server(provider) as upstream: + path: Final = _pre_call_progress_config(tmp_path, identity, policy.url, store_prompts=False) + with ( + owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + model: Final = scenario.model(model_info={"id": identity}, api_base=upstream.url + "/v1", num_retries=0) + eventually(lambda: owned.serves_only_on_every_worker(model, identity), bool, seconds=70) + with pytest.raises(APIError, match=message) as raised: + read_stream(candidate, model, identity) + assert "NoneType" not in str(raised.value), raised.value + assert len(upstream.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, metadata FROM "LiteLLM_SpendLogs" ' + "WHERE model_group=%s AND (request_id=%s OR metadata->>'litellm_call_id'=%s)", + (model, identity, identity), + ), + lambda values: len(values) == 1, + seconds=70, + ) + record_property("spend_request_id", rows[0]["request_id"]) + assert rows[0]["status"] == "failure", rows + error: Final = object_value(rows[0]["metadata"])["error_information"] + assert message in str(object_value(error)["error_message"]), rows + scans: Final = tuple(json.loads(request.body) for request in policy.drain()) + assert len(scans) == 1 and scans[0]["texts"] == [identity], scans + + +@pytest.mark.parametrize("index_fields", ({}, {"index": None}), ids=("omitted-index", "null-index")) +def test_disconnect_after_progress_without_choice_index_logs_only_the_received_prefix( + gateway: Gateway, + tmp_path: Path, + index_fields: dict[str, JsonValue], + record_property: Callable[[str, object], None], +) -> None: + identity: Final = "guardrail" + uuid.uuid4().hex + record_property("litellm_call_id", identity) + prefix: Final = "synthetic received prefix " + identity + tail: Final = "synthetic withheld tail " + identity + gate: Final = threading.Event() + frames: Final = ( + _chat_frame( + identity, + ( + { + **index_fields, + "delta": {"reasoning_content": "backend progress", "content": prefix}, + "finish_reason": None, + }, + ), + ), + _chat_frame(identity, (_chat_choice(0, {"content": tail}, "stop"),)), + b"data: [DONE]\n\n", + ) + + def provider(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + body: Final = json.loads(request.body) + assert body["stream"] is True, body + assert body["messages"] == [{"role": "user", "content": identity}], body + return Reply(content_type="text/event-stream", chunks=frames, gate_after_first=gate) + + with wire_server(_allowing_guardrail) as policy, wire_server(provider) as upstream: + path: Final = _pre_call_progress_config(tmp_path, identity, policy.url, store_prompts=True) + with ( + owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + model: Final = scenario.model(model_info={"id": identity}, api_base=upstream.url + "/v1", num_retries=0) + eventually(lambda: owned.serves_only_on_every_worker(model, identity), bool, seconds=70) + + try: + + async def disconnect() -> str: + async with AsyncOpenAI( + base_url=str(candidate.client.base_url) + "/v1", + api_key=candidate.key, + max_retries=0, + timeout=30, + ) as client: + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": identity}], + stream=True, + extra_headers={"x-litellm-call-id": identity}, + ) + async for chunk in stream: + for choice in chunk.choices: + if choice.delta.content == prefix: + await stream.close() + return choice.delta.content + raise AssertionError("The stream ended before the SDK received the provider prefix") + + assert asyncio.run(disconnect()) == prefix + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, response FROM "LiteLLM_SpendLogs" ' + "WHERE model_group=%s AND (request_id=%s OR metadata->>'litellm_call_id'=%s)", + (model, identity, identity), + ), + lambda values: len(values) == 1, + seconds=70, + ) + record_property("spend_request_id", rows[0]["request_id"]) + assert rows[0]["status"] == "success", rows + logged_choices: Final = object_value(rows[0]["response"])["choices"] + assert isinstance(logged_choices, list), rows + assert tuple(object_value(object_value(choice)["message"])["content"] for choice in logged_choices) == ( + prefix, + ), rows + assert len(upstream.drain()) == 1 + scans: Final = tuple(json.loads(request.body) for request in policy.drain()) + assert len(scans) == 1 and scans[0]["texts"] == [identity], scans + finally: + gate.set() + + def _cached_twin_rows(rig: _DisconnectRig) -> tuple[tuple[str, ...], ...]: first: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig, 0, stream=False)) second: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig, 0, stream=False)) diff --git a/tests/unit/litellm_core_utils/test_served_output_texts.py b/tests/unit/litellm_core_utils/test_served_output_texts.py index ade39b8fca9..637a20571c8 100644 --- a/tests/unit/litellm_core_utils/test_served_output_texts.py +++ b/tests/unit/litellm_core_utils/test_served_output_texts.py @@ -1,3 +1,7 @@ +from typing import Final + +import pytest + from litellm.litellm_core_utils.served_output_texts import ( SERVED_OUTPUT_TEXTS_KEY, overlay_served_output_texts, @@ -79,6 +83,29 @@ def test_served_stream_output_texts_keeps_every_choice_index_when_chunks_carry_o assert served_stream_output_texts(chunks) == ("first choice", MASKED) +@pytest.mark.parametrize("content", (None, "", "first")) +def test_null_choice_index_joins_choice_zero_without_changing_other_choices(content: str | None) -> None: + chunks: Final = ( + _stream_chunk("other", 1), + ModelResponseStream( + choices=[StreamingChoices(index=None, delta=Delta(content=content, reasoning_content="loading"))] + ), + _stream_chunk(" answer"), + ) + assert served_stream_output_texts(chunks) == ((content or "") + " answer", "other") + assert chunks[1].choices[0].index is None + + +@pytest.mark.parametrize("content, expected", ((None, None), ("", ("",)), ("answer", ("answer",)))) +def test_null_choice_index_preserves_empty_text_and_ignores_reasoning( + content: str | None, expected: tuple[str, ...] | None +) -> None: + chunk: Final = ModelResponseStream( + choices=[StreamingChoices(index=None, delta=Delta(content=content, reasoning_content="loading"))] + ) + assert served_stream_output_texts((chunk,)) == expected + + def test_blanked_output_is_served_as_empty_text_and_overlaid(): assert served_output_texts(_chat_response("")) == ("",) assert served_output_texts({"type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}]}) == ( From 44fcbc30b07792aceb24e312ea5c59e5cfbdb8a9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 14:29:59 -0700 Subject: [PATCH 49/64] test(mcp): replace the live Linear MCP OAuth e2e tests with integration coverage (#47154) * test(mcp): move the MCP OAuth happy path from a live Linear e2e to an integration test GitHub Actions now runs unit tests only, and the MCP OAuth happy path needed a captured Linear browser login to click through a real consent page. The OAuth flows it exercised are already integration tests against an owned authorization-server double, except one behavior: the stored per-user token must be served by a proxy process that never saw the consent, with no shared Redis to fall back on. That case is now in test_mcp_oauth_flows.py The e2e test, its owned-gateway helper, its workflow, the mcp_oauth_live marker and the JUnit gate script only that workflow used are removed * test(mcp): chat completions run an OAuth server tool with the caller's own stored token * test(e2e): drop the Linear chat-completion MCP OAuth e2e and its helpers The behavior it checked, a chat completion running an OAuth server's tool with the caller's stored per-user token, is now an integration test in test_mcp_llm_endpoints.py that needs no saved Linear browser login * test(mcp): decode the stored per-user credential as urlsafe base64 The plaintext check used b64decode, which drops the - and _ of a urlsafe value and raised on roughly one ciphertext in five * test(mcp): drop the chat OAuth integration test that duplicates test_mcp_slots F2-chat security/test_mcp_slots.py already sends a per-user OAuth token through /v1/chat/completions and asserts the peer receives it. Every mutant the new test killed was also killed by the existing suite, so it added no signal --- .github/e2e-stack/assert_tests_ran.py | 69 --- .github/workflows/test-linting.yml | 4 +- .github/workflows/test-mcp-oauth-e2e.yml | 180 ------- .../check_e2e_no_raw_requests.py | 1 - .../test_e2e_changed_gate.py | 156 ------- .../test_e2e_junit_report.py | 16 +- tests/e2e/AGENTS.md | 8 +- tests/e2e/CONTRIBUTING.md | 52 --- tests/e2e/conftest.py | 16 - tests/e2e/coverage_registry/mcp.yaml | 8 - tests/e2e/e2e_config.py | 5 - tests/e2e/mcp/linear_session_capture.py | 58 --- tests/e2e/mcp/oauth_chat_client.py | 442 ------------------ tests/e2e/mcp/oauth_gateway.py | 197 -------- .../mcp/test_mcp_chat_completion_oauth_e2e.py | 218 --------- .../e2e/mcp/test_mcp_oauth_happy_path_e2e.py | 216 --------- tests/e2e/models.py | 38 +- tests/e2e/pytest.ini | 1 - tests/integration/mcp/test_mcp_oauth_flows.py | 48 ++ tests/unit/test_circleci_path_filter.py | 2 +- 20 files changed, 61 insertions(+), 1674 deletions(-) delete mode 100644 .github/e2e-stack/assert_tests_ran.py delete mode 100644 .github/workflows/test-mcp-oauth-e2e.yml delete mode 100644 tests/code_coverage_tests/test_e2e_changed_gate.py delete mode 100644 tests/e2e/mcp/linear_session_capture.py delete mode 100644 tests/e2e/mcp/oauth_chat_client.py delete mode 100644 tests/e2e/mcp/oauth_gateway.py delete mode 100644 tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py delete mode 100644 tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py deleted file mode 100644 index 3af49007b1e..00000000000 --- a/.github/e2e-stack/assert_tests_ran.py +++ /dev/null @@ -1,69 +0,0 @@ -import os -import re -import sys -import xml.etree.ElementTree as ET -from pathlib import Path -from typing import Final - -sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e")) -from coverage_registry.management_cases import MANAGEMENT_CASES - - -def main() -> int: - selected: Final = tuple(sys.argv[2:]) - try: - report: Final = ET.parse(Path(sys.argv[1])).getroot() - except (ET.ParseError, OSError): - _ = sys.stdout.write("::error::could not read the test execution report\n") - return 1 - cases: Final = tuple(report.iter("testcase")) - expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT") - passed: Final = frozenset( - case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) - ) - missing: Final = tuple(path for path in selected if path not in passed) - required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected) - passed_nodes: Final = frozenset( - prop.get("value") - for case in cases - if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) - for prop in case.findall("./properties/property") - if prop.get("name") == "management_node" - ) - missing_nodes: Final = required_nodes - passed_nodes - for node in sorted(missing_nodes): - _ = sys.stdout.write(f"::error::required management case did not pass: {node}\n") - for path in selected: - collected: Final = sum(case.get("file") == path for case in cases) - skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases) - _ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n") - for case in cases: - if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")): - continue - outcome = "skipped" if case.find("skipped") is not None else "failed" - _ = sys.stdout.write(f" {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n") - for prop in case.findall("./properties/property"): - name = prop.get("name", "") - value = prop.get("value", "") - if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch( - r"[A-Za-z0-9_.:<>-]{1,240}", value - ): - _ = sys.stdout.write(f" {name}: {value}\n") - if expected_count is not None and ( - len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases) - ): - _ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n") - return 1 - if ( - selected - and not missing - and not missing_nodes - and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error")) - ): - return 0 - _ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n") - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 4c8aeb64e85..3b9a2614725 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -539,8 +539,8 @@ jobs: - name: test_unit_passed_gate run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_unit_passed_gate.py - - name: test_e2e_changed_gate - run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py + - name: test_e2e_idp_stack + run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_idp_stack.py - name: test_e2e_metadata env: diff --git a/.github/workflows/test-mcp-oauth-e2e.yml b/.github/workflows/test-mcp-oauth-e2e.yml deleted file mode 100644 index 5c07a68714e..00000000000 --- a/.github/workflows/test-mcp-oauth-e2e.yml +++ /dev/null @@ -1,180 +0,0 @@ -name: MCP OAuth happy path - -on: - pull_request: - paths: - - '.github/workflows/test-mcp-oauth-e2e.yml' - - '.github/e2e-stack/**' - - 'tests/e2e/*.py' - - 'tests/e2e/pytest.ini' - - 'tests/e2e/idp_realm.json' - - 'tests/e2e/mcp/**' - - 'litellm/experimental_mcp_client/**' - - 'litellm/proxy/_experimental/mcp_server/**' - - 'litellm/proxy/auth/**' - - 'litellm/proxy/management_endpoints/mcp_management_endpoints.py' - - 'litellm/proxy/_types.py' - - 'litellm/types/mcp_server/mcp_server_manager.py' - - 'litellm/proxy/management_endpoints/*sso*.py' - - 'litellm/proxy/management_endpoints/sso/**' - - 'litellm/proxy/common_utils/encrypt_decrypt_utils.py' - - 'litellm/proxy/proxy_server.py' - - 'litellm/proxy/schema.prisma' - - 'ui/litellm-dashboard/src/app/connect/**' - - 'ui/litellm-dashboard/src/app/mcp/oauth/**' - - 'pyproject.toml' - - 'uv.lock' - workflow_dispatch: - -permissions: {} - -concurrency: - group: mcp-oauth-${{ github.ref }} - cancel-in-progress: true - -jobs: - oauth: - if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository - runs-on: ubuntu-latest - environment: e2e-changed - timeout-minutes: 45 - permissions: - contents: read - id-token: write - services: - postgres: - image: postgres:16.6 - env: - POSTGRES_USER: litellm - POSTGRES_PASSWORD: dbpassword9090 - POSTGRES_DB: litellm - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready -U litellm" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - env: - DATABASE_HOST: 127.0.0.1 - DATABASE_PORT: '5432' - DATABASE_USER: litellm - DATABASE_PASSWORD: dbpassword9090 - DATABASE_NAME: litellm - DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm - E2E_KEYCLOAK_URL: http://127.0.0.1:8081 - E2E_KEYCLOAK_ADMIN_USER: admin - E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret - E2E_FIXTURE_MODE: live - E2E_PROVIDER_CACHE: '0' - E2E_MCP_OAUTH_LIVE: '1' - E2E_REQUIRED_TEST_COUNT: '4' - steps: - - name: Checkout the tested source - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - ref: ${{ github.event.pull_request.head.sha || github.sha }} - persist-credentials: false - - - name: Require and materialize the upstream login - env: - STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }} - run: | - umask 077 - python3 - <<'PY' - import base64 - import json - import os - import secrets - from pathlib import Path - encoded = os.environ.get("STORAGE_STATE", "") - if not encoded: - raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login") - state = json.loads(base64.b64decode(encoded, validate=True)) - if not isinstance(state, dict) or not state.get("cookies"): - raise SystemExit("The captured login must contain browser cookies") - directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private" - directory.mkdir(mode=0o700) - path = directory / "linear-state.json" - path.write_text(json.dumps(state)) - with open(os.environ["GITHUB_ENV"], "a") as output: - output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n") - for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"): - value = "sk-e2e-" + secrets.token_hex(24) - print(f"::add-mask::{value}") - output.write(f"{name}={value}\n") - PY - - - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: '3.13' - - uses: ./.github/actions/setup-uv-with-retries - with: - version: '0.10.9' - - uses: ./.github/actions/cache-cargo-build - - name: Install the frozen E2E environment - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev - uv run --no-sync python scripts/prisma_generate_if_needed.py - uv run --no-sync playwright install --with-deps chromium - - - name: Configure license access - id: aws - uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0 - with: - role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }} - aws-region: us-east-1 - role-session-name: mcp-oauth-${{ github.run_id }} - role-duration-seconds: 900 - output-env-credentials: false - output-credentials: true - - name: Load the E2E license - env: - AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }} - AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }} - AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }} - AWS_DEFAULT_REGION: us-east-1 - run: | - license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)" - test -n "${license}" - echo "::add-mask::${license}" - echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}" - - - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0 - with: - node-version-file: ui/litellm-dashboard/.nvmrc - - name: Build the gateway consent UI at the tested commit - run: | - cd ui/litellm-dashboard - ../../scripts/with_dashboard_node.sh npm ci - ../../scripts/with_dashboard_node.sh npm run build - mkdir -p ../../litellm/proxy/_experimental/out - cp -r out/. ../../litellm/proxy/_experimental/out/ - find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do - mkdir -p "${page%.html}" - mv "${page}" "${page%.html}/index.html" - done - - - name: Prepare the isolated database and IdP - run: | - umask 077 - bash .github/e2e-stack/start-idp.sh - uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1 - - - name: Run every required OAuth variant without retries - run: | - umask 077 - uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \ - --rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \ - --junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \ - > "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1 - - name: Report JUnit results and reject skipped or missing cases - if: always() - run: | - uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \ - "${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py - - name: Remove private login and logs - if: always() - run: | - docker rm -f e2e-keycloak >/dev/null 2>&1 || true - rm -rf "${RUNNER_TEMP}/mcp-oauth-private" diff --git a/tests/code_coverage_tests/check_e2e_no_raw_requests.py b/tests/code_coverage_tests/check_e2e_no_raw_requests.py index 4316f1b57b6..bea177d3b7d 100644 --- a/tests/code_coverage_tests/check_e2e_no_raw_requests.py +++ b/tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -25,7 +25,6 @@ ALLOWED_RAW_CLIENT_FILES = { "e2e/e2e_http.py": ("requests",), "e2e/conftest.py": ("requests",), "e2e/claude_code/pr_gate_version_resolver.py": ("urllib.request",), - "e2e/mcp/oauth_chat_client.py": ("httpx",), } EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"}) diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py deleted file mode 100644 index c4d12d1282c..00000000000 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ /dev/null @@ -1,156 +0,0 @@ -import os -import subprocess -import sys -import xml.etree.ElementTree as ET -from pathlib import Path -from typing import Final - -import pytest - -GATE: Final = Path(__file__).resolve().parents[2] / ".github/e2e-stack/assert_tests_ran.py" -SELECTED: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py") - - -@pytest.mark.parametrize( - ("second_outcome", "expected_status"), - (("passed", 0), ("skipped", 1), ("failure", 1), ("error", 1), ("deselected", 1)), -) -def test_each_changed_file_must_run(tmp_path: Path, second_outcome: str, expected_status: int) -> None: - suite: Final = ET.Element("testsuite") - _ = ET.SubElement(suite, "testcase", file=SELECTED[0]) - if second_outcome != "deselected": - second: Final = ET.SubElement(suite, "testcase", file=SELECTED[1]) - if second_outcome != "passed": - _ = ET.SubElement(second, second_outcome) - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - - result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) - - assert result.returncode == expected_status, result.stdout - - -@pytest.mark.parametrize("outcome", ("failure", "error")) -def test_passing_case_does_not_hide_a_failure_in_the_same_file(tmp_path: Path, outcome: str) -> None: - suite: Final = ET.Element("testsuite") - _ = ET.SubElement(suite, "testcase", file=SELECTED[0]) - failed: Final = ET.SubElement(suite, "testcase", file=SELECTED[0]) - _ = ET.SubElement(failed, outcome) - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - - result: Final = subprocess.run( - [sys.executable, str(GATE), str(report), SELECTED[0]], capture_output=True, text=True - ) - - assert result.returncode == 1 - - -def test_failed_cases_are_named_per_selected_file(tmp_path: Path) -> None: - suite: Final = ET.Element("testsuite") - _ = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_ok") - failed: Final = ET.SubElement( - suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_boom" - ) - _ = ET.SubElement(failed, "failure", message="secret-bearing message") - errored: Final = ET.SubElement( - suite, "testcase", file=SELECTED[1], classname="tests.e2e.access_control.test_b", name="test_setup" - ) - _ = ET.SubElement(errored, "error") - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - - result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) - - assert result.returncode == 1 - assert " failed: tests.e2e.access_control.test_a::test_boom\n" in result.stdout - assert " failed: tests.e2e.access_control.test_b::test_setup\n" in result.stdout - assert "test_ok" not in result.stdout - assert "secret-bearing message" not in result.stdout - - -@pytest.mark.parametrize("contents", ("", "')) -def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None: - report: Final = tmp_path / "report.xml" - _ = report.write_text(contents) - - result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) - - assert result.returncode == 1 - - -@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer")) -def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None: - suite: Final = ET.Element("testsuite") - path: Final = "tests/e2e/management/test_jwt_management_e2e.py" - case: Final = ET.SubElement(suite, "testcase", file=path) - properties: Final = ET.SubElement(case, "properties") - _ = ET.SubElement( - properties, - "property", - name="management_node", - value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]", - ) - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True) - assert result.returncode == 1 - assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout - - -@pytest.mark.parametrize("phase", ("setup", "call", "teardown")) -@pytest.mark.parametrize("required_count", ("1", "4")) -def test_oauth_failure_diagnostics_do_not_publish_private_payloads( - tmp_path: Path, phase: str, required_count: str -) -> None: - suite: Final = ET.Element("testsuite") - case: Final = ET.SubElement(suite, "testcase", file=SELECTED[0]) - private: Final = "private-token-in-exception-message" - failure: Final = ET.SubElement(case, "failure", message=private) - failure.text = private - properties: Final = ET.SubElement(case, "properties") - for name, value in ( - ("oauth_failure_phase", phase), - ("oauth_exception_type", "AssertionError"), - ("oauth_frame", "oauth_gateway.py:120:start"), - ("oauth_frame", f"injected\\n{private}"), - ("unrelated_property", private), - ): - _ = ET.SubElement(properties, "property", name=name, value=value) - report: Final = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - result: Final = subprocess.run( - [sys.executable, "-I", str(GATE), str(report), SELECTED[0]], - capture_output=True, - text=True, - env={**os.environ, "E2E_REQUIRED_TEST_COUNT": required_count}, - ) - assert result.returncode == 1 - assert f"oauth_failure_phase: {phase}" in result.stdout - assert "oauth_exception_type: AssertionError" in result.stdout - assert "oauth_frame: oauth_gateway.py:120:start" in result.stdout - assert private not in result.stdout + result.stderr - - -@pytest.mark.parametrize( - ("count", "skip", "expected"), ((0, False, 1), (3, False, 1), (4, False, 0), (5, False, 1), (4, True, 1)) -) -def test_required_count_reports_cases_before_rejecting(tmp_path: Path, count: int, skip: bool, expected: int) -> None: - suite = ET.Element("testsuite") - for index in range(count): - case = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="OAuth", name=f"variant{index}") - if skip and index == 0: - ET.SubElement(case, "skipped", message="private-skip-reason") - report = tmp_path / "report.xml" - ET.ElementTree(suite).write(report) - result = subprocess.run( - [sys.executable, "-I", str(GATE), str(report), SELECTED[0]], - env={**os.environ, "E2E_REQUIRED_TEST_COUNT": "4"}, - capture_output=True, - text=True, - ) - assert result.returncode == expected - assert f"{count} collected, {int(skip)} skipped" in result.stdout - if skip: - assert "skipped: OAuth::variant0" in result.stdout - assert "private-skip-reason" not in result.stdout + result.stderr diff --git a/tests/code_coverage_tests/test_e2e_junit_report.py b/tests/code_coverage_tests/test_e2e_junit_report.py index 140b9ba6dca..be5fcd245db 100644 --- a/tests/code_coverage_tests/test_e2e_junit_report.py +++ b/tests/code_coverage_tests/test_e2e_junit_report.py @@ -172,12 +172,7 @@ def open_consent() -> None: raise RuntimeError("consent page timed out") -@pytest.mark.mcp_oauth_live -def test_oauth_dies_on_consent() -> None: - open_consent() - - -def test_plain_dies_on_consent() -> None: +def test_dies_on_consent() -> None: open_consent() """ @@ -353,16 +348,13 @@ class TestStepsReachTheReport: def test_a_failed_phase_s_own_report_carries_the_steps(tmp_path: Path) -> None: """Plugins that read the failed setup or call report, not the teardown one - junitxml writes from, see where the test died too, oauth-live or not.""" + junitxml writes from, see where the test died too.""" write_suite(tmp_path, {"test_consent.py": FAILED_PHASE_SUITE, "report_spy.py": REPORT_SPY_PLUGIN}) - child: Final = run_child_pytest(tmp_path, "-p", "report_spy", env={"E2E_MCP_OAUTH_LIVE": "1"}) + child: Final = run_child_pytest(tmp_path, "-p", "report_spy") seen_path: Final = tmp_path / "failed-reports.jsonl" assert seen_path.exists(), f"no failed report reached the spy:\n{child.stdout}\n{child.stderr}" seen: Final = dict(map(FailedReport.validate_json, seen_path.read_text().splitlines())) - assert seen == { - "test_oauth_dies_on_consent": ("open the consent page",), - "test_plain_dies_on_consent": ("open the consent page",), - }, child.stdout + assert seen == {"test_dies_on_consent": ("open the consent page",)}, child.stdout class TestDeclaredPropertiesReachTheReport: diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 62ebfd12fe5..b32cfb25fa6 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -40,7 +40,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface, not Playwright) - `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) -- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below), plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token +- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below) - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) @@ -53,14 +53,14 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family ## MCP suite: real Datadog only -Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server, except the two Linear OAuth tests `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite +Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite - Register via `register_datadog_mcp` in `tests/e2e/mcp/datadog_mcp.py` (or extend that helper if you need a different `toolsets=` / `allowed_tools` slice of the same Datadog endpoint). That posts `/v1/mcp/server` with `url=datadog_mcp_url(...)` and static headers `DD-API-KEY` / `DD-APPLICATION-KEY` from the process env - Auth is Datadog's documented CI/header path, not a browser OAuth authorize/token dance. Hard-fail when `DD_API_KEY` or `DD_APP_KEY` is missing (`assert_dd_mcp_creds`). Never skip for a missing fake upstream - Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find. Dual-read with `dd_logs` from conftest when delivery matters - Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down - If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream. The default is always Datadog -- The two standing exceptions are `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so these tests drive a real Linear MCP server instead. They are still real remote upstreams, so the no-mock, no-fixture rule above holds unchanged. The direct OAuth test also uses the existing live provider edge to inspect forwarded headers without replay, and owns a separate source-built gateway for cold restarts +- Gateway-managed OAuth (an authorize/token dance and per-user stored tokens) is out of scope here, since Datadog only takes static headers. Those flows are integration tests against an owned authorization-server double, in `tests/integration/mcp/test_mcp_oauth_flows.py` and, for a chat completion calling the tool with the caller's stored token, the F2-chat cases of `tests/integration/security/test_mcp_slots.py`. Do not add a third-party OAuth upstream that needs a saved browser login ## Lay the pattern down in a class @@ -243,7 +243,7 @@ MCPs - endpoint features with the protocol op as the variant mcp... operation : list_tools | call_tool | list_resources | read_resource | list_prompts | get_prompt auth_family : none | api_key | bearer | oauth - assertion : succeeds | denied_without_permission | persists_across_processes + assertion : succeeds | denied_without_permission e.g. mcp.call_tool.oauth.succeeds ``` diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index a0059122ff0..b4379425854 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -261,55 +261,3 @@ The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthr Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The `python` job in `.github/workflows/test-linting.yml` runs them alongside the provider-edge and fixture tests at the root of `tests/e2e_harness/` with `--noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures`; they need only synthetic HTTP providers and temporary fixture storage - - -## MCP OAuth happy path - -`test_mcp_oauth_happy_path_e2e.py` runs one shared scenario with four variants: -aggregate gateway SSO and explicitly configured per-server JWT, each directly -against Linear and through the live provider edge. The edge forwards to real -Linear without replay and compares the forwarded bearer to the encrypted -canonical user/server credential. This observes the forwarding boundary, not -Linear's internal logs. Direct variants independently exercise discovery - -Use the existing database preparation, Prisma generation and Keycloak setup. -Build and stage the dashboard from the tested checkout as in the UI runner. -Provide `DATABASE_URL`, `LITELLM_MASTER_KEY`, `LITELLM_SALT_KEY`, `LITELLM_LICENSE`, -and the `E2E_KEYCLOAK_*` settings. Capture a test-account Linear login using -`mcp/linear_session_capture.py` and set `E2E_LINEAR_STORAGE_STATE` to that private -file. The test workspace must contain a team. Do not publish browser state or -raw test/proxy output - -```bash -E2E_MCP_OAUTH_LIVE=1 E2E_FIXTURE_MODE=live E2E_PROVIDER_CACHE=0 \ - uv run --no-sync pytest tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \ - --rootdir=. --reruns 0 -``` - -The test starts and restarts its own source-built proxy on a free loopback port, -retaining its database and SSO client but no Redis or process-local cache. It -does not restart an existing proxy or clear shared databases. Gateway login, -consent, immediate list/call and post-restart reconnect must all succeed. The -aggregate client never injects a gateway header; the explicitly labeled JWT -variant configures `x-litellm-api-key` for the first consent and reconnects with -only its gateway JWT after restart - -`.github/workflows/test-mcp-oauth-e2e.yml` automatically requests a run for -same-repository pull requests changing MCP, gateway authentication/SSO, consent -UI, dependencies or the relevant E2E harness/workflow paths. It retains manual -`workflow_dispatch` for targeted verification. The four cases run in the -protected `e2e-changed` environment after its normal deployment approval; -reviewers should approve and inspect this separate OAuth check when it appears. -Fork pull requests do not run this credentialed job; use a reviewed -same-repository branch for their verification. The workflow's path-filtered -check is not configured here as a globally required branch-protection check. -Provision `E2E_LINEAR_STORAGE_STATE_B64` as a secret there and retain the existing E2E license/AWS role configuration. A missing or -expired session fails the job; collection, deselection and skips are not passes. -No LLM call is needed - -Coverage remains limited to authorization-code OAuth over HTTP. M2M, OBO, -PKCE passthrough, static/BYOK, ID-JAG, forwarding, SigV4 and stdio are outside this -scenario; consult the registry and LIT-3559 for their existing coverage and gaps. -LIT-4506 owns broader isolation/failure regressions. LIT-7737 retains ownership -of dependency/Python compatibility and its matrix; this test reuses its delivered -environment and does not change dependency constraints or compatibility gates diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 3bb4ab0d3be..4c7c57ca18f 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -17,7 +17,6 @@ import functools import os from collections.abc import Generator, Iterator from datetime import datetime, timezone -from pathlib import Path from types import MappingProxyType from typing import Final @@ -29,7 +28,6 @@ from e2e_config import ( FIXTURE_DIR, FIXTURE_MODE_RAW, MANAGED_FILES_OPT_IN_ENV, - MCP_OAUTH_LIVE_OPT_IN_ENV, OTEL_TLS_OPT_IN_ENV, OTEL_V2_OPT_IN_ENV, OWNED_GATEWAY_OPT_IN_ENV, @@ -72,7 +70,6 @@ OPT_IN_MARKERS: Final = MappingProxyType( "redis_chaos": PERF_TESTS_OPT_IN_ENV, "perf": PERF_TESTS_OPT_IN_ENV, "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, - "mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV, "provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV, "owned_gateway": OWNED_GATEWAY_OPT_IN_ENV, "otel_v2": OTEL_V2_OPT_IN_ENV, @@ -170,11 +167,6 @@ def pytest_configure(config: pytest.Config) -> None: "quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; " "every other test waits for it to finish", ) - config.addinivalue_line( - "markers", - "mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless " - "E2E_MCP_OAUTH_LIVE is set", - ) config.addinivalue_line( "markers", "provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the " @@ -331,8 +323,6 @@ def pytest_runtest_setup(item: pytest.Item) -> None: item.user_properties.extend(item.config.stash[_IDLE_RSS].junit_properties) if not _reaches_proxy(item): return - if isinstance(item, pytest.Function) and "oauth_gateway" in item.fixturenames: - return reason = _proxy_fail_reason() if reason is not None: pytest.fail(reason) @@ -376,12 +366,6 @@ def pytest_runtest_makereport( report = yield if report.when in ("setup", "call"): attach_step_properties(item) - if item.get_closest_marker("mcp_oauth_live") is not None and call.excinfo is not None: - # Publish code locations only, never exception messages, source text or locals. - item.user_properties.append(("oauth_failure_phase", report.when)) - item.user_properties.append(("oauth_exception_type", call.excinfo.type.__name__)) - for entry in call.excinfo.traceback: - item.user_properties.append(("oauth_frame", f"{Path(entry.path).name}:{entry.lineno + 1}:{entry.name}")) if report.when == "call": item.stash[_CALL_PASSED] = report.passed report.user_properties = list(item.user_properties) diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index cdd194b267f..2ad815930a9 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -79,14 +79,6 @@ assertions: [succeeds] source: "db.py user_oauth_credential lookup" rationale: OAuth2 token passthrough; per-user credential storage -- id: mcp.call_tool.oauth.persists_across_processes - module: mcp - tier: P1 - operation: call_tool - auth_family: oauth - assertions: [persists_across_processes] - source: "outbound_credentials/per_user_oauth_store.py V2PerUserTokenStore" - rationale: Stored per-user token survives a verified restart of an owned gateway with no Redis cache - id: mcp.list_tools.none.succeeds module: mcp tier: P1 diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index efbc95966c4..724aa2362e7 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -111,10 +111,6 @@ CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haik CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5") S3_PARTITION_GRANULARITY = os.environ.get("E2E_S3_PARTITION_GRANULARITY", "day") -LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp") -LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "") -LINEAR_READONLY_TOOL: Final = "list_teams" # as listed by tools/list on mcp.linear.app when PR #33787 landed - # Jaeger query API of the compose stack's OTEL trace destination (the `jaeger` # service in docker-compose.yml maps it to host 16686). Trace-completeness tests # read exported spans back through it. @@ -207,7 +203,6 @@ MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" PERF_TESTS_OPT_IN_ENV: Final = "E2E_PERF_TESTS" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" -MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE" PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE" OWNED_GATEWAY_OPT_IN_ENV: Final = "E2E_OWNED_GATEWAY" OTEL_V2_OPT_IN_ENV: Final = "E2E_OTEL_V2" diff --git a/tests/e2e/mcp/linear_session_capture.py b/tests/e2e/mcp/linear_session_capture.py deleted file mode 100644 index 1c867e17e22..00000000000 --- a/tests/e2e/mcp/linear_session_capture.py +++ /dev/null @@ -1,58 +0,0 @@ -"""One-time helper to capture a logged-in Linear browser session for the -real-Linear MCP e2e test. - -The real-Linear test drives the genuine gateway-managed authorization_code -dance against ``mcp.linear.app``. The only step that cannot be scripted is -Linear's login (magic link / SSO), so a human authenticates once here and the -resulting session (cookies + local storage) is persisted to disk. The e2e test -then loads that session in a headless Playwright context and clicks Approve on -Linear's consent screen every run, with no human and no login automation. - -Run it with the e2e venv, log into Linear in the window that opens, then return -to the terminal and press Enter: - - LITELLM=~/litellm-mcpe2e - "$LITELLM"/.venv/bin/python "$LITELLM"/tests/e2e/mcp/linear_session_capture.py - -The session is written to ``E2E_LINEAR_STORAGE_STATE`` (default -``~/.litellm-e2e/linear_storage_state.json``), outside the repo. It is a -secret: never commit it. Re-run this whenever Linear expires the session. -""" - -from __future__ import annotations - -import os -from pathlib import Path - -from playwright.sync_api import sync_playwright - -DEFAULT_STATE_PATH = Path.home() / ".litellm-e2e" / "linear_storage_state.json" - - -def capture(state_path: Path) -> None: - """Open a headed browser at Linear, wait for the human to log in, then save - the authenticated session to ``state_path``.""" - state_path.parent.mkdir(parents=True, exist_ok=True) - with sync_playwright() as playwright: - browser = playwright.chromium.launch(headless=False) - context = browser.new_context() - page = context.new_page() - page.goto("https://linear.app/login", wait_until="domcontentloaded") - print("\n" + "=" * 72) - print("Log into Linear in the browser window that just opened.") - print("If Linear emails you a magic link, paste the link into THIS window's") - print("address bar (opening it in your default browser won't capture the") - print("session). Google SSO works too as long as you complete it here.") - print("When your Linear workspace has loaded, come back and press Enter.") - print("=" * 72) - input("Press Enter once you are logged in... ") - page.goto("https://mcp.linear.app/", wait_until="domcontentloaded") - context.storage_state(path=str(state_path)) - browser.close() - print(f"\nSaved Linear session to {state_path}") - print("Point the e2e test at it with:") - print(f' export E2E_LINEAR_STORAGE_STATE="{state_path}"') - - -if __name__ == "__main__": - capture(Path(os.environ.get("E2E_LINEAR_STORAGE_STATE", str(DEFAULT_STATE_PATH)))) diff --git a/tests/e2e/mcp/oauth_chat_client.py b/tests/e2e/mcp/oauth_chat_client.py deleted file mode 100644 index 68e838d8745..00000000000 --- a/tests/e2e/mcp/oauth_chat_client.py +++ /dev/null @@ -1,442 +0,0 @@ -"""Client for the mcp chat-completion OAuth e2e suite. - -Registers a gateway-managed OAuth (authorization_code) MCP server, seeds the -per-user upstream token by driving the interactive authorize dance with the -official mcp SDK's OAuthClientProvider (the browser leg is a headless Chromium -primed with a human's saved Linear session), then exercises the server through -/chat/completions, where the gateway lists and executes its tools with the -stored per-user token. - -Management routes (/v1/mcp/server CRUD, /chat/completions) go through the -shared ProxyClient transport. The MCP protocol used to seed the token goes through -the mcp SDK, the same library production MCP hosts run. -""" - -from __future__ import annotations - -import asyncio -import re -import time -from dataclasses import dataclass -from typing import TYPE_CHECKING, Final -from urllib.parse import parse_qsl - -import httpx -import httpx2 -import pytest -from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT -from e2e_http import AuthHeaders, NoBody, unwrap -from e2e_metadata import step -from idp import Identity -from mcp import ClientSession -from mcp.client.auth import OAuthClientProvider -from mcp.client.streamable_http import streamable_http_client -from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken -from mcp.types import TextContent -from models import ( - ChatBody, - ChatResponse, - McpOauthUserCredentialStatus, - McpServerCreateBody, - McpServerInfo, - McpServerUserCredentialListResponse, - McpServerUserCredentialRow, -) -from proxy_client import ProxyClient - -if TYPE_CHECKING: - from playwright.async_api import Route - -# Where the "browser" lands at the end of the authorize dance. Nothing listens -# here: the route interceptor short-circuits the final redirect and reads the -# code/state off its query string, exactly like a desktop MCP host intercepting -# its loopback redirect. -OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback" -BROWSER_CONSENT_TIMEOUT = 60.0 - - -def _mcp_url(alias: str, base_url: str = PROXY_BASE_URL) -> str: - return f"{base_url.rstrip('/')}/{alias}/mcp" - - -class InMemoryTokenStorage: - """The mcp SDK's TokenStorage protocol, in memory for one dance: the - DCR-registered client and the gateway tokens minted for it.""" - - def __init__(self) -> None: - self._tokens: OAuthToken | None = None - self._client_info: OAuthClientInformationFull | None = None - - async def get_tokens(self) -> OAuthToken | None: - return self._tokens - - async def set_tokens(self, tokens: OAuthToken) -> None: - self._tokens = tokens - - async def get_client_info(self) -> OAuthClientInformationFull | None: - return self._client_info - - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: - self._client_info = client_info - - -async def _browser_follow_authorize( - start_url: str, - storage_state_path: str, - identity: Identity | None = None, - server_alias: str | None = None, - allow_upstream_consent: bool = True, -) -> tuple[str, str | None]: - """Play the browser's role for a real upstream whose authorize endpoint - serves an interactive consent page (Linear). A headless Chromium primed - with a human's saved Linear session opens the gateway authorize URL and - clicks through Linear's consent screens (the mcp.linear.app Approve form, - then the linear.app workspace-selection page), riding the rest of the chain - (Linear -> gateway callback -> host redirect_uri). The final hop is - intercepted and short-circuited, since nothing listens there, and its - code/state are read off the query string.""" - from playwright.async_api import async_playwright - - captured: dict[str, str] = {} # mutable-ok: hand-off from the request listener - trail: list[str] = [] # mutable-ok: navigation diagnostics for a failed dance - - def _note_request(request: object) -> None: - url = getattr(request, "url", "") - host = httpx.URL(url).host - if not allow_upstream_consent and (host == "linear.app" or host.endswith(".linear.app")): - captured["upstream_consent"] = "seen" - if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured: - captured["url"] = url - - async def _swallow_redirect(route: Route) -> None: - await route.fulfill(status=200, content_type="text/plain", body="ok") - - async with async_playwright() as playwright: - browser = await playwright.chromium.launch(headless=True) - context = await browser.new_context(storage_state=storage_state_path) - await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect) - page = await context.new_page() - context.on("request", _note_request) - page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0])) - await page.goto(start_url, wait_until="domcontentloaded") - deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT - while "url" not in captured and time.monotonic() < deadline: - try: - await page.wait_for_load_state("networkidle", timeout=8000) - except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it - pass - if "upstream_consent" in captured or "url" in captured: - break - if await page.locator("#username").count() and identity is not None: - await page.locator("#username").fill(identity.username) - await page.locator("#password").fill(identity.password) - await page.locator("#kc-login").click() - continue - if "/ui/connect" in page.url and server_alias is not None: - card = page.locator("div.cursor-pointer").filter(has=page.get_by_text(server_alias, exact=True)) - if await card.count() != 1: - await asyncio.sleep(0.5) - continue - connect = card.get_by_text("Connect", exact=True) - if await connect.count(): - await connect.click() - continue - if not await card.locator("svg.text-success").count(): - await asyncio.sleep(0.5) - continue - finish = page.get_by_role("button", name="Finish connecting", exact=True) - if await finish.count() and await finish.is_enabled(): - await finish.click() - continue - control = page.locator( - 'button[name="action"][value="approve"], button:has-text("Authorize"), ' - 'button:has-text("Allow"), button:has-text("@"), a:has-text("@")' - ).first - try: - await control.click(timeout=5000) - except Exception: # noqa: BLE001 - nothing to advance yet; loop and re-check - await asyncio.sleep(0.5) - final_url = page.url - await browser.close() - - # A redirect chain can finish inside goto/networkidle before the loop checks the page. - assert "upstream_consent" not in captured, "cold reconnect required upstream consent" - landing = captured.get("url") - assert landing is not None, ( - f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; " - f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}" - ) - params = dict(parse_qsl(httpx.URL(landing).query.decode())) - assert "code" in params, "client redirect_uri carried no authorization code" - return params["code"], params.get("state") - - -def _oauth_provider( - url: str, - storage: InMemoryTokenStorage, - storage_state_path: str | None, - identity: Identity | None = None, - server_alias: str | None = None, - allow_upstream_consent: bool = True, -) -> OAuthClientProvider: - """The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR, - PKCE, token exchange) with the browser leg driven by Playwright against the - upstream's consent screen.""" - code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks - - async def _reject_redirect(_: str) -> None: - raise AssertionError("gateway demanded a fresh upstream consent; stored per-user token was not reused") - - async def _follow_redirect(authorize_url: str) -> None: - assert storage_state_path is not None - code, state = await _browser_follow_authorize( - authorize_url, storage_state_path, identity, server_alias, allow_upstream_consent - ) - code_holder["code"] = code - code_holder["state"] = state - - redirect_handler: Final = _reject_redirect if storage_state_path is None else _follow_redirect - - async def callback_handler() -> AuthorizationCodeResult: - code = code_holder.get("code") - assert code is not None, "callback_handler ran before the authorize redirect completed" - return AuthorizationCodeResult(code=code, state=code_holder.get("state")) - - return OAuthClientProvider( - server_url=url, - client_metadata=OAuthClientMetadata.model_validate( - { - "redirect_uris": [OAUTH_CLIENT_REDIRECT_URI], - "token_endpoint_auth_method": "none", - "grant_types": ["authorization_code", "refresh_token"], - "response_types": ["code"], - "client_name": "e2e-mcp-host", - } - ), - storage=storage, - redirect_handler=redirect_handler, - callback_handler=callback_handler, - ) - - -class _HeaderInjectingTransport(httpx2.AsyncBaseTransport): - """Adds the caller's LiteLLM key header to every outgoing SDK request - (discovery, DCR, token exchange), so the gateway resolves which user to - store the upstream token for from the key on the token exchange, exactly - like a production MCP host configured with a LiteLLM key header.""" - - def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None: - self._inner = inner - self._headers = headers - self._gateway_url = httpx2.URL(gateway_url) - - @staticmethod - def _port(url: httpx2.URL) -> int | None: - if url.port is not None: - return url.port - return {"http": 80, "https": 443}.get(url.scheme) - - async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response: - same_origin: Final = ( - request.url.scheme == self._gateway_url.scheme - and request.url.host == self._gateway_url.host - and self._port(request.url) == self._port(self._gateway_url) - ) - if same_origin: - for name, value in self._headers.items(): - if name not in request.headers: - request.headers[name] = value - else: - for name, value in self._headers.items(): - if request.headers.get(name) == value: - del request.headers[name] - return await self._inner.handle_async_request(request) - - async def aclose(self) -> None: - await self._inner.aclose() - - -def _oauth_http_client( - headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL -) -> httpx2.AsyncClient: - return httpx2.AsyncClient( - auth=auth, - timeout=httpx2.Timeout(REQUEST_TIMEOUT), - follow_redirects=True, - transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url), - ) - - -async def _seed_via_dance( - url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str -) -> tuple[str, ...]: - async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) as http_client: - async with streamable_http_client(url, http_client=http_client) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - listed = await session.list_tools() - return tuple(sorted(tool.name for tool in listed.tools)) - - -@dataclass(frozen=True, slots=True) -class OauthToolRun: - tools: tuple[str, ...] - is_error: bool - text: str - - -async def _list_and_call( - url: str, - headers: dict[str, str], - storage: InMemoryTokenStorage, - storage_state_path: str | None, - tool: str, - arguments: dict[str, str], - gateway_url: str = PROXY_BASE_URL, - identity: Identity | None = None, - server_alias: str | None = None, - allow_upstream_consent: bool = True, -) -> OauthToolRun: - async with _oauth_http_client( - headers, - _oauth_provider(url, storage, storage_state_path, identity, server_alias, allow_upstream_consent), - gateway_url, - ) as http_client: - async with streamable_http_client(url, http_client=http_client) as (read, write): - async with ClientSession(read, write) as session: - await session.initialize() - listed: Final = await session.list_tools() - result: Final = await session.call_tool(tool, arguments) - text: Final = "".join(content.text for content in result.content if isinstance(content, TextContent)) - return OauthToolRun( - tools=tuple(sorted(tool_item.name for tool_item in listed.tools)), - is_error=result.is_error, - text=text, - ) - - -@dataclass(frozen=True, slots=True) -class ChatMcpClient: - proxy: ProxyClient - - @step("Register the MCP server with the alias {body.alias}") - def create_server(self, body: McpServerCreateBody) -> McpServerInfo: - return unwrap( - self.proxy.transport.post( - "/v1/mcp/server", - headers=self.proxy.transport.master, - json=body, - response_type=McpServerInfo, - ) - ) - - @step("Read the MCP server back from /v1/mcp/server") - def server_info(self, server_id: str) -> McpServerInfo: - return unwrap( - self.proxy.transport.get( - f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, - params=NoBody(), - response_type=McpServerInfo, - ) - ) - - @step("Delete the MCP server") - def delete_server(self, server_id: str) -> None: - _ = self.proxy.transport.delete( - f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, - json=NoBody(), - response_type=NoBody, - ) - - @step("Sign the key's user in to the MCP server {alias} through the OAuth consent flow") - def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]: - """Drive the interactive authorize dance for `key`'s user so the gateway - stores their upstream token, retried to the shared deadline since the - just-created server and key propagate asynchronously. The LiteLLM key - rides x-litellm-api-key so the gateway binds the token to that user. - Returns the upstream tool names the dance listed, proof the token works.""" - headers = {"x-litellm-api-key": f"Bearer {key}"} - storage = InMemoryTokenStorage() - deadline = time.monotonic() + self.proxy.poll_timeout - last_error: Exception | None = None - while time.monotonic() < deadline: - try: - return asyncio.run(_seed_via_dance(_mcp_url(alias), headers, storage, storage_state_path)) - except Exception as exc: # noqa: BLE001 - retried to the deadline; the last error surfaces below - last_error = exc - time.sleep(self.proxy.poll_interval) - pytest.fail( - f"authorize dance for {alias!r} never completed within {self.proxy.poll_timeout}s; " - f"last error: {last_error!r}" - ) - - @step("List the tools on the MCP server {alias} and call {tool} over the MCP protocol") - def list_and_call( - self, - alias: str, - headers: dict[str, str], - storage: InMemoryTokenStorage, - storage_state_path: str | None, - tool: str, - arguments: dict[str, str], - base_url: str = PROXY_BASE_URL, - identity: Identity | None = None, - allow_upstream_consent: bool = True, - ) -> OauthToolRun: - return asyncio.run( - _list_and_call( - f"{base_url.rstrip('/')}/mcp" if identity is not None else _mcp_url(alias, base_url), - headers, - storage, - storage_state_path, - tool, - arguments, - base_url, - identity, - alias, - allow_upstream_consent, - ) - ) - - @step("List the users with a stored OAuth token for the MCP server") - def server_user_credentials(self, server_id: str) -> tuple[McpServerUserCredentialRow, ...]: - return unwrap( - self.proxy.transport.get( - f"/v1/mcp/server/{server_id}/user-credentials", - headers=self.proxy.transport.master, - params=NoBody(), - response_type=McpServerUserCredentialListResponse, - ) - ).root - - @step("Revoke the user's stored OAuth token for the MCP server") - def revoke_user_token(self, server_id: str, headers: AuthHeaders) -> None: - _ = unwrap( - self.proxy.transport.delete( - f"/v1/mcp/server/{server_id}/oauth-user-credential", - headers=headers, - json=NoBody(), - response_type=McpOauthUserCredentialStatus, - ) - ) - - @step("Send a /chat/completions request to {body.model} with an MCP server attached as a tool") - def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse: - """POST /chat/completions carrying the LiteLLM key in `headers` (either - ingress form) with an MCP server attached in `body.tools`. The gateway - resolves the user from the key and lists/executes the server's tools - with that user's stored upstream token.""" - return unwrap( - self.proxy.transport.post( - "/chat/completions", - headers=headers, - json=body, - response_type=ChatResponse, - ) - ) - - -def build_chat_client(proxy: ProxyClient) -> ChatMcpClient: - return ChatMcpClient(proxy=proxy) diff --git a/tests/e2e/mcp/oauth_gateway.py b/tests/e2e/mcp/oauth_gateway.py deleted file mode 100644 index 2abe8262b6e..00000000000 --- a/tests/e2e/mcp/oauth_gateway.py +++ /dev/null @@ -1,197 +0,0 @@ -"""An owned, source-built OAuth gateway with cold restarts and credential observations. - -Only this child process is restarted. Its database and SSO client survive while -its process-local caches do not; Redis is deliberately absent from its config. -The optional live edge measures headers without recording credentials or bodies. -""" - -from __future__ import annotations - -import os -import subprocess -import sys -import threading -import time -from collections.abc import Callable, Mapping -from contextlib import ExitStack -from dataclasses import dataclass, field -from pathlib import Path -from typing import Final - -import psycopg -from e2e_config import INHERITED_ENV_PREFIXES, available_port -from e2e_http import NoBody -from e2e_metadata import step -from idp import Keycloak, stop_process_group -from proxy_client import ProxyClient, build_proxy_client -from psycopg.rows import class_row -from pydantic import BaseModel, SecretStr, ValidationError - - -class StoredOAuth(BaseModel): - type: str - access_token: SecretStr - - -@dataclass(frozen=True, slots=True) -class CredentialRow: - credential_b64: str = field(repr=False) - - -@step("Read the user's stored OAuth credential for the MCP server from the database and decrypt it") -def stored_oauth(user_id: str, server_id: str) -> StoredOAuth: - """Read the encrypted credential because management APIs omit the plaintext token.""" - from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper - - with psycopg.Connection[CredentialRow].connect( - os.environ["DATABASE_URL"], row_factory=class_row(CredentialRow) - ) as conn: - row: Final = conn.execute( - 'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s', - (user_id, server_id), - ).fetchone() - assert row is not None, "canonical user/server has no persisted credential" - plaintext: Final = decrypt_value_helper( - row.credential_b64, "e2e_mcp_oauth", exception_type="debug", return_original_value=False - ) - assert plaintext is not None, "persisted credential must decrypt with the gateway salt" - assert plaintext != row.credential_b64, "persisted credential must be encrypted" - try: - credential: Final = StoredOAuth.model_validate_json(plaintext) - except ValidationError: - raise AssertionError("decrypted credential is not an OAuth payload") from None - assert credential.type == "oauth2" - assert bool(credential.access_token.get_secret_value()), "stored upstream token is empty" - return credential - - -class RpcMethod(BaseModel): - method: str = "" - - -@dataclass(slots=True) -class OAuthObservation: - gateway_token: str = field(default="", repr=False) - _seen: tuple[tuple[str, str, bool], ...] = field(default=(), init=False, repr=False) - _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) - - def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None: - if body is None or not url.endswith("/mcp"): - return - try: - operation: Final = RpcMethod.model_validate_json(body).method - except ValidationError: - return - if operation not in ("tools/list", "tools/call"): - return - received: Final = headers.get("authorization", "") - gateway_leaked: Final = any( - value in (self.gateway_token, f"Bearer {self.gateway_token}") for value in headers.values() - ) - with self._lock: - self._seen = (*self._seen, (operation, received, gateway_leaked)) - - def assert_forwarded(self, expected: StoredOAuth) -> None: - with self._lock: - snapshot: Final = self._seen - self._seen = () - assert {item[0] for item in snapshot} == {"tools/list", "tools/call"}, "missing upstream observations" - expected_header: Final = f"Bearer {expected.access_token.get_secret_value()}" - assert all(item[1] == expected_header for item in snapshot), "upstream bearer did not match the stored token" - assert all(not item[2] for item in snapshot), "gateway bearer leaked to the upstream" - - -@dataclass(slots=True) -class OAuthGateway: - base_url: str - proxy: ProxyClient - _environment: Mapping[str, str] = field(repr=False) - _command: tuple[str, ...] = field(repr=False) - _log_path: Path - _child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False) - - @step("Start the separate LiteLLM proxy and wait for /health/liveliness") - def start(self) -> None: - with self._log_path.open("ab") as log: - self._child = subprocess.Popen( - self._command, - env=self._environment, - stdout=log, - stderr=log, - start_new_session=True, - ) - deadline: Final = time.monotonic() + 120 - while time.monotonic() < deadline: - assert self._child.poll() is None, "owned OAuth gateway exited; inspect its private log" - result = self.proxy.transport.probe("/health/liveliness", params=NoBody()) - if result.status_code == 200: - return - time.sleep(0.5) - raise AssertionError("owned OAuth gateway did not become ready") - - @step("Stop the separate LiteLLM proxy") - def stop(self) -> None: - if self._child is not None: - stop_process_group(self._child) - assert self._child.poll() is not None, "old gateway process is still alive" - - @step("Restart the separate LiteLLM proxy process so its in-memory caches start empty") - def restart(self) -> None: - assert self._child is not None - previous: Final = self._child.pid - self.stop() - self.start() - assert self._child.pid != previous, "gateway restart did not create a new process" - - -@step("Start a separate LiteLLM proxy from source with JWT auth against Keycloak") -def owned_gateway(idp: Keycloak, directory: Path, cleanup: ExitStack) -> OAuthGateway: - for name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"): - assert os.environ.get(name), f"{name} is required for the owned OAuth gateway" - port: Final = available_port() - base_url: Final = f"http://127.0.0.1:{port}" - - def defer(callback: Callable[[], object]) -> None: - cleanup.callback(callback) - - browser: Final = idp.browser_client(callback_url=f"{base_url}/sso/callback", defer=defer) - config: Final = directory / "oauth-gateway.yaml" - config.write_text( - "model_list: []\n" - "general_settings:\n" - " master_key: os.environ/LITELLM_MASTER_KEY\n" - " database_url: os.environ/DATABASE_URL\n" - " enable_jwt_auth: true\n" - " litellm_jwtauth:\n" - " user_id_jwt_field: sub\n" - " user_email_jwt_field: email\n" - " team_ids_jwt_field: groups\n" - " user_id_upsert: true\n" - ) - environment: Final = { - **{key: value for key, value in os.environ.items() if not key.startswith(INHERITED_ENV_PREFIXES)}, - **browser.environment(idp.discovery()), - "PROXY_BASE_URL": base_url, - "JWT_PUBLIC_KEY_URL": idp.jwks_url, - "JWT_ISSUER": idp.issuer, - "JWT_AUDIENCE": "litellm-e2e", - "DISABLE_SCHEMA_UPDATE": "true", - "STORE_MODEL_IN_DB": "True", - "PYTHONPATH": str(Path(__file__).resolve().parents[3]), - } - gateway: Final = OAuthGateway( - base_url=base_url, - proxy=build_proxy_client( - base_url=base_url, - control_plane_base_url=base_url, - replica_urls=(base_url,), - control_replica_urls=(base_url,), - master_key=os.environ["LITELLM_MASTER_KEY"], - ), - _environment=environment, - _command=(sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config), "--port", str(port)), - _log_path=directory / "oauth-gateway.log", - ) - cleanup.callback(gateway.stop) - gateway.start() - return gateway diff --git a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py b/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py deleted file mode 100644 index 9d7e4713963..00000000000 --- a/tests/e2e/mcp/test_mcp_chat_completion_oauth_e2e.py +++ /dev/null @@ -1,218 +0,0 @@ -"""On-demand e2e: a chat completion drives a gateway-managed OAuth MCP server. - -The real end-user flow for MCP over an OAuth server: a user registers a Linear -authorization_code server, authorizes it once so the gateway stores their -upstream token, then sends a normal /chat/completions request with the Linear -MCP attached. The gateway resolves the user from the LiteLLM key, lists Linear's -tools with the stored per-user token, lets the model call one, executes it -upstream with that token, and returns the answer. This is proven against the -real Linear MCP server (mcp.linear.app) and a real Anthropic model, once per -documented ingress header (x-litellm-api-key and Authorization). - -The authorize dance is seeded through the mcp SDK's OAuthClientProvider; the one -step Linear cannot auto-approve is the human consent, so it is captured once out -of band (mcp/linear_session_capture.py) into a saved browser session and a -headless Chromium clicks Approve every run. The test therefore skips unless -E2E_LINEAR_STORAGE_STATE points at that session, so it never runs on the per-PR -CI path; it is a nightly/on-demand real-server smoke test. - -Fail-before-fix: without the stored per-user token the gateway lists no Linear -tools, so mcp_list_tools comes back empty, nothing is called, and the -assertions fail; a served, called, non-empty Linear tool proves the gateway -pulled and used the user's token. -""" - -from __future__ import annotations - -import os - -import pytest - -from e2e_config import CHEAP_ANTHROPIC_MODEL, LINEAR_MCP_URL, LINEAR_STORAGE_STATE, unique_marker -from e2e_http import AuthHeaders -from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta -from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, McpServerCreateBody, ObjectPermission -from proxy_client import ProxyClient - -pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`") -pytest.importorskip( - "playwright.async_api", - reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`", -) - -from oauth_chat_client import ChatMcpClient, build_chat_client # noqa: E402 # imports follow the importorskip guards - -pytestmark = [ - pytest.mark.e2e, - pytest.mark.skipif( - not LINEAR_STORAGE_STATE or not os.path.exists(LINEAR_STORAGE_STATE), - reason="set E2E_LINEAR_STORAGE_STATE to a Linear session captured via mcp/linear_session_capture.py", - ), -] - -# Pinned from a live dance during verification (never guessed); the gateway -# prefixes every upstream tool name with the server alias. list_teams is a -# read-only Linear tool that takes no arguments and returns the caller's teams. -LINEAR_READONLY_TOOL = "list_teams" -LINEAR_PROMPT = "Use the list_teams tool to list my Linear teams, then reply with the name of one of them." - - -@pytest.fixture(scope="session") -def chat_client(proxy: ProxyClient) -> ChatMcpClient: - return build_chat_client(proxy) - - -class TestMcpChatCompletionOauth: - """A scoped internal-user key on a real Linear authorization_code server, - used through /chat/completions once per ingress header: the gateway pulls - the user's stored upstream token, lists and executes Linear's tools during - the completion, and returns the answer.""" - - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - @meta( - Subject( - domain=Domain.MCP, - route=Route.CHAT_COMPLETIONS, - providers=(Provider.ANTHROPIC,), - models=(CHEAP_ANTHROPIC_MODEL,), - capabilities=(Capability.FUNCTION_CALLING,), - mode=Mode.NONSTREAM, - ) - ) - def test_chat_completion_uses_linear_with_x_litellm_api_key_header( - self, chat_client: ChatMcpClient, resources: ResourceManager - ) -> None: - marker = unique_marker() - alias = f"e2elinear{marker}" - created = chat_client.create_server( - McpServerCreateBody( - alias=alias, - url=LINEAR_MCP_URL, - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - ) - ) - resources.defer(lambda: chat_client.delete_server(created.server_id)) - - stored = chat_client.server_info(created.server_id) - assert stored.auth_type == "oauth2" - assert stored.oauth2_flow == "authorization_code" - assert stored.allow_all_keys is False - - key = chat_client.proxy.generate_key( - KeyGenerateBody( - user_id="e2e-test-user", - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - resources.defer(lambda: chat_client.proxy.delete_key(key)) - - seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) - assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( - f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" - ) - - response = chat_client.chat_with_mcp( - AuthHeaders.model_validate({"x-litellm-api-key": f"Bearer {key}"}), - ChatBody( - model=CHEAP_ANTHROPIC_MODEL, - messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], - tools=[ - McpChatTool( - server_url=f"litellm_proxy/mcp/{alias}", - server_label=alias, - require_approval="never", - ) - ], - ), - ) - - message = response.choices[0].message - assert message is not None and message.content, f"completion returned no answer: {response}" - meta = message.provider_specific_fields - assert meta is not None, f"no MCP metadata on the completion: {response}" - listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} - assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( - f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" - ) - results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] - assert results and results[0].result, ( - f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" - ) - - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - @meta( - Subject( - domain=Domain.MCP, - route=Route.CHAT_COMPLETIONS, - providers=(Provider.ANTHROPIC,), - models=(CHEAP_ANTHROPIC_MODEL,), - capabilities=(Capability.FUNCTION_CALLING,), - mode=Mode.NONSTREAM, - ) - ) - def test_chat_completion_uses_linear_with_authorization_bearer_header( - self, chat_client: ChatMcpClient, resources: ResourceManager - ) -> None: - marker = unique_marker() - alias = f"e2elinear{marker}" - created = chat_client.create_server( - McpServerCreateBody( - alias=alias, - url=LINEAR_MCP_URL, - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - ) - ) - resources.defer(lambda: chat_client.delete_server(created.server_id)) - - stored = chat_client.server_info(created.server_id) - assert stored.auth_type == "oauth2" - assert stored.oauth2_flow == "authorization_code" - assert stored.allow_all_keys is False - - key = chat_client.proxy.generate_key( - KeyGenerateBody( - user_id="e2e-test-user", - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - resources.defer(lambda: chat_client.proxy.delete_key(key)) - - seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE) - assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, ( - f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}" - ) - - response = chat_client.chat_with_mcp( - AuthHeaders.model_validate({"authorization": f"Bearer {key}"}), - ChatBody( - model=CHEAP_ANTHROPIC_MODEL, - messages=[ChatMessage(role="user", content=LINEAR_PROMPT)], - tools=[ - McpChatTool( - server_url=f"litellm_proxy/mcp/{alias}", - server_label=alias, - require_approval="never", - ) - ], - ), - ) - - message = response.choices[0].message - assert message is not None and message.content, f"completion returned no answer: {response}" - meta = message.provider_specific_fields - assert meta is not None, f"no MCP metadata on the completion: {response}" - listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function} - assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, ( - f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}" - ) - results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"] - assert results and results[0].result, ( - f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}" - ) diff --git a/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py b/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py deleted file mode 100644 index b57c941cb18..00000000000 --- a/tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Real OAuth consent, immediate MCP operations and cold-restart persistence. - -Aggregate SSO uses the SDK's normal authentication. The per-server variant is -explicitly a configured two-header client, not an Authorization-only OAuth host. -The observed variants forward to the same real Linear upstream and compare its -bearer at the forwarding boundary; direct variants retain unmodified discovery. -""" - -from __future__ import annotations - -import os -from collections.abc import Iterator -from contextlib import ExitStack -from pathlib import Path -from types import MappingProxyType -from typing import Final, Literal - -import pytest -from e2e_config import LINEAR_MCP_URL, LINEAR_READONLY_TOOL, LINEAR_STORAGE_STATE, unique_marker -from e2e_http import AuthHeaders, NoBody, get_external, unwrap -from e2e_metadata import Domain, Route, Subject, meta -from idp import Identity, Keycloak -from lifecycle import ResourceManager -from models import ( - McpOauthCredentials, - McpServerCreateBody, - ObjectPermission, - TeamMemberAddBody, - TeamMemberEntry, - TeamUpdateBody, -) -from oauth_chat_client import ChatMcpClient, InMemoryTokenStorage, OauthToolRun, build_chat_client -from oauth_gateway import OAuthGateway, OAuthObservation, owned_gateway, stored_oauth -from provider_edge import LiveEdge, start_provider_edge -from proxy_client import ProxyClient -from pydantic import BaseModel, ValidationError - -pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_live] - - -class OAuthMetadata(BaseModel): - issuer: str - authorization_endpoint: str - token_endpoint: str - registration_endpoint: str - - -class LinearTeam(BaseModel): - id: str - name: str - - -class LinearTeams(BaseModel): - teams: tuple[LinearTeam, ...] - - -def assert_tool_result(run: OauthToolRun, tool: str) -> None: - assert tool in run.tools - assert run.is_error is False - try: - result: Final = LinearTeams.model_validate_json(run.text) - except ValidationError: - raise AssertionError("list_teams did not return the expected teams payload") from None - assert result.teams, "the test workspace must contain at least one team" - assert all(team.id and team.name for team in result.teams), "team results must contain identifiers and names" - - -@pytest.fixture(scope="module") -def oauth_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OAuthGateway]: - assert LINEAR_STORAGE_STATE and Path(LINEAR_STORAGE_STATE).is_file(), ( - "E2E_LINEAR_STORAGE_STATE must name a captured Linear login; see mcp/linear_session_capture.py" - ) - assert os.environ.get("E2E_FIXTURE_MODE", "live") == "live", "OAuth acceptance cannot use replay" - with ExitStack() as cleanup: - yield owned_gateway(idp, tmp_path_factory.mktemp("mcp-oauth"), cleanup) - - -@pytest.fixture(scope="module") -def proxy(oauth_gateway: OAuthGateway) -> ProxyClient: - return oauth_gateway.proxy - - -@pytest.fixture(scope="module") -def client(proxy: ProxyClient) -> ChatMcpClient: - return build_chat_client(proxy) - - -class TestMcpOauthHappyPath: - @pytest.mark.covers("mcp.list_tools.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.succeeds") - @pytest.mark.covers("mcp.call_tool.oauth.persists_across_processes") - @meta( - Subject( - domain=Domain.MCP, - route=Route.MCP, - ) - ) - @pytest.mark.parametrize("route", ("aggregate_sso", "explicit_header_jwt")) - @pytest.mark.parametrize("observed", (False, True), ids=("direct", "observed")) - def test_consent_list_call_and_cold_restart( - self, - client: ChatMcpClient, - resources: ResourceManager, - jwt_identity: Identity, - idp: Keycloak, - oauth_gateway: OAuthGateway, - route: Literal["aggregate_sso", "explicit_header_jwt"], - observed: bool, - ) -> None: - alias: Final = f"e2elinear{unique_marker()}" - tool: Final = f"{alias}-{LINEAR_READONLY_TOOL}" - token: Final = idp.access_token(jwt_identity) - observation: Final = OAuthObservation(gateway_token=token) - edge: Final = ( - start_provider_edge( - LiveEdge(observe_request=observation.observe), - mounts=MappingProxyType( - {"linear": "https://mcp.linear.app", ".well-known": "https://mcp.linear.app/.well-known"} - ), - ) - if observed - else None - ) - if edge is not None: - resources.defer(edge.shutdown) - metadata: Final = ( - unwrap( - get_external( - "https://mcp.linear.app/.well-known/oauth-authorization-server", - response_type=OAuthMetadata, - ) - ) - if observed - else None - ) - created: Final = client.create_server( - McpServerCreateBody( - alias=alias, - server_name=alias, - url=f"{edge.edge.api_base('linear')}/mcp" if edge is not None else LINEAR_MCP_URL, - transport="http", - allow_all_keys=False, - auth_type="oauth2", - oauth2_flow="authorization_code", - per_server_oauth_discovery=route == "explicit_header_jwt", - issuer=metadata.issuer if metadata else None, - authorization_url=metadata.authorization_endpoint if metadata else None, - token_url=metadata.token_endpoint if metadata else None, - registration_url=metadata.registration_endpoint if metadata else None, - credentials=McpOauthCredentials(upstream_resource=LINEAR_MCP_URL) if observed else None, - ) - ) - resources.defer(lambda: client.delete_server(created.server_id)) - assert client.server_user_credentials(created.server_id) == (), ( - "scenario must start without upstream credentials" - ) - unwrap( - client.proxy.transport.post( - "/team/member_add", - headers=client.proxy.transport.master, - json=TeamMemberAddBody( - team_id=jwt_identity.group, member=TeamMemberEntry(user_id=jwt_identity.user_id, role="user") - ), - response_type=NoBody, - ) - ) - client.proxy.update_team( - TeamUpdateBody( - team_id=jwt_identity.group, - object_permission=ObjectPermission(mcp_servers=[created.server_id]), - ) - ) - headers: Final = {"x-litellm-api-key": f"Bearer {token}"} if route == "explicit_header_jwt" else {} - resources.defer( - lambda: client.revoke_user_token( - created.server_id, - AuthHeaders(authorization=f"Bearer {idp.access_token(jwt_identity)}"), - ) - ) - identity: Final = jwt_identity if route == "aggregate_sso" else None - first: Final = client.list_and_call( - alias, - headers, - InMemoryTokenStorage(), - LINEAR_STORAGE_STATE, - tool, - {}, - base_url=oauth_gateway.base_url, - identity=identity, - ) - assert_tool_result(first, tool) - credentials: Final = client.server_user_credentials(created.server_id) - assert len(credentials) == 1 - assert credentials[0].user_id == jwt_identity.user_id - assert credentials[0].credential_type == "oauth2" - first_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id) - if observed: - observation.assert_forwarded(first_stored_oauth) - oauth_gateway.restart() - fresh_token: Final = idp.access_token(jwt_identity) - observation.gateway_token = fresh_token - second: Final = client.list_and_call( - alias, - {"Authorization": f"Bearer {fresh_token}"} if identity is None else {}, - InMemoryTokenStorage(), - LINEAR_STORAGE_STATE if identity is not None else None, - tool, - {}, - base_url=oauth_gateway.base_url, - identity=identity, - allow_upstream_consent=False, - ) - assert_tool_result(second, tool) - second_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id) - if observed: - observation.assert_forwarded(second_stored_oauth) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index d7a10c6da1f..40e9e181982 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -706,28 +706,14 @@ class McpInfo(BaseModel): logo_url: str | None = None -class McpOauthCredentials(BaseModel): - upstream_resource: str - - class McpServerCreateBody(BaseModel): - """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is - `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints - are discovered and registered via DCR when left unset. `allow_all_keys` - false scopes the server to keys granted it through object_permission.""" + """POST /v1/mcp/server. `allow_all_keys` false scopes the server to keys + granted it through object_permission.""" alias: str url: str transport: str = "http" allow_all_keys: bool = True - auth_type: str | None = None - oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None - per_server_oauth_discovery: bool | None = None - issuer: str | None = None - authorization_url: str | None = None - token_url: str | None = None - registration_url: str | None = None - credentials: McpOauthCredentials | None = Field(default=None, repr=False) server_name: str | None = None description: str | None = None mcp_info: McpInfo | None = None @@ -767,26 +753,6 @@ class McpServerListResponse(RootModel[list[McpServerRow]]): """GET /v1/mcp/server answers with a bare array of servers.""" -class McpServerUserCredentialRow(BaseModel): - user_id: str - credential_type: Literal["oauth2", "byok"] - expires_at: str | None = None - connected_at: str | None = None - updated_at: str - - -class McpServerUserCredentialListResponse(RootModel[tuple[McpServerUserCredentialRow, ...]]): - """GET /v1/mcp/server/{server_id}/user-credentials answers with a bare array.""" - - -class McpOauthUserCredentialStatus(BaseModel): - server_id: str - has_credential: bool - expires_at: str | None = None - is_expired: bool = False - connected_at: str | None = None - - class ToolsetTool(BaseModel): server_id: str tool_name: str diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index da3dd72d795..f15cfd4e98b 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -14,7 +14,6 @@ markers = cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish - mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 6f59f694352..563d35bf794 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -697,6 +697,54 @@ def test_per_user_authorization_code_with_pkce_binds_the_token_to_the_authorizin ) +@pytest.mark.timeout(600) +def test_per_user_authorization_code_token_is_read_back_from_the_database_by_a_fresh_proxy_without_redis( + gateway: Gateway, tmp_path: Path +) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "cold" + uuid.uuid4().hex[:8] + identity: Final = _register_oauth( + scenario, + peer, + auth, + alias, + auth_type="oauth2", + oauth2_flow="authorization_code", + credentials={"client_id": "cold-client", "client_secret": "cold-secret", "scopes": ["tools.call"]}, + ) + owner_id: Final = scenario.user() + owner: Final = scenario.key(user_id=owner_id, object_permission={"mcp_servers": [identity]}) + registered: Final = gateway.client.post( + f"/{alias}/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"} + ) + assert registered.status_code in (200, 201), registered.text + client_id: Final = str(_response_object(registered)["client_id"]) + pkce: Final = _Pkce(secrets.token_urlsafe(32)) + code, _ = _authorize_through_gateway(gateway, auth, alias, owner, client_id, pkce) + redeemed: Final = _redeem(gateway, alias, owner, client_id, code, pkce) + assert redeemed.status_code == 200, redeemed.text + token: Final = _issued_token(redeemed.json()) + stored: Final = read_rows( + 'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s', + (owner_id, identity), + ) + assert len(stored) == 1, stored + stored_bytes: Final = base64.urlsafe_b64decode(str(stored[0]["credential_b64"])) + assert token.encode() not in stored_bytes, "stored token is plaintext" + config: Final = tmp_path / "no_redis.yaml" + config.write_text( + "general_settings:\n" + " master_key: os.environ/LITELLM_MASTER_KEY\n" + " database_url: os.environ/DATABASE_URL\n" + " store_model_in_db: true\n" + ) + peer.drain() + with owned_proxy(gateway, tmp_path, {}, config=config, remove_environment=("REDIS_HOST", "REDIS_PORT")) as cold: + called: Final = call_tool(cold, owner, identity, f"{alias}-add", ADD) + assert called.status_code == 200, called.text + assert _authorizations(peer) == (f"Bearer {token}".encode(),) + + def test_authorization_request_without_pkce_is_refused_before_reaching_the_authorization_server( gateway: Gateway, ) -> None: diff --git a/tests/unit/test_circleci_path_filter.py b/tests/unit/test_circleci_path_filter.py index 0856115e39b..f94b2fa863a 100644 --- a/tests/unit/test_circleci_path_filter.py +++ b/tests/unit/test_circleci_path_filter.py @@ -52,7 +52,7 @@ CI = [".github/workflows/test-unit.yml"] ("mcp-dependencies", ["pyproject.toml"], "run"), ("mcp-dependencies", ["uv.lock"], "run"), ("mcp-dependencies", ["litellm/experimental_mcp_client/client.py"], "run"), - ("mcp-dependencies", ["tests/e2e/mcp/oauth_chat_client.py"], "run"), + ("mcp-dependencies", ["tests/e2e/mcp/datadog_mcp.py"], "run"), ("mcp-dependencies", ["litellm-proxy-extras/pyproject.toml"], "run"), ("mcp-dependencies", ["scripts/check_mcp_sdk_install.py"], "run"), ("mcp-dependencies", [".github/workflows/test-mcp-dependency-resolution.yml"], "run"), From fc69f987dc05c7966d6d9d82384abf2f1cbe79cd Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Sat, 17 Oct 2026 14:33:26 -0700 Subject: [PATCH 50/64] fix(logging): keep root logger INFO lines when json_logs is on and LITELLM_LOG is set (#47175) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/_logging.py | 2 ++ tests/unit/test_logging.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/_logging.py b/litellm/_logging.py index 913e9a20101..d08254db3fd 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1000,6 +1000,8 @@ def _initialize_loggers_with_handler(handler: logging.Handler): lg.handlers.clear() # remove any existing handlers lg.addHandler(handler) # add JSON formatter handler lg.propagate = False # prevent bubbling to parent/root + if (configured_level := os.getenv("LITELLM_LOG")) is not None: + logging.getLogger().setLevel(resolve_log_level(configured_level)) def _get_uvicorn_json_log_config(): diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index 03b822411ac..cb4d3f82611 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -89,6 +89,40 @@ def test_json_mode_emits_one_record_per_logger(capfd): assert "timestamp" in obj, "`timestamp` key missing" +@pytest.mark.parametrize( + ("litellm_log", "want_info_lines"), + [("INFO", 2), ("DEBUG", 2), (None, 0)], +) +def test_json_mode_root_logger_honours_litellm_log(monkeypatch, capfd, litellm_log, want_info_lines): + root: Final = logging.getLogger() + saved_level: Final = root.level + saved_handlers: Final = list(root.handlers) + saved_propagate: Final = root.propagate + if litellm_log is None: + monkeypatch.delenv("LITELLM_LOG", raising=False) + else: + monkeypatch.setenv("LITELLM_LOG", litellm_log) + try: + turn_on_json() + logging.getLogger("some.third.party").info("library info line") + logging.info("root info line") + logging.getLogger("some.third.party").warning("library warning line") + out, err = capfd.readouterr() + finally: + root.handlers.clear() + root.handlers.extend(saved_handlers) + root.setLevel(saved_level) + root.propagate = saved_propagate + + records = [json.loads(line) for line in (out + err).splitlines() if line.strip()] + info_records = [r for r in records if r["level"] == "INFO"] + assert len(info_records) == want_info_lines, records + if want_info_lines: + assert {r["message"] for r in info_records} == {"library info line", "root info line"} + assert {r["component"] for r in info_records} == {"some.third.party", "root"} + assert [r["message"] for r in records if r["level"] == "WARNING"] == ["library warning line"] + + def test_json_formatter_parses_embedded_json_message(): """ Test that JsonFormatter parses embedded JSON in the message field and promotes From 897dda792047ba4eb9b3327615d101abc9b82e71 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:39:14 +0000 Subject: [PATCH 51/64] fix(router)!: stop a chat-mode Mantle deployment from rewriting the shared Responses routing row (#46263) * fix(router): stop a chat-mode Mantle deployment from rewriting the shared Responses routing row A deployment with an explicit custom_llm_provider and model_info.mode: chat used to overwrite the shared litellm.model_cost row, moving every other deployment and plain SDK calls onto Chat Completions until restart, because the downgrade guard read a double-prefixed key and never fired. The guard now reads the real key, and a deployment's own mode decides its routing in responses_api_bridge_check, except on chatgpt and github_copilot where Chat Completions is dead for responses-mode models. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(responses): keep the bridge-check mode tests in the sharded bridge flag file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): let a deployment mode pin replace only the shared-row mode so the bridge capability rules still run Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): keep Responses-only backends bridged under a deployment chat pin and satisfy the type gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): give the bridge casts a reason and drop the bare noqa so the type discipline gate passes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): keep a responses row with no supported_endpoints bridged under a deployment chat pin Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): read the downgrade guard's mode from the row registration writes A chat-mode openai/ deployment merged into the bare row while the guard read openai/, so the guard never fired and the shared responses row turned to chat. An explicit responses/ prefix now also outranks a deployment's chat pin. * fix(router): let an explicit responses/ prefix outrank a registered chat row The router now registers a chat-pinned openai/responses/ deployment under its own name, so the lookup found mode chat and the request went to Chat Completions as responses/. * fix(router): keep the shared row's mode for every deployment pin and honor the pin on /v1/messages A deployment's mode pin now never rewrites the shared cost-map row in either direction, including the responses/-stripped twin key, so siblings and SDK calls keep routing by the registered mode. The backend key formula goes back to main's, which keeps capability fields off rows that registration never wrote at base. The /v1/messages adapter and handler read the deployment's model_info, so a chat-pinned deployment gets a plain reasoning effort on chat completions. * fix(router): keep copilot responses opt-in, make unmapped pins order-free, send Mantle auto summary as auto A responses pin on a github_copilot model or a model missing from the cost map still marks the shared row, since that is how those models opt into the Responses API. A chat pin never marks a row, so registration order no longer matters. Mantle's OpenAI path gets summary auto in place of the auto summary detailed. The Azure o3-pro and o3-deep-research rows drop Chat Completions, which Azure does not serve for them. * fix(router): skip a non-string model_info.mode when keeping the registered mode * fix(router): keep chat-pinned Mantle tool calls and summaries on the Responses API A Mantle row that also serves the Responses API now follows the GPT-5 tool and reasoning summary bridge rules, since Mantle Chat Completions rejects function tools on gpt-5.4+ unless reasoning is off and rejects the summary aliases. Mantle chat gets a dict reasoning_effort as its effort level. The opt-in auto summary asks each Responses config for its level, auto on Mantle's OpenAI path, so a client's own detailed summary follows drop_params as at base. No mode replaces a registered responses mode, and chatgpt and github_copilot rows take any pin until they are responses, which restores the Copilot chat opt-out. * fix(responses): ask the target for its auto summary level on the chat bridge and reasoning_effort * test(integration): run the Mantle region-prefixed v1 base case again This PR routes that deployment by its own chat pin, so the case passes and no longer needs the skip * fix(bedrock_mantle): keep web search and reasoning summaries on the Responses API under a chat pin A chat-pinned Mantle GPT-5 deployment now sends web_search_options and a dict reasoning_effort carrying a summary to the Responses API, which serves them. Mantle Chat Completions drops n above 1 and a reasoning summary under drop_params and answers 400 without it. The /v1/messages adapter probes the bridge with only the client's own summary, so an auto-injected one never moves the request * fix(bedrock_mantle): check n and reasoning summaries only on the Chat Completions request * fix(router): keep a non-string deployment mode off the shared routing row A deployment whose model_info.mode was an int, a list, or an object still reached the shared cost-map row, so an unpinned sibling on the same backend stopped routing to the Responses API. The shared registration now drops a mode that is not a string, and the deployment keeps its own raw value. Adds integration cells for per-deployment routing across the OpenAI and Anthropic SDKs, raw httpx, wildcard routes, malformed modes, an upstream outage, config changes under traffic, and a worker kill plus restart. * test(integration): find each routing cell's spend row by the response it logged A bridged call logs its row under the wrapped Responses id, which the caller never sees on a streamed Messages answer, so look the row up by the answer's own ids and the upstream ids for its marker. Int and list modes 500 their own model group on main too, so those cells check only the sibling, and the summary cell sends the auto level Mantle's Responses path accepts * fix(bedrock_mantle): run extra_body fields through the chat checks A chat-pinned Mantle request that nests n or reasoning_effort in extra_body now gets the same 400, drop, and effort-level handling as the top-level fields. The bridge test injects a MockTransport client instead of patching HTTPHandler.post. * test(router): cover nested reasoning_effort routing, messages thinking summaries and per-target auto summary levels * fix(router): probe a pinned deployment's own API on /health and keep a requested summary beside an extra_body effort * fix(router): keep a custom-priced deployment's mode pin off the shared backend row on each request * fix(router): run per-request custom pricing through the shared mode guard so a health probe keeps the row's mode * fix(router): drop the debug log that wrote a request's model name from keep_registered_mode * test(router): cover custom-priced chat pins, their health probes and summary aliases beside an extra_body effort * fix(proxy): pass only the mode pin to a /health probe so a failing probe never cools its deployment down * test(router): cover failing health probes leaving their deployment out of cooldown, with and without background checks * test(router): kill the worker that holds burst calls, so the worker kill cell never depends on how the kernel spread them * fix(router): route a fine-tune the same whichever deployment registers first Main's register_model now seeds an unmapped fine-tune's row from its base model's built-in row, mode included, on the first registration. The mode guard read only the live row, so a responses pin marked the shared row only when it registered before its siblings. The guard now reads the row the registration merges onto, the live row or else the row it would seed. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 39 +- .../pass_through/adapters/handler.py | 10 +- .../pass_through/adapters/transformation.py | 11 +- .../pass_through/messages/handler.py | 25 +- .../responses_adapters/handler.py | 2 + .../responses_adapters/transformation.py | 8 +- .../llms/base_llm/responses/transformation.py | 8 +- .../bedrock_mantle/chat/transformation.py | 73 +- .../responses/transformation.py | 4 + litellm/main.py | 22 +- ...odel_prices_and_context_window_backup.json | 3 - litellm/proxy/health_check.py | 5 +- litellm/responses/bridge_check.py | 137 ++- .../transformation.py | 4 +- litellm/responses/main.py | 2 +- litellm/responses/utils.py | 18 +- litellm/router.py | 30 +- litellm/utils.py | 20 +- model_prices_and_context_window.json | 3 - .../test_bedrock_mantle_output_limit_wire.py | 22 +- .../test_deployment_mode_routing_wire.py | 1037 +++++++++++++++++ ...responses_transformation_transformation.py | 109 ++ ..._handler_reasoning_effort_normalization.py | 45 +- ...erimental_pass_through_messages_handler.py | 41 + .../test_responses_adapters_handler.py | 23 +- ...bedrock_mantle_responses_transformation.py | 32 + .../test_bedrock_mantle_transformation.py | 162 +++ .../proxy/test_health_check_max_tokens.py | 129 ++ .../test_responses_api_bridge_flag.py | 401 +++++-- .../unit/test_router_model_cost_isolation.py | 400 ++++++- tests/unit/test_xai_responses_auto_routing.py | 16 +- 31 files changed, 2657 insertions(+), 184 deletions(-) create mode 100644 tests/integration/routing/test_deployment_mode_routing_wire.py diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 8d23526b41a..98322aecf6a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -41,7 +41,11 @@ from litellm.responses.sse_output_recovery import ( record_output_item_chunk, record_output_text_chunk, ) -from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_responses_api_stream_options +from litellm.responses.utils import ( + ResponsesAPIRequestUtils, + auto_reasoning_summary, + normalize_responses_api_stream_options, +) from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionReasoningItem, @@ -176,6 +180,16 @@ def _without_audio_input_parts(input_items: list[object]) -> list[object]: return [_without_audio_input_parts_in_item(item) for item in input_items] +def _effort_with_requested_summary(nested_effort: object, top_level_effort: object) -> object: + match top_level_effort, nested_effort: + case {"summary": str() as summary}, str() as effort: + return {"effort": effort, "summary": summary} + case {"summary": str() as summary}, Mapping() as reasoning if "summary" not in reasoning: + return {**reasoning, "summary": summary} + case _: + return nested_effort + + def _served_provider(litellm_params: Mapping[str, object]) -> str | None: custom_llm_provider: Final = litellm_params.get("custom_llm_provider") return custom_llm_provider if isinstance(custom_llm_provider, str) else None @@ -953,6 +967,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): optional_params: dict, responses_api_request: "ResponsesAPIOptionalRequestParams", drop_params: bool = False, + model: str | None = None, + custom_llm_provider: str | None = None, ) -> None: """Map optional_params into responses_api_request (mutates in place).""" for key, value in optional_params.items(): @@ -984,7 +1000,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif key == "previous_response_id": responses_api_request["previous_response_id"] = value elif key == "reasoning_effort": - responses_api_request["reasoning"] = self.map_reasoning_effort(value) + responses_api_request["reasoning"] = self.map_reasoning_effort(value, model, custom_llm_provider) elif key == "web_search_options": self._add_web_search_tool(responses_api_request, value) @@ -1100,6 +1116,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): optional_params, responses_api_request, drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params), + model=model, + custom_llm_provider=_served_provider(litellm_params), ) stream: Final = optional_params.get("stream") or litellm_params.get("stream", False) @@ -1705,14 +1723,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Extract supported params from extra_body and merge into optional_params extra_body_copy: Final = extra_body.copy() + top_level_effort: Final = cast(object, optional_params.get("reasoning_effort")) # cast-ok: upcast for key, value in extra_body_copy.items(): if key in supported_responses_api_params: # Prefer extra_body value if it exists (may have more complete info like summary in reasoning_effort) - optional_params[key] = extra_body.pop(key) + extra_body.pop(key) + optional_params[key] = ( + _effort_with_requested_summary(cast(object, value), top_level_effort) # cast-ok: upcast + if key == "reasoning_effort" + else value + ) return optional_params - def map_reasoning_effort(self, reasoning_effort: object) -> Reasoning: + def map_reasoning_effort( + self, reasoning_effort: object, model: str | None = None, custom_llm_provider: str | None = None + ) -> Reasoning: # If dict is passed, convert it directly to Reasoning object if isinstance(reasoning_effort, dict): return Reasoning( @@ -1726,7 +1752,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) return ( - Reasoning(effort=reasoning_effort, summary="detailed") + Reasoning( + effort=reasoning_effort, + summary="detailed" if model is None else auto_reasoning_summary(model, custom_llm_provider), + ) if auto_summary_enabled else Reasoning(effort=reasoning_effort) ) diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 7b8c2550838..a9aa3de69ad 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -501,7 +501,11 @@ class LiteLLMMessagesToCompletionTransformationHandler: return None if resolved_provider == "litellm_proxy": return None - from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary, responses_api_bridge_check + from litellm.responses.bridge_check import ( + chat_target_accepts_reasoning_summary, + responses_api_bridge_check, + str_keyed_mapping, + ) if requested_summary is not None and chat_target_accepts_reasoning_summary(local_model, resolved_provider): return _ChatTargetReasoning(effort=effort, summary=requested_summary) @@ -510,6 +514,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: web_search_options: Final = completion_kwargs.get("web_search_options") tools: Final = completion_kwargs.get("tools") + client_summary: Final = thinking.get("summary") model_info, _ = responses_api_bridge_check( model=local_model, custom_llm_provider=resolved_provider, @@ -517,8 +522,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: cast(OpenAIWebSearchOptions, web_search_options) if isinstance(web_search_options, dict) else None ), tools=cast("list[dict[str, object]]", tools) if isinstance(tools, list) else None, - reasoning_effort=reasoning_effort, + reasoning_effort=effort if client_summary is None else {"effort": effort, "summary": client_summary}, api_base=resolved_api_base, + deployment_model_info=str_keyed_mapping(completion_kwargs.get("model_info")), ) return None if model_info.get("mode") == "responses" else _ChatTargetReasoning(effort=effort) diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index 3aad0be1913..0d3e48c49fd 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -14,6 +14,7 @@ from litellm.llms.anthropic.pass_through.utils import ( json_copy_or_original, prompt_cache_key_from_user_id, ) +from litellm.responses.utils import auto_reasoning_summary # OpenAI has a 64-character limit for function/tool names # Anthropic does not have this limit, so we need to truncate long names @@ -736,7 +737,7 @@ class LiteLLMAnthropicMessagesAdapter: if reasoning_effort: return { "reasoning_effort": LiteLLMAnthropicMessagesAdapter._apply_reasoning_summary_wrapping( - reasoning_effort, thinking + reasoning_effort, thinking, model, None ) } return {} @@ -745,6 +746,8 @@ class LiteLLMAnthropicMessagesAdapter: def _apply_reasoning_summary_wrapping( reasoning_effort: str, thinking: Mapping[str, object], + model: str, + custom_llm_provider: str | None, ) -> Any: """ Apply the reasoning_effort/summary wrapping rules shared by every @@ -762,7 +765,7 @@ class LiteLLMAnthropicMessagesAdapter: if summary: return {"effort": reasoning_effort, "summary": summary} if is_reasoning_auto_summary_enabled(): - return {"effort": reasoning_effort, "summary": "detailed"} + return {"effort": reasoning_effort, "summary": auto_reasoning_summary(model, custom_llm_provider)} return reasoning_effort def translate_anthropic_tool_choice_to_openai( @@ -1174,7 +1177,9 @@ class LiteLLMAnthropicMessagesAdapter: new_kwargs["reasoning_effort"] = ( reasoning_effort if is_claude_target - else self._apply_reasoning_summary_wrapping(reasoning_effort, cast(dict[str, object], thinking)) + else self._apply_reasoning_summary_wrapping( + reasoning_effort, cast(dict[str, object], thinking), model, custom_llm_provider + ) ) def _translate_output_format_to_openai( diff --git a/litellm/llms/anthropic/pass_through/messages/handler.py b/litellm/llms/anthropic/pass_through/messages/handler.py index 00bf1774dd0..88cfeaad66a 100644 --- a/litellm/llms/anthropic/pass_through/messages/handler.py +++ b/litellm/llms/anthropic/pass_through/messages/handler.py @@ -55,15 +55,19 @@ _SAFEGUARD_KWARGS: Final = TypeAdapter(Mapping[str, object]) _SAFEGUARD_MESSAGES: Final = TypeAdapter(Sequence[Mapping[str, object]]) -def _bridges_to_responses_api(model: str, custom_llm_provider: str) -> bool: - from litellm.responses.bridge_check import responses_api_bridge_check +def _bridges_to_responses_api(model: str, custom_llm_provider: str, deployment_model_info: object) -> bool: + from litellm.responses.bridge_check import responses_api_bridge_check, str_keyed_mapping - model_info, _ = responses_api_bridge_check(model=model, custom_llm_provider=custom_llm_provider) + model_info, _ = responses_api_bridge_check( + model=model, + custom_llm_provider=custom_llm_provider, + deployment_model_info=str_keyed_mapping(deployment_model_info), + ) return model_info.get("mode") == "responses" def _responses_mode_is_lost_by_prefix_strip( - requested_model: str, resolved_model: str, custom_llm_provider: str + requested_model: str, resolved_model: str, custom_llm_provider: str, deployment_model_info: object ) -> bool: """Whether a Responses-only deployment stops looking like one once its provider prefix is stripped. @@ -73,15 +77,16 @@ def _responses_mode_is_lost_by_prefix_strip( """ if requested_model == resolved_model: return False - return _bridges_to_responses_api(requested_model, custom_llm_provider) and not _bridges_to_responses_api( - resolved_model, custom_llm_provider - ) + return _bridges_to_responses_api( + requested_model, custom_llm_provider, deployment_model_info + ) and not _bridges_to_responses_api(resolved_model, custom_llm_provider, deployment_model_info) def _should_route_to_responses_api( custom_llm_provider: str | None, requested_model: str | None = None, resolved_model: str | None = None, + deployment_model_info: object = None, ) -> bool: """Return True when the request should use the Responses API path. @@ -94,7 +99,9 @@ def _should_route_to_responses_api( return True if custom_llm_provider is None or requested_model is None or resolved_model is None: return False - return _responses_mode_is_lost_by_prefix_strip(requested_model, resolved_model, custom_llm_provider) + return _responses_mode_is_lost_by_prefix_strip( + requested_model, resolved_model, custom_llm_provider, deployment_model_info + ) async def _with_native_safeguard_results( @@ -606,7 +613,7 @@ def anthropic_messages_handler( if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. if kwargs.get("compaction") is None and _should_route_to_responses_api( - custom_llm_provider, original_model, model + custom_llm_provider, original_model, model, kwargs.get("model_info") ): return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( max_tokens=max_tokens, diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index 692e5d4c175..05ea4e4301d 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -16,6 +16,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( ) from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.prompt_templates.common_utils import cache_control_as_prompt_cache_breakpoint +from litellm.responses.utils import auto_reasoning_summary from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic import ( AllAnthropicMessageValues, @@ -188,6 +189,7 @@ def _build_responses_kwargs( include_encrypted_reasoning=_provider_returns_encrypted_reasoning( model, forwarded_kwargs.get("custom_llm_provider") ), + auto_summary_level=auto_reasoning_summary(model, forwarded_kwargs.get("custom_llm_provider")), ) # Normalize reasoning effort based on model capabilities diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py index 97f2db36a85..fb611b64dda 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/transformation.py @@ -544,6 +544,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: def translate_thinking_to_reasoning( thinking: dict[str, object], output_config: dict[str, object] | None = None, + auto_summary_level: str = "detailed", ) -> dict[str, object] | None: """ Convert Anthropic thinking param to Responses API reasoning param. @@ -576,7 +577,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if summary: result["summary"] = summary elif auto_summary: - result["summary"] = "detailed" + result["summary"] = auto_summary_level return result @classmethod @@ -591,6 +592,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: self, anthropic_request: AnthropicMessagesRequest, include_encrypted_reasoning: bool = True, + auto_summary_level: str = "detailed", ) -> dict[str, Any]: """ Translate a full Anthropic /v1/messages request dict to @@ -599,7 +601,8 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ``include_encrypted_reasoning`` asks the provider for ``reasoning.encrypted_content`` on every call, so a reasoning model's items can be replayed intact next turn even when the client sent no ``thinking`` block; pass False for a provider whose - Responses API rejects ``include``. + Responses API rejects ``include``. ``auto_summary_level`` is the ``reasoning.summary`` the + opt-in auto summary injects when the client sent none. """ model: Final[str] = anthropic_request["model"] messages_list: Final = cast( @@ -670,6 +673,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: reasoning: Final = self.translate_thinking_to_reasoning( thinking, output_config=cast(dict[str, Any] | None, output_config), + auto_summary_level=auto_summary_level, ) if reasoning: responses_kwargs["reasoning"] = reasoning diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index ed8bfc8cb96..01baeed9bef 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -1,7 +1,7 @@ import types from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast import httpx @@ -26,6 +26,8 @@ else: LiteLLMLoggingObj = Any BaseLLMException = Any +ReasoningSummaryLevel: TypeAlias = Literal["auto", "concise", "detailed"] + class BaseResponsesAPIConfig(ABC): def __init__(self): @@ -91,6 +93,10 @@ class BaseResponsesAPIConfig(ABC): def get_supported_openai_params(self, model: str) -> list: pass + def auto_reasoning_summary(self) -> ReasoningSummaryLevel: + """The ``reasoning.summary`` LiteLLM's opt-in auto summary asks this provider for.""" + return "detailed" + @abstractmethod def map_openai_params( self, diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 1c791e994f2..61c7d51a889 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -10,19 +10,21 @@ Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the BedrockMantleAuthMixin in common_utils. """ -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx import litellm from litellm._logging import verbose_logger +from litellm.exceptions import UnsupportedParamsError from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock_mantle.common_utils import ( BEDROCK_MANTLE_DEFAULT_REGION, BedrockMantleAuthMixin, ) -from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name +from litellm.llms.openai.chat.gpt_5_transformation import get_effort_level, is_gpt_reasoning_series_name from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams @@ -39,6 +41,51 @@ if TYPE_CHECKING: ) +def _with_chat_reasoning_effort(non_default_params: Mapping[str, object]) -> Mapping[str, object]: + """Chat Completions takes ``reasoning_effort`` as a level, so a Responses-shaped ``{"effort", "summary"}`` + passes only its effort, the way the OpenAI GPT-5 chat config reads it.""" + reasoning_effort: Final = non_default_params.get("reasoning_effort") + if not isinstance(reasoning_effort, dict): + return non_default_params + effort: Final = get_effort_level(reasoning_effort) + rest: Final = {key: value for key, value in non_default_params.items() if key != "reasoning_effort"} + return rest if effort is None else {**rest, "reasoning_effort": effort} + + +def _unsupported_chat_values(non_default_params: Mapping[str, object]) -> tuple[str, ...]: + n: Final = non_default_params.get("n") + reasoning_effort: Final = non_default_params.get("reasoning_effort") + summary: Final = reasoning_effort.get("summary") if isinstance(reasoning_effort, Mapping) else None + return tuple( + value + for value, unsupported in ( + (f"n={n}", isinstance(n, int) and n > 1), + (f"reasoning summary {summary}", summary is not None), + ) + if unsupported + ) + + +def _drops_params(litellm_params: Mapping[str, object]) -> bool: + return litellm.drop_params is True or litellm_params.get("drop_params") is True + + +def _chat_request_params(optional_params: Mapping[str, object], drop_params: bool) -> Mapping[str, object]: + unsupported: Final = _unsupported_chat_values(optional_params) + if unsupported and not drop_params: + raise UnsupportedParamsError( + message=f"Bedrock Mantle Chat Completions does not support {', '.join(unsupported)}. To drop unsupported " + "openai params from the call, set `litellm.drop_params = True`", + llm_provider="bedrock_mantle", + status_code=400, + ) + servable: Final = _with_chat_reasoning_effort(optional_params) + n: Final = servable.get("n") + return MappingProxyType( + {key: value for key, value in servable.items() if not (key == "n" and isinstance(n, int) and n > 1)} + ) + + class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): """ Transformation config for Amazon Bedrock Mantle OpenAI-compatible API. @@ -158,6 +205,28 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): if param != "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: OpenAIGPTConfig.transform_request signature + request_params: Final = dict(_chat_request_params(optional_params, _drops_params(litellm_params))) + return super().transform_request(model, messages, request_params, litellm_params, headers) + + def transform_extra_body( + self, + extra_body: Mapping[str, object], + request: Mapping[str, object], + model: str, + litellm_params: Mapping[str, object], + ) -> Mapping[str, object]: + return _chat_request_params( + super().transform_extra_body(extra_body, request, model, litellm_params), _drops_params(litellm_params) + ) + def _supports_reasoning(self, model: str) -> bool: try: return litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index fa0da6a28b4..8a0ce58c8b2 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -24,6 +24,7 @@ import litellm from litellm._logging import verbose_logger from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items +from litellm.llms.base_llm.responses.transformation import ReasoningSummaryLevel from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock_mantle.common_utils import ( @@ -73,6 +74,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK_MANTLE + def auto_reasoning_summary(self) -> ReasoningSummaryLevel: + return "auto" if self.use_openai_path else super().auto_reasoning_summary() + def get_error_class( self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers ) -> BaseLLMException: diff --git a/litellm/main.py b/litellm/main.py index 3eae64d828e..85d00f28025 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1078,6 +1078,11 @@ def mock_completion( from litellm.responses.bridge_check import ( # noqa: E402 # re-export at the old definition site _resolve_openai_api_base as _resolve_openai_api_base, # pyright: ignore[reportPrivateUsage] # noqa: PLC0414 # re-export of moved helper ) +from litellm.responses.bridge_check import ( + bridge_reasoning_effort, + keep_registered_mode, + str_keyed_mapping, +) from litellm.responses.bridge_check import ( # noqa: E402 # re-export at the old definition site responses_api_bridge_check as responses_api_bridge_check, # pyright: ignore[reportUnknownVariableType] # noqa: PLC0414 # re-export of moved function ) @@ -1201,6 +1206,8 @@ def _register_custom_pricing_for_request( pricing overrides (e.g. a zero-cost wildcard) from clobbering built-in pricing used by sibling deployments of the same backend model. Direct SDK calls keep the legacy behavior of registering the shared key with pricing. + Either way the shared key's ``mode`` follows ``keep_registered_mode``, as on + deployment registration, since a call routes by its own ``model_info``. """ entry: Final = _build_custom_pricing_entry( custom_llm_provider=custom_llm_provider, @@ -1210,8 +1217,11 @@ def _register_custom_pricing_for_request( shared_key: Final = f"{custom_llm_provider}/{model}" deployment_id: Final = _get_router_deployment_id(kwargs) if deployment_id is None: - litellm.register_model({shared_key: entry}, persist_across_reloads=False) + litellm.register_model({shared_key: keep_registered_mode(shared_key, entry)}, persist_across_reloads=False) return + shared_entry: Final = keep_registered_mode( + shared_key, CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry) + ) litellm.register_model( {deployment_id: entry}, persist_across_reloads=False, @@ -1220,7 +1230,7 @@ def _register_custom_pricing_for_request( add_to_provider_model_sets=False, ) litellm.register_model( - {shared_key: CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry)}, + {shared_key: shared_entry}, persist_across_reloads=False, warning_display_name=shared_key, custom_llm_provider=custom_llm_provider, @@ -5539,6 +5549,9 @@ def completion( custom_llm_provider=custom_llm_provider, web_search_options=web_search_options, api_base=api_base, + deployment_model_info=str_keyed_mapping( + cast(object, model_info) # cast-ok: model_info is the untyped dict kwarg + ), ) if not _is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): @@ -5858,9 +5871,12 @@ def completion( custom_llm_provider=custom_llm_provider, web_search_options=web_search_options, tools=tools, - reasoning_effort=reasoning_effort, + reasoning_effort=bridge_reasoning_effort(reasoning_effort, optional_params.get("extra_body")), reasoning_summary=_reasoning_summary_for_bridge, api_base=api_base, + deployment_model_info=str_keyed_mapping( + cast(object, model_info) # cast-ok: model_info is the untyped dict kwarg + ), ) # Use base_model (the true underlying model) for Azure model-type diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6ebb54c0035..13210a7eb1c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10904,7 +10904,6 @@ }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ - "/v1/chat/completions", "/v1/batch", "/v1/responses" ], @@ -10973,7 +10972,6 @@ "output_cost_per_token_batches": 4e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ - "/v1/chat/completions", "/v1/batch", "/v1/responses" ], @@ -11005,7 +11003,6 @@ "output_cost_per_token_batches": 4e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ - "/v1/chat/completions", "/v1/batch", "/v1/responses" ], diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 14ec69a2059..0d2f8dc9a0c 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -551,10 +551,13 @@ async def _run_model_health_check(model: dict): model_info, litellm_params ) timeout: Final = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS + probe_params: Final = ( + {**litellm_params, "model_info": {"mode": model_info["mode"]}} if "mode" in model_info else litellm_params + ) return await run_with_timeout( litellm.ahealth_check( - litellm_params, + probe_params, mode=mode, prompt=DEFAULT_HEALTH_CHECK_PROMPT, input=["test from litellm"], diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 2552ee8c880..120a0d7c828 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -13,9 +13,52 @@ from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import OpenAIWebSearchOptions from litellm.types.utils import LlmProviders -from litellm.utils import ProviderConfigManager, get_model_info_helper +from litellm.utils import ProviderConfigManager, get_model_info_helper, model_cost_registration_row _OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1" +RESPONSES_ONLY_PROVIDERS: Final = frozenset({"chatgpt", "github_copilot"}) +ROUTABLE_MODES: Final = frozenset({"chat", "responses"}) + + +def deployment_routing_mode(deployment_model_info: Mapping[str, object] | None, custom_llm_provider: str) -> str | None: + """The mode a router deployment's own ``model_info`` pins its chat requests to, or ``None`` to defer + to the shared cost-map row. Lets two deployments of one backend model route differently without + either rewriting the shared row the SDK and the other deployments read.""" + if deployment_model_info is None or custom_llm_provider in RESPONSES_ONLY_PROVIDERS: + return None + mode: Final = deployment_model_info.get("mode") + return mode if isinstance(mode, str) and mode in ROUTABLE_MODES else None + + +def keep_registered_mode(backend_key: str, shared_model_info: Mapping[str, object]) -> Mapping[str, object]: + """A deployment's ``chat`` or ``responses`` pin routes only that deployment, so it never replaces the + routing mode registered for its shared backend row, which the SDK and every sibling deployment route by, + and no other mode replaces a registered ``responses`` mode either. + + A ``responses`` pin still marks a row that has no routing mode, since that is how a model missing from + the cost map opts into the Responses API. A mode the row inherits from the built-in entry its registration + falls back to counts as registered, since the first registration writes it whichever deployment that is. A ``chatgpt`` or ``github_copilot`` row takes any mode until it + is ``responses``, since those providers read it by name to pick their endpoint. Elsewhere a ``chat`` pin + never marks a row, so the outcome does not depend on which deployment registers first. A ``mode`` that + is not a string names no mode, so it never reaches the row.""" + deployment_mode: Final = shared_model_info.get("mode") + without_mode: Final = {key: value for key, value in shared_model_info.items() if key != "mode"} + if not isinstance(deployment_mode, str): + return without_mode + registered_mode: Final = model_cost_registration_row( + backend_key, shared_model_info.get("litellm_provider", "") + ).get("mode") + if registered_mode != "responses" and ( + deployment_mode not in ROUTABLE_MODES or backend_key.split("/", 1)[0] in RESPONSES_ONLY_PROVIDERS + ): + return shared_model_info + if registered_mode is None and deployment_mode == "chat": + return without_mode + if registered_mode is None: + return shared_model_info + return {**shared_model_info, "mode": registered_mode} + + _OPENAI_API_HOST: Final = "api.openai.com" @@ -38,6 +81,46 @@ def _resolve_openai_api_base(api_base: str | None) -> str: ) +def _mantle_responses_series_name(model: str) -> str | None: + """The bare OpenAI series name of a Bedrock Mantle model that also serves the Responses API, or ``None``.""" + from litellm.llms.bedrock_mantle.common_utils import mantle_supports_responses, split_mantle_region_prefix + + if not mantle_supports_responses(model, litellm.model_cost): + return None + return split_mantle_region_prefix(model)[1].removeprefix("openai.") + + +def str_keyed_mapping(value: object) -> Mapping[str, object] | None: + """A request field such as a deployment's ``model_info`` when it is a mapping, read without narrowing the + caller's variable (``litellm.completion`` reuses ``model_info`` for pricing registration).""" + if not isinstance(value, Mapping): + return None + return cast( # cast-ok: isinstance only narrows to Mapping[Unknown, Unknown]; request JSON keys are str + Mapping[str, object], value + ) + + +def bridge_reasoning_effort( + reasoning_effort: str | Mapping[str, object] | None, extra_body: object +) -> str | Mapping[str, object] | None: + nested_body: Final = str_keyed_mapping(extra_body) + if reasoning_effort is not None or nested_body is None: + return reasoning_effort + nested_effort: Final = nested_body.get("reasoning_effort") + return nested_effort if isinstance(nested_effort, str) else str_keyed_mapping(nested_effort) + + +def backend_serves_chat_completions(shared_model_info: Mapping[str, object]) -> bool: + """Whether the shared cost-map row lets a deployment pin ``mode: chat``. Responses-only backends + (``gpt-5.1-codex``, ``o3-pro``, ``gpt-5.4-pro``) list only ``/v1/responses`` and must stay bridged, and + a ``responses`` row with no ``supported_endpoints`` (regional Azure codex and pro, Databricks codex) is + read the same way.""" + supported_endpoints: Final = shared_model_info.get("supported_endpoints") + if not isinstance(supported_endpoints, Sequence) or isinstance(supported_endpoints, str): + return shared_model_info.get("mode") != "responses" + return "/v1/chat/completions" in supported_endpoints + + def chat_target_accepts_reasoning_summary(model: str, custom_llm_provider: str | None) -> bool: try: local_model, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) @@ -58,7 +141,8 @@ def responses_api_bridge_check( reasoning_effort: str | Mapping[str, object] | None = None, reasoning_summary: object | None = None, api_base: str | None = None, -) -> tuple[dict, str]: + deployment_model_info: Mapping[str, object] | None = None, +) -> tuple[dict[str, object], str]: model_info: dict[str, object] = {} # Global flag: route ALL OpenAI chat completions through Responses API. @@ -70,21 +154,21 @@ def responses_api_bridge_check( try: model_info = cast( - dict, + dict[str, object], get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider), ) - if model_info.get("mode") is None and model.startswith("responses/"): - model = model.replace("responses/", "") - mode = "responses" - model_info["mode"] = mode - except Exception as e: verbose_logger.debug("Error getting model info: %s", e) - if model.startswith("responses/"): # handle azure models - `azure/responses/` - model = model.replace("responses/", "") - mode = "responses" - model_info["mode"] = mode + explicit_responses_prefix: Final = model.startswith("responses/") + pinned_mode: Final = ( + None if explicit_responses_prefix else deployment_routing_mode(deployment_model_info, custom_llm_provider) + ) + if explicit_responses_prefix or pinned_mode == "responses": + model = model.replace("responses/", "") + model_info["mode"] = "responses" + elif pinned_mode == "chat" and backend_serves_chat_completions(model_info): + model_info["mode"] = "chat" if web_search_options is not None and custom_llm_provider == "xai": model_info["mode"] = "responses" @@ -112,6 +196,9 @@ def responses_api_bridge_check( # an explicit effort with function tools is rejected from gpt-5.6 on, and the unset # effort only from gpt-6 on (gpt-5.6 serves tools with reasoning silently off), so the # azure_ai gate keys on those measured boundaries instead of gpt-5.4+. + # - Bedrock Mantle serves the same API and enforces the gpt-5.4+ rule even with the effort + # unset (measured on gpt-5.6-terra), so a Mantle row that also serves Responses follows + # these rules, which keeps a deployment pinned to ``mode: chat`` from a 400 on them. # - Older GPT-5 names (e.g. ``gpt-5``, ``gpt-5.1``): bridge only when a reasoning # summary alias is present with ``reasoning_effort`` (tools alone stay on chat). has_function_tool: Final = any( @@ -140,8 +227,21 @@ def responses_api_bridge_check( on_foundry_openai_endpoint: Final = custom_llm_provider == "azure_ai" and azure_ai_supports_native_responses( model, api_base ) + mantle_series_name: Final = ( + _mantle_responses_series_name(model) if custom_llm_provider == "bedrock_mantle" else None + ) + on_mantle_responses_endpoint: Final = mantle_series_name is not None + summary_requested: Final = reasoning_summary is not None or ( + on_mantle_responses_endpoint + and isinstance(reasoning_effort, Mapping) + and reasoning_effort.get("summary") is not None + ) + gpt_series_name: Final = mantle_series_name or model on_constraint_enforcing_endpoint: Final = ( - custom_llm_provider == "azure" or resolved_api_base == "" or _is_openai_backed_api_base(resolved_api_base) + custom_llm_provider == "azure" + or on_mantle_responses_endpoint + or resolved_api_base == "" + or _is_openai_backed_api_base(resolved_api_base) ) chat_rejects_function_tools: Final = ( has_function_tool @@ -150,20 +250,21 @@ def responses_api_bridge_check( foundry_chat_rejects_function_tools_while_reasoning(model, reasoning_effort) if on_foundry_openai_endpoint else ( - OpenAIGPT5Config.is_model_gpt_5_4_plus_model(model) + OpenAIGPT5Config.is_model_gpt_5_4_plus_model(gpt_series_name) and (reasoning_effort is not None or on_constraint_enforcing_endpoint) ) ) ) if ( - (custom_llm_provider in ("openai", "azure") or on_foundry_openai_endpoint) + (custom_llm_provider in ("openai", "azure") or on_foundry_openai_endpoint or on_mantle_responses_endpoint) and model_info.get("mode") != "responses" - and OpenAIGPT5Config.is_model_gpt_5_model(model) - and not OpenAIGPT5Config.is_model_gpt_5_search_model(model) + and OpenAIGPT5Config.is_model_gpt_5_model(gpt_series_name) + and not OpenAIGPT5Config.is_model_gpt_5_search_model(gpt_series_name) and ( - (reasoning_effort is not None and reasoning_summary is not None) + (reasoning_effort is not None and summary_requested) or chat_rejects_function_tools or (carries_tool_search and on_constraint_enforcing_endpoint) + or (web_search_options is not None and on_mantle_responses_endpoint) ) ): model_info["mode"] = "responses" diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 35d4b8bf1ba..842230699b7 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -762,8 +762,8 @@ class LiteLLMCompletionResponsesConfig: """ Whether ``litellm.completion`` will route this model back onto the Responses API. - Delegates to the same check ``litellm.completion`` itself runs, and is asked with the - params this transform is about to emit, so the two cannot reach different answers. + Delegates to the same check ``litellm.completion`` itself runs, asked with the params this + transform is about to emit. A deployment's own ``model_info.mode`` pin is not passed here. """ from litellm.main import responses_api_bridge_check diff --git a/litellm/responses/main.py b/litellm/responses/main.py index d32d48c51ca..ec1786b529a 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1555,7 +1555,7 @@ def responses( ) reasoning_effort: Final[object] = local_vars.get("reasoning_effort") request_reasoning: Final = ( - LiteLLMResponsesTransformationHandler().map_reasoning_effort(reasoning_effort) + LiteLLMResponsesTransformationHandler().map_reasoning_effort(reasoning_effort, model, custom_llm_provider) if current_reasoning is None and reasoning_effort is not None else current_reasoning ) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index b16fe72b281..8e5e1c1cfee 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -10,7 +10,7 @@ from typing_extensions import TypeIs # noqa: TID251 # narrows untyped wire pay import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path -from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel from litellm.types.llms.openai import ( AllMessageValues, OutputTokensDetails, @@ -1294,3 +1294,19 @@ class ResponseAPILoggingUtils: return chat_usage _transform_response_api_usage_to_chat_usage = transform_response_api_usage_to_chat_usage + + +def auto_reasoning_summary(model: str, custom_llm_provider: object) -> ReasoningSummaryLevel: + """The summary level the auto summary injects for a target, read from its Responses API config.""" + from litellm.utils import ProviderConfigManager + + try: + provider: Final = ( + custom_llm_provider if isinstance(custom_llm_provider, str) else litellm.get_llm_provider(model=model)[1] + ) + except litellm.exceptions.BadRequestError: + return "detailed" + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config( + provider, model.removeprefix(f"{provider}/") + ) + return "detailed" if responses_config is None else responses_config.auto_reasoning_summary() diff --git a/litellm/router.py b/litellm/router.py index 8d77fd0dac5..e363f41ccde 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -139,6 +139,7 @@ from litellm.llms.openai_like.model_info import ( MODEL_INFO_REFRESH_SECONDS, get_openai_compatible_model_info, ) +from litellm.responses.bridge_check import keep_registered_mode from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.complexity_router.context_compaction import ( @@ -11010,7 +11011,6 @@ class Router: ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes backend_keys: Final = Router._backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) - backend_key: Final = backend_keys[0] # For the shared backend key, keep only cost-map schema fields # (minus custom pricing) so that one deployment's pricing overrides @@ -11019,34 +11019,8 @@ class Router: # name. Each deployment's full model_info is already stored under # its unique model_id above. shared_model_info: Final = shared_backend_model_info(model_info) - existing_shared_mode: Final = (cast(dict | None, litellm.model_cost.get(backend_key, {})) or {}).get("mode") - deployment_mode: Final = shared_model_info.get("mode") - # Keep the built-in bridge mode stable for shared backend keys. - # Multiple aliases can point at the same provider/model backend, - # but their deployment-level overrides should not downgrade the - # backend from responses -> chat via last-write-wins registration. - # Only preserve in that specific direction so legitimate upgrades - # (e.g. chat -> responses) and unrelated mode changes still apply, - # and so a missing deployment mode does not silently clear the - # existing shared backend mode. - is_responses_to_chat_downgrade: Final = existing_shared_mode == "responses" and deployment_mode == "chat" - would_clear_existing_mode: Final = existing_shared_mode is not None and deployment_mode is None - if is_responses_to_chat_downgrade or would_clear_existing_mode: - if deployment_mode is not None: - verbose_router_logger.warning( - "Router: preserving existing mode=%s for shared backend " - "key %s instead of the deployment-specified mode=%s " - "(prevents alias registration from downgrading the " - "shared backend mode).", - existing_shared_mode, - backend_key, - deployment_mode, - ) - shared_model_info["mode"] = existing_shared_mode - - # Always register the (possibly mode-preserved) shared backend info. litellm.register_model( - model_cost={_key: shared_model_info for _key in backend_keys}, + model_cost={_key: keep_registered_mode(_key, shared_model_info) for _key in backend_keys}, persist_across_reloads=False, ) diff --git a/litellm/utils.py b/litellm/utils.py index b37827a3b13..d6bdee4697b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3520,19 +3520,33 @@ def _get_builtin_model_info_for_registration(model: str, custom_llm_provider: st ) +def _skips_builtin_registration_lookup(key: str, provider: object) -> bool: + return provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO or any( + key.startswith(f"{p}/") for p in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO + ) + + @dataclass(frozen=True, slots=True) class _RegistrationWrite: key: str entry: Mapping[str, object] +def model_cost_registration_row(key: str, provider: object) -> Mapping[str, object]: + """The row a ``register_model`` write of a ``key`` entry with this ``litellm_provider`` merges onto: the + ``litellm.model_cost`` row it resolves to (``openai/o3-pro`` lands on ``o3-pro``), else the row it seeds ``key`` + with from the built-in entry it falls back to (an OpenAI fine-tune takes its base model's row).""" + base: Final = _registration_base( + key, {"litellm_provider": provider}, custom_llm_provider=None, warning_display_name=None + ) + return _builtin_cost_map().get(base.key) or base.entry + + def _registration_base( key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None ) -> _RegistrationWrite: provider: Final = str(value.get("litellm_provider") or "") - if provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO or any( - key.startswith(f"{p}/") for p in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO - ): + if _skips_builtin_registration_lookup(key, provider): return _RegistrationWrite(key=key, entry=dict(_builtin_cost_map().get(key, {}))) builtin_model_info: Final = _get_builtin_model_info_for_registration( model=key, custom_llm_provider=custom_llm_provider diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6ebb54c0035..13210a7eb1c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10904,7 +10904,6 @@ }, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ - "/v1/chat/completions", "/v1/batch", "/v1/responses" ], @@ -10973,7 +10972,6 @@ "output_cost_per_token_batches": 4e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ - "/v1/chat/completions", "/v1/batch", "/v1/responses" ], @@ -11005,7 +11003,6 @@ "output_cost_per_token_batches": 4e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supported_endpoints": [ - "/v1/chat/completions", "/v1/batch", "/v1/responses" ], diff --git a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py index 9b07cce2f3a..1e599103a27 100644 --- a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py @@ -250,6 +250,7 @@ def _deployment( *, base: str = _OPENAI_BASE, api_key: str = _API_KEY, + model_info: Mapping[str, JsonValue] | None = _CHAT_INFO, **extra: JsonValue, ) -> str: reasoning: Final[Mapping[str, JsonValue]] = ( @@ -260,7 +261,7 @@ def _deployment( custom_llm_provider="bedrock_mantle", api_base=wire.url + base, api_key=api_key, - model_info=_CHAT_INFO, + model_info=model_info, **reasoning, **extra, ) @@ -429,7 +430,7 @@ def test_stream_with_tools_carries_the_deployment_limit_as_max_completion_tokens text: Final = _stream_text( gateway, "/v1/chat/completions", - _chat_body(model, marker, stream=True, stream_options={"include_usage": True}, tools=_TOOLS), + _chat_body(model, marker, stream=True, stream_options={"include_usage": True}, tools=list(_TOOLS)), ) assert f"marker-{marker}" in text and text.rstrip().endswith("data: [DONE]"), text (attempt,) = _attempts(wire) @@ -471,6 +472,22 @@ async def test_async_anthropic_sdk_messages_stream_reaches_mantle_as_max_complet _assert_logged(message.id, model, "success") +def test_a_chat_mode_deployment_leaves_its_unpinned_sibling_on_the_responses_api(gateway: Gateway) -> None: + pinned_marker: Final = uuid.uuid4().hex + sibling_marker: Final = uuid.uuid4().hex + with wire_server(_mantle) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _GPT56) + sibling: Final = _deployment(scenario, wire, _GPT56, model_info=None) + for model, marker in ((pinned, pinned_marker), (sibling, sibling_marker)): + response = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert response.status_code == 200, response.text + assert _answer(marker) in response.text, response.text + pinned_attempt, sibling_attempt = _attempts(wire) + assert (pinned_attempt.target, _marker_in(pinned_attempt)) == (_CHAT_TARGET, pinned_marker), pinned_attempt + assert sibling_attempt.target in _RESPONSES_TARGETS, sibling_attempt + assert _marker_in(sibling_attempt) == sibling_marker, sibling_attempt + + def test_responses_max_output_tokens_is_forwarded_natively_on_both_legs(gateway: Gateway) -> None: unary: Final = uuid.uuid4().hex streamed: Final = uuid.uuid4().hex @@ -500,7 +517,6 @@ def test_responses_max_output_tokens_is_forwarded_natively_on_both_legs(gateway: _assert_logged(_managed_response_id(streamed_issued), model, "success") -@pytest.mark.skip(reason="unrealistic input") def test_region_prefixed_deployment_on_a_v1_base_sends_the_bare_model_with_max_completion_tokens( gateway: Gateway, ) -> None: diff --git a/tests/integration/routing/test_deployment_mode_routing_wire.py b/tests/integration/routing/test_deployment_mode_routing_wire.py new file mode 100644 index 00000000000..65e64883b18 --- /dev/null +++ b/tests/integration/routing/test_deployment_mode_routing_wire.py @@ -0,0 +1,1037 @@ +import asyncio +import itertools +import json +import re +import threading +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import anthropic +import httpx +import openai +import psutil +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.responses_vendor import same_response +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_API_KEY: Final = "synthetic-routing-bearer" +_MANTLE_GPT: Final = "bedrock_mantle/openai.gpt-5.6-terra" +_MANTLE_OSS: Final = "bedrock_mantle/openai.gpt-oss-120b" +_OPENAI_GPT: Final = "openai/gpt-5.4" +_PINNED: Final = "routing-chat-pinned" +_SIBLING: Final = "routing-unpinned-sibling" +_NO_RETRIES: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0}) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_SPEND_BY_MODEL_GROUP: Final = 'SELECT request_id, status, model_group FROM "LiteLLM_SpendLogs" WHERE model_group = %s' +_N_REJECTED_BY_MANTLE: Final = "'n' is not supported for values greater than 1" +_FUNCTION_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + } +) + +Api = Literal["chat", "responses"] +Client = Literal["openai-sdk", "openai-sdk-async-stream", "anthropic-sdk", "anthropic-sdk-async-stream", "httpx-stream"] +Endpoint = Literal["chat", "messages", "responses"] + + +@dataclass(frozen=True, slots=True) +class _Attempt: + api: Api + marker: str + body: Mapping[str, JsonValue] + + +@dataclass(frozen=True, slots=True) +class _Answer: + status: int + identity: str + call_id: str + text: str + + +@dataclass(frozen=True, slots=True) +class _Call: + model: str + endpoint: Endpoint + stream: bool + marker: str + + +def _question(marker: str) -> str: + return f"Question marker-{marker}" + + +def _answer(marker: str) -> str: + return f"answer marker-{marker}" + + +def _marker_in(text: str) -> str: + found: Final = _MARKER.search(text) + assert found is not None, text + return found.group(1) + + +def _api(target: str) -> Api: + if target.endswith("/chat/completions"): + return "chat" + assert target.endswith("/responses"), target + return "responses" + + +def _error(status: int, message: str, kind: str) -> Reply: + return Reply(status=status, body=json.dumps({"error": {"message": message, "type": kind, "code": None}}).encode()) + + +def _upstream(down: threading.Event | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(body=b'{"object": "list", "data": []}') + if request.headers.get("authorization") != f"Bearer {_API_KEY}": + return _error(401, "Invalid bearer token.", "authentication_error") + if down is not None and down.is_set(): + return _error(503, "The server is overloaded.", "server_error") + body: Final = _JSON_OBJECT.validate_json(request.body) + marker: Final = _marker_in(request.body.decode()) + stream: Final = body.get("stream") is True + backend: Final = string_value(body["model"]) + if _api(request.target) == "responses": + return responses_reply(f"resp_{marker}", backend, _answer(marker), stream=stream) + n: Final = body.get("n") + if isinstance(n, int) and n > 1: + return _error(400, _N_REJECTED_BY_MANTLE, "invalid_request_error") + return chat_reply(f"chatcmpl-{marker}", backend, _answer(marker), stream=stream) + + return respond + + +def _is_probe(request: Request) -> bool: + return request.method == "POST" and _MARKER.search(request.body.decode()) is None + + +def _attempts(wire: Wire, *, skip_probes: bool = False) -> tuple[_Attempt, ...]: + return tuple( + _Attempt(_api(request.target), _marker_in(request.body.decode()), _JSON_OBJECT.validate_json(request.body)) + for request in wire.drain() + if request.method == "POST" and not (skip_probes and _is_probe(request)) + ) + + +def _routes(attempts: tuple[_Attempt, ...]) -> Mapping[str, tuple[Api, ...]]: + markers: Final = frozenset(attempt.marker for attempt in attempts) + return {marker: tuple(attempt.api for attempt in attempts if attempt.marker == marker) for marker in markers} + + +def _api_base(wire: Wire, backend: str) -> str: + return wire.url + ("/openai/v1" if backend.startswith("bedrock_mantle/") else "/v1") + + +def _deployment( + scenario: Scenario, + wire: Wire, + backend: str, + mode: str | None, + *, + provider_pin: bool = True, + api_key: str = _API_KEY, +) -> str: + provider: Final[Mapping[str, JsonValue]] = {"custom_llm_provider": backend.split("/", 1)[0]} if provider_pin else {} + return scenario.model( + model=backend, + api_base=_api_base(wire, backend), + api_key=api_key, + model_info=None if mode is None else {"mode": mode}, + **provider, + ) + + +def _wildcard(scenario: Scenario, wire: Wire, mode: str | None) -> str: + prefix: Final = f"wild-{uuid.uuid4().hex}" + created: Final = scenario.gateway.post( + "/model/new", + { + "model_name": f"{prefix}/*", + "litellm_params": {"model": "bedrock_mantle/*", "api_base": wire.url + "/openai/v1", "api_key": _API_KEY}, + "model_info": {} if mode is None else {"mode": mode}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return f"{prefix}/{_MANTLE_GPT.removeprefix('bedrock_mantle/')}" + + +def _raw_model_info_deployment(scenario: Scenario, wire: Wire, raw_model_info: str) -> str: + name: Final = f"integration-{uuid.uuid4().hex}" + params: Final = json.dumps( + { + "model": _MANTLE_GPT, + "custom_llm_provider": "bedrock_mantle", + "api_base": wire.url + "/openai/v1", + "api_key": _API_KEY, + } + ) + response: Final = scenario.gateway.client.post( + "/model/new", + content=f'{{"model_name": "{name}", "litellm_params": {params}, "model_info": {raw_model_info}}}', + headers={"Authorization": f"Bearer {scenario.gateway.key}", "Content-Type": "application/json"}, + ) + assert response.status_code == 200, response.text + identity: Final = string_value(object_value(_JSON_OBJECT.validate_json(response.content)["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + return name + + +def _proxy_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _messages(marker: str) -> list[JsonValue]: + return [{"role": "user", "content": _question(marker)}] + + +def _identity(endpoint: Endpoint, stream: bool, text: str) -> str: + if not stream: + return string_value(_JSON_OBJECT.validate_json(text)["id"]) + first: Final = _JSON_OBJECT.validate_json( + next(line for line in text.splitlines() if line.startswith("data: {")).removeprefix("data: ") + ) + match endpoint: + case "chat": + return string_value(first["id"]) + case "messages": + return string_value(object_value(first["message"])["id"]) + case "responses": + return string_value(object_value(first["response"])["id"]) + + +def _body(call: _Call, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + match call.endpoint: + case "chat": + return {"model": call.model, "messages": _messages(call.marker), "stream": call.stream, **extra} + case "messages": + return {"model": call.model, "max_tokens": 64, "messages": _messages(call.marker), "stream": call.stream} + case "responses": + return {"model": call.model, "input": _question(call.marker), "stream": call.stream, **extra} + + +def _path(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +async def _send(client: httpx.AsyncClient, key: str, call: _Call, extra: Mapping[str, JsonValue]) -> _Answer: + async with client.stream( + "POST", + _path(call.endpoint), + json=_body(call, extra), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + text: Final = (await response.aread()).decode() + identity: Final = _identity(call.endpoint, call.stream, text) if response.status_code == 200 else "" + call_id: Final = response.headers["x-litellm-call-id"] if "x-litellm-call-id" in response.headers else "" + return _Answer(response.status_code, identity, call_id, text) + + +async def _http(gateway: Gateway, call: _Call, extra: Mapping[str, JsonValue] = _NO_RETRIES) -> _Answer: + async with httpx.AsyncClient(base_url=_proxy_url(gateway), timeout=60, trust_env=False) as client: + return await _send(client, gateway.key, call, extra) + + +async def _burst( + base_url: str, key: str, calls: tuple[_Call, ...], *, tolerate_drops: bool = False +) -> tuple[_Answer | None, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, call, _NO_RETRIES) for call in calls), return_exceptions=tolerate_drops + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result if isinstance(result, _Answer) else None for result in results) + + +def _openai_sdk(gateway: Gateway, model: str, marker: str) -> _Answer: + with openai.OpenAI(base_url=f"{_proxy_url(gateway)}/v1", api_key=gateway.key, max_retries=0) as sdk: + completion: Final = sdk.chat.completions.create( + model=model, messages=[{"role": "user", "content": _question(marker)}] + ) + return _Answer(200, completion.id, "", completion.choices[0].message.content or "") + + +async def _openai_sdk_async_stream(gateway: Gateway, model: str, marker: str) -> _Answer: + async with openai.AsyncOpenAI(base_url=f"{_proxy_url(gateway)}/v1", api_key=gateway.key, max_retries=0) as sdk: + stream: Final = await sdk.chat.completions.create( + model=model, messages=[{"role": "user", "content": _question(marker)}], stream=True + ) + chunks: Final = [chunk async for chunk in stream] + (identity,) = {chunk.id for chunk in chunks} + return _Answer( + 200, identity, "", "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + ) + + +def _anthropic_sdk(gateway: Gateway, model: str, marker: str) -> _Answer: + with anthropic.Anthropic(base_url=_proxy_url(gateway), api_key=gateway.key, max_retries=0) as sdk: + message: Final = sdk.messages.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": _question(marker)}] + ) + return _Answer(200, message.id, "", "".join(block.text for block in message.content if block.type == "text")) + + +async def _anthropic_sdk_async_stream(gateway: Gateway, model: str, marker: str) -> _Answer: + async with anthropic.AsyncAnthropic(base_url=_proxy_url(gateway), api_key=gateway.key, max_retries=0) as sdk: + async with sdk.messages.stream( + model=model, max_tokens=64, messages=[{"role": "user", "content": _question(marker)}] + ) as stream: + pieces: Final = [piece async for piece in stream.text_stream] + message: Final = await stream.get_final_message() + return _Answer(200, message.id, "", "".join(pieces)) + + +async def _ask(gateway: Gateway, client: Client, model: str, marker: str) -> _Answer: + match client: + case "openai-sdk": + return _openai_sdk(gateway, model, marker) + case "openai-sdk-async-stream": + return await _openai_sdk_async_stream(gateway, model, marker) + case "anthropic-sdk": + return _anthropic_sdk(gateway, model, marker) + case "anthropic-sdk-async-stream": + return await _anthropic_sdk_async_stream(gateway, model, marker) + case "httpx-stream": + return await _http(gateway, _Call(model, "chat", True, marker)) + + +def _logged_ids(answer: _Answer) -> tuple[str, ...]: + markers: Final = frozenset(_MARKER.findall(answer.text)) + upstream: Final = (*(f"resp_{marker}" for marker in markers), *(f"chatcmpl-{marker}" for marker in markers)) + return tuple(shown for shown in (answer.identity, answer.call_id, *upstream) if shown) + + +def _logs_answer(request_id: JsonValue, answer: _Answer) -> bool: + return any(same_response(string_value(request_id), shown) for shown in _logged_ids(answer)) + + +def _logged_rows(answer: _Answer, model: str) -> list[dict[str, JsonValue]]: + return [row for row in read_rows(_SPEND_BY_MODEL_GROUP, (model,)) if _logs_answer(row["request_id"], answer)] + + +def _assert_logged(answer: _Answer, model: str, status: str) -> None: + rows: Final = eventually(lambda: _logged_rows(answer, model), lambda found: len(found) >= 1, seconds=70) + assert len(rows) == 1, rows + assert (rows[0]["status"], rows[0]["model_group"]) == (status, model), rows + + +def _assert_answered(answer: _Answer | None, marker: str) -> _Answer: + assert answer is not None, marker + assert answer.status == 200, answer.text + assert _answer(marker) in answer.text, answer.text + assert set(_MARKER.findall(answer.text)) == {marker}, answer.text + return answer + + +def _waves(model: str, count: int) -> tuple[_Call, ...]: + endpoints: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") + return tuple( + _Call(model, endpoints[index % len(endpoints)], index % 2 == 0, uuid.uuid4().hex) for index in range(count) + ) + + +def _expected_route(call: _Call, pinned: str) -> Api: + return "chat" if call.model == pinned and call.endpoint != "responses" else "responses" + + +@pytest.mark.parametrize( + "client", ("openai-sdk", "openai-sdk-async-stream", "anthropic-sdk", "anthropic-sdk-async-stream", "httpx-stream") +) +async def test_a_chat_pin_leaves_its_unpinned_sibling_on_responses_on_every_client( + gateway: Gateway, client: Client +) -> None: + pinned_marker: Final = uuid.uuid4().hex + sibling_marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat") + sibling: Final = _deployment(scenario, wire, _MANTLE_GPT, None) + pinned_answer: Final = _assert_answered(await _ask(gateway, client, pinned, pinned_marker), pinned_marker) + sibling_answer: Final = _assert_answered(await _ask(gateway, client, sibling, sibling_marker), sibling_marker) + assert _routes(_attempts(wire)) == {pinned_marker: ("chat",), sibling_marker: ("responses",)} + _assert_logged(pinned_answer, pinned, "success") + _assert_logged(sibling_answer, sibling, "success") + + +@pytest.mark.parametrize("backend", (_MANTLE_OSS, _OPENAI_GPT)) +async def test_a_responses_pin_leaves_its_unpinned_sibling_on_chat(gateway: Gateway, backend: str) -> None: + pinned_marker: Final = uuid.uuid4().hex + sibling_marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, backend, "responses") + sibling: Final = _deployment(scenario, wire, backend, None) + pinned_answer: Final = _assert_answered( + await _http(gateway, _Call(pinned, "chat", False, pinned_marker)), pinned_marker + ) + sibling_answer: Final = _assert_answered( + await _http(gateway, _Call(sibling, "chat", False, sibling_marker)), sibling_marker + ) + assert _routes(_attempts(wire)) == {pinned_marker: ("responses",), sibling_marker: ("chat",)} + _assert_logged(pinned_answer, pinned, "success") + _assert_logged(sibling_answer, sibling, "success") + + +@pytest.mark.parametrize("stream", (False, True)) +async def test_the_responses_endpoint_serves_a_chat_pinned_deployment_natively(gateway: Gateway, stream: bool) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat") + answer: Final = _assert_answered(await _http(gateway, _Call(pinned, "responses", stream, marker)), marker) + assert _routes(_attempts(wire)) == {marker: ("responses",)} + _assert_logged(answer, pinned, "success") + + +async def test_a_wildcard_route_keeps_the_registered_mode_beside_a_chat_pin(gateway: Gateway) -> None: + pinned_marker: Final = uuid.uuid4().hex + wildcard_marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat") + wildcard: Final = _wildcard(scenario, wire, None) + _assert_answered(await _http(gateway, _Call(pinned, "chat", False, pinned_marker)), pinned_marker) + _assert_answered(await _http(gateway, _Call(wildcard, "chat", False, wildcard_marker)), wildcard_marker) + assert _routes(_attempts(wire)) == {pinned_marker: ("chat",), wildcard_marker: ("responses",)} + + +@pytest.mark.parametrize("shape", ("wildcard", "no-provider-pin")) +async def test_a_chat_pin_without_a_provider_pin_gets_chat_completions(gateway: Gateway, shape: str) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + model: Final = ( + _wildcard(scenario, wire, "chat") + if shape == "wildcard" + else _deployment(scenario, wire, _MANTLE_GPT, "chat", provider_pin=False) + ) + _assert_answered(await _http(gateway, _Call(model, "chat", False, marker)), marker) + assert _routes(_attempts(wire)) == {marker: ("chat",)} + + +async def test_an_explicit_responses_prefix_wins_over_a_chat_pin(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-5.4", api_base=wire.url + "/v1", api_key=_API_KEY, model_info={"mode": "chat"} + ) + answer: Final = _assert_answered(await _http(gateway, _Call(model, "chat", False, marker)), marker) + (attempt,) = _attempts(wire) + assert (attempt.api, attempt.body["model"]) == ("responses", "gpt-5.4"), attempt + _assert_logged(answer, model, "success") + + +_SUMMARY_REASONING: Final[Mapping[str, JsonValue]] = MappingProxyType({"effort": "low", "summary": "auto"}) + + +@pytest.mark.parametrize( + ("extra", "reasoning"), + ( + ({"tools": [dict(_FUNCTION_TOOL)]}, None), + ({"reasoning_effort": dict(_SUMMARY_REASONING)}, _SUMMARY_REASONING), + ({"extra_body": {"reasoning_effort": dict(_SUMMARY_REASONING)}}, _SUMMARY_REASONING), + ({"web_search_options": {"search_context_size": "low"}}, None), + ), + ids=("function-tools", "reasoning-summary", "extra-body-reasoning-summary", "web-search"), +) +async def test_a_mantle_chat_pin_sends_what_chat_cannot_serve_to_responses( + gateway: Gateway, extra: Mapping[str, JsonValue], reasoning: Mapping[str, JsonValue] | None +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat") + answer: Final = _assert_answered( + await _http(gateway, _Call(pinned, "chat", False, marker), {**_NO_RETRIES, **extra}), marker + ) + (attempt,) = _attempts(wire) + assert attempt.api == "responses", attempt + assert reasoning is None or attempt.body["reasoning"] == reasoning, attempt.body + _assert_logged(answer, pinned, "success") + + +def _placed(placement: str, params: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + return params if placement == "top-level" else {"extra_body": dict(params)} + + +async def _messages_with_thinking( + gateway: Gateway, model: str, marker: str, thinking: Mapping[str, JsonValue] +) -> _Answer: + async with httpx.AsyncClient(base_url=_proxy_url(gateway), timeout=60, trust_env=False) as client: + response: Final = await client.post( + "/v1/messages", + json={"model": model, "max_tokens": 2048, "thinking": dict(thinking), "messages": _messages(marker)}, + headers={"Authorization": f"Bearer {gateway.key}", "anthropic-version": "2023-06-01"}, + ) + identity: Final = _identity("messages", False, response.text) if response.status_code == 200 else "" + return _Answer(response.status_code, identity, response.headers.get("x-litellm-call-id", ""), response.text) + + +@pytest.mark.parametrize(("summary", "route"), ((None, "chat"), ("auto", "responses")), ids=("plain", "summary")) +async def test_a_mantle_chat_pin_moves_messages_thinking_to_responses_only_when_it_asks_for_a_summary( + gateway: Gateway, summary: str | None, route: Api +) -> None: + marker: Final = uuid.uuid4().hex + thinking: Final[Mapping[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024} | ( + {} if summary is None else {"summary": summary} + ) + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat") + answer: Final = _assert_answered(await _messages_with_thinking(gateway, pinned, marker, thinking), marker) + (attempt,) = _attempts(wire) + assert attempt.api == route, attempt + match route: + case "chat": + assert "summary" not in json.dumps(attempt.body), attempt.body + case "responses": + assert object_value(attempt.body["reasoning"])["summary"] == summary, attempt.body + _assert_logged(answer, pinned, "success") + + +def _sent_effort(attempt: _Attempt) -> JsonValue: + match attempt.api: + case "chat": + return attempt.body.get("reasoning_effort") + case "responses": + return object_value(attempt.body["reasoning"]).get("effort") + + +@pytest.mark.parametrize( + ("backend", "mode", "effort", "route"), + ((_OPENAI_GPT, None, "low", "responses"), (_MANTLE_GPT, "chat", "none", "chat")), + ids=("openai-low", "mantle-chat-pin-none"), +) +@pytest.mark.parametrize("placement", ("top-level", "extra-body")) +async def test_an_effort_inside_extra_body_routes_function_tools_like_a_top_level_one( + gateway: Gateway, backend: str, mode: str | None, effort: str, route: Api, placement: str +) -> None: + marker: Final = uuid.uuid4().hex + extra: Final[Mapping[str, JsonValue]] = { + **_NO_RETRIES, + "tools": [dict(_FUNCTION_TOOL)], + **_placed(placement, {"reasoning_effort": effort}), + } + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, backend, mode) + answer: Final = _assert_answered(await _http(gateway, _Call(model, "chat", False, marker), extra), marker) + (attempt,) = _attempts(wire) + assert (attempt.api, _sent_effort(attempt)) == (route, effort), attempt + _assert_logged(answer, model, "success") + + +_AUTO_SUMMARY_LEVELS: Final[Mapping[str, str]] = MappingProxyType( + {"auto-summary-mantle": "auto", "auto-summary-openai": "detailed"} +) +_PLAIN_THINKING: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "enabled", "budget_tokens": 1024}) + + +def _auto_summary_config(wire: Wire, tmp_path: Path) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + mantle: Final[dict[str, JsonValue]] = { + "model": _MANTLE_GPT, + "custom_llm_provider": "bedrock_mantle", + "api_base": wire.url + "/openai/v1", + "api_key": _API_KEY, + } + openai_codex: Final[dict[str, JsonValue]] = { + "model": "openai/gpt-5.3-codex", + "api_base": wire.url + "/v1", + "api_key": _API_KEY, + } + config: Final = { + **base, + "model_list": [ + {"model_name": "auto-summary-mantle", "litellm_params": mantle}, + {"model_name": "auto-summary-openai", "litellm_params": openai_codex}, + ], + "router_settings": {**object_value(base["router_settings"]), "num_retries": 0}, + } + path: Final = tmp_path / "deployment-mode-routing-auto-summary.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +async def _ask_for_reasoning(gateway: Gateway, call: _Call) -> _Answer: + if call.endpoint == "messages": + return await _messages_with_thinking(gateway, call.model, call.marker, _PLAIN_THINKING) + return await _http(gateway, call, {**_NO_RETRIES, "reasoning_effort": "low"}) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +async def test_the_auto_summary_asks_each_target_for_a_level_it_accepts_on_every_endpoint( + gateway: Gateway, tmp_path: Path +) -> None: + endpoints: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") + calls: Final = tuple( + _Call(model, endpoint, False, uuid.uuid4().hex) + for model, endpoint in itertools.product(_AUTO_SUMMARY_LEVELS, endpoints) + ) + with wire_server(_upstream()) as wire: + config: Final = _auto_summary_config(wire, tmp_path) + with owned_proxy_process(gateway, tmp_path, {"LITELLM_REASONING_AUTO_SUMMARY": "true"}, config=config) as owned: + for call in calls: + _assert_answered(await _ask_for_reasoning(owned.gateway, call), call.marker) + sent: Final = { + attempt.marker: (attempt.api, object_value(attempt.body["reasoning"]).get("summary")) + for attempt in _attempts(wire) + } + assert sent == {call.marker: ("responses", _AUTO_SUMMARY_LEVELS[call.model]) for call in calls} + + +def _health_upstream(request: Request) -> Reply: + if request.method != "POST": + return Reply(body=b'{"object": "list", "data": []}') + backend: Final = string_value(_JSON_OBJECT.validate_json(request.body)["model"]) + if _api(request.target) == "responses": + return responses_reply("resp_health", backend, "ok", stream=False) + return chat_reply("chatcmpl-health", backend, "ok", stream=False) + + +@pytest.mark.parametrize(("mode", "route"), (("chat", "chat"), (None, "responses")), ids=("chat-pin", "unpinned")) +def test_a_health_check_probes_the_api_its_deployment_sends_chat_to( + gateway: Gateway, mode: str | None, route: Api +) -> None: + with wire_server(_health_upstream) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, _MANTLE_GPT, mode) + response: Final = gateway.request("GET", "/health", params={"model": model}) + probed: Final = tuple(_api(request.target) for request in wire.drain() if request.method == "POST") + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["healthy_count"] == 1, response.text + assert probed == (route,), probed + + +_INPUT_PRICE: Final = 0.001 +_OUTPUT_PRICE: Final = 0.002 +_SPEND_OF_MODEL_GROUP: Final = 'SELECT request_id, spend FROM "LiteLLM_SpendLogs" WHERE model_group = %s' + + +def _priced_chat_pin(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=_MANTLE_GPT, + custom_llm_provider="bedrock_mantle", + api_base=_api_base(wire, _MANTLE_GPT), + api_key=_API_KEY, + input_cost_per_token=_INPUT_PRICE, + output_cost_per_token=_OUTPUT_PRICE, + model_info={"mode": "chat"}, + ) + + +def _probe_or_marked(request: Request) -> Reply: + if _is_probe(request): + return _health_upstream(request) + return _upstream()(request) + + +async def _answered_burst(gateway: Gateway, model: str, count: int) -> tuple[tuple[_Call, _Answer], ...]: + calls: Final = tuple(_Call(model, "chat", False, uuid.uuid4().hex) for _ in range(count)) + answers: Final = await _burst(_proxy_url(gateway), gateway.key, calls) + return tuple((call, _assert_answered(answer, call.marker)) for call, answer in zip(calls, answers, strict=True)) + + +async def test_a_custom_priced_chat_pin_bills_its_price_and_leaves_its_unpinned_sibling_on_responses( + gateway: Gateway, +) -> None: + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _priced_chat_pin(scenario, wire) + sibling: Final = _deployment(scenario, wire, _MANTLE_GPT, None) + pinned_calls: Final = await _answered_burst(gateway, pinned, 8) + sibling_calls: Final = await _answered_burst(gateway, sibling, 8) + assert _routes(_attempts(wire)) == {call.marker: ("chat",) for call, _ in pinned_calls} | { + call.marker: ("responses",) for call, _ in sibling_calls + } + _, answer = pinned_calls[0] + rows: Final = eventually( + lambda: [ + row for row in read_rows(_SPEND_OF_MODEL_GROUP, (pinned,)) if _logs_answer(row["request_id"], answer) + ], + lambda found: len(found) == 1, + seconds=70, + ) + assert float(str(rows[0]["spend"])) == pytest.approx(5 * _INPUT_PRICE + 3 * _OUTPUT_PRICE), rows + + +async def test_a_health_probe_of_a_custom_priced_chat_pin_leaves_its_unpinned_sibling_on_responses( + gateway: Gateway, +) -> None: + with wire_server(_probe_or_marked) as wire, gateway.scenario() as scenario: + pinned: Final = _priced_chat_pin(scenario, wire) + sibling: Final = _deployment(scenario, wire, _MANTLE_GPT, None) + async with httpx.AsyncClient(base_url=_proxy_url(gateway), timeout=60, trust_env=False) as client: + probes: Final = await asyncio.gather( + *( + client.get("/health", params={"model": pinned}, headers={"Authorization": f"Bearer {gateway.key}"}) + for _ in range(6) + ) + ) + probed: Final = tuple( + _api(request.target) + for request in wire.drain() + if request.method == "POST" and _MARKER.search(request.body.decode()) is None + ) + sibling_calls: Final = await _answered_burst(gateway, sibling, 8) + assert _routes(_attempts(wire)) == {call.marker: ("responses",) for call, _ in sibling_calls} + assert [(probe.status_code, _JSON_OBJECT.validate_json(probe.content)["healthy_count"]) for probe in probes] == [ + (200, 1) + ] * 6, [probe.text for probe in probes] + assert probed == ("chat",) * 6, probed + + +_PROBED: Final = "routing-probed-deployment" +_BACKGROUND_HEALTH_CHECKS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"background_health_checks": True, "health_check_interval": 1} +) + + +def _failing_probes(probes: SimpleQueue[Api]) -> Callable[[Request], Reply]: + respond: Final = _upstream() + + def fail_probe(request: Request) -> Reply: + if not _is_probe(request): + return respond(request) + probes.put(_api(request.target)) + return _error(404, "The model `openai.gpt-5.6-terra` does not exist.", "invalid_request_error") + + return fail_probe + + +def _probed_deployment_config(wire: Wire, tmp_path: Path, mode: str | None, general: Mapping[str, JsonValue]) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + params: Final[dict[str, JsonValue]] = { + "model": _MANTLE_GPT, + "custom_llm_provider": "bedrock_mantle", + "api_base": wire.url + "/openai/v1", + "api_key": _API_KEY, + } + pin: Final[dict[str, JsonValue]] = {} if mode is None else {"model_info": {"mode": mode}} + config: Final = { + **base, + "model_list": [{"model_name": _PROBED, "litellm_params": params, **pin}], + "general_settings": {**object_value(base["general_settings"]), "store_model_in_db": False, **general}, + "router_settings": { + **object_value(base["router_settings"]), + "disable_cooldowns": False, + "cooldown_time": 60, + "num_retries": 0, + }, + } + path: Final = tmp_path / "deployment-mode-routing-failing-probe.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _drained(probes: SimpleQueue[Api]) -> tuple[Api, ...]: + return tuple(probes.get() for _ in range(probes.qsize())) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize(("mode", "route"), (("chat", "chat"), (None, "responses")), ids=("chat-pin", "unpinned")) +async def test_a_failing_health_probe_leaves_its_deployment_serving_the_next_requests( + gateway: Gateway, tmp_path: Path, mode: str | None, route: Api +) -> None: + probes: Final[SimpleQueue[Api]] = SimpleQueue() + calls: Final = tuple(_Call(_PROBED, "chat", False, uuid.uuid4().hex) for _ in range(4)) + with wire_server(_failing_probes(probes)) as wire: + config: Final = _probed_deployment_config(wire, tmp_path, mode, {}) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + health: Final = owned.gateway.request("GET", "/health", params={"model": _PROBED}) + answers: Final = [await _http(owned.gateway, call) for call in calls] + served: Final = _routes(_attempts(wire, skip_probes=True)) + counts: Final = _JSON_OBJECT.validate_json(health.content) + assert (counts["healthy_count"], counts["unhealthy_count"]) == (0, 1), health.text + assert _drained(probes) == (route,) + for call, answer in zip(calls, answers, strict=True): + _assert_answered(answer, call.marker) + assert served == {call.marker: (route,) for call in calls} + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +async def test_failing_background_health_probes_leave_a_chat_pinned_deployment_serving_requests( + gateway: Gateway, tmp_path: Path +) -> None: + probes: Final[SimpleQueue[Api]] = SimpleQueue() + calls: Final = tuple(_Call(_PROBED, "chat", False, uuid.uuid4().hex) for _ in range(4)) + with wire_server(_failing_probes(probes)) as wire: + config: Final = _probed_deployment_config(wire, tmp_path, "chat", _BACKGROUND_HEALTH_CHECKS) + with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned: + await asyncio.to_thread(eventually, probes.qsize, lambda size: size >= 2, 60) + answers: Final = [await _http(owned.gateway, call) for call in calls] + served: Final = _routes(_attempts(wire, skip_probes=True)) + probed: Final = _drained(probes) + assert len(probed) >= 2 and set(probed) == {"chat"}, probed + for call, answer in zip(calls, answers, strict=True): + _assert_answered(answer, call.marker) + assert served == {call.marker: ("chat",) for call in calls} + + +@pytest.mark.parametrize( + "extra_body", + ( + {"reasoning_effort": "low", "reasoningSummary": "auto"}, + {"reasoning_effort": {"effort": "low"}, "reasoning_summary": "auto"}, + ), + ids=("string-effort", "object-effort"), +) +@pytest.mark.parametrize( + ("backend", "mode"), ((_OPENAI_GPT, None), (_MANTLE_GPT, "chat")), ids=("openai", "mantle-chat-pin") +) +async def test_a_summary_alias_beside_an_extra_body_effort_reaches_the_responses_api( + gateway: Gateway, extra_body: Mapping[str, JsonValue], backend: str, mode: str | None +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, backend, mode) + answer: Final = _assert_answered( + await _http(gateway, _Call(model, "chat", False, marker), {**_NO_RETRIES, "extra_body": dict(extra_body)}), + marker, + ) + (attempt,) = _attempts(wire) + assert (attempt.api, attempt.body.get("reasoning")) == ("responses", {"effort": "low", "summary": "auto"}), ( + attempt + ) + _assert_logged(answer, model, "success") + + +@pytest.mark.parametrize("placement", ("top-level", "extra-body")) +async def test_a_mantle_chat_pin_rejects_n_above_one_before_the_upstream(gateway: Gateway, placement: str) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat") + answer: Final = await _http( + gateway, _Call(pinned, "chat", False, marker), {**_NO_RETRIES, **_placed(placement, {"n": 2})} + ) + assert answer.status == 400, answer.text + assert "does not support n=2" in answer.text and "drop_params" in answer.text, answer.text + assert _attempts(wire) == () + _assert_logged(answer, pinned, "failure") + + +@pytest.mark.parametrize("placement", ("top-level", "extra-body")) +async def test_a_mantle_chat_pin_drops_n_above_one_under_drop_params(gateway: Gateway, placement: str) -> None: + marker: Final = uuid.uuid4().hex + extra: Final[Mapping[str, JsonValue]] = {**_NO_RETRIES, **_placed(placement, {"n": 2}), "drop_params": True} + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat") + answer: Final = _assert_answered(await _http(gateway, _Call(pinned, "chat", False, marker), extra), marker) + (attempt,) = _attempts(wire) + assert attempt.api == "chat" and "n" not in attempt.body, attempt + _assert_logged(answer, pinned, "success") + + +@pytest.mark.parametrize("placement", ("top-level", "extra-body")) +async def test_a_no_provider_mantle_chat_pin_sends_an_effort_object_as_its_level( + gateway: Gateway, placement: str +) -> None: + marker: Final = uuid.uuid4().hex + effort: Final[Mapping[str, JsonValue]] = {"reasoning_effort": {"effort": "low"}} + extra: Final[Mapping[str, JsonValue]] = {**_NO_RETRIES, **_placed(placement, effort)} + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat", provider_pin=False) + answer: Final = _assert_answered(await _http(gateway, _Call(pinned, "chat", False, marker), extra), marker) + (attempt,) = _attempts(wire) + assert (attempt.api, attempt.body.get("reasoning_effort")) == ("chat", "low"), attempt + _assert_logged(answer, pinned, "success") + + +@pytest.mark.parametrize( + ("raw_model_info", "own_route"), + ( + ('{"mode": ""}', "responses"), + ('{"mode": "embedding"}', "responses"), + ('{"mode": null}', "responses"), + ('{"mode": 5}', None), + ('{"mode": ["chat"]}', None), + ('{"mode": "' + "c" * 5000 + '"}', "responses"), + ('{"mode": "responses", "mode": "chat"}', "chat"), + ), + ids=("empty", "unroutable", "null", "int", "list", "5kb", "twice"), +) +async def test_a_malformed_mode_never_moves_an_unpinned_sibling( + gateway: Gateway, raw_model_info: str, own_route: Api | None +) -> None: + odd_marker: Final = uuid.uuid4().hex + sibling_markers: Final = tuple(uuid.uuid4().hex for _ in range(4)) + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + odd: Final = _raw_model_info_deployment(scenario, wire, raw_model_info) + sibling: Final = _deployment(scenario, wire, _MANTLE_GPT, None) + if own_route is not None: + _assert_answered(await _http(gateway, _Call(odd, "chat", False, odd_marker)), odd_marker) + for marker in sibling_markers: + _assert_answered(await _http(gateway, _Call(sibling, "chat", False, marker)), marker) + own: Final[dict[str, tuple[Api, ...]]] = {} if own_route is None else {odd_marker: (own_route,)} + assert _routes(_attempts(wire)) == own | {marker: ("responses",) for marker in sibling_markers} + + +async def test_a_bridged_sibling_reports_the_upstream_401_after_one_attempt(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + _deployment(scenario, wire, _MANTLE_GPT, "chat") + sibling: Final = _deployment(scenario, wire, _MANTLE_GPT, None, api_key="wrong-bearer") + answer: Final = await _http(gateway, _Call(sibling, "chat", False, marker)) + assert answer.status == 401, answer.text + assert "Invalid bearer token." in answer.text, answer.text + assert _routes(_attempts(wire)) == {marker: ("responses",)} + _assert_logged(answer, sibling, "failure") + + +async def test_adding_and_deleting_a_chat_pin_under_traffic_never_moves_the_sibling(gateway: Gateway) -> None: + with wire_server(_upstream()) as wire, gateway.scenario() as scenario: + sibling: Final = _deployment(scenario, wire, _MANTLE_GPT, None) + before: Final = _waves(sibling, 12) + during: Final = _waves(sibling, 12) + after: Final = _waves(sibling, 12) + first: Final = asyncio.create_task(_burst(_proxy_url(gateway), gateway.key, before)) + with gateway.scenario() as pin_scenario: + await asyncio.to_thread(_deployment, pin_scenario, wire, _MANTLE_GPT, "chat") + answered_before: Final = await first + answered_during: Final = await _burst(_proxy_url(gateway), gateway.key, during) + answered_after: Final = await _burst(_proxy_url(gateway), gateway.key, after) + calls: Final = (*before, *during, *after) + answers: Final = tuple( + _assert_answered(answer, call.marker) + for call, answer in zip(calls, (*answered_before, *answered_during, *answered_after), strict=True) + ) + assert _routes(_attempts(wire)) == {call.marker: ("responses",) for call in calls} + for answer in answers: + _assert_logged(answer, sibling, "success") + + +async def test_an_upstream_outage_mid_burst_reaches_callers_and_routing_recovers(gateway: Gateway) -> None: + down: Final = threading.Event() + with wire_server(_upstream(down)) as wire, gateway.scenario() as scenario: + pinned: Final = _deployment(scenario, wire, _MANTLE_GPT, "chat") + sibling: Final = _deployment(scenario, wire, _MANTLE_GPT, None) + outage: Final = (*_waves(pinned, 10), *_waves(sibling, 10)) + recovery: Final = (*_waves(pinned, 10), *_waves(sibling, 10)) + down.set() + failed: Final = await _burst(_proxy_url(gateway), gateway.key, outage) + down.clear() + recovered: Final = tuple( + _assert_answered(answer, call.marker) + for call, answer in zip(recovery, await _burst(_proxy_url(gateway), gateway.key, recovery), strict=True) + ) + for call, answer in zip(outage, failed, strict=True): + assert answer is not None and answer.status == 503, (call, answer) + assert "The server is overloaded." in answer.text, answer.text + _assert_logged(answer, call.model, "failure") + assert _routes(_attempts(wire)) == { + call.marker: (_expected_route(call, pinned),) for call in (*outage, *recovery) + } + for call, answer in zip(recovery, recovered, strict=True): + _assert_logged(answer, call.model, "success") + + +def _yaml_config(wire: Wire, tmp_path: Path, order: tuple[str, ...], label: str) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + params: Final[dict[str, JsonValue]] = { + "model": _MANTLE_GPT, + "custom_llm_provider": "bedrock_mantle", + "api_base": wire.url + "/openai/v1", + "api_key": _API_KEY, + } + entries: Final[Mapping[str, dict[str, JsonValue]]] = { + _PINNED: {"model_name": _PINNED, "litellm_params": params, "model_info": {"mode": "chat"}}, + _SIBLING: {"model_name": _SIBLING, "litellm_params": params}, + } + config: Final = { + **base, + "model_list": [entries[name] for name in order], + "router_settings": {**object_value(base["router_settings"]), "num_retries": 0}, + } + path: Final = tmp_path / f"deployment-mode-routing-{label}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _holding_upstream(held: SimpleQueue[str], release: threading.Event) -> Callable[[Request], Reply]: + respond: Final = _upstream() + + def hold(request: Request) -> Reply: + if request.method == "POST": + held.put(_marker_in(request.body.decode())) + assert release.wait(timeout=60), "The burst was never released" + return respond(request) + + return hold + + +def _held_upstream_calls(pid: int, wire: Wire) -> int: + port: Final = urlsplit(wire.url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +async def _follow_up(gateway: Gateway) -> tuple[_Call, ...]: + calls: Final = (*_waves(_PINNED, 6), *_waves(_SIBLING, 6)) + answers: Final = await _burst(_proxy_url(gateway), gateway.key, calls) + for call, answer in zip(calls, answers, strict=True): + _assert_logged(_assert_answered(answer, call.marker), call.model, "success") + return calls + + +@pytest.mark.timeout(5 * graceful_stop_seconds() + 120) +async def test_a_worker_kill_and_a_restart_in_either_order_keep_each_deployment_on_its_mode( + gateway: Gateway, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + burst_calls: Final = (*_waves(_PINNED, 10), *_waves(_SIBLING, 10)) + with wire_server(_holding_upstream(held, release)) as wire: + sibling_first: Final = _yaml_config(wire, tmp_path, (_SIBLING, _PINNED), "sibling-first") + with owned_proxy_process(gateway, tmp_path, {}, config=sibling_first, workers=2) as owned: + workers: Final = eventually( + lambda: tuple(int(found.group(1)) for found in _STARTED_WORKER.finditer(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst(_proxy_url(owned.gateway), owned.gateway.key, burst_calls, tolerate_drops=True) + ) + await asyncio.to_thread(eventually, held.qsize, lambda size: size == len(burst_calls), 60) + held_by_worker: Final = MappingProxyType({pid: _held_upstream_calls(pid, wire) for pid in workers}) + assert sum(held_by_worker.values()) == len(burst_calls), held_by_worker + victim: Final = max(workers, key=held_by_worker.__getitem__) + with owned.replacing_worker(victim) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + survivors: Final = tuple( + _assert_answered(answer, call.marker) + for call, answer in zip(burst_calls, served, strict=True) + if answer is not None + ) + assert len(survivors) == len(burst_calls) - held_by_worker[victim], (held_by_worker, served) + after_kill: Final = await _follow_up(owned.gateway) + pinned_first: Final = _yaml_config(wire, tmp_path, (_PINNED, _SIBLING), "pinned-first") + with owned_proxy_process(gateway, tmp_path, {}, config=pinned_first, workers=2) as restarted: + after_restart: Final = await _follow_up(restarted.gateway) + assert _routes(_attempts(wire)) == { + call.marker: (_expected_route(call, _PINNED),) for call in (*burst_calls, *after_kill, *after_restart) + } diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 83c59480d37..9aa494c1aab 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -82,6 +82,115 @@ def test_completion_over_responses_maps_anthropic_search_tool_to_openai_tool_sea assert request_body["tools"][0] == {"type": "tool_search", "execution": "server"} +_MANTLE_RESPONSES_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses" + + +@pytest.mark.parametrize( + "extra_body", + [ + {"reasoning_effort": "low", "reasoningSummary": "auto"}, + {"reasoning_effort": {"effort": "low"}, "reasoning_summary": "auto"}, + ], + ids=["string-effort", "dict-effort"], +) +@pytest.mark.parametrize( + ("model", "url", "model_info"), + [ + ("openai/gpt-5.4", "https://api.openai.com/v1/responses", None), + ("bedrock_mantle/openai.gpt-5.6-terra", _MANTLE_RESPONSES_URL, {"mode": "chat"}), + ], + ids=["openai", "mantle-chat-pin"], +) +def test_a_summary_alias_beside_an_extra_body_effort_reaches_the_responses_api( + monkeypatch: pytest.MonkeyPatch, + extra_body: Mapping[str, object], + model: str, + url: str, + model_info: Mapping[str, object] | None, +) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + response_body: Final = { + "id": "resp_nested_effort_summary", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": model.rsplit("/", 1)[-1], + "output": [ + { + "type": "message", + "id": "msg_nested_effort_summary", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + + with respx.mock(assert_all_called=True) as router: + upstream: Final = router.post(url).mock(return_value=httpx.Response(200, json=response_body)) + litellm.completion( + model=model, + messages=[{"role": "user", "content": "ping"}], + extra_body=dict(extra_body), + model_info=model_info, + api_key="test-key", + aws_region_name="us-east-1", + ) + + assert json.loads(upstream.calls.last.request.content)["reasoning"] == {"effort": "low", "summary": "auto"} + + +@pytest.mark.parametrize( + ("model", "url", "expected_summary"), + [ + ("bedrock_mantle/openai.gpt-5.6-terra", _MANTLE_RESPONSES_URL, "auto"), + ("openai/responses/gpt-5.5", "https://api.openai.com/v1/responses", "detailed"), + ], +) +@pytest.mark.parametrize("endpoint", ["chat_completions", "responses"]) +def test_auto_summary_sends_the_level_the_target_accepts( + monkeypatch: pytest.MonkeyPatch, model: str, url: str, expected_summary: str, endpoint: str +) -> None: + monkeypatch.setattr(litellm, "reasoning_auto_summary", True) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + response_body: Final = { + "id": "resp_auto_summary", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": model.rsplit("/", 1)[-1], + "output": [ + { + "type": "message", + "id": "msg_auto_summary", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + + with respx.mock(assert_all_called=True) as router: + upstream: Final = router.post(url).mock(return_value=httpx.Response(200, json=response_body)) + if endpoint == "chat_completions": + litellm.completion( + model=model, + messages=[{"role": "user", "content": "ping"}], + reasoning_effort="low", + api_key="test-key", + aws_region_name="us-east-1", + ) + else: + litellm.responses( + model=model, input="ping", reasoning_effort="low", api_key="test-key", aws_region_name="us-east-1" + ) + + request_body: Final = json.loads(upstream.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "low", "summary": expected_summary} + + def test_convert_chat_completion_messages_to_responses_api_image_input(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py index 4b96d2e7991..389398dcbf4 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py @@ -49,9 +49,12 @@ def _completion_kwargs_for_thinking( *, tools: list[dict[str, object]] | None = None, api_base: str | None = None, + model_info: dict[str, str] | None = None, ) -> Mapping[str, object]: extra_kwargs: Final = { - key: value for key, value in (("custom_llm_provider", provider), ("api_base", api_base)) if value is not None + key: value + for key, value in (("custom_llm_provider", provider), ("api_base", api_base), ("model_info", model_info)) + if value is not None } completion_kwargs, _ = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( max_tokens=1024, @@ -80,10 +83,11 @@ def _reasoning_effort_sent_for_thinking( *, tools: list[dict[str, object]] | None = None, api_base: str | None = None, + model_info: dict[str, str] | None = None, ) -> object: - return _completion_kwargs_for_thinking(model, provider, thinking, tools=tools, api_base=api_base).get( - "reasoning_effort" - ) + return _completion_kwargs_for_thinking( + model, provider, thinking, tools=tools, api_base=api_base, model_info=model_info + ).get("reasoning_effort") SUMMARIZED_THINKING = {"type": "enabled", "budget_tokens": 4096, "summary": "auto"} @@ -183,6 +187,22 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: assert sent == "high" + def test_a_chat_pinned_deployment_gets_the_plain_tier_on_a_responses_row( + self, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A ``mode: chat`` pin sends that deployment to chat completions even though its shared row + says ``responses``, so the adapter has to read the pin too or chat gets the Responses-only dict.""" + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + model: Final = "bedrock_mantle/openai.gpt-5.4" + + unpinned = _reasoning_effort_sent_for_thinking(model, "bedrock_mantle", PLAIN_THINKING) + pinned = _reasoning_effort_sent_for_thinking( + model, "bedrock_mantle", PLAIN_THINKING, model_info={"id": "pinned-id", "mode": "chat"} + ) + + assert isinstance(unpinned, dict) + assert pinned == "high" + @pytest.mark.parametrize( "model, provider", [ @@ -205,6 +225,23 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: assert sent == {"effort": "high", "summary": "detailed"} + @pytest.mark.parametrize( + "model, provider, summary", + [ + ("bedrock_mantle/openai.gpt-5.4", "bedrock_mantle", "auto"), + ("azure/responses/gpt-5-mini", "azure", "detailed"), + ], + ) + def test_auto_summary_asks_each_bridged_target_for_a_level_it_accepts( + self, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, model: str, provider: str, summary: str + ) -> None: + """Mantle's OpenAI path accepts only ``auto``, so the opt-in summary must not send it ``detailed``.""" + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + + sent = _reasoning_effort_sent_for_thinking(model, provider, PLAIN_THINKING) + + assert sent == {"effort": "high", "summary": summary} + @pytest.mark.parametrize( "api_base, expected", [ diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 0b68f54c504..e44cdf2080b 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1,10 +1,12 @@ import asyncio +import copy import json import os import uuid from typing import Any, Dict, Final, List import httpx +import respx import pytest from fastapi.testclient import TestClient from pydantic import ValidationError @@ -25,6 +27,7 @@ from litellm.types.utils import ( StandardLoggingPayloadErrorInformation, StreamingChoices, ) +from litellm.utils import _invalidate_model_cost_lowercase_map import unittest.mock @@ -1075,6 +1078,44 @@ async def test_messages_strips_provider_prefix_exactly_once(requested_model, exp assert captured["url"] == expected_url +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_info, expected_url", + [ + (None, "https://api.perplexity.ai/v1/responses"), + ({"id": "pinned-id", "mode": "chat"}, "https://api.perplexity.ai/chat/completions"), + ], +) +async def test_messages_honors_a_chat_pin_on_a_prefixed_responses_deployment( + monkeypatch: pytest.MonkeyPatch, model_info: dict[str, str] | None, expected_url: str +) -> None: + """A deployment id whose stripped form is a chat row is kept on the Responses API, unless the + deployment pins ``mode: chat``, which has to decide the route before the prefix is stripped.""" + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + litellm.model_cost["perplexity/perplexity/pin-test-model"] = { + "litellm_provider": "perplexity", + "mode": "responses", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + } + litellm.model_cost["perplexity/pin-test-model"] = {"litellm_provider": "perplexity", "mode": "chat"} + _invalidate_model_cost_lowercase_map() + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route(host="api.perplexity.ai").mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(litellm.exceptions.InternalServerError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="perplexity/perplexity/pin-test-model", + api_key="test-api-key", + model_info=model_info, + ) + + assert {str(call.request.url) for call in upstream.calls} == {expected_url} + + @pytest.mark.asyncio async def test_native_messages_strips_replayed_provider_specific_fields_from_wire(): captured = {} diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 453429df084..b6f831d4bef 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -3,7 +3,7 @@ import json import os import sys from collections.abc import Mapping -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, patch import pytest @@ -110,6 +110,27 @@ def test_build_responses_kwargs_skips_include_for_a_responses_provider_that_reje assert "reasoning" in responses_kwargs +@pytest.mark.parametrize( + "model, provider, summary", + [ + ("bedrock_mantle/openai.gpt-5.6-terra", "bedrock_mantle", "auto"), + ("openai/gpt-5.6-luna", "openai", "detailed"), + ], +) +def test_build_responses_kwargs_auto_summary_asks_each_target_for_a_level_it_accepts( + monkeypatch: pytest.MonkeyPatch, model: str, provider: str, summary: str +) -> None: + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + responses_kwargs: Final = _build_responses_kwargs( + max_tokens=1024, + messages=MESSAGES, + model=model, + thinking={"type": "enabled", "budget_tokens": 4096}, + extra_kwargs={"custom_llm_provider": provider}, + ) + assert cast(Mapping[str, object], responses_kwargs["reasoning"])["summary"] == summary + + def test_build_responses_kwargs_keeps_the_deployment_include_next_to_encrypted_reasoning(): responses_kwargs = _build_responses_kwargs( max_tokens=1024, diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 951911ac066..a2289167ad4 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -625,6 +625,38 @@ class TestBedrockMantleReasoningSummary: ) assert params["reasoning"] == {"effort": "medium", "summary": "auto"} + def test_auto_summary_on_still_drops_a_client_detailed_summary(self, monkeypatch): + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + cfg = BedrockMantleResponsesAPIConfig() + params = cfg.map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": "detailed"}}, + model="openai.gpt-5.6-sol", + drop_params=True, + ) + assert params["reasoning"] == {"effort": "medium"} + + def test_auto_summary_on_still_rejects_a_client_detailed_summary_without_drop_params(self, monkeypatch): + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + cfg = BedrockMantleResponsesAPIConfig() + with pytest.raises(litellm.UnsupportedParamsError): + cfg.map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": "detailed"}}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + + @pytest.mark.parametrize("use_openai_path, expected", [(True, "auto"), (False, "detailed")]) + def test_auto_reasoning_summary_is_a_level_the_path_accepts(self, use_openai_path, expected): + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=use_openai_path) + summary = cfg.auto_reasoning_summary() + assert summary == expected + params = cfg.map_openai_params( + response_api_optional_params={"reasoning": {"effort": "medium", "summary": summary}}, + model="openai.gpt-5.6-sol" if use_openai_path else "openai.gpt-oss-120b", + drop_params=False, + ) + assert params["reasoning"] == {"effort": "medium", "summary": summary} + def test_reasoning_summary_kept_on_standard_path(self): cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) params = cfg.map_openai_params( diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index a148ef1f57d..ee41af2110d 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -21,6 +21,7 @@ import litellm from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.utils import LlmProviders from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @@ -42,6 +43,63 @@ def local_cost_map(monkeypatch): litellm.get_model_info.cache_clear() +_RESPONSES_API_REPLY: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1733529600, + "status": "completed", + "model": "openai.gpt-5.6-terra", + "output": [ + { + "type": "message", + "id": "msg_test", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} + + +_CHAT_COMPLETIONS_REPLY: Final = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "openai.gpt-5.6-terra", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +def _use_mantle_bearer(monkeypatch: pytest.MonkeyPatch) -> None: + for var in ("AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_BASE", "BEDROCK_MANTLE_REGION"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "test-key") + + +def _reply_by_route(request: httpx.Request) -> httpx.Response: + reply: Final = _RESPONSES_API_REPLY if request.url.path.endswith("/responses") else _CHAT_COMPLETIONS_REPLY + return httpx.Response(status_code=200, json=reply, request=request) + + +def _recording_client(handler: Mock) -> HTTPHandler: + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + + +def _mantle_chat_request(drop_params: bool, **params: object) -> dict[str, object]: + optional_params: Final = litellm.get_optional_params( + model="openai.gpt-5.6-terra", custom_llm_provider="bedrock_mantle", drop_params=drop_params, **params + ) + return BedrockMantleChatConfig().transform_request( + model="openai.gpt-5.6-terra", + messages=[{"role": "user", "content": "hello"}], + optional_params=optional_params, + litellm_params={"drop_params": drop_params}, + headers={}, + ) + + class TestBedrockMantleProviderRegistration: def test_provider_enum_exists(self): assert LlmProviders.BEDROCK_MANTLE == "bedrock_mantle" @@ -273,6 +331,110 @@ class TestBedrockMantleConfig: assert optional_params["max_completion_tokens"] == 50 assert "max_tokens" not in optional_params + @pytest.mark.parametrize( + "reasoning_effort, expected", + [({"effort": "low", "summary": "auto"}, "low"), ("medium", "medium"), ({"summary": "auto"}, None)], + ) + def test_chat_completions_gets_reasoning_effort_as_a_level(self, reasoning_effort, expected) -> None: + """Mantle chat answers a Responses-shaped ``{"effort", "summary"}`` with ``invalid_type``.""" + request = _mantle_chat_request(drop_params=True, reasoning_effort=reasoning_effort) + assert request.get("reasoning_effort") == expected + + @pytest.mark.parametrize("params", [{"n": 2}, {"reasoning_effort": {"effort": "low", "summary": "auto"}}]) + def test_chat_completions_answers_400_for_what_mantle_chat_cannot_serve_without_drop_params(self, params) -> None: + with pytest.raises(litellm.UnsupportedParamsError): + _mantle_chat_request(drop_params=False, **params) + + @pytest.mark.parametrize("n, drop_params, expected", [(1, False, 1), (2, True, None)]) + def test_chat_completions_sends_n_only_where_mantle_chat_serves_it(self, n, drop_params, expected) -> None: + assert _mantle_chat_request(drop_params=drop_params, n=n).get("n") == expected + + @pytest.mark.parametrize( + "params, model_info", + [ + ({"n": 2}, None), + ({"reasoning_effort": {"effort": "low", "summary": "auto"}}, {"mode": "chat"}), + ({"web_search_options": {"search_context_size": "low"}}, {"mode": "chat"}), + ], + ) + def test_what_mantle_chat_cannot_serve_reaches_the_responses_api( + self, monkeypatch, local_cost_map, params, model_info + ) -> None: + _use_mantle_bearer(monkeypatch) + handler: Final = Mock(side_effect=_reply_by_route) + + response: Final = litellm.completion( + model="bedrock_mantle/openai.gpt-5.6-terra", + messages=[{"role": "user", "content": "hello"}], + aws_region_name="us-east-1", + model_info=model_info, + client=_recording_client(handler), + **params, + ) + + sent: Final = tuple(call.args[0] for call in handler.call_args_list) + assert response.choices[0].message.content == "ok" + assert [request.url.path.rsplit("/", 1)[-1] for request in sent] == ["responses"] + assert json.loads(sent[0].content).get("reasoning", {}).get("summary") == params.get( + "reasoning_effort", {} + ).get("summary") + + @pytest.mark.parametrize( + "extra_body, drop_params, expected", + [ + ({"reasoning_effort": {"effort": "low"}}, False, {"reasoning_effort": "low"}), + ({"n": 2}, True, {}), + ({"reasoning_effort": {}}, False, {}), + ({"reasoning_effort": {"effort": None}}, False, {}), + ({"n": 1, "reasoning_effort": "high"}, False, {"n": 1, "reasoning_effort": "high"}), + ], + ids=[ + "dict-effort", + "n-above-one-dropped", + "empty-effort-asks-for-the-default", + "null-effort-asks-for-the-default", + "servable", + ], + ) + def test_a_chat_pin_sends_extra_body_through_the_same_chat_checks( + self, monkeypatch, local_cost_map, extra_body, drop_params, expected + ) -> None: + _use_mantle_bearer(monkeypatch) + handler: Final = Mock(side_effect=_reply_by_route) + + litellm.completion( + model="bedrock_mantle/openai.gpt-5.6-terra", + messages=[{"role": "user", "content": "hello"}], + aws_region_name="us-east-1", + model_info={"mode": "chat"}, + client=_recording_client(handler), + extra_body=extra_body, + drop_params=drop_params, + ) + + (sent,) = (call.args[0] for call in handler.call_args_list) + body: Final = json.loads(sent.content) + assert sent.url.path.endswith("/chat/completions") + assert {key: body[key] for key in ("n", "reasoning_effort") if key in body} == expected + + def test_a_chat_pin_answers_400_for_n_above_one_in_extra_body_without_drop_params( + self, monkeypatch, local_cost_map + ) -> None: + _use_mantle_bearer(monkeypatch) + handler: Final = Mock(side_effect=_reply_by_route) + + with pytest.raises(litellm.BadRequestError, match="does not support n=2"): + litellm.completion( + model="bedrock_mantle/openai.gpt-5.6-terra", + messages=[{"role": "user", "content": "hello"}], + aws_region_name="us-east-1", + model_info={"mode": "chat"}, + client=_recording_client(handler), + extra_body={"n": 2}, + ) + + handler.assert_not_called() + def test_non_gpt_reasoning_models_keep_the_openai_compatible_max_tokens(self): optional_params = litellm.get_optional_params( model="openai.gpt-oss-120b", diff --git a/tests/unit/proxy/test_health_check_max_tokens.py b/tests/unit/proxy/test_health_check_max_tokens.py index cad919b09e3..24117f45106 100644 --- a/tests/unit/proxy/test_health_check_max_tokens.py +++ b/tests/unit/proxy/test_health_check_max_tokens.py @@ -1,4 +1,5 @@ import asyncio +import copy import json import logging from typing import Final @@ -607,6 +608,101 @@ async def test_run_model_health_check_honors_an_explicit_chat_mode_on_mantle_cla assert fake_ahealth_check.call_args.kwargs["mode"] == "chat" +_MANTLE_GPT_CHAT_REPLY: Final = { + "id": "chatcmpl-health", + "object": "chat.completion", + "created": 1, + "model": "openai.gpt-5.6-terra", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + +_MANTLE_GPT_RESPONSES_REPLY: Final = { + "id": "resp_health", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "openai.gpt-5.6-terra", + "output": [ + { + "type": "message", + "id": "msg_health", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_info", "probed_api"), + (({"mode": "chat"}, "chat"), ({}, "responses")), + ids=("chat-pin", "unpinned"), +) +async def test_run_model_health_check_probes_the_api_a_mantle_gpt_deployment_sends_chat_to( + monkeypatch, model_info, probed_api +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=False) as respx_mock: + chat_route: Final = respx_mock.post( + url__regex=r"^https://bedrock-mantle\.us-east-2\.api\.aws/.*/chat/completions$" + ).respond(json=_MANTLE_GPT_CHAT_REPLY) + responses_route: Final = respx_mock.post( + url__regex=r"^https://bedrock-mantle\.us-east-2\.api\.aws/.*/responses$" + ).respond(json=_MANTLE_GPT_RESPONSES_REPLY) + result: Final = await hc_module._run_model_health_check( + { + "litellm_params": { + "model": "bedrock_mantle/openai.gpt-5.6-terra", + "api_key": "test-bearer", + "aws_region_name": "us-east-2", + }, + "model_info": model_info, + } + ) + + assert "error" not in result, result + assert {"chat": chat_route.call_count, "responses": responses_route.call_count} == { + "chat": int(probed_api == "chat"), + "responses": int(probed_api == "responses"), + } + + +@pytest.mark.asyncio +async def test_a_health_probe_of_a_custom_priced_chat_pin_leaves_plain_calls_on_responses(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + litellm.in_memory_llm_clients_cache.flush_cache() + params: Final = { + "model": "bedrock_mantle/openai.gpt-5.6-terra", + "api_key": "test-bearer", + "aws_region_name": "us-east-2", + } + + with respx.mock(assert_all_called=False) as respx_mock: + chat_route: Final = respx_mock.post( + url__regex=r"^https://bedrock-mantle\.us-east-2\.api\.aws/.*/chat/completions$" + ).respond(json=_MANTLE_GPT_CHAT_REPLY) + responses_route: Final = respx_mock.post( + url__regex=r"^https://bedrock-mantle\.us-east-2\.api\.aws/.*/responses$" + ).respond(json=_MANTLE_GPT_RESPONSES_REPLY) + result: Final = await hc_module._run_model_health_check( + { + "litellm_params": {**params, "input_cost_per_token": 2e-05, "output_cost_per_token": 8e-05}, + "model_info": {"mode": "chat"}, + } + ) + await litellm.acompletion(messages=[{"role": "user", "content": "Say ok"}], **params) + + assert "error" not in result, result + assert (chat_route.call_count, responses_route.call_count) == (1, 1) + + def test_autodetected_embedding_skips_reasoning_effort(): """reasoning_effort must not leak into an embedding probe whose mode is auto-detected. @@ -1313,3 +1409,36 @@ async def test_background_health_check_timeout_does_not_cancel_a_sibling( assert fast_upstream.called assert healthy_model_ids == {"fast"} assert unhealthy_model_ids == {"slow"} + + +@pytest.mark.asyncio +async def test_a_failing_health_probe_leaves_the_deployment_serving_requests(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + litellm.in_memory_llm_clients_cache.flush_cache() + deployment: Final = { + "model_name": "probe-group", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://probe.example/v1", + }, + "model_info": {"id": "probe-deployment", "mode": "chat"}, + } + router: Final = litellm.Router(model_list=[deployment]) + + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.post("https://probe.example/v1/chat/completions").mock( + side_effect=( + httpx.Response(404, json={"error": {"message": "model not found", "type": "invalid_request_error"}}), + httpx.Response(200, json={**_MANTLE_GPT_CHAT_REPLY, "model": "gpt-4o-mini"}), + ) + ) + probe: Final = await hc_module._run_model_health_check(copy.deepcopy(deployment)) + reply: Final = await router.acompletion(model="probe-group", messages=[{"role": "user", "content": "Say ok"}]) + + assert "error" in probe + assert reply.choices[0].message.content == "ok" diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index 699144a31bb..63204114501 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -6,18 +6,28 @@ Includes file_search emulation: the flag must be forwarded on inner aresponses calls so routed requests do not hit a custom api_base /v1/responses endpoint. """ +import copy import json +from collections.abc import Mapping from importlib import import_module -from typing import Final +from typing import Final, cast from unittest.mock import MagicMock, patch import httpx import openai import pytest import respx +from pydantic import JsonValue import litellm +from litellm import Router from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.responses.bridge_check import ( + backend_serves_chat_completions, + bridge_reasoning_effort, + deployment_routing_mode, + responses_api_bridge_check, +) from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -48,12 +58,8 @@ class TestUseResponsesApiBridgeFlag: @patch.object( import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) - def test_bridge_used_when_use_chat_completions_api_true( - self, mock_get_config, mock_bridge_handler - ): + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") + def test_bridge_used_when_use_chat_completions_api_true(self, mock_get_config, mock_bridge_handler): """When use_chat_completions_api=True, the bridge handler should be called.""" mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig() mock_bridge_handler.return_value = MagicMock() @@ -70,9 +76,7 @@ class TestUseResponsesApiBridgeFlag: @patch.object( import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") def test_provider_affinity_header_is_forwarded_through_bridge(self, mock_get_config, mock_bridge_handler): mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig() mock_bridge_handler.return_value = MagicMock() @@ -94,12 +98,8 @@ class TestUseResponsesApiBridgeFlag: @patch.object( import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) - def test_bridge_used_when_model_uses_chat_completions_prefix( - self, mock_get_config, mock_bridge_handler - ): + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") + def test_bridge_used_when_model_uses_chat_completions_prefix(self, mock_get_config, mock_bridge_handler): """`openai/chat_completions/` normalizes to `openai/` and uses the bridge.""" mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig() mock_bridge_handler.return_value = MagicMock() @@ -115,12 +115,8 @@ class TestUseResponsesApiBridgeFlag: assert mock_bridge_handler.call_args.kwargs["model"].endswith("my-custom-model") @patch.object(import_module("litellm.responses.main").base_llm_http_handler, "response_api_handler") - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) - def test_native_forwarding_when_flag_absent( - self, mock_get_config, mock_native_handler - ): + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") + def test_native_forwarding_when_flag_absent(self, mock_get_config, mock_native_handler): """When use_chat_completions_api is not set, openai/ models should use native responses API forwarding (existing behavior).""" mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig() @@ -137,9 +133,7 @@ class TestUseResponsesApiBridgeFlag: @patch.object( import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") def test_flag_does_not_leak_into_kwargs(self, mock_get_config, mock_bridge_handler): """use_chat_completions_api should be popped and not passed to the bridge handler.""" mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig() @@ -159,12 +153,8 @@ class TestUseResponsesApiBridgeFlag: @patch.object( import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) - def test_bridge_used_when_provider_config_none( - self, mock_get_config, mock_bridge_handler - ): + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") + def test_bridge_used_when_provider_config_none(self, mock_get_config, mock_bridge_handler): """When the provider has no native responses API config (returns None), the bridge should be used regardless of the flag (existing behavior).""" mock_get_config.return_value = None @@ -179,12 +169,8 @@ class TestUseResponsesApiBridgeFlag: mock_bridge_handler.assert_called_once() @patch("litellm.acompletion") - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) - async def test_allowed_openai_params_forwarded_through_bridge( - self, mock_get_config, mock_acompletion - ): + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") + async def test_allowed_openai_params_forwarded_through_bridge(self, mock_get_config, mock_acompletion): """allowed_openai_params is a named param of responses(), so it must be explicitly forwarded to the bridge; otherwise litellm.acompletion raises UnsupportedParamsError for params the caller explicitly allowed.""" @@ -212,9 +198,7 @@ class TestUseResponsesApiBridgeFlag: ) mock_acompletion.assert_called_once() - assert mock_acompletion.call_args.kwargs.get("allowed_openai_params") == [ - "reasoning_effort" - ] + assert mock_acompletion.call_args.kwargs.get("allowed_openai_params") == ["reasoning_effort"] @pytest.mark.parametrize( ("model", "upstream_url", "use_chat_completions_api", "allowed_openai_params", "expected_chat_template_kwargs"), @@ -287,9 +271,7 @@ class TestUseResponsesApiBridgeFlag: assert request_body["messages"] == [{"role": "user", "content": "Hello"}] assert response.output[0].content[0].text == "Answer" - def test_bridge_drops_client_metadata_even_when_allowed_openai_params_names_it( - self, respx_mock: respx.MockRouter - ): + def test_bridge_drops_client_metadata_even_when_allowed_openai_params_names_it(self, respx_mock: respx.MockRouter): upstream: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( return_value=httpx.Response( status_code=200, @@ -441,12 +423,8 @@ class TestUseResponsesApiBridgeFlag: assert "chat_template_kwargs" not in json.loads(request.read()) @patch.object(import_module("litellm.responses.file_search.emulated_handler"), "_call_aresponses") - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) - async def test_bridge_flag_forwarded_to_file_search_emulation( - self, mock_get_config, mock_call_aresponses - ): + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") + async def test_bridge_flag_forwarded_to_file_search_emulation(self, mock_get_config, mock_call_aresponses): """When use_chat_completions_api=True and file_search tool is present, the flag should be forwarded to the inner aresponses call in the file_search emulation path.""" @@ -458,12 +436,8 @@ class TestUseResponsesApiBridgeFlag: id="resp_123", model="openai/my-custom-model", created_at=1234567890, - output=[ - {"type": "message", "content": [{"type": "text", "text": "Answer"}]} - ], - usage=ResponseAPIUsage( - input_tokens=10, output_tokens=5, total_tokens=15 - ), + output=[{"type": "message", "content": [{"type": "text", "text": "Answer"}]}], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15), ) mock_call_aresponses.return_value = mock_response @@ -478,17 +452,15 @@ class TestUseResponsesApiBridgeFlag: # Verify _call_aresponses was called with use_chat_completions_api=True mock_call_aresponses.assert_called_once() call_kwargs = mock_call_aresponses.call_args.kwargs - assert ( - call_kwargs.get("use_chat_completions_api") is True - ), "use_chat_completions_api should be forwarded to inner aresponses call" + assert call_kwargs.get("use_chat_completions_api") is True, ( + "use_chat_completions_api should be forwarded to inner aresponses call" + ) @patch.object( import_module("litellm.responses.main").litellm_completion_transformation_handler, "response_api_handler" ) @patch("litellm.vector_stores.main.asearch") - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") async def test_bridge_flag_prevents_native_responses_endpoint_call( self, mock_get_config, mock_asearch, mock_bridge_handler ): @@ -512,9 +484,7 @@ class TestUseResponsesApiBridgeFlag: "arguments": '{"queries": ["test query"]}', } ], - usage=ResponseAPIUsage( - input_tokens=10, output_tokens=5, total_tokens=15 - ), + usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15), ) second_response = ResponsesAPIResponse( id="resp_second", @@ -526,9 +496,7 @@ class TestUseResponsesApiBridgeFlag: "content": [{"type": "text", "text": "Final answer"}], } ], - usage=ResponseAPIUsage( - input_tokens=20, output_tokens=10, total_tokens=30 - ), + usage=ResponseAPIUsage(input_tokens=20, output_tokens=10, total_tokens=30), ) mock_bridge_handler.side_effect = [first_response, second_response] @@ -547,8 +515,7 @@ class TestUseResponsesApiBridgeFlag: ) assert mock_bridge_handler.call_count == 2, ( - "Bridge handler should be called twice: initial function-tool call " - "and follow-up with tool results" + "Bridge handler should be called twice: initial function-tool call and follow-up with tool results" ) for call in mock_bridge_handler.call_args_list: all_kwargs = call.kwargs if call.kwargs else {} @@ -558,12 +525,8 @@ class TestUseResponsesApiBridgeFlag: @patch.object(import_module("litellm.responses.main").base_llm_http_handler, "response_api_handler") @patch("litellm.vector_stores.main.asearch") - @patch.object( - import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config" - ) - async def test_without_bridge_flag_uses_native_endpoint( - self, mock_get_config, mock_asearch, mock_native_handler - ): + @patch.object(import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config") + async def test_without_bridge_flag_uses_native_endpoint(self, mock_get_config, mock_asearch, mock_native_handler): """Without the bridge flag, openai/ with native config uses the native handler.""" mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig() mock_asearch.return_value = [] @@ -577,9 +540,7 @@ class TestUseResponsesApiBridgeFlag: "content": [{"type": "text", "text": "Native response"}], } ], - usage=ResponseAPIUsage( - input_tokens=10, output_tokens=5, total_tokens=15 - ), + usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15), ) result = await litellm.aresponses( @@ -607,3 +568,287 @@ class TestUseResponsesApiBridgeFlag: assert exc_info.value.param == "stream_chunk_size" send.assert_not_called() + + +# Deployment-pinned routing mode (deployment_routing_mode / responses_api_bridge_check) +_CHAT: Final[Mapping[str, JsonValue]] = {"id": "dep-1", "mode": "chat"} +_RESPONSES: Final[Mapping[str, JsonValue]] = {"id": "dep-1", "mode": "responses"} + + +def _pin_shared_mode(monkeypatch: pytest.MonkeyPatch, model: str, mode: str) -> None: + cost_map: Final = cast(dict[str, dict[str, object]], litellm.model_cost) + monkeypatch.setitem(cost_map, model, {**cost_map.get(model, {}), "mode": mode}) + + +@pytest.mark.parametrize( + "deployment_model_info, custom_llm_provider, expected", + [ + (_CHAT, "bedrock_mantle", "chat"), + (_RESPONSES, "bedrock_mantle", "responses"), + (_CHAT, "openai", "chat"), + (_CHAT, "chatgpt", None), + (_CHAT, "github_copilot", None), + ({"id": "dep-1", "mode": "embedding"}, "bedrock_mantle", None), + ({"id": "dep-1"}, "bedrock_mantle", None), + (None, "bedrock_mantle", None), + ], +) +def test_deployment_routing_mode_pins_chat_or_responses_except_on_responses_only_backends( + deployment_model_info: Mapping[str, JsonValue] | None, custom_llm_provider: str, expected: str | None +) -> None: + assert deployment_routing_mode(deployment_model_info, custom_llm_provider) == expected + + +def test_a_chat_pinned_mantle_deployment_stays_on_chat_completions_while_the_shared_row_says_responses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model: Final = "bedrock_mantle/openai.gpt-5.6-sol" + _pin_shared_mode(monkeypatch, model, "responses") + shared_info, _ = responses_api_bridge_check(model=model, custom_llm_provider="bedrock_mantle") + pinned_info, pinned_model = responses_api_bridge_check( + model=model, custom_llm_provider="bedrock_mantle", deployment_model_info=_CHAT + ) + assert shared_info["mode"] == "responses" + assert (pinned_info["mode"], pinned_model) == ("chat", model) + + +def test_a_chat_pinned_chatgpt_alias_still_bridges_to_the_responses_api(monkeypatch: pytest.MonkeyPatch) -> None: + _pin_shared_mode(monkeypatch, "chatgpt/gpt-5.4", "responses") + info, _ = responses_api_bridge_check(model="gpt-5.4", custom_llm_provider="chatgpt", deployment_model_info=_CHAT) + assert info["mode"] == "responses" + + +_FUNCTION_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object", "properties": {}}}, +} + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure"]) +def test_a_chat_pin_cannot_pull_gpt_5_4_function_tool_calls_off_the_responses_api(custom_llm_provider: str) -> None: + """Chat Completions rejects function tools while reasoning is on for gpt-5.4+, so the capability rule + must still bridge the request even when the deployment pins ``mode: chat``.""" + info, model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider=custom_llm_provider, + tools=[_FUNCTION_TOOL], + api_base="https://example.openai.azure.com" if custom_llm_provider == "azure" else None, + deployment_model_info=_CHAT, + ) + assert (info["mode"], model) == ("responses", "gpt-5.4") + + +@pytest.mark.parametrize("model", ["openai.gpt-5.6-terra", "us-east-1/openai.gpt-5.6-terra"]) +@pytest.mark.parametrize("reasoning_effort", [None, "low"]) +def test_a_chat_pinned_mantle_deployment_sends_function_tools_to_the_responses_api( + model: str, reasoning_effort: str | None +) -> None: + """Mantle Chat Completions rejects function tools on gpt-5.4+ unless reasoning is off, so a ``mode: chat`` + pin cannot keep a tool call there.""" + info, routed_model = responses_api_bridge_check( + model=model, + custom_llm_provider="bedrock_mantle", + tools=[_FUNCTION_TOOL], + reasoning_effort=reasoning_effort, + deployment_model_info=_CHAT, + ) + assert (info["mode"], routed_model) == ("responses", model) + + +@pytest.mark.parametrize( + "reasoning_effort, reasoning_summary", + [("low", "auto"), ({"effort": "low", "summary": "auto"}, None)], +) +def test_a_chat_pinned_mantle_deployment_sends_a_reasoning_summary_to_the_responses_api( + reasoning_effort: str | Mapping[str, JsonValue], reasoning_summary: str | None +) -> None: + info, _ = responses_api_bridge_check( + model="openai.gpt-5.6-terra", + custom_llm_provider="bedrock_mantle", + reasoning_effort=reasoning_effort, + reasoning_summary=reasoning_summary, + deployment_model_info=_CHAT, + ) + assert info["mode"] == "responses" + + +def test_a_chat_pinned_mantle_deployment_sends_web_search_to_the_responses_api() -> None: + info, model = responses_api_bridge_check( + model="openai.gpt-5.6-terra", + custom_llm_provider="bedrock_mantle", + web_search_options={"search_context_size": "medium"}, + deployment_model_info=_CHAT, + ) + assert (info["mode"], model) == ("responses", "openai.gpt-5.6-terra") + + +@pytest.mark.parametrize( + "model, reasoning_effort", + [ + ("openai.gpt-5.6-terra", "none"), + ("openai.gpt-oss-120b", "low"), + ("openai.gpt-oss-safeguard-120b", "low"), + ], +) +def test_a_chat_pinned_mantle_tool_call_stays_on_chat_completions_when_chat_accepts_it( + model: str, reasoning_effort: str +) -> None: + info, routed_model = responses_api_bridge_check( + model=model, + custom_llm_provider="bedrock_mantle", + tools=[_FUNCTION_TOOL], + reasoning_effort=reasoning_effort, + deployment_model_info=_CHAT, + ) + assert (info["mode"], routed_model) == ("chat", model) + + +def test_a_chat_pin_cannot_pull_xai_web_search_off_the_responses_api() -> None: + info, model = responses_api_bridge_check( + model="grok-4", + custom_llm_provider="xai", + web_search_options={"search_context_size": "medium"}, + deployment_model_info=_CHAT, + ) + assert (info["mode"], model) == ("responses", "grok-4") + + +@pytest.mark.parametrize( + "model, custom_llm_provider, api_base", + [ + ("responses/my-azure-deployment", "azure", "https://example.openai.azure.com"), + ("responses/gpt-5.4", "openai", None), + ], +) +def test_a_chat_pin_does_not_override_an_explicit_responses_prefix( + model: str, custom_llm_provider: str, api_base: str | None +) -> None: + info, sent_model = responses_api_bridge_check( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + deployment_model_info=_CHAT, + ) + assert (info["mode"], sent_model) == ("responses", model.removeprefix("responses/")) + + +def test_a_chat_pinned_responses_prefixed_router_deployment_reaches_the_responses_api( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + """The router registers a ``mode: chat`` row under the deployment's own ``responses/gpt-5.6`` name, + so the prefix has to win over that row too, or the request goes to Chat Completions as ``responses/gpt-5.6``.""" + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + responses_route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + ) + router: Final = Router( + model_list=[ + { + "model_name": "prefixed-chat", + "litellm_params": { + "model": "openai/responses/gpt-5.6", + "custom_llm_provider": "openai", + "api_key": "sk-test", + }, + "model_info": {"id": "prefixed-chat-id", "mode": "chat"}, + } + ] + ) + + response: Final = router.completion(model="prefixed-chat", messages=[{"role": "user", "content": "Say ok"}]) + + assert json.loads(responses_route.calls[0].request.read())["model"] == "gpt-5.6" + assert response.choices[0].message.content == "ok" + + +def test_a_chat_pin_keeps_a_plain_gpt_5_4_chat_request_on_chat_completions() -> None: + """Without tools or a reasoning summary nothing forces the Responses API, so the deployment's own + mode decides and the shared ``responses`` row is left alone.""" + shared_mode = litellm.model_cost["gpt-5.4"]["mode"] + info, model = responses_api_bridge_check(model="gpt-5.4", custom_llm_provider="openai", deployment_model_info=_CHAT) + assert (info["mode"], model) == ("chat", "gpt-5.4") + assert litellm.model_cost["gpt-5.4"]["mode"] == shared_mode + + +@pytest.mark.parametrize( + "model, custom_llm_provider", + [ + ("gpt-5.1-codex", "azure"), + ("o3-pro", "openai"), + ("gpt-5.4-pro", "openai"), + ("eu/gpt-5-codex", "azure"), + ("us/gpt-5.4-pro", "azure"), + ("databricks-gpt-5-3-codex", "databricks"), + ], +) +def test_a_chat_pin_cannot_pull_a_responses_only_model_off_the_responses_api( + model: str, custom_llm_provider: str +) -> None: + """These rows list only ``/v1/responses`` in ``supported_endpoints``, or are ``responses`` rows with no + endpoint list at all, so a deployment pin to chat has nowhere to send the request and the bridge keeps it.""" + cost_map: Final = cast(dict[str, dict[str, object]], litellm.model_cost) + assert not backend_serves_chat_completions( + cost_map[model if custom_llm_provider == "openai" else f"{custom_llm_provider}/{model}"] + ) + info, routed_model = responses_api_bridge_check( + model=model, + custom_llm_provider=custom_llm_provider, + api_base="https://example.openai.azure.com" if custom_llm_provider == "azure" else None, + deployment_model_info=_CHAT, + ) + assert (info["mode"], routed_model) == ("responses", model) + + +@pytest.mark.parametrize( + "shared_model_info, expected", + [ + ({}, True), + ({"mode": "chat"}, True), + ({"mode": "responses"}, False), + ({"supported_endpoints": ["/v1/chat/completions", "/v1/responses"]}, True), + ({"supported_endpoints": ["/v1/responses", "/v1/batch"]}, False), + ({"supported_endpoints": "/v1/chat/completions"}, True), + ], +) +def test_backend_serves_chat_completions_reads_the_shared_row( + shared_model_info: Mapping[str, JsonValue], expected: bool +) -> None: + assert backend_serves_chat_completions(shared_model_info) is expected + + +@pytest.mark.parametrize( + "extra_body, expected", + [ + ({"reasoning_effort": {"effort": "low", "summary": "auto"}}, {"effort": "low", "summary": "auto"}), + ({"reasoning_effort": "low"}, "low"), + ({"reasoning_effort": 5}, None), + ({"reasoning_effort": ["low"]}, None), + ({}, None), + ("reasoning_effort=low", None), + (None, None), + ], + ids=["dict", "string", "int", "list", "absent", "not-a-mapping", "no-extra-body"], +) +def test_bridge_reasoning_effort_reads_only_a_string_or_mapping_effort_nested_in_extra_body( + extra_body: JsonValue, expected: JsonValue +) -> None: + assert bridge_reasoning_effort(None, extra_body) == expected diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index b353dc86d5f..b4a0837578b 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -10,18 +10,21 @@ should still use the built-in pricing. import asyncio import copy import itertools +import json import logging import os +import pathlib import queue import re import sys import threading from types import MappingProxyType -from typing import Final +from typing import Final, cast from unittest.mock import Mock, patch import httpx import pytest +import respx import litellm from litellm import Router @@ -44,6 +47,9 @@ from litellm.utils import ( ) +_COPILOT_TOKEN_EXPIRES_AT: Final = 4_102_444_800 + + def _simulate_price_data_reload(fetched_catalog): """Drive what a price data reload does to this process's litellm state. @@ -2130,6 +2136,398 @@ def test_backend_cost_map_keys_matches_what_registration_writes(model, custom_ll assert keys[0] == (model if custom_llm_provider is None else f"{custom_llm_provider}/{model}") +def test_mantle_chat_deployment_does_not_rewrite_the_shared_responses_row(): + """ + A Bedrock Mantle deployment that pins ``mode: chat`` routes its own chat + requests to Chat Completions, while the shared cost-map row keeps + ``mode: responses`` for every other deployment and plain SDK calls. + """ + from litellm.main import responses_api_bridge_check + + backend_model = "bedrock_mantle/openai.gpt-5.6-sol" + cost_map: Final = cast(dict[str, dict[str, object]], litellm.model_cost) + model_keys = { + key: copy.deepcopy(cost_map.get(key)) + for key in (backend_model, f"bedrock_mantle/{backend_model}", "mantle-chat-pinned") + } + + try: + litellm.model_cost[backend_model] = { + **(model_keys[backend_model] or {}), + "litellm_provider": "bedrock_mantle", + "mode": "responses", + } + _invalidate_model_cost_lowercase_map() + + Router( + model_list=[ + { + "model_name": "mantle-chat", + "litellm_params": { + "model": backend_model, + "custom_llm_provider": "bedrock_mantle", + "api_key": "test-bearer", + }, + "model_info": {"id": "mantle-chat-pinned", "mode": "chat"}, + } + ] + ) + + assert litellm.model_cost[backend_model]["mode"] == "responses" + + shared_info, _ = responses_api_bridge_check(model=backend_model, custom_llm_provider="bedrock_mantle") + assert shared_info["mode"] == "responses" + pinned_info, _ = responses_api_bridge_check( + model=backend_model, + custom_llm_provider="bedrock_mantle", + deployment_model_info={"id": "mantle-chat-pinned", "mode": "chat"}, + ) + assert pinned_info["mode"] == "chat" + finally: + _restore_model_cost_entries(model_keys) + + +def _register_routing_test_row(monkeypatch: pytest.MonkeyPatch, mode: str) -> None: + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + litellm.model_cost["routing-test-model"] = { + "litellm_provider": "openai", + "mode": mode, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "max_input_tokens": 1000, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize("registered_mode, pinned_mode", [("responses", "chat"), ("chat", "responses")]) +@pytest.mark.parametrize( + "model, custom_llm_provider", + [ + ("openai/routing-test-model", "openai"), + ("openai/routing-test-model", None), + ("routing-test-model", "openai"), + ("openai/responses/routing-test-model", None), + ], +) +def test_a_deployment_mode_never_rewrites_the_shared_row( + monkeypatch: pytest.MonkeyPatch, + registered_mode: str, + pinned_mode: str, + model: str, + custom_llm_provider: str | None, +) -> None: + """A ``mode`` pin routes only its own deployment. Every key registration merges into the shared + ```` row, the ``responses/``-stripped twin included, has to keep the mode the SDK and the + sibling deployments route by, whichever way the pin points.""" + from litellm.main import responses_api_bridge_check + + _register_routing_test_row(monkeypatch, registered_mode) + litellm_params: Final = {"model": model, "api_key": "sk-test"} | ( + {"custom_llm_provider": custom_llm_provider} if custom_llm_provider else {} + ) + + Router( + model_list=[ + { + "model_name": "pinned", + "litellm_params": litellm_params, + "model_info": {"id": "pinned-id", "mode": pinned_mode}, + }, + {"model_name": "sibling", "litellm_params": {"model": "openai/routing-test-model", "api_key": "sk-test"}}, + ] + ) + + assert litellm.model_cost["routing-test-model"]["mode"] == registered_mode + shared_info, _ = responses_api_bridge_check(model="routing-test-model", custom_llm_provider="openai") + assert shared_info["mode"] == registered_mode + + +def test_a_custom_priced_chat_pin_request_leaves_its_unpinned_sibling_on_responses( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + _register_routing_test_row(monkeypatch, "responses") + chat_route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "routing-test-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + responses_route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response( + 200, + json={ + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "routing-test-model", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + ) + router: Final = Router( + model_list=[ + { + "model_name": "priced-chat-pin", + "litellm_params": { + "model": "openai/routing-test-model", + "api_key": "sk-test", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 4e-06, + }, + "model_info": {"id": "priced-chat-pin-id", "mode": "chat"}, + }, + {"model_name": "sibling", "litellm_params": {"model": "openai/routing-test-model", "api_key": "sk-test"}}, + ] + ) + + router.completion(model="priced-chat-pin", messages=[{"role": "user", "content": "Say ok"}]) + router.completion(model="sibling", messages=[{"role": "user", "content": "Say ok"}]) + + assert (chat_route.call_count, responses_route.call_count) == (1, 1) + + +def test_a_prefixed_deployment_override_leaves_the_shared_row_limits_alone(monkeypatch: pytest.MonkeyPatch) -> None: + """A deployment that names its provider twice, in ``custom_llm_provider`` and in the model, registers + under its own key, so its ``max_input_tokens`` override stays off the row its siblings read.""" + _register_routing_test_row(monkeypatch, "responses") + shared_before: Final = copy.deepcopy(litellm.model_cost["routing-test-model"]) + + Router( + model_list=[ + { + "model_name": "small-window", + "litellm_params": { + "model": "openai/routing-test-model", + "custom_llm_provider": "openai", + "api_key": "sk-test", + }, + "model_info": {"id": "small-window-id", "max_input_tokens": 10}, + } + ] + ) + + assert litellm.model_cost["routing-test-model"] == shared_before + + +@pytest.mark.parametrize("pins", [("chat", "responses"), ("responses", "chat")], ids=["chat first", "responses first"]) +def test_an_unmapped_model_routes_the_same_whichever_pin_registers_first( + monkeypatch: pytest.MonkeyPatch, pins: tuple[str, str] +) -> None: + """A model missing from the cost map gets its shared row from the router, so conflicting pins must not + leave a row that depends on deployment order, and a ``responses`` pin keeps opting the model in.""" + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + Router( + model_list=[ + { + "model_name": f"pinned-{mode}", + "litellm_params": {"model": "openai/unmapped-routing-test-model", "api_key": "sk-test"}, + "model_info": {"id": f"pinned-{mode}", "mode": mode}, + } + for mode in pins + ] + ) + + shared_info, _ = responses_api_bridge_check(model="unmapped-routing-test-model", custom_llm_provider="openai") + assert shared_info.get("mode") == "responses" + + +_FINE_TUNE: Final = "openai/ft:gpt-4o-mini-2024-07-18:routing-test::abc123" + + +def _fine_tune_routing( + monkeypatch: pytest.MonkeyPatch, pristine_cost_map: dict[str, object], order: tuple[str | None, ...] +) -> tuple[object, ...]: + from litellm.main import responses_api_bridge_check + + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(pristine_cost_map)) + Router( + model_list=[ + { + "model_name": f"ft-{mode}", + "litellm_params": {"model": _FINE_TUNE, "api_key": "sk-test"}, + "model_info": {"id": f"ft-{mode}"} if mode is None else {"id": f"ft-{mode}", "mode": mode}, + } + for mode in order + ] + ) + return tuple( + responses_api_bridge_check( + model=_FINE_TUNE.removeprefix("openai/"), + custom_llm_provider="openai", + deployment_model_info=None if mode is None else {"mode": mode}, + )[0].get("mode") + for mode in (None, "chat", "responses") + ) + + +@pytest.mark.parametrize("order", list(itertools.permutations(("chat", "responses", None))), ids=str) +def test_a_fine_tune_routes_the_same_whichever_deployment_registers_first( + monkeypatch: pytest.MonkeyPatch, order: tuple[str | None, ...] +) -> None: + """A fine-tune missing from the cost map inherits its base model's row on its first registration, so a + ``responses`` pin must not opt the shared row in only when it happens to register before its siblings.""" + pristine_cost_map: Final = litellm.model_cost + + assert _fine_tune_routing(monkeypatch, pristine_cost_map, order) == _fine_tune_routing( + monkeypatch, pristine_cost_map, ("responses", "chat", None) + ) + + +def test_a_non_routing_mode_pin_still_reaches_the_shared_row(monkeypatch: pytest.MonkeyPatch) -> None: + """Only ``chat`` and ``responses`` pins are per-deployment routing. Other modes, such as a realtime + model pinned to ``audio_transcription``, are still read from the shared row by name.""" + _register_routing_test_row(monkeypatch, "chat") + + Router( + model_list=[ + { + "model_name": "transcriber", + "litellm_params": {"model": "openai/routing-test-model", "api_key": "sk-test"}, + "model_info": {"id": "transcriber-id", "mode": "audio_transcription"}, + } + ] + ) + + assert litellm.get_model_info(model="routing-test-model", custom_llm_provider="openai")["mode"] == ( + "audio_transcription" + ) + + +def test_a_non_routing_mode_pin_leaves_a_responses_row_on_the_responses_api(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.main import responses_api_bridge_check + + _register_routing_test_row(monkeypatch, "responses") + + Router( + model_list=[ + { + "model_name": "transcriber", + "litellm_params": {"model": "openai/routing-test-model", "api_key": "sk-test"}, + "model_info": {"id": "transcriber-id", "mode": "audio_transcription"}, + } + ] + ) + + shared_info, _ = responses_api_bridge_check(model="routing-test-model", custom_llm_provider="openai") + assert shared_info["mode"] == "responses" + + +@pytest.mark.parametrize("registered_mode", ["responses", "chat"]) +@pytest.mark.parametrize("malformed_mode", [5, ["chat"], {"mode": "chat"}], ids=["int", "list", "object"]) +def test_a_malformed_mode_never_reaches_the_shared_row( + monkeypatch: pytest.MonkeyPatch, registered_mode: str, malformed_mode: object +) -> None: + from litellm.main import responses_api_bridge_check + + _register_routing_test_row(monkeypatch, registered_mode) + + Router( + model_list=[ + { + "model_name": "malformed", + "litellm_params": {"model": "openai/routing-test-model", "api_key": "sk-test"}, + "model_info": {"id": "malformed-id", "mode": malformed_mode}, + } + ] + ) + + assert litellm.model_cost["routing-test-model"]["mode"] == registered_mode + shared_info, _ = responses_api_bridge_check(model="routing-test-model", custom_llm_provider="openai") + assert shared_info["mode"] == registered_mode + + +def test_a_github_copilot_chat_pin_opts_a_dual_endpoint_model_out_of_the_responses_api( + monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path +) -> None: + """``github_copilot`` sends a model listing both endpoints to ``/v1/responses`` unless its row says ``chat``, + so a ``chat`` pin has to reach the row to keep that model on chat completions.""" + from litellm.llms.github_copilot.responses.transformation import github_copilot_supports_responses_api + + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "api-key.json").write_text(json.dumps({"token": "tid=test", "expires_at": _COPILOT_TOKEN_EXPIRES_AT})) + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + litellm.model_cost["github_copilot/routing-test-model"] = { + "litellm_provider": "github_copilot", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "max_input_tokens": 1000, + } + _invalidate_model_cost_lowercase_map() + assert github_copilot_supports_responses_api("github_copilot/routing-test-model") is True + + Router( + model_list=[ + { + "model_name": "copilot", + "litellm_params": {"model": "github_copilot/routing-test-model"}, + "model_info": {"id": "copilot-chat-pinned", "mode": "chat"}, + } + ] + ) + + assert github_copilot_supports_responses_api("github_copilot/routing-test-model") is False + + +@pytest.mark.parametrize("registered_mode, pinned_mode", [("chat", "responses"), ("responses", "chat")]) +@pytest.mark.parametrize( + "model, custom_llm_provider", [("github_copilot/routing-test-model", None), ("routing-test-model", "github_copilot")] +) +def test_a_github_copilot_responses_pin_still_opts_the_model_into_the_responses_api( + monkeypatch: pytest.MonkeyPatch, + tmp_path: pathlib.Path, + registered_mode: str, + pinned_mode: str, + model: str, + custom_llm_provider: str | None, +) -> None: + """``github_copilot`` dispatches ``/v1/responses`` natively by the shared row's mode, and a ``responses`` + pin is its documented opt-in, while a ``chat`` pin must not move a Responses-only row off it.""" + from litellm.llms.github_copilot.responses.transformation import github_copilot_supports_responses_api + + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "api-key.json").write_text(json.dumps({"token": "tid=test", "expires_at": _COPILOT_TOKEN_EXPIRES_AT})) + monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) + litellm.model_cost["github_copilot/routing-test-model"] = { + "litellm_provider": "github_copilot", + "mode": registered_mode, + "supported_endpoints": ["/v1/responses"] if registered_mode == "responses" else ["/v1/chat/completions"], + "max_input_tokens": 1000, + } + _invalidate_model_cost_lowercase_map() + + Router( + model_list=[ + { + "model_name": "copilot", + "litellm_params": {"model": model} + | ({"custom_llm_provider": custom_llm_provider} if custom_llm_provider else {}), + "model_info": {"id": "copilot-pinned", "mode": pinned_mode}, + } + ] + ) + + assert github_copilot_supports_responses_api("github_copilot/routing-test-model") is True + + def test_a_discarded_router_stops_contributing_to_later_reloads(monkeypatch): """ `_route_user_config_request` builds a Router per request from caller-supplied diff --git a/tests/unit/test_xai_responses_auto_routing.py b/tests/unit/test_xai_responses_auto_routing.py index c31f599dd79..89b93df1891 100644 --- a/tests/unit/test_xai_responses_auto_routing.py +++ b/tests/unit/test_xai_responses_auto_routing.py @@ -4,14 +4,16 @@ Test automatic routing to xAI Responses API when tools are present import json from collections.abc import Mapping -from typing import Final +from typing import Final, cast from unittest.mock import MagicMock, patch import httpx import pytest + import litellm from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.main import responses_api_bridge_check +from litellm.types.llms.openai import OpenAIWebSearchOptions class _RecordingResponsesHandler: @@ -46,7 +48,6 @@ class TestXAIResponsesAutoRouting: assert model_info.get("mode") != "responses" assert updated_model == model - def test_responses_api_bridge_check_with_empty_tools(self): """Test that with empty tools list, xAI does not route to Responses API""" model = "grok-3" @@ -106,9 +107,6 @@ class TestXAIResponsesAutoRouting: assert model_info.get("mode") == "responses" assert updated_model == "grok-3" # prefix removed - - - def test_responses_api_bridge_check_with_web_search_options(self): """Test auto-routing with web_search_options""" model = "grok-4-1-fast" @@ -119,7 +117,7 @@ class TestXAIResponsesAutoRouting: model_info, updated_model = responses_api_bridge_check( model=model, custom_llm_provider=custom_llm_provider, - web_search_options=web_search_options, + web_search_options=cast(OpenAIWebSearchOptions | None, web_search_options), ) # Should auto-route with web_search_options @@ -136,7 +134,7 @@ class TestXAIResponsesAutoRouting: model_info, updated_model = responses_api_bridge_check( model=model, custom_llm_provider=custom_llm_provider, - web_search_options=web_search_options, + web_search_options=cast(OpenAIWebSearchOptions | None, web_search_options), ) # Should auto-route with both present @@ -155,9 +153,7 @@ class TestXAIResponsesAutoRouting: assert updated_model == "grok-not-in-cost-map" @patch("litellm.completion_extras.responses_api_bridge.completion") - def test_completion_with_tools_routes_to_responses_api( - self, mock_responses_completion - ): + def test_completion_with_tools_routes_to_responses_api(self, mock_responses_completion): """Test that completion() with tools routes to Responses API""" # Mock the responses_api_bridge.completion to avoid actual API calls mock_responses_completion.return_value = MagicMock() From 4bafebbcdcdc515eca916dcc9082cbb53f2f3883 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 17 Oct 2026 14:39:33 -0700 Subject: [PATCH 52/64] fix(ui): take the dashboard's user role and identity from the server (#47025) * fix(ui): take the dashboard's user role and identity from the server Adds an internal GET /management/v1/users/me that returns the user id, role and email the proxy resolved for the calling key. The dashboard now reads identity from it instead of the session cookie's claims, so its navigation and role gating always match what the backend authorizes. The layout waits on it behind the existing loading screen, and a failed lookup clears the session and returns to login. Sessions that must change their password can also read it * fix(ui): sign out when the current user lookup fails and wait for it on chat The dashboard layout now holds the loading screen and signs out when /management/v1/users/me fails, instead of mounting the shell and looping back to loading as each consumer refetched the failed query. The chat layout also waits for the identity so its panels never query without a user id. The lookup goes through the typed client with the regenerated schema.d.ts * fix(ui): send the session's auth header on the first user lookup and keep chat closed until authorized The first lookup runs before AuthProvider registers the custom header name, so it now reads auth_header_name from the session cookie. Chat also waits for isAuthorized, so a failed lookup no longer mounts the shell with an empty user id * fix(ui): type the dashboard callers of the now typed user id useAuthorized used to hand back the cookie's user id untyped, which hid six strict type errors in its callers. The id is now string | null, so the production build failed. Callers handle the null case, and query key filters accept the boolean blocked filter * fix(ui): look up the current user only after the proxy URL is discovered Pages outside the dashboard shell, like chat and connect, could start the lookup before UI config discovery set the proxy base URL. On a split-origin deployment both attempts then hit the UI origin and a valid session was signed out * fix(ui): keep a loaded identity while UI config discovery reruns Pausing the lookup by passing a null token switched the query key, so on the dashboard the identity the layout had already loaded vanished until the shell's own UI config request finished. The key now stays the same and only the fetch waits for discovery --- litellm/proxy/_types.py | 1 + litellm/proxy/auth/login_utils.py | 2 +- litellm/proxy/management/users/endpoints.py | 20 ++ litellm/proxy/management/users/exceptions.py | 13 + litellm/proxy/management/users/schemas.py | 11 + .../management_v1/__init__.py | 2 + tests/unit/proxy/auth/test_login_utils.py | 4 +- tests/unit/proxy/auth/test_route_checks.py | 25 +- .../proxy/management/users/test_endpoints.py | 67 +++++ .../hooks/common/queryKeysFactory.ts | 2 +- .../hooks/currentUser/useCurrentUser.ts | 28 ++ .../app/(dashboard)/hooks/redirectToLogin.ts | 8 + .../(dashboard)/hooks/useAuthorized.test.ts | 273 +++++++++++++++++- .../app/(dashboard)/hooks/useAuthorized.ts | 36 ++- .../(dashboard)/layout.integration.test.tsx | 2 + .../src/app/(dashboard)/layout.test.tsx | 117 ++++++++ .../src/app/(dashboard)/layout.tsx | 20 +- .../components/AllModelsTab.tsx | 2 +- .../policies/_components/add_policy_form.tsx | 2 +- .../src/app/chat/layout.test.tsx | 30 ++ ui/litellm-dashboard/src/app/chat/layout.tsx | 14 +- .../DeploymentHealth.integration.test.tsx | 6 + .../src/components/add_model/AddModelForm.tsx | 2 +- .../liteadmin/LiteAdmin.integration.test.tsx | 9 +- .../organisms/create_key_button.tsx | 4 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 45 +++ 26 files changed, 702 insertions(+), 43 deletions(-) create mode 100644 litellm/proxy/management/users/endpoints.py create mode 100644 litellm/proxy/management/users/exceptions.py create mode 100644 litellm/proxy/management/users/schemas.py create mode 100644 tests/unit/proxy/management/users/test_endpoints.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/currentUser/useCurrentUser.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/redirectToLogin.ts diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dc7383569dc..a7bd16804c1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -563,6 +563,7 @@ class LiteLLMRoutes(enum.Enum): "/team/metadata_schema", "/user/info", "/v2/user/info", + "/management/v1/users/me", "/model/info", "/v1/model/info", "/v2/model/info", diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 629b31024e2..9783f44b41f 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -57,7 +57,7 @@ if TYPE_CHECKING: from prisma import types as prisma_types BREACH_RECHECK_INTERVAL: Final = timedelta(hours=24) -PASSWORD_RESET_ALLOWED_ROUTES: Final = ("/user/password/change", "/session/logout") +PASSWORD_RESET_ALLOWED_ROUTES: Final = ("/user/password/change", "/session/logout", "/management/v1/users/me") PASSWORD_SESSION_METADATA: Final = MappingProxyType({"login_method": "username_password"}) diff --git a/litellm/proxy/management/users/endpoints.py b/litellm/proxy/management/users/endpoints.py new file mode 100644 index 00000000000..62dc3a2a518 --- /dev/null +++ b/litellm/proxy/management/users/endpoints.py @@ -0,0 +1,20 @@ +from typing import Annotated, Final + +from fastapi import APIRouter, Depends + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management.users.exceptions import current_user_not_found +from litellm.proxy.management.users.schemas import CurrentUserResponse +from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX + +router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX) + + +@router.get("/users/me", include_in_schema=False) +async def get_current_user( + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> CurrentUserResponse: + if caller.user_id is None: + raise current_user_not_found() + return CurrentUserResponse(user_id=caller.user_id, user_role=caller.user_role, user_email=caller.user_email) diff --git a/litellm/proxy/management/users/exceptions.py b/litellm/proxy/management/users/exceptions.py new file mode 100644 index 00000000000..2c1e58c9283 --- /dev/null +++ b/litellm/proxy/management/users/exceptions.py @@ -0,0 +1,13 @@ +from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem +from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail + + +def current_user_not_found() -> ManagementProblem: + return ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}user-not-found", + title="User not found", + status=404, + detail="The key used for this request is not bound to a user.", + ) + ) diff --git a/litellm/proxy/management/users/schemas.py b/litellm/proxy/management/users/schemas.py new file mode 100644 index 00000000000..243e970c3a1 --- /dev/null +++ b/litellm/proxy/management/users/schemas.py @@ -0,0 +1,11 @@ +from pydantic import BaseModel, ConfigDict + +from litellm.proxy._types import LitellmUserRoles + + +class CurrentUserResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + user_id: str + user_role: LitellmUserRoles | None + user_email: str | None diff --git a/litellm/proxy/management_endpoints/management_v1/__init__.py b/litellm/proxy/management_endpoints/management_v1/__init__.py index b29b7fe5dd5..a02d989cc28 100644 --- a/litellm/proxy/management_endpoints/management_v1/__init__.py +++ b/litellm/proxy/management_endpoints/management_v1/__init__.py @@ -4,6 +4,7 @@ from typing import Final from fastapi import APIRouter +from litellm.proxy.management.users.endpoints import router as current_user_router from litellm.proxy.management_endpoints.management_v1.budgets import ( router as budgets_router, ) @@ -22,5 +23,6 @@ router.include_router(budgets_router) router.include_router(spend_logs_router) router.include_router(teams_router) router.include_router(users_router) +router.include_router(current_user_router) __all__ = ["router"] diff --git a/tests/unit/proxy/auth/test_login_utils.py b/tests/unit/proxy/auth/test_login_utils.py index 1fb648fe5e6..ec9b9a04f72 100644 --- a/tests/unit/proxy/auth/test_login_utils.py +++ b/tests/unit/proxy/auth/test_login_utils.py @@ -2144,7 +2144,7 @@ class TestPasswordResetRequiredSessionMinting: row = _db_user_row(password="Str0ng!Passw0rd", password_reset_required=True) result, key_kwargs = await self._login(_prisma_with_user(row)) - assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout"] + assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout", "/management/v1/users/me"] assert key_kwargs["metadata"] == {"login_method": "username_password", "password_reset_required": True} assert result.password_reset_required is True @@ -2205,7 +2205,7 @@ class TestPasswordResetRequiredSessionMinting: result, key_kwargs, _ = await self._login_with_screen_result(mock_prisma_client, breached=True) - assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout"] + assert key_kwargs["allowed_routes"] == ["/user/password/change", "/session/logout", "/management/v1/users/me"] assert key_kwargs["metadata"] == {"login_method": "username_password", "password_reset_required": True} assert result.password_reset_required is True diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index 7bdcb8afd37..a519a4e576a 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -17,6 +17,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import _is_api_route_allowed from litellm.proxy.auth.auth_checks_organization import user_is_org_admin +from litellm.proxy.auth.login_utils import PASSWORD_RESET_ALLOWED_ROUTES from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router @@ -3992,20 +3993,38 @@ def _password_reset_session_token() -> UserAPIKeyAuth: `password_reset_required`.""" return UserAPIKeyAuth( user_id="flagged_user", - allowed_routes=["/user/password/change"], + allowed_routes=list(PASSWORD_RESET_ALLOWED_ROUTES), metadata={"password_reset_required": True}, ) -def test_password_reset_session_can_reach_change_password(): +@pytest.mark.parametrize("route", ["/user/password/change", "/management/v1/users/me"]) +def test_password_reset_session_can_reach_change_password_and_its_own_identity(route): result = RouteChecks.is_virtual_key_allowed_to_call_route( - route="/user/password/change", + route=route, valid_token=_password_reset_session_token(), ) assert result is True +@pytest.mark.parametrize( + "role", + [LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value], +) +def test_non_admin_user_can_read_its_own_identity(role): + allowed = RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=LiteLLM_UserTable(user_id="test_user", user_role=role), + _user_role=role, + route="/management/v1/users/me", + request=MagicMock(spec=Request, query_params={}), + valid_token=UserAPIKeyAuth(user_id="test_user", user_role=role), + request_data={}, + ) + + assert allowed is None + + @pytest.mark.parametrize( "route", [ diff --git a/tests/unit/proxy/management/users/test_endpoints.py b/tests/unit/proxy/management/users/test_endpoints.py new file mode 100644 index 00000000000..7ec0e5b7f53 --- /dev/null +++ b/tests/unit/proxy/management/users/test_endpoints.py @@ -0,0 +1,67 @@ +from collections.abc import Iterator +from typing import Final + +import pytest +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient + +from litellm.proxy._types import LitellmUserRoles +from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth +from litellm.proxy.list_api.common import ManagementProblem, problem_response +from litellm.proxy.management_endpoints.management_v1 import router + +app: Final = FastAPI() + + +@app.exception_handler(ManagementProblem) +async def management_problem_exception_handler(request: Request, exc: ManagementProblem): + return problem_response(exc.problem) + + +app.include_router(router) +client: Final = TestClient(app) + +USERS_ME_PATH: Final = "/management/v1/users/me" + + +@pytest.fixture +def caller() -> Iterator[None]: + yield + app.dependency_overrides.clear() + + +def authenticate_as(resolved: UserAPIKeyAuth) -> None: + app.dependency_overrides[user_api_key_auth] = lambda: resolved + + +@pytest.mark.parametrize( + "role", + [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], +) +def test_returns_the_identity_auth_resolved_for_the_caller(caller, role): + authenticate_as(UserAPIKeyAuth(user_id="u-1", user_role=role, user_email="u1@example.com")) + + response = client.get(USERS_ME_PATH, headers={"Authorization": "Bearer k"}) + + assert response.status_code == 200 + assert response.json() == {"user_id": "u-1", "user_role": role.value, "user_email": "u1@example.com"} + + +def test_key_without_a_user_gets_a_404_problem_document(caller): + authenticate_as(UserAPIKeyAuth(team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER)) + + response = client.get(USERS_ME_PATH, headers={"Authorization": "Bearer k"}) + + assert response.status_code == 404 + assert response.headers["content-type"] == "application/problem+json" + assert response.json()["type"] == "urn:litellm:error:user-not-found" + assert response.json()["status"] == 404 + + +def test_route_is_served_but_hidden_from_the_openapi_schema(caller): + authenticate_as(UserAPIKeyAuth(user_id="u-1", user_role=LitellmUserRoles.INTERNAL_USER)) + paths: Final = app.openapi()["paths"] + + assert client.get(USERS_ME_PATH).status_code == 200 + assert "/management/v1/users/bulk" in paths + assert USERS_ME_PATH not in paths diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.ts index eceaae9d422..a7367069014 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.ts @@ -3,7 +3,7 @@ type ListParams = { page?: number; limit?: number; - filters?: Record; + filters?: Record; }; /** diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/currentUser/useCurrentUser.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/currentUser/useCurrentUser.ts new file mode 100644 index 00000000000..6b688ec2668 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/currentUser/useCurrentUser.ts @@ -0,0 +1,28 @@ +import { skipToken, useQuery, type UseQueryOptions } from "@tanstack/react-query"; +import { fetchClient } from "@/lib/http/api"; +import { getAuthHeaderName } from "@/lib/http/runtime"; +import type { components } from "@/lib/http/schema"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +export type CurrentUser = components["schemas"]["CurrentUserResponse"]; + +const currentUserKeys = createQueryKeys("currentUser"); + +const getCurrentUser = async (accessToken: string, authHeaderName: string): Promise => { + const { data } = await fetchClient.GET("/management/v1/users/me", { + headers: { [authHeaderName]: `Bearer ${accessToken}` }, + }); + if (!data) throw new Error("Current user response is empty"); + return data; +}; + +export const useCurrentUser = (accessToken: string | null, authHeaderName?: string, canFetch = true) => { + const options: UseQueryOptions = { + queryKey: currentUserKeys.detail(accessToken ?? ""), + queryFn: + accessToken && canFetch ? () => getCurrentUser(accessToken, authHeaderName ?? getAuthHeaderName()) : skipToken, + staleTime: Infinity, + retry: 1, + }; + return useQuery(options); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/redirectToLogin.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/redirectToLogin.ts new file mode 100644 index 00000000000..ec53589a8b2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/redirectToLogin.ts @@ -0,0 +1,8 @@ +import { getProxyBaseUrl, getUiPageUrl } from "@/components/networking"; +import { buildLoginUrlWithReturn, getLoginUrl, storeReturnUrl } from "@/utils/returnUrlUtils"; + +export const redirectToLogin = () => { + storeReturnUrl(); + const baseLoginUrl = getUiPageUrl(getLoginUrl(getProxyBaseUrl())); + window.location.replace(buildLoginUrlWithReturn(baseLoginUrl)); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts index 1ca0581a391..5d7825b310a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -1,9 +1,11 @@ /* @vitest-environment jsdom */ import React from "react"; -import { renderHook, waitFor } from "@testing-library/react"; +import { act, renderHook, waitFor } from "@testing-library/react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import useAuthorized from "./useAuthorized"; +import { useUIConfig } from "./uiConfig/useUIConfig"; +import { useCurrentUser } from "./currentUser/useCurrentUser"; // Unmock useAuthorized to test the actual implementation vi.unmock("@/app/(dashboard)/hooks/useAuthorized"); @@ -60,6 +62,32 @@ vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { storeReturnUrl: vi.fn(), }; }); +const fetchMock = vi.fn<(request: Request) => Promise>(); + +const CURRENT_USER_PATH = "/management/v1/users/me"; + +const currentUserRequests = () => + fetchMock.mock.calls + .map(([request]) => request) + .filter((request) => new URL(request.url).pathname === CURRENT_USER_PATH); + +type CurrentUserBody = { user_id: string; user_email: string | null; user_role: string | null }; + +const serveCurrentUser = (reply: CurrentUserBody | ((request: Request) => CurrentUserBody)) => { + fetchMock.mockImplementation(async (request) => Response.json(typeof reply === "function" ? reply(request) : reply)); + vi.stubGlobal("fetch", fetchMock); +}; + +const stallCurrentUser = () => { + fetchMock.mockImplementation(() => new Promise(() => {})); + vi.stubGlobal("fetch", fetchMock); +}; + +const failCurrentUser = (status: number) => { + fetchMock.mockImplementation(async () => Response.json({ detail: "rejected" }, { status })); + vi.stubGlobal("fetch", fetchMock); +}; + const createQueryClient = () => new QueryClient({ defaultOptions: { @@ -114,10 +142,12 @@ describe("useAuthorized", () => { decodeTokenMock.mockReset(); checkTokenValidityMock.mockReset(); buildLoginUrlWithReturnMock.mockClear(); + vi.unstubAllGlobals(); + fetchMock.mockReset(); clearCookie(); }); - it("should decode the token and expose user details", async () => { + it("should take user details from the server and the rest from the token", async () => { getUiConfigMock.mockResolvedValue({ server_root_path: "/", proxy_base_url: null, @@ -138,6 +168,7 @@ describe("useAuthorized", () => { decodeTokenMock.mockReturnValue(decodedPayload); checkTokenValidityMock.mockReturnValue(true); + serveCurrentUser({ user_id: "user-1", user_email: "user@example.com", user_role: "proxy_admin" }); const token = createJwt(decodedPayload); document.cookie = `token=${token}; path=/;`; @@ -145,9 +176,13 @@ describe("useAuthorized", () => { const { result } = renderHook(() => useAuthorized(), { wrapper }); await waitFor(() => { - expect(result.current.token).toBe(token); + expect(result.current.userId).toBe("user-1"); }); + expect(currentUserRequests().map((request) => request.headers.get("Authorization"))).toEqual([ + "Bearer api-key-123", + ]); + expect(result.current.token).toBe(token); expect(result.current.accessToken).toBe("api-key-123"); expect(result.current.userId).toBe("user-1"); expect(result.current.userEmail).toBe("user@example.com"); @@ -161,6 +196,39 @@ describe("useAuthorized", () => { expect(clearTokenCookiesMock).not.toHaveBeenCalled(); }); + it("sends the first lookup on the auth header named in the session cookie", async () => { + const uiConfig = { + server_root_path: "/", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: false, + sso_configured: false, + }; + getUiConfigMock.mockResolvedValue(uiConfig); + + const decodedPayload = { + key: "api-key-custom-header", + user_id: "user-3", + user_role: "internal_user", + auth_header_name: "X-Litellm-Key", + }; + + decodeTokenMock.mockReturnValue(decodedPayload); + checkTokenValidityMock.mockReturnValue(true); + serveCurrentUser({ user_id: "user-3", user_email: null, user_role: "internal_user" }); + document.cookie = `token=${createJwt(decodedPayload)}; path=/;`; + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(result.current.userId).toBe("user-3"); + }); + + expect(currentUserRequests().map((request) => request.headers.get("X-Litellm-Key"))).toEqual([ + "Bearer api-key-custom-header", + ]); + }); + it("should present proxy_admin_viewer as Admin while flagging it view-only", async () => { getUiConfigMock.mockResolvedValue({ server_root_path: "/", @@ -182,6 +250,11 @@ describe("useAuthorized", () => { decodeTokenMock.mockReturnValue(decodedPayload); checkTokenValidityMock.mockReturnValue(true); + serveCurrentUser({ + user_id: "user-2", + user_email: "viewer@example.com", + user_role: "proxy_admin_viewer", + }); const token = createJwt(decodedPayload); document.cookie = `token=${token}; path=/;`; @@ -189,7 +262,7 @@ describe("useAuthorized", () => { const { result } = renderHook(() => useAuthorized(), { wrapper }); await waitFor(() => { - expect(result.current.token).toBe(token); + expect(result.current.userId).toBe("user-2"); }); expect(result.current.userRole).toBe("Admin"); @@ -222,6 +295,7 @@ describe("useAuthorized", () => { expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login/"); expect(result.current.accessToken).toBeNull(); expect(result.current.userRole).toBe("Undefined Role"); + expect(currentUserRequests()).toHaveLength(0); }); it("should redirect even with valid token if admin_ui_disabled is true", async () => { @@ -245,6 +319,7 @@ describe("useAuthorized", () => { decodeTokenMock.mockReturnValue(decodedPayload); checkTokenValidityMock.mockReturnValue(true); + serveCurrentUser({ user_id: "user-1", user_email: "user@example.com", user_role: "proxy_admin" }); const token = createJwt(decodedPayload); document.cookie = `token=${token}; path=/;`; @@ -256,8 +331,6 @@ describe("useAuthorized", () => { }); expect(result.current.accessToken).toBe("api-key-123"); - expect(result.current.userId).toBe("user-1"); - expect(result.current.userEmail).toBe("user@example.com"); }); it("should redirect when token is missing", async () => { @@ -313,5 +386,193 @@ describe("useAuthorized", () => { expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login/"); expect(checkTokenValidityMock).toHaveBeenCalledWith(token); + expect(currentUserRequests()).toHaveLength(0); + }); + + describe("identity from the server", () => { + const uiConfig = { + server_root_path: "/", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: false, + sso_configured: false, + }; + + const setSessionCookie = (claims: Record) => { + decodeTokenMock.mockReturnValue(claims); + checkTokenValidityMock.mockReturnValue(true); + const token = createJwt(claims); + document.cookie = `token=${token}; path=/;`; + return token; + }; + + it("uses the role the server resolved, not the role claimed in the cookie", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "internal-key", user_id: "internal-1", user_role: "proxy_admin" }); + serveCurrentUser({ + user_id: "internal-1", + user_email: "internal@example.com", + user_role: "internal_user", + }); + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(result.current.userRole).toBe("Internal User"); + }); + expect(result.current.userRoleLabel).toBe("Internal User"); + expect(result.current.isViewOnly).toBe(false); + expect(result.current.userEmail).toBe("internal@example.com"); + expect(replaceMock).not.toHaveBeenCalled(); + }); + + it("stays loading with no role until the server answers", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "admin-claim-key", user_id: "u-1", user_role: "proxy_admin" }); + stallCurrentUser(); + + const { result } = renderHook(() => ({ auth: useAuthorized(), uiConfig: useUIConfig() }), { wrapper }); + + await waitFor(() => { + expect(result.current.uiConfig.isSuccess).toBe(true); + }); + expect(result.current.auth.isLoading).toBe(true); + expect(result.current.auth.userRole).toBe("Undefined Role"); + expect(replaceMock).not.toHaveBeenCalled(); + }); + + it("waits for backend discovery before looking up the current user", async () => { + let finishDiscovery!: (config: typeof uiConfig) => void; + getUiConfigMock.mockReturnValue(new Promise((resolve) => (finishDiscovery = resolve))); + setSessionCookie({ key: "discovery-key", user_id: "u-9", user_role: "internal_user" }); + serveCurrentUser({ user_id: "u-9", user_email: null, user_role: "internal_user" }); + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => expect(getUiConfigMock).toHaveBeenCalled()); + await act(() => new Promise((resolve) => setTimeout(resolve, 50))); + expect(currentUserRequests()).toHaveLength(0); + expect(result.current.isLoading).toBe(true); + + await act(async () => finishDiscovery(uiConfig)); + + await waitFor(() => expect(result.current.userId).toBe("u-9")); + expect(currentUserRequests()).toHaveLength(1); + expect(replaceMock).not.toHaveBeenCalled(); + }); + + it("keeps an identity the page already loaded while backend discovery reruns", async () => { + getUiConfigMock.mockReturnValue(new Promise(() => {})); + setSessionCookie({ key: "cached-key", user_id: "u-7", user_role: "proxy_admin" }); + serveCurrentUser({ user_id: "u-7", user_email: null, user_role: "internal_user" }); + const queryClient = createQueryClient(); + const sharedWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + const { result: layoutLookup } = renderHook(() => useCurrentUser("cached-key"), { wrapper: sharedWrapper }); + await waitFor(() => expect(layoutLookup.current.isSuccess).toBe(true)); + + const { result } = renderHook(() => useAuthorized(), { wrapper: sharedWrapper }); + + expect(result.current.userId).toBe("u-7"); + expect(result.current.userRole).toBe("Internal User"); + expect(currentUserRequests()).toHaveLength(1); + }); + + it("clears the cookie and redirects to login when the current user cannot be loaded", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "orphan-key", user_id: "u-1", user_role: "internal_user" }); + failCurrentUser(404); + + renderHook(() => useAuthorized(), { wrapper }); + + await waitFor( + () => { + expect(replaceMock).toHaveBeenCalledWith("http://proxy.example/ui/login/"); + }, + { timeout: 4000 }, + ); + expect(clearTokenCookiesMock).toHaveBeenCalled(); + expect(currentUserRequests()).toHaveLength(2); + }); + + it("refetches the identity when a different session key logs in", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + serveCurrentUser((request) => + request.headers.get("Authorization") === "Bearer first-key" + ? { user_id: "first", user_email: null, user_role: "proxy_admin" } + : { user_id: "second", user_email: null, user_role: "internal_user" }, + ); + const queryClient = createQueryClient(); + const sharedClientWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + setSessionCookie({ key: "first-key", user_id: "first" }); + + const { result, rerender } = renderHook(() => useAuthorized(), { wrapper: sharedClientWrapper }); + await waitFor(() => { + expect(result.current.userId).toBe("first"); + }); + + setSessionCookie({ key: "second-key", user_id: "second" }); + act(() => rerender()); + + await waitFor(() => { + expect(result.current.userId).toBe("second"); + }); + expect(result.current.userRole).toBe("Internal User"); + }); + + it.each([ + { cookieRole: "internal_user", serverRole: "proxy_admin_viewer", isViewOnly: true }, + { cookieRole: "proxy_admin_viewer", serverRole: "proxy_admin", isViewOnly: false }, + ])("derives view-only from the server role ($serverRole), not the cookie ($cookieRole)", async (scenario) => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: `view-only-${scenario.serverRole}`, user_id: "u-1", user_role: scenario.cookieRole }); + serveCurrentUser({ user_id: "u-1", user_email: null, user_role: scenario.serverRole }); + + const { result } = renderHook(() => useAuthorized(), { wrapper }); + + await waitFor(() => { + expect(result.current.userId).toBe("u-1"); + }); + expect(result.current.isViewOnly).toBe(scenario.isViewOnly); + }); + + it("does not refetch the identity when a consumer mounts after it has loaded", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "late-key", user_id: "u-1" }); + serveCurrentUser({ user_id: "u-1", user_email: null, user_role: "internal_user" }); + const queryClient = createQueryClient(); + const sharedClientWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + const first = renderHook(() => useAuthorized(), { wrapper: sharedClientWrapper }); + await waitFor(() => { + expect(first.result.current.userId).toBe("u-1"); + }); + const second = renderHook(() => useAuthorized(), { wrapper: sharedClientWrapper }); + + expect(second.result.current.userId).toBe("u-1"); + await act(async () => {}); + expect(currentUserRequests()).toHaveLength(1); + }); + + it("fetches the identity once for many consumers on the page", async () => { + getUiConfigMock.mockResolvedValue(uiConfig); + setSessionCookie({ key: "shared-key", user_id: "u-1" }); + serveCurrentUser({ user_id: "u-1", user_email: null, user_role: "internal_user" }); + const queryClient = createQueryClient(); + const sharedClientWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + const { result } = renderHook(() => [useAuthorized(), useAuthorized(), useAuthorized()], { + wrapper: sharedClientWrapper, + }); + + await waitFor(() => { + expect(result.current.map((auth) => auth.userId)).toEqual(["u-1", "u-1", "u-1"]); + }); + expect(currentUserRequests()).toHaveLength(1); + }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 0d240c02f45..47716fccfeb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -1,11 +1,11 @@ "use client"; -import { getProxyBaseUrl, getUiPageUrl } from "@/components/networking"; import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { checkTokenValidity, decodeToken } from "@/utils/jwtUtils"; -import { buildLoginUrlWithReturn, getLoginUrl, storeReturnUrl } from "@/utils/returnUrlUtils"; -import { useCallback, useEffect, useMemo } from "react"; +import { useEffect, useMemo } from "react"; import { effectiveSessionRole, formatUserRole, isViewOnlySessionRole } from "@/utils/roles"; +import { useCurrentUser } from "./currentUser/useCurrentUser"; +import { redirectToLogin } from "./redirectToLogin"; import { useUIConfig } from "./uiConfig/useUIConfig"; const useAuthorized = () => { @@ -15,16 +15,14 @@ const useAuthorized = () => { const decoded = useMemo(() => decodeToken(token), [token]); const isTokenValid = useMemo(() => checkTokenValidity(token), [token]); - const isLoading = isUIConfigLoading; - const isAuthorized = isTokenValid && !uiConfig?.admin_ui_disabled; - - // Helper function to redirect to login while preserving the current URL - const redirectToLogin = useCallback(() => { - storeReturnUrl(); - const baseLoginUrl = getUiPageUrl(getLoginUrl(getProxyBaseUrl())); - const loginUrlWithReturn = buildLoginUrlWithReturn(baseLoginUrl); - window.location.replace(loginUrlWithReturn); - }, []); + const { + data: currentUser, + isLoading: isCurrentUserLoading, + isError: isCurrentUserError, + } = useCurrentUser(isTokenValid ? decoded?.key ?? null : null, decoded?.auth_header_name, !isUIConfigLoading); + const serverRole = currentUser?.user_role ?? ""; + const isLoading = isUIConfigLoading || isCurrentUserLoading; + const isAuthorized = isTokenValid && !uiConfig?.admin_ui_disabled && !isCurrentUserError; // Single useEffect for all redirect logic useEffect(() => { @@ -36,18 +34,18 @@ const useAuthorized = () => { } redirectToLogin(); } - }, [isLoading, isAuthorized, token, redirectToLogin]); + }, [isLoading, isAuthorized, token]); return { isLoading, isAuthorized, token: isAuthorized ? token : null, accessToken: decoded?.key ?? null, - userId: decoded?.user_id ?? null, - userEmail: decoded?.user_email ?? null, - userRole: effectiveSessionRole(decoded?.user_role), - userRoleLabel: formatUserRole(decoded?.user_role), - isViewOnly: isViewOnlySessionRole(decoded?.user_role), + userId: currentUser?.user_id ?? null, + userEmail: currentUser?.user_email ?? null, + userRole: effectiveSessionRole(serverRole), + userRoleLabel: formatUserRole(serverRole), + isViewOnly: isViewOnlySessionRole(serverRole), premiumUser: decoded?.premium_user ?? null, disabledPersonalKeyCreation: decoded?.disabled_non_admin_personal_key_creation ?? null, loginMethod: decoded?.login_method ?? null, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.integration.test.tsx index aa2069eb74e..c4c32d5f7f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.integration.test.tsx @@ -81,6 +81,8 @@ beforeEach(() => { }; return Response.json(usage); } + case "/management/v1/users/me": + return Response.json({ user_id: "layout-admin", user_role: "proxy_admin", user_email: null }); case "/sso/get/ui_settings": return Response.json({}); case "/get/user_banner": diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index 154f72719e2..0fdbd047786 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -7,6 +7,17 @@ import Layout from "./layout"; const { replaceMock } = vi.hoisted(() => ({ replaceMock: vi.fn() })); +const CURRENT_USER_PATH = "/management/v1/users/me"; +const fetchMock = vi.fn<(input: RequestInfo | URL) => Promise>(); +const pathOf = (input: RequestInfo | URL) => + new URL(input instanceof Request ? input.url : String(input), "http://localhost").pathname; +let currentUserReply: () => Promise; + +const currentUserRequests = () => + fetchMock.mock.calls + .map(([input]) => input) + .filter((input): input is Request => input instanceof Request && pathOf(input) === CURRENT_USER_PATH); + let searchParamsValue = new URLSearchParams(); vi.mock("@/lib/telemetry/uiEvents", () => ({ @@ -79,6 +90,17 @@ describe("(dashboard) Layout", () => { pendingUiConfig = createDeferred(); searchParamsValue = new URLSearchParams(); vi.mocked(usePathname).mockReturnValue("/ui/guardrails"); + currentUserReply = async () => + Response.json({ user_id: "normal-user", user_role: "internal_user", user_email: null }); + fetchMock.mockReset(); + fetchMock.mockImplementation((input) => + pathOf(input) === CURRENT_USER_PATH ? currentUserReply() : Promise.resolve(Response.json({}, { status: 404 })), + ); + vi.stubGlobal("fetch", fetchMock); + }); + + afterEach(() => { + vi.unstubAllGlobals(); }); it("starts mobile navigation closed, opens a modal drawer and closes it after choosing a page", async () => { @@ -220,6 +242,7 @@ describe("(dashboard) Layout", () => { afterEach(() => { document.cookie = "token=; Max-Age=0; Path=/"; + document.cookie = "litellm_return_url=; Max-Age=0; Path=/"; }); it("routes a session flagged password_reset_required to the change-password page", async () => { @@ -264,5 +287,99 @@ describe("(dashboard) Layout", () => { expect(await screen.findByTestId("page-content")).toBeInTheDocument(); expect(replaceMock).not.toHaveBeenCalled(); }); + + it("shows the loading screen instead of the shell until the current user resolves", async () => { + let resolveCurrentUser!: (response: Response) => void; + currentUserReply = () => new Promise((resolve) => (resolveCurrentUser = resolve)); + document.cookie = `token=${sessionCookie({ user_id: "normal-user", key: "sk-pending-session", user_role: "proxy_admin" })}; Path=/`; + + render( + + +
+ + , + ); + pendingUiConfig.resolve(); + + await waitFor(() => + expect(currentUserRequests().map((request) => request.headers.get("Authorization"))).toEqual([ + "Bearer sk-pending-session", + ]), + ); + expect(screen.getByTestId("loading-screen")).toBeInTheDocument(); + expect(screen.queryByTestId("sidebar")).not.toBeInTheDocument(); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + + resolveCurrentUser(Response.json({ user_id: "normal-user", user_role: "internal_user", user_email: null })); + + expect(await screen.findByTestId("page-content")).toBeInTheDocument(); + expect(screen.getByTestId("sidebar")).toBeInTheDocument(); + }); + + it("signs out instead of mounting the shell when the current user cannot be loaded", async () => { + const originalLocation = window.location; + const locationReplace = vi.fn(); + Object.defineProperty(window, "location", { + value: { ...originalLocation, replace: locationReplace }, + writable: true, + }); + currentUserReply = async () => Response.json({ detail: "rejected" }, { status: 401 }); + document.cookie = `token=${sessionCookie({ user_id: "gone-user", key: "sk-rejected-session" })}; Path=/`; + + try { + render( + + +
+ + , + ); + pendingUiConfig.resolve(); + + await waitFor(() => expect(locationReplace).toHaveBeenCalledWith(expect.stringContaining("/login")), { + timeout: 4000, + }); + expect(document.cookie).not.toContain("token="); + expect(document.cookie).toContain("litellm_return_url="); + expect(screen.queryByTestId("sidebar")).not.toBeInTheDocument(); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + expect(currentUserRequests()).toHaveLength(2); + } finally { + Object.defineProperty(window, "location", { value: originalLocation, writable: true }); + } + }); + + it("keeps the session of an invitation link whose current user cannot be loaded", async () => { + const originalLocation = window.location; + const locationReplace = vi.fn(); + Object.defineProperty(window, "location", { + value: { ...originalLocation, replace: locationReplace }, + writable: true, + }); + searchParamsValue = new URLSearchParams("invitation_id=abc123"); + currentUserReply = async () => Response.json({ detail: "rejected" }, { status: 401 }); + document.cookie = `token=${sessionCookie({ user_id: "invited-user", key: "sk-invited-session" })}; Path=/`; + + try { + render( + + +
+ + , + ); + pendingUiConfig.resolve(); + + await waitFor(() => expect(currentUserRequests()).toHaveLength(2), { timeout: 4000 }); + await waitFor(() => + expect(replaceMock).toHaveBeenCalledWith(expect.stringContaining("/onboarding?invitation_id=abc123")), + ); + expect(locationReplace).not.toHaveBeenCalled(); + expect(document.cookie).toContain("token="); + } finally { + Object.defineProperty(window, "location", { value: originalLocation, writable: true }); + } + }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 083118bc65d..92c8806df07 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -7,6 +7,10 @@ import Navbar from "@/components/navbar"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import { ThemeProvider } from "@/contexts/ThemeContext"; import { useAuth } from "@/contexts/AuthContext"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/currentUser/useCurrentUser"; +import { redirectToLogin } from "@/app/(dashboard)/hooks/redirectToLogin"; +import { clearTokenCookies } from "@/utils/cookieUtils"; +import { decodeToken } from "@/utils/jwtUtils"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import { usePathname, useRouter, useSearchParams } from "next/navigation"; import { @@ -218,8 +222,14 @@ function LayoutContent({ children }: { children: React.ReactNode }) { const router = useRouter(); const searchParams = useSearchParams(); const pathname = usePathname(); - const { accessToken, authLoading, passwordResetRequired } = useAuth(); + const { accessToken, token, authLoading, passwordResetRequired } = useAuth(); + const session = decodeToken(token); + const { isLoading: currentUserLoading, isError: currentUserFailed } = useCurrentUser( + session?.key ?? null, + session?.auth_header_name, + ); const isInvitationFlow = Boolean(searchParams.get("invitation_id")); + const identityPending = currentUserLoading || currentUserFailed; // Legacy invitation links point at /ui/?invitation_id=; the onboarding form now lives at its own // /onboarding route. Redirect once ui-config has loaded so uiHref resolves the SERVER_ROOT_PATH base. @@ -237,7 +247,13 @@ function LayoutContent({ children }: { children: React.ReactNode }) { } }, [authLoading, passwordResetRequired, pathname, router]); - if (authLoading || isInvitationFlow) { + useEffect(() => { + if (!currentUserFailed || isInvitationFlow) return; + clearTokenCookies(); + redirectToLogin(); + }, [currentUserFailed, isInvitationFlow]); + + if (authLoading || identityPending || isInvitationFlow) { return ; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 478ddcd2458..92890031ddd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -337,7 +337,7 @@ const AllModelsTab = ({ availableModelGroups={availableModelGroups} availableModelAccessGroups={availableModelAccessGroups} userRole={userRole} - userID={userId} + userID={userId ?? ""} isViewOnly={isViewOnly} onModelIdClick={setSelectedModelId} onTeamIdClick={setSelectedTeamId} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx index b43bcba7bd0..0d1fc4f7aba 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/add_policy_form.tsx @@ -247,7 +247,7 @@ const AddPolicyForm: React.FC = ({ }, [visible, accessToken]); const loadAvailableModels = async () => { - if (!accessToken) return; + if (!accessToken || !userId) return; try { const response = await modelAvailableCall(accessToken, userId, userRole); if (response?.data) { diff --git a/ui/litellm-dashboard/src/app/chat/layout.test.tsx b/ui/litellm-dashboard/src/app/chat/layout.test.tsx index ac4fe2068fb..d455b84afa3 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.test.tsx @@ -7,6 +7,8 @@ const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, state } = const state = { enableChatUI: false, isUISettingsLoading: false, + isIdentityLoading: false, + isAuthorized: true, }; return { state, @@ -18,6 +20,8 @@ const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, state } = userId: "user-1", userEmail: "user@example.com", premiumUser: false, + isLoading: state.isIdentityLoading, + isAuthorized: state.isAuthorized, })), mockUseUISettings: vi.fn(() => ({ data: { values: { enable_chat_ui: state.enableChatUI } }, @@ -47,6 +51,8 @@ describe("ChatLayout", () => { afterEach(() => { state.enableChatUI = false; state.isUISettingsLoading = false; + state.isIdentityLoading = false; + state.isAuthorized = true; mockReplace.mockClear(); mockUiHref.mockClear(); }); @@ -84,4 +90,28 @@ describe("ChatLayout", () => { expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(mockReplace).not.toHaveBeenCalled(); }); + + it("renders nothing until the signed-in user's identity has loaded", () => { + state.enableChatUI = true; + state.isIdentityLoading = true; + render( + +
+ , + ); + expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + }); + + it("renders nothing when the signed-in user's identity could not be confirmed", () => { + state.enableChatUI = true; + state.isAuthorized = false; + render( + +
+ , + ); + expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); + expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/chat/layout.tsx b/ui/litellm-dashboard/src/app/chat/layout.tsx index 0bed5bda4a8..bfe494d2234 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.tsx @@ -12,7 +12,15 @@ import { uiHref } from "@/utils/uiHref"; // The nuqs Next adapter uses useSearchParams, so keep the chat tree behind Suspense. function ChatLayoutContent({ children }: { children: React.ReactNode }) { - const { accessToken, userRole, userId, userEmail, premiumUser } = useAuthorized(); + const { + accessToken, + userRole, + userId, + userEmail, + premiumUser, + isLoading: isIdentityLoading, + isAuthorized, + } = useAuthorized(); const { data: uiSettings, isLoading: isUISettingsLoading } = useUISettings(); const router = useRouter(); @@ -23,7 +31,9 @@ function ChatLayoutContent({ children }: { children: React.ReactNode }) { if (blocked) router.replace(uiHref("")); }, [blocked, router]); - if (isUISettingsLoading || blocked) return null; + const identityReady = !isIdentityLoading && isAuthorized; + + if (!identityReady || isUISettingsLoading || blocked) return null; return ( diff --git a/ui/litellm-dashboard/src/components/DeploymentHealth/DeploymentHealth.integration.test.tsx b/ui/litellm-dashboard/src/components/DeploymentHealth/DeploymentHealth.integration.test.tsx index 1b4d96c352d..87c34165fd9 100644 --- a/ui/litellm-dashboard/src/components/DeploymentHealth/DeploymentHealth.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/DeploymentHealth/DeploymentHealth.integration.test.tsx @@ -36,7 +36,10 @@ const snoozedBy = (user: string, until: string): DeploymentWarning["ignore"] => ignored_until: until, }); +let signedInRole = "proxy_admin"; + const signIn = (userRole: string) => { + signedInRole = userRole; const claims = { key: "sk-health-test", user_id: "admin-1", user_role: userRole, exp: NOW.getTime() / 1000 + 3600 }; const token = `${btoa(JSON.stringify({ alg: "none" }))}.${btoa(JSON.stringify(claims))}.test`; document.cookie = `token=${token}; Path=/`; @@ -47,6 +50,9 @@ const handleRequest = async (request: Request): Promise => { if (pathname === "/litellm/.well-known/litellm-ui-config") { return Response.json({ server_root_path: "", proxy_base_url: "", admin_ui_disabled: false }); } + if (pathname === "/management/v1/users/me") { + return Response.json({ user_id: "admin-1", user_role: signedInRole, user_email: null }); + } if (pathname === "/health/deployment_warnings" && request.method === "GET") { proxy.gets += 1; return Response.json({ warnings: proxy.warnings, ignore_available: proxy.ignoreAvailable }); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 7e90991523c..beab81c16a2 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -202,7 +202,7 @@ const AddModelForm: React.FC = ({ : null; const isAdmin = all_admin_roles.includes(userRole); - const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId); + const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId ?? ""); // Same owner the Auto-Routers tab uses, so the two creation forms cannot disagree about // who has to name a team. This form is only reachable when creation is allowed at all. const createScope = modelCreationScope( diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx index 61c3dbbc1f1..b0ff9871570 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx @@ -64,15 +64,18 @@ function deferred() { return { promise, resolve }; } +let signedInUser = { user_id: "first-admin", user_role: "proxy_admin" }; + function session( role = "proxy_admin", user = "first-admin", license: { premium_user?: boolean | null | string } = { premium_user: true }, + key = `sk-session-${user}`, ) { const encode = (value: object) => btoa(JSON.stringify(value)).replaceAll("=", "").replaceAll("+", "-").replaceAll("/", "_"); const claims = { - key: `sk-session-${user}`, + key, user_id: user, user_role: role, ...license, @@ -81,6 +84,7 @@ function session( }; const token = `${encode({ alg: "none" })}.${encode(claims)}.test`; document.cookie = `token=${token}; Path=/`; + signedInUser = { user_id: user, user_role: role }; return token; } @@ -134,6 +138,7 @@ function gateway(replies: (ModelReply | Promise)[], options: Gateway const recorded = { url: request.url, headers: request.headers, body, signal: request.signal }; const path = new URL(request.url).pathname; if (path.endsWith("/ui/litellm-runtime-config.json")) return json({}); + if (path.endsWith("/management/v1/users/me")) return json({ ...signedInUser, user_email: null }); requests.push(recorded); if (path.endsWith("/litellm-ui-config")) return json({ proxy_base_url: MANAGEMENT, server_root_path: "", admin_ui_disabled: false }); @@ -1018,7 +1023,7 @@ describe("LiteAdmin in the gateway", () => { settings.target = `${INFERENCE}/updated`; await act(async () => view.client.invalidateQueries({ queryKey: ["proxySettings"] })); } else { - session("proxy_admin_viewer"); + session("proxy_admin_viewer", "first-admin", { premium_user: true }, "sk-session-first-admin-viewer"); view.refresh(); expect(screen.queryByRole("button", { name: "LiteAdmin" })).not.toBeInTheDocument(); expect(screen.queryByRole("region", { name: "Ask LiteAdmin", hidden: true })).not.toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index a15f6c6e31c..30c7a80c456 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -479,7 +479,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const response = endpoint === "service_account" ? await keyCreateServiceAccountCall(accessToken, payload) - : await keyCreateCall(accessToken, userID, payload); + : await keyCreateCall(accessToken, userID ?? "", payload); // Add the data to the state in the parent component // Also directly update the keys list in VirtualKeysTable without an API call @@ -1902,7 +1902,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp Create New User Date: Sat, 17 Oct 2026 22:08:02 +0000 Subject: [PATCH 53/64] fix(router)!: bound fake-streamed calls by the request timeout, not stream_timeout (#47162) * fix(router): bound fake-streamed calls by the request timeout, not stream_timeout A stream=True call that the provider serves as one non-streaming request (Bedrock Converse for json_schema on models without native structured output, tools on some Llama models, ai21) gets no byte until generation ends, so stream_timeout acted as a total-duration cap. The Router now passes timeout and stream_timeout separately to chat completion calls, and completion() applies stream_timeout only when it reads a live stream. * test(timeout): cover stream_timeout on live streams only and the router's timeout split * fix(router): cover every fake-stream site and keep the target group's stream_timeout under prompt management * fix(router): never shorten a fake stream below stream_timeout, and split timeouts on every surface A faked stream now takes the longer of timeout and stream_timeout, so a caller who raised stream_timeout to cover a slow fake stream keeps that headroom. Azure o-series dispatch is mirrored when deciding whether a stream is faked. text_completion, /v1/messages and /v1/responses through the Router now get timeout and stream_timeout separately, like chat calls. * fix(router): report the effective stream timeout, split adapter streams, and cover M365 and Rust messages * fix(router): report and log the timeout each stream was bounded by Native /v1/messages and /v1/responses now record the timeout they applied, so x-litellm-timeout and the cooldown guard read the real limit instead of the raw request value. A prompt-managed model also passes on only the caller's or its own timeout, so the target deployment's timeout still wins. * fix(router): keep the caller's request_timeout on a prompt-managed handoff and report the applied stream limit * fix(router): give streams answered in one piece the longer limit and accept text timeouts on responses * fix(router): give /v1/messages and /v1/responses streams answered in one piece the longer limit * test(integration): cover faked, live and converted stream timeout limits per endpoint * test(router): regression tests for faked-stream timeouts --- .../litellm_core_utils/completion_timeout.py | 124 ++++++-- litellm/litellm_core_utils/litellm_logging.py | 8 +- litellm/llms/custom_httpx/llm_http_handler.py | 22 +- .../chat/transformation.py | 5 + litellm/main.py | 93 +++++- litellm/messages/dispatch.py | 9 +- litellm/responses/main.py | 18 +- litellm/router.py | 103 +++++-- litellm/types/router.py | 3 +- litellm/utils.py | 19 +- .../test_bedrock_stream_timeout_wire.py | 57 +++- ...t_interception_converted_stream_options.py | 71 ++++- .../messages/test_dispatch.py | 22 ++ .../chat/test_azure_base_model_routing.py | 39 +++ .../chat/test_bedrock_converse_handler.py | 277 ++++++++++++++++- tests/unit/llms/bedrock/slow_upstream.py | 35 ++- .../chat/test_handler.py | 27 ++ .../test_vertex_gemma_transformation.py | 27 +- .../test_responses_api_request_body.py | 43 ++- .../test_completion_timeout_resolution.py | 56 +++- tests/unit/test_main.py | 76 +++++ tests/unit/test_router/test_router.py | 4 +- tests/unit/test_router/test_router_timeout.py | 285 ++++++++++++++++++ 23 files changed, 1324 insertions(+), 99 deletions(-) diff --git a/litellm/litellm_core_utils/completion_timeout.py b/litellm/litellm_core_utils/completion_timeout.py index e6f18257138..a2e7780ded0 100644 --- a/litellm/litellm_core_utils/completion_timeout.py +++ b/litellm/litellm_core_utils/completion_timeout.py @@ -2,15 +2,25 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping +from enum import Enum +from functools import partial +from types import MappingProxyType from typing import Final import httpx from openai import Timeout as SDKTimeout +from typing_extensions import assert_never from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS +class StreamKind(Enum): + NOT_STREAMED = "not_streamed" + LIVE = "live" + FAKED = "faked" + + class CompletionTimeout: """Resolves HTTP timeout for ``completion()`` from model vs global settings.""" @@ -36,17 +46,91 @@ class CompletionTimeout: return COMPLETION_HTTP_FALLBACK_SECONDS return float(global_timeout) + @staticmethod + def _longer_limit(first: float | None, second: float | None) -> float | None: + return None if first is None or second is None else max(first, second) + + @staticmethod + def _longer(first: float | httpx.Timeout, second: float | httpx.Timeout) -> float | httpx.Timeout: + if not isinstance(first, httpx.Timeout) and not isinstance(second, httpx.Timeout): + return max(first, second) + first_limits: Final = httpx.Timeout(first) + second_limits: Final = httpx.Timeout(second) + longer: Final = CompletionTimeout._longer_limit + return httpx.Timeout( + connect=longer(first_limits.connect, second_limits.connect), + read=longer(first_limits.read, second_limits.read), + write=longer(first_limits.write, second_limits.write), + pool=longer(first_limits.pool, second_limits.pool), + ) + + @staticmethod + def for_stream( + stream_kind: StreamKind, request_timeout: float | httpx.Timeout, stream_timeout: float | httpx.Timeout | None + ) -> float | httpx.Timeout: + """ + ``stream_timeout`` bounds each read of a live stream, so it replaces the request timeout there. A fake + stream is one non-streaming request whose first byte arrives when generation ends, so it takes the longer + of each limit and never fails sooner than the same call would with either one alone. + """ + if stream_timeout is None: + return request_timeout + match stream_kind: + case StreamKind.NOT_STREAMED: + return request_timeout + case StreamKind.LIVE: + return stream_timeout + case StreamKind.FAKED: + return CompletionTimeout._longer(request_timeout, stream_timeout) + case _: + assert_never(stream_kind) + + @staticmethod + def as_limit(timeout: object) -> float | httpx.Timeout: + if isinstance(timeout, (httpx.Timeout, SDKTimeout)): + return CompletionTimeout.normalize(timeout) + if isinstance(timeout, (float, int, str)): + return float(timeout) + raise TypeError(f"timeout must be a number or an httpx.Timeout, got {type(timeout).__name__}") + + @staticmethod + def stream_request_kwargs(non_stream_timeout: object, stream_timeout: object) -> Mapping[str, object]: + if stream_timeout is None: + return MappingProxyType({"timeout": non_stream_timeout}) + return MappingProxyType( + { + "timeout": None + if non_stream_timeout is None + else CompletionTimeout.for_stream( + StreamKind.FAKED, + CompletionTimeout.as_limit(non_stream_timeout), + CompletionTimeout.as_limit(stream_timeout), + ), + "stream_timeout": stream_timeout, + } + ) + + @staticmethod + def _coerce( + timeout: object, custom_llm_provider: str, supports_httpx_timeout: Callable[[str], bool] + ) -> float | httpx.Timeout: + limit: Final = CompletionTimeout.as_limit(timeout) + if not isinstance(limit, httpx.Timeout) or supports_httpx_timeout(custom_llm_provider): + return limit + return float(limit.read) if limit.read is not None else COMPLETION_HTTP_FALLBACK_SECONDS + @staticmethod def resolve( model_timeout: float | str | httpx.Timeout | SDKTimeout | None, - kwargs: dict, + kwargs: Mapping[str, object], custom_llm_provider: str, *, global_timeout: float | str | None, supports_httpx_timeout: Callable[[str], bool], + stream_kind: StreamKind, ) -> float | httpx.Timeout: """ - Resolution order (first non-None wins): + Request timeout resolution order (first non-None wins): 1. ``model_timeout`` (call argument / merged ``litellm_params``) 2. ``kwargs["timeout"]`` @@ -54,24 +138,20 @@ class CompletionTimeout: 4. ``global_timeout`` (the explicitly-configured ``litellm.request_timeout``), or 600 when nothing was configured. + A stream then applies ``kwargs["stream_timeout"]`` per :meth:`for_stream`. Coerce :class:`httpx.Timeout` when the provider does not support it. """ - resolved: float | str | httpx.Timeout | SDKTimeout - if model_timeout is not None: - resolved = model_timeout - elif kwargs.get("timeout") is not None: - resolved = kwargs["timeout"] - elif kwargs.get("request_timeout") is not None: - resolved = kwargs["request_timeout"] - else: - resolved = CompletionTimeout._fallback_when_no_explicit_timeout(global_timeout) - - if isinstance(resolved, (httpx.Timeout, SDKTimeout)): - normalized: Final = CompletionTimeout.normalize(resolved) - if supports_httpx_timeout(custom_llm_provider): - return normalized - read_timeout: Final = normalized.read - return ( - float(read_timeout) if read_timeout is not None else COMPLETION_HTTP_FALLBACK_SECONDS - ) # default 10 min timeout - return float(resolved) + coerce: Final = partial( + CompletionTimeout._coerce, + custom_llm_provider=custom_llm_provider, + supports_httpx_timeout=supports_httpx_timeout, + ) + non_stream_candidates: Final = (model_timeout, kwargs.get("timeout"), kwargs.get("request_timeout")) + explicit: Final = next((candidate for candidate in non_stream_candidates if candidate is not None), None) + request_timeout: Final = coerce( + explicit if explicit is not None else CompletionTimeout._fallback_when_no_explicit_timeout(global_timeout) + ) + stream_timeout: Final = None if stream_kind is StreamKind.NOT_STREAMED else kwargs.get("stream_timeout") + return CompletionTimeout.for_stream( + stream_kind, request_timeout, None if stream_timeout is None else coerce(stream_timeout) + ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 99ebf479b20..b6bdc21bc05 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,7 +19,7 @@ from functools import lru_cache from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast -from httpx import Response +from httpx import Response, Timeout from pydantic import BaseModel, JsonValue, TypeAdapter import litellm @@ -671,6 +671,7 @@ class Logging(LiteLLMLoggingBaseClass): self.message_copy_duration_ms: float = (time.time() - _copy_start) * 1000 self.callback_duration_ms: float = 0.0 self.stream = stream + self.applied_timeout: float | Timeout | None = None self.start_time = start_time # log the call start time self.call_type = call_type self.relays_llm_passthrough: Final = call_type in ( @@ -1056,6 +1057,11 @@ class Logging(LiteLLMLoggingBaseClass): if "custom_llm_provider" in self.model_call_details: self.custom_llm_provider = self.model_call_details["custom_llm_provider"] + def record_applied_timeout(self, timeout: float | Timeout) -> None: + self.applied_timeout = timeout + self.litellm_params["timeout"] = timeout + self.model_call_details["litellm_params"] = self.litellm_params + def update_from_kwargs( self, kwargs: dict, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 1bccef7dd9e..3aa925db70e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2239,24 +2239,25 @@ class BaseLLMHTTPHandler: stream: bool, custom_llm_provider: str, ) -> float | httpx.Timeout | None: - from litellm.litellm_core_utils.completion_timeout import CompletionTimeout + from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) from litellm.utils import supports_httpx_timeout stream_timeout: Final = litellm_params.get("stream_timeout") if stream else None - model_timeout: Final = stream_timeout if stream_timeout is not None else litellm_params.get("timeout") + model_timeout: Final = litellm_params.get("timeout") request_timeout: Final = litellm_params.get("request_timeout") global_timeout: Final = get_configured_request_timeout() - if model_timeout is None and request_timeout is None and global_timeout is None: + if stream_timeout is None and model_timeout is None and request_timeout is None and global_timeout is None: return None return CompletionTimeout.resolve( model_timeout, - {"request_timeout": request_timeout}, + {"stream_timeout": stream_timeout, "request_timeout": request_timeout}, custom_llm_provider, global_timeout=global_timeout, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.LIVE if stream else StreamKind.NOT_STREAMED, ) async def async_anthropic_messages_handler( @@ -2410,6 +2411,13 @@ class BaseLLMHTTPHandler: }, ) + applied_timeout: Final = self._resolve_anthropic_messages_timeout( + litellm_params=litellm_params, + stream=stream or False, + custom_llm_provider=custom_llm_provider, + ) + if applied_timeout is not None: + logging_obj.record_applied_timeout(applied_timeout) response: Final = await self._async_post_anthropic_messages_with_http_error_retry( async_httpx_client=async_httpx_client, request_url=request_url, @@ -2422,11 +2430,7 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, api_key=api_key, model=model, - timeout=self._resolve_anthropic_messages_timeout( - litellm_params=litellm_params, - stream=stream or False, - custom_llm_provider=custom_llm_provider, - ), + timeout=applied_timeout, ) # used for logging + cost tracking diff --git a/litellm/llms/microsoft_365_copilot/chat/transformation.py b/litellm/llms/microsoft_365_copilot/chat/transformation.py index 6b7366685fd..0a40790001c 100644 --- a/litellm/llms/microsoft_365_copilot/chat/transformation.py +++ b/litellm/llms/microsoft_365_copilot/chat/transformation.py @@ -109,6 +109,11 @@ class Microsoft365CopilotChatConfig(BaseConfig): def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig requires a list return return ["stream", "max_tokens", "max_completion_tokens"] + def should_fake_stream( + self, model: str | None, stream: bool | None, custom_llm_provider: str | None = None + ) -> bool: + return True + def transform_request( self, model: str, diff --git a/litellm/main.py b/litellm/main.py index 85d00f28025..9cad1f82c7e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -77,7 +77,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( from litellm.litellm_core_utils.chat_completion_agentic_loop import ( maybe_run_chat_completion_agentic_loop, ) -from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_litellm_params import ( AWS_CREDENTIAL_KWARGS_KEYS, @@ -530,7 +530,14 @@ async def acompletion( mock_timeout = kwargs.get("mock_timeout", None) if mock_timeout is True: - await _handle_mock_timeout_async(mock_timeout, timeout, model) + stream_timeout: Final = kwargs.get("stream_timeout") + await _handle_mock_timeout_async( + mock_timeout, + stream_timeout + if stream is True and isinstance(stream_timeout, int | float | str | httpx.Timeout) + else timeout, + model, + ) loop: Final = asyncio.get_event_loop() custom_llm_provider = kwargs.get("custom_llm_provider", None) @@ -1135,6 +1142,52 @@ def _drop_anthropic_only_tool_keys(tools: list[dict[str, object]] | None) -> lis return [_without_anthropic_only_tool_keys(tool) if isinstance(tool, dict) else tool for tool in tools] +def _azure_completion_fakes_stream(model: str, base_model: str | None) -> bool: + if not litellm.AzureOpenAIO1Config().is_o_series_model(model=base_model or model): + return False + o_series_handler_config: Final = ProviderConfigManager.get_provider_chat_config( + model=model, provider=LlmProviders.AZURE + ) + return o_series_handler_config is not None and o_series_handler_config.should_fake_stream( + model=model, stream=True, custom_llm_provider=LlmProviders.AZURE.value + ) + + +def _completion_fakes_stream( + model: str, + custom_llm_provider: str | None, + provider_config: BaseConfig | None, + optional_params: Mapping[str, object], + *, + base_model: str | None, + bridges_to_responses_api: bool, +) -> bool: + if optional_params.get("fake_stream") is True: + return True + if custom_llm_provider is None: + return False + if bridges_to_responses_api: + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config( + provider=custom_llm_provider, model=model + ) + return responses_config is not None and responses_config.should_fake_stream( + model=model, stream=True, custom_llm_provider=custom_llm_provider + ) + if ( + custom_llm_provider == LlmProviders.VERTEX_AI.value + and get_vertex_ai_model_route(model=model, litellm_params={"base_model": base_model}) + == VertexAIModelRoute.GEMMA + ): + return True + if isinstance(provider_config, litellm.AmazonAI21Config): + return True + if custom_llm_provider == LlmProviders.AZURE.value: + return _azure_completion_fakes_stream(model, base_model) + return provider_config is not None and provider_config.should_fake_stream( + model=model, stream=True, custom_llm_provider=custom_llm_provider + ) + + class _ProxyAuthHeadersProvider(Protocol): def get_auth_headers(self) -> Mapping[str, str]: ... @@ -5574,15 +5627,6 @@ def completion( "aws_region_name", None ) # support region-based pricing for bedrock - ### TIMEOUT LOGIC ### - timeout = CompletionTimeout.resolve( - timeout, - kwargs, - custom_llm_provider, - global_timeout=get_configured_request_timeout(), - supports_httpx_timeout=supports_httpx_timeout, - ) - ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### if ( (input_cost_per_token is not None and output_cost_per_token is not None) @@ -5713,6 +5757,33 @@ def completion( model_info=_MODEL_INFO_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else None, **non_default_params, ) + stream_kind: Final = ( + ( + StreamKind.FAKED + if _completion_fakes_stream( + model, + custom_llm_provider, + provider_config, + optional_params, + base_model=base_model if isinstance(base_model, str) else None, + bridges_to_responses_api=bridges_to_responses_api is True, + ) + else StreamKind.LIVE + ) + if stream is True + else StreamKind.NOT_STREAMED + ) + + ### TIMEOUT LOGIC ### + timeout = CompletionTimeout.resolve( + timeout, + kwargs, + custom_llm_provider, + global_timeout=get_configured_request_timeout(), + supports_httpx_timeout=supports_httpx_timeout, + stream_kind=stream_kind, + ) + logging.record_applied_timeout(timeout) processed_non_default_params: Final = pre_process_non_default_params( model=model, passed_params=optional_param_args, diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index cbf3ff6a04d..cb6ed8d5c75 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -50,6 +50,13 @@ _PYTHON_AMESSAGES: Final = _python_amessages() _AMESSAGES: Final = signature(_PYTHON_AMESSAGES) +def _native_kwargs(fields: Mapping[str, object], kwargs: Mapping[str, object]) -> Mapping[str, object]: + stream_timeout: Final = kwargs.get("stream_timeout") + if fields.get("stream") is not True or stream_timeout is None: + return kwargs + return {**kwargs, "timeout": stream_timeout} + + def _public_request( legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object] ) -> NativeCall | None: @@ -61,7 +68,7 @@ def _public_request( max_tokens: Final = fields.get("max_tokens") if not isinstance(model, str) or messages is None or not isinstance(max_tokens, int): return None - return native_call(legacy, args, kwargs) + return native_call(legacy, args, _native_kwargs(fields, kwargs)) def _resolved_provider(request: NativeCall) -> str | None: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index ec1786b529a..e775e155dea 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -28,6 +28,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( configured_injection_points, ) from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -1654,6 +1655,17 @@ def responses( if custom_llm_provider is None: raise ValueError("custom_llm_provider is required but passed as None") + fake_stream: Final = responses_api_provider_config.should_fake_stream( + model=model, stream=stream, custom_llm_provider=custom_llm_provider + ) + applied_timeout: Final = CompletionTimeout.for_stream( + (StreamKind.FAKED if fake_stream else StreamKind.LIVE) if stream is True else StreamKind.NOT_STREAMED, + CompletionTimeout.as_limit(timeout or request_timeout), + None + if stream is not True or litellm_params.stream_timeout is None + else float(litellm_params.stream_timeout), + ) + litellm_logging_obj.record_applied_timeout(applied_timeout) response = base_llm_http_handler.response_api_handler( model=model, input=input, @@ -1664,12 +1676,10 @@ def responses( logging_obj=litellm_logging_obj, extra_headers=effective_extra_headers, extra_body=extra_body, - timeout=timeout or request_timeout, + timeout=applied_timeout, _is_async=_is_async, client=kwargs.get("client"), - fake_stream=responses_api_provider_config.should_fake_stream( - model=model, stream=stream, custom_llm_provider=custom_llm_provider - ), + fake_stream=fake_stream, litellm_metadata=kwargs.get("litellm_metadata", {}), shared_session=kwargs.get("shared_session"), ) diff --git a/litellm/router.py b/litellm/router.py index e363f41ccde..c30108b4e38 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -81,6 +81,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.routing import routing_decision_attributes from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, get_litellm_metadata_from_kwargs, @@ -347,6 +348,7 @@ from litellm.types.router import ( DiscoveredDeploymentModelInfo, FallbackAccessCheck, FallbackBudgetCheck, + GenericLiteLLMParams, GuardrailTypedDict, LiteLLM_Params, LiteLLMParamsTypedDict, @@ -511,6 +513,7 @@ _CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" _PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) +_STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS: Final = frozenset({"anthropic_messages", "aresponses", "responses"}) _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( { "prompt_caching": PromptCachingDeploymentCheck, @@ -549,6 +552,10 @@ _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) _RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None) _ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) _FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_DEPLOYMENT_PARAMS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | GenericLiteLLMParams]] = TypeAdapter( + Mapping[str, object] | GenericLiteLLMParams +) _OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) _DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" _CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) @@ -2946,7 +2953,7 @@ class Router: ).start() kwargs.setdefault("messages", messages) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) kwargs.pop("silent_model", None) # Ensure it's not in kwargs either model_name = litellm_params["model"] potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs) @@ -4289,7 +4296,7 @@ class Router: ) kwargs.setdefault("messages", messages) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) kwargs.pop("silent_model", None) # Ensure it's not in kwargs either model_name = litellm_params["model"] @@ -4554,6 +4561,7 @@ class Router: deployment: dict, kwargs: dict, function_name: str | None = None, + splits_stream_timeout: bool = False, ) -> None: """ 4 jobs: @@ -4671,14 +4679,35 @@ class Router: if self.request_timeout is not None else self.default_litellm_params.get("stream_timeout") ) - kwargs["timeout"] = resolve_llm_passthrough_timeout( - kwargs=kwargs, + passthrough_timeout: Final = partial( + resolve_llm_passthrough_timeout, litellm_params=deployment["litellm_params"], router_timeout=_router_timeout, router_stream_timeout=_router_stream_timeout, ) + kwargs.update( + CompletionTimeout.stream_request_kwargs( + passthrough_timeout(kwargs={**kwargs, "stream": False}), passthrough_timeout(kwargs=kwargs) + ) + if splits_stream_timeout and kwargs.get("stream") + else {"timeout": passthrough_timeout(kwargs=kwargs)} + ) + elif splits_stream_timeout: + request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) + deployment_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) + kwargs.update( + CompletionTimeout.stream_request_kwargs( + self._get_non_stream_timeout(kwargs=request_kwargs, data=deployment_params), + self._get_stream_timeout(kwargs=request_kwargs, data=deployment_params) + if request_kwargs.get("stream") + else None, + ) + ) else: - kwargs["timeout"] = self._get_timeout(kwargs=kwargs, data=deployment["litellm_params"]) + kwargs["timeout"] = self._get_timeout( + kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + data=_DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]), + ) self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) @@ -4715,7 +4744,9 @@ class Router: return model_client - def _get_stream_timeout(self, kwargs: dict, data: dict) -> float | int | None: + def _get_stream_timeout( + self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams + ) -> object: """Helper to get stream timeout from kwargs or deployment params""" return ( kwargs.get("stream_timeout", None) # the params dynamically set by user @@ -4725,9 +4756,11 @@ class Router: or self.default_litellm_params.get("stream_timeout", None) ) - def _get_non_stream_timeout(self, kwargs: dict, data: dict) -> float | int | None: + def _get_non_stream_timeout( + self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams + ) -> object: """Helper to get non-stream timeout from kwargs or deployment params""" - timeout: Final = ( + return ( kwargs.get("timeout", None) # the params dynamically set by user or kwargs.get("request_timeout", None) # the params dynamically set by user or data.get("timeout", None) # timeout set on litellm_params for this deployment @@ -4736,18 +4769,11 @@ class Router: or self.timeout # timeout set on router (router_settings.timeout) or self.default_litellm_params.get("timeout", None) ) - return timeout - def _get_timeout(self, kwargs: dict, data: dict) -> float | int | None: + def _get_timeout(self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams) -> object: """Helper to get timeout from kwargs or deployment params""" - timeout: float | int | None = None - if kwargs.get("stream", False): - timeout = self._get_stream_timeout(kwargs=kwargs, data=data) - if timeout is None: - timeout = self._get_non_stream_timeout( - kwargs=kwargs, data=data - ) # default to this if no stream specific timeout set - return timeout + stream_timeout: Final = self._get_stream_timeout(kwargs=kwargs, data=data) if kwargs.get("stream") else None + return stream_timeout if stream_timeout is not None else self._get_non_stream_timeout(kwargs=kwargs, data=data) async def abatch_completion( self, @@ -5089,6 +5115,8 @@ class Router: } ) litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) + caller_timeout: Final = kwargs.get("timeout") or kwargs.get("request_timeout") + caller_stream_timeout: Final = kwargs.get("stream_timeout") specific_deployment: Final = kwargs.pop("specific_deployment", None) prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, @@ -5097,7 +5125,9 @@ class Router: request_kwargs=kwargs, ) - self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs) + self._update_kwargs_with_deployment( + deployment=prompt_management_deployment, kwargs=kwargs, splits_stream_timeout=True + ) data: Final = prompt_management_deployment["litellm_params"].copy() litellm_model: Final = data.get("model", None) @@ -5159,7 +5189,21 @@ class Router: kwargs.pop("original_function") return await litellm.acompletion(**kwargs) - return await self.async_function_with_fallbacks(**kwargs) + target_group_timeouts: Final = { + key: value + for key, value in ( + ("timeout", caller_timeout if caller_timeout is not None else filtered_data.get("timeout")), + ( + "stream_timeout", + caller_stream_timeout if caller_stream_timeout is not None else filtered_data.get("stream_timeout"), + ), + ) + if value is not None + } + return await self.async_function_with_fallbacks( + **{key: value for key, value in kwargs.items() if key not in ("timeout", "stream_timeout")}, + **target_group_timeouts, + ) def image_generation(self, prompt: str, model: str, **kwargs): try: @@ -5596,7 +5640,7 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) data: Final = deployment["litellm_params"].copy() model_name: Final = data["model"] @@ -5666,7 +5710,7 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) data: Final = deployment["litellm_params"].copy() model_name: Final = data["model"] @@ -5901,7 +5945,13 @@ class Router: return await original_generic_function(model=model, **kwargs) raise e - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=function_name) + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name=function_name, + splits_stream_timeout=getattr(original_generic_function, "__name__", None) + in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + ) data: Final = deployment["litellm_params"].copy() model_name: Final = data["model"] @@ -6613,7 +6663,12 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name="generic_api_call", + splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + ) data: Final = deployment["litellm_params"].copy() model_name: Final = data["model"] diff --git a/litellm/types/router.py b/litellm/types/router.py index 96add8bb2a1..1bd1c2e5502 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -502,7 +502,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): itpm: int | None = None otpm: int | None = None timeout: float | str | httpx.Timeout | SDKTimeout | None = None # if str, pass in as os.environ/ - stream_timeout: float | str | None = None # timeout when making stream=True calls, if str, pass in as os.environ/ + # read timeout of a live stream, a faked stream takes the longer of this and timeout. If str, pass in as os.environ/ + stream_timeout: float | str | None = None max_retries: int | None = None drop_params: bool | str | None = None organization: str | None = None # for openai orgs diff --git a/litellm/utils.py b/litellm/utils.py index d6bdee4697b..9136e3f12d6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1482,15 +1482,16 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu return num_retries, kwargs -def _get_wrapper_timeout(kwargs: dict[str, object], exception: Exception) -> float | int | httpx.Timeout | None: - """ - Get the timeout from the kwargs - Used for the wrapper functions. - """ +def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: + from litellm.litellm_core_utils.litellm_logging import Logging - timeout: Final = cast(float | int | httpx.Timeout | None, kwargs.get("timeout", None)) - - return timeout + if kwargs.get("stream") is not True: + return kwargs.get("timeout") + logging_obj: Final = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, Logging) and logging_obj.applied_timeout is not None: + return logging_obj.applied_timeout + stream_timeout: Final = kwargs.get("stream_timeout") + return kwargs.get("timeout") if stream_timeout is None else stream_timeout def check_coroutine(value) -> bool: @@ -2353,7 +2354,7 @@ def client(original_function): if deployment_num_retries is not None: setattr(e, "num_retries", deployment_num_retries) - timeout: Final = _get_wrapper_timeout(kwargs=kwargs, exception=e) + timeout: Final = _get_wrapper_timeout(kwargs=kwargs) setattr(e, "timeout", timeout) raise e diff --git a/tests/integration/providers/test_bedrock_stream_timeout_wire.py b/tests/integration/providers/test_bedrock_stream_timeout_wire.py index 631232a0420..66267ca80ce 100644 --- a/tests/integration/providers/test_bedrock_stream_timeout_wire.py +++ b/tests/integration/providers/test_bedrock_stream_timeout_wire.py @@ -22,6 +22,8 @@ _CONVERSE_MODEL_ID: Final = "global.moonshotai.kimi-k3" _CONVERSE_MODEL: Final = f"bedrock/converse/{_CONVERSE_MODEL_ID}" _INVOKE_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0" _INVOKE_MODEL: Final = f"bedrock/invoke/{_INVOKE_MODEL_ID}" +_FAKED_MODEL_ID: Final = "ai21.jamba-1-5-mini-v1:0" +_FAKED_MODEL: Final = f"bedrock/converse/{_FAKED_MODEL_ID}" _EVENT_STREAM: Final = "application/vnd.amazon.eventstream" _ANSWER_HEAD: Final = "scripted answer part one " _ANSWER_TAIL: Final = "scripted answer part two" @@ -31,6 +33,9 @@ _USER_TURN: Final[dict[str, JsonValue]] = {"role": "user", "content": _PROMPT} _TIMEOUT_SECONDS: Final = 1 _PAUSE_SECONDS: Final = 3.0 _HOLD_SECONDS: Final = 40.0 +_GENERATION_SECONDS: Final = 3.0 +_LONG_TIMEOUT_SECONDS: Final = 6 +_BOTH_LIMITS_TIMEOUT_SECONDS: Final = 2 _CLIENT_WINDOW: Final = 10.0 _RETRY_WINDOW: Final = 25.0 _SHORT_WINDOW: Final = 3.0 @@ -65,8 +70,8 @@ _INVOKE_RESPONSE: Final = json.dumps( } ).encode() -Endpoint = Literal["chat", "messages", "responses"] -Mode = Literal["fast", "stall", "pause", "cut"] +Endpoint = Literal["chat", "messages", "responses", "completions"] +Mode = Literal["fast", "stall", "pause", "cut", "slow"] def _frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: @@ -136,6 +141,8 @@ class _Peer: target: Final = unquote(request.target) if self.mode == "stall": self.release.wait(timeout=_HOLD_SECONDS) + if self.mode == "slow": + self.release.wait(timeout=_GENERATION_SECONDS) if not _is_stream(target): return Reply(body=_json_for(target)) frames: Final = _frames_for(target) @@ -163,6 +170,10 @@ def _converse(scenario: Scenario, wire: Wire, **extra: JsonValue) -> str: return scenario.model(model=_CONVERSE_MODEL, api_base=wire.url, **_AWS, **extra) +def _faked(scenario: Scenario, wire: Wire, **extra: JsonValue) -> str: + return scenario.model(model=_FAKED_MODEL, api_base=wire.url, **_AWS, **extra) + + def _invoke(scenario: Scenario, wire: Wire, **extra: JsonValue) -> str: return scenario.model( model=_INVOKE_MODEL, api_base=wire.url, aws_bedrock_runtime_endpoint=wire.url, **_AWS, **extra @@ -205,6 +216,8 @@ def _path(endpoint: Endpoint) -> str: return "/v1/messages" case "responses": return "/v1/responses" + case "completions": + return "/v1/completions" def _body(endpoint: Endpoint, model: str, *, stream: bool = True, **extra: JsonValue) -> dict[str, JsonValue]: @@ -215,6 +228,8 @@ def _body(endpoint: Endpoint, model: str, *, stream: bool = True, **extra: JsonV return {"model": model, "messages": [_USER_TURN], "max_tokens": 16, "stream": stream, **_EXTRA, **extra} case "responses": return {"model": model, "input": _PROMPT, "stream": stream, **_EXTRA, **extra} + case "completions": + return {"model": model, "prompt": _PROMPT, "stream": stream, **_EXTRA, **extra} @dataclass(frozen=True, slots=True) @@ -318,6 +333,11 @@ def _received(wire: Wire, expected: int) -> None: assert len(wire.drain()) == expected +def _received_targets(wire: Wire, expected: int) -> tuple[str, ...]: + eventually(wire.received.qsize, lambda size: size == expected, 10) + return tuple(unquote(request.target) for request in wire.drain()) + + def _model_identity(gateway: Gateway, model: str) -> str: entries: Final = _MODELS.validate_python(gateway.get("/model/info")["data"]) (entry,) = tuple(item for item in entries if item["model_name"] == model) @@ -441,6 +461,39 @@ def test_h13_an_invoke_stream_cut_off_mid_answer_reports_the_timeout_in_band(gat _received(wire, 1) +@pytest.mark.parametrize("endpoint", ("chat", "messages", "responses", "completions")) +def test_f1_to_f4_a_faked_stream_slower_than_its_stream_timeout_is_answered_within_the_deployment_timeout( + gateway: Gateway, endpoint: Endpoint +) -> None: + with _peer("slow") as wire, gateway.scenario() as scenario: + model: Final = _faked(scenario, wire, timeout=_LONG_TIMEOUT_SECONDS, stream_timeout=_TIMEOUT_SECONDS) + _assert_answered(_streamed(gateway, endpoint, _body(endpoint, model))) + (target,) = _received_targets(wire, 1) + assert target.endswith("/converse"), target + + +def test_f5_a_faked_stream_slower_than_both_limits_times_out_at_the_deployment_timeout(gateway: Gateway) -> None: + with _peer("stall") as wire, gateway.scenario() as scenario: + model: Final = _faked(scenario, wire, timeout=_BOTH_LIMITS_TIMEOUT_SECONDS, stream_timeout=_TIMEOUT_SECONDS) + served: Final = _streamed(gateway, "chat", _body("chat", model)) + _assert_timed_out(served.status, served.text, _BOTH_LIMITS_TIMEOUT_SECONDS) + assert served.headers.get("x-litellm-timeout") == f"{_BOTH_LIMITS_TIMEOUT_SECONDS:.1f}", served.headers + _received(wire, 1) + _failure_rows(model, 1) + + +@pytest.mark.parametrize("endpoint", ("chat", "messages", "responses")) +def test_c11_to_c13_a_live_stream_slower_than_its_stream_timeout_still_fails_at_the_stream_timeout( + gateway: Gateway, endpoint: Endpoint +) -> None: + with _peer("slow") as wire, gateway.scenario() as scenario: + model: Final = _converse(scenario, wire, timeout=_LONG_TIMEOUT_SECONDS, stream_timeout=_TIMEOUT_SECONDS) + served: Final = _streamed(gateway, endpoint, _body(endpoint, model)) + _assert_timed_out(served.status, served.text) + (target,) = _received_targets(wire, 1) + assert _is_stream(target), target + + def test_c1_chat_without_streaming_already_fails_at_the_deployment_timeout(gateway: Gateway) -> None: with _peer("stall") as wire, gateway.scenario() as scenario: model: Final = _converse(scenario, wire, timeout=_TIMEOUT_SECONDS) diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py index b9d39b02afd..83a3af06355 100644 --- a/tests/integration/streaming/test_interception_converted_stream_options.py +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from pathlib import Path from queue import SimpleQueue from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit import anthropic @@ -549,6 +549,75 @@ def test_converted_chat_stream_with_usage_replays_the_single_call_usage( _assert_landed_once(deployment.model, frozenset({identity}), rows=1) +_SLOW_REPLY_SECONDS: Final = 3.0 +_SHORT_TIMEOUT_SECONDS: Final = 1 +_LONG_STREAM_TIMEOUT_SECONDS: Final = 10 +_MESSAGES_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"name": iv.SEARCH_TOOL, "description": "Search the web", "input_schema": dict(_PARAMETERS)} +) +_MESSAGES_LOOKUP_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"name": "lookup", "description": "Look up an order", "input_schema": dict(_PARAMETERS)} +) + +_ConvertedEndpoint: TypeAlias = Literal["chat", "messages"] + + +def _slowly(request: Request) -> Reply: + threading.Event().wait(_SLOW_REPLY_SECONDS) + return iv.respond(request) + + +def _messages_text(stream: _Stream) -> str: + if stream.text.startswith("{"): + blocks: Final = rv.ITEMS.validate_python(iv.JSON_OBJECT.validate_json(stream.text)["content"]) + return "".join(string_value(block["text"]) for block in blocks if block.get("type") == "text") + deltas: Final = tuple( + object_value(frame["delta"]) for frame in stream.frames if frame.get("type") == "content_block_delta" + ) + return "".join(string_value(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + + +def _converted_reply_text(rig: _Rig, endpoint: _ConvertedEndpoint, deployment: _Deployment) -> str: + if endpoint == "chat": + stream: Final = _chat(rig, _chat_body(deployment)) + _assert_restreamed(stream, iv.plain(deployment.name), usage=None) + return stream.content + messages: Final = _post( + rig, + "/v1/messages", + { + "model": deployment.model, + "max_tokens": 64, + "messages": [{"role": "user", "content": f"Search for {deployment.name}"}], + "tools": [dict(_MESSAGES_SEARCH_TOOL), dict(_MESSAGES_LOOKUP_TOOL)], + "stream": True, + **_NO_CACHE, + }, + ) + assert messages.status == 200, messages.text + return _messages_text(messages) + + +@pytest.mark.parametrize("endpoint", ("chat", "messages")) +def test_a_converted_stream_slower_than_the_deployment_timeout_is_answered_within_its_stream_timeout( + rig: _Rig, scenario: Scenario, endpoint: _ConvertedEndpoint +) -> None: + with wire_server(_slowly) as slow: + name: Final = iv.deployment("plain") + model: Final = scenario.model( + model=f"azure/{name}", + api_base=slow.url, + api_key=iv.AZURE_KEY, + api_version=iv.API_VERSION, + timeout=_SHORT_TIMEOUT_SECONDS, + stream_timeout=_LONG_STREAM_TIMEOUT_SECONDS, + ) + assert _converted_reply_text(rig, endpoint, _Deployment(model, name)) == iv.plain(name) + received: Final = iv.received(slow.drain(), name) + _assert_upstream(received, model_calls=1, searches=0) + assert not any(call.get("stream") for call in received.model_calls), received.model_calls + + def test_anthropic_sdk_messages_stream_with_only_web_search_on_azure_answers_from_the_search( rig: _Rig, scenario: Scenario ) -> None: diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index c38104495a2..2a32e31c250 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -172,6 +172,28 @@ def test_native_receives_normalized_request_and_original_call_shape() -> None: assert call_kwargs["litellm_metadata"] is metadata +@pytest.mark.parametrize("stream, expected_timeout", ((True, 0.5), (False, 30.0))) +def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_timeout: float) -> None: + captured: Final[list[NativeCall]] = [] + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("Required Rust dispatch must not call Python") + + def native(request: NativeCall) -> AnthropicMessagesResponse: + captured.append(request) + return response() + + _ = _DISPATCH.run( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + {"stream": stream, "timeout": 30.0, "stream_timeout": 0.5}, + python=python, + binding=messages_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=RUST_RULES, + ) + assert captured[0].resolved["timeout"] == expected_timeout + + def test_internal_async_marker_bypasses_native() -> None: args: Final[tuple[object, ...]] = (16, MESSAGES, "claude-sonnet-4-5") kwargs: Final[Mapping[str, object]] = {"is_async": True} diff --git a/tests/unit/llms/azure/chat/test_azure_base_model_routing.py b/tests/unit/llms/azure/chat/test_azure_base_model_routing.py index 1e8e23c38ca..aae2587085f 100644 --- a/tests/unit/llms/azure/chat/test_azure_base_model_routing.py +++ b/tests/unit/llms/azure/chat/test_azure_base_model_routing.py @@ -5,12 +5,51 @@ setting ``base_model`` should drive model-type detection (o-series, gpt-5, etc.) so the correct config, supported params, and param mapping are used. """ +from typing import Final + +import httpx +import openai import pytest import litellm from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.chat.o_series_transformation import AzureOpenAIO1Config from litellm.utils import ProviderConfigManager, get_optional_params +from tests.unit.llms.bedrock.slow_upstream import STREAM_TIMEOUT_SECONDS, slow_upstream_transport + +_O3_STREAM_ANSWER: Final = ( + b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":0,"model":"o3",' + b'"choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},"finish_reason":"stop"}]}\n\n' + b"data: [DONE]\n\n" +) + + +@pytest.mark.parametrize( + ("model", "base_model"), + [("azure/chat", "azure/o3"), ("azure/gpt4o-pro1", "azure/gpt-4o")], + ids=["o3_base_model_streams_natively", "o_series_looking_deployment_of_a_non_o_series_base_model"], +) +def test_live_stream_on_a_deployment_named_apart_from_its_base_model_keeps_its_stream_timeout( + model: str, base_model: str +) -> None: + client: Final = openai.AzureOpenAI( + api_key="test-key", + api_version="2024-12-01-preview", + azure_endpoint="https://example.openai.azure.com", + max_retries=0, + http_client=httpx.Client(transport=slow_upstream_transport(_O3_STREAM_ANSWER)), + ) + + with pytest.raises(litellm.Timeout): + litellm.completion( + model=model, + base_model=base_model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + stream_timeout=STREAM_TIMEOUT_SECONDS, + timeout=30.0, + client=client, + ) # --------------------------------------------------------------------------- diff --git a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index dec32b3325e..3c96163a7b8 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -6,6 +6,8 @@ AWS credential resolution is stubbed so nothing reaches STS. from __future__ import annotations import asyncio +import json +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from datetime import datetime, timedelta, timezone from itertools import chain from typing import Final @@ -13,18 +15,24 @@ from unittest.mock import MagicMock, patch import boto3 import httpx +import openai import pytest +import respx from botocore.credentials import Credentials from botocore.exceptions import ClientError import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import configuration -from litellm.types.utils import ModelResponse +from litellm.types.utils import AdapterCompletionStreamWrapper, ModelResponse from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from tests.unit.llms.bedrock.slow_upstream import ( + CONVERSE_ANSWER, STREAM_TIMEOUT_SECONDS, + UPSTREAM_FIRST_BYTE_SECONDS, + slow_upstream_answer, slow_upstream_async_client, slow_upstream_sync_client, ) @@ -305,6 +313,273 @@ def test_sync_converse_streaming_fails_at_the_request_timeout_not_the_upstreams_ litellm.completion(client=slow_upstream_sync_client(), **_converse_streaming_kwargs()) +@pytest.mark.asyncio +async def test_live_stream_with_a_text_stream_timeout_reports_the_stream_timeout_it_was_cut_at() -> None: + with pytest.raises(litellm.Timeout) as excinfo: + await litellm.acompletion( + client=slow_upstream_async_client(), + **{ + **_converse_streaming_kwargs(), + "timeout": UPSTREAM_FIRST_BYTE_SECONDS * 2, + "stream_timeout": str(STREAM_TIMEOUT_SECONDS), + }, + ) + assert getattr(excinfo.value, "timeout", None) == STREAM_TIMEOUT_SECONDS + + + +_FAKE_STREAMED_REQUESTS: Final = pytest.mark.parametrize( + ("model", "request_params"), + [ + ("bedrock/converse/ai21.jamba-1-5-large-v1:0", {}), + ( + "bedrock/converse/meta.llama3-3-70b-instruct-v1:0", + {"tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}]}, + ), + ( + "bedrock/converse/meta.llama3-3-70b-instruct-v1:0", + {"response_format": {"type": "json_schema", "json_schema": {"name": "answer", "schema": {"type": "object"}}}}, + ), + ("bedrock/invoke/ai21.j2-ultra-v1", {}), + ], + ids=[ + "model_without_streaming", + "request_shape_without_streaming", + "json_schema_without_native_structured_output", + "invoke_route_without_streaming", + ], +) + + +def _bedrock_router(model: str, *, timeout: float, stream_timeout: float) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "bedrock", + "litellm_params": { + "model": model, + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + }, + } + ], + timeout=timeout, + stream_timeout=stream_timeout, + num_retries=0, + ) + + +def _router_with_a_stream_timeout_shorter_than_the_upstreams_first_byte(model: str) -> litellm.Router: + return _bedrock_router(model, timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, stream_timeout=STREAM_TIMEOUT_SECONDS) + + +@_FAKE_STREAMED_REQUESTS +def test_sync_fake_stream_is_bounded_by_the_request_timeout_not_the_stream_timeout( + model: str, request_params: Mapping[str, object] +) -> None: + router: Final = _router_with_a_stream_timeout_shorter_than_the_upstreams_first_byte(model) + response: Final = router.completion( + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=slow_upstream_sync_client(CONVERSE_ANSWER), + **request_params, + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in response) == "hi" + + +@pytest.mark.asyncio +@_FAKE_STREAMED_REQUESTS +async def test_async_fake_stream_is_bounded_by_the_request_timeout_not_the_stream_timeout( + model: str, request_params: Mapping[str, object] +) -> None: + router: Final = _router_with_a_stream_timeout_shorter_than_the_upstreams_first_byte(model) + response: Final = await router.acompletion( + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=slow_upstream_async_client(CONVERSE_ANSWER), + **request_params, + ) + + assert "".join([chunk.choices[0].delta.content or "" async for chunk in response]) == "hi" + + +@pytest.mark.asyncio +async def test_fake_stream_keeps_a_stream_timeout_longer_than_the_request_timeout() -> None: + router: Final = _bedrock_router( + "bedrock/converse/ai21.jamba-1-5-large-v1:0", + timeout=STREAM_TIMEOUT_SECONDS, + stream_timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + ) + response: Final = await router.acompletion( + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=slow_upstream_async_client(CONVERSE_ANSWER), + ) + + assert "".join([chunk.choices[0].delta.content or "" async for chunk in response]) == "hi" + + +_BEDROCK_RUNTIME_URL: Final = r"https://bedrock-runtime\.us-east-1\.amazonaws\.com/.*" + + +async def _text_completion_text(router: litellm.Router) -> str: + response: Final = await router.atext_completion(model="bedrock", prompt="hi", stream=True) + return "".join([chunk.choices[0].text or "" async for chunk in response]) + + +async def _anthropic_messages_text(router: litellm.Router) -> str: + response: Final = await router.aanthropic_messages( + model="bedrock", messages=[{"role": "user", "content": "hi"}], max_tokens=10, stream=True + ) + sse: Final = b"".join([chunk async for chunk in response]).decode() + events: Final = [json.loads(line.removeprefix("data: ")) for line in sse.splitlines() if line.startswith("data: ")] + return "".join(event["delta"]["text"] for event in events if event["type"] == "content_block_delta") + + +async def _responses_text(router: litellm.Router) -> str: + response: Final = await router.aresponses(model="bedrock", input="hi", stream=True) + return "".join([event.delta async for event in response if event.type == "response.output_text.delta"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stream_text", + [_text_completion_text, _anthropic_messages_text, _responses_text], + ids=["text_completion", "anthropic_messages", "responses"], +) +async def test_fake_stream_on_each_router_surface_is_bounded_by_the_request_timeout( + stream_text: Callable[[litellm.Router], Awaitable[str]], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=slow_upstream_answer(CONVERSE_ANSWER)) + router: Final = _router_with_a_stream_timeout_shorter_than_the_upstreams_first_byte( + "bedrock/converse/ai21.jamba-1-5-large-v1:0" + ) + + assert await stream_text(router) == "hi" + + +class _ChatCompletionAdapter(CustomLogger): + def translate_completion_input_params(self, kwargs: Mapping[str, object]) -> Mapping[str, object]: + return kwargs + + def translate_completion_output_params_streaming(self, completion_stream: object) -> AdapterCompletionStreamWrapper: + return AdapterCompletionStreamWrapper(completion_stream=completion_stream) + + +@pytest.mark.asyncio +async def test_fake_stream_through_a_router_adapter_is_bounded_by_the_request_timeout( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "adapters", [{"id": "chat", "adapter": _ChatCompletionAdapter()}]) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=slow_upstream_answer(CONVERSE_ANSWER)) + router: Final = _router_with_a_stream_timeout_shorter_than_the_upstreams_first_byte( + "bedrock/converse/ai21.jamba-1-5-large-v1:0" + ) + + response: Final = await router.aadapter_completion( + adapter_id="chat", model="bedrock", messages=[{"role": "user", "content": "hi"}], stream=True + ) + + assert isinstance(response, AsyncIterator) + contents: Final = [ + chunk.choices[0].delta.content or "" async for chunk in response if isinstance(chunk, litellm.ModelResponseStream) + ] + assert "".join(contents) == "hi" + + +@pytest.mark.asyncio +async def test_live_anthropic_messages_stream_through_the_router_fails_at_the_stream_timeout( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=slow_upstream_answer()) + router: Final = _router_with_a_stream_timeout_shorter_than_the_upstreams_first_byte( + "bedrock/invoke/anthropic.claude-sonnet-4-5-v1:0" + ) + + with pytest.raises(litellm.Timeout) as excinfo: + await _anthropic_messages_text(router) + assert getattr(excinfo.value, "timeout", None) == STREAM_TIMEOUT_SECONDS + + +def test_sync_live_stream_through_the_router_fails_at_the_stream_timeout() -> None: + router: Final = _router_with_a_stream_timeout_shorter_than_the_upstreams_first_byte( + "bedrock/anthropic.claude-sonnet-4-5-v1:0" + ) + with pytest.raises(litellm.Timeout): + router.completion( + model="bedrock", messages=[{"role": "user", "content": "hi"}], stream=True, client=slow_upstream_sync_client() + ) + + +@pytest.mark.asyncio +async def test_async_live_stream_through_the_router_fails_at_the_stream_timeout() -> None: + router: Final = _router_with_a_stream_timeout_shorter_than_the_upstreams_first_byte( + "bedrock/anthropic.claude-sonnet-4-5-v1:0" + ) + with pytest.raises(litellm.Timeout) as excinfo: + await router.acompletion( + model="bedrock", messages=[{"role": "user", "content": "hi"}], stream=True, client=slow_upstream_async_client() + ) + assert getattr(excinfo.value, "timeout", None) == STREAM_TIMEOUT_SECONDS + + +@pytest.mark.asyncio +async def test_failed_fake_stream_reports_the_request_timeout_it_was_bounded_by() -> None: + request_timeout: Final = UPSTREAM_FIRST_BYTE_SECONDS / 2 + router: Final = _bedrock_router( + "bedrock/converse/ai21.jamba-1-5-large-v1:0", timeout=request_timeout, stream_timeout=STREAM_TIMEOUT_SECONDS + ) + with pytest.raises(litellm.Timeout) as excinfo: + await router.acompletion( + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=slow_upstream_async_client(CONVERSE_ANSWER), + ) + assert getattr(excinfo.value, "timeout", None) == request_timeout + + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stream_text", + [_text_completion_text, _anthropic_messages_text, _responses_text], + ids=["text_completion", "anthropic_messages", "responses"], +) +@pytest.mark.parametrize( + ("model", "upstream_answer", "bounding_timeout"), + [ + ("bedrock/converse/ai21.jamba-1-5-large-v1:0", CONVERSE_ANSWER, UPSTREAM_FIRST_BYTE_SECONDS / 2), + ("bedrock/anthropic.claude-sonnet-4-5-v1:0", b"", STREAM_TIMEOUT_SECONDS), + ], + ids=["fake_stream", "live_stream"], +) +async def test_failed_stream_on_each_router_surface_reports_the_timeout_it_was_bounded_by( + stream_text: Callable[[litellm.Router], Awaitable[str]], + model: str, + upstream_answer: bytes, + bounding_timeout: float, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=slow_upstream_answer(upstream_answer)) + router: Final = _bedrock_router(model, timeout=UPSTREAM_FIRST_BYTE_SECONDS / 2, stream_timeout=STREAM_TIMEOUT_SECONDS) + + with pytest.raises(openai.APIError) as excinfo: + await stream_text(router) + assert getattr(excinfo.value, "timeout", None) == bounding_timeout + RECORDED_REASONING: Final = "Tokyo proper has about 14 million people." RECORDED_THINKING_TOKENS: Final = 309 JSON_MODE_THINKING_CONVERSE_RESPONSE: Final = { diff --git a/tests/unit/llms/bedrock/slow_upstream.py b/tests/unit/llms/bedrock/slow_upstream.py index ed262d0f3ed..f23a947105c 100644 --- a/tests/unit/llms/bedrock/slow_upstream.py +++ b/tests/unit/llms/bedrock/slow_upstream.py @@ -7,6 +7,8 @@ gets the answer. from __future__ import annotations +import json +from collections.abc import Callable from typing import Final import httpx @@ -16,20 +18,35 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler STREAM_TIMEOUT_SECONDS: Final = 0.5 UPSTREAM_FIRST_BYTE_SECONDS: Final = 2.0 +CONVERSE_ANSWER: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + "metrics": {"latencyMs": 1}, + } +).encode() _TIMEOUT_EXTENSION: Final = TypeAdapter(dict[str, float | None]) -def _answer_once_the_first_byte_is_due(request: httpx.Request) -> httpx.Response: - read_timeout: Final = _TIMEOUT_EXTENSION.validate_python(request.extensions["timeout"])["read"] - if read_timeout is not None and read_timeout < UPSTREAM_FIRST_BYTE_SECONDS: - raise httpx.ReadTimeout(f"no byte arrived within {read_timeout}s", request=request) - return httpx.Response(200, content=b"", request=request) +def slow_upstream_answer(body: bytes = b"") -> Callable[[httpx.Request], httpx.Response]: + def answer(request: httpx.Request) -> httpx.Response: + read_timeout: Final = _TIMEOUT_EXTENSION.validate_python(request.extensions["timeout"])["read"] + if read_timeout is not None and read_timeout < UPSTREAM_FIRST_BYTE_SECONDS: + raise httpx.ReadTimeout(f"no byte arrived within {read_timeout}s", request=request) + return httpx.Response(200, content=body, request=request) + + return answer -def slow_upstream_async_client() -> AsyncHTTPHandler: - return AsyncHTTPHandler(transport=httpx.MockTransport(_answer_once_the_first_byte_is_due)) +def slow_upstream_transport(body: bytes = b"") -> httpx.MockTransport: + return httpx.MockTransport(slow_upstream_answer(body)) -def slow_upstream_sync_client() -> HTTPHandler: - return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_answer_once_the_first_byte_is_due))) +def slow_upstream_async_client(body: bytes = b"") -> AsyncHTTPHandler: + return AsyncHTTPHandler(transport=slow_upstream_transport(body)) + + +def slow_upstream_sync_client(body: bytes = b"") -> HTTPHandler: + return HTTPHandler(client=httpx.Client(transport=slow_upstream_transport(body))) diff --git a/tests/unit/llms/microsoft_365_copilot/chat/test_handler.py b/tests/unit/llms/microsoft_365_copilot/chat/test_handler.py index 7be5db74af5..6c41a6dcad5 100644 --- a/tests/unit/llms/microsoft_365_copilot/chat/test_handler.py +++ b/tests/unit/llms/microsoft_365_copilot/chat/test_handler.py @@ -17,6 +17,11 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.utils import ModelResponse from litellm.utils import token_counter +from tests.unit.llms.bedrock.slow_upstream import ( + STREAM_TIMEOUT_SECONDS, + UPSTREAM_FIRST_BYTE_SECONDS, + slow_upstream_answer, +) _MODEL: Final = "microsoft_365_copilot/chat" _MESSAGES_ADAPTER: Final = TypeAdapter(list[AllMessageValues]) @@ -552,6 +557,28 @@ def test_streaming_wraps_full_reply_as_fake_stream() -> None: assert usage_chunk.usage.total_tokens == (usage_chunk.usage.prompt_tokens + usage_chunk.usage.completion_tokens) +def test_streaming_is_bounded_by_the_request_timeout_not_the_stream_timeout() -> None: + slow_chat: Final = slow_upstream_answer(json.dumps(_DEFAULT_CHAT_RESPONSE).encode()) + + def respond(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/chat"): + return slow_chat(request) + return httpx.Response(status_code=201, json={"id": "conversation/42", "state": "active"}, request=request) + + response: Final = litellm.completion( + model=_MODEL, + messages=_messages(), + api_key="delegated-stream-token-9306", + stream=True, + timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + stream_timeout=STREAM_TIMEOUT_SECONDS, + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + ) + + assert isinstance(response, CustomStreamWrapper) + assert "".join(chunk.choices[0].delta.content or "" for chunk in response if chunk.choices) == "Copilot response" + + @pytest.mark.asyncio async def test_async_streaming_wraps_full_reply_as_fake_stream() -> None: client, _ = _async_client() diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index 1110b53624a..fb1cf4ced2d 100644 --- a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -6,9 +6,10 @@ Maps to: litellm/llms/vertex_ai/vertex_gemma_models/transformation.py import json from collections.abc import AsyncIterator -from typing import cast +from typing import Final, cast from unittest.mock import AsyncMock, Mock, patch +import httpx import pytest import litellm @@ -17,6 +18,11 @@ from litellm.types.llms.openai import ( ResponseCompletedEvent, ResponsesAPIStreamingResponse, ) +from tests.unit.llms.bedrock.slow_upstream import ( + STREAM_TIMEOUT_SECONDS, + UPSTREAM_FIRST_BYTE_SECONDS, + slow_upstream_transport, +) @pytest.fixture(autouse=True) @@ -1354,6 +1360,25 @@ def test_sync_gemma_stream(_gemma_cached_access_token): assert chunks[1].choices[0].finish_reason == "stop" +def test_gemma_stream_is_bounded_by_the_request_timeout_not_the_stream_timeout(_gemma_cached_access_token: None) -> None: + stream: Final = litellm.completion( + model="vertex_ai/gemma/test-model", + messages=[{"role": "user", "content": "Reply exactly READY"}], + stream=True, + stream_timeout=STREAM_TIMEOUT_SECONDS, + timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + api_base=_GEMMA_VERTEX_URL, + vertex_project="test", + vertex_location="us-central1", + vertex_credentials=_FAKE_GEMMA_CREDENTIALS, + client=httpx.Client( + transport=slow_upstream_transport(json.dumps(_make_gemma_vertex_response(content="READY")).encode()) + ), + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in stream) == "READY" + + @pytest.mark.asyncio async def test_async_gemma_responses_stream(_gemma_cached_access_token): import httpx diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index 8c8b347be82..4a2fc0627ab 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -16,7 +16,12 @@ import pytest import respx import litellm -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from tests.unit.llms.bedrock.slow_upstream import ( + STREAM_TIMEOUT_SECONDS, + UPSTREAM_FIRST_BYTE_SECONDS, + slow_upstream_sync_client, +) def _expected_dir() -> Path: @@ -1035,3 +1040,39 @@ async def test_a_native_responses_provider_places_every_point_itself(): assert user_only == [0, 2, 4, 6] assert with_unmatchable_system == user_only + + +def test_non_streamed_response_ignores_a_stream_timeout_it_never_uses() -> None: + answer: Final = json.dumps(_minimal_responses_api_payload("resp_unstreamed", "gpt-5.6")).encode() + client: Final = HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, content=answer))) + ) + + response: Final = litellm.responses( + model="openai/gpt-5.6", input="hi", stream_timeout="unset", api_key="fake-key", client=client + ) + + assert response.output[0].content[0].text == "Done." + + +def test_faked_stream_accepts_a_text_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem( + litellm.model_cost, + "slow-faked-responses-model", + {"litellm_provider": "openai", "mode": "responses", "supports_native_streaming": False}, + ) + answer: Final = json.dumps(_minimal_responses_api_payload("resp_faked", "slow-faked-responses-model")).encode() + + response: Final = litellm.responses( + model="openai/slow-faked-responses-model", + input="hi", + stream=True, + timeout=str(UPSTREAM_FIRST_BYTE_SECONDS * 2), + stream_timeout=STREAM_TIMEOUT_SECONDS, + api_key="fake-key", + client=slow_upstream_sync_client(answer), + ) + + assert [event.response.output[0].content[0].text for event in response if event.type == "response.completed"] == [ + "Done." + ] diff --git a/tests/unit/test_completion_timeout_resolution.py b/tests/unit/test_completion_timeout_resolution.py index 77166d6e947..37b5cceef5a 100644 --- a/tests/unit/test_completion_timeout_resolution.py +++ b/tests/unit/test_completion_timeout_resolution.py @@ -10,7 +10,7 @@ from openai import Timeout as SDKTimeout sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))) -from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.utils import supports_httpx_timeout @@ -22,6 +22,7 @@ def test_explicit_timeout_wins(): "openai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == 12.5 ) @@ -35,6 +36,7 @@ def test_kwargs_timeout_when_param_none(): "azure_ai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == 21.0 ) @@ -48,6 +50,7 @@ def test_request_timeout_alias_in_kwargs(): "bedrock", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == 33.0 ) @@ -61,6 +64,7 @@ def test_global_timeout_from_litellm_settings(): "vertex_ai", global_timeout=360.0, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == 360.0 ) @@ -76,6 +80,7 @@ def test_explicit_global_timeout_6000_is_preserved(): "openai", global_timeout=6000.0, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == 6000.0 ) @@ -90,6 +95,7 @@ def test_explicit_request_timeout_6000_preserved(): "openai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == 6000.0 ) @@ -103,6 +109,7 @@ def test_explicit_model_timeout_6000_preserved(): "openai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == 6000.0 ) @@ -116,6 +123,7 @@ def test_fallback_600_when_no_global_timeout(): "azure_ai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == 600.0 ) @@ -129,6 +137,7 @@ def test_httpx_timeout_coerced_for_provider_without_httpx_timeout_support(): "azure_ai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) assert out == 50.0 assert not isinstance(out, httpx.Timeout) @@ -142,6 +151,7 @@ def test_httpx_timeout_preserved_for_openai(): "openai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) assert out is t assert isinstance(out, httpx.Timeout) @@ -158,6 +168,7 @@ def test_sdk_timeout_preserves_components_for_httpx_providers( "openai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) assert isinstance(resolved, httpx.Timeout) assert resolved.as_dict() == timeout.as_dict() @@ -174,6 +185,49 @@ def test_sdk_timeout_coerces_read_timeout_for_providers_without_httpx_support( "azure_ai", global_timeout=None, supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.NOT_STREAMED, ) == expected ) + + +@pytest.mark.parametrize( + "stream_kind, timeout, stream_timeout, expected", + ( + (StreamKind.NOT_STREAMED, 120.0, 15.0, 120.0), + (StreamKind.LIVE, 120.0, 15.0, 15.0), + (StreamKind.LIVE, 1.0, 15.0, 15.0), + (StreamKind.FAKED, 120.0, 15.0, 120.0), + (StreamKind.FAKED, 1.0, 15.0, 15.0), + (StreamKind.FAKED, httpx.Timeout(5.0, read=None), 15.0, httpx.Timeout(15.0, read=None)), + (StreamKind.FAKED, httpx.Timeout(1.0, read=120.0), 15.0, httpx.Timeout(15.0, read=120.0)), + ), +) +def test_stream_timeout_bounds_a_live_stream_and_never_shortens_a_fake_one( + stream_kind: StreamKind, timeout: float | httpx.Timeout, stream_timeout: float, expected: float | httpx.Timeout +) -> None: + assert ( + CompletionTimeout.resolve( + timeout, + {"stream_timeout": stream_timeout}, + "openai", + global_timeout=None, + supports_httpx_timeout=supports_httpx_timeout, + stream_kind=stream_kind, + ) + == expected + ) + + +def test_fake_stream_with_only_a_stream_timeout_keeps_the_global_request_timeout() -> None: + assert ( + CompletionTimeout.resolve( + None, + {"stream_timeout": 15.0}, + "bedrock", + global_timeout=None, + supports_httpx_timeout=supports_httpx_timeout, + stream_kind=StreamKind.FAKED, + ) + == 600.0 + ) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 0293af05c84..7c292ea8f08 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -48,6 +48,11 @@ from litellm.types.litellm_params import ControlOptions from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.utils import Delta, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, Usage +from tests.unit.llms.bedrock.slow_upstream import ( + STREAM_TIMEOUT_SECONDS, + UPSTREAM_FIRST_BYTE_SECONDS, + slow_upstream_sync_client, +) @pytest.fixture(autouse=True) @@ -1937,6 +1942,60 @@ def test_responses_bridge_preserves_reasoning_effort_with_drop_params( assert request_body["reasoning"] == {"effort": "high"} +_BRIDGED_RESPONSE_ANSWER: Final = json.dumps( + { + "id": "resp_slow", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "slow-responses-model", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + "tool_choice": "auto", + "tools": [], + } +).encode() + + +def test_responses_bridge_fake_stream_is_bounded_by_the_request_timeout_not_the_stream_timeout( + restore_model_registry: None, +) -> None: + model: Final = "openai/slow-responses-model" + litellm.register_model( + { + model: { + "litellm_provider": "openai", + "mode": "responses", + "supports_native_streaming": False, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + + response: Final = litellm.completion( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + stream_timeout=STREAM_TIMEOUT_SECONDS, + timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + api_key="fake-key", + client=slow_upstream_sync_client(_BRIDGED_RESPONSE_ANSWER), + ) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in response) == "hi" + + _FOUNDRY_RESPONSES_FUNCTION_CALL_BODY: Final = { "id": "resp_foundry", "object": "response", @@ -6978,6 +7037,23 @@ def test_mock_request_with_mock_timeout_raises_litellm_timeout(): assert error.value.model == "gpt-5.6" +@pytest.mark.asyncio +async def test_async_mock_timeout_on_a_stream_waits_only_for_the_stream_timeout() -> None: + with pytest.raises(litellm.Timeout): + await asyncio.wait_for( + litellm.acompletion( + model="gpt-5.6", + messages=[{"role": "user", "content": "hello"}], + stream=True, + timeout=30.0, + stream_timeout=0.01, + mock_timeout=True, + num_retries=0, + ), + timeout=5, + ) + + def test_router_mock_request_with_mock_timeout_raises_litellm_timeout(): router: Final = litellm.Router( model_list=[ diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index bd3cf738af8..1dec4de45f4 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -2023,7 +2023,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): ] ) - def inject_alias_into_kwargs(deployment, kwargs, function_name=None): + def inject_alias_into_kwargs( + deployment: dict, kwargs: dict, function_name: str | None = None, splits_stream_timeout: bool = False + ) -> None: # Simulate the alias leaking into kwargs (as happens when # _ageneric_api_call_with_fallbacks sets kwargs["model"] = alias before # calling the helper through async_function_with_fallbacks). diff --git a/tests/unit/test_router/test_router_timeout.py b/tests/unit/test_router/test_router_timeout.py index 3c3f0692043..3265179a513 100644 --- a/tests/unit/test_router/test_router_timeout.py +++ b/tests/unit/test_router/test_router_timeout.py @@ -3,6 +3,8 @@ from __future__ import annotations import asyncio import json import os +from datetime import datetime, timedelta +from collections.abc import Awaitable, Callable, Mapping from typing import Final from unittest.mock import AsyncMock, Mock @@ -13,9 +15,22 @@ import respx import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router_utils.cooldown_handlers import is_caller_timeout_408 +from litellm.types.integrations.custom_logger import WEBSEARCH_INTERCEPTION_PREFIX, as_converted_stream +from litellm.types.llms.openai import AllMessageValues import time from unittest.mock import patch, MagicMock +from tests.unit.llms.bedrock.slow_upstream import ( + CONVERSE_ANSWER, + STREAM_TIMEOUT_SECONDS, + UPSTREAM_FIRST_BYTE_SECONDS, + slow_upstream_answer, + slow_upstream_async_client, +) + @pytest.fixture def restore_request_timeout(monkeypatch: pytest.MonkeyPatch) -> None: @@ -138,6 +153,276 @@ def test_unit_test_streaming_timeout(stream): assert stream_timeout_val == normal_timeout +async def _drain_responses_stream(router: Router) -> None: + response: Final = await router.aresponses(model="gpt", input="hi", stream=True) + _ = [event async for event in response] + + +@pytest.mark.asyncio +async def test_live_responses_stream_through_the_router_fails_at_the_stream_timeout( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=slow_upstream_answer()) + router: Final = Router( + model_list=[{"model_name": "gpt", "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}}], + timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + stream_timeout=STREAM_TIMEOUT_SECONDS, + num_retries=0, + ) + + with pytest.raises(litellm.Timeout): + await _drain_responses_stream(router) + + +@pytest.mark.asyncio +async def test_prompt_managed_stream_keeps_the_target_deployments_own_stream_timeout() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "claude-agent", + "litellm_params": {"model": "litellm_agent/claude-target"}, + }, + { + "model_name": "claude-target", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-4-5-v1:0", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + "stream_timeout": STREAM_TIMEOUT_SECONDS, + }, + }, + ], + timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + stream_timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + num_retries=0, + ) + messages: Final[list[AllMessageValues]] = [{"role": "user", "content": "hi"}] + + with pytest.raises(litellm.Timeout): + await router.acompletion( + model="claude-agent", + messages=messages, + stream=True, + client=slow_upstream_async_client(), + litellm_logging_obj=LiteLLMLogging( + model="claude-agent", + messages=messages, + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="prompt-managed-stream", + function_id="f", + ), + ) + + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("agent_params", "target_params", "request_params"), + [ + ({}, {"timeout": UPSTREAM_FIRST_BYTE_SECONDS * 2}, {}), + ({"timeout": STREAM_TIMEOUT_SECONDS}, {}, {"request_timeout": UPSTREAM_FIRST_BYTE_SECONDS * 2}), + ], + ids=["target_deployment_timeout", "caller_request_timeout"], +) +async def test_prompt_managed_fake_stream_keeps_the_timeout_its_caller_or_target_set( + agent_params: Mapping[str, float], target_params: Mapping[str, float], request_params: Mapping[str, float] +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "jamba-agent", + "litellm_params": {"model": "litellm_agent/jamba-target", **agent_params}, + }, + { + "model_name": "jamba-target", + "litellm_params": { + "model": "bedrock/converse/ai21.jamba-1-5-large-v1:0", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + **target_params, + }, + }, + ], + timeout=STREAM_TIMEOUT_SECONDS, + stream_timeout=STREAM_TIMEOUT_SECONDS, + num_retries=0, + ) + messages: Final[list[AllMessageValues]] = [{"role": "user", "content": "hi"}] + + response: Final = await router.acompletion( + model="jamba-agent", + messages=messages, + stream=True, + client=slow_upstream_async_client(CONVERSE_ANSWER), + **request_params, + litellm_logging_obj=LiteLLMLogging( + model="jamba-agent", + messages=messages, + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="prompt-managed-fake-stream", + function_id="f", + ), + ) + + assert isinstance(response, litellm.CustomStreamWrapper) + assert "".join([chunk.choices[0].delta.content or "" async for chunk in response]) == "hi" + + +class _AnswersStreamsInOnePiece(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs: dict, call_type: object) -> dict: + return as_converted_stream(kwargs, WEBSEARCH_INTERCEPTION_PREFIX) if kwargs.get("stream") else kwargs + + +_INVOKE_ANSWER: Final = json.dumps( + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } +).encode() + + +async def _streamed_chat_text(router: Router) -> str: + response: Final = await router.acompletion(model="claude", messages=[{"role": "user", "content": "hi"}], stream=True) + return "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + + +async def _streamed_messages_text(router: Router) -> str: + response: Final = await router.aanthropic_messages( + model="claude", messages=[{"role": "user", "content": "hi"}], max_tokens=10, stream=True + ) + events: Final = [ + json.loads(line.removeprefix(b"data: ")) + for line in b"".join([chunk async for chunk in response]).splitlines() + if line.startswith(b"data: ") + ] + return "".join(event["delta"]["text"] for event in events if event["type"] == "content_block_delta") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "answer", "streamed_text"), + [ + ("bedrock/converse/anthropic.claude-sonnet-4-5-v1:0", CONVERSE_ANSWER, _streamed_chat_text), + ("bedrock/invoke/anthropic.claude-sonnet-4-5-v1:0", _INVOKE_ANSWER, _streamed_messages_text), + ], + ids=["chat_completions", "anthropic_messages"], +) +async def test_stream_a_deployment_hook_answers_in_one_piece_gets_the_longer_limit( + model: str, + answer: bytes, + streamed_text: Callable[[Router], Awaitable[str]], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [_AnswersStreamsInOnePiece()]) + respx_mock.post(url__regex=r"https://bedrock-runtime\.us-east-1\.amazonaws\.com/.*").mock( + side_effect=slow_upstream_answer(answer) + ) + router: Final = Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": { + "model": model, + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + "timeout": STREAM_TIMEOUT_SECONDS, + "stream_timeout": UPSTREAM_FIRST_BYTE_SECONDS * 2, + }, + } + ], + num_retries=0, + ) + + assert await streamed_text(router) == "hi" + + +class _FailureDetails(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: Final = asyncio.Event() + self.details: Mapping[str, object] = {} + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.details = kwargs + self.logged.set() + + +async def _drain_messages_stream(router: Router) -> None: + response: Final = await router.aanthropic_messages( + model="native", messages=[{"role": "user", "content": "hi"}], max_tokens=10, stream=True, client_side_timeout=True + ) + _ = [chunk async for chunk in response] + + +async def _drain_native_responses_stream(router: Router) -> None: + response: Final = await router.aresponses(model="native", input="hi", stream=True, client_side_timeout=True) + _ = [event async for event in response] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "drain"), + [ + ("bedrock/invoke/anthropic.claude-sonnet-4-5-v1:0", _drain_messages_stream), + ("openai/gpt-4.1", _drain_native_responses_stream), + ], + ids=["anthropic_messages", "responses"], +) +async def test_native_live_stream_cut_at_the_stream_timeout_is_not_blamed_on_the_deployment( + model: str, + drain: Callable[[Router], Awaitable[None]], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + failure_details: Final = _FailureDetails() + monkeypatch.setattr(litellm, "callbacks", [failure_details]) + respx_mock.post(url__regex=r"https://(bedrock-runtime\.us-east-1\.amazonaws\.com|api\.openai\.com)/.*").mock( + side_effect=slow_upstream_answer() + ) + router: Final = Router( + model_list=[ + { + "model_name": "native", + "litellm_params": { + "model": model, + "api_key": "fake", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + }, + } + ], + timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + stream_timeout=STREAM_TIMEOUT_SECONDS, + num_retries=0, + ) + + with pytest.raises(litellm.Timeout): + await drain(router) + await asyncio.wait_for(failure_details.logged.wait(), timeout=10) + + started: Final = failure_details.details.get("api_call_start_time") + assert isinstance(started, datetime) + assert is_caller_timeout_408(failure_details.details, 408, ended=started + timedelta(seconds=STREAM_TIMEOUT_SECONDS)) + @pytest.mark.asyncio async def test_bedrock_timeout_preserves_the_provider_timeout_error( respx_mock: respx.MockRouter, From 8b442f00fc2ccd57797e6405acbd36a5258adac3 Mon Sep 17 00:00:00 2001 From: Misbah Syed <35427888+misbahsy@users.noreply.github.com> Date: Sat, 17 Oct 2026 15:49:57 -0700 Subject: [PATCH 54/64] fix(proxy): count the whole minor line in the upgrade notice for patch releases (#46942) * fix(proxy): count the whole minor line in the upgrade notice for patch releases A patch release's notes only list its backports, so v1.105.1 showed "0 new features, 1 fix, and 0 other updates" even though upgrading from v1.104.x brings everything in v1.105.0 too. When the latest release is a patch, also fetch the recent release list and sum the stable releases on the same major.minor line (skipping rc/dev prereleases and drafts). Co-Authored-By: Claude Opus 5.5 * fix(proxy): use tuples for release lists and sync the dashboard API types Satisfies the LIT001 type-discipline gate and regenerates the latest_release_info description in schema.d.ts. Co-Authored-By: Claude Opus 5.5 * test(proxy): cover the bad release-list payload and non-semver latest tag Co-Authored-By: Claude Opus 5.5 * fix(proxy): type the release-line Counters so basedpyright sees no unknowns Co-Authored-By: Claude Opus 5.5 * fix(proxy): fetch earlier patches of the minor line by tag Reading the first page of /releases could miss the line's .0 release when prereleases push it off the page, and one unrelated release with a null body failed validation for the whole list. Fetch each earlier patch tag directly instead (v1.105.0 .. v1.105.N-1), skip tags with no stable release, and treat a null release body as empty notes. --------- Co-authored-by: Claude Opus 5.5 --- .../latest_release_endpoints.py | 75 +++++++++-- .../test_latest_release_endpoints.py | 120 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 3 +- 3 files changed, 183 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py index 720b4dd8d32..7411f1b4947 100644 --- a/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py @@ -18,6 +18,7 @@ from litellm.types.llms.base import LiteLLMBaseModel router: Final = APIRouter() LATEST_RELEASE_URL: Final = "https://api.github.com/repos/BerriAI/litellm/releases/latest" +RELEASE_BY_TAG_URL: Final = "https://api.github.com/repos/BerriAI/litellm/releases/tags/{tag}" LATEST_RELEASE_FETCH_TIMEOUT_SECONDS: Final = 5 LATEST_RELEASE_CACHE_TTL_SECONDS: Final = 60 * 60 LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS: Final = 5 * 60 @@ -25,6 +26,7 @@ LATEST_RELEASE_CACHE_KEY: Final = "latest_release_info" _RELEASE_BULLET_PATTERN: Final = re.compile(r"^\*\s+(?:([A-Za-z]+)(?:\([^)]*\))?!?:\s)?\S") _NEW_CONTRIBUTOR_PATTERN: Final = re.compile(r"^\*\s+@\S+ made their first contribution\b") +_STABLE_TAG_PATTERN: Final = re.compile(r"^v?(\d+)\.(\d+)\.(\d+)$") _Bucket: TypeAlias = Literal["new_features", "bug_fixes", "other_updates"] _PREFIX_BUCKETS: Final[Mapping[str, _Bucket]] = MappingProxyType({"feat": "new_features", "fix": "bug_fixes"}) @@ -46,7 +48,9 @@ class LatestReleaseUnavailable: class _GitHubRelease(LiteLLMBaseModel): tag_name: str html_url: str - body: str + body: str | None + prerelease: bool = False + draft: bool = False class _AsyncGetClient(Protocol): @@ -87,29 +91,71 @@ def count_release_bullets(body: str) -> Mapping[_Bucket, int]: return MappingProxyType(Counter(bucket for line in body.splitlines() if (bucket := _bucket_for(line)) is not None)) -def parse_latest_release(response: httpx.Response) -> LatestReleaseInfo | LatestReleaseUnavailable: - if response.status_code != 200: - return LatestReleaseUnavailable(reason=f"GitHub responded with status {response.status_code}") - try: - release: Final = _GitHubRelease.model_validate_json(response.content) - except ValidationError as e: - return LatestReleaseUnavailable(reason=f"GitHub release payload was not the expected shape: {e}") - counts: Final = count_release_bullets(release.body) +def _stable_version(tag_name: str) -> tuple[int, int, int] | None: + match: Final = _STABLE_TAG_PATTERN.match(tag_name) + return None if match is None else (int(match.group(1)), int(match.group(2)), int(match.group(3))) + + +def earlier_patch_tags(version: tuple[int, int, int]) -> tuple[str, ...]: + """Tags of the patches that precede ``version`` on its minor line (v1.105.0 for v1.105.1).""" + major, minor, patch = version + return tuple(f"v{major}.{minor}.{earlier}" for earlier in range(patch)) + + +def summarize_release_line(latest: _GitHubRelease, earlier: tuple[_GitHubRelease, ...]) -> LatestReleaseInfo: + """A patch release's notes only list its backports, so count the whole minor line it belongs to.""" + counts: Final = sum( + (Counter[_Bucket](count_release_bullets(release.body or "")) for release in (latest, *earlier)), + Counter[_Bucket](), + ) return LatestReleaseInfo( - version=release.tag_name.removeprefix("v"), + version=latest.tag_name.removeprefix("v"), new_features=counts.get("new_features", 0), bug_fixes=counts.get("bug_fixes", 0), other_updates=counts.get("other_updates", 0), - release_url=release.html_url, + release_url=latest.html_url, ) +def parse_github_release(response: httpx.Response) -> _GitHubRelease | LatestReleaseUnavailable: + if response.status_code != 200: + return LatestReleaseUnavailable(reason=f"GitHub responded with status {response.status_code}") + try: + return _GitHubRelease.model_validate_json(response.content) + except ValidationError as e: + return LatestReleaseUnavailable(reason=f"GitHub release payload was not the expected shape: {e}") + + +async def _fetch_earlier_release(client: _AsyncGetClient, tag: str) -> _GitHubRelease | LatestReleaseUnavailable | None: + """The stable release for ``tag``, or None when that patch was never published as one.""" + response: Final = await client.get(RELEASE_BY_TAG_URL.format(tag=tag), timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) + if response.status_code == 404: + return None + release: Final = parse_github_release(response) + if isinstance(release, _GitHubRelease) and (release.prerelease or release.draft): + return None + return release + + async def fetch_latest_release(client: _AsyncGetClient) -> LatestReleaseInfo | LatestReleaseUnavailable: try: - response: Final = await client.get(LATEST_RELEASE_URL, timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) + latest: Final = parse_github_release( + await client.get(LATEST_RELEASE_URL, timeout=LATEST_RELEASE_FETCH_TIMEOUT_SECONDS) + ) + if isinstance(latest, LatestReleaseUnavailable): + return latest + version: Final = _stable_version(latest.tag_name) + if version is None: + return summarize_release_line(latest, ()) + earlier: Final = await asyncio.gather( + *(_fetch_earlier_release(client, tag) for tag in earlier_patch_tags(version)) + ) except httpx.HTTPError as e: return LatestReleaseUnavailable(reason=f"{type(e).__name__}: {e}") - return parse_latest_release(response) + unavailable: Final = next((release for release in earlier if isinstance(release, LatestReleaseUnavailable)), None) + if unavailable is not None: + return unavailable + return summarize_release_line(latest, tuple(release for release in earlier if isinstance(release, _GitHubRelease))) async def get_latest_release_info( @@ -144,7 +190,8 @@ async def latest_release_info( fetch_lock: Annotated[asyncio.Lock, Depends(_default_fetch_lock)], ) -> LatestReleaseInfo | None: """ - Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. + Latest stable LiteLLM GitHub release with the PR count of its whole minor line (e.g. v1.105.0 + v1.105.1) + split into new features, bug fixes and other updates. Returns null when GitHub can't be reached so the dashboard upgrade banner simply doesn't render. """ result: Final = await get_latest_release_info(client=client, cache=cache, fetch_lock=fetch_lock) diff --git a/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py index ae0ac32c62d..25f1301dd9d 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_latest_release_endpoints.py @@ -16,12 +16,14 @@ from litellm.proxy.ui_crud_endpoints.latest_release_endpoints import ( LATEST_RELEASE_CACHE_TTL_SECONDS, LATEST_RELEASE_UNAVAILABLE_CACHE_TTL_SECONDS, LATEST_RELEASE_URL, + RELEASE_BY_TAG_URL, LatestReleaseInfo, LatestReleaseUnavailable, _default_cache, _default_client, _default_fetch_lock, count_release_bullets, + earlier_patch_tags, get_latest_release_info, ) from tests._master_key import MASTER_KEY @@ -233,6 +235,124 @@ class TestGetLatestReleaseInfo: assert len(client.calls) == 1 +def _release(tag: str, body: str | None, *, prerelease: bool = False) -> dict[str, object]: + return { + "tag_name": tag, + "html_url": f"https://github.com/BerriAI/litellm/releases/tag/{tag}", + "body": body, + "prerelease": prerelease, + "draft": False, + } + + +def _tag_url(tag: str) -> str: + return RELEASE_BY_TAG_URL.format(tag=tag) + + +class _GitHubByUrl: + """Answers each GitHub URL from a fixed map; unknown URLs are 404, like a tag with no release.""" + + def __init__(self, routes: dict[str, httpx.Response | Exception]) -> None: + self._routes = routes + self.urls: list[str] = [] + + async def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: + self.urls.append(url) + outcome = self._routes.get(url, httpx.Response(404, content=b"{}")) + if isinstance(outcome, Exception): + raise outcome + return outcome + + +PATCH_RELEASE: Final = _release("v1.102.2", "* fix(proxy): second backport by @k in https://x/pull/10\n") +FIRST_PATCH: Final = _release("v1.102.1", "* fix(proxy): first backport by @k in https://x/pull/9\n") + + +def _minor_line(**overrides: httpx.Response | Exception) -> dict[str, httpx.Response | Exception]: + routes: dict[str, httpx.Response | Exception] = { + LATEST_RELEASE_URL: _github_response(payload=PATCH_RELEASE), + _tag_url("v1.102.0"): _github_response(payload=SAMPLE_RELEASE), + _tag_url("v1.102.1"): _github_response(payload=FIRST_PATCH), + } + routes.update({_tag_url(tag): outcome for tag, outcome in overrides.items()}) + return routes + + +async def _info(client: _GitHubByUrl) -> LatestReleaseInfo | LatestReleaseUnavailable: + return await get_latest_release_info(client=client, cache=_fresh_cache(), fetch_lock=asyncio.Lock()) + + +class TestPatchReleaseAggregatesMinorLine: + def test_earlier_patch_tags_cover_the_whole_line(self): + assert earlier_patch_tags((1, 105, 2)) == ("v1.105.0", "v1.105.1") + assert earlier_patch_tags((1, 105, 0)) == () + + @pytest.mark.asyncio + async def test_patch_release_counts_include_every_earlier_patch_on_its_line(self): + client = _GitHubByUrl(_minor_line()) + result = await _info(client) + assert isinstance(result, LatestReleaseInfo) + assert result.model_dump() == { + "version": "1.102.2", + "new_features": 2, + "bug_fixes": 4, + "other_updates": 4, + "release_url": PATCH_RELEASE["html_url"], + } + assert sorted(client.urls) == sorted([LATEST_RELEASE_URL, _tag_url("v1.102.0"), _tag_url("v1.102.1")]) + + @pytest.mark.asyncio + async def test_minor_release_makes_a_single_request(self): + client = _RecordingClient([_github_response()]) + await get_latest_release_info(client=client, cache=_fresh_cache(), fetch_lock=asyncio.Lock()) + assert client.calls == [(LATEST_RELEASE_URL, 5)] + + @pytest.mark.asyncio + async def test_non_semver_latest_tag_counts_only_its_own_notes(self): + client = _GitHubByUrl({LATEST_RELEASE_URL: _github_response(payload=_release("v1.102.1-hotfix", "* fix: x\n"))}) + result = await _info(client) + assert isinstance(result, LatestReleaseInfo) + assert (result.new_features, result.bug_fixes, result.other_updates) == (0, 1, 0) + assert client.urls == [LATEST_RELEASE_URL] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "skipped", + [ + httpx.Response(404, content=b"{}"), + _github_response(payload=_release("v1.102.1", "* feat: rc only\n", prerelease=True)), + _github_response(payload=_release("v1.102.1", None)), + ], + ids=["never_published", "prerelease", "null_body"], + ) + async def test_missing_prerelease_or_empty_patch_adds_nothing(self, skipped: httpx.Response): + result = await _info(_GitHubByUrl(_minor_line(**{"v1.102.1": skipped}))) + assert isinstance(result, LatestReleaseInfo) + assert (result.new_features, result.bug_fixes, result.other_updates) == (2, 3, 4) + + @pytest.mark.asyncio + async def test_latest_release_with_null_body_still_counts_its_line(self): + routes = _minor_line() + routes[LATEST_RELEASE_URL] = _github_response(payload=_release("v1.102.2", None)) + result = await _info(_GitHubByUrl(routes)) + assert isinstance(result, LatestReleaseInfo) + assert (result.new_features, result.bug_fixes, result.other_updates) == (2, 3, 4) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "outcome", + [ + httpx.ConnectError("boom"), + _github_response(status=403, payload={"message": "rate limited"}), + httpx.Response(200, content=b"not json"), + ], + ids=["unreachable", "rate_limited", "not_json"], + ) + async def test_earlier_patch_failure_is_unavailable(self, outcome: httpx.Response | Exception): + result = await _info(_GitHubByUrl(_minor_line(**{"v1.102.0": outcome}))) + assert isinstance(result, LatestReleaseUnavailable) + + class TestLatestReleaseInfoEndpoint: def test_returns_release_stats_for_authenticated_user(self, http_client): _override_dependencies(_RecordingClient([_github_response()]), _fresh_cache(), LitellmUserRoles.INTERNAL_USER) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f56fa87e93b..468811966de 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -5832,7 +5832,8 @@ export interface paths { }; /** * Latest Release Info - * @description Latest stable LiteLLM GitHub release with its PR count split into new features, bug fixes and other updates. + * @description Latest stable LiteLLM GitHub release with the PR count of its whole minor line (e.g. v1.105.0 + v1.105.1) + * split into new features, bug fixes and other updates. * Returns null when GitHub can't be reached so the dashboard upgrade banner simply doesn't render. */ get: operations["latest_release_info_get_latest_release_info_get"]; From e2295b596fe4f879786b4c8ce8f92cee175d8458 Mon Sep 17 00:00:00 2001 From: Aayush Jain Date: Sun, 18 Oct 2026 10:22:19 +1100 Subject: [PATCH 55/64] feat(openinfer): add OpenInfer as an OpenAI-compatible provider (#41034) * feat(openinfer): add OpenInfer as OpenAI-compatible provider Register OpenInfer in the JSON provider registry with catalog models, proxy create fields, and unit coverage so requests route via openinfer/ Co-authored-by: Cursor * fix(openinfer): assert catalog invariants instead of vendor rates Pinning per-token prices made CI fail when the vendor list changed. Keep the backup agreement and positivity checks our code owns. Co-authored-by: Cursor * fix(openinfer): add openinfer to the Rust OpenAI-compatible provider mirror * fix(openinfer): align catalog with OpenInfer's live /v1/models Replace the four catalog rows with the two models OpenInfer serves and publishes a price for (@oi/GLM-5.3-Flash, @oi/Qwen3.5-27B), using its listed per-token rates and context windows, and price-check them in tests * fix(openinfer): use OpenInfer's native Responses and Messages APIs, allow tools on unpriced models, price Qwen3.5-9B Declare /v1/responses and /v1/messages for openinfer, since OpenInfer serves both natively. Let JSON providers declare provider-wide capabilities in providers.json, read through get_provider_info, so tool calls reach OpenInfer for models without a cost map row. Send max_completion_tokens as is, set output caps to each model's context window, mark the three rows as reasoning models, add Qwen3.5-9B at its listed rate, and wire openinfer into models_by_provider so the openinfer/* wildcard lists its models. * fix(openinfer): let a model's own capability row win over provider defaults in get_model_info Provider defaults now only fill capability fields the cost-map row leaves unset, so a row or deployment that turns function calling off stays off. Also marks openinfer_models as Final. --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../src/exception_mapping_utils/original.rs | 1 + litellm/__init__.py | 5 + litellm/constants.py | 2 + litellm/llms/openai_like/README.md | 6 + litellm/llms/openai_like/json_loader.py | 8 + litellm/llms/openai_like/providers.json | 10 + ...odel_prices_and_context_window_backup.json | 36 +++ .../provider_endpoints_support_backup.json | 17 ++ .../provider_create_fields.json | 28 ++ litellm/types/utils.py | 1 + litellm/utils.py | 10 +- model_prices_and_context_window.json | 36 +++ provider_endpoints_support.json | 17 ++ .../openai_like/test_openinfer_provider.py | 247 ++++++++++++++++++ 14 files changed, 420 insertions(+), 4 deletions(-) create mode 100644 tests/unit/llms/openai_like/test_openinfer_provider.py diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs index 82392cbd7ee..4de67e6e7b5 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs @@ -94,6 +94,7 @@ const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[ "docker_model_runner", "ragflow", "pinstripes", + "openinfer", "darkbloom", "meta", "cognition", diff --git a/litellm/__init__.py b/litellm/__init__.py index 9cc8965b628..cc5fbc1e72e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -676,6 +676,7 @@ qwen_ai_platform_models: Set = set() moonshot_models: Set = set() publicai_models: Set = set() darkbloom_models: Set = set() +openinfer_models: Final[Set[str]] = set() v0_models: Set = set() morph_models: Set = set() lambda_ai_models: Set = set() @@ -945,6 +946,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje publicai_models.add(key) elif value.get("litellm_provider") == "darkbloom": darkbloom_models.add(key) + elif value.get("litellm_provider") == "openinfer": + openinfer_models.add(key) elif value.get("litellm_provider") == "v0": v0_models.add(key) elif value.get("litellm_provider") == "morph": @@ -1113,6 +1116,7 @@ model_list = list( | moonshot_models | publicai_models | darkbloom_models + | openinfer_models | v0_models | morph_models | lambda_ai_models @@ -1229,6 +1233,7 @@ def _build_models_by_provider() -> dict: "moonshot": moonshot_models, "publicai": publicai_models, "darkbloom": darkbloom_models, + "openinfer": openinfer_models, "v0": v0_models, "morph": morph_models, "lambda_ai": lambda_ai_models, diff --git a/litellm/constants.py b/litellm/constants.py index 94ff02a8c26..da95b97036a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1030,6 +1030,7 @@ openai_compatible_endpoints: Final[list[str]] = [ "https://api.prisminference.com/v1", "https://api.reka.ai/v1", "https://gigachat.devices.sberbank.ru/api/v1", + "https://api.openinfer.ai/v1", ] @@ -1098,6 +1099,7 @@ openai_compatible_providers: Final[list] = [ "docker_model_runner", "ragflow", "pinstripes", # Pinstripes - JSON-configured provider + "openinfer", "darkbloom", "meta", # Meta Model API (Muse Spark) - JSON-configured provider "cognition", diff --git a/litellm/llms/openai_like/README.md b/litellm/llms/openai_like/README.md index e1409b81c35..c56f58b24c9 100644 --- a/litellm/llms/openai_like/README.md +++ b/litellm/llms/openai_like/README.md @@ -63,6 +63,12 @@ That's it! The provider will be automatically loaded and available. // Optional: Special handling flags "special_handling": { "convert_content_list_to_string": true + }, + + // Optional: Capabilities every model of this provider has, used when a + // model has no cost map row or its row leaves the field unset + "provider_info": { + "supports_function_calling": true } } } diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index 5cdaff90d24..23b63177137 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -6,7 +6,12 @@ import json from pathlib import Path from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_logger +from litellm.types.utils import ProviderSpecificModelInfo + +_PROVIDER_INFO_ADAPTER: Final = TypeAdapter(ProviderSpecificModelInfo) class SimpleProviderConfig: @@ -22,6 +27,9 @@ class SimpleProviderConfig: self.constraints = data.get("constraints", {}) self.special_handling = data.get("special_handling", {}) self.supported_endpoints = data.get("supported_endpoints", []) + self.provider_info: ProviderSpecificModelInfo = ( + _PROVIDER_INFO_ADAPTER.validate_python(data["provider_info"]) if "provider_info" in data else {} + ) class JSONProviderRegistry: diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index aa8a17c3824..36337fddbdc 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -230,6 +230,16 @@ "api_key_env": "REKA_API_KEY", "api_base_env": "REKA_API_BASE" }, + "openinfer": { + "base_url": "https://api.openinfer.ai/v1", + "api_key_env": "OPENINFER_API_KEY", + "api_base_env": "OPENINFER_API_BASE", + "base_class": "openai_gpt", + "provider_info": { + "supports_function_calling": true + }, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + }, "inferenceapis": { "base_url": "https://api.inferenceapis.com/v1", "api_key_env": "INFERENCEAPIS_API_KEY", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 13210a7eb1c..6ad071ad2ec 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -75526,6 +75526,42 @@ "supports_reasoning": true, "supports_vision": true }, + "openinfer/@oi/GLM-5.3-Flash": { + "max_tokens": 393216, + "max_input_tokens": 393216, + "max_output_tokens": 393216, + "input_cost_per_token": 1.95e-07, + "output_cost_per_token": 1.56e-06, + "litellm_provider": "openinfer", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "source": "https://api.openinfer.ai/v1/models" + }, + "openinfer/@oi/Qwen3.5-27B": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1.95e-07, + "output_cost_per_token": 1.56e-06, + "litellm_provider": "openinfer", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "source": "https://api.openinfer.ai/v1/models" + }, + "openinfer/@oi/Qwen3.5-9B": { + "max_tokens": 100000, + "max_input_tokens": 100000, + "max_output_tokens": 100000, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "openinfer", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "source": "https://api.openinfer.ai/v1/models" + }, "xai/grok-voice-3.0": { "input_cost_per_second": 0.001333333333333, "input_cost_per_text_message": 0.004, diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index a10f2f93d48..f95b553c2b4 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -2731,6 +2731,23 @@ "messages": true, "responses": true } + }, + "openinfer": { + "display_name": "OpenInfer (`openinfer`)", + "url": "https://docs.litellm.ai/docs/providers/openinfer", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } } }, "endpoints": { diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index cb9e193fa48..a794854ce19 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2735,6 +2735,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "OPENINFER", + "provider_display_name": "OpenInfer", + "litellm_provider": "openinfer", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.openinfer.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "openinfer/@oi/Qwen3.5-27B" + }, { "provider": "OPENAI_LIKE", "provider_display_name": "Openai Like", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8e56ebdc39b..6c47ac7fa32 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4389,6 +4389,7 @@ class LlmProviders(str, Enum): CORTECS = "cortecs" SCX_AI = "scx-ai" PRISM = "prism" + OPENINFER = "openinfer" INFERENCEAPIS = "inferenceapis" DARKBLOOM = "darkbloom" META = "meta" diff --git a/litellm/utils.py b/litellm/utils.py index 9136e3f12d6..a283ea0e1e5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6265,11 +6265,13 @@ def get_provider_info(model: str, custom_llm_provider: str | None) -> ProviderSp model=model, provider=LlmProviders(custom_llm_provider) ) - model_info: ProviderSpecificModelInfo | None = None if provider_config: - model_info = provider_config.get_provider_info(model=model) + return provider_config.get_provider_info(model=model) - return model_info + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + json_provider: Final = JSONProviderRegistry.get(custom_llm_provider) if custom_llm_provider else None + return json_provider.provider_info if json_provider is not None and json_provider.provider_info else None def _is_potential_model_name_in_model_cost( @@ -6726,7 +6728,7 @@ def _build_model_info( provider_info: Final = get_provider_info(model=model, custom_llm_provider=custom_llm_provider) if provider_info: for key, value in provider_info.items(): - if value is not None: + if value is not None and _model_info.get(key) is None: _model_info[key] = value # if verbose_logger.isEnabledFor(logging.DEBUG): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 13210a7eb1c..6ad071ad2ec 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -75526,6 +75526,42 @@ "supports_reasoning": true, "supports_vision": true }, + "openinfer/@oi/GLM-5.3-Flash": { + "max_tokens": 393216, + "max_input_tokens": 393216, + "max_output_tokens": 393216, + "input_cost_per_token": 1.95e-07, + "output_cost_per_token": 1.56e-06, + "litellm_provider": "openinfer", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "source": "https://api.openinfer.ai/v1/models" + }, + "openinfer/@oi/Qwen3.5-27B": { + "max_tokens": 32768, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "input_cost_per_token": 1.95e-07, + "output_cost_per_token": 1.56e-06, + "litellm_provider": "openinfer", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "source": "https://api.openinfer.ai/v1/models" + }, + "openinfer/@oi/Qwen3.5-9B": { + "max_tokens": 100000, + "max_input_tokens": 100000, + "max_output_tokens": 100000, + "input_cost_per_token": 5e-06, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "openinfer", + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "source": "https://api.openinfer.ai/v1/models" + }, "xai/grok-voice-3.0": { "input_cost_per_second": 0.001333333333333, "input_cost_per_text_message": 0.004, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index ad38b064903..c2cec2c5490 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -3254,6 +3254,23 @@ "rerank": false, "a2a": false } + }, + "openinfer": { + "display_name": "OpenInfer (`openinfer`)", + "url": "https://docs.litellm.ai/docs/providers/openinfer", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } } }, "endpoints": { diff --git a/tests/unit/llms/openai_like/test_openinfer_provider.py b/tests/unit/llms/openai_like/test_openinfer_provider.py new file mode 100644 index 00000000000..7bcc489fcf8 --- /dev/null +++ b/tests/unit/llms/openai_like/test_openinfer_provider.py @@ -0,0 +1,247 @@ +import json +from pathlib import Path +from typing import Final + +import pytest +import respx + +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.proxy.auth.model_checks import get_provider_models + +_REPO_ROOT: Final = Path(__file__).resolve().parents[4] +_PRIMARY_PRICES: Final = json.loads((_REPO_ROOT / "model_prices_and_context_window.json").read_text()) +_BACKUP_PRICES: Final = json.loads( + (_REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json").read_text() +) +_CATALOG_MODELS: Final = tuple( + key for key, row in _PRIMARY_PRICES.items() if isinstance(row, dict) and row.get("litellm_provider") == "openinfer" +) +_REASONING_MODELS: Final = tuple(key for key in _CATALOG_MODELS if _PRIMARY_PRICES[key].get("supports_reasoning")) +_CHAT_URL: Final = "https://api.openinfer.ai/v1/chat/completions" +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +_CHAT_RESPONSE: Final = { + "id": "chatcmpl_openinfer", + "object": "chat.completion", + "created": 1_792_000_000, + "model": "@oi/Qwen3.5-27B", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello from OpenInfer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}, +} + + +def _sent_chat_request(model: str, **params: object) -> tuple[str, str, dict[str, object]]: + with respx.mock() as upstream: + route: Final = upstream.post(_CHAT_URL).respond(200, json=_CHAT_RESPONSE) + litellm.completion( + model=f"openinfer/{model}", + messages=[{"role": "user", "content": "Say hello"}], + api_key="sk-test", + **params, + ) + request: Final = route.calls.last.request + return str(request.url), request.headers["authorization"], json.loads(request.content) + + +def test_provider_prefixed_model_routes_to_openinfer(): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, api_key, api_base = get_llm_provider( + model="openinfer/@oi/Qwen3.5-27B", + custom_llm_provider=None, + api_base=None, + api_key="sk-test", + ) + + assert model == "@oi/Qwen3.5-27B" + assert provider == "openinfer" + assert api_key == "sk-test" + assert api_base == "https://api.openinfer.ai/v1" + + +def test_api_key_and_base_resolved_from_env(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("OPENINFER_API_KEY", "sk-env-key") + monkeypatch.setenv("OPENINFER_API_BASE", "https://proxy.internal/v1") + + _, provider, api_key, api_base = get_llm_provider( + model="openinfer/@oi/Qwen3.5-9B", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert provider == "openinfer" + assert api_key == "sk-env-key" + assert api_base == "https://proxy.internal/v1" + + +def test_url_autodetection_from_api_base(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("OPENINFER_API_KEY", "sk-env-key") + + _, provider, api_key, _ = get_llm_provider( + model="@oi/Qwen3.5-27B", + custom_llm_provider=None, + api_base="https://api.openinfer.ai/v1", + api_key=None, + ) + + assert provider == "openinfer" + assert api_key == "sk-env-key" + + +def test_chat_request_reaches_openinfer_with_the_bare_model_id(): + url, authorization, body = _sent_chat_request("@oi/Qwen3.5-27B") + + assert url == _CHAT_URL + assert authorization == "Bearer sk-test" + assert body["model"] == "@oi/Qwen3.5-27B" + + +def test_tools_reach_openinfer_for_a_model_without_a_cost_map_row(): + _, _, body = _sent_chat_request("@oi/model-without-a-catalog-row", tools=[_WEATHER_TOOL]) + + assert body["tools"] == [_WEATHER_TOOL] + + +def _register_openinfer_row(model: str, **capabilities: bool) -> None: + litellm.register_model( + {f"openinfer/{model}": {"litellm_provider": "openinfer", "mode": "chat", **capabilities}}, + persist_across_reloads=False, + add_to_provider_model_sets=False, + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_model_info_keeps_a_row_that_turns_tools_off(): + model: Final = "@oi/row-without-tools" + _register_openinfer_row(model, supports_function_calling=False) + + assert litellm.get_model_info(model, custom_llm_provider="openinfer")["supports_function_calling"] is False + assert litellm.supports_function_calling(model, custom_llm_provider="openinfer") is False + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_model_info_fills_tool_support_from_the_openinfer_default(): + model: Final = "@oi/row-silent-on-tools" + _register_openinfer_row(model) + + assert litellm.get_model_info(model, custom_llm_provider="openinfer")["supports_function_calling"] is True + assert litellm.supports_function_calling(model, custom_llm_provider="openinfer") is True + + +def test_max_tokens_and_max_completion_tokens_both_reach_openinfer(): + _, _, body = _sent_chat_request("@oi/Qwen3.5-27B", max_tokens=64, max_completion_tokens=128) + + assert body["max_tokens"] == 64 + assert body["max_completion_tokens"] == 128 + + +@pytest.mark.parametrize("model", _REASONING_MODELS) +def test_reasoning_effort_reaches_openinfer_for_reasoning_rows(model: str): + _, _, body = _sent_chat_request(model.removeprefix("openinfer/"), reasoning_effort="low") + + assert body["reasoning_effort"] == "low" + + +def test_responses_request_goes_to_openinfer_natively(): + with respx.mock() as upstream: + route: Final = upstream.post("https://api.openinfer.ai/v1/responses").respond( + 200, + json={ + "id": "resp_openinfer", + "object": "response", + "created_at": 1_792_000_000, + "model": "@oi/Qwen3.5-27B", + "status": "completed", + "output": [ + { + "id": "msg_openinfer", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Hello from OpenInfer", "annotations": []}], + } + ], + "usage": {"input_tokens": 4, "output_tokens": 3, "total_tokens": 7}, + }, + ) + response: Final = litellm.responses(model="openinfer/@oi/Qwen3.5-27B", input="Say hello", api_key="sk-test") + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert request.headers["authorization"] == "Bearer sk-test" + assert body["model"] == "@oi/Qwen3.5-27B" + assert body["input"] == "Say hello" + assert response.output[0].content[0].text == "Hello from OpenInfer" + + +@pytest.mark.asyncio +async def test_anthropic_messages_request_goes_to_openinfer_natively(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + with respx.mock() as upstream: + route: Final = upstream.post("https://api.openinfer.ai/v1/messages").respond( + 200, + json={ + "id": "msg_openinfer", + "type": "message", + "role": "assistant", + "model": "@oi/Qwen3.5-27B", + "content": [{"type": "text", "text": "Hello from OpenInfer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 4, "output_tokens": 3}, + }, + ) + response: Final = await litellm.anthropic.messages.acreate( + model="openinfer/@oi/Qwen3.5-27B", + messages=[{"role": "user", "content": "Say hello"}], + max_tokens=32, + api_key="sk-test", + ) + + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert request.headers["authorization"] == "Bearer sk-test" + assert body["model"] == "@oi/Qwen3.5-27B" + assert body["messages"] == [{"role": "user", "content": "Say hello"}] + assert response["content"][0]["text"] == "Hello from OpenInfer" + + +def test_wildcard_lists_every_catalog_model(): + assert sorted(get_provider_models("openinfer") or []) == sorted(_CATALOG_MODELS) + + +@pytest.mark.parametrize("model", _CATALOG_MODELS) +def test_catalog_row_matches_its_backup_and_fits_its_context(model: str): + row: Final = _PRIMARY_PRICES[model] + + assert row == _BACKUP_PRICES[model] + assert row["mode"] == "chat" + assert row["max_tokens"] == row["max_output_tokens"] + assert 0 < row["max_output_tokens"] <= row["max_input_tokens"] + assert row["input_cost_per_token"] > 0 + assert row["output_cost_per_token"] > 0 + + +@pytest.mark.parametrize("model", _CATALOG_MODELS) +def test_chat_response_is_priced_at_its_catalog_row(model: str): + row: Final = litellm.model_cost[model] + response: Final = litellm.ModelResponse( + model=model.removeprefix("openinfer/"), + usage=litellm.Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200), + ) + + cost: Final = litellm.completion_cost(completion_response=response, model=model) + + assert cost == pytest.approx(1000 * row["input_cost_per_token"] + 200 * row["output_cost_per_token"]) From c58fc0c99638d279f2c0f091ed738a24ecf7b182 Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Sat, 17 Oct 2026 16:51:48 -0700 Subject: [PATCH 56/64] fix(rag)!: use team provider credentials and allow viewer and JWT team access (#45497) * fix(rag): use team provider credentials and allow viewer and JWT team access Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rag): assert persisted managed store data excludes team credentials Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(jwt): run real team model checks in the route denial test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(rag): scope team provider credentials to openai vector stores Bedrock and other providers kept their pre-PR credential and region resolution: a team bedrock deployment's region no longer overrides the caller's region on RAG query. Query tests now fake the provider at the HTTP boundary with a real vector store registry. * fix(rag): use only the team's own OpenAI deployment for RAG credentials * fix(rag): ignore caller credential fields the query search never uses when picking team credentials * fix(rag): treat blank credential fields as absent when picking team credentials * test(rag): delete the stores the integration ingests create before the owned proxy exits * fix(proxy): let a restricted key use the team deployment its wildcard or model covers A key restricted to gpt-4o on a team whose own OpenAI deployment is openai/* never resolved that deployment's credentials, because the key allowlist was matched against the deployment's public name literally. The allowlist now matches a wildcard on either side the way request authorization does, through the auth layer's own helper, so /rag, /files and /batches resolve the team deployment for every key that may call a model on it. * fix(batches): read a restricted key's batch input through its team deployment The batch access check downloaded the input JSONL with the proxy's OPENAI_API_KEY while the upload and the create used the team deployment, so a key restricted under a team openai/* deployment could not create a batch when the two keys belong to different OpenAI projects. The read now resolves the team deployment the way the create does when no model pins it Adds the audit cells for keys restricted under the team wildcard on /rag, /files and /batches, and for the rag-vs-chat split of a key restricted to a global exact-name deployment * fix(rag): load deployments oldest first, reject file_id on stores that cannot take one, and show ingest errors in the dashboard * fix(rag): pick the team deployment with the lowest id, return 400 for caller mistakes, keep the ingest id in errors, and show partial success in the dashboard * fix(rag): pick the team's oldest OpenAI deployment so adding one never moves the team * fix(rag): keep config file team deployments first and try JWT teams in sorted order A team's config file OpenAI deployments now win in file order, as on the base, before its DB ones oldest first, so upgrading never moves a team that has both. A JWT in several teams tries them in sorted order, so every worker resolves the same team. * fix(rag): return 400 for Bedrock, S3 Vectors and Gemini store config mistakes on ingest * fix(rag): keep forwarded provider keys on /rag/query and mixed JWT team claims working A forwarded x-api-key now counts as a caller credential, so the /rag/query search uses it over the team deployment as at the merge base. Sorting the JWT team claim compares entries as text, so a stray number or null no longer turns into a 401. A Gemini store with no proxy key is a 500 again, and a non-string vector_store_id keeps the merge base's 400 message. * fix(rag): let a forwarded provider key win over team credentials on files, batches and ingest * fix(rag): keep team credentials when a request carries its own key or api_base A forwarded x-api-key or a body api_key now replaces only the key of the team deployment, keeping its api_base, as chat completions does. A body litellm_credential_name no longer turns team credentials off on /v1/batches. * fix(batches): keep the team key on the team host when a request sends only api_base A request api_base now replaces the team deployment's api_base only when the request also brings its own api_key, and a non-string api_key or api_base is ignored, as on the merge base. * fix(files): ignore a blank request api_key or api_base next to team credentials * fix(files): answer 400 for a request api_key that cannot be sent as an HTTP header * fix(rag): honor a body drop_params for a forwarded key that cannot be an HTTP header * fix(rag): read the request field for drop_params only on multipart ingest bodies * fix(rag): type the ingest body read for drop_params * fix(batches): count a model-routed batch's input with the deployment key it is created with * fix(batches): reject a batch create whose x-litellm-model headers name different models * fix(rag): treat a null provider and whitespace store credentials as unset when picking team credentials * refactor(files): type the team deployment the credential resolver ranks --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/_types.py | 4 + litellm/proxy/auth/auth_checks.py | 4 +- litellm/proxy/auth/handle_jwt.py | 17 +- litellm/proxy/batches_endpoints/endpoints.py | 9 + litellm/proxy/hooks/batch_rate_limiter.py | 71 +- .../openai_files_endpoints/common_utils.py | 141 +- litellm/proxy/rag_endpoints/endpoints.py | 183 +- litellm/rag/ingestion/base_ingestion.py | 23 +- litellm/rag/ingestion/bedrock_ingestion.py | 6 +- litellm/rag/ingestion/openai_ingestion.py | 4 +- litellm/rag/ingestion/s3_vectors_ingestion.py | 6 +- litellm/rag/ingestion/vertex_ai_ingestion.py | 6 +- litellm/types/rag.py | 11 +- .../test_rag_query_vector_store_allowlist.py | 26 + .../test_rag_team_provider_credentials.py | 2316 +++++++++++++++++ tests/unit/proxy/auth/test_handle_jwt.py | 136 + .../proxy/batches_endpoints/test_endpoints.py | 47 +- .../proxy/hooks/test_batch_rate_limiter.py | 151 +- .../test_files_common_utils.py | 205 +- .../proxy/rag_endpoints/test_rag_endpoints.py | 756 +++++- .../test_batch_x_litellm_model_encoding.py | 5 +- .../unit/rag/ingestion/test_base_ingestion.py | 164 ++ .../_components/CreateVectorStore.test.tsx | 39 + .../_components/CreateVectorStore.tsx | 25 +- .../src/components/networking.test.ts | 24 + .../src/components/networking.tsx | 2 +- 26 files changed, 4173 insertions(+), 208 deletions(-) create mode 100644 tests/integration/authorization/test_rag_team_provider_credentials.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a7bd16804c1..5970ede5c7e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4829,6 +4829,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): "mcp_routes", "/v1/messages", "/v1/messages/count_tokens", + "/rag/ingest", + "/v1/rag/ingest", + "/rag/query", + "/v1/rag/query", ] team_id_default: str | None = Field( default=None, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b395a826625..d0708fcda2f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4784,7 +4784,7 @@ def check_model_access_helper( if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases): return True - if _model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models): + if model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models): return True all_model_access: bool = False @@ -7434,7 +7434,7 @@ def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool: return False -def _model_matches_any_wildcard_pattern_in_list(model: str, allowed_model_list: list) -> bool: +def model_matches_any_wildcard_pattern_in_list(model: str, allowed_model_list: Sequence[str]) -> bool: """ Returns True if a model matches any wildcard pattern in a list. diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 05b8cca19af..9d800a1fd63 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1767,6 +1767,7 @@ class JWTAuthManager: from litellm.proxy.proxy_server import llm_router denied_pass_through_team: LiteLLM_TeamTable | None = None + route_denied_team_id: str | None = None # rebind-ok: retain route-denied team while checking all claims if not team_ids: if ( @@ -1780,7 +1781,7 @@ class JWTAuthManager: return None, None any_claim_team_resolved = False - for team_id in team_ids: + for team_id in sorted(team_ids, key=str): try: team_object = await get_team_object( team_id=team_id, @@ -1809,7 +1810,9 @@ class JWTAuthManager: user_route=route, litellm_proxy_roles=jwt_handler.litellm_jwtauth, ) - if is_allowed and not JWTAuthManager._team_has_passthrough_route_access( + if not is_allowed: + route_denied_team_id = team_id + elif not JWTAuthManager._team_has_passthrough_route_access( team_object=team_object, route=route, request_method=request_method, @@ -1832,7 +1835,13 @@ class JWTAuthManager: # Claim resolved but no model access, or fallback disabled — deny. raise HTTPException( status_code=403, - detail=f"No team has access to the requested model: {requested_model}. Checked teams={team_ids}. Check `/models` to see all available models.", + detail=( + f"Team {route_denied_team_id} can access model {requested_model} but route {route} " + "is not in litellm_jwtauth.team_allowed_routes" + if route_denied_team_id is not None + else f"No team has access to the requested model: {requested_model}. " + f"Checked teams={team_ids}. Check `/models` to see all available models." + ), ) # No claim team resolved and fallback enabled — defer to fallback. @@ -2003,7 +2012,7 @@ class JWTAuthManager: status_code=403, detail=( f"x-litellm-team-id '{header_value}' does not resolve to a team id or a unique team alias in your " - f"JWT's allowed teams. Allowed team ids: {sorted(allowed_team_ids)}" + f"JWT's allowed teams. Allowed team ids: {sorted(allowed_team_ids, key=str)}" ), ) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index ead1cd8e2ca..763737191ef 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -149,6 +149,14 @@ async def _raise_when_input_file_must_be_managed(model: str, credentials: Mappin ) +def _raise_when_routing_headers_disagree(request: Request, body_model: object) -> None: + if (isinstance(body_model, str) and body_model) or request.query_params.get("model"): + return + if len(frozenset(request.headers.getlist("x-litellm-model"))) < 2: + return + raise batch_error(400, "Send one x-litellm-model header. This request has several naming different models") + + def _litellm_metadata_of(data: MutableMapping[str, object]) -> MutableMapping[str, object]: """The request's litellm_metadata mapping, created on the request when it carries none. @@ -274,6 +282,7 @@ async def create_batch( data: dict = {} try: data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line + _raise_when_routing_headers_disagree(request, data.get("model")) verbose_proxy_logger.debug( "Request received by LiteLLM:\n%s", json.dumps(data, indent=4), diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 5f214816578..94762232e0b 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -22,6 +22,7 @@ from collections.abc import Callable, Iterable, Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias +from urllib.parse import parse_qs, urlsplit from fastapi import HTTPException from pydantic import Field, TypeAdapter, ValidationError @@ -99,6 +100,28 @@ else: _BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object]) _WINDOW_START_ADAPTER: Final[TypeAdapter[int | float | str | None]] = TypeAdapter(int | float | str | None) + +def _batch_create_routing_model(data: Mapping[str, object]) -> str | None: + body_model: Final = data.get("model") + if isinstance(body_model, str) and body_model: + return body_model + proxy_server_request: Final = data.get("proxy_server_request") + if not isinstance(proxy_server_request, Mapping): + return None + request_line: Final = _BATCH_BODY_ADAPTER.validate_python(proxy_server_request) + url: Final = request_line.get("url") + query_model: Final = ( + parse_qs(urlsplit(url).query, keep_blank_values=True).get("model", [None])[-1] if isinstance(url, str) else None + ) + if query_model: + return query_model + headers: Final = request_line.get("headers") + header_model: Final = ( + _BATCH_BODY_ADAPTER.validate_python(headers).get("x-litellm-model") if isinstance(headers, Mapping) else None + ) + return header_model if isinstance(header_model, str) and header_model else None + + IncrementAmounts: TypeAlias = dict[Literal["requests", "tokens"], int] @@ -509,6 +532,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_id: str, custom_llm_provider: str, data: dict, + user_api_key_dict: UserAPIKeyAuth | None = None, ) -> tuple[str, dict[str, Any]]: """ Map proxy-facing file IDs to provider file IDs and credentials. @@ -519,8 +543,10 @@ class _PROXY_BatchRateLimiter(CustomLogger): """ from litellm.proxy.openai_files_endpoints.common_utils import ( decode_model_from_file_id, + explicit_request_credentials, get_credentials_for_model, get_original_file_id, + get_team_provider_credentials, ) from litellm.proxy.proxy_server import llm_router @@ -546,22 +572,36 @@ class _PROXY_BatchRateLimiter(CustomLogger): pass return get_original_file_id(file_id), fetch_kwargs - request_model: Final = data.get("model") - if isinstance(request_model, str) and request_model and llm_router is not None: - try: - credentials = get_credentials_for_model( - llm_router=llm_router, - model_id=request_model, - operation_context="batch input file read (rate limiting)", - ) - fetch_kwargs.update(extract_file_access_credentials(credentials)) - fetch_kwargs["model"] = request_model - provider = credentials.get("custom_llm_provider") - if provider: - fetch_kwargs["custom_llm_provider"] = provider - except HTTPException: - pass + request_model: Final = _batch_create_routing_model(data) + if request_model is not None: + if llm_router is not None: + try: + credentials = get_credentials_for_model( + llm_router=llm_router, + model_id=request_model, + operation_context="batch input file read (rate limiting)", + ) + fetch_kwargs.update(extract_file_access_credentials(credentials)) + fetch_kwargs["model"] = request_model + provider = credentials.get("custom_llm_provider") + if provider: + fetch_kwargs["custom_llm_provider"] = provider + except HTTPException: + pass + return file_id, fetch_kwargs + team_credentials: Final = ( + get_team_provider_credentials( + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) + if user_api_key_dict is not None + else None + ) + fetch_kwargs.update( + extract_file_access_credentials({**(team_credentials or {}), **explicit_request_credentials(data)}) + ) return file_id, fetch_kwargs async def _reserve_batch_enqueued_tokens( @@ -912,6 +952,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_id=file_id, custom_llm_provider=custom_llm_provider, data=data or {}, + user_api_key_dict=user_api_key_dict, ) # For non-managed files, use the standard litellm.afile_content file_content = await litellm.afile_content( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 48306cc8d7a..2e07694876b 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -4,6 +4,7 @@ import mimetypes import re from collections.abc import Mapping, Sequence from dataclasses import dataclass, field +from datetime import datetime, timezone from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -17,6 +18,7 @@ from typing import ( ) from pydantic import TypeAdapter +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.batches.batch_utils import batch_cost_is_final from litellm.constants import MAX_FILE_LIST_LIMIT @@ -43,6 +45,7 @@ FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500 BATCH_CREATE_HIDDEN_PARAM: Final = "batch_create" LITELLM_EXECUTED_BATCH_ID_PREFIX: Final = "litellm_batch_" +EXPLICIT_CREDENTIAL_KEYS: Final = frozenset({"api_key", "api_base", "litellm_credential_name"}) def validate_file_list_limit(limit: int | None) -> None: @@ -474,10 +477,47 @@ async def get_authorized_credentials_for_model( ) +def _key_allowlist_covers(public_model_name: str, key_model_allowlist: Sequence[str]) -> bool: + from litellm.proxy.auth.auth_checks import model_matches_any_wildcard_pattern_in_list + + if public_model_name in key_model_allowlist: + return True + if model_matches_any_wildcard_pattern_in_list(model=public_model_name, allowed_model_list=key_model_allowlist): + return True + return any( + model_matches_any_wildcard_pattern_in_list(model=key_model, allowed_model_list=(public_model_name,)) + for key_model in key_model_allowlist + ) + + +_CONFIG_FILE_RANK: Final = datetime.min.replace(tzinfo=timezone.utc) +_UNKNOWN_CREATION_TIME: Final = datetime.max.replace(tzinfo=timezone.utc) + + +class _RouterDeployment(TypedDict): + model_name: ReadOnly[NotRequired[str]] + model_info: ReadOnly[Mapping[str, object]] + + +def _team_deployment_rank(model_info: Mapping[str, object], deployment_id: str) -> tuple[datetime, str]: + if model_info.get("db_model") is not True: + return _CONFIG_FILE_RANK, "" + created_at: Final = model_info.get("created_at") + if not isinstance(created_at, datetime): + return _UNKNOWN_CREATION_TIME, deployment_id + return (created_at if created_at.tzinfo else created_at.replace(tzinfo=timezone.utc)), deployment_id + + +def _by_team_deployment_rank(order: tuple[tuple[datetime, str], str]) -> tuple[datetime, str]: + return order[0] + + def get_team_provider_credentials( llm_router: Optional["Router"], user_api_key_dict: "UserAPIKeyAuth", custom_llm_provider: str, + *, + include_accessible_deployments: bool = True, ) -> dict | None: """ Resolve upstream credentials for a provider-scoped file operation @@ -490,12 +530,15 @@ def get_team_provider_credentials( global one. 2. Fallback: any deployment the caller is granted access to for this provider, expanding wildcard routes and the all-proxy-models sentinel. + Skipped when ``include_accessible_deployments`` is False. Credential lookup is scoped to both the team's allowlist and the key's own model allowlist (``user_api_key_dict.models``), so neither a team nor a restricted key within a team can resolve a provider key for a deployment it isn't authorized to use. A key restricted to an explicit model list - only narrows the team scope; sentinel-bearing keys (all-proxy-models / + only narrows the team scope, matching a wildcard on either side the way + request authorization does (a key allowed ``gpt-4o`` may use the team's + ``openai/*`` deployment); sentinel-bearing keys (all-proxy-models / all-team-models) defer to the team scope instead of widening past it. Returns None when the router is unavailable or no authorized deployment matches, so the caller can fall back to default credential resolution. @@ -531,33 +574,52 @@ def get_team_provider_credentials( if key_is_restricted else () ) - key_model_allowlist_set: Final = frozenset(key_model_allowlist) - def _key_may_use(public_model_name: str | None) -> bool: - if not key_model_allowlist_set: + def _key_may_use(public_model_name: object) -> bool: + if not key_model_allowlist: return True - return public_model_name is not None and public_model_name in key_model_allowlist_set + return isinstance(public_model_name, str) and _key_allowlist_covers(public_model_name, key_model_allowlist) - def _provider_credentials(model_id: str) -> dict | None: + def _provider_credentials(model_id: str) -> dict[str, object] | None: credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id) if credentials is not None and credentials.get("custom_llm_provider") == custom_llm_provider: return {key: value for key, value in credentials.items() if key != "model"} return None - # 1. Prefer the team's own BYOK deployment, matched by model_info.team_id. + def _usable_team_deployment_order(deployment: _RouterDeployment) -> tuple[tuple[datetime, str], str] | None: + model_info: Final[Mapping[str, object]] = deployment["model_info"] + if model_info.get("team_id") != team_id: + return None + if not _key_may_use(model_info.get("team_public_model_name") or deployment.get("model_name")): + return None + deployment_id: Final = model_info.get("id") + if not isinstance(deployment_id, str): + return None + return _team_deployment_rank(model_info, deployment_id), deployment_id + + # 1. Prefer the team's own BYOK deployment, matched by model_info.team_id. Config file + # deployments win, kept in file order by the stable sort, then DB ones oldest first, then + # by lowest id, so adding a deployment never moves the team off the one it already uses, + # and every worker picks the same one in any DB load or edit order. if team_id is not None: - for deployment in llm_router.model_list or []: - model_info = deployment.get("model_info") or {} - if model_info.get("team_id") != team_id: - continue - deployment_id = model_info.get("id") - if deployment_id is None: - continue - if not _key_may_use(model_info.get("team_public_model_name") or deployment.get("model_name")): - continue - credentials = _provider_credentials(deployment_id) - if credentials is not None: - return credentials + team_deployment_orders: Final = sorted( + ( + order + for deployment in llm_router.model_list or [] + if (order := _usable_team_deployment_order(deployment)) is not None + ), + key=_by_team_deployment_rank, + ) + team_deployment_ids: Final = tuple(deployment_id for _, deployment_id in team_deployment_orders) + team_credentials: Final = next( + (credentials for credentials in map(_provider_credentials, team_deployment_ids) if credentials is not None), + None, + ) + if team_credentials is not None: + return team_credentials + + if not include_accessible_deployments: + return None # 2. Fall back to deployments the caller is allowed to access. The key's # effective allowlist (sentinels and access groups already expanded by @@ -594,6 +656,42 @@ def get_team_provider_credentials( return None +HTTP_HEADER_VALUE: Final = re.compile(r"[\x21-\x7e]+(?:[ \t]+[\x21-\x7e]+)*") + + +def _non_blank_string(value: object) -> str | None: + return value if isinstance(value, str) and value.strip() else None + + +def _request_api_key(data: Mapping[str, object]) -> str | None: + from fastapi import HTTPException + + import litellm + + api_key: Final = _non_blank_string(data.get("api_key")) + if api_key is None or HTTP_HEADER_VALUE.fullmatch(api_key): + return api_key + if data.get("drop_params") is True or litellm.drop_params is True: + return None + raise HTTPException( + status_code=400, + detail={ + "error": "The request's provider api_key is not a valid HTTP header value. Remove surrounding " + "whitespace and non-ASCII characters, or set drop_params to send the request without it." + }, + ) + + +def explicit_request_credentials(data: Mapping[str, object]) -> Mapping[str, str]: + api_key: Final = _request_api_key(data) + if api_key is None: + return MappingProxyType({}) + api_base: Final = _non_blank_string(data.get("api_base")) + if api_base is None: + return MappingProxyType({"api_key": api_key}) + return MappingProxyType({"api_key": api_key, "api_base": api_base}) + + def apply_team_provider_credentials( data: dict, # mutable-ok: credentials are merged into the request payload in place, same contract as prepare_data_with_credentials llm_router: Optional["Router"], @@ -605,6 +703,9 @@ def apply_team_provider_credentials( ``get_team_provider_credentials`` and merge them into ``data`` in-place. Leaves ``data`` untouched when no authorized deployment matches, so the caller falls back to environment-variable credentials exactly as before. + An ``api_key`` the request carries itself (a provider key forwarded with + ``forward_llm_provider_auth_headers``) replaces the deployment's, and the + request's ``api_base`` replaces the deployment's only alongside that key. """ credentials: Final = get_team_provider_credentials( llm_router=llm_router, @@ -613,7 +714,7 @@ def apply_team_provider_credentials( ) if credentials is None: return - prepare_data_with_credentials(data=data, credentials=credentials) + prepare_data_with_credentials(data=data, credentials={**credentials, **explicit_request_credentials(data)}) def add_internal_model_credentials( diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index a1a7d62899f..a7146f13035 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -13,8 +13,9 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import orjson -from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import ORJSONResponse, StreamingResponse +from pydantic import TypeAdapter from starlette.datastructures import UploadFile import litellm @@ -38,6 +39,11 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, safe_get_request_headers, ) +from litellm.proxy.openai_files_endpoints.common_utils import ( + EXPLICIT_CREDENTIAL_KEYS, + explicit_request_credentials, + get_team_provider_credentials, +) from litellm.proxy.rag_endpoints.upload_security import ( MAX_UPLOAD_SIZE_BYTES, EicarTestMalwareScanner, @@ -57,9 +63,13 @@ from litellm.repositories.table_repositories import ManagedVectorStoresRepositor from litellm.types.utils import ModelResponse if TYPE_CHECKING: + from litellm import Router from litellm.proxy.utils import PrismaClient router: Final = APIRouter() +_TEAM_PROVIDER_CREDENTIALS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_EMPTY_TEAM_PROVIDER_CREDENTIALS: Final[Mapping[str, object]] = MappingProxyType({}) +_TEAM_CREDENTIAL_PROVIDER: Final = "openai" def _as_string_keyed_mapping(value: object) -> Mapping[str, object] | None: @@ -186,21 +196,91 @@ def _caller_vector_store_options( managed_store: LiteLLM_ManagedVectorStore | None, ) -> Mapping[str, object]: if managed_store is None: - return request_vector_store_config + return _without_blank_credentials(request_vector_store_config) return MappingProxyType( {key: value for key, value in request_vector_store_config.items() if key in _MANAGED_STORE_CALLER_OPTIONS} ) +def _is_blank(value: object) -> bool: + return value is None or (isinstance(value, str) and not value.strip()) + + +def _without_blank_credentials(config: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType( + {key: value for key, value in config.items() if key not in EXPLICIT_CREDENTIAL_KEYS or not _is_blank(value)} + ) + + +def _store_provider(config: Mapping[str, object]) -> object: + provider: Final = config.get("custom_llm_provider") + return _TEAM_CREDENTIAL_PROVIDER if provider is None else provider + + def _managed_store_overrides(managed_store: LiteLLM_ManagedVectorStore | None) -> Mapping[str, object]: if managed_store is None: return MappingProxyType({}) + stored: Final = build_request_data_from_managed_vector_store(managed_store) + return _without_blank_credentials({key: value for key, value in stored.items() if value is not None}) + + +def _forwarded_store_credentials( + vector_store_config: Mapping[str, object], request_data: Mapping[str, object] +) -> Mapping[str, object]: + if _store_provider(vector_store_config) != _TEAM_CREDENTIAL_PROVIDER: + return _EMPTY_TEAM_PROVIDER_CREDENTIALS + return explicit_request_credentials(request_data) + + +def _with_body_drop_params(request_data: Mapping[str, object], body: Mapping[str, object]) -> Mapping[str, object]: + if body.get("drop_params") is not True: + return request_data + return MappingProxyType({**request_data, "drop_params": True}) + + +async def _rag_ingest_body(request: Request) -> Mapping[str, object]: + body: Final = _TEAM_PROVIDER_CREDENTIALS_ADAPTER.validate_python(await read_request_body(request)) + if "multipart/form-data" not in request.headers.get("content-type", ""): + return body + form_request: Final = body.get("request") + parsed: Final[object] = ( + orjson.loads(form_request) if isinstance(form_request, (str, bytes)) and form_request else None + ) + return ( + _TEAM_PROVIDER_CREDENTIALS_ADAPTER.validate_python(parsed) if isinstance(parsed, dict) else MappingProxyType({}) + ) + + +def _with_request_credentials( + team_credentials: Mapping[str, object], request_data: Mapping[str, object] +) -> Mapping[str, object]: + if not team_credentials: + return team_credentials + return MappingProxyType({**team_credentials, **explicit_request_credentials(request_data)}) + + +def _team_provider_credentials( + effective_config: Mapping[str, object], + llm_router: "Router | None", + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + if not EXPLICIT_CREDENTIAL_KEYS.isdisjoint(effective_config): + return _EMPTY_TEAM_PROVIDER_CREDENTIALS + + if _store_provider(effective_config) != _TEAM_CREDENTIAL_PROVIDER: + return _EMPTY_TEAM_PROVIDER_CREDENTIALS + + validated_credentials: Final = _TEAM_PROVIDER_CREDENTIALS_ADAPTER.validate_python( + get_team_provider_credentials( + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=_TEAM_CREDENTIAL_PROVIDER, + include_accessible_deployments=False, + ) + or {} + ) return MappingProxyType( - { - key: value - for key, value in build_request_data_from_managed_vector_store(managed_store).items() - if value is not None - } + {key: value for key, value in validated_credentials.items() if key not in {"custom_llm_provider", "model"}} ) @@ -591,18 +671,6 @@ async def rag_ingest( request, scanner=EicarTestMalwareScanner() ) - # INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores only - if user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value and not ingest_options.get( - "vector_store", {} - ).get("vector_store_id"): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": "internal_user_viewer role can only ingest files to an existing vector store. " - "Provide 'vector_store_id' in ingest_options.vector_store." - }, - ) - resolved_stores: Final = await _authorize_nested_vector_store_ids( payload=ingest_options, user_api_key_dict=user_api_key_dict, @@ -620,10 +688,26 @@ async def rag_ingest( raise HTTPException(status_code=400, detail={"error": str(e)}) managed_store: Final = resolved_stores.get(request_vector_store_config.get("vector_store_id")) - merged_vector_store_config: Final = { + base: Final = { **_caller_vector_store_options(request_vector_store_config, managed_store), **_managed_store_overrides(managed_store), } + request_data: Final[dict[str, Any]] = await add_litellm_data_to_request( + data={}, + request=request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_config=proxy_config, + ) + forwarded_credentials: Final = _forwarded_store_credentials( + base, _with_body_drop_params(request_data, await _rag_ingest_body(request)) + ) + merged_vector_store_config: Final = { + **_team_provider_credentials(base, llm_router, user_api_key_dict), + **forwarded_credentials, + **base, + } merged_ingest_options: Final = { **ingest_options, "vector_store": merged_vector_store_config, @@ -636,17 +720,6 @@ async def rag_ingest( detail={"error": provider_error}, ) - # Add litellm data - request_data: dict[str, Any] = {} - request_data = await add_litellm_data_to_request( - data=request_data, - request=request, - general_settings=general_settings, - user_api_key_dict=user_api_key_dict, - version=version, - proxy_config=proxy_config, - ) - verbose_proxy_logger.debug( "RAG Ingest - options: %s, custom_llm_provider: %s", ingest_options, @@ -662,6 +735,11 @@ async def rag_ingest( router=llm_router, **request_data, ) + if response.get("status") == "failed": + raise HTTPException( + status_code=response.get("error_status_code") or 500, + detail={"id": response.get("id"), "status": "failed", "error": response.get("error")}, + ) # Save vector store to database if it was newly created and prisma_client is available verbose_proxy_logger.debug( @@ -777,7 +855,7 @@ async def rag_query( # Extract required fields model: Final = data.get("model") messages: Final = data.get("messages") - retrieval_config: Final = data.get("retrieval_config") + raw_retrieval_config: Final = data.get("retrieval_config") rerank: Final = data.get("rerank") stream: Final = data.get("stream", False) @@ -792,21 +870,28 @@ async def rag_query( status_code=400, detail={"error": "messages is required"}, ) - if not retrieval_config: + if not raw_retrieval_config: raise HTTPException( status_code=400, detail={"error": "retrieval_config is required"}, ) - if not isinstance(retrieval_config, dict): + if not isinstance(raw_retrieval_config, dict): raise HTTPException( status_code=400, detail={"error": "retrieval_config must be an object"}, ) + retrieval_config: Final = _TEAM_PROVIDER_CREDENTIALS_ADAPTER.validate_python(raw_retrieval_config) if "vector_store_id" not in retrieval_config: raise HTTPException( status_code=400, detail={"error": "retrieval_config must contain 'vector_store_id'"}, ) + vector_store_id: Final = retrieval_config.get("vector_store_id") + if not isinstance(vector_store_id, str): + raise HTTPException( + status_code=400, + detail={"error": "vector_store_id must be a non-empty string"}, + ) reject_caller_embedding_selection_params(payload=retrieval_config, source="retrieval_config") resolved_stores: Final = await _authorize_nested_vector_store_ids( payload=retrieval_config, @@ -817,27 +902,39 @@ async def rag_query( # model, credentials, ...) from the registry: the same source the direct # /vector_stores/{id}/search endpoint uses. Store-managed keys win on # conflict so callers cannot override the store's provider or credentials. - managed_store: Final = resolved_stores.get(retrieval_config["vector_store_id"]) - store_data: Final = ( + managed_store: Final = resolved_stores.get(vector_store_id) + store_data: Final[Mapping[str, object]] = ( build_request_data_from_managed_vector_store(managed_store) if managed_store is not None - else MappingProxyType({}) + else _EMPTY_TEAM_PROVIDER_CREDENTIALS ) + managed_store_params: Final = _managed_store_overrides(managed_store) merged_retrieval_config: Final = { **retrieval_config, **store_data, } - - # Add litellm data - request_data: dict[str, object] = {} - request_data = await add_litellm_data_to_request( - data=request_data, + effective_vector_store_config: Final = ( + managed_store_params + if managed_store is not None + else MappingProxyType( + {key: value for key, value in retrieval_config.items() if key not in EXPLICIT_CREDENTIAL_KEYS} + ) + ) + request_data: Final[dict[str, object]] = await add_litellm_data_to_request( + data={}, request=request, general_settings=general_settings, user_api_key_dict=user_api_key_dict, version=version, proxy_config=proxy_config, ) + team_credentials: Final = _with_request_credentials( + _team_provider_credentials(effective_vector_store_config, llm_router, user_api_key_dict), + _with_body_drop_params(request_data, data), + ) + vector_store_params: Final = ( + {**team_credentials, **managed_store_params} if managed_store is not None else team_credentials + ) verbose_proxy_logger.debug( "RAG Query - model: %s, vector_store_id: %s, custom_llm_provider: %s", @@ -851,7 +948,7 @@ async def rag_query( model=model, messages=messages, retrieval_config=merged_retrieval_config, - vector_store_params=store_data, + vector_store_params=vector_store_params, rerank=rerank, stream=stream, router=llm_router, diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index f9e39e042f1..8b2cf6e0a4e 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -34,6 +34,10 @@ if TYPE_CHECKING: from litellm import Router +class RAGIngestBadRequestError(ValueError): + status_code: Final = 400 + + class BaseRAGIngestion(ABC): """ Base class for RAG ingestion. @@ -319,13 +323,18 @@ class BaseRAGIngestion(ABC): file_id=file_id, ) - try: - if existing_file_id and not self.supports_existing_file_id: - raise ValueError( - f"{self.__class__.__name__} does not support ingesting an existing file_id. " - "Upload file data or provide file_url instead." - ) + if existing_file_id and not self.supports_existing_file_id: + return RAGIngestResponse( + id=self.ingest_id, + status="failed", + vector_store_id="", + file_id=None, + error=f"{self.__class__.__name__} does not support ingesting an existing file_id. " + "Upload file data or provide file_url instead.", + error_status_code=400, + ) + try: # Step 2: OCR (optional) extracted_text: Final = await self.ocr( file_content=file_content, @@ -361,10 +370,12 @@ class BaseRAGIngestion(ABC): except Exception as e: verbose_logger.exception("RAG Pipeline failed: %s", e) + error_status_code: Final[object] = getattr(e, "status_code", None) return RAGIngestResponse( id=self.ingest_id, status="failed", vector_store_id="", file_id=None, error=str(e), + error_status_code=error_status_code if isinstance(error_status_code, int) else None, ) diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index e412b7bcc8d..1e7aa6ac1c6 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_arn_prefix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM -from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError if TYPE_CHECKING: from litellm import Router @@ -149,7 +149,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): data_sources: Final = ds_response.get("dataSourceSummaries", []) if not data_sources: - raise ValueError( + raise RAGIngestBadRequestError( f"No data sources found for Knowledge Base {self.knowledge_base_id}. " "Please create a data source first or provide data_source_id and s3_bucket." ) @@ -181,7 +181,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): self.s3_prefix = prefixes[0] else: if not self._s3_bucket: - raise ValueError( + raise RAGIngestBadRequestError( f"Could not auto-detect S3 bucket for data source {self.data_source_id}. " "Please provide s3_bucket in config." ) diff --git a/litellm/rag/ingestion/openai_ingestion.py b/litellm/rag/ingestion/openai_ingestion.py index 6c7c619237d..e6a3f0bb74f 100644 --- a/litellm/rag/ingestion/openai_ingestion.py +++ b/litellm/rag/ingestion/openai_ingestion.py @@ -10,7 +10,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Final import litellm -from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError from litellm.vector_store_files.main import acreate as vector_store_file_acreate from litellm.vector_stores.main import acreate as vector_store_acreate @@ -87,7 +87,7 @@ class OpenAIRAGIngestion(BaseRAGIngestion): api_base: Final = self.vector_store_config.get("api_base") if existing_file_id and not vector_store_id: - raise ValueError("vector_store_id is required when ingesting an existing file_id") + raise RAGIngestBadRequestError("vector_store_id is required when ingesting an existing file_id") # Create vector store if not provided if not vector_store_id: diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index e3eafd680de..f31b068e709 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -37,7 +37,7 @@ from litellm.llms.s3_vectors.vector_stores.transformation import ( s3_vectors_ingest_embedding_options, s3_vectors_ingest_target, ) -from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError if TYPE_CHECKING: import httpx @@ -269,13 +269,13 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Validate bucket name (AWS S3 naming rules) if len(self.vector_bucket_name) < 3: - raise ValueError( + raise RAGIngestBadRequestError( f"Invalid vector_bucket_name '{self.vector_bucket_name}': " f"AWS S3 bucket names must be at least 3 characters long. " f"Please provide a valid bucket name (e.g., 'my-vector-bucket')." ) if not self.vector_bucket_name.replace("-", "").replace(".", "").isalnum(): - raise ValueError( + raise RAGIngestBadRequestError( f"Invalid vector_bucket_name '{self.vector_bucket_name}': " f"AWS S3 bucket names can only contain lowercase letters, numbers, hyphens, and periods. " f"Please provide a valid bucket name (e.g., 'my-vector-bucket')." diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index eff8ad1b8cb..cceddf31117 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -22,7 +22,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError if TYPE_CHECKING: from litellm import Router @@ -134,7 +134,9 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): Tuple of (rag_corpus_id, file_id) """ if not self.project_id: - raise ValueError("vertex_project is required for Vertex AI RAG ingestion. Set it in vector_store config.") + raise RAGIngestBadRequestError( + "vertex_project is required for Vertex AI RAG ingestion. Set it in vector_store config." + ) # Get or create RAG corpus rag_corpus_id = self.vector_store_config.get("vector_store_id") diff --git a/litellm/types/rag.py b/litellm/types/rag.py index 00c070b0996..2f912d21587 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -219,11 +219,12 @@ class RAGIngestOptions(TypedDict, total=False): class RAGIngestResponse(TypedDict, total=False): """Response from RAG ingest API.""" - id: str # Unique ingest job ID - status: Literal["completed", "in_progress", "failed"] - vector_store_id: str # The vector store ID (created or existing) - file_id: str | None # The file ID in the vector store - error: str | None # Error message if status is "failed" + id: ReadOnly[str] # Unique ingest job ID + status: ReadOnly[Literal["completed", "in_progress", "failed"]] + vector_store_id: ReadOnly[str] # The vector store ID (created or existing) + file_id: ReadOnly[str | None] # The file ID in the vector store + error: ReadOnly[str | None] # Error message if status is "failed" + error_status_code: ReadOnly[int | None] class RAGIngestRequest(LiteLLMBaseModel): diff --git a/tests/integration/authorization/test_rag_query_vector_store_allowlist.py b/tests/integration/authorization/test_rag_query_vector_store_allowlist.py index 39f0ad1c7e9..f4b5063441f 100644 --- a/tests/integration/authorization/test_rag_query_vector_store_allowlist.py +++ b/tests/integration/authorization/test_rag_query_vector_store_allowlist.py @@ -321,6 +321,32 @@ def test_deny_by_default_searches_explicitly_granted_store(strict_gateway: Stric assert marker in str(object_value(searches[0]["body"])["query"]), searches +def test_jwt_team_route_outside_operator_allowed_routes_is_denied_before_search( + strict_gateway: StrictGateway, +) -> None: + with strict_gateway.gateway.scenario() as scenario: + model: Final = scenario.model() + store_id: Final = f"vs_unregistered_{uuid.uuid4().hex}" + team: Final = scenario.team(models=_json_array(model), object_permission=_permission_for_stores(store_id)) + member: Final = scenario.member(team) + marker: Final = f"lit9370 alias {uuid.uuid4().hex}" + + response: Final = _rag_query( + strict_gateway.gateway, + model, + marker, + strict_gateway.jwt(member, (team,)), + store_id=store_id, + path="/rag/query", + ) + + assert response.status_code == 403, response.text + assert ( + f"Team {team} can access model {model} but route /rag/query is not in litellm_jwtauth.team_allowed_routes" + ) in response.text, response.text + assert _searches_for_marker(strict_gateway.upstream, marker, store_id) == () + + @pytest.mark.parametrize( ("scope", "error_type"), (("key", "key_vector_store_access_denied"), ("team", "team_vector_store_access_denied")), diff --git a/tests/integration/authorization/test_rag_team_provider_credentials.py b/tests/integration/authorization/test_rag_team_provider_credentials.py new file mode 100644 index 00000000000..8d556b89107 --- /dev/null +++ b/tests/integration/authorization/test_rag_team_provider_credentials.py @@ -0,0 +1,2316 @@ +from __future__ import annotations + +import asyncio +import json +import re +import threading +import time +import uuid +from collections.abc import Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import anthropic +import httpx +import jwt +import openai +import psutil +import pytest +import yaml +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.ports import allocate_port +from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +TEAM_API_KEY: Final = "integration-rag-team-key" +WILDCARD_API_KEY: Final = "integration-rag-wildcard-key" +SHARED_API_KEY: Final = "integration-rag-shared-key" +ENVIRONMENT_API_KEY: Final = "integration-rag-environment-key" +CALLER_API_KEY: Final = "integration-rag-caller-key" +FORWARDED_API_KEY: Final = "integration-rag-forwarded-key" +CREDENTIAL_API_KEY: Final = "integration-rag-credential-key" +GEMINI_API_KEY: Final = "integration-rag-gemini-key" +AZURE_API_KEY: Final = "integration-rag-azure-key" +JWT_KEY_ID: Final = "integration-rag-team-provider-jwt" +WILDCARD_MODEL: Final = "openai/gpt-4o-mini" +BARE_WILDCARD_MODEL: Final = "gpt-4o-mini" +OPENAI: Final[Mapping[str, JsonValue]] = MappingProxyType({"custom_llm_provider": "openai"}) +DOCUMENT: Final = ("document.txt", b"team provider credentials", "text/plain") +BATCH_INPUT_ROW: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "custom_id": "integration-rag-batch", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": BARE_WILDCARD_MODEL, "messages": [{"role": "user", "content": "batch input"}]}, + } +) +BATCH_INPUT: Final = ("batch.jsonl", (json.dumps(dict(BATCH_INPUT_ROW)) + "\n").encode(), "application/jsonl") +FAILING_STORE_NAMES: Final = MappingProxyType({"rag-upstream-401": 401, "rag-upstream-429": 429}) +OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 +STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +PREFIXED: Final = re.compile(r"^/(team|wildcard|shared|environment|credential|caller|gemini|azure|bedrock)(/[^?]*)") +STORE_FILES: Final = re.compile(r"^/v1/vector_stores/([^/]+)/files$") +FILE_CONTENT: Final = re.compile(r"^/v1/files/([^/]+)/content$") +FILE_ITEM: Final = re.compile(r"^/v1/files/([^/]+)$") +BATCH_ITEM: Final = re.compile(r"^/v1/batches/([^/]+)$") +BATCH_CANCEL: Final = re.compile(r"^/v1/batches/([^/]+)/cancel$") +BEDROCK_DATA_SOURCES: Final = re.compile(r"^/knowledgebases/([^/]+)/datasources/$") +BEDROCK_DATA_SOURCE: Final = re.compile(r"^/knowledgebases/([^/]+)/datasources/([^/]+)$") +EMPTY_KNOWLEDGE_BASE: Final = "EMPTYKB000" +BUCKETLESS_KNOWLEDGE_BASE: Final = "NOBUCKET00" +WEB_DATA_SOURCE: Final = "WEBSOURCE0" +STORE_SEARCH: Final = re.compile(r"^(?:/v1|/openai)/vector_stores/([^/]+)/search$") +GEMINI_UPLOAD: Final = re.compile(r"^/upload/v1beta/fileSearchStores/([^:]+):uploadToFileSearchStore$") + +JWT_TEAM_COUNT: Final = 10 +AWS_AMBIENT_CREDENTIALS: Final = ( + "AWS_PROFILE", + "AWS_SESSION_TOKEN", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_ROLE_ARN", + "AWS_ROLE_NAME", + "AWS_SESSION_NAME", + "AWS_WEB_IDENTITY_TOKEN_FILE", +) + +pytestmark: Final = pytest.mark.timeout(OWNED_PROXY_CELL_SECONDS) + +Signature = tuple[str, str, str, str | None] + + +@dataclass(frozen=True, slots=True) +class Call: + prefix: str + method: str + route: str + authorization: str | None + google_key: str | None + azure_key: str | None + + @property + def signature(self) -> Signature: + return (self.prefix, self.method, self.route, self.authorization) + + +@dataclass(frozen=True, slots=True) +class Tenants: + shared_model: str + global_model: str + team_a: str + team_a_model: str + team_a_key: str + restricted_key: str + viewer_key: str + jwt_member: str + team_w: str + team_w_key: str + wildcard_bare_key: str + wildcard_prefixed_key: str + wildcard_global_key: str + team_s: str + team_s_key: str + teamless_key: str + jwt_teams: tuple[str, ...] + jwt_team_model: str + jwt_multi_member: str + azure_restricted_key: str + credential_name: str + unreachable_credential_name: str + + +@dataclass(frozen=True, slots=True) +class RagRig: + gateway: Gateway + upstream: Wire + outage: threading.Event + tenants: Tenants + signing_key: rsa.RSAPrivateKey + scenario: Scenario + + def jwt(self, subject: str, *team_ids: JsonValue) -> str: + issued_at: Final = int(time.time()) + return jwt.encode( + {"sub": subject, "groups": list(team_ids), "iat": issued_at, "exp": issued_at + 300}, + self.signing_key, + algorithm="RS256", + headers={"kid": JWT_KEY_ID}, + ) + + def calls(self) -> tuple[Call, ...]: + return tuple(_call(request) for request in self.upstream.drain()) + + def signatures(self) -> tuple[Signature, ...]: + return tuple(call.signature for call in self.calls()) + + +def _call(request: Request) -> Call: + matched: Final = PREFIXED.match(request.target) + assert matched is not None, request.target + return Call( + prefix=matched[1], + method=request.method, + route=matched[2], + authorization=request.headers.get("authorization"), + google_key=request.headers.get("x-goog-api-key"), + azure_key=request.headers.get("api-key"), + ) + + +def _json_body(request: Request) -> Mapping[str, JsonValue]: + if not request.body or not request.headers.get("content-type", "").startswith("application/json"): + return MappingProxyType({}) + return object_value(json.loads(request.body)) + + +def _json_reply(payload: Mapping[str, JsonValue], status: int = 200) -> Reply: + return Reply(status=status, body=json.dumps(payload).encode()) + + +def _error_reply(status: int) -> Reply: + return _json_reply( + {"error": {"message": f"scripted provider {status}", "type": "invalid_request_error", "code": str(status)}}, + status, + ) + + +def _vector_store_reply(request: Request, outage: threading.Event) -> Reply: + if outage.is_set(): + return _error_reply(503) + name: Final = str(_json_body(request).get("name")) + failing: Final = FAILING_STORE_NAMES.get(name) + if failing is not None: + return _error_reply(failing) + return _json_reply( + { + "id": f"vs_{uuid.uuid4().hex}", + "object": "vector_store", + "created_at": int(time.time()), + "name": name, + "usage_bytes": 0, + "status": "completed", + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + } + ) + + +def _owned_by(prefix: str, identity: str) -> bool: + return identity.startswith((f"file-{prefix}-", f"batch_{prefix}_")) + + +def _file_object(identity: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "file", + "bytes": len(DOCUMENT[1]), + "created_at": int(time.time()), + "filename": DOCUMENT[0], + "purpose": "assistants", + "status": "processed", + } + + +def _file_reply(prefix: str) -> Reply: + return _json_reply(_file_object(f"file-{prefix}-{uuid.uuid4().hex}")) + + +def _batch_input_reply() -> Reply: + return Reply(status=200, body=BATCH_INPUT[1]) + + +def _batch_object(identity: str, input_file_id: str, status: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": input_file_id, + "completion_window": "24h", + "status": status, + "created_at": int(time.time()), + "errors": None, + "output_file_id": None, + "error_file_id": None, + "request_counts": {"total": 0, "completed": 0, "failed": 0}, + "metadata": None, + } + + +def _batch_reply(request: Request, prefix: str) -> Reply: + body: Final = _json_body(request) + return _json_reply( + _batch_object(f"batch_{prefix}_{uuid.uuid4().hex}", str(body.get("input_file_id")), "validating") + ) + + +def _list_reply() -> Reply: + return _json_reply({"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}) + + +def _owned_reply(prefix: str, identity: str, payload: Mapping[str, JsonValue]) -> Reply: + return _json_reply(payload) if _owned_by(prefix, identity) else _error_reply(404) + + +def _openai_get(route: str, prefix: str) -> Reply: + if route in {"/v1/files", "/v1/batches"}: + return _list_reply() + content: Final = FILE_CONTENT.match(route) + if content is not None: + return _batch_input_reply() if _owned_by(prefix, content[1]) else _error_reply(404) + file_item: Final = FILE_ITEM.match(route) + if file_item is not None: + return _owned_reply(prefix, file_item[1], _file_object(file_item[1])) + batch_item: Final = BATCH_ITEM.match(route) + if batch_item is not None: + return _owned_reply(prefix, batch_item[1], _batch_object(batch_item[1], f"file-{prefix}-input", "validating")) + return _error_reply(404) + + +def _openai_delete(route: str, prefix: str) -> Reply: + file_item: Final = FILE_ITEM.match(route) + if file_item is None: + return _error_reply(404) + return _owned_reply(prefix, file_item[1], {"id": file_item[1], "object": "file", "deleted": True}) + + +def _embedding_reply(request: Request) -> Reply: + body: Final = _json_body(request) + inputs: Final = body.get("input") + count: Final = len(inputs) if isinstance(inputs, list) else 1 + return _json_reply( + { + "object": "list", + "data": [{"object": "embedding", "index": index, "embedding": [0.1, 0.2, 0.3]} for index in range(count)], + "model": str(body.get("model")), + "usage": {"prompt_tokens": count, "total_tokens": count}, + } + ) + + +def _bedrock_reply(request: Request, route: str) -> Reply: + listed: Final = BEDROCK_DATA_SOURCES.match(route) + if request.method == "POST" and listed is not None: + summary: Final[dict[str, JsonValue]] = { + "knowledgeBaseId": listed[1], + "dataSourceId": WEB_DATA_SOURCE, + "name": "integration-web-source", + "status": "AVAILABLE", + "updatedAt": "2026-01-01T00:00:00Z", + } + return _json_reply({"dataSourceSummaries": [] if listed[1] == EMPTY_KNOWLEDGE_BASE else [summary]}) + detail: Final = BEDROCK_DATA_SOURCE.match(route) + if request.method == "GET" and detail is not None: + return _json_reply( + { + "dataSource": { + "knowledgeBaseId": detail[1], + "dataSourceId": detail[2], + "name": "integration-web-source", + "status": "AVAILABLE", + "dataSourceConfiguration": {"type": "WEB"}, + "createdAt": "2026-01-01T00:00:00Z", + "updatedAt": "2026-01-01T00:00:00Z", + } + } + ) + return _error_reply(404) + + +def _attach_reply(request: Request, store_id: str) -> Reply: + return _json_reply( + { + "id": str(_json_body(request).get("file_id")), + "object": "vector_store.file", + "created_at": int(time.time()), + "vector_store_id": store_id, + "status": "completed", + "last_error": None, + "usage_bytes": 0, + } + ) + + +def _search_reply(request: Request) -> Reply: + return _json_reply( + { + "object": "vector_store.search_results.page", + "search_query": _json_body(request).get("query", ""), + "data": [], + "has_more": False, + "next_page": None, + } + ) + + +def _sse(payload: Mapping[str, JsonValue]) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def _chunk(identity: str, model: str, delta: Mapping[str, JsonValue], finish_reason: str | None) -> bytes: + return _sse( + { + "id": identity, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + } + ) + + +def _chat_reply(request: Request) -> Reply: + body: Final = _json_body(request) + model: Final = str(body.get("model")) + identity: Final = f"chatcmpl-{uuid.uuid4().hex}" + if body.get("stream") is True: + return Reply( + content_type="text/event-stream", + chunks=( + _chunk(identity, model, {"role": "assistant", "content": "integration response"}, None), + _chunk(identity, model, {}, "stop"), + b"data: [DONE]\n\n", + ), + ) + return _json_reply( + { + "id": identity, + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "integration response"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ) + + +def _responses_reply(request: Request) -> Reply: + identity: Final = uuid.uuid4().hex + completed: Final = _responses_object(identity, str(_json_body(request).get("model"))) + if _json_body(request).get("stream") is not True: + return _json_reply(completed) + events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = ( + ( + "response.created", + { + "type": "response.created", + "sequence_number": 0, + "response": {**completed, "status": "in_progress", "output": []}, + }, + ), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": "integration response", + "logprobs": [], + }, + ), + ("response.completed", {"type": "response.completed", "sequence_number": 2, "response": completed}), + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events), + ) + + +def _responses_object(identity: str, model: str) -> dict[str, JsonValue]: + return { + "id": f"resp_{identity}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": model, + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "integration response", "annotations": []}], + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + + +def _gemini_reply(request: Request, route: str) -> Reply: + if request.method == "POST" and route == "/v1beta/fileSearchStores": + return _json_reply({"name": f"fileSearchStores/rag-{uuid.uuid4().hex}"}) + upload: Final = GEMINI_UPLOAD.match(route) + if request.method == "POST" and upload is not None: + session_url: Final = f"http://{request.headers['host']}/gemini/upload-session/{upload[1]}" + return Reply(body=b"{}", headers={"x-goog-upload-url": session_url}) + if request.method == "PUT" and route.startswith("/upload-session/"): + return _json_reply({"name": f"documents/{uuid.uuid4().hex}"}) + return _error_reply(404) + + +def _openai_reply(request: Request, prefix: str, route: str, outage: threading.Event) -> Reply: + if request.method == "GET": + return _openai_get(route, prefix) + if request.method == "DELETE": + return _openai_delete(route, prefix) + if request.method != "POST": + return _error_reply(404) + if route == "/v1/vector_stores": + return _vector_store_reply(request, outage) + if route in {"/v1/files", "/openai/files"}: + return _file_reply(prefix) + if route == "/v1/batches": + return _batch_reply(request, prefix) + cancelled: Final = BATCH_CANCEL.match(route) + if cancelled is not None: + return _owned_reply(prefix, cancelled[1], _batch_object(cancelled[1], f"file-{prefix}-input", "cancelling")) + if route == "/v1/chat/completions": + return _chat_reply(request) + if route == "/v1/responses": + return _responses_reply(request) + if route == "/v1/embeddings": + return _embedding_reply(request) + attach: Final = STORE_FILES.match(route) + if attach is not None: + return _attach_reply(request, attach[1]) + if STORE_SEARCH.match(route) is not None: + return _search_reply(request) + return _error_reply(404) + + +def _responder(outage: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + matched: Final = PREFIXED.match(request.target) + if matched is None: + return _error_reply(404) + if matched[1] == "gemini": + return _gemini_reply(request, matched[2]) + if matched[1] == "bedrock": + return _bedrock_reply(request, matched[2]) + return _openai_reply(request, matched[1], matched[2], outage) + + return respond + + +def _proxy_config(directory: Path) -> Path: + base: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + general_settings: Final = object_value(base["general_settings"]) + config: Final = { + **base, + "general_settings": { + **general_settings, + "enable_jwt_auth": True, + "litellm_jwtauth": {"user_id_jwt_field": "sub", "team_ids_jwt_field": "groups"}, + "forward_llm_provider_auth_headers": True, + }, + "files_settings": [{"custom_llm_provider": "openai"}, {"custom_llm_provider": "azure"}], + } + path: Final = directory / "rag_team_provider_credentials.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _provider_environment(upstream: str) -> Mapping[str, str]: + return MappingProxyType( + { + "OPENAI_API_KEY": ENVIRONMENT_API_KEY, + "OPENAI_API_BASE": f"{upstream}/environment/v1", + "OPENAI_BASE_URL": f"{upstream}/environment/v1", + "GEMINI_API_KEY": GEMINI_API_KEY, + "GEMINI_API_BASE": f"{upstream}/gemini", + "AZURE_API_KEY": AZURE_API_KEY, + "AZURE_API_BASE": f"{upstream}/azure", + "AZURE_API_VERSION": "2025-04-01-preview", + } + ) + + +def _aws_environment(upstream: str) -> Mapping[str, str]: + return MappingProxyType( + { + "AWS_ENDPOINT_URL_BEDROCK_AGENT": f"{upstream}/bedrock", + "AWS_ACCESS_KEY_ID": "integration-rag-aws-access-key", + "AWS_SECRET_ACCESS_KEY": "integration-rag-aws-secret-key", + "AWS_REGION_NAME": "us-east-1", + "AWS_DEFAULT_REGION": "us-east-1", + "AWS_EC2_METADATA_DISABLED": "true", + } + ) + + +def _jwt_team_key(team_id: str) -> str: + return f"integration-rag-jwt-{team_id}" + + +def _jwt_team(gateway: Gateway, scenario: Scenario, upstream: str, model: str, member: str) -> str: + team_id: Final = scenario.team(models=[model]) + _deployment( + gateway, + scenario, + name=model, + model="openai/gpt-4o-mini", + api_key=_jwt_team_key(team_id), + api_base=f"{upstream}/team/v1", + team_id=team_id, + ) + gateway.post("/team/member_add", {"team_id": team_id, "member": {"user_id": member, "role": "user"}}) + return team_id + + +def _deployment( + gateway: Gateway, + scenario: Scenario, + *, + name: str, + model: str, + api_key: str, + api_base: str, + team_id: str | None, + model_id: str | None = None, +) -> str: + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": {"model": model, "api_key": api_key, "api_base": api_base}, + "model_info": ({} if team_id is None else {"team_id": team_id}) + | ({} if model_id is None else {"id": model_id}), + }, + ) + model_id: Final = string_value(object_value(created["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, model_id) + return model_id + + +def _delete_credential(gateway: Gateway, name: str) -> None: + response: Final = gateway.request("DELETE", f"/credentials/{name}") + assert response.status_code == 200, response.text + + +def _credential(gateway: Gateway, scenario: Scenario, api_base: str) -> str: + name: Final = f"integration-rag-credential-{uuid.uuid4().hex}" + gateway.post( + "/credentials", + { + "credential_name": name, + "credential_values": {"api_key": CREDENTIAL_API_KEY, "api_base": api_base}, + "credential_info": {}, + }, + ) + scenario.cleanups.callback(_delete_credential, gateway, name) + return name + + +def _tenants(gateway: Gateway, scenario: Scenario, upstream: str) -> Tenants: + shared_model: Final = f"integration-rag-shared-{uuid.uuid4().hex}" + _deployment( + gateway, + scenario, + name=shared_model, + model="openai/gpt-4o-mini", + api_key=SHARED_API_KEY, + api_base=f"{upstream}/shared/v1", + team_id=None, + ) + global_model: Final = f"openai/integration-rag-global-{uuid.uuid4().hex}" + _deployment( + gateway, + scenario, + name=global_model, + model="openai/gpt-4o-mini", + api_key=SHARED_API_KEY, + api_base=f"{upstream}/shared/v1", + team_id=None, + ) + team_a_model: Final = f"integration-rag-team-{uuid.uuid4().hex}" + team_a: Final = scenario.team(models=[team_a_model, shared_model]) + _deployment( + gateway, + scenario, + name=team_a_model, + model="openai/gpt-4o-mini", + api_key=TEAM_API_KEY, + api_base=f"{upstream}/team/v1", + team_id=team_a, + ) + viewer: Final = scenario.user(user_role="internal_user_viewer") + gateway.post("/team/member_add", {"team_id": team_a, "member": {"user_id": viewer, "role": "user"}}) + team_w: Final = scenario.team(models=[shared_model]) + _deployment( + gateway, + scenario, + name="openai/*", + model="openai/*", + api_key=WILDCARD_API_KEY, + api_base=f"{upstream}/wildcard/v1", + team_id=team_w, + ) + team_s: Final = scenario.team(models=[shared_model]) + team_z: Final = scenario.team(models=[shared_model]) + _deployment( + gateway, + scenario, + name="azure/*", + model="azure/*", + api_key=WILDCARD_API_KEY, + api_base=f"{upstream}/azure", + team_id=team_z, + ) + jwt_team_model: Final = f"integration-rag-jwt-{uuid.uuid4().hex}" + jwt_multi_member: Final = scenario.user(user_role="internal_user") + return Tenants( + shared_model=shared_model, + global_model=global_model, + team_a=team_a, + team_a_model=team_a_model, + team_a_key=scenario.key(team_id=team_a), + restricted_key=scenario.key(team_id=team_a, models=[shared_model]), + viewer_key=scenario.key(user_id=viewer, team_id=team_a), + jwt_member=scenario.member(team_a), + team_w=team_w, + team_w_key=scenario.key(team_id=team_w), + wildcard_bare_key=scenario.key(team_id=team_w, models=[BARE_WILDCARD_MODEL]), + wildcard_prefixed_key=scenario.key(team_id=team_w, models=[WILDCARD_MODEL]), + wildcard_global_key=scenario.key(team_id=team_w, models=[global_model]), + team_s=team_s, + team_s_key=scenario.key(team_id=team_s), + teamless_key=scenario.key(), + jwt_teams=tuple( + _jwt_team(gateway, scenario, upstream, jwt_team_model, jwt_multi_member) for _ in range(JWT_TEAM_COUNT) + ), + jwt_team_model=jwt_team_model, + jwt_multi_member=jwt_multi_member, + azure_restricted_key=scenario.key(team_id=team_z, models=["azure/gpt-4o-mini"]), + credential_name=_credential(gateway, scenario, f"{upstream}/credential/v1"), + unreachable_credential_name=_credential(gateway, scenario, _closed_url()), + ) + + +@pytest.fixture(scope="module") +def rag_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[RagRig]: + signing_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwks_body: Final = json.dumps( + { + "keys": [ + { + **json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(signing_key.public_key())), + "kid": JWT_KEY_ID, + "use": "sig", + "alg": "RS256", + } + ] + } + ).encode() + + def jwks_reply(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=jwks_body) + + outage: Final = threading.Event() + directory: Final = tmp_path_factory.mktemp("rag_team_provider_credentials") + with ( + gateway_from_environment() as admin_gateway, + admin_gateway.scenario() as scenario, + wire_server(_responder(outage)) as upstream, + wire_server(jwks_reply) as jwks, + ): + tenants: Final = _tenants(admin_gateway, scenario, upstream.url) + with ( + owned_proxy( + admin_gateway, + directory, + { + **_provider_environment(upstream.url), + **_aws_environment(upstream.url), + "JWT_PUBLIC_KEY_URL": jwks.url, + }, + config=_proxy_config(directory), + remove_environment=("GOOGLE_API_KEY", *AWS_AMBIENT_CREDENTIALS), + workers=2, + ) as gateway, + gateway.scenario() as owned_scenario, + ): + yield RagRig(gateway, upstream, outage, tenants, signing_key, owned_scenario) + + +@pytest.fixture(autouse=True) +def _drain_calls_an_earlier_failed_cell_left(request: pytest.FixtureRequest) -> None: + if "rag_rig" in request.fixturenames: + request.getfixturevalue("rag_rig").upstream.drain() + + +def _ingest_request(name: str, vector_store: Mapping[str, JsonValue]) -> Mapping[str, str]: + return MappingProxyType( + {"request": json.dumps({"ingest_options": {"name": name, "vector_store": dict(vector_store)}})} + ) + + +def _ingest( + gateway: Gateway, + key: str, + vector_store: Mapping[str, JsonValue] = OPENAI, + *, + name: str = "litellm-rag-ingest", + path: str = "/v1/rag/ingest", +) -> httpx.Response: + return gateway.request_multipart(path, _ingest_request(name, vector_store), {"file": DOCUMENT}, key=key) + + +def _query( + gateway: Gateway, + key: str, + model: str, + retrieval_config: Mapping[str, JsonValue], + *, + stream: bool = False, + path: str = "/v1/rag/query", +) -> httpx.Response: + return gateway.request( + "POST", + path, + { + "model": model, + "messages": [{"role": "user", "content": f"query team vector store {uuid.uuid4().hex}"}], + "retrieval_config": {"top_k": 1, **retrieval_config}, + **({"stream": True} if stream else {}), + }, + key=key, + ) + + +def _ingest_signatures(prefix: str, api_key: str, store_id: str) -> tuple[Signature, ...]: + bearer: Final = f"Bearer {api_key}" + return ( + (prefix, "POST", "/v1/vector_stores", bearer), + (prefix, "POST", "/v1/files", bearer), + (prefix, "POST", f"/v1/vector_stores/{store_id}/files", bearer), + ) + + +def _search_signature(prefix: str, api_key: str, store_id: str) -> Signature: + return (prefix, "POST", f"/v1/vector_stores/{store_id}/search", f"Bearer {api_key}") + + +def _chat_signature(prefix: str, api_key: str) -> Signature: + return (prefix, "POST", "/v1/chat/completions", f"Bearer {api_key}") + + +def _completed_store(rig: RagRig, response: httpx.Response) -> str: + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + assert body["status"] == "completed", response.text + store_id: Final = string_value(body["vector_store_id"]) + rig.scenario.cleanups.callback(rig.gateway.post, "/vector_store/delete", {"vector_store_id": store_id}) + return store_id + + +def _store_rows(store_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT vector_store_id, team_id, litellm_params, vector_store_metadata FROM "LiteLLM_ManagedVectorStoresTable" ' + "WHERE vector_store_id = %s", + (store_id,), + ) + + +def _closed_url() -> str: + return f"http://127.0.0.1:{allocate_port()}/v1" + + +def test_viewer_ingest_and_team_and_jwt_query_use_team_provider_credentials(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, tenants.viewer_key)) + assert rag_rig.signatures() == _ingest_signatures("team", TEAM_API_KEY, store_id) + + team_query: Final = _query( + rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, {**OPENAI, "vector_store_id": store_id} + ) + assert team_query.status_code == 200, team_query.text + assert rag_rig.signatures() == ( + _search_signature("team", TEAM_API_KEY, store_id), + _chat_signature("team", TEAM_API_KEY), + ) + + jwt_query: Final = _query( + rag_rig.gateway, + rag_rig.jwt(tenants.jwt_member, tenants.team_a), + tenants.team_a_model, + {**OPENAI, "vector_store_id": store_id}, + ) + assert jwt_query.status_code == 200, jwt_query.text + assert rag_rig.signatures() == ( + _search_signature("team", TEAM_API_KEY, store_id), + _chat_signature("team", TEAM_API_KEY), + ) + + +@pytest.mark.parametrize( + ("name", "status"), (("rag-upstream-401", 401), ("rag-upstream-429", 429)), ids=("provider_401", "provider_429") +) +def test_failed_ingest_returns_the_provider_status_without_saving_a_store( + rag_rig: RagRig, name: str, status: int +) -> None: + before: Final = _store_rows("") + + response: Final = _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, name=name) + + assert response.status_code == status, response.text + assert f"scripted provider {status}" in response.text, response.text + assert rag_rig.signatures() == (("team", "POST", "/v1/vector_stores", f"Bearer {TEAM_API_KEY}"),) + assert _store_rows("") == before + + +def test_unreachable_credential_api_base_fails_ingest_without_saving_a_store(rag_rig: RagRig) -> None: + before: Final = _store_rows("") + + response: Final = _ingest( + rag_rig.gateway, + rag_rig.tenants.team_a_key, + {**OPENAI, "litellm_credential_name": rag_rig.tenants.unreachable_credential_name}, + ) + + assert response.status_code == 500, response.text + assert rag_rig.calls() == () + assert _store_rows("") == before + + +def test_rag_alias_ingest_and_query_use_the_team_wildcard_deployment(rag_rig: RagRig) -> None: + store_id: Final = _completed_store( + rag_rig, _ingest(rag_rig.gateway, rag_rig.tenants.team_w_key, path="/rag/ingest") + ) + assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id) + rows: Final = _store_rows(store_id) + assert len(rows) == 1, rows + assert WILDCARD_API_KEY not in json.dumps(rows), rows + + response: Final = _query( + rag_rig.gateway, + rag_rig.tenants.team_w_key, + WILDCARD_MODEL, + {**OPENAI, "vector_store_id": store_id}, + path="/rag/query", + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature("wildcard", WILDCARD_API_KEY, store_id), + _chat_signature("wildcard", WILDCARD_API_KEY), + ) + + +def test_jwt_team_member_ingest_uses_team_provider_credentials(rag_rig: RagRig) -> None: + token: Final = rag_rig.jwt(rag_rig.tenants.jwt_member, rag_rig.tenants.team_a) + + store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, token)) + + assert rag_rig.signatures() == _ingest_signatures("team", TEAM_API_KEY, store_id) + rows: Final = _store_rows(store_id) + assert [row["team_id"] for row in rows] == [rag_rig.tenants.team_a], rows + + +def test_caller_credential_name_wins_over_the_team_deployment(rag_rig: RagRig) -> None: + store_id: Final = _completed_store( + rag_rig, + _ingest( + rag_rig.gateway, + rag_rig.tenants.team_a_key, + {**OPENAI, "litellm_credential_name": rag_rig.tenants.credential_name}, + ), + ) + + assert rag_rig.signatures() == _ingest_signatures("credential", CREDENTIAL_API_KEY, store_id) + + +@pytest.mark.parametrize("field", ("api_key", "api_base")) +@pytest.mark.parametrize("value", ("set", "null", "empty")) +def test_caller_api_key_or_api_base_on_ingest_is_rejected_before_any_upstream_call( + rag_rig: RagRig, field: Literal["api_key", "api_base"], value: Literal["set", "null", "empty"] +) -> None: + set_values: Final = MappingProxyType({"api_key": CALLER_API_KEY, "api_base": f"{rag_rig.upstream.url}/caller/v1"}) + sent: Final[JsonValue] = MappingProxyType({"set": set_values[field], "null": None, "empty": ""})[value] + + response: Final = _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, {**OPENAI, field: sent}) + + assert response.status_code == 400, response.text + assert f"'{field}' cannot be set in ingest_options.vector_store" in response.text, response.text + assert rag_rig.calls() == () + + +@pytest.mark.parametrize("blank", (None, ""), ids=("null", "empty")) +def test_blank_caller_credential_name_falls_back_to_the_team_deployment(rag_rig: RagRig, blank: str | None) -> None: + tenants: Final = rag_rig.tenants + + store_id: Final = _completed_store( + rag_rig, _ingest(rag_rig.gateway, tenants.team_a_key, {**OPENAI, "litellm_credential_name": blank}) + ) + assert rag_rig.signatures() == _ingest_signatures("team", TEAM_API_KEY, store_id) + + response: Final = _query( + rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, {**OPENAI, "vector_store_id": store_id} + ) + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature("team", TEAM_API_KEY, store_id), + _chat_signature("team", TEAM_API_KEY), + ) + + +@pytest.mark.parametrize("caller", ("teamless_key", "team_s_key", "restricted_key")) +def test_ingest_without_a_usable_team_deployment_uses_the_environment_credentials( + rag_rig: RagRig, caller: Literal["teamless_key", "team_s_key", "restricted_key"] +) -> None: + key: Final = getattr(rag_rig.tenants, caller) + + store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, key)) + + assert rag_rig.signatures() == _ingest_signatures("environment", ENVIRONMENT_API_KEY, store_id) + + +@pytest.mark.parametrize( + ("caller", "prefix", "api_key"), + (("team_a_key", "team", TEAM_API_KEY), ("team_s_key", "shared", SHARED_API_KEY)), + ids=("team_deployment", "shared_deployment"), +) +def test_files_upload_keeps_the_accessible_deployment_fallback( + rag_rig: RagRig, caller: Literal["team_a_key", "team_s_key"], prefix: str, api_key: str +) -> None: + response: Final = rag_rig.gateway.request_multipart( + "/v1/files", {"purpose": "assistants"}, {"file": DOCUMENT}, key=getattr(rag_rig.tenants, caller) + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ((prefix, "POST", "/v1/files", f"Bearer {api_key}"),) + + +@pytest.mark.parametrize( + ("caller", "model"), + (("wildcard_bare_key", BARE_WILDCARD_MODEL), ("wildcard_prefixed_key", WILDCARD_MODEL)), + ids=("bare_model", "prefixed_model"), +) +def test_key_restricted_to_a_model_under_the_team_wildcard_uses_the_team_deployment( + rag_rig: RagRig, caller: Literal["wildcard_bare_key", "wildcard_prefixed_key"], model: str +) -> None: + key: Final = getattr(rag_rig.tenants, caller) + + store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, key)) + assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id) + rows: Final = _store_rows(store_id) + assert [row["team_id"] for row in rows] == [rag_rig.tenants.team_w], rows + assert WILDCARD_API_KEY not in json.dumps(rows), rows + + response: Final = _query(rag_rig.gateway, key, model, {**OPENAI, "vector_store_id": store_id}) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature("wildcard", WILDCARD_API_KEY, store_id), + _chat_signature("wildcard", WILDCARD_API_KEY), + ) + + +def test_key_restricted_to_a_global_model_under_the_team_wildcard_splits_rag_and_chat(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + + store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, tenants.wildcard_global_key)) + assert rag_rig.signatures() == _ingest_signatures("wildcard", WILDCARD_API_KEY, store_id) + + response: Final = _query( + rag_rig.gateway, tenants.wildcard_global_key, tenants.global_model, {**OPENAI, "vector_store_id": store_id} + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature("wildcard", WILDCARD_API_KEY, store_id), + _chat_signature("shared", SHARED_API_KEY), + ) + + +def test_files_upload_by_a_key_restricted_under_the_team_wildcard_uses_the_team_deployment(rag_rig: RagRig) -> None: + response: Final = rag_rig.gateway.request_multipart( + "/v1/files", {"purpose": "assistants"}, {"file": DOCUMENT}, key=rag_rig.tenants.wildcard_bare_key + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == (("wildcard", "POST", "/v1/files", f"Bearer {WILDCARD_API_KEY}"),) + + +@pytest.mark.parametrize("batch_options", ({}, {"litellm_credential_name": "no-such-credential"}), ids=("plain", "credential_name")) +@pytest.mark.parametrize( + ("caller", "prefix", "api_key", "reads_input"), + (("team_a_key", "team", TEAM_API_KEY, False), ("wildcard_bare_key", "wildcard", WILDCARD_API_KEY, True)), + ids=("team_key", "restricted_wildcard_key"), +) +def test_batch_create_reads_the_input_and_creates_the_batch_through_the_team_deployment( + rag_rig: RagRig, + caller: Literal["team_a_key", "wildcard_bare_key"], + prefix: str, + api_key: str, + reads_input: bool, + batch_options: Mapping[str, JsonValue], +) -> None: + key: Final = getattr(rag_rig.tenants, caller) + uploaded: Final = rag_rig.gateway.request_multipart( + "/v1/files", {"purpose": "batch"}, {"file": BATCH_INPUT}, key=key + ) + assert uploaded.status_code == 200, uploaded.text + file_id: Final = string_value(object_value(uploaded.json())["id"]) + + created: Final = rag_rig.gateway.request( + "POST", + "/v1/batches", + {"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h", **batch_options}, + key=key, + ) + + assert created.status_code == 200, created.text + assert object_value(created.json())["input_file_id"] == file_id, created.text + bearer: Final = f"Bearer {api_key}" + input_read: Final = ((prefix, "GET", f"/v1/files/{file_id}/content", bearer),) if reads_input else () + assert rag_rig.signatures() == ( + (prefix, "POST", "/v1/files", bearer), + *input_read, + (prefix, "POST", "/v1/batches", bearer), + ) + + +def test_gemini_ingest_never_receives_the_team_openai_key(rag_rig: RagRig) -> None: + store_id: Final = _completed_store( + rag_rig, _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, {"custom_llm_provider": "gemini"}) + ) + + store_name: Final = store_id.removeprefix("fileSearchStores/") + assert tuple( + (call.prefix, call.method, call.route, call.google_key, call.authorization) for call in rag_rig.calls() + ) == ( + ("gemini", "POST", "/v1beta/fileSearchStores", GEMINI_API_KEY, None), + ( + "gemini", + "POST", + f"/upload/v1beta/fileSearchStores/{store_name}:uploadToFileSearchStore", + GEMINI_API_KEY, + None, + ), + ("gemini", "PUT", f"/upload-session/{store_name}", None, None), + ) + + +@pytest.mark.parametrize("provider", (7, ["openai"], "", "x" * 5120), ids=("int", "list", "empty", "five_kb")) +def test_malformed_ingest_provider_is_rejected_before_any_upstream_call(rag_rig: RagRig, provider: JsonValue) -> None: + response: Final = _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, {"custom_llm_provider": provider}) + + assert response.status_code == 400, response.text + assert rag_rig.calls() == () + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +def test_unmanaged_store_query_searches_with_the_team_deployment(rag_rig: RagRig, stream: bool) -> None: + tenants: Final = rag_rig.tenants + store_id: Final = f"vs_unmanaged_{uuid.uuid4().hex}" + + response: Final = _query( + rag_rig.gateway, + tenants.team_a_key, + tenants.team_a_model, + {**OPENAI, "vector_store_id": store_id}, + stream=stream, + ) + + assert response.status_code == 200, response.text + assert "integration response" in response.text, response.text + assert rag_rig.signatures() == ( + _search_signature("team", TEAM_API_KEY, store_id), + _chat_signature("team", TEAM_API_KEY), + ) + + +@pytest.mark.parametrize("case", ("api_key", "api_base", "null_api_key", "empty_api_key")) +def test_caller_credentials_in_retrieval_config_never_reach_an_unmanaged_search( + rag_rig: RagRig, case: Literal["api_key", "api_base", "null_api_key", "empty_api_key"] +) -> None: + tenants: Final = rag_rig.tenants + store_id: Final = f"vs_unmanaged_{uuid.uuid4().hex}" + caller_fields: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "api_key": {"api_key": CALLER_API_KEY}, + "api_base": {"api_base": f"{rag_rig.upstream.url}/caller/v1"}, + "null_api_key": {"api_key": None}, + "empty_api_key": {"api_key": ""}, + }[case] + ) + + response: Final = _query( + rag_rig.gateway, + tenants.team_a_key, + tenants.team_a_model, + {**OPENAI, **caller_fields, "vector_store_id": store_id}, + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature("team", TEAM_API_KEY, store_id), + _chat_signature("team", TEAM_API_KEY), + ) + + +def _managed_store(rig: RagRig, scenario: Scenario, litellm_params: Mapping[str, JsonValue]) -> str: + store_id: Final = f"vs_managed_{uuid.uuid4().hex}" + rig.gateway.post( + "/vector_store/new", + {"vector_store_id": store_id, "custom_llm_provider": "openai", "litellm_params": dict(litellm_params)}, + ) + scenario.cleanups.callback(rig.gateway.post, "/vector_store/delete", {"vector_store_id": store_id}) + return store_id + + +@pytest.mark.parametrize( + ("stored", "prefix", "api_key"), + (("own_key", "caller", CALLER_API_KEY), ("null_key", "team", TEAM_API_KEY), ("empty_key", "team", TEAM_API_KEY)), + ids=("own_key", "null_key", "empty_key"), +) +def test_managed_store_credentials_decide_the_search_key( + rag_rig: RagRig, stored: Literal["own_key", "null_key", "empty_key"], prefix: str, api_key: str +) -> None: + tenants: Final = rag_rig.tenants + litellm_params: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "own_key": {"api_key": CALLER_API_KEY, "api_base": f"{rag_rig.upstream.url}/caller/v1"}, + "null_key": {"api_key": None}, + "empty_key": {"api_key": ""}, + }[stored] + ) + with rag_rig.gateway.scenario() as scenario: + store_id: Final = _managed_store(rag_rig, scenario, litellm_params) + + response: Final = _query( + rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, {**OPENAI, "vector_store_id": store_id} + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature(prefix, api_key, store_id), + _chat_signature("team", TEAM_API_KEY), + ) + + +def test_credential_name_in_retrieval_config_is_rejected_before_search(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + + response: Final = _query( + rag_rig.gateway, + tenants.team_a_key, + tenants.team_a_model, + { + **OPENAI, + "vector_store_id": f"vs_unmanaged_{uuid.uuid4().hex}", + "litellm_credential_name": tenants.credential_name, + }, + ) + + assert response.status_code == 400, response.text + assert rag_rig.calls() == () + + +def test_azure_query_never_receives_the_team_openai_key(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + store_id: Final = f"vs_unmanaged_{uuid.uuid4().hex}" + + response: Final = _query( + rag_rig.gateway, + tenants.team_a_key, + tenants.team_a_model, + {"custom_llm_provider": "azure", "vector_store_id": store_id}, + ) + + assert response.status_code == 200, response.text + assert tuple((call.prefix, call.route, call.azure_key, call.authorization) for call in rag_rig.calls()) == ( + ("azure", f"/openai/vector_stores/{store_id}/search", AZURE_API_KEY, None), + ("team", "/v1/chat/completions", None, f"Bearer {TEAM_API_KEY}"), + ) + + +def test_another_team_cannot_query_a_team_store(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + store_id: Final = _completed_store(rag_rig, _ingest(rag_rig.gateway, tenants.team_a_key)) + assert rag_rig.signatures() == _ingest_signatures("team", TEAM_API_KEY, store_id) + + response: Final = _query( + rag_rig.gateway, tenants.team_w_key, WILDCARD_MODEL, {**OPENAI, "vector_store_id": store_id} + ) + + assert response.status_code == 403, response.text + assert rag_rig.calls() == () + + +@pytest.mark.parametrize("vector_store_id", (7, ["vs_a"], None), ids=("int", "list", "null")) +def test_non_string_vector_store_id_is_rejected_before_search(rag_rig: RagRig, vector_store_id: JsonValue) -> None: + tenants: Final = rag_rig.tenants + + response: Final = _query( + rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, {**OPENAI, "vector_store_id": vector_store_id} + ) + + assert response.status_code == 400, response.text + assert "vector_store_id must be a non-empty string" in response.text, response.text + assert rag_rig.calls() == () + + +def test_empty_vector_store_id_is_rejected_before_search(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + + response: Final = _query( + rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, {**OPENAI, "vector_store_id": ""} + ) + + assert (response.status_code, rag_rig.signatures()) == (400, ()), response.text + + +def test_five_kb_vector_store_id_searches_with_the_team_deployment(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + store_id: Final = f"vs_{'v' * 5120}" + + response: Final = _query( + rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, {**OPENAI, "vector_store_id": store_id} + ) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + _search_signature("team", TEAM_API_KEY, store_id), + _chat_signature("team", TEAM_API_KEY), + ) + + +@pytest.mark.parametrize("provider", (7, ["openai"], "", "x" * 5120), ids=("int", "list", "empty", "five_kb")) +def test_malformed_query_provider_fails_without_any_upstream_call(rag_rig: RagRig, provider: JsonValue) -> None: + tenants: Final = rag_rig.tenants + + response: Final = _query( + rag_rig.gateway, + tenants.team_a_key, + tenants.team_a_model, + {"custom_llm_provider": provider, "vector_store_id": f"vs_unmanaged_{uuid.uuid4().hex}"}, + ) + + assert (response.status_code, rag_rig.signatures()) == (500, ()), response.text + + +def test_jwt_route_outside_the_team_allowed_routes_names_the_route(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + + response: Final = rag_rig.gateway.request( + "POST", + "/v1beta/interactions", + {"model": tenants.team_a_model, "input": "hello"}, + key=rag_rig.jwt(tenants.jwt_member, tenants.team_a), + ) + + assert response.status_code == 403, response.text + assert ( + f"Team {tenants.team_a} can access model {tenants.team_a_model} but route /v1beta/interactions " + "is not in litellm_jwtauth.team_allowed_routes" + ) in response.text, response.text + assert rag_rig.calls() == () + + +def _llm_body(route: str, model: str) -> dict[str, JsonValue]: + prompt: Final = f"hello {uuid.uuid4().hex}" + if route == "/v1/responses": + return {"model": model, "input": prompt} + if route == "/v1/messages": + return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]} + return {"model": model, "messages": [{"role": "user", "content": prompt}]} + + +@pytest.mark.parametrize("route", ("/v1/chat/completions", "/v1/messages", "/v1/responses")) +def test_jwt_team_member_keeps_its_llm_routes(rag_rig: RagRig, route: str) -> None: + tenants: Final = rag_rig.tenants + + response: Final = rag_rig.gateway.request( + "POST", + route, + _llm_body(route, tenants.team_a_model), + key=rag_rig.jwt(tenants.jwt_member, tenants.team_a), + headers={"anthropic-version": "2023-06-01"}, + ) + + assert response.status_code == 200, response.text + calls: Final = rag_rig.calls() + assert [(call.prefix, call.authorization) for call in calls] == [("team", f"Bearer {TEAM_API_KEY}")], calls + + +async def _ingest_burst(rig: RagRig, names: tuple[str, ...]) -> tuple[httpx.Response, ...]: + async with httpx.AsyncClient(base_url=str(rig.gateway.client.base_url), timeout=60, trust_env=False) as client: + return tuple( + await asyncio.gather( + *( + client.post( + "/v1/rag/ingest", + data=_ingest_request(name, OPENAI), + files={"file": DOCUMENT}, + headers={"Authorization": f"Bearer {rig.tenants.team_a_key}"}, + ) + for name in names + ) + ) + ) + + +def _burst_names() -> tuple[str, ...]: + return tuple(f"rag-burst-{uuid.uuid4().hex}" for _ in range(10)) + + +def _all_ingest_signatures(store_ids: Iterable[str]) -> list[Signature]: + return sorted(chain.from_iterable(_ingest_signatures("team", TEAM_API_KEY, store_id) for store_id in store_ids)) + + +async def test_ingest_burst_during_a_provider_outage_fails_each_request_then_recovers(rag_rig: RagRig) -> None: + before: Final = _store_rows("") + rag_rig.outage.set() + try: + failed: Final = await _ingest_burst(rag_rig, _burst_names()) + finally: + rag_rig.outage.clear() + + assert [response.status_code for response in failed] == [503] * 10, [response.text for response in failed] + outage_calls: Final = rag_rig.signatures() + assert outage_calls == (("team", "POST", "/v1/vector_stores", f"Bearer {TEAM_API_KEY}"),) * 10, outage_calls + assert _store_rows("") == before + + recovered: Final = await _ingest_burst(rag_rig, _burst_names()) + + store_ids: Final = tuple(_completed_store(rag_rig, response) for response in recovered) + assert len(frozenset(store_ids)) == 10, store_ids + assert sorted(rag_rig.signatures()) == _all_ingest_signatures(store_ids) + assert [[row["team_id"] for row in _store_rows(store_id)] for store_id in store_ids] == [ + [rag_rig.tenants.team_a] + ] * 10 + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +async def _query_burst(gateway: Gateway, key: str, model: str, count: int) -> tuple[httpx.Response, ...]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *( + client.post( + "/v1/rag/query", + json={ + "model": model, + "messages": [{"role": "user", "content": f"held query {uuid.uuid4().hex}"}], + "retrieval_config": {**OPENAI, "vector_store_id": f"vs_burst_{uuid.uuid4().hex}", "top_k": 1}, + }, + headers={"Authorization": f"Bearer {key}"}, + ) + for _ in range(count) + ), + return_exceptions=True, + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, httpx.Response)) + + +async def test_worker_sigkill_mid_query_burst_leaves_the_sibling_on_team_credentials(tmp_path: Path) -> None: + release: Final = threading.Event() + held_searches: Final[SimpleQueue[str]] = SimpleQueue() + + def respond(request: Request) -> Reply: + matched: Final = PREFIXED.match(request.target) + if matched is None: + return _error_reply(404) + if STORE_SEARCH.match(matched[2]) is not None: + held_searches.put(request.target) + assert release.wait(timeout=60), "The burst was never released" + return _search_reply(request) + return _openai_reply(request, matched[1], matched[2], threading.Event()) + + with ( + gateway_from_environment() as admin_gateway, + admin_gateway.scenario() as scenario, + wire_server(respond) as upstream, + ): + model: Final = f"integration-rag-kill-{uuid.uuid4().hex}" + team_id: Final = scenario.team(models=[model]) + _deployment( + admin_gateway, + scenario, + name=model, + model="openai/gpt-4o-mini", + api_key=TEAM_API_KEY, + api_base=f"{upstream.url}/team/v1", + team_id=team_id, + ) + key: Final = scenario.key(team_id=team_id) + with owned_proxy_process( + admin_gateway, + tmp_path, + _provider_environment(upstream.url), + remove_environment=("GOOGLE_API_KEY",), + workers=2, + ) as owned: + workers: Final = eventually( + lambda: tuple(int(pid) for pid in STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task(_query_burst(owned.gateway, key, model, 20)) + await asyncio.to_thread(eventually, held_searches.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, upstream.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert [response.status_code for response in served] == [200] * held_by[survivor_pid], held_by + (follow_up,) = await _query_burst(owned.gateway, key, model, 1) + assert follow_up.status_code == 200, follow_up.text + calls: Final = tuple(_call(request) for request in upstream.drain()) + assert {(call.prefix, call.authorization) for call in calls} == {("team", f"Bearer {TEAM_API_KEY}")}, calls + assert sum(1 for call in calls if call.route == "/v1/chat/completions") == held_by[survivor_pid] + 1, calls + + +def _ingest_and_drain( + gateway: Gateway, scenario: Scenario, key: str, upstream: Wire +) -> tuple[str, tuple[Signature, ...]]: + upstream.drain() + response: Final = _ingest(gateway, key) + assert response.status_code == 200, response.text + store_id: Final = string_value(object_value(response.json())["vector_store_id"]) + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": store_id}) + calls: Final = tuple(_call(request) for request in upstream.drain()) + return store_id, tuple(call.signature for call in calls if call.route != "/v1/models") + + +def _ingests_after_editing( + gateway: Gateway, scenario: Scenario, key: str, upstream: Wire, model_id: str +) -> tuple[tuple[str, tuple[Signature, ...]], ...]: + edited: Final = gateway.request("PATCH", f"/model/{model_id}/update", {"litellm_params": {"timeout": 30}}) + assert edited.status_code == 200, edited.text + return tuple(_ingest_and_drain(gateway, scenario, key, upstream) for _ in range(4)) + + +def test_a_team_with_two_openai_deployments_keeps_ingesting_through_the_older_one_after_either_is_edited( + tmp_path: Path, +) -> None: + with ( + gateway_from_environment() as admin_gateway, + admin_gateway.scenario() as scenario, + wire_server(_responder(threading.Event())) as upstream, + ): + team_model: Final = f"integration-rag-team-{uuid.uuid4().hex}" + wildcard_model: Final = f"integration-rag-wildcard-{uuid.uuid4().hex}" + team_id: Final = scenario.team(models=[team_model, wildcard_model]) + team_deployment: Final = ( + _deployment( + admin_gateway, + scenario, + name=team_model, + model="openai/gpt-4o-mini", + api_key=TEAM_API_KEY, + api_base=f"{upstream.url}/team/v1", + team_id=team_id, + model_id=f"z-{uuid.uuid4().hex}", + ), + "team", + TEAM_API_KEY, + ) + wildcard_deployment: Final = ( + _deployment( + admin_gateway, + scenario, + name=wildcard_model, + model="openai/gpt-4o-mini", + api_key=WILDCARD_API_KEY, + api_base=f"{upstream.url}/wildcard/v1", + team_id=team_id, + model_id=f"a-{uuid.uuid4().hex}", + ), + "wildcard", + WILDCARD_API_KEY, + ) + (picked_id, picked_prefix, picked_key), (other_id, _, _) = team_deployment, wildcard_deployment + key: Final = scenario.key(team_id=team_id) + with ( + owned_proxy( + admin_gateway, + tmp_path, + _provider_environment(upstream.url), + remove_environment=("GOOGLE_API_KEY",), + workers=2, + ) as gateway, + gateway.scenario() as owned_scenario, + ): + ingests: Final = ( + _ingest_and_drain(gateway, owned_scenario, key, upstream), + *_ingests_after_editing(gateway, owned_scenario, key, upstream, picked_id), + *_ingests_after_editing(gateway, owned_scenario, key, upstream, other_id), + ) + + assert [calls for _, calls in ingests] == [ + _ingest_signatures(picked_prefix, picked_key, store_id) for store_id, _ in ingests + ] + + +@pytest.mark.parametrize( + ("stored", "prefix", "api_key"), + (("own_key", "caller", CALLER_API_KEY), ("null_key", "team", TEAM_API_KEY)), + ids=("own_key", "null_key"), +) +def test_ingest_into_a_managed_store_uses_its_own_key_else_the_team_key( + rag_rig: RagRig, stored: Literal["own_key", "null_key"], prefix: str, api_key: str +) -> None: + litellm_params: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "own_key": {"api_key": CALLER_API_KEY, "api_base": f"{rag_rig.upstream.url}/caller/v1"}, + "null_key": {"api_key": None}, + }[stored] + ) + with rag_rig.gateway.scenario() as scenario: + store_id: Final = _managed_store(rag_rig, scenario, litellm_params) + + response: Final = _ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, {**OPENAI, "vector_store_id": store_id}) + + assert response.status_code == 200, response.text + assert object_value(response.json())["vector_store_id"] == store_id, response.text + bearer: Final = f"Bearer {api_key}" + assert rag_rig.signatures() == ( + (prefix, "POST", "/v1/files", bearer), + (prefix, "POST", f"/v1/vector_stores/{store_id}/files", bearer), + ) + + +def _file_id_ingest(gateway: Gateway, key: str, vector_store: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request( + "POST", + "/v1/rag/ingest", + { + "file_id": "file-team-existing", + "ingest_options": {"name": "litellm-rag-ingest", "vector_store": dict(vector_store)}, + }, + key=key, + ) + + +@pytest.mark.parametrize( + ("vector_store", "ingestion"), + ( + ({"custom_llm_provider": "gemini"}, "GeminiRAGIngestion"), + ({"custom_llm_provider": "bedrock", "vector_store_id": BUCKETLESS_KNOWLEDGE_BASE}, "BedrockRAGIngestion"), + ({"custom_llm_provider": "s3_vectors", "vector_bucket_name": "integration-bucket"}, "S3VectorsRAGIngestion"), + ({"custom_llm_provider": "vertex_ai"}, "VertexAIRAGIngestion"), + ), + ids=("gemini", "bedrock", "s3_vectors", "vertex_ai"), +) +def test_existing_file_id_on_a_store_that_cannot_take_one_is_rejected_before_any_upstream_call( + rag_rig: RagRig, vector_store: Mapping[str, JsonValue], ingestion: str +) -> None: + before: Final = _store_rows("") + + response: Final = _file_id_ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, vector_store) + + assert response.status_code == 400, response.text + assert f"{ingestion} does not support ingesting an existing file_id" in response.text, response.text + assert rag_rig.calls() == () + assert _store_rows("") == before + + +def test_openai_file_id_without_a_vector_store_id_is_rejected_before_any_upstream_call(rag_rig: RagRig) -> None: + before: Final = _store_rows("") + + response: Final = _file_id_ingest(rag_rig.gateway, rag_rig.tenants.team_a_key, OPENAI) + + assert response.status_code == 400, response.text + assert "vector_store_id is required when ingesting an existing file_id" in response.text, response.text + assert rag_rig.calls() == () + assert _store_rows("") == before + + +def test_openai_file_id_attaches_to_an_existing_store_with_the_team_key(rag_rig: RagRig) -> None: + store_id: Final = f"vs_existing_{uuid.uuid4().hex}" + + response: Final = _file_id_ingest( + rag_rig.gateway, rag_rig.tenants.team_a_key, {**OPENAI, "vector_store_id": store_id} + ) + + assert _completed_store(rag_rig, response) == store_id + assert rag_rig.signatures() == (("team", "POST", f"/v1/vector_stores/{store_id}/files", f"Bearer {TEAM_API_KEY}"),) + + +@pytest.mark.parametrize( + ("knowledge_base", "error", "routes"), + ( + ( + EMPTY_KNOWLEDGE_BASE, + f"No data sources found for Knowledge Base {EMPTY_KNOWLEDGE_BASE}", + (("POST", f"/knowledgebases/{EMPTY_KNOWLEDGE_BASE}/datasources/"),), + ), + ( + BUCKETLESS_KNOWLEDGE_BASE, + f"Could not auto-detect S3 bucket for data source {WEB_DATA_SOURCE}", + ( + ("POST", f"/knowledgebases/{BUCKETLESS_KNOWLEDGE_BASE}/datasources/"), + ("GET", f"/knowledgebases/{BUCKETLESS_KNOWLEDGE_BASE}/datasources/{WEB_DATA_SOURCE}"), + ), + ), + ), + ids=("no_data_source", "no_s3_bucket"), +) +def test_bedrock_knowledge_base_the_caller_named_without_a_usable_data_source_fails_ingest_with_400( + rag_rig: RagRig, knowledge_base: str, error: str, routes: tuple[tuple[str, str], ...] +) -> None: + before: Final = _store_rows("") + + response: Final = _ingest( + rag_rig.gateway, + rag_rig.tenants.team_a_key, + {"custom_llm_provider": "bedrock", "vector_store_id": knowledge_base}, + ) + + assert response.status_code == 400, response.text + assert error in response.text, response.text + assert tuple((call.prefix, call.method, call.route) for call in rag_rig.calls()) == tuple( + ("bedrock", method, route) for method, route in routes + ) + assert _store_rows("") == before + + +def test_s3_vectors_bucket_name_the_caller_sent_is_rejected_with_400_after_embedding(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + before: Final = _store_rows("") + request: Final = { + "request": json.dumps( + { + "ingest_options": { + "name": "litellm-rag-ingest", + "embedding": {"model": tenants.shared_model}, + "vector_store": {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "ab", "dimension": 3}, + } + } + ) + } + + response: Final = rag_rig.gateway.request_multipart( + "/v1/rag/ingest", request, {"file": DOCUMENT}, key=tenants.team_a_key + ) + + assert response.status_code == 400, response.text + assert "Invalid vector_bucket_name 'ab'" in response.text, response.text + assert rag_rig.signatures() == (("shared", "POST", "/v1/embeddings", f"Bearer {SHARED_API_KEY}"),) + assert _store_rows("") == before + + +@dataclass(frozen=True, slots=True) +class KeylessRig: + gateway: Gateway + upstream: Wire + key: str + + +@pytest.fixture(scope="module") +def keyless_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[KeylessRig]: + directory: Final = tmp_path_factory.mktemp("rag_keyless_providers") + with ( + gateway_from_environment() as admin_gateway, + admin_gateway.scenario() as scenario, + wire_server(_responder(threading.Event())) as upstream, + ): + key: Final = scenario.key() + with owned_proxy( + admin_gateway, + directory, + {"OPENAI_API_KEY": ENVIRONMENT_API_KEY, "OPENAI_API_BASE": f"{upstream.url}/environment/v1"}, + remove_environment=("GEMINI_API_KEY", "GOOGLE_API_KEY", "VERTEXAI_PROJECT"), + workers=2, + ) as gateway: + yield KeylessRig(gateway, upstream, key) + + +@pytest.mark.parametrize( + ("provider", "status", "error"), + ( + ("gemini", 500, "GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search"), + ("vertex_ai", 400, "vertex_project is required for Vertex AI RAG ingestion"), + ), + ids=("gemini_without_a_proxy_key", "vertex_ai_without_a_project"), +) +def test_ingest_status_says_whether_the_caller_or_the_proxy_must_fix_the_missing_setting( + keyless_rig: KeylessRig, provider: str, status: int, error: str +) -> None: + before: Final = _store_rows("") + + response: Final = _ingest(keyless_rig.gateway, keyless_rig.key, {"custom_llm_provider": provider}) + + assert response.status_code == status, response.text + assert error in response.text, response.text + assert tuple(keyless_rig.upstream.drain()) == () + assert _store_rows("") == before + + +async def _jwt_chat_burst(rig: RagRig, token: str, model: str, count: int) -> tuple[httpx.Response, ...]: + async with httpx.AsyncClient(base_url=str(rig.gateway.client.base_url), timeout=60, trust_env=False) as client: + return tuple( + await asyncio.gather( + *( + client.post( + "/v1/chat/completions", + json=_llm_body("/v1/chat/completions", model), + headers={"Authorization": f"Bearer {token}"}, + ) + for _ in range(count) + ) + ) + ) + + +async def test_jwt_in_ten_teams_that_hold_the_model_always_resolves_the_lowest_team_id(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + token: Final = rag_rig.jwt(tenants.jwt_multi_member, *reversed(tenants.jwt_teams)) + + responses: Final = await _jwt_chat_burst(rag_rig, token, tenants.jwt_team_model, 16) + + assert [response.status_code for response in responses] == [200] * 16, [response.text for response in responses] + assert ( + rag_rig.signatures() + == (("team", "POST", "/v1/chat/completions", f"Bearer {_jwt_team_key(min(tenants.jwt_teams))}"),) * 16 + ) + + +def _prompt() -> str: + return f"hello {uuid.uuid4().hex}" + + +def _sync_sdk_text(route: str, base_url: str, token: str, model: str, stream: bool) -> str: + if route == "/v1/messages": + messages_client: Final = anthropic.Anthropic(base_url=base_url, api_key=token, max_retries=0) + request: Final = {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": _prompt()}]} + if stream: + return "".join( + event.delta.text + for event in messages_client.messages.create(**request, stream=True) + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ) + message: Final = messages_client.messages.create(**request) + return "".join(block.text for block in message.content if block.type == "text") + client: Final = openai.OpenAI(base_url=f"{base_url}/v1", api_key=token, max_retries=0) + if route == "/v1/responses": + if stream: + return "".join( + event.delta + for event in client.responses.create(model=model, input=_prompt(), stream=True) + if event.type == "response.output_text.delta" + ) + return client.responses.create(model=model, input=_prompt()).output_text + chat_messages: Final = [{"role": "user", "content": _prompt()}] + if stream: + return "".join( + chunk.choices[0].delta.content or "" + for chunk in client.chat.completions.create(model=model, messages=chat_messages, stream=True) + if chunk.choices + ) + return client.chat.completions.create(model=model, messages=chat_messages).choices[0].message.content or "" + + +async def _async_sdk_text(route: str, base_url: str, token: str, model: str, stream: bool) -> str: + if route == "/v1/messages": + messages_client: Final = anthropic.AsyncAnthropic(base_url=base_url, api_key=token, max_retries=0) + request: Final = {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": _prompt()}]} + if stream: + return "".join( + [ + event.delta.text + async for event in await messages_client.messages.create(**request, stream=True) + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ] + ) + message: Final = await messages_client.messages.create(**request) + return "".join(block.text for block in message.content if block.type == "text") + client: Final = openai.AsyncOpenAI(base_url=f"{base_url}/v1", api_key=token, max_retries=0) + if route == "/v1/responses": + if stream: + return "".join( + [ + event.delta + async for event in await client.responses.create(model=model, input=_prompt(), stream=True) + if event.type == "response.output_text.delta" + ] + ) + return (await client.responses.create(model=model, input=_prompt())).output_text + chat_messages: Final = [{"role": "user", "content": _prompt()}] + if stream: + return "".join( + [ + chunk.choices[0].delta.content or "" + async for chunk in await client.chat.completions.create( + model=model, messages=chat_messages, stream=True + ) + if chunk.choices + ] + ) + completion: Final = await client.chat.completions.create(model=model, messages=chat_messages) + return completion.choices[0].message.content or "" + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +@pytest.mark.parametrize("client", ("sync", "async")) +@pytest.mark.parametrize("route", ("/v1/chat/completions", "/v1/messages", "/v1/responses")) +async def test_jwt_team_member_reaches_its_llm_routes_through_the_sdks( + rag_rig: RagRig, route: str, client: Literal["sync", "async"], stream: bool +) -> None: + tenants: Final = rag_rig.tenants + base_url: Final = str(rag_rig.gateway.client.base_url).rstrip("/") + token: Final = rag_rig.jwt(tenants.jwt_member, tenants.team_a) + + text: Final = ( + await asyncio.to_thread(_sync_sdk_text, route, base_url, token, tenants.team_a_model, stream) + if client == "sync" + else await _async_sdk_text(route, base_url, token, tenants.team_a_model, stream) + ) + + assert text == "integration response", text + calls: Final = rag_rig.calls() + assert [(call.prefix, call.authorization) for call in calls] == [("team", f"Bearer {TEAM_API_KEY}")], calls + + +def _batch_for(rig: RagRig, headers: Mapping[str, str]) -> tuple[str, str]: + uploaded: Final = rig.gateway.client.post( + "/v1/files", data={"purpose": "batch"}, files={"file": BATCH_INPUT}, headers=dict(headers) + ) + assert uploaded.status_code == 200, uploaded.text + file_id: Final = string_value(object_value(uploaded.json())["id"]) + created: Final = rig.gateway.request( + "POST", + "/v1/batches", + {"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + headers=headers, + ) + assert created.status_code == 200, created.text + rig.upstream.drain() + return file_id, string_value(object_value(created.json())["id"]) + + +FOLLOW_UP_CALLS: Final = ( + ("GET", "/v1/files/{file_id}", "/v1/files/{file_id}"), + ("GET", "/v1/files/{file_id}/content", "/v1/files/{file_id}/content"), + ("DELETE", "/v1/files/{file_id}", "/v1/files/{file_id}"), + ("GET", "/v1/files?purpose=batch&custom_llm_provider=openai", "/v1/files"), + ("GET", "/v1/batches/{batch_id}", "/v1/batches/{batch_id}"), + ("POST", "/v1/batches/{batch_id}/cancel", "/v1/batches/{batch_id}/cancel"), +) +FOLLOW_UP_CALL_IDS: Final = ( + "file_retrieve", + "file_content", + "file_delete", + "file_list", + "batch_retrieve", + "batch_cancel", +) + + +def _bearer(key: str) -> Mapping[str, str]: + return MappingProxyType({"Authorization": f"Bearer {key}"}) + + +@pytest.mark.parametrize(("method", "path", "route"), FOLLOW_UP_CALLS, ids=FOLLOW_UP_CALL_IDS) +def test_files_and_batches_calls_by_a_key_restricted_under_the_team_wildcard_use_the_team_deployment( + rag_rig: RagRig, method: str, path: str, route: str +) -> None: + key: Final = rag_rig.tenants.wildcard_bare_key + file_id, batch_id = _batch_for(rag_rig, _bearer(key)) + + response: Final = rag_rig.gateway.request(method, path.format(file_id=file_id, batch_id=batch_id), key=key) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + ("wildcard", method, route.format(file_id=file_id, batch_id=batch_id), f"Bearer {WILDCARD_API_KEY}"), + ) + + +def test_batch_list_by_a_key_restricted_under_the_team_wildcard_is_answered_by_the_proxy_without_an_upstream_call( + rag_rig: RagRig, +) -> None: + key: Final = rag_rig.tenants.wildcard_bare_key + _batch_for(rag_rig, _bearer(key)) + + response: Final = rag_rig.gateway.request("GET", "/v1/batches?custom_llm_provider=openai", key=key) + + assert response.status_code == 200, response.text + assert object_value(response.json())["object"] == "list", response.text + assert rag_rig.signatures() == () + + +def _upload_batch_input(rig: RagRig, key: str) -> str: + uploaded: Final = rig.gateway.request_multipart("/v1/files", {"purpose": "batch"}, {"file": BATCH_INPUT}, key=key) + assert uploaded.status_code == 200, uploaded.text + return string_value(object_value(uploaded.json())["id"]) + + +def _create_batch(rig: RagRig, key: str, file_id: str) -> httpx.Response: + return rig.gateway.request( + "POST", + "/v1/batches", + {"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + key=key, + ) + + +@pytest.mark.parametrize( + ("team", "prefix", "api_key"), + (("team_a", "team", TEAM_API_KEY), ("team_s", "shared", SHARED_API_KEY)), + ids=("team_deployment", "accessible_deployment"), +) +def test_batch_over_the_key_tpm_limit_is_rejected_after_reading_its_input_where_it_was_uploaded( + rag_rig: RagRig, team: Literal["team_a", "team_s"], prefix: str, api_key: str +) -> None: + key: Final = rag_rig.scenario.key(team_id=getattr(rag_rig.tenants, team), tpm_limit=1) + file_id: Final = _upload_batch_input(rag_rig, key) + + response: Final = _create_batch(rag_rig, key, file_id) + + assert response.status_code == 429, response.text + bearer: Final = f"Bearer {api_key}" + assert rag_rig.signatures() == ( + (prefix, "POST", "/v1/files", bearer), + (prefix, "GET", f"/v1/files/{file_id}/content", bearer), + ) + + +def test_batch_whose_input_names_a_model_outside_the_key_allowlist_is_rejected_before_creation( + rag_rig: RagRig, +) -> None: + key: Final = rag_rig.tenants.restricted_key + file_id: Final = _upload_batch_input(rag_rig, key) + + response: Final = _create_batch(rag_rig, key, file_id) + + assert response.status_code == 403, response.text + assert BARE_WILDCARD_MODEL in response.text, response.text + bearer: Final = f"Bearer {SHARED_API_KEY}" + assert rag_rig.signatures() == ( + ("shared", "POST", "/v1/files", bearer), + ("shared", "GET", f"/v1/files/{file_id}/content", bearer), + ) + + +def test_azure_files_upload_by_a_key_restricted_under_the_team_azure_wildcard_uses_the_team_deployment( + rag_rig: RagRig, +) -> None: + response: Final = rag_rig.gateway.request_multipart( + "/v1/files", + {"purpose": "assistants", "custom_llm_provider": "azure"}, + {"file": DOCUMENT}, + key=rag_rig.tenants.azure_restricted_key, + ) + + assert response.status_code == 200, response.text + assert tuple((call.prefix, call.method, call.route, call.azure_key) for call in rag_rig.calls()) == ( + ("azure", "POST", "/openai/files", WILDCARD_API_KEY), + ) + + +def test_query_without_a_vector_store_id_is_rejected_before_search(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + + response: Final = _query(rag_rig.gateway, tenants.team_a_key, tenants.team_a_model, OPENAI) + + assert response.status_code == 400, response.text + assert "retrieval_config must contain 'vector_store_id'" in response.text, response.text + assert rag_rig.calls() == () + + +def _forwarding(key: str) -> Mapping[str, str]: + return MappingProxyType({"Authorization": f"Bearer {key}", "x-api-key": FORWARDED_API_KEY}) + + +@pytest.mark.parametrize( + ("stored", "stream", "search"), + ( + (None, False, ("team", FORWARDED_API_KEY)), + (None, True, ("team", FORWARDED_API_KEY)), + ("own_key", False, ("caller", CALLER_API_KEY)), + ("null_key", False, ("team", FORWARDED_API_KEY)), + ), + ids=("unmanaged_json", "unmanaged_stream", "managed_own_key", "managed_null_key"), +) +def test_forwarded_provider_key_on_query_wins_over_team_provider_credentials( + rag_rig: RagRig, stored: Literal["own_key", "null_key"] | None, stream: bool, search: tuple[str, str] +) -> None: + tenants: Final = rag_rig.tenants + with rag_rig.gateway.scenario() as scenario: + store_id: Final = ( + f"vs_unmanaged_{uuid.uuid4().hex}" + if stored is None + else _managed_store( + rag_rig, + scenario, + {"api_key": CALLER_API_KEY, "api_base": f"{rag_rig.upstream.url}/caller/v1"} + if stored == "own_key" + else {"api_key": None}, + ) + ) + + response: Final = rag_rig.gateway.request( + "POST", + "/v1/rag/query", + { + "model": tenants.team_a_model, + "messages": [{"role": "user", "content": f"query forwarded key {uuid.uuid4().hex}"}], + "retrieval_config": {**OPENAI, "top_k": 1, "vector_store_id": store_id}, + **({"stream": True} if stream else {}), + }, + headers=_forwarding(tenants.team_a_key), + ) + + assert response.status_code == 200, response.text + assert "integration response" in response.text, response.text + assert rag_rig.signatures() == ( + _search_signature(search[0], search[1], store_id), + _chat_signature("team", FORWARDED_API_KEY), + ) + + +def test_forwarded_provider_key_on_ingest_wins_over_team_provider_credentials(rag_rig: RagRig) -> None: + tenants: Final = rag_rig.tenants + + response: Final = rag_rig.gateway.client.post( + "/v1/rag/ingest", + data=dict(_ingest_request("litellm-rag-forwarded-key", OPENAI)), + files={"file": DOCUMENT}, + headers=dict(_forwarding(tenants.team_a_key)), + ) + + store_id: Final = _completed_store(rag_rig, response) + assert rag_rig.signatures() == _ingest_signatures("team", FORWARDED_API_KEY, store_id) + + +@pytest.mark.parametrize( + ("caller", "prefix"), + (("team_a_key", "team"), ("wildcard_bare_key", "wildcard")), + ids=("team_key", "restricted_wildcard_key"), +) +def test_forwarded_provider_key_on_files_and_batches_wins_over_team_provider_credentials( + rag_rig: RagRig, caller: Literal["team_a_key", "wildcard_bare_key"], prefix: str +) -> None: + headers: Final = _forwarding(getattr(rag_rig.tenants, caller)) + uploaded: Final = rag_rig.gateway.client.post( + "/v1/files", data={"purpose": "batch"}, files={"file": BATCH_INPUT}, headers=dict(headers) + ) + assert uploaded.status_code == 200, uploaded.text + file_id: Final = string_value(object_value(uploaded.json())["id"]) + + created: Final = rag_rig.gateway.request( + "POST", + "/v1/batches", + {"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + headers=headers, + ) + + assert created.status_code == 200, created.text + calls: Final = rag_rig.calls() + assert [(call.prefix, call.authorization) for call in calls] == [ + (prefix, f"Bearer {FORWARDED_API_KEY}") + ] * len(calls), calls + assert {(call.method, call.route) for call in calls} >= {("POST", "/v1/files"), ("POST", "/v1/batches")}, calls + + +@pytest.mark.parametrize("stray_claim", (7, None), ids=("int", "null")) +@pytest.mark.parametrize("route", ("/v1/chat/completions", "/v1/rag/query")) +def test_jwt_team_claim_with_a_non_string_entry_still_resolves_the_string_team( + rag_rig: RagRig, route: str, stray_claim: JsonValue +) -> None: + tenants: Final = rag_rig.tenants + token: Final = rag_rig.jwt(tenants.jwt_member, tenants.team_a, stray_claim) + store_id: Final = f"vs_unmanaged_{uuid.uuid4().hex}" + + response: Final = ( + _query(rag_rig.gateway, token, tenants.team_a_model, {**OPENAI, "vector_store_id": store_id}) + if route == "/v1/rag/query" + else rag_rig.gateway.request("POST", route, _llm_body(route, tenants.team_a_model), key=token) + ) + + assert response.status_code == 200, response.text + search: Final = (_search_signature("team", TEAM_API_KEY, store_id),) if route == "/v1/rag/query" else () + assert rag_rig.signatures() == (*search, _chat_signature("team", TEAM_API_KEY)) + + +@pytest.mark.parametrize("stray_claim", (7, None), ids=("int", "null")) +def test_jwt_header_team_outside_a_claim_with_a_non_string_entry_is_denied_with_403( + rag_rig: RagRig, stray_claim: JsonValue +) -> None: + tenants: Final = rag_rig.tenants + header_team: Final = f"integration-rag-unknown-team-{uuid.uuid4().hex}" + + response: Final = rag_rig.gateway.request( + "POST", + "/v1/chat/completions", + _llm_body("/v1/chat/completions", tenants.team_a_model), + key=rag_rig.jwt(tenants.jwt_member, tenants.team_a, stray_claim), + headers={"x-litellm-team-id": header_team}, + ) + + assert response.status_code == 403, response.text + assert f"x-litellm-team-id '{header_team}' does not resolve to a team id or a unique team alias" in response.text, ( + response.text + ) + assert f"'{tenants.team_a}'" in response.text, response.text + assert rag_rig.calls() == () + + +@pytest.mark.parametrize(("method", "path", "route"), FOLLOW_UP_CALLS, ids=FOLLOW_UP_CALL_IDS) +@pytest.mark.parametrize( + ("caller", "prefix"), + (("team_a_key", "team"), ("wildcard_bare_key", "wildcard")), + ids=("team_key", "restricted_wildcard_key"), +) +def test_forwarded_provider_key_on_files_and_batches_follow_up_calls_replaces_the_team_deployment_key( + rag_rig: RagRig, + caller: Literal["team_a_key", "wildcard_bare_key"], + prefix: str, + method: str, + path: str, + route: str, +) -> None: + headers: Final = _forwarding(getattr(rag_rig.tenants, caller)) + file_id, batch_id = _batch_for(rag_rig, headers) + + response: Final = rag_rig.gateway.request(method, path.format(file_id=file_id, batch_id=batch_id), headers=headers) + + assert response.status_code == 200, response.text + assert rag_rig.signatures() == ( + (prefix, method, route.format(file_id=file_id, batch_id=batch_id), f"Bearer {FORWARDED_API_KEY}"), + ) + + +@pytest.mark.parametrize( + ("x_api_key", "prefix", "api_key"), + ( + (("",), "wildcard", WILDCARD_API_KEY), + (("f" * 5120,), "wildcard", "f" * 5120), + ((FORWARDED_API_KEY, FORWARDED_API_KEY), "wildcard", FORWARDED_API_KEY), + ), + ids=("empty", "five_kb", "repeated"), +) +def test_blank_oversized_or_repeated_forwarded_provider_key_on_files_and_batches( + rag_rig: RagRig, x_api_key: tuple[str, ...], prefix: str, api_key: str +) -> None: + headers: Final = ( + ("Authorization", f"Bearer {rag_rig.tenants.wildcard_bare_key}"), + *(("x-api-key", value) for value in x_api_key), + ) + uploaded: Final = rag_rig.gateway.client.post( + "/v1/files", data={"purpose": "batch"}, files={"file": BATCH_INPUT}, headers=headers + ) + assert uploaded.status_code == 200, uploaded.text + file_id: Final = string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]) + + created: Final = rag_rig.gateway.client.post( + "/v1/batches", + json={"input_file_id": file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + headers=headers, + ) + + assert created.status_code == 200, created.text + calls: Final = rag_rig.calls() + assert [(call.prefix, call.authorization) for call in calls] == [(prefix, f"Bearer {api_key}")] * len(calls), calls + assert {(call.method, call.route) for call in calls} == { + ("POST", "/v1/files"), + ("GET", f"/v1/files/{file_id}/content"), + ("POST", "/v1/batches"), + }, calls + + +def _forwarded_ingest(rig: RagRig, vector_store: Mapping[str, JsonValue]) -> httpx.Response: + return rig.gateway.client.post( + "/v1/rag/ingest", + data=dict(_ingest_request("litellm-rag-forwarded-key", vector_store)), + files={"file": DOCUMENT}, + headers=dict(_forwarding(rig.tenants.team_a_key)), + ) + + +def test_forwarded_provider_key_on_ingest_never_overrides_a_named_credential(rag_rig: RagRig) -> None: + store_id: Final = _completed_store( + rag_rig, _forwarded_ingest(rag_rig, {**OPENAI, "litellm_credential_name": rag_rig.tenants.credential_name}) + ) + + assert rag_rig.signatures() == _ingest_signatures("credential", CREDENTIAL_API_KEY, store_id) + + +@pytest.mark.parametrize( + ("stored", "prefix", "api_key"), + (("own_key", "caller", CALLER_API_KEY), ("null_key", "team", FORWARDED_API_KEY)), + ids=("own_key", "null_key"), +) +def test_forwarded_provider_key_on_ingest_into_a_managed_store_yields_only_to_the_store_key( + rag_rig: RagRig, stored: Literal["own_key", "null_key"], prefix: str, api_key: str +) -> None: + with rag_rig.gateway.scenario() as scenario: + store_id: Final = _managed_store( + rag_rig, + scenario, + {"api_key": CALLER_API_KEY, "api_base": f"{rag_rig.upstream.url}/caller/v1"} + if stored == "own_key" + else {"api_key": None}, + ) + + response: Final = _forwarded_ingest(rag_rig, {**OPENAI, "vector_store_id": store_id}) + + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["vector_store_id"] == store_id, response.text + bearer: Final = f"Bearer {api_key}" + assert rag_rig.signatures() == ( + (prefix, "POST", "/v1/files", bearer), + (prefix, "POST", f"/v1/vector_stores/{store_id}/files", bearer), + ) + + +def test_forwarded_provider_key_never_reaches_a_gemini_ingest(rag_rig: RagRig) -> None: + store_id: Final = _completed_store(rag_rig, _forwarded_ingest(rag_rig, {"custom_llm_provider": "gemini"})) + + requests: Final = rag_rig.upstream.drain() + store_name: Final = store_id.removeprefix("fileSearchStores/") + assert tuple( + (call.prefix, call.method, call.route, call.google_key, call.authorization) + for call in (_call(request) for request in requests) + ) == ( + ("gemini", "POST", "/v1beta/fileSearchStores", GEMINI_API_KEY, None), + ( + "gemini", + "POST", + f"/upload/v1beta/fileSearchStores/{store_name}:uploadToFileSearchStore", + GEMINI_API_KEY, + None, + ), + ("gemini", "PUT", f"/upload-session/{store_name}", None, None), + ) + assert [request.headers.get("x-api-key") for request in requests] == [None] * len(requests), requests diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index e4a3935a58d..a0ea6ce0403 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -1,7 +1,9 @@ import asyncio import re +import sys import time from collections.abc import Mapping, Sequence +from types import ModuleType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -296,6 +298,84 @@ async def test_find_team_with_model_access_uses_request_method_for_passthrough_a assert "allowed_passthrough_routes" in exc_info.value.detail + +class _TeamClaims(set[str]): + def __init__(self, claim_order: tuple[str, ...]): + super().__init__(claim_order) + self.claim_order = claim_order + + def __iter__(self): + return iter(self.claim_order) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("claim_order", [("team-a", "team-b"), ("team-b", "team-a")]) +async def test_a_jwt_in_two_teams_resolves_the_same_team_in_any_claim_order(claim_order: tuple[str, ...]): + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + teams = {team_id: LiteLLM_TeamTable(team_id=team_id, models=["gpt-4"], metadata={}) for team_id in claim_order} + + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + side_effect=lambda team_id, **_: teams[team_id], + ), + patch( + "litellm.proxy.auth.handle_jwt.can_team_access_model", + new_callable=AsyncMock, + return_value=True, + ), + ): + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids=_TeamClaims(claim_order), + requested_model="gpt-4", + route="/v1/rag/query", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert team_id == "team-a" + assert team_obj == teams["team-a"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stray_claim", [1, None]) +async def test_a_jwt_team_claim_with_a_non_string_entry_still_resolves_the_string_team(stray_claim: object): + jwt_handler: Final = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") + team: Final = LiteLLM_TeamTable(team_id="aa-team", models=["gpt-4"], metadata={}) + + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + side_effect=lambda team_id, **_: {"aa-team": team}[team_id], + ), + patch( + "litellm.proxy.auth.handle_jwt.can_team_access_model", + new_callable=AsyncMock, + return_value=True, + ), + ): + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids=set(jwt_handler.get_team_ids_from_jwt({"groups": ["aa-team", stray_claim]})), + requested_model="gpt-4", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert team_id == "aa-team" + assert team_obj == team + + _AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = { "test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": { "endpoint_id": "test-uuid-1", @@ -1483,8 +1563,13 @@ async def test_find_team_with_model_access_v1_messages_default_routes(monkeypatc [ ("/v1/messages", True), ("/v1/messages/count_tokens", True), + ("/rag/ingest", True), + ("/v1/rag/ingest", True), + ("/rag/query", True), + ("/v1/rag/query", True), ("/v1/skills", False), ("/v1/skills/skill_abc123", False), + ("/lens", False), ], ) def test_default_team_allowed_routes_cover_messages_but_not_skills(route, expected): @@ -1500,6 +1585,40 @@ def test_default_team_allowed_routes_cover_messages_but_not_skills(route, expect ) +@pytest.mark.asyncio +async def test_find_team_with_model_access_reports_team_route_allowlist_denial(): + jwt_handler: Final = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_allowed_routes=["/chat/completions"]) + team: Final = LiteLLM_TeamTable(team_id="team-route-denied", models=["gpt-4o-mini"]) + user_api_key_cache: Final = DualCache() + user_api_key_cache.set_cache(key="team_id:team-route-denied", value=team) + prisma_client: Final = MagicMock() + prisma_team_table: Final = prisma_client.db.litellm_teamtable + prisma_team_table.find_unique = AsyncMock() + proxy_server_module: Final = ModuleType("proxy_server") + proxy_server_module.llm_router = None + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": proxy_server_module}): + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids={"team-route-denied"}, + requested_model="gpt-4o-mini", + route="/rag/query", + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + prisma_team_table.find_unique.assert_not_awaited() + assert exc_info.value.status_code == 403 + assert ( + exc_info.value.detail == "Team team-route-denied can access model gpt-4o-mini but route /rag/query " + "is not in litellm_jwtauth.team_allowed_routes" + ) + + @pytest.mark.asyncio async def test_auth_builder_returns_team_membership_object(): """ @@ -7011,6 +7130,23 @@ async def test_resolve_team_from_header_denies_aliases_of_teams_the_jwt_does_not ) +@pytest.mark.asyncio +@pytest.mark.parametrize("stray_claim", [1, None]) +async def test_resolve_team_from_header_denies_an_unknown_team_with_403_when_the_jwt_claim_holds_a_non_string( + stray_claim: object, +): + jwt_handler: Final = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") + allowed_team_ids: Final = set(jwt_handler.get_team_ids_from_jwt({"groups": ["team_a", stray_claim]})) + + with pytest.raises(HTTPException) as denied: + await _resolve_header("no_such", allowed_team_ids, False, _team_lookup_404, _teams_by_alias({})) + + assert denied.value.status_code == 403 + assert "'no_such' does not resolve" in denied.value.detail + assert "'team_a'" in denied.value.detail + + @pytest.mark.asyncio async def test_resolve_team_from_header_under_db_fallback_tries_the_id_before_the_alias(): """Under fallback_to_db_teams a claimless JWT's header is provisional: a value diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 15a254dd7a1..9ef35444419 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -36,7 +36,7 @@ from contextlib import ExitStack from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -59,6 +59,7 @@ from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import CredentialItem, LiteLLMBatch, SpecialEnums from fastapi import HTTPException, Request, Response +from starlette.datastructures import Headers, QueryParams # --------------------------------------------------------------------------- # # Fixtures: distinguishable credentials per model so a wrong/hardcoded model_id @@ -124,8 +125,8 @@ class FakeRequest: headers: Optional[Dict[str, str]] = None, query: Optional[Dict[str, str]] = None, ): - self.headers = headers or {} - self.query_params = query or {} + self.headers = Headers(headers=headers or {}) + self.query_params = QueryParams(query or {}) @pytest.fixture @@ -937,6 +938,42 @@ async def test_create__raw_file_with_executed_model_400_with_upload_guidance(har harness.router_acreate.assert_not_called() +async def _create_with_model_headers(harness: Harness, *models: str): + harness.upstream_files_route.mock(return_value=httpx.Response(200, json={"object": "list", "data": []})) + set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/batches", + "query_string": b"", + "headers": [(b"x-litellm-model", model.encode()) for model in models], + } + ) + return await endpoints.create_batch( + request=request, fastapi_response=Response(), provider=None, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test") + ) + + +@pytest.mark.asyncio +async def test_create__model_headers_naming_different_models_are_rejected_before_anything_runs(harness): + with pytest.raises(ProxyException) as exc: + await _create_with_model_headers(harness, "vertex-model", "my-vllm") + + assert exc.value.code == "400" + assert "x-litellm-model" in exc.value.message + harness.pre_call.assert_not_called() + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__the_same_model_header_twice_still_routes_to_that_model(harness): + await _create_with_model_headers(harness, "my-vllm", "my-vllm") + + assert harness.acreate_kwargs()["api_base"] == CREDS["my-vllm"]["api_base"] + + @pytest.mark.asyncio @pytest.mark.parametrize( "upstream_answer", @@ -1059,8 +1096,8 @@ def _raw_batches_request(body: Dict[str, Any]) -> MagicMock: request.url.path = "/v1/batches" request.method = "POST" request.scope = {"type": "http", "method": "POST", "path": "/v1/batches"} - request.query_params = {} - request.headers = {"Content-Type": "application/json"} + request.query_params = QueryParams({}) + request.headers = Headers(headers={"Content-Type": "application/json"}) request.client = MagicMock() request.client.host = "127.0.0.1" request.body = AsyncMock(return_value=json.dumps(body).encode()) diff --git a/tests/unit/proxy/hooks/test_batch_rate_limiter.py b/tests/unit/proxy/hooks/test_batch_rate_limiter.py index b4e972ea34f..b10cd9e7107 100644 --- a/tests/unit/proxy/hooks/test_batch_rate_limiter.py +++ b/tests/unit/proxy/hooks/test_batch_rate_limiter.py @@ -8,7 +8,7 @@ are charged against a 24h token window instead of their minute counters. import json import time -from collections.abc import Iterator, Sequence +from collections.abc import Callable, Iterator, Sequence from datetime import datetime, timezone from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -482,6 +482,155 @@ async def test_batch_rate_limiter_without_user_context(openai_files: respx.MockR assert usage_without_context.total_tokens == usage_with_context.total_tokens == _token_counter_total(rows) +@pytest.mark.asyncio +async def test_batch_input_without_a_model_is_read_through_the_team_deployment( + openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm import Router + from litellm.proxy import proxy_server + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_key": "sk-team", + "api_base": "https://team.example/v1", + }, + "model_info": {"id": "team-openai", "team_id": "team-1"}, + } + ] + ), + ) + _, _, batch_limiter = _make_limiters() + rows: Final = _batch_rows(("Hello",)) + environment_route: Final = _serve_file(openai_files, "file-abc123", rows) + team_route: Final = openai_files.get("https://team.example/v1/files/file-abc123/content").mock( + return_value=httpx.Response(200, content="\n".join(json.dumps(row) for row in rows).encode()) + ) + team_key: Final = UserAPIKeyAuth(api_key="sk", team_id="team-1", team_models=["openai/*"], models=[_BATCH_MODEL]) + teamless_key: Final = UserAPIKeyAuth(api_key="sk") + + team_usage: Final = await batch_limiter.count_input_file_usage( + file_id="file-abc123", custom_llm_provider="openai", user_api_key_dict=team_key, data={} + ) + teamless_usage: Final = await batch_limiter.count_input_file_usage( + file_id="file-abc123", custom_llm_provider="openai", user_api_key_dict=teamless_key, data={} + ) + + assert team_route.call_count == 1 + assert team_route.calls.last.request.headers["authorization"] == "Bearer sk-team" + assert environment_route.call_count == 1 + assert environment_route.calls.last.request.headers["authorization"] == "Bearer sk-test" + assert team_usage.total_tokens == teamless_usage.total_tokens == _token_counter_total(rows) + + forwarded_usage: Final = await batch_limiter.count_input_file_usage( + file_id="file-abc123", + custom_llm_provider="openai", + user_api_key_dict=team_key, + data={"api_key": "sk-forwarded"}, + ) + + assert team_route.call_count == 2 + assert team_route.calls.last.request.headers["authorization"] == "Bearer sk-forwarded" + assert environment_route.call_count == 1 + assert forwarded_usage.total_tokens == _token_counter_total(rows) + + +_ROUTED_FILE_URL: Final = "https://routed.example/v1/files/file-abc123/content" + + +def _serve_file_to_the_deployment_key(rows: Sequence[_BatchLine]) -> Callable[[httpx.Request], httpx.Response]: + jsonl: Final = "\n".join(json.dumps(row) for row in rows).encode() + + def respond(request: httpx.Request) -> httpx.Response: + if request.headers.get("authorization") != "Bearer sk-deployment": + return httpx.Response(401, json={"error": {"message": "Incorrect API key provided"}}) + return httpx.Response(200, content=jsonl) + + return respond + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "query_string, routing_headers", + [(b"model=routed-batch-model", ()), (b"", ((b"x-litellm-model", b"routed-batch-model"),))], + ids=["query", "header"], +) +async def test_model_routed_batch_is_counted_with_the_deployment_key_not_the_request_key( + openai_files: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + query_string: bytes, + routing_headers: tuple[tuple[bytes, bytes], ...], +): + from starlette.requests import Request + + from litellm import Router + from litellm.proxy import proxy_server + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "routed-batch-model", + "litellm_params": { + "model": f"openai/{_BATCH_MODEL}", + "api_key": "sk-deployment", + "api_base": "https://routed.example/v1", + }, + "model_info": {"id": "routed-batch-deployment"}, + } + ] + ), + ) + file_route: Final = openai_files.get(_ROUTED_FILE_URL).mock( + side_effect=_serve_file_to_the_deployment_key(_batch_rows(("Hello " * 50,))) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("sk-batch"), tpm_limit=1) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/batches", + "raw_path": b"/v1/batches", + "root_path": "", + "scheme": "http", + "server": ("proxy.test", 80), + "client": ("127.0.0.1", 50000), + "query_string": query_string, + "headers": [(b"content-type", b"application/json"), *routing_headers], + } + ) + data: Final = await add_litellm_data_to_request( + data={ + "input_file_id": "file-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "api_key": "sk-not-the-deployment-key", + }, + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + ) + _, _, batch_limiter = _make_limiters() + + with pytest.raises(HTTPException) as exc: + await batch_limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type="acreate_batch" + ) + + assert exc.value.status_code == 429 + assert file_route.call_count == 1 + + async def _install_managed_file_hook( openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rows: str ) -> tuple[respx.Route, respx.Route]: diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index 180ece77a6e..f6c6b51fc93 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -1,14 +1,20 @@ +from datetime import datetime, timezone from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles +from litellm import Router from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( + apply_team_provider_credentials, apply_unified_file_ids, get_credentials_for_model, + get_team_provider_credentials, is_litellm_executed_batch, map_raw_file_ids_to_unified, ) @@ -461,9 +467,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( ) -def _completed_batch_for_retire( - output_file_id: str | None, counts: BatchRequestCounts | None = None -) -> LiteLLMBatch: +def _completed_batch_for_retire(output_file_id: str | None, counts: BatchRequestCounts | None = None) -> LiteLLMBatch: kwargs = dict( id="batch-1", completion_window="24h", @@ -652,3 +656,198 @@ async def test_batch_cancel_updates_database(caplog: pytest.LogCaptureFixture): assert "batch_processed" not in written assert json.loads(written["file_object"])["cancelled_at"] == 1234567999 assert f"Updating batch {batch_id} status to cancelled after cancel" in caplog.messages + + +def _team_openai_router(team_public_model_name: str, deployment_model: str) -> Router: + return Router( + model_list=[ + { + "model_name": team_public_model_name, + "litellm_params": {"model": deployment_model, "api_key": "team-openai-key"}, + "model_info": { + "id": "team-a-openai", + "team_id": "team-a", + "team_public_model_name": team_public_model_name, + }, + }, + { + "model_name": "claude-opus-4-6", + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "anthropic-key"}, + }, + ] + ) + + +def _team_only_openai_credentials( + team_public_model_name: str, deployment_model: str, key_models: list[str] +) -> dict | None: + return get_team_provider_credentials( + llm_router=_team_openai_router(team_public_model_name, deployment_model), + user_api_key_dict=UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="team-a", + team_models=[team_public_model_name, "claude-opus-4-6"], + models=key_models, + ), + custom_llm_provider="openai", + include_accessible_deployments=False, + ) + + +@pytest.mark.parametrize( + "team_public_model_name, deployment_model, key_models", + [ + ("openai/*", "openai/*", ["gpt-4o"]), + ("openai/*", "openai/*", ["openai/gpt-4o"]), + ("gpt-4o", "openai/gpt-4o", ["openai/*"]), + ], +) +def test_a_restricted_key_resolves_the_team_deployment_its_wildcard_or_model_covers( + team_public_model_name: str, deployment_model: str, key_models: list[str] +): + credentials = _team_only_openai_credentials(team_public_model_name, deployment_model, key_models) + + assert credentials is not None + assert credentials["api_key"] == "team-openai-key" + + +@pytest.mark.parametrize("key_models", [["claude-opus-4-6"], ["anthropic/*"]]) +def test_a_restricted_key_outside_the_team_wildcard_never_resolves_its_credentials(key_models: list[str]): + assert _team_only_openai_credentials("openai/*", "openai/*", key_models) is None + + +@pytest.mark.parametrize( + "request_credentials, sent_api_key, sent_api_base", + [ + ({}, "team-openai-key", "https://team.example/v1"), + ({"api_key": ""}, "team-openai-key", "https://team.example/v1"), + ({"api_key": " "}, "team-openai-key", "https://team.example/v1"), + ({"api_key": "\n"}, "team-openai-key", "https://team.example/v1"), + ({"api_key": "sk-forwarded"}, "sk-forwarded", "https://team.example/v1"), + ({"api_base": "https://caller.example/v1"}, "team-openai-key", "https://team.example/v1"), + ( + {"api_key": "sk-forwarded", "api_base": "https://caller.example/v1"}, + "sk-forwarded", + "https://caller.example/v1", + ), + ({"api_key": 123}, "team-openai-key", "https://team.example/v1"), + ( + {"api_key": "sk-forwarded", "api_base": ["https://caller.example/v1"]}, + "sk-forwarded", + "https://team.example/v1", + ), + ({"api_key": "sk-forwarded", "api_base": " "}, "sk-forwarded", "https://team.example/v1"), + ({"litellm_credential_name": "no-such-credential"}, "team-openai-key", "https://team.example/v1"), + ({"api_key": "sk-forwarded\n", "drop_params": True}, "team-openai-key", "https://team.example/v1"), + ({"api_key": "sk-for warded"}, "sk-for warded", "https://team.example/v1"), + ], +) +def test_the_team_key_only_goes_to_the_team_host_unless_the_request_brings_its_own_key( + request_credentials: dict[str, object], sent_api_key: str, sent_api_base: str +): + data: Final[dict[str, object]] = dict(request_credentials) + + _apply_team_a_openai_credentials(data) + + assert (data.get("api_key"), data.get("api_base")) == (sent_api_key, sent_api_base) + + +@pytest.mark.parametrize("api_key", ["sk-forwarded\n", " sk-forwarded", "\u200b", "sk-forwarded\u00a0"]) +def test_a_request_key_that_cannot_be_an_http_header_is_a_bad_request(api_key: str): + data: Final[dict[str, object]] = {"api_key": api_key} + + with pytest.raises(HTTPException) as refused: + _apply_team_a_openai_credentials(data) + + assert refused.value.status_code == 400 + assert api_key not in str(refused.value.detail) + + +def _apply_team_a_openai_credentials(data: dict[str, object]) -> None: + apply_team_provider_credentials( + data=data, + llm_router=Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_key": "team-openai-key", + "api_base": "https://team.example/v1", + }, + "model_info": {"id": "team-a-openai", "team_id": "team-a"}, + } + ] + ), + user_api_key_dict=UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="team-a", team_models=["openai/*"], models=["gpt-4o"] + ), + custom_llm_provider="openai", + ) + + +def _team_openai_deployment( + deployment_id: str, api_key: str, created_at: datetime | None = None, db_model: bool = True +) -> dict[str, object]: + return { + "model_name": f"team-gpt-{deployment_id}", + "litellm_params": {"model": "openai/gpt-4o", "api_key": api_key}, + "model_info": {"id": deployment_id, "team_id": "team-a", "db_model": db_model} + | ({} if created_at is None else {"created_at": created_at}), + } + + +@pytest.mark.parametrize("deployment_ids", [("deployment-a", "deployment-b"), ("deployment-b", "deployment-a")]) +def test_a_team_with_two_openai_deployments_resolves_the_same_one_in_any_router_order( + deployment_ids: tuple[str, str], +): + router: Final = Router( + model_list=[_team_openai_deployment(deployment_id, f"{deployment_id}-key") for deployment_id in deployment_ids] + ) + + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="team-a"), + custom_llm_provider="openai", + include_accessible_deployments=False, + ) + + assert credentials is not None + assert credentials["api_key"] == "deployment-a-key" + + +@pytest.mark.parametrize("newer_first", [True, False]) +def test_a_team_keeps_its_older_openai_deployment_when_a_newer_one_has_a_lower_id(newer_first: bool): + older: Final = _team_openai_deployment("deployment-z", "older-key", datetime(2026, 1, 1, tzinfo=timezone.utc)) + newer: Final = _team_openai_deployment("deployment-a", "newer-key", datetime(2026, 6, 1, tzinfo=timezone.utc)) + router: Final = Router(model_list=[newer, older] if newer_first else [older, newer]) + + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="team-a"), + custom_llm_provider="openai", + include_accessible_deployments=False, + ) + + assert credentials is not None + assert credentials["api_key"] == "older-key" + + +def test_a_team_keeps_its_first_config_file_openai_deployment_over_older_db_ones(): + db_deployment: Final = _team_openai_deployment("deployment-a", "db-key", datetime(2020, 1, 1, tzinfo=timezone.utc)) + config_deployments: Final = tuple( + _team_openai_deployment(deployment_id, f"{deployment_id}-key", db_model=False) + for deployment_id in ("deployment-z", "deployment-b") + ) + router: Final = Router(model_list=[db_deployment, *config_deployments]) + + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="team-a"), + custom_llm_provider="openai", + include_accessible_deployments=False, + ) + + assert credentials is not None + assert credentials["api_key"] == "deployment-z-key" diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 2d654ea28ec..d9f3ceda1ed 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -1,21 +1,37 @@ """ Tests for RAG proxy endpoints. - -Covers: -- internal_user_viewer restriction: can only ingest to existing vector stores (must provide vector_store_id) """ import io import json +from collections.abc import Iterator, Mapping, Sequence +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi.testclient import TestClient import litellm +from litellm import Router from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore +from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + +_TEAM_PROVIDER_KEY: Final = "team-openai-key" +_REGISTRY_PROVIDER_KEY: Final = "registry-openai-key" +_ENVIRONMENT_OPENAI_KEY: Final = "environment-openai-key" +_TEAM_ID: Final = "team-a" +_TEAM_OPENAI_BASE: Final = "https://team-openai.example/v1" +_REGISTRY_OPENAI_BASE: Final = "https://registry-openai.example/v1" +_OPENAI_SEARCH_PAGE: Final = {"object": "vector_store.search_results.page", "search_query": "q", "data": []} +_BEDROCK_KNOWLEDGE_BASE_ID: Final = "KB12345678" +_CALLER_AWS_ACCESS_KEY: Final = "caller-access-key" +_TEAM_AWS_ACCESS_KEY: Final = "team-access-key" +_HEADER_UNSAFE_PROVIDER_KEY: Final = "sk-forwarded\u00a0".encode("latin-1") @pytest.fixture @@ -48,31 +64,32 @@ def client_internal_user(): app.dependency_overrides = original_overrides -def test_internal_user_viewer_rag_ingest_without_vector_store_id_rejected( - client_internal_user_viewer, -): - """ - internal_user_viewer cannot create new vector stores - must provide vector_store_id. - """ - # Form upload without vector_store_id (would create new store) - response = client_internal_user_viewer.post( - "/v1/rag/ingest", - files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, - data={ - "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' - }, - ) +@pytest.fixture +def client_team_a() -> Iterator[TestClient]: + mock_auth: Final = UserAPIKeyAuth(user_id="test_team_user", team_id=_TEAM_ID) + original_overrides: Final = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + try: + yield TestClient(app) + finally: + app.dependency_overrides = original_overrides - assert response.status_code == 403 - detail = response.json() - assert "detail" in detail - error_msg = ( - detail["detail"]["error"] - if isinstance(detail["detail"], dict) - else str(detail["detail"]) - ) - assert "internal_user_viewer" in error_msg - assert "vector_store_id" in error_msg + +def test_internal_user_viewer_rag_ingest_without_vector_store_id_allowed(client_internal_user_viewer): + with patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new_callable=AsyncMock, + return_value={"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) as mock_aingest: + response: Final = client_internal_user_viewer.post( + "/v1/rag/ingest", + files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, + data={"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'}, + ) + + assert response.status_code == 200, response.json() + assert response.json()["status"] == "completed" + mock_aingest.assert_awaited_once() def test_internal_user_viewer_rag_ingest_with_vector_store_id_passes_check( @@ -97,8 +114,7 @@ def test_internal_user_viewer_rag_ingest_with_vector_store_id_passes_check( # Should not be 403 (role check passed) assert response.status_code != 403, ( - f"internal_user_viewer with vector_store_id should pass role check. " - f"Response: {response.json()}" + f"internal_user_viewer with vector_store_id should pass role check. Response: {response.json()}" ) @@ -114,15 +130,12 @@ def test_internal_user_rag_ingest_without_vector_store_id_allowed(client_interna response = client_internal_user.post( "/v1/rag/ingest", files={"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, - data={ - "request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' - }, + data={"request": '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}'}, ) # Should not be 403 assert response.status_code != 403, ( - f"internal_user should be allowed to create new vector stores. " - f"Response: {response.json()}" + f"internal_user should be allowed to create new vector stores. Response: {response.json()}" ) @@ -170,13 +183,13 @@ def test_rag_ingest_blocks_clientside_credentials(client_internal_user, blocked_ }, }, ) - assert ( - response.status_code == 400 - ), f"Expected 400 when '{blocked_field}' is set clientside, got {response.status_code}: {response.json()}" + assert response.status_code == 400, ( + f"Expected 400 when '{blocked_field}' is set clientside, got {response.status_code}: {response.json()}" + ) body = response.json() - assert blocked_field in str( - body - ), f"Response should mention '{blocked_field}': {body}" + assert blocked_field in str(body), f"Response should mention '{blocked_field}': {body}" + + class TestRagIngestSSRFBlocked: """ aws_sts_endpoint and related credential-redirect fields must be rejected @@ -193,9 +206,7 @@ class TestRagIngestSSRFBlocked: ("aws_bedrock_runtime_endpoint", "https://attacker.example/bedrock"), ], ) - def test_ssrf_field_in_vector_store_config_rejected( - self, field, value, client_internal_user - ): + def test_ssrf_field_in_vector_store_config_rejected(self, field, value, client_internal_user): payload = { "file_url": "https://example.com/doc.pdf", "ingest_options": { @@ -215,9 +226,7 @@ class TestRagIngestSSRFBlocked: ) body = response.json() detail = body.get("detail", {}) - error_text = ( - detail.get("error", "") if isinstance(detail, dict) else str(detail) - ) + error_text = detail.get("error", "") if isinstance(detail, dict) else str(detail) assert field in error_text, f"Error should name the offending field: {error_text}" def test_clean_bedrock_ingest_options_not_rejected(self, client_internal_user): @@ -230,14 +239,10 @@ class TestRagIngestSSRFBlocked: "/v1/rag/ingest", json={ "file_url": "https://example.com/doc.pdf", - "ingest_options": { - "vector_store": {"custom_llm_provider": "bedrock"} - }, + "ingest_options": {"vector_store": {"custom_llm_provider": "bedrock"}}, }, ) - assert response.status_code != 400, ( - f"Clean Bedrock ingest_options should not be rejected: {response.json()}" - ) + assert response.status_code != 400, f"Clean Bedrock ingest_options should not be rejected: {response.json()}" S3_REGISTRY_STORE = { @@ -292,10 +297,20 @@ def _registry_with(store): return registry -def _ingest_form(vector_store): +def _ingest_form(vector_store, **request_fields): return { "files": {"file": ("sample.txt", io.BytesIO(b"test content"), "text/plain")}, - "data": {"request": json.dumps({"ingest_options": {"vector_store": vector_store}})}, + "data": {"request": json.dumps({"ingest_options": {"vector_store": vector_store}, **request_fields})}, + } + + +def _ingest_json(vector_store, **request_fields): + return { + "json": { + "file_url": "https://example.com/doc.txt", + "ingest_options": {"vector_store": vector_store}, + **request_fields, + } } @@ -320,6 +335,557 @@ def _patched_prisma_client(prisma_client): ) +def _team_provider_router() -> Router: + return Router( + model_list=[ + { + "model_name": "openai/team-provider", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": _TEAM_PROVIDER_KEY, + "api_base": _TEAM_OPENAI_BASE, + }, + "model_info": {"team_id": _TEAM_ID}, + }, + _completion_deployment(), + ] + ) + + +def _completion_deployment() -> dict[str, object]: + return { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "completion-key", + "mock_response": "hello", + }, + } + + +def _team_bedrock_router() -> Router: + return Router( + model_list=[ + { + "model_name": "bedrock/team-provider", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_region_name": "us-west-2", + "aws_access_key_id": _TEAM_AWS_ACCESS_KEY, + "aws_secret_access_key": "team-secret-key", + }, + "model_info": {"team_id": _TEAM_ID}, + }, + _completion_deployment(), + ] + ) + + +def _managed_openai_store(vector_store_id: str, litellm_params: Mapping[str, object]) -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore( + vector_store_id=vector_store_id, + custom_llm_provider="openai", + litellm_credential_name=None, + litellm_params=dict(litellm_params), + ) + + +def _post_rag_query( + client: TestClient, + *, + retrieval_config: Mapping[str, object], + managed_stores: Sequence[LiteLLM_ManagedVectorStore], + router: Router, + headers: Mapping[str, str | bytes] | None = None, + body: Mapping[str, object] | None = None, +) -> httpx.Response: + with ( + patch.object( # test-quality-ok: seeds the real managed-store registry the query path reads + litellm, + "vector_store_registry", + VectorStoreRegistry(vector_stores=list(managed_stores)), + ), + patch("litellm.proxy.proxy_server.llm_router", router), + _patched_prisma_client(None), + ): + return client.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": dict(retrieval_config), + **(body or {}), + }, + headers=dict(headers or {}), + ) + + +def test_rag_ingest_forwards_team_provider_credentials_without_persisting_them(client_team_a): + router: Final = _team_provider_router() + prisma_client: Final = MagicMock() + managed_vector_store_table: Final = prisma_client.db.litellm_managedvectorstorestable + managed_vector_store_table.find_unique = AsyncMock(return_value=None) + created_row: Final = MagicMock() + created_row.model_dump.return_value = { + "vector_store_id": "vs_new", + "custom_llm_provider": "openai", + "litellm_params": "{}", + "vector_store_metadata": "{}", + } + managed_vector_store_table.create = AsyncMock(return_value=created_row) + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", router), + _patched_prisma_client(prisma_client), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + **_ingest_form({"custom_llm_provider": "openai"}), + ) + + assert response.status_code == 200, response.json() + forwarded: Final = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["api_key"] == _TEAM_PROVIDER_KEY + assert forwarded["api_base"] == _TEAM_OPENAI_BASE + managed_vector_store_table.create.assert_awaited_once() + saved_data: Final = managed_vector_store_table.create.await_args.kwargs["data"] + assert saved_data["vector_store_id"] == "vs_new" + assert _TEAM_PROVIDER_KEY not in json.dumps(saved_data, default=str) + + +def test_rag_ingest_uses_the_forwarded_provider_key_over_team_provider_credentials(client_team_a): + forwarded_provider_key: Final = "sk-forwarded-provider-key" + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + patch.dict("litellm.proxy.proxy_server.general_settings", {"forward_llm_provider_auth_headers": True}), + _patched_prisma_client(None), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + headers={"Authorization": "Bearer sk-litellm-virtual-key", "x-api-key": forwarded_provider_key}, + **_ingest_form({"custom_llm_provider": "openai"}), + ) + + assert response.status_code == 200, response.json() + forwarded: Final = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["api_key"] == forwarded_provider_key + assert forwarded["api_base"] == _TEAM_OPENAI_BASE + + +@pytest.mark.parametrize("ingest_request", [_ingest_form, _ingest_json]) +def test_rag_ingest_drop_params_drops_a_forwarded_key_that_cannot_be_an_http_header(client_team_a, ingest_request): + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + patch.dict("litellm.proxy.proxy_server.general_settings", {"forward_llm_provider_auth_headers": True}), + _patched_prisma_client(None), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + headers={"Authorization": "Bearer sk-litellm-virtual-key", "x-api-key": _HEADER_UNSAFE_PROVIDER_KEY}, + **ingest_request({"custom_llm_provider": "openai"}, drop_params=True), + ) + + assert response.status_code == 200, response.json() + forwarded: Final = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["api_key"] == _TEAM_PROVIDER_KEY + assert forwarded["api_base"] == _TEAM_OPENAI_BASE + + +@pytest.mark.parametrize("unread_request_field", ["not json", json.dumps({"drop_params": False})]) +def test_rag_ingest_json_body_ignores_a_request_field_it_does_not_read(client_team_a, unread_request_field): + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + patch.dict("litellm.proxy.proxy_server.general_settings", {"forward_llm_provider_auth_headers": True}), + _patched_prisma_client(None), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + headers={"Authorization": "Bearer sk-litellm-virtual-key", "x-api-key": _HEADER_UNSAFE_PROVIDER_KEY}, + **_ingest_json({"custom_llm_provider": "openai"}, drop_params=True, request=unread_request_field), + ) + + assert response.status_code == 200, response.json() + assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"]["api_key"] == _TEAM_PROVIDER_KEY + + +def test_rag_ingest_caller_credential_name_suppresses_team_provider_credentials(client_team_a): + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + _patched_prisma_client(None), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + **_ingest_form({"custom_llm_provider": "openai", "litellm_credential_name": "caller-openai"}), + ) + + assert response.status_code == 200, response.json() + forwarded: Final = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["litellm_credential_name"] == "caller-openai" + assert "api_key" not in forwarded + + +def test_rag_ingest_managed_store_credentials_override_team_provider_credentials(client_team_a): + managed_store: Final = _managed_openai_store("registry-store-with-key", {"api_key": _REGISTRY_PROVIDER_KEY}) + aingest_patch, registry_patch = _patched_ingest_boundary( + managed_store, + {"status": "completed", "vector_store_id": "registry-store-with-key", "file_id": "file_123"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + patch( + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), + _patched_prisma_client(None), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + **_ingest_form({"vector_store_id": "registry-store-with-key"}), + ) + + assert response.status_code == 200, response.json() + forwarded: Final = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["api_key"] == _REGISTRY_PROVIDER_KEY + assert forwarded["api_key"] != _TEAM_PROVIDER_KEY + + +def test_rag_ingest_managed_store_without_credentials_uses_team_provider_credentials(client_team_a): + managed_store: Final = _managed_openai_store("registry-store-without-key", {}) + aingest_patch, registry_patch = _patched_ingest_boundary( + managed_store, + {"status": "completed", "vector_store_id": "registry-store-without-key", "file_id": "file_123"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + patch( + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), + _patched_prisma_client(None), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + **_ingest_form({"vector_store_id": "registry-store-without-key"}), + ) + + assert response.status_code == 200, response.json() + forwarded: Final = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["api_key"] == _TEAM_PROVIDER_KEY + + +@pytest.mark.parametrize("blank_credential_name", [None, "", " "]) +def test_rag_ingest_blank_caller_credential_name_uses_team_provider_credentials(client_team_a, blank_credential_name): + aingest_patch, registry_patch = _patched_ingest_boundary( + None, + {"status": "completed", "vector_store_id": "vs_new", "file_id": "file_123"}, + ) + + with ( + aingest_patch as mock_aingest, + registry_patch, + patch("litellm.proxy.proxy_server.llm_router", _team_provider_router()), + _patched_prisma_client(None), + ): + response: Final = client_team_a.post( + "/v1/rag/ingest", + **_ingest_form({"custom_llm_provider": "openai", "litellm_credential_name": blank_credential_name}), + ) + + assert response.status_code == 200, response.json() + forwarded: Final = mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] + assert forwarded["api_key"] == _TEAM_PROVIDER_KEY + assert forwarded["api_base"] == _TEAM_OPENAI_BASE + + +def test_rag_ingest_failed_upstream_status_is_returned_without_persisting_store(client_internal_user): + prisma_client: Final = MagicMock() + managed_vector_store_table: Final = prisma_client.db.litellm_managedvectorstorestable + managed_vector_store_table.find_unique = AsyncMock(return_value=None) + managed_vector_store_table.create = AsyncMock() + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new=AsyncMock( + return_value={ + "id": "ingest_abc", + "status": "failed", + "error": "unauthorized", + "error_status_code": 401, + } + ), + ), + patch("litellm.vector_store_registry", None), + _patched_prisma_client(prisma_client), + ): + response: Final = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form({"custom_llm_provider": "openai"}), + ) + + assert response.status_code == 401, response.json() + assert response.json() == {"detail": {"id": "ingest_abc", "status": "failed", "error": "unauthorized"}} + managed_vector_store_table.create.assert_not_called() + + +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + "caller_fields", + [ + {}, + {"api_key": None}, + {"api_base": None}, + {"api_key": "sk-caller-key-the-search-never-uses"}, + {"custom_llm_provider": None}, + ], +) +def test_rag_query_forwards_team_provider_credentials_to_unmanaged_search( + client_team_a: TestClient, respx_mock: respx.MockRouter, caller_fields: Mapping[str, object] +) -> None: + search: Final = respx_mock.post(f"{_TEAM_OPENAI_BASE}/vector_stores/provider-native-store/search").respond( + json=_OPENAI_SEARCH_PAGE + ) + + response: Final = _post_rag_query( + client_team_a, + retrieval_config={"vector_store_id": "provider-native-store", "custom_llm_provider": "openai", **caller_fields}, + managed_stores=(), + router=_team_provider_router(), + ) + + assert response.status_code == 200, response.json() + assert search.call_count == 1 + assert search.calls.last.request.headers["authorization"] == f"Bearer {_TEAM_PROVIDER_KEY}" + + +@pytest.mark.usefixtures("httpx_transport") +def test_rag_query_managed_store_credentials_override_team_provider_credentials( + client_team_a: TestClient, respx_mock: respx.MockRouter +) -> None: + search: Final = respx_mock.post(f"{_REGISTRY_OPENAI_BASE}/vector_stores/registry-store-with-key/search").respond( + json=_OPENAI_SEARCH_PAGE + ) + + response: Final = _post_rag_query( + client_team_a, + retrieval_config={"vector_store_id": "registry-store-with-key"}, + managed_stores=( + _managed_openai_store( + "registry-store-with-key", + {"api_key": _REGISTRY_PROVIDER_KEY, "api_base": _REGISTRY_OPENAI_BASE}, + ), + ), + router=_team_provider_router(), + ) + + assert response.status_code == 200, response.json() + assert search.call_count == 1 + assert search.calls.last.request.headers["authorization"] == f"Bearer {_REGISTRY_PROVIDER_KEY}" + + +@pytest.mark.usefixtures("httpx_transport") +def test_rag_query_managed_store_without_credentials_uses_team_provider_credentials( + client_team_a: TestClient, respx_mock: respx.MockRouter +) -> None: + search: Final = respx_mock.post(f"{_TEAM_OPENAI_BASE}/vector_stores/registry-store-without-key/search").respond( + json=_OPENAI_SEARCH_PAGE + ) + + response: Final = _post_rag_query( + client_team_a, + retrieval_config={"vector_store_id": "registry-store-without-key"}, + managed_stores=(_managed_openai_store("registry-store-without-key", {}),), + router=_team_provider_router(), + ) + + assert response.status_code == 200, response.json() + assert search.call_count == 1 + assert search.calls.last.request.headers["authorization"] == f"Bearer {_TEAM_PROVIDER_KEY}" + + +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize("blank_key", ["", " "]) +def test_rag_query_managed_store_with_an_empty_key_uses_team_provider_credentials( + client_team_a: TestClient, respx_mock: respx.MockRouter, blank_key: str +) -> None: + search: Final = respx_mock.post(f"{_TEAM_OPENAI_BASE}/vector_stores/registry-store-blank-key/search").respond( + json=_OPENAI_SEARCH_PAGE + ) + + response: Final = _post_rag_query( + client_team_a, + retrieval_config={"vector_store_id": "registry-store-blank-key"}, + managed_stores=(_managed_openai_store("registry-store-blank-key", {"api_key": blank_key}),), + router=_team_provider_router(), + ) + + assert response.status_code == 200, response.json() + assert search.call_count == 1 + assert search.calls.last.request.headers["authorization"] == f"Bearer {_TEAM_PROVIDER_KEY}" + + +@pytest.mark.usefixtures("httpx_transport") +def test_rag_query_without_a_team_deployment_keeps_the_environment_openai_credentials( + client_internal_user: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", _ENVIRONMENT_OPENAI_KEY) + for unset in ("OPENAI_BASE_URL", "OPENAI_API_BASE"): + monkeypatch.delenv(unset, raising=False) + search: Final = respx_mock.post("https://api.openai.com/v1/vector_stores/provider-native-store/search").respond( + json=_OPENAI_SEARCH_PAGE + ) + + response: Final = _post_rag_query( + client_internal_user, + retrieval_config={"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}, + managed_stores=(), + router=_team_provider_router(), + ) + + assert response.status_code == 200, response.json() + assert search.call_count == 1 + assert search.calls.last.request.headers["authorization"] == f"Bearer {_ENVIRONMENT_OPENAI_KEY}" + + +@pytest.mark.usefixtures("httpx_transport") +def test_rag_query_search_uses_the_forwarded_provider_key_over_team_provider_credentials( + client_team_a: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + forwarded_provider_key: Final = "sk-forwarded-provider-key" + monkeypatch.setenv("OPENAI_API_KEY", _ENVIRONMENT_OPENAI_KEY) + for unset in ("OPENAI_BASE_URL", "OPENAI_API_BASE"): + monkeypatch.delenv(unset, raising=False) + search: Final = respx_mock.post(f"{_TEAM_OPENAI_BASE}/vector_stores/provider-native-store/search").respond( + json=_OPENAI_SEARCH_PAGE + ) + + with patch.dict("litellm.proxy.proxy_server.general_settings", {"forward_llm_provider_auth_headers": True}): + response: Final = _post_rag_query( + client_team_a, + retrieval_config={"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}, + managed_stores=(), + router=_team_provider_router(), + headers={"Authorization": "Bearer sk-litellm-virtual-key", "x-api-key": forwarded_provider_key}, + ) + + assert response.status_code == 200, response.json() + assert search.call_count == 1 + assert search.calls.last.request.headers["authorization"] == f"Bearer {forwarded_provider_key}" + + +@pytest.mark.usefixtures("httpx_transport") +def test_rag_query_drop_params_drops_a_forwarded_key_that_cannot_be_an_http_header( + client_team_a: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", _ENVIRONMENT_OPENAI_KEY) + for unset in ("OPENAI_BASE_URL", "OPENAI_API_BASE"): + monkeypatch.delenv(unset, raising=False) + search: Final = respx_mock.post(f"{_TEAM_OPENAI_BASE}/vector_stores/provider-native-store/search").respond( + json=_OPENAI_SEARCH_PAGE + ) + + with patch.dict("litellm.proxy.proxy_server.general_settings", {"forward_llm_provider_auth_headers": True}): + response: Final = _post_rag_query( + client_team_a, + retrieval_config={"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}, + managed_stores=(), + router=_team_provider_router(), + headers={"Authorization": "Bearer sk-litellm-virtual-key", "x-api-key": _HEADER_UNSAFE_PROVIDER_KEY}, + body={"drop_params": True}, + ) + + assert response.status_code == 200, response.json() + assert search.call_count == 1 + assert search.calls.last.request.headers["authorization"] == f"Bearer {_TEAM_PROVIDER_KEY}" + + +@pytest.mark.parametrize( + "retrieval_config, error", + [ + ({"custom_llm_provider": "openai"}, "retrieval_config must contain 'vector_store_id'"), + ({"vector_store_id": 7, "custom_llm_provider": "openai"}, "vector_store_id must be a non-empty string"), + ], +) +def test_rag_query_names_what_is_wrong_with_the_vector_store_id( + client_team_a: TestClient, retrieval_config: Mapping[str, object], error: str +) -> None: + response: Final = _post_rag_query( + client_team_a, retrieval_config=retrieval_config, managed_stores=(), router=_team_provider_router() + ) + + assert response.status_code == 400, response.json() + assert error in response.text + + +@pytest.mark.usefixtures("httpx_transport") +def test_rag_query_bedrock_search_keeps_the_callers_region_and_credentials_over_team_deployment( + client_team_a: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("AWS_ACCESS_KEY_ID", _CALLER_AWS_ACCESS_KEY) + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "caller-secret-key") + for unset in ("AWS_BEARER_TOKEN_BEDROCK", "AWS_SESSION_TOKEN", "AWS_PROFILE", "AWS_REGION_NAME", "AWS_REGION"): + monkeypatch.delenv(unset, raising=False) + retrieve: Final = respx_mock.post( + url__regex=rf"^https://bedrock-agent-runtime\.[a-z0-9-]+\.amazonaws\.com/knowledgebases/{_BEDROCK_KNOWLEDGE_BASE_ID}/retrieve$" + ).respond(json={"retrievalResults": []}) + + response: Final = _post_rag_query( + client_team_a, + retrieval_config={ + "vector_store_id": _BEDROCK_KNOWLEDGE_BASE_ID, + "custom_llm_provider": "bedrock", + "aws_region_name": "us-east-1", + }, + managed_stores=(), + router=_team_bedrock_router(), + ) + + assert response.status_code == 200, response.json() + assert retrieve.call_count == 1 + sent: Final = retrieve.calls.last.request + assert sent.url.host == "bedrock-agent-runtime.us-east-1.amazonaws.com" + assert f"Credential={_CALLER_AWS_ACCESS_KEY}/" in sent.headers["authorization"] + + def test_rag_ingest_resolves_registry_store_provider_and_params(client_internal_user): aingest_patch, registry_patch = _patched_ingest_boundary( S3_REGISTRY_STORE, {"vector_store_id": "s3-store", "file_id": "file_123"} @@ -703,12 +1269,14 @@ def test_rag_query_returns_response_cost_header(client_internal_user): ) mock_response._hidden_params["response_cost"] = 3.45e-06 - with patch( - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new_callable=AsyncMock, - return_value=mock_response, - ), patch("litellm.vector_store_registry", None), patch( - "litellm.proxy.proxy_server.prisma_client", None + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ), + patch("litellm.vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), ): response = client_internal_user.post( "/v1/rag/query", @@ -742,7 +1310,9 @@ def test_rag_query_surfaces_upstream_status_code(client_internal_user, upstream_ new=AsyncMock(side_effect=upstream_error), ), patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", None + ), # test-quality-ok: proxy module global, no injection seam ): response = client_internal_user.post( "/v1/rag/query", @@ -779,10 +1349,14 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): api_key="test-key", ) - with patch( - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new=AsyncMock(side_effect=fake_aquery), - ), patch("litellm.vector_store_registry", None), patch("litellm.proxy.proxy_server.prisma_client", None): + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(side_effect=fake_aquery), + ), + patch("litellm.vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): response = client_internal_user.post( "/v1/rag/query", json={ @@ -825,7 +1399,9 @@ def test_rag_query_stream_pings_while_retrieval_is_still_running(client_internal new=AsyncMock(side_effect=slow_aquery), ), patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", None + ), # test-quality-ok: proxy module global, no injection seam ): response = client_internal_user.post( "/v1/rag/query", @@ -849,9 +1425,7 @@ def test_rag_query_stream_pings_while_retrieval_is_still_running(client_internal assert response.text.endswith("data: [DONE]\n\n") -def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepalive( - client_internal_user, monkeypatch -): +def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepalive(client_internal_user, monkeypatch): import litellm as litellm_module monkeypatch.setattr(litellm_module, "sse_keepalive_ping_interval_seconds", 5) @@ -873,7 +1447,9 @@ def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepal new=AsyncMock(side_effect=fast_aquery), ), patch("litellm.vector_store_registry", None), # test-quality-ok: proxy module global, no injection seam - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: proxy module global, no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", None + ), # test-quality-ok: proxy module global, no injection seam ): response = client_internal_user.post( "/v1/rag/query", @@ -925,13 +1501,17 @@ def test_rag_query_merges_managed_store_params(client_internal_user): model="gpt-4o-mini", ) - with patch( # test-quality-ok: aquery is the endpoint's downstream boundary; the forwarded config is what the test asserts - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new_callable=AsyncMock, - return_value=mock_response, - ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs - "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", - new=AsyncMock(return_value=True), + with ( + patch( # test-quality-ok: aquery is the endpoint's downstream boundary; the forwarded config is what the test asserts + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_aquery, + patch.object(litellm, "vector_store_registry", mock_registry), + patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), ): response = client_internal_user.post( "/v1/rag/query", @@ -971,13 +1551,17 @@ def test_rag_query_store_params_win_over_user_retrieval_config(client_internal_u model="gpt-4o-mini", ) - with patch( # test-quality-ok: aquery is the endpoint's downstream boundary; the forwarded config is what the test asserts - "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", - new_callable=AsyncMock, - return_value=mock_response, - ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs - "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", - new=AsyncMock(return_value=True), + with ( + patch( # test-quality-ok: aquery is the endpoint's downstream boundary; the forwarded config is what the test asserts + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_aquery, + patch.object(litellm, "vector_store_registry", mock_registry), + patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), ): response = client_internal_user.post( "/v1/rag/query", @@ -1164,9 +1748,7 @@ class TestVectorStoreUploadControls: content=b"benign document text\n", content_type="text/plain", ) - _options, file_data, _url, _file_id = await parse_rag_ingest_request( - request, scanner=EicarTestMalwareScanner() - ) + _options, file_data, _url, _file_id = await parse_rag_ingest_request(request, scanner=EicarTestMalwareScanner()) assert file_data is not None server_filename, content_bytes, secured_content_type = file_data assert server_filename != "../../etc/passwd" diff --git a/tests/unit/proxy/test_batch_x_litellm_model_encoding.py b/tests/unit/proxy/test_batch_x_litellm_model_encoding.py index f8c734c4efd..22fada05aa7 100644 --- a/tests/unit/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/unit/proxy/test_batch_x_litellm_model_encoding.py @@ -10,6 +10,7 @@ from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from starlette.datastructures import Headers, QueryParams from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -23,8 +24,8 @@ from litellm.types.utils import LiteLLMBatch def _make_mock_request(headers: dict) -> MagicMock: """Create a mock FastAPI Request with the given headers.""" mock_request = MagicMock() - mock_request.headers = headers - mock_request.query_params = {} + mock_request.headers = Headers(headers=headers) + mock_request.query_params = QueryParams({}) mock_request.url = MagicMock() mock_request.url.port = 4000 mock_request.method = "POST" diff --git a/tests/unit/rag/ingestion/test_base_ingestion.py b/tests/unit/rag/ingestion/test_base_ingestion.py index d20b46d356b..63a0f0d7da5 100644 --- a/tests/unit/rag/ingestion/test_base_ingestion.py +++ b/tests/unit/rag/ingestion/test_base_ingestion.py @@ -1,4 +1,5 @@ from collections import UserString +from types import SimpleNamespace from typing import Final import httpx @@ -6,10 +7,18 @@ import pytest import respx import litellm +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion +from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion +from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion +from litellm.rag.ingestion.vertex_ai_ingestion import VertexAIRAGIngestion from litellm.types.utils import CredentialItem _FILE_URL: Final = "https://files.example/docs/report.pdf" +_GEMINI_API_BASE: Final = "https://gemini.example" +_GEMINI_STORE: Final = "fileSearchStores/docs-1" +_GEMINI_UPLOAD_SESSION_URL: Final = "https://gemini.example/upload/session-1" +_TEXT_FILE: Final = ("report.txt", b"plain text", "text/plain") _STORED_CREDENTIAL: Final = CredentialItem( credential_name="gemini-prod", credential_info={}, @@ -17,6 +26,32 @@ _STORED_CREDENTIAL: Final = CredentialItem( ) +class _UnauthorizedIngestionError(Exception): + status_code: Final = 401 + + +class _FailingStoreIngestion(BaseRAGIngestion): + def __init__(self, failure: Exception) -> None: + super().__init__(ingest_options={"vector_store": {}}) + self._failure = failure + + async def store( + self, + file_content: bytes | None, + filename: str | None, + content_type: str | None, + chunks: list[str], + embeddings: list[list[float]] | None, + existing_file_id: str | None = None, + ) -> tuple[str | None, str | None]: + raise self._failure + + +class _EmbeddingRouter: + async def aembedding(self, model: str, input: list[str]) -> SimpleNamespace: + return SimpleNamespace(data=[{"embedding": [0.1, 0.2]} for _ in input]) + + def _vector_store_options(litellm_credential_name: object) -> dict[str, object]: return { "custom_llm_provider": "gemini", @@ -99,3 +134,132 @@ async def test_upload_from_a_url_takes_the_content_type_from_the_response_header uploaded: Final = await ingestion.upload(file_url=_FILE_URL) assert uploaded == ("report.pdf", b"%PDF-1.7", expected_content_type, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "expected_status_code"), + [ + (_UnauthorizedIngestionError("unauthorized"), 401), + (RuntimeError("ingestion failed"), None), + ], +) +async def test_ingest_response_preserves_integer_exception_status_code( + failure: Exception, + expected_status_code: int | None, +) -> None: + response: Final = await _FailingStoreIngestion(failure).ingest( + file_data=("report.txt", b"plain text", "text/plain") + ) + + assert response["error_status_code"] == expected_status_code + + +async def test_an_existing_file_id_on_a_store_that_cannot_take_one_fails_as_a_bad_request() -> None: + response: Final = await _FailingStoreIngestion(RuntimeError("store must not run")).ingest(file_id="file-abc") + + assert response["status"] == "failed" + assert response["error_status_code"] == 400 + assert "does not support ingesting an existing file_id" in str(response["error"]) + + +@pytest.mark.parametrize( + ("ingestion_class", "ingest_kwargs"), + [ + (OpenAIRAGIngestion, {"file_id": "file-abc"}), + (VertexAIRAGIngestion, {"file_data": ("report.txt", b"plain text", "text/plain")}), + ], +) +async def test_a_vector_store_config_missing_a_caller_field_fails_as_a_bad_request( + ingestion_class: type[BaseRAGIngestion], + ingest_kwargs: dict[str, object], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + ingestion: Final = ingestion_class(ingest_options={"vector_store": {"api_key": "caller-key"}}) + + response: Final = await ingestion.ingest(**ingest_kwargs) + + assert response["status"] == "failed" + assert response["error_status_code"] == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("vector_store_id", "upstream_statuses", "expected_status_code"), + [ + (None, {"create_store": 401}, 401), + (_GEMINI_STORE, {"start_upload": 429}, 429), + (_GEMINI_STORE, {"upload": 503}, 503), + ], +) +async def test_a_gemini_upstream_error_keeps_its_status_on_the_failed_ingest( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + vector_store_id: str | None, + upstream_statuses: dict[str, int], + expected_status_code: int, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{_GEMINI_API_BASE}/v1beta/fileSearchStores").mock( + return_value=httpx.Response(upstream_statuses.get("create_store", 200), json={"name": _GEMINI_STORE}) + ) + respx_mock.post(f"{_GEMINI_API_BASE}/upload/v1beta/{_GEMINI_STORE}:uploadToFileSearchStore").mock( + return_value=httpx.Response( + upstream_statuses.get("start_upload", 200), headers={"x-goog-upload-url": _GEMINI_UPLOAD_SESSION_URL} + ) + ) + respx_mock.put(_GEMINI_UPLOAD_SESSION_URL).mock( + return_value=httpx.Response(upstream_statuses.get("upload", 200), json={"name": f"{_GEMINI_STORE}/documents/d"}) + ) + vector_store: Final = { + "custom_llm_provider": "gemini", + "api_key": "test-key", + "api_base": _GEMINI_API_BASE, + **({"vector_store_id": vector_store_id} if vector_store_id else {}), + } + + response: Final = await GeminiRAGIngestion(ingest_options={"vector_store": vector_store}).ingest( + file_data=_TEXT_FILE + ) + + assert response["status"] == "failed" + assert response["error_status_code"] == expected_status_code + + +@pytest.mark.asyncio +async def test_a_gemini_store_with_no_api_key_anywhere_fails_without_a_caller_error_status(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + ingestion: Final = GeminiRAGIngestion(ingest_options={"vector_store": {"custom_llm_provider": "gemini"}}) + + response: Final = await ingestion.ingest(file_data=_TEXT_FILE) + + assert response["status"] == "failed" + assert response.get("error_status_code") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "vector_bucket_name", + ["ab", "bucket_with_underscores"], +) +async def test_an_invalid_s3_vectors_bucket_name_fails_as_a_bad_request(vector_bucket_name: str) -> None: + ingestion: Final = S3VectorsRAGIngestion( + ingest_options={ + "embedding": {"model": "text-embedding-3-small"}, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "aws_region_name": "us-west-2", + "vector_bucket_name": vector_bucket_name, + "dimension": 2, + }, + }, + router=_EmbeddingRouter(), + ) + + response: Final = await ingestion.ingest(file_data=_TEXT_FILE) + + assert response["status"] == "failed" + assert response["error_status_code"] == 400 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx index ff0706b24cb..9b7e60aa12a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx @@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import CreateVectorStore from "./CreateVectorStore"; import * as networking from "@/components/networking"; +import { toast } from "@/lib/toast"; // Mock the networking module vi.mock("@/components/networking", () => ({ @@ -215,6 +216,44 @@ describe("CreateVectorStore", () => { }); }); + it("should report the created store and the failed file when a later file fails to ingest", async () => { + const firstIngest = { id: "ingest-1", status: "completed", vector_store_id: "vs_123", file_id: "file_1" }; + vi.mocked(networking.ragIngestCall) + .mockResolvedValueOnce(firstIngest) + .mockRejectedValueOnce(new Error("upstream rejected the file")); + const onSuccess = vi.fn(); + render(); + + const uploadInput = document.querySelector('input[type="file"]') as HTMLInputElement; + await act(async () => { + fireEvent.change(uploadInput, { + target: { + files: [ + new File(["first"], "first.txt", { type: "text/plain" }), + new File(["second"], "second.txt", { type: "text/plain" }), + ], + }, + }); + }); + await waitFor(() => { + expect(screen.getByText("Uploaded Documents (2)")).toBeInTheDocument(); + }); + + await act(async () => { + fireEvent.click(screen.getByRole("button", { name: /Create Vector Store/i })); + }); + + await waitFor(() => { + expect(toast.warning).toHaveBeenCalledWith( + "Vector store vs_123 was created with 1 document(s), but second.txt failed to ingest", + { description: "upstream rejected the file" }, + ); + }); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(onSuccess).toHaveBeenCalledWith("vs_123"); + expect(screen.getByText("Vector Store Created Successfully")).toBeInTheDocument(); + }); + it("should exclude valkey from the provider dropdown since it has no RAG ingestion", async () => { render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx index 8f24e47340c..2422e83355d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx @@ -175,6 +175,7 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu setIsCreating(true); const results: RAGIngestResponse[] = []; let vectorStoreId: string | undefined; + let failure: { name: string; error: unknown } | undefined; try { // Ingest each document @@ -208,19 +209,35 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu console.error(`Error ingesting ${doc.name}:`, error); // Update document status to error setDocuments((prev) => prev.map((d) => (d.uid === doc.uid ? { ...d, status: "error" as const } : d))); - throw error; // Stop processing on first error + failure = { name: doc.name, error }; + break; } } + if (failure && !vectorStoreId) { + throw failure.error; + } + setIngestResults(results); - toast.success( - `Successfully created vector store with ${results.length} document(s). Vector Store ID: ${vectorStoreId}`, - ); + if (failure) { + toast.warning( + `Vector store ${vectorStoreId} was created with ${results.length} document(s), but ${failure.name} failed to ingest`, + { description: failure.error instanceof Error ? failure.error.message : String(failure.error) }, + ); + } else { + toast.success( + `Successfully created vector store with ${results.length} document(s). Vector Store ID: ${vectorStoreId}`, + ); + } if (onSuccess && vectorStoreId) { onSuccess(vectorStoreId); } + if (failure) { + return; + } + // Clear documents after successful creation setTimeout(() => { setDocuments([]); diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index dc582da6833..e7a55870345 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -975,6 +975,30 @@ describe("modelCostMap", () => { }); }); +describe("ragIngestCall", () => { + let currentFetch: typeof global.fetch; + + beforeEach(() => { + currentFetch = global.fetch; + }); + + afterEach(() => { + global.fetch = currentFetch; + }); + + it("rejects with the proxy's error text when the detail is an object", async () => { + const reason = "Error code: 404 - No vector store found with id 'vs_missing'"; + global.fetch = vi.fn().mockResolvedValue({ + ok: false, + json: async () => ({ detail: { error: reason } }), + } as any) as any; + + await expect( + Networking.ragIngestCall("sk-test", new File(["text"], "doc.txt"), "openai", "vs_missing"), + ).rejects.toThrow(reason); + }); +}); + describe("getUiConfig runtime backend discovery", () => { const discoveryConfig = () => ({ server_root_path: "/", diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 0f7d76a7156..4df34275923 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -5786,7 +5786,7 @@ export const ragIngestCall = async ( if (!response.ok) { const error = await response.json(); - throw new Error(error.error?.message || error.detail || "Failed to ingest document"); + throw new Error(deriveErrorMessage(error)); } return await response.json(); From 52f5a286794645863098ec670108f69f473db1eb Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:06:07 +0000 Subject: [PATCH 57/64] fix(bedrock): send native structured output first on Opus 4.8 and Sonnet 5 and fall back to the json tool when Bedrock rejects it (#47100) * fix(bedrock): send native structured output first on Opus 4.8 and Sonnet 5 and fall back to the json tool when Bedrock rejects it * fix(bedrock): retry with the json tool when Bedrock refuses the native output schema Count a 400 under output_config.format.schema: as a native output rejection, so a schema Bedrock cannot compile (like a pattern with a nested {n,m} group) gets the json tool retry instead of reaching the caller * fix(bedrock): make the native output retry send what an opted-out deployment sends The retry after a native output rejection now reruns the same optional param mapping completion() ran, with the deployment model_info opted out of supports_bedrock_native_structured_output, so tool_choice, effort and the json tool schema match an opted-out deployment instead of a hand-built tool. Each config keeps only the predicate that decides when native output came from the opt-in key. Bedrock batch records for Anthropic models map with the opt-in turned off, since a batch job has no retry. Models on supports_native_structured_output keep native output in batch * fix(bedrock): build the messages retry from the caller's request and keep other errors unchanged * fix(bedrock): give sync converse stream errors the upstream body as text * fix(bedrock): let a malformed caller outputConfig reach Bedrock instead of raising A Converse json_schema request whose caller outputConfig failed litellm's shape check raised a pydantic error before the request was sent, so the caller got a 500 where the merge base passed Bedrock's 400 through. The opt-in check now treats that outputConfig as carrying no native output. Also opts the vertex wire test's pinned Opus 4.8 json tool deployment out of native output, and extends the Bedrock native output integration tests. * fix(bedrock): send a refused caller outputConfig once and give opt-out advice only when native output is unsupported A caller can send its own outputConfig on Converse. When Bedrock refused that schema on an opt-in deployment, the retry still carried the caller's textFormat, so Bedrock got the same request twice and the warning told the operator to turn off an opt-in that was not the cause. The retry is now skipped when the opted-out request would still carry a native textFormat, so Bedrock's 400 reaches the caller after one call as on the merge base. The opt-out advice is logged only for the account-level "Extra inputs are not permitted" rejection. * fix(bedrock): keep the caller's continuation message on the native output retry The Converse transform pops user_continue_message from litellm_params on the first attempt, so a retry built from those params afterwards dropped it. Both Bedrock chat routes now hand the fallback a copy taken before the first attempt. * test(bedrock): pin additional_drop_params on the /v1/messages native output retry * fix(bedrock): give the Converse native output retry the fake stream timeout --- .../anthropic_messages/transformation.py | 21 +- litellm/llms/bedrock/chat/converse_handler.py | 24 +- .../bedrock/chat/converse_transformation.py | 32 +- .../anthropic_claude3_transformation.py | 6 + .../bedrock/chat/native_output_fallback.py | 62 +++ litellm/llms/bedrock/common_utils.py | 81 +++- litellm/llms/bedrock/files/transformation.py | 6 +- .../anthropic_claude3_transformation.py | 27 ++ litellm/llms/custom_httpx/llm_http_handler.py | 86 ++-- litellm/main.py | 199 ++++++--- ...odel_prices_and_context_window_backup.json | 20 + litellm/types/completion.py | 2 + model_prices_and_context_window.json | 20 + .../test_bedrock_batch_files_wire.py | 55 ++- ...t_bedrock_native_structured_output_wire.py | 390 +++++++++++++++++- ...vertex_anthropic_structured_output_wire.py | 9 +- ...drock_converse_stream_sync_decoder_wire.py | 14 + ...test_bedrock_native_output_fallback_sdk.py | 336 +++++++++++++++ ...ations_anthropic_claude3_transformation.py | 60 +++ .../chat/test_bedrock_converse_handler.py | 113 +++++ .../chat/test_converse_transformation.py | 148 +++++++ .../chat/test_native_output_fallback.py | 123 ++++++ .../test_bedrock_files_transformation.py | 25 ++ .../test_anthropic_claude3_transformation.py | 160 +++++++ .../llms/bedrock/test_bedrock_common_utils.py | 72 ++++ tests/unit/types/test_completion.py | 2 + 26 files changed, 1985 insertions(+), 108 deletions(-) create mode 100644 litellm/llms/bedrock/chat/native_output_fallback.py create mode 100644 tests/integration/sdk/test_bedrock_native_output_fallback_sdk.py create mode 100644 tests/unit/llms/bedrock/chat/test_native_output_fallback.py diff --git a/litellm/llms/base_llm/anthropic_messages/transformation.py b/litellm/llms/base_llm/anthropic_messages/transformation.py index bd6c53a0453..06717207dc0 100644 --- a/litellm/llms/base_llm/anthropic_messages/transformation.py +++ b/litellm/llms/base_llm/anthropic_messages/transformation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -180,7 +180,7 @@ class BaseAnthropicMessagesConfig(ABC): return 2 def should_retry_anthropic_messages_on_http_error( - self, e: httpx.HTTPStatusError, litellm_params: dict, request_body: object = None + self, e: httpx.HTTPStatusError, litellm_params: Mapping[str, object], request_body: object = None ) -> bool: """ When True, async_anthropic_messages_handler will transform the request body @@ -202,7 +202,9 @@ class BaseAnthropicMessagesConfig(ABC): ) and is_anthropic_thinking_binding_mismatch_error(error_text) return not caller_asked_for_this_error and is_anthropic_invalid_thinking_block_error(error_text) - def transform_anthropic_messages_request_on_http_error(self, e: httpx.HTTPStatusError, request_data: dict) -> dict: + def transform_anthropic_messages_request_on_http_error( + self, e: httpx.HTTPStatusError, request_data: dict[str, object] + ) -> dict[str, object]: """ Mutates request_data in place when retrying after a recoverable HTTP error. """ @@ -214,3 +216,16 @@ class BaseAnthropicMessagesConfig(ABC): if e.response.status_code == 400 and is_anthropic_invalid_thinking_block_error(e.response.text): strip_thinking_blocks_from_anthropic_messages_request_dict(request_data) return request_data + + def native_output_fallback_litellm_params( + self, + model: str, + anthropic_messages_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + ) -> GenericLiteLLMParams | None: + return None + + def should_retry_anthropic_messages_without_native_output( + self, e: httpx.HTTPStatusError, request_body: Mapping[str, object], model: str + ) -> bool: + return False diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 5a64e3fb78e..7d02b9ef230 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -42,14 +42,22 @@ def make_sync_call( if client is None: client = get_httpx_client() # Create a new client if none provided - response: Final = client.post( - api_base, - headers=headers, - data=data, - stream=not fake_stream, - logging_obj=logging_obj, - timeout=timeout, - ) + try: + response: Final = client.post( + api_base, + headers=headers, + data=data, + stream=not fake_stream, + logging_obj=logging_obj, + timeout=timeout, + ) + except httpx.HTTPStatusError as err: + raise BedrockError( + status_code=err.response.status_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) from err if response.status_code != 200: raise BedrockError( diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 215db3a4ee4..9deed66d647 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -12,7 +12,7 @@ from itertools import chain from typing import TYPE_CHECKING, Final, Literal, cast, overload import httpx -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger @@ -61,8 +61,10 @@ from litellm.llms.bedrock.common_utils import ( bedrock_cleaned_output_schema, bedrock_model_supports_regex_lookaround, bedrock_native_output_format, + bedrock_native_output_is_opt_in, bedrock_supports_tool_search, deployment_model_info, + requests_output_config_format, ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, @@ -170,6 +172,18 @@ def _tools_the_model_accepts( return [tool_with_sanitized_parameters(tool, drop_lookaround_regex_patterns) for tool in tools] +_OUTPUT_CONFIG_BLOCK: Final = TypeAdapter(OutputConfigBlock) + + +def _has_native_text_format(output_config: object) -> bool: + if not isinstance(output_config, Mapping): + return False + try: + return _OUTPUT_CONFIG_BLOCK.validate_python(output_config).get("textFormat") is not None + except ValidationError: + return False + + class AmazonConverseConfig(BaseConfig): """ Reference - https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html @@ -1333,6 +1347,22 @@ class AmazonConverseConfig(BaseConfig): optional_params["json_mode"] = True return optional_params + @staticmethod + def sends_native_text_format(optional_params: Mapping[str, object]) -> bool: + return _has_native_text_format(optional_params.get("outputConfig")) + + def uses_opt_in_native_output( + self, model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object] + ) -> bool: + native_output_requested: Final = ( + optional_params.get("json_mode") is True and self.sends_native_text_format(optional_params) + ) or requests_output_config_format(optional_params) + return native_output_requested and bedrock_native_output_is_opt_in( + bedrock_capability_model(model, litellm_params.get("base_model")), + self.custom_llm_provider or "bedrock_converse", + deployment_model_info(litellm_params), + ) + def update_optional_params_with_thinking_tokens(self, non_default_params: dict, optional_params: dict): """ Handles scenario where max tokens is not specified. For anthropic models (anthropic api/bedrock/vertex ai), this requires having the max tokens being set and being greater than the thinking token budget. diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index ba0bb607975..240a85ce327 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -31,6 +31,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, + requests_opt_in_native_output, strip_unsupported_bedrock_invoke_output_config_keys, tools_without_eager_input_streaming, ) @@ -102,6 +103,11 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model, self.is_thinking_requested(request_params), model_info=self._deployment_model_info ) and (json_schema is None or bedrock_cleaned_output_schema(model, json_schema) is not None) + def uses_opt_in_native_output( + self, model: str, optional_params: Mapping[str, object], litellm_params: Mapping[str, object] + ) -> bool: + return requests_opt_in_native_output(model, optional_params, deployment_model_info(litellm_params)) + def map_openai_params( self, non_default_params: dict, diff --git a/litellm/llms/bedrock/chat/native_output_fallback.py b/litellm/llms/bedrock/chat/native_output_fallback.py new file mode 100644 index 00000000000..1b8ce5fcda4 --- /dev/null +++ b/litellm/llms/bedrock/chat/native_output_fallback.py @@ -0,0 +1,62 @@ +from collections.abc import Callable, Coroutine +from typing import Final, TypeAlias + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.bedrock.common_utils import ( + NativeOutputFallback, + is_bedrock_native_output_rejection, + warn_bedrock_native_output_rejected, +) +from litellm.types.utils import ModelResponse +from litellm.utils import CustomStreamWrapper + +NativeOutputResponse: TypeAlias = ModelResponse | CustomStreamWrapper +NativeOutputAttemptResult: TypeAlias = NativeOutputResponse | Coroutine[object, object, NativeOutputResponse] +NativeOutputAttempt: TypeAlias = Callable[[NativeOutputFallback | None], NativeOutputAttemptResult] +NativeOutputFallbackFactory: TypeAlias = Callable[[], NativeOutputFallback | None] + + +def _fallback_for( + error: BaseLLMException, fallback: NativeOutputFallbackFactory, model: str +) -> NativeOutputFallback | None: + if not is_bedrock_native_output_rejection(error.status_code, error.message): + return None + retry: Final = fallback() + if retry is not None: + warn_bedrock_native_output_rejected(model, error.message) + return retry + + +async def _await_with_native_output_fallback( + first_attempt: Coroutine[object, object, NativeOutputResponse], + attempt: NativeOutputAttempt, + fallback: NativeOutputFallbackFactory, + model: str, +) -> NativeOutputResponse: + try: + return await first_attempt + except BaseLLMException as error: + retry: Final = _fallback_for(error, fallback, model) + if retry is None: + raise + second_attempt: Final = attempt(retry) + return await second_attempt if isinstance(second_attempt, Coroutine) else second_attempt + + +def call_with_native_output_fallback( + attempt: NativeOutputAttempt, + fallback: NativeOutputFallbackFactory | None, + model: str, +) -> NativeOutputAttemptResult: + if fallback is None: + return attempt(None) + try: + first_attempt: Final = attempt(None) + except BaseLLMException as error: + retry: Final = _fallback_for(error, fallback, model) + if retry is None: + raise + return attempt(retry) + if isinstance(first_attempt, Coroutine): + return _await_with_native_output_fallback(first_attempt, attempt, fallback, model) + return first_attempt diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 17a8192ca81..bbbc929fad3 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -11,6 +11,7 @@ import os import re import warnings from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias @@ -184,7 +185,7 @@ def _get_local_model_cost_map() -> dict: return GetModelCostMap.load_local_model_cost_map() -def pop_bedrock_invoke_output_config_format(request_body: dict) -> dict | None: +def pop_bedrock_invoke_output_config_format(request_body: dict[str, object]) -> dict[str, object] | None: """ Remove and return Anthropic's nested ``output_config.format`` field. @@ -205,8 +206,8 @@ def pop_bedrock_invoke_output_config_format(request_body: dict) -> dict | None: def convert_bedrock_invoke_output_format_to_inline_schema( - output_format: dict, - request_body: dict, + output_format: Mapping[str, object], + request_body: dict[str, object], ) -> None: """ Embed an Anthropic structured-output schema into the last user message. @@ -292,6 +293,77 @@ def bedrock_chat_uses_native_structured_output( ) +def bedrock_native_output_is_opt_in( + model: str, custom_llm_provider: str = "bedrock", model_info: Mapping[str, object] | None = None +) -> bool: + return not _bedrock_model_supports( + model, "supports_native_structured_output", custom_llm_provider + ) and _bedrock_native_output_opt_in(model, custom_llm_provider, model_info) + + +def without_bedrock_native_output_opt_in(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + return { + **litellm_params, + "model_info": {**deployment_model_info(litellm_params), "supports_bedrock_native_structured_output": False}, + } + + +_BEDROCK_NATIVE_OUTPUT_UNSUPPORTED: Final = "output_config.format: Extra inputs are not permitted" +_BEDROCK_NATIVE_OUTPUT_REJECTIONS: Final = ( + _BEDROCK_NATIVE_OUTPUT_UNSUPPORTED, + "compilation exceeded limits: Schema is too complex", + "output_config.format.schema:", +) + + +def is_bedrock_native_output_rejection(status_code: int, error_text: str) -> bool: + return status_code == 400 and any(rejection in error_text for rejection in _BEDROCK_NATIVE_OUTPUT_REJECTIONS) + + +def warn_bedrock_native_output_rejected(model: str, error_text: str) -> None: + opt_out_advice: Final = ( + " Set `supports_bedrock_native_structured_output: false` in this deployment's model_info to skip the " + "rejected call." + if _BEDROCK_NATIVE_OUTPUT_UNSUPPORTED in error_text + else "" + ) + verbose_logger.warning( + "Bedrock rejected native structured output for model=%s, so litellm sent the request again without it.%s", + model, + opt_out_advice, + ) + + +@dataclass(frozen=True, slots=True) +class NativeOutputFallback: + optional_params: Mapping[str, object] + litellm_params: Mapping[str, object] + + +def requests_output_config_format(optional_params: Mapping[str, object]) -> bool: + match optional_params: + case {"output_config": {"format": object()}}: + return True + case _: + return False + + +def requests_opt_in_native_output( + model: str, anthropic_optional_params: Mapping[str, object], model_info: Mapping[str, object] | None +) -> bool: + native_output_requested: Final = isinstance( + anthropic_optional_params.get("output_format"), Mapping + ) or requests_output_config_format(anthropic_optional_params) + return native_output_requested and bedrock_native_output_is_opt_in(model, model_info=model_info) + + +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) + + +def json_object_or_none(value: object) -> Mapping[str, object] | None: + return _JSON_OBJECT.validate_python(value) if isinstance(value, Mapping) else None + + BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES: Final = 64 @@ -358,7 +430,8 @@ def apply_bedrock_invoke_structured_output( """ legacy_output_format: Final = request_body.pop("output_format", None) output_config_format: Final = pop_bedrock_invoke_output_config_format(request_body) - schema_format: Final = legacy_output_format if isinstance(legacy_output_format, dict) else output_config_format + legacy_format: Final = json_object_or_none(legacy_output_format) + schema_format: Final = legacy_format if legacy_format is not None else output_config_format if schema_format is None: return diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 61e084fe544..ce2c8f15d66 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -66,8 +66,10 @@ from litellm.utils import get_llm_provider, get_optional_params from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( BedrockError, + deployment_model_info, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id, + without_bedrock_native_output_opt_in, ) S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" @@ -989,10 +991,12 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) config: Final = AmazonAnthropicClaudeConfig() + litellm_params: Final = without_bedrock_native_output_opt_in({}) mapped_params = get_optional_params( model=_invoke_route_model(model), custom_llm_provider="bedrock", messages=messages, + model_info=deployment_model_info(litellm_params), **MappingProxyType( { k: v @@ -1007,7 +1011,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): optional_params={ k: v for k, v in mapped_params.items() if k not in _MAPPED_PARAMS_THE_REQUEST_HANDLER_STRIPS }, - litellm_params={}, + litellm_params=dict(litellm_params), headers={}, ) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 96a7186b136..df85c72ecdc 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -40,12 +40,16 @@ from litellm.llms.bedrock.common_utils import ( bedrock_supports_tool_search, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, + is_bedrock_native_output_rejection, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, + requests_opt_in_native_output, + requests_output_config_format, strip_unsupported_bedrock_invoke_output_config_keys, supports_bedrock_prompt_cache_ttl, tools_without_eager_input_streaming, + warn_bedrock_native_output_rejected, ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, @@ -111,6 +115,29 @@ class AmazonAnthropicClaudeMessagesConfig( BaseAnthropicMessagesConfig.__init__(self, **kwargs) AmazonInvokeConfig.__init__(self, **kwargs) + def native_output_fallback_litellm_params( + self, + model: str, + anthropic_messages_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + ) -> GenericLiteLLMParams | None: + model_info: Final = _MODEL_INFO.validate_python(litellm_params.model_info) + if not requests_opt_in_native_output(model, anthropic_messages_optional_request_params, model_info): + return None + return litellm_params.model_copy( + update={"model_info": {**(model_info or {}), "supports_bedrock_native_structured_output": False}} + ) + + def should_retry_anthropic_messages_without_native_output( + self, e: httpx.HTTPStatusError, request_body: Mapping[str, object], model: str + ) -> bool: + if not requests_output_config_format(request_body) or not is_bedrock_native_output_rejection( + e.response.status_code, e.response.text + ): + return False + warn_bedrock_native_output_rejected(model, e.response.text) + return True + def validate_anthropic_messages_environment( self, headers: dict, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3aa925db70e..bbda4f57ec8 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,8 +1,9 @@ import asyncio +import copy import inspect import json import ssl -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager from functools import lru_cache, partial, reduce from types import MappingProxyType @@ -469,6 +470,8 @@ async def _aiter_bytes_then_close(response: httpx.Response, *, chunk_size: int) _DECODED_BODY_STALE_HEADERS: Final[frozenset[str]] = frozenset({"content-encoding", "content-length"}) _OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object]) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, object]]) def _decoded_body_headers(response: httpx.Response) -> httpx.Headers: @@ -2182,11 +2185,13 @@ class BaseLLMHTTPHandler: api_key: str | None, model: str, timeout: float | httpx.Timeout | None = None, + rebuild_without_native_output: Callable[[], Mapping[str, object]] | None = None, ) -> httpx.Response: max_attempts: Final = max(provider_config.max_retry_on_anthropic_messages_http_error, 1) litellm_params_dict: Final = dict(litellm_params) optional_params_dict: Final = dict(litellm_params) - for attempt_idx in range(max_attempts): + recovered_errors: tuple[httpx.HTTPStatusError, ...] = () # rebind-ok: a rebuilt body replays every recovery + while True: try: response = await async_httpx_client.post( url=request_url, @@ -2199,40 +2204,54 @@ class BaseLLMHTTPHandler: response.raise_for_status() return response except httpx.HTTPStatusError as e: - hit_max_attempt = attempt_idx + 1 == max_attempts + rebuild_body = ( + rebuild_without_native_output + if rebuild_without_native_output is not None + and provider_config.should_retry_anthropic_messages_without_native_output( + e=e, request_body=request_body, model=model + ) + else None + ) + hit_max_attempt = len(recovered_errors) + 1 == max_attempts should_retry = provider_config.should_retry_anthropic_messages_on_http_error( e=e, litellm_params=litellm_params_dict, request_body=request_body ) - if should_retry and not hit_max_attempt: - if logging_obj.baseline_cache_context is not None: - await logging_obj.invalidate_baseline_cache_estimate("retried_request") + if rebuild_body is None and (hit_max_attempt or not should_retry): + raise self._handle_error(e=e, provider_config=provider_config) + if logging_obj.baseline_cache_context is not None: + await logging_obj.invalidate_baseline_cache_estimate("retried_request") + if rebuild_body is not None: + request_body.clear() + request_body.update(rebuild_body()) + for recovered_error in recovered_errors: + provider_config.transform_anthropic_messages_request_on_http_error( + e=recovered_error, request_data=request_body + ) + else: verbose_logger.debug( "Anthropic /v1/messages: invalid thinking signature; " "stripping thinking blocks and retrying (attempt %s/%s).", - attempt_idx + 2, + len(recovered_errors) + 2, max_attempts, ) provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body) - headers, signed_json_body = await sign_request_off_loop_if_aws( - provider_config, - provider_config.sign_request, - headers=headers, - optional_params=optional_params_dict, - request_data=request_body, - api_base=request_url, - api_key=api_key, - stream=stream, - fake_stream=False, - model=model, - ) - logging_obj.model_call_details.update(request_body) - continue - raise self._handle_error(e=e, provider_config=provider_config) + recovered_errors = (*recovered_errors, e) + headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, + headers=headers, + optional_params=optional_params_dict, + request_data=request_body, + api_base=request_url, + api_key=api_key, + stream=stream, + fake_stream=False, + model=model, + ) + logging_obj.model_call_details.update(request_body) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) - raise RuntimeError("unreachable: anthropic messages HTTP retry loop exited without return") - @staticmethod def _resolve_anthropic_messages_timeout( litellm_params: GenericLiteLLMParams, @@ -2355,6 +2374,24 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, ) + native_output_fallback_params: Final = anthropic_messages_provider_config.native_output_fallback_litellm_params( + model=model, + anthropic_messages_optional_request_params=_OPTIONAL_PARAMS.validate_python(request_params), + litellm_params=litellm_params, + ) + rebuild_without_native_output: Final = ( + None + if native_output_fallback_params is None + else partial[Mapping[str, object]]( + anthropic_messages_provider_config.transform_anthropic_messages_request, + model=model, + messages=copy.deepcopy(_JSON_OBJECTS.validate_python(messages)), + anthropic_messages_optional_request_params=copy.deepcopy(_JSON_OBJECT.validate_python(request_params)), + litellm_params=native_output_fallback_params, + headers=_JSON_OBJECT.validate_python(headers), + ) + ) + # Prepare request body request_body: Final = anthropic_messages_provider_config.transform_anthropic_messages_request( model=model, @@ -2431,6 +2468,7 @@ class BaseLLMHTTPHandler: api_key=api_key, model=model, timeout=applied_timeout, + rebuild_without_native_output=rebuild_without_native_output, ) # used for logging + cost tracking diff --git a/litellm/main.py b/litellm/main.py index 9cad1f82c7e..2ac65b738a4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -119,7 +119,10 @@ from litellm.llms.base_llm.base_model_iterator import ( from litellm.llms.base_llm.chat.transformation import with_attribution_headers from litellm.llms.bedrock.common_utils import ( BedrockModelInfo, + NativeOutputFallback, bedrock_route_for_request, + deployment_model_info, + without_bedrock_native_output_opt_in, without_bedrock_route_prefix, ) from litellm.llms.bedrock_mantle.chat.claude_transformation import bedrock_mantle_chat_config @@ -223,6 +226,9 @@ from .llms.azure.completion.handler import AzureTextCompletion from .llms.azure_ai.anthropic.handler import AzureAnthropicChatCompletion from .llms.azure_ai.embed import AzureAIEmbedding from .llms.bedrock.chat import BedrockConverseLLM +from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import AmazonAnthropicClaudeConfig +from .llms.bedrock.chat.native_output_fallback import NativeOutputAttemptResult, call_with_native_output_fallback from .llms.bedrock.embed.embedding import BedrockEmbedding from .llms.bedrock.image_edit.handler import BedrockImageEdit from .llms.bedrock.image_generation.image_handler import BedrockImageGeneration @@ -4200,6 +4206,53 @@ _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) _BASE_URL_HINT: Final = TypeAdapter(BaseUrlRoutingHint) +_AWS_BEDROCK_CLIENT_CREDENTIAL_KEYS: Final = ( + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_region_name", +) + + +def _native_output_fallback( + ctx: CompletionDispatchContext, + first_attempt_params: Mapping[str, object], + first_attempt_litellm_params: Mapping[str, object], +) -> NativeOutputFallback: + litellm_params: Final = without_bedrock_native_output_opt_in(first_attempt_litellm_params) + remapped_params: Final = ctx.optional_params_for_model_info(deployment_model_info(litellm_params)) + return NativeOutputFallback( + optional_params={ + **{key: value for key, value in remapped_params.items() if key != "aws_bedrock_client"}, + **{ + key: first_attempt_params[key] + for key in _AWS_BEDROCK_CLIENT_CREDENTIAL_KEYS + if key in first_attempt_params + }, + }, + litellm_params=litellm_params, + ) + + +def _converse_native_output_fallback( + ctx: CompletionDispatchContext, + first_attempt_params: Mapping[str, object], + first_attempt_litellm_params: Mapping[str, object], +) -> NativeOutputFallback | None: + fallback: Final = _native_output_fallback(ctx, first_attempt_params, first_attempt_litellm_params) + return None if AmazonConverseConfig.sends_native_text_format(fallback.optional_params) else fallback + + +def _native_output_attempt_timeout( + ctx: CompletionDispatchContext, fallback: NativeOutputFallback | None +) -> float | str | httpx.Timeout | None: + if fallback is None: + return ctx.timeout + retry_timeout: Final = ctx.timeout_for_optional_params(fallback.optional_params) + ctx.logging.record_applied_timeout(retry_timeout) + return retry_timeout + + def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base @@ -4271,22 +4324,36 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu ) elif bedrock_route == "converse": model = without_bedrock_route_prefix(model) - response = bedrock_converse_chat_completion.completion( - model=model, - messages=messages, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - extra_headers=headers, # Use merged headers instead of original extra_headers - timeout=timeout, - acompletion=acompletion, - client=client, - api_base=api_base, - api_key=api_key, + converse_model: Final = model + + def converse_attempt(fallback: NativeOutputFallback | None) -> NativeOutputAttemptResult: + return bedrock_converse_chat_completion.completion( + model=converse_model, + messages=messages, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + optional_params=optional_params if fallback is None else dict(fallback.optional_params), + litellm_params=litellm_params if fallback is None else dict(fallback.litellm_params), + logger_fn=logger_fn, + encoding=_get_encoding(), + logging_obj=logging, + extra_headers=headers, # Use merged headers instead of original extra_headers + timeout=_native_output_attempt_timeout(ctx, fallback), + acompletion=acompletion, + client=client, + api_base=api_base, + api_key=api_key, + ) + + converse_params: Final = _OPTIONAL_PARAMS_ADAPTER.validate_python(optional_params) + response = call_with_native_output_fallback( + converse_attempt, + ( + partial(_converse_native_output_fallback, ctx, converse_params, dict(litellm_params)) + if AmazonConverseConfig().uses_opt_in_native_output(converse_model, converse_params, litellm_params) + else None + ), + converse_model, ) elif bedrock_route == "converse_like": model = model.replace("converse_like/", "") @@ -4308,22 +4375,36 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu client=client, ) else: - response = base_llm_http_handler.completion( - model=model, - stream=stream, - messages=messages, - acompletion=acompletion, - api_base=api_base, - model_response=model_response, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider="bedrock", - timeout=timeout, - headers=headers, - encoding=_get_encoding(), - api_key=api_key, - logging_obj=logging, - client=client, + + def invoke_attempt(fallback: NativeOutputFallback | None) -> NativeOutputAttemptResult: + return base_llm_http_handler.completion( + model=model, + stream=stream, + messages=messages, + acompletion=acompletion, + api_base=api_base, + model_response=model_response, + optional_params=optional_params if fallback is None else dict(fallback.optional_params), + litellm_params=litellm_params if fallback is None else dict(fallback.litellm_params), + custom_llm_provider="bedrock", + timeout=_native_output_attempt_timeout(ctx, fallback), + headers=headers, + encoding=_get_encoding(), + api_key=api_key, + logging_obj=logging, + client=client, + ) + + invoke_params: Final = _OPTIONAL_PARAMS_ADAPTER.validate_python(optional_params) + response = call_with_native_output_fallback( + invoke_attempt, + ( + partial(_native_output_fallback, ctx, invoke_params, dict(litellm_params)) + if isinstance(provider_config, AmazonAnthropicClaudeConfig) + and provider_config.uses_opt_in_native_output(model, invoke_params, litellm_params) + else None + ), + model, ) return response @@ -5752,37 +5833,41 @@ def completion( accepts_summary=provider_config is not None and provider_config.accepts_reasoning_summary_in_reasoning_effort, ) - optional_params = get_optional_params( - **{**optional_param_args, "reasoning_effort": provider_reasoning_effort}, - model_info=_MODEL_INFO_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else None, - **non_default_params, - ) - stream_kind: Final = ( - ( - StreamKind.FAKED - if _completion_fakes_stream( - model, - custom_llm_provider, - provider_config, - optional_params, - base_model=base_model if isinstance(base_model, str) else None, - bridges_to_responses_api=bridges_to_responses_api is True, - ) - else StreamKind.LIVE - ) - if stream is True - else StreamKind.NOT_STREAMED - ) - ### TIMEOUT LOGIC ### - timeout = CompletionTimeout.resolve( + mapping_args: Final = MappingProxyType({**optional_param_args, "reasoning_effort": provider_reasoning_effort}) + mapping_special_params: Final = MappingProxyType(dict(non_default_params)) + + def optional_params_for_model_info(deployment_info: Mapping[str, object] | None) -> dict[str, object]: + return get_optional_params(**mapping_args, model_info=deployment_info, **mapping_special_params) + + optional_params = optional_params_for_model_info( + _MODEL_INFO_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else None + ) + fakes_stream: Final = partial( + _completion_fakes_stream, + model, + custom_llm_provider, + provider_config, + base_model=base_model if isinstance(base_model, str) else None, + bridges_to_responses_api=bridges_to_responses_api is True, + ) + resolve_timeout: Final = partial( + CompletionTimeout.resolve, timeout, kwargs, custom_llm_provider, global_timeout=get_configured_request_timeout(), supports_httpx_timeout=supports_httpx_timeout, - stream_kind=stream_kind, ) + streams: Final = stream is True + + def timeout_for_optional_params(attempt_params: Mapping[str, object]) -> float | httpx.Timeout: + if not streams: + return resolve_timeout(stream_kind=StreamKind.NOT_STREAMED) + return resolve_timeout(stream_kind=StreamKind.FAKED if fakes_stream(attempt_params) else StreamKind.LIVE) + + ### TIMEOUT LOGIC ### + timeout = timeout_for_optional_params(optional_params) logging.record_applied_timeout(timeout) processed_non_default_params: Final = pre_process_non_default_params( model=model, @@ -6020,6 +6105,7 @@ def completion( model=model, model_response=model_response, optional_params=optional_params, + optional_params_for_model_info=optional_params_for_model_info, organization=organization, provider_config=provider_config, request_params=MappingProxyType( @@ -6030,6 +6116,7 @@ def completion( temperature=temperature, text_completion=text_completion, timeout=timeout, + timeout_for_optional_params=timeout_for_optional_params, top_p=top_p, ) if custom_llm_provider == "azure": diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6ad071ad2ec..51a535736f6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -2323,6 +2323,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2362,6 +2363,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2401,6 +2403,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2440,6 +2443,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2479,6 +2483,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2518,6 +2523,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2594,6 +2600,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2633,6 +2640,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2672,6 +2680,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2711,6 +2720,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2750,6 +2760,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2789,6 +2800,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -48468,6 +48480,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -48502,6 +48515,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -67143,6 +67157,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -67176,6 +67191,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -67413,6 +67429,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -67446,6 +67463,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -72798,6 +72816,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72930,6 +72949,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", diff --git a/litellm/types/completion.py b/litellm/types/completion.py index 97b77b4ca01..a6693030e89 100644 --- a/litellm/types/completion.py +++ b/litellm/types/completion.py @@ -229,6 +229,7 @@ class CompletionDispatchContext: model: str model_response: ModelResponse optional_params: dict + optional_params_for_model_info: Callable[[Mapping[str, object] | None], dict[str, object]] organization: str | None provider_config: BaseConfig | None request_params: Mapping[str, object] @@ -237,6 +238,7 @@ class CompletionDispatchContext: temperature: float | None text_completion: bool timeout: float | str | httpx.Timeout | None + timeout_for_optional_params: Callable[[Mapping[str, object]], float | httpx.Timeout] top_p: float | None diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6ad071ad2ec..51a535736f6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -2323,6 +2323,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2362,6 +2363,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2401,6 +2403,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2440,6 +2443,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2479,6 +2483,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2518,6 +2523,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2594,6 +2600,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2633,6 +2640,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2672,6 +2680,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2711,6 +2720,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2750,6 +2760,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -2789,6 +2800,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -48468,6 +48480,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -48502,6 +48515,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -67143,6 +67157,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -67176,6 +67191,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -67413,6 +67429,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -67446,6 +67463,7 @@ "supports_max_reasoning_effort": true, "supports_mid_conversation_system": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, "supports_output_config": true, "supports_parallel_tool_use_config": true, "supports_pdf_input": true, @@ -72798,6 +72816,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", @@ -72930,6 +72949,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": true, "supports_max_reasoning_effort": true, "supports_output_config": true, "bedrock_output_config_effort_ceiling": "xhigh", diff --git a/tests/integration/providers/test_bedrock_batch_files_wire.py b/tests/integration/providers/test_bedrock_batch_files_wire.py index 1a834fc2f8c..823bcda7e06 100644 --- a/tests/integration/providers/test_bedrock_batch_files_wire.py +++ b/tests/integration/providers/test_bedrock_batch_files_wire.py @@ -2,7 +2,7 @@ import json from typing import Final import pytest -from integration._support.client import Gateway +from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value from integration._support.wire import Reply, Request, wire_server MODEL: Final = "bedrock/anthropic.claude-3-haiku-20240307-v1:0" @@ -76,3 +76,56 @@ def test_completions_and_responses_batch_records_upload_as_anthropic_user_messag assert len(uploads) == 1, f"Expected exactly one S3 PUT, saw {[upload.target for upload in uploads]}" stored: Final = tuple(json.loads(line) for line in uploads[0].body.decode().splitlines() if line.strip()) assert stored == EXPECTED_S3_OBJECT + + +_JSON_TOOL: Final = "json_tool_call" +_SCHEMA: Final = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, +} + + +@pytest.mark.parametrize( + "backend", ["us.anthropic.claude-haiku-5-5", "us.anthropic.claude-opus-4-8"], ids=["haiku_5_5", "opus_4_8"] +) +def test_json_schema_batch_record_with_thinking_uploads_the_json_tool_and_no_native_output( + gateway: Gateway, backend: str +) -> None: + record: Final = { + "custom_id": "schema-record", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": f"bedrock/{backend}", + "messages": [{"role": "user", "content": "Name a city."}], + "max_tokens": 2048, + "thinking": {"type": "adaptive"}, + "response_format": {"type": "json_schema", "json_schema": {"name": "city", "schema": _SCHEMA}}, + }, + } + with wire_server(s3_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/{backend}", + api_key=None, + api_base=None, + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_region_name="us-east-1", + s3_bucket_name=BUCKET, + s3_endpoint_url=wire.url, + ) + response: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "model": model}, + {"file": ("in.jsonl", (json.dumps(record) + "\n").encode(), "application/jsonl")}, + ) + assert response.status_code == 200, response.text + (upload,) = wire.drain() + (stored,) = tuple(JSON_OBJECT.validate_json(line) for line in upload.body.decode().splitlines() if line.strip()) + model_input: Final = object_value(stored["modelInput"]) + assert "format" not in object_value(model_input.get("output_config", {})), model_input + (tool,) = (object_value(entry) for entry in list_value(model_input["tools"])) + assert tool["name"] == _JSON_TOOL, model_input + assert object_value(tool["input_schema"])["properties"] == _SCHEMA["properties"], model_input diff --git a/tests/integration/providers/test_bedrock_native_structured_output_wire.py b/tests/integration/providers/test_bedrock_native_structured_output_wire.py index f389e53d132..5803d1dee44 100644 --- a/tests/integration/providers/test_bedrock_native_structured_output_wire.py +++ b/tests/integration/providers/test_bedrock_native_structured_output_wire.py @@ -1,14 +1,16 @@ +import base64 import json import re import uuid -from collections.abc import Iterable, Iterator, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping from itertools import chain from typing import Final from urllib.parse import unquote import httpx import pytest -from integration._support.client import Gateway, Scenario, list_value, object_value, string_value +from integration._support.client import Gateway, Scenario, held, list_value, object_value, string_value +from integration._support.database import read_rows from integration._support.upstream import aws_event_stream_frame from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -17,7 +19,8 @@ _HAIKU: Final = "us.anthropic.claude-haiku-5-5" _FABLE: Final = "us.anthropic.claude-fable-5" _SONNET: Final = "us.anthropic.claude-sonnet-4-6" _OPUS: Final = "us.anthropic.claude-opus-5-5" -_OPT_IN: Final = "au.anthropic.claude-opus-4-8" +_OPT_IN: Final = _OPUS +_OPUS_4_8: Final = "us.anthropic.claude-opus-4-8" _DEEPSEEK: Final = "deepseek.v3.2" _AWS: Final[dict[str, JsonValue]] = { "aws_access_key_id": "AKIASCRIPTEDPROVIDER", @@ -38,6 +41,20 @@ _JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _FABLE_REJECTION: Final = ( "The model returned the following errors: output_config.format: Extra inputs are not permitted" ) +_SCHEMA_TOO_COMPLEX: Final = ( + "The model returned the following errors: compilation exceeded limits: Schema is too complex." +) +_PATTERN_TOO_COMPLEX: Final = ( + "The model returned the following errors: output_config.format.schema: Unsupported regex feature in pattern " + "field: pattern is too complex for structured output: reduce the {n,m} upper bound, narrow the character class " + "range, or avoid nesting quantified groups" +) +_MAX_TOKENS_400: Final = ( + "The model returned the following errors: max_tokens: 2048 > 1024, the maximum for this request" +) +_SIGNATURE_400: Final = ( + "The model returned the following errors: messages.1.content.0: Invalid `signature` in `thinking` block" +) _UNSUPPORTED_KEYWORDS: Final = ( "minimum", "maximum", @@ -327,6 +344,54 @@ def _converse_frames(blocks: tuple[dict[str, JsonValue], ...], stop_reason: str) ) +def _invoke_chunk(event: Mapping[str, JsonValue]) -> bytes: + return aws_event_stream_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, + json.dumps({"bytes": base64.b64encode(json.dumps(event).encode()).decode()}).encode(), + ) + + +def _invoke_block_events(index: int, block: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + def start(content_block: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "content_block_start", "index": index, "content_block": content_block} + + def delta(payload: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "content_block_delta", "index": index, "delta": payload} + + stop: Final[dict[str, JsonValue]] = {"type": "content_block_stop", "index": index} + match block.get("type"): + case "thinking": + return ( + start({"type": "thinking", "thinking": ""}), + delta({"type": "thinking_delta", "thinking": block["thinking"]}), + delta({"type": "signature_delta", "signature": block["signature"]}), + stop, + ) + case "tool_use": + return ( + start({"type": "tool_use", "id": block["id"], "name": block["name"], "input": {}}), + delta({"type": "input_json_delta", "partial_json": json.dumps(block["input"])}), + stop, + ) + case _: + return (start({"type": "text", "text": ""}), delta({"type": "text_delta", "text": block["text"]}), stop) + + +def _invoke_frames(message: Mapping[str, JsonValue]) -> bytes: + blocks: Final = tuple(object_value(block) for block in list_value(message["content"])) + events: Final[tuple[Mapping[str, JsonValue], ...]] = ( + {"type": "message_start", "message": {**message, "content": [], "stop_reason": None}}, + *chain.from_iterable(_invoke_block_events(index, block) for index, block in enumerate(blocks)), + { + "type": "message_delta", + "delta": {"stop_reason": message["stop_reason"], "stop_sequence": None}, + "usage": {"output_tokens": 7}, + }, + {"type": "message_stop"}, + ) + return b"".join(_invoke_chunk(event) for event in events) + + def _called_tool(tools: Iterable[str], native: bool) -> str | None: names: Final = frozenset(tools) if _USER_TOOL in names: @@ -365,6 +430,8 @@ def _bedrock(request: Request) -> Reply: "stop_sequence": None, "usage": {"input_tokens": 11, "output_tokens": 7}, } + if _is_stream(request): + return Reply(body=_invoke_frames(message), content_type=_EVENT_STREAM) return Reply(body=json.dumps(message).encode()) if _is_stream(request): return Reply(body=_converse_frames(blocks, stop_reason), content_type=_EVENT_STREAM) @@ -446,6 +513,20 @@ def _messages_text(blocks: tuple[dict[str, JsonValue], ...]) -> JsonValue: return _JSON_VALUE.validate_json(string_value(text["text"])) +def _messages_stream_text(response: httpx.Response) -> JsonValue: + assert response.status_code == 200, response.text + events: Final = tuple( + _JSON.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + deltas: Final = tuple( + object_value(event["delta"]) for event in events if event.get("type") == "content_block_delta" + ) + text: Final = "".join(string_value(delta["text"]) for delta in deltas if delta.get("type") == "text_delta") + return _JSON_VALUE.validate_json(text) + + def _refs(schema: JsonValue) -> frozenset[str]: return frozenset(string_value(node["$ref"]) for node in _schema_nodes(schema) if "$ref" in node) @@ -551,15 +632,18 @@ def test_bedrock_native_model_without_thinking_keeps_the_forced_json_tool(gatewa @pytest.mark.parametrize( - ("backend", "model_info"), - [(_OPUS, None), (_OPT_IN, {"supports_bedrock_native_structured_output": True})], - ids=["neither_key", "deployment_opt_in"], + ("backend", "model_info", "native"), + [ + (_OPUS, None, False), + (_OPT_IN, {"supports_bedrock_native_structured_output": True}, True), + (_OPUS_4_8, {"supports_bedrock_native_structured_output": False}, False), + ], + ids=["neither_key", "deployment_opt_in", "deployment_opt_out"], ) def test_thinking_request_follows_the_bedrock_native_key_the_model_resolves_to( - gateway: Gateway, backend: str, model_info: Mapping[str, JsonValue] | None + gateway: Gateway, backend: str, model_info: Mapping[str, JsonValue] | None, native: bool ) -> None: marker, marker_hex = _marker() - native: Final = model_info is not None with wire_server(_bedrock) as wire, gateway.scenario() as scenario: model: Final = _deployment(scenario, wire, "converse", backend, model_info) response: Final = gateway.request( @@ -818,3 +902,293 @@ def test_numeric_formats_and_wide_schemas_follow_the_bedrock_native_key( tools: Final = _tools(body, invoke) assert (_JSON_TOOL in tools) is not native, body assert native or tools[_JSON_TOOL]["properties"] == schema["properties"], body + + +def _rejecting_native_output(message: str) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + native: Final = _native_schema(_JSON.validate_json(request.body), _is_invoke(request)) + return _bedrock(request) if native is None else _error_reply(message) + + return reply + + +def _fallback_answer(gateway: Gateway, entry: str, model: str, marker: str) -> JsonValue: + match entry: + case "converse_chat_stream" | "invoke_chat_stream": + streamed: Final = _chat_body( + model, marker, stream=True, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA) + ) + return _content_json(_stream_message(gateway.request("POST", "/v1/chat/completions", streamed))) + case "invoke_messages_stream": + streamed_messages: Final = _messages_body( + model, marker, stream=True, output_config={"format": {"type": "json_schema", "schema": _SCHEMA}} + ) + return _messages_stream_text(gateway.request("POST", "/v1/messages", streamed_messages)) + case "converse_output_config": + chat: Final = _chat_body( + model, marker, output_config={"format": {"type": "json_schema", "schema": _SCHEMA}} + ) + return _content_json(_chat_message(gateway.request("POST", "/v1/chat/completions", chat))) + case _: + return _opt_in_answer(gateway, entry, model, marker) + + +@pytest.mark.parametrize( + ("entry", "rejection"), + [ + ("converse_chat", _FABLE_REJECTION), + ("converse_chat", _PATTERN_TOO_COMPLEX), + ("converse_chat_stream", _SCHEMA_TOO_COMPLEX), + ("converse_responses", _FABLE_REJECTION), + ("converse_messages", _FABLE_REJECTION), + ("converse_output_config", _FABLE_REJECTION), + ("invoke_chat", _FABLE_REJECTION), + ("invoke_chat_stream", _FABLE_REJECTION), + ("invoke_messages", _SCHEMA_TOO_COMPLEX), + ("invoke_messages_stream", _FABLE_REJECTION), + ], +) +def test_bedrock_rejecting_opt_in_native_output_gets_one_retry_without_it( + gateway: Gateway, entry: str, rejection: str +) -> None: + marker, marker_hex = _marker() + route: Final = entry.split("_", 1)[0] + invoke: Final = route == "invoke" + with wire_server(_rejecting_native_output(rejection)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _OPUS_4_8) + assert _fallback_answer(gateway, entry, model, marker) == _payload(marker_hex) + rejected, retried = (_JSON.validate_json(request.body) for request in wire.drain()) + assert _native_schema(rejected, invoke) is not None, rejected + assert _native_schema(retried, invoke) is None, retried + tools: Final = _tools(retried, invoke) + match entry: + case "converse_output_config": + assert not tools, retried + case "invoke_messages" | "invoke_messages_stream": + assert not tools, retried + prompt: Final = json.dumps(retried["messages"]) + assert all(name in prompt for name in object_value(_SCHEMA["properties"])), retried + case _: + assert tools[_JSON_TOOL]["properties"] == _SCHEMA["properties"], retried + + +@pytest.mark.parametrize( + ("backend", "rejection"), + [(_OPUS_4_8, _MAX_TOKENS_400), (_SONNET, _FABLE_REJECTION)], + ids=["other_400", "native_model"], +) +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_bedrock_400_is_not_retried_unless_it_rejects_opt_in_native_output( + gateway: Gateway, route: str, backend: str, rejection: str +) -> None: + marker, marker_hex = _marker() + with wire_server(_rejecting_native_output(rejection)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, backend) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body(model, marker, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA)), + ) + assert response.status_code == 400, response.text + assert rejection.split(": ", 1)[1] in response.text, response.text + _, body = _one_call(wire, marker_hex) + assert _native_schema(body, route == "invoke") == _SCHEMA, body + + +def _always_rejecting(message: str) -> Callable[[Request], Reply]: + def reply(_: Request) -> Reply: + return _error_reply(message) + + return reply + + +@pytest.mark.parametrize( + ("backend", "rejection"), + [(_OPUS_4_8, _MAX_TOKENS_400), (_SONNET, _FABLE_REJECTION)], + ids=["other_400", "native_model"], +) +def test_messages_400_on_bedrock_invoke_reaches_the_caller_after_one_call_unless_it_rejects_opt_in_native_output( + gateway: Gateway, backend: str, rejection: str +) -> None: + marker, marker_hex = _marker() + with wire_server(_always_rejecting(rejection)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "invoke", backend) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, output_config={"format": {"type": "json_schema", "schema": _SCHEMA}}), + ) + assert response.status_code == 400, response.text + assert rejection.split(": ", 1)[1] in response.text, response.text + _one_call(wire, marker_hex) + + +def test_a_messages_retry_without_native_output_still_omits_the_deployments_dropped_params(gateway: Gateway) -> None: + marker, marker_hex = _marker() + with wire_server(_rejecting_native_output(_FABLE_REJECTION)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/invoke/{_OPUS_4_8}", + api_base=wire.url, + model_info=None, + additional_drop_params=["stop_sequences"], + **_AWS, + ) + request: Final = _messages_body( + model, marker, stop_sequences=["END"], output_config={"format": {"type": "json_schema", "schema": _SCHEMA}} + ) + assert _messages_text(_messages_blocks(gateway.request("POST", "/v1/messages", request))) == _payload( + marker_hex + ) + rejected, retried = (_JSON.validate_json(call.body) for call in wire.drain()) + assert _native_schema(rejected, True) is not None, rejected + assert _native_schema(retried, True) is None, retried + assert "stop_sequences" not in rejected, rejected + assert "stop_sequences" not in retried, retried + + +def _carries_thinking_block(message: JsonValue) -> bool: + content: Final = object_value(message).get("content") + return isinstance(content, list) and any(object_value(block).get("type") == "thinking" for block in content) + + +def _replays_thinking(body: Mapping[str, JsonValue]) -> bool: + return any(_carries_thinking_block(message) for message in list_value(body["messages"])) + + +def _rejecting_native_output_and_foreign_signatures(signature_first: bool) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + body: Final = _JSON.validate_json(request.body) + checks: Final = ( + (_replays_thinking(body), _SIGNATURE_400), + (_native_schema(body, True) is not None, _FABLE_REJECTION), + ) + failing: Final = tuple(message for failed, message in (checks if signature_first else checks[::-1]) if failed) + return _error_reply(failing[0]) if failing else _bedrock(request) + + return reply + + +@pytest.mark.parametrize("signature_first", [False, True], ids=["rejection_first", "signature_first"]) +def test_messages_retry_without_native_output_keeps_the_thinking_signature_fix( + gateway: Gateway, signature_first: bool +) -> None: + marker, marker_hex = _marker() + history: Final[list[JsonValue]] = [ + {"role": "user", "content": [{"type": "text", "text": f"first question {marker}"}]}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "from another deployment", "signature": _SIGNATURE}, + {"type": "text", "text": "first answer"}, + ], + }, + {"role": "user", "content": [{"type": "text", "text": "second question"}]}, + ] + with ( + wire_server(_rejecting_native_output_and_foreign_signatures(signature_first)) as wire, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, wire, "invoke", _OPUS_4_8) + request: Final = { + **_messages_body( + model, marker, thinking=_ADAPTIVE, output_config={"format": {"type": "json_schema", "schema": _SCHEMA}} + ), + "messages": history, + } + assert _messages_text(_messages_blocks(gateway.request("POST", "/v1/messages", request))) == _payload( + marker_hex + ) + sent: Final = tuple(_JSON.validate_json(call.body) for call in wire.drain()) + assert len(sent) == 3, sent + assert _native_schema(sent[0], True) is not None and _replays_thinking(sent[0]), sent[0] + assert _native_schema(sent[-1], True) is None and not _replays_thinking(sent[-1]), sent[-1] + assert all(name in json.dumps(sent[-1]["messages"]) for name in object_value(_SCHEMA["properties"])), sent[-1] + + +@pytest.mark.parametrize("entry", ["converse_chat", "invoke_messages"]) +def test_a_request_retried_without_native_output_lands_one_success_spend_row(gateway: Gateway, entry: str) -> None: + marker, marker_hex = _marker() + invoke: Final = entry.startswith("invoke") + with wire_server(_rejecting_native_output(_FABLE_REJECTION)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, entry.split("_", 1)[0], _OPUS_4_8) + assert _fallback_answer(gateway, entry, model, marker) == _payload(marker_hex) + sent: Final = tuple(_JSON.validate_json(call.body) for call in wire.drain()) + assert len(sent) == (1 if _native_schema(sent[0], invoke) is None else 2), sent + rows: Final = held( + lambda: read_rows('SELECT status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) == 1, + holding=2, + seconds=70, + ) + assert rows[0]["status"] == "success", rows + + +_CALLER_OUTPUT_CONFIG_WITH_A_SCHEMA_OBJECT: Final[dict[str, JsonValue]] = { + "textFormat": {"type": "json_schema", "structure": {"jsonSchema": {"name": "risk", "schema": {"type": "object"}}}} +} +_BEDROCK_SCHEMA_OBJECT_400: Final = "Start of structure or map found where not expected" + + +def _converse_json_schema_field(body: Mapping[str, JsonValue]) -> JsonValue: + text_format: Final = object_value(body.get("outputConfig") or {}).get("textFormat") or {} + structure: Final = object_value(text_format).get("structure") or {} + return object_value(object_value(structure).get("jsonSchema") or {}).get("schema") + + +def _deserializing_like_bedrock(request: Request) -> Reply: + schema: Final = _converse_json_schema_field(_JSON.validate_json(request.body)) + if schema is None or isinstance(schema, str): + return _bedrock(request) + return Reply(status=400, body=json.dumps({"Message": _BEDROCK_SCHEMA_OBJECT_400}).encode()) + + +def test_a_malformed_caller_output_config_on_an_opt_in_model_gets_bedrocks_own_400_after_one_call( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker() + with wire_server(_deserializing_like_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _OPUS_4_8) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body( + model, + marker, + thinking=_ADAPTIVE, + response_format=_response_format(_SCHEMA), + outputConfig=_CALLER_OUTPUT_CONFIG_WITH_A_SCHEMA_OBJECT, + ), + ) + assert response.status_code == 400, response.text + assert _BEDROCK_SCHEMA_OBJECT_400 in response.text, response.text + _, body = _one_call(wire, marker_hex) + assert body["outputConfig"] == _CALLER_OUTPUT_CONFIG_WITH_A_SCHEMA_OBJECT, body + + +_CALLER_OUTPUT_CONFIG: Final[dict[str, JsonValue]] = { + "textFormat": {"type": "json_schema", "structure": {"jsonSchema": {"name": "risk", "schema": json.dumps(_SCHEMA)}}} +} + + +@pytest.mark.parametrize("thinking", [None, _ADAPTIVE], ids=["no_thinking", "adaptive_thinking"]) +def test_a_caller_output_config_bedrock_refuses_is_sent_once_on_an_opt_in_model( + gateway: Gateway, thinking: dict[str, JsonValue] | None +) -> None: + marker, marker_hex = _marker() + with wire_server(_rejecting_native_output(_PATTERN_TOO_COMPLEX)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _OPUS_4_8) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body( + model, + marker, + response_format=_response_format(_SCHEMA), + outputConfig=_CALLER_OUTPUT_CONFIG, + **({} if thinking is None else {"thinking": thinking}), + ), + ) + assert response.status_code == 400, response.text + assert "pattern is too complex for structured output" in response.text, response.text + _, body = _one_call(wire, marker_hex) + assert body["outputConfig"] == _CALLER_OUTPUT_CONFIG, body diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index b2aba7fb100..bba79049283 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -757,12 +757,15 @@ def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, back ) -def _bedrock_invoke_deployment(scenario: Scenario, api_base: str, backend: str) -> str: +def _bedrock_invoke_deployment( + scenario: Scenario, api_base: str, backend: str, model_info: Mapping[str, JsonValue] | None = None +) -> str: return scenario.model( model=f"bedrock/invoke/{backend}", api_base=api_base, api_key=_BEDROCK_TOKEN, aws_region_name="us-east-1", + model_info=model_info, ) @@ -1521,7 +1524,9 @@ def _json_tool_deployment(gateway: Gateway, scenario: Scenario, api_base: str, p case "vertex_ai": return _vertex_deployment(gateway, scenario, api_base, _TOOL_MODEL) case "bedrock_invoke": - return _bedrock_invoke_deployment(scenario, api_base, _BEDROCK_JSON_TOOL_MODEL) + return _bedrock_invoke_deployment( + scenario, api_base, _BEDROCK_JSON_TOOL_MODEL, {"supports_bedrock_native_structured_output": False} + ) case _: raise AssertionError(provider) diff --git a/tests/integration/sdk/test_bedrock_converse_stream_sync_decoder_wire.py b/tests/integration/sdk/test_bedrock_converse_stream_sync_decoder_wire.py index 90acdd774d0..87efbc5529e 100644 --- a/tests/integration/sdk/test_bedrock_converse_stream_sync_decoder_wire.py +++ b/tests/integration/sdk/test_bedrock_converse_stream_sync_decoder_wire.py @@ -98,3 +98,17 @@ def test_k03_sync_stream_with_normal_frames_delivers_the_text_and_a_stop() -> No assert consumed.text == _ANSWER, consumed assert consumed.finish_reasons[-1] == "stop", consumed assert len(wire.drain()) == 1 + + +def test_k04_sync_stream_open_400_carries_the_bedrock_body_as_text() -> None: + body: Final = '{"message":"The request is invalid."}' + + def reject(_: Request) -> Reply: + return Reply(status=400, body=body.encode()) + + with wire_server(reject) as wire: + with pytest.raises(BadRequestError) as raised: + _consume_sync_stream(wire) + assert raised.value.status_code == 400, raised.value + assert f"BedrockException - {body}" in raised.value.message, raised.value.message + assert len(wire.drain()) == 1 diff --git a/tests/integration/sdk/test_bedrock_native_output_fallback_sdk.py b/tests/integration/sdk/test_bedrock_native_output_fallback_sdk.py new file mode 100644 index 00000000000..745c1ceb603 --- /dev/null +++ b/tests/integration/sdk/test_bedrock_native_output_fallback_sdk.py @@ -0,0 +1,336 @@ +import json +import re +import uuid +from collections.abc import Mapping +from typing import Final + +import boto3 +import pytest +from integration._support.client import list_value, object_value, string_value +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm import CustomStreamWrapper, ModelResponse +from litellm.types.utils import Choices, StreamingChoices + +_OPUS_4_8: Final = "us.anthropic.claude-opus-4-8" +_JSON_TOOL: Final = "json_tool_call" +_MARKER: Final = re.compile(r"probe ([0-9a-f]{32})") +_REJECTION: Final = "The model returned the following errors: output_config.format: Extra inputs are not permitted" +_SCHEMA: Final = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, +} +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_ROUTES: Final = pytest.mark.parametrize( + ("route", "stream"), + [("converse", False), ("converse", True), ("invoke", False)], + ids=["converse", "converse_stream", "invoke"], +) + + +@pytest.fixture +def marker_hex() -> str: + return uuid.uuid4().hex + + +def _is_invoke(request: Request) -> bool: + return "/invoke" in request.target + + +def _native(body: Mapping[str, JsonValue], invoke: bool) -> bool: + output_config: Final = body.get("output_config" if invoke else "outputConfig") + return isinstance(output_config, dict) and ("format" if invoke else "textFormat") in output_config + + +def _tool_names(body: Mapping[str, JsonValue], invoke: bool) -> frozenset[str]: + if invoke: + return frozenset(string_value(object_value(tool)["name"]) for tool in list_value(body.get("tools", []))) + tools: Final = list_value(object_value(body.get("toolConfig") or {"tools": []})["tools"]) + return frozenset(string_value(object_value(object_value(tool)["toolSpec"])["name"]) for tool in tools) + + +def _tool_reply(invoke: bool, answer: Mapping[str, JsonValue]) -> Reply: + if invoke: + message: Final = { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": [{"type": "tool_use", "id": "toolu_scripted", "name": _JSON_TOOL, "input": answer}], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 7}, + } + return Reply(body=json.dumps(message).encode()) + converse: Final = { + "output": { + "message": { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "tooluse_scripted", "name": _JSON_TOOL, "input": answer}}], + } + }, + "stopReason": "tool_use", + "usage": {"inputTokens": 11, "outputTokens": 7, "totalTokens": 18}, + "metrics": {"latencyMs": 1}, + } + return Reply(body=json.dumps(converse).encode()) + + +def _bedrock_without_native_output(request: Request) -> Reply: + body: Final = _JSON.validate_json(request.body) + invoke: Final = _is_invoke(request) + if _native(body, invoke): + return Reply(status=400, body=json.dumps({"message": _REJECTION}).encode()) + found: Final = _MARKER.search(json.dumps(body)) + assert found is not None and _JSON_TOOL in _tool_names(body, invoke), body + return _tool_reply(invoke, {"city": found.group(1)}) + + +def _request(route: str, stream: bool, wire: Wire, marker_hex: str) -> Mapping[str, object]: + return { + "model": f"bedrock/{route}/{_OPUS_4_8}", + "messages": [{"role": "user", "content": f"Name a city. probe {marker_hex}"}], + "api_base": wire.url, + "api_key": "synthetic-bedrock-bearer", + "aws_region_name": "us-east-1", + "max_tokens": 2048, + "thinking": {"type": "adaptive"}, + "response_format": {"type": "json_schema", "json_schema": {"name": "city", "schema": _SCHEMA}}, + "num_retries": 0, + "stream": stream, + } + + +def _message_content(response: ModelResponse) -> str: + (choice,) = response.choices + assert isinstance(choice, Choices), response + return string_value(choice.message.content) + + +def _delta_content(chunk: ModelResponse) -> str: + (choice,) = chunk.choices + assert isinstance(choice, StreamingChoices), chunk + return choice.delta.content or "" + + +def _assert_one_retry_without_native_output(wire: Wire, route: str) -> None: + rejected, retried = (_JSON.validate_json(request.body) for request in wire.drain()) + invoke: Final = route == "invoke" + assert _native(rejected, invoke), rejected + assert not _native(retried, invoke), retried + assert _JSON_TOOL in _tool_names(retried, invoke), retried + + +@_ROUTES +def test_sync_call_retries_a_native_output_rejection_once_with_the_json_tool( + route: str, stream: bool, marker_hex: str +) -> None: + with wire_server(_bedrock_without_native_output) as wire: + response: Final = litellm.completion(**_request(route, stream, wire, marker_hex)) + content: Final = ( + "".join(_delta_content(chunk) for chunk in response) + if isinstance(response, CustomStreamWrapper) + else _message_content(response) + ) + assert json.loads(content) == {"city": marker_hex} + _assert_one_retry_without_native_output(wire, route) + + +@_ROUTES +async def test_async_call_retries_a_native_output_rejection_once_with_the_json_tool( + route: str, stream: bool, marker_hex: str +) -> None: + with wire_server(_bedrock_without_native_output) as wire: + response: Final = await litellm.acompletion(**_request(route, stream, wire, marker_hex)) + content: Final = ( + "".join([_delta_content(chunk) async for chunk in response]) + if isinstance(response, CustomStreamWrapper) + else _message_content(response) + ) + assert json.loads(content) == {"city": marker_hex} + _assert_one_retry_without_native_output(wire, route) + + +_CONSTRAINED_SCHEMA: Final = { + "type": "object", + "properties": { + "city": {"type": "string", "maxLength": 20}, + "population": {"type": "integer", "minimum": 0}, + "districts": {"type": "array", "items": {"type": "string"}, "minItems": 2, "uniqueItems": True}, + "code": {"oneOf": [{"type": "string"}, {"type": "integer"}]}, + }, + "required": ["city", "population", "districts", "code"], + "additionalProperties": False, +} + + +def _response_format(schema: Mapping[str, object]) -> Mapping[str, object]: + return {"type": "json_schema", "json_schema": {"name": "city", "schema": schema}} + + +def _consume(response: ModelResponse | CustomStreamWrapper) -> None: + if isinstance(response, CustomStreamWrapper): + tuple(response) + + +def _sent_bodies( + route: str, + stream: bool, + marker_hex: str, + request_params: Mapping[str, object], + model_info: Mapping[str, object] | None, +) -> tuple[Mapping[str, JsonValue], ...]: + with wire_server(_bedrock_without_native_output) as wire: + request: Final = { + **_request(route, stream, wire, marker_hex), + **{"thinking": None, "response_format": None}, + **request_params, + **({} if model_info is None else {"model_info": model_info}), + } + _consume(litellm.completion(**{key: value for key, value in request.items() if value is not None})) + return tuple(_JSON.validate_json(sent.body) for sent in wire.drain()) + + +@_ROUTES +@pytest.mark.parametrize( + "request_params", + [ + {"reasoning_effort": "low", "response_format": _response_format(_SCHEMA)}, + {"thinking": {"type": "adaptive"}, "response_format": _response_format(_CONSTRAINED_SCHEMA)}, + { + "thinking": {"type": "adaptive"}, + "response_format": _response_format(_SCHEMA), + "output_config": {"format": {"type": "json_schema", "schema": _SCHEMA}}, + }, + { + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "low"}, + "response_format": _response_format(_SCHEMA), + }, + ], + ids=[ + "reasoning_effort", + "constrained_schema", + "response_format_and_output_config_format", + "adaptive_thinking_with_effort", + ], +) +def test_the_retry_sends_what_a_deployment_without_native_output_sends( + route: str, stream: bool, marker_hex: str, request_params: Mapping[str, object] +) -> None: + rejected, retried = _sent_bodies(route, stream, marker_hex, request_params, None) + (opted_out,) = _sent_bodies( + route, stream, marker_hex, request_params, {"supports_bedrock_native_structured_output": False} + ) + + assert _native(rejected, route == "invoke"), rejected + assert retried == opted_out + + +@_ROUTES +def test_the_retry_keeps_the_callers_continuation_message(route: str, stream: bool, marker_hex: str) -> None: + request_params: Final = { + "thinking": {"type": "adaptive"}, + "response_format": _response_format(_SCHEMA), + "messages": [ + {"role": "user", "content": f"Name a city. probe {marker_hex}"}, + {"role": "assistant", "content": "Which country should it be in?"}, + ], + "user_continue_message": {"role": "user", "content": "Any country in South America"}, + } + + rejected, retried = _sent_bodies(route, stream, marker_hex, request_params, None) + (opted_out,) = _sent_bodies( + route, stream, marker_hex, request_params, {"supports_bedrock_native_structured_output": False} + ) + + assert _native(rejected, route == "invoke"), rejected + assert retried == opted_out + + +_MAX_LENGTH_ONE_OF_SCHEMA: Final = { + "type": "object", + "properties": { + "city": {"type": "string", "maxLength": 20}, + "code": {"oneOf": [{"type": "string"}, {"type": "integer"}]}, + }, + "required": ["city", "code"], + "additionalProperties": False, +} + + +def _messages_bedrock_without_native_output(request: Request) -> Reply: + if not _is_invoke(request) or _native(_JSON.validate_json(request.body), invoke=True): + return _bedrock_without_native_output(request) + message: Final = { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": [{"type": "text", "text": '{"city": "Paris"}'}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 7}, + } + return Reply(body=json.dumps(message).encode()) + + +async def _messages_bodies( + route: str, marker_hex: str, model_info: Mapping[str, object] | None +) -> tuple[Mapping[str, JsonValue], ...]: + with wire_server(_messages_bedrock_without_native_output) as wire: + await litellm.anthropic_messages( + model=f"bedrock/{route}/{_OPUS_4_8}", + messages=[{"role": "user", "content": f"Name a city. probe {marker_hex}"}], + api_base=wire.url, + api_key="synthetic-bedrock-bearer", + aws_region_name="us-east-1", + max_tokens=2048, + thinking={"type": "adaptive"}, + output_config={"effort": "low", "format": {"type": "json_schema", "schema": _MAX_LENGTH_ONE_OF_SCHEMA}}, + **({} if model_info is None else {"model_info": model_info}), + ) + return tuple(_JSON.validate_json(sent.body) for sent in wire.drain()) + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +async def test_a_messages_retry_sends_what_a_deployment_without_native_output_sends( + route: str, marker_hex: str +) -> None: + rejected, retried = await _messages_bodies(route, marker_hex, None) + (opted_out,) = await _messages_bodies(route, marker_hex, {"supports_bedrock_native_structured_output": False}) + + assert _native(rejected, invoke=route == "invoke"), rejected + assert retried == opted_out + + +@_ROUTES +def test_the_retry_signs_with_the_credentials_of_a_deprecated_aws_bedrock_client( + route: str, stream: bool, marker_hex: str +) -> None: + aws_bedrock_client: Final = boto3.client( + "bedrock-runtime", + region_name="us-west-2", + aws_access_key_id="AKIASYNTHETICCLIENT", + aws_secret_access_key="synthetic-client-secret", + aws_session_token="synthetic-client-session", + ) + with wire_server(_bedrock_without_native_output) as wire: + request: Final = { + key: value + for key, value in _request(route, stream, wire, marker_hex).items() + if key not in {"api_key", "aws_region_name"} + } + _consume(litellm.completion(**request, aws_bedrock_client=aws_bedrock_client)) + sent: Final = wire.drain() + + assert len(sent) == 2 + for attempt in sent: + headers = {name.lower(): value for name, value in attempt.headers.items()} + assert "Credential=AKIASYNTHETICCLIENT/" in headers["authorization"] + assert "/us-west-2/bedrock/aws4_request" in headers["authorization"] + assert headers["x-amz-security-token"] == "synthetic-client-session" diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index ca34a4eee06..bd98e4eb4de 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1948,3 +1948,63 @@ def test_refusal_stop_details_kept_through_bedrock_invoke_transform(): provider_fields: Final = result.choices[0].message.provider_specific_fields assert provider_fields["stop_reason"] == "refusal" assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + + +_CITY_SCHEMA: Final = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, +} +_CITY_RESPONSE_FORMAT: Final = {"type": "json_schema", "json_schema": {"name": "city", "schema": _CITY_SCHEMA}} + + +def _invoke_chat_params( + model: str, model_info: Mapping[str, object] | None, non_default_params: Mapping[str, object] +) -> Mapping[str, object]: + return AmazonAnthropicClaudeConfig(model_info=model_info).map_openai_params( + non_default_params=dict(non_default_params), optional_params={}, model=model, drop_params=False + ) + + +@pytest.mark.parametrize( + "thinking_params", + [{"thinking": {"type": "adaptive"}}, {"reasoning_effort": "low"}], + ids=["adaptive", "reasoning_effort"], +) +@pytest.mark.parametrize("model", ["us.anthropic.claude-opus-4-8", "global.anthropic.claude-sonnet-5"]) +def test_invoke_uses_opt_in_native_output_when_the_opt_in_drove_native_output( + local_model_cost_map, model, thinking_params +): + optional_params: Final = _invoke_chat_params( + model, None, {"response_format": _CITY_RESPONSE_FORMAT, **thinking_params} + ) + + assert AmazonAnthropicClaudeConfig().uses_opt_in_native_output(model, optional_params, {}) + + +@pytest.mark.parametrize( + ("model", "model_info", "request_params"), + [ + ("us.anthropic.claude-opus-4-8", None, {"response_format": _CITY_RESPONSE_FORMAT}), + ( + "us.anthropic.claude-opus-4-8", + {"supports_bedrock_native_structured_output": False}, + {"response_format": _CITY_RESPONSE_FORMAT, "thinking": {"type": "adaptive"}}, + ), + ( + "us.anthropic.claude-sonnet-4-6", + None, + {"response_format": _CITY_RESPONSE_FORMAT, "thinking": {"type": "adaptive"}}, + ), + ], + ids=["no_thinking", "deployment_opt_out", "native_model"], +) +def test_invoke_uses_opt_in_native_output_is_false_when_native_output_did_not_come_from_the_opt_in( + local_model_cost_map, model, model_info, request_params +): + optional_params: Final = _invoke_chat_params(model, model_info, request_params) + + assert not AmazonAnthropicClaudeConfig().uses_opt_in_native_output( + model, optional_params, {"model_info": model_info} + ) diff --git a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py index 3c96163a7b8..1fa83254822 100644 --- a/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/unit/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -549,6 +549,73 @@ async def test_failed_fake_stream_reports_the_request_timeout_it_was_bounded_by( assert getattr(excinfo.value, "timeout", None) == request_timeout +_NATIVE_OUTPUT_REJECTION: Final = json.dumps( + {"message": "The model returned the following errors: output_config.format: Extra inputs are not permitted"} +).encode() + + +_ANSWER_SCHEMA: Final = { + "type": "object", + "properties": {"answer": {"type": "string"}}, + "required": ["answer"], + "additionalProperties": False, +} + + +def _rejects_native_output_then_answers_slowly(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/converse-stream"): + return httpx.Response(400, content=_NATIVE_OUTPUT_REJECTION, request=request) + return slow_upstream_answer(CONVERSE_ANSWER)(request) + + +async def _stream_json_schema_with_thinking(router: litellm.Router) -> str: + response: Final = await router.acompletion( + model="bedrock", + messages=[{"role": "user", "content": "hi"}], + stream=True, + thinking={"type": "adaptive"}, + response_format={"type": "json_schema", "json_schema": {"name": "answer", "schema": _ANSWER_SCHEMA}}, + ) + return "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + + +def _sent_paths(respx_mock: respx.MockRouter) -> list[str]: + return [call.request.url.path.rsplit("/", 1)[-1] for call in respx_mock.calls] + + +@pytest.mark.asyncio +async def test_native_output_retry_is_bounded_by_the_request_timeout_not_the_stream_timeout( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=_rejects_native_output_then_answers_slowly) + router: Final = _bedrock_router( + "bedrock/converse/us.anthropic.claude-opus-4-8", + timeout=UPSTREAM_FIRST_BYTE_SECONDS * 2, + stream_timeout=STREAM_TIMEOUT_SECONDS, + ) + + assert await _stream_json_schema_with_thinking(router) == "hi" + assert _sent_paths(respx_mock) == ["converse-stream", "converse"] + + +@pytest.mark.asyncio +async def test_failed_native_output_retry_reports_the_request_timeout_it_was_bounded_by( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_BEDROCK_RUNTIME_URL).mock(side_effect=_rejects_native_output_then_answers_slowly) + request_timeout: Final = UPSTREAM_FIRST_BYTE_SECONDS / 2 + router: Final = _bedrock_router( + "bedrock/converse/us.anthropic.claude-opus-4-8", timeout=request_timeout, stream_timeout=STREAM_TIMEOUT_SECONDS + ) + + with pytest.raises(litellm.Timeout) as excinfo: + await _stream_json_schema_with_thinking(router) + assert getattr(excinfo.value, "timeout", None) == request_timeout + assert _sent_paths(respx_mock) == ["converse-stream", "converse"] + + @pytest.mark.asyncio @pytest.mark.parametrize( @@ -805,3 +872,49 @@ def test_fake_streamed_json_mode_answer_closes_merged_reasoning_before_the_answe assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == ( f"{answer.choices[0].message.reasoning_content}{answer.choices[0].message.content}" ) + + +_UPSTREAM_ERROR_BODY: Final = '{"message":"The request is invalid."}' + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "model", ["us.anthropic.claude-sonnet-4-5-20250929-v1:0", "us.anthropic.claude-opus-4-8"], ids=["sonnet", "opus"] +) +@pytest.mark.parametrize( + ("upstream_status", "error_class", "status_code"), + [ + (400, litellm.BadRequestError, 400), + (429, litellm.RateLimitError, 429), + (500, litellm.ServiceUnavailableError, 503), + (503, litellm.ServiceUnavailableError, 503), + ], +) +def test_sync_converse_streaming_upstream_error_carries_the_upstream_body_as_text( + model: str, + upstream_status: int, + error_class: type[litellm.BadRequestError | litellm.RateLimitError | litellm.ServiceUnavailableError], + status_code: int, +) -> None: + client: Final = HTTPHandler( + client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(upstream_status, content=_UPSTREAM_ERROR_BODY.encode()) + ) + ) + ) + + with pytest.raises(error_class) as raised: + litellm.completion( + model=f"bedrock/converse/{model}", + messages=[{"role": "user", "content": "hi"}], + stream=True, + num_retries=0, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + client=client, + ) + + assert raised.value.status_code == status_code + assert raised.value.message == f"litellm.{error_class.__name__}: BedrockException - {_UPSTREAM_ERROR_BODY}" diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 6092934fc7a..f0dd40d4bc9 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -10581,3 +10581,151 @@ def test_map_openai_params_drops_the_openai_tool_search_tool_and_every_deferral_ ) assert [tool["function"]["name"] for tool in mapped["tools"]] == ["get_weather", "get_stock_price"] assert all("defer_loading" not in tool and "defer_loading" not in tool["function"] for tool in mapped["tools"]) + + +def _converse_params( + model: str, model_info: Mapping[str, object] | None, non_default_params: Mapping[str, object] +) -> Mapping[str, object]: + return AmazonConverseConfig().map_openai_params( + non_default_params=dict(non_default_params), + optional_params={}, + model=model, + drop_params=False, + model_info=model_info, + ) + + +_OUTPUT_CONFIG_FORMAT: Final = { + "output_config": {"format": {"type": "json_schema", "schema": CAPITAL_RESPONSE_FORMAT["json_schema"]["schema"]}} +} + + +@pytest.mark.parametrize( + "request_params", + [ + {"response_format": CAPITAL_RESPONSE_FORMAT, "thinking": ADAPTIVE_THINKING}, + {"response_format": CAPITAL_RESPONSE_FORMAT, "reasoning_effort": "low"}, + {"response_format": CAPITAL_RESPONSE_FORMAT, "thinking": {"type": "enabled", "budget_tokens": 2048}}, + _OUTPUT_CONFIG_FORMAT, + ], + ids=["adaptive", "reasoning_effort", "budget", "output_config_format"], +) +@pytest.mark.parametrize("model", ["us.anthropic.claude-opus-4-8", "global.anthropic.claude-sonnet-5"]) +def test_uses_opt_in_native_output_when_the_opt_in_drove_native_output(local_model_cost_map, model, request_params): + optional_params: Final = _converse_params(model, None, request_params) + + assert AmazonConverseConfig().uses_opt_in_native_output(model, optional_params, {"model_info": {"id": "d1"}}) + + +@pytest.mark.parametrize( + ("model", "model_info", "request_params"), + [ + ("us.anthropic.claude-opus-4-8", None, {"response_format": CAPITAL_RESPONSE_FORMAT}), + ( + "us.anthropic.claude-opus-4-8", + {"supports_bedrock_native_structured_output": False}, + {"response_format": CAPITAL_RESPONSE_FORMAT, "thinking": ADAPTIVE_THINKING}, + ), + ( + "us.anthropic.claude-sonnet-4-6", + {"supports_bedrock_native_structured_output": True}, + {"response_format": CAPITAL_RESPONSE_FORMAT, "thinking": ADAPTIVE_THINKING}, + ), + ], + ids=["no_thinking", "deployment_opt_out", "native_model"], +) +def test_uses_opt_in_native_output_is_false_when_native_output_did_not_come_from_the_opt_in( + local_model_cost_map, model, model_info, request_params +): + optional_params: Final = _converse_params(model, model_info, request_params) + + assert not AmazonConverseConfig().uses_opt_in_native_output(model, optional_params, {"model_info": model_info}) + + +def test_uses_opt_in_native_output_reads_the_base_model_of_an_application_inference_profile(local_model_cost_map): + profile: Final = "arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/profile-id" + litellm_params: Final = {"base_model": "bedrock/us.anthropic.claude-opus-4-8"} + optional_params: Final = _converse_params( + "us.anthropic.claude-opus-4-8", + None, + {"response_format": CAPITAL_RESPONSE_FORMAT, "thinking": ADAPTIVE_THINKING}, + ) + + assert AmazonConverseConfig().uses_opt_in_native_output(profile, optional_params, litellm_params) + + +_CALLER_OUTPUT_CONFIG_WITH_A_SCHEMA_OBJECT: Final = { + "textFormat": {"type": "json_schema", "structure": {"jsonSchema": {"name": "capital", "schema": {"type": "object"}}}} +} +_BEDROCK_SCHEMA_TYPE_ERROR: Final = "outputConfig.textFormat.structure.jsonSchema.schema must be a string" + + +@pytest.mark.parametrize("model", [_OPUS, _LLAMA], ids=["opt_in", "not_opt_in"]) +@pytest.mark.respx(assert_all_called=True) +def test_a_malformed_caller_output_config_gets_bedrocks_own_400( + model: str, respx_mock: respx.Router, local_model_cost_map: None +) -> None: + route: Final = respx_mock.post(url__regex=r"/converse$").mock( + return_value=httpx.Response(400, json={"message": _BEDROCK_SCHEMA_TYPE_ERROR}) + ) + + with pytest.raises(litellm.BadRequestError, match=re.escape(_BEDROCK_SCHEMA_TYPE_ERROR)): + completion( + model=f"bedrock/converse/{model}", + messages=[{"role": "user", "content": "What is the capital of France?"}], + response_format=CAPITAL_RESPONSE_FORMAT, + outputConfig=_CALLER_OUTPUT_CONFIG_WITH_A_SCHEMA_OBJECT, + num_retries=0, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + client=HTTPHandler(), + ) + + assert route.call_count == 1 + + +_CALLER_OUTPUT_CONFIG: Final = { + "textFormat": { + "type": "json_schema", + "structure": { + "jsonSchema": {"name": "capital", "schema": json.dumps(CAPITAL_RESPONSE_FORMAT["json_schema"]["schema"])} + }, + } +} +_BEDROCK_SCHEMA_REFUSAL: Final = "output_config.format.schema: pattern is too complex for structured output" + + +@pytest.mark.parametrize( + ("caller_params", "expected_calls"), + [ + ({"thinking": ADAPTIVE_THINKING}, 2), + ({"outputConfig": _CALLER_OUTPUT_CONFIG}, 1), + ({"outputConfig": _CALLER_OUTPUT_CONFIG, "thinking": ADAPTIVE_THINKING}, 1), + ], + ids=["litellm_native_output", "caller_output_config", "caller_output_config_thinking"], +) +@pytest.mark.respx(assert_all_called=True) +def test_a_refused_native_schema_is_sent_again_only_when_the_retry_drops_native_output( + caller_params: dict, expected_calls: int, respx_mock: respx.Router, local_model_cost_map: None +) -> None: + route: Final = respx_mock.post(url__regex=r"/converse$").mock( + return_value=httpx.Response(400, json={"message": _BEDROCK_SCHEMA_REFUSAL}) + ) + + with pytest.raises(litellm.BadRequestError, match=re.escape(_BEDROCK_SCHEMA_REFUSAL)): + completion( + model=f"bedrock/converse/{_OPUS}", + messages=[{"role": "user", "content": "What is the capital of France?"}], + response_format=CAPITAL_RESPONSE_FORMAT, + num_retries=0, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + client=HTTPHandler(), + **caller_params, + ) + + assert route.call_count == expected_calls + assert all("outputConfig" in json.loads(call.request.content) for call in route.calls[:1]) + assert all("outputConfig" not in json.loads(call.request.content) for call in route.calls[1:]) diff --git a/tests/unit/llms/bedrock/chat/test_native_output_fallback.py b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py new file mode 100644 index 00000000000..301fdaeb245 --- /dev/null +++ b/tests/unit/llms/bedrock/chat/test_native_output_fallback.py @@ -0,0 +1,123 @@ +import asyncio +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final + +import pytest + +from litellm.llms.bedrock.chat.native_output_fallback import NativeOutputAttemptResult, call_with_native_output_fallback +from litellm.llms.bedrock.common_utils import BedrockError, NativeOutputFallback +from litellm.types.utils import ModelResponse + +_MODEL: Final = "us.anthropic.claude-opus-4-8" +_FALLBACK: Final = NativeOutputFallback( + optional_params={"tools": "json_tool_call"}, litellm_params={"model_info": {"opted_out": True}} +) +_REJECTION: Final = BedrockError( + status_code=400, + message='{"message":"The model returned the following errors: output_config.format: Extra inputs are not permitted"}', +) +_SCHEMA_REFUSAL: Final = BedrockError( + status_code=400, message='{"message":"output_config.format.schema: pattern is too complex for structured output"}' +) +_OTHER_400: Final = BedrockError(status_code=400, message='{"message":"messages.0.content: Field required"}') + + +@dataclass(slots=True) +class ScriptedBedrock: + errors: tuple[BedrockError | None, ...] + asynchronous: bool + sent: tuple[NativeOutputFallback | None, ...] = () + + def _answer(self, error: BedrockError | None) -> ModelResponse: + if error is not None: + raise error + return ModelResponse(id=f"answer-{len(self.sent)}") + + async def _answer_later(self, error: BedrockError | None) -> ModelResponse: + await asyncio.sleep(0) + return self._answer(error) + + def __call__(self, fallback: NativeOutputFallback | None) -> NativeOutputAttemptResult: + self.sent = (*self.sent, fallback) + error: Final = self.errors[len(self.sent) - 1] + return self._answer_later(error) if self.asynchronous else self._answer(error) + + +def _resolve(result: NativeOutputAttemptResult) -> object: + if isinstance(result, Coroutine): + return asyncio.run(result) + return result + + +def _call(bedrock: ScriptedBedrock, fallback: NativeOutputFallback | None) -> object: + return _resolve(call_with_native_output_fallback(bedrock, None if fallback is None else lambda: fallback, _MODEL)) + + +_MODES: Final = pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) + + +@_MODES +def test_a_native_output_rejection_is_sent_again_once_with_the_fallback_params(asynchronous: bool) -> None: + bedrock: Final = ScriptedBedrock(errors=(_REJECTION, None), asynchronous=asynchronous) + + response: Final = _call(bedrock, _FALLBACK) + + assert isinstance(response, ModelResponse) and response.id == "answer-2" + assert bedrock.sent == (None, _FALLBACK) + + +@_MODES +def test_a_second_native_output_rejection_reaches_the_caller(asynchronous: bool) -> None: + bedrock: Final = ScriptedBedrock(errors=(_REJECTION, _REJECTION), asynchronous=asynchronous) + + with pytest.raises(BedrockError): + _call(bedrock, _FALLBACK) + + assert len(bedrock.sent) == 2 + + +@_MODES +@pytest.mark.parametrize( + ("error", "fallback"), + [(_OTHER_400, _FALLBACK), (_REJECTION, None)], + ids=["other_400", "no_fallback"], +) +def test_no_second_request_without_both_a_rejection_and_a_fallback( + asynchronous: bool, error: BedrockError, fallback: NativeOutputFallback | None +) -> None: + bedrock: Final = ScriptedBedrock(errors=(error,), asynchronous=asynchronous) + + with pytest.raises(BedrockError) as raised: + _call(bedrock, fallback) + + assert raised.value is error + assert bedrock.sent == (None,) + + +@_MODES +def test_a_rejection_reaches_the_caller_when_the_retry_would_send_native_output_again(asynchronous: bool) -> None: + bedrock: Final = ScriptedBedrock(errors=(_REJECTION,), asynchronous=asynchronous) + + with pytest.raises(BedrockError) as raised: + _resolve(call_with_native_output_fallback(bedrock, lambda: None, _MODEL)) + + assert raised.value is _REJECTION + assert bedrock.sent == (None,) + + +@pytest.mark.parametrize( + ("error", "advises_opting_out"), + [(_REJECTION, True), (_SCHEMA_REFUSAL, False)], + ids=["native_output_unsupported", "schema_refused"], +) +def test_a_rejection_logs_the_retry_and_advises_opting_out_only_when_native_output_is_unsupported( + error: BedrockError, advises_opting_out: bool, caplog: pytest.LogCaptureFixture +) -> None: + bedrock: Final = ScriptedBedrock(errors=(error, None), asynchronous=False) + + with caplog.at_level("WARNING", logger="LiteLLM"): + _call(bedrock, _FALLBACK) + + assert _MODEL in caplog.text + assert ("supports_bedrock_native_structured_output: false" in caplog.text) is advises_opting_out diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index b050bb7625b..d64a9ee845a 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -2047,6 +2047,31 @@ class TestBedrockBatchAnthropicRowParams: assert "response_format" not in model_input assert ("output_config" in model_input) == (response_format["type"] == "json_schema") + @pytest.mark.parametrize( + ("model", "native_output"), + [("bedrock/us.anthropic.claude-opus-4-8", False), (MODEL, True)], + ids=["opt_in_native_output_model", "native_output_model"], + ) + def test_thinking_json_schema_rows_use_native_output_only_on_models_that_always_take_it( + self, local_model_cost_map, model, native_output + ): + model_input: Final = self._transform( + "/v1/chat/completions", + { + "messages": [{"role": "user", "content": "Weather in Paris?"}], + "thinking": {"type": "adaptive"}, + "response_format": { + "type": "json_schema", + "json_schema": {"name": "weather", "schema": self.PARAMETERS}, + }, + }, + model=model, + ) + tool_names: Final = frozenset(tool["name"] for tool in model_input.get("tools", [])) + + assert ("format" in model_input.get("output_config", {})) is native_output + assert (tool_names == {"json_tool_call"}) is not native_output + def test_provider_native_params_still_pass_through(self): model_input = self._transform( "/v1/chat/completions", diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index eb6b4fdb063..6fe32785281 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -6,6 +6,7 @@ import os import struct import zlib from collections.abc import AsyncIterator, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from types import SimpleNamespace from typing import Final @@ -4009,3 +4010,162 @@ def test_bedrock_removed_tool_change_does_not_add_beta(explicit_beta: bool) -> N assert result["messages"] == [{"role": "user", "content": "Reply with OK"}] assert result.get("anthropic_beta", []).count(beta) == int(explicit_beta) + + +_NATIVE_OUTPUT_REJECTION: Final = { + "message": "The model returned the following errors: compilation exceeded limits: Schema is too complex." +} +_STALE_SIGNATURE_ERROR: Final = {"message": "messages.1.content.0: Invalid `signature` in `thinking` block"} +_OPTED_OUT: Final = {"supports_bedrock_native_structured_output": False} +_CONSTRAINED_SCHEMA: Final = { + "type": "object", + "properties": { + "summary": {"type": "string", "maxLength": 50}, + "kind": {"oneOf": [{"type": "string"}, {"type": "integer"}]}, + }, + "required": ["summary", "kind"], + "additionalProperties": False, +} +_ANTHROPIC_MESSAGE: Final = { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": [{"type": "text", "text": '{"summary": "sunny", "kind": 1}'}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} +_THINKING_HISTORY: Final = ( + {"role": "user", "content": "Describe the weather."}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "checking", "signature": "stale"}, + {"type": "text", "text": "Sunny."}, + ], + }, + {"role": "user", "content": "Summarise it as JSON."}, +) + + +@dataclass(slots=True) +class ScriptedBedrockInvoke: + replies: tuple[tuple[int, Mapping[str, object]], ...] + bodies: tuple[dict[str, object], ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.bodies = (*self.bodies, REASONING_BODY_ADAPTER.validate_json(request.content)) + status_code, body = self.replies[len(self.bodies) - 1] + return httpx.Response(status_code, json=body) + + +async def _send_bedrock_messages( + model: str, + model_info: Mapping[str, object] | None, + replies: Sequence[tuple[int, Mapping[str, object]]], + messages: Sequence[Mapping[str, object]] = ({"role": "user", "content": "Summarise the weather as JSON."},), +) -> tuple[ScriptedBedrockInvoke, Exception | None]: + bedrock: Final = ScriptedBedrockInvoke(replies=tuple(replies)) + try: + await litellm.anthropic_messages( + model=f"bedrock/invoke/{model}", + messages=[dict(message) for message in messages], + max_tokens=4000, + system="Answer in French.", + thinking={"type": "adaptive"}, + output_config={"format": {"type": "json_schema", "schema": _CONSTRAINED_SCHEMA}}, + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=AsyncHTTPHandler(transport=httpx.MockTransport(bedrock)), + **({} if model_info is None else {"model_info": model_info}), + ) + except Exception as error: + return bedrock, error + return bedrock, None + + +_OPT_IN_DEPLOYMENTS: Final = pytest.mark.parametrize( + ("model", "model_info"), + [ + ("us.anthropic.claude-opus-4-8", None), + ("us.anthropic.claude-opus-5-5", {"supports_bedrock_native_structured_output": True}), + ], + ids=["mapped_opt_in", "deployment_opt_in"], +) + + +@pytest.mark.usefixtures("local_model_cost_map") +@_OPT_IN_DEPLOYMENTS +async def test_bedrock_messages_native_output_rejection_resends_the_body_an_opted_out_deployment_sends( + model: str, model_info: Mapping[str, object] | None +) -> None: + opt_in, opt_in_error = await _send_bedrock_messages( + model, model_info, [(400, _NATIVE_OUTPUT_REJECTION), (200, _ANTHROPIC_MESSAGE)] + ) + opted_out, opted_out_error = await _send_bedrock_messages(model, _OPTED_OUT, [(200, _ANTHROPIC_MESSAGE)]) + + assert opt_in_error is None and opted_out_error is None + native, retried = opt_in.bodies + assert "format" in REASONING_BODY_ADAPTER.validate_python(native["output_config"]) + assert opted_out.bodies == (retried,) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("model_info", [None, _OPTED_OUT], ids=["native_output", "opted_out"]) +async def test_bedrock_messages_thinking_signature_400_that_persists_reaches_bedrock_twice_as_before( + model_info: Mapping[str, object] | None, +) -> None: + bedrock, error = await _send_bedrock_messages( + "us.anthropic.claude-opus-4-8", model_info, [(400, _STALE_SIGNATURE_ERROR)] * 3, _THINKING_HISTORY + ) + + assert isinstance(error, litellm.BadRequestError) + assert len(bedrock.bodies) == 2 + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "first_error, second_error", + [(_NATIVE_OUTPUT_REJECTION, _STALE_SIGNATURE_ERROR), (_STALE_SIGNATURE_ERROR, _NATIVE_OUTPUT_REJECTION)], + ids=["rejection_first", "signature_first"], +) +async def test_bedrock_messages_native_output_rejection_and_signature_400_each_get_one_retry_like_an_opted_out_deployment( + first_error: Mapping[str, object], second_error: Mapping[str, object] +) -> None: + opt_in, opt_in_error = await _send_bedrock_messages( + "us.anthropic.claude-opus-4-8", + None, + [(400, first_error), (400, second_error), (200, _ANTHROPIC_MESSAGE)], + _THINKING_HISTORY, + ) + opted_out, opted_out_error = await _send_bedrock_messages( + "us.anthropic.claude-opus-4-8", + _OPTED_OUT, + [(400, _STALE_SIGNATURE_ERROR), (200, _ANTHROPIC_MESSAGE)], + _THINKING_HISTORY, + ) + + assert opt_in_error is None and opted_out_error is None + assert len(opt_in.bodies) == 3 + assert opt_in.bodies[-1] == opted_out.bodies[-1] + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("model", "model_info", "error"), + [ + ("us.anthropic.claude-opus-4-8", None, {"message": "messages.0.content: Field required"}), + ("us.anthropic.claude-sonnet-4-6", None, _NATIVE_OUTPUT_REJECTION), + ("us.anthropic.claude-opus-4-8", _OPTED_OUT, _NATIVE_OUTPUT_REJECTION), + ], + ids=["other_400", "native_model", "deployment_opt_out"], +) +async def test_bedrock_messages_400_is_not_retried_unless_it_rejects_opt_in_native_output( + model: str, model_info: Mapping[str, object] | None, error: Mapping[str, object] +) -> None: + bedrock, raised = await _send_bedrock_messages(model, model_info, [(400, error), (200, _ANTHROPIC_MESSAGE)]) + + assert isinstance(raised, litellm.BadRequestError) + assert len(bedrock.bodies) == 1 diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 9aab5052cfe..d5bd7d44223 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -1,15 +1,20 @@ import asyncio, importlib, litellm, litellm.litellm_core_utils.get_model_cost_map as bedrock_govcloud_model_cost_map, pytest +from typing import Final from litellm.llms.bedrock.common_utils import( AmazonBedrockGlobalConfig, BedrockModelInfo, + bedrock_native_output_is_opt_in, + deployment_model_info, extract_model_name_from_bedrock_arn, get_bedrock_base_model, get_bedrock_cross_region_inference_regions, + is_bedrock_native_output_rejection, strip_bedrock_routing_prefix, strip_bedrock_throughput_suffix, + without_bedrock_native_output_opt_in, ) from collections.abc import Iterator from litellm import completion @@ -2054,3 +2059,70 @@ def test_bedrock_capability_model_reads_base_model_only_for_application_profile_ from litellm.llms.bedrock.common_utils import bedrock_capability_model assert bedrock_capability_model(model, base_model) == expected + + +_FORMAT_REJECTION: Final = "The model returned the following errors: output_config.format: Extra inputs are not permitted" +_COMPILE_REJECTION: Final = "The model returned the following errors: compilation exceeded limits: Schema is too complex." +_SCHEMA_PATTERN_REJECTION: Final = ( + "The model returned the following errors: output_config.format.schema: Unsupported regex feature in pattern " + "field: pattern is too complex for structured output: reduce the {n,m} upper bound, narrow the character class " + "range, or avoid nesting quantified groups" +) + + +@pytest.mark.parametrize( + ("status_code", "error_text", "expected"), + [ + (400, f'{{"message":"{_FORMAT_REJECTION}"}}', True), + (400, f'{{"message":"{_COMPILE_REJECTION}"}}', True), + (400, f'{{"message":"{_SCHEMA_PATTERN_REJECTION}"}}', True), + (400, "Malformed input request: #: extraneous key [output_config] is not permitted", False), + (400, "The model returned the following errors: messages.0.content: Field required", False), + (500, _FORMAT_REJECTION, False), + (429, _COMPILE_REJECTION, False), + ], + ids=[ + "format_not_permitted", + "schema_too_complex", + "schema_pattern_too_complex", + "other_key_not_permitted", + "other_400", + "500", + "429", + ], +) +def test_is_bedrock_native_output_rejection_matches_only_the_400s_that_reject_native_output( + status_code, error_text, expected +): + assert is_bedrock_native_output_rejection(status_code, error_text) is expected + + +@pytest.mark.parametrize( + ("model", "model_info", "expected"), + [ + ("us.anthropic.claude-opus-4-8", None, True), + ("global.anthropic.claude-sonnet-5", None, True), + ("us-gov.anthropic.claude-opus-4-8", None, False), + ("us.anthropic.claude-opus-4-8", {"supports_bedrock_native_structured_output": False}, False), + ("us.anthropic.claude-opus-5-5", {"supports_bedrock_native_structured_output": True}, True), + ("us.anthropic.claude-sonnet-4-6", {"supports_bedrock_native_structured_output": True}, False), + ], + ids=["opus_4_8", "sonnet_5", "govcloud_opus_4_8", "deployment_opt_out", "deployment_opt_in", "native_model"], +) +def test_bedrock_native_output_is_opt_in_only_where_the_opt_in_key_turns_native_output_on( + local_model_cost_map, model, model_info, expected +): + assert bedrock_native_output_is_opt_in(model, model_info=model_info) is expected + + +def test_without_bedrock_native_output_opt_in_turns_the_opt_in_off_and_keeps_the_deployment(local_model_cost_map): + litellm_params: Final = {"api_base": "https://bedrock.example", "model_info": {"id": "deployment-1"}} + + turned_off: Final = without_bedrock_native_output_opt_in(litellm_params) + + assert not bedrock_native_output_is_opt_in( + "us.anthropic.claude-opus-4-8", model_info=deployment_model_info(turned_off) + ) + assert turned_off["api_base"] == "https://bedrock.example" + assert deployment_model_info(turned_off)["id"] == "deployment-1" + assert litellm_params["model_info"] == {"id": "deployment-1"} diff --git a/tests/unit/types/test_completion.py b/tests/unit/types/test_completion.py index 28d13b7da4f..d6e5801bc0b 100644 --- a/tests/unit/types/test_completion.py +++ b/tests/unit/types/test_completion.py @@ -179,6 +179,7 @@ def _build_dispatch_context() -> CompletionDispatchContext: model="gpt-4o", model_response=None, # type: ignore[arg-type] optional_params={}, + optional_params_for_model_info=lambda model_info: {}, organization=None, provider_config=None, request_params={}, @@ -187,6 +188,7 @@ def _build_dispatch_context() -> CompletionDispatchContext: temperature=None, text_completion=False, timeout=None, + timeout_for_optional_params=lambda optional_params: 600.0, top_p=None, ) From 6c68c2cbef17d132ec86ebaa32290a328cef1670 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:08:20 -0700 Subject: [PATCH 58/64] fix(bedrock): add the July 7, 2026 EOL date to the Llama 3.1 405B and Llama 3.2 rows (#47196) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++++++++ model_prices_and_context_window.json | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 51a535736f6..3828957f2fb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24758,6 +24758,7 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.meta.llama3-2-1b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.3e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -24765,10 +24766,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.3e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-1b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "eu.meta.llama3-2-3b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.9e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -24776,6 +24779,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.9e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-3b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, @@ -40529,6 +40533,7 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "meta.llama3-1-405b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 2.4e-06, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -40571,6 +40576,7 @@ "supports_tool_choice": false }, "meta.llama3-2-11b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.6e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -40584,6 +40590,7 @@ "supports_vision": true }, "meta.llama3-2-1b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -40591,10 +40598,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-1b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "meta.llama3-2-3b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -40602,10 +40611,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-3b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "meta.llama3-2-90b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 7.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -48996,6 +49007,7 @@ "supports_vision": false }, "us.meta.llama3-1-405b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 2.4e-06, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -49038,6 +49050,7 @@ "supports_tool_choice": false }, "us.meta.llama3-2-11b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.6e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -49051,6 +49064,7 @@ "supports_vision": true }, "us.meta.llama3-2-1b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -49058,10 +49072,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-1b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "us.meta.llama3-2-3b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -49069,10 +49085,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-3b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "us.meta.llama3-2-90b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 7.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 51a535736f6..3828957f2fb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24758,6 +24758,7 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "eu.meta.llama3-2-1b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.3e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -24765,10 +24766,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.3e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-1b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "eu.meta.llama3-2-3b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.9e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -24776,6 +24779,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.9e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-3b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, @@ -40529,6 +40533,7 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "meta.llama3-1-405b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 2.4e-06, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -40571,6 +40576,7 @@ "supports_tool_choice": false }, "meta.llama3-2-11b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.6e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -40584,6 +40590,7 @@ "supports_vision": true }, "meta.llama3-2-1b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -40591,10 +40598,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-1b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "meta.llama3-2-3b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -40602,10 +40611,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-3b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "meta.llama3-2-90b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 7.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -48996,6 +49007,7 @@ "supports_vision": false }, "us.meta.llama3-1-405b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 2.4e-06, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -49038,6 +49050,7 @@ "supports_tool_choice": false }, "us.meta.llama3-2-11b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.6e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -49051,6 +49064,7 @@ "supports_vision": true }, "us.meta.llama3-2-1b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -49058,10 +49072,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-1b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "us.meta.llama3-2-3b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 1.5e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, @@ -49069,10 +49085,12 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-meta-llama-3-2-3b-instruct.md", "supports_function_calling": true, "supports_tool_choice": false }, "us.meta.llama3-2-90b-instruct-v1:0": { + "deprecation_date": "2026-07-07", "input_cost_per_token": 7.2e-07, "litellm_provider": "bedrock", "max_input_tokens": 128000, From 3fb95c8fe385ebe422d82d8cb10c9e52fac266fb Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sat, 17 Oct 2026 19:10:22 -0500 Subject: [PATCH 59/64] fix(batches): preserve spend across malformed output and metadata failures (#47184) * fix(batches): preserve recoverable spend and terminal reconciliation state * fix(batches): narrow validated envelope types * fix(batches): serialize stored metadata when releasing claims * fix(batches): sanitize metadata before claiming completed batches * refactor(batches): expose existing metadata serializer for shared use --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 68 ++++--- .../proxy/hooks/managed_files.py | 8 +- litellm/batches/batch_utils.py | 16 +- tests/unit/batches/test_batch_utils.py | 28 +++ .../proxy/test_managed_files_access_check.py | 21 +- .../common_utils/test_check_batch_cost.py | 187 +++++++++++------- 6 files changed, 209 insertions(+), 119 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 94c3674945d..2d2bc92dbd2 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,6 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t """ import asyncio +import json from collections.abc import Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone @@ -227,7 +228,7 @@ class CheckBatchCost: "so it will no longer be polled" ) - async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool: + async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool: """ Atomically flip batch_processed from false to true, returning whether this pod won the row. Every pod and uvicorn worker schedules its own poller against the shared @@ -241,10 +242,16 @@ class CheckBatchCost: """ if not self._has_batch_processed_column: return True + from litellm_enterprise.proxy.hooks.managed_files import postgres_storable_json + try: claimed: Final = await _managed_object_table(self.prisma_client).update_many( where={"id": job.id, "batch_processed": False}, - data={"batch_processed": True}, + data={ + "batch_processed": True, + "status": response.status if response.status != "completed" else "complete", + "file_object": postgres_storable_json(response), + }, ) except Exception as db_err: verbose_proxy_logger.error( @@ -253,7 +260,9 @@ class CheckBatchCost: return False return claimed > 0 - async def _release_job_claim(self, job: "ManagedObjectRow") -> None: + async def _release_job_claim( + self, job: "ManagedObjectRow", original_status: str | None, original_file_object: object + ) -> None: """Give a claimed row back once billing it failed, so a later poll cycle retries it. Safe to match on batch_processed=True: while this poller is active the retrieve @@ -265,7 +274,11 @@ class CheckBatchCost: try: await _managed_object_table(self.prisma_client).update_many( where={"id": job.id, "batch_processed": True}, - data={"batch_processed": False}, + data={ + "batch_processed": False, + "status": original_status, + "file_object": json.dumps(original_file_object), + }, ) except Exception as db_err: verbose_proxy_logger.error( @@ -564,8 +577,6 @@ class CheckBatchCost: @staticmethod def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]: - import json - from litellm.types.utils import LiteLLMBatch file_object: Final = job.file_object @@ -718,8 +729,7 @@ class CheckBatchCost: ) except Exception as _e: verbose_proxy_logger.warning( - f"CheckBatchCost: failed to create managed file ID for " - f"{_file_attr}={_raw_file_id!r}: {_e}" + f"CheckBatchCost: failed to create managed file ID for {_file_attr}={_raw_file_id!r}: {_e}" ) # Pass the deployment's router-registered pricing (litellm_params custom @@ -788,7 +798,9 @@ class CheckBatchCost: custom_llm_provider=str(llm_provider) if llm_provider else None, ) - if not await self._claim_job_for_costing(job): + original_status: Final = job.status + original_file_object: Final = job.file_object + if not await self._claim_job_for_costing(job, response.model_copy(update={"usage": batch_result.usage})): verbose_proxy_logger.info( f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod " "in this window, so its cost is already being tracked there" @@ -807,7 +819,7 @@ class CheckBatchCost: batch_completion_cost=batch_result.completion_cost, ) except Exception: - await self._release_job_claim(job) + await self._release_job_claim(job, original_status, original_file_object) raise # Record batch duration (completed_at - created_at) @@ -1100,10 +1112,7 @@ class CheckBatchCost: continue ## RETRIEVE THE BATCH JOB OUTPUT FILE - if ( - response.status in PROVIDER_TERMINAL_BATCH_STATUSES - and response.output_file_id is not None - ): + if response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None: try: tracked = await self._track_completed_batch_cost( job=job, @@ -1134,24 +1143,19 @@ class CheckBatchCost: # Track this job for the final metrics summary processed_models.append(tracked) - # mark the job as complete - try: - update_data: dict = { - "status": response.status if response.status != "completed" else "complete", - "file_object": response.model_dump_json(), - } - if self._has_batch_processed_column: - update_data["batch_processed"] = True - await _managed_object_table(self.prisma_client).update( - where={"id": job.id}, - data=update_data, - ) - except Exception as db_err: - verbose_proxy_logger.error( - f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" - ) - if not self._has_batch_processed_column: - # nothing claimed the row, so it stays in the poll page and its lines wait for that retry + if not self._has_batch_processed_column: + try: + await _managed_object_table(self.prisma_client).update( + where={"id": job.id}, + data={ + "status": response.status if response.status != "completed" else "complete", + "file_object": response.model_dump_json(), + }, + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" + ) continue self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index a33b33662ab..439770390f8 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -416,7 +416,7 @@ def _storable_json_value(value: JsonValue) -> JsonValue: return value -def _postgres_storable_json(file_object: BaseModel) -> str: +def postgres_storable_json(file_object: BaseModel) -> str: dumped: Final = file_object.model_dump_json() if "\\u0000" not in dumped: return dumped @@ -650,7 +650,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) # FIX: Update status and file_object on every operation to keep state in sync update_columns: Final = { - "file_object": _postgres_storable_json(stored_object), + "file_object": postgres_storable_json(stored_object), "status": file_object.status, "updated_by": user_api_key_dict.user_id, } @@ -666,7 +666,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): data={ "create": { "unified_object_id": unified_object_id, - "file_object": _postgres_storable_json(stored_object), + "file_object": postgres_storable_json(stored_object), "model_object_id": model_object_id, "file_purpose": file_purpose, "created_by": resolve_resource_owner_id(user_api_key_dict), @@ -702,7 +702,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "unified_object_id": response_id, "model_object_id": decrypted_response_id(response_id), "file_purpose": "response", - "file_object": _postgres_storable_json(_without_model_output(response)), + "file_object": postgres_storable_json(_without_model_output(response)), "status": PENDING_RESPONSE_STATUSES[0], "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index a50cdf75d33..45b0eaa0b57 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -794,11 +794,17 @@ def _batch_response_was_successful( message batch results lines report ``result.type == "succeeded"``; Bedrock batch output lines report ``modelOutput`` (and no ``error``). """ - if custom_llm_provider == "anthropic": - return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded" if custom_llm_provider == "bedrock": return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None - _response: Final[dict[str, object]] = cast( # cast-ok: batch response bodies come from provider JSON - dict[str, object], batch_job_output_file.get("response", None) or {} + envelope_key: Final = "result" if custom_llm_provider == "anthropic" else "response" + envelope: Final = batch_job_output_file.get(envelope_key) + if not isinstance(envelope, Mapping): + if envelope is not None: + verbose_logger.warning("batch output %s must be an object, treating the line as unsuccessful", envelope_key) + return False + response: Final = cast(Mapping[object, object], envelope) # cast-ok: mapping shape validated above + return ( + response.get("type") == "succeeded" + if custom_llm_provider == "anthropic" + else response.get("status_code") == 200 ) - return _response.get("status_code", None) == 200 diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index 912d4cfe84c..630b00e172b 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable. import asyncio, json, time import logging from types import MappingProxyType +from typing import Final, Literal import httpx import pytest @@ -3347,3 +3348,30 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): assert mock_batch._hidden_params["batch_successful_requests"] == 8 assert mock_batch._hidden_params["batch_failed_requests"] == 0 assert mock_batch.usage == expected_usage + + +@pytest.mark.parametrize("provider", ("openai", "anthropic")) +@pytest.mark.parametrize("envelope", (1, "invalid", [1], None, {})) +@pytest.mark.parametrize("include_valid", (True, False)) +def test_invalid_batch_envelopes_preserve_recoverable_cost( + provider: Literal["openai", "anthropic"], envelope: object, include_valid: bool +) -> None: + invalid: Final[dict[str, object]] = {"response" if provider == "openai" else "result": envelope} + valid: Final[dict[str, object]] = ( + _success_row(model="batch-accounting-control", usage=_usage(100, 20)) + if provider == "openai" + else { + "result": { + "type": "succeeded", + "message": {"model": "batch-accounting-control", "usage": {"input_tokens": 100, "output_tokens": 20}}, + } + } + ) + entries: Final = (valid, invalid, valid) if include_valid else (invalid,) + result: Final = bu._aggregate_batch_cost_usage_models( + entries, provider, model_info={"input_cost_per_token_batches": 0.001, "output_cost_per_token_batches": 0.002} + ) + assert result.cost == pytest.approx(0.28 if include_valid else 0) + assert result.usage.prompt_tokens == (200 if include_valid else 0) + assert result.usage.completion_tokens == (40 if include_valid else 0) + assert (result.successful_requests, result.failed_requests) == (2 if include_valid else 0, 1) diff --git a/tests/unit/enterprise/proxy/test_managed_files_access_check.py b/tests/unit/enterprise/proxy/test_managed_files_access_check.py index 62302446a29..0cdc212502c 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/unit/enterprise/proxy/test_managed_files_access_check.py @@ -9,6 +9,8 @@ with deployment credentials, bypassing the managed files access-control hooks. """ import base64 +import json +from typing import Final import pytest from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -367,13 +369,14 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti # managed_files_obj.afile_content should NOT have been called mock_managed_files_hook.afile_content.assert_not_called() - # Verify the DB update writes batch_processed, status, and file_object - mock_prisma.db.litellm_managedobjecttable.update.assert_called_once() - update_call_kwargs = ( - mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs + claim_data: Final = next( + call.kwargs["data"] + for call in mock_prisma.db.litellm_managedobjecttable.update_many.call_args_list + if call.kwargs["where"].get("id") == mock_job.id ) - assert update_call_kwargs["data"]["batch_processed"] is True - assert update_call_kwargs["data"]["status"] == "complete" - assert ( - "file_object" in update_call_kwargs["data"] - ), "file_object must be written to DB so list_batches reads updated status" + assert claim_data["batch_processed"] is True + assert claim_data["status"] == "complete" + persisted_batch: Final = json.loads(claim_data["file_object"]) + assert persisted_batch["status"] == "completed" + assert persisted_batch["usage"]["prompt_tokens"] == 10 + assert persisted_batch["usage"]["completion_tokens"] == 5 diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index c974e55c0c9..eb2a5b7f66c 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -17,6 +17,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from litellm.types.utils import LiteLLMBatch + if TYPE_CHECKING: from litellm.batches.batch_utils import BatchCostUsageResult @@ -668,10 +670,16 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) - mock_response = MagicMock() - mock_response.status = "completed" - mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response: Final = LiteLLMBatch( + id="batch-1", + status="completed", + output_file_id="file-output-123", + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) @@ -730,10 +738,10 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( "Expected update() to be called exactly once for the completed job" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True, ( "update() must include batch_processed=True when column is present" ) @@ -863,7 +871,7 @@ class TestCheckBatchCost: mock_update_database.assert_awaited_once() assert mock_update_database.call_args.kwargs["response_cost"] == 0.01 - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( "the job must still be marked processed once cost tracking succeeds" ) @@ -1285,11 +1293,16 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) - mock_response = MagicMock() - mock_response.status = terminal_status - mock_response.output_file_id = "file-output-123" - mock_response.error_file_id = None - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' + mock_response: Final = LiteLLMBatch( + id="batch-1", + status=terminal_status, + output_file_id="file-output-123", + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) @@ -1352,8 +1365,8 @@ class TestCheckBatchCost: f"{terminal_status} batch with an output file must fetch results and be billed" ) mock_logging_obj.async_success_handler.assert_awaited_once() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3 + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == terminal_status, ( f"billed {terminal_status} batch must keep its real terminal status in the DB" @@ -1847,13 +1860,16 @@ class TestUnmanagedVertexRouting: router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash" router.get_model_ids.return_value = ["deploy-1"] - mock_response = MagicMock() - mock_response.status = "completed" - mock_response.output_file_id = "gs://bucket/out/predictions.jsonl" - mock_response.error_file_id = None - mock_response.completed_at = None - mock_response.created_at = None - mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' + mock_response: Final = LiteLLMBatch( + id="8823717160934178816", + status="completed", + output_file_id="gs://bucket/out/predictions.jsonl", + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1921,8 +1937,8 @@ class TestUnmanagedVertexRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01 - assert prisma.db.litellm_managedobjecttable.update.call_count == 1 - update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2065,13 +2081,16 @@ class TestUnmanagedBedrockRouting: router.resolve_model_name_from_model_id.return_value = "claude-sonnet-4" router.get_model_ids.return_value = ["deploy-1"] - mock_response = MagicMock() - mock_response.status = "completed" - mock_response.output_file_id = "s3://bucket/out/predictions.jsonl" - mock_response.error_file_id = None - mock_response.completed_at = None - mock_response.created_at = None - mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' + mock_response: Final = LiteLLMBatch( + id=self._ARN, + status="completed", + output_file_id="s3://bucket/out/predictions.jsonl", + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) @@ -2135,8 +2154,8 @@ class TestUnmanagedBedrockRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02 - assert prisma.db.litellm_managedobjecttable.update.call_count == 1 - update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -3003,13 +3022,20 @@ class TestMultiPodBatchCostClaim: @staticmethod def _router(): - response = MagicMock() - response.status = "completed" - response.output_file_id = _CLAIM_OUTPUT_FILE_ID - response.error_file_id = None - response.created_at = 1 - response.completed_at = 2 - response.model_dump_json.return_value = '{"id":"batch-456","status":"completed"}' + from litellm.types.utils import LiteLLMBatch + + response: Final = LiteLLMBatch( + id="batch-456", + status="completed", + output_file_id=_CLAIM_OUTPUT_FILE_ID, + error_file_id=None, + created_at=1, + completed_at=2, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "openai" @@ -3089,8 +3115,9 @@ class TestMultiPodBatchCostClaim: ("openai/batch-test-model", 2, 0, "batch-test-model"), ), ) + @pytest.mark.parametrize("failure_envelope", (None, 1, "invalid", [1])) async def test_batch_accounting_preserves_models_and_mixed_results( - self, deployment_model: str, successes: int, failures: int, expected_model: str + self, deployment_model: str, successes: int, failures: int, expected_model: str, failure_envelope: object ) -> None: import base64 from datetime import datetime @@ -3125,7 +3152,7 @@ class TestMultiPodBatchCostClaim: }, "error": None, } - failure: Final = {"response": None, "error": {"code": "batch_expired", "message": "expired"}} + failure: Final = {"response": failure_envelope, "error": {"code": "batch_expired", "message": "expired"}} rows: Final = (success,) * successes + (failure,) * failures content: Final = b"\n".join(json.dumps(row).encode() for row in rows) @@ -3217,6 +3244,30 @@ class TestMultiPodBatchCostClaim: with patch("litellm.proxy.proxy_server.scheduler", scheduler): await guard._check_file_deletion_allowed(file_id) + @pytest.mark.asyncio + @pytest.mark.parametrize("status", ("completed", "failed", "expired", "cancelled")) + @pytest.mark.parametrize("metadata_text", ("before\x00after", r"\u0000")) + async def test_cost_claim_persists_terminal_metadata_without_a_later_write( + self, status: str, metadata_text: str + ) -> None: + row: Final = _FakeManagedObjectRow() + journal: Final = [] + prisma: Final = self._prisma(row, journal) + router: Final = self._router() + router.aretrieve_batch.return_value = router.aretrieve_batch.return_value.model_copy( + update={"status": status, "metadata": {"annotation": metadata_text}} + ) + prisma.db.litellm_managedobjecttable.update.side_effect = RuntimeError("final write unavailable") + worker: Final = self._instance(prisma, router) + with self._billing_patches(journal) as logging_obj: + await worker.check_batch_cost() + await worker.check_batch_cost() + assert row.batch_processed is True + assert row.status == ("complete" if status == "completed" else status) + assert json.loads(row.file_object)["status"] == status + assert json.loads(row.file_object)["metadata"] == {"annotation": metadata_text.replace("\x00", "")} + assert logging_obj.async_success_handler.await_count == 1 + @pytest.mark.asyncio async def test_winning_pod_claims_the_row_between_fetching_and_billing(self): """The claim flips batch_processed false -> true after the results are in hand and @@ -3228,13 +3279,9 @@ class TestMultiPodBatchCostClaim: with self._billing_patches(journal) as logging_obj: await self._instance(prisma, self._router()).check_batch_cost() - assert journal == ["fetch", "claim", "bill", "finalize"] - assert self._claim_calls(prisma) == [ - { - "where": {"id": "job-claim-1", "batch_processed": False}, - "data": {"batch_processed": True}, - } - ] + assert journal == ["fetch", "claim", "bill"] + assert row.status == "complete" + assert json.loads(row.file_object)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5} logging_obj.async_success_handler.assert_awaited_once() assert row.batch_processed is True @@ -3254,18 +3301,17 @@ class TestMultiPodBatchCostClaim: assert journal == ["fetch"] logging_obj.async_success_handler.assert_not_awaited() - assert self._claim_calls(prisma) == [ - { - "where": {"id": "job-claim-1", "batch_processed": False}, - "data": {"batch_processed": True}, - } - ] + assert row.status == "in_progress" + assert json.loads(row.file_object)["status"] == "in_progress" @pytest.mark.asyncio - async def test_a_failed_spend_log_write_releases_the_claim(self): + @pytest.mark.parametrize("stored_as_json_string", (True, False)) + async def test_a_failed_spend_log_write_releases_the_claim(self, stored_as_json_string: bool) -> None: """A transient failure while billing a claimed batch must hand the row back, or its spend is silently lost instead of being retried on the next cycle.""" row = _FakeManagedObjectRow() + original_file_object: Final = row.file_object if stored_as_json_string else json.loads(row.file_object) + row.file_object = original_file_object journal = [] prisma = self._prisma(row, journal) @@ -3274,10 +3320,16 @@ class TestMultiPodBatchCostClaim: assert journal == ["fetch", "claim", "bill", "release"] assert row.batch_processed is False - assert self._claim_calls(prisma)[-1] == { - "where": {"id": "job-claim-1", "batch_processed": True}, - "data": {"batch_processed": False}, - } + assert row.status == "in_progress" + assert json.loads(row.file_object) == original_file_object + + with self._billing_patches(journal) as logging_obj: + await self._instance(prisma, self._router()).check_batch_cost() + await self._instance(prisma, self._router()).check_batch_cost() + + assert row.batch_processed is True + assert row.status == "complete" + logging_obj.async_success_handler.assert_awaited_once() @pytest.mark.asyncio async def test_a_worker_interrupted_mid_costing_leaves_the_batch_billable(self): @@ -3309,7 +3361,7 @@ class TestMultiPodBatchCostClaim: with self._billing_patches(survivor_journal) as survivor_logging: await self._instance(survivor_prisma, self._router()).check_batch_cost() - assert survivor_journal == ["fetch", "claim", "bill", "finalize"] + assert survivor_journal == ["fetch", "claim", "bill"] survivor_logging.async_success_handler.assert_awaited_once() assert row.batch_processed is True @@ -3340,7 +3392,7 @@ class TestMultiPodBatchCostClaim: finish_fetch.set() await asyncio.wait_for(costing, timeout=5) - assert journal == ["fetch", "claim", "bill", "finalize"] + assert journal == ["fetch", "claim", "bill"] assert row.batch_processed is True await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) @@ -3465,7 +3517,6 @@ class TestBatchLineItemsFromThePoller: f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", - "finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines", @@ -3539,7 +3590,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal[-4:] == ["finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines"] + assert journal[-4:] == ["bill", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines"] ((line,),) = [payload["lines"] for payload in logger.payloads] assert line["request"] == json.loads(self._INPUT) assert line["response"] == json.loads(self._OUTPUT) @@ -3558,7 +3609,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "finalize", "lines"] + assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "lines"] assert row.batch_processed is True assert row.status == "complete" @@ -3574,7 +3625,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"] @pytest.mark.asyncio async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is_retired(self, monkeypatch): @@ -3751,7 +3802,6 @@ class TestBatchLineItemsFromThePoller: f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", - "finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines started", @@ -3834,7 +3884,6 @@ class TestBatchLineItemsFromThePoller: f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", - "finalize", "store batch-earlier", "store batch-456", ] @@ -3863,7 +3912,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"] assert logger.payloads == [] assert (row.batch_processed, row.status) == (True, "complete") assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text @@ -3887,7 +3936,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"] assert logger.payloads == [] assert (row.batch_processed, row.status) == (True, "complete") assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text From 75faaeab641e161a3d6bc38a42fcaa71c5602d90 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:11:27 +0000 Subject: [PATCH 60/64] feat(count_tokens)!: configurable token_count_framing and token_count_model, count source in the response (#46777) * feat(count_tokens): configurable token_count_framing and token_count_fallback_model, count source in the response Count endpoints now say where each number came from in litellm_count_source (provider, fallback_model, local_estimate) and name the fallback model in litellm_count_fallback_model. When the provider answers 400 or 404 for a model it cannot count, the proxy counts the same request with token_count_fallback_model. token_count_framing adds a per-model constant to local estimates. Both settings ship in the cost map for Bedrock Claude Sonnet/Haiku 5.5, Bedrock Kimi K3, and Databricks Kimi K3 and DeepSeek V4 Flash, and can be overridden in model_info. * test(count_tokens): classify the new cost map keys in the schema generator and expect count_source in exact-shape tests * test(count_tokens): e2e matrix against real inference and integration wire tests for count source, fallback and framing * fix(count_tokens): keep the provider error when local counting is disabled, name the fallback error, keep settings in cooldown Moves the cost map lookup off the event loop * ci(cost map): bound token_count_framing to non-negative and token_count_fallback_model to non-empty * fix(count_tokens): keep a failed fallback's own status, ship fallbacks only where they count, add system.ai framing A fallback model that answers 404 returns the requested model's error unchanged, and any other fallback failure answers with its own status and message. The eu, apac, au and jp Sonnet and Haiku 5.5 entries no longer ship the Opus 5.5 fallback, which AWS serves only in us-east-1. A configured empty token_count_fallback_model turns the shipped fallback off. Config framing accepts whole-number floats and rejects negatives. The Databricks system.ai Kimi K3 and DeepSeek V4 Flash names carry the same framing as their serving endpoint names, and the e2e matrix runs on Bedrock Haiku 5.5 too. * test(databricks): list the system.ai Kimi K3 and DeepSeek V4 Flash names with their published DBU rates * fix(count_tokens): skip a token count fallback model for an hour after it answers 404 on that deployment A us. or global. Sonnet 5.5 deployment outside us-east-1 paid two extra Bedrock calls and about 200 ms on every count, since Opus 5.5 cannot count there either. The Kimi K3 e2e also runs the large system prompt shape now. * test(count_tokens): use the shared integration port allocator in the fallback wire test * fix(cost map): keep the Databricks system.ai entries to framing only, so their spend stays as today * test(count_tokens): cover the fallback 404 skip, inference profile base_model, shipped framing, and deployment-scoped token count settings Adds integration cells for a fallback model that answers 404, an application inference profile ARN deployment whose base_model carries the shipped fallback, and the shipped Kimi K3 framing with token_count_framing 0 turning it off. Adds unit tests that keep token count settings off the shared backend cost map key and that store a token count setting on PATCH only when the operator changed it * fix(count_tokens): keep the counter's own 400 and the provider's error when a framed estimate cannot read the body A 400 the OpenAI or Gemini counter returns before calling the provider is the caller's answer, so it no longer tries the fallback model or a framed estimate. When the framed local estimate cannot read a content block, the client gets the provider's error instead of a 500. Adds /audit round 2's integration tests for the fallback leg's converted body, the Anthropic envelope on fallback failures, and both framing cases. * fix(count_tokens): answer the provider's error whenever a framed estimate raises A tool schema the local estimate cannot format, such as a boolean subschema or an array with no items, raised AttributeError or KeyError and answered a 500 on a framed deployment whose provider had already answered 400. Any estimate failure now answers the provider's own error. * fix(count_tokens): load the custom tokenizer inside the guarded framed estimate A framed deployment whose custom_tokenizer could not load answered a 500 or 503 where the provider had already answered 400, because the tokenizer was selected outside the estimate's error handling. Selecting it inside answers the provider's error, logs a warning with the estimate's error, and still raises the Rust extension's fork-guard errors. * fix(count_tokens): look up framing under the deployment's custom_llm_provider * fix(count_tokens): share one tokenizer download per worker and bound the wait after a provider error A framed deployment whose custom_tokenizer is not cached locally made every count wait out the Hugging Face hub's whole retry cycle when the hub could not be reached, and a burst of those loads held the default thread limiter that every count's cost map lookup also needs. Counts on a worker now share one download per minute, a failed download is reused for the rest of that minute, and after a provider error a count waits at most 5 s before answering the provider's error. The swallowed-estimate warning logs only the error type. * fix(count_tokens): keep one tokenizer download per key running and load the local route per request * test(count_tokens): type the scripted counter and wait on events instead of wall time where possible * fix(count_tokens): run framed tokenizer loads on their own capped threads and answer a tokenizer panic with the provider error Framed loads used anyio's default 40 threads, so a hub that drops packets let about 40 hanging loads stall every count on the worker. Loads now run on a per-worker pool of 4, and a count that finds 4 loads running answers the provider's error at once. Any finished load is reused for a minute after it finishes, so a load that fell back to tiktoken is not retried on every count. A pyo3 PanicException from the tokenizer, while loading or counting, now answers the provider's error instead of a plain 500 * feat(proxy): count tokens with token_count_model in its pinned region token_count_model {model, region} replaces token_count_fallback_model. Bedrock Sonnet and Haiku 5.5 count with Opus 5.5 in us-east-1, or Melbourne for au. profiles, and a failure names the region to enable * fix(count_tokens): use the renamed count model constant in the e2e test and list token_count_model with the object keys * fix(proxy): keep the upstream error envelope when adding the region hint * test(proxy): annotate the token count setting echo test * ci(cost map): require non-blank token_count_model strings and test the region rules * fix(count_tokens): answer the token_count_model's error even when framing is set * fix(count_tokens): keep the provider error for Google contents instead of a framed estimate * fix(count_tokens): estimate Google contents the route translated to messages * fix(count_tokens): keep the provider error for Google contents with non-text parts * test(proxy): drive tokenizer wait and cost map tests by events and a skippable loop clock * fix(count_tokens): check Google text-only parts without a KeysView comparison basedpyright flags * feat(proxy): count Bedrock Kimi K3 tokens exactly with its pinned tokenizer * test(bedrock): keep the Kimi K3 image check in-process * fix(count_tokens): join the region hint to an upstream message that already ends with a period Also pins that a blank token_count_model in config falls through to the cost map default * fix(count_tokens): count Gemini tool-call turns and Responses tools, and stop tokenizer retries from stalling counts The Google countTokens route estimates contents with functionCall, functionResponse and thoughtSignature parts again, and still answers the provider error for media, thought and unknown parts The Kimi K3 exact count renders Responses API and Anthropic tools through the chat tool translation the Vertex count already used, now shared from the prompt factory, and falls back to the estimate when the Converse body cannot carry a tool A failed Kimi K3 tokenizer load is retried in the background without making counts wait, and a load still running after the wait logs a warning * fix(count_tokens): count Kimi K3 tools the way Bedrock inference sends them A flat Responses API function or custom tool becomes the chat tool the Responses route sends to inference, and every other tool goes to the Converse body as is, so hosted tools Bedrock drops are left out of the count instead of being counted or turning it into an estimate. * test(count_tokens): annotate new token counter test locals with Final * fix(count_tokens): keep shipped count regions inside the deployment's geography A plain Bedrock Sonnet or Haiku 5.5 deployment in a non-US region no longer sends its count prompt to us-east-1. Global count models and configured token_count_model keep their pinned region. Also expect litellm_count_source in the Vertex Llama count test main added * fix(count_tokens): resolve count settings off the event loop --- .../prompt_templates/factory.py | 22 + litellm/llms/bedrock/count_tokens/handler.py | 2 +- litellm/llms/bedrock/count_tokens/kimi_k3.py | 502 ++++++ .../bedrock/count_tokens/mantle_handler.py | 2 +- .../responses/count_tokens/token_counter.py | 1 + .../llms/vertex_ai/count_tokens/context.py | 25 +- ...odel_prices_and_context_window_backup.json | 52 + .../proxy/anthropic_endpoints/endpoints.py | 3 +- .../common_utils/token_count_settings.py | 164 ++ litellm/proxy/google_endpoints/endpoints.py | 11 +- litellm/proxy/proxy_server.py | 461 +++++- .../proxy/response_api_endpoints/endpoints.py | 5 +- litellm/types/llms/vertex_ai.py | 6 +- litellm/types/token_count.py | 16 + litellm/types/utils.py | 9 + litellm/utils.py | 3 + model_prices_and_context_window.json | 52 + model_prices_and_context_window.schema.json | 30 + .../test_count_tokens_source_e2e.py | 362 +++++ tests/e2e/models.py | 34 +- tests/e2e/proxy_client.py | 11 + .../test_anthropic_count_tokens_wire.py | 4 +- .../anthropic/test_anthropic_gateway_wire.py | 5 +- ...anthropic_count_tokens_system_lift_wire.py | 24 +- .../test_bedrock_mantle_count_tokens_wire.py | 89 +- .../providers/test_count_tokens_model_wire.py | 719 +++++++++ .../providers/test_malformed_tools_wire.py | 10 +- ...onses_bridge_message_cache_control_wire.py | 2 + .../test_vertex_partner_count_tokens_wire.py | 28 +- .../integration/spend/test_cache_and_quota.py | 8 +- .../count_tokens/test_token_counter.py | 5 + .../count_tokens/test_bedrock_kimi_k3.py | 524 +++++++ .../count_tokens/test_token_counter.py | 5 + .../common_utils/test_token_count_settings.py | 209 +++ .../test_google_api_endpoints.py | 42 +- .../test_model_management_endpoints.py | 33 + .../proxy/proxy_server/test_routes_utils.py | 1 + tests/unit/proxy/test_proxy_token_counter.py | 1352 ++++++++++++++++- .../unit/test_router_model_cost_isolation.py | 38 +- tests/unit/test_utils.py | 7 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +- 41 files changed, 4732 insertions(+), 158 deletions(-) create mode 100644 litellm/llms/bedrock/count_tokens/kimi_k3.py create mode 100644 litellm/proxy/common_utils/token_count_settings.py create mode 100644 litellm/types/token_count.py create mode 100644 tests/e2e/llm_translation/test_count_tokens_source_e2e.py create mode 100644 tests/integration/providers/test_count_tokens_model_wire.py create mode 100644 tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py create mode 100644 tests/unit/proxy/common_utils/test_token_count_settings.py diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0e9206b422e..0549b0af11e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -6166,3 +6166,25 @@ def resolve_structured_messages( if structured: return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured] return None + + +def is_flat_responses_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return tool_type in ("function", "custom") and tool_type not in tool and "input_schema" not in tool + + +def chat_completion_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + if is_flat_responses_tool(tool): + chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in a flat Responses tool shape + ) + return tuple(chat_tools) + anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai( + [dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in an Anthropic or chat tool shape + ) + return tuple(anthropic_tools) diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index be8243e48ef..3bae1f22490 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -48,7 +48,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): verbose_logger.debug("Processing CountTokens request for resolved model: %s", resolved_model) # Get AWS region using existing LiteLLM function - aws_region_name: Final = self._get_aws_region_name( + aws_region_name: Final = self.get_aws_region_name( optional_params=litellm_params, model=resolved_model, model_id=None, diff --git a/litellm/llms/bedrock/count_tokens/kimi_k3.py b/litellm/llms/bedrock/count_tokens/kimi_k3.py new file mode 100644 index 00000000000..5e3fd56e4ba --- /dev/null +++ b/litellm/llms/bedrock/count_tokens/kimi_k3.py @@ -0,0 +1,502 @@ +""" +Exact Bedrock token counts for Kimi K3, which has no CountTokens support. + +Renders the Converse body litellm sends through Moonshot's Kimi K3 chat template and counts it with the +official tokenizer, downloaded once from Hugging Face at a pinned revision and checked against its SHA-256. +""" + +import asyncio +import base64 +import copy +import hashlib +import json +import threading +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from concurrent.futures import Future +from dataclasses import dataclass +from enum import Enum +from itertools import chain, groupby +from pathlib import Path +from typing import Annotated, Final, Literal, TypeAlias + +import regex +import tiktoken +from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator +from typing_extensions import Self + +import litellm +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.factory import ( + chat_completion_tools_for, + is_flat_responses_tool, +) +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam + +KIMI_K3_REPO: Final = "moonshotai/Kimi-K3" +KIMI_K3_REVISION: Final = "f831ab66814297da540d832a5235f8e904f29d06" +KIMI_K3_RANKS_FILE: Final = "tiktoken.model" +KIMI_K3_RANKS_SHA256: Final = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103" +_PATTERN: Final = "|".join( + ( + r"[\p{Han}]+", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?", + r"\p{N}{1,3}", + r" ?[^\s\p{L}\p{N}]+[\r\n]*", + r"\s*[\r\n]+", + r"\s+(?!\S)", + r"\s+", + ) +) +# Ends chunks only where the pattern above always splits, so their counts sum to the whole text's count while BPE +# stays linear. Only a run of over 1024 characters with no such split, like 2000 spaces, is cut mid-piece. +_CHUNK: Final = regex.compile( + r"(?s).{1,1024}(?:(?<=\S)(?= \S)|(?<=\p{N})(?=\P{N})|(?<=[^\s\p{N}])(?=\p{N})|(?<=\p{L})(?=[^\p{L}\p{M}'])" + r"|(?<=[\r\n])(?=\S)|\Z)|.{1024}\p{N}*" +) + + +@dataclass(frozen=True, slots=True) +class KimiK3TokenizerSource: + ranks_path: Callable[[], str] + sha256: str + + +def _hugging_face_ranks_path() -> str: + import huggingface_hub + + return huggingface_hub.hf_hub_download( + repo_id=KIMI_K3_REPO, + filename=KIMI_K3_RANKS_FILE, + revision=KIMI_K3_REVISION, + local_files_only=litellm.disable_hf_tokenizer_download is True, + ) + + +HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256) + + +def _loaded(source: KimiK3TokenizerSource) -> tiktoken.Encoding | str: + try: + data: Final = Path(source.ranks_path()).read_bytes() + except Exception as e: # noqa: BLE001 # a download or read failure of any kind falls back to the estimate + return type(e).__name__ + if hashlib.sha256(data).hexdigest() != source.sha256: + return "SHA-256 mismatch" + rows: Final = (line.split() for line in data.splitlines() if line) + return tiktoken.Encoding( + name="kimi_k3", + pat_str=_PATTERN, + mergeable_ranks={base64.b64decode(token): int(rank) for token, rank in rows}, + special_tokens={}, + ) + + +def load_kimi_k3_encoding(source: KimiK3TokenizerSource) -> tiktoken.Encoding | None: + loaded: Final = _loaded(source) + if isinstance(loaded, tiktoken.Encoding): + return loaded + verbose_logger.warning( + "Kimi K3 tokenizer unavailable (%s), so Kimi K3 token counts are estimated until a later count loads it. " + "Allow downloads of %s from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache", + loaded, + KIMI_K3_RANKS_FILE, + KIMI_K3_REPO, + KIMI_K3_REVISION, + ) + return None + + +def _started_load(source: KimiK3TokenizerSource) -> "Future[tiktoken.Encoding | None]": + encoding: Final = Future[tiktoken.Encoding | None]() + threading.Thread( + target=lambda: encoding.set_result(load_kimi_k3_encoding(source)), name="litellm-kimi-k3-tokenizer", daemon=True + ).start() + return encoding + + +@dataclass(frozen=True, slots=True) +class _EncodingLoad: + encoding: "Future[tiktoken.Encoding | None]" + started_at: float + is_retry: bool + + def retry_due(self, now: float, retry_after_seconds: float) -> bool: + return self.encoding.done() and self.encoding.result() is None and now - self.started_at >= retry_after_seconds + + +@dataclass(slots=True) +class KimiK3Encodings: + source: KimiK3TokenizerSource + retry_after_seconds: float + clock: Callable[[], float] = time.monotonic + current: _EncodingLoad | None = None + + def _load(self) -> _EncodingLoad: + now: Final = self.clock() + current: Final = self.current + if current is not None and not current.retry_due(now, self.retry_after_seconds): + return current + load: Final = _EncodingLoad(encoding=_started_load(self.source), started_at=now, is_retry=current is not None) + self.current = load + return load + + async def encoding(self, wait_seconds: float) -> tiktoken.Encoding | None: + load: Final = self._load() + remaining: Final = 0 if load.is_retry else load.started_at + wait_seconds - self.clock() + if load.encoding.done() or remaining <= 0: + return load.encoding.result() if load.encoding.done() else None + await asyncio.wait((asyncio.wrap_future(load.encoding),), timeout=remaining) + if load.encoding.done(): + return load.encoding.result() + verbose_logger.warning( + "Kimi K3 tokenizer still loading after %s s, so Kimi K3 token counts are estimated until it loads. " + "Check that %s can be downloaded from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache", + wait_seconds, + KIMI_K3_RANKS_FILE, + KIMI_K3_REPO, + KIMI_K3_REVISION, + ) + return None + + +class _Converse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class _Text(_Converse): + text: str + + +class _CachePoint(_Converse): + cache_point: JsonValue = Field(alias="cachePoint") + + +class _ToolUse(_Converse): + tool_use_id: str = Field(alias="toolUseId") + name: str + arguments: Mapping[str, JsonValue] = Field(alias="input") + + +class _ToolUseBlock(_Converse): + tool_use: _ToolUse = Field(alias="toolUse") + + +class _ToolResult(_Converse): + tool_use_id: str = Field(alias="toolUseId") + content: tuple[_Text, ...] + status: str | None = None + + +class _ToolResultBlock(_Converse): + tool_result: _ToolResult = Field(alias="toolResult") + + +class _ReasoningText(_Converse): + text: str + signature: str | None = None + + +class _Reasoning(_Converse): + reasoning_text: _ReasoningText = Field(alias="reasoningText") + + +class _ReasoningBlock(_Converse): + reasoning_content: _Reasoning = Field(alias="reasoningContent") + + +class _UserMessage(_Converse): + role: Literal["user"] + content: tuple[_Text | _ToolResultBlock | _CachePoint, ...] + + +class _AssistantMessage(_Converse): + role: Literal["assistant"] + content: tuple[_Text | _ToolUseBlock | _ReasoningBlock | _CachePoint, ...] + + def calls(self) -> tuple[_ToolUse, ...]: + return tuple(block.tool_use for block in self.content if isinstance(block, _ToolUseBlock)) + + +class _InputSchema(_Converse): + definition: JsonValue = Field(alias="json") + + +class _ToolSpec(_Converse): + name: str + description: str | None = None + input_schema: _InputSchema = Field(alias="inputSchema") + + def declared(self) -> Mapping[str, JsonValue]: + described: Final[Mapping[str, JsonValue]] = ( + {} if self.description is None else {"description": self.description} + ) + return {"name": self.name, **described, "parameters": self.input_schema.definition} + + +class _ToolSpecBlock(_Converse): + tool_spec: _ToolSpec = Field(alias="toolSpec") + + +class _ToolConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + tools: tuple[_ToolSpecBlock | _CachePoint, ...] + + +_Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")] + + +def _result_ids(message: _Message) -> frozenset[str]: + if isinstance(message, _AssistantMessage): + return frozenset() + return frozenset(block.tool_result.tool_use_id for block in message.content if isinstance(block, _ToolResultBlock)) + + +def _call_ids(message: _Message | None) -> frozenset[str]: + if not isinstance(message, _AssistantMessage): + return frozenset() + return frozenset(call.tool_use_id for call in message.calls()) + + +class _ConverseRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + messages: tuple[_Message, ...] + system: tuple[_Text | _CachePoint, ...] = () + tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig") + + @model_validator(mode="after") + def _results_answer_the_previous_calls(self) -> Self: + pairs: Final = zip((None, *self.messages), self.messages, strict=False) + if not all(_result_ids(message) <= _call_ids(previous) for previous, message in pairs): + raise ValueError("a toolResult does not answer a toolUse of the previous assistant message") + return self + + +class _Control(Enum): + OPEN = "<|open|>" + CLOSE = "<|close|>" + SEP = "<|sep|>" + END_OF_MSG = "<|end_of_msg|>" + + +_Segment: TypeAlias = str | _Control +_Segments: TypeAlias = tuple[_Segment, ...] + + +def _open(tag: str, attrs: Iterable[tuple[str, str]] = ()) -> _Segments: + escaped: Final = ( + (f" {key}", '="', value.replace("&", "&").replace('"', """), '"') for key, value in attrs + ) + return (_Control.OPEN, tag, *chain.from_iterable(escaped), _Control.SEP) + + +def _close(tag: str) -> _Segments: + return (_Control.CLOSE, tag, _Control.SEP) + + +def _message(attrs: Iterable[tuple[str, str]], body: Iterable[_Segment]) -> _Segments: + return (*_open("message", attrs), *body, *_close("message"), _Control.END_OF_MSG) + + +def _tool_declare(specs: Sequence[_ToolSpec]) -> _Segments: + tools: Final = json.dumps( + [{"type": "function", "function": spec.declared()} for spec in specs], + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + return _message( + (("role", "system"), ("type", "tool-declare")), + (f"# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n{tools}\n```",), + ) + + +def _text(blocks: Iterable[_Converse]) -> str: + return "".join(block.text for block in blocks if isinstance(block, _Text)) + + +def _json_type(value: JsonValue) -> str: + match value: + case bool(): + return "boolean" + case None: + return "null" + case int() | float(): + return "number" + case str(): + return "string" + case dict(): + return "object" + case _: + return "array" + + +def _argument(key: str, value: JsonValue) -> _Segments: + text: Final = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) + return (*_open("argument", (("key", key), ("type", _json_type(value)))), text, *_close("argument")) + + +def _call(index: int, call: _ToolUse) -> _Segments: + arguments: Final = chain.from_iterable(_argument(key, value) for key, value in call.arguments.items()) + return (*_open("call", (("tool", call.name), ("index", str(index)))), *arguments, *_close("call")) + + +def _tool_calls(calls: Sequence[_ToolUse]) -> _Segments: + if not calls: + return () + rendered: Final = chain.from_iterable(_call(index, call) for index, call in enumerate(calls, start=1)) + return (*_open("tools"), *rendered, *_close("tools")) + + +def _assistant_message(message: _AssistantMessage) -> _Segments: + reasoning: Final = "".join( + block.reasoning_content.reasoning_text.text for block in message.content if isinstance(block, _ReasoningBlock) + ) + return _message( + (("role", "assistant"),), + ( + *_open("think"), + reasoning, + *_close("think"), + *_open("response"), + _text(message.content), + *_close("response"), + *_tool_calls(message.calls()), + ), + ) + + +def _tool_messages(results: Sequence[_ToolResult], previous: _AssistantMessage) -> _Segments: + names: Final = {call.tool_use_id: call.name for call in previous.calls()} + return tuple( + chain.from_iterable( + _message( + (("role", "tool"), ("tool", names[result.tool_use_id]), ("index", str(index))), + (_text(result.content),), + ) + for index, result in enumerate(results, start=1) + ) + ) + + +def _user_run(blocks: Sequence[_Text | _ToolResultBlock], previous: _Message | None) -> _Segments: + results: Final = tuple(block.tool_result for block in blocks if isinstance(block, _ToolResultBlock)) + if results and isinstance(previous, _AssistantMessage): + return _tool_messages(results, previous) + return _message((("role", "user"),), (_text(blocks),)) + + +def _user_messages(message: _UserMessage, previous: _Message | None) -> _Segments: + blocks: Final = tuple(block for block in message.content if not isinstance(block, _CachePoint)) + runs: Final = groupby(blocks, key=lambda block: isinstance(block, _ToolResultBlock)) + return tuple(chain.from_iterable(_user_run(tuple(run), previous) for _, run in runs)) + + +def _turn(previous: _Message | None, message: _Message) -> _Segments: + if isinstance(message, _AssistantMessage): + return _assistant_message(message) + return _user_messages(message, previous) + + +def _rendered(request: _ConverseRequest) -> _Segments: + specs: Final = ( + () + if request.tool_config is None + else tuple(block.tool_spec for block in request.tool_config.tools if isinstance(block, _ToolSpecBlock)) + ) + system: Final = _text(request.system) + turns: Final = chain.from_iterable( + _turn(previous, message) for previous, message in zip((None, *request.messages), request.messages, strict=False) + ) + return ( + *(_tool_declare(specs) if specs else ()), + *(_message((("role", "system"),), (system,)) if system else ()), + *turns, + *_open("message", (("role", "assistant"),)), + *_open("think"), + ) + + +def _text_tokens(encoding: tiktoken.Encoding, text: str) -> int: + spans: Final = (match.span() for match in _CHUNK.finditer(text)) + return sum(len(encoding.encode_ordinary(text[start:end])) for start, end in spans) + + +class _TextPart(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + type: Literal["text"] + + +class _TextOnlyMessage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + content: str | tuple[_TextPart, ...] | None = None + + +def _text_only(message: AllMessageValues) -> bool: + try: + _TextOnlyMessage.model_validate(message) + except ValidationError: + return False + return True + + +def _as_converse_input(tool: ChatCompletionToolParam) -> tuple[Mapping[str, object], ...]: + return chat_completion_tools_for(tool) if is_flat_responses_tool(tool) else (tool,) + + +def _converse_request( + model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None +) -> _ConverseRequest | None: + if not all(_text_only(message) for message in messages): + return None + try: + chat_tools: Final = tuple( + chat_tool for tool in tools or () for chat_tool in _as_converse_input(tool) + ) # comprehension-ok: flattens each tool's translation + return _ConverseRequest.model_validate( + litellm.AmazonConverseConfig().transform_request( + model=model, + messages=copy.deepcopy(list(messages)), + optional_params={"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {}, + litellm_params={}, + headers={}, + ) + ) + except Exception: # noqa: BLE001 # a body litellm cannot send to Bedrock, or cannot render, falls back to the estimate + return None + + +def kimi_k3_tokens( + encoding: tiktoken.Encoding, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: + request: Final = _converse_request(model, messages, tools) + if request is None: + return None + return sum( + 1 if isinstance(segment, _Control) else _text_tokens(encoding, segment) for segment in _rendered(request) + ) + + +def kimi_k3_anthropic_tokens(encoding: tiktoken.Encoding, model: str, fields: Mapping[str, object]) -> int | None: + body: Final = {**fields, "model": model, "max_tokens": 1} + try: + translated, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + copy.deepcopy(body), # pyright: ignore[reportArgumentType] # raw /v1/messages fields + custom_llm_provider="bedrock", + ) + except Exception: # noqa: BLE001 # a request the adapter cannot translate falls back to the estimate + return None + match translated: + case {"tools": tools}: + return kimi_k3_tokens(encoding, model, translated["messages"], tools) + case _: + return kimi_k3_tokens(encoding, model, translated["messages"], None) + + +def kimi_k3_template_tokens(message_count: int, has_tools: bool) -> int: + return 14 + 13 * max(message_count - 1, 0) + (40 if has_tools else 0) diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 23c3dc54c95..4ea4440a3c4 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -52,7 +52,7 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): ) -> dict[str, JsonValue]: try: request: Final = _COUNT_REQUEST.validate_python(request_data) - aws_region_name: Final = self._get_aws_region_name( + aws_region_name: Final = self.get_aws_region_name( optional_params=litellm_params, model=resolved_model, model_id=None ) endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index c7c997bc961..633f154c4cd 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -124,6 +124,7 @@ class OpenAITokenCounter(BaseTokenCounter): error=True, error_message=request_error, status_code=400, + rejected_before_provider_call=True, ) deployment = deployment or {} diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index 871532b7af5..c218b4710fd 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import assert_never +from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for from litellm.types.llms.anthropic import AnthropicMessagesTextParam from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ContentType @@ -75,27 +76,6 @@ def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool: return isinstance(tool.get("name"), str) or "input_schema" in tool or "function" in tool -def _is_responses_function_tool(tool: Mapping[str, JsonValue]) -> bool: - return tool.get("type") == "function" and "function" not in tool and "input_schema" not in tool - - -def _chat_tool(tool: Mapping[str, JsonValue]) -> tuple[Mapping[str, object], ...]: - from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - if _is_responses_function_tool(tool): - chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - [dict(tool)] # pyright: ignore[reportArgumentType] # validated JSON dict in the Responses function tool shape - ) - return tuple(chat_tools) - anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai( - [dict(tool)] # pyright: ignore[reportArgumentType] # validated JSON dict in an Anthropic or chat tool shape - ) - return tuple(anthropic_tools) - - def _map_function_tools( tools: object | None, ) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError: @@ -112,7 +92,7 @@ def _map_function_tools( from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig chat_tools: Final = tuple( - chat_tool for tool in function_tools for chat_tool in _chat_tool(tool) + chat_tool for tool in function_tools for chat_tool in chat_completion_tools_for(tool) ) # comprehension-ok: flattens each tool's translation try: mapped_tools: Final = VertexGeminiConfig().map_function( @@ -163,6 +143,7 @@ def count_context_error_response( error=True, error_message=error.error_message, status_code=error.status_code, + rejected_before_provider_call=True, ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3828957f2fb..cc931c89a31 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21144,6 +21144,7 @@ "input_cost_per_token": 1.4e-07, "input_dbu_cost_per_token": 2e-06, "litellm_provider": "databricks", + "token_count_framing": 97, "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, @@ -21166,6 +21167,11 @@ "supports_tool_choice": true, "supports_vision": false }, + "databricks/system.ai.deepseek-v4-flash-0731": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 97 + }, "databricks/databricks-deepseek-v4-pro-0813": { "cache_creation_input_token_cost": 1.31999e-06, "cache_read_input_token_cost": 1.3202e-07, @@ -22134,6 +22140,7 @@ "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", + "token_count_framing": 100, "max_input_tokens": 1000000, "max_output_tokens": 1048576, "max_tokens": 1048576, @@ -22157,6 +22164,11 @@ "supports_tool_choice": true, "supports_vision": true }, + "databricks/system.ai.kimi-k3": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 100 + }, "databricks/databricks-llama-4-maverick": { "cache_creation_input_token_cost": 5.0001e-07, "cache_read_input_token_cost": 5.0001e-07, @@ -71540,6 +71552,8 @@ "input_cost_per_token_flex": 1.65e-06, "input_cost_per_token_priority": 5.775e-06, "litellm_provider": "bedrock_converse", + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, @@ -71570,6 +71584,8 @@ "input_cost_per_token_flex": 1.5e-06, "input_cost_per_token_priority": 5.25e-06, "litellm_provider": "bedrock_converse", + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, @@ -71600,6 +71616,8 @@ "input_cost_per_token_flex": 1.65e-06, "input_cost_per_token_priority": 5.775e-06, "litellm_provider": "bedrock_converse", + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, @@ -73100,6 +73118,10 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -73173,6 +73195,10 @@ "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "au.anthropic.claude-sonnet-5-5": { + "token_count_model": { + "model": "au.anthropic.claude-opus-5-5", + "region": "ap-southeast-4" + }, "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, @@ -73373,6 +73399,10 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "global.anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -73582,6 +73612,10 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "us.anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -74382,6 +74416,8 @@ "input_cost_per_token_flex": 1.65e-06, "input_cost_per_token_priority": 5.775e-06, "litellm_provider": "bedrock_converse", + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, @@ -74555,6 +74591,10 @@ "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 1e-07, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -74633,6 +74673,10 @@ "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "au.anthropic.claude-haiku-5-5": { + "token_count_model": { + "model": "au.anthropic.claude-opus-5-5", + "region": "ap-southeast-4" + }, "input_cost_per_token_above_100k_tokens": 5.5e-07, "output_cost_per_token_above_100k_tokens": 2.75e-06, "cache_creation_input_token_cost_above_100k_tokens": 6.875e-07, @@ -74863,6 +74907,10 @@ "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 1e-07, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "global.anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -75010,6 +75058,10 @@ "cache_read_input_token_cost": 1.1e-08, "input_cost_per_token": 1.1e-07, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "us.anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1fa276d9ab2..b12676e08f3 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -48,6 +48,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.litellm_pre_call_utils import should_auto_drop_params_for_agentic_cli from litellm.router import Router @@ -420,7 +421,7 @@ 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} + return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)} @router.post( diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py new file mode 100644 index 00000000000..6f805616820 --- /dev/null +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -0,0 +1,164 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, ValidationError + +import litellm +from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.proxy._types import ProxyException +from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer +from litellm.types.utils import TokenCountResponse + +CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) +TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0 +TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0 +TOKENIZER_LOAD_MAX_RUNNING: Final = 4 + + +class _UpstreamErrorDetail(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + message: str + + +class _UpstreamErrorEnvelope(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + type: Literal["error"] + error: _UpstreamErrorDetail + + +def _joined(message: str, hint: str) -> str: + return f"{message.rstrip().rstrip('.')}. {hint}" + + +def _with_hint(message: str, hint: str) -> str: + try: + envelope: Final = _UpstreamErrorEnvelope.model_validate_json(message) + except ValidationError: + return _joined(AnthropicExceptionMapping.extract_error_message(message), hint) + detail: Final = envelope.error.model_copy(update={"message": _joined(envelope.error.message, hint)}) + return envelope.model_copy(update={"error": detail}).model_dump_json() + + +@dataclass(frozen=True, slots=True) +class TokenCountModel: + model: str + region: str | None + + +@dataclass(frozen=True, slots=True) +class TokenCountSettings: + framing: int | None + count_model: TokenCountModel | None + tokenizer: TokenCountTokenizer | None = None + + +@dataclass(frozen=True, slots=True) +class ProviderCountFailure: + status_code: int + message: str + rejected_before_provider_call: bool = False + + @property + def cannot_count_model(self) -> bool: + return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES + + def with_hint(self, hint: str) -> "ProviderCountFailure": + return replace(self, message=_with_hint(self.message, hint)) + + def to_proxy_exception(self) -> ProxyException: + return ProxyException( + message=self.message, + type="token_counting_error", + param="model", + code=self.status_code, + ) + + +def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models + return MappingProxyType({}) + + +def _framing(value: object) -> int | None: + match value: + case bool(): + return None + case int() if value >= 0: + return value + case float() if value.is_integer() and value >= 0: + return int(value) + case _: + return None + + +def _non_blank(value: object) -> str | None: + return value if isinstance(value, str) and value.strip() else None + + +def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: + model: Final = _non_blank(value.get("model")) + return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region"))) + + +def _geography(region: str) -> str: + return region.rsplit("-", 2)[0] + + +def _within_deployment_geography( + count_model: TokenCountModel | None, deployment_region: Callable[[], str] +) -> TokenCountModel | None: + if count_model is None or count_model.region is None or count_model.model.startswith("global."): + return count_model + if _geography(deployment_region()) == _geography(count_model.region): + return count_model + return replace(count_model, region=None) + + +def _resolved_count_model( + configured: object, default: object, deployment_region: Callable[[], str] +) -> TokenCountModel | None: + configured_mapping: Final = as_str_mapping(configured) + if configured_mapping is not None: + return _count_model(configured_mapping) + default_mapping: Final = as_str_mapping(default) + return ( + None + if default_mapping is None + else _within_deployment_geography(_count_model(default_mapping), deployment_region) + ) + + +def _tokenizer(value: object) -> TokenCountTokenizer | None: + match value: + case "kimi_k3": + return "kimi_k3" + case _: + return None + + +def resolve_token_count_settings( + configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] +) -> TokenCountSettings: + framing: Final = _framing(configured.get("token_count_framing")) + return TokenCountSettings( + framing=framing if framing is not None else _framing(default.get("token_count_framing")), + count_model=_resolved_count_model( + configured.get("token_count_model"), default.get("token_count_model"), deployment_region + ), + tokenizer=_tokenizer(default.get("token_count_tokenizer")), + ) + + +def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields: + match response.count_source: + case None: + return {} + case "token_count_model": + return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used} + case source: + return {"litellm_count_source": source} diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index d382ff53ffb..3d1a9ab4cfe 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -21,6 +21,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, ) +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.types.llms.vertex_ai import TokenCountDetailsResponse @@ -309,10 +310,12 @@ async def google_count_tokens(request: Request, model_name: str): ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), ) original_response: Final[dict] = token_response.original_response or {} - return TokenCountDetailsResponse( - totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, - promptTokensDetails=original_response.get("promptTokensDetails", []), - ) + counted: Final[TokenCountDetailsResponse] = { + "totalTokens": original_response.get("totalTokens") or token_response.total_tokens or 0, + "promptTokensDetails": original_response.get("promptTokensDetails", []), + **count_source_fields(token_response), + } + return counted # ============================================================ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dba4b987725..0c7840c92fd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -55,6 +55,7 @@ from typing import ( ) import anyio +import anyio.lowlevel import websockets import websockets.exceptions from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config @@ -99,6 +100,13 @@ from litellm.litellm_core_utils.litellm_logging import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + HUGGING_FACE_SOURCE, + KimiK3Encodings, + kimi_k3_anthropic_tokens, + kimi_k3_template_tokens, + kimi_k3_tokens, +) from litellm.proxy._types import ( UI_TEAM_ID, AwsErrorTypeProxyException, @@ -149,6 +157,15 @@ from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 # _realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export realtime_request_body, ) +from litellm.proxy.common_utils.token_count_settings import ( + TOKENIZER_LOAD_MAX_RUNNING, + TOKENIZER_LOAD_REUSE_SECONDS, + TOKENIZER_LOAD_WAIT_SECONDS, + ProviderCountFailure, + TokenCountSettings, + cost_map_entry, + resolve_token_count_settings, +) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.router_utils.access_windows import access_windows_config_error from litellm.router_utils.add_retry_fallback_headers import ( @@ -304,6 +321,7 @@ from litellm.constants import ( CLI_SSO_SESSION_TTL_SECONDS, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, + DEFAULT_MAX_LRU_CACHE_SIZE, DEFAULT_MODEL_CREATED_AT_TIME, GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -337,6 +355,7 @@ from litellm.litellm_core_utils.core_helpers import ( drop_params_flag, get_litellm_metadata_from_kwargs, get_parent_otel_span_from_kwargs, + iter_items, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -993,7 +1012,13 @@ from litellm.types.secret_managers.main import ( KeyManagementSettings, KeyManagementSystem, ) -from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload +from litellm.types.utils import ( + CredentialItem, + CustomHuggingfaceTokenizer, + RawRequestTypedDict, + SelectTokenizerResponse, + StandardLoggingPayload, +) from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.utils import add_custom_logger_callback_to_specific_event @@ -14673,25 +14698,36 @@ async def _try_provider_token_count( provider_counter: "BaseTokenCounter", custom_llm_provider: str | None, model_to_use: str, + region: str | None, messages: list | None, contents: list | None, - deployment: dict[str, object] | None, + deployment: Mapping[str, object] | None, request_model: str, tools: list | None = None, system: str | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: CountContextManagement | None = None, -) -> Optional["TokenCountResponse"]: - """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.""" +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + """Attempt provider-specific token counting. Returns None to fall through to local counting.""" if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider): return None + deployment_in_region: Final = ( + None + if deployment is None + else dict(deployment) + if region is None + else { + **deployment, + "litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), "aws_region_name": region}, + } + ) try: result: Final = await provider_counter.count_tokens( model_to_use=model_to_use or "", messages=messages, contents=contents, - deployment=deployment, + deployment=deployment_in_region, request_model=request_model, tools=tools, system=system, @@ -14700,24 +14736,325 @@ async def _try_provider_token_count( context_management=context_management, ) except httpx.HTTPStatusError as e: - error_message: Final = getattr(e, "message", None) or str(e) - status_code: Final = getattr(e, "status_code", None) or e.response.status_code - raise ProxyException( - message=error_message, - type="token_counting_error", - param="model", - code=status_code, + return ProviderCountFailure( + status_code=getattr(e, "status_code", None) or e.response.status_code, + message=getattr(e, "message", None) or str(e), ) if result is not None and result.error is True: - raise ProxyException( + return ProviderCountFailure( + status_code=result.status_code or 500, message=result.error_message or "Token counting failed", - type="token_counting_error", - param="model", - code=result.status_code or 500, + rejected_before_provider_call=result.rejected_before_provider_call, ) return result +async def _provider_token_count( + count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]], + model_to_use: str, + settings: TokenCountSettings, + local_estimate_allowed: bool, +) -> Union["TokenCountResponse", ProviderCountFailure, None]: + """Count with the provider, using the token_count_model in its region when one is set. + + Returns None when the count should be estimated locally, or the provider's failure when the framed local + estimate stands in for it and that failure is the answer if the estimate cannot read the request.""" + count_model: Final = settings.count_model + counted: Final = await ( + count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region) + ) + if counted is None: + return None + if isinstance(counted, TokenCountResponse): + if count_model is None: + return counted.model_copy(update={"count_source": "provider"}) + return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model}) + failure: Final = ( + counted + if count_model is None or count_model.region is None + else counted.with_hint( + f"This model counts tokens with {count_model.model} in {count_model.region}. " + "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" + ) + ) + if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed: + return failure + raise failure.to_proxy_exception() + + +@dataclass(frozen=True, slots=True) +class _LocalTokenEstimate: + total_tokens: int + tokenizer_type: str + + +def _is_fork_guard_error(error: BaseException) -> bool: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking + + return isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)) + + +def _is_tokenizer_panic(error: BaseException) -> bool: + return type(error).__module__ == "pyo3_runtime" and type(error).__name__ == "PanicException" + + +def _answers_provider_error(error: BaseException) -> bool: + return (isinstance(error, Exception) or _is_tokenizer_panic(error)) and not _is_fork_guard_error(error) + + +def _provider_error_after_failed_estimate( + provider_failure: ProviderCountFailure, error: BaseException +) -> ProxyException: + verbose_proxy_logger.warning( + "Local token estimate failed after the provider's count error %s, answering that error: %s", + provider_failure.status_code, + type(error).__name__, + ) + return provider_failure.to_proxy_exception() + + +async def _estimate_or_provider_error(estimate: Awaitable[int], provider_failure: ProviderCountFailure) -> int: + try: + return await estimate + except BaseException as e: + if not _answers_provider_error(e): + raise + raise _provider_error_after_failed_estimate(provider_failure, e) from e + + +@dataclass(frozen=True, slots=True) +class _CustomTokenizerKey: + identifier: str + revision: str + auth_token: str | None + + +class _TokenizerLoadsBusy(Exception): + pass + + +@dataclass(frozen=True, slots=True) +class _TokenizerLoadOutcome: + tokenizer: SelectTokenizerResponse | BaseException + finished_at: float + + +@dataclass(frozen=True, slots=True) +class _SharedTokenizerLoad: + outcome: asyncio.Task[_TokenizerLoadOutcome] + started_at: float + + def reusable(self, now: float) -> bool: + return not self.outcome.done() or now - self.outcome.result().finished_at < TOKENIZER_LOAD_REUSE_SECONDS + + +async def _loaded_tokenizer( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter +) -> SelectTokenizerResponse | BaseException: + try: + return await asyncify(litellm.utils.select_tokenizer, limiter=threads)( + model=model, custom_tokenizer=custom_tokenizer + ) + except asyncio.CancelledError: + raise + except BaseException as e: # noqa: BLE001 # every count sharing this load decides what it answers, pyo3 panics included + return e + + +async def _tokenizer_load_outcome( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter +) -> _TokenizerLoadOutcome: + tokenizer: Final = await _loaded_tokenizer(model, custom_tokenizer, threads) + return _TokenizerLoadOutcome(tokenizer=tokenizer, finished_at=asyncio.get_running_loop().time()) + + +@dataclass(slots=True) +class _TokenizerLoads: + threads: anyio.CapacityLimiter + running: int = 0 + + def start( + self, model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, loop: asyncio.AbstractEventLoop + ) -> _SharedTokenizerLoad | None: + if self.running >= self.threads.total_tokens: + return None + self.running += 1 + outcome: Final = loop.create_task(_tokenizer_load_outcome(model, custom_tokenizer, self.threads)) + outcome.add_done_callback(self._finished) + return _SharedTokenizerLoad(outcome=outcome, started_at=loop.time()) + + def _finished(self, _outcome: asyncio.Task[_TokenizerLoadOutcome]) -> None: + self.running -= 1 + + +_TOKENIZER_LOADS: Final = anyio.lowlevel.RunVar[_TokenizerLoads]("litellm_proxy_tokenizer_loads") + + +def _tokenizer_loads_for_this_loop() -> _TokenizerLoads: + existing: Final = _TOKENIZER_LOADS.get(None) + if existing is not None: + return existing + created: Final = _TokenizerLoads(threads=anyio.CapacityLimiter(TOKENIZER_LOAD_MAX_RUNNING)) + _TOKENIZER_LOADS.set(created) + return created + + +@dataclass(slots=True) +class _TokenizerLoadSlot: + model: str + custom_tokenizer: CustomHuggingfaceTokenizer | None + current: _SharedTokenizerLoad | None = None + + def load(self, loads: _TokenizerLoads, loop: asyncio.AbstractEventLoop) -> _SharedTokenizerLoad | None: + current: Final = self.current + if current is not None and current.reusable(loop.time()): + return current + self.current = loads.start(self.model, self.custom_tokenizer, loop) + return self.current + + +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def _tokenizer_load_slot( + model: str, custom_tokenizer: _CustomTokenizerKey | None, loop: asyncio.AbstractEventLoop +) -> _TokenizerLoadSlot: + return _TokenizerLoadSlot( + model=model, + custom_tokenizer=None + if custom_tokenizer is None + else CustomHuggingfaceTokenizer( + identifier=custom_tokenizer.identifier, + revision=custom_tokenizer.revision, + auth_token=custom_tokenizer.auth_token, + ), + ) + + +async def _tokenizer_after_provider_error( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None +) -> SelectTokenizerResponse | BaseException: + loop: Final = asyncio.get_running_loop() + try: + load: Final = _tokenizer_load_slot( + model, + None + if custom_tokenizer is None + else _CustomTokenizerKey( + identifier=custom_tokenizer["identifier"], + revision=custom_tokenizer["revision"], + auth_token=custom_tokenizer["auth_token"], + ), + loop, + ).load(_tokenizer_loads_for_this_loop(), loop) + if load is None: + return _TokenizerLoadsBusy() + outcome: Final = await asyncio.wait_for( + asyncio.shield(load.outcome), + timeout=max(load.started_at + TOKENIZER_LOAD_WAIT_SECONDS - loop.time(), 0), + ) + except Exception as e: # noqa: BLE001 # a tokenizer that cannot be read, loaded in time, or loaded at all answers the provider's error + return e + return outcome.tokenizer + + +async def _tokenizer_or_provider_error( + model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, provider_failure: ProviderCountFailure +) -> SelectTokenizerResponse: + tokenizer: Final = await _tokenizer_after_provider_error(model, custom_tokenizer) + if not isinstance(tokenizer, BaseException): + return tokenizer + if not _answers_provider_error(tokenizer): + raise tokenizer.with_traceback(None) + raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer + + +async def _local_token_estimate( + model: str, + prompt: str | None, + messages: Sequence[AllMessageValues] | None, + tools: list[ChatCompletionToolParam] | None, + custom_tokenizer: CustomHuggingfaceTokenizer | None, + provider_failure: ProviderCountFailure | None, +) -> _LocalTokenEstimate: + tokenizer: Final = ( + await asyncify(litellm.utils.select_tokenizer)(model=model, custom_tokenizer=custom_tokenizer) + if provider_failure is None + else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure) + ) + count: Final = offload_token_count(litellm.token_counter)( + model=model, + text=prompt, + messages=messages, + tools=tools, + custom_tokenizer=tokenizer, + ) + total_tokens: Final = await ( + count if provider_failure is None else _estimate_or_provider_error(count, provider_failure) + ) + return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=str(tokenizer["type"])) + + +_KIMI_K3_ENCODINGS: Final = KimiK3Encodings( + source=HUGGING_FACE_SOURCE, retry_after_seconds=TOKENIZER_LOAD_REUSE_SECONDS +) + + +async def _kimi_k3_count( + request: TokenCountRequest, + messages_format: TokenCountMessagesFormat, + model: str, + messages: Sequence[AllMessageValues], + tools: Sequence[ChatCompletionToolParam] | None, +) -> int | None: + encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS) + if encoding is None: + return None + match messages_format: + case "anthropic": + return await offload_token_count(kimi_k3_anthropic_tokens)( + encoding, model, request.model_dump(include={"messages", "system", "tools"}, exclude_none=True) + ) + case "openai": + return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools) + + +class _TokenCountDeployment(TypedDict, total=False): + litellm_params: ReadOnly[Mapping[str, object]] + model_info: ReadOnly[Mapping[str, object]] + + +_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment) + + +def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None: + group: Final = () if router is None else router.get_all_deployments(model_name=model_group) + return group[0] if group else None + + +def _non_empty_str(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _deployment_aws_region(deployment: _TokenCountDeployment) -> str: + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + params: Final = deployment.get("litellm_params", {}) + return BaseAWSLLM().get_aws_region_name( + {"aws_region_name": params.get("aws_region_name")}, model=_non_empty_str(params.get("model")) + ) + + +def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]: + model_info: Final = deployment.get("model_info", {}) + params: Final = deployment.get("litellm_params", {}) + base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model")) + if base_model is not None: + return base_model, None + model: Final = _non_empty_str(params.get("model")) + if model is None: + return request_model, None + return model, _non_empty_str(params.get("custom_llm_provider")) + + _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} ) @@ -14728,6 +15065,14 @@ def _drops_params(deployment: Mapping[str, object] | None) -> bool: return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) +_ESTIMABLE_GOOGLE_PART_KEYS: Final = frozenset({"text", "functionCall", "functionResponse", "thoughtSignature"}) + + +def _estimable_contents(contents: Sequence[Mapping[str, object]]) -> bool: + parts: Final = chain.from_iterable(iter_items(content.get("parts") or ()) for content in contents) + return all(set(as_str_mapping(part) or ()) <= _ESTIMABLE_GOOGLE_PART_KEYS for part in parts) + + def _router_default_litellm_params(router: Router) -> Mapping[str, object]: defaults: Final[Mapping[str, object] | None] = getattr(router, "default_litellm_params", None) if defaults is None: @@ -14831,6 +15176,14 @@ async def count_request_tokens( model_to_use: str = ( litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model ) + deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python( + deployment or _model_group_deployment(llm_router, request.model) or {} + ) + settings: Final = await asyncify(resolve_token_count_settings)( + configured=deployment_view.get("model_info", {}), + default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)), + deployment_region=partial(_deployment_aws_region, deployment_view), + ) # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -14850,24 +15203,33 @@ async def count_request_tokens( code=400, ) prompt_only: Final = bool(prompt) and conversation_empty - if provider_counter is not None: - prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES - result: Final = await _try_provider_token_count( - provider_counter=provider_counter, - custom_llm_provider=custom_llm_provider, + contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents)) + prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES + ladder: Final = ( + None + if provider_counter is None + else await _provider_token_count( + count_with=partial( + _try_provider_token_count, + provider_counter, + custom_llm_provider, + messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, + contents=None if prompt_as_message else contents, + deployment=deployment, + request_model=request.model, + tools=tools, + system=system, + messages_format=messages_format, + output_format=output_format, + context_management=_count_context_management(context_management, request.model, request_litellm_params), + ), model_to_use=model_to_use, - messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, - contents=None if prompt_as_message else contents, - deployment=deployment, - request_model=request.model, - tools=tools, - system=system, - messages_format=messages_format, - output_format=output_format, - context_management=_count_context_management(context_management, request.model, request_litellm_params), + settings=settings, + local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, ) - if result is not None: - return result + ) + if isinstance(ladder, TokenCountResponse): + return ladder # Check if token counter is disabled before fallback if litellm.disable_token_counter is True: @@ -14885,11 +15247,6 @@ async def count_request_tokens( CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await asyncify(litellm.utils.select_tokenizer)( - model=model_to_use, custom_tokenizer=custom_tokenizer - ) - - tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes @@ -14902,18 +15259,38 @@ async def count_request_tokens( counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - total_tokens: Final = await offload_token_count(litellm.token_counter)( + framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 + kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None + exact: Final = ( + await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools) + if kimi_k3_messages is not None and custom_tokenizer is None + else None + ) + if exact is not None: + return TokenCountResponse( + total_tokens=exact + framing, + request_model=request.model, + model_used=model_to_use, + tokenizer_type="huggingface_tokenizer", + count_source="local_estimate", + ) + estimate: Final = await _local_token_estimate( model=model_to_use, - text=prompt if counted_messages is None else None, + prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, - custom_tokenizer=_tokenizer_used, + custom_tokenizer=custom_tokenizer, + provider_failure=ladder, + ) + template: Final = ( + 0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(counted_tools)) ) return TokenCountResponse( - total_tokens=total_tokens, + total_tokens=estimate.total_tokens + template + framing, request_model=request.model, model_used=model_to_use, - tokenizer_type=tokenizer_used, + tokenizer_type=estimate.tokenizer_type, + count_source="local_estimate", ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 3c727a6414c..f032cd79783 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -67,6 +67,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, openai_error_type_for_status, ) +from litellm.proxy.common_utils.token_count_settings import count_source_fields from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id from litellm.proxy.litellm_pre_call_utils import ( @@ -86,6 +87,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.responses.main import DeleteResponseResult +from litellm.types.token_count import TokenCountSourceFields from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: @@ -1373,7 +1375,7 @@ class _ResponsesApiErrorBody(TypedDict): error: ReadOnly[_ResponsesApiErrorDetail] -class _ResponsesInputTokensResult(TypedDict): +class _ResponsesInputTokensResult(TokenCountSourceFields): object: ReadOnly[str] input_tokens: ReadOnly[int] @@ -1610,6 +1612,7 @@ async def responses_input_tokens( result: Final[_ResponsesInputTokensResult] = { "object": "response.input_tokens", "input_tokens": token_response.total_tokens, + **count_source_fields(token_response), } return result diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 3dc7a56239d..0480d6b1f3e 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -9,6 +9,7 @@ from typing_extensions import ( ) from litellm.types.llms.openai import ChatCompletionFileObject, EmbeddingInput +from litellm.types.token_count import TokenCountSourceFields # Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit # opt-in for combined embeddings — a provider-specific extension of the @@ -320,12 +321,13 @@ class UsageMetadata(TypedDict, total=False): candidatesTokensDetails: list[PromptTokensDetails] # Alternative key name used in some responses -class TokenCountDetailsResponse(TypedDict): +class TokenCountDetailsResponse(TokenCountSourceFields): """ Response structure for token count details with modality breakdown. Example: - {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}], + 'litellm_count_source': 'provider'} """ totalTokens: int diff --git a/litellm/types/token_count.py b/litellm/types/token_count.py new file mode 100644 index 00000000000..3b21401066c --- /dev/null +++ b/litellm/types/token_count.py @@ -0,0 +1,16 @@ +from typing import Literal, TypeAlias + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"] +TokenCountTokenizer: TypeAlias = Literal["kimi_k3"] + + +class TokenCountModelInfo(TypedDict): + model: ReadOnly[str] + region: ReadOnly[NotRequired[str]] + + +class TokenCountSourceFields(TypedDict, total=False): + litellm_count_source: ReadOnly[TokenCountSource] + litellm_count_model: ReadOnly[str] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6c47ac7fa32..8b2b319d1e0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -58,6 +58,7 @@ from litellm.types.llms.base import ( LiteLLMPydanticObjectBase, ) from litellm.types.mcp import MCPServerCostInfo +from litellm.types.token_count import TokenCountModelInfo, TokenCountSource, TokenCountTokenizer from litellm.types.workload_identity import ( ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS, @@ -230,6 +231,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False): bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None bedrock_converse_supports_strict_tools: bool | None supports_regex_lookaround: ReadOnly[bool | None] + token_count_framing: ReadOnly[int | None] + token_count_model: ReadOnly[TokenCountModelInfo | None] + token_count_tokenizer: ReadOnly[TokenCountTokenizer | None] class SearchContextCostPerQuery(TypedDict, total=False): @@ -4051,11 +4055,14 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( {"supports_regex_lookaround", "supports_assistant_prefill", "supports_bedrock_native_structured_output"} ) +DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) + SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) - frozenset(CustomPricingLiteLLMParams.model_fields) - DEPLOYMENT_SCOPED_PRICING_FIELDS - DEPLOYMENT_SCOPED_CAPABILITY_FIELDS + - DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS ) @@ -4492,6 +4499,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase): request_model: str model_used: str tokenizer_type: str + count_source: TokenCountSource | None = None original_response: dict | None = None """ Original Response from upstream API call - if an API call was made for token counting @@ -4502,6 +4510,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase): HTTP status code from the token counting API (e.g., 200 for success, 429 for rate limit, 400 for bad request) """ status_code: int | None = None + rejected_before_provider_call: bool = Field(default=False, exclude=True) class CustomHuggingfaceTokenizer(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index a283ea0e1e5..142a1745143 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6694,6 +6694,9 @@ def get_model_info_helper( uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), supported_audio_formats=_model_info.get("supported_audio_formats", None), + token_count_framing=_model_info.get("token_count_framing", None), + token_count_model=_model_info.get("token_count_model", None), + token_count_tokenizer=_model_info.get("token_count_tokenizer", None), vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), ) for cost_key, cost_value in _model_info.items(): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3828957f2fb..cc931c89a31 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21144,6 +21144,7 @@ "input_cost_per_token": 1.4e-07, "input_dbu_cost_per_token": 2e-06, "litellm_provider": "databricks", + "token_count_framing": 97, "max_input_tokens": 1000000, "max_output_tokens": 393216, "max_tokens": 393216, @@ -21166,6 +21167,11 @@ "supports_tool_choice": true, "supports_vision": false }, + "databricks/system.ai.deepseek-v4-flash-0731": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 97 + }, "databricks/databricks-deepseek-v4-pro-0813": { "cache_creation_input_token_cost": 1.31999e-06, "cache_read_input_token_cost": 1.3202e-07, @@ -22134,6 +22140,7 @@ "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", + "token_count_framing": 100, "max_input_tokens": 1000000, "max_output_tokens": 1048576, "max_tokens": 1048576, @@ -22157,6 +22164,11 @@ "supports_tool_choice": true, "supports_vision": true }, + "databricks/system.ai.kimi-k3": { + "litellm_provider": "databricks", + "mode": "chat", + "token_count_framing": 100 + }, "databricks/databricks-llama-4-maverick": { "cache_creation_input_token_cost": 5.0001e-07, "cache_read_input_token_cost": 5.0001e-07, @@ -71540,6 +71552,8 @@ "input_cost_per_token_flex": 1.65e-06, "input_cost_per_token_priority": 5.775e-06, "litellm_provider": "bedrock_converse", + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, @@ -71570,6 +71584,8 @@ "input_cost_per_token_flex": 1.5e-06, "input_cost_per_token_priority": 5.25e-06, "litellm_provider": "bedrock_converse", + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, @@ -71600,6 +71616,8 @@ "input_cost_per_token_flex": 1.65e-06, "input_cost_per_token_priority": 5.775e-06, "litellm_provider": "bedrock_converse", + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, @@ -73100,6 +73118,10 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -73173,6 +73195,10 @@ "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "au.anthropic.claude-sonnet-5-5": { + "token_count_model": { + "model": "au.anthropic.claude-opus-5-5", + "region": "ap-southeast-4" + }, "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, @@ -73373,6 +73399,10 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "global.anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -73582,6 +73612,10 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "us.anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -74382,6 +74416,8 @@ "input_cost_per_token_flex": 1.65e-06, "input_cost_per_token_priority": 5.775e-06, "litellm_provider": "bedrock_converse", + "token_count_framing": 64, + "token_count_tokenizer": "kimi_k3", "max_input_tokens": 1000000, "max_output_tokens": 131072, "max_tokens": 131072, @@ -74555,6 +74591,10 @@ "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 1e-07, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -74633,6 +74673,10 @@ "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "au.anthropic.claude-haiku-5-5": { + "token_count_model": { + "model": "au.anthropic.claude-opus-5-5", + "region": "ap-southeast-4" + }, "input_cost_per_token_above_100k_tokens": 5.5e-07, "output_cost_per_token_above_100k_tokens": 2.75e-06, "cache_creation_input_token_cost_above_100k_tokens": 6.875e-07, @@ -74863,6 +74907,10 @@ "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 1e-07, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "global.anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -75010,6 +75058,10 @@ "cache_read_input_token_cost": 1.1e-08, "input_cost_per_token": 1.1e-07, "litellm_provider": "bedrock_converse", + "token_count_model": { + "model": "us.anthropic.claude-opus-5-5", + "region": "us-east-1" + }, "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 33a784e2dae..de57fc0616c 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -1210,6 +1210,36 @@ "additionalProperties": false } }, + "token_count_framing": { + "type": "integer", + "minimum": 0, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing." + }, + "token_count_model": { + "type": "object", + "required": [ + "model" + ], + "additionalProperties": false, + "properties": { + "model": { + "type": "string", + "pattern": "\\S" + }, + "region": { + "type": "string", + "pattern": "\\S" + } + }, + "description": "Model that counts tokens for this one, and the AWS region its count API runs in." + }, + "token_count_tokenizer": { + "type": "string", + "description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.", + "enum": [ + "kimi_k3" + ] + }, "tpm": { "type": "integer", "minimum": 0, diff --git a/tests/e2e/llm_translation/test_count_tokens_source_e2e.py b/tests/e2e/llm_translation/test_count_tokens_source_e2e.py new file mode 100644 index 00000000000..9df4f14ae6e --- /dev/null +++ b/tests/e2e/llm_translation/test_count_tokens_source_e2e.py @@ -0,0 +1,362 @@ +from __future__ import annotations + +import base64 +import io +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from anthropic import Anthropic, BadRequestError, Omit, omit +from anthropic.types import DocumentBlockParam, MessageParam, TextBlockParam, ToolParam +from e2e_config import unique_marker +from e2e_http import unwrap +from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from models import ( + ChatMessage, + CountTokensResponse, + LiteLLMParamsBody, + ModelInfoBody, + ModelNewBody, + ResponsesInputTokensResponse, + TokenCounterBody, + TokenCountModelBody, +) +from proxy_client import ProxyClient +from reportlab.lib.pagesizes import LETTER +from reportlab.pdfgen.canvas import Canvas +from sdk_clients import NO_PROXY_CACHE, SdkClients + +pytestmark = pytest.mark.e2e + +BEDROCK_SONNET: Final = "bedrock/us.anthropic.claude-sonnet-5-5" +BEDROCK_HAIKU: Final = "bedrock/us.anthropic.claude-haiku-5-5" +COST_MAP_COUNT_MODEL: Final = "us.anthropic.claude-opus-5-5" +CONFIGURED_COUNT_MODEL: Final = "us.anthropic.claude-sonnet-5" +ANTHROPIC_SONNET: Final = "anthropic/claude-sonnet-5-5" +BEDROCK_KIMI: Final = "bedrock/us.moonshotai.kimi-k3" +BEDROCK_NOVA: Final = "bedrock/us.amazon.nova-pro-v1:0" +OPENAI_GPT: Final = "openai/gpt-5.6" + +COUNT_TOLERANCE: Final = 0.05 +MAX_TOKENS: Final = 16 +PDF_LINES_PER_PAGE: Final = 45 + +LARGE_SYSTEM: Final = " ".join( + f"Rule {n}: answer in plain English, cite clause {n} when it applies, and never invent a figure." + for n in range(1, 301) +) +LARGE_TOOL: Final[ToolParam] = { + "name": "search_orders", + "description": " ".join( + f"Filter {n} narrows the order search by the attribute recorded in field_{n}." for n in range(1, 151) + ), + "input_schema": { + "type": "object", + "properties": { + f"field_{n}": {"type": "string", "description": f"Order attribute number {n}, matched exactly."} + for n in range(1, 61) + }, + "required": ["field_1"], + }, +} + + +@dataclass(frozen=True, slots=True) +class Served: + model: str + key: str + + +@dataclass(frozen=True, slots=True) +class PromptShape: + name: str + system: str | None = None + tool: ToolParam | None = None + pdf_pages: int = 0 + + +SHAPES: Final = ( + PromptShape(name="baseline"), + PromptShape(name="large_system", system=LARGE_SYSTEM), + PromptShape(name="large_tool", tool=LARGE_TOOL), + PromptShape(name="system_and_tool", system=LARGE_SYSTEM, tool=LARGE_TOOL), + PromptShape(name="pdf_1_page", pdf_pages=1), + PromptShape(name="pdf_12_pages", pdf_pages=12), +) +SHAPE_IDS: Final = tuple(shape.name for shape in SHAPES) +TEXT_SHAPES: Final = SHAPES[:2] +TEXT_SHAPE_IDS: Final = SHAPE_IDS[:2] + + +def _bedrock(model: str) -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=model, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ) + + +def _serve(proxy: ProxyClient, backend: LiteLLMParamsBody, model_info: ModelInfoBody | None = None) -> Iterator[Served]: + resources: Final = ResourceManager(client=proxy) + try: + model: Final = f"e2e-count-source-{unique_marker()}" + model_id: Final = proxy.register_model( + ModelNewBody(model_name=model, litellm_params=backend, model_info=model_info or ModelInfoBody()) + ) + resources.defer(lambda: proxy.delete_model(model_id)) + yield Served(model=model, key=resources.key()) + finally: + resources.teardown() + + +@pytest.fixture(scope="module") +def bedrock_sonnet(proxy: ProxyClient) -> Iterator[Served]: + yield from _serve(proxy, _bedrock(BEDROCK_SONNET)) + + +@pytest.fixture(scope="module") +def bedrock_haiku(proxy: ProxyClient) -> Iterator[Served]: + yield from _serve(proxy, _bedrock(BEDROCK_HAIKU)) + + +@pytest.fixture(scope="module") +def bedrock_sonnet_with_configured_count_model(proxy: ProxyClient) -> Iterator[Served]: + yield from _serve( + proxy, + _bedrock(BEDROCK_SONNET), + ModelInfoBody(token_count_model=TokenCountModelBody(model=CONFIGURED_COUNT_MODEL)), + ) + + +@pytest.fixture(scope="module") +def anthropic_sonnet(proxy: ProxyClient) -> Iterator[Served]: + yield from _serve(proxy, LiteLLMParamsBody(model=ANTHROPIC_SONNET, api_key="os.environ/ANTHROPIC_API_KEY")) + + +@pytest.fixture(scope="module") +def bedrock_kimi(proxy: ProxyClient) -> Iterator[Served]: + yield from _serve(proxy, _bedrock(BEDROCK_KIMI)) + + +@pytest.fixture(scope="module") +def bedrock_nova(proxy: ProxyClient) -> Iterator[Served]: + yield from _serve(proxy, _bedrock(BEDROCK_NOVA)) + + +@pytest.fixture(scope="module") +def openai_gpt(proxy: ProxyClient) -> Iterator[Served]: + yield from _serve(proxy, LiteLLMParamsBody(model=OPENAI_GPT, api_key="os.environ/OPENAI_API_KEY")) + + +def _pdf(pages: int) -> str: + buffer: Final = io.BytesIO() + canvas: Final = Canvas(buffer, pagesize=LETTER, invariant=True) + for page in range(1, pages + 1): + for line in range(1, PDF_LINES_PER_PAGE + 1): + canvas.drawString( + 72, 740 - 14 * line, f"Page {page}, line {line}: revenue, costs and headcount for region {line}." + ) + canvas.showPage() + canvas.save() + return base64.b64encode(buffer.getvalue()).decode() + + +def _user_turn(shape: PromptShape) -> MessageParam: + question: Final[TextBlockParam] = {"type": "text", "text": f"Reply with one word. {unique_marker()}"} + if shape.pdf_pages == 0: + return {"role": "user", "content": [question]} + document: Final[DocumentBlockParam] = { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": _pdf(shape.pdf_pages)}, + } + return {"role": "user", "content": [document, question]} + + +def _count_and_bill(client: Anthropic, model: str, shape: PromptShape) -> tuple[CountTokensResponse, int]: + messages: Final = [_user_turn(shape)] + system: Final[str | Omit] = omit if shape.system is None else shape.system + tools: Final[list[ToolParam] | Omit] = omit if shape.tool is None else [shape.tool] + counted: Final = CountTokensResponse.model_validate_json( + client.messages.with_raw_response.count_tokens(model=model, messages=messages, system=system, tools=tools).text + ) + usage: Final = client.messages.create( + model=model, + max_tokens=MAX_TOKENS, + messages=messages, + system=system, + tools=tools, + extra_body=NO_PROXY_CACHE, + ).usage + billed: Final = usage.input_tokens + (usage.cache_creation_input_tokens or 0) + (usage.cache_read_input_tokens or 0) + return counted, billed + + +def _within_tolerance(counted: int, billed: int) -> bool: + return abs(counted - billed) <= billed * COUNT_TOLERANCE + + +def _assert_counted_by_the_cost_map_count_model(served: Served, sdk: SdkClients, shape: PromptShape) -> None: + counted, billed = _count_and_bill(sdk.anthropic(served.key), served.model, shape) + + assert counted.litellm_count_source == "token_count_model", counted + assert counted.litellm_count_model == COST_MAP_COUNT_MODEL, counted + assert _within_tolerance(counted.input_tokens, billed), (counted, billed) + + +class TestBedrockClaudeCountsWithCostMapCountModel: + @pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS) + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.BEDROCK,), + models=(BEDROCK_SONNET, COST_MAP_COUNT_MODEL), + capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT), + mode=Mode.NONSTREAM, + ) + ) + def test_sonnet_count_names_the_count_model_and_matches_billed_input( + self, bedrock_sonnet: Served, sdk: SdkClients, shape: PromptShape + ) -> None: + _assert_counted_by_the_cost_map_count_model(bedrock_sonnet, sdk, shape) + + @pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS) + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.BEDROCK,), + models=(BEDROCK_HAIKU, COST_MAP_COUNT_MODEL), + capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT), + mode=Mode.NONSTREAM, + ) + ) + def test_haiku_count_names_the_count_model_and_matches_billed_input( + self, bedrock_haiku: Served, sdk: SdkClients, shape: PromptShape + ) -> None: + _assert_counted_by_the_cost_map_count_model(bedrock_haiku, sdk, shape) + + +class TestAnthropicClaudeCountsWithProvider: + @pytest.mark.parametrize("shape", SHAPES, ids=SHAPE_IDS) + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.ANTHROPIC,), + models=(ANTHROPIC_SONNET,), + capabilities=(Capability.FUNCTION_CALLING, Capability.PDF_INPUT), + mode=Mode.NONSTREAM, + ) + ) + def test_count_comes_from_the_provider_and_matches_billed_input( + self, anthropic_sonnet: Served, sdk: SdkClients, shape: PromptShape + ) -> None: + counted, billed = _count_and_bill(sdk.anthropic(anthropic_sonnet.key), anthropic_sonnet.model, shape) + + assert counted.litellm_count_source == "provider", counted + assert counted.litellm_count_model is None, counted + assert _within_tolerance(counted.input_tokens, billed), (counted, billed) + + +class TestBedrockKimiCountsWithFramedLocalEstimate: + @pytest.mark.parametrize("shape", TEXT_SHAPES, ids=TEXT_SHAPE_IDS) + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.BEDROCK,), + models=(BEDROCK_KIMI,), + mode=Mode.NONSTREAM, + ) + ) + def test_text_estimate_matches_billed_input(self, bedrock_kimi: Served, sdk: SdkClients, shape: PromptShape) -> None: + counted, billed = _count_and_bill(sdk.anthropic(bedrock_kimi.key), bedrock_kimi.model, shape) + + assert counted.litellm_count_source == "local_estimate", counted + assert _within_tolerance(counted.input_tokens, billed), (counted, billed) + + +class TestConfiguredCountModelOverridesCostMap: + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.BEDROCK,), + models=(BEDROCK_SONNET, CONFIGURED_COUNT_MODEL), + ) + ) + def test_count_names_the_model_info_count_model( + self, bedrock_sonnet_with_configured_count_model: Served, sdk: SdkClients + ) -> None: + served: Final = bedrock_sonnet_with_configured_count_model + raw: Final = sdk.anthropic(served.key).messages.with_raw_response.count_tokens( + model=served.model, messages=[_user_turn(SHAPES[0])] + ) + counted: Final = CountTokensResponse.model_validate_json(raw.text) + + assert counted.litellm_count_source == "token_count_model", counted + assert counted.litellm_count_model == CONFIGURED_COUNT_MODEL, counted + assert counted.input_tokens > 0, counted + + +class TestProviderCountErrorWithoutCountModel: + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.BEDROCK,), + models=(BEDROCK_NOVA,), + ) + ) + def test_count_keeps_the_provider_400_instead_of_estimating(self, bedrock_nova: Served, sdk: SdkClients) -> None: + with pytest.raises(BadRequestError) as rejected: + _ = sdk.anthropic(bedrock_nova.key).messages.count_tokens( + model=bedrock_nova.model, messages=[_user_turn(SHAPES[0])] + ) + + assert rejected.value.status_code == 400, rejected.value + + +class TestCountSourceOnOtherCountEndpoints: + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.OPENAI,), + models=(OPENAI_GPT,), + ) + ) + def test_responses_input_tokens_says_the_provider_counted(self, openai_gpt: Served, sdk: SdkClients) -> None: + raw: Final = sdk.openai(openai_gpt.key).responses.input_tokens.with_raw_response.count( + model=openai_gpt.model, input=f"Reply with one word. {unique_marker()}" + ) + counted: Final = ResponsesInputTokensResponse.model_validate_json(raw.text) + + assert counted.litellm_count_source == "provider", counted + assert counted.input_tokens > 0, counted + + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.COUNT_TOKENS, + providers=(Provider.BEDROCK,), + models=(BEDROCK_KIMI,), + ) + ) + def test_token_counter_says_kimi_was_estimated_locally(self, proxy: ProxyClient, bedrock_kimi: Served) -> None: + counted: Final = unwrap( + proxy.token_counter( + bedrock_kimi.key, + TokenCounterBody( + model=bedrock_kimi.model, + messages=[ChatMessage(role="user", content=f"Reply with one word. {unique_marker()}")], + ), + ) + ) + + assert counted.count_source == "local_estimate", counted + assert counted.total_tokens > 0, counted diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 40e9e181982..f0a8d67b444 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -677,13 +677,39 @@ class AnthropicMessagesResponse(BaseModel): usage: Usage | None = None -class CountTokensResponse(BaseModel): +TokenCountSource = Literal["provider", "token_count_model", "local_estimate"] + + +class LitellmCountSource(BaseModel, frozen=True): + litellm_count_source: TokenCountSource | None = None + litellm_count_model: str | None = None + + +class CountTokensResponse(LitellmCountSource, frozen=True): """`/v1/messages/count_tokens` answer. `input_tokens` is required so a 200 whose body lacks it fails validation instead of passing vacuously.""" input_tokens: int +class ResponsesInputTokensResponse(LitellmCountSource, frozen=True): + object: str + input_tokens: int + + +class TokenCounterBody(BaseModel, frozen=True): + model: str + messages: list[ChatMessage] + + +class TokenCounterResponse(BaseModel, frozen=True): + total_tokens: int + request_model: str + model_used: str + tokenizer_type: str + count_source: TokenCountSource | None = None + + class AnthropicErrorBody(BaseModel): type: str message: str @@ -1392,6 +1418,11 @@ class LiteLLMParamsBody(BaseModel): ModelMode = Literal["chat", "batch", "realtime", "image_generation"] +class TokenCountModelBody(BaseModel, frozen=True): + model: str + region: str | None = None + + class ModelInfoBody(BaseModel): # id is left unset so the proxy assigns a unique model_id per deployment. # Pinning it to the model_name made re-registrations of a fixed-name model @@ -1404,6 +1435,7 @@ class ModelInfoBody(BaseModel): team_id: str | None = None allowed_fails: int | None = None allowed_fails_policy: dict[str, int] | None = None + token_count_model: TokenCountModelBody | None = None class ModelNewBody(BaseModel): diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 23ab6487889..8cc7bf0c2b4 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -100,6 +100,8 @@ from models import ( TeamNewBody, TeamNewResponse, TeamUpdateBody, + TokenCounterBody, + TokenCounterResponse, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, @@ -1049,6 +1051,15 @@ class ProxyClient: response_type=CountTokensResponse, ) + @step("Count tokens with /utils/token_counter for {body.model}") + def token_counter(self, key: str, body: TokenCounterBody) -> Result[TokenCounterResponse]: + return self.transport.post( + "/utils/token_counter", + headers=self.transport.bearer(key), + json=body, + response_type=TokenCounterResponse, + ) + @step("Send a /v1/messages request to {body.model}") def messages( self, key: str, body: AnthropicMessagesBody, *, session_id: str | None = None diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py index 92512a304bc..1009a954bee 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py @@ -104,7 +104,7 @@ def test_anthropic_count_tokens_forwards_body_and_refuses_a_key_without_model_ac client: Final = _client(gateway, gateway.key) result: Final = _sdk_count(client, model, spelling, _COUNT_BETA) - assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS} + assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS, "litellm_count_source": "provider"} requests: Final = wire.drain() assert len(requests) == 1, requests @@ -236,5 +236,5 @@ def test_anthropic_count_tokens_merges_the_client_beta_with_the_token_counting_b num_retries=0, ) result: Final = _sdk_count(_client(gateway, gateway.key), model, spelling, _CLIENT_BETA) - assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS} + assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS, "litellm_count_source": "provider"} assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py index 8a79d5c6a7a..b13c34fe2d6 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py @@ -245,7 +245,10 @@ def test_enabled_gateway_count_tokens_alias_matches_native_count_tokens_route(en "/claude_code_gateway/v1/messages/count_tokens?beta=true", json=body, headers=headers ) assert gateway_response.status_code == 200, gateway_response.text - assert _JSON_OBJECT.validate_json(gateway_response.content) == {"input_tokens": 41}, gateway_response.text + assert _JSON_OBJECT.validate_json(gateway_response.content) == { + "input_tokens": 41, + "litellm_count_source": "provider", + }, gateway_response.text native_response: Final = enabled_gateway.client.post( "/v1/messages/count_tokens?beta=true", json=body, headers=headers ) diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index c3cd4c5b0e7..afd25fcfa04 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -274,7 +274,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run( ) -> None: with _listening(deployment) as peer: response: Final = _count(gateway, count_request(deployment.model, case)) - assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text assert _bodies(peer, deployment.provider) == (expected_count_body(_MODEL, case),) @@ -329,7 +329,11 @@ def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gate "/v1/responses/input_tokens", {"model": deployment.model, "input": USER_TEXT, "instructions": INSTRUCTION}, ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _COUNT, + "litellm_count_source": "provider", + }, response.text assert _bodies(peer, deployment.provider) == (_STRING_BODY,) @@ -342,7 +346,11 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite "/v1/responses/input_tokens", {"model": deployment.model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION}, ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _COUNT, + "litellm_count_source": "provider", + }, response.text assert _bodies(peer, deployment.provider) == ( {"model": _MODEL, "messages": [USER], "system": [LIFTED, {"type": "text", "text": REMINDER}]}, ) @@ -373,7 +381,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system( ) -> None: with _listening(deployment) as peer: response: Final = _count(gateway, {**count_request(deployment.model, STRING_CASE), "tools": TOOLS}) - assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text assert _bodies(peer, deployment.provider) == ({**_STRING_BODY, "tools": TOOLS},) @@ -386,7 +394,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place( "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], } response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text assert _bodies(peer, deployment.provider) == ( {"model": _MODEL, "messages": [USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], "system": [LIFTED]}, ) @@ -418,7 +426,7 @@ def test_messages_count_tokens_repeated_request_lifts_each_time(deployment: _Dep answers: Final = tuple( _payload(_count(gateway, count_request(deployment.model, STRING_CASE))) for _ in range(2) ) - assert answers == ({"input_tokens": _COUNT},) * 2 + assert answers == ({"input_tokens": _COUNT, "litellm_count_source": "provider"},) * 2 assert _bodies(peer, deployment.provider) == (_STRING_BODY,) * 2 @@ -432,7 +440,7 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_ assert "input_tokens" not in response.text, response.text assert peer.drain() == () follow_up: Final = _count(gateway, count_request(deployment.model, STRING_CASE)) - assert _payload(follow_up) == {"input_tokens": _COUNT}, follow_up.text + assert _payload(follow_up) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, follow_up.text assert _bodies(peer, deployment.provider) == (_STRING_BODY,) @@ -447,7 +455,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value( content=f'{{"model": "{deployment.model}", "messages": {first}, "messages": {last}}}', headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) - assert _payload(response) == {"input_tokens": _COUNT}, response.text + assert _payload(response) == {"input_tokens": _COUNT, "litellm_count_source": "provider"}, response.text assert _bodies(peer, deployment.provider) == (_STRING_BODY,) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 701d185b241..ab0c35d0ca4 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -378,7 +378,7 @@ def test_messages_count_tokens_counts_through_mantle_when_the_runtime_cannot( ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) @@ -394,7 +394,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_a_bare_body_to_mant ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) @@ -429,7 +429,11 @@ def test_gemini_count_tokens_route_reports_the_mantle_total(counting_proxy: Owne ): model: Final = _deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) - assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert _payload(response) == { + "totalTokens": _MANTLE_COUNT, + "promptTokensDetails": [], + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == ({"model": _OPUS_BASE, "messages": _GEMINI_MESSAGES},) @@ -445,7 +449,11 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock ): model: Final = _deployment(scenario, runtime.url, _SONNET) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", _GEMINI_BODY) - assert _payload(response) == {"totalTokens": _RUNTIME_COUNT, "promptTokensDetails": []}, response.text + assert _payload(response) == { + "totalTokens": _RUNTIME_COUNT, + "promptTokensDetails": [], + "litellm_count_source": "provider", + }, response.text assert _runtime_count_targets(runtime) == (f"/model/{_SONNET}/count-tokens",) assert mantle.drain() == () @@ -460,7 +468,7 @@ def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) @@ -520,7 +528,11 @@ def test_google_count_tokens_forwards_system_and_tools_to_mantle( ): model: Final = _deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) - assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert _payload(response) == { + "totalTokens": _MANTLE_COUNT, + "promptTokensDetails": [], + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == ( _mantle_body( @@ -565,7 +577,11 @@ def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle( ): model: Final = _deployment(scenario, runtime.url) response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) - assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert _payload(response) == { + "totalTokens": _MANTLE_COUNT, + "promptTokensDetails": [], + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == ( _mantle_body( @@ -595,7 +611,11 @@ def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy response: Final = gateway.request( "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _MANTLE_COUNT, + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) @@ -640,7 +660,10 @@ def test_model_the_runtime_counts_never_reaches_mantle(counting_proxy: OwnedProx gateway.scenario() as scenario, ): model: Final = _deployment(scenario, runtime.url, _SONNET) - assert _payload(_count(gateway, _full(model))) == {"input_tokens": _RUNTIME_COUNT} + assert _payload(_count(gateway, _full(model))) == { + "input_tokens": _RUNTIME_COUNT, + "litellm_count_source": "provider", + } detailed: Final = gateway.request( "POST", "/utils/token_counter", _full(model), params={"call_endpoint": "true"} ) @@ -728,7 +751,11 @@ def test_responses_input_tokens_with_instructions_counts_through_mantle( "/v1/responses/input_tokens", {"model": model, "input": "Count this message", "instructions": "Be terse"}, ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _MANTLE_COUNT, + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=[{"type": "text", "text": "Be terse"}]),) @@ -745,7 +772,7 @@ def test_messages_count_tokens_lifts_the_leading_system_run_for_mantle( ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, count_request(model, case)) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (expected_count_body(_OPUS_BASE, case),) @@ -858,7 +885,11 @@ def test_responses_input_tokens_lifts_instructions_ahead_of_a_leading_system_ite "/v1/responses/input_tokens", {"model": model, "input": [MID_SYSTEM, USER], "instructions": INSTRUCTION}, ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _MANTLE_COUNT, + "litellm_count_source": "provider", + }, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED, {"type": "text", "text": REMINDER}]),) @@ -874,7 +905,7 @@ def test_messages_count_tokens_forwards_tools_beside_the_lifted_system_to_mantle ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**count_request(model, STRING_CASE), "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=[LIFTED], tools=_TOOLS),) @@ -894,7 +925,7 @@ def test_messages_count_tokens_keeps_a_mid_conversation_system_in_place_for_mant "messages": [LEADING, USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], } response: Final = _count(gateway, body) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == ( _mantle_body(messages=[USER, MID_SYSTEM, ASSISTANT, FOLLOW_UP], system=[LIFTED]), @@ -947,7 +978,7 @@ def test_messages_count_tokens_repeated_request_lifts_the_leading_system_each_ti ): model: Final = _deployment(scenario, runtime.url) answers: Final = tuple(_payload(_count(gateway, count_request(model, STRING_CASE))) for _ in range(2)) - assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2 + assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) * 2 @@ -968,7 +999,9 @@ def test_messages_count_tokens_surfaces_the_counting_error_for_a_leading_system_ assert "input_tokens" not in response.text, response.text assert mantle.drain() == () follow_up: Final = _count(gateway, count_request(model, STRING_CASE)) - assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT}, follow_up.text + assert _payload(follow_up) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, ( + follow_up.text + ) assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) @@ -990,7 +1023,7 @@ def test_messages_count_tokens_duplicate_messages_key_lifts_the_last_value_for_m content=f'{{"model": "{model}", "messages": {first}, "messages": {last}}}', headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_LIFTED,) @@ -1015,7 +1048,9 @@ def test_disabled_token_counter_counts_a_leading_system_through_mantle( ) with wire_server(_mantle(_strict), port=mantle_port) as strict: counted: Final = _count(owned.gateway, count_request(_OWNED_OPUS, STRING_CASE)) - assert _payload(counted) == {"input_tokens": _MANTLE_COUNT}, counted.text + assert _payload(counted) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, ( + counted.text + ) assert _mantle_bodies(strict) == (_MANTLE_LIFTED,) assert len(_runtime_count_targets(runtime)) == 1 @@ -1220,7 +1255,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim_to_mantle( ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "system": system}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_mantle_body(system=system),) @@ -1241,7 +1276,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 (sent,) = mantle.drain() assert _mantle_requests((sent,)) == (_MANTLE_FULL,) @@ -1259,7 +1294,7 @@ def test_messages_count_tokens_leaves_empty_tools_and_system_out_of_the_mantle_b ): model: Final = _deployment(scenario, runtime.url) response: Final = _count(gateway, {**_bare(model), "tools": [], "system": ""}) - assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert _payload(response) == {"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"}, response.text assert len(_runtime_count_targets(runtime)) == 1 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) @@ -1275,7 +1310,7 @@ def test_messages_count_tokens_repeated_request_reaches_both_peers_each_time( ): model: Final = _deployment(scenario, runtime.url) answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) - assert answers == ({"input_tokens": _MANTLE_COUNT},) * 2 + assert answers == ({"input_tokens": _MANTLE_COUNT, "litellm_count_source": "provider"},) * 2 assert _runtime_count_targets(runtime) == (f"/model/{_OPUS_BASE}/count-tokens",) * 2 assert _mantle_bodies(mantle) == (_MANTLE_BARE,) * 2 @@ -1304,9 +1339,15 @@ def test_disabled_token_counter_surfaces_the_mantle_error_instead_of_counting_lo assert _REJECTION in refused.text and "input_tokens" not in refused.text, refused.text assert _mantle_bodies(refusing) == (_MANTLE_FULL,) with wire_server(_mantle(), port=mantle_port) as counting: - assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == {"input_tokens": _MANTLE_COUNT} + assert _payload(_count(owned.gateway, _full(_OWNED_OPUS))) == { + "input_tokens": _MANTLE_COUNT, + "litellm_count_source": "provider", + } assert _mantle_bodies(counting) == (_MANTLE_FULL,) - assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == {"input_tokens": _RUNTIME_COUNT} + assert _payload(_count(owned.gateway, _full(_OWNED_SONNET))) == { + "input_tokens": _RUNTIME_COUNT, + "litellm_count_source": "provider", + } rejected: Final = _count(owned.gateway, _full(_OWNED_NOVA)) assert rejected.status_code == 400, rejected.text assert _UNSUPPORTED in rejected.text and "input_tokens" not in rejected.text, rejected.text diff --git a/tests/integration/providers/test_count_tokens_model_wire.py b/tests/integration/providers/test_count_tokens_model_wire.py new file mode 100644 index 00000000000..46e5396ef64 --- /dev/null +++ b/tests/integration/providers/test_count_tokens_model_wire.py @@ -0,0 +1,719 @@ +import json +import re +from collections.abc import Callable, Iterator, Mapping +from itertools import chain +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import anthropic +import httpx +import pytest +from anthropic.types import MessageCountTokensToolParam, MessageTokensCount +from integration._support.bedrock_runtime_peer import target_of +from integration._support.client import Gateway, Scenario, gateway_from_environment +from integration._support.ports import allocate_port +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_SONNET: Final = "us.anthropic.claude-sonnet-5-5" +_SONNET_BASE: Final = "anthropic.claude-sonnet-5-5" +_OPUS: Final = "us.anthropic.claude-opus-5-5" +_OPUS_BASE: Final = "anthropic.claude-opus-5-5" +_NOVA: Final = "amazon.nova-lite-v1:0" +_KIMI: Final = "us.moonshotai.kimi-k3" +_KIMI_BASE: Final = "moonshotai.kimi-k3" +_PROFILE_ID: Final = "lit9563count" +_PROFILE_ARN: Final = f"arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/{_PROFILE_ID}" +_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" +_REGION: Final = "us-west-2" +_COUNT_REGION: Final = "us-east-1" +_ACCESS_KEY: Final = "AKIAINTEGRATIONCOUNT" +_SECRET_KEY: Final = "integration-count-secret" +_PROVIDER_COUNT: Final = 1345 +_COUNT_MODEL_COUNT: Final = 2718 +_FRAMING: Final = 50 +_UNSUPPORTED: Final = "The provided model doesn't support counting tokens." +_COUNT_TARGET: Final = re.compile(r"^/model/(.+)/count-tokens$") +_SCRIPTED_STATUS: Final = re.compile(r"status=(\d{3})") +_SCRIPTED_COUNT_MODEL_STATUS: Final = re.compile(r"count_model=(\d{3})") +_CREDENTIAL: Final = re.compile(r"Credential=([^/]+)/\d+/([^/]+)/") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, JsonValue]]) + +_SYSTEM: Final = "You are a terse assistant that answers in one sentence" +_WEATHER_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string", "description": "City to look up"}}, + "required": ["city"], +} +_SDK_TOOLS: Final[list[MessageCountTokensToolParam]] = [ + { + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City to look up"}}, + "required": ["city"], + }, + } +] +_COUNT_MODEL_SOURCE: Final = {"litellm_count_source": "token_count_model", "litellm_count_model": _OPUS} +_OPENAI_PRIMARY: Final = "gpt-count-primary" +_OPENAI_COUNT_MODEL: Final = "gpt-count-model" +_OPENAI_COUNT_TARGET: Final = "/v1/responses/input_tokens" +_COUNTED_FIELDS: Final = ("input", "instructions", "tools", "text") +_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}} +_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = { + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": _WEATHER_SCHEMA, +} +_TOOL_TURNS: Final[list[JsonValue]] = [ + {"role": "user", "content": "What is the weather in Paris?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "Sunny"}]}, +] +_LOCAL_SOURCE: Final = {"litellm_count_source": "local_estimate"} + + +def _prompt(status: int, count_model_status: int | None = None) -> str: + scripted: Final = "" if count_model_status is None else f" count_model={count_model_status}" + return f"Count this message status={status}{scripted}" + + +def _converse_body(prompt: str) -> dict[str, JsonValue]: + return { + "input": { + "converse": { + "messages": [{"role": "user", "content": [{"text": prompt}]}], + "system": [{"text": _SYSTEM}], + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": "get_weather", + "description": "Look up the current weather for a city", + "inputSchema": {"json": _WEATHER_SCHEMA}, + } + } + ] + }, + } + } + } + + +def _json_reply(status: int, payload: Mapping[str, JsonValue]) -> Reply: + return Reply(status=status, body=json.dumps(payload).encode()) + + +def _counted_model(request: Request) -> str: + counted: Final = _COUNT_TARGET.match(target_of(request)) + return target_of(request) if counted is None else counted.group(1) + + +def _scripted(pattern: re.Pattern[str], request: Request, default: int) -> int: + scripted: Final = pattern.search(request.body.decode(errors="replace")) + return default if scripted is None else int(scripted.group(1)) + + +def _runtime(request: Request) -> Reply: + model: Final = _counted_model(request) + if model == _OPUS_BASE: + count_model_status: Final = _scripted(_SCRIPTED_COUNT_MODEL_STATUS, request, 200) + if count_model_status == 200: + return _json_reply(200, {"inputTokens": _COUNT_MODEL_COUNT}) + return _json_reply(count_model_status, {"message": f"scripted runtime {count_model_status} for {model}"}) + status: Final = _scripted(_SCRIPTED_STATUS, request, 500) + if status == 200: + return _json_reply(200, {"inputTokens": _PROVIDER_COUNT}) + return _json_reply(status, {"message": _UNSUPPORTED if status == 400 else f"scripted runtime {status}"}) + + +def _mantle(request: Request) -> Reply: + model: Final = _JSON_OBJECT.validate_json(request.body).get("model") + return _json_reply( + 404, {"type": "error", "error": {"type": "not_found_error", "message": f"scripted mantle 404 for {model}"}} + ) + + +def _signer(request: Request) -> tuple[str, str]: + credential: Final = _CREDENTIAL.search(request.headers.get("authorization", "")) + return ("", "") if credential is None else (credential.group(1), credential.group(2)) + + +def _runtime_counts(wire: Wire) -> tuple[tuple[str, tuple[str, str], dict[str, JsonValue]], ...]: + return tuple( + (_counted_model(request), _signer(request), _JSON_OBJECT.validate_json(request.body)) + for request in wire.drain() + ) + + +def _runtime_models(wire: Wire) -> tuple[str, ...]: + return tuple(_counted_model(request) for request in wire.drain()) + + +def _mantle_models(wire: Wire) -> tuple[JsonValue, ...]: + return tuple(_JSON_OBJECT.validate_json(request.body).get("model") for request in wire.drain()) + + +def _url(port: int) -> str: + return f"http://127.0.0.1:{port}" + + +def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]: + return { + "model": f"bedrock/{model}", + "api_base": api_base, + "aws_access_key_id": _ACCESS_KEY, + "aws_secret_access_key": _SECRET_KEY, + "aws_region_name": _REGION, + } + + +@pytest.fixture(scope="module") +def runtime_port() -> int: + return allocate_port() + + +@pytest.fixture(scope="module") +def mantle_port() -> int: + return allocate_port() + + +@pytest.fixture(scope="module") +def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-model"), + { + "BEDROCK_MANTLE_API_BASE": _url(mantle_port), + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "HF_HUB_OFFLINE": "1", + }, + workers=1, + remove_environment=("AWS_BEARER_TOKEN_BEDROCK",), + ) as owned, + ): + yield owned + + +def _deployment( + scenario: Scenario, runtime_port: int, model: str = _SONNET, model_info: Mapping[str, JsonValue] | None = None +) -> str: + return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port))) + + +def _count(gateway: Gateway, model: str, status: int, count_model_status: int | None = None) -> MessageTokensCount: + with anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) as client: + return client.messages.count_tokens( + model=model, + messages=[{"role": "user", "content": _prompt(status, count_model_status)}], + system=_SYSTEM, + tools=_SDK_TOOLS, + ) + + +def _messages_estimate(gateway: Gateway, model: str) -> tuple[int, JsonValue]: + counted: Final = _count(gateway, model, 400) + return counted.input_tokens, (counted.model_extra or {}).get("litellm_count_source") + + +def _utils_estimate(gateway: Gateway, model: str) -> tuple[int, JsonValue]: + response: Final = gateway.request( + "POST", "/utils/token_counter", {"model": model, "messages": [{"role": "user", "content": _prompt(400)}]} + ) + assert response.status_code == 200, response.text + body: Final = _JSON_OBJECT.validate_json(response.content) + total: Final = body["total_tokens"] + assert isinstance(total, int), body + return total, body.get("count_source") + + +def _shipped_framing(model: str) -> JsonValue: + return _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())[model].get("token_count_framing") + + +def test_shipped_token_count_model_counts_the_request_in_its_region_without_trying_the_model( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + counted: Final = _count(gateway, _deployment(scenario, runtime_port), 404) + assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted + assert _runtime_counts(runtime) == ((_OPUS_BASE, (_ACCESS_KEY, _COUNT_REGION), _converse_body(_prompt(404))),) + assert _mantle_models(mantle) == () + + +@pytest.mark.parametrize( + ("count_model_status", "error", "message"), + [ + (401, anthropic.AuthenticationError, f"scripted runtime 401 for {_OPUS_BASE}"), + (429, anthropic.RateLimitError, f"scripted runtime 429 for {_OPUS_BASE}"), + (404, anthropic.NotFoundError, f"scripted mantle 404 for {_OPUS_BASE}"), + ], +) +def test_token_count_model_failure_answers_its_error_and_names_the_region_to_enable( + counting_proxy: OwnedProxy, + runtime_port: int, + mantle_port: int, + count_model_status: int, + error: type[anthropic.APIStatusError], + message: str, +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + with pytest.raises(error) as refused: + _count(gateway, _deployment(scenario, runtime_port), 200, count_model_status=count_model_status) + assert message in str(refused.value) and f"{_OPUS} in {_COUNT_REGION}" in str(refused.value), refused.value + assert _runtime_models(runtime) == (_OPUS_BASE,) + assert _mantle_models(mantle) == ((_OPUS_BASE,) if count_model_status == 404 else ()) + + +def test_framing_is_added_to_the_local_estimate_when_the_runtime_cannot_count( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + framed: Final = _count( + gateway, _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": _FRAMING}), 400 + ) + unframed: Final = _count(gateway, _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": 0}), 400) + assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed) + assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed) + assert _runtime_models(runtime) == (_NOVA, _NOVA) + assert _mantle_models(mantle) == () + + +def test_framing_on_one_deployment_leaves_a_sibling_of_the_same_model_raising_the_400( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + framed: Final = _deployment(scenario, runtime_port, _NOVA, {"token_count_framing": _FRAMING}) + sibling: Final = _deployment(scenario, runtime_port, _NOVA) + assert _count(gateway, framed, 400).model_extra == _LOCAL_SOURCE + with pytest.raises(anthropic.BadRequestError, match=re.escape(_UNSUPPORTED)): + _count(gateway, sibling, 400) + assert _runtime_models(runtime) == (_NOVA, _NOVA) + assert _mantle_models(mantle) == () + + +def test_empty_configured_token_count_model_counts_the_model_itself_in_the_deployment_region( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + counted: Final = _count(gateway, _deployment(scenario, runtime_port, model_info={"token_count_model": {}}), 200) + assert (counted.input_tokens, counted.model_extra) == (_PROVIDER_COUNT, {"litellm_count_source": "provider"}) + assert _runtime_counts(runtime) == ((_SONNET_BASE, (_ACCESS_KEY, _REGION), _converse_body(_prompt(200))),) + assert _mantle_models(mantle) == () + + +@pytest.mark.parametrize("base_model_in", ["model_info", "litellm_params"]) +def test_inference_profile_arn_counts_with_the_token_count_model_of_its_base_model( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int, base_model_in: str +) -> None: + gateway: Final = counting_proxy.gateway + base_model: Final[dict[str, JsonValue]] = {"base_model": f"bedrock/{_SONNET}"} + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model_info=base_model if base_model_in == "model_info" else None, + api_key=None, + **_litellm_params(_PROFILE_ARN, _url(runtime_port)), + **(base_model if base_model_in == "litellm_params" else {}), + ) + counted: Final = _count(gateway, model, 404) + assert (counted.input_tokens, counted.model_extra) == (_COUNT_MODEL_COUNT, _COUNT_MODEL_SOURCE), counted + assert _runtime_models(runtime) == (_OPUS_BASE,) + assert _mantle_models(mantle) == () + + +@pytest.mark.parametrize( + ("estimate", "runtime_calls"), + [(_messages_estimate, (_KIMI_BASE, _KIMI_BASE)), (_utils_estimate, ())], + ids=["messages_count_tokens", "utils_token_counter"], +) +def test_shipped_framing_is_added_to_the_estimate_and_a_configured_zero_turns_it_off( + counting_proxy: OwnedProxy, + runtime_port: int, + mantle_port: int, + estimate: Callable[[Gateway, str], tuple[int, JsonValue]], + runtime_calls: tuple[str, ...], +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + shipped: Final = estimate(gateway, _deployment(scenario, runtime_port, _KIMI)) + unframed: Final = estimate(gateway, _deployment(scenario, runtime_port, _KIMI, {"token_count_framing": 0})) + assert (shipped[1], unframed[1]) == ("local_estimate", "local_estimate"), (shipped, unframed) + framing: Final = _shipped_framing(_KIMI) + assert isinstance(framing, int) and framing > 0, framing + assert shipped[0] - unframed[0] == framing, (shipped, unframed) + assert _runtime_models(runtime) == runtime_calls + assert _mantle_models(mantle) == () + + +def _openai_count(request: Request) -> Reply: + model: Final = _JSON_OBJECT.validate_json(request.body).get("model") + if request.target == _OPENAI_COUNT_TARGET and model == _OPENAI_COUNT_MODEL: + return _json_reply(200, {"object": "response.input_tokens", "input_tokens": _COUNT_MODEL_COUNT}) + return _json_reply( + 404, + {"error": {"message": f"scripted openai 404 for {model}", "type": "invalid_request_error", "code": None}}, + ) + + +def _counted_fields(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {key: body[key] for key in _COUNTED_FIELDS if key in body} + + +def _openai_counts(wire: Wire) -> tuple[tuple[str, JsonValue, dict[str, JsonValue]], ...]: + bodies: Final = tuple((request.target, _JSON_OBJECT.validate_json(request.body)) for request in wire.drain()) + return tuple((target, body.get("model"), _counted_fields(body)) for target, body in bodies) + + +def _openai_deployment(scenario: Scenario, wire: Wire, model_info: Mapping[str, JsonValue]) -> str: + return scenario.model(model=f"openai/{_OPENAI_PRIMARY}", api_base=f"{wire.url}/v1", model_info=model_info) + + +def test_token_count_model_counts_the_anthropic_body_converted_like_the_provider_count( + counting_proxy: OwnedProxy, +) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_openai_count) as openai, gateway.scenario() as scenario: + model: Final = _openai_deployment(scenario, openai, {"token_count_model": {"model": _OPENAI_COUNT_MODEL}}) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": model, + "system": _SYSTEM, + "tools": [_ANTHROPIC_TOOL], + "output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}}, + "messages": _TOOL_TURNS, + }, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content) == { + "input_tokens": _COUNT_MODEL_COUNT, + "litellm_count_source": "token_count_model", + "litellm_count_model": _OPENAI_COUNT_MODEL, + } + counts: Final = _openai_counts(openai) + assert tuple((target, counted) for target, counted, _ in counts) == ( + (_OPENAI_COUNT_TARGET, _OPENAI_COUNT_MODEL), + ), counts + primary: Final = counts[0][2] + assert primary["instructions"] == _SYSTEM, primary + text: Final = primary.get("text") + assert isinstance(text, dict) and isinstance(text.get("format"), dict), primary + input_items: Final = primary["input"] + assert isinstance(input_items, list), primary + item_types: Final = {item.get("type") for item in input_items if isinstance(item, dict)} + assert {"function_call", "function_call_output"} <= item_types, input_items + + +@pytest.mark.parametrize( + ("body", "named"), + [ + ( + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + { + "type": "search_result", + "source": "https://example.com/notes", + "title": "Notes", + "content": [], + }, + ], + } + ] + }, + "search_result", + ), + ( + {"messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]}, + "messages.1.role: Field required", + ), + ], + ids=["unconvertible_block", "missing_role"], +) +def test_framing_keeps_the_400_for_a_body_the_openai_count_rejects_before_calling_openai( + counting_proxy: OwnedProxy, body: Mapping[str, JsonValue], named: str +) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_openai_count) as openai, gateway.scenario() as scenario: + model: Final = _openai_deployment(scenario, openai, {"token_count_framing": _FRAMING}) + rejected: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **body}) + assert rejected.status_code == 400, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict) and named in str(error.get("message")), error + assert openai.drain() == () + estimated: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": "hello"}]} + ) + assert estimated.status_code == 200, estimated.text + assert _JSON_OBJECT.validate_json(estimated.content).get("litellm_count_source") == "local_estimate" + assert tuple(counted for _, counted, _ in _openai_counts(openai)) == (_OPENAI_PRIMARY,) + + +def _runtime_cannot_count(_request: Request) -> Reply: + return _json_reply(400, {"message": _UNSUPPORTED}) + + +_REPLAYED_WEB_SEARCH: Final[list[JsonValue]] = [ + {"role": "user", "content": "What is the newest LiteLLM release?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "LiteLLM"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + } + ], + }, + {"type": "text", "text": "The newest release is in the release notes."}, + ], + }, + {"role": "user", "content": "Thanks."}, +] +_BOOLEAN_SUBSCHEMA_TOOL: Final[dict[str, JsonValue]] = { + "name": "lookup", + "description": "Look up a value", + "input_schema": {"type": "object", "properties": {"query": True}}, +} +_UNREADABLE_BODIES: Final[dict[str, dict[str, JsonValue]]] = { + "web_search_replay": {"messages": _REPLAYED_WEB_SEARCH}, + "boolean_subschema_tool": { + "messages": [{"role": "user", "content": "Look up the value."}], + "tools": [_BOOLEAN_SUBSCHEMA_TOOL], + }, +} +_COUNT_ENDPOINTS: Final[dict[str, tuple[str, str]]] = { + "messages_count_tokens": ("/v1/messages/count_tokens", "litellm_count_source"), + "utils_token_counter": ("/utils/token_counter?call_endpoint=true", "count_source"), +} + + +def _count_answer(response: httpx.Response, source_key: str) -> tuple[int, str]: + if response.status_code not in (200, 400): + return response.status_code, response.text + body: Final = _JSON_OBJECT.validate_json(response.content) + error: Final = body.get("error") + if response.status_code == 400 and isinstance(error, dict): + return 400, str(error.get("message")) + return response.status_code, str(body.get(source_key)) + + +@pytest.mark.parametrize("endpoint", sorted(_COUNT_ENDPOINTS)) +@pytest.mark.parametrize("unreadable", sorted(_UNREADABLE_BODIES)) +def test_shipped_framing_never_answers_a_500_for_a_body_the_local_count_cannot_read( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int, unreadable: str, endpoint: str +) -> None: + gateway: Final = counting_proxy.gateway + path, source_key = _COUNT_ENDPOINTS[endpoint] + with ( + wire_server(_runtime_cannot_count, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + response: Final = gateway.request( + "POST", path, {"model": _deployment(scenario, runtime_port, _KIMI), **_UNREADABLE_BODIES[unreadable]} + ) + status, answer = _count_answer(response, source_key) + assert (status, answer) == (200, "local_estimate") or (status == 400 and _UNSUPPORTED in answer), ( + response.status_code, + response.text, + ) + assert _runtime_models(runtime) == (_KIMI_BASE,) + assert _mantle_models(mantle) == () + + +_GEMINI_MODEL: Final = "gemini-4-argon" +_GEMINI_COUNT_MODEL: Final = "gemini-3.8-flash" +_VERTEX_PROJECT: Final = "scripted-project" +_COMPACT_UNDER_THE_MINIMUM_TRIGGER: Final[dict[str, JsonValue]] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] +} +_CLEAR_OLDER_TOOL_RESULTS: Final[dict[str, JsonValue]] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] +} +_READ_FILE_TOOL: Final[dict[str, JsonValue]] = { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_FILE_TEXT: Final = "line of a long file that the agent read earlier. " * 180 +_FILES: Final = 3 + + +def _google_cannot_count(_request: Request) -> Reply: + return _json_reply(400, {"error": {"code": 400, "message": "scripted count 400", "status": "INVALID_ARGUMENT"}}) + + +def _google_counts_only_the_count_model(request: Request) -> Reply: + if request.target.endswith(f"/{_GEMINI_COUNT_MODEL}:countTokens"): + return _json_reply(200, {"totalTokens": _COUNT_MODEL_COUNT}) + return _google_cannot_count(request) + + +def _google_deployment( + gateway: Gateway, scenario: Scenario, wire: Wire, provider: str, model_info: dict[str, JsonValue] +) -> str: + if provider == "gemini": + return scenario.model( + model=f"gemini/{_GEMINI_MODEL}", api_base=wire.url, api_key="scripted-gemini-key", model_info=model_info + ) + return scenario.model( + model=f"vertex_ai/{_GEMINI_MODEL}", + api_base=wire.url, + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location="us-east5", + vertex_credentials=service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")), + model_info=model_info, + ) + + +def _google_count_targets(wire: Wire) -> tuple[bool, ...]: + return tuple(request.target.endswith(":countTokens") for request in wire.drain()) + + +def _read_file_round(marker: str, index: int) -> tuple[JsonValue, ...]: + return ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}.txt"}} + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": f"toolu_{index}", + "content": f"{_FILE_TEXT}", + } + ], + }, + ) + + +def _agent_messages(marker: str) -> list[JsonValue]: + rounds: Final = tuple(chain.from_iterable(_read_file_round(marker, index) for index in range(_FILES))) + return [*rounds, {"role": "user", "content": "Summarize what you read."}] + + +def _counted_files(request: Request, marker: str) -> tuple[bool, tuple[bool, ...]]: + sent: Final = json.dumps(_JSON_OBJECT.validate_json(request.body)) + return ( + request.target.endswith(f"/{_GEMINI_COUNT_MODEL}:countTokens"), + tuple(f"" in sent for index in range(_FILES)), + ) + + +@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"]) +def test_framing_keeps_the_400_for_a_context_edit_the_gemini_count_rejects_before_calling_google( + counting_proxy: OwnedProxy, provider: str +) -> None: + gateway: Final = counting_proxy.gateway + hello: Final[list[JsonValue]] = [{"role": "user", "content": "hello"}] + with wire_server(_google_cannot_count) as google, gateway.scenario() as scenario: + model: Final = _google_deployment(gateway, scenario, google, provider, {"token_count_framing": _FRAMING}) + rejected: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": model, "messages": hello, "context_management": _COMPACT_UNDER_THE_MINIMUM_TRIGGER}, + ) + assert rejected.status_code == 400, rejected.text + assert _google_count_targets(google) == () + estimated: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, "messages": hello}) + assert estimated.status_code == 200, estimated.text + assert _JSON_OBJECT.validate_json(estimated.content).get("litellm_count_source") == "local_estimate" + assert _google_count_targets(google) == (True,) + + +@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"]) +def test_token_count_model_counts_the_body_left_after_the_context_management_edits( + counting_proxy: OwnedProxy, provider: str +) -> None: + gateway: Final = counting_proxy.gateway + marker: Final = uuid4().hex + with wire_server(_google_counts_only_the_count_model) as google, gateway.scenario() as scenario: + model: Final = _google_deployment( + gateway, scenario, google, provider, {"token_count_model": {"model": _GEMINI_COUNT_MODEL}} + ) + counted: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": model, + "tools": [_READ_FILE_TOOL], + "messages": _agent_messages(marker), + "context_management": _CLEAR_OLDER_TOOL_RESULTS, + }, + ) + assert counted.status_code == 200, counted.text + assert _JSON_OBJECT.validate_json(counted.content) == { + "input_tokens": _COUNT_MODEL_COUNT, + "litellm_count_source": "token_count_model", + "litellm_count_model": _GEMINI_COUNT_MODEL, + } + assert tuple(_counted_files(request, marker) for request in google.drain()) == ((True, (False, False, True)),) diff --git a/tests/integration/providers/test_malformed_tools_wire.py b/tests/integration/providers/test_malformed_tools_wire.py index 25a86f7fc7c..706babac398 100644 --- a/tests/integration/providers/test_malformed_tools_wire.py +++ b/tests/integration/providers/test_malformed_tools_wire.py @@ -46,6 +46,7 @@ _CLAUDE_BACKEND: Final = "claude-haiku-4-5" _GEMINI_BACKEND: Final = "gemini-2.5-flash" _PROVIDER_KEY: Final = "scripted-malformed-tools-key" _COUNT: Final = 37 +_PROVIDER_SOURCE: Final[Mapping[str, JsonValue]] = MappingProxyType({"litellm_count_source": "provider"}) _CONFIG_STORE: Final = "vs_integration_config_store" _OTHER_STORE: Final = "vs_malformed_tools_other_store" _CLAUDE_CODE_AGENT: Final = "claude-cli/2.1.0 (external, cli)" @@ -821,7 +822,7 @@ def test_count_tokens_forwards_well_formed_tools(rig: _Rig) -> None: "POST", "/v1/messages/count_tokens", _count_body(rig, marker, tools=[MESSAGES.tool]) ) assert response.status_code == 200, response.text - assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text + assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT, **_PROVIDER_SOURCE}, response.text assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool] @@ -853,6 +854,7 @@ def test_responses_input_tokens_counts_well_formed_tools(rig: _Rig) -> None: assert _JSON_OBJECT.validate_json(response.content) == { "object": "response.input_tokens", "input_tokens": _COUNT, + **_PROVIDER_SOURCE, }, response.text sent: Final = _one_delivery(rig, marker, "/v1/messages/count_tokens") assert "get_weather" in json.dumps(sent["tools"]), sent @@ -1396,7 +1398,7 @@ def test_count_tokens_drops_the_bad_entry_a_deployment_list_names(rig: _Rig, fie {**_count_body(rig, marker, tools=[MESSAGES.tool, 5]), "model": field_groups.anthropic}, ) assert response.status_code == 200, response.text - assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT}, response.text + assert _JSON_OBJECT.validate_json(response.content) == {"input_tokens": _COUNT, **_PROVIDER_SOURCE}, response.text assert _one_delivery(rig, marker, "/v1/messages/count_tokens")["tools"] == [MESSAGES.tool] @@ -1656,7 +1658,7 @@ _CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType( named_as=MESSAGES, send=_count_tokens, target="/v1/messages/count_tokens", - answered=_counted({"input_tokens": _COUNT}), + answered=_counted({"input_tokens": _COUNT, **_PROVIDER_SOURCE}), kept_by_a_tools_list=("get_weather",), bad_entry=_MESSAGES_BAD_ENTRY, ), @@ -1666,7 +1668,7 @@ _CHECKED: Final[Mapping[str, _Checked]] = MappingProxyType( named_as=MESSAGES, send=_input_tokens, target="/v1/messages/count_tokens", - answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT}), + answered=_counted({"object": "response.input_tokens", "input_tokens": _COUNT, **_PROVIDER_SOURCE}), kept_by_a_tools_list=("get_weather",), bad_entry=_RESPONSES_BAD_ENTRY, ), diff --git a/tests/integration/providers/test_responses_bridge_message_cache_control_wire.py b/tests/integration/providers/test_responses_bridge_message_cache_control_wire.py index ea55552c331..4b25275285e 100644 --- a/tests/integration/providers/test_responses_bridge_message_cache_control_wire.py +++ b/tests/integration/providers/test_responses_bridge_message_cache_control_wire.py @@ -1118,6 +1118,7 @@ def test_input_tokens_counts_marked_items_without_their_marks(gateway: Gateway) assert _JSON_OBJECT.validate_json(response.text) == { "object": "response.input_tokens", "input_tokens": _COUNTED_TOKENS, + "litellm_count_source": "provider", } assert counted.target.endswith("/count_tokens"), counted.target assert anthropic_labels(counted) == [] @@ -1199,6 +1200,7 @@ def test_input_tokens_reads_each_input_shape_without_item_marks(gateway: Gateway assert _JSON_OBJECT.validate_json(response.text) == { "object": "response.input_tokens", "input_tokens": _COUNTED_TOKENS, + "litellm_count_source": "provider", } assert [anthropic_labels(request) for request in counted] == [[]] diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index 8d520a038af..1956301f1be 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -396,7 +396,7 @@ def test_messages_count_tokens_forwards_system_and_tools(gateway: Gateway, syste with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system, "tools": _TOOLS},) @@ -404,7 +404,7 @@ def test_messages_count_tokens_without_system_or_tools_sends_bare_body(gateway: with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == (_PEER_BARE,) @@ -443,7 +443,11 @@ def test_responses_input_tokens_counts_through_the_partner_peer(gateway: Gateway response: Final = gateway.request( "POST", "/v1/responses/input_tokens", {"model": model, "input": "Count this message"} ) - assert _payload(response) == {"object": "response.input_tokens", "input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == { + "object": "response.input_tokens", + "input_tokens": _PEER_COUNT, + "litellm_count_source": "provider", + }, response.text assert _counted_bodies(wire) == (_PEER_BARE,) @@ -498,7 +502,7 @@ def test_llama_count_still_reaches_the_meta_count_route(gateway: Gateway) -> Non with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url, model=f"vertex_ai/{_LLAMA_BACKEND}") response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text (sent,) = _count_bodies(wire.drain(), _LLAMA_COUNT_TARGET) assert sent["messages"] == _MESSAGES, sent @@ -616,7 +620,7 @@ def test_messages_count_tokens_forwards_an_empty_tools_list(gateway: Gateway) -> with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "tools": []}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "tools": []},) @@ -624,7 +628,7 @@ def test_messages_count_tokens_forwards_an_empty_system_string(gateway: Gateway) with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": ""}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "system": ""},) @@ -632,7 +636,7 @@ def test_messages_count_tokens_leaves_null_system_and_tools_out(gateway: Gateway with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": None, "tools": None}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == (_PEER_BARE,) @@ -650,7 +654,7 @@ def test_messages_count_tokens_forwards_a_5kb_system_verbatim(gateway: Gateway) with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) response: Final = _count(gateway, {**_bare(model), "system": system}) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text assert _counted_bodies(wire) == ({**_PEER_BARE, "system": system},) @@ -663,7 +667,7 @@ def test_messages_count_tokens_duplicate_system_and_tools_keys_forward_one_value content=f'{{"model": "{model}", "messages": {json.dumps(_MESSAGES)}, {fields}, {fields}}}', headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, ) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text (sent,) = _count_requests(wire.drain()) assert _count_bodies((sent,)) == (_PEER_FULL,) assert (sent.body.count(b'"system"'), sent.body.count(b'"tools"')) == (1, 1), sent.body @@ -696,7 +700,7 @@ def test_count_tokens_location_override_targets_the_count_region(gateway: Gatewa gateway, scenario, wire.url, vertex_location="global", vertex_count_tokens_location="europe-west1" ) response: Final = _count(gateway, _bare(model)) - assert _payload(response) == {"input_tokens": _PEER_COUNT}, response.text + assert _payload(response) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, response.text target: Final = _COUNT_TARGET.replace(f"/locations/{_LOCATION}/", "/locations/europe-west1/") assert _count_bodies(wire.drain(), target) == (_PEER_BARE,) @@ -705,7 +709,7 @@ def test_messages_count_tokens_repeated_request_reaches_the_peer_each_time(gatew with wire_server(_peer()) as wire, gateway.scenario() as scenario: model: Final = _deployment(gateway, scenario, wire.url) answers: Final = tuple(_payload(_count(gateway, _bare(model))) for _ in range(2)) - assert answers == ({"input_tokens": _PEER_COUNT},) * 2 + assert answers == ({"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"},) * 2 assert _counted_bodies(wire) == (_PEER_BARE,) * 2 @@ -722,7 +726,7 @@ def test_disabled_token_counter_surfaces_provider_failures_instead_of_counting_l ) with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: counted: Final = _count(owned.gateway, _full(_OWNED_MODEL)) - assert _payload(counted) == {"input_tokens": _PEER_COUNT}, counted.text + assert _payload(counted) == {"input_tokens": _PEER_COUNT, "litellm_count_source": "provider"}, counted.text rejected: Final = _count( owned.gateway, {**_full(_OWNED_MODEL), "messages": [{"role": "user", "content": _REJECT_TEXT}]} ) diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 801a14c87d3..511a812e895 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -453,9 +453,11 @@ def test_in_flight_count_tokens_does_not_reserve_key_budget_away_from_a_completi [{"role": "user", "content": prompt}] ] assert count.status_code == 200, count.text - assert count.json() == {"totalTokens": 12, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}]}, ( - count.text - ) + assert count.json() == { + "totalTokens": 12, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}], + "litellm_count_source": "provider", + }, count.text provider_calls: Final = wire.drain() assert [(call.method, call.target) for call in provider_calls] == [ ("POST", "/v1beta/models/gemini-3.8-flash:countTokens") diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 86ad6c8d9d9..4913d2b0d1d 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -54,6 +54,7 @@ async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": payload, "error": False, "error_message": None, @@ -98,6 +99,7 @@ async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_v "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": None, "error": True, "error_message": str(rejection.value), @@ -126,6 +128,7 @@ async def test_count_payload_that_is_not_an_object_becomes_an_error_response_nam "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": None, "error": True, "error_message": str(failed_lookup.value), @@ -156,6 +159,7 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": None, "error": True, "error_message": upstream_text, @@ -180,6 +184,7 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "request_model": "azure_ai/claude-sonnet-4-5", "model_used": "claude-sonnet-4-5", "tokenizer_type": "azure_ai_anthropic_api", + "count_source": None, "original_response": None, "error": True, "error_message": f"CountTokens processing error: {undecodable.value}", diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py new file mode 100644 index 00000000000..ffddb09cca7 --- /dev/null +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py @@ -0,0 +1,524 @@ +import asyncio +import base64 +import hashlib +import json +import logging +import re +import threading +from collections.abc import Iterable +from itertools import product +from pathlib import Path +from typing import Final + +import httpx +import pytest +import respx +import tiktoken + +import litellm +from litellm._logging import verbose_logger +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_anthropic_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) +from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig + +MODEL: Final = "us.moonshotai.kimi-k3" +_CONTROL: Final = re.compile(r"<\|(?:open|close|sep|end_of_msg)\|>") +_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather in °C", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}, "days": {"type": "integer"}}, + "required": ["city"], + }, + }, +} + + +def _call(call_id: str, arguments: dict[str, object]) -> dict[str, object]: + return {"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": json.dumps(arguments)}} + + +def _written_ranks(path: Path, tokens: Iterable[bytes]) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(token) + b" %d" % rank for rank, token in enumerate(tokens))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _single_bytes() -> tuple[bytes, ...]: + return tuple(bytes((value,)) for value in range(256)) + + +@pytest.fixture(scope="module") +def byte_source(tmp_path_factory: pytest.TempPathFactory) -> KimiK3TokenizerSource: + return _written_ranks(tmp_path_factory.mktemp("bytes") / "tiktoken.model", _single_bytes()) + + +@pytest.fixture(scope="module") +def byte_encoding(byte_source: KimiK3TokenizerSource) -> tiktoken.Encoding: + loaded: Final = load_kimi_k3_encoding(byte_source) + assert loaded is not None + return loaded + + +@pytest.fixture(scope="module") +def pair_encoding(tmp_path_factory: pytest.TempPathFactory) -> tiktoken.Encoding: + pairs: Final = (bytes(pair) for pair in product(range(256), repeat=2)) + source: Final = _written_ranks(tmp_path_factory.mktemp("pairs") / "tiktoken.model", (*_single_bytes(), *pairs)) + loaded: Final = load_kimi_k3_encoding(source) + assert loaded is not None + return loaded + + +def _template_count(encoding: tiktoken.Encoding, rendered: str) -> int: + texts: Final = _CONTROL.split(rendered) + return len(_CONTROL.findall(rendered)) + sum(len(encoding.encode_ordinary(text)) for text in texts) + + +_GENERATION_PROMPT: Final = '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>' + + +# Expected renders follow the Kimi K3 chat template in moonshotai/Kimi-K3 at revision f831ab6, which matched +# Bedrock's prompt_tokens for us.moonshotai.kimi-k3 in 19 live calls on 2026-10-16 +@pytest.mark.parametrize( + ("messages", "tools", "rendered"), + [ + pytest.param( + [{"role": "user", "content": "Hello there"}], + None, + '<|open|>message role="user"<|sep|>Hello there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="user", + ), + pytest.param( + [ + { + "role": "system", + "content": [{"type": "text", "text": "Be terse."}, {"type": "text", "text": "Answer in French."}], + }, + {"role": "user", "content": [{"type": "text", "text": "Hi"}, {"type": "text", "text": " there"}]}, + ], + None, + '<|open|>message role="system"<|sep|>Be terse.Answer in French.<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Hi there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="text_parts_join", + ), + pytest.param( + [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "And 3+3?"}, + ], + None, + '<|open|>message role="user"<|sep|>What is 2+2?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|>4<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="user"<|sep|>And 3+3?<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT, + id="assistant_turn", + ), + pytest.param( + [ + {"role": "user", "content": "Weather in Paris and Oslo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [_call("c1", {"city": "Paris", "days": 3}), _call("c2", {"city": "Oslo"})], + }, + {"role": "tool", "tool_call_id": "c2", "content": "Snow"}, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ], + [_WEATHER], + '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n' + "Here are the available tools, described in JSONSchema.\n\n```json\n" + '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":' + '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},' + '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Weather in Paris and Oslo?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>" + '<|open|>call tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|>' + '<|open|>argument key="days" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|>' + '<|open|>call tool="get_weather" index="2"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Oslo<|close|>argument<|sep|><|close|>call<|sep|>' + "<|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="tool" tool="get_weather" index="2"<|sep|>Snow<|close|>message<|sep|><|end_of_msg|>' + + _GENERATION_PROMPT, + id="tool_results_out_of_call_order", + ), + pytest.param( + [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + _call( + "c1", {"a&b": 'x"y', "flag": True, "none": None, "list": [1, "zé"], "object": {"k": True}} + ) + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Rain"}, + ], + [_WEATHER], + '<|open|>message role="system" type="tool-declare"<|sep|># Tools\n' + "Here are the available tools, described in JSONSchema.\n\n```json\n" + '[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":' + '{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},' + '"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="user"<|sep|>Weather?<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>' + "<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>" + '<|open|>call tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="a&b" type="string"<|sep|>x"y<|close|>argument<|sep|>' + '<|open|>argument key="flag" type="boolean"<|sep|>true<|close|>argument<|sep|>' + '<|open|>argument key="none" type="null"<|sep|>null<|close|>argument<|sep|>' + '<|open|>argument key="list" type="array"<|sep|>[1, "zé"]<|close|>argument<|sep|>' + '<|open|>argument key="object" type="object"<|sep|>{"k": true}<|close|>argument<|sep|>' + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>" + '<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Rain<|close|>message<|sep|><|end_of_msg|>' + + _GENERATION_PROMPT, + id="argument_types_and_escaping", + ), + ], +) +@pytest.mark.parametrize("encoding_fixture", ["byte_encoding", "pair_encoding"]) +def test_count_is_the_kimi_k3_chat_template_render_of_the_bedrock_request( + request: pytest.FixtureRequest, + encoding_fixture: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + rendered: str, +): + encoding: Final[tiktoken.Encoding] = request.getfixturevalue(encoding_fixture) + + assert kimi_k3_tokens(encoding, MODEL, messages, tools) == _template_count(encoding, rendered) + + +def _parallel_calls(count: int) -> list[dict[str, object]]: + ids: Final = tuple(f"c{number}" for number in range(1, count + 1)) + return [ + {"role": "user", "content": "Weather?"}, + {"role": "assistant", "content": None, "tool_calls": [_call(call_id, {"city": "Paris"}) for call_id in ids]}, + *({"role": "tool", "tool_call_id": call_id, "content": "Sun"} for call_id in ids), + ] + + +def test_tenth_parallel_call_and_its_result_are_numbered_ten(byte_encoding: tiktoken.Encoding): + nine: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(9), [_WEATHER]) + ten: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(10), [_WEATHER]) + assert nine is not None and ten is not None + + assert ten - nine == _template_count( + byte_encoding, + '<|open|>call tool="get_weather" index="10"<|sep|>' + '<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|>' + '<|open|>message role="tool" tool="get_weather" index="10"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>', + ) + + +@pytest.mark.parametrize( + "text", + [ + pytest.param("The quick brown fox, jumps over 12345 lazy dogs. " * 60, id="prose"), + pytest.param("a" + "1234567890" * 250 + "1", id="long_digit_run"), + pytest.param("--- " * 600, id="symbols"), + pytest.param("def f(x):\n return x # it's 42\n\n" * 80, id="code"), + pytest.param("word next\t\ttab\n \n" * 120, id="whitespace"), + pytest.param("naïve café, “quotes” & 中文汉字，日本語のテキスト 😀🎉\r\n" * 40, id="unicode"), + pytest.param("We're sure they'd say it's DON'T or I'M " * 60, id="contractions"), + pytest.param("abc123def4567ghi89 " * 120, id="letters_and_digits"), + ], +) +def test_long_text_counts_what_the_tokenizer_counts_for_the_whole_text(pair_encoding: tiktoken.Encoding, text: str): + def counted(content: str) -> int | None: + return kimi_k3_tokens(pair_encoding, MODEL, [{"role": "user", "content": content}], None) + + short: Final = counted("x") + assert short is not None + + assert counted(text) == short - len(pair_encoding.encode_ordinary("x")) + len(pair_encoding.encode_ordinary(text)) + + +@pytest.mark.parametrize( + ("messages", "tools"), + [ + pytest.param( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + ], + } + ], + None, + id="image", + ), + pytest.param( + [{"role": "user", "content": "Weather?"}, {"role": "tool", "tool_call_id": "c1", "content": "Sun"}], + [_WEATHER], + id="tool_result_without_its_call", + ), + ], +) +def test_request_the_template_cannot_render_is_not_counted( + pair_encoding: tiktoken.Encoding, messages: list[dict[str, object]], tools: list[dict[str, object]] | None +): + assert kimi_k3_tokens(pair_encoding, MODEL, messages, tools) is None + + +@pytest.mark.parametrize( + "tool", + [ + pytest.param( + { + "type": "function", + "name": "get_weather", + "description": "Get the weather in °C", + "parameters": _WEATHER["function"]["parameters"], + }, + id="responses_api", + ), + pytest.param( + { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], + }, + id="anthropic_messages", + ), + ], +) +def test_tool_in_another_api_shape_counts_like_the_same_chat_completions_tool( + pair_encoding: tiktoken.Encoding, tool: dict[str, object] +): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [tool]) + + assert counted is not None + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + +@pytest.mark.parametrize( + "tool", + [ + pytest.param({"type": "web_search_20250305", "name": "web_search"}, id="anthropic_web_search"), + pytest.param({"type": "web_search"}, id="responses_web_search"), + pytest.param( + { + "type": "mcp", + "server_label": "docs", + "server_url": "https://mcp.example/mcp", + "require_approval": "never", + }, + id="mcp", + ), + ], +) +def test_tool_bedrock_inference_drops_counts_like_the_function_tools_alone( + pair_encoding: tiktoken.Encoding, tool: dict[str, object] +): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, tool]) + + assert counted is not None + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + +def test_responses_custom_tool_counts_like_the_function_tool_inference_sends(pair_encoding: tiktoken.Encoding): + messages: Final = [{"role": "user", "content": "Weather in Paris?"}] + custom_tool: Final = {"type": "custom", "name": "grammar", "description": "Free-form SQL"} + sent_to_inference, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [dict(custom_tool)] + ) + + without_it: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER]) + + counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, custom_tool]) + + assert counted is not None and without_it is not None + assert counted > without_it + assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, *sent_to_inference]) + + +def test_counting_a_request_with_an_image_url_never_downloads_the_image( + pair_encoding: tiktoken.Encoding, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "user_url_validation", False) + image: Final = respx_mock.get("https://images.example/cat.png").mock(return_value=httpx.Response(200)) + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + {"type": "image_url", "image_url": {"url": "https://images.example/cat.png"}}, + ], + } + ] + + assert kimi_k3_tokens(pair_encoding, MODEL, messages, None) is None + assert not image.called + + +def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: tiktoken.Encoding): + anthropic: Final = kimi_k3_anthropic_tokens( + pair_encoding, + MODEL, + { + "system": "Be terse.", + "messages": [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]}, + ], + "tools": [ + { + "name": "get_weather", + "description": "Get the weather in °C", + "input_schema": _WEATHER["function"]["parameters"], + } + ], + }, + ) + openai: Final = kimi_k3_tokens( + pair_encoding, + MODEL, + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [_call("c1", {"city": "Paris"})]}, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ], + [_WEATHER], + ) + + assert openai is not None + assert anthropic == openai + + +@pytest.mark.parametrize("failure", ["checksum", "missing_file"]) +def test_tokenizer_that_fails_its_checksum_or_cannot_be_read_is_not_loaded( + byte_source: KimiK3TokenizerSource, tmp_path: Path, caplog: pytest.LogCaptureFixture, failure: str +): + source: Final = ( + KimiK3TokenizerSource(ranks_path=byte_source.ranks_path, sha256=hashlib.sha256(b"other").hexdigest()) + if failure == "checksum" + else KimiK3TokenizerSource(ranks_path=lambda: str(tmp_path / "missing"), sha256=byte_source.sha256) + ) + verbose_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + loaded: Final = load_kimi_k3_encoding(source) + finally: + verbose_logger.removeHandler(caplog.handler) + + assert loaded is None + assert any("Kimi K3 tokenizer unavailable" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_slow_tokenizer_load_is_not_waited_on_past_the_wait_and_is_shared_once_done( + byte_source: KimiK3TokenizerSource, +): + released: Final = threading.Event() + + def slow_ranks_path() -> str: + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=slow_ranks_path, sha256=byte_source.sha256), + retry_after_seconds=60, + clock=lambda: 0.0, + ) + + assert await encodings.encoding(wait_seconds=0) is None + released.set() + loaded: Final = await encodings.encoding(wait_seconds=5) + assert loaded is not None + assert await encodings.encoding(wait_seconds=0) is loaded + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + +async def _loaded_eventually(encodings: KimiK3Encodings) -> tiktoken.Encoding | None: + for _ in range(500): + if (loaded := await encodings.encoding(wait_seconds=0)) is not None: + return loaded + await asyncio.sleep(0.01) + return None + + +@pytest.mark.asyncio +async def test_failed_tokenizer_load_is_retried_in_the_background_only_after_the_retry_window( + byte_source: KimiK3TokenizerSource, tmp_path: Path +): + clock: Final = _Clock() + reads: Final[list[float]] = [] + released: Final = threading.Event() + + def ranks_path() -> str: + reads.append(clock.now) + if len(reads) == 1: + return str(tmp_path / "missing") + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=ranks_path, sha256=byte_source.sha256), + retry_after_seconds=60, + clock=clock, + ) + + assert await encodings.encoding(wait_seconds=5) is None + clock.now = 59.0 + assert await encodings.encoding(wait_seconds=5) is None + clock.now = 60.0 + assert await asyncio.wait_for(encodings.encoding(wait_seconds=5), timeout=1) is None + released.set() + assert await _loaded_eventually(encodings) is not None + assert reads == [0.0, 60.0] + + +@pytest.mark.asyncio +async def test_tokenizer_load_still_running_after_the_wait_is_logged( + byte_source: KimiK3TokenizerSource, caplog: pytest.LogCaptureFixture +): + released: Final = threading.Event() + + def hung_ranks_path() -> str: + released.wait(10) + return byte_source.ranks_path() + + encodings: Final = KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=hung_ranks_path, sha256=byte_source.sha256), retry_after_seconds=60 + ) + verbose_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + loaded: Final = await encodings.encoding(wait_seconds=0.05) + finally: + verbose_logger.removeHandler(caplog.handler) + released.set() + + assert loaded is None + assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records) diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index befdc5ee9fd..be394b71ec7 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -54,6 +54,7 @@ async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload( "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": payload, "error": False, "error_message": None, @@ -98,6 +99,7 @@ async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_v "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": None, "error": True, "error_message": str(rejection.value), @@ -126,6 +128,7 @@ async def test_count_payload_that_is_not_an_object_becomes_an_error_response_nam "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": None, "error": True, "error_message": str(failed_lookup.value), @@ -156,6 +159,7 @@ async def test_upstream_error_status_becomes_an_error_response_with_that_status_ "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": None, "error": True, "error_message": upstream_text, @@ -382,6 +386,7 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "request_model": "openai/gpt-4o", "model_used": "gpt-4o", "tokenizer_type": "openai_api", + "count_source": None, "original_response": None, "error": True, "error_message": f"CountTokens processing error: {undecodable.value}", diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py new file mode 100644 index 00000000000..8bd5ced1b58 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -0,0 +1,209 @@ +import json +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + count_source_fields, + resolve_token_count_settings, +) +from litellm.types.utils import TokenCountResponse + +_OPUS_US_EAST_1: Final = TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1") +_DEFAULT: Final = { + "token_count_framing": 97, + "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}, +} +_COST_MAP: Final = Path(__file__).parents[4] / "model_prices_and_context_window.json" + + +def _in_us_west_2() -> str: + return "us-west-2" + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}}, + TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")), + ), + ( + {"token_count_model": {"model": "custom-model"}}, + TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + ), + ( + {"token_count_model": {"model": "custom-model", "region": " "}}, + TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + ), + ({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_framing": None, "token_count_model": None}, + TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1), + ), + ( + {"token_count_framing": True, "token_count_model": "custom-model"}, + TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1), + ), + ({"token_count_framing": 78.0}, TokenCountSettings(framing=78, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": 1.5}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": -1}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_framing": "78"}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ({"token_count_model": {}}, TokenCountSettings(framing=97, count_model=None)), + ({"token_count_model": ""}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), + ( + {"token_count_model": {"model": " ", "region": "us-east-1"}}, + TokenCountSettings(framing=97, count_model=None), + ), + ], +) +def test_config_overrides_the_cost_map_default_key_by_key( + configured: dict[str, object], expected: TokenCountSettings +) -> None: + assert ( + resolve_token_count_settings(configured=configured, default=_DEFAULT, deployment_region=_in_us_west_2) + == expected + ) + + +@pytest.mark.parametrize( + ("deployment_region", "configured", "expected_region"), + [ + ("us-west-2", {}, "us-east-1"), + ("eu-central-1", {}, None), + ("ca-central-1", {}, None), + ("us-gov-west-1", {}, None), + ( + "eu-central-1", + {"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}}, + "us-east-1", + ), + ], +) +def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography( + deployment_region: str, configured: dict[str, object], expected_region: str | None +) -> None: + settings: Final = resolve_token_count_settings( + configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region + ) + assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) + + +def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None: + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}}, + deployment_region=lambda: "eu-west-1", + ) + assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1") + + +def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None: + settings: Final = resolve_token_count_settings( + configured={}, + default={"token_count_model": {"model": "au.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + deployment_region=lambda: "ap-southeast-2", + ) + assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4") + + +def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None: + assert resolve_token_count_settings( + configured={}, default={}, deployment_region=_in_us_west_2 + ) == TokenCountSettings(framing=None, count_model=None) + + +@pytest.mark.parametrize( + ("value", "expected"), + [("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)], +) +def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object, expected: str | None) -> None: + settings: Final = resolve_token_count_settings( + configured={}, default={"token_count_tokenizer": value}, deployment_region=_in_us_west_2 + ) + assert settings.tokenizer == expected + + +def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse: + return TokenCountResponse( + total_tokens=7, + request_model="sonnet", + model_used=model_used, + tokenizer_type="bedrock_api", + count_source=source, + ) + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + (None, {}), + ("provider", {"litellm_count_source": "provider"}), + ("local_estimate", {"litellm_count_source": "local_estimate"}), + ( + "token_count_model", + {"litellm_count_source": "token_count_model", "litellm_count_model": "us.anthropic.claude-opus-5-5"}, + ), + ], +) +def test_count_source_fields_name_the_counting_model_only_when_token_count_model_counted( + source: str | None, expected: dict[str, str] +) -> None: + assert count_source_fields(_counted(source)) == expected + + +def test_every_shipped_token_count_model_is_another_mapped_model_of_the_same_provider_in_a_named_region() -> None: + cost_map: Final = json.loads(_COST_MAP.read_text()) + count_models: Final = { + name: entry["token_count_model"] + for name, entry in cost_map.items() + if isinstance(entry, dict) and "token_count_model" in entry + } + assert count_models + assert { + name: count_model + for name, count_model in count_models.items() + if cost_map.get(count_model["model"], {}).get("litellm_provider") != cost_map[name]["litellm_provider"] + or count_model["model"] == name + or not count_model.get("region") + } == {} + + +_UPSTREAM_404: Final = { + "type": "error", + "request_id": "req_upstream", + "error": {"type": "not_found_error", "message": "The model 'anthropic.claude-opus-5-5' does not exist"}, +} + + +@pytest.mark.parametrize( + ("message", "expected"), + [ + ( + json.dumps(_UPSTREAM_404), + { + **_UPSTREAM_404, + "error": { + "type": "not_found_error", + "message": "The model 'anthropic.claude-opus-5-5' does not exist. Enable eu-west-1", + }, + }, + ), + (json.dumps({"message": "Too many tokens"}), "Too many tokens. Enable eu-west-1"), + ("scripted 404", "scripted 404. Enable eu-west-1"), + ( + "Credential should be scoped to a valid region. ", + "Credential should be scoped to a valid region. Enable eu-west-1", + ), + ], +) +def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message: str, expected: object) -> None: + hinted: Final = ProviderCountFailure(status_code=404, message=message).with_hint("Enable eu-west-1") + + assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected + assert hinted.status_code == 404 diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index 114bbb184c6..2777d18a399 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -4,10 +4,11 @@ Test to verify the Google GenAI proxy API endpoints """ from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest +from litellm.types.utils import TokenCountResponse def _build_test_client(): @@ -88,9 +89,7 @@ def test_google_stream_generate_content_endpoint(): assert init_kwargs["data"]["_litellm_raw_sse_stream"] is True assert init_kwargs["data"]["_litellm_skip_openai_stream_done"] is True assert init_kwargs["data"]["model"] == "test-model" - assert init_kwargs["data"]["contents"] == [ - {"role": "user", "parts": [{"text": "Hello"}]} - ] + assert init_kwargs["data"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] def test_google_generate_content_data_flows_through_processor(): @@ -123,9 +122,7 @@ def test_google_generate_content_data_flows_through_processor(): data = mock_init.call_args.kwargs["data"] assert data["model"] == "test-model" assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] - assert data["systemInstruction"] == { - "parts": [{"text": "Your name is Doodle."}] - } + assert data["systemInstruction"] == {"parts": [{"text": "Your name is Doodle."}]} # generationConfig arrives intact here; the rename to `config` is # done downstream in route_request (see test_route_llm_request). assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] @@ -160,12 +157,17 @@ def test_google_count_tokens_unchanged(): except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - fake_response = MagicMock() - fake_response.original_response = { - "totalTokens": 7, - "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], - } - fake_response.total_tokens = 7 + fake_response = TokenCountResponse( + total_tokens=7, + request_model="test-model", + model_used="test-model", + tokenizer_type="gemini_api", + original_response={ + "totalTokens": 7, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], + }, + count_source="provider", + ) with patch( "litellm.proxy.proxy_server.token_counter", @@ -180,6 +182,7 @@ def test_google_count_tokens_unchanged(): assert response.status_code == 200 body = response.json() assert body["totalTokens"] == 7 + assert body["litellm_count_source"] == "provider" def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope(): @@ -212,9 +215,14 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - fake_response: Final = MagicMock() - fake_response.original_response = {"input_tokens": 2167} - fake_response.total_tokens = 2167 + fake_response: Final = TokenCountResponse( + total_tokens=2167, + request_model="claude-opus-4-8", + model_used="us.anthropic.claude-opus-5-5", + tokenizer_type="bedrock_api", + original_response={"input_tokens": 2167}, + count_source="token_count_model", + ) with patch( "litellm.proxy.proxy_server.token_counter", @@ -230,3 +238,5 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs() body: Final = response.json() assert body["totalTokens"] == 2167 assert body["promptTokensDetails"] == [] + assert body["litellm_count_source"] == "token_count_model" + assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5" diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index ad75344afba..7cd1107abce 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4184,6 +4184,39 @@ class TestModelInfoCostMapEchoFilter: assert "litellm_provider" not in info assert "supported_openai_params" not in info + @pytest.mark.parametrize( + ("model", "field", "edited"), + [ + ("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}), + ("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0), + ], + ) + def test_token_count_setting_is_stored_only_when_the_operator_changed_it( + self, model: str, field: str, edited: object + ) -> None: + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + entry: Final = litellm.get_model_info(model) + assert entry[field] not in (None, edited), entry + db_model: Final = Deployment( + model_name="count-settings", + litellm_params=LiteLLM_Params(model=model), + model_info=ModelInfo(id="dep-echo-count"), + ) + + echoed: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count"})), + ) + changed: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count", field: edited})), + ) + + assert field not in json.loads(echoed["model_info"]) + assert json.loads(changed["model_info"])[field] == edited + def test_metadata_without_a_cost_map_key_is_persisted(self): import litellm diff --git a/tests/unit/proxy/proxy_server/test_routes_utils.py b/tests/unit/proxy/proxy_server/test_routes_utils.py index 118b80da31e..f93c23c8ce4 100644 --- a/tests/unit/proxy/proxy_server/test_routes_utils.py +++ b/tests/unit/proxy/proxy_server/test_routes_utils.py @@ -54,6 +54,7 @@ def test_token_counter_happy_path(client, auth_as, patched_token_counter): "request_model": "gpt-4", "model_used": "gpt-4", "tokenizer_type": "openai_tokenizer", + "count_source": "local_estimate", "original_response": None, "error": False, "error_message": None, diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index f7d1970fdae..720739a730b 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,15 +2,32 @@ # 1. Generate a Key, and use it to make a call +import asyncio +import base64 +import hashlib +import itertools import json -from collections.abc import Callable, Mapping -from typing import Final +import logging +import threading +import time +import traceback +from collections.abc import Callable, Coroutine, Mapping +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final, NoReturn from unittest.mock import AsyncMock, MagicMock, patch +import anyio import httpx import pytest +import tiktoken from dotenv import load_dotenv from pydantic import BaseModel, ConfigDict, JsonValue +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from tokenizers.processors import TemplateProcessing load_dotenv() @@ -22,10 +39,20 @@ from fastapi.responses import JSONResponse import litellm from litellm import Router +from litellm._logging import verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler +from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.bedrock.count_tokens.kimi_k3 import ( + KIMI_K3_RANKS_SHA256, + KimiK3Encodings, + KimiK3TokenizerSource, + kimi_k3_template_tokens, + kimi_k3_tokens, + load_kimi_k3_encoding, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter @@ -33,8 +60,39 @@ from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, ) -from litellm.proxy.proxy_server import token_counter -from litellm.types.utils import TokenCountResponse +from litellm.proxy.common_utils.token_count_settings import ( + ProviderCountFailure, + TokenCountModel, + TokenCountSettings, + cost_map_entry, +) +from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request +from litellm.proxy.proxy_server import ( + _estimate_or_provider_error, + _provider_token_count, + _try_provider_token_count, + count_request_tokens, + token_counter, +) +from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from litellm.types.utils import SelectTokenizerResponse, TokenCountMessagesFormat, TokenCountResponse + + +def _hugging_face_offline() -> NoReturn: + raise OSError("unit tests never download the Kimi K3 tokenizer") + + +@pytest.fixture(autouse=True) +def _kimi_k3_tokenizer_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "_KIMI_K3_ENCODINGS", + KimiK3Encodings( + source=KimiK3TokenizerSource(ranks_path=_hugging_face_offline, sha256=KIMI_K3_RANKS_SHA256), + retry_after_seconds=60, + ), + ) @pytest.mark.asyncio @@ -1217,6 +1275,1282 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = original_token_counter +class _ScriptedCounter(BaseTokenCounter): + def __init__(self, replies: dict[str, TokenCountResponse], burst: int = 1) -> None: + self._replies = replies + self._burst = burst + self._burst_arrived = asyncio.Event() + self.calls: list[tuple[str, object, object, object]] = [] + self.regions: list[object] = [] + + def should_use_token_counting_api(self, custom_llm_provider: str | None = None) -> bool: + return True + + async def count_tokens( + self, + model_to_use: str, + messages: object, + contents: object, + deployment: object = None, + request_model: str = "", + tools: object = None, + system: object = None, + messages_format: TokenCountMessagesFormat = "openai", + output_format: JsonValue = None, + context_management: object = None, + ) -> TokenCountResponse | None: + self.calls.append((model_to_use, messages, tools, system)) + self.regions.append(None if deployment is None else deployment["litellm_params"].get("aws_region_name")) + if len(self.calls) >= self._burst: + self._burst_arrived.set() + await self._burst_arrived.wait() + return self._replies[model_to_use] + + +def _provider_reply( + model: str, total_tokens: int = 0, status_code: int | None = None, rejected_before_provider_call: bool = False +) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=total_tokens, + request_model="sonnet-5-5", + model_used=model, + tokenizer_type="bedrock_api", + error=status_code is not None, + error_message=None if status_code is None else f"scripted {status_code} for {model}", + status_code=status_code, + rejected_before_provider_call=rejected_before_provider_call, + ) + + +_COUNT_REQUEST = TokenCountRequest( + model="sonnet-5-5", + messages=[{"role": "user", "content": "hello"}], + tools=[{"name": "get_quote", "input_schema": {"type": "object"}}], + system="be brief", +) + + +_SONNET_DEPLOYMENT: Final = MappingProxyType( + {"litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5", "aws_region_name": "us-west-2"}} +) +_OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5", region="us-east-1") + + +async def _ladder( + counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True +) -> TokenCountResponse | None: + return await _provider_token_count( + count_with=partial( + _try_provider_token_count, + counter, + "bedrock", + messages=_COUNT_REQUEST.messages, + contents=_COUNT_REQUEST.contents, + deployment=_SONNET_DEPLOYMENT, + request_model=_COUNT_REQUEST.model, + tools=_COUNT_REQUEST.tools, + system=_COUNT_REQUEST.system, + ), + model_to_use="us.anthropic.claude-sonnet-5-5", + settings=settings, + local_estimate_allowed=local_estimate_allowed, + ) + + +@pytest.mark.asyncio +async def test_model_without_token_count_model_counts_itself_in_the_deployment_region(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + + assert result is not None + assert (result.total_tokens, result.count_source) == (447, "provider") + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-sonnet-5-5"], ["us-west-2"]) + + +@pytest.mark.parametrize( + ("count_model", "region"), + [ + (_OPUS_IN_US_EAST_1, "us-east-1"), + (TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None), "us-west-2"), + ], +) +@pytest.mark.asyncio +async def test_token_count_model_counts_the_same_request_in_its_own_region_without_trying_the_model( + count_model: TokenCountModel, region: str +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=None, count_model=count_model)) + + assert result is not None + assert (result.total_tokens, result.count_source, result.model_used) == ( + 447, + "token_count_model", + "us.anthropic.claude-opus-5-5", + ) + assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-opus-5-5"], [region]) + assert counter.calls[0][1:] == (_COUNT_REQUEST.messages, _COUNT_REQUEST.tools, _COUNT_REQUEST.system) + assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" + + +@pytest.mark.parametrize("status_code", [401, 403, 429, 500]) +@pytest.mark.asyncio +async def test_errors_other_than_cannot_count_raise_even_with_framing(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + + assert (raised.value.code, raised.value.message) == (str(status_code), f"scripted {status_code} for sonnet") + + +@pytest.mark.parametrize("status_code", [400, 403, 404, 500]) +@pytest.mark.asyncio +async def test_token_count_model_failure_names_its_region_so_the_user_can_enable_it(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=None, count_model=_OPUS_IN_US_EAST_1)) + + assert raised.value.code == str(status_code) + assert raised.value.message.startswith(f"scripted {status_code} for opus") + assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + assert len(counter.calls) == 1 + + +@pytest.mark.asyncio +async def test_token_count_model_without_a_region_answers_its_error_unchanged(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) + + with pytest.raises(ProxyException) as raised: + await _ladder( + counter, + TokenCountSettings( + framing=None, count_model=TokenCountModel(model="us.anthropic.claude-opus-5-5", region=None) + ), + ) + + assert (raised.value.code, raised.value.message) == ("404", "scripted 404 for opus") + + +@pytest.mark.asyncio +async def test_uncountable_model_without_framing_raises(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=None, count_model=None)) + + assert raised.value.code == "404" + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None)) + + assert isinstance(result, ProviderCountFailure) + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize("status_code", [400, 404]) +@pytest.mark.asyncio +async def test_token_count_model_error_stands_even_with_framing(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + + assert raised.value.code == str(status_code) + assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + + +@pytest.mark.asyncio +async def test_uncountable_model_with_framing_raises_when_local_counting_is_disabled(): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=10, count_model=None), local_estimate_allowed=False) + + assert raised.value.code == "400" + + +@pytest.mark.asyncio +async def test_request_the_counter_rejects_before_calling_the_provider_answers_400_without_framing(): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply( + "sonnet", status_code=400, rejected_before_provider_call=True + ) + } + ) + + with pytest.raises(ProxyException) as raised: + await _ladder(counter, TokenCountSettings(framing=100, count_model=None)) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for sonnet") + + +@pytest.mark.parametrize("model", ["openai/gpt-5.4-mini", "gemini/gemini-3.8-flash"]) +@pytest.mark.asyncio +async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_convert( + monkeypatch: pytest.MonkeyPatch, model: str +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": model, "api_key": "sk-unused-framed"}, + "model_info": {"token_count_framing": 100}, + } + ] + ), + ) + search_result: Final = {"type": "search_result", "source": "https://example.com", "title": "Notes", "content": []} + + with pytest.raises(ProxyException) as raised: + await count_request_tokens( + request=TokenCountRequest( + model="framed", + messages=[{"role": "user", "content": [{"type": "text", "text": "Summarize."}, search_result]}], + ), + call_endpoint=True, + messages_format="anthropic", + ) + + assert raised.value.code == "400" + assert "search_result" in raised.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unreadable_body", + [ + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_01", + "name": "web_search", + "input": {"query": "new"}, + }, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_01", "content": []}, + ], + }, + ] + }, + id="web_search_blocks", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [{"name": "search", "input_schema": {"type": "object", "properties": {"query": True}}}], + }, + id="boolean_subschema_tool", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "What is new?"}], + "tools": [ + {"name": "search", "input_schema": {"type": "object", "properties": {"tags": {"type": "array"}}}} + ], + }, + id="array_schema_without_items", + ), + pytest.param( + {"messages": [], "contents": [{"role": "user", "parts": [{"text": "What is new?"}]}]}, + id="google_contents", + ), + pytest.param( + { + "messages": [{"role": "user", "content": "Describe this"}], + "contents": [ + { + "role": "user", + "parts": [ + {"text": "Describe this"}, + {"fileData": {"mimeType": "image/png", "fileUri": "gs://b/c.png"}}, + ], + } + ], + }, + id="google_contents_with_media_the_translation_dropped", + ), + pytest.param( + { + "messages": [ + {"role": "user", "content": "What is new?"}, + {"role": "assistant", "content": "Weighing the release notes"}, + ], + "contents": [ + {"role": "user", "parts": [{"text": "What is new?"}]}, + {"role": "model", "parts": [{"text": "Weighing the release notes", "thought": True}]}, + ], + }, + id="google_contents_with_a_thought_the_translation_counts_as_text", + ), + ], +) +async def test_framed_estimate_that_cannot_read_the_body_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, unreadable_body: dict[str, list[dict[str, object]]] +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter(request=TokenCountRequest(model="kimi", **unreadable_body), call_endpoint=True) + readable: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + assert readable.count_source == "local_estimate" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "contents", + [ + pytest.param([{"role": "user", "parts": [{"text": "What is new in the release notes?"}]}], id="text"), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Weather in Paris?"}]}, + { + "role": "model", + "parts": [ + {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, "thoughtSignature": "c2ln"} + ], + }, + { + "role": "user", + "parts": [{"functionResponse": {"name": "get_weather", "response": {"forecast": "Sun"}}}], + }, + ], + id="sdk_agent_loop", + ), + pytest.param( + [ + {"role": "user", "parts": [{"text": "Hi"}]}, + {"role": "model", "parts": [{"text": "Hello!", "thoughtSignature": "c2ln"}]}, + {"role": "user", "parts": [{"text": "What is new?"}]}, + ], + id="sdk_chat_with_a_thought_signature", + ), + pytest.param([], id="empty"), + ], +) +async def test_framed_estimate_counts_google_contents_like_the_same_conversation_sent_as_messages( + monkeypatch: pytest.MonkeyPatch, contents: list[dict[str, object]] +): + _framed_kimi(monkeypatch) + google_request: Final = _google_count_tokens_request("kimi", {"contents": contents}) + + from_google: Final = await token_counter(request=google_request, call_endpoint=True) + from_messages: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=google_request.messages), call_endpoint=True + ) + + assert from_google.count_source == "local_estimate" + assert from_google.total_tokens == from_messages.total_tokens + + +@pytest.mark.asyncio +async def test_framed_deployment_whose_custom_tokenizer_cannot_load_answers_the_provider_error( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"custom_tokenizer": {"identifier": "moonshotai/kimi-k3", "revision": "main"}}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + + assert (raised.value.code, raised.value.message) == ("400", "scripted 400 for kimi") + + +def _framed_kimis_with_custom_tokenizers( + monkeypatch: pytest.MonkeyPatch, tokenizers: Mapping[str, tuple[str, str]], burst: int = 1 +) -> None: + counter: Final = _ScriptedCounter( + { + "us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400), + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=12), + }, + burst=burst, + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": {"model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + "model_info": {"token_count_model": {}}, + }, + *( + { + "model_name": name, + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": { + "custom_tokenizer": {"identifier": identifier, "revision": revision, "auth_token": None} + }, + } + for name, (identifier, revision) in tokenizers.items() + ), + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +def _framed_kimi_with_custom_tokenizer(monkeypatch: pytest.MonkeyPatch, identifier: str) -> None: + _framed_kimis_with_custom_tokenizers(monkeypatch, {"kimi": (identifier, "main")}) + + +async def _count_status(model: str = "kimi") -> str: + try: + response: Final = await token_counter( + request=TokenCountRequest(model=model, messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + except ProxyException as e: + return e.code + return response.count_source or "" + + +class _SkippableClockLoop(asyncio.SelectorEventLoop): + skipped_seconds: float = 0.0 + + def time(self) -> float: + return super().time() + self.skipped_seconds + + def skip(self, seconds: float) -> None: + self.skipped_seconds += seconds + + +def _run_on_skippable_clock( + scenario: Callable[[_SkippableClockLoop], Coroutine[object, object, tuple[str, ...]]], +) -> tuple[str, ...]: + loop: Final = _SkippableClockLoop() + try: + return loop.run_until_complete(scenario(loop)) + finally: + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +async def _until_finished(tasks: tuple[asyncio.Future[str], ...], count: int) -> None: + for finished in itertools.islice(asyncio.as_completed(tasks), count): + await finished + + +def _tokenizer_that_panics_while_encoding() -> SelectTokenizerResponse: + template: Final = Tokenizer(WordLevel({"[UNK]": 0, "": 1}, unk_token="[UNK]")) + template.pre_tokenizer = Whitespace() + template.post_processor = TemplateProcessing(single=" $A", special_tokens=[("", 1)]) + spec: Final = json.loads(template.to_str()) + return litellm.utils.create_tokenizer( + json.dumps({**spec, "post_processor": {**spec["post_processor"], "special_tokens": {}}}) + ) + + +def test_hanging_tokenizer_download_answers_the_provider_error_after_the_wait_with_one_download_while_it_runs( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/hanging-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def hanging_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "select_tokenizer", hanging_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + + async def burst_past_the_wait_then_one_more(loop: _SkippableClockLoop) -> tuple[str, ...]: + try: + burst: Final = asyncio.gather(*(_count_status() for _ in range(60))) + await asyncio.to_thread(download_started.wait, 10) + loop.skip(31) + answered: Final = await asyncio.wait_for(burst, 8) + later: Final = await asyncio.wait_for(_count_status(), 8) + finally: + hub_released.set() + return (*answered, later) + + assert set(_run_on_skippable_clock(burst_past_the_wait_then_one_more)) == {"400"} + assert len(downloads) == 1 + + +@pytest.mark.asyncio +async def test_framed_count_waits_for_a_tokenizer_download_that_finishes_within_the_wait( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-tokenizer") + real_select_tokenizer: Final = litellm.utils.select_tokenizer + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + + def slow_download(model: str, custom_tokenizer: object = None): + download_started.set() + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "select_tokenizer", slow_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + count: Final = asyncio.create_task(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + finally: + hub_released.set() + + assert await asyncio.wait_for(count, 10) == "local_estimate" + + +@pytest.mark.asyncio +async def test_failed_tokenizer_download_is_reused_until_the_retry_window_passes(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/flaky-tokenizer") + real_select_tokenizer: Final = litellm.utils.select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + within_window: Final = [await _count_status() for _ in range(3)] + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e-9) + after_window: Final = await _count_status() + + assert within_window == ["400", "400", "400"] + assert after_window == "local_estimate" + assert len(downloads) == 2 + + +def test_failed_tokenizer_load_is_reused_for_the_window_after_it_fails(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/slow-failing-tokenizer") + download_started: Final = threading.Event() + hub_released: Final = threading.Event() + downloads: Final[list[str]] = [] + + def slow_failing_download(model: str, custom_tokenizer: object = None) -> NoReturn: + downloads.append(model) + download_started.set() + hub_released.wait(10) + raise OSError("hub unreachable") + + monkeypatch.setattr(litellm.utils, "select_tokenizer", slow_failing_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 1000) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 60) + + async def fail_after_longer_than_the_window_then_retry(loop: _SkippableClockLoop) -> tuple[str, ...]: + first: Final = asyncio.ensure_future(_count_status()) + try: + await asyncio.to_thread(download_started.wait, 10) + loop.skip(120) + finally: + hub_released.set() + return await asyncio.wait_for(first, 8), await asyncio.wait_for(_count_status(), 8) + + assert _run_on_skippable_clock(fail_after_longer_than_the_window_then_retry) == ("400", "400") + assert len(downloads) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "second_tokenizer", + [("lit-tests/tokenizer-b", "main"), ("lit-tests/tokenizer-a", "v2")], + ids=["identifier", "revision"], +) +async def test_framed_deployments_with_different_custom_tokenizers_each_load_their_own( + monkeypatch: pytest.MonkeyPatch, second_tokenizer: tuple[str, str] +): + _framed_kimis_with_custom_tokenizers( + monkeypatch, {"kimi-a": ("lit-tests/tokenizer-a", "main"), "kimi-b": second_tokenizer} + ) + real_select_tokenizer: Final = litellm.utils.select_tokenizer + loaded: Final[list[object]] = [] + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + loaded.append(custom_tokenizer) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "select_tokenizer", load) + + statuses: Final = [await _count_status(name) for name in ("kimi-a", "kimi-b")] + + assert statuses == ["local_estimate", "local_estimate"] + assert loaded == [ + {"identifier": "lit-tests/tokenizer-a", "revision": "main", "auth_token": None}, + {"identifier": second_tokenizer[0], "revision": second_tokenizer[1], "auth_token": None}, + ] + + +@pytest.mark.asyncio +async def test_tokenizer_loads_past_the_cap_answer_the_provider_error_without_stalling_other_counts( + monkeypatch: pytest.MonkeyPatch, +): + tokenizers: Final = {f"kimi-{n}": (f"lit-tests/unreachable-tokenizer-{n}", "main") for n in range(5)} + _framed_kimis_with_custom_tokenizers(monkeypatch, tokenizers, burst=len(tokenizers)) + real_select_tokenizer: Final = litellm.utils.select_tokenizer + hub_released: Final = threading.Event() + downloads: Final[list[object]] = [] + + def held_download(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + downloads.append(custom_tokenizer) + hub_released.wait(10) + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "select_tokenizer", held_download) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 2) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_WAIT_SECONDS", 30) + monkeypatch.setattr(anyio.to_thread.current_default_thread_limiter(), "total_tokens", 2) + framed: Final = tuple(asyncio.ensure_future(_count_status(name)) for name in tokenizers) + try: + await asyncio.wait_for(_until_finished(framed, 3), 8) + unrelated: Final = await asyncio.wait_for(_count_status("sonnet"), 5) + finally: + hub_released.set() + statuses: Final = await asyncio.wait_for(asyncio.gather(*framed), 8) + + assert sorted(statuses) == ["400"] * 3 + ["local_estimate"] * 2 + assert unrelated == "provider" + assert len(downloads) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("panics_while", ["loading", "counting"]) +async def test_framed_count_answers_the_provider_error_when_the_tokenizer_panics( + monkeypatch: pytest.MonkeyPatch, panics_while: str +): + _framed_kimi_with_custom_tokenizer(monkeypatch, f"lit-tests/panicking-tokenizer-{panics_while}") + panicking: Final = _tokenizer_that_panics_while_encoding() + + def load(model: str, custom_tokenizer: object = None) -> SelectTokenizerResponse: + if panics_while == "loading": + panicking["tokenizer"].encode("hello world") + return panicking + + monkeypatch.setattr(litellm.utils, "select_tokenizer", load) + + assert await _count_status() == "400" + + +def _framed_kimi(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + +@pytest.mark.asyncio +async def test_local_count_loads_the_tokenizer_again_after_a_failed_load(monkeypatch: pytest.MonkeyPatch): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/local-route-tokenizer") + real_select_tokenizer: Final = litellm.utils.select_tokenizer + downloads: Final[list[str]] = [] + + def hub_down_once(model: str, custom_tokenizer: object = None): + downloads.append(model) + if len(downloads) == 1: + raise OSError("hub unreachable") + return real_select_tokenizer(model="gpt-5.4-mini") + + monkeypatch.setattr(litellm.utils, "select_tokenizer", hub_down_once) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", 1e9) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + with pytest.raises(OSError, match="hub unreachable"): + await token_counter(request=request) + retried: Final = await token_counter(request=request) + + assert retried.count_source == "local_estimate" and retried.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reuse_seconds", "tokenizer_types"), + [(1e9, ("openai_tokenizer", "openai_tokenizer")), (1e-9, ("openai_tokenizer", "huggingface_tokenizer"))], + ids=["within_the_window", "after_the_window"], +) +async def test_framed_count_loads_the_tokenizer_again_only_after_the_window_of_a_load_that_fell_back( + monkeypatch: pytest.MonkeyPatch, reuse_seconds: float, tokenizer_types: tuple[str, str] +): + _framed_kimi(monkeypatch) + real_select_tokenizer: Final = litellm.utils.select_tokenizer + hub_tokenizer: Final = litellm.utils.create_tokenizer( + Tokenizer(WordLevel({"[UNK]": 0}, unk_token="[UNK]")).to_str() + ) + loads: Final[list[str]] = [] + + def hub_up_on_the_second_load(model: str, custom_tokenizer: object = None): + loads.append(model) + return real_select_tokenizer(model="gpt-5.4-mini") if len(loads) == 1 else hub_tokenizer + + monkeypatch.setattr(litellm.utils, "select_tokenizer", hub_up_on_the_second_load) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_REUSE_SECONDS", reuse_seconds) + monkeypatch.setattr(litellm.proxy.proxy_server, "TOKENIZER_LOAD_MAX_RUNNING", 1) + request: Final = TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]) + + first: Final = await token_counter(request=request, call_endpoint=True) + second: Final = await token_counter(request=request, call_endpoint=True) + + assert (first.tokenizer_type, second.tokenizer_type) == tokenizer_types + + +async def _fork_guard_traceback_frames() -> int: + with pytest.raises(ForkedAfterNativeRuntimeStarted) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "What is new?"}]), + call_endpoint=True, + ) + return len(traceback.extract_tb(raised.value.__traceback__)) + + +@pytest.mark.asyncio +async def test_fork_guard_error_while_loading_the_tokenizer_is_raised_to_every_count_sharing_the_load( + monkeypatch: pytest.MonkeyPatch, +): + _framed_kimi(monkeypatch) + + def forked_load(model: str, custom_tokenizer: object = None) -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + monkeypatch.setattr(litellm.utils, "select_tokenizer", forked_load) + + frames: Final = [await _fork_guard_traceback_frames() for _ in range(3)] + + assert len(set(frames)) == 1 + + +@pytest.mark.asyncio +async def test_swallowed_estimate_error_is_logged_by_type_without_the_request_content( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "kimi", "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}}]), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name), pytest.raises(ProxyException): + await token_counter( + request=TokenCountRequest( + model="kimi", + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok", "tool_calls": "private-marker-in-the-prompt"}, + ], + ), + call_endpoint=True, + ) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("count error 400" in message and "ValueError" in message for message in warnings) + assert not any("private-marker-in-the-prompt" in message for message in warnings) + + +@pytest.mark.asyncio +async def test_fork_guard_error_in_the_estimate_is_never_answered_as_the_provider_error(): + async def forked_estimate() -> NoReturn: + raise ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + with pytest.raises(ForkedAfterNativeRuntimeStarted): + await _estimate_or_provider_error( + forked_estimate(), provider_failure=ProviderCountFailure(status_code=400, message="scripted 400") + ) + + +@pytest.mark.asyncio +async def test_local_estimate_adds_the_configured_framing_once_per_request(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": f"sk-unused-{name}"}, + "model_info": model_info, + } + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[deployment("framed", {"token_count_framing": 100}), deployment("plain", {})]), + ) + messages: Final = [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hello there"}] + + plain: Final = await token_counter(request=TokenCountRequest(model="plain", messages=messages)) + framed: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + framed_prompt: Final = await token_counter(request=TokenCountRequest(model="framed", prompt="hello there")) + + assert plain.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + assert framed.total_tokens == plain.total_tokens + 100 + assert (plain.count_source, framed.count_source) == ("local_estimate", "local_estimate") + assert framed_prompt.total_tokens == litellm.token_counter(model="gpt-5.4-mini", text="hello there") + + +@pytest.mark.asyncio +async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_names(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, litellm_params: dict[str, object], model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {**litellm_params, "api_key": "unused", "api_base": "https://unused.invalid"}, + "model_info": model_info, + } + + kimi: Final = "databricks/system.ai.kimi-k3" + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + deployment("prefixed", {"model": kimi}, {}), + deployment("split", {"model": "system.ai.kimi-k3", "custom_llm_provider": "databricks"}, {}), + deployment("aliased", {"model": "gpt-5.4-mini", "custom_llm_provider": "openai"}, {"base_model": kimi}), + ] + ), + ) + messages: Final = [{"role": "user", "content": "hello there"}] + framing: Final = cost_map_entry(kimi)["token_count_framing"] + + prefixed: Final = await token_counter(request=TokenCountRequest(model="prefixed", messages=messages)) + split: Final = await token_counter(request=TokenCountRequest(model="split", messages=messages)) + aliased: Final = await token_counter(request=TokenCountRequest(model="aliased", messages=messages)) + + assert prefixed.total_tokens == litellm.token_counter(model="system.ai.kimi-k3", messages=messages) + framing + assert split.total_tokens == prefixed.total_tokens + assert aliased.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + framing + + +_KIMI_WEATHER: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +def _byte_level_kimi_k3(path: Path) -> KimiK3TokenizerSource: + path.write_bytes(b"\n".join(base64.b64encode(bytes((value,))) + b" %d" % value for value in range(256))) + return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + + +def _available_kimi_k3(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> tiktoken.Encoding: + source: Final = _byte_level_kimi_k3(tmp_path / "tiktoken.model") + monkeypatch.setattr( + litellm.proxy.proxy_server, "_KIMI_K3_ENCODINGS", KimiK3Encodings(source=source, retry_after_seconds=60) + ) + encoding: Final = load_kimi_k3_encoding(source) + assert encoding is not None + return encoding + + +def _shipped_kimi_k3_framing() -> int: + framing: Final = cost_map_entry("bedrock/us.moonshotai.kimi-k3")["token_count_framing"] + assert isinstance(framing, int) + return framing + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_is_counted_with_its_tokenizer_plus_the_shipped_framing_in_both_request_formats( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi(monkeypatch) + encoding: Final = _available_kimi_k3(monkeypatch, tmp_path) + question: Final = {"role": "user", "content": "Weather in Paris?"} + openai_messages: Final = [ + {"role": "system", "content": "Be terse."}, + question, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}} + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "Sun"}, + ] + + openai: Final = await count_request_tokens( + request=TokenCountRequest(model="kimi", messages=openai_messages, tools=[_KIMI_WEATHER]), call_endpoint=True + ) + anthropic: Final = await count_request_tokens( + request=TokenCountRequest( + model="kimi", + system="Be terse.", + messages=[ + question, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]}, + ], + tools=[ + { + "name": "get_weather", + "description": "Get the weather", + "input_schema": _KIMI_WEATHER["function"]["parameters"], + } + ], + ), + call_endpoint=True, + messages_format="anthropic", + ) + exact: Final = kimi_k3_tokens(encoding, "us.moonshotai.kimi-k3", openai_messages, [_KIMI_WEATHER]) + + assert exact is not None + assert (openai.total_tokens, openai.tokenizer_type, openai.count_source) == ( + exact + _shipped_kimi_k3_framing(), + "huggingface_tokenizer", + "local_estimate", + ) + assert (anthropic.total_tokens, anthropic.tokenizer_type) == (openai.total_tokens, "huggingface_tokenizer") + + +@pytest.mark.asyncio +async def test_bedrock_kimi_k3_deployment_with_a_custom_tokenizer_keeps_counting_with_it( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + _framed_kimi_with_custom_tokenizer(monkeypatch, "lit-tests/kimi-custom-tokenizer") + _available_kimi_k3(monkeypatch, tmp_path) + custom: Final = litellm.utils.select_tokenizer(model="gpt-5.4-mini") + monkeypatch.setattr(litellm.utils, "select_tokenizer", lambda model, custom_tokenizer=None: custom) + messages: Final = [{"role": "user", "content": "What is new?"}] + + counted: Final = await token_counter(request=TokenCountRequest(model="kimi", messages=messages), call_endpoint=True) + + assert counted.total_tokens == ( + litellm.token_counter(model=counted.model_used, messages=messages, custom_tokenizer=custom) + + kimi_k3_template_tokens(len(messages), False) + + _shipped_kimi_k3_framing() + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("messages", "tools", "template_tokens"), + [ + pytest.param([{"role": "user", "content": "What is new?"}], None, 14, id="one_message"), + pytest.param( + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": "Sunny."}, + {"role": "user", "content": "And Oslo?"}, + ], + [_KIMI_WEATHER], + 14 + 3 * 13 + 40, + id="turns_and_tools", + ), + ], +) +async def test_bedrock_kimi_k3_without_its_tokenizer_adds_the_template_offsets_to_the_local_estimate( + monkeypatch: pytest.MonkeyPatch, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + template_tokens: int, +): + _framed_kimi(monkeypatch) + + counted: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=messages, tools=tools), call_endpoint=True + ) + estimate: Final = litellm.token_counter( + model=counted.model_used, + messages=messages, + tools=tools, + custom_tokenizer=litellm.utils.select_tokenizer(model=counted.model_used), + ) + + assert counted.total_tokens == estimate + template_tokens + _shipped_kimi_k3_framing() + + +@pytest.mark.asyncio +async def test_slow_cost_map_lookup_does_not_stall_other_requests_on_the_worker(monkeypatch: pytest.MonkeyPatch): + real_get_model_info: Final = litellm.get_model_info + lookup_started: Final = threading.Event() + other_request_ran: Final = threading.Event() + released_by_the_other_request: Final[list[bool]] = [] + + def slow_get_model_info(model: str, *args: object, **kwargs: object): + if model == "slow-lookup-model": + lookup_started.set() + released_by_the_other_request.append(other_request_ran.wait(5)) + return real_get_model_info(model, *args, **kwargs) + + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) + monkeypatch.setattr(litellm, "get_model_info", slow_get_model_info) + counting: Final = asyncio.create_task( + token_counter( + request=TokenCountRequest(model="slow-lookup-model", messages=[{"role": "user", "content": "hi"}]) + ) + ) + await asyncio.to_thread(lookup_started.wait, 5) + other_request_ran.set() + await counting + + assert set(released_by_the_other_request) == {True} + + +@pytest.mark.asyncio +async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is_disabled( + monkeypatch: pytest.MonkeyPatch, +): + counter: Final = _ScriptedCounter({"us.moonshotai.kimi-k3": _provider_reply("kimi", status_code=400)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "kimi", + "litellm_params": {"model": "bedrock/us.moonshotai.kimi-k3"}, + "model_info": {"token_count_framing": 78}, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + monkeypatch.setattr(litellm, "disable_token_counter", True) + + with pytest.raises(ProxyException) as raised: + await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert raised.value.code == "400" + assert "scripted 400 for kimi" in raised.value.message + + +@pytest.mark.parametrize( + ("model_info", "counted_with"), + [ + ({}, ("us.anthropic.claude-opus-5-5", "us-east-1", 447)), + ( + {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}}, + ("us.anthropic.claude-opus-5-5", "ap-southeast-4", 447), + ), + ({"token_count_model": {}}, ("us.anthropic.claude-sonnet-5-5", "us-west-2", 300)), + ], +) +@pytest.mark.asyncio +async def test_configured_token_count_model_overrides_the_shipped_one_and_an_empty_one_counts_the_model_itself( + monkeypatch: pytest.MonkeyPatch, model_info: dict[str, object], counted_with: tuple[str, str, int] +): + counter: Final = _ScriptedCounter( + { + "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=300), + "us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447), + } + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_region_name": "us-west-2", + }, + "model_info": model_info, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + counted: Final = await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0], counted.total_tokens) == counted_with + assert len(counter.calls) == 1 + + +@pytest.mark.parametrize( + ("aws_region_name", "counted_in"), + [("eu-central-1", "eu-central-1"), ("us-west-2", "us-east-1")], +) +@pytest.mark.asyncio +async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography( + monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str +): + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "sonnet", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-5-5", + "aws_region_name": aws_region_name, + }, + } + ] + ), + ) + monkeypatch.setattr( + litellm.proxy.proxy_server, "_get_provider_token_counter", lambda deployment, model: (counter, model, "bedrock") + ) + + await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True + ) + + assert (counter.calls[0][0], counter.regions[0]) == ("anthropic.claude-opus-5-5", counted_in) + + +class _SlowVault: + def sync_read_secret(self, secret_name: str) -> None: + time.sleep(0.5) + + +async def _longest_stall(done: asyncio.Event) -> float: + if done.is_set(): + return 0.0 + started: Final = time.monotonic() + await asyncio.sleep(0.01) + return max(time.monotonic() - started, await _longest_stall(done)) + + +@pytest.mark.asyncio +async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + + def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]: + return counter, model, "bedrock" + + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[{"model_name": "sonnet", "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-5-5"}}] + ), + ) + monkeypatch.setattr(proxy_server, "_get_provider_token_counter", bedrock_counter) + monkeypatch.setattr(litellm, "secret_manager_client", _SlowVault()) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + done: Final = asyncio.Event() + + async def count() -> TokenCountResponse: + try: + return await token_counter( + request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), + call_endpoint=True, + ) + finally: + done.set() + + started: Final = time.monotonic() + stall, response = await asyncio.gather(_longest_stall(done), count()) + + assert time.monotonic() - started >= 1.0 + assert stall < 0.4 + assert response.total_tokens == 447 + + +@pytest.mark.asyncio +async def test_framing_still_applies_while_the_deployment_is_cooling_down(monkeypatch: pytest.MonkeyPatch): + router: Final = Router( + model_list=[ + { + "model_name": "framed", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-unused-framed"}, + "model_info": {"id": "framed-deployment", "token_count_framing": 100}, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + messages: Final = [{"role": "user", "content": "hello there"}] + warm: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + router.cooldown_cache.add_deployment_to_cooldown( + model_id="framed-deployment", + original_exception=Exception("rate limited"), + exception_status=429, + cooldown_time=60, + ) + cooling: Final = await token_counter(request=TokenCountRequest(model="framed", messages=messages)) + + assert warm.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + 100 + assert cooling.total_tokens - 100 == litellm.token_counter(model="framed", messages=messages) + assert cooling.count_source == "local_estimate" + + _GEMINI_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 @@ -1360,7 +2694,7 @@ async def test_gemini_messages_count_tokens_applies_the_context_edits_inference_ router, {**request_options, "model": requested_model} ) - assert response == {"input_tokens": 41} + assert response == {"input_tokens": 41, "litellm_count_source": "provider"} assert len(sent_bodies) == 1 assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent @@ -1426,7 +2760,7 @@ async def test_gemini_messages_count_tokens_keeps_tool_results_when_a_routable_d Router(model_list=model_list), {"model": requested_model} ) - assert response == {"input_tokens": 41} + assert response == {"input_tokens": 41, "litellm_count_source": "provider"} assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 @@ -1493,7 +2827,7 @@ async def test_gemini_messages_count_tokens_reads_the_opt_out_of_the_deployment_ make_router(), {**request_options, "model": requested_model} ) - assert response == {"input_tokens": 41} + assert response == {"input_tokens": 41, "litellm_count_source": "provider"} assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent @@ -1517,7 +2851,7 @@ async def test_gemini_messages_count_tokens_reads_the_opt_out_of_the_completion_ with patch.object(proxy_server, "general_settings", {"completion_model": completion_model}): response, sent_bodies = await _count_through_the_messages_endpoint(router, {"model": requested_model}) - assert response == {"input_tokens": 41} + assert response == {"input_tokens": 41, "litellm_count_source": "provider"} assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == expected_tool_outputs_sent @@ -1591,7 +2925,7 @@ async def test_gemini_messages_count_tokens_counts_the_unedited_body_when_drop_p router, {**request_options, "model": "gemini-agent", "context_management": _INVALID_COMPACT_EDIT} ) - assert response == {"input_tokens": 41} + assert response == {"input_tokens": 41, "litellm_count_source": "provider"} assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index b4a0837578b..094a13c37b2 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -584,6 +584,40 @@ def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: _restore_model_cost_entries(model_keys) +def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: + """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on + the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them.""" + from litellm.proxy.common_utils.token_count_settings import cost_map_entry + + backend_model: Final = "bedrock/us.anthropic.claude-sonnet-5-5" + deploy_id: Final = "sonnet-deploy-keep-token-count" + configured: Final = { + "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"}, + "token_count_framing": 7, + } + builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured} + assert all(builtin[field] != value for field, value in configured.items()), builtin + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in (deploy_id, backend_model, backend_model.removeprefix("bedrock/")) + } + try: + Router( + model_list=[ + { + "model_name": "sonnet-keep-token-count", + "litellm_params": {"model": backend_model}, + "model_info": {"id": deploy_id, **configured}, + } + ], + ) + + assert {field: litellm.model_cost[deploy_id][field] for field in configured} == configured + assert {field: cost_map_entry(backend_model).get(field) for field in configured} == builtin + finally: + _restore_model_cost_entries(model_keys) + + def test_should_store_full_pricing_under_deployment_model_id(): """ Per-deployment pricing (including zero) should be stored and @@ -3884,9 +3918,7 @@ def test_router_capacity_widens_to_a_larger_base_model_window( ) _invalidate_model_cost_lowercase_map() router: Final = _capacity_test_router(route, base, base_location, explicit_limit) - expected: Final = ( - explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192) - ) + expected: Final = explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192) assert _effective_max_input_tokens(router) == expected assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) _invalidate_model_cost_lowercase_map() diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a8e90fded83..b4150aad1c7 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1138,6 +1138,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_box_questions": {"type": "boolean"}, "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, + "token_count_framing": {"type": "integer"}, + "token_count_model": { + "type": "object", + "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, + "required": ["model"], + }, + "token_count_tokenizer": {"type": "string", "enum": ["kimi_k3"]}, "supports_bedrock_runtime_chat_completions_tools_with_reasoning": {"type": "boolean"}, "supports_bedrock_runtime_chat_completions_response_format": {"type": "boolean"}, "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 468811966de..c33ef86af72 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -46847,9 +46847,17 @@ export interface components { * @description Response structure for token count details with modality breakdown. * * Example: - * {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]} + * {'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}], + * 'litellm_count_source': 'provider'} */ TokenCountDetailsResponse: { + /** Litellm Count Model */ + litellm_count_model?: string; + /** + * Litellm Count Source + * @enum {string} + */ + litellm_count_source?: "provider" | "token_count_model" | "local_estimate"; /** Prompttokensdetails */ promptTokensDetails: components["schemas"]["PromptTokensDetails"][]; /** Totaltokens */ @@ -46878,6 +46886,8 @@ export interface components { }; /** TokenCountResponse */ TokenCountResponse: { + /** Count Source */ + count_source?: ("provider" | "token_count_model" | "local_estimate") | null; /** * Error * @default false From 13335266f59a822dcafaadf7329c6a938867c265 Mon Sep 17 00:00:00 2001 From: Yousef Joe Date: Sun, 18 Oct 2026 02:16:03 +0200 Subject: [PATCH 61/64] fix(baseten): forward reasoning_effort to Baseten Model API (#40152) * fix(baseten): forward reasoning_effort to Baseten Model API Baseten's Model API accepts a top-level reasoning_effort parameter and validates it server-side (https://docs.baseten.co/inference/model-apis/reasoning), but BasetenConfig.get_supported_openai_params did not list it, so LiteLLM silently dropped it when drop_params=True and raised UnsupportedParamsError otherwise. Add reasoning_effort to the supported params so map_openai_params forwards it verbatim, mirroring the hosted_vllm and deepseek providers. Baseten returns a clear 400 for values outside a model's supported set and ignores the parameter on non-reasoning models, so verbatim passthrough is safe for the whole catalog. [Claude_Code] * fix(baseten): forward reasoning_effort to Baseten Model API Baseten's Model API accepts a top-level reasoning_effort parameter and validates it server-side (https://docs.baseten.co/inference/model-apis/reasoning), but BasetenConfig.get_supported_openai_params did not list it, so LiteLLM silently dropped it when drop_params=True and raised UnsupportedParamsError otherwise. Add reasoning_effort to the supported params so map_openai_params forwards it verbatim, mirroring the hosted_vllm and deepseek providers. Baseten returns a clear 400 for values outside a model's supported set and ignores the parameter on non-reasoning models, so verbatim passthrough is safe for the whole catalog. --------- Co-authored-by: ryan-crabbe-berri --- litellm/llms/baseten/chat.py | 1 + .../baseten/chat/test_baseten_completions.py | 36 +++++++++++++++---- 2 files changed, 30 insertions(+), 7 deletions(-) diff --git a/litellm/llms/baseten/chat.py b/litellm/llms/baseten/chat.py index 0374646f2f3..e92f47c41c0 100644 --- a/litellm/llms/baseten/chat.py +++ b/litellm/llms/baseten/chat.py @@ -54,6 +54,7 @@ class BasetenConfig(OpenAIGPTConfig): return [ "max_tokens", "max_completion_tokens", + "reasoning_effort", "response_format", "seed", "stop", diff --git a/tests/unit/llms/baseten/chat/test_baseten_completions.py b/tests/unit/llms/baseten/chat/test_baseten_completions.py index caf96987933..9195df13242 100644 --- a/tests/unit/llms/baseten/chat/test_baseten_completions.py +++ b/tests/unit/llms/baseten/chat/test_baseten_completions.py @@ -18,10 +18,7 @@ class TestBasetenRouting: ) # Model API (non-8-character) - assert ( - config.get_api_base_for_model("openai/gpt-oss-120b") - == "https://inference.baseten.co/v1" - ) + assert config.get_api_base_for_model("openai/gpt-oss-120b") == "https://inference.baseten.co/v1" class TestBasetenModelAPI: @@ -47,12 +44,37 @@ class TestBasetenModelAPI: assert result["top_p"] == 0.9 # Test provider info - api_base, api_key = config._get_openai_compatible_provider_info( - None, "test-key" - ) + api_base, api_key = config._get_openai_compatible_provider_info(None, "test-key") assert api_base == "https://inference.baseten.co/v1" assert api_key == "test-key" +class TestBasetenReasoningEffort: + def test_reasoning_effort_passthrough_with_drop_params(self): + config = BasetenConfig() + + result = config.map_openai_params( + non_default_params={"reasoning_effort": "high", "max_tokens": 100}, + optional_params={}, + model="deepseek-ai/DeepSeek-V4-Flash-0731", + drop_params=True, + ) + + assert result["reasoning_effort"] == "high" + assert result["max_tokens"] == 100 + + def test_reasoning_effort_none_passthrough(self): + config = BasetenConfig() + + result = config.map_openai_params( + non_default_params={"reasoning_effort": "none"}, + optional_params={}, + model="moonshotai/Kimi-K3", + drop_params=True, + ) + + assert result["reasoning_effort"] == "none" + + if __name__ == "__main__": pytest.main([__file__]) From d5a8c5fd5120b987157671abe054fc8dd3f8e865 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:27:12 -0700 Subject: [PATCH 62/64] fix(bedrock): take the 200K context window from the Claude Sonnet 4 model card (#47202) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 10 +++++----- model_prices_and_context_window.json | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cc931c89a31..da625869e9a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3028,7 +3028,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -3313,7 +3313,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -24707,7 +24707,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -33913,7 +33913,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -48940,7 +48940,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cc931c89a31..da625869e9a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3028,7 +3028,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -3313,7 +3313,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -24707,7 +24707,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -33913,7 +33913,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -48940,7 +48940,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", - "max_input_tokens": 1000000, + "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", From 4a8b47542cfa64db706cbd3d0a4e8fa2cea2a5f1 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:34:03 -0700 Subject: [PATCH 63/64] feat(ui): show gateway URL and agent setup instructions after creating a key (#47199) * feat(ui): show gateway URL and agent setup instructions after creating a key Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(ui): make the selected agent button visibly distinct in dark mode Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(ui): hide agent instructions by default and use a generic title Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(ui): use one shared agent instructions card with Claude Code and Codex logos Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../add_agent_form.integration.test.tsx | 1 + .../_components/add_agent_form.test.tsx | 1 + .../organisms/create_key_button.tsx | 2 +- .../shared/CreatedKeyDisplay.test.tsx | 115 +++++++++++++++++- .../components/shared/CreatedKeyDisplay.tsx | 80 ++++++++++++ 5 files changed, 197 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx index 67ac770a65f..53ce661baaf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx @@ -10,6 +10,7 @@ import { chooseSelectOption, renderWithProviders as render } from "../../../../. vi.mock("@/components/networking", () => ({ apiClient: { get: vi.fn() }, createAgentCall: vi.fn(), + getProxyBaseUrl: vi.fn(() => "http://gw.test:4000/"), getAgentCreateMetadata: vi.fn(), getAgentsList: vi.fn(), keyCreateForAgentCall: vi.fn(), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx index 5e8ba145396..5b8cb652510 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx @@ -8,6 +8,7 @@ import type { AgentCreateInfo } from "@/components/networking"; vi.mock("@/components/networking", () => ({ createAgentCall: vi.fn(), + getProxyBaseUrl: vi.fn(() => "http://gw.test:4000/"), getAgentCreateMetadata: vi.fn(), getAgentsList: vi.fn(), keyCreateForAgentCall: vi.fn(), diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 30c7a80c456..30ff097a87c 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1914,7 +1914,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp {apiKey && ( !open && handleCancel()}> - +
Save your Key {apiKey != null ? ( diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx index 0e10d32f8d1..310f69d28bc 100644 --- a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx @@ -1,13 +1,43 @@ +import { cloneElement, type MouseEventHandler, type ReactElement, type ReactNode } from "react"; import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { render, screen, act } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import CreatedKeyDisplay from "./CreatedKeyDisplay"; +import CreatedKeyDisplay, { buildAgentInstructions } from "./CreatedKeyDisplay"; import { toast } from "@/lib/toast"; +const mockCopy = vi.hoisted(() => vi.fn<(text: string) => void>()); + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => "http://gw.test:4000/"), +})); + +vi.mock("react-copy-to-clipboard", () => ({ + CopyToClipboard: ({ + text, + onCopy, + children, + }: { + text: string; + onCopy?: (text: string, result: boolean) => void; + children: ReactNode; + }) => { + const child = children as ReactElement<{ onClick?: MouseEventHandler }>; + return cloneElement(child, { + onClick: (event) => { + child.props.onClick?.(event); + mockCopy(text); + onCopy?.(text, true); + }, + }); + }, +})); + describe("CreatedKeyDisplay", () => { beforeEach(() => { vi.useFakeTimers({ shouldAdvanceTime: true }); + mockCopy.mockReset(); + vi.mocked(toast.success).mockClear(); }); afterEach(() => { @@ -69,4 +99,87 @@ describe("CreatedKeyDisplay", () => { expect(screen.getByRole("button", { name: /copy virtual key/i })).toBeInTheDocument(); }); + + it("shows the gateway URL without its trailing slash and copies it", async () => { + const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); + render(); + + expect(screen.getByText("http://gw.test:4000")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Copy URL" })); + + expect(mockCopy).toHaveBeenCalledWith("http://gw.test:4000"); + expect(toast.success).toHaveBeenCalledWith("Gateway URL copied to clipboard"); + }); + + it("shows one shared instruction card with both agent logos and expands on demand", async () => { + const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); + render(); + + expect(screen.getByRole("img", { name: "Claude Code" })).toBeInTheDocument(); + expect(screen.getByRole("img", { name: "Codex" })).toBeInTheDocument(); + expect(screen.getByText("Agent instructions")).toBeInTheDocument(); + expect(screen.queryByText("Set up your agent")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Claude Code" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Codex" })).not.toBeInTheDocument(); + + const showButton = screen.getByRole("button", { name: "Show" }); + expect(showButton).toHaveAttribute("aria-expanded", "false"); + expect(screen.queryByText(/POST http:\/\/gw\.test:4000\/v1\/responses/)).not.toBeInTheDocument(); + + await user.click(showButton); + + expect(screen.getByRole("button", { name: "Hide" })).toHaveAttribute("aria-expanded", "true"); + expect(screen.getByText(/POST http:\/\/gw\.test:4000\/v1\/responses/)).toBeInTheDocument(); + expect(screen.getByText(/ANTHROPIC_BASE_URL=http:\/\/gw\.test:4000/)).toBeInTheDocument(); + expect(screen.getByText(/OPENAI_BASE_URL=http:\/\/gw\.test:4000\/v1/)).toBeInTheDocument(); + }); + + it("copies complete shared instructions while the card is collapsed", async () => { + const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime }); + render(); + + expect(screen.getByRole("button", { name: "Show" })).toHaveAttribute("aria-expanded", "false"); + expect(screen.queryByText(/POST http:\/\/gw\.test:4000\/v1\/responses/)).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Copy instructions" })); + + const copiedInstructions = mockCopy.mock.calls.at(-1)?.[0] ?? ""; + expect(copiedInstructions).toContain("Gateway URL: http://gw.test:4000"); + expect(copiedInstructions).toContain("Virtual key: sk-test-123"); + expect(copiedInstructions).toContain("POST http://gw.test:4000/v1/chat/completions"); + expect(copiedInstructions).toContain("POST http://gw.test:4000/v1/responses"); + expect(copiedInstructions).toContain("POST http://gw.test:4000/v1/messages"); + expect(copiedInstructions).toContain("GET http://gw.test:4000/v1/models"); + expect(copiedInstructions).toContain( + "To run Claude Code through this gateway, set ANTHROPIC_BASE_URL=http://gw.test:4000 and ANTHROPIC_AUTH_TOKEN=sk-test-123.", + ); + expect(copiedInstructions).toContain( + "To run Codex through this gateway, set OPENAI_BASE_URL=http://gw.test:4000/v1 and OPENAI_API_KEY=sk-test-123.", + ); + expect(screen.getByRole("button", { name: "Copied!" })).toBeInTheDocument(); + expect(toast.success).toHaveBeenCalledWith("Instructions copied"); + }); +}); + +describe("buildAgentInstructions", () => { + it("builds shared instructions with both agent settings and a normalized URL", () => { + expect(buildAgentInstructions("http://gw.test:4000/", "sk-test-123")).toBe( + [ + "Use this LiteLLM AI gateway for model inference.", + "", + "Gateway URL: http://gw.test:4000", + "Virtual key: sk-test-123", + "", + "LiteLLM is a gateway that exposes OpenAI and Anthropic compatible APIs. Send the virtual key as a Bearer token (Authorization: Bearer sk-test-123) to:", + "- POST http://gw.test:4000/v1/chat/completions (OpenAI Chat Completions)", + "- POST http://gw.test:4000/v1/responses (OpenAI Responses)", + "- POST http://gw.test:4000/v1/messages (Anthropic Messages)", + "- GET http://gw.test:4000/v1/models (models this key can use)", + "", + "To run Claude Code through this gateway, set ANTHROPIC_BASE_URL=http://gw.test:4000 and ANTHROPIC_AUTH_TOKEN=sk-test-123.", + "To run Codex through this gateway, set OPENAI_BASE_URL=http://gw.test:4000/v1 and OPENAI_API_KEY=sk-test-123.", + "", + "Docs for everything else LiteLLM supports: https://docs.litellm.ai/docs/", + ].join("\n"), + ); + }); }); diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx index 5dbc17b13bf..d58a1864575 100644 --- a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx +++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx @@ -1,5 +1,9 @@ import React, { useState } from "react"; import { CopyToClipboard } from "react-copy-to-clipboard"; +import { ChevronDown, ChevronUp } from "lucide-react"; +import anthropicLogo from "../../../public/assets/logos/anthropic.svg"; +import openaiLogo from "../../../public/assets/logos/openai_small.svg"; +import { getProxyBaseUrl } from "@/components/networking"; import { Button } from "@/components/ui/button"; import { toast } from "@/lib/toast"; @@ -7,12 +11,41 @@ interface CreatedKeyDisplayProps { apiKey: string; } +export const buildAgentInstructions = (baseUrl: string, apiKey: string): string => { + const normalizedBaseUrl = baseUrl.replace(/\/+$/, ""); + const endpointIntroduction = + "LiteLLM is a gateway that exposes OpenAI and Anthropic compatible APIs. " + + `Send the virtual key as a Bearer token (Authorization: Bearer ${apiKey}) to:`; + + return [ + "Use this LiteLLM AI gateway for model inference.", + "", + `Gateway URL: ${normalizedBaseUrl}`, + `Virtual key: ${apiKey}`, + "", + endpointIntroduction, + `- POST ${normalizedBaseUrl}/v1/chat/completions (OpenAI Chat Completions)`, + `- POST ${normalizedBaseUrl}/v1/responses (OpenAI Responses)`, + `- POST ${normalizedBaseUrl}/v1/messages (Anthropic Messages)`, + `- GET ${normalizedBaseUrl}/v1/models (models this key can use)`, + "", + `To run Claude Code through this gateway, set ANTHROPIC_BASE_URL=${normalizedBaseUrl} and ANTHROPIC_AUTH_TOKEN=${apiKey}.`, + `To run Codex through this gateway, set OPENAI_BASE_URL=${normalizedBaseUrl}/v1 and OPENAI_API_KEY=${apiKey}.`, + "", + "Docs for everything else LiteLLM supports: https://docs.litellm.ai/docs/", + ].join("\n"); +}; + /** * Shared component for displaying a newly-created virtual key. * Used on the Virtual Keys page and in the Add Agent wizard. */ const CreatedKeyDisplay: React.FC = ({ apiKey }) => { const [copied, setCopied] = useState(false); + const [instructionsCopied, setInstructionsCopied] = useState(false); + const [showInstructions, setShowInstructions] = useState(false); + const baseUrl = getProxyBaseUrl().replace(/\/+$/, ""); + const instructions = buildAgentInstructions(baseUrl, apiKey); const handleCopy = () => { setCopied(true); @@ -20,6 +53,12 @@ const CreatedKeyDisplay: React.FC = ({ apiKey }) => { setTimeout(() => setCopied(false), 2000); }; + const handleCopyInstructions = () => { + setInstructionsCopied(true); + toast.success("Instructions copied"); + setTimeout(() => setInstructionsCopied(false), 2000); + }; + return (

@@ -36,6 +75,47 @@ const CreatedKeyDisplay: React.FC = ({ apiKey }) => { + +

Gateway URL:

+
+ 
{baseUrl}
+ toast.success("Gateway URL copied to clipboard")}> + + +
+ +
+
+
+

Agent instructions

+ Claude Code + Codex +
+
+ + + + +
+
+ {showInstructions && ( + 
{instructions}
+ )} +
); }; From 52023601a872d9ae4b14d8af3f35ad6da43258d1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 17 Oct 2026 17:47:15 -0700 Subject: [PATCH 64/64] test(baseten): annotate test locals with Final (#47205) Follow-up to #40152. Greptile flagged the two new reasoning_effort tests for assigning config and result without Final. Annotates every local in the file so it passes the LIT010 check cleanly --- .../baseten/chat/test_baseten_completions.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/unit/llms/baseten/chat/test_baseten_completions.py b/tests/unit/llms/baseten/chat/test_baseten_completions.py index 9195df13242..7cac16cca1c 100644 --- a/tests/unit/llms/baseten/chat/test_baseten_completions.py +++ b/tests/unit/llms/baseten/chat/test_baseten_completions.py @@ -1,4 +1,5 @@ import os +from typing import Final import pytest from unittest.mock import patch from litellm.llms.baseten.chat import BasetenConfig @@ -9,7 +10,7 @@ class TestBasetenRouting: def test_routing_logic(self): """Test routing between Model API and dedicated deployments""" - config = BasetenConfig() + config: Final = BasetenConfig() # Dedicated deployment (8-character alphanumeric) assert ( @@ -27,12 +28,12 @@ class TestBasetenModelAPI: @patch.dict(os.environ, {"BASETEN_API_KEY": "test-key"}) def test_model_api_inference(self): """Test Model API inference with basic parameters""" - config = BasetenConfig() + config: Final = BasetenConfig() # Test parameter mapping - non_default_params = {"max_tokens": 100, "temperature": 0.7, "top_p": 0.9} + non_default_params: Final = {"max_tokens": 100, "temperature": 0.7, "top_p": 0.9} - result = config.map_openai_params( + result: Final = config.map_openai_params( non_default_params=non_default_params, optional_params={}, model="openai/gpt-oss-120b", @@ -51,9 +52,9 @@ class TestBasetenModelAPI: class TestBasetenReasoningEffort: def test_reasoning_effort_passthrough_with_drop_params(self): - config = BasetenConfig() + config: Final = BasetenConfig() - result = config.map_openai_params( + result: Final = config.map_openai_params( non_default_params={"reasoning_effort": "high", "max_tokens": 100}, optional_params={}, model="deepseek-ai/DeepSeek-V4-Flash-0731", @@ -64,9 +65,9 @@ class TestBasetenReasoningEffort: assert result["max_tokens"] == 100 def test_reasoning_effort_none_passthrough(self): - config = BasetenConfig() + config: Final = BasetenConfig() - result = config.map_openai_params( + result: Final = config.map_openai_params( non_default_params={"reasoning_effort": "none"}, optional_params={}, model="moonshotai/Kimi-K3",