diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 14d47a15c6d..b217c46a28f 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -7,7 +7,8 @@ import json import mimetypes import re from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence -from itertools import groupby, islice +from graphlib import CycleError, TopologicalSorter +from itertools import chain, groupby, islice from os import PathLike from pathlib import Path from types import MappingProxyType @@ -1371,6 +1372,123 @@ def _schema_levels(schema: Mapping[str, object]) -> Iterator[tuple[Mapping[str, frontier = tuple(child for node in frontier for child in _subschemas(node)) +def schema_fragment_ref(ref: str) -> str: + return f"#{ref}" if ref.startswith(("/$defs/", "/definitions/")) else ref + + +def schema_has_unresolvable_ref(schema: Mapping[str, object]) -> bool: + definitions: Final = MappingProxyType(dict(_root_definitions(schema))) + if not _refs_in(schema) <= definitions.keys(): + return True + try: + tuple(TopologicalSorter({pointer: _refs_in(body) for pointer, body in definitions.items()}).static_order()) + except CycleError: + return True + return False + + +def schema_union_property_count(schema: Mapping[str, object]) -> int: + return sum(1 for _, node in _property_schemas(schema) if _is_union(node)) + + +def schema_optional_property_count(schema: Mapping[str, object]) -> int: + return sum(len(_optional_property_names(node)) for node in _schema_nodes(schema)) + + +def schema_property_count(schema: Mapping[str, object]) -> int: + return sum(len(properties) for node in _schema_nodes(schema) if (properties := _schema_at(node, "properties"))) + + +def schema_grammar_branch_count(schema: Mapping[str, object]) -> int: + return sum(_grammar_branches(node) for node in _schema_nodes(schema)) + + +def schema_has_open_object(schema: Mapping[str, object]) -> bool: + return any(_is_open_object(node) for node in _schema_nodes(schema)) + + +def schema_has_lookaround_pattern(schema: Mapping[str, object]) -> bool: + return any( + isinstance(pattern := node.get("pattern"), str) and _uses_regex_lookaround(pattern) + for node in _schema_nodes(schema) + ) + + +def schema_string_formats(schema: Mapping[str, object]) -> frozenset[str]: + return frozenset( + value for node in _schema_nodes(schema) if _may_be_string(node) and isinstance(value := node.get("format"), str) + ) + + +def _may_be_string(node: Mapping[str, object]) -> bool: + node_type: Final = node.get("type") + return node_type is None or node_type == "string" or "string" in _list_at(node, "type") + + +def is_object_schema(node: Mapping[str, object]) -> bool: + node_type: Final = node.get("type") + return node_type == "object" or (isinstance(node_type, list) and "object" in node_type) + + +def _schema_nodes(schema: Mapping[str, object]) -> Iterator[Mapping[str, object]]: + return chain.from_iterable(islice(_schema_levels(schema), _MAX_SCHEMA_NESTING)) + + +_SCHEMA_LIST: Final = TypeAdapter(tuple[object, ...]) + + +def _list_at(container: Mapping[str, object], key: str) -> tuple[object, ...]: + value: Final = container.get(key) + return _SCHEMA_LIST.validate_python(value) if isinstance(value, list) else () + + +def _named_schemas(container: Mapping[str, object] | None) -> Iterator[tuple[str, Mapping[str, object]]]: + if container is None: + return iter(()) + return ((key, value) for key in container if (value := _schema_at(container, key)) is not None) + + +def _root_definitions(schema: Mapping[str, object]) -> Iterator[tuple[str, Mapping[str, object]]]: + for defs_key in ("$defs", "definitions"): + yield from ((f"#/{defs_key}/{name}", body) for name, body in _named_schemas(_schema_at(schema, defs_key))) + + +def _refs_in(schema: Mapping[str, object]) -> frozenset[str]: + return frozenset( + schema_fragment_ref(ref) for node in _schema_nodes(schema) if isinstance(ref := node.get("$ref"), str) + ) + + +def _property_schemas(schema: Mapping[str, object]) -> Iterator[tuple[str, Mapping[str, object]]]: + return chain.from_iterable(_named_schemas(_schema_at(node, "properties")) for node in _schema_nodes(schema)) + + +def _is_union(node: Mapping[str, object]) -> bool: + return isinstance(node.get("anyOf"), list) or isinstance(node.get("oneOf"), list) or len(_list_at(node, "type")) > 1 + + +def _grammar_branches(node: Mapping[str, object]) -> int: + properties: Final = _schema_at(node, "properties") + if properties is None: + return 0 + optional: Final = len(_optional_property_names(node)) + return 2**optional * (len(properties) - optional + 1) + + +def _optional_property_names(node: Mapping[str, object]) -> frozenset[str]: + properties: Final = _schema_at(node, "properties") + if properties is None: + return frozenset() + return frozenset(properties) - frozenset(name for name in _list_at(node, "required") if isinstance(name, str)) + + +def _is_open_object(node: Mapping[str, object]) -> bool: + additional_properties: Final = node.get("additionalProperties") + if additional_properties is True or isinstance(additional_properties, dict) or "patternProperties" in node: + return True + return is_object_schema(node) and "properties" not in node and additional_properties is not False + + def _subschemas(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]: for key, value in node.items(): if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict): diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 26d219a5ebc..101eca56c16 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5129,14 +5129,14 @@ def _bedrock_tools_pt(tools: list, model: str | None = None) -> list[BedrockTool description = name defs = parameters.pop("$defs", {}) - defs_copy = copy.deepcopy(defs) + defs_copy = copy.deepcopy(defs) if isinstance(defs, dict) else {} # Expand $ref references in parameters using the definitions # Note: We don't pre-flatten defs as that causes exponential memory growth # with circular references (see issue #19098). unpack_defs handles nested # refs recursively and correctly detects/skips circular references. unpack_defs(parameters, defs_copy) normalize_json_schema_custom_types_to_object(parameters) - if parameters.get("type") not in _valid_json_schema_root_types: + if not (isinstance(root_type := parameters.get("type"), str) and root_type in _valid_json_schema_root_types): parameters["type"] = "object" tool_block = cast( BedrockToolBlock, diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 3bffee48d6a..6c26a2797ce 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1,3 +1,4 @@ +import copy import json import re import time @@ -24,7 +25,18 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.prompt_templates.common_utils import ( + DEFS_MAX_INLINED_BYTES, + is_object_schema, sanitize_input_schema_for_anthropic, + schema_fragment_ref, + schema_grammar_branch_count, + schema_has_lookaround_pattern, + schema_has_open_object, + schema_has_unresolvable_ref, + schema_optional_property_count, + schema_string_formats, + schema_union_property_count, + unpack_defs, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( RemoteMedia, @@ -36,7 +48,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import split_leading_system_run, ) from litellm.llms.base_llm.base_utils import type_to_response_format_param -from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.chat.transformation import REQUEST_PARAMS, BaseConfig, BaseLLMException from litellm.types.llms.anthropic import ( ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_BETA_HEADER_VALUES, @@ -137,6 +149,12 @@ else: # response side. _ANTHROPIC_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 +_NATIVE_OUTPUT_STRING_FORMATS: Final = frozenset( + ("date-time", "time", "date", "duration", "email", "hostname", "uri", "ipv4", "ipv6", "uuid") +) +NATIVE_OUTPUT_MAX_UNION_PARAMETERS: Final = 16 +NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS: Final = 24 +NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES: Final = 2048 class _AnthropicUsageIteration(TypedDict, total=False): @@ -536,7 +554,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): grammar Anthropic compiles the schema into, so the API rejects them with a 400 ``invalid_request_error`` (e.g. "output_format.schema: For 'array' type, property 'uniqueItems' is not supported"): - - maxItems/minItems/uniqueItems/contains/minContains/maxContains/prefixItems: array constraints + - maxItems/minItems above 1/uniqueItems/contains/minContains/maxContains/prefixItems: array constraints - minimum/maximum/exclusiveMinimum/exclusiveMaximum/multipleOf: numeric constraints - minLength/maxLength: string constraints - minProperties/maxProperties/patternProperties/propertyNames: object constraints @@ -593,7 +611,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): "else": "conditional (else): {}", "not": "must not match: {}", } - unsupported_fields: Final = set(constraint_labels) + min_items: Final = schema.get("minItems") + enforced_min_items: Final = ( + isinstance(min_items, int) and not isinstance(min_items, bool) and min_items in (0, 1) + ) + unsupported_fields: Final = set(constraint_labels) - ({"minItems"} if enforced_min_items else set()) # Build description additions from removed constraints. Iterating # constraint_labels (not the set) keeps the note order deterministic across @@ -601,7 +623,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # PYTHONHASHSEED and stay cache-friendly. constraint_descriptions: Final[list] = [] for field, label in constraint_labels.items(): - if field not in schema: + if field not in schema or field not in unsupported_fields: continue value = schema[field] # A falsy boolean constraint (e.g. ``uniqueItems: false``) imposes no @@ -636,7 +658,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Already handled above continue - if key == "properties" and isinstance(value, dict): + if key in ("properties", "$defs", "definitions") and isinstance(value, dict): result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()} elif key == "items" and isinstance(value, dict): result[key] = AnthropicConfig.filter_anthropic_output_schema(value) @@ -648,6 +670,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ] elif key == "allOf" and isinstance(value, list): result[key] = [AnthropicConfig.filter_anthropic_output_schema(item) for item in value] + elif key == "$ref" and isinstance(value, str): + result[key] = schema_fragment_ref(value) elif key == "oneOf" and isinstance(value, list): # Anthropic rejects oneOf ("Schema type 'oneOf' is not supported"); # the Anthropic SDK rewrites it to anyOf, so do the same. @@ -659,7 +683,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Anthropic requires additionalProperties=false for object schemas # See: https://docs.anthropic.com/en/docs/build-with-claude/structured-outputs - if result.get("type") == "object": + if is_object_schema(result): result["additionalProperties"] = False return result @@ -1336,32 +1360,55 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return json_schema + @staticmethod + def _defs_unpacked(json_schema: Mapping[str, object]) -> Mapping[str, object]: + unpacked: Final = copy.deepcopy(dict(json_schema)) + defs: Final = unpacked.pop("$defs", unpacked.pop("definitions", None)) + if isinstance(defs, dict) and defs: + unpack_defs(unpacked, defs, max_inlined_bytes=DEFS_MAX_INLINED_BYTES) + return unpacked + + @staticmethod + def native_output_schema(json_schema: Mapping[str, object]) -> Mapping[str, object] | None: + if ( + schema_has_unresolvable_ref(json_schema) + or schema_has_open_object(json_schema) + or schema_has_lookaround_pattern(json_schema) + or not schema_string_formats(json_schema) <= _NATIVE_OUTPUT_STRING_FORMATS + or schema_union_property_count(json_schema) > NATIVE_OUTPUT_MAX_UNION_PARAMETERS + or schema_optional_property_count(json_schema) > NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS + or schema_grammar_branch_count(json_schema) > NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES + ): + return None + return AnthropicConfig.filter_anthropic_output_schema(dict(json_schema)) + + def uses_native_structured_output_for_schema( + self, + model: str, + json_schema: Mapping[str, object] | None, + request_params: Mapping[str, object], + ) -> bool: + open_object_takes_forced_json_tool: Final = ( + json_schema is not None + and schema_has_open_object(json_schema) + and not self.is_thinking_requested(request_params) + and not self.is_thinking_enabled(dict(request_params)) + and not request_params.get("tools") + and not AnthropicModelInfo.forced_tool_use_unsupported(model) + ) + return self.uses_native_structured_output(model) and not open_object_takes_forced_json_tool + + def native_output_format_schema(self, json_schema: Mapping[str, object]) -> Mapping[str, object] | None: + return self.filter_anthropic_output_schema(dict(self._defs_unpacked(json_schema))) + def map_response_format_to_anthropic_output_format(self, value: dict | None) -> AnthropicOutputSchema | None: - json_schema: dict | None = self._extract_json_schema_from_response_format(value) + json_schema: Final = self._extract_json_schema_from_response_format(value) if json_schema is None: return None - - # Resolve $ref/$defs before filtering — Anthropic doesn't support - # external schema references (e.g., /$defs/CalendarEvent). - import copy - - from litellm.litellm_core_utils.prompt_templates.common_utils import ( - DEFS_MAX_INLINED_BYTES, - unpack_defs, - ) - - json_schema = copy.deepcopy(json_schema) - defs: Final = json_schema.pop("$defs", json_schema.pop("definitions", {})) - if defs: - unpack_defs(json_schema, defs, max_inlined_bytes=DEFS_MAX_INLINED_BYTES) - - # Filter out unsupported fields for Anthropic's output_format API - filtered_schema: Final = self.filter_anthropic_output_schema(json_schema) - - return AnthropicOutputSchema( - type="json_schema", - schema=filtered_schema, - ) + native_schema: Final = self.native_output_format_schema(json_schema) + if native_schema is None: + return None + return AnthropicOutputSchema(type="json_schema", schema=dict(native_schema)) def map_response_format_to_anthropic_tool( self, value: dict | None, optional_params: dict, is_thinking_enabled: bool @@ -1526,10 +1573,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): output_key=param, ) elif param == "response_format" and isinstance(value, dict): - if AnthropicConfig._supports_model_capability( + if self.uses_native_structured_output_for_schema( model, - "supports_native_structured_output", - self._resolved_provider, + self._extract_json_schema_from_response_format(value), + REQUEST_PARAMS.validate_python(non_default_params), ): _output_format = self.map_response_format_to_anthropic_output_format(value) if _output_format is not None: @@ -2665,6 +2712,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if json_mode_message is not None: completion_response["stop_reason"] = "stop" + if thinking_blocks: + json_mode_message.thinking_blocks = thinking_blocks + json_mode_message.reasoning_content = reasoning_content _message = json_mode_message model_response.choices[0].message = _message diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 7decf1b4186..7d6109c84bc 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Final, Union import httpx -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm.constants import DEFAULT_MAX_TOKENS, RESPONSE_FORMAT_TOOL_NAME from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -38,6 +38,9 @@ else: LiteLLMLoggingObj = Any +REQUEST_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + class BaseLLMException(Exception): def __init__( self, @@ -97,6 +100,17 @@ class BaseConfig(ABC): thinking_type: Final = thinking.get("type") if isinstance(thinking, dict) else None return thinking is True or thinking_type == "enabled" or non_default_params.get("reasoning_effort") is not None + def is_thinking_requested(self, non_default_params: Mapping[str, object]) -> bool: + match non_default_params: + case {"thinking": True | {"type": "enabled" | "adaptive"}}: + return True + case {"reasoning_effort": None | "none"}: + return False + case {"reasoning_effort": object()}: + return True + case _: + return False + def is_max_tokens_in_request(self, non_default_params: dict) -> bool: """ OpenAI spec allows max_tokens or max_completion_tokens to be specified. diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 2a2f3052b2a..51bde66d660 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( _parse_content_for_reasoning, + is_object_schema, ) from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockConverseMessagesProcessor, @@ -48,7 +49,13 @@ from litellm.llms.anthropic.chat.transformation import ( AnthropicConfig, ) from litellm.llms.anthropic.common_utils import AnthropicModelInfo -from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.base_llm.chat.transformation import REQUEST_PARAMS, BaseConfig, BaseLLMException +from litellm.llms.bedrock.common_utils import ( + bedrock_chat_uses_native_structured_output, + bedrock_cleaned_output_schema, + bedrock_native_output_format, + deployment_model_info, +) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, bedrock_request_metadata_is_owned, @@ -906,17 +913,22 @@ class AmazonConverseConfig(BaseConfig): return _tool @staticmethod - def _supports_native_structured_outputs(model: str, custom_llm_provider: str | None = None) -> bool: - """Check if the Bedrock model supports native structured outputs (outputConfig.textFormat). + def _supports_native_structured_outputs( + model: str, + custom_llm_provider: str | None, + thinking_requested: bool, + model_info: Mapping[str, object] | None, + ) -> bool: + """Check if the Bedrock model takes native structured outputs (outputConfig.textFormat) for this request. - Delegates to the standard ``supports_native_structured_output`` utility - which looks up the flag in ``litellm.model_cost`` via - ``_get_model_info_helper``. + ``supports_native_structured_output`` in ``litellm.model_cost`` turns it on for every request, and + ``supports_bedrock_native_structured_output`` (the deployment's ``model_info`` first) only when the + request asks for thinking, which the json tool path would drop. Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html """ - from litellm.utils import supports_native_structured_output - - return supports_native_structured_output(model=model, custom_llm_provider=custom_llm_provider) + return bedrock_chat_uses_native_structured_output( + model, thinking_requested, custom_llm_provider or "bedrock_converse", model_info + ) @staticmethod def _add_additional_properties_to_schema(schema: dict) -> dict: @@ -931,9 +943,9 @@ class AmazonConverseConfig(BaseConfig): if not isinstance(schema, dict): return schema - result: Final = dict(schema) + result: Final[dict[str, object]] = dict(schema) - if result.get("type") == "object" and "additionalProperties" not in result: + if is_object_schema(result) and "additionalProperties" not in result: result["additionalProperties"] = False # Recurse into nested schemas @@ -944,16 +956,46 @@ class AmazonConverseConfig(BaseConfig): if "items" in result and isinstance(result["items"], dict): result["items"] = AmazonConverseConfig._add_additional_properties_to_schema(result["items"]) for defs_key in ("$defs", "definitions"): - if defs_key in result and isinstance(result[defs_key], dict): + if isinstance(definitions := result.get(defs_key), dict): result[defs_key] = { - k: AmazonConverseConfig._add_additional_properties_to_schema(v) for k, v in result[defs_key].items() + k: AmazonConverseConfig._add_additional_properties_to_schema(v) for k, v in definitions.items() } for key in ("anyOf", "allOf", "oneOf"): - if key in result and isinstance(result[key], list): - result[key] = [AmazonConverseConfig._add_additional_properties_to_schema(item) for item in result[key]] + if isinstance(variants := result.get(key), list): + result[key] = [AmazonConverseConfig._add_additional_properties_to_schema(item) for item in variants] return result + def _native_output_schema( + self, + model: str, + json_schema: Mapping[str, object], + thinking_requested: bool, + model_info: Mapping[str, object] | None, + ) -> Mapping[str, object] | None: + if not self._supports_native_structured_outputs( + model, self.custom_llm_provider, thinking_requested, model_info + ): + return None + if BedrockModelInfo.get_base_model(model).startswith("anthropic"): + return bedrock_cleaned_output_schema(model, json_schema, self.custom_llm_provider or "bedrock_converse") + return json_schema + + def _output_config_for_anthropic_format( + self, model: str, output_config_format: Mapping[str, object], model_info: Mapping[str, object] | None + ) -> OutputConfigBlock | None: + native_format: Final = bedrock_native_output_format( + model, output_config_format, self.custom_llm_provider or "bedrock_converse", model_info + ) + if native_format is None: + return None + schema, name, description = (native_format.get(key) for key in ("schema", "name", "description")) + return AmazonConverseConfig._create_output_config_for_response_format( + json_schema={**schema} if isinstance(schema, Mapping) else None, + name=name if isinstance(name, str) else None, + description=description if isinstance(description, str) else None, + ) + @staticmethod def _create_output_config_for_response_format( json_schema: dict | None = None, @@ -1013,6 +1055,7 @@ class AmazonConverseConfig(BaseConfig): optional_params: dict, model: str, drop_params: bool, + model_info: Mapping[str, object] | None = None, ) -> dict: is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) base_model: Final = BedrockModelInfo.get_base_model(model) @@ -1029,6 +1072,7 @@ class AmazonConverseConfig(BaseConfig): optional_params=optional_params, non_default_params=non_default_params, is_thinking_enabled=is_thinking_enabled, + model_info=model_info, ) if param == "max_tokens" or param == "max_completion_tokens": optional_params["maxTokens"] = ( @@ -1205,6 +1249,7 @@ class AmazonConverseConfig(BaseConfig): optional_params: dict, non_default_params: dict, is_thinking_enabled: bool, + model_info: Mapping[str, object] | None = None, ) -> dict: """ Handles translation of response_format parameter to Bedrock format. @@ -1228,13 +1273,23 @@ class AmazonConverseConfig(BaseConfig): if "type" in value and value["type"] == "text": return optional_params - if self._supports_native_structured_outputs(model, self.custom_llm_provider) and json_schema is not None: + native_schema: Final = ( + self._native_output_schema( + model, + json_schema, + self.is_thinking_requested(REQUEST_PARAMS.validate_python(non_default_params)), + model_info, + ) + if json_schema is not None + else None + ) + if native_schema is not None: # Use Bedrock's native structured outputs API (outputConfig.textFormat) # No synthetic tool injection, no fake_stream needed. # Requires an explicit schema — json_object with no schema falls through # to the tool-call path below. output_config: Final = self._create_output_config_for_response_format( - json_schema=json_schema, + json_schema=dict(native_schema), name=name, description=description, ) @@ -1566,7 +1621,11 @@ class AmazonConverseConfig(BaseConfig): tool_choice.pop("type", None) def _prepare_request_params( - self, optional_params: dict, model: str, drop_params: bool = False + self, + optional_params: dict, + model: str, + drop_params: bool = False, + model_info: Mapping[str, object] | None = None, ) -> tuple[dict, dict, dict, OutputConfigBlock | None]: """Prepare and separate request parameters.""" # Consume the internal ``_output_config_normalized`` marker set by @@ -1609,17 +1668,12 @@ class AmazonConverseConfig(BaseConfig): and output_config_format is not None and output_config_format.get("type") == "json_schema" and base_model.startswith("anthropic") - and self._supports_native_structured_outputs(model, self.custom_llm_provider) ): - output_config = self._create_output_config_for_response_format( - json_schema=output_config_format.get("schema"), - name=output_config_format.get("name"), - description=output_config_format.get("description"), - ) - elif output_config is None and output_config_format is not None: + output_config = self._output_config_for_anthropic_format(model, output_config_format, model_info) + if output_config is None and output_config_format is not None: litellm.verbose_logger.warning( "Bedrock Converse: dropping `output_config.format` for model=%s — " - "model does not advertise `supports_native_structured_output` in " + "model does not take native structured output for this schema per " "model_prices_and_context_window.json. The schema will not be " "enforced; pass `response_format` to use the synthetic tool-call " "fallback.", @@ -1898,7 +1952,7 @@ class AmazonConverseConfig(BaseConfig): additional_request_params, request_metadata, output_config, - ) = self._prepare_request_params(optional_params, model, drop_params) + ) = self._prepare_request_params(optional_params, model, drop_params, deployment_model_info(litellm_params)) original_tools: Final = inference_params.pop("tools", []) 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 8d1ff1d2bd9..b94435345d3 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py @@ -92,8 +92,11 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): optional_params: dict, model: str, drop_params: bool, + model_info: Mapping[str, object] | None = None, ) -> dict: - return AmazonConverseConfig.map_openai_params(self, non_default_params, optional_params, model, drop_params) + return AmazonConverseConfig.map_openai_params( + self, non_default_params, optional_params, model, drop_params, model_info + ) def transform_request( self, 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 a8b94fb5703..a46c67bcdec 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -18,7 +18,10 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation ) from litellm.llms.bedrock.common_utils import ( apply_bedrock_invoke_structured_output, + bedrock_chat_uses_native_structured_output, + bedrock_cleaned_output_schema, bedrock_supports_tool_search, + deployment_model_info, get_anthropic_beta_from_headers, normalize_bedrock_opus_output_config_effort, normalize_custom_field_on_tools, @@ -41,6 +44,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) + class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): """ @@ -62,6 +67,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_version: str = "bedrock-2023-05-31" + def __init__(self, model_info: Mapping[str, object] | None = None) -> None: + super().__init__() + self._deployment_model_info: Final = model_info + @property def custom_llm_provider(self) -> str | None: return "bedrock" @@ -72,6 +81,19 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): def get_supported_openai_params(self, model: str) -> list[str]: return AnthropicConfig.get_supported_openai_params(self, model) + def native_output_format_schema(self, json_schema: Mapping[str, object]) -> Mapping[str, object] | None: + return AnthropicConfig.native_output_schema(json_schema) + + def uses_native_structured_output_for_schema( + self, + model: str, + json_schema: Mapping[str, object] | None, + request_params: Mapping[str, object], + ) -> bool: + return bedrock_chat_uses_native_structured_output( + 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 map_openai_params( self, non_default_params: dict, @@ -229,6 +251,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): apply_bedrock_invoke_structured_output( model=model, request_body=anthropic_request, + model_info=deployment_model_info(_LITELLM_PARAMS.validate_python(litellm_params)), ) strip_unsupported_bedrock_invoke_output_config_keys( model=model, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5f044897b2c..853abc27a11 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -23,6 +23,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.litellm_core_utils.prompt_templates.common_utils import schema_property_count from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -234,15 +235,99 @@ def convert_bedrock_invoke_output_format_to_inline_schema( request_body["messages"] = new_messages -def _bedrock_model_supports(model: str, key: str) -> bool: +def _bedrock_model_supports(model: str, key: str, custom_llm_provider: str = "bedrock") -> bool: from litellm.utils import _supports_factory - return _supports_factory(model=model, custom_llm_provider="bedrock", key=key) + return _supports_factory(model=model, custom_llm_provider=custom_llm_provider, key=key) + + +_DEPLOYMENT_MODEL_INFO: Final = TypeAdapter(dict[str, object]) + + +def deployment_model_info(litellm_params: Mapping[str, object] | None) -> Mapping[str, object]: + model_info: Final = (litellm_params or {}).get("model_info") + return _DEPLOYMENT_MODEL_INFO.validate_python(model_info) if isinstance(model_info, Mapping) else {} + + +def _bedrock_native_output_opt_in( + model: str, custom_llm_provider: str, model_info: Mapping[str, object] | None +) -> bool: + deployment_flag: Final = (model_info or {}).get("supports_bedrock_native_structured_output") + if isinstance(deployment_flag, bool): + return deployment_flag + return _bedrock_model_supports(model, "supports_bedrock_native_structured_output", custom_llm_provider) + + +def bedrock_supports_native_structured_output( + model: str, custom_llm_provider: str = "bedrock", model_info: Mapping[str, object] | None = None +) -> bool: + return _bedrock_model_supports( + model, "supports_native_structured_output", custom_llm_provider + ) or _bedrock_native_output_opt_in(model, custom_llm_provider, model_info) + + +def bedrock_chat_uses_native_structured_output( + model: str, + thinking_requested: bool, + custom_llm_provider: str = "bedrock", + model_info: Mapping[str, object] | None = None, +) -> bool: + return _bedrock_model_supports(model, "supports_native_structured_output", custom_llm_provider) or ( + thinking_requested and _bedrock_native_output_opt_in(model, custom_llm_provider, model_info) + ) + + +BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES: Final = 64 + + +def bedrock_cleaned_output_schema( + model: str, schema: Mapping[str, object], custom_llm_provider: str = "bedrock" +) -> Mapping[str, object] | None: + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + if ( + not _bedrock_model_supports(model, "supports_native_structured_output", custom_llm_provider) + and schema_property_count(schema) > BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES + ): + return None + return AnthropicConfig.native_output_schema(schema) + + +def bedrock_native_output_schema( + model: str, + schema: Mapping[str, object], + custom_llm_provider: str = "bedrock", + model_info: Mapping[str, object] | None = None, +) -> Mapping[str, object] | None: + native_schema: Final = bedrock_cleaned_output_schema(model, schema, custom_llm_provider) + if _bedrock_model_supports(model, "supports_native_structured_output", custom_llm_provider): + return schema if native_schema is None else native_schema + if _bedrock_native_output_opt_in(model, custom_llm_provider, model_info): + return native_schema + return None + + +def bedrock_native_output_format( + model: str, + output_format: Mapping[str, object], + custom_llm_provider: str = "bedrock", + model_info: Mapping[str, object] | None = None, +) -> Mapping[str, object] | None: + schema: Final = output_format.get("schema") + if not isinstance(schema, Mapping): + return ( + output_format if bedrock_supports_native_structured_output(model, custom_llm_provider, model_info) else None + ) + native_schema: Final = bedrock_native_output_schema(model, schema, custom_llm_provider, model_info) + if native_schema is None: + return None + return {**output_format, "schema": native_schema} def apply_bedrock_invoke_structured_output( model: str, request_body: dict[str, object], # mutable-ok: edited in place like siblings + model_info: Mapping[str, object] | None = None, ) -> None: """ Route Anthropic structured-output params to what the Bedrock model supports. @@ -262,17 +347,18 @@ def apply_bedrock_invoke_structured_output( if schema_format is None: return - if _bedrock_model_supports(model, "supports_native_structured_output"): + native_format: Final = bedrock_native_output_format(model, schema_format, model_info=model_info) + if native_format is not None: existing_output_config: Final = request_body.get("output_config") if isinstance(existing_output_config, dict): - existing_output_config["format"] = schema_format + existing_output_config["format"] = native_format else: - request_body["output_config"] = {"format": schema_format} # rebind-ok: out-param # mutable-ok: json + request_body["output_config"] = {"format": native_format} # rebind-ok: out-param return verbose_logger.warning( - "Bedrock Invoke: model=%s does not advertise `supports_native_structured_output` " - "in model_prices_and_context_window.json, so the JSON schema was inlined into " + "Bedrock Invoke: model=%s does not take native structured output for this schema " + "per model_prices_and_context_window.json, so the JSON schema was inlined into " "the last user message and is NOT enforced by the model.", model, ) 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 14bd2bee6cf..5575f2808ab 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -3,6 +3,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast import httpx +from pydantic import TypeAdapter import litellm from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers @@ -64,6 +65,9 @@ else: LiteLLMLoggingObj = Any +_MODEL_INFO: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) + + class AmazonAnthropicClaudeMessagesConfig( AnthropicMessagesConfig, AmazonInvokeConfig, @@ -706,6 +710,7 @@ class AmazonAnthropicClaudeMessagesConfig( apply_bedrock_invoke_structured_output( model=model, request_body=anthropic_messages_request, + model_info=_MODEL_INFO.validate_python(getattr(litellm_params, "model_info", None)), ) normalize_bedrock_opus_output_config_effort( model=model, diff --git a/litellm/main.py b/litellm/main.py index 12854db15d0..280e717952e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4148,6 +4148,10 @@ def _complete_sagemaker(ctx: _CompletionDispatchContext) -> _CompletionDispatchR ) +_OPTIONAL_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + def _complete_bedrock(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base @@ -5603,7 +5607,11 @@ def completion( "allowed_openai_params": allowed_openai_params, "base_model": base_model, } - optional_params = get_optional_params(**optional_param_args, **non_default_params) + optional_params = get_optional_params( + **optional_param_args, + model_info=_MODEL_INFO_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else None, + **non_default_params, + ) processed_non_default_params: Final = pre_process_non_default_params( model=model, passed_params=optional_param_args, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2cce951ca60..a6581edfc5c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -188,6 +188,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_pdf_input: bool | None supports_native_streaming: bool | None supports_native_structured_output: bool | None + supports_bedrock_native_structured_output: ReadOnly[bool | None] supports_parallel_function_calling: bool | None supports_web_search: bool | None supports_reasoning: bool | None @@ -3830,6 +3831,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) +DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( + {"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]] = ( diff --git a/litellm/utils.py b/litellm/utils.py index 64ecfcac98c..8c00bdaf9f4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4469,7 +4469,8 @@ def get_optional_params( store: bool | None = None, prompt_cache_key: str | None = None, base_model: str | None = None, - **kwargs, + model_info: Mapping[str, object] | None = None, + **kwargs: object, ): drop_params = normalize_drop_params(drop_params) # rebind-ok: config and DB deployments pass "true" as a string passed_params: Final = locals().copy() @@ -4478,6 +4479,7 @@ def get_optional_params( # non_default_params / _check_valid_arg — it's a routing hint, not an # OpenAI param. passed_params.pop("base_model", None) + passed_params.pop("model_info", None) provider_config: BaseConfig | None = None if custom_llm_provider is not None and custom_llm_provider in [provider.value for provider in LlmProviders]: provider_config = ProviderConfigManager.get_provider_chat_config( @@ -4710,6 +4712,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, drop_params=bool(drop_params), + model_info=model_info, ) elif bedrock_route == "openai": optional_params = litellm.AmazonBedrockOpenAIConfig().map_openai_params( @@ -4727,7 +4730,7 @@ def get_optional_params( drop_params=bool(drop_params), ) else: - optional_params = litellm.AmazonAnthropicClaudeConfig().map_openai_params( + optional_params = litellm.AmazonAnthropicClaudeConfig(model_info=model_info).map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, @@ -6293,6 +6296,9 @@ def _get_model_info_helper( supports_embedding_image_input=_model_info.get("supports_embedding_image_input", None), supports_native_streaming=_model_info.get("supports_native_streaming", None), supports_native_structured_output=_model_info.get("supports_native_structured_output", None), + supports_bedrock_native_structured_output=_model_info.get( + "supports_bedrock_native_structured_output", None + ), supports_web_search=_model_info.get("supports_web_search", None), supports_url_context=_model_info.get("supports_url_context", None), supports_reasoning=_model_info.get("supports_reasoning", None), diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 5b2d52f11ee..5fbd45e612e 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -918,6 +918,9 @@ "supports_audio_output": { "type": "boolean" }, + "supports_bedrock_native_structured_output": { + "type": "boolean" + }, "supports_computer_use": { "type": "boolean" }, diff --git a/tests/integration/providers/test_bedrock_native_structured_output_wire.py b/tests/integration/providers/test_bedrock_native_structured_output_wire.py new file mode 100644 index 00000000000..f389e53d132 --- /dev/null +++ b/tests/integration/providers/test_bedrock_native_structured_output_wire.py @@ -0,0 +1,820 @@ +import json +import re +import uuid +from collections.abc import 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.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_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" +_DEEPSEEK: Final = "deepseek.v3.2" +_AWS: Final[dict[str, JsonValue]] = { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", +} +_NO_CACHE: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}, "num_retries": 0} +_JSON_TOOL: Final = "json_tool_call" +_USER_TOOL: Final = "lookup" +_SIGNATURE: Final = "c2NyaXB0ZWQtc2lnbmF0dXJl" +_ADAPTIVE: Final[dict[str, JsonValue]] = {"type": "adaptive"} +_THINKING_TYPES: Final = frozenset({"enabled", "adaptive"}) +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_USAGE: Final[dict[str, JsonValue]] = {"inputTokens": 11, "outputTokens": 7, "totalTokens": 18} +_MARKER: Final = re.compile(r"probe ([0-9a-f]{32})") +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_FABLE_REJECTION: Final = ( + "The model returned the following errors: output_config.format: Extra inputs are not permitted" +) +_UNSUPPORTED_KEYWORDS: Final = ( + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "multipleOf", + "maxItems", + "uniqueItems", + "minProperties", + "maxProperties", + "patternProperties", + "oneOf", +) +_SCHEMA_MAPS: Final = ("properties", "$defs", "definitions") +_SCHEMA_LISTS: Final = ("anyOf", "allOf", "oneOf") +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], + "additionalProperties": False, +} +_LIMITS_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "tier": {"type": "integer", "minimum": 1, "maximum": 5}, + "summary": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}, "minItems": 1}, + "owner": {"anyOf": [{"$ref": "#/$defs/Person"}, {"type": "null"}]}, + "backup": {"type": ["object", "null"], "properties": {"name": {"type": "string"}}, "required": ["name"]}, + }, + "required": ["tier", "summary", "tags", "owner", "backup"], + "$defs": {"Person": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}}, +} +_OPEN_COUNTS: Final[dict[str, JsonValue]] = {"type": "object", "additionalProperties": {"type": "integer"}} +_DICT_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"store": {"type": "string"}, "counts": _OPEN_COUNTS}, + "required": ["store", "counts"], + "additionalProperties": False, +} +_NULLABLE_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "backup": {"type": ["object", "null"], "properties": {"name": {"type": "string"}}, "required": ["name"]}, + }, + "required": ["summary", "backup"], + "additionalProperties": False, +} +_CLOSED_NULLABLE_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "backup": { + "type": ["object", "null"], + "properties": {"name": {"type": "string"}}, + "required": ["name"], + "additionalProperties": False, + }, + }, + "required": ["summary", "backup"], + "additionalProperties": False, +} +_UNCLEANED_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer", "minimum": 1, "maximum": 5}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], + "additionalProperties": False, +} +_NUMERIC_FORMATS_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "tier": {"type": "integer", "format": "int32"}, + "score": {"type": "number", "format": "double"}, + "summary": {"type": "string"}, + }, + "required": ["tier", "score", "summary"], + "additionalProperties": False, +} +_WIDE_FIELDS: Final = tuple(f"field_{index:02d}" for index in range(65)) +_WIDE_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {name: {"type": "string"} for name in _WIDE_FIELDS}, + "required": list(_WIDE_FIELDS), + "additionalProperties": False, +} +_USER_TOOL_SPEC: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": _USER_TOOL, + "description": "Look up a record", + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}, "required": ["q"]}, + }, +} + + +def _payload(marker_hex: str) -> dict[str, JsonValue]: + return {"tier": 2, "summary": marker_hex} + + +def _marker() -> tuple[str, str]: + marker_hex: Final = uuid.uuid4().hex + return f"probe {marker_hex}", marker_hex + + +def _response_format(schema: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "json_schema", "json_schema": {"name": "risk", "schema": dict(schema)}} + + +def _model_id(request: Request) -> str: + return unquote(request.target).split("/model/", 1)[1].rsplit("/", 1)[0] + + +def _is_invoke(request: Request) -> bool: + return unquote(request.target).rsplit("/", 1)[1].startswith("invoke") + + +def _is_stream(request: Request) -> bool: + return unquote(request.target).endswith(("/converse-stream", "/invoke-with-response-stream")) + + +def _thinking_of(body: Mapping[str, JsonValue], invoke: bool) -> JsonValue: + if invoke: + return body.get("thinking") + return object_value(body.get("additionalModelRequestFields") or {}).get("thinking") + + +def _thinking_requested(body: Mapping[str, JsonValue], invoke: bool) -> bool: + thinking: Final = _thinking_of(body, invoke) + return isinstance(thinking, dict) and thinking.get("type") in _THINKING_TYPES + + +def _native_schema(body: Mapping[str, JsonValue], invoke: bool) -> dict[str, JsonValue] | None: + if invoke: + output_config: Final = body.get("output_config") + if not isinstance(output_config, dict) or "format" not in output_config: + return None + return object_value(object_value(object_value(output_config)["format"])["schema"]) + text_format: Final = object_value(body.get("outputConfig") or {}).get("textFormat") + if text_format is None: + return None + json_schema: Final = object_value(object_value(object_value(text_format)["structure"])["jsonSchema"]) + return _JSON.validate_json(string_value(json_schema["schema"])) + + +def _tools(body: Mapping[str, JsonValue], invoke: bool) -> dict[str, dict[str, JsonValue]]: + if invoke: + return { + string_value(object_value(tool)["name"]): object_value(object_value(tool)["input_schema"]) + for tool in list_value(body.get("tools", [])) + } + specs: Final = tuple( + object_value(object_value(tool)["toolSpec"]) + for tool in list_value(object_value(body.get("toolConfig") or {"tools": []})["tools"]) + ) + return {string_value(spec["name"]): object_value(object_value(spec["inputSchema"])["json"]) for spec in specs} + + +def _is_object_type(node: Mapping[str, JsonValue]) -> bool: + node_type: Final = node.get("type") + return node_type == "object" or (isinstance(node_type, list) and "object" in node_type) + + +def _child_schemas(node: Mapping[str, JsonValue]) -> Iterable[JsonValue]: + mapped: Final = (node.get(key) for key in _SCHEMA_MAPS) + listed: Final = (node.get(key) for key in _SCHEMA_LISTS) + single: Final = (node.get(key) for key in ("items", "additionalProperties")) + return chain( + chain.from_iterable(children.values() for children in mapped if isinstance(children, dict)), + chain.from_iterable(variants for variants in listed if isinstance(variants, list)), + (child for child in single if isinstance(child, dict)), + ) + + +def _schema_nodes(schema: JsonValue) -> Iterator[dict[str, JsonValue]]: + if not isinstance(schema, dict): + return + yield schema + yield from chain.from_iterable(_schema_nodes(child) for child in _child_schemas(schema)) + + +def _node_fault(node: Mapping[str, JsonValue]) -> str | None: + node_type: Final = node.get("type") + unsupported: Final = tuple(key for key in _UNSUPPORTED_KEYWORDS if key in node) + if unsupported: + return f"For '{node_type}' type, property '{unsupported[0]}' is not supported" + min_items: Final = node.get("minItems") + if isinstance(min_items, int) and min_items > 1: + return "For 'array' type, property 'minItems' is not supported" + if _is_object_type(node) and node.get("additionalProperties") is not False: + return "For 'object' type, 'additionalProperties' must be explicitly set to false" + return None + + +def _bedrock_schema_fault(schema: JsonValue) -> str | None: + faults: Final = tuple(fault for node in _schema_nodes(schema) if (fault := _node_fault(node)) is not None) + if not faults: + return None + return f"The model returned the following errors: output_config.format.schema: {faults[0]}" + + +def _rejection(model_id: str, native: dict[str, JsonValue] | None) -> str | None: + if native is None: + return None + if model_id == _FABLE: + return _FABLE_REJECTION + if "anthropic." not in model_id: + return None + return _bedrock_schema_fault(native) + + +def _reasoning_block(marker_hex: str, invoke: bool) -> dict[str, JsonValue]: + if invoke: + return {"type": "thinking", "thinking": f"reasoning {marker_hex}", "signature": _SIGNATURE} + return {"reasoningContent": {"reasoningText": {"text": f"reasoning {marker_hex}", "signature": _SIGNATURE}}} + + +def _answer_block(name: str | None, marker_hex: str, invoke: bool) -> dict[str, JsonValue]: + tool_input: Final[dict[str, JsonValue]] = {"q": marker_hex} if name == _USER_TOOL else _payload(marker_hex) + if name is None and invoke: + return {"type": "text", "text": json.dumps(_payload(marker_hex))} + if name is None: + return {"text": json.dumps(_payload(marker_hex))} + if invoke: + return {"type": "tool_use", "id": f"toolu_{marker_hex}", "name": name, "input": tool_input} + return {"toolUse": {"toolUseId": f"tooluse_{marker_hex}", "name": name, "input": tool_input}} + + +def _converse_frames(blocks: tuple[dict[str, JsonValue], ...], stop_reason: str) -> bytes: + def frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return aws_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload).encode(), + ) + + def block_frames(index: int, block: dict[str, JsonValue]) -> bytes: + if "reasoningContent" in block: + reasoning: Final = object_value(object_value(block["reasoningContent"])["reasoningText"]) + return b"".join( + ( + frame( + "contentBlockDelta", + {"delta": {"reasoningContent": {"text": reasoning["text"]}}, "contentBlockIndex": index}, + ), + frame( + "contentBlockDelta", + { + "delta": {"reasoningContent": {"signature": reasoning["signature"]}}, + "contentBlockIndex": index, + }, + ), + frame("contentBlockStop", {"contentBlockIndex": index}), + ) + ) + if "toolUse" in block: + tool_use: Final = object_value(block["toolUse"]) + return b"".join( + ( + frame( + "contentBlockStart", + { + "start": {"toolUse": {"toolUseId": tool_use["toolUseId"], "name": tool_use["name"]}}, + "contentBlockIndex": index, + }, + ), + frame( + "contentBlockDelta", + {"delta": {"toolUse": {"input": json.dumps(tool_use["input"])}}, "contentBlockIndex": index}, + ), + frame("contentBlockStop", {"contentBlockIndex": index}), + ) + ) + return b"".join( + ( + frame("contentBlockDelta", {"delta": {"text": block["text"]}, "contentBlockIndex": index}), + frame("contentBlockStop", {"contentBlockIndex": index}), + ) + ) + + return b"".join( + ( + frame("messageStart", {"role": "assistant"}), + *(block_frames(index, block) for index, block in enumerate(blocks)), + frame("messageStop", {"stopReason": stop_reason}), + frame("metadata", {"usage": _USAGE, "metrics": {"latencyMs": 1}}), + ) + ) + + +def _called_tool(tools: Iterable[str], native: bool) -> str | None: + names: Final = frozenset(tools) + if _USER_TOOL in names: + return _USER_TOOL + if native or _JSON_TOOL not in names: + return None + return _JSON_TOOL + + +def _error_reply(message: str) -> Reply: + return Reply(status=400, body=json.dumps({"message": message}).encode()) + + +def _bedrock(request: Request) -> Reply: + body: Final = _JSON.validate_json(request.body) + invoke: Final = _is_invoke(request) + found: Final = _MARKER.search(json.dumps(body)) + assert found is not None, body + marker_hex: Final = found.group(1) + native: Final = _native_schema(body, invoke) + rejection: Final = _rejection(_model_id(request), native) + if rejection is not None: + return _error_reply(rejection) + called: Final = _called_tool(_tools(body, invoke).keys(), native is not None) + reasoning: Final = (_reasoning_block(marker_hex, invoke),) if _thinking_requested(body, invoke) else () + blocks: Final = (*reasoning, _answer_block(called, marker_hex, invoke)) + stop_reason: Final = "tool_use" if called is not None else "end_turn" + if invoke: + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{marker_hex}", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": list(blocks), + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 7}, + } + return Reply(body=json.dumps(message).encode()) + if _is_stream(request): + return Reply(body=_converse_frames(blocks, stop_reason), content_type=_EVENT_STREAM) + converse: Final[dict[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": list(blocks)}}, + "stopReason": stop_reason, + "usage": _USAGE, + "metrics": {"latencyMs": 1}, + } + return Reply(body=json.dumps(converse).encode()) + + +def _deployment( + scenario: Scenario, wire: Wire, route: str, backend: str, model_info: Mapping[str, JsonValue] | None = None +) -> str: + return scenario.model(model=f"bedrock/{route}/{backend}", api_base=wire.url, model_info=model_info, **_AWS) + + +def _chat_body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], "max_tokens": 2048, **_NO_CACHE, **extra} + + +def _messages_body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "max_tokens": 2048, "messages": [{"role": "user", "content": marker}], **_NO_CACHE, **extra} + + +def _one_call(wire: Wire, marker_hex: str) -> tuple[Request, dict[str, JsonValue]]: + (request,) = wire.drain() + body: Final = _JSON.validate_json(request.body) + assert marker_hex in json.dumps(body), body + return request, body + + +def _chat_message(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + choices: Final = list_value(_JSON.validate_json(response.content)["choices"]) + return object_value(object_value(choices[0])["message"]) + + +def _stream_message(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + deltas: Final = tuple( + object_value(object_value(list_value(_JSON.validate_json(line[6:])["choices"])[0])["delta"]) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + return { + "content": "".join(string_value(delta["content"]) for delta in deltas if delta.get("content")), + "reasoning_content": "".join( + string_value(delta["reasoning_content"]) for delta in deltas if delta.get("reasoning_content") + ), + } + + +def _content_json(message: Mapping[str, JsonValue]) -> JsonValue: + return _JSON_VALUE.validate_json(string_value(message["content"])) + + +def _responses_items(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + assert response.status_code == 200, response.text + return tuple(object_value(item) for item in list_value(_JSON.validate_json(response.content)["output"])) + + +def _responses_text(items: tuple[dict[str, JsonValue], ...]) -> JsonValue: + (message,) = tuple(item for item in items if item.get("type") == "message") + (text,) = tuple( + object_value(part) for part in list_value(message["content"]) if object_value(part).get("type") == "output_text" + ) + return _JSON_VALUE.validate_json(string_value(text["text"])) + + +def _messages_blocks(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + assert response.status_code == 200, response.text + return tuple(object_value(block) for block in list_value(_JSON.validate_json(response.content)["content"])) + + +def _messages_text(blocks: tuple[dict[str, JsonValue], ...]) -> JsonValue: + (text,) = tuple(block for block in blocks if block.get("type") == "text") + return _JSON_VALUE.validate_json(string_value(text["text"])) + + +def _refs(schema: JsonValue) -> frozenset[str]: + return frozenset(string_value(node["$ref"]) for node in _schema_nodes(schema) if "$ref" in node) + + +def _resolves(schema: Mapping[str, JsonValue], ref: str) -> bool: + prefix: Final = "#/$defs/" + return ref.startswith(prefix) and ref.removeprefix(prefix) in object_value(schema.get("$defs") or {}) + + +def _assert_native_without_tools(body: Mapping[str, JsonValue], invoke: bool, schema: Mapping[str, JsonValue]) -> None: + assert _native_schema(body, invoke) == schema, body + assert _JSON_TOOL not in _tools(body, invoke), body + + +@pytest.mark.parametrize( + ("route", "stream"), + [("converse", False), ("converse", True), ("invoke", False)], + ids=["converse", "converse_stream", "invoke"], +) +def test_thinking_request_on_a_bedrock_native_model_gets_native_output_and_keeps_reasoning( + gateway: Gateway, route: str, stream: bool +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body(model, marker, stream=stream, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA)), + ) + message: Final = _stream_message(response) if stream else _chat_message(response) + assert _content_json(message) == _payload(marker_hex), response.text + assert message["reasoning_content"] == f"reasoning {marker_hex}", response.text + request, body = _one_call(wire, marker_hex) + invoke: Final = route == "invoke" + assert _is_stream(request) is stream, request.target + _assert_native_without_tools(body, invoke, _SCHEMA) + assert _thinking_requested(body, invoke), body + + +def test_responses_reasoning_effort_on_a_bedrock_native_model_gets_native_output_and_keeps_reasoning( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/responses", + { + "model": model, + "input": marker, + "max_output_tokens": 2048, + "reasoning": {"effort": "low"}, + "text": {"format": {"type": "json_schema", "name": "risk", "schema": _SCHEMA}}, + **_NO_CACHE, + }, + ) + items: Final = _responses_items(response) + assert _responses_text(items) == _payload(marker_hex), response.text + (reasoning,) = tuple(item for item in items if item.get("type") == "reasoning") + assert f"reasoning {marker_hex}" in json.dumps(reasoning), response.text + _, body = _one_call(wire, marker_hex) + _assert_native_without_tools(body, False, _SCHEMA) + assert _thinking_requested(body, False), body + + +def test_messages_thinking_on_a_bedrock_converse_native_model_gets_native_output_and_keeps_reasoning( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, thinking=_ADAPTIVE, output_format={"type": "json_schema", "schema": _SCHEMA}), + ) + blocks: Final = _messages_blocks(response) + assert _messages_text(blocks) == _payload(marker_hex), response.text + (thinking,) = tuple(block for block in blocks if block.get("type") == "thinking") + assert thinking["thinking"] == f"reasoning {marker_hex}", response.text + _, body = _one_call(wire, marker_hex) + _assert_native_without_tools(body, False, _SCHEMA) + assert _thinking_requested(body, False), body + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_bedrock_native_model_without_thinking_keeps_the_forced_json_tool(gateway: Gateway, route: str) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _HAIKU) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_SCHEMA)) + ) + message: Final = _chat_message(response) + assert _content_json(message) == _payload(marker_hex), response.text + assert not message.get("tool_calls"), response.text + _, body = _one_call(wire, marker_hex) + invoke: Final = route == "invoke" + assert _native_schema(body, invoke) is None, body + assert _tools(body, invoke)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], body + + +@pytest.mark.parametrize( + ("backend", "model_info"), + [(_OPUS, None), (_OPT_IN, {"supports_bedrock_native_structured_output": True})], + ids=["neither_key", "deployment_opt_in"], +) +def test_thinking_request_follows_the_bedrock_native_key_the_model_resolves_to( + gateway: Gateway, backend: str, model_info: Mapping[str, JsonValue] | None +) -> 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( + "POST", + "/v1/chat/completions", + _chat_body(model, marker, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA)), + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + assert (_native_schema(body, False) == _SCHEMA) is native, body + assert (_JSON_TOOL in _tools(body, False)) is not native, body + + +def _opt_in_answer(gateway: Gateway, entry: str, model: str, marker: str) -> JsonValue: + output_format: Final[dict[str, JsonValue]] = {"type": "json_schema", "schema": _SCHEMA} + match entry: + case "converse_chat" | "invoke_chat": + chat: Final = _chat_body(model, marker, thinking=_ADAPTIVE, response_format=_response_format(_SCHEMA)) + return _content_json(_chat_message(gateway.request("POST", "/v1/chat/completions", chat))) + case "converse_responses": + responses: Final[dict[str, JsonValue]] = { + "model": model, + "input": marker, + "max_output_tokens": 2048, + "reasoning": {"effort": "low"}, + "text": {"format": {**output_format, "name": "risk"}}, + **_NO_CACHE, + } + return _responses_text(_responses_items(gateway.request("POST", "/v1/responses", responses))) + case "converse_messages": + adapted: Final = _messages_body(model, marker, thinking=_ADAPTIVE, output_format=output_format) + return _messages_text(_messages_blocks(gateway.request("POST", "/v1/messages", adapted))) + case _: + invoked: Final = _messages_body(model, marker, output_config={"format": output_format}) + return _messages_text(_messages_blocks(gateway.request("POST", "/v1/messages", invoked))) + + +def _sent_native_schema(gateway: Gateway, wire: Wire, entry: str, model: str) -> dict[str, JsonValue] | None: + marker, marker_hex = _marker() + assert _opt_in_answer(gateway, entry, model, marker) == _payload(marker_hex) + _, body = _one_call(wire, marker_hex) + invoke: Final = entry.startswith("invoke") + native: Final = _native_schema(body, invoke) + assert native is None or _JSON_TOOL not in _tools(body, invoke), body + return native + + +@pytest.mark.parametrize( + "entry", ["converse_chat", "converse_responses", "converse_messages", "invoke_chat", "invoke_messages"] +) +def test_a_deployment_native_output_opt_in_stays_on_the_deployment_that_sets_it(gateway: Gateway, entry: str) -> None: + route: Final = entry.split("_", 1)[0] + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + opted_in: Final = _deployment( + scenario, wire, route, _OPT_IN, {"supports_bedrock_native_structured_output": True} + ) + sibling: Final = _deployment(scenario, wire, route, _OPT_IN) + assert _sent_native_schema(gateway, wire, entry, sibling) is None + assert _sent_native_schema(gateway, wire, entry, opted_in) == _SCHEMA + + +def test_fable_5_json_schema_takes_the_json_tool_instead_of_the_native_output_bedrock_rejects( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _FABLE) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_SCHEMA)) + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + assert "outputConfig" not in body, body + assert _tools(body, False)[_JSON_TOOL]["properties"] == _SCHEMA["properties"], body + + +def _limits_answer(gateway: Gateway, entry: str, model: str, marker: str) -> JsonValue: + output_config: Final[dict[str, JsonValue]] = {"format": {"type": "json_schema", "schema": _LIMITS_SCHEMA}} + match entry: + case "converse_chat" | "invoke_chat": + body: Final = _chat_body(model, marker, response_format=_response_format(_LIMITS_SCHEMA)) + return _content_json(_chat_message(gateway.request("POST", "/v1/chat/completions", body))) + case "converse_output_config": + chat: Final = _chat_body(model, marker, output_config=output_config) + return _content_json(_chat_message(gateway.request("POST", "/v1/chat/completions", chat))) + case _: + messages: Final = _messages_body(model, marker, output_config=output_config) + return _messages_text(_messages_blocks(gateway.request("POST", "/v1/messages", messages))) + + +@pytest.mark.parametrize("entry", ["converse_chat", "converse_output_config", "invoke_chat", "invoke_messages"]) +def test_native_schema_reaches_bedrock_with_limits_in_descriptions_and_every_object_closed( + gateway: Gateway, entry: str +) -> None: + marker, marker_hex = _marker() + route: Final = entry.split("_", 1)[0] + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _SONNET) + assert _limits_answer(gateway, entry, model, marker) == _payload(marker_hex) + _, body = _one_call(wire, marker_hex) + schema: Final = _native_schema(body, route == "invoke") + assert schema is not None, body + assert _bedrock_schema_fault(schema) is None, schema + properties: Final = object_value(schema["properties"]) + assert object_value(properties["tags"])["minItems"] == 1, schema + assert {"1", "5"} <= set(re.findall(r"\d+", string_value(object_value(properties["tier"])["description"]))) + assert all(_resolves(schema, ref) for ref in _refs(schema)), schema + assert _JSON_TOOL not in _tools(body, route == "invoke"), body + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_open_object_schema_takes_the_json_tool_on_a_native_model_and_keeps_the_dict_open( + gateway: Gateway, route: str +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, route, _SONNET) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(_DICT_SCHEMA)) + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + invoke: Final = route == "invoke" + assert _native_schema(body, invoke) is None, body + tool_schema: Final = _tools(body, invoke)[_JSON_TOOL] + assert object_value(tool_schema["properties"])["counts"] == _OPEN_COUNTS, tool_schema + + +def test_thinking_with_a_caller_tool_on_a_bedrock_native_model_sends_only_the_callers_tool(gateway: Gateway) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body( + model, + marker, + thinking=_ADAPTIVE, + tools=[_USER_TOOL_SPEC], + response_format=_response_format(_SCHEMA), + ), + ) + (call,) = list_value(_chat_message(response)["tool_calls"]) + function: Final = object_value(object_value(call)["function"]) + assert function["name"] == _USER_TOOL, response.text + assert json.loads(string_value(function["arguments"])) == {"q": marker_hex}, response.text + _, body = _one_call(wire, marker_hex) + assert _tools(body, False).keys() == {_USER_TOOL}, body + assert _native_schema(body, False) == _SCHEMA, body + assert _thinking_requested(body, False), body + + +@pytest.mark.parametrize( + ("schema", "expected"), + [(_UNCLEANED_SCHEMA, _UNCLEANED_SCHEMA), (_NULLABLE_SCHEMA, _CLOSED_NULLABLE_SCHEMA)], + ids=["limits_kept", "nullable_object_closed"], +) +def test_non_anthropic_native_model_gets_the_callers_schema_with_objects_closed( + gateway: Gateway, schema: dict[str, JsonValue], expected: dict[str, JsonValue] +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _DEEPSEEK) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(schema)) + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + assert _native_schema(body, False) == expected, body + + +@pytest.mark.parametrize( + ("schema", "native"), + [ + (_SCHEMA, True), + (_DICT_SCHEMA, False), + ], + ids=["closed_schema_native", "open_object_inlined"], +) +def test_messages_on_bedrock_invoke_follows_the_bedrock_native_key_without_thinking( + gateway: Gateway, schema: dict[str, JsonValue], native: bool +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "invoke", _HAIKU) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, output_config={"format": {"type": "json_schema", "schema": schema}}), + ) + assert _messages_text(_messages_blocks(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + sent: Final = _native_schema(body, True) + assert (sent is not None) is native, body + prompt: Final = json.dumps(body["messages"]) + assert all((name in prompt) is not native for name in object_value(schema["properties"])), body + + +@pytest.mark.parametrize( + "schema", + [ + {"type": "object", "properties": ["tier", "summary"]}, + {"type": "object", "properties": {"tier": {"$ref": 5}}, "additionalProperties": False}, + {"type": "object", "$defs": [1, 2], "properties": {"tier": {"$ref": "#/$defs/Tier"}}}, + {"type": ["object", "null"]}, + {"type": "object", "properties": {"tier": {"type": "integer"}}, "required": "tier"}, + ], + ids=["properties_list", "ref_not_string", "defs_list", "type_list_without_properties", "required_string"], +) +def test_malformed_schema_on_a_native_model_reaches_bedrock_once_without_a_proxy_error( + gateway: Gateway, schema: dict[str, JsonValue] +) -> None: + marker, marker_hex = _marker() + with wire_server(_bedrock) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, "converse", _SONNET) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_response_format(schema)) + ) + assert response.status_code < 500, response.text + _, body = _one_call(wire, marker_hex) + native: Final = _native_schema(body, False) + fault: Final = None if native is None else _bedrock_schema_fault(native) + assert response.status_code == (200 if fault is None else 400), response.text + assert fault is None or fault.split(": ", 1)[1] in response.text, response.text + assert (native is None) is (_JSON_TOOL in _tools(body, False)), body + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize( + ("route", "backend", "thinking", "schema", "native"), + [ + ("converse", _SONNET, False, _NUMERIC_FORMATS_SCHEMA, True), + ("invoke", _HAIKU, True, _NUMERIC_FORMATS_SCHEMA, True), + ("converse", _SONNET, False, _WIDE_SCHEMA, True), + ("converse", _HAIKU, True, _WIDE_SCHEMA, False), + ], + ids=["numeric_formats_old_key", "numeric_formats_new_key_invoke", "wide_old_key", "wide_new_key"], +) +def test_numeric_formats_and_wide_schemas_follow_the_bedrock_native_key( + gateway: Gateway, route: str, backend: str, thinking: bool, schema: dict[str, JsonValue], native: bool +) -> None: + marker, marker_hex = _marker() + extra: Final[dict[str, JsonValue]] = {"thinking": _ADAPTIVE} if thinking else {} + with wire_server(_bedrock) 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, response_format=_response_format(schema), **extra), + ) + assert _content_json(_chat_message(response)) == _payload(marker_hex), response.text + _, body = _one_call(wire, marker_hex) + invoke: Final = route == "invoke" + assert (_native_schema(body, invoke) == schema) is native, body + tools: Final = _tools(body, invoke) + assert (_JSON_TOOL in tools) is not native, body + assert native or tools[_JSON_TOOL]["properties"] == schema["properties"], 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 new file mode 100644 index 00000000000..b2aba7fb100 --- /dev/null +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -0,0 +1,2161 @@ +import asyncio +import base64 +import functools +import json +import re +import signal +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import anthropic +import httpx +import openai +import psutil +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.upstream import aws_event_stream_frame +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 + +_JSON_TOOL: Final = "json_tool_call" +_USER_TOOL: Final = "lookup" +_STRUCTURED_OUTPUTS_BETA: Final = "structured-outputs-2025-11-13" +_TOOL_MODEL: Final = "claude-sonnet-5" +_EFFORT_GATED_MODEL: Final = "claude-haiku-4-5" +_NATIVE_MODEL: Final = "claude-opus-5-5" +_FORCED_TOOL_NATIVE_MODEL: Final = "claude-haiku-5-5" +_UNMAPPED_MODEL: Final = "claude-scripted-9" +_BEDROCK_INVOKE_MODEL: Final = "anthropic.claude-sonnet-5" +_BEDROCK_JSON_TOOL_MODEL: Final = "us.anthropic.claude-opus-4-8" +_THINKING_SIGNATURE: Final = "c2NyaXB0ZWQtc2lnbmF0dXJl" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-east5" +_VERTEX_MODELS: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/" +_VERTEX_STREAM_SUFFIX: Final = ":streamRawPredict?alt=sse" +_ANTHROPIC_KEY: Final = "synthetic-anthropic-key" +_AZURE_KEY: Final = "synthetic-azure-ai-key" +_BEDROCK_TOKEN: Final = "synthetic-bedrock-bearer" +_PLATFORM_KEY: Final = "synthetic-platform-key" +_MANTLE_KEY: Final = "synthetic-bedrock-mantle-key" +_SCRIPTED_KEY: Final = "scripted" +_NIM_GROUP: Final = "scripted-nim" +_WATSONX_VERSION: Final = "2024-03-13" +_OWNED_MODEL: Final = "partner-claude" +_DEPRECATED_OUTPUT_FORMAT: Final = "output_format: this field is deprecated, use output_config.format" +_BAD_OUTPUT_CONFIG: Final = "output_config: Input should be a valid dictionary" +_BAD_FORMAT: Final = "output_config.format: Input should be a json_schema object" +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}}, + "required": ["tier", "summary"], + "additionalProperties": False, +} +_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "risk", "schema": _SCHEMA}, +} +_TEXT_FORMAT: Final[dict[str, JsonValue]] = { + "format": {"type": "json_schema", "name": "risk", "schema": _SCHEMA}, +} +_OUTPUT_FORMAT: Final[dict[str, JsonValue]] = {"type": "json_schema", "schema": _SCHEMA} +_NAME_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"name": {"type": "string"}, "tier": {"type": "integer"}}, + "required": ["name", "tier"], + "additionalProperties": False, +} +_NAME_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "named", "schema": _NAME_SCHEMA}, +} +_OPTIONAL_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": _SCHEMA["properties"]} +_VALUES_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"values": {"type": "object"}}, + "required": ["values"], +} +_VALUES_RESPONSE_FORMAT: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, +} +_ADDITIONAL_OBJECT_SCHEMA: Final[dict[str, JsonValue]] = { + **_OPTIONAL_SCHEMA, + "additionalProperties": {"type": "object"}, +} +_PATTERN_SCHEMA: Final[dict[str, JsonValue]] = { + **_OPTIONAL_SCHEMA, + "patternProperties": {"^[a-z_]+$": {"type": "object"}}, +} +_ADDITIONAL_TRUE_SCHEMA: Final[dict[str, JsonValue]] = {**_OPTIONAL_SCHEMA, "additionalProperties": True} +_CLOSED_OPTIONAL_SCHEMA: Final[dict[str, JsonValue]] = {**_OPTIONAL_SCHEMA, "additionalProperties": False} +_CALLERS_FORMAT: Final[dict[str, JsonValue]] = {"type": "json_schema", "schema": {**_SCHEMA, "title": "callers"}} +_ANTHROPIC_JSON_TOOL: Final[dict[str, JsonValue]] = {"name": _JSON_TOOL, "input_schema": _SCHEMA} +_FORCED_JSON_TOOL: Final[dict[str, JsonValue]] = {"type": "tool", "name": _JSON_TOOL} +_REASONING_FIELDS: Final = ("thinking", "output_config", "max_tokens") +_USER_TOOL_SPEC: Final[dict[str, JsonValue]] = { + "type": "function", + "function": { + "name": _USER_TOOL, + "parameters": {"type": "object", "properties": {"q": {"type": "string"}}, "required": ["q"]}, + }, +} +_MARKER: Final = re.compile(r"\b([a-z_]+) ([0-9a-f]{32})\b") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 +_BURST_SHAPES: Final = ("wrapped", "plain", "named", "values") +_STREAM_SHAPES: Final = ( + "plain", + "wrapped", + "named", + "values", + "large", + "two_candidates", + "inner_incomplete", + "self_shaped", + "scalar", + "empty", + "list_input", + "deep", + "other_tool", + "named_partial", + "parameter_label", + "function_label", + "label_unshared", + "label_two_objects", + "label_encoded", + "named_encoded", + "undeclared_encoded", + "encoded_deep", + "number_first", + "bool_first", + "null_first", +) +_HELD_STREAM_SHAPES: Final = frozenset( + { + "wrapped", + "named", + "values", + "large", + "two_candidates", + "inner_incomplete", + "deep", + "other_tool", + "named_partial", + "parameter_label", + "function_label", + "label_unshared", + "label_two_objects", + "label_encoded", + "named_encoded", + "undeclared_encoded", + "encoded_deep", + } +) +_AS_WRITTEN_SHAPES: Final = ("label_non_ascii", "label_two_objects_non_ascii") +_TEXT_ANSWER_SHAPES: Final = frozenset({"undeclared_wrapper"}) +_LABEL: Final = "RiskAssessment" +_NON_ASCII_LABEL: Final = "Évaluation des risques 风险" +_ENCODED_DEPTH: Final = 200_000 +_FAMILY: Final = ( + "anthropic_unmapped", + "azure_ai_unmapped", + "bedrock_invoke", + "bedrock_mantle", + "claude_platform_unmapped", + "vertex_haiku", +) +_FRAGMENT_SIZE: Final = 20 +_TIMING_BYTES: Final = 512 * 1024 +_TIMING_FRAGMENT: Final = 4096 +_TIMING_SMALL_BYTES: Final = 128 * 1024 +_TIMING_GROWTH_LIMIT: Final = 6.0 +_HEALTH_LATENCY_SECONDS: Final = 1.0 +_SSE: Final = "text/event-stream" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_RELAY_CHUNKS: Final = ( + b'data: {"choices":[{"delta":{"content":"relay one"}}]}\n\n', + b'data: {"choices":[{"delta":{"content":"relay two"}}]}\n\n', + b"data: [DONE]\n\n", +) +_PASS_THROUGH_ROUTES: Final = ( + "gemini", + "cohere", + "vllm", + "mistral", + "typesafe", + "openrouter", + "azure", + "openai", + "gigachat", + "watsonx", + "tinyfish", + "nvidia_nim", +) + + +@dataclass(frozen=True, slots=True) +class _Delivery: + target: str + beta: str | None + body: dict[str, JsonValue] + + +def _payload(marker_hex: str) -> dict[str, JsonValue]: + return {"tier": 2, "summary": marker_hex} + + +def _nest(payload: JsonValue, depth: int) -> JsonValue: + return functools.reduce(lambda inner, _: {"parameters": inner}, range(depth), payload) + + +def _tool_input(shape: str, marker_hex: str) -> JsonValue: + payload: Final = _payload(marker_hex) + match shape: + case "plain": + return payload + case "wrapped": + return {"tool_name": _JSON_TOOL, "parameters": payload} + case "named": + return {"name": _JSON_TOOL, "parameters": payload} + case "values": + return {"values": payload} + case "large": + return {"tool_name": _JSON_TOOL, "parameters": {**payload, "summary": marker_hex + "x" * 5000}} + case "two_candidates": + return {"parameters": payload, "alternate": {**payload, "tier": 3}} + case "inner_incomplete": + return {"parameters": {"tier": 2}} + case "self_shaped": + return {**payload, "nested": payload} + case "scalar": + return marker_hex + case "empty": + return {} + case "list_input": + return [payload] + case "deep": + return _nest(payload, 60) + case "other_tool": + return {"tool_name": "other_tool", "parameters": payload} + case "named_input": + return {"name": _JSON_TOOL, "input": payload} + case "named_two_objects": + return {"tool_name": _JSON_TOOL, "parameters": payload, "extra": {"note": marker_hex}} + case "named_string": + return {"tool_name": _JSON_TOOL, "parameters": marker_hex} + case "named_partial": + return {"tool_name": _JSON_TOOL, "parameters": {"tier": 2}} + case "name_property": + return {"name": _JSON_TOOL, "parameters": {"name": marker_hex, "tier": 1}} + case "name_value_plain": + return {"name": marker_hex, "tier": 1} + case "parameter_label": + return {"parameter_name": _LABEL, "parameter_value": payload} + case "function_label": + return {"$FUNCTION_NAME": _JSON_TOOL, "PARAMETER_VALUE": payload} + case "label_unshared": + return {"parameter_name": _LABEL, "parameter_value": {"verdict": marker_hex}} + case "label_two_objects": + return {"parameter_name": _LABEL, "parameter_value": payload, "alternate": {"tier": 3}} + case "label_encoded": + return {"parameter_name": _LABEL, "parameter_value": json.dumps(payload)} + case "label_then_declared": + return {"parameter_name": _LABEL, **payload} + case "named_encoded": + return {"name": _JSON_TOOL, "parameters": json.dumps(payload)} + case "undeclared_encoded": + return {"parameters": json.dumps(payload)} + case "encoded_deep": + return {"tool_name": _JSON_TOOL, "parameters": "[" * _ENCODED_DEPTH + "]" * _ENCODED_DEPTH} + case "number_first": + return {"rank": 3, "parameter_value": payload} + case "bool_first": + return {"rank": True, "parameter_value": payload} + case "null_first": + return {"rank": None, "parameter_value": payload} + case "undeclared_wrapper": + return {"parameters": payload} + case "label_non_ascii": + return {"parameter_name": _NON_ASCII_LABEL, "parameter_value": {"verdict": f"résumé ✓ {marker_hex}"}} + case "label_two_objects_non_ascii": + return {"parameter_name": _NON_ASCII_LABEL, "parameter_value": payload, "alternate": {"note": "风险 ✓"}} + raise AssertionError(shape) + + +def _tool_text(shape: str, marker_hex: str) -> str: + tool_input: Final = _tool_input(shape, marker_hex) + if shape in _AS_WRITTEN_SHAPES: + return json.dumps(tool_input, ensure_ascii=False, separators=(",", ":")) + return json.dumps(tool_input) + + +def _structured_text(shape: str, marker_hex: str) -> str: + return _tool_text(shape, marker_hex) if shape in _TEXT_ANSWER_SHAPES else json.dumps(_payload(marker_hex)) + + +def _timing_digits(size: int) -> str: + return "".join(f"{index:08d}" for index in range(size // 8)) + + +def _timing_text(shape: str) -> str | None: + match shape: + case "unclosed_key": + return '{"' + _timing_digits(_TIMING_BYTES) + case "unclosed_key_small": + return '{"' + _timing_digits(_TIMING_SMALL_BYTES) + case "unclosed_name": + return '{"name": "' + _timing_digits(_TIMING_BYTES) + case "unclosed_name_small": + return '{"name": "' + _timing_digits(_TIMING_SMALL_BYTES) + case "unclosed_label": + return '{"parameter_name": "' + _timing_digits(_TIMING_BYTES) + case "unclosed_label_small": + return '{"parameter_name": "' + _timing_digits(_TIMING_SMALL_BYTES) + case "unclosed_label_key": + return f'{{"parameter_name": "{_LABEL}", "' + _timing_digits(_TIMING_BYTES) + case "unclosed_label_key_small": + return f'{{"parameter_name": "{_LABEL}", "' + _timing_digits(_TIMING_SMALL_BYTES) + return None + + +def _expected(shape: str, marker_hex: str) -> JsonValue: + match shape: + case ( + "plain" + | "wrapped" + | "named" + | "values" + | "named_input" + | "deep" + | "other_tool" + | "parameter_label" + | "function_label" + | "named_encoded" + ): + return _payload(marker_hex) + case "large": + return {**_payload(marker_hex), "summary": marker_hex + "x" * 5000} + case "named_partial" | "inner_incomplete": + return {"tier": 2} + case "name_property": + return {"name": marker_hex, "tier": 1} + return _tool_input(shape, marker_hex) + + +def _fragments(text: str, size: int) -> tuple[str, ...]: + return tuple(text[start : start + size] for start in range(0, len(text), size)) + + +def _fragment_size(shape: str) -> int: + return _TIMING_FRAGMENT if shape == "encoded_deep" or _timing_text(shape) is not None else _FRAGMENT_SIZE + + +def _fragment_count(shape: str, marker_hex: str) -> int: + return len(_fragments(_tool_text(shape, marker_hex), _fragment_size(shape))) + + +def _message(message_id: str, content: list[JsonValue], stop_reason: str) -> dict[str, JsonValue]: + return { + "id": message_id, + "type": "message", + "role": "assistant", + "model": _TOOL_MODEL, + "content": content, + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + } + + +def _events( + message: dict[str, JsonValue], block: dict[str, JsonValue], deltas: tuple[dict[str, JsonValue], ...] +) -> tuple[tuple[str, dict[str, JsonValue]], ...]: + return ( + ( + "message_start", + { + "type": "message_start", + "message": { + **message, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + }, + }, + ), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": block}), + *(("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": delta}) for delta in deltas), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": message["stop_reason"], "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + ), + ("message_stop", {"type": "message_stop"}), + ) + + +def _aws_chunk(data: Mapping[str, JsonValue]) -> bytes: + payload: Final = json.dumps({"bytes": base64.b64encode(json.dumps(data).encode()).decode()}).encode() + return aws_event_stream_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, payload + ) + + +def _stream_reply( + message: dict[str, JsonValue], block: dict[str, JsonValue], deltas: tuple[dict[str, JsonValue], ...], framing: str +) -> Reply: + events: Final = _events(message, block, deltas) + if framing == "aws": + return Reply(content_type=_EVENT_STREAM, chunks=tuple(_aws_chunk(data) for _, data in events)) + return Reply( + content_type=_SSE, + chunks=tuple(f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() for name, data in events), + ) + + +def _error(status: int, message: str) -> Reply: + body: Final[dict[str, JsonValue]] = { + "type": "error", + "error": {"type": "invalid_request_error", "message": message}, + } + return Reply(status=status, body=json.dumps(body).encode()) + + +def _format_fault(body: Mapping[str, JsonValue], vertex: bool) -> str | None: + if vertex and "output_format" in body: + return _DEPRECATED_OUTPUT_FORMAT + output_config: Final = body.get("output_config") + if output_config is None: + return None + if not isinstance(output_config, dict): + return _BAD_OUTPUT_CONFIG + output_format: Final = output_config.get("format") + if output_format is None: + return None + if not isinstance(output_format, dict) or output_format.get("type") != "json_schema": + return _BAD_FORMAT + return None if isinstance(output_format.get("schema"), dict) else _BAD_FORMAT + + +def _marker_of(body: Mapping[str, JsonValue]) -> tuple[str, str]: + found: Final = _MARKER.search(json.dumps(body.get("messages"))) + assert found is not None, body + return found.group(1), found.group(2) + + +def _tool_names(body: Mapping[str, JsonValue]) -> frozenset[str]: + return frozenset(string_value(object_value(tool)["name"]) for tool in list_value(body.get("tools", []))) + + +def _tool_use(block_id: str, name: str, tool_input: JsonValue) -> dict[str, JsonValue]: + return {"type": "tool_use", "id": block_id, "name": name, "input": tool_input} + + +def _tool_reply(message_id: str, shape: str, marker_hex: str, user_tool: bool, framing: str | None) -> Reply: + timing: Final = _timing_text(shape) + tool_input: Final = {} if timing is not None else _tool_input(shape, marker_hex) + json_block: Final = _tool_use(f"toolu_{marker_hex}", _JSON_TOOL, tool_input) + user_blocks: Final[list[JsonValue]] = ( + [_tool_use(f"toolu_user_{marker_hex}", _USER_TOOL, {"q": marker_hex})] if user_tool else [] + ) + message: Final = _message(message_id, [*user_blocks, json_block], "tool_use") + if framing is None: + return Reply(body=json.dumps(message).encode()) + partial: Final = _tool_text(shape, marker_hex) if timing is None else timing + deltas: Final = tuple( + {"type": "input_json_delta", "partial_json": piece} for piece in _fragments(partial, _fragment_size(shape)) + ) + return _stream_reply(message, {**json_block, "input": {}}, deltas, framing) + + +def _text_reply(message_id: str, text: str, framing: str | None) -> Reply: + message: Final = _message(message_id, [{"type": "text", "text": text}], "end_turn") + if framing is None: + return Reply(body=json.dumps(message).encode()) + return _stream_reply(message, {"type": "text", "text": ""}, ({"type": "text_delta", "text": text},), framing) + + +def _scripted_target(request: Request) -> bool: + return ( + request.target.startswith(_VERTEX_MODELS) + or request.target.endswith("/v1/messages") + or request.target.endswith("/invoke") + or request.target.endswith("/invoke-with-response-stream") + ) + + +def _framing(request: Request, body: Mapping[str, JsonValue]) -> str | None: + if request.target.endswith("/invoke-with-response-stream"): + return "aws" + if request.target.endswith(_VERTEX_STREAM_SUFFIX) or body.get("stream") is True: + return "sse" + return None + + +def _structured(body: Mapping[str, JsonValue]) -> bool: + output_config: Final = body.get("output_config") + return "output_format" in body or (isinstance(output_config, dict) and "format" in output_config) + + +def _peer(request: Request) -> Reply: + vertex: Final = request.target.startswith(_VERTEX_MODELS) + if vertex: + assert request.headers["authorization"] == "Bearer scripted-token", request.headers + elif not _scripted_target(request): + return _error(404, f"unscripted target {request.target}") + body: Final = _JSON_OBJECT.validate_json(request.body) + fault: Final = _format_fault(body, vertex) + if fault is not None: + return _error(400, fault) + shape, marker_hex = _marker_of(body) + framing: Final = _framing(request, body) + tools: Final = _tool_names(body) + message_id: Final = f"msg_{marker_hex}" + if _JSON_TOOL in tools: + return _tool_reply(message_id, shape, marker_hex, _USER_TOOL in tools, framing) + if _structured(body): + return _text_reply(message_id, _structured_text(shape, marker_hex), framing) + return _text_reply(message_id, f"plain text {marker_hex}", framing) + + +def _relay_peer(request: Request) -> Reply: + assert request.body, request.target + return Reply(content_type=_SSE, chunks=_RELAY_CHUNKS) + + +def _failing(status: int) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + _, marker_hex = _marker_of(_JSON_OBJECT.validate_json(request.body)) + return _error(status, f"scripted {status} {marker_hex}") + + return respond + + +def _holding(held: SimpleQueue[str], release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + held.put(request.target) + assert release.wait(timeout=60), "Held request was never released" + return _peer(request) + + return respond + + +def _delivered(wire: Wire) -> tuple[_Delivery, ...]: + return tuple( + _Delivery(request.target, request.headers.get("anthropic-beta"), _JSON_OBJECT.validate_json(request.body)) + for request in wire.drain() + ) + + +def _one_delivery(wire: Wire, marker_hex: str) -> _Delivery: + (sent,) = _delivered(wire) + assert marker_hex in json.dumps(sent.body), sent.body + return sent + + +def _json_tool_schema(delivery: _Delivery) -> dict[str, JsonValue]: + tools: Final = tuple(object_value(tool) for tool in list_value(delivery.body["tools"])) + (json_tool,) = tuple(tool for tool in tools if tool["name"] == _JSON_TOOL) + assert json_tool.keys() == {"name", "input_schema"}, json_tool + return object_value(json_tool["input_schema"]) + + +def _folded_format(delivery: _Delivery) -> dict[str, JsonValue]: + assert "output_format" not in delivery.body, delivery.body + return object_value(object_value(delivery.body["output_config"])["format"]) + + +def _marker(shape: str) -> tuple[str, str]: + marker_hex: Final = uuid.uuid4().hex + return f"{shape} {marker_hex}", marker_hex + + +def _chat_body( + model: str, + marker: str, + *, + stream: bool = False, + response_format: JsonValue = _RESPONSE_FORMAT, + tools: JsonValue = None, +) -> dict[str, JsonValue]: + body: Final[dict[str, JsonValue]] = { + "model": model, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + "response_format": response_format, + } + return body if tools is None else {**body, "tools": tools} + + +def _effort_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], "reasoning_effort": "low"} + + +def _reasoning_fields(delivery: _Delivery) -> dict[str, JsonValue]: + return {field: delivery.body.get(field) for field in _REASONING_FIELDS} + + +def _messages_body(model: str, marker: str, *, stream: bool = False, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + **extra, + } + + +def _responses_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "input": marker, "text": _TEXT_FORMAT} + + +def _chat_message(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + choice: Final = object_value(list_value(_JSON_OBJECT.validate_json(response.content)["choices"])[0]) + assert choice["finish_reason"] == "stop", response.text + return object_value(choice["message"]) + + +def _chat_content(response: httpx.Response) -> JsonValue: + return json.loads(string_value(_chat_message(response)["content"])) + + +def _response_id(response: httpx.Response) -> str: + return string_value(_JSON_OBJECT.validate_json(response.content)["id"]) + + +def _delta_content(frame: str) -> str: + choices: Final = list_value(_JSON_OBJECT.validate_json(frame).get("choices", [])) + if not choices: + return "" + content: Final = object_value(object_value(choices[0])["delta"]).get("content") + return "" if content is None else string_value(content) + + +def _content_frames(response: httpx.Response) -> tuple[str, ...]: + assert response.status_code == 200, response.text + assert response.text.rstrip().endswith("data: [DONE]"), response.text + contents: Final = tuple( + _delta_content(line.removeprefix("data: ")) for line in response.text.splitlines() if line.startswith("data: {") + ) + return tuple(content for content in contents if content) + + +def _timing_frames(response: httpx.Response) -> tuple[str, int]: + assert response.status_code == 200, response.text + frames: Final = tuple( + line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: {") + ) + error_frames: Final = sum(1 for frame in frames if "error" in _JSON_OBJECT.validate_json(frame)) + return "".join(_delta_content(frame) for frame in frames), error_frames + + +def _timed_stream( + client: httpx.Client, + key: str, + wire: Wire, + model: str, + shape: str, + response_format: dict[str, JsonValue], + record_property: Callable[[str, object], None], +) -> float: + marker, marker_hex = _marker(shape) + raw: Final = _timing_text(shape) + assert raw is not None + started: Final = time.perf_counter() + response: Final = _post( + client, key, "/v1/chat/completions", _chat_body(model, marker, stream=True, response_format=response_format) + ) + wall: Final = time.perf_counter() - started + text, error_frames = _timing_frames(response) + record_property(f"{shape}_seconds", round(wall, 3)) + record_property(f"{shape}_error_frames", error_frames) + assert text == raw, (shape, len(text), len(raw), text[:80]) + _one_delivery(wire, marker_hex) + return wall + + +def _stream_text(response: httpx.Response) -> str: + return "".join(_content_frames(response)) + + +def _stream_content(response: httpx.Response) -> JsonValue: + return json.loads(_stream_text(response)) + + +def _responses_stream_text(response: httpx.Response) -> JsonValue: + assert response.status_code == 200, response.text + frames: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + assert any(frame["type"] == "response.completed" for frame in frames), response.text + return json.loads( + "".join(string_value(frame["delta"]) for frame in frames if frame["type"] == "response.output_text.delta") + ) + + +def _responses_text(response: httpx.Response) -> JsonValue: + assert response.status_code == 200, response.text + items: Final = tuple( + object_value(item) for item in list_value(_JSON_OBJECT.validate_json(response.content)["output"]) + ) + (message,) = tuple(item for item in items if item["type"] == "message") + return json.loads(string_value(object_value(list_value(message["content"])[0])["text"])) + + +def _text_block_json(message: Mapping[str, JsonValue]) -> JsonValue: + block: Final = object_value(list_value(message["content"])[0]) + return json.loads(string_value(block["text"])) + + +def _messages_text(response: httpx.Response) -> JsonValue: + assert response.status_code == 200, response.text + return _text_block_json(_JSON_OBJECT.validate_json(response.content)) + + +def _messages_stream_text(response: httpx.Response) -> JsonValue: + assert response.status_code == 200, response.text + assert "event: message_stop" in response.text, response.text + frames: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: {") + ) + deltas: Final = tuple(object_value(frame["delta"]) for frame in frames if frame["type"] == "content_block_delta") + return json.loads("".join(string_value(delta["text"]) for delta in deltas if delta["type"] == "text_delta")) + + +def _vertex_deployment(gateway: Gateway, scenario: Scenario, api_base: str, backend: str) -> str: + return scenario.model( + model=f"vertex_ai/{backend}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + +def _bedrock_invoke_deployment(scenario: Scenario, api_base: str, backend: str) -> str: + return scenario.model( + model=f"bedrock/invoke/{backend}", + api_base=api_base, + api_key=_BEDROCK_TOKEN, + aws_region_name="us-east-1", + ) + + +def _rejects_forced_tools_without_native_output(row: Mapping[str, JsonValue]) -> bool: + return row.get("supports_forced_tool_use") is False and row.get("supports_native_structured_output") is not True + + +def _claude_rejecting_forced_tools(gateway: Gateway, namespace: str, family: str) -> str: + rows: Final = gateway.get("/public/litellm_model_cost_map") + names: Final = sorted( + name.removeprefix(namespace) + for name, row in rows.items() + if name.startswith(namespace + family) and _rejects_forced_tools_without_native_output(object_value(row)) + ) + assert names, f"The cost map has no {namespace}{family} row rejecting forced tools without native output" + return names[0] + + +def _family_deployment(gateway: Gateway, scenario: Scenario, api_base: str, member: str) -> str: + match member: + case "anthropic_unmapped": + return scenario.model(model=f"anthropic/{_UNMAPPED_MODEL}", api_base=api_base, api_key=_ANTHROPIC_KEY) + case "azure_ai_unmapped": + return scenario.model(model=f"azure_ai/{_UNMAPPED_MODEL}", api_base=api_base, api_key=_AZURE_KEY) + case "bedrock_invoke": + return _bedrock_invoke_deployment(scenario, api_base, _BEDROCK_INVOKE_MODEL) + case "bedrock_mantle": + return scenario.model( + model=f"bedrock_mantle/{_BEDROCK_INVOKE_MODEL}", + api_base=api_base, + api_key=_MANTLE_KEY, + aws_region_name="us-east-1", + ) + case "claude_platform_unmapped": + return scenario.model( + model=f"bedrock/claude_platform/{_UNMAPPED_MODEL}", + api_base=api_base, + api_key=_PLATFORM_KEY, + aws_region_name="us-east-1", + workspace_id="scripted-workspace", + ) + case "vertex_haiku": + return _vertex_deployment(gateway, scenario, api_base, _EFFORT_GATED_MODEL) + raise AssertionError(member) + + +def _owned_config(path: Path, gateway: Gateway, api_base: str) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + { + "model_name": _OWNED_MODEL, + "litellm_params": { + "model": f"vertex_ai/{_TOOL_MODEL}", + "api_base": api_base, + "vertex_project": _PROJECT, + "vertex_location": _LOCATION, + "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + }, + }, + { + "model_name": _NIM_GROUP, + "litellm_params": { + "model": f"nvidia_nim/{_NIM_GROUP}", + "api_base": api_base, + "api_key": _SCRIPTED_KEY, + }, + }, + ], + } + ) + ) + return path + + +def _pass_through_environment(api_base: str) -> dict[str, str]: + bases: Final = ( + "GEMINI_API_BASE", + "COHERE_API_BASE", + "VLLM_API_BASE", + "MISTRAL_API_BASE", + "TYPESAFE_API_BASE", + "OPENROUTER_API_BASE", + "AZURE_API_BASE", + "OPENAI_API_BASE", + "GIGACHAT_API_BASE", + "WATSONX_API_BASE", + "TINYFISH_AGENT_API_BASE", + ) + keys: Final = ( + "GEMINI_API_KEY", + "COHERE_API_KEY", + "VLLM_API_KEY", + "MISTRAL_API_KEY", + "TYPESAFE_API_KEY", + "OPENROUTER_API_KEY", + "AZURE_API_KEY", + "OPENAI_API_KEY", + "GIGACHAT_ACCESS_TOKEN", + "WATSONX_TOKEN", + "TINYFISH_API_KEY", + ) + return {**{name: api_base for name in bases}, **{name: _SCRIPTED_KEY for name in keys}} + + +@dataclass(frozen=True, slots=True) +class _PassThroughCall: + path: str + body: dict[str, JsonValue] + upstream_target: str + headers: dict[str, str] + + +def _pass_through_call(route: str, marker: str, key: str) -> _PassThroughCall: + chat: Final[dict[str, JsonValue]] = { + "model": "scripted-chat", + "messages": [{"role": "user", "content": marker}], + "stream": True, + } + bearer: Final = {"Authorization": f"Bearer {key}"} + match route: + case "gemini": + return _PassThroughCall( + "/gemini/v1beta/models/scripted-model:streamGenerateContent?alt=sse", + {"contents": [{"parts": [{"text": marker}]}]}, + "/v1beta/models/scripted-model:streamGenerateContent?alt=sse&key=scripted", + {"x-goog-api-key": key}, + ) + case "cohere": + return _PassThroughCall("/cohere/v2/chat", chat, "/v2/chat", bearer) + case "vllm": + return _PassThroughCall("/vllm/v1/chat/completions", chat, "/v1/chat/completions", bearer) + case "mistral": + return _PassThroughCall("/mistral/v1/chat/completions", chat, "/v1/chat/completions", bearer) + case "typesafe": + return _PassThroughCall("/typesafe/v1/chat/completions", chat, "/v1/chat/completions", bearer) + case "openrouter": + return _PassThroughCall("/openrouter/v1/chat/completions", chat, "/v1/chat/completions", bearer) + case "azure": + return _PassThroughCall( + "/azure/openai/deployments/scripted-deployment/chat/completions?api-version=2024-10-21", + chat, + "/openai/deployments/scripted-deployment/chat/completions?api-version=2024-10-21", + bearer, + ) + case "openai": + return _PassThroughCall("/openai/v1/chat/completions", chat, "/v1/chat/completions", bearer) + case "gigachat": + return _PassThroughCall("/gigachat/chat/completions", chat, "/chat/completions", bearer) + case "watsonx": + return _PassThroughCall( + "/watsonx/ml/v1/text/chat_stream", + {**chat, "model_id": "scripted-chat", "project_id": "scripted-project"}, + f"/ml/v1/text/chat_stream?version={_WATSONX_VERSION}", + bearer, + ) + case "tinyfish": + return _PassThroughCall( + "/tinyfish/v1/automation/run-sse", + {"url": "https://example.com", "goal": marker}, + "/v1/automation/run-sse", + bearer, + ) + case "nvidia_nim": + return _PassThroughCall( + f"/nvidia_nim/{_NIM_GROUP}/v1/chat/completions", chat, "/v1/chat/completions", bearer + ) + raise AssertionError(route) + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=60, trust_env=False)) for _ in range(count) + ) + + +def _post(client: httpx.Client, key: str, path: str, body: Mapping[str, JsonValue]) -> httpx.Response: + return client.post(path, json=dict(body), headers={"Authorization": f"Bearer {key}"}) + + +def _post_or_dropped(client: httpx.Client, key: str, path: str, body: Mapping[str, JsonValue]) -> httpx.Response | None: + try: + return _post(client, key, path, body) + except httpx.TransportError: + return None + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + port: Final = int(response.extensions["network_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 + ) + + +def _seen_once(delivered: tuple[_Delivery, ...], marker_hexes: tuple[str, ...]) -> None: + bodies: Final = tuple(json.dumps(sent.body) for sent in delivered) + seen: Final = tuple(sum(marker_hex in body for body in bodies) for marker_hex in marker_hexes) + assert seen == (1,) * len(marker_hexes), seen + + +def _pass_through_spend_rows(key: str) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT request_id, model, call_type FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s', + (sha256(key.encode()).hexdigest(), "pass_through_endpoint"), + ), + lambda rows: len(rows) == 1, + seconds=90, + ) + + +def _spend_rows(request_id: str) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows('SELECT request_id, model FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)), + lambda rows: len(rows) == 1, + seconds=70, + ) + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +@pytest.fixture(scope="module") +def pass_through_proxy( + session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[tuple[OwnedProxy, Wire]]: + directory: Final = tmp_path_factory.mktemp("structured-output-pass-through") + with wire_server(_relay_peer) as wire: + config: Final = _owned_config(directory / "pass-through.yaml", session_gateway, wire.url) + with owned_proxy_process( + session_gateway, directory, _pass_through_environment(wire.url), config=config + ) as owned: + yield owned, wire + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +@pytest.mark.parametrize( + "shape", + [ + "plain", + "wrapped", + "named", + "values", + "large", + "two_candidates", + "inner_incomplete", + "self_shaped", + "scalar", + "empty", + "list_input", + "deep", + "parameter_label", + "function_label", + "label_unshared", + "label_two_objects", + "label_encoded", + "label_then_declared", + "named_encoded", + "undeclared_encoded", + "encoded_deep", + "number_first", + "bool_first", + "null_first", + "label_non_ascii", + "label_two_objects_non_ascii", + ], +) +def test_vertex_chat_json_schema_content_is_the_schema_object(gateway: Gateway, shape: str) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert _chat_content(response) == _expected(shape, marker_hex), response.text + assert "tool_calls" not in _chat_message(response) or _chat_message(response)["tool_calls"] is None + sent: Final = _one_delivery(wire, marker_hex) + assert sent.target == f"{_VERTEX_MODELS}{_TOOL_MODEL}:rawPredict", sent.target + assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + assert "output_config" not in sent.body and "output_format" not in sent.body, sent.body + + +def test_vertex_chat_through_the_openai_sdk_unwraps_the_envelope(gateway: Gateway) -> None: + marker, marker_hex = _marker("wrapped") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + with _openai_client(gateway) as client: + reply: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": marker}], + response_format={"type": "json_schema", "json_schema": {"name": "risk", "schema": _SCHEMA}}, + ) + assert json.loads(reply.choices[0].message.content or "") == _payload(marker_hex), reply + _one_delivery(wire, marker_hex) + + +def test_vertex_chat_through_the_async_openai_sdk_unwraps_the_envelope(gateway: Gateway) -> None: + marker, marker_hex = _marker("wrapped") + + async def generate(model: str) -> str | None: + async with openai.AsyncOpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) as client: + reply: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": marker}], + response_format={"type": "json_schema", "json_schema": {"name": "risk", "schema": _SCHEMA}}, + ) + return reply.choices[0].message.content + + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + content: Final = asyncio.run(generate(model)) + assert json.loads(content or "") == _payload(marker_hex), content + _one_delivery(wire, marker_hex) + + +def test_vertex_responses_json_schema_output_text_is_the_schema_object(gateway: Gateway) -> None: + marker, marker_hex = _marker("wrapped") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/responses", _responses_body(model, marker)) + assert _responses_text(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + + +@pytest.mark.parametrize("shape", _STREAM_SHAPES) +def test_vertex_chat_stream_json_schema_content_is_the_schema_object(gateway: Gateway, shape: str) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) + frames: Final = _content_frames(response) + assert json.loads("".join(frames)) == _expected(shape, marker_hex), response.text + if shape in _HELD_STREAM_SHAPES: + assert len(frames) == 1, frames + else: + assert len(frames) >= _fragment_count(shape, marker_hex), frames + sent: Final = _one_delivery(wire, marker_hex) + assert sent.target == f"{_VERTEX_MODELS}{_TOOL_MODEL}{_VERTEX_STREAM_SUFFIX}", sent.target + assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + + +def test_vertex_chat_stream_label_before_a_declared_key_keeps_streaming(gateway: Gateway) -> None: + marker, marker_hex = _marker("label_then_declared") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) + frames: Final = _content_frames(response) + assert json.loads("".join(frames)) == _tool_input("label_then_declared", marker_hex), response.text + assert len(frames) > 1, frames + _one_delivery(wire, marker_hex) + + +@pytest.mark.parametrize("shape", _AS_WRITTEN_SHAPES) +def test_vertex_chat_stream_held_answer_without_an_envelope_is_sent_as_the_model_wrote_it( + gateway: Gateway, shape: str +) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) + assert _content_frames(response) == (_tool_text(shape, marker_hex),), response.text + _one_delivery(wire, marker_hex) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize("schema", [_OPTIONAL_SCHEMA, {"type": "object"}], ids=["no_required", "no_properties"]) +@pytest.mark.parametrize( + ("shape", "unwrapped"), [("wrapped", True), ("parameter_label", False)], ids=["wrapped", "label"] +) +def test_vertex_chat_json_schema_without_required_keys_unwraps_only_the_tool_named_envelope( + gateway: Gateway, schema: dict[str, JsonValue], shape: str, unwrapped: bool, stream: bool +) -> None: + marker, marker_hex = _marker(shape) + response_format: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "risk", "schema": schema}, + } + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format) + ) + content: Final = _stream_content(response) if stream else _chat_content(response) + assert content == (_payload(marker_hex) if unwrapped else _tool_input(shape, marker_hex)), response.text + if stream: + assert (len(_content_frames(response)) == 1) is unwrapped, response.text + _one_delivery(wire, marker_hex) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_vertex_chat_json_schema_that_declares_the_wrapper_key_keeps_the_answer_whole( + gateway: Gateway, stream: bool +) -> None: + marker, marker_hex = _marker("values") + response_format: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "valued", "schema": _VALUES_SCHEMA}, + } + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format) + ) + content: Final = _stream_content(response) if stream else _chat_content(response) + assert content == _tool_input("values", marker_hex), response.text + if stream: + assert len(_content_frames(response)) >= _fragment_count("values", marker_hex), response.text + _one_delivery(wire, marker_hex) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize( + ("schema", "wrapper_unwrapped"), + [ + (_ADDITIONAL_OBJECT_SCHEMA, False), + (_PATTERN_SCHEMA, False), + (_ADDITIONAL_TRUE_SCHEMA, False), + (_CLOSED_OPTIONAL_SCHEMA, True), + ], + ids=["additional_object", "pattern_properties", "additional_true", "additional_false"], +) +@pytest.mark.parametrize("shape", ["undeclared_wrapper", "wrapped"]) +def test_vertex_chat_json_schema_that_allows_undeclared_keys_keeps_an_undeclared_wrapper_whole( + gateway: Gateway, shape: str, schema: dict[str, JsonValue], wrapper_unwrapped: bool, stream: bool +) -> None: + marker, marker_hex = _marker(shape) + unwrapped: Final = shape == "wrapped" or wrapper_unwrapped + response_format: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "dynamic", "schema": schema}, + } + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream, response_format=response_format) + ) + content: Final = _stream_content(response) if stream else _chat_content(response) + assert content == (_payload(marker_hex) if unwrapped else _tool_input(shape, marker_hex)), response.text + if stream and unwrapped: + assert len(_content_frames(response)) == 1, response.text + elif stream: + assert len(_content_frames(response)) >= _fragment_count(shape, marker_hex), response.text + _one_delivery(wire, marker_hex) + + +def test_vertex_responses_stream_json_schema_output_text_is_the_schema_object(gateway: Gateway) -> None: + marker, marker_hex = _marker("wrapped") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/responses", {**_responses_body(model, marker), "stream": True}) + assert _responses_stream_text(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert sent.target == f"{_VERTEX_MODELS}{_TOOL_MODEL}{_VERTEX_STREAM_SUFFIX}", sent.target + + +@pytest.mark.parametrize("shape", ["other_tool", "named_input", "named_two_objects", "named_string", "named_partial"]) +def test_vertex_chat_json_schema_unwraps_a_tool_named_envelope_only_around_one_object( + gateway: Gateway, shape: str +) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert _chat_content(response) == _expected(shape, marker_hex), response.text + _one_delivery(wire, marker_hex) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize("shape", ["name_property", "name_value_plain"]) +def test_vertex_chat_json_schema_with_a_name_property_unwraps_only_the_tool_named_wrapper( + gateway: Gateway, shape: str, stream: bool +) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body(model, marker, stream=stream, response_format=_NAME_RESPONSE_FORMAT), + ) + content: Final = _stream_content(response) if stream else _chat_content(response) + assert content == _expected(shape, marker_hex), response.text + if stream and shape == "name_property": + assert len(_content_frames(response)) == 1, response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _json_tool_schema(sent)["properties"] == _NAME_SCHEMA["properties"], sent.body + + +@pytest.mark.parametrize("shape", ["wrapped", "inner_incomplete"]) +def test_vertex_chat_json_schema_validation_runs_on_the_unwrapped_object(gateway: Gateway, shape: str) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "enable_json_schema_validation": True} + ) + if shape == "wrapped": + assert _chat_content(response) == _payload(marker_hex), response.text + else: + assert response.status_code >= 400, response.text + assert "JSONSchemaValidationError" in response.text, response.text + _one_delivery(wire, marker_hex) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def test_vertex_chat_json_schema_next_to_a_user_tool_unwraps_the_envelope_and_keeps_the_tool_call( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker("wrapped") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, tools=[_USER_TOOL_SPEC]) + ) + assert response.status_code == 200, response.text + message: Final = object_value( + object_value(list_value(_JSON_OBJECT.validate_json(response.content)["choices"])[0])["message"] + ) + assert json.loads(string_value(message["content"])) == _payload(marker_hex), response.text + (tool_call,) = tuple(object_value(call) for call in list_value(message["tool_calls"])) + function: Final = object_value(tool_call["function"]) + assert function["name"] == _USER_TOOL, response.text + assert json.loads(string_value(function["arguments"])) == {"q": marker_hex}, response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _tool_names(sent.body) == frozenset({_USER_TOOL, _JSON_TOOL}), sent.body + + +@pytest.mark.parametrize("member", _FAMILY) +def test_tool_path_family_unwraps_the_envelope(gateway: Gateway, member: str) -> None: + marker, marker_hex = _marker("wrapped") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _family_deployment(gateway, scenario, wire.url, member) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert _chat_content(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + + +@pytest.mark.parametrize("member", _FAMILY) +def test_tool_path_family_stream_unwraps_the_envelope(gateway: Gateway, member: str) -> None: + marker, marker_hex = _marker("wrapped") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _family_deployment(gateway, scenario, wire.url, member) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=True)) + frames: Final = _content_frames(response) + assert frames, response.text + assert json.loads("".join(frames)) == _payload(marker_hex), response.text + assert len(frames) == 1, frames + sent: Final = _one_delivery(wire, marker_hex) + assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + + +@pytest.mark.parametrize("accepts_forced_tools", [True, False], ids=["accepts_forced_tools", "rejects_forced_tools"]) +def test_vertex_chat_json_schema_forces_the_json_tool_only_when_the_model_accepts_forced_tools( + gateway: Gateway, accepts_forced_tools: bool +) -> None: + marker, marker_hex = _marker("plain") + backend: Final = ( + _TOOL_MODEL if accepts_forced_tools else _claude_rejecting_forced_tools(gateway, "vertex_ai/", "claude") + ) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, backend) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert _chat_content(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + assert sent.body.get("tool_choice") == (_FORCED_JSON_TOOL if accepts_forced_tools else None), sent.body + assert "output_config" not in sent.body and "output_format" not in sent.body, sent.body + + +def test_vertex_chat_json_schema_on_a_model_rejecting_forced_tools_maps_reasoning_effort_like_plain_chat( + gateway: Gateway, +) -> None: + backend: Final = _claude_rejecting_forced_tools(gateway, "vertex_ai/", "claude") + structured_marker, structured_hex = _marker("plain") + plain_marker, plain_hex = _marker("plain") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, backend) + structured: Final = gateway.request( + "POST", + "/v1/chat/completions", + {**_effort_body(model, structured_marker), "response_format": _RESPONSE_FORMAT}, + ) + assert _chat_content(structured) == _payload(structured_hex), structured.text + structured_sent: Final = _one_delivery(wire, structured_hex) + plain: Final = gateway.request("POST", "/v1/chat/completions", _effort_body(model, plain_marker)) + assert plain.status_code == 200, plain.text + plain_sent: Final = _one_delivery(wire, plain_hex) + assert "thinking" in plain_sent.body, plain_sent.body + assert _reasoning_fields(structured_sent) == _reasoning_fields(plain_sent), (structured_sent.body, plain_sent.body) + + +def test_bedrock_invoke_json_schema_sends_the_json_tool_unforced_when_the_model_rejects_forced_tools( + gateway: Gateway, +) -> None: + marker, marker_hex = _marker("plain") + backend: Final = _claude_rejecting_forced_tools(gateway, "", "anthropic.claude") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_invoke_deployment(scenario, wire.url, backend) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert _chat_content(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert sent.target == f"/model/{backend}/invoke", sent.target + assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + assert "tool_choice" not in sent.body, sent.body + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_vertex_chat_json_schema_native_folds_output_format_into_output_config(gateway: Gateway, stream: bool) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=stream)) + content: Final = _stream_content(response) if stream else _chat_content(response) + assert content == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _folded_format(sent) == _OUTPUT_FORMAT, sent.body + assert "tools" not in sent.body, sent.body + + +def test_vertex_chat_native_through_the_openai_sdk_folds_output_format(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + with _openai_client(gateway) as client: + reply: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": marker}], + response_format={"type": "json_schema", "json_schema": {"name": "risk", "schema": _SCHEMA}}, + ) + assert json.loads(reply.choices[0].message.content or "") == _payload(marker_hex), reply + assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT + + +def test_vertex_responses_native_folds_output_format(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request("POST", "/v1/responses", _responses_body(model, marker)) + assert _responses_text(response) == _payload(marker_hex), response.text + assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT + + +def _native_deployment(gateway: Gateway, scenario: Scenario, api_base: str, provider: str, backend: str) -> str: + match provider: + case "vertex_ai": + return _vertex_deployment(gateway, scenario, api_base, backend) + case "anthropic": + return scenario.model(model=f"anthropic/{backend}", api_base=api_base, api_key=_ANTHROPIC_KEY) + case "azure_ai": + return scenario.model(model=f"azure_ai/{backend}", api_base=api_base, api_key=_AZURE_KEY) + case _: + raise AssertionError(provider) + + +@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"]) +def test_native_model_json_schema_with_an_open_object_takes_the_json_tool_and_keeps_the_object_open( + gateway: Gateway, provider: str +) -> None: + marker, marker_hex = _marker("values") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _native_deployment(gateway, scenario, wire.url, provider, _FORCED_TOOL_NATIVE_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=_VALUES_RESPONSE_FORMAT) + ) + assert _chat_content(response) == _tool_input("values", marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert not _structured(sent.body), sent.body + assert _json_tool_schema(sent)["properties"] == _VALUES_SCHEMA["properties"], sent.body + + +@pytest.mark.parametrize( + ("provider", "backend", "extra"), + [ + ("anthropic", _NATIVE_MODEL, {}), + ("vertex_ai", _NATIVE_MODEL, {}), + ("anthropic", _FORCED_TOOL_NATIVE_MODEL, {"thinking": {"type": "adaptive"}}), + ("anthropic", _FORCED_TOOL_NATIVE_MODEL, {"tools": [_USER_TOOL_SPEC]}), + ("anthropic", _FORCED_TOOL_NATIVE_MODEL, {"reasoning_effort": "none"}), + ("azure_ai", _FORCED_TOOL_NATIVE_MODEL, {"reasoning_effort": "none"}), + ], + ids=[ + "forced_tool_refused-anthropic", + "forced_tool_refused-vertex_ai", + "thinking", + "caller_tools", + "effort_none-anthropic", + "effort_none-azure_ai", + ], +) +def test_open_object_schema_keeps_native_output_when_the_json_tool_cannot_be_forced( + gateway: Gateway, provider: str, backend: str, extra: dict[str, JsonValue] +) -> None: + marker, marker_hex = _marker("values") + expected_tools: Final[frozenset[str]] = frozenset({_USER_TOOL}) if "tools" in extra else frozenset() + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _native_deployment(gateway, scenario, wire.url, provider, backend) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(model, marker, response_format=_VALUES_RESPONSE_FORMAT), **extra}, + ) + assert _chat_content(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _structured(sent.body), sent.body + assert _tool_names(sent.body) == expected_tools, sent.body + assert "tool_choice" not in sent.body, sent.body + assert ("thinking" in sent.body) is ("thinking" in extra), sent.body + + +def _thinking_then_json_tool(request: Request) -> Reply: + body: Final = _JSON_OBJECT.validate_json(request.body) + _, marker_hex = _marker_of(body) + assert _JSON_TOOL in _tool_names(body), body + thinking: Final[dict[str, JsonValue]] = { + "type": "thinking", + "thinking": f"thinking {marker_hex}", + "signature": _THINKING_SIGNATURE, + } + answer: Final = _tool_use(f"toolu_{marker_hex}", _JSON_TOOL, _payload(marker_hex)) + return Reply(body=json.dumps(_message(f"msg_{marker_hex}", [thinking, answer], "tool_use")).encode()) + + +def _json_tool_deployment(gateway: Gateway, scenario: Scenario, api_base: str, provider: str) -> str: + match provider: + 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) + case _: + raise AssertionError(provider) + + +@pytest.mark.parametrize("provider", ["vertex_ai", "bedrock_invoke"]) +def test_json_mode_answer_keeps_the_thinking_that_preceded_the_json_tool_call(gateway: Gateway, provider: str) -> None: + marker, marker_hex = _marker("plain") + with wire_server(_thinking_then_json_tool) as wire, gateway.scenario() as scenario: + model: Final = _json_tool_deployment(gateway, scenario, wire.url, provider) + response: Final = gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "reasoning_effort": "high"} + ) + message: Final = _chat_message(response) + assert json.loads(string_value(message["content"])) == _payload(marker_hex), response.text + assert message.get("reasoning_content") == f"thinking {marker_hex}", response.text + (block,) = tuple(object_value(item) for item in list_value(message.get("thinking_blocks") or [])) + assert (block["thinking"], block["signature"]) == (f"thinking {marker_hex}", _THINKING_SIGNATURE), block + sent: Final = _one_delivery(wire, marker_hex) + assert "thinking" in sent.body, sent.body + assert sent.body.get("tool_choice") != _FORCED_JSON_TOOL, sent.body + + +def test_vertex_chat_native_schema_keeps_min_items_one_and_closes_a_nullable_object(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + schema: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}, "minItems": 1}, + "backup": {"type": ["object", "null"], "properties": {"name": {"type": "string"}}, "required": ["name"]}, + }, + "required": ["summary", "tags", "backup"], + "additionalProperties": False, + } + expected: Final[dict[str, JsonValue]] = { + **schema, + "properties": { + "summary": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}, "minItems": 1}, + "backup": { + "type": ["object", "null"], + "properties": {"name": {"type": "string"}}, + "required": ["name"], + "additionalProperties": False, + }, + }, + } + response_format: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "tagged", "schema": schema}, + } + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker, response_format=response_format) + ) + assert _chat_content(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _folded_format(sent) == {"type": "json_schema", "schema": expected}, sent.body + assert "tools" not in sent.body, sent.body + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_vertex_messages_folds_output_format_into_output_config(gateway: Gateway, stream: bool) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, stream=stream, output_format=_OUTPUT_FORMAT), + headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA}, + ) + text: Final = _messages_stream_text(response) if stream else _messages_text(response) + assert text == _payload(marker_hex), response.text + assert f"msg_{marker_hex}" in response.text, response.text + sent: Final = _one_delivery(wire, marker_hex) + assert _folded_format(sent) == _OUTPUT_FORMAT, sent.body + assert sent.target.endswith(_VERTEX_STREAM_SUFFIX if stream else ":rawPredict"), sent.target + + +def test_vertex_messages_through_the_anthropic_sdk_folds_output_format(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + with _anthropic_client(gateway) as client: + reply: Final = client.beta.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + output_format={"type": "json_schema", "schema": _SCHEMA}, + betas=[_STRUCTURED_OUTPUTS_BETA], + ) + assert reply.id == f"msg_{marker_hex}", reply + assert _text_block_json(_JSON_OBJECT.validate_json(reply.model_dump_json())) == _payload(marker_hex), reply + assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT + + +def test_vertex_messages_through_the_async_anthropic_sdk_folds_output_format(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + + async def generate(model: str) -> tuple[str, str]: + async with anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=60), + ) as client: + reply: Final = await client.beta.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + output_format={"type": "json_schema", "schema": _SCHEMA}, + betas=[_STRUCTURED_OUTPUTS_BETA], + ) + return reply.id, reply.model_dump_json() + + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + reply_id, dumped = asyncio.run(generate(model)) + assert reply_id == f"msg_{marker_hex}", dumped + assert _text_block_json(_JSON_OBJECT.validate_json(dumped)) == _payload(marker_hex), dumped + assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize( + ("callers_config", "expected_output_config"), + [ + ({"format": _CALLERS_FORMAT}, {"format": _OUTPUT_FORMAT}), + ({"effort": "high", "format": _CALLERS_FORMAT}, {"effort": "high", "format": _OUTPUT_FORMAT}), + ], + ids=["format", "effort_and_format"], +) +def test_vertex_messages_output_format_wins_over_the_callers_output_config_format( + gateway: Gateway, callers_config: dict[str, JsonValue], expected_output_config: dict[str, JsonValue], stream: bool +) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, stream=stream, output_format=_OUTPUT_FORMAT, output_config=callers_config), + ) + text: Final = _messages_stream_text(response) if stream else _messages_text(response) + assert text == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert "output_format" not in sent.body, sent.body + assert sent.body["output_config"] == expected_output_config, sent.body + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.parametrize( + "schema", + [_ADDITIONAL_OBJECT_SCHEMA, _PATTERN_SCHEMA, _ADDITIONAL_TRUE_SCHEMA], + ids=["additional_object", "pattern_properties", "additional_true"], +) +def test_vertex_messages_json_schema_that_allows_undeclared_keys_keeps_the_answer_whole( + gateway: Gateway, schema: dict[str, JsonValue], stream: bool +) -> None: + marker, marker_hex = _marker("undeclared_wrapper") + output_format: Final[dict[str, JsonValue]] = {"type": "json_schema", "schema": schema} + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, stream=stream, output_format=output_format), + headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA}, + ) + text: Final = _messages_stream_text(response) if stream else _messages_text(response) + assert text == _tool_input("undeclared_wrapper", marker_hex), response.text + assert _folded_format(_one_delivery(wire, marker_hex)) == output_format + + +def test_vertex_messages_replaces_a_non_dict_output_config_with_the_folded_format(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, output_format=_OUTPUT_FORMAT, output_config="bogus"), + ) + assert _messages_text(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert sent.body["output_config"] == {"format": _OUTPUT_FORMAT}, sent.body + + +def test_vertex_messages_drops_effort_but_keeps_the_folded_format_on_haiku(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _EFFORT_GATED_MODEL) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, output_format=_OUTPUT_FORMAT, output_config={"effort": "high"}), + ) + assert _messages_text(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert sent.body["output_config"] == {"format": _OUTPUT_FORMAT}, sent.body + + +def test_vertex_messages_forwards_an_explicit_output_config_format_unchanged(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request( + "POST", "/v1/messages", _messages_body(model, marker, output_config={"format": _OUTPUT_FORMAT}) + ) + assert _messages_text(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert sent.body["output_config"] == {"format": _OUTPUT_FORMAT}, sent.body + + +@pytest.mark.parametrize( + ("extra", "expected_output_config"), + [ + ({"output_format": _OUTPUT_FORMAT, "output_config": None}, {"format": _OUTPUT_FORMAT}), + ({"output_format": _OUTPUT_FORMAT, "output_config": {"format": None}}, {"format": _OUTPUT_FORMAT}), + ({"output_format": _OUTPUT_FORMAT, "output_config": {}}, {"format": _OUTPUT_FORMAT}), + ( + {"output_format": _OUTPUT_FORMAT, "output_config": {"effort": "high"}}, + {"effort": "high", "format": _OUTPUT_FORMAT}, + ), + ({"output_format": None, "output_config": {"format": _OUTPUT_FORMAT}}, {"format": _OUTPUT_FORMAT}), + ({"output_format": None}, None), + ({"output_config": None}, None), + ], + ids=[ + "null_config", + "null_format_key", + "empty_config", + "effort_config", + "null_output_format_beside_callers", + "null_output_format_alone", + "null_config_alone", + ], +) +def test_vertex_messages_output_format_beside_an_explicit_output_config( + gateway: Gateway, extra: dict[str, JsonValue], expected_output_config: JsonValue +) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, marker, **extra)) + if expected_output_config is None: + assert response.status_code == 200, response.text + assert f"plain text {marker_hex}" in response.text, response.text + else: + assert _messages_text(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert "output_format" not in sent.body, sent.body + if expected_output_config is None: + assert "output_config" not in sent.body, sent.body + else: + assert sent.body["output_config"] == expected_output_config, sent.body + + +def test_anthropic_messages_forwards_top_level_output_format_verbatim(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_NATIVE_MODEL}", api_base=wire.url, api_key=_ANTHROPIC_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + _messages_body(model, marker, output_format=_OUTPUT_FORMAT), + headers={"anthropic-beta": _STRUCTURED_OUTPUTS_BETA}, + ) + assert _messages_text(response) == _payload(marker_hex), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert sent.target == "/v1/messages", sent.target + assert sent.body["output_format"] == _OUTPUT_FORMAT, sent.body + assert "output_config" not in sent.body, sent.body + + +def test_vertex_chat_without_response_format_sends_no_tool_and_no_output_config(gateway: Gateway) -> None: + marker, marker_hex = _marker("plain_text") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + assert response.status_code == 200, response.text + assert f"plain text {marker_hex}" in response.text, response.text + sent: Final = _one_delivery(wire, marker_hex) + assert not {"tools", "output_config", "output_format"} & sent.body.keys(), sent.body + + +@pytest.mark.parametrize( + "output_format", + [5, [], {}, "x" * 5000], + ids=["int", "list", "empty", "5kb"], +) +def test_vertex_messages_malformed_output_format_reaches_the_caller_as_the_peer_400( + gateway: Gateway, output_format: JsonValue +) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request( + "POST", "/v1/messages", _messages_body(model, marker, output_format=output_format) + ) + assert response.status_code == 400, response.text + assert _BAD_FORMAT in response.text, response.text + sent: Final = _one_delivery(wire, marker_hex) + assert "output_format" not in sent.body, sent.body + assert object_value(sent.body["output_config"])["format"] == output_format, sent.body + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +def test_vertex_messages_duplicate_output_format_keys_fold_the_last_value(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + raw: Final = ( + json.dumps(_messages_body(model, marker, output_format=5))[:-1] + + f', "output_format": {json.dumps(_OUTPUT_FORMAT)}}}' + ) + response: Final = gateway.client.post( + "/v1/messages", + content=raw, + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + assert _messages_text(response) == _payload(marker_hex), response.text + assert _folded_format(_one_delivery(wire, marker_hex)) == _OUTPUT_FORMAT + + +@pytest.mark.parametrize( + "schema", + [ + {"type": "object", "properties": ["tier", "summary"]}, + {"type": "object", "properties": "tier"}, + {"type": "object", "properties": None}, + ], + ids=["properties_list", "properties_string", "properties_null"], +) +@pytest.mark.parametrize("shape", ["wrapped", "other_tool"]) +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_vertex_chat_malformed_response_format_schema_keeps_the_tool_input( + gateway: Gateway, schema: JsonValue, shape: str, stream: bool +) -> None: + marker, marker_hex = _marker(shape) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body( + model, + marker, + stream=stream, + response_format={"type": "json_schema", "json_schema": {"name": "risk", "schema": schema}}, + ), + ) + content: Final = _stream_content(response) if stream else _chat_content(response) + assert content == _tool_input(shape, marker_hex), response.text + _one_delivery(wire, marker_hex) + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize("status", [400, 403, 404]) +def test_vertex_chat_peer_error_reaches_the_caller_and_the_deployment_keeps_serving( + gateway: Gateway, status: int +) -> None: + marker, marker_hex = _marker("wrapped") + after, after_hex = _marker("wrapped") + port: Final = allocate_port() + with gateway.scenario() as scenario: + with wire_server(_failing(status), port=port) as failing: + model: Final = _vertex_deployment(gateway, scenario, failing.url, _TOOL_MODEL) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert response.status_code == status, response.text + assert f"scripted {status} {marker_hex}" in response.text, response.text + _one_delivery(failing, marker_hex) + with wire_server(_peer, port=port) as wire: + recovered: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, after)) + assert _chat_content(recovered) == _payload(after_hex), recovered.text + _one_delivery(wire, after_hex) + + +def test_vertex_chat_unauthenticated_request_never_reaches_the_peer(gateway: Gateway) -> None: + marker, _ = _marker("wrapped") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, marker), key="sk-not-a-key-this-proxy-issued" + ) + assert response.status_code == 401, response.text + assert wire.drain() == () + + +def test_vertex_chat_repeated_requests_each_reach_the_peer_and_log_spend_once(gateway: Gateway) -> None: + markers: Final = tuple(_marker("wrapped") for _ in range(2)) + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + responses: Final = tuple( + gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) for marker, _ in markers + ) + for response, (_, marker_hex) in zip(responses, markers, strict=True): + assert _chat_content(response) == _payload(marker_hex), response.text + delivered: Final = _delivered(wire) + assert len(delivered) == 2, delivered + _seen_once(delivered, tuple(marker_hex for _, marker_hex in markers)) + request_ids: Final = tuple(_response_id(response) for response in responses) + assert len(set(request_ids)) == 2, request_ids + for request_id in request_ids: + assert _spend_rows(request_id)[0]["request_id"] == request_id + + +def test_concurrent_burst_across_both_workers_unwraps_every_envelope(gateway: Gateway) -> None: + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_peer)) + scenario: Final = stack.enter_context(gateway.scenario()) + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + shapes: Final = tuple(_BURST_SHAPES[index % len(_BURST_SHAPES)] for index in range(12)) + markers: Final = tuple(_marker(shape) for shape in shapes) + streams: Final = tuple(index % 3 == 2 for index in range(len(markers))) + bodies: Final = tuple( + _chat_body(model, marker, stream=stream) for (marker, _), stream in zip(markers, streams, strict=True) + ) + clients: Final = _clients(stack, str(gateway.client.base_url), len(bodies)) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(bodies))) + responses: Final = tuple( + pool.map( + lambda pair: _post(pair[0], gateway.key, "/v1/chat/completions", pair[1]), + zip(clients, bodies, strict=True), + ) + ) + for response, shape, (_, marker_hex), stream in zip(responses, shapes, markers, streams, strict=True): + observed: Final = _stream_content(response) if stream else _chat_content(response) + expected: Final = _expected(shape, marker_hex) + assert observed == expected, (shape, stream, response.text) + _seen_once(_delivered(wire), tuple(marker_hex for _, marker_hex in markers)) + + +def test_upstream_outage_mid_run_recovers_with_envelopes_still_unwrapped(gateway: Gateway) -> None: + with ExitStack() as stack: + clients: Final = _clients(stack, str(gateway.client.base_url), 8) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + scenario: Final = stack.enter_context(gateway.scenario()) + + streams: Final = tuple(index % 2 == 1 for index in range(len(clients))) + + def wave() -> tuple[tuple[httpx.Response | None, str, bool], ...]: + markers: Final = tuple(_marker("wrapped") for _ in clients) + responses: Final = tuple( + pool.map( + lambda item: _post_or_dropped( + item[0], gateway.key, "/v1/chat/completions", _chat_body(model, item[1][0], stream=item[2]) + ), + zip(clients, markers, streams, strict=True), + ) + ) + return tuple(zip(responses, tuple(marker_hex for _, marker_hex in markers), streams, strict=True)) + + def unwrapped(wire: Wire, results: tuple[tuple[httpx.Response | None, str, bool], ...]) -> None: + for response, marker_hex, stream in results: + assert response is not None, results + observed: Final = _stream_content(response) if stream else _chat_content(response) + assert observed == _payload(marker_hex), response.text + _seen_once(_delivered(wire), tuple(marker_hex for _, marker_hex, _ in results)) + + port: Final = allocate_port() + with wire_server(_peer, port=port) as wire: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + unwrapped(wire, wave()) + outage: Final = wave() + assert all(response is not None and response.status_code >= 500 for response, _, _ in outage), outage + assert gateway.request("GET", "/health/liveliness").status_code == 200 + with wire_server(_peer, port=port) as revived: + unwrapped(revived, wave()) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_worker_sigkill_mid_burst_leaves_the_replacement_unwrapping_envelopes(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(held, release))) + stack.callback(release.set) + config: Final = _owned_config(tmp_path / "structured-output-worker-kill.yaml", gateway, wire.url) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2)) + proxy_url: Final = owned.gateway.client.base_url + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + clients: Final = _clients(stack, str(proxy_url), 12) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + ports: Final = tuple(_local_port(client) for client in clients) + markers: Final = tuple(_marker("wrapped") for _ in clients) + futures: Final = tuple( + pool.submit(_post_or_dropped, client, gateway.key, "/v1/chat/completions", _chat_body(_OWNED_MODEL, 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])) + psutil.Process(victim).send_signal(signal.SIGKILL) + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result, (_, marker_hex) in zip(ports, results, markers, strict=True): + if port in shares[victim]: + assert result is None, (port, result, shares) + else: + assert result is not None and _chat_content(result) == _payload(marker_hex), (port, result, shares) + eventually(lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda started: started >= 3, 120) + second_wave: Final = tuple(_marker("wrapped") for _ in range(6)) + for marker, marker_hex in second_wave: + after: Final = _post( + owned.gateway.client, gateway.key, "/v1/chat/completions", _chat_body(_OWNED_MODEL, marker) + ) + assert _chat_content(after) == _payload(marker_hex), after.text + _seen_once(_delivered(wire), tuple(marker_hex for _, marker_hex in markers + second_wave)) + assert owned.process.poll() is None + + +@pytest.mark.parametrize( + ("shape", "response_format"), + [ + ("unclosed_key", _RESPONSE_FORMAT), + ("unclosed_name", _NAME_RESPONSE_FORMAT), + ("unclosed_label", _RESPONSE_FORMAT), + ("unclosed_label_key", _RESPONSE_FORMAT), + ], + ids=["unclosed_key", "unclosed_name", "unclosed_label", "unclosed_label_key"], +) +def test_vertex_chat_stream_never_closing_tool_input_stays_linear_and_keeps_health_responsive( + gateway: Gateway, shape: str, response_format: dict[str, JsonValue], record_property: Callable[[str, object], None] +) -> None: + latencies: Final[SimpleQueue[float]] = SimpleQueue() + done: Final = threading.Event() + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_peer)) + scenario: Final = stack.enter_context(gateway.scenario()) + model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) + streaming, probing = _clients(stack, str(gateway.client.base_url), 2) + _local_port(probing) + + def probe() -> None: + while not done.wait(0.05): + started: Final = time.perf_counter() + assert probing.get("/health/liveliness").status_code == 200 + latencies.put(time.perf_counter() - started) + + prober: Final = threading.Thread(target=probe) + prober.start() + stack.callback(prober.join, 30) + stack.callback(done.set) + small, large = ( + _timed_stream(streaming, gateway.key, wire, model, size_shape, response_format, record_property) + for size_shape in (f"{shape}_small", shape) + ) + done.set() + prober.join(timeout=30) + probes: Final = tuple(latencies.get_nowait() for _ in range(latencies.qsize())) + slowest: Final = max(probes, default=0.0) + record_property("slowest_health_seconds", round(slowest, 3)) + record_property("health_probes", len(probes)) + assert large <= _TIMING_GROWTH_LIMIT * small, (shape, small, large) + assert slowest <= _HEALTH_LATENCY_SECONDS, (shape, slowest, len(probes)) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_anthropic_pass_through_stream_relays_the_tool_stream_verbatim_and_logs_spend( + gateway: Gateway, tmp_path: Path +) -> None: + marker, marker_hex = _marker("wrapped") + overrides: Final = {"ANTHROPIC_API_BASE": "", "ANTHROPIC_API_KEY": _ANTHROPIC_KEY} + with wire_server(_peer) as wire, gateway.scenario() as scenario: + key: Final = scenario.key() + with owned_proxy_process(gateway, tmp_path, {**overrides, "ANTHROPIC_API_BASE": wire.url}) as owned: + response: Final = _post( + owned.gateway.client, + key, + "/anthropic/v1/messages", + _messages_body(_UNMAPPED_MODEL, marker, stream=True, tools=[_ANTHROPIC_JSON_TOOL]), + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith(_SSE), response.headers + expected: Final = _tool_reply(f"msg_{marker_hex}", "wrapped", marker_hex, False, "sse").chunks + assert expected is not None + assert response.content == b"".join(expected), response.text + sent: Final = _one_delivery(wire, marker_hex) + assert sent.target == "/v1/messages", sent.target + assert _json_tool_schema(sent) == _SCHEMA, sent.body + (row,) = _pass_through_spend_rows(key) + assert row["model"] == _UNMAPPED_MODEL, row + + +def test_bedrock_pass_through_invoke_stream_relays_the_event_stream_verbatim(gateway: Gateway) -> None: + marker, marker_hex = _marker("wrapped") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + deployment: Final = scenario.model( + model=f"bedrock/{_BEDROCK_INVOKE_MODEL}", + api_base=wire.url, + aws_access_key_id="AKIASCRIPTEDPROVIDER", + aws_secret_access_key="scripted-secret", + aws_region_name="us-east-1", + ) + response: Final = gateway.request( + "POST", + f"/bedrock/model/{deployment}/invoke-with-response-stream", + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 64, + "messages": [{"role": "user", "content": marker}], + "tools": [_ANTHROPIC_JSON_TOOL], + }, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"] == _EVENT_STREAM, response.headers + expected: Final = _tool_reply(f"msg_{marker_hex}", "wrapped", marker_hex, False, "aws").chunks + assert expected is not None + assert response.content == b"".join(expected), response.content[:200] + sent: Final = _one_delivery(wire, marker_hex) + assert sent.target == f"/model/{_BEDROCK_INVOKE_MODEL}/invoke-with-response-stream", sent.target + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +@pytest.mark.parametrize("route", _PASS_THROUGH_ROUTES) +def test_pass_through_stream_relays_the_scripted_frames_verbatim( + pass_through_proxy: tuple[OwnedProxy, Wire], session_gateway: Gateway, route: str +) -> None: + owned, wire = pass_through_proxy + marker_hex: Final = uuid.uuid4().hex + call: Final = _pass_through_call(route, f"relay {marker_hex}", session_gateway.key) + response: Final = owned.gateway.client.post(call.path, json=call.body, headers=call.headers) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith(_SSE), response.headers + assert response.content == b"".join(_RELAY_CHUNKS), response.text + (seen,) = tuple(request for request in wire.drain() if marker_hex in request.body.decode()) + assert seen.target == call.upstream_target, seen.target diff --git a/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py b/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py new file mode 100644 index 00000000000..1c419e5b11c --- /dev/null +++ b/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py @@ -0,0 +1,139 @@ +import json +import re +import uuid +from collections.abc import Callable, Iterator, Mapping +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm import ModelResponse, Router +from litellm.utils import _invalidate_model_cost_lowercase_map # pyright: ignore[reportPrivateUsage] # required cache invalidation + +_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})") +_SCHEMA: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, +} +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]]) +_OPTIONS: Final[dict[str, JsonValue]] = { + "max_tokens": 2048, + "thinking": {"type": "adaptive"}, + "response_format": {"type": "json_schema", "json_schema": {"name": "city", "schema": _SCHEMA}}, + "num_retries": 0, +} + + +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 _invoke_reply(answer: dict[str, JsonValue], native: bool) -> Reply: + block: Final[dict[str, JsonValue]] = ( + {"type": "text", "text": json.dumps(answer)} + if native + else {"type": "tool_use", "id": "toolu_scripted", "name": _JSON_TOOL, "input": answer} + ) + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": "claude-scripted", + "content": [block], + "stop_reason": "end_turn" if native else "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 7}, + } + return Reply(body=json.dumps(message).encode()) + + +def _converse_reply(answer: dict[str, JsonValue], native: bool) -> Reply: + block: Final[dict[str, JsonValue]] = ( + {"text": json.dumps(answer)} + if native + else {"toolUse": {"toolUseId": "tooluse_scripted", "name": _JSON_TOOL, "input": answer}} + ) + converse: Final[dict[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [block]}}, + "stopReason": "end_turn" if native else "tool_use", + "usage": {"inputTokens": 11, "outputTokens": 7, "totalTokens": 18}, + "metrics": {"latencyMs": 1}, + } + return Reply(body=json.dumps(converse).encode()) + + +def _bedrock(request: Request) -> Reply: + body: Final = _JSON.validate_json(request.body) + found: Final = _MARKER.search(json.dumps(body)) + assert found is not None, body + invoke: Final = "/invoke" in request.target + reply: Final[Callable[[dict[str, JsonValue], bool], Reply]] = _invoke_reply if invoke else _converse_reply + return reply({"city": found.group(1)}, _native(body, invoke)) + + +@pytest.fixture +def opus_4_8_opted_in_for_the_process(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "model_cost", _COST_MAP.validate_python(litellm.model_cost)) # pyright: ignore[reportUnknownMemberType] # the cost map is untyped + litellm.register_model( + {_OPUS_4_8: {"supports_bedrock_native_structured_output": True}}, persist_across_reloads=False + ) + yield + monkeypatch.undo() + _invalidate_model_cost_lowercase_map() + + +def _model(route: str) -> str: + return f"bedrock/{route}/{_OPUS_4_8}" + + +def _connection(wire: Wire) -> dict[str, JsonValue]: + return {"api_base": wire.url, "api_key": "synthetic-bedrock-bearer", "aws_region_name": "us-east-1"} + + +def _sent_native(wire: Wire, route: str, call: Callable[[list[dict[str, str]]], object]) -> bool: + marker_hex: Final = uuid.uuid4().hex + response: Final = call([{"role": "user", "content": f"Name a city. probe {marker_hex}"}]) + assert isinstance(response, ModelResponse), response + assert json.loads(str(response.choices[0].message.content)) == {"city": marker_hex}, response # pyright: ignore[reportAttributeAccessIssue] # choices also types streaming chunks + (request,) = wire.drain() + return _native(_JSON.validate_json(request.body), route == "invoke") + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +@pytest.mark.usefixtures("opus_4_8_opted_in_for_the_process") +def test_a_register_model_native_output_opt_in_reaches_every_call_but_a_deployment_that_opts_out(route: str) -> None: + with wire_server(_bedrock) as wire: + router: Final = Router( + model_list=[ + { + "model_name": "opted-out", + "litellm_params": {"model": _model(route), **_connection(wire)}, + "model_info": {"supports_bedrock_native_structured_output": False}, + }, + {"model_name": "inherits", "litellm_params": {"model": _model(route), **_connection(wire)}}, + ], + num_retries=0, + ) + assert _sent_native( + wire, + route, + lambda messages: router.completion(model="inherits", messages=messages, **_OPTIONS), # pyright: ignore[reportUnknownMemberType] # Router.completion leaves **kwargs untyped + ) + assert not _sent_native( + wire, + route, + lambda messages: router.completion(model="opted-out", messages=messages, **_OPTIONS), # pyright: ignore[reportUnknownMemberType] # Router.completion leaves **kwargs untyped + ) + assert _sent_native( + wire, + route, + lambda messages: litellm.completion(_model(route), messages, **_connection(wire), **_OPTIONS), + ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index e91d42e1cbf..73491fbf035 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -2,10 +2,10 @@ import asyncio import base64 import importlib import os -from collections.abc import Coroutine, Iterator +from collections.abc import Callable, Coroutine, Iterator from dataclasses import dataclass, field from pathlib import Path -from typing import Final +from typing import Final, Literal, TypeAlias import boto3 import httpx @@ -299,6 +299,22 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() +NativeStructuredOutputKey: TypeAlias = Literal["legacy", "bedrock"] | None + + +@pytest.fixture +def set_native_structured_output( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> Callable[[str, NativeStructuredOutputKey], None]: + def set_flags(model_cost_key: str, native_key: NativeStructuredOutputKey) -> None: + entry: Final = litellm.model_cost[model_cost_key] + monkeypatch.setitem(entry, "supports_native_structured_output", native_key == "legacy") + monkeypatch.setitem(entry, "supports_bedrock_native_structured_output", native_key == "bedrock") + litellm.get_model_info.cache_clear() + + return set_flags + + @pytest.fixture def local_beta_headers_config(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: monkeypatch.setenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", "True") diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 79c50bf2369..ce005fc9464 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -3,6 +3,7 @@ import functools import json import os import sys +from collections.abc import Mapping from typing import Final import pytest @@ -21,6 +22,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( is_encrypted_reasoning_block, merge_consecutive_system_messages, responses_reasoning_items_from_thinking_blocks, + salvage_concatenated_tool_arguments, + schema_grammar_branch_count, + schema_has_open_object, + schema_has_unresolvable_ref, + schema_property_count, split_concatenated_json_objects, strip_encrypted_reasoning_from_messages, system_messages_first, @@ -1952,3 +1958,76 @@ class TestMergeConsecutiveSystemMessages: merged = merge_consecutive_system_messages([{"role": "system"}, {"role": "system"}, {"role": "user", "content": "Hi"}]) assert merged == [{"role": "system"}, {"role": "user", "content": "Hi"}] + + +@pytest.mark.parametrize( + ("schema", "is_open"), + [ + ({"type": "object", "properties": {"a": {"type": "string"}}}, False), + ({"type": "object", "properties": {}}, False), + ({"type": "object", "additionalProperties": False}, False), + ({"type": "string"}, False), + ({"type": "object"}, True), + ({"type": ["object", "null"]}, True), + ({"type": "object", "properties": {"a": {"type": "string"}}, "additionalProperties": True}, True), + ({"type": "object", "additionalProperties": {"type": "integer"}}, True), + ({"type": "object", "properties": {}, "patternProperties": {"^x": {"type": "string"}}}, True), + ({"type": "object", "properties": {"inner": {"anyOf": [{"type": "object"}, {"type": "null"}]}}}, True), + ({"type": "array", "items": {"type": "object", "additionalProperties": {"type": "string"}}}, True), + ], +) +def test_schema_has_open_object_flags_any_node_that_accepts_unlisted_keys(schema, is_open): + assert schema_has_open_object(schema) is is_open + + +def _object_with(properties: Mapping[str, object], **defs: Mapping[str, object]) -> Mapping[str, object]: + return {"type": "object", "properties": properties, **defs} + + +@pytest.mark.parametrize( + ("schema", "unresolvable"), + [ + (_object_with({"a": {"type": "string"}}), False), + (_object_with({"a": {"$ref": "#/$defs/A"}, "b": {"$ref": "#/$defs/A"}}, **{"$defs": {"A": {"type": "string"}}}), False), + (_object_with({"a": {"$ref": "#/definitions/A"}}, definitions={"A": {"$ref": "#/definitions/B"}, "B": {"type": "integer"}}), False), + (_object_with({"a": {"$ref": "/$defs/A"}}, **{"$defs": {"A": {"items": {"$ref": "/$defs/B"}}, "B": {"type": "string"}}}), False), + (_object_with({"a": {"$ref": "#/$defs/Missing"}}, **{"$defs": {"A": {"type": "string"}}}), True), + (_object_with({"a": {"$ref": "other.json#/A"}}), True), + (_object_with({"parent": {"$ref": "#"}}), True), + (_object_with({"a": {"$ref": "#/$defs/A"}}, **{"$defs": {"A": {"type": "array", "items": {"$ref": "#/$defs/A"}}}}), True), + (_object_with({"a": {"$ref": "/$defs/A"}}, **{"$defs": {"A": {"type": "array", "items": {"$ref": "/$defs/A"}}}}), True), + ( + _object_with( + {"a": {"$ref": "#/$defs/A"}}, + **{"$defs": {"A": {"properties": {"b": {"$ref": "#/$defs/B"}}}, "B": {"properties": {"a": {"$ref": "#/$defs/A"}}}}}, + ), + True, + ), + ], + ids=["no-refs", "shared-def", "chained-definitions", "root-path-defs", "missing-def", "external", "root", "self-cycle", "root-path-self-cycle", "mutual-cycle"], +) +def test_schema_has_unresolvable_ref_flags_refs_that_cannot_be_inlined(schema, unresolvable): + assert schema_has_unresolvable_ref(schema) is unresolvable + + +@pytest.mark.parametrize( + ("schema", "properties", "branches"), + [ + (_object_with({"a": {"type": "string"}, "b": {"type": "string"}, "c": {"type": "string"}}, required=["a"]), 3, 8), + ( + _object_with( + {"x": {"$ref": "#/$defs/Pair"}, "y": {"$ref": "#/$defs/Pair"}}, + required=["x", "y"], + **{"$defs": {"Pair": _object_with({"l": {"type": "string"}, "r": {"type": "string"}}, required=["l"])}}, + ), + 4, + 7, + ), + (_object_with({"inner": _object_with({"z": {"type": "string"}})}, required=["inner"]), 2, 4), + ], + ids=["flat", "shared-def-counted-once", "nested"], +) +def test_schema_size_counts_every_written_object_once(schema, properties, branches): + assert schema_property_count(schema) == properties + assert schema_grammar_branch_count(schema) == branches + diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 26124ac24de..f84086138dc 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1012,6 +1012,25 @@ def test_bedrock_tools_unpack_defs(): _bedrock_tools_pt(tools=tools) +@pytest.mark.parametrize( + "parameters", + [ + {"type": ["object", "null"]}, + {"type": "object", "properties": {"tier": {"type": "integer"}}, "$defs": [{"type": "string"}]}, + ], + ids=["root_type_list", "defs_list"], +) +def test_bedrock_tools_pt_sends_malformed_parameters_as_an_object_schema(parameters: dict[str, object]) -> None: + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_tools_pt + + tools: Final = [{"type": "function", "function": {"name": "json_tool_call", "parameters": parameters}}] + + input_schema: Final = _bedrock_tools_pt(tools=tools)[0]["toolSpec"]["inputSchema"]["json"] + + assert input_schema["type"] == "object" + assert "$defs" not in input_schema + + def test_bedrock_tools_pt_strict_parameter(): """Regression for strict tools on the Bedrock Converse path. diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 729f46ec57f..a8b929ffa20 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1302,6 +1302,74 @@ def test_opus_uses_native_structured_output(model_name): assert optional_params.get("json_mode") is True +@pytest.mark.parametrize( + ("field_schema", "native"), + [ + ({"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, True), + ({"type": "object", "additionalProperties": {"type": "integer"}}, False), + ({"anyOf": [{"type": "object"}, {"type": "null"}]}, False), + ], + ids=["closed-object", "dict-field", "nullable-dict-field"], +) +def test_dict_fields_take_the_json_tool_on_a_native_structured_output_model(field_schema, native): + optional_params: Final = AnthropicConfig().map_openai_params( + non_default_params={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "out", + "schema": {"type": "object", "properties": {"field": field_schema}, "required": ["field"]}, + }, + } + }, + optional_params={}, + model="claude-opus-4-8", + drop_params=False, + ) + + assert ("output_format" in optional_params) is native + assert ([tool["name"] for tool in optional_params.get("tools", [])] == ["json_tool_call"]) is not native + + + +@pytest.mark.parametrize( + ("model", "thinking_params"), + [ + ("claude-opus-4-8", {"thinking": {"type": "adaptive"}}), + ("claude-opus-4-8", {"reasoning_effort": "low"}), + ("claude-opus-4-8", {"reasoning_effort": "none"}), + ("claude-sonnet-5-5", {}), + ( + "claude-opus-4-8", + {"tools": [{"type": "function", "function": {"name": "lookup", "parameters": {"type": "object"}}}]}, + ), + ], + ids=["adaptive-thinking", "reasoning-effort", "reasoning-effort-none", "model-without-forced-tools", "caller-tools"], +) +def test_dict_fields_keep_native_output_when_the_json_tool_would_not_be_forced(model, thinking_params): + optional_params: Final = AnthropicConfig().map_openai_params( + non_default_params={ + **thinking_params, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "out", + "schema": { + "type": "object", + "properties": {"field": {"type": "object", "additionalProperties": {"type": "integer"}}}, + "required": ["field"], + }, + }, + }, + }, + optional_params={}, + model=model, + drop_params=False, + ) + + assert "output_format" in optional_params + assert "json_tool_call" not in [tool["name"] for tool in optional_params.get("tools", [])] + def test_native_structured_output_uses_bundled_capability_when_remote_map_lags( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -6798,3 +6866,15 @@ def test_chat_dummy_tool_result_for_an_orphaned_tool_call_replays_a_byte_identic _assert_prefix_stable(requests) assert [m["role"] for m in requests[0]["messages"]] == ["user", "assistant", "user"] assert requests[0]["messages"][2]["content"][0]["type"] == "tool_result" + + +def test_json_mode_keeps_the_thinking_that_preceded_the_json_tool_call(): + thinking_block: Final = {"type": "thinking", "thinking": "The user wants a greeting.", "signature": "sig-1"} + body: Final = _json_mode_message_body("tool_use", None) + result: Final = _transform_json_mode_body({**body, "content": [thinking_block, *body["content"]]}) + + message: Final = result.choices[0].message + assert json.loads(message.content) == {"answer": "Hello"} + assert message.reasoning_content == thinking_block["thinking"] + assert [block["signature"] for block in message.thinking_blocks] == [thinking_block["signature"]] + diff --git a/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py b/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py index b2850192f91..0ed3b907a6f 100644 --- a/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py +++ b/tests/unit/llms/anthropic/test_anthropic_output_format_filter.py @@ -12,6 +12,18 @@ from litellm.llms.anthropic.chat.transformation import AnthropicConfig class TestOutputFormatArrayObjectConstraints: + def test_closes_an_object_whose_type_is_a_union_with_null(self): + schema = { + "type": "object", + "properties": { + "address": {"type": ["object", "null"], "properties": {"city": {"type": "string"}}}, + }, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["properties"]["address"]["additionalProperties"] is False + def test_removes_uniqueitems(self): schema = { "type": "object", @@ -129,7 +141,7 @@ class TestOutputFormatArrayObjectConstraints: schema = { "type": "array", "items": {"type": "string"}, - "minItems": 1, + "minItems": 2, "maxItems": 10, "uniqueItems": True, "minContains": 2, @@ -139,7 +151,7 @@ class TestOutputFormatArrayObjectConstraints: result = AnthropicConfig.filter_anthropic_output_schema(schema) assert result["description"] == ( - "Note: minimum number of items: 1, maximum number of items: 10, " + "Note: minimum number of items: 2, maximum number of items: 10, " "all array items must be unique, minimum number of matching items: 2, " "maximum number of matching items: 3." ) diff --git a/tests/unit/llms/anthropic/test_anthropic_schema_filter.py b/tests/unit/llms/anthropic/test_anthropic_schema_filter.py index 32423d6679f..407a52c183e 100644 --- a/tests/unit/llms/anthropic/test_anthropic_schema_filter.py +++ b/tests/unit/llms/anthropic/test_anthropic_schema_filter.py @@ -69,7 +69,7 @@ class TestFilterAnthropicOutputSchema: schema = { "type": "array", "items": {"type": "string"}, - "minItems": 1, + "minItems": 2, "maxItems": 10, } @@ -80,9 +80,27 @@ class TestFilterAnthropicOutputSchema: assert result["type"] == "array" assert result["items"] == {"type": "string"} # Description should contain constraint info - assert "minimum number of items: 1" in result["description"] + assert "minimum number of items: 2" in result["description"] assert "maximum number of items: 10" in result["description"] + def test_keeps_min_items_of_zero_or_one(self): + for min_items in (0, 1): + schema = {"type": "array", "items": {"type": "string"}, "minItems": min_items} + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["minItems"] == min_items + assert "description" not in result + + def test_moves_a_min_items_that_is_not_an_integer_into_the_description(self): + for min_items in (True, 1.0): + schema = {"type": "array", "items": {"type": "string"}, "minItems": min_items} + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "minItems" not in result + assert "minimum number of items" in result["description"] + def test_handles_nested_schemas(self): """Test that nested schemas are also filtered.""" schema = { @@ -100,7 +118,7 @@ class TestFilterAnthropicOutputSchema: } }, }, - "minItems": 1, + "minItems": 2, } }, } @@ -336,7 +354,7 @@ class TestFilterAnthropicOutputSchema: schema = { "type": "array", "items": {"type": "string"}, - "minItems": 1, + "minItems": 2, "maxItems": 10, "uniqueItems": True, "minContains": 2, @@ -346,7 +364,7 @@ class TestFilterAnthropicOutputSchema: result = AnthropicConfig.filter_anthropic_output_schema(schema) assert result["description"] == ( - "Note: minimum number of items: 1, maximum number of items: 10, " + "Note: minimum number of items: 2, maximum number of items: 10, " "all array items must be unique, minimum number of matching items: 2, " "maximum number of matching items: 3." ) 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 5d97beeb3fc..f855a846475 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 @@ -4,11 +4,13 @@ import copy import json import uuid from types import SimpleNamespace +from collections.abc import Mapping from typing import Final from unittest.mock import patch import httpx import pytest +from pydantic import BaseModel, ConfigDict # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. @@ -521,7 +523,7 @@ def test_output_config_format_converted_for_bedrock_chat_invoke_request(): with patch( # test-quality-ok: pin non-native path "litellm.llms.bedrock.common_utils._bedrock_model_supports", - side_effect=lambda _model, key: key == "supports_output_config", + side_effect=lambda _model, key, *_: key == "supports_output_config", ): result = config.transform_request( model="anthropic.claude-opus-4-7", @@ -543,10 +545,12 @@ def test_output_config_format_converted_for_bedrock_chat_invoke_request(): @pytest.mark.parametrize("model", ["anthropic.claude-opus-4-7", "us.anthropic.claude-opus-4-8"]) -def test_output_config_format_inlined_for_bedrock_chat_invoke_opus_4_7_and_4_8(local_model_cost_map, model): - """Bedrock rejects ``output_config.format`` on Claude Opus 4.7 and 4.8, so the - Invoke chat path inlines the schema into the last user message and keeps effort, - driven by the cost map alone (no capability stub).""" +def test_output_config_format_inlined_for_bedrock_chat_invoke_profile_without_native_output( + set_native_structured_output, model +): + """On a profile without native structured output the Invoke chat path inlines the + schema into the last user message and keeps effort.""" + set_native_structured_output(model, None) schema = {"type": "object", "properties": {"answer": {"type": "string"}}} result = AmazonAnthropicClaudeConfig().transform_request( @@ -564,12 +568,155 @@ def test_output_config_format_inlined_for_bedrock_chat_invoke_opus_4_7_and_4_8(l assert json.loads(result["messages"][-1]["content"][-1]["text"]) == schema +TREE_SCHEMA: Final = { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Node"}}, + "$defs": { + "Node": { + "type": "object", + "properties": {"children": {"type": "array", "items": {"$ref": "#/$defs/Node"}}}, + } + }, +} +OPEN_OBJECT_SCHEMA: Final = { + "type": "object", + "properties": {"counts": {"type": "object", "additionalProperties": {"type": "integer"}}}, +} + + +def _invoke_params_for_schema( + model: str, schema: Mapping[str, object], thinking_params: Mapping[str, object] | None = None +) -> Mapping[str, object]: + return AmazonAnthropicClaudeConfig().map_openai_params( + non_default_params={ + "response_format": {"type": "json_schema", "json_schema": {"name": "out", "schema": schema}}, + **({"thinking": {"type": "adaptive"}} if thinking_params is None else thinking_params), + }, + optional_params={}, + model=model, + drop_params=False, + ) + + +@pytest.mark.parametrize( + ("native_key", "thinking_params", "native"), + [ + ("legacy", {}, True), + ("bedrock", {}, False), + ("bedrock", {"thinking": {"type": "adaptive"}}, True), + ("bedrock", {"reasoning_effort": "low"}, True), + ], + ids=["legacy-key", "bedrock-key-no-thinking", "bedrock-key-adaptive", "bedrock-key-effort"], +) +def test_bedrock_native_output_key_applies_only_when_thinking_is_requested( + set_native_structured_output, native_key, thinking_params, native +): + model = "us.anthropic.claude-haiku-5-5" + set_native_structured_output(model, native_key) + schema = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]} + + result = _invoke_params_for_schema(model, schema, thinking_params) + + assert ("output_format" in result) is native + assert ("tools" in result) is not native + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +@pytest.mark.parametrize("schema", [TREE_SCHEMA, OPEN_OBJECT_SCHEMA], ids=["recursive", "open-object"]) +def test_schema_native_output_cannot_hold_keeps_the_tool_path_on_a_native_profile( + set_native_structured_output, schema, native_key +): + model = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + result = _invoke_params_for_schema(model, schema) + + assert "output_format" not in result + assert [tool["name"] for tool in result["tools"]] == ["json_tool_call"] + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_closed_schema_goes_native_on_either_native_key(set_native_structured_output, native_key): + model = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + schema = {"type": "object", "properties": {"age": {"type": "integer", "minimum": 0}}, "required": ["age"]} + + result = _invoke_params_for_schema(model, schema) + + assert result["output_format"]["schema"]["properties"]["age"]["type"] == "integer" + assert "minimum" not in result["output_format"]["schema"]["properties"]["age"] + assert "tools" not in result + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_shared_defs_reach_native_output_once_instead_of_inlined_per_use(set_native_structured_output, native_key): + model = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + party_ref = {"$ref": "#/$defs/Party"} + schema = { + "type": "object", + "properties": {"buyer": party_ref, "seller": party_ref}, + "required": ["buyer", "seller"], + "$defs": { + "Party": { + "type": "object", + "properties": {"email": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + "required": ["email"], + } + }, + } + + result = _invoke_params_for_schema(model, schema) + + sent = result["output_format"]["schema"] + assert sent["properties"]["buyer"] == party_ref + assert json.dumps(sent).count('"anyOf"') == 1 + assert "tools" not in result + + + +class Relabel(BaseModel): + model_config = ConfigDict(frozen=True) + + box: str + new_label: str + + +class PuzzleAnswer(BaseModel): + model_config = ConfigDict(frozen=True) + + relabels: list[Relabel] + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_nested_pydantic_response_format_reaches_native_output_with_refs_that_resolve( + set_native_structured_output, native_key +): + model = "us.anthropic.claude-haiku-5-5" + set_native_structured_output(model, native_key) + config: Final = AmazonAnthropicClaudeConfig() + + result: Final = config.map_openai_params( + non_default_params={ + "response_format": config.get_json_schema_from_pydantic_object(PuzzleAnswer), + "thinking": {"type": "adaptive"}, + }, + optional_params={}, + model=model, + drop_params=False, + ) + + sent: Final = result["output_format"]["schema"] + item_ref: Final = sent["properties"]["relabels"]["items"]["$ref"] + assert item_ref.removeprefix("#/$defs/") in sent["$defs"] + assert "tools" not in result + def test_output_config_format_forwarded_for_bedrock_chat_invoke_request(local_model_cost_map): """Bedrock Invoke chat path forwards ``output_config.format`` alongside effort for models with native structured-output support (Claude Sonnet 4.6).""" schema_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"answer": {"type": "string"}}}, + "schema": {"type": "object", "properties": {"answer": {"type": "string"}}, "additionalProperties": False}, } result = AmazonAnthropicClaudeConfig().transform_request( @@ -712,7 +859,7 @@ def test_bedrock_chat_invoke_drop_params_keeps_native_output_config_format(local monkeypatch.setattr(litellm, "drop_params", True) schema_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}, "additionalProperties": False}, } result = AmazonAnthropicClaudeConfig().transform_request( @@ -754,10 +901,11 @@ def test_bedrock_chat_invoke_drop_params_still_inlines_for_non_native(local_mode "model", ["us.anthropic.claude-fable-5-1", "anthropic.claude-fable-5-1"], ) -def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice(local_model_cost_map, model): - """Regression: Bedrock rejects both native ``output_config.format`` and forced - tool_choice for Fable 5.1, so invoke must use the tool-based path without a +def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice(set_native_structured_output, model): + """Regression: Bedrock rejects forced tool_choice for Fable 5.1, so on a profile + without native structured output invoke must use the tool-based path without a forced ``tool_choice``.""" + set_native_structured_output(model, None) result = AmazonAnthropicClaudeConfig().map_openai_params( non_default_params={ "response_format": { @@ -779,7 +927,10 @@ def test_bedrock_chat_invoke_fable_5_1_response_format_avoids_forced_tool_choice @pytest.mark.parametrize("model", ["us.anthropic.claude-sonnet-5", "us.anthropic.claude-fable-5-1"]) -def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_thinking(local_model_cost_map, model): +def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_thinking( + set_native_structured_output, model +): + set_native_structured_output(model, None) result = AmazonAnthropicClaudeConfig().map_openai_params( non_default_params={ "response_format": { @@ -802,6 +953,9 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th assert result["output_config"] == {"effort": "high"} + + + def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking(local_model_cost_map): """Regression: the tool-based ``response_format`` path swaps in a Claude 3 stub model before the shared Anthropic mapping, which hid the adaptive-only model diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 499096621c5..46f2d11402d 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -1,6 +1,7 @@ import copy import json import os +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock, patch @@ -8,8 +9,19 @@ import httpx import pytest import litellm -from litellm import ModelResponse +from litellm import ModelResponse, Router, completion +from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + _bedrock_tools_pt, +) from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import CONVERTED_SYSTEM_NOTE +from litellm.llms.anthropic.chat.transformation import ( + NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES, + NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS, + NATIVE_OUTPUT_MAX_UNION_PARAMETERS, +) +from litellm.llms.bedrock.common_utils import BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig from litellm.types.llms.bedrock import ConverseTokenUsageBlock @@ -637,6 +649,9 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} + + + def test_reasoning_effort_requests_summarized_display_converse(): """Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must request the summarized display, otherwise the provider returns a blank thinking @@ -831,32 +846,25 @@ def test_output_config_format_dropped_on_unsupported_converse_model_warns(caplog """When Converse model lacks native structured-output support, the silently dropped ``output_config.format`` must surface as a warning so callers can diagnose plain-text responses.""" - from unittest.mock import patch - config = AmazonConverseConfig() schema = { "type": "object", "properties": {"answer": {"type": "string"}}, } - with patch.object( - AmazonConverseConfig, - "_supports_native_structured_outputs", - return_value=False, - ): - with caplog.at_level("WARNING"): - result = config._transform_request( - model="bedrock/converse/us.anthropic.claude-3-haiku-20240307-v1:0", - messages=[{"role": "user", "content": "hi"}], - optional_params={ - "maxTokens": 256, - "output_config": { - "format": {"type": "json_schema", "schema": schema}, - }, + with caplog.at_level("WARNING"): + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-3-haiku-20240307-v1:0", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "output_config": { + "format": {"type": "json_schema", "schema": schema}, }, - litellm_params={}, - headers={}, - ) + }, + litellm_params={}, + headers={}, + ) assert "outputConfig" not in result assert any( @@ -4293,24 +4301,6 @@ def test_translate_response_format_native_output_config(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) -BEDROCK_OPUS_4_7_AND_4_8_MODELS: Final = ( - "anthropic.claude-opus-4-7", - "global.anthropic.claude-opus-4-7", - "us.anthropic.claude-opus-4-7", - "eu.anthropic.claude-opus-4-7", - "au.anthropic.claude-opus-4-7", - "jp.anthropic.claude-opus-4-7", - "anthropic.claude-opus-4-8", - "global.anthropic.claude-opus-4-8", - "us.anthropic.claude-opus-4-8", - "eu.anthropic.claude-opus-4-8", - "au.anthropic.claude-opus-4-8", - "jp.anthropic.claude-opus-4-8", - "us-gov.anthropic.claude-opus-4-8", - "us-gov-west-1/anthropic.claude-opus-4-8", - "us-gov-east-1/anthropic.claude-opus-4-8", -) - CAPITAL_RESPONSE_FORMAT: Final = { "type": "json_schema", "json_schema": { @@ -4325,10 +4315,15 @@ CAPITAL_RESPONSE_FORMAT: Final = { } -def _converse_request_for_json_schema(model: str, stream: bool) -> tuple[dict, dict]: +ADAPTIVE_THINKING: Final = {"type": "adaptive"} + + +def _converse_request_for_json_schema( + model: str, stream: bool, thinking_params: Mapping[str, object] | None = None +) -> tuple[Mapping[str, object], RequestObject]: config = AmazonConverseConfig() optional_params = config.map_openai_params( - non_default_params={"response_format": CAPITAL_RESPONSE_FORMAT, "stream": stream}, + non_default_params={"response_format": CAPITAL_RESPONSE_FORMAT, "stream": stream, **(thinking_params or {})}, optional_params={}, model=model, drop_params=False, @@ -4343,21 +4338,452 @@ def _converse_request_for_json_schema(model: str, stream: bool) -> tuple[dict, d return optional_params, request -@pytest.mark.parametrize("model", BEDROCK_OPUS_4_7_AND_4_8_MODELS) +@pytest.mark.parametrize( + ("native_profile", "tool_profile"), + [ + ("us.anthropic.claude-opus-4-8", "anthropic.claude-opus-4-8"), + ("global.anthropic.claude-opus-4-7", "us.anthropic.claude-opus-4-7"), + ("anthropic.claude-opus-4-8", "us-gov.anthropic.claude-opus-4-8"), + ], +) @pytest.mark.parametrize("stream", [False, True]) -def test_opus_4_7_and_4_8_json_schema_sent_as_forced_tool_not_output_config(monkeypatch, model, stream): - """Regression for issue #27846: Bedrock rejects outputConfig on Opus 4.7 and 4.8 - (``output_config.format: Extra inputs are not permitted``), so json_schema has to - go out as the forced json_tool_call tool, streamed through fake_stream.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) +def test_json_schema_shape_follows_each_inference_profile_own_native_output_flag( + set_native_structured_output, native_profile, tool_profile, stream +): + set_native_structured_output(native_profile, "bedrock") + set_native_structured_output(tool_profile, None) - optional_params, request = _converse_request_for_json_schema(model=model, stream=stream) + thinking_params: Final = {"thinking": ADAPTIVE_THINKING} + native_params, native_request = _converse_request_for_json_schema(native_profile, stream, thinking_params) + tool_params, tool_request = _converse_request_for_json_schema(tool_profile, stream, thinking_params) + + assert native_request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] == "capital" + assert "toolConfig" not in native_request + assert "fake_stream" not in native_params + assert "outputConfig" not in tool_request + assert [tool["toolSpec"]["name"] for tool in tool_request["toolConfig"]["tools"]] == ["json_tool_call"] + assert tool_request["toolConfig"]["toolChoice"] == {"tool": {"name": "json_tool_call"}} + assert tool_params.get("fake_stream", False) is stream + + +@pytest.mark.parametrize( + ("native_key", "thinking_params", "native"), + [ + ("legacy", {}, True), + ("bedrock", {}, False), + ("bedrock", {"thinking": ADAPTIVE_THINKING}, True), + ("bedrock", {"thinking": {"type": "enabled", "budget_tokens": 1024}}, True), + ("bedrock", {"reasoning_effort": "low"}, True), + ("bedrock", {"reasoning_effort": "none"}, False), + ("bedrock", {"thinking": {"type": "disabled"}}, False), + ], + ids=[ + "legacy-key", + "bedrock-key-no-thinking", + "bedrock-key-adaptive", + "bedrock-key-enabled", + "bedrock-key-effort", + "bedrock-key-effort-none", + "bedrock-key-thinking-disabled", + ], +) +def test_bedrock_native_output_key_applies_only_when_thinking_is_requested( + set_native_structured_output, native_key, thinking_params, native +): + model: Final = "us.anthropic.claude-haiku-5-5" + set_native_structured_output(model, native_key) + + _, request = _converse_request_for_json_schema(model, False, thinking_params) + + assert ("outputConfig" in request) is native + assert ("toolConfig" in request) is not native + + +def test_adaptive_thinking_is_kept_next_to_json_schema_on_a_native_output_profile(set_native_structured_output): + model: Final = "us.anthropic.claude-opus-5" + set_native_structured_output(model, "bedrock") + config: Final = AmazonConverseConfig() + thinking: Final = {"type": "adaptive", "display": "summarized"} + + optional_params: Final = config.map_openai_params( + non_default_params={"response_format": CAPITAL_RESPONSE_FORMAT, "thinking": thinking}, + optional_params={}, + model=model, + drop_params=False, + ) + request: Final = config._transform_request( + model=model, + messages=[{"role": "user", "content": "Name the capital of France."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert request["additionalModelRequestFields"]["thinking"] == thinking + assert request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["name"] == "capital" + assert "toolConfig" not in request + + +_OPUS_4_8_INVOKE_RESPONSE: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": '{"city": "Paris", "country": "France"}'}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, +} + + +def _routed_request_body(router: Router, model_name: str, route: str) -> Mapping[str, object]: + response: Final = _mock_converse_response() + if route == "invoke": + response.json.return_value = _OPUS_4_8_INVOKE_RESPONSE + response.text = json.dumps(_OPUS_4_8_INVOKE_RESPONSE) + client: Final = HTTPHandler() + with patch.object(client, "post", return_value=response) as mock_post: + router.completion( + model=model_name, + messages=[{"role": "user", "content": "Name the capital of France."}], + response_format=CAPITAL_RESPONSE_FORMAT, + thinking=ADAPTIVE_THINKING, + client=client, + ) + return json.loads(mock_post.call_args.kwargs["data"]) + + +def _sends_native_output(body: Mapping[str, object], route: str) -> bool: + if route == "converse": + return "outputConfig" in body and "toolConfig" not in body + output_config: Final = body.get("output_config") + return isinstance(output_config, dict) and "format" in output_config and "tools" not in body + + +@pytest.mark.parametrize("route", ["converse", "invoke"]) +def test_a_deployment_native_output_opt_in_stays_on_that_deployment(set_native_structured_output, route): + set_native_structured_output("us.anthropic.claude-opus-4-8", None) + litellm_params: Final = { + "model": f"bedrock/{route}/us.anthropic.claude-opus-4-8", + "aws_access_key_id": "fake-access-key", + "aws_secret_access_key": "fake-secret-key", + "aws_region_name": "us-east-1", + } + router: Final = Router( + model_list=[ + { + "model_name": "opted-in", + "litellm_params": litellm_params, + "model_info": {"supports_bedrock_native_structured_output": True}, + }, + {"model_name": "sibling", "litellm_params": litellm_params}, + ] + ) + + assert _sends_native_output(_routed_request_body(router, "opted-in", route), route) + assert not _sends_native_output(_routed_request_body(router, "sibling", route), route) + + +NATIVE_OUTPUT_REJECTED_KEYWORDS: Final = frozenset({"minimum", "maximum", "uniqueItems", "oneOf"}) +CONSTRAINED_PERSON_SCHEMA: Final = { + "type": "object", + "properties": { + "age": {"type": "integer", "minimum": 0, "maximum": 150}, + "tags": {"type": "array", "items": {"type": "string"}, "minItems": 2, "uniqueItems": True}, + "aliases": {"type": "array", "items": {"type": "string"}, "minItems": 1}, + "contact": {"oneOf": [{"type": "string"}, {"type": "integer"}]}, + "address": {"type": ["object", "null"], "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, + "required": ["age", "tags", "aliases", "contact", "address"], +} +RECURSIVE_TREE_SCHEMA: Final = { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Node"}}, + "required": ["root"], + "$defs": { + "Node": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "children": {"type": "array", "items": {"$ref": "#/$defs/Node"}}, + }, + "required": ["name", "children"], + } + }, +} +OPEN_OBJECT_SCHEMA: Final = { + "type": "object", + "properties": {"counts": {"type": "object", "additionalProperties": {"type": "integer"}}}, + "required": ["counts"], +} +LOOKAROUND_PATTERN_SCHEMA: Final = { + "type": "object", + "properties": {"amount": {"type": "string", "pattern": "^(?!^[-+.]*$)[+-]?0*\\d*\\.?\\d*$"}}, + "required": ["amount"], +} +UNKNOWN_FORMAT_SCHEMA: Final = { + "type": "object", + "properties": {"blob": {"type": "string", "format": "x-litellm-unlisted-format"}}, + "required": ["blob"], +} +ROOT_REF_SCHEMA: Final = { + "type": "object", + "properties": {"name": {"type": "string"}, "parent": {"anyOf": [{"$ref": "#"}, {"type": "null"}]}}, + "required": ["name", "parent"], +} + + +def _schema_keys(node: object) -> frozenset[str]: + if isinstance(node, dict): + return frozenset(node).union(*(_schema_keys(value) for value in node.values())) + if isinstance(node, list): + return frozenset().union(*(_schema_keys(item) for item in node)) + return frozenset() + + +def _converse_request_for_schema(model: str, schema: Mapping[str, object]) -> RequestObject: + config: Final = AmazonConverseConfig() + optional_params: Final = config.map_openai_params( + non_default_params={ + "response_format": {"type": "json_schema", "json_schema": {"name": "out", "schema": schema}}, + "thinking": ADAPTIVE_THINKING, + }, + optional_params={}, + model=model, + drop_params=False, + ) + return config._transform_request( + model=model, + messages=[{"role": "user", "content": "Fill in a plausible example."}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_native_output_schema_drops_keywords_bedrock_rejects_and_closes_objects( + set_native_structured_output, native_key +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = _converse_request_for_schema(model, CONSTRAINED_PERSON_SCHEMA) + + sent: Final = json.loads(request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert not _schema_keys(sent) & NATIVE_OUTPUT_REJECTED_KEYWORDS + assert sent["additionalProperties"] is False + assert sent["properties"]["contact"]["anyOf"] == [{"type": "string"}, {"type": "integer"}] + assert "minimum value: 0" in sent["properties"]["age"]["description"] + assert "minItems" not in sent["properties"]["tags"] + assert sent["properties"]["aliases"]["minItems"] == 1 + assert sent["properties"]["address"]["additionalProperties"] is False + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +@pytest.mark.parametrize( + ("schema", "kept_keyword"), + [ + (RECURSIVE_TREE_SCHEMA, "$ref"), + (ROOT_REF_SCHEMA, "$ref"), + (OPEN_OBJECT_SCHEMA, "additionalProperties"), + (LOOKAROUND_PATTERN_SCHEMA, "pattern"), + (UNKNOWN_FORMAT_SCHEMA, "format"), + ], + ids=["recursive", "root-ref", "open-object", "lookaround-pattern", "unlisted-format"], +) +def test_schema_native_output_cannot_hold_keeps_the_tool_path_on_a_native_profile( + set_native_structured_output, schema, kept_keyword, native_key +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = _converse_request_for_schema(model, schema) assert "outputConfig" not in request - assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["json_tool_call"] - assert request["toolConfig"]["toolChoice"] == {"tool": {"name": "json_tool_call"}} - assert optional_params.get("fake_stream", False) is stream + tool_spec: Final = request["toolConfig"]["tools"][0]["toolSpec"] + assert tool_spec["name"] == "json_tool_call" + assert kept_keyword in _schema_keys(tool_spec["inputSchema"]["json"]) + + +def _shared_defs_schema(defs_key: str) -> Mapping[str, object]: + party_ref: Final = {"$ref": f"#/{defs_key}/Party"} + return { + "type": "object", + "properties": {"buyer": party_ref, "seller": party_ref, "broker": {"anyOf": [party_ref, {"type": "null"}]}}, + "required": ["buyer", "seller", "broker"], + defs_key: { + "Party": { + "type": ["object", "null"], + "properties": {"name": {"type": "string"}, "email": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + "required": ["name", "email"], + } + }, + } + + +@pytest.mark.parametrize("defs_key", ["$defs", "definitions"]) +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_shared_defs_reach_native_output_once_instead_of_inlined_per_use( + set_native_structured_output, native_key, defs_key +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = _converse_request_for_schema(model, _shared_defs_schema(defs_key)) + + sent: Final = json.loads(request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert sent["properties"]["buyer"] == {"$ref": f"#/{defs_key}/Party"} + assert json.dumps(sent).count('"anyOf"') == 2 + assert sent[defs_key]["Party"]["additionalProperties"] is False + + +def _flat_schema(property_schema: Mapping[str, object], count: int, required: bool) -> Mapping[str, object]: + names: Final = tuple(f"field_{index}" for index in range(count)) + return { + "type": "object", + "properties": dict.fromkeys(names, property_schema), + "required": list(names) if required else [], + } + + +def _optional_groups_schema(count: int, group_size: int = 6) -> Mapping[str, object]: + groups: Final = { + f"group_{start}": _flat_schema({"type": "string"}, min(group_size, count - start), required=False) + for start in range(0, count, group_size) + } + return {"type": "object", "properties": groups, "required": list(groups)} + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +@pytest.mark.parametrize( + ("schema_for", "limit"), + [ + (lambda count: _flat_schema({"type": ["string", "null"]}, count, True), NATIVE_OUTPUT_MAX_UNION_PARAMETERS), + ( + lambda count: _flat_schema({"anyOf": [{"type": "string"}, {"type": "integer"}]}, count, True), + NATIVE_OUTPUT_MAX_UNION_PARAMETERS, + ), + (_optional_groups_schema, NATIVE_OUTPUT_MAX_OPTIONAL_PARAMETERS), + ( + lambda count: _flat_schema({"type": "string"}, count, False), + NATIVE_OUTPUT_MAX_GRAMMAR_BRANCHES.bit_length() - 1, + ), + ], + ids=["type-list-unions", "any-of-unions", "optional", "optional-grammar-branches"], +) +def test_schema_past_a_native_output_size_limit_keeps_the_tool_path( + set_native_structured_output, native_key, schema_for, limit +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + at_limit: Final = _converse_request_for_schema(model, schema_for(limit)) + past_limit: Final = _converse_request_for_schema(model, schema_for(limit + 1)) + + sent: Final = json.loads(at_limit["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert sent["properties"].keys() == schema_for(limit)["properties"].keys() + assert "outputConfig" not in past_limit + assert past_limit["toolConfig"]["tools"][0]["toolSpec"]["name"] == "json_tool_call" + + +@pytest.mark.parametrize( + ("native_key", "native_past_limit"), [("legacy", True), ("bedrock", False)], ids=["legacy-key", "bedrock-key"] +) +def test_field_count_limit_sends_the_json_tool_only_on_the_bedrock_native_key( + set_native_structured_output, native_key, native_past_limit +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + at_limit: Final = _converse_request_for_schema( + model, _flat_schema({"type": "string"}, BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES, True) + ) + past_limit: Final = _converse_request_for_schema( + model, _flat_schema({"type": "string"}, BEDROCK_NATIVE_OUTPUT_MAX_PROPERTIES + 1, True) + ) + + assert "outputConfig" in at_limit + assert ("outputConfig" in past_limit) is native_past_limit + + +NUMERIC_FORMAT_SCHEMA: Final = { + "type": "object", + "properties": { + "count": {"type": "integer", "format": "int32"}, + "total": {"type": "integer", "format": "int64"}, + "ratio": {"type": "number", "format": "double"}, + }, + "required": ["count", "total", "ratio"], +} + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_numeric_formats_keep_native_output(set_native_structured_output, native_key): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = _converse_request_for_schema(model, NUMERIC_FORMAT_SCHEMA) + + sent: Final = json.loads(request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert {name: field["format"] for name, field in sent["properties"].items()} == { + "count": "int32", + "total": "int64", + "ratio": "double", + } + + +def _doubling_defs_schema(depth: int) -> Mapping[str, object]: + levels: Final = { + f"Level{index}": { + "type": "object", + "properties": { + "left": {"$ref": f"#/$defs/Level{index + 1}"}, + "right": {"$ref": f"#/$defs/Level{index + 1}"}, + }, + "required": ["left", "right"], + } + for index in range(depth) + } + return { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Level0"}}, + "required": ["root"], + "$defs": {**levels, f"Level{depth}": {"type": "string"}}, + } + + +@pytest.mark.timeout(30) +def test_defs_that_double_at_every_level_reach_native_output_without_being_expanded(set_native_structured_output): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, "legacy") + schema: Final = _doubling_defs_schema(30) + + request: Final = _converse_request_for_schema(model, schema) + + sent: Final = json.loads(request["outputConfig"]["textFormat"]["structure"]["jsonSchema"]["schema"]) + assert sent["$defs"].keys() == schema["$defs"].keys() + assert sent["properties"]["root"] == {"$ref": "#/$defs/Level0"} + + +@pytest.mark.parametrize( + ("native_key", "forwarded"), [("legacy", True), ("bedrock", False)], ids=["legacy-key", "bedrock-key"] +) +def test_recursive_output_config_format_follows_the_profile_native_key( + set_native_structured_output, native_key, forwarded +): + model: Final = "us.anthropic.claude-opus-4-8" + set_native_structured_output(model, native_key) + + request: Final = AmazonConverseConfig()._transform_request( + model=model, + messages=[{"role": "user", "content": "Fill in a plausible example."}], + optional_params={"output_config": {"format": {"type": "json_schema", "schema": RECURSIVE_TREE_SCHEMA}}}, + litellm_params={}, + headers={}, + ) + + assert ("outputConfig" in request) is forwarded + def test_sonnet_4_6_json_schema_still_uses_native_output_config(monkeypatch): @@ -7399,11 +7825,12 @@ def test_unforced_tool_choice_unaffected_on_fable_5_1_converse(local_model_cost_ ["anthropic.claude-fable-5-1", "us.anthropic.claude-fable-5-1"], ) def test_response_format_avoids_native_and_forced_tool_choice_on_fable_5_1_converse( - local_model_cost_map, model + set_native_structured_output, model ): - """Regression: Bedrock rejects both ``outputConfig`` structured output and forced - tool_choice for Fable 5.1, so response_format must map to a tool without a forced - tool_choice.""" + """Regression: Bedrock rejects forced tool_choice for Fable 5.1, so on a profile + without native structured output response_format must map to a tool without a + forced tool_choice.""" + set_native_structured_output(model, None) config = AmazonConverseConfig() result = config.map_openai_params( 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 f4d51d975bb..414573cc4e4 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 @@ -1071,6 +1071,7 @@ def test_bedrock_messages_forwards_output_config_with_output_format(): "schema": { "type": "object", "properties": {"answer": {"type": "string"}}, + "additionalProperties": False, }, } optional_params = { @@ -1119,7 +1120,7 @@ def test_bedrock_messages_converts_output_config_format_to_inline_schema(): with patch( "litellm.llms.bedrock.common_utils._bedrock_model_supports", - side_effect=lambda _model, key: key == "supports_output_config", + side_effect=lambda _model, key, *_: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -1193,7 +1194,7 @@ def test_bedrock_messages_does_not_mutate_callers_messages_when_embedding_schema with patch( "litellm.llms.bedrock.common_utils._bedrock_model_supports", - side_effect=lambda _model, key: key == "supports_output_config", + side_effect=lambda _model, key, *_: key == "supports_output_config", ): result = cfg.transform_anthropic_messages_request( model="anthropic.claude-opus-4-7", @@ -3227,11 +3228,12 @@ def test_bedrock_messages_forwards_output_config_format_natively(local_model_cos assert "zebra_count" not in json.dumps(result["messages"]) -def test_bedrock_messages_inlines_schema_for_claude_5(local_model_cost_map): - """Bedrock rejects ``output_config.format`` for the Claude 5 family, so the - schema falls back to the inline-text path instead of a deterministic 400.""" +def test_bedrock_messages_inlines_schema_for_a_profile_without_native_output(set_native_structured_output): + """On a profile without native structured output the schema falls back to the + inline-text path instead of a deterministic 400.""" from litellm.types.router import GenericLiteLLMParams + set_native_structured_output("us.anthropic.claude-sonnet-5", None) cfg = AmazonAnthropicClaudeMessagesConfig() schema = { "type": "object", @@ -3254,6 +3256,115 @@ def test_bedrock_messages_inlines_schema_for_claude_5(local_model_cost_map): assert json.loads(last_content[-1]["text"]) == schema +def _bedrock_messages_request_with_schema(model: str, schema: Mapping[str, object]) -> Mapping[str, object]: + from litellm.types.router import GenericLiteLLMParams + + return AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": [{"type": "text", "text": "say hello"}]}], + anthropic_messages_optional_request_params={ + "max_tokens": 100, + "output_config": {"format": {"type": "json_schema", "schema": schema}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + +def test_bedrock_messages_native_output_drops_schema_keywords_bedrock_rejects(set_native_structured_output): + model: Final = "us.anthropic.claude-sonnet-5" + set_native_structured_output(model, "bedrock") + schema: Final = { + "type": "object", + "properties": { + "count": {"type": "integer", "minimum": 1}, + "ids": {"type": "array", "items": {"type": "string"}, "uniqueItems": True}, + }, + "required": ["count", "ids"], + } + + result: Final = _bedrock_messages_request_with_schema(model, schema) + + sent: Final = result["output_config"]["format"]["schema"] + assert "minimum" not in sent["properties"]["count"] + assert "uniqueItems" not in sent["properties"]["ids"] + assert sent["additionalProperties"] is False + + +@pytest.mark.parametrize("native_key", ["legacy", "bedrock"]) +def test_bedrock_messages_shared_defs_reach_native_output_once_instead_of_inlined_per_use( + set_native_structured_output, native_key +): + model: Final = "us.anthropic.claude-sonnet-5" + set_native_structured_output(model, native_key) + party_ref: Final = {"$ref": "#/$defs/Party"} + schema: Final = { + "type": "object", + "properties": {"buyer": party_ref, "seller": party_ref}, + "required": ["buyer", "seller"], + "additionalProperties": False, + "$defs": { + "Party": { + "type": "object", + "properties": {"email": {"anyOf": [{"type": "string"}, {"type": "null"}]}}, + "required": ["email"], + "additionalProperties": False, + } + }, + } + + result: Final = _bedrock_messages_request_with_schema(model, schema) + + sent: Final = result["output_config"]["format"]["schema"] + assert sent["properties"]["buyer"] == party_ref + assert json.dumps(sent).count('"anyOf"') == 1 + + +MESSAGES_TREE_SCHEMA: Final = { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Node"}}, + "$defs": { + "Node": { + "type": "object", + "properties": {"children": {"type": "array", "items": {"$ref": "#/$defs/Node"}}}, + } + }, +} +MESSAGES_OPEN_OBJECT_SCHEMA: Final = { + "type": "object", + "properties": {"counts": {"type": "object", "additionalProperties": {"type": "integer"}}}, +} + + +@pytest.mark.parametrize( + "schema", [MESSAGES_TREE_SCHEMA, MESSAGES_OPEN_OBJECT_SCHEMA], ids=["recursive", "open-object"] +) +def test_bedrock_messages_schema_native_output_cannot_hold_is_inlined_on_a_bedrock_native_profile( + set_native_structured_output, schema +): + model: Final = "us.anthropic.claude-sonnet-5" + set_native_structured_output(model, "bedrock") + + result: Final = _bedrock_messages_request_with_schema(model, schema) + + assert "output_config" not in result + assert json.loads(result["messages"][-1]["content"][-1]["text"]) == schema + + +@pytest.mark.parametrize( + "schema", [MESSAGES_TREE_SCHEMA, MESSAGES_OPEN_OBJECT_SCHEMA], ids=["recursive", "open-object"] +) +def test_bedrock_messages_schema_native_output_cannot_hold_is_forwarded_on_a_legacy_native_profile( + set_native_structured_output, schema +): + model: Final = "us.anthropic.claude-sonnet-4-6" + set_native_structured_output(model, "legacy") + + result: Final = _bedrock_messages_request_with_schema(model, schema) + + assert result["output_config"]["format"]["schema"] == schema + + def test_bedrock_messages_legacy_output_format_wins_over_output_config_format(local_model_cost_map): """When a request carries both schema forms, the legacy top-level ``output_format`` keeps winning, matching the pre-existing precedence.""" @@ -3262,7 +3373,11 @@ def test_bedrock_messages_legacy_output_format_wins_over_output_config_format(lo cfg = AmazonAnthropicClaudeMessagesConfig() legacy_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"legacy_field": {"type": "string"}}}, + "schema": { + "type": "object", + "properties": {"legacy_field": {"type": "string"}}, + "additionalProperties": False, + }, } newer_format = { "type": "json_schema", @@ -3296,7 +3411,7 @@ def test_bedrock_messages_drop_params_keeps_native_output_config_format(local_mo cfg = AmazonAnthropicClaudeMessagesConfig() schema_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}, "additionalProperties": False}, } result = cfg.transform_anthropic_messages_request( @@ -3321,7 +3436,7 @@ def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_mo cfg = AmazonAnthropicClaudeMessagesConfig() schema_format = { "type": "json_schema", - "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}}, + "schema": {"type": "object", "properties": {"zebra_count": {"type": "integer"}}, "additionalProperties": False}, } result = cfg.transform_anthropic_messages_request( diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index e5118f90e44..6f225840f99 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -557,8 +557,14 @@ def test_apply_structured_output_prefers_legacy_output_format(local_model_cost_m apply_bedrock_invoke_structured_output, ) - legacy = {"type": "json_schema", "schema": {"type": "object", "properties": {"a": {"type": "string"}}}} - newer = {"type": "json_schema", "schema": {"type": "object", "properties": {"b": {"type": "string"}}}} + legacy = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"a": {"type": "string"}}, "additionalProperties": False}, + } + newer = { + "type": "json_schema", + "schema": {"type": "object", "properties": {"b": {"type": "string"}}, "additionalProperties": False}, + } body = { "messages": [{"role": "user", "content": "hi"}], "output_format": legacy, diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index e80ca2a5820..4bd30764197 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1056,6 +1056,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_native_streaming": {"type": "boolean"}, "supports_image_size": {"type": "boolean"}, "supports_native_structured_output": {"type": "boolean"}, + "supports_bedrock_native_structured_output": {"type": "boolean"}, "use_openai_responses_path": {"type": "boolean"}, "off_peak_pricing": { "type": "object", diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py new file mode 100644 index 00000000000..c3aa223b41f --- /dev/null +++ b/tests/unit/test_utils_get_model_info.py @@ -0,0 +1,522 @@ +import json +import re +from collections.abc import Collection, Iterator, Mapping +from pathlib import Path +from typing import Final, Literal, cast + +import httpx +import pytest + +import litellm +from litellm import get_model_info +from litellm.llms.bedrock.common_utils import BedrockModelInfo +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.utils import ModelInfoBase +from litellm.utils import _invalidate_model_cost_lowercase_map +import os +from unittest.mock import MagicMock, patch + + +@pytest.fixture(autouse=True) +def isolate_model_info_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + monkeypatch.setattr(litellm, "custom_provider_map", list(litellm.custom_provider_map)) + yield + _invalidate_model_cost_lowercase_map() + + +def test_get_model_info_simple_model_name(): + """ + tests if model name given, and model exists in model info - the object is returned + """ + model = "claude-opus-5-5" + litellm.get_model_info(model) + + +def test_get_model_info_custom_llm_with_model_name(): + """ + Tests if {custom_llm_provider}/{model_name} name given, and model exists in model info, the object is returned + """ + model = "anthropic/claude-opus-5-5" + litellm.get_model_info(model) + + +def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch): + """ + Tests if {custom_llm_provider}/{model_name} name given, and model exists in model info, the object is returned + """ + model = "command-r-plus" + provider = "openai" # vllm is openai-compatible + litellm.register_model( + { + "openai/command-r-plus": { + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + } + ) + model_info = litellm.get_model_info(model, custom_llm_provider=provider) + print("model_info", model_info) + assert model_info["input_cost_per_token"] == 0.0 + + +def test_get_model_info_ollama_chat(): + from litellm.llms.ollama.completion.transformation import OllamaConfig + + with patch.object( + litellm.module_level_client, + "post", + return_value=MagicMock( + json=lambda: { + "model_info": {"llama.context_length": 32768}, + "template": "tools", + } + ), + ) as mock_client: + info = OllamaConfig().get_model_info("unknown-model") + assert info["supports_function_calling"] is True + + info = get_model_info("ollama/unknown-model") + print("info", info) + assert info["supports_function_calling"] is True + + mock_client.assert_called() + + print(mock_client.call_args.kwargs) + + assert mock_client.call_args.kwargs["json"]["name"] == "unknown-model" + + +def test_get_model_info_bedrock_region(monkeypatch): + regional_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + model_cost_without_regional_entry = { + key: value for key, value in litellm.get_model_cost_map(url="").items() if key != regional_model + } + monkeypatch.setattr(litellm, "model_cost", model_cost_without_regional_entry) + _invalidate_model_cost_lowercase_map() + info = litellm.get_model_info(model=regional_model, custom_llm_provider="bedrock") + print("info", info) + assert info["key"] == "anthropic.claude-haiku-4-5-20251001-v1:0" + assert info["litellm_provider"] == "bedrock_converse" + + +@pytest.mark.parametrize( + "model", + [ + "ft:gpt-3.5-turbo:my-org:custom_suffix:id", + "ft:gpt-4-0613:my-org:custom_suffix:id", + "ft:davinci-002:my-org:custom_suffix:id", + "ft:babbage-002:my-org:custom_suffix:id", + "gpt-35-turbo", + "ada", + ], +) +def test_get_model_info_completion_cost_unit_tests(model): + info = litellm.get_model_info(model) + print("info", info) + + +def test_get_model_info_ft_model_with_provider_prefix(): + args = { + "model": "openai/ft:gpt-3.5-turbo:my-org:custom_suffix:id", + "custom_llm_provider": "openai", + } + info = litellm.get_model_info(**args) + print("info", info) + assert info["key"] == "ft:gpt-3.5-turbo" + + +def _enforce_bedrock_converse_models( + model_cost: Mapping[str, ModelInfoBase], whitelist_models: Collection[str] +) -> None: + for model, info in model_cost.items(): + if ( + info.get("litellm_provider") == "bedrock" + and info.get("mode") == "chat" + and model not in whitelist_models + and not ( + (base_model := BedrockModelInfo.get_base_model(model)) != model + and model_cost.get(base_model, {}).get("litellm_provider") == "bedrock_converse" + and BedrockModelInfo.get_bedrock_route(model) == "converse" + ) + ): + raise AssertionError(f"Unlisted Bedrock chat model does not route to Converse: {model}") + + +def _read_whitelisted_bedrock_models() -> tuple[str, ...]: + path: Final = Path(__file__).resolve().parents[2] / "whitelisted_bedrock_models.txt" + return tuple(path.read_text().splitlines()) + + +def _normalize_bedrock_model_key(model_key: str) -> str: + without_wildcard: Final = model_key.replace("*/", "") + return re.sub(r"(?:1-month-commitment|3-month-commitment|6-month-commitment)/", "", without_wildcard) + + +def test_model_info_bedrock_converse(monkeypatch): + """ + Assert unlisted Bedrock chat models declare or inherit Converse routing. + + This ensures they are automatically routed to the converse endpoint. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + try: + # Load whitelist models from file + with open("whitelisted_bedrock_models.txt", "r") as file: + whitelist_models = [line.strip() for line in file.readlines()] + except FileNotFoundError: + pytest.skip("whitelisted_bedrock_models.txt not found") + + _enforce_bedrock_converse_models( + model_cost=litellm.model_cost, whitelist_models=whitelist_models + ) + + +@pytest.mark.flaky(retries=6, delay=2) +def test_model_info_bedrock_converse_enforcement(monkeypatch): + """ + Test the enforcement of the whitelist by adding a fake model and ensuring the test fails. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + + # Add a fake unwhitelisted model + litellm.model_cost["fake.bedrock-chat-model"] = { + "litellm_provider": "bedrock", + "mode": "chat", + } + + try: + # Load whitelist models from file + with open("whitelisted_bedrock_models.txt", "r") as file: + whitelist_models = [line.strip() for line in file.readlines()] + + # Check for unwhitelisted models + with pytest.raises(AssertionError, match=r"fake\.bedrock-chat-model"): + _enforce_bedrock_converse_models( + model_cost=litellm.model_cost, whitelist_models=whitelist_models + ) + except FileNotFoundError as e: + pytest.skip("whitelisted_bedrock_models.txt not found") + + +@pytest.mark.parametrize("region", ("us-gov-east-1", "us-gov-west-1")) +@pytest.mark.parametrize("base_provider", ("bedrock_converse", "bedrock")) +def test_regional_bedrock_alias_requires_canonical_converse_metadata( + region: str, base_provider: Literal["bedrock_converse", "bedrock"] +) -> None: + base_model: Final = next( + model for model in sorted(litellm.bedrock_converse_models) if BedrockModelInfo.get_base_model(model) == model + ) + model: Final = f"bedrock/{region}/{base_model}" + model_cost: Final[Mapping[str, ModelInfoBase]] = { + model: {"litellm_provider": "bedrock", "mode": "chat"}, + base_model: {"litellm_provider": base_provider, "mode": "chat"}, + } + assert BedrockModelInfo.get_bedrock_route(model) == "converse" + if base_provider == "bedrock": + with pytest.raises(AssertionError, match=re.escape(model)): + _enforce_bedrock_converse_models(model_cost, ()) + return + _enforce_bedrock_converse_models(model_cost, ()) + + +GOVCLOUD_LAGGING_CAPABILITIES: Final = frozenset({"supports_bedrock_native_structured_output"}) + + +def _govcloud_may_lag(model_key: str, capability: str) -> bool: + is_govcloud: Final = model_key.startswith("us-gov.") or any( + part.startswith("us-gov-") for part in model_key.split("/") + ) + return is_govcloud and capability in GOVCLOUD_LAGGING_CAPABILITIES + + +def test_get_model_info_bedrock_models(): + """ + Check for drift in base model info for bedrock models and regional model info for bedrock models. + GovCloud regions may lag their commercial base on Bedrock native structured output, never lead it. + """ + from litellm.llms.bedrock.common_utils import BedrockModelInfo + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + for k, v in litellm.model_cost.items(): + if v["litellm_provider"] == "bedrock": + k = k.replace("*/", "") + potential_commitments = [ + "1-month-commitment", + "3-month-commitment", + "6-month-commitment", + ] + if any(commitment in k for commitment in potential_commitments): + for commitment in potential_commitments: + k = k.replace(f"{commitment}/", "") + base_model = BedrockModelInfo.get_base_model(k) + # get_base_model() returns model id without "bedrock/" prefix; cost map keys use "bedrock/" + base_model_key = ( + base_model + if base_model in litellm.model_cost + else f"bedrock/{base_model}" + ) + if base_model_key not in litellm.model_cost: + continue + base_model_info = litellm.model_cost[base_model_key] + for base_model_key, base_model_value in base_model_info.items(): + if "invoke/" in k: + continue + if base_model_key.startswith("supports_"): + assert ( + base_model_key in v + ), f"{base_model_key} is not in model cost map for {k}" + if _govcloud_may_lag(k, base_model_key): + assert base_model_value or not v[base_model_key], f"{base_model_key} is on {k} but not its base" + continue + assert ( + v[base_model_key] == base_model_value + ), f"{base_model_key} is not equal to {base_model_value} for model {k}" + + +def _cross_region_base_model_key( + model_key: str, + prefixes: tuple[str, ...], + model_cost: Mapping[str, Mapping[str, object]], +) -> str | None: + matched_prefix: Final = next((prefix for prefix in prefixes if model_key.startswith(prefix)), None) + if matched_prefix is None: + return None + base_model_key: Final = model_key[len(matched_prefix) :] + return base_model_key if base_model_key in model_cost else None + + +def _cross_region_profiles( + model_cost: Mapping[str, Mapping[str, object]], + prefixes: tuple[str, ...], +) -> tuple[tuple[str, Mapping[str, object], str], ...]: + return tuple( + (model_key, model_info, base_model_key) + for model_key, model_info in model_cost.items() + if str(model_info.get("litellm_provider", "")).startswith("bedrock") + and ( + base_model_key := _cross_region_base_model_key(model_key, prefixes, model_cost) + ) + is not None + ) + + +def test_get_model_info_bedrock_cross_region_capability_parity(): + """ + Cross-region inference profiles carry litellm_provider "bedrock_converse", so the + regional drift check above (which filters on "bedrock") never reaches them. + GovCloud is a separate AWS partition that gets features later, so a us-gov profile + may lack Bedrock native structured output its commercial base has, never the reverse. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + prefixes = ("us.", "eu.", "apac.", "us-gov.") + checked = 0 + + for k, v in litellm.model_cost.items(): + if not str(v.get("litellm_provider", "")).startswith("bedrock"): + continue + base_model_key = next( + (k[len(p) :] for p in prefixes if k.startswith(p)), + None, + ) + if base_model_key is None or base_model_key not in litellm.model_cost: + continue + checked += 1 + for cap, base_value in litellm.model_cost[base_model_key].items(): + if not cap.startswith("supports_"): + continue + assert cap in v, f"{cap} is on {base_model_key} but missing from {k}" + if _govcloud_may_lag(k, cap): + assert base_value or not v[cap], f"{cap} is on {k} but not on {base_model_key}" + continue + assert ( + v[cap] == base_value + ), f"{cap} is {v[cap]} on {k} but {base_value} on {base_model_key}" + + assert checked > 0, "no cross-region bedrock profiles found - the filter is inert" + + +def _is_positive_cost(value: object) -> bool: + return isinstance(value, (int, float)) and value > 0 + + +def test_get_model_info_bedrock_priced_cross_region_profile_has_priced_base(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + prefixes = ("us.", "eu.", "apac.", "us-gov.", "au.", "global.") + checked = 0 + + for k, v in litellm.model_cost.items(): + if not str(v.get("litellm_provider", "")).startswith("bedrock"): + continue + base_model_key = next( + (k[len(p) :] for p in prefixes if k.startswith(p)), + None, + ) + if base_model_key is None or base_model_key not in litellm.model_cost: + continue + checked += 1 + base = litellm.model_cost[base_model_key] + for cost_key in ("input_cost_per_token", "output_cost_per_token"): + if (v.get(cost_key) or 0) > 0: + assert ( + base.get(cost_key) or 0 + ) > 0, f"{k} charges {cost_key} but its base {base_model_key} is free" + + assert checked > 0, "no cross-region bedrock profiles found - the filter is inert" + + +def test_get_model_info_case_insensitive_lookup(monkeypatch): + """ + Test that model info lookup is case-insensitive. + + This ensures that users can use lowercase model names even when the model cost + map has mixed-case keys (e.g., "Qwen/Qwen3-Next-80B-A3B-Thinking"). + + Related Slack discussion: Users were getting "does not support parameters: ['tools']" + errors when using lowercase model names like "qwen/qwen3-next-80b-a3b-thinking" + because the lookup was case-sensitive. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + + # Register a test model with mixed-case name + litellm.register_model( + { + "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "input_cost_per_token": 0.0001, + "output_cost_per_token": 0.0002, + "litellm_provider": "together_ai", + "supports_function_calling": True, + } + } + ) + + # Test 1: Exact case should work + info = litellm.get_model_info( + model="Qwen/Qwen3-Next-80B-A3B-Thinking", custom_llm_provider="together_ai" + ) + assert info is not None + assert info["supports_function_calling"] is True + + # Test 2: Lowercase should also work (case-insensitive lookup) + info_lower = litellm.get_model_info( + model="qwen/qwen3-next-80b-a3b-thinking", custom_llm_provider="together_ai" + ) + assert info_lower is not None + assert info_lower["supports_function_calling"] is True + + # Test 3: Mixed case should also work + info_mixed = litellm.get_model_info( + model="QWEN/qwen3-NEXT-80b-a3b-thinking", custom_llm_provider="together_ai" + ) + assert info_mixed is not None + assert info_mixed["supports_function_calling"] is True + + +def test_get_model_info_case_insensitive_supports_function_calling(monkeypatch): + """ + Test that supports_function_calling check works with case-insensitive model lookup. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + + # Register a model with mixed-case name that supports function calling + litellm.register_model( + { + "test_provider/TestModel-ABC": { + "input_cost_per_token": 0.0001, + "output_cost_per_token": 0.0002, + "litellm_provider": "test_provider", + "supports_function_calling": True, + } + } + ) + + # Test that supports_function_calling works with lowercase model name + from litellm.utils import supports_function_calling + + # Exact case + assert ( + supports_function_calling("TestModel-ABC", custom_llm_provider="test_provider") + is True + ) + + # Lowercase (should now work with case-insensitive lookup) + assert ( + supports_function_calling("testmodel-abc", custom_llm_provider="test_provider") + is True + ) + + +def test_get_model_info_custom_model_router(): + from litellm import Router + from litellm import get_model_info + + litellm.turn_on_debug() + + router = Router( + model_list=[ + { + "model_name": "ma-summary", + "litellm_params": { + "api_base": "http://ma-mix-llm-serving.cicero.svc.cluster.local/v1", + "input_cost_per_token": 1, + "output_cost_per_token": 1, + "model": "openai/meta-llama/Meta-Llama-3-8B-Instruct", + }, + "model_info": { + "id": "c20d603e-1166-4e0f-aa65-ed9c476ad4ca", + }, + } + ] + ) + info = get_model_info("c20d603e-1166-4e0f-aa65-ed9c476ad4ca") + print("info", info) + assert info is not None + + +def test_get_model_info_custom_provider(): + # Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server: + import litellm + from litellm import CustomLLM, completion + + class MyCustomLLM(CustomLLM): + def completion(self, *args, **kwargs) -> litellm.ModelResponse: + return litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello world"}], + mock_response="Hi!", + ) # type: ignore + + my_custom_llm = MyCustomLLM() + + litellm.custom_provider_map = [ # 👈 KEY STEP - REGISTER HANDLER + {"provider": "my-custom-llm", "custom_handler": my_custom_llm} + ] + + resp = completion( + model="my-custom-llm/my-fake-model", + messages=[{"role": "user", "content": "Hello world!"}], + ) + + assert resp.choices[0].message.content == "Hi!" + + # Register model info + model_info = {"my-custom-llm/my-fake-model": {"max_tokens": 2048}} + litellm.register_model(model_info) + + # Get registered model info + from litellm import get_model_info + + get_model_info( + model="my-custom-llm/my-fake-model" + ) # 💥 "Exception: This model isn't mapped yet." in v1.56.10